Move types to their own file

Make it so I can specify the payload message

1 files changed, 191 insertions, 0 deletions

diff --git a/src/saprus_message.zig b/src/saprus_message.zig
new file mode 100644
index 0000000..443bbfb
--- /dev/null
+++ b/src/saprus_message.zig
@@ -0,0 +1,191 @@
+const base64Enc = std.base64.Base64Encoder.init(std.base64.standard_alphabet_chars, '=');
+const base64Dec = std.base64.Base64Decoder.init(std.base64.standard_alphabet_chars, '=');
+
+/// Type tag for SaprusMessage union.
+/// This is the first value in the actual packet sent over the network.
+pub const SaprusPacketType = enum(u16) {
+    relay = 0x003C,
+    file_transfer = 0x8888,
+    connection = 0x00E9,
+    _,
+};
+
+/// Reserved option values.
+/// Currently unused.
+pub const SaprusConnectionOptions = packed struct(u8) {
+    opt1: bool = false,
+    opt2: bool = false,
+    opt3: bool = false,
+    opt4: bool = false,
+    opt5: bool = false,
+    opt6: bool = false,
+    opt7: bool = false,
+    opt8: bool = false,
+};
+
+pub const SaprusError = error{
+    NotImplementedSaprusType,
+    UnknownSaprusType,
+};
+
+/// All Saprus messages
+pub const SaprusMessage = union(SaprusPacketType) {
+    pub const Relay = struct {
+        pub const Header = packed struct {
+            dest: @Vector(4, u8),
+        };
+        header: Header,
+        payload: []const u8,
+    };
+    pub const Connection = struct {
+        pub 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,
+    file_transfer: void, // unimplemented
+    connection: Connection,
+
+    /// Should be called for any SaprusMessage that was declared using a function that you pass an allocator to.
+    pub fn deinit(self: SaprusMessage, allocator: Allocator) void {
+        switch (self) {
+            .relay => |r| allocator.free(r.payload),
+            .connection => |c| allocator.free(c.payload),
+            else => unreachable,
+        }
+    }
+
+    inline fn toBytesAux(
+        Header: type,
+        header: Header,
+        payload: []const u8,
+        w: std.ArrayList(u8).Writer,
+        allocator: Allocator,
+    ) !void {
+        // Create a growable string to store the base64 bytes in.
+        // Doing this first so I can use the length of the encoded bytes for the length field.
+        var payload_list = std.ArrayList(u8).init(allocator);
+        defer payload_list.deinit();
+        const buf_w = payload_list.writer();
+
+        // Write the payload bytes as base64 to the growable string.
+        try base64Enc.encodeWriter(buf_w, payload);
+
+        // Write the packet body to the output writer.
+        try w.writeStructEndian(header, .big);
+        try w.writeInt(u16, @intCast(payload_list.items.len), .big);
+        try w.writeAll(payload_list.items);
+    }
+
+    /// Caller is responsible for freeing the returned bytes.
+    pub fn toBytes(self: SaprusMessage, allocator: Allocator) ![]u8 {
+        // Create a growable list of bytes to store the output in.
+        var buf = std.ArrayList(u8).init(allocator);
+        // Create a writer for an easy interface to append arbitrary bytes.
+        const w = buf.writer();
+
+        // Start with writing the message type, which is the first 16 bits of every Saprus message.
+        try w.writeInt(u16, @intFromEnum(self), .big);
+
+        // Write the proper header and payload for the given packet type.
+        switch (self) {
+            .relay => |r| try toBytesAux(Relay.Header, r.header, r.payload, w, allocator),
+            .connection => |c| try toBytesAux(Connection.Header, c.header, c.payload, w, allocator),
+            .file_transfer => return SaprusError.NotImplementedSaprusType,
+        }
+
+        // Collect the growable list as a slice and return it.
+        return buf.toOwnedSlice();
+    }
+
+    inline fn fromBytesAux(
+        packet: SaprusPacketType,
+        Header: type,
+        r: std.io.FixedBufferStream([]const u8).Reader,
+        allocator: Allocator,
+    ) !SaprusMessage {
+        // Read the header for the current message type.
+        const header = try r.readStructEndian(Header, .big);
+        // Read the length of the base64 encoded payload.
+        const len = try r.readInt(u16, .big);
+
+        // Read the base64 bytes into a list to be able to call the decoder on it.
+        var payload_buf = std.ArrayList(u8).init(allocator);
+        defer payload_buf.deinit();
+        try r.readAllArrayList(&payload_buf, len);
+
+        // Create a buffer to store the payload in, and decode the base64 bytes into the payload field.
+        const payload = try allocator.alloc(u8, try base64Dec.calcSizeForSlice(payload_buf.items));
+        try base64Dec.decode(payload, payload_buf.items);
+
+        // Return the type of SaprusMessage specified by the `packet` argument.
+        return @unionInit(SaprusMessage, @tagName(packet), .{
+            .header = header,
+            .payload = payload,
+        });
+    }
+
+    /// Caller is responsible for calling .deinit on the returned value.
+    pub fn fromBytes(bytes: []const u8, allocator: Allocator) !SaprusMessage {
+        var s = std.io.fixedBufferStream(bytes);
+        const r = s.reader();
+
+        const packet_type = @as(SaprusPacketType, @enumFromInt(try r.readInt(u16, .big)));
+
+        switch (packet_type) {
+            .relay => return fromBytesAux(.relay, Relay.Header, r, allocator),
+            .connection => return fromBytesAux(.connection, Connection.Header, r, allocator),
+            .file_transfer => return SaprusError.NotImplementedSaprusType,
+            else => return SaprusError.UnknownSaprusType,
+        }
+    }
+};
+
+const std = @import("std");
+const Allocator = std.mem.Allocator;
+
+test "Round trip Relay 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 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);
+}