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,
    InvalidMessage,
};

// ZERO COPY STUFF
// &payload could be a void value that is treated as a pointer to a [*]u8
pub const ZeroCopyMessage = packed struct {
    const Relay = packed struct {
        dest: @Vector(4, u8),
        payload: void,

        pub fn getPayload(self: *align(1) Relay) []u8 {
            const len: *u16 = @ptrFromInt(@intFromPtr(self) - @sizeOf(u16));
            return @as([*]u8, @ptrCast(&self.payload))[0 .. len.* - @sizeOf(Relay)];
        }
    };
    const Connection = 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 = .{},
        payload: void,

        pub fn getPayload(self: *align(1) Connection) []u8 {
            const len: *u16 = @ptrFromInt(@intFromPtr(self) - @sizeOf(u16));
            return @as([*]u8, @ptrCast(&self.payload))[0 .. len.* - @sizeOf(Connection)];
        }

        fn nativeFromNetworkEndian(self: *align(1) Connection) Error!void {
            self.src_port = bigToNative(@TypeOf(self.src_port), self.src_port);
            self.dest_port = bigToNative(@TypeOf(self.dest_port), self.dest_port);
            self.seq_num = bigToNative(@TypeOf(self.seq_num), self.seq_num);
            self.msg_id = bigToNative(@TypeOf(self.msg_id), self.msg_id);
        }

        fn networkFromNativeEndian(self: *align(1) Connection) Error!void {
            self.src_port = nativeToBig(@TypeOf(self.src_port), self.src_port);
            self.dest_port = nativeToBig(@TypeOf(self.dest_port), self.dest_port);
            self.seq_num = nativeToBig(@TypeOf(self.seq_num), self.seq_num);
            self.msg_id = nativeToBig(@TypeOf(self.msg_id), self.msg_id);
        }
    };

    const Self = @This();

    type: PacketType,
    length: u16,
    bytes: void = {},

    pub fn init(allocator: Allocator, comptime @"type": PacketType, payload_len: u16) !*Self {
        const header_size = @sizeOf(switch (@"type") {
            .relay => Relay,
            .connection => Connection,
            else => return error.Bad,
        });
        const size = payload_len + @sizeOf(Self) + header_size;
        const bytes = try allocator.alignedAlloc(u8, @alignOf(Self), size);
        const res: *Self = @ptrCast(bytes.ptr);
        res.type = @"type";
        res.length = payload_len + header_size;
        return res;
    }

    pub fn deinit(self: *Self, allocator: Allocator) void {
        allocator.free(self.toBytes());
    }

    fn getRelay(self: *Self) *align(1) Relay {
        return std.mem.bytesAsValue(Relay, &self.bytes);
    }
    fn getConnection(self: *Self) *align(1) Connection {
        return std.mem.bytesAsValue(Connection, &self.bytes);
    }

    pub fn getSaprusTypePayload(self: *Self) Error!(union(PacketType) {
        relay: *align(1) Relay,
        file_transfer: void,
        connection: *align(1) Connection,
    }) {
        return switch (self.type) {
            .relay => .{ .relay = self.getRelay() },
            .connection => .{ .connection = self.getConnection() },
            .file_transfer => Error.NotImplementedSaprusType,
            else => Error.UnknownSaprusType,
        };
    }

    pub fn nativeFromNetworkEndian(self: *Self) Error!void {
        self.type = @enumFromInt(bigToNative(
            @typeInfo(@TypeOf(self.type)).@"enum".tag_type,
            @intFromEnum(self.type),
        ));
        self.length = bigToNative(@TypeOf(self.length), self.length);
        switch (try self.getSaprusTypePayload()) {
            .relay => {},
            .connection => |*con| try con.*.nativeFromNetworkEndian(),
            // We know other values are unreachable,
            // because they would have returned an error from the switch condition.
            else => unreachable,
        }
    }

    pub fn networkFromNativeEndian(self: *Self) Error!void {
        try switch (try self.getSaprusTypePayload()) {
            .relay => {},
            .connection => |*con| con.*.networkFromNativeEndian(),
            .file_transfer => Error.NotImplementedSaprusType,
            else => Error.UnknownSaprusType,
        };
        self.type = @enumFromInt(nativeToBig(
            @typeInfo(@TypeOf(self.type)).@"enum".tag_type,
            @intFromEnum(self.type),
        ));
        self.length = nativeToBig(@TypeOf(self.length), self.length);
    }

    pub fn fromBytesUnchecked(bytes: []align(@alignOf(Self)) u8) *Self {
        return @ptrCast(@alignCast(bytes.ptr));
    }

    pub fn fromBytes(bytes: []align(@alignOf(Self)) u8) !*Self {
        const res: *Self = @ptrCast(@alignCast(bytes.ptr));
        return if (bytes.len == res.length + @sizeOf(Self)) res else Error.InvalidMessage;
    }

    pub fn toBytes(self: *Self) []align(@alignOf(Self)) u8 {
        return @as([*]align(@alignOf(Self)) u8, @ptrCast(self))[0 .. @sizeOf(Self) + self.length];
    }
};

test "testing variable length zero copy struct" {
    const gpa = std.testing.allocator;
    const payload = "Hello darkness my old friend";

    // Create a view of the byte slice as a ZeroCopyMessage
    const zcm: *ZeroCopyMessage = try .init(gpa, .relay, payload.len);
    defer zcm.deinit(gpa);

    {
        // Set the message values
        {
            // These are both set by the init call.
            // zcm.type = .relay;
            // zcm.length = payload_len;
        }
        const relay = (try zcm.getSaprusTypePayload()).relay;
        relay.dest = .{ 1, 2, 3, 4 };
        @memcpy(relay.getPayload(), payload);
    }

    {
        const bytes = zcm.toBytes();

        // Print the message as hex using the network byte order
        try zcm.networkFromNativeEndian();
        // We know the error from nativeFromNetworkEndian is unreachable because
        // it would have returned an error from networkFromNativeEndian.
        defer zcm.nativeFromNetworkEndian() catch unreachable;
        std.debug.print("network bytes: {s}\n", .{bytes});
        std.debug.print("bytes len: {d}\n", .{bytes.len});
    }

    if (false) {
        // Illegal behavior
        std.debug.print("{any}\n", .{(try zcm.getSaprusTypePayload()).connection});
    }
}

/// 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;
const bigToNative = std.mem.bigToNative;

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);
}

test {
    std.testing.refAllDeclsRecursive(@This());
}