zaprus/src/message.zig

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 Message union.
/// This is the first value in the actual packet sent over the network.
pub const PacketType = enum(u16) {
    relay = 0x003C,
    file_transfer = 0x8888,
    connection = 0x00E9,
    _,
};

/// Reserved option values.
/// Currently unused.
pub const ConnectionOptions = 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 Error = error{
    NotImplementedSaprusType,
    UnknownSaprusType,
};

/// All Saprus messages
pub const Message = union(PacketType) {
    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, // random number > 1024
            dest_port: u16, // random number > 1024
            seq_num: u32 = 0,
            msg_id: u32 = 0,
            reserved: u8 = 0,
            options: ConnectionOptions = .{},
        };
        header: Header,
        payload: []const u8,
    };
    relay: Relay,
    file_transfer: void, // unimplemented
    connection: Connection,

    /// Should be called for any Message that was declared using a function that you pass an allocator to.
    pub fn deinit(self: Message, allocator: Allocator) void {
        switch (self) {
            .relay => |r| allocator.free(r.payload),
            .connection => |c| allocator.free(c.payload),
            else => unreachable,
        }
    }

    fn toBytesAux(
        header: anytype,
        payload: []const u8,
        buf: *std.ArrayList(u8),
        allocator: Allocator,
    ) !void {
        const Header = @TypeOf(header);
        // 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);

        // At this point, payload_list contains the base64 encoded payload.

        // Add the payload length to the output buf.
        try buf.*.appendSlice(
            asBytes(&nativeToBig(u16, @intCast(payload_list.items.len + @bitSizeOf(Header) / 8))),
        );

        // Add the header bytes to the output buf.
        var header_buf: [@sizeOf(Header)]u8 = undefined;
        var header_buf_stream = std.io.fixedBufferStream(&header_buf);
        try header_buf_stream.writer().writeStructEndian(header, .big);
        // Add the exact number of bits in the header without padding.
        try buf.*.appendSlice(header_buf[0 .. @bitSizeOf(Header) / 8]);

        try buf.*.appendSlice(payload_list.items);
    }

    /// Caller is responsible for freeing the returned bytes.
    pub fn toBytes(self: Message, allocator: Allocator) ![]u8 {
        // Create a growable list of bytes to store the output in.
        var buf = std.ArrayList(u8).init(allocator);
        errdefer buf.deinit();

        // Start with writing the message type, which is the first 16 bits of every Saprus message.
        try buf.appendSlice(asBytes(&nativeToBig(u16, @intFromEnum(self))));

        // Write the proper header and payload for the given packet type.
        switch (self) {
            .relay => |r| try toBytesAux(r.header, r.payload, &buf, allocator),
            .connection => |c| try toBytesAux(c.header, c.payload, &buf, allocator),
            .file_transfer => return Error.NotImplementedSaprusType,
        }

        // Collect the growable list as a slice and return it.
        return buf.toOwnedSlice();
    }

    fn fromBytesAux(
        comptime packet: PacketType,
        len: u16,
        r: std.io.FixedBufferStream([]const u8).Reader,
        allocator: Allocator,
    ) !Message {
        const Header = @field(@FieldType(Message, @tagName(packet)), "Header");

        // Read the header for the current message type.
        var header_bytes: [@sizeOf(Header)]u8 = undefined;
        _ = try r.read(header_bytes[0 .. @bitSizeOf(Header) / 8]);
        var header_stream = std.io.fixedBufferStream(&header_bytes);
        const header = try header_stream.reader().readStructEndian(Header, .big);

        // Read the base64 bytes into a list to be able to call the decoder on it.
        const payload_buf = try allocator.alloc(u8, len - @bitSizeOf(Header) / 8);
        defer allocator.free(payload_buf);
        _ = try r.readAll(payload_buf);

        // 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));
        try base64Dec.decode(payload, payload_buf);

        // Return the type of Message specified by the `packet` argument.
        return @unionInit(Message, @tagName(packet), .{
            .header = header,
            .payload = payload,
        });
    }

    /// Caller is responsible for calling .deinit on the returned value.
    pub fn fromBytes(bytes: []const u8, allocator: Allocator) !Message {
        var s = std.io.fixedBufferStream(bytes);
        const r = s.reader();

        // Read packet type
        const packet_type = @as(PacketType, @enumFromInt(try r.readInt(u16, .big)));

        // Read the length of the header + base64 encoded payload.
        const len = try r.readInt(u16, .big);

        switch (packet_type) {
            .relay => return fromBytesAux(.relay, len, r, allocator),
            .connection => return fromBytesAux(.connection, len, r, allocator),
            .file_transfer => return Error.NotImplementedSaprusType,
            else => return Error.UnknownSaprusType,
        }
    }
};

const std = @import("std");
const Allocator = std.mem.Allocator;

const asBytes = std.mem.asBytes;
const nativeToBig = std.mem.nativeToBig;

test "Round trip Relay toBytes and fromBytes" {
    const gpa = std.testing.allocator;
    const msg = Message{
        .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 Message.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 = Message{
        .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 Message.fromBytes(to_bytes, gpa);
    defer from_bytes.deinit(gpa);

    try std.testing.expectEqualDeep(msg, from_bytes);
}