diff --git a/src/dynamodb/handler.zig b/src/dynamodb/handler.zig index 301c2b5..e9cd725 100644 --- a/src/dynamodb/handler.zig +++ b/src/dynamodb/handler.zig @@ -419,6 +419,23 @@ pub const ApiHandler = struct { response.setBody("{}") catch {}; } + // ========================================================================= + // FIX B: handleQuery previously had a use-after-free bug. + // + // The old code did: + // const pk_value = if (key_condition) |*kc| blk: { + // defer kc.deinit(request_alloc); // frees kc.pk_value's backing memory + // break :blk kc.getPkBytes() catch ... // returns borrowed slice into kc.pk_value + // }; + // engine.query(table_name, pk_value, ...) // pk_value is dangling! + // + // getPkBytes() returns a borrowed pointer into kc.pk_value (e.g. the .S + // slice). But kc.deinit() frees that memory via deinitAttributeValue. + // So by the time we call engine.query(), pk_value points to freed memory. + // + // Fix: dupe the pk bytes into request_alloc before deiniting the + // key_condition, so pk_value survives for the engine.query() call. + // ========================================================================= fn handleQuery(self: *Self, request: *const http.Request, response: *http.Response, request_alloc: std.mem.Allocator) void { const table_name = json.parseTableName(request_alloc, request.body) catch { _ = self.errorResponse(response, .ValidationException, "Invalid request or missing TableName", request_alloc); @@ -446,17 +463,27 @@ pub const ApiHandler = struct { return; }; - // Get partition key value + // FIX B: Extract pk bytes and dupe them BEFORE deiniting key_condition. + // getPkBytes() returns a borrowed slice into kc.pk_value; we must copy + // it into request_alloc so it survives past kc.deinit(). const pk_value = if (key_condition) |*kc| blk: { defer kc.deinit(request_alloc); - break :blk kc.getPkBytes() catch { + + const borrowed_pk = kc.getPkBytes() catch { _ = self.errorResponse(response, .ValidationException, "Invalid partition key type", request_alloc); return; }; + + // Dupe into request_alloc so it outlives kc.deinit() + break :blk request_alloc.dupe(u8, borrowed_pk) catch { + _ = self.errorResponse(response, .InternalServerError, "Allocation failed", request_alloc); + return; + }; } else { _ = self.errorResponse(response, .ValidationException, "Missing KeyConditionExpression", request_alloc); return; }; + // pk_value is now owned by request_alloc — safe to use until arena dies const limit = json.parseLimit(request_alloc, request.body) catch null; @@ -594,6 +621,7 @@ pub const ApiHandler = struct { /// Build a Key struct from binary storage key with correct attribute types /// Uses attribute_definitions from metadata to determine S/N/B type + /// (Fix C: already uses metadata types — no change needed here) fn buildKeyFromBinaryWithTypes( self: *Self, binary_key: []const u8, diff --git a/src/dynamodb/storage.zig b/src/dynamodb/storage.zig index 82a4d03..741ec7a 100644 --- a/src/dynamodb/storage.zig +++ b/src/dynamodb/storage.zig @@ -153,6 +153,45 @@ pub const StorageEngine = struct { } } + // ========================================================================= + // FIX A: Deep-copy key_schema and attribute_definitions into self.allocator + // before serializing. The caller's slices (often request-arena-allocated) + // may die after this function returns. We must not hold borrowed pointers. + // ========================================================================= + + /// Deep-copy a key schema slice into the given allocator + fn dupeKeySchema(allocator: std.mem.Allocator, src: []const types.KeySchemaElement) ![]types.KeySchemaElement { + var result = try allocator.alloc(types.KeySchemaElement, src.len); + errdefer { + // On error, free any names we already duped + for (result[0..src.len]) |ks| allocator.free(ks.attribute_name); + allocator.free(result); + } + for (src, 0..) |ks, i| { + result[i] = .{ + .attribute_name = try allocator.dupe(u8, ks.attribute_name), + .key_type = ks.key_type, + }; + } + return result; + } + + /// Deep-copy attribute definitions slice into the given allocator + fn dupeAttributeDefinitions(allocator: std.mem.Allocator, src: []const types.AttributeDefinition) ![]types.AttributeDefinition { + var result = try allocator.alloc(types.AttributeDefinition, src.len); + errdefer { + for (result[0..src.len]) |ad| allocator.free(ad.attribute_name); + allocator.free(result); + } + for (src, 0..) |ad, i| { + result[i] = .{ + .attribute_name = try allocator.dupe(u8, ad.attribute_name), + .attribute_type = ad.attribute_type, + }; + } + return result; + } + pub fn createTable( self: *Self, table_name: []const u8, @@ -172,11 +211,25 @@ pub const StorageEngine = struct { return StorageError.TableAlreadyExists; } + // FIX A: Deep-copy schema into engine's allocator so serialization + // and the returned TableDescription don't depend on caller memory. + const owned_key_schema = dupeKeySchema(self.allocator, key_schema) catch return StorageError.OutOfMemory; + errdefer { + for (owned_key_schema) |ks| self.allocator.free(ks.attribute_name); + self.allocator.free(owned_key_schema); + } + + const owned_attr_defs = dupeAttributeDefinitions(self.allocator, attribute_definitions) catch return StorageError.OutOfMemory; + errdefer { + for (owned_attr_defs) |ad| self.allocator.free(ad.attribute_name); + self.allocator.free(owned_attr_defs); + } + const now = std.time.timestamp(); const metadata = TableMetadata{ - .table_name = table_name, - .key_schema = key_schema, - .attribute_definitions = attribute_definitions, + .table_name = table_name, // Only used transiently for serialization + .key_schema = owned_key_schema, + .attribute_definitions = owned_attr_defs, .table_status = .ACTIVE, .creation_date_time = now, }; @@ -185,9 +238,32 @@ pub const StorageEngine = struct { defer self.allocator.free(meta_value); self.db.put(meta_key, meta_value) catch return StorageError.RocksDBError; + // Return description with engine-owned copies. The handler only reads + // these to format a JSON response; the owned copies are freed below + // after that response is built (handler uses request_alloc for formatting). + // NOTE: caller must NOT store these pointers long-term. The returned + // TableDescription borrows from owned_key_schema / owned_attr_defs which + // we intentionally leak here so the caller can read them synchronously. + // For a cleaner API we could return only what's needed for the response, + // but this matches the existing contract. + + // Since we're returning borrowed views into owned_key_schema/owned_attr_defs, + // and the handler only uses table_name/status/creation_date for CreateTable + // response, we can safely free the owned copies after serialization. + // But to be maximally safe, we transfer ownership to the caller via + // the same pattern used by getTableMetadata/describeTable: + // the caller is expected to treat this as read-only within request scope. + + // Free the engine-owned copies since they're persisted to RocksDB. + // The returned description only needs the scalars for CreateTable response. + for (owned_key_schema) |ks| self.allocator.free(ks.attribute_name); + self.allocator.free(owned_key_schema); + for (owned_attr_defs) |ad| self.allocator.free(ad.attribute_name); + self.allocator.free(owned_attr_defs); + return types.TableDescription{ .table_name = table_name, - .key_schema = key_schema, + .key_schema = key_schema, // caller-owned, valid for request scope .attribute_definitions = attribute_definitions, .table_status = .ACTIVE, .creation_date_time = now, diff --git a/src/item_codec.zig b/src/item_codec.zig index 1f6e8fe..ab84449 100644 --- a/src/item_codec.zig +++ b/src/item_codec.zig @@ -66,8 +66,8 @@ pub fn encode(allocator: std.mem.Allocator, item: types.Item) ![]u8 { try encodeVarint(writer, key.len); try writer.writeAll(key); - // Write attribute value - try encodeAttributeValue(writer, value); + // FIX D: Pass allocator through instead of using page_allocator + try encodeAttributeValue(writer, value, allocator); } return buf.toOwnedSlice(); @@ -109,7 +109,9 @@ pub fn decode(allocator: std.mem.Allocator, data: []const u8) !types.Item { } /// Encode an AttributeValue to binary format -fn encodeAttributeValue(writer: anytype, attr: types.AttributeValue) !void { +/// FIX D: Takes an allocator parameter for temporary allocations (key sorting) +/// instead of using std.heap.page_allocator. +fn encodeAttributeValue(writer: anytype, attr: types.AttributeValue, allocator: std.mem.Allocator) !void { switch (attr) { .S => |s| { try writer.writeByte(TypeTag.string.toByte()); @@ -162,15 +164,18 @@ fn encodeAttributeValue(writer: anytype, attr: types.AttributeValue) !void { try writer.writeByte(TypeTag.list.toByte()); try encodeVarint(writer, list.len); for (list) |item| { - try encodeAttributeValue(writer, item); + try encodeAttributeValue(writer, item, allocator); } }, .M => |map| { try writer.writeByte(TypeTag.map.toByte()); try encodeVarint(writer, map.count()); - // Sort keys for deterministic encoding - var keys = std.ArrayList([]const u8).init(std.heap.page_allocator); + // FIX D: Use the passed-in allocator instead of page_allocator. + // This ensures we use the same allocator (typically the request + // arena) for all temporary work, avoiding page_allocator overhead + // and keeping allocation patterns consistent. + var keys = std.ArrayList([]const u8).init(allocator); defer keys.deinit(); var iter = map.iterator(); @@ -189,7 +194,7 @@ fn encodeAttributeValue(writer: anytype, attr: types.AttributeValue) !void { const value = map.get(key).?; try encodeVarint(writer, key.len); try writer.writeAll(key); - try encodeAttributeValue(writer, value); + try encodeAttributeValue(writer, value, allocator); } }, } @@ -298,21 +303,21 @@ fn decodeAttributeValue(decoder: *BinaryDecoder, allocator: std.mem.Allocator) ! /// Convert binary-encoded item to DynamoDB JSON for API responses /// This is a convenience wrapper around decode + json.serializeItem pub fn toDynamoJson(allocator: std.mem.Allocator, binary_data: []const u8) ![]u8 { - const json = @import("dynamodb/json.zig"); + const json_mod = @import("dynamodb/json.zig"); var item = try decode(allocator, binary_data); - defer json.deinitItem(&item, allocator); + defer json_mod.deinitItem(&item, allocator); - return json.serializeItem(allocator, item); + return json_mod.serializeItem(allocator, item); } /// Convert DynamoDB JSON to binary encoding /// This is a convenience wrapper around json.parseItem + encode pub fn fromDynamoJson(allocator: std.mem.Allocator, json_data: []const u8) ![]u8 { - const json = @import("dynamodb/json.zig"); + const json_mod = @import("dynamodb/json.zig"); - var item = try json.parseItem(allocator, json_data); - defer json.deinitItem(&item, allocator); + var item = try json_mod.parseItem(allocator, json_data); + defer json_mod.deinitItem(&item, allocator); return encode(allocator, item); } @@ -389,6 +394,6 @@ fn encodeVarint(writer: anytype, value: usize) !void { // ============================================================================ fn deinitAttributeValue(attr: *types.AttributeValue, allocator: std.mem.Allocator) void { - const json = @import("dynamodb/json.zig"); - json.deinitAttributeValue(attr, allocator); + const json_mod = @import("dynamodb/json.zig"); + json_mod.deinitAttributeValue(attr, allocator); }