Merge pull request #1 from Dynom/WhiteListParams

White-listing parameters

argumentlist.go

@@ -0,0 +1,14 @@
+package main
+
+import "strings"
+
+type argumentList []string
+
+func (l argumentList) String() string {
+	return strings.Join(l, ",")
+}
+
+func (a *argumentList) Set(value string) error {
+	*a = append(*a, strings.Split(value, ",")...)
+	return nil
+}

handlers/allowedparams.go

@@ -0,0 +1,36 @@
+package handlers
+
+import (
+	"net/http"
+
+	"strings"
+
+	"github.com/go-kit/kit/log"
+)
+
+func NewAllowedParams(l log.Logger, allowedParams []string) func(h http.Handler) http.Handler {
+	var params = make(map[string]bool, len(allowedParams))
+	for _, p := range allowedParams {
+		if p == "" {
+			continue
+		}
+
+		params[p] = true
+	}
+
+	return func(h http.Handler) http.Handler {
+		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			requestParams := r.URL.Query()
+
+			for p := range requestParams {
+				if _, exists := params[p]; !exists {
+					l.Log("error", "parameter is not white-listed", "parameter", p, "allowed", strings.Join(allowedParams, ","))
+					http.Error(w, "Unregisterd parameter", http.StatusBadRequest)
+					return
+				}
+			}
+
+			h.ServeHTTP(w, r)
+		})
+	}
+}

main.go

@@ -7,8 +7,6 @@ import (
 	"os"
 	"time"
 
-	"strings"
-
 	"flag"
 
 	"fmt"
@@ -19,11 +17,12 @@ import (
 )
 
 var (
-	allowedHosts argumentList
-	imaginaryURL string
-	listenPort   int64
-	bucketRate   float64
-	bucketSize   int64
+	allowedHosts           argumentList
+	allowedImaginaryParams argumentList
+	imaginaryURL           string
+	listenPort             int64
+	bucketRate             float64
+	bucketSize             int64
 
 	Version = "dev"
 	logger  = log.With(
@@ -33,19 +32,10 @@ var (
 	)
 )
 
-type argumentList []string
-
-func (l argumentList) String() string {
-	return strings.Join(l, ",")
-}
-
-func (l *argumentList) Set(value string) error {
-	*l = append(*l, value)
-	return nil
-}
-
 func init() {
-	flag.Var(&allowedHosts, "allow-host", "Repeatable flag for hosts to allow for the URL parameter (e.g. \"d2dktr6aauwgqs.cloudfront.net\")")
+	flag.Var(&allowedHosts, "allow-hosts", "Repeatable flag (or a comma-separated list) for hosts to allow for the URL parameter (e.g. \"d2dktr6aauwgqs.cloudfront.net\")")
+	flag.Var(&allowedImaginaryParams, "allowed-params", "A comma seperated list of parameters allows to be sent upstream. If empty, everything is allowed.")
+
 	flag.StringVar(&imaginaryURL, "imaginary-url", "http://localhost:9000", "URL to imaginary (default: http://localhost:9000)")
 	flag.Int64Var(&listenPort, "listen-port", 8080, "Port to listen on")
 	flag.Float64Var(&bucketRate, "bucket-rate", 20, "Rate limiter bucket fill rate (req/s)")
@@ -60,6 +50,7 @@ func main() {
 		"msg", "Starting.",
 		"version", Version,
 		"allowed_hosts", allowedHosts.String(),
+		"allowed_params", allowedImaginaryParams.String(),
 		"imaginary_backend", imaginaryURL,
 	)
 
@@ -97,10 +88,25 @@ type httpHandler func(h http.Handler) http.Handler
 
 func decorateHandler(h http.Handler, b *ratelimit.Bucket) http.Handler {
 	decorators := []httpHandler{
-		handlers.NewRateLimitHandler(b, logger),
-		handlers.NewIgnoreFaviconRequests(),
 		handlers.NewValidateURLParameter(logger, allowedHosts),
 	}
+
+	if len(allowedImaginaryParams) > 0 {
+		decorators = append(
+			decorators,
+			handlers.NewAllowedParams(
+				logger,
+				allowedImaginaryParams,
+			))
+	}
+
+	// Defining early needed handlers last
+	decorators = append(
+		decorators,
+		handlers.NewIgnoreFaviconRequests(),
+		handlers.NewRateLimitHandler(b, logger),
+	)
+
 	var handler http.Handler = h
 	for _, d := range decorators {
 		handler = d(handler)