From d317dfc06923318b1033c4dbdada7bf8f7776e6f Mon Sep 17 00:00:00 2001
From: Jimmy Zelinskie <jimmyzelinskie@gmail.com>
Date: Sat, 4 Apr 2015 15:41:58 -0400
Subject: [PATCH] udp: implement BEP45 response

---
 udp/protocol.go |  1 +
 udp/writer.go   | 54 ++++++++++++++++++++++++++++++++++++++++---------
 2 files changed, 45 insertions(+), 10 deletions(-)

diff --git a/udp/protocol.go b/udp/protocol.go
index ae40632..c236e79 100644
--- a/udp/protocol.go
+++ b/udp/protocol.go
@@ -18,6 +18,7 @@ const (
 	announceActionID
 	scrapeActionID
 	errorActionID
+	announceDualStackActionID
 )
 
 var (
diff --git a/udp/writer.go b/udp/writer.go
index b28fd8a..86e0d50 100644
--- a/udp/writer.go
+++ b/udp/writer.go
@@ -28,14 +28,33 @@ func (w *Writer) WriteError(err error) error {
 	return nil
 }
 
-// WriteAnnounce encodes an announce response according to the UDP spec.
-func (w *Writer) WriteAnnounce(res *models.AnnounceResponse) error {
-	w.writeHeader(announceActionID)
-	binary.Write(w.buf, binary.BigEndian, uint32(res.Interval/time.Second))
-	binary.Write(w.buf, binary.BigEndian, uint32(res.Incomplete))
-	binary.Write(w.buf, binary.BigEndian, uint32(res.Complete))
+// WriteAnnounce encodes an announce response by selecting the proper announce
+// format based on the BitTorrent spec.
+func (w *Writer) WriteAnnounce(resp *models.AnnounceResponse) (err error) {
+	if resp.Announce.HasIPv6() {
+		err = w.WriteAnnounceIPv6(resp)
+	} else {
+		err = w.WriteAnnounceIPv4(resp)
+	}
 
-	for _, peer := range res.IPv4Peers {
+	return
+}
+
+// WriteAnnounceIPv6 encodes an announce response according to BEP45.
+func (w *Writer) WriteAnnounceIPv6(resp *models.AnnounceResponse) error {
+	w.writeHeader(announceDualStackActionID)
+	binary.Write(w.buf, binary.BigEndian, uint32(resp.Interval/time.Second))
+	binary.Write(w.buf, binary.BigEndian, uint32(resp.Incomplete))
+	binary.Write(w.buf, binary.BigEndian, uint32(resp.Complete))
+	binary.Write(w.buf, binary.BigEndian, uint32(len(resp.IPv4Peers)))
+	binary.Write(w.buf, binary.BigEndian, uint32(len(resp.IPv6Peers)))
+
+	for _, peer := range resp.IPv4Peers {
+		w.buf.Write(peer.IP)
+		binary.Write(w.buf, binary.BigEndian, peer.Port)
+	}
+
+	for _, peer := range resp.IPv6Peers {
 		w.buf.Write(peer.IP)
 		binary.Write(w.buf, binary.BigEndian, peer.Port)
 	}
@@ -43,11 +62,26 @@ func (w *Writer) WriteAnnounce(res *models.AnnounceResponse) error {
 	return nil
 }
 
-// WriteScrape encodes a scrape response according to the UDP spec.
-func (w *Writer) WriteScrape(res *models.ScrapeResponse) error {
+// WriteAnnounceIPv4 encodes an announce response according to BEP15.
+func (w *Writer) WriteAnnounceIPv4(resp *models.AnnounceResponse) error {
+	w.writeHeader(announceActionID)
+	binary.Write(w.buf, binary.BigEndian, uint32(resp.Interval/time.Second))
+	binary.Write(w.buf, binary.BigEndian, uint32(resp.Incomplete))
+	binary.Write(w.buf, binary.BigEndian, uint32(resp.Complete))
+
+	for _, peer := range resp.IPv4Peers {
+		w.buf.Write(peer.IP)
+		binary.Write(w.buf, binary.BigEndian, peer.Port)
+	}
+
+	return nil
+}
+
+// WriteScrape encodes a scrape response according to BEP15.
+func (w *Writer) WriteScrape(resp *models.ScrapeResponse) error {
 	w.writeHeader(scrapeActionID)
 
-	for _, torrent := range res.Files {
+	for _, torrent := range resp.Files {
 		binary.Write(w.buf, binary.BigEndian, uint32(torrent.Seeders.Len()))
 		binary.Write(w.buf, binary.BigEndian, uint32(torrent.Snatches))
 		binary.Write(w.buf, binary.BigEndian, uint32(torrent.Leechers.Len()))