diff --git a/zig/build/cuda.zig b/zig/build/cuda.zig index c224600140..e3b603eac0 100644 --- a/zig/build/cuda.zig +++ b/zig/build/cuda.zig @@ -382,7 +382,7 @@ pub fn hostTests(b: *std.Build, draft_ids: *std.Build.Module, build_options: *st @import("qwen_cost_open.zig").tests(b, mods, native_modules, step); const windows_port = @import("windows_host.zig").steps(b, cuda, native_modules.api); const adapter_tests = b.createModule(.{ .root_source_file = b.path("zig/src/native/slide_adapter_test.zig"), .target = host, .link_libc = true, .imports = &.{.{ .name = "engine_api", .module = native_modules.api }} }); - b.step("test-slide-adapter-host", "Initial learner progress before its first long unit").dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = adapter_tests })).step); + b.step("test-slide-adapter-host", "Slide adapter: initial learner progress, and --slide's startup refusals").dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = adapter_tests })).step); @import("drama3_host.zig").native(b, step, drama3, mods.core, native_modules.api, cuda, "drama_cuda_native_host"); @import("drama3_host.zig").metal(b, mods.core, native_modules.api); const safety = b.step("test-kv-safety-host", "Actual native KV startup budget and paged constructors under allocation/completion faults"); diff --git a/zig/src/core/checkpoint_index.zig b/zig/src/core/checkpoint_index.zig index 7636a31357..1d60cf3114 100644 --- a/zig/src/core/checkpoint_index.zig +++ b/zig/src/core/checkpoint_index.zig @@ -15,16 +15,31 @@ pub fn plainName(name: []const u8) bool { return true; } +/// An index's shard: a plain name, or plain names joined by '/' with none starting with '.', inside the model folder. +pub fn shardName(name: []const u8) bool { + if (std.mem.indexOfScalar(u8, name, '/') == null) return plainName(name); + var parts = std.mem.splitScalar(u8, name, '/'); + while (parts.next()) |part| if (!plainName(part) or part[0] == '.') return false; + return true; +} + pub fn weightMap(root: std.json.Value) error{BadIndex}!std.json.ObjectMap { if (root != .object) return error.BadIndex; const map = root.object.get("weight_map") orelse return error.BadIndex; if (map != .object) return error.BadIndex; for (map.object.values()) |value| { - if (value != .string or !plainName(value.string)) return error.BadIndex; + if (value != .string or !shardName(value.string)) return error.BadIndex; } return map.object; } +/// The weight map with each shard directly in the model folder, as writes need: shard_edit.linked sees only those. +pub fn plainWeightMap(root: std.json.Value) error{BadIndex}!std.json.ObjectMap { + const map = try weightMap(root); + for (map.values()) |value| if (!plainName(value.string)) return error.BadIndex; + return map; +} + test "portable shard names exclude path and Windows device aliases" { const bad = [_][]const u8{ "", ".", "..", ".. ", "model.", "model ", "model\x00suffix", "model\nname", @@ -35,3 +50,15 @@ test "portable shard names exclude path and Windows device aliases" { for ([_][]const u8{ "model.safetensors", "model-00001-of-00002.safetensors", "model.v1", "unicode-λ" }) |name| try std.testing.expect(plainName(name)); } + +test "a shard name may name a folder inside the model folder, each part plain and none starting with '.'" { + const bad = [_][]const u8{ + "", "/model.safetensors", "sub/", "sub//model", "./model", "sub/./model", "sub/../model", + "../sub", "sub/..", "sub\\model", "C:/model", "CON/model", "sub/CON", "sub./model", + "sub /model", "sub/model ", "sub/m\x00x", "a/b/../../x", ".cache/model", "sub/.hid/model", "sub/.model", + "tensors/native-QB.", + }; + for (bad) |name| try std.testing.expect(!shardName(name)); + for ([_][]const u8{ "model.safetensors", "tensors/native-QB.safetensors", "a/b/c.safetensors" }) |name| + try std.testing.expect(shardName(name)); +} diff --git a/zig/src/core/checkpoint_security_test.zig b/zig/src/core/checkpoint_security_test.zig index 220a3fb2b4..643eb6e60a 100644 --- a/zig/src/core/checkpoint_security_test.zig +++ b/zig/src/core/checkpoint_security_test.zig @@ -33,10 +33,10 @@ test "security: checkpoint shard names cannot escape their model folder" { const dir = try tmp.dir.realPathFileAlloc(io, ".", a); defer a.free(dir); const invalid = [_][]const u8{ - "../outside.safetensors", "/tmp/outside.safetensors", "sub/../outside.safetensors", - "sub/model.safetensors", "..\\outside.safetensors", "C:\\outside.safetensors", - "C:outside.safetensors", "model.safetensors\x00extra", "", - ".", "..", + "../outside.safetensors", "/tmp/outside.safetensors", "sub/../outside.safetensors", + "..\\outside.safetensors", "C:\\outside.safetensors", "C:outside.safetensors", + "model.safetensors\x00extra", "", ".", + "..", }; for (invalid) |name| { const json = try std.json.Stringify.valueAlloc(a, .{ .weight_map = .{ .a = name } }, .{}); @@ -46,6 +46,30 @@ test "security: checkpoint shard names cannot escape their model folder" { } } +test "security: a checkpoint shard may sit in a folder inside its model folder" { + const a = std.testing.allocator; + const io = std.testing.io; + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + try tmp.dir.createDirPath(io, "sub"); + const dir = try tmp.dir.realPathFileAlloc(io, ".", a); + defer a.free(dir); + const shard = try std.fs.path.join(a, &.{ dir, "sub", "model.safetensors" }); + defer a.free(shard); + try edit.create(a, shard, &.{}, &.{.{ .name = "a", .dtype = .u8, .shape = &.{4}, .bytes = "here" }}); + try tmp.dir.writeFile(io, .{ + .sub_path = "model.safetensors.index.json", + .data = "{\"weight_map\":{\"a\":\"sub/model.safetensors\"}}", + }); + const files = try ck.shardFiles(a, io, dir); + defer ck.freeShardFiles(a, files); + try std.testing.expectEqual(@as(usize, 1), files.len); + try std.testing.expectEqualStrings("sub/model.safetensors", files[0][dir.len + 1 ..]); + var model = try ck.Checkpoint.openModel(a, io, dir); + defer model.close(); + try std.testing.expectEqualStrings("here", (try model.get("a")).bytes); +} + test "security: an unreadable checkpoint index is not silently ignored" { const a = std.testing.allocator; const io = std.testing.io; @@ -84,6 +108,61 @@ test "security: a checkpoint index cannot redirect Living Weights into another f try std.testing.expectError(error.BadIndex, changed); } +test "security: Living Weights refuse an index with a shard in a folder before writing to it" { + const a = std.testing.allocator; + const io = std.testing.io; + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + try tmp.dir.createDirPath(io, "sub"); + const dir = try tmp.dir.realPathFileAlloc(io, ".", a); + defer a.free(dir); + const shard = try std.fs.path.join(a, &.{ dir, "sub", "model.safetensors" }); + defer a.free(shard); + try edit.create(a, shard, &.{}, &.{.{ .name = "a.weight", .dtype = .u8, .shape = &.{4}, .bytes = "safe" }}); + const before = try tmp.dir.readFileAlloc(io, "sub/model.safetensors", a, .limited(4096)); + defer a.free(before); + try tmp.dir.writeFile(io, .{ + .sub_path = "model.safetensors.index.json", + .data = "{\"weight_map\":{\"a.weight\":\"sub/model.safetensors\"}}", + }); + const evil: []const edit.Replacement = &.{.{ .name = "a.weight", .dtype = .u8, .shape = &.{4}, .bytes = "evil" }}; + try std.testing.expectError(error.BadIndex, edit.overwrite(a, io, dir, evil)); + try std.testing.expectError(error.BadIndex, edit.grow(a, io, dir, evil)); + try std.testing.expectError(error.BadIndex, edit.bake(a, io, dir, evil)); + const after = try tmp.dir.readFileAlloc(io, "sub/model.safetensors", a, .limited(4096)); + defer a.free(after); + try std.testing.expectEqualSlices(u8, before, after); +} + +test "security: Living Weights keep to the model folder's own shards, so a linked folder cannot redirect them" { + if (@import("builtin").os.tag == .windows) return error.SkipZigTest; + const a = std.testing.allocator; + const io = std.testing.io; + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + try tmp.dir.createDirPath(io, "model"); + try tmp.dir.createDirPath(io, "outside"); + const root = try tmp.dir.realPathFileAlloc(io, ".", a); + defer a.free(root); + const dir = try std.fs.path.join(a, &.{ root, "model" }); + defer a.free(dir); + const victim = try std.fs.path.join(a, &.{ root, "outside", "victim.safetensors" }); + defer a.free(victim); + try edit.create(a, victim, &.{}, &.{.{ .name = "a.weight", .dtype = .u8, .shape = &.{4}, .bytes = "safe" }}); + const before = try tmp.dir.readFileAlloc(io, "outside/victim.safetensors", a, .limited(4096)); + defer a.free(before); + try tmp.dir.symLink(io, "../outside", "model/sub", .{ .is_directory = true }); + try tmp.dir.writeFile(io, .{ + .sub_path = "model/model.safetensors.index.json", + .data = "{\"weight_map\":{\"a.weight\":\"sub/victim.safetensors\"}}", + }); + const changed = edit.overwrite(a, io, dir, &.{.{ .name = "a.weight", .dtype = .u8, .shape = &.{4}, .bytes = "evil" }}); + const after = try tmp.dir.readFileAlloc(io, "outside/victim.safetensors", a, .limited(4096)); + defer a.free(after); + try std.testing.expectEqualSlices(u8, before, after); + try std.testing.expectError(error.BadIndex, changed); +} + test "security: shard and snapshot atomic writes cannot follow preexisting temporary symlinks" { if (@import("builtin").os.tag == .windows) return error.SkipZigTest; const a = std.testing.allocator; diff --git a/zig/src/core/shard_edit.zig b/zig/src/core/shard_edit.zig index 9bd4caecde..499bb91540 100644 --- a/zig/src/core/shard_edit.zig +++ b/zig/src/core/shard_edit.zig @@ -54,7 +54,7 @@ pub fn bake(gpa: Allocator, io: Io, dir: []const u8, edits: []const Replacement) for (edits) |r| if (r.bytes.len != numel(r.shape) * r.dtype.size() or !std.mem.endsWith(u8, r.name, ".weight")) return error.BadReplacement; const index_path = try std.fs.path.join(arena, &.{ dir, index_name }); const index: ?Value = readJson(arena, io, index_path) catch |e| if (e == error.FileNotFound) null else return e; - if (index) |ix| _ = try index_format.weightMap(ix); + if (index) |ix| _ = try index_format.plainWeightMap(ix); var shards: std.StringArrayHashMapUnmanaged(std.ArrayList(Replacement)) = .empty; for (edits) |r| { const shard = if (index) |ix| str((try field(try field(ix, "weight_map"), r.name))) orelse return error.BadIndex else single_name; @@ -95,7 +95,7 @@ pub fn grow(gpa: Allocator, io: Io, dir: []const u8, edits: []const Replacement) for (edits) |r| if (r.bytes.len != numel(r.shape) * r.dtype.size()) return error.BadReplacement; const index_path = try std.fs.path.join(arena, &.{ dir, index_name }); const index: ?Value = readJson(arena, io, index_path) catch |e| if (e == error.FileNotFound) null else return e; - if (index) |ix| _ = try index_format.weightMap(ix); + if (index) |ix| _ = try index_format.plainWeightMap(ix); var shards: std.StringArrayHashMapUnmanaged(std.ArrayList(Replacement)) = .empty; for (edits) |r| { const shard = if (index) |ix| str((try field(try field(ix, "weight_map"), r.name))) orelse return error.BadIndex else single_name; @@ -167,7 +167,7 @@ pub fn overwrite(gpa: Allocator, io: Io, dir: []const u8, edits: []const Replace const arena = arena_state.allocator(); const index_path = try std.fs.path.join(arena, &.{ dir, index_name }); const index: ?Value = readJson(arena, io, index_path) catch |e| if (e == error.FileNotFound) null else return e; - if (index) |ix| _ = try index_format.weightMap(ix); + if (index) |ix| _ = try index_format.plainWeightMap(ix); for (edits) |r| { const shard = if (index) |ix| str((try field(try field(ix, "weight_map"), r.name))) orelse return error.BadIndex else single_name; const path = try std.fs.path.join(arena, &.{ dir, shard }); diff --git a/zig/src/families/radio/weights.zig b/zig/src/families/radio/weights.zig index f52db83bea..2ba0c1d2b4 100644 --- a/zig/src/families/radio/weights.zig +++ b/zig/src/families/radio/weights.zig @@ -191,7 +191,7 @@ fn open(gpa: Allocator, io: std.Io, dir: []const u8) !std.ArrayList(st.File) { var named = try indexed(gpa, io, dir); defer if (named) |*n| freeNames(gpa, n); for (paths) |path| { - if (named) |n| if (!n.contains(std.fs.path.basename(path))) continue; + if (named) |n| if (!n.contains(path[dir.len + 1 ..])) continue; var f = try st.File.open(gpa, io, path); const has = for (f.names.keys()) |k| { if (ours(k)) break true; @@ -242,6 +242,29 @@ test "a layer's tensors are known; a layer scale is not" { try std.testing.expect(!Planner.known(tower ++ "blocks.3.ls1.gamma")); } +test "the image tensors' shard is found in the folder the index names" { + const gpa = std.testing.allocator; + const io = std.testing.io; + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + try tmp.dir.createDirPath(io, "tensors"); + const dir = try tmp.dir.realPathFileAlloc(io, ".", gpa); + defer gpa.free(dir); + const path = try std.fs.path.join(gpa, &.{ dir, "tensors", "vision.safetensors" }); + defer gpa.free(path); + try core.shard_edit.create(gpa, path, &.{}, &.{.{ .name = "vision_model.x", .dtype = .u8, .shape = &.{4}, .bytes = "seen" }}); + try tmp.dir.writeFile(io, .{ + .sub_path = "model.safetensors.index.json", + .data = "{\"weight_map\":{\"vision_model.x\":\"tensors/vision.safetensors\"}}", + }); + var files = try open(gpa, io, dir); + defer { + for (files.items) |*f| f.close(io); + files.deinit(gpa); + } + try std.testing.expectEqual(@as(usize, 1), files.items.len); +} + const Failed = struct { gpa: Allocator, w: Weights, diff --git a/zig/src/native/slide_adapter.zig b/zig/src/native/slide_adapter.zig index 24f63f76aa..f8cef10659 100644 --- a/zig/src/native/slide_adapter.zig +++ b/zig/src/native/slide_adapter.zig @@ -10,6 +10,21 @@ const max_neurons = 64; pub fn refusal(comptime Checks: type, a: Allocator, io: std.Io, dir: []const u8, copy: []const u8) !?[]const u8 { if (@hasDecl(Checks, "plainShared")) if (try Checks.plainShared(a, io, dir)) |name| return try std.fmt.allocPrint(a, "--slide adds learned neurons beside every shared expert, and {s} in {s} is unquantized, as TensorFold 1.0.3 to 1.0.6 saved their lessons: serve this folder without --slide (what it learned stays), or teach a fresh copy", .{ name, dir }); if (try Checks.linked(a, io, dir)) |name| return try std.fmt.allocPrint(a, "--slide rewrites the model's own files, and {s} in {s} is a link to data another file shares (a Hugging Face cache's blob, or a hard link), so learning would change that file too: serve a copy made with {s}", .{ name, dir, copy }); + if (try folderShard(a, io, dir)) |name| return try std.fmt.allocPrint(a, "--slide rewrites the model's own shards, and {s} in {s} sits in a folder, where the checks for links do not look: serve this folder without --slide", .{ name, dir }); + return null; +} + +/// The first shard the index places in a folder (shard_edit writes only the shards beside it), or null. +fn folderShard(a: Allocator, io: std.Io, dir: []const u8) !?[]const u8 { + const path = try std.fs.path.join(a, &.{ dir, "model.safetensors.index.json" }); + const text = std.Io.Dir.cwd().readFileAlloc(io, path, a, .limited(1 << 26)) catch |e| switch (e) { + error.FileNotFound => return null, + else => return e, + }; + const index = try std.json.parseFromSliceLeaky(std.json.Value, a, text, .{}); + const map = if (index == .object) index.object.get("weight_map") orelse return null else return null; + if (map != .object) return null; + for (map.object.values()) |v| if (v == .string and std.mem.indexOfScalar(u8, v.string, '/') != null) return v.string; return null; } diff --git a/zig/src/native/slide_adapter_test.zig b/zig/src/native/slide_adapter_test.zig index 081e6de5db..de7912158a 100644 --- a/zig/src/native/slide_adapter_test.zig +++ b/zig/src/native/slide_adapter_test.zig @@ -1,7 +1,8 @@ -//! Initial reader progress must be visible before its first potentially long step begins. +//! The slide adapter: initial reader progress before its first potentially long step, and --slide's startup refusals. const std = @import("std"); const api = @import("engine_api"); const Adapter = @import("slide_adapter.zig").Adapter; +const refusal = @import("slide_adapter.zig").refusal; const Fixture = struct { pub const Example = api.Example; pub const Learner = struct { @@ -44,3 +45,26 @@ test "lesson adapter: publish reading zero before the first long unit" { try hook.begin(hook.ctx, &.{}, .{ .ctx = &seen, .event = Seen.on }); try std.testing.expectEqual(@as(usize, 1), seen.reads); } +const NoLinks = struct { + pub fn linked(_: std.mem.Allocator, _: std.Io, _: []const u8) !?[]const u8 { + return null; + } +}; +test "slide refusal: a shard the index places in a folder is refused at startup, before any save" { + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + const io = std.testing.io; + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + const dir = try tmp.dir.realPathFileAlloc(io, ".", a); + try std.testing.expect((try refusal(NoLinks, a, io, dir, "a copy")) == null); + try tmp.dir.writeFile(io, .{ .sub_path = "model.safetensors.index.json", .data = "{\"weight_map\":{\"a\":\"model.safetensors\"}}" }); + try std.testing.expect((try refusal(NoLinks, a, io, dir, "a copy")) == null); + try tmp.dir.writeFile(io, .{ + .sub_path = "model.safetensors.index.json", + .data = "{\"weight_map\":{\"a\":\"model.safetensors\",\"b\":\"sub/b.safetensors\"}}", + }); + const why = (try refusal(NoLinks, a, io, dir, "a copy")) orelse return error.TestUnexpectedResult; + try std.testing.expect(std.mem.indexOf(u8, why, "sub/b.safetensors") != null); +}