tracker/udp-server_test.go

   1 package tracker
   2
   3 import (
   4         "bytes"
   5         "encoding"
   6         "encoding/binary"
   7         "fmt"
   8         "math/rand"
   9         "net"
  10
  11         "github.com/anacrolix/dht/v2/krpc"
  12         "github.com/anacrolix/missinggo/v2"
  13         "github.com/anacrolix/torrent/tracker/udp"
  14 )
  15
  16 type torrent struct {
  17         Leechers int32
  18         Seeders  int32
  19         Peers    []krpc.NodeAddr
  20 }
  21
  22 type server struct {
  23         pc    net.PacketConn
  24         conns map[udp.ConnectionId]struct{}
  25         t     map[[20]byte]torrent
  26 }
  27
  28 func marshal(parts ...interface{}) (ret []byte, err error) {
  29         var buf bytes.Buffer
  30         for _, p := range parts {
  31                 err = binary.Write(&buf, binary.BigEndian, p)
  32                 if err != nil {
  33                         return
  34                 }
  35         }
  36         ret = buf.Bytes()
  37         return
  38 }
  39
  40 func (s *server) respond(addr net.Addr, rh udp.ResponseHeader, parts ...interface{}) (err error) {
  41         b, err := marshal(append([]interface{}{rh}, parts...)...)
  42         if err != nil {
  43                 return
  44         }
  45         _, err = s.pc.WriteTo(b, addr)
  46         return
  47 }
  48
  49 func (s *server) newConn() (ret udp.ConnectionId) {
  50         ret = rand.Uint64()
  51         if s.conns == nil {
  52                 s.conns = make(map[udp.ConnectionId]struct{})
  53         }
  54         s.conns[ret] = struct{}{}
  55         return
  56 }
  57
  58 func (s *server) serveOne() (err error) {
  59         b := make([]byte, 0x10000)
  60         n, addr, err := s.pc.ReadFrom(b)
  61         if err != nil {
  62                 return
  63         }
  64         r := bytes.NewReader(b[:n])
  65         var h udp.RequestHeader
  66         err = udp.Read(r, &h)
  67         if err != nil {
  68                 return
  69         }
  70         switch h.Action {
  71         case udp.ActionConnect:
  72                 if h.ConnectionId != udp.ConnectRequestConnectionId {
  73                         return
  74                 }
  75                 connId := s.newConn()
  76                 err = s.respond(addr, udp.ResponseHeader{
  77                         udp.ActionConnect,
  78                         h.TransactionId,
  79                 }, udp.ConnectionResponse{
  80                         connId,
  81                 })
  82                 return
  83         case udp.ActionAnnounce:
  84                 if _, ok := s.conns[h.ConnectionId]; !ok {
  85                         s.respond(addr, udp.ResponseHeader{
  86                                 TransactionId: h.TransactionId,
  87                                 Action:        udp.ActionError,
  88                         }, []byte("not connected"))
  89                         return
  90                 }
  91                 var ar AnnounceRequest
  92                 err = udp.Read(r, &ar)
  93                 if err != nil {
  94                         return
  95                 }
  96                 t := s.t[ar.InfoHash]
  97                 bm := func() encoding.BinaryMarshaler {
  98                         ip := missinggo.AddrIP(addr)
  99                         if ip.To4() != nil {
 100                                 return krpc.CompactIPv4NodeAddrs(t.Peers)
 101                         }
 102                         return krpc.CompactIPv6NodeAddrs(t.Peers)
 103                 }()
 104                 b, err = bm.MarshalBinary()
 105                 if err != nil {
 106                         panic(err)
 107                 }
 108                 err = s.respond(addr, udp.ResponseHeader{
 109                         TransactionId: h.TransactionId,
 110                         Action:        udp.ActionAnnounce,
 111                 }, udp.AnnounceResponseHeader{
 112                         Interval: 900,
 113                         Leechers: t.Leechers,
 114                         Seeders:  t.Seeders,
 115                 }, b)
 116                 return
 117         default:
 118                 err = fmt.Errorf("unhandled action: %d", h.Action)
 119                 s.respond(addr, udp.ResponseHeader{
 120                         TransactionId: h.TransactionId,
 121                         Action:        udp.ActionError,
 122                 }, []byte("unhandled action"))
 123                 return
 124         }
 125 }