Add packet decoding

Big bug where I was setting the payload length do be the pre-encoded size rather than the encoded size.

1 files changed, 65 insertions, 9 deletions

diff --git a/src/main.zig b/src/main.zig
index 1de51f4..b7f8d54 100644
--- a/src/main.zig
+++ b/src/main.zig
@@ -1,10 +1,12 @@
 const is_debug = builtin.mode == .Debug;
-const base64 = std.base64.Base64Encoder.init(std.base64.standard_alphabet_chars, '=');
+const base64Enc = std.base64.Base64Encoder.init(std.base64.standard_alphabet_chars, '=');
+const base64Dec = std.base64.Base64Decoder.init(std.base64.standard_alphabet_chars, '=');
 
 const SaprusPacketType = enum(u16) {
     relay = 0x003C,
     file_transfer = 0x8888,
     connection = 0x00E9,
+    _,
 };
 
 const SaprusConnectionOptions = packed struct(u8) {
@@ -20,20 +22,22 @@ const SaprusConnectionOptions = packed struct(u8) {
 
 const SaprusMessage = union(SaprusPacketType) {
     const Relay = struct {
-        header: packed struct {
+        const Header = packed struct {
             dest: @Vector(4, u8),
-        },
+        };
+        header: Header,
         payload: []const u8,
     };
     const Connection = struct {
-        header: packed struct {
+        const Header = packed struct {
             src_port: u16,
             dest_port: u16,
             seq_num: u32 = 0,
             msg_id: u32 = 0,
             reserved: u8 = 0,
             options: SaprusConnectionOptions = .{},
-        },
+        };
+        header: Header,
         payload: []const u8,
     };
     relay: Relay,
@@ -47,20 +51,54 @@ const SaprusMessage = union(SaprusPacketType) {
 
         switch (self) {
             .relay => |r| {
+                var b64_payload_buf = std.ArrayList(u8).init(allocator);
+                defer b64_payload_buf.deinit();
+                const buf_w = b64_payload_buf.writer();
+                try base64Enc.encodeWriter(buf_w, r.payload);
+
                 try w.writeStructEndian(r.header, .big);
-                try w.writeInt(u16, @intCast(r.payload.len), .big);
-                try base64.encodeWriter(w, r.payload);
+                try w.writeInt(u16, @intCast(b64_payload_buf.items.len), .big);
+                try w.writeAll(b64_payload_buf.items);
             },
             .file_transfer => unreachable,
             .connection => |c| {
+                var b64_payload_buf = std.ArrayList(u8).init(allocator);
+                defer b64_payload_buf.deinit();
+                const buf_w = b64_payload_buf.writer();
+                try base64Enc.encodeWriter(buf_w, c.payload);
+
                 try w.writeStructEndian(c.header, .big);
-                try w.writeInt(u16, @intCast(c.payload.len), .big);
-                try base64.encodeWriter(w, c.payload);
+                try w.writeInt(u16, @intCast(b64_payload_buf.items.len), .big);
+                try w.writeAll(b64_payload_buf.items);
             },
         }
 
         return buf.toOwnedSlice();
     }
+
+    fn fromBytes(bytes: []const u8, allocator: Allocator) !SaprusMessage {
+        var s = std.io.fixedBufferStream(bytes);
+        const r = s.reader();
+
+        switch (@as(SaprusPacketType, @enumFromInt(try r.readInt(u16, .big)))) {
+            .relay => {
+                const header = try r.readStructEndian(Relay.Header, .big);
+                const len = try r.readInt(u16, .big);
+                var payload_b64_buf = std.ArrayList(u8).init(allocator);
+                defer payload_b64_buf.deinit();
+                try r.readAllArrayList(&payload_b64_buf, len);
+                const payload = try allocator.alloc(u8, try base64Dec.calcSizeForSlice(payload_b64_buf.items));
+                try base64Dec.decode(payload, payload_b64_buf.items);
+                return .{
+                    .relay = .{
+                        .header = header,
+                        .payload = payload,
+                    },
+                };
+            },
+            else => unreachable,
+        }
+    }
 };
 
 pub fn main() !void {
@@ -111,3 +149,21 @@ const Allocator = std.mem.Allocator;
 const DebugAllocator = std.heap.DebugAllocator(.{});
 
 const network = @import("network");
+
+test "Round trip toBytes and fromBytes" {
+    const gpa = std.testing.allocator;
+    const msg = SaprusMessage{
+        .relay = .{
+            .header = .{ .dest = .{ 255, 255, 255, 255 } },
+            .payload = "Hello darkness my old friend",
+        },
+    };
+
+    const to_bytes = try msg.toBytes(gpa);
+    defer gpa.free(to_bytes);
+
+    const from_bytes = try SaprusMessage.fromBytes(to_bytes, gpa);
+    defer gpa.free(from_bytes.relay.payload);
+
+    try std.testing.expectEqualDeep(msg, from_bytes);
+}