From 039f25f57194196d0ca604396256f93f9bcb7f93 Mon Sep 17 00:00:00 2001
From: Jimmy Zelinskie <jimmy.zelinskie+git@gmail.com>
Date: Wed, 3 May 2017 22:12:08 -0400
Subject: [PATCH] storage/memorybysubnet: add peerSubnet type

---
 storage/memorybysubnet/peer_store.go | 76 ++++++++++++++--------------
 1 file changed, 39 insertions(+), 37 deletions(-)

diff --git a/storage/memorybysubnet/peer_store.go b/storage/memorybysubnet/peer_store.go
index d5f1b47..22de2e6 100644
--- a/storage/memorybysubnet/peer_store.go
+++ b/storage/memorybysubnet/peer_store.go
@@ -120,15 +120,39 @@ func New(cfg Config) (storage.PeerStore, error) {
 
 type serializedPeer string
 
+func newPeerKey(p bittorrent.Peer) serializedPeer {
+	b := make([]byte, 20+2+len(p.IP.IP))
+	copy(b[:20], p.ID[:])
+	binary.BigEndian.PutUint16(b[20:22], p.Port)
+	copy(b[22:], p.IP.IP)
+
+	return serializedPeer(b)
+}
+
+type peerSubnet string
+
+func newPeerSubnet(ip bittorrent.IP, ipv4Mask, ipv6Mask net.IPMask) peerSubnet {
+	var maskedIP net.IP
+	switch ip.AddressFamily {
+	case bittorrent.IPv4:
+		maskedIP = ip.Mask(ipv4Mask)
+	case bittorrent.IPv6:
+		maskedIP = ip.Mask(ipv6Mask)
+	default:
+		panic("IP is neither v4 nor v6")
+	}
+
+	return peerSubnet(maskedIP.String())
+}
+
 type peerShard struct {
 	swarms map[bittorrent.InfoHash]swarm
 	sync.RWMutex
 }
 
 type swarm struct {
-	// map serialized peer to mtime
-	seeders  map[string]map[serializedPeer]int64
-	leechers map[string]map[serializedPeer]int64
+	seeders  map[peerSubnet]map[serializedPeer]int64
+	leechers map[peerSubnet]map[serializedPeer]int64
 }
 
 func (s swarm) lenSeeders() (i int) {
@@ -165,15 +189,6 @@ func (s *peerStore) shardIndex(infoHash bittorrent.InfoHash, af bittorrent.Addre
 	return idx
 }
 
-func newPeerKey(p bittorrent.Peer) serializedPeer {
-	b := make([]byte, 20+2+len(p.IP.IP))
-	copy(b[:20], p.ID[:])
-	binary.BigEndian.PutUint16(b[20:22], p.Port)
-	copy(b[22:], p.IP.IP)
-
-	return serializedPeer(b)
-}
-
 func decodePeerKey(pk serializedPeer) bittorrent.Peer {
 	peer := bittorrent.Peer{
 		ID:   bittorrent.PeerIDFromString(string(pk[:20])),
@@ -192,19 +207,6 @@ func decodePeerKey(pk serializedPeer) bittorrent.Peer {
 	return peer
 }
 
-func (s *peerStore) mask(p bittorrent.Peer) string {
-	var maskedIP net.IP
-	switch p.IP.AddressFamily {
-	case bittorrent.IPv4:
-		maskedIP = p.IP.IP.Mask(s.ipv4Mask)
-	case bittorrent.IPv6:
-		maskedIP = p.IP.IP.Mask(s.ipv6Mask)
-	default:
-		panic("IP is neither v4 nor v6")
-	}
-	return maskedIP.String()
-}
-
 func (s *peerStore) PutSeeder(ih bittorrent.InfoHash, p bittorrent.Peer) error {
 	select {
 	case <-s.closed:
@@ -219,13 +221,13 @@ func (s *peerStore) PutSeeder(ih bittorrent.InfoHash, p bittorrent.Peer) error {
 
 	if _, ok := shard.swarms[ih]; !ok {
 		shard.swarms[ih] = swarm{
-			seeders:  make(map[string]map[serializedPeer]int64),
-			leechers: make(map[string]map[serializedPeer]int64),
+			seeders:  make(map[peerSubnet]map[serializedPeer]int64),
+			leechers: make(map[peerSubnet]map[serializedPeer]int64),
 		}
 		recordInfohashesDelta(1)
 	}
 
-	preferredSubnet := s.mask(p)
+	preferredSubnet := newPeerSubnet(p.IP, s.ipv4Mask, s.ipv6Mask)
 	if shard.swarms[ih].seeders[preferredSubnet] == nil {
 		shard.swarms[ih].seeders[preferredSubnet] = make(map[serializedPeer]int64)
 	}
@@ -252,7 +254,7 @@ func (s *peerStore) DeleteSeeder(ih bittorrent.InfoHash, p bittorrent.Peer) erro
 		return storage.ErrResourceDoesNotExist
 	}
 
-	preferredSubnet := s.mask(p)
+	preferredSubnet := newPeerSubnet(p.IP, s.ipv4Mask, s.ipv6Mask)
 	if _, ok := shard.swarms[ih].seeders[preferredSubnet][pk]; !ok {
 		shard.Unlock()
 		return storage.ErrResourceDoesNotExist
@@ -283,13 +285,13 @@ func (s *peerStore) PutLeecher(ih bittorrent.InfoHash, p bittorrent.Peer) error
 
 	if _, ok := shard.swarms[ih]; !ok {
 		shard.swarms[ih] = swarm{
-			seeders:  make(map[string]map[serializedPeer]int64),
-			leechers: make(map[string]map[serializedPeer]int64),
+			seeders:  make(map[peerSubnet]map[serializedPeer]int64),
+			leechers: make(map[peerSubnet]map[serializedPeer]int64),
 		}
 		recordInfohashesDelta(1)
 	}
 
-	preferredSubnet := s.mask(p)
+	preferredSubnet := newPeerSubnet(p.IP, s.ipv4Mask, s.ipv6Mask)
 	if shard.swarms[ih].leechers[preferredSubnet] == nil {
 		shard.swarms[ih].leechers[preferredSubnet] = make(map[serializedPeer]int64)
 	}
@@ -316,7 +318,7 @@ func (s *peerStore) DeleteLeecher(ih bittorrent.InfoHash, p bittorrent.Peer) err
 		return storage.ErrResourceDoesNotExist
 	}
 
-	preferredSubnet := s.mask(p)
+	preferredSubnet := newPeerSubnet(p.IP, s.ipv4Mask, s.ipv6Mask)
 	if _, ok := shard.swarms[ih].leechers[preferredSubnet][pk]; !ok {
 		shard.Unlock()
 		return storage.ErrResourceDoesNotExist
@@ -347,13 +349,13 @@ func (s *peerStore) GraduateLeecher(ih bittorrent.InfoHash, p bittorrent.Peer) e
 
 	if _, ok := shard.swarms[ih]; !ok {
 		shard.swarms[ih] = swarm{
-			seeders:  make(map[string]map[serializedPeer]int64),
-			leechers: make(map[string]map[serializedPeer]int64),
+			seeders:  make(map[peerSubnet]map[serializedPeer]int64),
+			leechers: make(map[peerSubnet]map[serializedPeer]int64),
 		}
 		recordInfohashesDelta(1)
 	}
 
-	preferredSubnet := s.mask(p)
+	preferredSubnet := newPeerSubnet(p.IP, s.ipv4Mask, s.ipv6Mask)
 	delete(shard.swarms[ih].leechers[preferredSubnet], pk)
 
 	if shard.swarms[ih].seeders[preferredSubnet] == nil {
@@ -380,7 +382,7 @@ func (s *peerStore) AnnouncePeers(ih bittorrent.InfoHash, seeder bool, numWant i
 		return nil, storage.ErrResourceDoesNotExist
 	}
 
-	preferredSubnet := s.mask(announcer)
+	preferredSubnet := newPeerSubnet(announcer.IP, s.ipv4Mask, s.ipv6Mask)
 
 	if seeder {
 		// Append as many close leechers as possible.