lib/machine/src/checkpoint/roots/branch.zig

daab053ee43316e1809a84551d573ddd1e5bf3d2

  1 const os = @import("os");
  2 const std = @import("std");
  3 const owner = @import("owner.zig");
  4 const tree = @import("tree.zig");
  5 const types = @import("types.zig");
  6 
  7 const authentication_word_bits: usize = @bitSizeOf(u64);
  8 pub const authentication_word_count: usize =
  9     (types.page_count + authentication_word_bits - 1) /
 10     authentication_word_bits;
 11 
 12 const BranchError = error{
 13     BranchAddressOverflow,
 14     BranchOutOfBounds,
 15     BranchPageCapacityExceeded,
 16     BranchStorageMismatch,
 17 };
 18 
 19 /// Combines stored-capture failures together with branch range, capacity, and
 20 /// storage failures.
 21 pub const Error = types.Error || BranchError;
 22 
 23 /// Sizes the storage a branch needs for its private pages exactly. The byte
 24 /// counts cover one two-byte page index and one aligned 4096-byte page for
 25 /// every slot.
 26 pub const Capacity = struct {
 27     pages: u16,
 28     index_storage_bytes: u32,
 29     page_storage_bytes: u32,
 30 
 31     /// Derives the storage counts for `pages` slots. A count above 16,384
 32     /// returns `BranchPageCapacityExceeded`.
 33     pub fn derive(pages: usize) error{BranchPageCapacityExceeded}!Capacity {
 34         if (pages > types.page_count) return error.BranchPageCapacityExceeded;
 35         const index_bytes = std.math.mul(
 36             usize,
 37             pages,
 38             @sizeOf(u16),
 39         ) catch return error.BranchPageCapacityExceeded;
 40         const page_storage_bytes = std.math.mul(
 41             usize,
 42             pages,
 43             types.page_bytes,
 44         ) catch return error.BranchPageCapacityExceeded;
 45         return .{
 46             .pages = @intCast(pages),
 47             .index_storage_bytes = @intCast(index_bytes),
 48             .page_storage_bytes = @intCast(page_storage_bytes),
 49         };
 50     }
 51 };
 52 
 53 pub const maximum_capacity = Capacity.derive(types.page_count) catch unreachable;
 54 
 55 /// Holds the caller-owned memory a branch uses. The page pool holds private
 56 /// writes. A bitmap of fixed size and a digest array remember which shared
 57 /// pages have been authenticated.
 58 pub const Storage = struct {
 59     capacity: Capacity,
 60     indices: []u16,
 61     pages: []align(types.page_bytes) u8,
 62     authenticated: *[authentication_word_count]u64,
 63     digests: *[types.page_count]os.abi.Digest,
 64 
 65     /// Pairs each aligned private page with the index slot that names it. The
 66     /// authentication bitmap holds 256 words of 64 bits and the digest array
 67     /// holds one digest per page, both covering all 16,384 pages. More than
 68     /// 16,384 indices returns `BranchPageCapacityExceeded`. Any other
 69     /// page-storage length returns `BranchStorageMismatch`.
 70     pub fn init(
 71         indices: []u16,
 72         pages: []align(types.page_bytes) u8,
 73         authenticated: *[authentication_word_count]u64,
 74         digests: *[types.page_count]os.abi.Digest,
 75     ) Error!Storage {
 76         const capacity = try Capacity.derive(indices.len);
 77         if (pages.len != capacity.page_storage_bytes) {
 78             return error.BranchStorageMismatch;
 79         }
 80         return .{
 81             .capacity = capacity,
 82             .indices = indices,
 83             .pages = pages,
 84             .authenticated = authenticated,
 85             .digests = digests,
 86         };
 87     }
 88 };
 89 
 90 /// Provides a mutable memory view over one stored manifest root. A read
 91 /// authenticates a shared page the first time that page is touched. A page is
 92 /// copied into private storage the caller supplied when a write first touches
 93 /// it. Ownership of the provider and of every storage region stays with the
 94 /// caller.
 95 pub const Branch = struct {
 96     root_storage: types.Storage,
 97     root: types.ManifestRoot,
 98     manifest: types.Manifest,
 99     storage: Storage,
100     private_count: u16 = 0,
101     authenticated_count: u16 = 0,
102 
103     /// Fills a caller's output buffer from a byte range that lies inside
104     /// memory. A shared page is checked against its digest and the tree path
105     /// above it the first time the branch touches it. A read consumes no pool
106     /// slot. A range reaching past 67,108,864 bytes returns
107     /// `BranchOutOfBounds`, and an address whose sum overflows returns
108     /// `BranchAddressOverflow`.
109     pub fn read(
110         self: *@This(),
111         address: usize,
112         output: []u8,
113     ) Error!void {
114         const range = try checkedRange(address, output.len);
115         var cursor = range.start;
116         var output_offset: usize = 0;
117         while (cursor < range.end) {
118             const page_index: u16 = @intCast(cursor / types.page_bytes);
119             const page_offset = cursor % types.page_bytes;
120             const count = @min(types.page_bytes - page_offset, range.end - cursor);
121             var page: [types.page_bytes]u8 align(types.page_bytes) = undefined;
122             try self.readPage(page_index, &page);
123             @memcpy(output[output_offset..][0..count], page[page_offset..][0..count]);
124             cursor += count;
125             output_offset += count;
126         }
127         std.debug.assert(output_offset == output.len);
128     }
129 
130     /// Sends a caller's bytes into private pages over a range that lies inside
131     /// memory. Too little pool capacity is refused before any input byte is
132     /// written.
133     pub fn write(
134         self: *@This(),
135         address: usize,
136         input: []const u8,
137     ) Error!void {
138         const range = try checkedRange(address, input.len);
139         try self.preflightPrivate(range);
140         try self.materializeRange(range);
141         var cursor = range.start;
142         var input_offset: usize = 0;
143         while (cursor < range.end) {
144             const page_index: u16 = @intCast(cursor / types.page_bytes);
145             const page_offset = cursor % types.page_bytes;
146             const count = @min(types.page_bytes - page_offset, range.end - cursor);
147             const page = self.privatePage(page_index) orelse unreachable;
148             @memcpy(page[page_offset..][0..count], input[input_offset..][0..count]);
149             cursor += count;
150             input_offset += count;
151         }
152         std.debug.assert(input_offset == input.len);
153     }
154 
155     /// Sets every byte of a range inside memory to one value, working through
156     /// private pages. Too little pool capacity is refused before any
157     /// destination byte is changed.
158     pub fn fill(
159         self: *@This(),
160         address: usize,
161         byte_count: usize,
162         value: u8,
163     ) Error!void {
164         const range = try checkedRange(address, byte_count);
165         try self.preflightPrivate(range);
166         try self.materializeRange(range);
167         var cursor = range.start;
168         while (cursor < range.end) {
169             const page_index: u16 = @intCast(cursor / types.page_bytes);
170             const page_offset = cursor % types.page_bytes;
171             const count = @min(types.page_bytes - page_offset, range.end - cursor);
172             const page = self.privatePage(page_index) orelse unreachable;
173             @memset(page[page_offset..][0..count], value);
174             cursor += count;
175         }
176     }
177 
178     /// Delivers a whole page, named by its zero-based index, into an aligned
179     /// buffer the caller owns. A private page is copied from the pool, and any
180     /// other page is read and authenticated through the provider. An index of
181     /// 16,384 or more returns `BranchOutOfBounds`.
182     pub fn readPage(
183         self: *@This(),
184         page_index: u16,
185         output: *align(types.page_bytes) [types.page_bytes]u8,
186     ) Error!void {
187         if (page_index >= types.page_count) return error.BranchOutOfBounds;
188         if (self.privatePage(page_index)) |page| {
189             @memcpy(output, page);
190             return;
191         }
192         return self.readSharedPage(page_index, output);
193     }
194 
195     /// Hands back a page the caller can write into directly, named by its
196     /// zero-based index. The first access authenticates the shared page, copies
197     /// it into the pool, and consumes one pool slot. A full pool returns
198     /// `BranchPageCapacityExceeded`.
199     pub fn writablePage(
200         self: *@This(),
201         page_index: u16,
202     ) Error!*align(types.page_bytes) [types.page_bytes]u8 {
203         if (page_index >= types.page_count) return error.BranchOutOfBounds;
204         if (self.privatePage(page_index)) |page| return page;
205         if (self.private_count == self.storage.capacity.pages) {
206             return error.BranchPageCapacityExceeded;
207         }
208         var page: [types.page_bytes]u8 align(types.page_bytes) = undefined;
209         try self.readSharedPage(page_index, &page);
210         const slot = self.private_count;
211         const destination = pageAt(self.storage.pages, slot);
212         @memcpy(destination, &page);
213         self.storage.indices[slot] = page_index;
214         self.private_count += 1;
215         return destination;
216     }
217 
218     /// Makes every distinct page index in the list private. The whole list is
219     /// sized against the pool before a single slot is taken, and an index that
220     /// appears twice costs one slot. A caller reserves pages up front so a
221     /// later write cannot run out of pool.
222     pub fn preparePages(self: *@This(), pages: []const u16) Error!void {
223         var missing: usize = 0;
224         for (pages, 0..) |page_index, index| {
225             if (page_index >= types.page_count) return error.BranchOutOfBounds;
226             if (self.privatePage(page_index) != null) continue;
227             var duplicate = false;
228             for (pages[0..index]) |previous| {
229                 duplicate = duplicate or previous == page_index;
230             }
231             if (!duplicate) missing += 1;
232         }
233         const available = self.storage.capacity.pages - self.private_count;
234         if (missing > available) return error.BranchPageCapacityExceeded;
235         for (pages) |page_index| _ = try self.writablePage(page_index);
236     }
237 
238     /// Gives the branch a private copy of every page a byte range touches. The
239     /// range is sized against the pool before a single slot is taken.
240     pub fn prepareWrite(
241         self: *@This(),
242         address: usize,
243         byte_count: usize,
244     ) Error!void {
245         const range = try checkedRange(address, byte_count);
246         try self.preflightPrivate(range);
247         try self.materializeRange(range);
248     }
249 
250     /// Reports how many pool slots the branch has spent on private pages.
251     pub fn privatePageCount(self: *const @This()) u16 {
252         return self.private_count;
253     }
254 
255     /// Returns the branch metadata together with the authenticated-digest and
256     /// private-page bytes. The total counts the branch value, the whole
257     /// authentication bitmap, one digest for each authenticated page, and 4098
258     /// bytes for each private page.
259     pub fn residentBytes(self: *const @This()) u64 {
260         const authenticated_bytes = @as(u64, self.authenticated_count) *
261             @sizeOf(os.abi.Digest);
262         const private_bytes = @as(u64, self.private_count) *
263             (types.page_bytes + @sizeOf(u16));
264         return @sizeOf(Branch) + @sizeOf(@TypeOf(self.storage.authenticated.*)) +
265             authenticated_bytes + private_bytes;
266     }
267 
268     /// Answers whether a byte range touches the branch value itself or any of
269     /// the four caller-owned regions: the index slots, the page pool, the
270     /// authentication bitmap, and the digest cache. A caller checks this to
271     /// keep a buffer clear of everything the branch borrows.
272     pub fn aliases(self: *const @This(), bytes: []const u8) bool {
273         return buffersOverlap(bytes, std.mem.asBytes(self)) or
274             buffersOverlap(bytes, std.mem.sliceAsBytes(self.storage.indices)) or
275             buffersOverlap(bytes, self.storage.pages) or
276             buffersOverlap(bytes, std.mem.asBytes(self.storage.authenticated)) or
277             buffersOverlap(bytes, std.mem.asBytes(self.storage.digests));
278     }
279 
280     fn readSharedPage(
281         self: *@This(),
282         page_index: u16,
283         output: *align(types.page_bytes) [types.page_bytes]u8,
284     ) Error!void {
285         if (self.authenticated(page_index)) {
286             return tree.readAuthenticatedPage(
287                 self.root_storage,
288                 self.storage.digests[page_index],
289                 output,
290             );
291         }
292         const digest = try tree.readPageAt(
293             self.root_storage,
294             self.manifest.pages,
295             page_index,
296             output,
297         );
298         self.storage.digests[page_index] = digest;
299         self.markAuthenticated(page_index);
300     }
301 
302     fn privatePage(
303         self: *const @This(),
304         page_index: u16,
305     ) ?*align(types.page_bytes) [types.page_bytes]u8 {
306         for (self.storage.indices[0..self.private_count], 0..) |index, slot| {
307             if (index == page_index) return pageAt(self.storage.pages, slot);
308         }
309         return null;
310     }
311 
312     fn preflightPrivate(self: *const @This(), range: Range) Error!void {
313         if (range.start == range.end) return;
314         const first = range.start / types.page_bytes;
315         const last = (range.end - 1) / types.page_bytes;
316         var missing: usize = 0;
317         for (first..last + 1) |page_index| {
318             if (self.privatePage(@intCast(page_index)) == null) missing += 1;
319         }
320         const available = self.storage.capacity.pages - self.private_count;
321         if (missing > available) return error.BranchPageCapacityExceeded;
322     }
323 
324     fn materializeRange(self: *@This(), range: Range) Error!void {
325         if (range.start == range.end) return;
326         const first = range.start / types.page_bytes;
327         const last = (range.end - 1) / types.page_bytes;
328         for (first..last + 1) |page_index| {
329             _ = try self.writablePage(@intCast(page_index));
330         }
331     }
332 
333     fn authenticated(self: *const @This(), page_index: u16) bool {
334         const word = page_index / authentication_word_bits;
335         const bit: u6 = @intCast(page_index % authentication_word_bits);
336         return self.storage.authenticated[word] & (@as(u64, 1) << bit) != 0;
337     }
338 
339     fn markAuthenticated(self: *@This(), page_index: u16) void {
340         const word = page_index / authentication_word_bits;
341         const bit: u6 = @intCast(page_index % authentication_word_bits);
342         std.debug.assert(!self.authenticated(page_index));
343         self.storage.authenticated[word] |= @as(u64, 1) << bit;
344         self.authenticated_count += 1;
345     }
346 };
347 
348 const Range = struct {
349     start: usize,
350     end: usize,
351 };
352 
353 /// Opens one stored manifest root for shared-page access. The function
354 /// authenticates the manifest chain and its delta relationships, leaving each
355 /// page for the read that first touches it. The call validates the private pool
356 /// and clears the authentication bitmap, so no page starts out trusted. The
357 /// returned branch borrows every supplied resource.
358 pub fn restore(
359     root_storage: types.Storage,
360     root: types.ManifestRoot,
361     storage: Storage,
362 ) Error!Branch {
363     try validateStorage(storage);
364     const value = try owner.openShared(root_storage, root);
365     @memset(storage.authenticated, 0);
366     return .{
367         .root_storage = root_storage,
368         .root = root,
369         .manifest = value,
370         .storage = storage,
371     };
372 }
373 
374 fn validateStorage(storage: Storage) Error!void {
375     const actual = try Capacity.derive(storage.indices.len);
376     if (!std.meta.eql(actual, storage.capacity) or
377         storage.pages.len != storage.capacity.page_storage_bytes)
378     {
379         return error.BranchStorageMismatch;
380     }
381 }
382 
383 fn checkedRange(address: usize, byte_count: usize) Error!Range {
384     const end = std.math.add(usize, address, byte_count) catch
385         return error.BranchAddressOverflow;
386     const ram_bytes = types.page_count * types.page_bytes;
387     if (address > ram_bytes or end > ram_bytes) return error.BranchOutOfBounds;
388     return .{ .start = address, .end = end };
389 }
390 
391 fn pageAt(
392     pages: []align(types.page_bytes) u8,
393     slot: usize,
394 ) *align(types.page_bytes) [types.page_bytes]u8 {
395     const start = slot * types.page_bytes;
396     std.debug.assert(start + types.page_bytes <= pages.len);
397     return @ptrCast(@alignCast(pages[start..][0..types.page_bytes]));
398 }
399 
400 fn buffersOverlap(left: []const u8, right: []const u8) bool {
401     if (left.len == 0 or right.len == 0) return false;
402     const left_start = @intFromPtr(left.ptr);
403     const right_start = @intFromPtr(right.ptr);
404     const left_end = std.math.add(usize, left_start, left.len) catch return true;
405     const right_end = std.math.add(usize, right_start, right.len) catch return true;
406     return left_start < right_end and right_start < left_end;
407 }
408 
409 test "branch capacity rejects max plus one" {
410     try std.testing.expectEqual(
411         @as(u16, types.page_count),
412         maximum_capacity.pages,
413     );
414     try std.testing.expectError(
415         error.BranchPageCapacityExceeded,
416         Capacity.derive(types.page_count + 1),
417     );
418 }
419 
420 comptime {
421     std.debug.assert(types.page_count % authentication_word_bits == 0);
422     std.debug.assert(types.page_count <= std.math.maxInt(u16));
423     std.debug.assert(authentication_word_count == 256);
424 }