-
Notifications
You must be signed in to change notification settings - Fork 3
fix: comprehensive audit — fix 14 critical and 8 high-severity bugs #99
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from all commits
9b3a1e2
82dc5cd
ac9f366
95d54b3
ab8d669
54addd2
2f16de2
5190777
806fd36
cd77279
e2bdf26
d9b3f7f
3fdffe2
9d7675a
491e7b0
5b09556
ea3e173
f3b26de
547e3de
9ccd975
98a84f2
9d5efc7
42f5366
09a0f24
811538d
bc0ba24
6277cfa
ddb6727
c894d48
22c366d
fcd8047
1b35cf7
cdb6663
eeaf44e
81dcab9
4b75ce0
f628ede
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,91 @@ | ||
| """WebSocket streaming bulk insert for TurboDB. | ||
|
|
||
| Fire-and-forget protocol: client streams NDJSON frames without waiting | ||
| for per-frame responses. ~10-20x faster than HTTP batches over high-latency | ||
| connections (SSH tunnels, WAN). | ||
|
|
||
| Usage:: | ||
|
|
||
| from turbodb.stream import StreamInserter | ||
|
|
||
| with StreamInserter("ws://host:port", "mycollection") as s: | ||
| for doc in docs: | ||
| s.add(doc["key"], doc["value"]) | ||
| print(s.result) # {"inserted": 12000, "errors": 0, "frames": 6} | ||
|
|
||
| Requires: pip install websocket-client | ||
| """ | ||
|
|
||
| import json | ||
|
|
||
| try: | ||
| import websocket | ||
| except ImportError: | ||
| websocket = None | ||
|
|
||
|
|
||
| class StreamInserter: | ||
| """Fire-and-forget WebSocket streaming bulk inserter.""" | ||
|
|
||
| def __init__(self, url, collection, tenant=None, batch_size=2000): | ||
| if websocket is None: | ||
| raise ImportError("pip install websocket-client") | ||
| self.url = url | ||
| self.collection = collection | ||
| self.tenant = tenant | ||
| self.batch_size = batch_size | ||
| self._ws = None | ||
| self._buf = [] | ||
| self._sent_frames = 0 | ||
| self.result = None | ||
|
|
||
| def open(self): | ||
| """Connect and enter STREAM mode.""" | ||
| self._ws = websocket.create_connection(self.url) | ||
| init = f"STREAM /db/{self.collection}/bulk" | ||
| if self.tenant: | ||
| init += f" tenant={self.tenant}" | ||
| self._ws.send(init) | ||
| ack = json.loads(self._ws.recv()) | ||
| if "error" in ack: | ||
| raise RuntimeError(f"Stream init failed: {ack['error']}") | ||
| return self | ||
|
|
||
| def add(self, key, value): | ||
| """Buffer a document. Auto-flushes at batch_size.""" | ||
| if isinstance(value, (dict, list)): | ||
| value = json.dumps(value) | ||
| self._buf.append(json.dumps({"key": key, "value": value})) | ||
| if len(self._buf) >= self.batch_size: | ||
| self.flush() | ||
|
|
||
| def send_ndjson(self, ndjson_str): | ||
| """Send a pre-built NDJSON string as one frame (fire-and-forget).""" | ||
| self._ws.send(ndjson_str) | ||
| self._sent_frames += 1 | ||
|
|
||
| def flush(self): | ||
| """Send buffered docs as one NDJSON frame.""" | ||
| if not self._buf: | ||
| return | ||
| self._ws.send("\n".join(self._buf)) | ||
| self._sent_frames += 1 | ||
| self._buf.clear() | ||
|
|
||
| def close(self): | ||
| """Flush remaining, close WS, return final summary.""" | ||
| self.flush() | ||
| self._ws.send_close() | ||
| try: | ||
| self.result = json.loads(self._ws.recv()) | ||
| except Exception: | ||
| self.result = {"inserted": 0, "errors": 0, "frames": self._sent_frames} | ||
| self._ws.close() | ||
| return self.result | ||
|
|
||
| def __enter__(self): | ||
| self.open() | ||
| return self | ||
|
|
||
| def __exit__(self, *args): | ||
| self.close() | ||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -22,13 +22,12 @@ pub const Branch = struct { | |
| agent_id: [64]u8, // which agent owns this branch | ||
| agent_id_len: u8, | ||
|
|
||
| // Branch-local storage: key_hash -> value (only modified keys) | ||
| // Branch-local storage: actual key string -> value (only modified keys) | ||
| // This is the CoW layer — unmodified keys fall through to main | ||
| writes: std.AutoHashMap(u64, BranchWrite), | ||
| writes: std.StringHashMap(BranchWrite), | ||
| allocator: Allocator, | ||
|
|
||
| pub const BranchWrite = struct { | ||
| key: []const u8, // owned copy | ||
| value: []const u8, // owned copy | ||
| deleted: bool, // true = tombstone (deleted on branch) | ||
| epoch: u64, // when this write happened | ||
|
|
@@ -45,25 +44,30 @@ pub const Branch = struct { | |
| pub fn deinit(self: *Branch) void { | ||
| var it = self.writes.iterator(); | ||
| while (it.next()) |entry| { | ||
| if (entry.value_ptr.key.len > 0) self.allocator.free(entry.value_ptr.key); | ||
| if (entry.value_ptr.value.len > 0) self.allocator.free(entry.value_ptr.value); | ||
| self.allocator.free(@constCast(entry.key_ptr.*)); | ||
| } | ||
| self.writes.deinit(); | ||
| } | ||
|
|
||
| /// Write a key-value pair on this branch (CoW — only stores the delta) | ||
| pub fn write(self: *Branch, key: []const u8, value: []const u8, epoch: u64) !void { | ||
| const key_hash = fnv1a(key); | ||
| // Free old write if exists | ||
| if (self.writes.getPtr(key_hash)) |old| { | ||
| if (old.key.len > 0) self.allocator.free(old.key); | ||
| if (old.value.len > 0) self.allocator.free(old.value); | ||
| } | ||
| // Allocate new copies BEFORE freeing old ones — if alloc fails, | ||
| // the existing entry stays valid. | ||
| const owned_key = try self.allocator.dupe(u8, key); | ||
| errdefer self.allocator.free(owned_key); | ||
| const owned_val = try self.allocator.dupe(u8, value); | ||
| try self.writes.put(key_hash, .{ | ||
| .key = owned_key, | ||
| errdefer self.allocator.free(owned_val); | ||
| // Free old write if exists (safe — new copies already allocated) | ||
| if (self.writes.getPtr(key)) |old| { | ||
| if (old.value.len > 0) self.allocator.free(old.value); | ||
| // Free the old key that was used as the map key. | ||
| const old_map_key = self.writes.getKey(key).?; | ||
| self.allocator.free(@constCast(old_map_key)); | ||
| // Remove old entry so we can insert with new owned key. | ||
| _ = self.writes.remove(key); | ||
|
Comment on lines
+65
to
+68
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
In Useful? React with 👍 / 👎.
Comment on lines
+65
to
+68
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
When updating an existing branch key, this code frees Useful? React with 👍 / 👎. |
||
| } | ||
| try self.writes.put(owned_key, .{ | ||
| .value = owned_val, | ||
| .deleted = false, | ||
| .epoch = epoch, | ||
|
|
@@ -72,14 +76,14 @@ pub const Branch = struct { | |
|
|
||
| /// Mark a key as deleted on this branch | ||
| pub fn delete(self: *Branch, key: []const u8, epoch: u64) !void { | ||
| const key_hash = fnv1a(key); | ||
| if (self.writes.getPtr(key_hash)) |old| { | ||
| if (old.key.len > 0) self.allocator.free(old.key); | ||
| if (self.writes.getPtr(key)) |old| { | ||
| if (old.value.len > 0) self.allocator.free(old.value); | ||
| const old_map_key = self.writes.getKey(key).?; | ||
| self.allocator.free(@constCast(old_map_key)); | ||
| _ = self.writes.remove(key); | ||
| } | ||
| const owned_key = try self.allocator.dupe(u8, key); | ||
| try self.writes.put(key_hash, .{ | ||
| .key = owned_key, | ||
| try self.writes.put(owned_key, .{ | ||
| .value = &.{}, | ||
| .deleted = true, | ||
| .epoch = epoch, | ||
|
|
@@ -88,8 +92,7 @@ pub const Branch = struct { | |
|
|
||
| /// Read a key on this branch. Returns branch-local value or null (fall through to main). | ||
| pub fn read(self: *const Branch, key: []const u8) ?BranchRead { | ||
| const key_hash = fnv1a(key); | ||
| if (self.writes.get(key_hash)) |w| { | ||
| if (self.writes.get(key)) |w| { | ||
| if (w.deleted) return .{ .deleted = true, .value = null }; | ||
| return .{ .deleted = false, .value = w.value }; | ||
| } | ||
|
|
@@ -108,7 +111,7 @@ pub const Branch = struct { | |
| while (it.next()) |entry| { | ||
| const w = entry.value_ptr.*; | ||
| try entries.append(alloc, .{ | ||
| .key = w.key, | ||
| .key = entry.key_ptr.*, | ||
| .value = w.value, | ||
| .deleted = w.deleted, | ||
| .epoch = w.epoch, | ||
|
|
@@ -165,7 +168,7 @@ pub fn compareBranches(branches: []*const Branch, alloc: Allocator) !CompareResu | |
| for (branches) |br| { | ||
| var it = br.writes.iterator(); | ||
| while (it.next()) |entry| { | ||
| try all_keys.put(entry.value_ptr.key, {}); | ||
| try all_keys.put(entry.key_ptr.*, {}); | ||
| } | ||
| } | ||
|
|
||
|
|
@@ -183,7 +186,7 @@ pub fn compareBranches(branches: []*const Branch, alloc: Allocator) !CompareResu | |
| .agent_id = br.getAgentId(), | ||
| .value = r.value, | ||
| .deleted = r.deleted, | ||
| .epoch = if (br.writes.get(fnv1a(key))) |w| w.epoch else 0, | ||
| .epoch = if (br.writes.get(key)) |w| w.epoch else 0, | ||
| }); | ||
| } | ||
| } | ||
|
|
@@ -432,7 +435,7 @@ pub const BranchManager = struct { | |
| .status = .active, | ||
| .agent_id = undefined, | ||
| .agent_id_len = aid_len, | ||
| .writes = std.AutoHashMap(u64, Branch.BranchWrite).init(self.allocator), | ||
| .writes = std.StringHashMap(Branch.BranchWrite).init(self.allocator), | ||
| .allocator = self.allocator, | ||
| }; | ||
| @memcpy(branch.name[0..name_len], name_arg[0..name_len]); | ||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
When
valueis adict/list,add()serializes it first and then serializes the outer document again, producing NDJSON like"value":"{...}"instead of"value":{...}. On the server side this gets stored as a JSON string literal rather than the original object/array, so streamed inserts have different semantics than normal client inserts for the same Python input.Useful? React with 👍 / 👎.