mirror of
https://github.com/buger/goreplay.git
synced 2024-04-21 12:32:02 +00:00
Add utilities for modifying HTTP payload
This commit is contained in:
Leonid Bugaev
@@ -14,7 +14,7 @@ dbuild:
|
||||
docker build -t gor .
|
||||
|
||||
dtest:
|
||||
docker run -v `pwd`:$(SOURCE_PATH) -t -i --env GORACE="halt_on_error=1" gor go test $(ARGS) -race -v
|
||||
docker run -v `pwd`:$(SOURCE_PATH) -t -i --env GORACE="halt_on_error=1" gor go test ./... $(ARGS) -race -v
|
||||
|
||||
dfmt:
|
||||
docker run -v `pwd`:$(SOURCE_PATH) -t -i gor go fmt
|
||||
|
||||
@@ -0,0 +1,37 @@
|
||||
package byteutils
|
||||
|
||||
func Cut(a []byte, from, to int) []byte {
|
||||
copy(a[from:], a[to:])
|
||||
a = a[:len(a)-to+from]
|
||||
|
||||
return a
|
||||
}
|
||||
|
||||
func Insert(a []byte, i int, b []byte) []byte {
|
||||
a = append(a, make([]byte, len(b))...)
|
||||
copy(a[i+len(b):], a[i:])
|
||||
copy(a[i:i+len(b)], b)
|
||||
|
||||
return a
|
||||
}
|
||||
|
||||
// Unlike bytes.Replace it allows you to specify range
|
||||
func Replace(a []byte, from, to int, new []byte) []byte {
|
||||
lenDiff := len(new) - (to - from)
|
||||
|
||||
if lenDiff > 0 {
|
||||
// Extend if new segment bigger
|
||||
a = append(a, make([]byte, lenDiff)...)
|
||||
copy(a[to+lenDiff:], a[to:])
|
||||
copy(a[from:from+len(new)], new)
|
||||
|
||||
return a
|
||||
} else if lenDiff < 0 {
|
||||
copy(a[from:], new)
|
||||
copy(a[from+len(new):],a[to:])
|
||||
return a[:len(a) + lenDiff]
|
||||
} else { // same size
|
||||
copy(a[from:], new)
|
||||
return a
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,32 @@
|
||||
package byteutils
|
||||
|
||||
import (
|
||||
"testing"
|
||||
"bytes"
|
||||
)
|
||||
|
||||
func TestCut(t *testing.T) {
|
||||
if !bytes.Equal(Cut([]byte("123456"), 2, 4), []byte("1256")) {
|
||||
t.Error("Should properly cut")
|
||||
}
|
||||
}
|
||||
|
||||
func TestInsert(t *testing.T) {
|
||||
if !bytes.Equal(Insert([]byte("123456"), 2, []byte("abcd")), []byte("12abcd3456")) {
|
||||
t.Error("Should insert into middle of slice")
|
||||
}
|
||||
}
|
||||
|
||||
func TestReplace(t *testing.T) {
|
||||
if !bytes.Equal(Replace([]byte("123456"), 2, 4, []byte("ab")), []byte("12ab56")) {
|
||||
t.Error("Should replace when same length")
|
||||
}
|
||||
|
||||
if !bytes.Equal(Replace([]byte("123456"), 2, 4, []byte("abcd")), []byte("12abcd56")) {
|
||||
t.Error("Should replace when replacement length bigger")
|
||||
}
|
||||
|
||||
if !bytes.Equal(Replace([]byte("123456"), 2, 5, []byte("ab")), []byte("12ab6")) {
|
||||
t.Error("Should replace when replacement length bigger")
|
||||
}
|
||||
}
|
||||
+7
-27
@@ -7,9 +7,7 @@ import (
|
||||
"net/url"
|
||||
"strings"
|
||||
"time"
|
||||
"bytes"
|
||||
"bufio"
|
||||
"errors"
|
||||
"github.com/buger/gor/proto"
|
||||
)
|
||||
|
||||
var defaultPorts = map[string]string{
|
||||
@@ -85,26 +83,6 @@ func (c *HTTPClient) isAlive() bool {
|
||||
return true
|
||||
}
|
||||
|
||||
func header(payload []byte, name []byte) ([]byte, error) {
|
||||
buf := bytes.NewBuffer(payload)
|
||||
reader := bufio.NewReader(buf)
|
||||
|
||||
// Skip status line
|
||||
reader.ReadLine()
|
||||
|
||||
for {
|
||||
line, _, err := reader.ReadLine()
|
||||
|
||||
if err != nil {
|
||||
return nil, errors.New("Header not found")
|
||||
}
|
||||
|
||||
if bytes.HasPrefix(line, name) {
|
||||
return bytes.Split(line, []byte(": "))[1], nil
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (c *HTTPClient) Send(data []byte) (response []byte, err error) {
|
||||
if c.conn == nil || !c.isAlive() {
|
||||
Debug("Connecting:", c.baseURL)
|
||||
@@ -132,18 +110,20 @@ func (c *HTTPClient) Send(data []byte) (response []byte, err error) {
|
||||
return
|
||||
}
|
||||
|
||||
payload := c.respBuf[:n]
|
||||
|
||||
if c.config.Debug {
|
||||
Debug("Received:", string(c.respBuf[:n]))
|
||||
Debug("Received:", string(payload))
|
||||
}
|
||||
|
||||
if c.config.FollowRedirects > 0 && c.redirectsCount < c.config.FollowRedirects {
|
||||
status := c.respBuf[9:12]
|
||||
status := payload[9:12]
|
||||
|
||||
// 3xx requests
|
||||
if status[0] == '3' {
|
||||
c.redirectsCount += 1
|
||||
|
||||
location, _ := header(c.respBuf[:n], []byte("Location:"))
|
||||
location, _, _, _ := proto.Header(payload, []byte("Location"))
|
||||
redirectPayload := []byte("GET " + string(location) + " HTTP/1.1\r\n\r\n")
|
||||
|
||||
if c.config.Debug {
|
||||
@@ -156,5 +136,5 @@ func (c *HTTPClient) Send(data []byte) (response []byte, err error) {
|
||||
|
||||
c.redirectsCount = 0
|
||||
|
||||
return c.respBuf[:n], err
|
||||
return payload, err
|
||||
}
|
||||
|
||||
@@ -0,0 +1,81 @@
|
||||
// Low-level interaction with HTTP request payload
|
||||
package proto
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"github.com/buger/gor/byteutils"
|
||||
_ "log"
|
||||
)
|
||||
|
||||
var CLRF = []byte("\r\n")
|
||||
var EMPTY_LINE = []byte("\r\n\r\n")
|
||||
var HEADER_DELIM = []byte(": ")
|
||||
|
||||
// Headers should end with empty line
|
||||
func MIMEHeadersEndPos(payload []byte) int {
|
||||
return bytes.Index(payload, EMPTY_LINE)
|
||||
}
|
||||
|
||||
func MIMEHeadersStartPos(payload []byte) int {
|
||||
return bytes.Index(payload, CLRF) + 2 // Find first line end
|
||||
}
|
||||
|
||||
// Find header value or return error
|
||||
// Do not support multi-line headers
|
||||
func Header(payload []byte, name []byte) (value []byte, headerStart, valueStart, headerEnd int) {
|
||||
headerStart = bytes.Index(payload, name)
|
||||
|
||||
if headerStart == -1 {
|
||||
return
|
||||
}
|
||||
|
||||
valueStart = headerStart + len(name) + 1 // Skip ":" after header name
|
||||
if payload[valueStart] == ' ' { // Ignore empty space after ':'
|
||||
valueStart += 1
|
||||
}
|
||||
headerEnd = valueStart + bytes.IndexByte(payload[valueStart:], '\r')
|
||||
value = payload[valueStart:headerEnd]
|
||||
|
||||
return
|
||||
}
|
||||
|
||||
func SetHeader(payload, name, value []byte) []byte {
|
||||
_, hs, vs, he := Header(payload, name)
|
||||
|
||||
// If header found
|
||||
if hs != -1 {
|
||||
return byteutils.Replace(payload, vs, he, value)
|
||||
} else {
|
||||
return AddHeader(payload, name, value)
|
||||
}
|
||||
}
|
||||
|
||||
func AddHeader(payload, name, value []byte) []byte {
|
||||
header := make([]byte, len(name) + 2 + len(value) + 2)
|
||||
copy(header[0:], name)
|
||||
copy(header[len(name):], HEADER_DELIM)
|
||||
copy(header[len(name)+2:], value)
|
||||
copy(header[len(header)-2:], CLRF)
|
||||
|
||||
mimeStart := MIMEHeadersStartPos(payload)
|
||||
|
||||
return byteutils.Insert(payload, mimeStart, header)
|
||||
}
|
||||
|
||||
func Path(payload []byte) []byte {
|
||||
start := bytes.IndexByte(payload, ' ')
|
||||
start += 1
|
||||
|
||||
end := bytes.IndexByte(payload[start:], ' ')
|
||||
|
||||
return payload[start:start+end]
|
||||
}
|
||||
|
||||
func SetPath(payload, path []byte) []byte {
|
||||
start := bytes.IndexByte(payload, ' ')
|
||||
start += 1
|
||||
|
||||
end := bytes.IndexByte(payload[start:], ' ')
|
||||
|
||||
return byteutils.Replace(payload, start, start+end, path)
|
||||
}
|
||||
@@ -0,0 +1,90 @@
|
||||
package proto
|
||||
|
||||
import (
|
||||
"testing"
|
||||
"bytes"
|
||||
)
|
||||
|
||||
func TestHeader(t *testing.T) {
|
||||
var payload, val []byte
|
||||
var headerStart int
|
||||
|
||||
payload = []byte("POST /post HTTP/1.1\r\nContent-Length: 7\r\nHost: www.w3.org\r\n\r\na=1&b=2")
|
||||
|
||||
if val, _, _, _ = Header(payload, []byte("Content-Length")); !bytes.Equal(val, []byte("7")) {
|
||||
t.Error("Should find header value")
|
||||
}
|
||||
|
||||
payload = []byte("POST /post HTTP/1.1\r\nContent-Length:7\r\nHost: www.w3.org\r\n\r\na=1&b=2")
|
||||
|
||||
if val, _, _, _ = Header(payload, []byte("Content-Length")); !bytes.Equal(val, []byte("7")) {
|
||||
t.Error("Should find header value without space after :")
|
||||
}
|
||||
|
||||
if _, headerStart, _, _ = Header(payload, []byte("Not-Found")); headerStart != -1 {
|
||||
t.Error("Should not found header")
|
||||
}
|
||||
}
|
||||
|
||||
func TestMIMEHeadersEndPos(t *testing.T) {
|
||||
head := []byte("POST /post HTTP/1.1\r\nContent-Length: 7\r\nHost: www.w3.org")
|
||||
payload := []byte("POST /post HTTP/1.1\r\nContent-Length: 7\r\nHost: www.w3.org\r\n\r\na=1&b=2")
|
||||
|
||||
end := MIMEHeadersEndPos(payload)
|
||||
|
||||
if !bytes.Equal(payload[:end], head) {
|
||||
t.Error("Wrong headers end position:", end)
|
||||
}
|
||||
}
|
||||
|
||||
func TestMIMEHeadersStartPos(t *testing.T) {
|
||||
headers := []byte("Content-Length: 7\r\nHost: www.w3.org")
|
||||
payload := []byte("POST /post HTTP/1.1\r\nContent-Length: 7\r\nHost: www.w3.org\r\n\r\na=1&b=2")
|
||||
|
||||
start := MIMEHeadersStartPos(payload)
|
||||
end := MIMEHeadersEndPos(payload)
|
||||
|
||||
if !bytes.Equal(payload[start:end], headers) {
|
||||
t.Error("Wrong headers end position:", start, end)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSetHeader(t *testing.T) {
|
||||
var payload, payload_after []byte
|
||||
|
||||
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\nContent-Length: 14\r\nHost: www.w3.org\r\n\r\na=1&b=2")
|
||||
|
||||
if payload = SetHeader(payload, []byte("Content-Length"), []byte("14")); !bytes.Equal(payload, payload_after) {
|
||||
t.Error("Should update header if it exists", 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 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 = SetHeader(payload, []byte("User-Agent"), []byte("Gor")); !bytes.Equal(payload, payload_after) {
|
||||
t.Error("Should add header if not found", string(payload))
|
||||
}
|
||||
}
|
||||
|
||||
func TestPath(t *testing.T) {
|
||||
var path, payload []byte
|
||||
|
||||
payload = []byte("POST /post HTTP/1.1\r\nContent-Length: 7\r\nHost: www.w3.org\r\n\r\na=1&b=2")
|
||||
|
||||
if path = Path(payload); !bytes.Equal(path, []byte("/post")) {
|
||||
t.Error("Should find path", string(path))
|
||||
}
|
||||
}
|
||||
|
||||
func TestSetPath(t *testing.T) {
|
||||
var payload, payload_after []byte
|
||||
|
||||
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 /new_path HTTP/1.1\r\nContent-Length: 7\r\nHost: www.w3.org\r\n\r\na=1&b=2")
|
||||
|
||||
if payload = SetPath(payload, []byte("/new_path")); !bytes.Equal(payload, payload_after) {
|
||||
t.Error("Should replace path", string(payload))
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user