wstracker.go

   1 package torrent
   2
   3 import (
   4         "context"
   5         "fmt"
   6         "net"
   7         netHttp "net/http"
   8         "net/url"
   9         "sync"
  10
  11         "github.com/anacrolix/log"
  12         "github.com/gorilla/websocket"
  13         "github.com/pion/datachannel"
  14         "github.com/pion/webrtc/v3"
  15
  16         "github.com/anacrolix/torrent/tracker"
  17         httpTracker "github.com/anacrolix/torrent/tracker/http"
  18         "github.com/anacrolix/torrent/webtorrent"
  19 )
  20
  21 type websocketTrackerStatus struct {
  22         url url.URL
  23         tc  *webtorrent.TrackerClient
  24 }
  25
  26 func (me websocketTrackerStatus) statusLine() string {
  27         return fmt.Sprintf("%+v", me.tc.Stats())
  28 }
  29
  30 func (me websocketTrackerStatus) URL() *url.URL {
  31         return &me.url
  32 }
  33
  34 type refCountedWebtorrentTrackerClient struct {
  35         webtorrent.TrackerClient
  36         refCount int
  37 }
  38
  39 type websocketTrackers struct {
  40         PeerId                     [20]byte
  41         Logger                     log.Logger
  42         GetAnnounceRequest         func(event tracker.AnnounceEvent, infoHash [20]byte) (tracker.AnnounceRequest, error)
  43         OnConn                     func(datachannel.ReadWriteCloser, webtorrent.DataChannelContext)
  44         mu                         sync.Mutex
  45         clients                    map[string]*refCountedWebtorrentTrackerClient
  46         Proxy                      httpTracker.ProxyFunc
  47         DialContext                func(ctx context.Context, network, addr string) (net.Conn, error)
  48         WebsocketTrackerHttpHeader func() netHttp.Header
  49         ICEServers                 []webrtc.ICEServer
  50 }
  51
  52 func (me *websocketTrackers) Get(url string, infoHash [20]byte) (*webtorrent.TrackerClient, func()) {
  53         me.mu.Lock()
  54         defer me.mu.Unlock()
  55         value, ok := me.clients[url]
  56         if !ok {
  57                 dialer := &websocket.Dialer{Proxy: me.Proxy, NetDialContext: me.DialContext, HandshakeTimeout: websocket.DefaultDialer.HandshakeTimeout}
  58                 value = &refCountedWebtorrentTrackerClient{
  59                         TrackerClient: webtorrent.TrackerClient{
  60                                 Dialer:             dialer,
  61                                 Url:                url,
  62                                 GetAnnounceRequest: me.GetAnnounceRequest,
  63                                 PeerId:             me.PeerId,
  64                                 OnConn:             me.OnConn,
  65                                 Logger: me.Logger.WithText(func(m log.Msg) string {
  66                                         return fmt.Sprintf("tracker client for %q: %v", url, m)
  67                                 }),
  68                                 WebsocketTrackerHttpHeader: me.WebsocketTrackerHttpHeader,
  69                                 ICEServers:                 me.ICEServers,
  70                         },
  71                 }
  72                 value.TrackerClient.Start(func(err error) {
  73                         if err != nil {
  74                                 me.Logger.Printf("error running tracker client for %q: %v", url, err)
  75                         }
  76                 })
  77                 if me.clients == nil {
  78                         me.clients = make(map[string]*refCountedWebtorrentTrackerClient)
  79                 }
  80                 me.clients[url] = value
  81         }
  82         value.refCount++
  83         return &value.TrackerClient, func() {
  84                 me.mu.Lock()
  85                 defer me.mu.Unlock()
  86                 value.TrackerClient.CloseOffersForInfohash(infoHash)
  87                 value.refCount--
  88                 if value.refCount == 0 {
  89                         value.TrackerClient.Close()
  90                         delete(me.clients, url)
  91                 }
  92         }
  93 }