From 3bc2276fb3c4719330b8c4bc7d7459a8d1462846 Mon Sep 17 00:00:00 2001
From: "Lawrence, Rendall" <rlawrence@sot-te.ch>
Date: Thu, 21 Apr 2022 21:57:18 +0300
Subject: [PATCH] (tested) Refactor code

* merge 9d04e4c from https://github.com/jzelinskie/chihaya
* sanitize ip address on Scrape requests
* remove NewConnectionID and ValidConnectionID functions from production code
---
 bittorrent/sanitize.go             |  5 +++++
 frontend/http/parser.go            |  2 +-
 frontend/http/writer.go            |  1 -
 frontend/udp/connection_id.go      | 15 ---------------
 frontend/udp/connection_id_test.go | 15 +++++++++++++++
 storage/memory/storage.go          |  2 +-
 6 files changed, 22 insertions(+), 18 deletions(-)

diff --git a/bittorrent/sanitize.go b/bittorrent/sanitize.go
index 7cffaaf..c6bb42d 100644
--- a/bittorrent/sanitize.go
+++ b/bittorrent/sanitize.go
@@ -46,6 +46,11 @@ func SanitizeScrape(r *ScrapeRequest, maxScrapeInfoHashes uint32) error {
 		r.InfoHashes = r.InfoHashes[:maxScrapeInfoHashes]
 	}
 
+	r.AddrPort = netip.AddrPortFrom(r.Addr(), r.Port())
+	if !r.Addr().IsValid() || r.Addr().IsUnspecified() {
+		return ErrInvalidIP
+	}
+
 	log.Debug("sanitized scrape", r, log.Fields{
 		"maxScrapeInfoHashes": maxScrapeInfoHashes,
 	})
diff --git a/frontend/http/parser.go b/frontend/http/parser.go
index a2e655f..43fbb5a 100644
--- a/frontend/http/parser.go
+++ b/frontend/http/parser.go
@@ -178,5 +178,5 @@ func requestedIP(r *http.Request, p bittorrent.Params, opts ParseOptions) (netip
 	}
 
 	addrPort, err := netip.ParseAddrPort(r.RemoteAddr)
-	return addrPort.Addr().Unmap(), false, err
+	return addrPort.Addr(), false, err
 }
diff --git a/frontend/http/writer.go b/frontend/http/writer.go
index bc452ae..b88fade 100644
--- a/frontend/http/writer.go
+++ b/frontend/http/writer.go
@@ -22,7 +22,6 @@ func WriteError(w http.ResponseWriter, err error) {
 		log.Error("http: internal error", log.Err(err))
 	}
 
-	w.WriteHeader(http.StatusOK)
 	if err = bencode.NewEncoder(w).Encode(map[string]any{
 		"failure reason": message,
 	}); err != nil {
diff --git a/frontend/udp/connection_id.go b/frontend/udp/connection_id.go
index 74b8ec0..f8bcd5a 100644
--- a/frontend/udp/connection_id.go
+++ b/frontend/udp/connection_id.go
@@ -15,21 +15,6 @@ import (
 // ttl is the duration a connection ID should be valid according to BEP 15.
 const ttl = 2 * time.Minute
 
-// NewConnectionID creates an 8-byte connection identifier for UDP packets as
-// described by BEP 15.
-// This is a wrapper around creating a new ConnectionIDGenerator and generating
-// an ID. It is recommended to use the generator for performance.
-func NewConnectionID(ip netip.Addr, now time.Time, key string) []byte {
-	return NewConnectionIDGenerator(key).Generate(ip, now)
-}
-
-// ValidConnectionID determines whether a connection identifier is legitimate.
-// This is a wrapper around creating a new ConnectionIDGenerator and validating
-// the ID. It is recommended to use the generator for performance.
-func ValidConnectionID(connectionID []byte, ip netip.Addr, now time.Time, maxClockSkew time.Duration, key string) bool {
-	return NewConnectionIDGenerator(key).Validate(connectionID, ip, now, maxClockSkew)
-}
-
 // A ConnectionIDGenerator is a reusable generator and validator for connection
 // IDs as described in BEP 15.
 // It is not thread safe, but is safe to be pooled and reused by other
diff --git a/frontend/udp/connection_id_test.go b/frontend/udp/connection_id_test.go
index 5289242..edc44b2 100644
--- a/frontend/udp/connection_id_test.go
+++ b/frontend/udp/connection_id_test.go
@@ -27,6 +27,21 @@ var golden = []struct {
 	{0, 0, "::1", "", true},
 }
 
+// NewConnectionID creates an 8-byte connection identifier for UDP packets as
+// described by BEP 15.
+// This is a wrapper around creating a new ConnectionIDGenerator and generating
+// an ID. It is recommended to use the generator for performance.
+func NewConnectionID(ip netip.Addr, now time.Time, key string) []byte {
+	return NewConnectionIDGenerator(key).Generate(ip, now)
+}
+
+// ValidConnectionID determines whether a connection identifier is legitimate.
+// This is a wrapper around creating a new ConnectionIDGenerator and validating
+// the ID. It is recommended to use the generator for performance.
+func ValidConnectionID(connectionID []byte, ip netip.Addr, now time.Time, maxClockSkew time.Duration, key string) bool {
+	return NewConnectionIDGenerator(key).Validate(connectionID, ip, now, maxClockSkew)
+}
+
 // simpleNewConnectionID generates a new connection ID the explicit way.
 // This is used to verify correct behaviour of the generator.
 func simpleNewConnectionID(ip netip.Addr, now time.Time, key string) []byte {
diff --git a/storage/memory/storage.go b/storage/memory/storage.go
index 018a50b..dd27c38 100644
--- a/storage/memory/storage.go
+++ b/storage/memory/storage.go
@@ -229,7 +229,7 @@ func (ps *peerStore) shardIndex(infoHash bittorrent.InfoHash, addr netip.Addr) u
 	// half is dedicated to IPv4 swarms and the second half is dedicated to
 	// IPv6 swarms.
 	idx := binary.BigEndian.Uint32([]byte(infoHash[:4])) % (uint32(len(ps.shards)) / 2)
-	if addr.Is6() && !addr.Is4In6() {
+	if addr.Is6() {
 		idx += uint32(len(ps.shards) / 2)
 	}
 	return idx