github.com/anacrolix/torrent@v1.61.0/tracker/udp-server_test.go (about)

     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  }