wstracker.go

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