Add download rate limiting

[btrtrc.git] / client.go

diff --git a/client.go b/client.go
index 51aaf193238a6504e4b87ffeadf7fd9618a087fb..d28f9b9db784e080eaeafb068081182f3d836108 100644 (file)
--- a/client.go
+++ b/client.go
@@ -74,8 +74,9 @@ type Client struct {
        // Our BitTorrent protocol extension bytes, sent in our BT handshakes.
        extensionBytes peerExtensionBytes
        // The net.Addr.String part that should be common to all active listeners.
-       listenAddr  string
-       uploadLimit *rate.Limiter
+       listenAddr    string
+       uploadLimit   *rate.Limiter
+       downloadLimit *rate.Limiter
 
        // Set of addresses that have our client ID. This intentionally will
        // include ourselves if we end up trying to connect to our own address
@@ -263,6 +264,11 @@ func NewClient(cfg *Config) (cl *Client, err error) {
        } else {
                cl.uploadLimit = cfg.UploadRateLimiter
        }
+       if cfg.DownloadRateLimiter == nil {
+               cl.downloadLimit = rate.NewLimiter(rate.Inf, 0)
+       } else {
+               cl.downloadLimit = cfg.DownloadRateLimiter
+       }
        missinggo.CopyExact(&cl.extensionBytes, defaultExtensionBytes)
        cl.event.L = &cl.mu
        storageImpl := cfg.DefaultStorage
@@ -1583,5 +1589,6 @@ func (cl *Client) newConnection(nc net.Conn) (c *connection) {
                PeerMaxRequests: 250,
        }
        c.setRW(connStatsReadWriter{nc, &cl.mu, c})
+       c.r = rateLimitedReader{cl.downloadLimit, c.r}
        return
 }