codecrafters submit [skip ci]

2 files changed, 87 insertions, 17 deletions

diff --git a/cmd/mybittorrent/client.go b/cmd/mybittorrent/client.go
index f3988b6..ef3d860 100644
--- a/cmd/mybittorrent/client.go
+++ b/cmd/mybittorrent/client.go
@@ -1,9 +1,13 @@
 package main
 
 import (
+	"bytes"
 	"encoding/binary"
+	"encoding/hex"
 	"fmt"
+	"io"
 	"log/slog"
+	"net"
 	"net/http"
 	"net/url"
 	"os"
@@ -30,6 +34,8 @@ type PeerResponse struct {
 	Peers    []string
 }
 
+type HandshakeMessage []byte
+
 func NewClient(peerId string, port int) *Client {
 	return &Client{
 		PeerId:   peerId,
@@ -134,3 +140,51 @@ func (c *Client) GetPeers(filename string) (PeerResponse, error) {
 		Peers:    DecodePeers([]byte(resp.Peers)),
 	}, err
 }
+
+func (m HandshakeMessage) PeerIdHex() string {
+
+	return hex.EncodeToString(m[48:])
+
+}
+
+func (c *Client) Handshake(filename, peerAddr string) (HandshakeMessage, error) {
+	ct, ok := c.Torrents[filename]
+	if !ok {
+		return nil, fmt.Errorf("missing torrent file: %s", filename)
+	}
+
+	buf := new(bytes.Buffer)
+	buf.WriteByte(19)
+	buf.WriteString("BitTorrent protocol")
+	buf.Write([]byte{0, 0, 0, 0, 0, 0, 0, 0})
+
+	hash, err := ct.Meta.InfoHash()
+	if err != nil {
+		return nil, err
+	}
+
+	buf.Write(hash)
+	buf.WriteString(c.PeerId)
+
+	conn, err := net.Dial("tcp", peerAddr)
+	defer func(conn net.Conn) {
+		err := conn.Close()
+		if err != nil {
+			slog.Error("failed to close connection", "peerAddr", peerAddr)
+		}
+	}(conn)
+
+	if err != nil {
+		return nil, err
+	}
+
+	_, err = buf.WriteTo(conn)
+	if err != nil {
+		return nil, err
+	}
+
+	respBuf := make([]byte, 68)
+	_, err = io.LimitReader(conn, 68).Read(respBuf)
+
+	return respBuf, nil
+}
diff --git a/cmd/mybittorrent/main.go b/cmd/mybittorrent/main.go
index c460c35..f000aa5 100644
--- a/cmd/mybittorrent/main.go
+++ b/cmd/mybittorrent/main.go
@@ -9,6 +9,18 @@ import (
 	bencode "github.com/jackpal/bencode-go" // Available if you need it!
 )
 
+func createClient() *Client {
+	fn := os.Args[2]
+	c := NewClient("00112233445566778899", 6881)
+
+	_, err := c.AddTorrentFile(fn)
+	if err != nil {
+		panic(err)
+	}
+
+	return c
+}
+
 func main() {
 	command := os.Args[1]
 
@@ -25,46 +37,50 @@ func main() {
 		fmt.Println(string(jsonOutput))
 
 	case "info":
-		c := NewClient("00112233445566778899", 6881)
-		tf, err := c.AddTorrentFile(os.Args[2])
-		if err != nil {
-			panic(err)
-		}
+		fn := os.Args[2]
+		c := createClient()
+		meta := c.Torrents[fn].Meta
+		fmt.Printf("Tracker URL: %s\n", meta.Announce)
+		fmt.Printf("Length: %d\n", meta.Info.Length)
 
-		fmt.Printf("Tracker URL: %s\n", tf.Meta.Announce)
-		fmt.Printf("Length: %d\n", tf.Meta.Info.Length)
-
-		infoHash, err := tf.Meta.InfoHash()
+		infoHash, err := meta.InfoHash()
 		if err != nil {
 			panic(err)
 		}
 
 		fmt.Printf("Info Hash: %x", infoHash)
 
-		fmt.Printf("Piece Length: %d\n", tf.Meta.Info.PieceLength)
+		fmt.Printf("Piece Length: %d\n", meta.Info.PieceLength)
 		fmt.Println("Piece Hashes:")
 
-		for _, h := range tf.Meta.PieceHashes() {
+		for _, h := range meta.PieceHashes() {
 			fmt.Printf("%x\n", h)
 		}
 
 	case "peers":
-		c := NewClient("00112233445566778899", 6881)
+		fn := os.Args[2]
+		c := createClient()
 
-		_, err := c.AddTorrentFile(os.Args[2])
+		pr, err := c.GetPeers(fn)
 		if err != nil {
 			panic(err)
 		}
 
-		pr, err := c.GetPeers(os.Args[2])
+		for _, peer := range pr.Peers {
+			fmt.Println(peer)
+		}
+
+	case "handshake":
+		fn := os.Args[2]
+		c := createClient()
+		peerAddr := os.Args[3]
 
+		hs, err := c.Handshake(fn, peerAddr)
 		if err != nil {
 			panic(err)
 		}
 
-		for _, peer := range pr.Peers {
-			fmt.Println(peer)
-		}
+		fmt.Printf("Peer ID: %s\n", hs.PeerIdHex())
 
 	default:
 		fmt.Println("Unknown command: " + command)