biondizzle/zyna-db
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix allocator again

Browse Source
This commit is contained in:
biondizzle
2026-02-15 07:20:28 -05:00
parent 3f7f4b6a6e
commit e4b82c247a
3 changed files with 130 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files

32
src/dynamodb/handler.zig
View File

@@ -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,

84
src/dynamodb/storage.zig
View File

@@ -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,

35
src/item_codec.zig
View File

@@ -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);
}