diff --git a/src/float.zig b/src/float.zig index 4858bbc..0544df1 100644 --- a/src/float.zig +++ b/src/float.zig @@ -4,6 +4,7 @@ const hdrs = @import("headers.zig"); const NonOptional = @import("utils.zig").NonOptional; const maybePackNull = @import("null.zig").maybePackNull; const packHeaderAndInt = @import("utils.zig").packHeaderAndInt; +const takeInt = @import("utils.zig").takeInt; const maybeUnpackNull = @import("null.zig").maybeUnpackNull; inline fn assertFloatType(comptime T: type) type { @@ -49,12 +50,8 @@ pub fn packFloat(writer: *std.Io.Writer, comptime T: type, value_or_maybe_null: } pub fn readFloatValue(reader: *std.Io.Reader, comptime SourceFloat: type, comptime TargetFloat: type) !TargetFloat { - const size = @sizeOf(SourceFloat); - var buf: [size]u8 = undefined; - try reader.readSliceAll(&buf); - const IntType = std.meta.Int(.unsigned, @bitSizeOf(SourceFloat)); - const int_value = std.mem.readInt(IntType, &buf, .big); + const int_value = try takeInt(reader, IntType); return @floatCast(@as(SourceFloat, @bitCast(int_value))); } diff --git a/src/int.zig b/src/int.zig index 1b61ec8..3491e41 100644 --- a/src/int.zig +++ b/src/int.zig @@ -4,6 +4,7 @@ const hdrs = @import("headers.zig"); const NonOptional = @import("utils.zig").NonOptional; const maybePackNull = @import("null.zig").maybePackNull; const packHeaderAndInt = @import("utils.zig").packHeaderAndInt; +const takeInt = @import("utils.zig").takeInt; const maybeUnpackNull = @import("null.zig").maybeUnpackNull; inline fn assertIntType(comptime T: type) type { @@ -155,10 +156,7 @@ pub fn unpackShortIntValue(header: u8, min_value: u8, max_value: u8, comptime Ta } pub fn unpackIntValue(reader: *std.Io.Reader, comptime SourceType: type, comptime TargetType: type) !TargetType { - const size = @divExact(@bitSizeOf(SourceType), 8); - var buf: [size]u8 = undefined; - try reader.readSliceAll(&buf); - const value = std.mem.readInt(SourceType, &buf, .big); + const value = try takeInt(reader, SourceType); const source_type_info = @typeInfo(SourceType).int; const target_type_info = @typeInfo(TargetType).int; diff --git a/src/msgpack.zig b/src/msgpack.zig index 50aebdf..22a3632 100644 --- a/src/msgpack.zig +++ b/src/msgpack.zig @@ -43,6 +43,7 @@ pub const packStringLiteral = @import("string.zig").packStringLiteral; pub const unpackStringHeader = @import("string.zig").unpackStringHeader; pub const unpackString = @import("string.zig").unpackString; pub const unpackStringInto = @import("string.zig").unpackStringInto; +pub const unpackStringBorrowed = @import("string.zig").unpackStringBorrowed; pub const packBinaryHeader = @import("binary.zig").packBinaryHeader; pub const packBinary = @import("binary.zig").packBinary; @@ -287,6 +288,120 @@ test { _ = std.testing.refAllDecls(@This()); } +/// A reader that yields one byte per fill, out of a caller-sized buffer. Models +/// a trickling network stream, where a value can straddle a fill and the +/// reader's buffer can be smaller than the value being decoded. `Reader.fixed` +/// cannot express either case, since its buffer is the whole payload. +const TrickleReader = struct { + data: []const u8, + pos: usize = 0, + reader: std.Io.Reader, + + fn init(buffer: []u8, data: []const u8) TrickleReader { + return .{ + .data = data, + .reader = .{ + .vtable = &.{ + .stream = stream, + .discard = discard, + .readVec = readVec, + .rebase = std.Io.Reader.defaultRebase, + }, + .buffer = buffer, + .seek = 0, + .end = 0, + }, + }; + } + + fn readVec(r: *std.Io.Reader, data: [][]u8) std.Io.Reader.Error!usize { + _ = data; + const self: *TrickleReader = @fieldParentPtr("reader", r); + if (self.pos >= self.data.len) return error.EndOfStream; + if (r.end >= r.buffer.len) return 0; + r.buffer[r.end] = self.data[self.pos]; + r.end += 1; + self.pos += 1; + return 1; + } + + fn stream(_: *std.Io.Reader, _: *std.Io.Writer, _: std.Io.Limit) std.Io.Reader.StreamError!usize { + return error.EndOfStream; + } + + fn discard(_: *std.Io.Reader, _: std.Io.Limit) std.Io.Reader.Error!usize { + return error.EndOfStream; + } +}; + +fn encodeToBuffer(value: anytype, buffer: []u8) ![]const u8 { + var writer = std.Io.Writer.fixed(buffer); + try encode(value, &writer); + return writer.buffered(); +} + +test "decode from a streaming reader whose buffer holds the largest value" { + const Numeric = struct { id: u64, seq: i32, ratio: f64, ok: bool }; + const value = Numeric{ .id = 0xdead_beef_cafe_1234, .seq = -4242, .ratio = 3.14159, .ok = true }; + + var encoded: [64]u8 = undefined; + const bytes = try encodeToBuffer(value, &encoded); + + // 8 bytes is the largest fixed-size payload (u64/f64), and every field name + // here is a fixstr short enough to fit alongside its header. + for ([_]usize{ 8, 9, 16, 64 }) |buffer_len| { + var buffer: [64]u8 = undefined; + var trickle = TrickleReader.init(buffer[0..buffer_len], bytes); + const decoded = try decodeLeaky(Numeric, std.testing.allocator, &trickle.reader); + try std.testing.expectEqualDeep(value, decoded); + } +} + +test "decode from a streaming reader with too small a buffer errors, never panics" { + const Numeric = struct { id: u64, seq: i32, ratio: f64, ok: bool }; + const value = Numeric{ .id = 0xdead_beef_cafe_1234, .seq = -4242, .ratio = 3.14159, .ok = true }; + + var encoded: [64]u8 = undefined; + const bytes = try encodeToBuffer(value, &encoded); + + // The u64 payload needs 8 buffered bytes; anything smaller cannot rebase. + for ([_]usize{ 1, 2, 4, 7 }) |buffer_len| { + var buffer: [8]u8 = undefined; + var trickle = TrickleReader.init(buffer[0..buffer_len], bytes); + try std.testing.expectError( + error.ReaderBufferTooSmall, + decodeLeaky(Numeric, std.testing.allocator, &trickle.reader), + ); + } +} + +test "decode a key too long for the reader buffer errors, never panics" { + // 33 characters, so the key is a str8 rather than a fixstr. + const LongKey = struct { + this_field_name_is_longer_than_32: u8, + }; + const value = LongKey{ .this_field_name_is_longer_than_32 = 7 }; + + var encoded: [64]u8 = undefined; + const bytes = try encodeToBuffer(value, &encoded); + + { + var buffer: [16]u8 = undefined; + var trickle = TrickleReader.init(&buffer, bytes); + try std.testing.expectError( + error.ReaderBufferTooSmall, + decodeLeaky(LongKey, std.testing.allocator, &trickle.reader), + ); + } + + { + var buffer: [64]u8 = undefined; + var trickle = TrickleReader.init(&buffer, bytes); + const decoded = try decodeLeaky(LongKey, std.testing.allocator, &trickle.reader); + try std.testing.expectEqualDeep(value, decoded); + } +} + test "encode/decode" { const Message = struct { name: []const u8, diff --git a/src/string.zig b/src/string.zig index 7331b29..8506590 100644 --- a/src/string.zig +++ b/src/string.zig @@ -101,6 +101,35 @@ pub fn unpackString(reader: *std.Io.Reader, allocator: std.mem.Allocator) ![]u8 return data; } +/// Unpacks a string without copying it, borrowing the reader's buffer. The +/// result is only valid until the next read, so it suits keys that are compared +/// and discarded, not values that are kept. The reader's buffer must be able to +/// hold the whole string. +pub fn unpackStringBorrowed(reader: *std.Io.Reader) ![]const u8 { + // A buffered fixstr, which is what a map key almost always is, needs no + // header parse and no copy. + const buffered = reader.buffer[reader.seek..reader.end]; + if (buffered.len > 0 and buffered[0] >= hdrs.FIXSTR_MIN and buffered[0] <= hdrs.FIXSTR_MAX) { + @branchHint(.likely); + const len = buffered[0] - hdrs.FIXSTR_MIN; + if (1 + len <= buffered.len) { + reader.seek += 1 + len; + return buffered[1..][0..len]; + } + } + + const len = try unpackStringHeader(reader, u32); + + // `take` rebases, which asserts the reader's buffer can hold `len`; report + // that as an error rather than letting the assert fire. + if (len > reader.buffer.len) { + @branchHint(.unlikely); + return error.ReaderBufferTooSmall; + } + + return reader.take(len); +} + pub fn unpackStringInto(reader: *std.Io.Reader, buf: []u8) ![]u8 { const len = try unpackStringHeader(reader, u32); diff --git a/src/struct.zig b/src/struct.zig index 4793385..183febd 100644 --- a/src/struct.zig +++ b/src/struct.zig @@ -17,7 +17,8 @@ const packInt = @import("int.zig").packInt; const unpackInt = @import("int.zig").unpackInt; const packStringLiteral = @import("string.zig").packStringLiteral; -const unpackStringInto = @import("string.zig").unpackStringInto; +const unpackStringBorrowed = @import("string.zig").unpackStringBorrowed; +const eqlLiteral = @import("utils.zig").eqlLiteral; const packArrayHeader = @import("array.zig").packArrayHeader; const unpackArrayHeader = @import("array.zig").unpackArrayHeader; @@ -164,8 +165,6 @@ pub fn unpackStructFromMapBody(reader: *std.Io.Reader, allocator: std.mem.Alloca var fields_seen = std.bit_set.StaticBitSet(fields.len).initEmpty(); - var field_name_buffer: [256]u8 = undefined; - var result: Type = undefined; for (0..field_count) |_| { @@ -184,9 +183,9 @@ pub fn unpackStructFromMapBody(reader: *std.Io.Reader, allocator: std.mem.Alloca } }, .field_name => { - const field_name = try unpackStringInto(reader, &field_name_buffer); + const field_name = try unpackStringBorrowed(reader); inline for (fields, 0..) |field, i| { - if (std.mem.eql(u8, field.name, field_name)) { + if (eqlLiteral(field.name, field_name)) { fields_seen.set(i); @field(result, field.name) = try unpackAny(reader, allocator, field.type); break; @@ -196,7 +195,7 @@ pub fn unpackStructFromMapBody(reader: *std.Io.Reader, allocator: std.mem.Alloca } }, .field_name_prefix => |prefix| { - const field_name = try unpackStringInto(reader, &field_name_buffer); + const field_name = try unpackStringBorrowed(reader); inline for (fields, 0..) |field, i| { if (std.mem.startsWith(u8, field.name, strPrefix(field_name, prefix))) { fields_seen.set(i); diff --git a/src/union.zig b/src/union.zig index 1f8cc77..9f5c958 100644 --- a/src/union.zig +++ b/src/union.zig @@ -16,7 +16,8 @@ const packInt = @import("int.zig").packInt; const unpackInt = @import("int.zig").unpackInt; const packStringLiteral = @import("string.zig").packStringLiteral; -const unpackStringInto = @import("string.zig").unpackStringInto; +const unpackStringBorrowed = @import("string.zig").unpackStringBorrowed; +const eqlLiteral = @import("utils.zig").eqlLiteral; const packArrayHeader = @import("array.zig").packArrayHeader; const unpackArrayHeader = @import("array.zig").unpackArrayHeader; @@ -150,7 +151,7 @@ pub fn packUnion(writer: *std.Io.Writer, comptime T: type, value_or_maybe_null: } } -pub fn unpackUnionAsMap(reader: *std.Io.Reader, allocator: std.mem.Allocator, comptime T: type, opts: UnionAsMapOptions) !T { +pub fn unpackUnionAsMap(reader: *std.Io.Reader, allocator: std.mem.Allocator, comptime T: type, comptime opts: UnionAsMapOptions) !T { const len = if (@typeInfo(T) == .optional) try unpackMapHeader(reader, ?u16) orelse return null else @@ -164,8 +165,6 @@ pub fn unpackUnionAsMap(reader: *std.Io.Reader, allocator: std.mem.Allocator, co const type_info = @typeInfo(Type); const fields = type_info.@"union".fields; - var field_name_buffer: [256]u8 = undefined; - var result: Type = undefined; switch (opts.key) { @@ -182,9 +181,9 @@ pub fn unpackUnionAsMap(reader: *std.Io.Reader, allocator: std.mem.Allocator, co } }, .field_name => { - const field_name = try unpackStringInto(reader, &field_name_buffer); + const field_name = try unpackStringBorrowed(reader); inline for (fields) |field| { - if (std.mem.eql(u8, field.name, field_name)) { + if (eqlLiteral(field.name, field_name)) { const value = try unpackAny(reader, allocator, field.type); result = @unionInit(Type, field.name, value); break; @@ -194,7 +193,7 @@ pub fn unpackUnionAsMap(reader: *std.Io.Reader, allocator: std.mem.Allocator, co } }, .field_name_prefix => |prefix| { - const field_name = try unpackStringInto(reader, &field_name_buffer); + const field_name = try unpackStringBorrowed(reader); inline for (fields) |field| { if (std.mem.startsWith(u8, field.name, strPrefix(field_name, prefix))) { const value = try unpackAny(reader, allocator, field.type); @@ -210,7 +209,7 @@ pub fn unpackUnionAsMap(reader: *std.Io.Reader, allocator: std.mem.Allocator, co return result; } -pub fn unpackUnionAsTagged(reader: *std.Io.Reader, allocator: std.mem.Allocator, comptime T: type, opts: UnionAsTaggedOptions) !T { +pub fn unpackUnionAsTagged(reader: *std.Io.Reader, allocator: std.mem.Allocator, comptime T: type, comptime opts: UnionAsTaggedOptions) !T { const len = if (@typeInfo(T) == .optional) try unpackMapHeader(reader, ?u16) orelse return null else @@ -224,11 +223,8 @@ pub fn unpackUnionAsTagged(reader: *std.Io.Reader, allocator: std.mem.Allocator, const type_info = @typeInfo(Type); const fields = type_info.@"union".fields; - var tag_field_buffer: [256]u8 = undefined; - var tag_value_buffer: [256]u8 = undefined; - - const tag_field_name = try unpackStringInto(reader, &tag_field_buffer); - if (!std.mem.eql(u8, tag_field_name, opts.tag_field)) { + const tag_field_name = try unpackStringBorrowed(reader); + if (!eqlLiteral(opts.tag_field, tag_field_name)) { return error.InvalidTagField; } @@ -241,11 +237,11 @@ pub fn unpackUnionAsTagged(reader: *std.Io.Reader, allocator: std.mem.Allocator, union_field_index = field_index; }, .field_name => { - const field_name = try unpackStringInto(reader, &tag_value_buffer); + const field_name = try unpackStringBorrowed(reader); union_field_name = field_name; }, .field_name_prefix => { - const field_name = try unpackStringInto(reader, &tag_value_buffer); + const field_name = try unpackStringBorrowed(reader); union_field_name = field_name; }, } @@ -253,7 +249,7 @@ pub fn unpackUnionAsTagged(reader: *std.Io.Reader, allocator: std.mem.Allocator, inline for (fields, 0..) |field, i| { const is_match = switch (opts.tag_value) { .field_index => union_field_index == i, - .field_name => if (union_field_name) |name| std.mem.eql(u8, field.name, name) else false, + .field_name => if (union_field_name) |name| eqlLiteral(field.name, name) else false, .field_name_prefix => |prefix| if (union_field_name) |name| std.mem.startsWith(u8, field.name, strPrefix(name, prefix)) else false, }; @@ -281,7 +277,7 @@ pub fn unpackUnionAsTagged(reader: *std.Io.Reader, allocator: std.mem.Allocator, pub fn unpackUnion(reader: *std.Io.Reader, allocator: std.mem.Allocator, comptime T: type) !T { const Type = NonOptional(T); - const format = if (std.meta.hasFn(Type, "msgpackFormat")) Type.msgpackFormat() else default_union_format; + const format = comptime if (std.meta.hasFn(Type, "msgpackFormat")) Type.msgpackFormat() else default_union_format; switch (format) { .as_map => |opts| { return try unpackUnionAsMap(reader, allocator, T, opts); diff --git a/src/utils.zig b/src/utils.zig index bcd2cdb..552d78b 100644 --- a/src/utils.zig +++ b/src/utils.zig @@ -35,6 +35,25 @@ pub inline fn reserveArray(writer: *std.Io.Writer, comptime len: usize) ?*[len]u return null; } +/// Reads a big-endian `T` from the reader. `Reader.takeInt` rebases, which +/// asserts the reader's buffer can hold `@sizeOf(T)`; report that as an error +/// rather than letting the assert fire. +pub inline fn takeInt(reader: *std.Io.Reader, comptime T: type) !T { + if (reader.buffer.len < @divExact(@bitSizeOf(T), 8)) { + @branchHint(.unlikely); + return error.ReaderBufferTooSmall; + } + return reader.takeInt(T, .big); +} + +/// Compares `value` against a comptime-known `name`. The length test is a +/// compare against a constant, and the byte compare that follows has a +/// comptime-known length, so it lowers to inline compares rather than a call. +pub inline fn eqlLiteral(comptime name: []const u8, value: []const u8) bool { + if (value.len != name.len) return false; + return std.mem.eql(u8, value[0..name.len], name); +} + /// Writes a header byte followed by a big-endian `T`, as a single contiguous /// store when the writer's buffer has room, and as a single `writeAll` /// otherwise. Never splits the header from its payload across a drain.