From 3999e35453b8691264e524ade5a98e4d4aa63e33 Mon Sep 17 00:00:00 2001 From: Jimmy Zelinskie Date: Sat, 28 Jun 2014 20:44:43 -0400 Subject: [PATCH] Refactor writing peers & IPv6 Compact support. Closes #14. --- models/models.go | 32 ++++---- models/query/query.go | 36 ++++++--- server/serve_announce.go | 170 ++++++++++++++++++++++++++------------- 3 files changed, 159 insertions(+), 79 deletions(-) diff --git a/models/models.go b/models/models.go index 9647272..24cf67e 100644 --- a/models/models.go +++ b/models/models.go @@ -8,6 +8,7 @@ package models import ( "errors" + "net" "net/http" "path" "strconv" @@ -29,7 +30,7 @@ type Peer struct { UserID uint64 `json:"user_id"` TorrentID uint64 `json:"torrent_id"` - IP string `json:"ip"` + IP net.IP `json:"ip"` Port uint64 `json:"port"` Uploaded uint64 `json:"uploaded"` @@ -92,15 +93,15 @@ type Torrent struct { } // InSeederPool returns true if a peer is within a Torrent's pool of seeders. -func (t *Torrent) InSeederPool(p *Peer) bool { - _, exists := t.Seeders[p.Key()] - return exists +func (t *Torrent) InSeederPool(p *Peer) (exists bool) { + _, exists = t.Seeders[p.Key()] + return } // InLeecherPool returns true if a peer is within a Torrent's pool of leechers. -func (t *Torrent) InLeecherPool(p *Peer) bool { - _, exists := t.Leechers[p.Key()] - return exists +func (t *Torrent) InLeecherPool(p *Peer) (exists bool) { + _, exists = t.Leechers[p.Key()] + return } // User is a registered user for private trackers. @@ -121,7 +122,7 @@ type Announce struct { Compact bool `json:"compact"` Downloaded uint64 `json:"downloaded"` Event string `json:"event"` - IP string `json:"ip"` + IP net.IP `json:"ip"` Infohash string `json:"infohash"` Left uint64 `json:"left"` NumWant int `json:"numwant"` @@ -139,15 +140,18 @@ func NewAnnounce(r *http.Request, conf *config.Config) (*Announce, error) { } compact := q.Params["compact"] != "0" - downloaded, downloadedErr := q.Uint64("downloaded") event, _ := q.Params["event"] infohash, _ := q.Params["info_hash"] - ip, _ := q.RequestedIP(r) - left, leftErr := q.Uint64("left") - numWant := q.RequestedPeerCount(conf.NumWantFallback) - dir, _ := path.Split(r.URL.Path) peerID, _ := q.Params["peer_id"] + + dir, _ := path.Split(r.URL.Path) + numWant := q.RequestedPeerCount(conf.NumWantFallback) + + ip, ipErr := q.RequestedIP(r) port, portErr := q.Uint64("port") + + left, leftErr := q.Uint64("left") + downloaded, downloadedErr := q.Uint64("downloaded") uploaded, uploadedErr := q.Uint64("uploaded") if downloadedErr != nil || @@ -156,7 +160,7 @@ func NewAnnounce(r *http.Request, conf *config.Config) (*Announce, error) { peerID == "" || portErr != nil || uploadedErr != nil || - ip == "" || + ipErr != nil || len(dir) != 34 { return nil, ErrMalformedRequest } diff --git a/models/query/query.go b/models/query/query.go index 93c51e6..6c74ae5 100644 --- a/models/query/query.go +++ b/models/query/query.go @@ -7,6 +7,7 @@ package query import ( "errors" + "net" "net/http" "net/url" "strconv" @@ -120,21 +121,35 @@ func (q Query) RequestedPeerCount(fallback int) int { } // RequestedIP returns the requested IP address from a Query. -func (q Query) RequestedIP(r *http.Request) (string, error) { - if ip, ok := q.Params["ip"]; ok { - return ip, nil +func (q Query) RequestedIP(r *http.Request) (net.IP, error) { + if ipstr, ok := q.Params["ip"]; ok { + if ip := net.ParseIP(ipstr); ip != nil { + return ip, nil + } } - if ip, ok := q.Params["ipv4"]; ok { - return ip, nil + if ipstr, ok := q.Params["ipv4"]; ok { + if ip := net.ParseIP(ipstr); ip != nil { + return ip, nil + } + } + + if ipstr, ok := q.Params["ipv6"]; ok { + if ip := net.ParseIP(ipstr); ip != nil { + return ip, nil + } } if xRealIPs, ok := q.Params["X-Real-Ip"]; ok { - return string(xRealIPs[0]), nil + if ip := net.ParseIP(string(xRealIPs[0])); ip != nil { + return ip, nil + } } if r.RemoteAddr == "" { - return "127.0.0.1", nil + if ip := net.ParseIP("127.0.0.1"); ip != nil { + return ip, nil + } } portIndex := len(r.RemoteAddr) - 1 @@ -145,8 +160,11 @@ func (q Query) RequestedIP(r *http.Request) (string, error) { } if portIndex != -1 { - return r.RemoteAddr[0:portIndex], nil + ipstr := r.RemoteAddr[0:portIndex] + if ip := net.ParseIP(ipstr); ip != nil { + return ip, nil + } } - return "", errors.New("failed to parse IP address") + return nil, errors.New("failed to parse IP address") } diff --git a/server/serve_announce.go b/server/serve_announce.go index 42066c1..00791b9 100644 --- a/server/serve_announce.go +++ b/server/serve_announce.go @@ -6,7 +6,6 @@ package server import ( "io" - "net" "net/http" "strconv" @@ -74,14 +73,12 @@ func (s Server) serveAnnounce(w http.ResponseWriter, r *http.Request) { w.(http.Flusher).Flush() - if log.V(5) { - log.Infof( - "announce: ip: %s, user: %s, torrent: %s", - announce.IP, - user.ID, - torrent.ID, - ) - } + log.V(5).Infof( + "announce: ip: %s, user: %s, torrent: %s", + announce.IP, + user.ID, + torrent.ID, + ) } func updateTorrent(c tracker.Conn, a *models.Announce, p *models.Peer, t *models.Torrent) (created bool, err error) { @@ -165,10 +162,17 @@ func handleEvent(c tracker.Conn, a *models.Announce, p *models.Peer, u *models.U } func writeAnnounceResponse(w io.Writer, a *models.Announce, u *models.User, t *models.Torrent) { - bencoder := bencode.NewEncoder(w) seedCount := len(t.Seeders) leechCount := len(t.Leechers) + var peerCount int + if a.Left == 0 { + peerCount = minInt(a.NumWant, leechCount) + } else { + peerCount = minInt(a.NumWant, leechCount+seedCount-1) + } + + bencoder := bencode.NewEncoder(w) bencoder.Encode("d") bencoder.Encode("complete") bencoder.Encode(seedCount) @@ -180,71 +184,125 @@ func writeAnnounceResponse(w io.Writer, a *models.Announce, u *models.User, t *m bencoder.Encode(a.Config.MinAnnounce.Duration) if a.NumWant > 0 && a.Event != "stopped" && a.Event != "paused" { - bencoder.Encode("peers") - - var peerCount int - if a.Left == 0 { - peerCount = minInt(a.NumWant, leechCount) - } else { - peerCount = minInt(a.NumWant, leechCount+seedCount-1) - } - if a.Compact { - // 6 is the number of bytes 1 compact peer takes up. - bencoder.Encode(strconv.Itoa(peerCount * 6)) - bencoder.Encode(":") + writePeersCompact(w, a, u, t, peerCount) } else { - bencoder.Encode("l") - } - - if a.Left == 0 { - // If they're seeding, give them only leechers - writePeers(w, u, t.Leechers, peerCount, a.Compact) - } else { - // If they're leeching, prioritize giving them seeders - count := writePeers(w, u, t.Seeders, peerCount, a.Compact) - writePeers(w, u, t.Leechers, peerCount-count, a.Compact) - } - - if !a.Compact { - bencoder.Encode("e") + writePeersList(w, a, u, t, peerCount) } } + bencoder.Encode("e") } -func writePeers(w io.Writer, user *models.User, peers map[string]models.Peer, numWant int, compact bool) (count int) { +func writePeersCompact(w io.Writer, a *models.Announce, u *models.User, t *models.Torrent, peerCount int) { + ipv4s, ipv6s := getPeers(a, u, t, peerCount) bencoder := bencode.NewEncoder(w) - for _, peer := range peers { - if count >= numWant { - break - } - if peer.UserID == user.ID { - continue - } - - if compact { - if ip := net.ParseIP(peer.IP); ip != nil { + if len(ipv4s) > 0 { + // 6 is the number of bytes that represents 1 compact IPv4 address. + bencoder.Encode("peers") + bencoder.Encode(strconv.Itoa(len(ipv4s) * 6)) + bencoder.Encode(":") + for _, peer := range ipv4s { + if ip := peer.IP.To4(); ip != nil { w.Write(ip) w.Write([]byte{byte(peer.Port >> 8), byte(peer.Port & 0xff)}) } - } else { - bencoder.Encode("d") - bencoder.Encode("ip") - bencoder.Encode(peer.IP) - bencoder.Encode("peer id") - bencoder.Encode(peer.ID) - bencoder.Encode("port") - bencoder.Encode(peer.Port) - bencoder.Encode("e") } + } + + if len(ipv6s) > 0 { + // 18 is the number of bytes that represents 1 compact IPv6 address. + bencoder.Encode("peers6") + bencoder.Encode(strconv.Itoa(len(ipv6s) * 18)) + for _, peer := range ipv6s { + if ip := peer.IP.To16(); ip != nil { + w.Write(ip) + w.Write([]byte{byte(peer.Port >> 8), byte(peer.Port & 0xff)}) + } + } + } +} + +func getPeers(a *models.Announce, u *models.User, t *models.Torrent, peerCount int) (ipv4s, ipv6s []*models.Peer) { + if a.Left == 0 { + // If they're seeding, give them only leechers. + splitPeers(ipv4s, ipv6s, u, t.Leechers, peerCount) + } else { + // If they're leeching, prioritize giving them seeders. + count := splitPeers(ipv4s, ipv6s, u, t.Seeders, peerCount) + splitPeers(ipv4s, ipv6s, u, t.Leechers, peerCount-count) + } + + return +} + +func splitPeers(ipv4s, ipv6s []*models.Peer, u *models.User, peers map[string]models.Peer, peerCount int) (count int) { + for _, peer := range peers { + if count >= peerCount { + break + } + + if peer.UserID == u.ID { + continue + } + + if ip := peer.IP.To4(); ip != nil { + ipv4s = append(ipv4s, &peer) + } else { + ipv6s = append(ipv6s, &peer) + } + count++ } return } +func writePeersList(w io.Writer, a *models.Announce, u *models.User, t *models.Torrent, peerCount int) { + bencoder := bencode.NewEncoder(w) + bencoder.Encode("peers") + bencoder.Encode("l") + + if a.Left == 0 { + // If they're seeding, give them only leechers + writePeerDicts(w, u, t.Leechers, peerCount) + } else { + // If they're leeching, prioritize giving them seeders + count := writePeerDicts(w, u, t.Seeders, peerCount) + writePeerDicts(w, u, t.Leechers, peerCount-count) + } + + bencoder.Encode("e") +} + +func writePeerDicts(w io.Writer, u *models.User, peers map[string]models.Peer, peerCount int) (count int) { + bencoder := bencode.NewEncoder(w) + + for _, peer := range peers { + if count >= peerCount { + break + } + + if peer.UserID == u.ID { + continue + } + + bencoder.Encode("d") + bencoder.Encode("ip") + bencoder.Encode(peer.IP) + bencoder.Encode("peer id") + bencoder.Encode(peer.ID) + bencoder.Encode("port") + bencoder.Encode(peer.Port) + bencoder.Encode("e") + + count++ + } + + return count +} + func minInt(a, b int) int { if a < b { return a