Refactor writing peers & IPv6 Compact support. Closes #14.

This commit is contained in:
Jimmy Zelinskie 2014-06-28 20:44:43 -04:00
parent e0ead67737
commit 3999e35453
3 changed files with 159 additions and 79 deletions

View file

@ -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
}

View file

@ -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")
}

View file

@ -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