IT'S RIGHT

1 files changed, 67 insertions, 68 deletions

diff --git a/src/NetWriter.zig b/src/NetWriter.zig
index b08ccfe..42d3c91 100644
--- a/src/NetWriter.zig
+++ b/src/NetWriter.zig
@@ -5,8 +5,8 @@ rand: Random,
 wrapped: *Writer,
 interface: Writer,
 
-pub fn init(w: *Writer) !NetWriter {
-    std.debug.assert(w.buffer.len > @sizeOf(EthernetHeaders) + @sizeOf(IpHeaders) + @sizeOf(UdpHeaders));
+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;
@@ -20,33 +20,34 @@ pub fn init(w: *Writer) !NetWriter {
         .interface = .{
             .vtable = &.{
                 .drain = drain,
-                .flush = flush,
+                // .flush = flush,
             },
-            .buffer = &.{},
+            .buffer = 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 headers_byte_len = comptime (EthernetHeaders.byte_len + IpHeaders.byte_len + UdpHeaders.byte_len);
+    const headers: [headers_byte_len]u8 = blk: {
         const ether_headers: EthernetHeaders = .{
             .dest_mac = .{ 0xff, 0xff, 0xff, 0xff, 0xfe, 0xff },
-            .src_mac = blk: {
+            .src_mac = src_blk: {
                 var output_bytes: [6]u8 = undefined;
                 output_bytes[0] = 0xee;
                 w.rand.bytes(output_bytes[1..]);
-                break :blk output_bytes;
+                break :src_blk output_bytes;
             },
             .ether_type = 0x0800,
         };
 
+        const total_len = Writer.countSplat(data, splat) + w.interface.end;
+
         const ip_headers: IpHeaders = .{
-            .total_length = @intCast(res - 92),
-            .ttl = 0x64,
+            // 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 },
@@ -55,23 +56,35 @@ fn drain(io_w: *Writer, data: []const []const u8, splat: usize) Writer.Error!usi
         const udp_headers: UdpHeaders = .{
             .src_port = 0xbbbb,
             .dest_port = 8888,
-            .length = @intCast(res),
+            .length = @intCast(total_len + UdpHeaders.byte_len),
         };
 
-        res += try ether_headers.write(w.wrapped);
-        res += try ip_headers.write(w.wrapped);
-        res += try udp_headers.write(w.wrapped);
-    }
+        var buf: [headers_byte_len]u8 = undefined;
+        var buf_w = Writer.fixed(&buf);
 
-    res += try w.wrapped.writeSplat(data, splat);
-    return res;
-}
+        _ = try ether_headers.write(&buf_w);
+        _ = try ip_headers.write(&buf_w);
+        std.debug.print("after ip: {x}\n", .{buf_w.buffered()});
+        _ = try udp_headers.write(&buf_w);
+        std.debug.print("after udp: {x}\n", .{buf_w.buffered()});
+
+        break :blk buf;
+    };
+
+    _ = try w.wrapped.write(&headers);
+    const total_len = try w.wrapped.writeSplatHeader(w.interface.buffered(), data, splat);
+
+    std.debug.print("total splat: {}\theaders.len: {}\tsplat: {}\n", .{ total_len, headers.len, splat });
 
-fn flush(io_w: *Writer) Writer.Error!void {
-    const w: *NetWriter = @alignCast(@fieldParentPtr("interface", io_w));
     try w.wrapped.flush();
+    return total_len - w.interface.consumeAll();
 }
 
+// fn flush(io_w: *Writer) Writer.Error!void {
+//     const w: *NetWriter = @alignCast(@fieldParentPtr("interface", io_w));
+//     try io_w.defaultFlush();
+// }
+
 const EthernetHeaders = struct {
     dest_mac: @Vector(6, u8),
 
@@ -80,21 +93,19 @@ const EthernetHeaders = struct {
     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;
+        return byte_len;
     }
 
-    const byte_len = @bitSizeOf(EthernetHeaders) / 8;
+    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;
@@ -120,42 +131,33 @@ const IpHeaders = struct {
     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;
+        return byte_len;
     }
 
-    const byte_len = @bitSizeOf(IpHeaders) / 8;
+    const byte_len = blk: {
+        var res: usize = 0;
+        res += @sizeOf(u8);
+        res += @sizeOf(u8);
+        res += @sizeOf(u16);
+        res += @sizeOf(u16);
+        res += @sizeOf(u16);
+        res += @sizeOf(u8);
+        res += @sizeOf(u8);
+        res += @sizeOf(u16);
+        res += @sizeOf(u32);
+        res += @sizeOf(u32);
+        break :blk res;
+    };
 
     fn bytes(hdr: IpHeaders) [byte_len]u8 {
         var res: [byte_len]u8 = undefined;
@@ -171,24 +173,21 @@ const UdpHeaders = packed struct {
     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;
+        return byte_len;
     }
 
-    const byte_len = @bitSizeOf(UdpHeaders) / 8;
+    const byte_len = blk: {
+        var res: usize = 0;
+        res += @sizeOf(u16);
+        res += @sizeOf(u16);
+        res += @sizeOf(u16);
+        res += @sizeOf(u16);
+        break :blk res;
+    };
 
     fn bytes(hdr: UdpHeaders) [byte_len]u8 {
         var res: [byte_len]u8 = undefined;