Merge pull request #197 from mrd0ll4r/refactor-udp

Refactor udp

frontend/udp/frontend.go

@@ -6,7 +6,6 @@ import (
 	"bytes"
 	"context"
 	"encoding/binary"
-	"log"
 	"net"
 	"sync"
 	"time"
@@ -94,13 +93,12 @@ func (t *Frontend) ListenAndServe() error {
 	}
 	defer t.socket.Close()
 
-	pool := bytepool.New(256, 2048)
+	pool := bytepool.New(2048, 2048)
 
 	for {
 		// Check to see if we need to shutdown.
 		select {
 		case <-t.closing:
-			t.wg.Wait()
 			return nil
 		default:
 		}
@@ -124,7 +122,6 @@ func (t *Frontend) ListenAndServe() error {
 			continue
 		}
 
-		log.Println("Got UDP Request")
 		t.wg.Add(1)
 		go func() {
 			defer t.wg.Done()
@@ -132,11 +129,10 @@ func (t *Frontend) ListenAndServe() error {
 
 			// Handle the request.
 			start := time.Now()
-			response, action, err := t.handleRequest(
+			action, err := t.handleRequest(
 				Request{buffer[:n], addr.IP},
 				ResponseWriter{t.socket, addr},
 			)
-			log.Printf("Handled UDP Request: %s, %s, %s\n", response, action, err)
 			recordResponseDuration(action, err, time.Since(start))
 		}()
 	}
@@ -162,7 +158,7 @@ func (w ResponseWriter) Write(b []byte) (int, error) {
 }
 
 // handleRequest parses and responds to a UDP Request.
-func (t *Frontend) handleRequest(r Request, w ResponseWriter) (response []byte, actionName string, err error) {
+func (t *Frontend) handleRequest(r Request, w ResponseWriter) (actionName string, err error) {
 	if len(r.Packet) < 16 {
 		// Malformed, no client packets are less than 16 bytes.
 		// We explicitly return nothing in case this is a DoS attempt.

frontend/udp/parser.go

@@ -12,7 +12,6 @@ const (
 	announceActionID
 	scrapeActionID
 	errorActionID
-	announceDualStackActionID
 )
 
 // Option-Types as described in BEP 41 and BEP 45.

frontend/udp/writer.go

@@ -26,32 +26,44 @@ func WriteError(w io.Writer, txID []byte, err error) {
 
 // WriteAnnounce encodes an announce response according to BEP 15.
 func WriteAnnounce(w io.Writer, txID []byte, resp *bittorrent.AnnounceResponse) {
-	writeHeader(w, txID, announceActionID)
+	var buf bytes.Buffer
-	binary.Write(w, binary.BigEndian, uint32(resp.Interval/time.Second))
+
-	binary.Write(w, binary.BigEndian, uint32(resp.Incomplete))
+	writeHeader(&buf, txID, announceActionID)
-	binary.Write(w, binary.BigEndian, uint32(resp.Complete))
+	binary.Write(&buf, binary.BigEndian, uint32(resp.Interval/time.Second))
+	binary.Write(&buf, binary.BigEndian, uint32(resp.Incomplete))
+	binary.Write(&buf, binary.BigEndian, uint32(resp.Complete))
 
 	for _, peer := range resp.IPv4Peers {
-		w.Write(peer.IP)
+		buf.Write(peer.IP)
-		binary.Write(w, binary.BigEndian, peer.Port)
+		binary.Write(&buf, binary.BigEndian, peer.Port)
 	}
+
+	w.Write(buf.Bytes())
 }
 
 // WriteScrape encodes a scrape response according to BEP 15.
 func WriteScrape(w io.Writer, txID []byte, resp *bittorrent.ScrapeResponse) {
-	writeHeader(w, txID, scrapeActionID)
+	var buf bytes.Buffer
+
+	writeHeader(&buf, txID, scrapeActionID)
 
 	for _, scrape := range resp.Files {
-		binary.Write(w, binary.BigEndian, scrape.Complete)
+		binary.Write(&buf, binary.BigEndian, scrape.Complete)
-		binary.Write(w, binary.BigEndian, scrape.Snatches)
+		binary.Write(&buf, binary.BigEndian, scrape.Snatches)
-		binary.Write(w, binary.BigEndian, scrape.Incomplete)
+		binary.Write(&buf, binary.BigEndian, scrape.Incomplete)
 	}
+
+	w.Write(buf.Bytes())
 }
 
 // WriteConnectionID encodes a new connection response according to BEP 15.
 func WriteConnectionID(w io.Writer, txID, connID []byte) {
-	writeHeader(w, txID, connectActionID)
+	var buf bytes.Buffer
-	w.Write(connID)
+
+	writeHeader(&buf, txID, connectActionID)
+	buf.Write(connID)
+
+	w.Write(buf.Bytes())
 }
 
 // writeHeader writes the action and transaction ID to the provided response