zaprus/src/NetWriter.zig

//! Wraps a writer with UDP headers.
//! This is useful for wrapping RawSocket Writer with appropriate headers.

rand: Random,
wrapped: *Writer,
interface: Writer,

pub fn init(w: *Writer, buffer: []u8) !NetWriter {
    std.debug.assert(buffer.len > @sizeOf(EthernetHeaders) + @sizeOf(IpHeaders) + @sizeOf(UdpHeaders));

    var prng = Random.DefaultPrng.init(blk: {
        var seed: u64 = undefined;
        try posix.getrandom(mem.asBytes(&seed));
        break :blk seed;
    });

    return .{
        .rand = prng.random(),
        .wrapped = w,
        .interface = .{
            .vtable = &.{
                .drain = drain,
                // .flush = flush,
            },
            .buffer = buffer,
        },
    };
}

fn drain(io_w: *Writer, data: []const []const u8, splat: usize) Writer.Error!usize {
    const w: *NetWriter = @alignCast(@fieldParentPtr("interface", io_w));
    const headers_byte_len = comptime (EthernetHeaders.byte_len + (@bitSizeOf(IpHeaders) / 8) + @sizeOf(UdpHeaders));
    const headers: [headers_byte_len]u8 = blk: {
        const ether_headers: EthernetHeaders = .{
            .dest_mac = .{ 0xff, 0xff, 0xff, 0xff, 0xfe, 0xff },
            .src_mac = src_blk: {
                var output_bytes: [6]u8 = undefined;
                output_bytes[0] = 0xee;
                w.rand.bytes(output_bytes[1..]);
                break :src_blk output_bytes;
            },
            .ether_type = 0x0800,
        };

        const total_len = Writer.countSplat(data, splat) + w.interface.end;

        const ip_headers: IpHeaders = .{
            // length of the packet minus eth header
            .total_length = @intCast(headers_byte_len + total_len - EthernetHeaders.byte_len), //@intCast(total_len),
            .ttl = 64,
            .protocol = 0x11,
            .src_ip = .{ 0xff, 0x02, 0x03, 0x04 },
            .dest_ip = .{ 0xff, 0xff, 0xff, 0xff },
        };

        const udp_headers: UdpHeaders = .{
            .src_port = 0xbbbb,
            .dest_port = 8888,
            .length = @intCast(total_len + @sizeOf(UdpHeaders)),
        };

        var buf: [headers_byte_len]u8 = undefined;
        var buf_w = Writer.fixed(&buf);

        _ = try ether_headers.write(&buf_w);
        try ip_headers.write(&buf_w);
        try buf_w.writeStruct(udp_headers, .big);

        break :blk buf;
    };

    _ = try w.wrapped.write(&headers);
    const total_len = try w.wrapped.writeSplatHeader(w.interface.buffered(), data, splat);

    try w.wrapped.flush();
    return total_len - w.interface.consumeAll();
}

const EthernetHeaders = struct {
    dest_mac: @Vector(6, u8),

    src_mac: @Vector(6, u8),

    ether_type: u16,

    fn write(hdr: EthernetHeaders, writer: *std.Io.Writer) Writer.Error!usize {
        try writer.writeInt(u48, @bitCast(hdr.dest_mac), .big);
        try writer.writeInt(u48, @bitCast(hdr.src_mac), .big);
        try writer.writeInt(u16, hdr.ether_type, .big);
        return byte_len;
    }

    const byte_len = blk: {
        var res: usize = 0;
        res += @bitSizeOf(u48) / 8;
        res += @bitSizeOf(u48) / 8;
        res += @bitSizeOf(u16) / 8;
        break :blk res;
    };

    fn bytes(hdr: EthernetHeaders) [byte_len]u8 {
        var res: [byte_len]u8 = undefined;
        hdr.write(Writer.fixed(&res)) catch unreachable;
    }
};

