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