diff --git a/bittorrent/bittorrent.go b/bittorrent/peer.go similarity index 91% rename from bittorrent/bittorrent.go rename to bittorrent/peer.go index 8a62a29..701413b 100644 --- a/bittorrent/bittorrent.go +++ b/bittorrent/peer.go @@ -117,22 +117,6 @@ func (i InfoHash) RawString() string { return string(i) } -// Scrape represents the state of a swarm that is returned in a scrape response. -type Scrape struct { - InfoHash InfoHash - Snatches uint32 - Complete uint32 - Incomplete uint32 -} - -// MarshalZerologObject writes fields into zerolog event -func (s Scrape) MarshalZerologObject(e *zerolog.Event) { - e.Stringer("infoHash", s.InfoHash). - Uint32("snatches", s.Snatches). - Uint32("complete", s.Complete). - Uint32("incomplete", s.Incomplete) -} - // Peer represents the connection details of a peer that is returned in an // announce response. type Peer struct { diff --git a/bittorrent/request.go b/bittorrent/request.go index 9ea096d..b5cd4a0 100644 --- a/bittorrent/request.go +++ b/bittorrent/request.go @@ -35,19 +35,19 @@ func (a RequestAddress) MarshalZerologObject(e *zerolog.Event) { // connection information about peer type RequestAddresses []RequestAddress -func (aa RequestAddresses) Len() int { - return len(aa) +func (aa *RequestAddresses) Len() int { + return len(*aa) } // Less returns true only if i-th RequestAddress is marked as // RequestAddress.Provided and j-th is not (provided address has // higher priority) -func (aa RequestAddresses) Less(i, j int) bool { - return aa[i].Provided && !aa[j].Provided +func (aa *RequestAddresses) Less(i, j int) bool { + return (*aa)[i].Provided && !(*aa)[j].Provided } -func (aa RequestAddresses) Swap(i, j int) { - aa[i], aa[j] = aa[j], aa[i] +func (aa *RequestAddresses) Swap(i, j int) { + (*aa)[i], (*aa)[j] = (*aa)[j], (*aa)[i] } // Add checks if provided RequestAddress is valid and adds unmapped @@ -82,25 +82,27 @@ func (aa *RequestAddresses) Sanitize(ignorePrivate bool) bool { *aa = append(*aa, RequestAddress{a, p}) } if len(*aa) > 1 { - sort.Sort(*aa) + sort.Sort(aa) } return len(uniqueAddresses) > 0 } // GetFirst returns first address from array // or empty netip.Addr if array is empty -func (aa RequestAddresses) GetFirst() netip.Addr { +func (aa *RequestAddresses) GetFirst() netip.Addr { var a netip.Addr - if len(aa) > 0 { - a = aa[0].Addr + if len(*aa) > 0 { + a = (*aa)[0].Addr } return a } // MarshalZerologArray writes array elements to zerolog event -func (aa RequestAddresses) MarshalZerologArray(a *zerolog.Array) { - for _, addr := range aa { - a.Object(addr) +func (aa *RequestAddresses) MarshalZerologArray(a *zerolog.Array) { + if aa != nil { + for _, addr := range *aa { + a.Object(addr) + } } } @@ -137,7 +139,7 @@ func (rp RequestPeer) Peers() (peers Peers) { // MarshalZerologObject writes fields into zerolog event func (rp RequestPeer) MarshalZerologObject(e *zerolog.Event) { e.Stringer("id", rp.ID). - Array("addresses", rp.RequestAddresses). + Array("addresses", &rp.RequestAddresses). Uint16("port", rp.Port) } @@ -216,18 +218,48 @@ type ScrapeRequest struct { // MarshalZerologObject writes fields into zerolog event func (r ScrapeRequest) MarshalZerologObject(e *zerolog.Event) { - e.Array("addresses", r.RequestAddresses). + e.Array("addresses", &r.RequestAddresses). Array("infoHashes", r.InfoHashes). Object("params", r.Params) } +// Scrape represents the state of a swarm that is returned in a scrape response. +type Scrape struct { + InfoHash InfoHash + Snatches uint32 + Complete uint32 + Incomplete uint32 +} + +// MarshalZerologObject writes fields into zerolog event +func (s Scrape) MarshalZerologObject(e *zerolog.Event) { + e.Stringer("infoHash", s.InfoHash). + Uint32("snatches", s.Snatches). + Uint32("complete", s.Complete). + Uint32("incomplete", s.Incomplete) +} + // Scrapes wrapper of array of Scrape-s type Scrapes []Scrape +func (s *Scrapes) Len() int { + return len(*s) +} + +func (s *Scrapes) Less(i, j int) bool { + return (*s)[i].InfoHash < (*s)[j].InfoHash +} + +func (s *Scrapes) Swap(i, j int) { + (*s)[i], (*s)[j] = (*s)[j], (*s)[i] +} + // MarshalZerologArray writes array elements to zerolog event -func (s Scrapes) MarshalZerologArray(a *zerolog.Array) { - for _, scrape := range s { - a.Object(scrape) +func (s *Scrapes) MarshalZerologArray(a *zerolog.Array) { + if s != nil { + for _, scrape := range *s { + a.Object(scrape) + } } } @@ -241,5 +273,5 @@ type ScrapeResponse struct { // MarshalZerologObject writes fields into zerolog event func (sr ScrapeResponse) MarshalZerologObject(e *zerolog.Event) { - e.Array("scrapes", sr.Files) + e.Array("scrapes", &sr.Files) } diff --git a/frontend/http/frontend.go b/frontend/http/frontend.go index e5a51b6..8fee959 100644 --- a/frontend/http/frontend.go +++ b/frontend/http/frontend.go @@ -6,14 +6,11 @@ import ( "context" "crypto/tls" "errors" - "net" "net/http" "net/netip" - "sync" "time" "github.com/julienschmidt/httprouter" - "github.com/libp2p/go-reuseport" "github.com/sot-tech/mochi/bittorrent" "github.com/sot-tech/mochi/frontend" @@ -26,7 +23,11 @@ import ( const Name = "http" -var logger = log.NewLogger(Name) +var ( + logger = log.NewLogger(Name) + errTLSNotProvided = errors.New("tls certificate/key not provided") + errRoutesNotProvided = errors.New("routes not provided") +) func init() { frontend.RegisterBuilder(Name, newFrontend) @@ -34,241 +35,135 @@ func init() { // Config represents all configurable options for an HTTP BitTorrent Frontend type Config struct { - Addr string - HTTPSAddr string `cfg:"https_addr"` - ReadTimeout time.Duration `cfg:"read_timeout"` - WriteTimeout time.Duration `cfg:"write_timeout"` - IdleTimeout time.Duration `cfg:"idle_timeout"` - EnableKeepAlive bool `cfg:"enable_keepalive"` - TLSCertPath string `cfg:"tls_cert_path"` - TLSKeyPath string `cfg:"tls_key_path"` - ReusePort bool `cfg:"reuse_port"` - AnnounceRoutes []string `cfg:"announce_routes"` - ScrapeRoutes []string `cfg:"scrape_routes"` - PingRoutes []string `cfg:"ping_routes"` - EnableRequestTiming bool `cfg:"enable_request_timing"` + frontend.ListenOptions + IdleTimeout time.Duration `cfg:"idle_timeout"` + EnableKeepAlive bool `cfg:"enable_keepalive"` + UseTLS bool `cfg:"tls"` + TLSCertPath string `cfg:"tls_cert_path"` + TLSKeyPath string `cfg:"tls_key_path"` + AnnounceRoutes []string `cfg:"announce_routes"` + ScrapeRoutes []string `cfg:"scrape_routes"` + PingRoutes []string `cfg:"ping_routes"` ParseOptions } -// Default config constants. -const ( - defaultReadTimeout = 2 * time.Second - defaultWriteTimeout = 2 * time.Second - defaultIdleTimeout = 30 * time.Second -) +const defaultIdleTimeout = 30 * time.Second // Validate sanity checks values set in a config and returns a new config with // default values replacing anything that is invalid. // // This function warns to the logger when a value is changed. -func (cfg Config) Validate() Config { - validcfg := cfg - - if cfg.ReadTimeout <= 0 { - validcfg.ReadTimeout = defaultReadTimeout - logger.Warn(). - Str("name", "http.ReadTimeout"). - Dur("provided", cfg.ReadTimeout). - Dur("default", validcfg.ReadTimeout). - Msg("falling back to default configuration") +func (cfg Config) Validate() (validCfg Config, err error) { + if validCfg.ListenOptions, err = cfg.ListenOptions.Validate(); err != nil { + return } - - if cfg.WriteTimeout <= 0 { - validcfg.WriteTimeout = defaultWriteTimeout - logger.Warn(). - Str("name", "http.WriteTimeout"). - Dur("provided", cfg.WriteTimeout). - Dur("default", validcfg.WriteTimeout). - Msg("falling back to default configuration") + if cfg.UseTLS && (len(cfg.TLSCertPath) == 0 || len(cfg.TLSKeyPath) == 0) { + err = errTLSNotProvided + return } - if cfg.IdleTimeout <= 0 { - validcfg.IdleTimeout = defaultIdleTimeout - + validCfg.IdleTimeout = defaultIdleTimeout if cfg.EnableKeepAlive { // If keepalive is disabled, this configuration isn't used anyway. logger.Warn(). Str("name", "http.IdleTimeout"). Dur("provided", cfg.IdleTimeout). - Dur("default", validcfg.IdleTimeout). + Dur("default", validCfg.IdleTimeout). Msg("falling back to default configuration") } } - - validcfg.ParseOptions.ParseOptions = cfg.ParseOptions.ParseOptions.Validate() - - return validcfg + validCfg.ParseOptions.ParseOptions = cfg.ParseOptions.ParseOptions.Validate() + return } type httpFE struct { - srv *http.Server - srvMu sync.Mutex - tlsSrv *http.Server - tlsSrvMu sync.Mutex - tlsCfg *tls.Config - - logic *middleware.Logic - Config + srv *http.Server + logic *middleware.Logic + collectTimings bool + ParseOptions } -func newFrontend(c conf.MapConfig, logic *middleware.Logic) (frontend.Frontend, error) { - var provided Config - if err := c.Unmarshal(&provided); err != nil { - return nil, err +func newFrontend(c conf.MapConfig, logic *middleware.Logic) (_ frontend.Frontend, err error) { + var cfg Config + if err = c.Unmarshal(&cfg); err != nil { + return + } + if cfg, err = cfg.Validate(); err != nil { + return + } + if len(cfg.AnnounceRoutes) < 1 || len(cfg.ScrapeRoutes) < 1 { + err = errRoutesNotProvided + return } - cfg := provided.Validate() f := &httpFE{ - logic: logic, - Config: cfg, - srvMu: sync.Mutex{}, - tlsSrvMu: sync.Mutex{}, - } - - if cfg.Addr == "" && cfg.HTTPSAddr == "" { - return nil, errors.New("must specify addr or https_addr or both") - } - - if len(cfg.AnnounceRoutes) < 1 || len(cfg.ScrapeRoutes) < 1 { - return nil, errors.New("must specify routes") + logic: logic, + srv: &http.Server{ + ReadTimeout: cfg.ReadTimeout, + ReadHeaderTimeout: cfg.ReadTimeout, + WriteTimeout: cfg.WriteTimeout, + IdleTimeout: cfg.IdleTimeout, + }, + collectTimings: cfg.EnableRequestTiming, + ParseOptions: cfg.ParseOptions, } // If TLS is enabled, create a key pair. - if cfg.TLSCertPath != "" && cfg.TLSKeyPath != "" { - var err error - f.tlsCfg = &tls.Config{ - Certificates: make([]tls.Certificate, 1), + if cfg.UseTLS { + var cert tls.Certificate + if cert, err = tls.LoadX509KeyPair(cfg.TLSCertPath, cfg.TLSKeyPath); err != nil { + return + } + f.srv.TLSConfig = &tls.Config{ + Certificates: []tls.Certificate{cert}, MinVersion: tls.VersionTLS12, } - f.tlsCfg.Certificates[0], err = tls.LoadX509KeyPair(cfg.TLSCertPath, cfg.TLSKeyPath) - if err != nil { - return nil, err - } - } - - if (cfg.HTTPSAddr == "") != (f.tlsCfg == nil) { - return nil, errors.New("must specify both https_addr, tls_cert_path and tls_key_path") } router := httprouter.New() - for _, route := range f.AnnounceRoutes { + for _, route := range cfg.AnnounceRoutes { router.GET(route, f.announceRoute) } - for _, route := range f.ScrapeRoutes { + for _, route := range cfg.ScrapeRoutes { router.GET(route, f.scrapeRoute) } + for _, route := range cfg.PingRoutes { + router.GET(route, f.ping) + router.HEAD(route, f.ping) + } + f.srv.Handler = router - if len(f.PingRoutes) > 0 { - for _, route := range f.PingRoutes { - router.GET(route, f.ping) - router.HEAD(route, f.ping) + f.srv.SetKeepAlivesEnabled(cfg.EnableKeepAlive) + + go func() { + ln, err := cfg.ListenTCP() + defer logger.Err(ln.Close()).Msg("closing listener") + if err == nil { + if f.srv.TLSConfig == nil { + err = f.srv.Serve(ln) + } else { + err = f.srv.ServeTLS(ln, "", "") + } } - } - - if cfg.Addr != "" { - go func() { - if err := f.serveHTTP(router, false); err != nil { - logger.Fatal().Err(err).Str("proto", "http").Msg("failed while serving") - } - }() - } - - if cfg.HTTPSAddr != "" { - go func() { - if err := f.serveHTTP(router, true); err != nil { - logger.Fatal().Err(err).Str("proto", "https").Msg("failed while serving") - } - }() - } + if errors.Is(err, http.ErrServerClosed) { + err = nil + } else { + logger.Fatal().Err(err).Msg("server failed") + } + }() return f, nil } // Stop provides a thread-safe way to shut down a currently running Frontend. func (f *httpFE) Stop() stop.Result { - stopGroup := stop.NewGroup() - - f.srvMu.Lock() + c := make(stop.Channel) if f.srv != nil { - stopGroup.AddFunc(f.makeStopFunc(f.srv)) - } - f.srvMu.Unlock() - - f.tlsSrvMu.Lock() - if f.tlsSrv != nil { - stopGroup.AddFunc(f.makeStopFunc(f.tlsSrv)) - } - f.tlsSrvMu.Unlock() - - return stopGroup.Stop() -} - -func (f *httpFE) makeStopFunc(stopSrv *http.Server) stop.Func { - return func() stop.Result { - c := make(stop.Channel) go func() { - c.Done(stopSrv.Shutdown(context.Background())) + c.Done(f.srv.Shutdown(context.Background())) }() - return c.Result() } -} - -// serveHTTP blocks while listening and serving non-TLS HTTP BitTorrent -// requests until Stop() is called or an error is returned. -func (f *httpFE) serveHTTP(handler http.Handler, tls bool) error { - srv := &http.Server{ - Handler: handler, - ReadTimeout: f.ReadTimeout, - ReadHeaderTimeout: f.ReadTimeout, - WriteTimeout: f.WriteTimeout, - IdleTimeout: f.IdleTimeout, - } - - srv.SetKeepAlivesEnabled(f.EnableKeepAlive) - - var err error - if tls { - if f.ReusePort { - var ln net.Listener - if ln, err = reuseport.Listen("tcp", f.HTTPSAddr); err == nil { - defer ln.Close() - srv.TLSConfig = f.tlsCfg - f.tlsSrvMu.Lock() - f.tlsSrv = srv - f.tlsSrvMu.Unlock() - err = srv.ServeTLS(ln, "", "") - } - } else { - srv.Addr = f.HTTPSAddr - srv.TLSConfig = f.tlsCfg - f.tlsSrvMu.Lock() - f.tlsSrv = srv - f.tlsSrvMu.Unlock() - err = srv.ListenAndServeTLS("", "") - } - } else { - if f.ReusePort { - var ln net.Listener - if ln, err = reuseport.Listen("tcp", f.Addr); err == nil { - defer ln.Close() - f.srvMu.Lock() - f.srv = srv - f.srvMu.Unlock() - err = srv.Serve(ln) - } - } else { - srv.Addr = f.Addr - f.srvMu.Lock() - f.srv = srv - f.srvMu.Unlock() - err = srv.ListenAndServe() - } - } - // Start the HTTP server. - if errors.Is(err, http.ErrServerClosed) { - err = nil - } - return err + return c.Result() } func injectRouteParamsToContext(ctx context.Context, ps httprouter.Params) context.Context { @@ -285,7 +180,7 @@ func (f *httpFE) announceRoute(w http.ResponseWriter, r *http.Request, ps httpro var start time.Time var addr netip.Addr var req *bittorrent.AnnounceRequest - if f.EnableRequestTiming && metrics.Enabled() { + if f.collectTimings && metrics.Enabled() { start = time.Now() defer func() { recordResponseDuration("announce", addr, err, time.Since(start)) @@ -321,7 +216,7 @@ func (f *httpFE) scrapeRoute(w http.ResponseWriter, r *http.Request, ps httprout var err error var start time.Time var addr netip.Addr - if f.EnableRequestTiming && metrics.Enabled() { + if f.collectTimings && metrics.Enabled() { start = time.Now() defer func() { recordResponseDuration("scrape", addr, err, time.Since(start)) diff --git a/frontend/http/writer.go b/frontend/http/writer.go index 8b194be..c78ab76 100644 --- a/frontend/http/writer.go +++ b/frontend/http/writer.go @@ -1,16 +1,20 @@ package http import ( + "bytes" "errors" - "net" "net/http" + "strconv" "time" "github.com/anacrolix/torrent/bencode" "github.com/sot-tech/mochi/bittorrent" + "github.com/sot-tech/mochi/pkg/bytepool" ) +var respBufferPool = bytepool.NewBufferPool() + // WriteError communicates an error to a BitTorrent client over HTTP. func WriteError(w http.ResponseWriter, err error) { message := "internal server error" @@ -20,17 +24,15 @@ func WriteError(w http.ResponseWriter, err error) { } else { logger.Error().Err(err).Msg("http: internal error") } - - if err = bencode.NewEncoder(w).Encode(map[string]any{ - "failure reason": message, - }); err != nil { - logger.Error().Err(err).Msg("unable to encode message") - } + _, _ = w.Write([]byte("d14:failure reason" + strconv.Itoa(len(message)) + ":" + message + "e")) } // WriteAnnounceResponse communicates the results of an Announce to a // BitTorrent client over HTTP. func WriteAnnounceResponse(w http.ResponseWriter, resp *bittorrent.AnnounceResponse) error { + bb := respBufferPool.Get() + defer respBufferPool.Put(bb) + if resp.Interval > 0 { resp.Interval /= time.Second } @@ -39,54 +41,79 @@ func WriteAnnounceResponse(w http.ResponseWriter, resp *bittorrent.AnnounceRespo resp.MinInterval /= time.Second } - bdict := map[string]any{ - "complete": resp.Complete, - "incomplete": resp.Incomplete, - "interval": resp.Interval, - "min interval": resp.MinInterval, - } + bb.WriteString("d8:completei") + bb.WriteString(strconv.FormatUint(uint64(resp.Complete), 10)) + bb.WriteString("e10:incompletei") + bb.WriteString(strconv.FormatUint(uint64(resp.Incomplete), 10)) + bb.WriteString("e8:intervali") + bb.WriteString(strconv.FormatUint(uint64(resp.Interval), 10)) + bb.WriteString("e12:min intervali") + bb.WriteString(strconv.FormatUint(uint64(resp.MinInterval), 10)) + bb.WriteByte('e') // Add the peers to the dictionary in the compact format. if resp.Compact { // Add the IPv4 peers to the dictionary. - compactAddresses := make([]byte, 0, (net.IPv4len+2)*len(resp.IPv4Peers)) + bb.WriteString("5:peersl") for _, peer := range resp.IPv4Peers { - compactAddresses = append(compactAddresses, compactAddress(peer)...) - } - if len(compactAddresses) > 0 { - bdict["peers"] = compactAddresses + compactAddress(bb, peer) } + bb.WriteByte('e') // Add the IPv6 peers to the dictionary. - compactAddresses = make([]byte, 0, (net.IPv6len+2)*len(resp.IPv6Peers)) // IP + port + bb.WriteString("6:peers6l") for _, peer := range resp.IPv6Peers { - compactAddresses = append(compactAddresses, compactAddress(peer)...) - } - if len(compactAddresses) > 0 { - bdict["peers6"] = compactAddresses + compactAddress(bb, peer) } + bb.WriteByte('e') } else { // Add the peers to the dictionary. - peers := make([]map[string]any, 0, len(resp.IPv4Peers)+len(resp.IPv6Peers)) // IP + port + bb.WriteString("5:peersl") for _, peer := range resp.IPv4Peers { - peers = append(peers, dict(peer)) + dictAddress(bb, peer) } for _, peer := range resp.IPv6Peers { - peers = append(peers, dict(peer)) + dictAddress(bb, peer) } - bdict["peers"] = peers + bb.WriteByte('e') } + bb.WriteByte('e') - return bencode.NewEncoder(w).Encode(bdict) + _, err := bb.WriteTo(w) + return err +} + +func compactAddress(bb *bytes.Buffer, peer bittorrent.Peer) { + addr, port := peer.Addr().AsSlice(), peer.Port() + bb.WriteString(strconv.Itoa(len(addr) + 2)) + bb.WriteByte(':') + bb.Write(addr) + bb.WriteByte(byte(port >> 8)) + bb.WriteByte(byte(port)) + return +} + +func dictAddress(bb *bytes.Buffer, peer bittorrent.Peer) { + bb.WriteString("d2:ip") + addr := peer.Addr().String() + bb.WriteString(strconv.Itoa(len(addr))) + bb.WriteByte(':') + bb.WriteString(addr) + bb.WriteString("7:peer id20:") + bb.WriteString(peer.ID.RawString()) + bb.WriteString("4:porti") + bb.WriteString(strconv.FormatUint(uint64(peer.Port()), 10)) + bb.WriteString("ee") } // WriteScrapeResponse communicates the results of a Scrape to a BitTorrent // client over HTTP. func WriteScrapeResponse(w http.ResponseWriter, resp *bittorrent.ScrapeResponse) error { - filesDict := make(map[string]any, len(resp.Files)) + filesDict := make(map[bittorrent.InfoHash]any, len(resp.Files)) for _, scrape := range resp.Files { - filesDict[string(scrape.InfoHash[:])] = map[string]any{ + filesDict[scrape.InfoHash] = map[string]any{ "complete": scrape.Complete, + "downloaded": scrape.Snatches, "incomplete": scrape.Incomplete, } } @@ -95,18 +122,3 @@ func WriteScrapeResponse(w http.ResponseWriter, resp *bittorrent.ScrapeResponse) "files": filesDict, }) } - -func compactAddress(peer bittorrent.Peer) (buf []byte) { - buf = append(buf, peer.Addr().AsSlice()...) - port := peer.Port() - buf = append(buf, byte(port>>8), byte(port)) - return -} - -func dict(peer bittorrent.Peer) map[string]any { - return map[string]any{ - "peer id": peer.ID.RawString(), - "ip": peer.Addr(), - "port": peer.Port(), - } -} diff --git a/frontend/options.go b/frontend/options.go index c77f84c..3b83a5d 100644 --- a/frontend/options.go +++ b/frontend/options.go @@ -1,5 +1,94 @@ package frontend +import ( + "errors" + "net" + "time" + + "github.com/libp2p/go-reuseport" +) + +const ( + defaultReadTimeout = 2 * time.Second + defaultWriteTimeout = 2 * time.Second +) + +var ( + errAddressNotProvided = errors.New("address not provided") + errUnexpectedListenerType = errors.New("unexpected listener type") +) + +// ListenOptions is the base configuration which may be used in net listeners +type ListenOptions struct { + Addr string + ReusePort bool `cfg:"reuse_port"` + ReadTimeout time.Duration `cfg:"read_timeout"` + WriteTimeout time.Duration `cfg:"write_timeout"` + EnableRequestTiming bool `cfg:"enable_request_timing"` +} + +func (lo ListenOptions) Validate() (validOptions ListenOptions, err error) { + validOptions = lo + if len(lo.Addr) == 0 { + err = errAddressNotProvided + } else { + if lo.ReadTimeout <= 0 { + validOptions.ReadTimeout = defaultReadTimeout + logger.Warn(). + Str("name", "http.ReadTimeout"). + Dur("provided", lo.ReadTimeout). + Dur("default", validOptions.ReadTimeout). + Msg("falling back to default configuration") + } + + if lo.WriteTimeout <= 0 { + validOptions.WriteTimeout = defaultWriteTimeout + logger.Warn(). + Str("name", "http.WriteTimeout"). + Dur("provided", lo.WriteTimeout). + Dur("default", validOptions.WriteTimeout). + Msg("falling back to default configuration") + } + } + return +} + +func (lo ListenOptions) ListenTCP() (conn *net.TCPListener, err error) { + if lo.ReusePort { + var ln net.Listener + if ln, err = reuseport.Listen("tcp", lo.Addr); err == nil { + var ok bool + if conn, ok = ln.(*net.TCPListener); !ok { + err = errUnexpectedListenerType + } + } + } else { + var addr *net.TCPAddr + if addr, err = net.ResolveTCPAddr("tcp", lo.Addr); err == nil { + conn, err = net.ListenTCP("tcp", addr) + } + } + return +} + +func (lo ListenOptions) ListenUDP() (conn *net.UDPConn, err error) { + if lo.ReusePort { + var ln net.PacketConn + if ln, err = reuseport.ListenPacket("udp", lo.Addr); err == nil { + var ok bool + if conn, ok = ln.(*net.UDPConn); !ok { + err = errUnexpectedListenerType + } + } + } else { + var addr *net.UDPAddr + if addr, err = net.ResolveUDPAddr("udp", lo.Addr); err == nil { + conn, err = net.ListenUDP("udp", addr) + } + } + return +} + // ParseOptions is the configuration used to parse an Announce Request. // // If AllowIPSpoofing is true, IPs provided via params will be used. diff --git a/frontend/udp/frontend.go b/frontend/udp/frontend.go index c7b1ec8..c0f5a40 100644 --- a/frontend/udp/frontend.go +++ b/frontend/udp/frontend.go @@ -17,8 +17,8 @@ import ( "github.com/sot-tech/mochi/bittorrent" "github.com/sot-tech/mochi/frontend" - "github.com/sot-tech/mochi/frontend/udp/bytepool" "github.com/sot-tech/mochi/middleware" + "github.com/sot-tech/mochi/pkg/bytepool" "github.com/sot-tech/mochi/pkg/conf" "github.com/sot-tech/mochi/pkg/log" "github.com/sot-tech/mochi/pkg/metrics" @@ -41,11 +41,9 @@ func init() { // Config represents all of the configurable options for a UDP BitTorrent // Tracker. type Config struct { - Addr string - ReusePort bool `cfg:"reuse_port"` - PrivateKey string `cfg:"private_key"` - MaxClockSkew time.Duration `cfg:"max_clock_skew"` - EnableRequestTiming bool `cfg:"enable_request_timing"` + frontend.ListenOptions + PrivateKey string `cfg:"private_key"` + MaxClockSkew time.Duration `cfg:"max_clock_skew"` frontend.ParseOptions } @@ -113,7 +111,7 @@ func newFrontend(c conf.MapConfig, logic *middleware.Logic) (frontend.Frontend, f.wg.Add(1) go func() { if err := f.serve(); err != nil { - logger.Fatal().Err(err).Str("proto", "udp").Msg("failed while serving") + logger.Fatal().Err(err).Msg("failed while serving") } }() @@ -162,7 +160,7 @@ func (t *udpFE) listen() (err error) { // serve blocks while listening and serving UDP BitTorrent requests // until Stop() is called or an error is returned. func (t *udpFE) serve() error { - pool := bytepool.New(2048) + pool := bytepool.NewBytePool(2048) defer t.wg.Done() for { diff --git a/frontend/udp/parser.go b/frontend/udp/parser.go index f3c0aaf..91e8571 100644 --- a/frontend/udp/parser.go +++ b/frontend/udp/parser.go @@ -1,15 +1,14 @@ package udp import ( - "bytes" "encoding/binary" "fmt" "net" "net/netip" - "sync" "github.com/sot-tech/mochi/bittorrent" "github.com/sot-tech/mochi/frontend" + "github.com/sot-tech/mochi/pkg/bytepool" ) const ( @@ -48,6 +47,8 @@ var ( errUnknownOptionType = bittorrent.ClientError("unknown option type") errInvalidInfoHash = bittorrent.ClientError("invalid info hash") errInvalidPeerID = bittorrent.ClientError("invalid info hash") + + reqRespBufferPool = bytepool.NewBufferPool() ) // ParseAnnounce parses an AnnounceRequest from a UDP request. @@ -111,23 +112,6 @@ func ParseAnnounce(r Request, v6Action bool, opts frontend.ParseOptions) (*bitto return request, err } -type buffer struct { - bytes.Buffer -} - -var bufferFree = sync.Pool{ - New: func() any { return new(buffer) }, -} - -func newBuffer() *buffer { - return bufferFree.Get().(*buffer) -} - -func (b *buffer) free() { - b.Reset() - bufferFree.Put(b) -} - // handleOptionalParameters parses the optional parameters as described in BEP // 41 and updates an announce with the values parsed. func handleOptionalParameters(packet []byte) (bittorrent.Params, error) { @@ -135,8 +119,8 @@ func handleOptionalParameters(packet []byte) (bittorrent.Params, error) { return bittorrent.ParseURLData("") } - buf := newBuffer() - defer buf.free() + buf := reqRespBufferPool.Get() + defer reqRespBufferPool.Put(buf) for i := 0; i < len(packet); { option := packet[i] diff --git a/frontend/udp/writer.go b/frontend/udp/writer.go index fe860d3..c6d8e34 100644 --- a/frontend/udp/writer.go +++ b/frontend/udp/writer.go @@ -18,8 +18,8 @@ func WriteError(w io.Writer, txID []byte, err error) { err = fmt.Errorf("internal error occurred: %w", err) } - buf := newBuffer() - defer buf.free() + buf := reqRespBufferPool.Get() + defer reqRespBufferPool.Put(buf) writeHeader(buf, txID, errorActionID) _, _ = buf.WriteString(err.Error()) _, _ = buf.WriteRune('\000') @@ -32,8 +32,8 @@ func WriteError(w io.Writer, txID []byte, err error) { // If v6Action is set, the action will be 4, according to // https://web.archive.org/web/20170503181830/http://opentracker.blog.h3q.com/2007/12/28/the-ipv6-situation/ func WriteAnnounce(w io.Writer, txID []byte, resp *bittorrent.AnnounceResponse, v6Action, v6Peers bool) { - buf := newBuffer() - defer buf.free() + buf := reqRespBufferPool.Get() + defer reqRespBufferPool.Put(buf) if v6Action { writeHeader(buf, txID, announceV6ActionID) @@ -59,8 +59,8 @@ func WriteAnnounce(w io.Writer, txID []byte, resp *bittorrent.AnnounceResponse, // WriteScrape encodes a scrape response according to BEP 15. func WriteScrape(w io.Writer, txID []byte, resp *bittorrent.ScrapeResponse) { - buf := newBuffer() - defer buf.free() + buf := reqRespBufferPool.Get() + defer reqRespBufferPool.Put(buf) writeHeader(buf, txID, scrapeActionID) @@ -75,8 +75,8 @@ func WriteScrape(w io.Writer, txID []byte, resp *bittorrent.ScrapeResponse) { // WriteConnectionID encodes a new connection response according to BEP 15. func WriteConnectionID(w io.Writer, txID, connID []byte) { - buf := newBuffer() - defer buf.free() + buf := reqRespBufferPool.Get() + defer reqRespBufferPool.Put(buf) writeHeader(buf, txID, connectActionID) _, _ = buf.Write(connID) diff --git a/go.sum b/go.sum index 1b09602..58b07d7 100644 --- a/go.sum +++ b/go.sum @@ -37,6 +37,8 @@ crawshaw.io/sqlite v0.3.2/go.mod h1:igAO5JulrQ1DbdZdtVq48mnZUBAPOeFzer7VhDWNtW4= dmitri.shuralyov.com/gpu/mtl v0.0.0-20190408044501-666a987793e9/go.mod h1:H6x//7gZCb22OMCxBHrMx7a5I7Hp++hsVxbQ4BYO7hU= github.com/BurntSushi/toml v0.3.1/go.mod h1:xHWCNGjB5oqiDr8zfno3MHue2Ht5sIBksp03qcyfWMU= github.com/BurntSushi/xgb v0.0.0-20160522181843-27f122750802/go.mod h1:IVnqGOEym/WlBOVXweHU+Q+/VP0lqqI8lqeDx9IjBqo= +github.com/IncSW/go-bencode v0.2.2 h1:RmkviUMnINqHhmBKVgrSJaHbPDw3hczN1weiX9UEoZA= +github.com/IncSW/go-bencode v0.2.2/go.mod h1:WPQp/z0JCQPy8cXJCRi/x7F7n/U0o9CIdBCpfkHGQY0= github.com/MicahParks/keyfunc v1.5.1 h1:RlyyYgKQI/adkIw1yXYtPvTAOb7hBhSX42aH23d8N0Q= github.com/MicahParks/keyfunc v1.5.1/go.mod h1:IdnCilugA0O/99dW+/MkvlyrsX8+L8+x95xuVNtM5jw= github.com/RoaringBitmap/roaring v0.4.7/go.mod h1:8khRDP4HmeXns4xIj9oGrKSz7XTQiJx2zgh7AcNke4w= diff --git a/middleware/logic.go b/middleware/logic.go index b876250..458c156 100644 --- a/middleware/logic.go +++ b/middleware/logic.go @@ -2,6 +2,7 @@ package middleware import ( "context" + "sort" "time" "github.com/sot-tech/mochi/bittorrent" @@ -88,6 +89,7 @@ func (l *Logic) HandleScrape(ctx context.Context, req *bittorrent.ScrapeRequest) return nil, nil, err } } + sort.Sort(&resp.Files) logger.Debug().Object("response", resp).Msg("generated scrape response") return ctx, resp, nil diff --git a/pkg/bytepool/bufferpool.go b/pkg/bytepool/bufferpool.go new file mode 100644 index 0000000..3fdcb72 --- /dev/null +++ b/pkg/bytepool/bufferpool.go @@ -0,0 +1,27 @@ +package bytepool + +import ( + "bytes" + "sync" +) + +// ByteBufferPool is a cached pool of reusable byte buffers. +type ByteBufferPool struct { + sync.Pool +} + +// NewBufferPool allocates a new ByteBufferPool. +func NewBufferPool() *ByteBufferPool { + return &ByteBufferPool{sync.Pool{New: func() any { return new(bytes.Buffer) }}} +} + +// Get returns a bytes.Buffer from the pool. +func (bbp *ByteBufferPool) Get() *bytes.Buffer { + return bbp.Pool.Get().(*bytes.Buffer) +} + +// Put returns bytes.Buffer to the pool. +func (bbp *ByteBufferPool) Put(b *bytes.Buffer) { + b.Reset() + bbp.Pool.Put(b) +} diff --git a/frontend/udp/bytepool/bytepool.go b/pkg/bytepool/bytepool.go similarity index 86% rename from frontend/udp/bytepool/bytepool.go rename to pkg/bytepool/bytepool.go index 8482803..bb05bfa 100644 --- a/frontend/udp/bytepool/bytepool.go +++ b/pkg/bytepool/bytepool.go @@ -8,8 +8,8 @@ type BytePool struct { sync.Pool } -// New allocates a new BytePool with slices of equal length and capacity. -func New(length int) *BytePool { +// NewBytePool allocates a new BytePool with slices of equal length and capacity. +func NewBytePool(length int) *BytePool { var bp BytePool bp.Pool.New = func() any { // This avoids allocations for the slice metadata, see: