wstracker.go

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