Added a server side ratelimiter for client uploads. (Velocidex#345)

Ensures total server traffic remains below a specified rate. Helpful for slow links.

config/proto/config.proto

@@ -406,6 +406,10 @@ message FrontendConfig {
     // The services that will run on this frontend. If not set, all
     // services will run on the primary frontend.
     ServerServicesConfig server_services = 20;
+
+    // Per client rate in bytes/sec
+    int64 per_client_upload_rate = 21;
+    int64 global_upload_rate = 22;
 }
 
 

go.mod

@@ -58,6 +58,7 @@ require (
 	github.com/jehiah/go-strftime v0.0.0-20171201141054-1d33003b3869 // indirect
 	github.com/jmoiron/sqlx v1.2.0
 	github.com/jstemmer/go-junit-report v0.9.1 // indirect
+	github.com/juju/ratelimit v1.0.1
 	github.com/kierdavis/dateparser v0.0.0-20171227112021-81e70b820720
 	github.com/konsorten/go-windows-terminal-sequences v1.0.2 // indirect
 	github.com/kr/pretty v0.2.0 // indirect
@@ -108,7 +109,7 @@ require (
 	gopkg.in/sourcemap.v1 v1.0.5 // indirect
 	gopkg.in/yaml.v2 v2.2.8 // indirect
 	howett.net/plist v0.0.0-20181124034731-591f970eefbb
-	www.velocidex.com/golang/evtx v0.0.2-0.20200504101835-a9fb49a9a1a0
+	www.velocidex.com/golang/evtx v0.0.2-0.20200505002428-02e98e3e472b
 	www.velocidex.com/golang/go-ese v0.0.0-20200111070159-4b7484475321
 	www.velocidex.com/golang/go-ntfs v0.0.0-20200110083657-950cbe916617
 	www.velocidex.com/golang/go-pe v0.1.1-0.20191103232346-ac12e8190bb6

go.sum

@@ -253,6 +253,8 @@ github.com/jstemmer/go-junit-report v0.0.0-20190106144839-af01ea7f8024 h1:rBMNdl
 github.com/jstemmer/go-junit-report v0.0.0-20190106144839-af01ea7f8024/go.mod h1:6v2b51hI/fHJwM22ozAgKL4VKDeJcHhJFhtBdhmNjmU=
 github.com/jstemmer/go-junit-report v0.9.1 h1:6QPYqodiu3GuPL+7mfx+NwDdp2eTkp9IfEUpgAwUN0o=
 github.com/jstemmer/go-junit-report v0.9.1/go.mod h1:Brl9GWCQeLvo8nXZwPNNblvFj/XSXhF0NWZEnDohbsk=
+github.com/juju/ratelimit v1.0.1 h1:+7AIFJVQ0EQgq/K9+0Krm7m530Du7tIz0METWzN0RgY=
+github.com/juju/ratelimit v1.0.1/go.mod h1:qapgC/Gy+xNh9UxzV13HGGl/6UXNN+ct+vwSgWNm/qk=
 github.com/julienschmidt/httprouter v1.2.0/go.mod h1:SYymIcj16QtmaHHD7aYtjjsJG7VTCxuUUipMqKk8s4w=
 github.com/kballard/go-shellquote v0.0.0-20180428030007-95032a82bc51 h1:Z9n2FFNUXsshfwJMBgNA0RU6/i7WVaAegv3PtuIHPMs=
 github.com/kballard/go-shellquote v0.0.0-20180428030007-95032a82bc51/go.mod h1:CzGEWj7cYgsdH8dAjBGEr58BoE7ScuLd+fwFZ44+/x8=
@@ -600,8 +602,8 @@ rsc.io/quote/v3 v3.1.0/go.mod h1:yEA65RcK8LyAZtP9Kv3t0HmxON59tX3rD+tICJqUlj0=
 rsc.io/sampler v1.3.0/go.mod h1:T1hPZKmBbMNahiBKFy5HrXp6adAjACjK9JXDnKaTXpA=
 www.velocidex.com/golang/binparsergen v0.1.0 h1:oNsMHGnlb4jrGwxKxqqmsics6FgYin3HR5UNtLXc8S0=
 www.velocidex.com/golang/binparsergen v0.1.0/go.mod h1:UC43Ecj0mjsidlClTYZ3H4dXdyv7CVI0HsYi4yY3qtc=
-www.velocidex.com/golang/evtx v0.0.2-0.20200504101835-a9fb49a9a1a0 h1:5dbv0fNN1wvWviSNHxs/uscVPv7eMkC+ydC+281Cq1Y=
-www.velocidex.com/golang/evtx v0.0.2-0.20200504101835-a9fb49a9a1a0/go.mod h1:+u26IeGeVIwL9j5V0I/UafWFaMV61pQNwXZK/VQksLQ=
+www.velocidex.com/golang/evtx v0.0.2-0.20200505002428-02e98e3e472b h1:WCN8TCDUdt6oimdAkmRL0JxW6776jGUxF1sxFaoiYvk=
+www.velocidex.com/golang/evtx v0.0.2-0.20200505002428-02e98e3e472b/go.mod h1:+u26IeGeVIwL9j5V0I/UafWFaMV61pQNwXZK/VQksLQ=
 www.velocidex.com/golang/go-ese v0.0.0-20200111070159-4b7484475321 h1:0FELGb4G9dT2y6rrBL2B5nbEw2foiPXtEO7jHppQGe0=
 www.velocidex.com/golang/go-ese v0.0.0-20200111070159-4b7484475321/go.mod h1:d3PHzQhyhe+AO9RYBnDKZ40As15T+38zr++Dnv4ufuc=
 www.velocidex.com/golang/go-ntfs v0.0.0-20200110083657-950cbe916617 h1:nETlXizhx7Mkjr6lIl3bge7rp9ml2ORP50tnsF7jBCQ=

server/comms.go

@@ -36,6 +36,7 @@ import (
 	"www.velocidex.com/golang/velociraptor/utils"
 
 	"github.com/golang/protobuf/proto"
+	"github.com/juju/ratelimit"
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/client_golang/prometheus/promauto"
 	"golang.org/x/crypto/acme/autocert"
@@ -269,7 +270,6 @@ func maybeRedirectFrontend(handler string, w http.ResponseWriter, r *http.Reques
 	redirect_url, ok := frontend.GetFrontendURL()
 	if ok {
 		// We should redirect to another frontend.
-		fmt.Printf("Redirecting to %v\n", redirect_url)
 		http.Redirect(w, r, redirect_url, 301)
 		return true
 	}
@@ -307,9 +307,22 @@ func control(server_obj *Server) http.Handler {
 			defer server_obj.EndConcurrencyControl()
 		}
 
-		body, err := ioutil.ReadAll(
-			io.LimitReader(req.Body, int64(server_obj.config.
-				Frontend.MaxUploadSize*2)))
+		reader := io.LimitReader(req.Body, int64(server_obj.config.
+			Frontend.MaxUploadSize*2))
+
+		if server_obj.config.Frontend.PerClientUploadRate > 0 {
+			bucket := ratelimit.NewBucketWithRate(
+				float64(server_obj.config.Frontend.PerClientUploadRate),
+				100*1024)
+			reader = ratelimit.Reader(reader, bucket)
+		}
+
+		if server_obj.Bucket != nil {
+			reader = ratelimit.Reader(reader, server_obj.Bucket)
+		}
+
+		body, err := ioutil.ReadAll(reader)
+
 		if err != nil {
 			logger.Debug("Unable to read body from %v: %+v (read %v)",
 				req.RemoteAddr, err, len(body))

server/server.go

@@ -22,6 +22,7 @@ import (
 	"context"
 	"time"
 
+	"github.com/juju/ratelimit"
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/client_golang/prometheus/promauto"
 	actions_proto "www.velocidex.com/golang/velociraptor/actions/proto"
@@ -49,7 +50,9 @@ type Server struct {
 	db      datastore.DataStore
 
 	// Limit concurrency for processing messages.
-	concurrency      chan bool
+	concurrency chan bool
+
+	Bucket           *ratelimit.Bucket
 	APIClientFactory grpc_client.APIClientFactory
 }
 
@@ -98,6 +101,15 @@ func NewServer(
 		concurrency:      make(chan bool, concurrency),
 		APIClientFactory: grpc_client.GRPCAPIClient{},
 	}
+
+	if config_obj.Frontend.GlobalUploadRate > 0 {
+		result.logger.Info("Global upload rate set to %v bytes per second",
+			config_obj.Frontend.GlobalUploadRate)
+		result.Bucket = ratelimit.NewBucketWithRate(
+			float64(config_obj.Frontend.GlobalUploadRate),
+			1024*1024)
+	}
+
 	return &result, nil
 }