diff --git a/middleware/hooks.go b/middleware/hooks.go index 0cfbeac..5269ba4 100644 --- a/middleware/hooks.go +++ b/middleware/hooks.go @@ -4,6 +4,7 @@ import ( "context" "github.com/chihaya/chihaya/bittorrent" + "github.com/chihaya/chihaya/storage" ) // Hook abstracts the concept of anything that needs to interact with a @@ -12,3 +13,54 @@ 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.Event == bittorrent.Stopped: + err = h.store.DeleteSeeder(req.InfoHash, req.Peer) + if err != nil && err != storage.ErrResourceDoesNotExist { + return ctx, err + } + + err = h.store.DeleteLeecher(req.InfoHash, req.Peer) + if err != nil && 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 +} diff --git a/middleware/middleware.go b/middleware/middleware.go index 56cff22..dec1e9f 100644 --- a/middleware/middleware.go +++ b/middleware/middleware.go @@ -26,7 +26,7 @@ func NewLogic(cfg Config, peerStore storage.PeerStore, preHooks, postHooks []Hoo announceInterval: cfg.AnnounceInterval, peerStore: peerStore, preHooks: preHooks, - postHooks: postHooks, + postHooks: append(postHooks, &swarmInteractionHook{store: peerStore}), } return l