Start breaking out net logic to NetWriter

1 files changed, 205 insertions, 0 deletions

diff --git a/src/NetWriter.zig b/src/NetWriter.zig
new file mode 100644
index 0000000..b08ccfe
--- /dev/null
+++ b/src/NetWriter.zig
@@ -0,0 +1,205 @@
+//! 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) !NetWriter {
+    std.debug.assert(w.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 = &.{},
+        },
+    };
+}
+
+fn drain(io_w: *Writer, data: []const []const u8, splat: usize) Writer.Error!usize {
+    const w: *NetWriter = @alignCast(@fieldParentPtr("interface", io_w));
+
+    var res: usize = 0;
+
+    if (io_w.end == 0) {
+        const ether_headers: EthernetHeaders = .{
+            .dest_mac = .{ 0xff, 0xff, 0xff, 0xff, 0xfe, 0xff },
+            .src_mac = blk: {
+                var output_bytes: [6]u8 = undefined;
+                output_bytes[0] = 0xee;
+                w.rand.bytes(output_bytes[1..]);
+                break :blk output_bytes;
+            },
+            .ether_type = 0x0800,
+        };
+
+        const ip_headers: IpHeaders = .{
+            .total_length = @intCast(res - 92),
+            .ttl = 0x64,
+            .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(res),
+        };
+
+        res += try ether_headers.write(w.wrapped);
+        res += try ip_headers.write(w.wrapped);
+        res += try udp_headers.write(w.wrapped);
+    }
+
+    res += try w.wrapped.writeSplat(data, splat);
+    return res;
+}
+
+fn flush(io_w: *Writer) Writer.Error!void {
+    const w: *NetWriter = @alignCast(@fieldParentPtr("interface", io_w));
+    try w.wrapped.flush();
+}
+
+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 {
+        comptime var res: usize = 0;
+
+        res += @sizeOf(u48);
+        try writer.writeInt(u48, @bitCast(hdr.dest_mac), .big);
+
+        res += @sizeOf(u48);
+        try writer.writeInt(u48, @bitCast(hdr.src_mac), .big);
+
+        res += @sizeOf(u16);
+        try writer.writeInt(u16, hdr.ether_type, .big);
+
+        return res;
+    }
+
+    const byte_len = @bitSizeOf(EthernetHeaders) / 8;
+
+    fn bytes(hdr: EthernetHeaders) [byte_len]u8 {
+        var res: [byte_len]u8 = undefined;
+        hdr.write(Writer.fixed(&res)) catch unreachable;
+    }
+};
+
+const IpHeaders = struct {
+    // ip_version: u4,
+    // header_length: u4 = 0,
+    type_of_service: u8 = 0,
+    total_length: u16 = 0x04,
+
+    identification: u16 = 0,
+    // ethernet_flags: u3 = 0,
+    // fragment_offset: u13 = 0,
+    ttl: u8 = 0,
+    protocol: u8 = 0,
+
+    header_checksum: @Vector(2, u8) = .{ 0, 0 },
+    src_ip: @Vector(4, u8),
+
+    dest_ip: @Vector(4, u8),
+
+    fn write(hdr: @This(), writer: *std.Io.Writer) Writer.Error!usize {
+        comptime var res: usize = 0;
+
+        res += @sizeOf(u8);
+        try writer.writeInt(u8, 0x45, .big); // ip version and header length
+
+        res += @sizeOf(u8);
+        try writer.writeByte(hdr.type_of_service);
+
+        res += @sizeOf(u16);
+        try writer.writeInt(u16, hdr.total_length, .big);
+
+        res += @sizeOf(u16);
+        try writer.writeInt(u16, hdr.identification, .big);
+
+        res += @sizeOf(u16);
+        try writer.writeInt(u16, 0x00, .big); // ethernet flags and fragment offset
+
+        res += @sizeOf(u8);
+        try writer.writeByte(hdr.ttl);
+
+        res += @sizeOf(u8);
+        try writer.writeByte(hdr.protocol);
+
+        res += @sizeOf(u16);
+        try writer.writeInt(u16, @bitCast(hdr.header_checksum), .big);
+
+        res += @sizeOf(u32);
+        try writer.writeInt(u32, @bitCast(hdr.src_ip), .big);
+
+        res += @sizeOf(u32);
+        try writer.writeInt(u32, @bitCast(hdr.dest_ip), .big);
+
+        return res;
+    }
+
+    const byte_len = @bitSizeOf(IpHeaders) / 8;
+
+    fn bytes(hdr: IpHeaders) [byte_len]u8 {
+        var res: [byte_len]u8 = undefined;
+        hdr.write(Writer.fixed(&res)) catch unreachable;
+    }
+};
+
+const UdpHeaders = packed struct {
+    src_port: u16,
+
+    dest_port: u16,
+    length: u16,
+    checksum: @Vector(2, u8) = .{ 0, 0 },
+
+    fn write(hdr: @This(), writer: *std.Io.Writer) Writer.Error!usize {
+        comptime var res: usize = 0;
+
+        res += @sizeOf(u16);
+        try writer.writeInt(u16, hdr.src_port, .big);
+
+        res += @sizeOf(u16);
+        try writer.writeInt(u16, hdr.dest_port, .big);
+
+        res += @sizeOf(u16);
+        try writer.writeInt(u16, hdr.length, .big);
+
+        res += @sizeOf(u16);
+        try writer.writeInt(u16, @bitCast(hdr.checksum), .big);
+
+        return res;
+    }
+
+    const byte_len = @bitSizeOf(UdpHeaders) / 8;
+
+    fn bytes(hdr: UdpHeaders) [byte_len]u8 {
+        var res: [byte_len]u8 = undefined;
+        hdr.write(Writer.fixed(&res)) catch unreachable;
+    }
+};
+
+const std = @import("std");
+const Random = std.Random;
+const posix = std.posix;
+const Writer = std.Io.Writer;
+const mem = std.mem;
+
+const NetWriter = @This();