zbt

main.zig (5430B)

---

 const std = @import("std");
 const MetaInfo = @import("metainfo.zig");
 const bencode = @import("bencode.zig");
 const AnyWriter = @import("anywriter.zig");
 const peerproto = @import("peer_protocol.zig");
 const trackproto = @import("tracker_protocol.zig");
 
 // 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.
 const peer_id = [20]u8{ 0x30, 0x30, 0x31, 0x31, 0x32, 0x32, 0x33, 0x33, 0x34, 0x34, 0x35, 0x35, 0x36, 0x36, 0x37, 0x37, 0x38, 0x38, 0x39, 0x39 };
 // To make this a practical torrent app...
 // Add some UI? or keep it a console app? I might keep it as a console app for now
 const Args = struct {
     torrent_file: []const u8,
 };
 
 fn read_args() !Args {
     var args = std.process.args();
     if (!args.skip()) return error.NotEnoughArgs;
     const nxt = args.next() orelse return error.NotEnoughArgs;
     if (std.mem.eql(u8, nxt, "-h") or std.mem.eql(u8, nxt, "--help")) {
         try printUsage(false);
         return error.Help;
     }
     return .{
         .torrent_file = nxt,
     };
 }
 
 fn do_main() !void {
     var gpa = std.heap.GeneralPurposeAllocator(.{}){};
     defer _ = gpa.deinit();
     const a = gpa.allocator();
 
     const args = read_args() catch |e| switch (e) {
         error.Help => return,
         else => |err| return err,
     };
 
     const f = try std.fs.cwd().openFile(args.torrent_file, .{});
     defer f.close();
     var fr = f.reader();
     var mib = try bencode.bdecode(a, fr);
     defer mib.deinit(a);
     var mi = try MetaInfo.parse(a, mib);
     defer mi.deinit(a);
     const info_hash = try mi.info.hash(a);
 
     var c = std.http.Client{
         .allocator = a,
     };
     defer c.deinit();
 
     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) {
         return error.TrackerHttpError;
     }
 
     var trb = try bencode.bdecodeBuf(a, res.body.?);
     defer trb.deinit(a);
     var tr = try trackproto.TrackerResp.parse(a, trb);
     defer tr.deinit(a);
 
     if (tr.peers.len == 0) {
         std.log.info("no peers", .{});
         return;
     }
 
     for (tr.peers) |peer| {
         std.log.info("peer: {}", .{peer});
     }
 
     // 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) });
 
     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", .{});
 }
 
 pub fn main() !void {
     do_main() catch |e| {
         try printUsage(true);
         return e;
     };
 }
 
 fn printUsage(err: bool) !void {
     var f = if (err) std.io.getStdErr() else std.io.getStdOut();
     try f.writeAll(
         \\ Usage: zbt TORRENTFILE 
         \\ Download the data specfied by TORRENTFILE to the current directory
         \\ 
     );
 }
 
 test {
     _ = bencode;
     _ = MetaInfo;
     _ = peerproto;
 }