package middleware import ( "context" "errors" "github.com/chihaya/chihaya/bittorrent" "github.com/chihaya/chihaya/storage" ) // Hook abstracts the concept of anything that needs to interact with a // BitTorrent client's request and response to a BitTorrent tracker. // PreHooks and PostHooks both use the same interface. // // A Hook can implement stop.Stopper if clean shutdown is required. type Hook interface { HandleAnnounce(context.Context, *bittorrent.AnnounceRequest, *bittorrent.AnnounceResponse) (context.Context, error) HandleScrape(context.Context, *bittorrent.ScrapeRequest, *bittorrent.ScrapeResponse) (context.Context, error) } type skipSwarmInteraction struct{} // SkipSwarmInteractionKey is a key for the context of an Announce to control // whether the swarm interaction middleware should run. // Any non-nil value set for this key will cause the swarm interaction // middleware to skip. var SkipSwarmInteractionKey = skipSwarmInteraction{} type swarmInteractionHook struct { store storage.PeerStore } func (h *swarmInteractionHook) HandleAnnounce(ctx context.Context, req *bittorrent.AnnounceRequest, resp *bittorrent.AnnounceResponse) (_ context.Context, err error) { if ctx.Value(SkipSwarmInteractionKey) != nil { return ctx, nil } switch { case req.Port < 100: return ctx, nil case req.Event == bittorrent.Stopped: err = h.store.DeleteSeeder(req.InfoHash, req.Peer) if err != nil && !errors.Is(err, storage.ErrResourceDoesNotExist) { return ctx, err } err = h.store.DeleteLeecher(req.InfoHash, req.Peer) if err != nil && !errors.Is(err, storage.ErrResourceDoesNotExist) { return ctx, err } case req.Event == bittorrent.Completed: err = h.store.GraduateLeecher(req.InfoHash, req.Peer) return ctx, err case req.Left == 0: // Completed events will also have Left == 0, but by making this // an extra case we can treat "old" seeders differently from // graduating leechers. (Calling PutSeeder is probably faster // than calling GraduateLeecher.) err = h.store.PutSeeder(req.InfoHash, req.Peer) return ctx, err default: err = h.store.PutLeecher(req.InfoHash, req.Peer) return ctx, err } return ctx, nil } func (h *swarmInteractionHook) HandleScrape(ctx context.Context, _ *bittorrent.ScrapeRequest, _ *bittorrent.ScrapeResponse) (context.Context, error) { // Scrapes have no effect on the swarm. return ctx, nil } type skipResponseHook struct{} // SkipResponseHookKey is a key for the context of an Announce or Scrape to // control whether the response middleware should run. // Any non-nil value set for this key will cause the response middleware to // skip. var SkipResponseHookKey = skipResponseHook{} type scrapeAddressType struct{} // ScrapeIsIPv6Key is the key under which to store whether or not the // address used to request a scrape was an IPv6 address. // The value is expected to be of type bool. // A missing value or a value that is not a bool for this key is equivalent to // it being set to false. var ScrapeIsIPv6Key = scrapeAddressType{} type responseHook struct { store storage.PeerStore } func (h *responseHook) HandleAnnounce(ctx context.Context, req *bittorrent.AnnounceRequest, resp *bittorrent.AnnounceResponse) (_ context.Context, err error) { if ctx.Value(SkipResponseHookKey) != nil { return ctx, nil } // Add the Scrape data to the response. s := h.store.ScrapeSwarm(req.InfoHash, req.IP.AddressFamily) resp.Incomplete = s.Incomplete resp.Complete = s.Complete err = h.appendPeers(req, resp) return ctx, err } func (h *responseHook) appendPeers(req *bittorrent.AnnounceRequest, resp *bittorrent.AnnounceResponse) error { seeding := req.Left == 0 peers, err := h.store.AnnouncePeers(req.InfoHash, seeding, int(req.NumWant), req.Peer) if err != nil && !errors.Is(err, storage.ErrResourceDoesNotExist) { return err } // Some clients expect a minimum of their own peer representation returned to // them if they are the only peer in a swarm. if len(peers) == 0 { if seeding { resp.Complete++ } else { resp.Incomplete++ } peers = append(peers, req.Peer) } switch req.IP.AddressFamily { case bittorrent.IPv4: resp.IPv4Peers = append(resp.IPv4Peers, peers...) case bittorrent.IPv6: resp.IPv6Peers = append(resp.IPv6Peers, peers...) default: panic("attempted to append peer that is neither IPv4 nor IPv6") } return nil } func (h *responseHook) HandleScrape(ctx context.Context, req *bittorrent.ScrapeRequest, resp *bittorrent.ScrapeResponse) (context.Context, error) { if ctx.Value(SkipResponseHookKey) != nil { return ctx, nil } for _, infoHash := range req.InfoHashes { resp.Files = append(resp.Files, h.store.ScrapeSwarm(infoHash, req.AddressFamily)) } return ctx, nil }