From 72b15db93ff18884165468964d2e648ee6132c4d Mon Sep 17 00:00:00 2001 From: Leonid Bugaev Date: Sun, 5 Jul 2015 09:42:59 +0500 Subject: [PATCH] Add utilities for modifying HTTP payload --- Makefile | 2 +- byteutils/byteutils.go | 37 +++++++++++++++ byteutils/byteutils_test.go | 32 +++++++++++++ http_client.go | 34 +++----------- proto/proto.go | 81 +++++++++++++++++++++++++++++++++ proto/proto_test.go | 90 +++++++++++++++++++++++++++++++++++++ 6 files changed, 248 insertions(+), 28 deletions(-) create mode 100644 byteutils/byteutils.go create mode 100644 byteutils/byteutils_test.go create mode 100644 proto/proto.go create mode 100644 proto/proto_test.go diff --git a/Makefile b/Makefile index 12db096..9ff6068 100644 --- a/Makefile +++ b/Makefile @@ -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 diff --git a/byteutils/byteutils.go b/byteutils/byteutils.go new file mode 100644 index 0000000..33ea49e --- /dev/null +++ b/byteutils/byteutils.go @@ -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 + } +} \ No newline at end of file diff --git a/byteutils/byteutils_test.go b/byteutils/byteutils_test.go new file mode 100644 index 0000000..5b7bb04 --- /dev/null +++ b/byteutils/byteutils_test.go @@ -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") + } +} \ No newline at end of file diff --git a/http_client.go b/http_client.go index 125a8a5..4f78e9d 100644 --- a/http_client.go +++ b/http_client.go @@ -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 } diff --git a/proto/proto.go b/proto/proto.go new file mode 100644 index 0000000..8a68949 --- /dev/null +++ b/proto/proto.go @@ -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) +} diff --git a/proto/proto_test.go b/proto/proto_test.go new file mode 100644 index 0000000..bc7dcaf --- /dev/null +++ b/proto/proto_test.go @@ -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)) + } +} \ No newline at end of file