zaprus/src/message.zig

/// 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 MessageTypeError = error{
    NotImplementedSaprusType,
    UnknownSaprusType,
};
pub const MessageParseError = MessageTypeError || error{
    InvalidMessage,
};

// ZERO COPY STUFF
// &payload could be a void value that is treated as a pointer to a [*]u8
/// All Saprus messages
pub const Message = 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) 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) 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();
    const SelfBytes = []align(@alignOf(Self)) u8;

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

    /// Takes a byte slice, and returns a Message struct backed by the slice.
    /// This properly initializes the top level headers within the slice.
    pub fn init(@"type": PacketType, bytes: []align(@alignOf(Self)) u8) *Self {
        std.debug.assert(bytes.len >= @sizeOf(Self));
        const res: *Self = @ptrCast(bytes.ptr);
        res.type = @"type";
        res.length = @intCast(bytes.len - @sizeOf(Self));
        return res;
    }

    pub fn lengthForPayloadLength(comptime @"type": PacketType, payload_len: usize) MessageTypeError!u16 {
        std.debug.assert(payload_len < std.math.maxInt(u16));
        const header_size = @sizeOf(switch (@"type") {
            .relay => Relay,
            .connection => Connection,
            .file_transfer => return MessageTypeError.NotImplementedSaprusType,
            else => return MessageTypeError.UnknownSaprusType,
        });
        return @intCast(payload_len + @sizeOf(Self) + header_size);
    }

    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) MessageTypeError!(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 => MessageTypeError.NotImplementedSaprusType,
            else => MessageTypeError.UnknownSaprusType,
        };
    }

    pub fn nativeFromNetworkEndian(self: *Self) MessageTypeError!void {
        self.type = @enumFromInt(bigToNative(
            @typeInfo(@TypeOf(self.type)).@"enum".tag_type,
            @intFromEnum(self.type),
        ));
        self.length = bigToNative(@TypeOf(self.length), self.length);
        errdefer {
            // If the payload specific headers fail, revert the top level header values
            self.type = @enumFromInt(nativeToBig(
                @typeInfo(@TypeOf(self.type)).@"enum".tag_type,
                @intFromEnum(self.type),
            ));
            self.length = nativeToBig(@TypeOf(self.length), self.length);
        }
        switch (try self.getSaprusTypePayload()) {
            .relay => {},
            .connection => |*con| 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) MessageTypeError!void {
        try switch (try self.getSaprusTypePayload()) {
            .relay => {},
            .connection => |*con| con.*.networkFromNativeEndian(),
            .file_transfer => MessageTypeError.NotImplementedSaprusType,
            else => MessageTypeError.UnknownSaprusType,
        };
        self.type = @enumFromInt(nativeToBig(
            @typeInfo(@TypeOf(self.type)).@"enum".tag_type,
            @intFromEnum(self.type),
        ));
        self.length = nativeToBig(@TypeOf(self.length), self.length);
    }

    /// Deprecated.
    /// If I need the bytes, I should just pass around the slice that is backing this to begin with.
    pub fn asBytes(self: *Self) SelfBytes {
        const size = @sizeOf(Self) + self.length;
        return @as([*]align(@alignOf(Self)) u8, @ptrCast(self))[0..size];
    }
};

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 Message
    const msg: *Message = try .init(gpa, .relay, payload.len);
    defer msg.deinit(gpa);

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

    {
        const bytes = msg.asBytes();

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

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

    try std.testing.expectEqualDeep(msg, try Message.bytesAsValue(msg.asBytes()));
}

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