const IpHeaders = extern struct {
    dest_ip: @Vector(4, u8) align(1),
    src_ip: @Vector(4, u8) align(1),
    header_checksum: @Vector(2, u8) align(1) = .{ 0, 0 },
    protocol: u8 align(1) = 0,
    ttl: u8 align(1) = 0,
    fragment_and_flags: packed struct(u16) {
        fragment_offset: u13 = 0,
        ethernet_flags: u3 = 0,
    } = .{},
    // fragment_offset: u13 = 0,
    // ethernet_flags: u3 = 0,
    identification: u16 align(1) = 0,
    total_length: u16 align(1) = 0x04,
    type_of_service: u8 align(1) = 0,
    header_and_version: packed struct(u8) {
        header_length: u4 = 5,
        ip_version: u4 = 4,
    } = .{},
    // header_length: u4 = 5,
    // ip_version: u4 = 4,

    fn write(hdr: IpHeaders, writer: *std.Io.Writer) Writer.Error!void {
        try writer.writeInt(u8, 0x45, .big); // ip version and header length
        try writer.writeByte(hdr.type_of_service);
        try writer.writeInt(u16, hdr.total_length, .big);
        try writer.writeInt(u16, hdr.identification, .big);
        try writer.writeInt(u16, 0x00, .big); // ethernet flags and fragment offset
        try writer.writeByte(hdr.ttl);
        try writer.writeByte(hdr.protocol);
        try writer.writeInt(u16, @bitCast(hdr.header_checksum), .big);
        try writer.writeInt(u32, @bitCast(hdr.src_ip), .big);
        try writer.writeInt(u32, @bitCast(hdr.dest_ip), .big);
    }

    fn bytes(hdr: IpHeaders) [@bitSizeOf(IpHeaders) / 8]u8 {
        var res: [@bitSizeOf(IpHeaders) / 8]u8 = undefined;
        var w = Writer.fixed(&res);
        _ = hdr.write(&w) catch unreachable;
        return res;
    }
};

test IpHeaders {
    var a: IpHeaders = .{
        .dest_ip = .{ 1, 2, 3, 4 },
        .src_ip = .{ 5, 6, 7, 8 },
        // .header_checksum = .{ 0x1, 0x1 },
        // .protocol = 0,
        // .ttl = 0x64,
        // .identification = 3,
        // .type_of_service = 7,
    };
    // a.dest_ip = .{ 1, 2, 3, 4 };
    // a.src_ip = .{ 5, 6, 7, 8 };
    std.debug.print("ip struct: {}\n", .{a});
    std.debug.print("dest ip: {}\n", .{a.dest_ip});
    std.debug.print("src ip: {}\n", .{a.src_ip});

    std.debug.print("raw b   :  {x}\n", .{@as([@bitSizeOf(IpHeaders) / 8]u8, @bitCast(a))});

    std.debug.print("ip bytes:  {x}\n", .{a.bytes()});
    var a_struct: [@bitSizeOf(IpHeaders) / 8]u8 = undefined;
    var a_struct_w = Writer.fixed(&a_struct);
    _ = a_struct_w.writeStruct(a, .big) catch unreachable;
    std.debug.print("ip struct: {x}\n", .{a_struct});
    try std.testing.expectEqual(a, a);
}

const UdpHeaders = packed struct {
    checksum: @Vector(2, u8) = .{ 0, 0 },
    length: u16,
    dest_port: u16,
    src_port: u16,

    fn write(hdr: UdpHeaders, writer: *std.Io.Writer) Writer.Error!void {
        try writer.writeStruct(hdr, .big);
    }

    fn bytes(hdr: UdpHeaders) [@sizeOf(UdpHeaders)]u8 {
        var res: [@sizeOf(UdpHeaders)]u8 = undefined;
        var w = Writer.fixed(&res);
        _ = hdr.write(&w) catch unreachable;
        return res;
    }
};

test UdpHeaders {
    const a: UdpHeaders = .{
        .src_port = 1,
        .dest_port = 2,
        .length = 3,
        .checksum = .{ 0x04, 0x05 },
    };
    var a_struct: [@sizeOf(UdpHeaders)]u8 = undefined;
    var a_struct_w = Writer.fixed(&a_struct);
    _ = a_struct_w.writeStruct(a, .big) catch unreachable;
    try std.testing.expectEqual(a.bytes(), a_struct);
}

const std = @import("std");
const Random = std.Random;
const posix = std.posix;
const Writer = std.Io.Writer;
const mem = std.mem;

const NetWriter = @This();