diff --git a/bittorrent/bittorrent.go b/bittorrent/bittorrent.go index 8f0e98c..241c08a 100644 --- a/bittorrent/bittorrent.go +++ b/bittorrent/bittorrent.go @@ -84,8 +84,8 @@ type AnnounceRequest struct { // response. type AnnounceResponse struct { Compact bool - Complete int32 - Incomplete int32 + Complete uint32 + Incomplete uint32 Interval time.Duration MinInterval time.Duration IPv4Peers []Peer diff --git a/frontend/udp/writer.go b/frontend/udp/writer.go index d8fc87d..f0800ed 100644 --- a/frontend/udp/writer.go +++ b/frontend/udp/writer.go @@ -37,8 +37,8 @@ func WriteAnnounce(w io.Writer, txID []byte, resp *bittorrent.AnnounceResponse, writeHeader(buf, txID, announceActionID) } binary.Write(buf, binary.BigEndian, uint32(resp.Interval/time.Second)) - binary.Write(buf, binary.BigEndian, uint32(resp.Incomplete)) - binary.Write(buf, binary.BigEndian, uint32(resp.Complete)) + binary.Write(buf, binary.BigEndian, resp.Incomplete) + binary.Write(buf, binary.BigEndian, resp.Complete) peers := resp.IPv4Peers if v6 { diff --git a/middleware/hooks.go b/middleware/hooks.go index 5269ba4..769fb7d 100644 --- a/middleware/hooks.go +++ b/middleware/hooks.go @@ -2,6 +2,8 @@ package middleware import ( "context" + "errors" + "net" "github.com/chihaya/chihaya/bittorrent" "github.com/chihaya/chihaya/storage" @@ -64,3 +66,62 @@ func (h *swarmInteractionHook) HandleScrape(ctx context.Context, _ *bittorrent.S // Scrapes have no effect on the swarm. return ctx, nil } + +// ErrInvalidIP indicates an invalid IP for an Announce. +var ErrInvalidIP = errors.New("invalid IP") + +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 + } + + s := h.store.ScrapeSwarm(req.InfoHash, len(req.IP) == net.IPv6len) + resp.Incomplete = s.Incomplete + resp.Complete = s.Complete + + switch len(req.IP) { + case net.IPv4len: + resp.IPv4Peers, err = h.store.AnnouncePeers(req.InfoHash, req.Left == 0, int(req.NumWant), req.Peer) + case net.IPv6len: + resp.IPv6Peers, err = h.store.AnnouncePeers(req.InfoHash, req.Left == 0, int(req.NumWant), req.Peer) + default: + return ctx, ErrInvalidIP + } + + return ctx, err +} + +func (h *responseHook) HandleScrape(ctx context.Context, req *bittorrent.ScrapeRequest, resp *bittorrent.ScrapeResponse) (context.Context, error) { + if ctx.Value(SkipResponseHookKey) != nil { + return ctx, nil + } + + v6, _ := ctx.Value(ScrapeIsIPv6Key).(bool) + + for _, infoHash := range req.InfoHashes { + resp.Files[infoHash] = h.store.ScrapeSwarm(infoHash, v6) + } + + return ctx, nil +} diff --git a/middleware/middleware.go b/middleware/middleware.go index dec1e9f..77b246f 100644 --- a/middleware/middleware.go +++ b/middleware/middleware.go @@ -25,7 +25,7 @@ func NewLogic(cfg Config, peerStore storage.PeerStore, preHooks, postHooks []Hoo l := &Logic{ announceInterval: cfg.AnnounceInterval, peerStore: peerStore, - preHooks: preHooks, + preHooks: append(preHooks, &responseHook{store: peerStore}), postHooks: append(postHooks, &swarmInteractionHook{store: peerStore}), } @@ -44,7 +44,9 @@ type Logic struct { // HandleAnnounce generates a response for an Announce. func (l *Logic) HandleAnnounce(ctx context.Context, req *bittorrent.AnnounceRequest) (resp *bittorrent.AnnounceResponse, err error) { resp = &bittorrent.AnnounceResponse{ - Interval: l.announceInterval, + Interval: l.announceInterval, + MinInterval: l.announceInterval, + Compact: req.Compact, } for _, h := range l.preHooks { if ctx, err = h.HandleAnnounce(ctx, req, resp); err != nil {