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