Move tracker protocol into it's own file

Change piece_length -> u32, protocol constrains that anyway
annnnd, successfully download a file (albeit from one torrent only with
no checking of a ton of stuff)

1 files changed, 87 insertions, 85 deletions

diff --git a/src/main.zig b/src/main.zig
index d50cc90..f27735d 100644
--- a/src/main.zig
+++ b/src/main.zig
@@ -3,18 +3,19 @@ const MetaInfo = @import("metainfo.zig");
 const bencode = @import("bencode.zig");
 const AnyWriter = @import("anywriter.zig");
 const peerproto = @import("peer_protocol.zig");
-
-const prr = "00112233445566778899";
+const trackproto = @import("tracker_protocol.zig");
 
 pub fn main() !void {
-    var peer_id: [20]u8 = undefined;
-    @memcpy(&peer_id, prr);
-
     var gpa = std.heap.GeneralPurposeAllocator(.{}){};
     defer _ = gpa.deinit();
     const a = gpa.allocator();
 
-    // open & parse the torrent file
+    // TODO figure this out. It's not that important, I think, unless
+    // other clients have special handling for different patterns. 
+    // Spec looks like a bit of a free-for-all here.
+    var peer_id: [20]u8 = undefined;
+    @memcpy(&peer_id, "00112233445566778899");
+
     const f = try std.fs.cwd().openFile("src/sample.torrent", .{});
     defer f.close();
     var fr = f.reader();
@@ -22,105 +23,106 @@ pub fn main() !void {
     defer mib.deinit(a);
     var mi = try MetaInfo.parse(a, mib);
     defer mi.deinit(a);
+    const info_hash = try mi.info.hash(a);
 
-    // call the tracker...
     var c = std.http.Client{
         .allocator = a,
     };
     defer c.deinit();
 
-    var q = std.StringHashMap([]const u8).init(a);
-    defer q.deinit();
-    var info_hash = try mi.info.hash(a);
-    var buf_left = [_]u8{0} ** 1024;
-    try q.put("info_hash", &info_hash);
-    try q.put("peer_id", &peer_id);
-    try q.put("port", "6881");
-    try q.put("uploaded", "0");
-    try q.put("downloaded", "0");
-    try q.put("left", try std.fmt.bufPrint(&buf_left, "{}", .{mi.info.files[0].length}));
-    try q.put("compact", "1");
-    var qs = try toqs(a, q);
-    defer a.free(qs);
-    var url = try std.fmt.allocPrint(a, "{s}?{s}", .{ mi.announce, qs });
+    const url = try trackproto.trackerRequestUrl(a, info_hash, peer_id, mi.info.files[0].length, mi.announce);
     defer a.free(url);
     var res = try c.fetch(a, .{ .location = .{ .url = url } });
     defer res.deinit();
-    if (res.status == .ok) {
-        var trb = try bencode.bdecodeBuf(a, res.body.?);
-        defer trb.deinit(a);
-        var tr = try TrackerResp.parse(a, trb);
-        defer tr.deinit(a);
-
-        for (tr.peers) |peer| {
-            std.log.info("peer: {}", .{peer});
-        }
-
-        // Handle peers... 
-
+    if (res.status != .ok) {
+        return error.TrackerHttpError;
+    }
 
-        const p = tr.peers[0];
-        var ps = try std.net.tcpConnectToAddress(p);
-        defer ps.close();
-        var pw = ps.writer();
-        var pr = ps.reader();
-        var hs: peerproto.Handshake = .{
-            .info_hash = info_hash,
-            .peer_id = peer_id,
-        };
+    var trb = try bencode.bdecodeBuf(a, res.body.?);
+    defer trb.deinit(a);
+    var tr = try trackproto.TrackerResp.parse(a, trb);
+    defer tr.deinit(a);
 
-        try hs.write(pw);
-        var phs = try peerproto.Handshake.read(pr);
-        std.log.info("peer at {} peer_id {s}", .{ p, std.fmt.fmtSliceHexLower(&phs.peer_id) });
+    if (tr.peers.len == 0) {
+        std.log.info("no peers", .{});
+        return;
     }
-}
 
-fn toqs(a: std.mem.Allocator, hm: std.StringHashMap([]const u8)) ![]const u8 {
-    var al = std.ArrayList(u8).init(a);
-    defer al.deinit();
-    var w = al.writer();
-    var it = hm.iterator();
-    var first = true;
-    while (it.next()) |entry| {
-        if (!first) try w.writeByte('&');
-        try std.Uri.writeEscapedQuery(w, entry.key_ptr.*);
-        try w.writeByte('=');
-        try std.Uri.writeEscapedQuery(w, entry.value_ptr.*);
-        first = false;
+    for (tr.peers) |peer| {
+        std.log.info("peer: {}", .{peer});
     }
-    return try al.toOwnedSlice();
-}
 
-const TrackerResp = struct {
-    // interval: u64,
-    peers: []std.net.Address,
-
-    pub fn parse(a: std.mem.Allocator, b: bencode.BValue) !TrackerResp {
-        var ipl = std.ArrayList(std.net.Address).init(a);
-        defer ipl.deinit();
-        var d = b.asDict() catch return error.Malformatted;
-        var pb = d.get("peers") orelse return error.Malformatted;
-        var ps = pb.asString() catch return error.Malformatted;
-        if ((ps.len % 6) != 0) return error.Malformatted;
-        for (0..ps.len / 6) |ix| {
-            const start = ix * 6;
-            const port = std.mem.readInt(u16, ps[start + 4 .. start + 6][0..2], .big);
-            var ip = [_]u8{0} ** 4;
-            @memcpy(&ip, ps[start .. start + 4]);
-            try ipl.append(std.net.Address.initIp4(ip, port));
-        }
-        return .{
-            .peers = try ipl.toOwnedSlice(),
-        };
-    }
+    // Handle peers, PoC we're just going to handle 1 peer and download everything from them very simplistically.
+    const p = tr.peers[0];
+    const file = mi.info.files[0];
+    var ps = try std.net.tcpConnectToAddress(p);
+    defer ps.close();
+    var pw = ps.writer();
+    var pr = ps.reader();
+    var hs: peerproto.Handshake = .{
+        .info_hash = info_hash,
+        .peer_id = peer_id,
+    };
+
+    try hs.write(pw);
+    var phs = try peerproto.Handshake.read(pr);
+    std.log.info("peer at {} peer_id {s}", .{ p, std.fmt.fmtSliceHexLower(&phs.peer_id) });
+
 
-    pub fn deinit(self: *TrackerResp, a: std.mem.Allocator) void {
-        a.free(self.peers);
+    var bf = try peerproto.readMessage(a, pr, peerproto.Bitfield);
+    _ = bf; // ignore it for now.
+    try peerproto.Interested.write(pw);
+    _ = try peerproto.readMessage(a, pr, peerproto.Unchoke);
+
+    var of = try std.fs.cwd().createFile(file.name, .{});
+    defer of.close();
+    errdefer {
+        // try to truncate the now-bad file...
+        of.setEndPos(0) catch {};
     }
-};
 
+    // Read the piece into memory, we'll check the hash before it goes to disk...
+    var piece_buf = try a.alloc(u8, mi.info.piece_length);
+    defer a.free(piece_buf);
+
+    for (0..mi.info.pieceCount()) |pi| {
+        const piece_length = @min(mi.info.piece_length, file.length - (pi * mi.info.piece_length));
+        var s1 = std.crypto.hash.Sha1.init(.{});
+
+        // Send a request message for each 16KiB block of the first piece
+        const blklen: u32 = 16*1024;
+        var blkcount = try std.math.divCeil(u32, piece_length, blklen);
+        for (0..blkcount) |i| {
+            const begin = std.math.cast(u32, i*blklen).?; 
+            const len = @min(blklen, piece_length - begin);
+            const req = peerproto.Request{
+                .index = @intCast(pi),
+                .begin = begin,
+                .length = len,
+            };
+            std.log.info("Request {any}", .{req});
+            try req.write(pw);
+            var piece = try peerproto.readMessage(a, pr, peerproto.Piece);
+            defer piece.deinit(a);
+            if (piece.index != req.index) return error.ProtocolError;
+            if (piece.begin != req.begin) return error.ProtocolError;
+            if (piece.block.len != req.length) return error.ProtocolError;
+            s1.update(piece.block);
+            @memcpy(piece_buf[piece.begin..piece.begin+piece.block.len], piece.block);
+        }
+        var ah = s1.finalResult();
+        var ph0 = mi.info.pieceHash(pi).?;
+        if (std.mem.eql(u8, &ah, &ph0)) {
+            try of.writeAll(piece_buf[0..piece_length]);
+        } else {
+            return error.BadHash;        
+        }
+    }
+    std.log.info("fin", .{});
+}
 
 test {
     _ = bencode;
     _ = MetaInfo;
+    _ = peerproto;
 }