diff --git a/http_modifier.go b/http_modifier.go index b6f994c..3aa0476 100644 --- a/http_modifier.go +++ b/http_modifier.go @@ -15,6 +15,7 @@ func NewHTTPModifier(config *HTTPModifierConfig) *HTTPModifier { len(config.urlRewrite) == 0 && len(config.headerFilters) == 0 && len(config.headerHashFilters) == 0 && + len(config.paramHashFilters) == 0 && len(config.headers) == 0 && len(config.methods) == 0 { return nil @@ -60,15 +61,28 @@ func (m *HTTPModifier) Rewrite(payload []byte) (response []byte) { for _, f := range m.config.headerHashFilters { value, s, _, _ := proto.Header(payload, f.name) - if s == -1 { - return + if s != -1 { + hasher := fnv.New32a() + hasher.Write(value) + + if (hasher.Sum32() % 100) >= f.percent { + return + } } + } + } - hasher := fnv.New32a() - hasher.Write(value) + if len(m.config.paramHashFilters) > 0 { + for _, f := range m.config.paramHashFilters { + value, s, _ := proto.PathParam(payload, f.name) - if (hasher.Sum32() % 100) >= f.percent { - return + if s != -1 { + hasher := fnv.New32a() + hasher.Write(value) + + if (hasher.Sum32() % 100) >= f.percent { + return + } } } } diff --git a/http_modifier_settings.go b/http_modifier_settings.go index 5ea8601..391854a 100644 --- a/http_modifier_settings.go +++ b/http_modifier_settings.go @@ -14,7 +14,8 @@ type HTTPModifierConfig struct { urlRegexp HTTPUrlRegexp urlRewrite UrlRewriteMap headerFilters HTTPHeaderFilters - headerHashFilters HTTPHeaderHashFilters + headerHashFilters HTTPHashFilters + paramHashFilters HTTPHashFilters headers HTTPHeaders methods HTTPMethods @@ -50,20 +51,20 @@ func (h *HTTPHeaderFilters) Set(value string) error { } // -// Handling of --http-allow-header-hash options +// Handling of --http-allow-header-hash and --http-allow-param-hash options // -type headerHashFilter struct { +type hashFilter struct { name []byte percent uint32 } -type HTTPHeaderHashFilters []headerHashFilter +type HTTPHashFilters []hashFilter -func (h *HTTPHeaderHashFilters) String() string { +func (h *HTTPHashFilters) String() string { return fmt.Sprint(*h) } -func (h *HTTPHeaderHashFilters) Set(value string) error { +func (h *HTTPHashFilters) Set(value string) error { valArr := strings.SplitN(value, ":", 2) if len(valArr) < 2 { return errors.New("need both header and value, colon-delimited (ex. user_id:1/2).") @@ -78,7 +79,7 @@ func (h *HTTPHeaderHashFilters) Set(value string) error { num, _ = strconv.ParseUint(fracArr[0], 10, 64) den, _ = strconv.ParseUint(fracArr[1], 10, 64) - var f headerHashFilter + var f hashFilter f.name = []byte(valArr[0]) f.percent = uint32((float64(num) / float64(den)) * 100) *h = append(*h, f) diff --git a/http_modifier_settings_test.go b/http_modifier_settings_test.go index badf30e..0b1701d 100644 --- a/http_modifier_settings_test.go +++ b/http_modifier_settings_test.go @@ -24,8 +24,8 @@ func TestHTTPHeaderFilters(t *testing.T) { } } -func TestHTTPHeaderHashFilters(t *testing.T) { - filters := HTTPHeaderHashFilters{} +func TestHTTPHashFilters(t *testing.T) { + filters := HTTPHashFilters{} err := filters.Set("Header1:1/2") if err != nil { diff --git a/http_modifier_test.go b/http_modifier_test.go index 4689aca..80a3bb3 100644 --- a/http_modifier_test.go +++ b/http_modifier_test.go @@ -69,7 +69,7 @@ func TestHTTPModifierURLRewrite(t *testing.T) { } func TestHTTPModifierHeaderHashFilters(t *testing.T) { - filters := HTTPHeaderHashFilters{} + filters := HTTPHashFilters{} filters.Set("Header2:1/2") modifier := NewHTTPModifier(&HTTPModifierConfig{ @@ -80,8 +80,8 @@ func TestHTTPModifierHeaderHashFilters(t *testing.T) { return []byte("POST / HTTP/1.1\r\n" + string(header) + "Content-Length: 7\r\nHost: www.w3.org\r\n\r\na=1&b=2") } - if p := modifier.Rewrite(payload([]byte(""))); len(p) > 0 { - t.Error("Request should not pass filters, Header2 does not exist") + if p := modifier.Rewrite(payload([]byte(""))); len(p) == 0 { + t.Error("Request should pass filters if Header does not exist") } if p := modifier.Rewrite(payload([]byte("Header2: 3\r\n"))); len(p) > 0 { @@ -93,6 +93,32 @@ func TestHTTPModifierHeaderHashFilters(t *testing.T) { } } + +func TestHTTPModifierParamHashFilters(t *testing.T) { + filters := HTTPHashFilters{} + filters.Set("user_id:1/2") + + modifier := NewHTTPModifier(&HTTPModifierConfig{ + paramHashFilters: filters, + }) + + payload := func(value []byte) []byte { + return []byte("POST /" + string(value) + " HTTP/1.1\r\nContent-Length: 7\r\nHost: www.w3.org\r\n\r\na=1&b=2") + } + + if p := modifier.Rewrite(payload([]byte(""))); len(p) == 0 { + t.Error("Request should pass filters if param does not exist") + } + + if p := modifier.Rewrite(payload([]byte("?user_id=3"))); len(p) > 0 { + t.Error("Request should not pass filters", string(p)) + } + + if p := modifier.Rewrite(payload([]byte("?user_id=1"))); len(p) == 0 { + t.Error("Request should pass filters") + } +} + func TestHTTPModifierHeaders(t *testing.T) { headers := HTTPHeaders{} headers.Set("Header1:1") diff --git a/proto/proto.go b/proto/proto.go index 330115b..9b0e5cf 100644 --- a/proto/proto.go +++ b/proto/proto.go @@ -86,6 +86,24 @@ func SetPath(payload, path []byte) []byte { return byteutils.Replace(payload, start, start+end, path) } +func PathParam(payload, name []byte) (value []byte, valueStart, valueEnd int) { + path := Path(payload) + + if paramStart := bytes.Index(path, append(name, '=')); paramStart != -1 { + valueStart := paramStart + len(name) + 1 + paramEnd := bytes.IndexByte(path[valueStart:], '&') + if paramEnd == -1 { // It is final param + paramEnd = len(path) + } else { + paramEnd += valueStart + } + + return path[valueStart:paramEnd], valueStart, paramEnd + } else { + return []byte(""), -1, -1 + } +} + func SetHost(payload, url, host []byte) []byte { // If this is HTTP 1.0 traffic or proxy traffic it may include host right into path variable, so instead of setting Host header we rewrite Path // Fix for https://github.com/buger/gor/issues/156 diff --git a/proto/proto_test.go b/proto/proto_test.go index 3ada427..12227e2 100644 --- a/proto/proto_test.go +++ b/proto/proto_test.go @@ -89,6 +89,21 @@ func TestSetPath(t *testing.T) { } } +func TestPathParam(t *testing.T) { + var payload []byte + + payload = []byte("POST /post?param=test&user_id=1 HTTP/1.1\r\nContent-Length: 7\r\nHost: www.w3.org\r\n\r\na=1&b=2") + + if val, _, _ := PathParam(payload, []byte("param")); !bytes.Equal(val, []byte("test")) { + t.Error("Should detect attribute", string(val)) + } + + if val, _, _ := PathParam(payload, []byte("user_id")); !bytes.Equal(val, []byte("1")) { + t.Error("Should detect attribute", string(val)) + } +} + + func TestSetHostHTTP10(t *testing.T) { var payload, payload_after []byte diff --git a/settings.go b/settings.go index 58b8400..aa9cb6d 100644 --- a/settings.go +++ b/settings.go @@ -110,6 +110,8 @@ func init() { flag.Var(&Settings.modifierConfig.headerHashFilters, "http-allow-header-hash", "Takes a fraction of requests, consistently taking or rejecting a request based on the FNV32-1A hash of a specific header:\n\t gor --input-raw :8080 --output-http staging.com --http-allow-header-hash user-id:1/4") flag.Var(&Settings.modifierConfig.headerHashFilters, "output-http-header-hash-filter", "WARNING: `output-http-header-hash-filter` DEPRECATED, use `--http-allow-header-hash` instead") + + flag.Var(&Settings.modifierConfig.paramHashFilters, "http-allow-param-hash", "Takes a fraction of requests, consistently taking or rejecting a request based on the FNV32-1A hash of a specific GET param:\n\t gor --input-raw :8080 --output-http staging.com --http-allow-param-hash user_id:1/4") } func Debug(args ...interface{}) {