Allow url param rewriting

This commit is contained in:
Leonid Bugaev
2015-07-09 16:58:42 +05:00
parent 29035d4518
commit 48492dc8cb
7 changed files with 162 additions and 34 deletions
+30 -8
View File
@@ -1,6 +1,7 @@
package main
import (
"bytes"
"github.com/buger/gor/proto"
"hash/fnv"
)
@@ -17,6 +18,7 @@ func NewHTTPModifier(config *HTTPModifierConfig) *HTTPModifier {
len(config.headerFilters) == 0 &&
len(config.headerHashFilters) == 0 &&
len(config.paramHashFilters) == 0 &&
len(config.params) == 0 &&
len(config.headers) == 0 &&
len(config.methods) == 0 {
return nil
@@ -26,8 +28,33 @@ func NewHTTPModifier(config *HTTPModifierConfig) *HTTPModifier {
}
func (m *HTTPModifier) Rewrite(payload []byte) (response []byte) {
if len(m.config.methods) > 0 && !m.config.methods.Contains(proto.Method(payload)) {
return
if len(m.config.methods) > 0 {
method := proto.Method(payload)
matched := false
for _, m := range m.config.methods {
if bytes.Equal(method, m) {
matched = true
break
}
}
if !matched {
return
}
}
if len(m.config.headers) > 0 {
for _, header := range m.config.headers {
payload = proto.SetHeader(payload, []byte(header.Name), []byte(header.Value))
}
}
if len(m.config.params) > 0 {
for _, param := range m.config.params {
payload = proto.SetPathParam(payload, param.Name, param.Value)
}
}
if len(m.config.urlRegexp) > 0 {
@@ -38,6 +65,7 @@ func (m *HTTPModifier) Rewrite(payload []byte) (response []byte) {
for _, f := range m.config.urlRegexp {
if f.regexp.Match(path) {
matched = true
break
}
}
@@ -109,11 +137,5 @@ func (m *HTTPModifier) Rewrite(payload []byte) (response []byte) {
}
}
if len(m.config.headers) > 0 {
for _, header := range m.config.headers {
payload = proto.SetHeader(payload, []byte(header.Name), []byte(header.Value))
}
}
return payload
}
+29 -10
View File
@@ -1,7 +1,6 @@
package main
import (
"bytes"
"errors"
"fmt"
"regexp"
@@ -17,6 +16,7 @@ type HTTPModifierConfig struct {
headerHashFilters HTTPHashFilters
paramHashFilters HTTPHashFilters
params HTTPParams
headers HTTPHeaders
methods HTTPMethods
}
@@ -121,6 +121,34 @@ func (h *HTTPHeaders) Set(value string) error {
return nil
}
//
// Handling of --http-set-param option
//
type HTTPParams []HTTPParam
type HTTPParam struct {
Name []byte
Value []byte
}
func (h *HTTPParams) String() string {
return fmt.Sprint(*h)
}
func (h *HTTPParams) Set(value string) error {
v := strings.SplitN(value, "=", 2)
if len(v) != 2 {
return errors.New("Expected `Key=Value`")
}
param := HTTPParam{
[]byte(strings.TrimSpace(v[0])),
[]byte(strings.TrimSpace(v[1])),
}
*h = append(*h, param)
return nil
}
//
// Handling of --http-allow-method option
//
@@ -135,15 +163,6 @@ func (h *HTTPMethods) Set(value string) error {
return nil
}
func (h *HTTPMethods) Contains(value []byte) bool {
for _, method := range *h {
if bytes.Equal(value, method) {
return true
}
}
return false
}
//
// Handling of --http-rewrite-url option
//
-15
View File
@@ -51,21 +51,6 @@ func TestHTTPHashFilters(t *testing.T) {
}
}
func TestHTTPMethods(t *testing.T) {
methods := HTTPMethods{}
methods.Set("GET")
methods.Set("POST")
if !methods.Contains([]byte("GET")) {
t.Error("Does not contain GET")
}
if !methods.Contains([]byte("POST")) {
t.Error("Does not contain POST")
}
}
func TestUrlRewriteMap(t *testing.T) {
var err error
rewrites := UrlRewriteMap{}
+32
View File
@@ -186,3 +186,35 @@ func TestHTTPModifierURLNegativeRegexp(t *testing.T) {
t.Error("Should not pass url")
}
}
func TestHTTPModifierSetHeader(t *testing.T) {
filters := HTTPHeaders{}
filters.Set("User-Agent:Gor")
modifier := NewHTTPModifier(&HTTPModifierConfig{
headers: filters,
})
payload := []byte("POST /post HTTP/1.1\r\nContent-Length: 7\r\nHost: www.w3.org\r\n\r\na=1&b=2")
payload_after := []byte("POST /post HTTP/1.1\r\nUser-Agent: Gor\r\nContent-Length: 7\r\nHost: www.w3.org\r\n\r\na=1&b=2")
if payload = modifier.Rewrite(payload); !bytes.Equal(payload_after, payload) {
t.Error("Should add new header", string(payload))
}
}
func TestHTTPModifierSetParam(t *testing.T) {
filters := HTTPParams{}
filters.Set("api_key=1")
modifier := NewHTTPModifier(&HTTPModifierConfig{
params: filters,
})
payload := []byte("POST /post?api_key=1234 HTTP/1.1\r\nContent-Length: 7\r\nHost: www.w3.org\r\n\r\na=1&b=2")
payload_after := []byte("POST /post?api_key=1 HTTP/1.1\r\nContent-Length: 7\r\nHost: www.w3.org\r\n\r\na=1&b=2")
if payload = modifier.Rewrite(payload); !bytes.Equal(payload_after, payload) {
t.Error("Should override param", string(payload))
}
}
+32
View File
@@ -104,6 +104,38 @@ func PathParam(payload, name []byte) (value []byte, valueStart, valueEnd int) {
}
}
func SetPathParam(payload, name, value []byte) []byte {
path := Path(payload)
_, vs, ve := PathParam(payload, name)
if vs != -1 {
newPath := make([]byte, len(path))
copy(newPath, path)
newPath = byteutils.Replace(newPath, vs, ve, value)
return SetPath(payload, newPath)
} else { // if param not found append to end of url
// Adding 2 because of '?' or '&' at start, and '=' in middle
newParam := make([]byte, len(name) + len(value) + 2)
if bytes.IndexByte(path, '?') == -1 {
newParam[0] = '?'
} else {
newParam[0] = '&'
}
copy(newParam[1:], name)
newParam[1+len(name)] = '='
copy(newParam[2+len(name):], value)
newPath := make([]byte, len(path) + len(newParam))
copy(newPath, path)
copy(newPath[len(path):], newParam)
return SetPath(payload, newPath)
}
}
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
+36
View File
@@ -104,6 +104,42 @@ func TestPathParam(t *testing.T) {
}
func TestSetPathParam(t *testing.T) {
var payload, payload_after []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")
payload_after = []byte("POST /post?param=new&user_id=1 HTTP/1.1\r\nContent-Length: 7\r\nHost: www.w3.org\r\n\r\na=1&b=2")
if payload = SetPathParam(payload, []byte("param"), []byte("new")); !bytes.Equal(payload, payload_after) {
t.Error("Should replace existing value", string(payload))
}
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")
payload_after = []byte("POST /post?param=test&user_id=2 HTTP/1.1\r\nContent-Length: 7\r\nHost: www.w3.org\r\n\r\na=1&b=2")
if payload = SetPathParam(payload, []byte("user_id"), []byte("2")); !bytes.Equal(payload, payload_after) {
t.Error("Should replace existing value", string(payload))
}
payload = []byte("POST /post HTTP/1.1\r\nContent-Length: 7\r\nHost: www.w3.org\r\n\r\na=1&b=2")
payload_after = []byte("POST /post?param=test HTTP/1.1\r\nContent-Length: 7\r\nHost: www.w3.org\r\n\r\na=1&b=2")
if payload = SetPathParam(payload, []byte("param"), []byte("test")); !bytes.Equal(payload, payload_after) {
t.Error("Should set param if url have no params", string(payload))
}
payload = []byte("POST /post?param=test HTTP/1.1\r\nContent-Length: 7\r\nHost: www.w3.org\r\n\r\na=1&b=2")
payload_after = []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 payload = SetPathParam(payload, []byte("user_id"), []byte("1")); !bytes.Equal(payload, payload_after) {
t.Error("Should set param at the end if url params", string(payload))
}
}
func TestSetHostHTTP10(t *testing.T) {
var payload, payload_after []byte
+3 -1
View File
@@ -89,6 +89,8 @@ func init() {
flag.Var(&Settings.modifierConfig.headers, "http-set-header", "Inject additional headers to http reqest:\n\tgor --input-raw :8080 --output-http staging.com --http-set-header 'User-Agent: Gor'")
flag.Var(&Settings.modifierConfig.headers, "output-http-header", "WARNING: `--output-http-header` DEPRECATED, use `--http-set-header` instead")
flag.Var(&Settings.modifierConfig.params, "http-set-param", "Set request url param, if param already exists it will be overwritten:\n\tgor --input-raw :8080 --output-http staging.com --http-set-param api_key=1")
flag.Var(&Settings.modifierConfig.methods, "http-allow-method", "Whitelist of HTTP methods to replay. Anything else will be dropped:\n\tgor --input-raw :8080 --output-http staging.com --http-allow-method GET --http-allow-method OPTIONS")
flag.Var(&Settings.modifierConfig.methods, "output-http-method", "WARNING: `--output-http-method` DEPRECATED, use `--http-allow-method` instead")
@@ -97,7 +99,7 @@ func init() {
flag.Var(&Settings.modifierConfig.urlNegativeRegexp, "http-diallow-url", "A regexp to match requests against. Filter get matched agains full url with domain. Anything else will be dropped:\n\t gor --input-raw :8080 --output-http staging.com --http-disallow-url ^www.")
flag.Var(&Settings.modifierConfig.urlRewrite, "http-rewrite-url", "Rewrite the requst url based on a mapping:\n\tgor --input-raw :8080 --output-http staging.com --http-rewrite-url /v1/user/([^\\/]+)/ping:/v2/user/$1/ping")
flag.Var(&Settings.modifierConfig.urlRewrite, "http-rewrite-url", "Rewrite the request url based on a mapping:\n\tgor --input-raw :8080 --output-http staging.com --http-rewrite-url /v1/user/([^\\/]+)/ping:/v2/user/$1/ping")
flag.Var(&Settings.modifierConfig.urlRewrite, "output-http-rewrite-url", "WARNING: `--output-http-rewrite-url` DEPRECATED, use `--http-rewrite-url` instead")
flag.Var(&Settings.modifierConfig.headerFilters, "http-allow-header", "A regexp to match a specific header against. Requests with non-matching headers will be dropped:\n\t gor --input-raw :8080 --output-http staging.com --http-allow-header api-version:^v1")