Add decoding connection messages

1 files changed, 60 insertions, 4 deletions

diff --git a/src/main.zig b/src/main.zig
index b7f8d54..bce3ca6 100644
--- a/src/main.zig
+++ b/src/main.zig
@@ -20,6 +20,11 @@ const SaprusConnectionOptions = packed struct(u8) {
     opt8: bool = false,
 };
 
+const SaprusError = error{
+    NotImplementedSaprusType,
+    UnknownSaprusType,
+};
+
 const SaprusMessage = union(SaprusPacketType) {
     const Relay = struct {
         const Header = packed struct {
@@ -44,6 +49,14 @@ const SaprusMessage = union(SaprusPacketType) {
     file_transfer: void, // unimplemented
     connection: Connection,
 
+    fn deinit(self: SaprusMessage, allocator: Allocator) void {
+        switch (self) {
+            .relay => |r| allocator.free(r.payload),
+            .connection => |c| allocator.free(c.payload),
+            else => unreachable,
+        }
+    }
+
     fn toBytes(self: SaprusMessage, allocator: Allocator) ![]u8 {
         var buf = std.ArrayList(u8).init(allocator);
         const w = buf.writer();
@@ -60,7 +73,7 @@ const SaprusMessage = union(SaprusPacketType) {
                 try w.writeInt(u16, @intCast(b64_payload_buf.items.len), .big);
                 try w.writeAll(b64_payload_buf.items);
             },
-            .file_transfer => unreachable,
+            .file_transfer => return SaprusError.NotImplementedSaprusType,
             .connection => |c| {
                 var b64_payload_buf = std.ArrayList(u8).init(allocator);
                 defer b64_payload_buf.deinit();
@@ -84,11 +97,14 @@ const SaprusMessage = union(SaprusPacketType) {
             .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,
@@ -96,7 +112,26 @@ const SaprusMessage = union(SaprusPacketType) {
                     },
                 };
             },
-            else => unreachable,
+            .file_transfer => return SaprusError.NotImplementedSaprusType,
+            .connection => {
+                const header = try r.readStructEndian(Connection.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 .{
+                    .connection = .{
+                        .header = header,
+                        .payload = payload,
+                    },
+                };
+            },
+            else => return SaprusError.UnknownSaprusType,
         }
     }
 };
@@ -150,7 +185,7 @@ const DebugAllocator = std.heap.DebugAllocator(.{});
 
 const network = @import("network");
 
-test "Round trip toBytes and fromBytes" {
+test "Round trip Relay toBytes and fromBytes" {
     const gpa = std.testing.allocator;
     const msg = SaprusMessage{
         .relay = .{
@@ -163,7 +198,28 @@ test "Round trip toBytes and fromBytes" {
     defer gpa.free(to_bytes);
 
     const from_bytes = try SaprusMessage.fromBytes(to_bytes, gpa);
-    defer gpa.free(from_bytes.relay.payload);
+    defer from_bytes.deinit(gpa);
+
+    try std.testing.expectEqualDeep(msg, from_bytes);
+}
+
+test "Round trip Connection toBytes and fromBytes" {
+    const gpa = std.testing.allocator;
+    const msg = SaprusMessage{
+        .connection = .{
+            .header = .{
+                .src_port = 0,
+                .dest_port = 0,
+            },
+            .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 from_bytes.deinit(gpa);
 
     try std.testing.expectEqualDeep(msg, from_bytes);
 }