Merge commit '0dbb8856e873c22f9a8ec7fcd12d8e189cd8ce7e'

Author: mindreframer

src/github.com/PuerkitoBio/throttled/varyby.go

@@ -30,6 +30,11 @@ type VaryBy struct {
 	// Use this separator string to concatenate the various criteria of the VaryBy struct.
 	// Defaults to a newline character if empty (\n).
 	Separator string
+
+	// Custom specifies the custom-generated key to use for this request.
+	// If not nil, the value returned by this function is used instead of any
+	// VaryBy criteria.
+	Custom func(r *http.Request) string
 }
 
 // Key returns the key for this request based on the criteria defined by the VaryBy struct.
@@ -39,6 +44,10 @@ func (vb *VaryBy) Key(r *http.Request) string {
 	if vb == nil {
 		return "" // Special case for no vary-by option
 	}
+	if vb.Custom != nil {
+		// A custom key generator is specified
+		return vb.Custom(r)
+	}
 	sep := vb.Separator
 	if sep == "" {
 		sep = "\n" // Separator defaults to newline

src/github.com/PuerkitoBio/throttled/varyby_test.go

@@ -39,6 +39,13 @@ func TestVaryBy(t *testing.T) {
 			&http.Request{Header: http.Header{"Cookie": []string{ck.String()}}},
 			"test\n",
 		},
+		6: {
+			&VaryBy{Cookies: []string{"ssn"}, RemoteAddr: true, Custom: func(r *http.Request) string {
+				return "blah"
+			}},
+			&http.Request{Header: http.Header{"Cookie": []string{ck.String()}}},
+			"blah",
+		},
 	}
 	for i, c := range cases {
 		got := c.vb.Key(c.r)