From 6647481fe07061bc90be6db55611dd34179d7c6c Mon Sep 17 00:00:00 2001 From: Leonid Bugaev Date: Mon, 13 Jul 2015 09:31:11 +0500 Subject: [PATCH] Improve style and docs for proto package --- Makefile | 2 +- elasticsearch.go | 30 +++---- gor.go | 2 +- http_client.go | 4 +- http_modifier.go | 8 +- input_raw_test.go | 2 +- proto/proto.go | 134 +++++++++++++++++++----------- proto/proto_test.go | 46 +++++----- raw_socket_listener/listener.go | 40 ++++----- raw_socket_listener/tcp_packet.go | 2 +- 10 files changed, 154 insertions(+), 116 deletions(-) diff --git a/Makefile b/Makefile index 2f8b95b..cb4bb17 100644 --- a/Makefile +++ b/Makefile @@ -14,7 +14,7 @@ dbuild: docker build -t gor . dlint: - docker run -v `pwd`:$(SOURCE_PATH) -t -i --env GORACE="halt_on_error=1" gor golint $(ARGS) + docker run -v `pwd`:$(SOURCE_PATH) -t -i --env GORACE="halt_on_error=1" gor golint $(PKG) drace: docker run -v `pwd`:$(SOURCE_PATH) -t -i --env GORACE="halt_on_error=1" gor go test ./... $(ARGS) -v -race -timeout 15s diff --git a/elasticsearch.go b/elasticsearch.go index 801e75c..a930ac3 100644 --- a/elasticsearch.go +++ b/elasticsearch.go @@ -122,24 +122,24 @@ func (p *ESPlugin) ResponseAnalyze(req, resp []byte, start, stop time.Time) { esResp := ESRequestResponse{ ReqUrl: proto.Path(req), ReqMethod: proto.Method(req), - ReqUserAgent: proto.GetHeader(req, "User-Agent"), - ReqAcceptLanguage: proto.GetHeader(req, "Accept-Language"), - ReqAccept: proto.GetHeader(req, "Accept"), - ReqAcceptEncoding: proto.GetHeader(req, "Accept-Encoding"), - ReqIfModifiedSince: proto.GetHeader(req, "If-Modified-Since"), - ReqConnection: proto.GetHeader(req, "Connection"), - ReqCookies: proto.GetHeader(req, "Cookie"), + ReqUserAgent: proto.Header(req, []byte("User-Agent")), + ReqAcceptLanguage: proto.Header(req, []byte("Accept-Language")), + ReqAccept: proto.Header(req, []byte("Accept")), + ReqAcceptEncoding: proto.Header(req, []byte("Accept-Encoding")), + ReqIfModifiedSince: proto.Header(req, []byte("If-Modified-Since")), + ReqConnection: proto.Header(req, []byte("Connection")), + ReqCookies: proto.Header(req, []byte("Cookie")), RespStatus: proto.Status(resp), RespStatusCode: proto.Status(resp), RespProto: proto.Method(resp), - RespContentLength: proto.GetHeader(resp, "Content-Length"), - RespContentType: proto.GetHeader(resp, "Content-Type"), - RespTransferEncoding: proto.GetHeader(resp, "Transfer-Encoding"), - RespContentEncoding: proto.GetHeader(resp, "Content-Encoding"), - RespExpires: proto.GetHeader(resp, "Expires"), - RespCacheControl: proto.GetHeader(resp, "Cache-Control"), - RespVary: proto.GetHeader(resp, "Vary"), - RespSetCookie: proto.GetHeader(resp, "Set-Cookie"), + RespContentLength: proto.Header(resp, []byte("Content-Length")), + RespContentType: proto.Header(resp, []byte("Content-Type")), + RespTransferEncoding: proto.Header(resp, []byte("Transfer-Encoding")), + RespContentEncoding: proto.Header(resp, []byte("Content-Encoding")), + RespExpires: proto.Header(resp, []byte("Expires")), + RespCacheControl: proto.Header(resp, []byte("Cache-Control")), + RespVary: proto.Header(resp, []byte("Vary")), + RespSetCookie: proto.Header(resp, []byte("Set-Cookie")), Rtt: rtt, Timestamp: t, } diff --git a/gor.go b/gor.go index c3c30bb..af9bf5f 100644 --- a/gor.go +++ b/gor.go @@ -7,9 +7,9 @@ import ( "fmt" "log" "os" + "runtime" "runtime/debug" "runtime/pprof" - "runtime" "time" ) diff --git a/http_client.go b/http_client.go index dd6a739..c54e6a8 100644 --- a/http_client.go +++ b/http_client.go @@ -147,7 +147,7 @@ func (c *HTTPClient) Send(data []byte) (response []byte, err error) { if status[0] == '3' { c.redirectsCount += 1 - location, _, _, _ := proto.Header(payload, []byte("Location")) + location := proto.Header(payload, []byte("Location")) redirectPayload := []byte("GET " + string(location) + " HTTP/1.1\r\n\r\n") if c.config.Debug { @@ -167,4 +167,4 @@ func (c *HTTPClient) Get(path string) (response []byte, err error) { payload := "GET " + path + " HTTP/1.1\r\n\r\n" return c.Send([]byte(payload)) -} \ No newline at end of file +} diff --git a/http_modifier.go b/http_modifier.go index d820df6..96be68b 100644 --- a/http_modifier.go +++ b/http_modifier.go @@ -86,9 +86,9 @@ func (m *HTTPModifier) Rewrite(payload []byte) (response []byte) { if len(m.config.headerFilters) > 0 { for _, f := range m.config.headerFilters { - value, s, _, _ := proto.Header(payload, f.name) + value := proto.Header(payload, f.name) - if s != -1 && !f.regexp.Match(value) { + if len(value) > 0 && !f.regexp.Match(value) { return } } @@ -96,9 +96,9 @@ func (m *HTTPModifier) Rewrite(payload []byte) (response []byte) { if len(m.config.headerHashFilters) > 0 { for _, f := range m.config.headerHashFilters { - value, s, _, _ := proto.Header(payload, f.name) + value := proto.Header(payload, f.name) - if s != -1 { + if len(value) > 0 { hasher := fnv.New32a() hasher.Write(value) diff --git a/input_raw_test.go b/input_raw_test.go index 2f8652b..6262667 100644 --- a/input_raw_test.go +++ b/input_raw_test.go @@ -31,7 +31,7 @@ func TestRAWInput(t *testing.T) { address := strings.Replace(listener.Addr().String(), "[::]", "127.0.0.1", -1) - client := NewHTTPClient(address, &HTTPClientConfig{}) + client := NewHTTPClient(address, &HTTPClientConfig{}) go Start(quit) diff --git a/proto/proto.go b/proto/proto.go index ee4a1a0..f595234 100644 --- a/proto/proto.go +++ b/proto/proto.go @@ -1,28 +1,50 @@ -// Low-level interaction with HTTP request payload +/* +Package proto provides byte-level interaction with HTTP request payload. + +Example of HTTP payload for future references, new line symbols escaped: + + POST /upload HTTP/1.1\r\n + User-Agent: Gor\r\n + Content-Length: 11\r\n + \r\n + Hello world + + GET /index.html HTTP/1.1\r\n + User-Agent: Gor\r\n + \r\n + \r\n +*/ package proto import ( "bytes" "github.com/buger/gor/byteutils" - _ "log" ) +// In HTTP newline defined by 2 bytes (for both windows and *nix support) var CLRF = []byte("\r\n") -var EMPTY_LINE = []byte("\r\n\r\n") -var HEADER_DELIM = []byte(": ") -// Headers should end with empty line +// New line acts as separator: end of Headers or Body (in some cases) +var EmptyLine = []byte("\r\n\r\n") + +// Separator for Header line. Header looks like: `HeaderName: value` +var HeaderDelim = []byte(": ") + +// MIMEHeadersEndPos finds end of the Headers section, which should end with empty line. func MIMEHeadersEndPos(payload []byte) int { - return bytes.Index(payload, EMPTY_LINE) + return bytes.Index(payload, EmptyLine) } +// MIMEHeadersStartPos finds start of Headers section +// It just finds position of second line (first contains location and method). 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) { +// header return value and positions of header/value start/end. +// If not found, value will be blank, and headerStart will be -1 +// 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 { @@ -31,7 +53,7 @@ func Header(payload []byte, name []byte) (value []byte, headerStart, valueStart, valueStart = headerStart + len(name) + 1 // Skip ":" after header name if payload[valueStart] == ' ' { // Ignore empty space after ':' - valueStart += 1 + valueStart++ } headerEnd = valueStart + bytes.IndexByte(payload[valueStart:], '\r') value = payload[valueStart:headerEnd] @@ -39,27 +61,32 @@ func Header(payload []byte, name []byte) (value []byte, headerStart, valueStart, return } -func GetHeader(payload []byte, name string) []byte { - val, _, _, _ := Header(payload, []byte(name)) +// Header returns header value, if header not found, value will be blank +func Header(payload, name []byte) []byte { + val, _, _, _ := header(payload, name) return val } +// SetHeader sets header value. If header not found it creates new one. +// Returns modified request payload func SetHeader(payload, name, value []byte) []byte { - _, hs, vs, he := Header(payload, name) + _, hs, vs, he := header(payload, name) - // If header found if hs != -1 { + // If header found we just repace its value return byteutils.Replace(payload, vs, he, value) - } else { - return AddHeader(payload, name, value) } + + return AddHeader(payload, name, value) } +// AddHeader takes http payload and appends new header to the start of headers section +// Returns modified request payload 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):], HeaderDelim) copy(header[len(name)+2:], value) copy(header[len(header)-2:], CLRF) @@ -68,30 +95,31 @@ func AddHeader(payload, name, value []byte) []byte { return byteutils.Insert(payload, mimeStart, header) } +// Path takes payload and retuns request path: Split(firstLine, ' ')[1] func Path(payload []byte) []byte { - start := bytes.IndexByte(payload, ' ') - start += 1 - + start := bytes.IndexByte(payload, ' ') + 1 end := bytes.IndexByte(payload[start:], ' ') return payload[start : start+end] } +// SetPath takes payload, sets new path and returns modified payload func SetPath(payload, path []byte) []byte { - start := bytes.IndexByte(payload, ' ') - start += 1 - + start := bytes.IndexByte(payload, ' ') + 1 end := bytes.IndexByte(payload[start:], ' ') return byteutils.Replace(payload, start, start+end, path) } +// PathParam returns URL query attribute by given name, if no found: valueStart will be -1 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:], '&') + + // Param can end with '&' (another param), or end of line if paramEnd == -1 { // It is final param paramEnd = len(path) } else { @@ -99,43 +127,51 @@ func PathParam(payload, name []byte) (value []byte, valueStart, valueEnd int) { } return path[valueStart:paramEnd], valueStart, paramEnd - } else { - return []byte(""), -1, -1 } + + return []byte(""), -1, -1 } +// SetPathParam takes payload and updates path Query attribute +// If query param not found, it will append new +// Returns modified payload func SetPathParam(payload, name, value []byte) []byte { path := Path(payload) _, vs, ve := PathParam(payload, name) - if vs != -1 { + if vs != -1 { // If param found, replace its value and set new Path 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) } + + // 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 "param=value" into buffer, after it looks like "?param=value" + copy(newParam[1:], name) + newParam[1+len(name)] = '=' + copy(newParam[2+len(name):], value) + + // Append param to the end of path + newPath := make([]byte, len(path)+len(newParam)) + copy(newPath, path) + copy(newPath[len(path):], newParam) + + return SetPath(payload, newPath) } +// SetHost updates Host header for HTTP/1.1 or updates host in path for HTTP/1.0 or Proxy requests +// Returns modified payload 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 @@ -149,18 +185,20 @@ func SetHost(payload, url, host []byte) []byte { newPath = byteutils.Replace(newPath, 0, hostEnd, url) return SetPath(payload, newPath) - } else { - return SetHeader(payload, []byte("Host"), host) } + + return SetHeader(payload, []byte("Host"), host) } +// Method returns HTTP method func Method(payload []byte) []byte { end := bytes.IndexByte(payload, ' ') return payload[:end] } -// Status in response have same position as Path in request +// Status returns response status. +// It happend to be in same position as request payload path func Status(payload []byte) []byte { return Path(payload) } diff --git a/proto/proto_test.go b/proto/proto_test.go index f232ee3..80adbe2 100644 --- a/proto/proto_test.go +++ b/proto/proto_test.go @@ -11,17 +11,17 @@ func TestHeader(t *testing.T) { 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")) { + 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")) { + 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 { + if _, headerStart, _, _ = header(payload, []byte("Not-Found")); headerStart != -1 { t.Error("Should not found header") } } @@ -50,19 +50,19 @@ func TestMIMEHeadersStartPos(t *testing.T) { } func TestSetHeader(t *testing.T) { - var payload, payload_after []byte + var payload, payloadAfter []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") + payloadAfter = []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) { + if payload = SetHeader(payload, []byte("Content-Length"), []byte("14")); !bytes.Equal(payload, payloadAfter) { 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") + payloadAfter = []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) { + if payload = SetHeader(payload, []byte("User-Agent"), []byte("Gor")); !bytes.Equal(payload, payloadAfter) { t.Error("Should add header if not found", string(payload)) } } @@ -78,12 +78,12 @@ func TestPath(t *testing.T) { } func TestSetPath(t *testing.T) { - var payload, payload_after []byte + var payload, payloadAfter []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") + payloadAfter = []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) { + if payload = SetPath(payload, []byte("/new_path")); !bytes.Equal(payload, payloadAfter) { t.Error("Should replace path", string(payload)) } } @@ -103,44 +103,44 @@ func TestPathParam(t *testing.T) { } func TestSetPathParam(t *testing.T) { - var payload, payload_after []byte + var payload, payloadAfter []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") + payloadAfter = []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) { + if payload = SetPathParam(payload, []byte("param"), []byte("new")); !bytes.Equal(payload, payloadAfter) { 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") + payloadAfter = []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) { + if payload = SetPathParam(payload, []byte("user_id"), []byte("2")); !bytes.Equal(payload, payloadAfter) { 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") + payloadAfter = []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) { + if payload = SetPathParam(payload, []byte("param"), []byte("test")); !bytes.Equal(payload, payloadAfter) { 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") + payloadAfter = []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) { + if payload = SetPathParam(payload, []byte("user_id"), []byte("1")); !bytes.Equal(payload, payloadAfter) { t.Error("Should set param at the end if url params", string(payload)) } } func TestSetHostHTTP10(t *testing.T) { - var payload, payload_after []byte + var payload, payloadAfter []byte payload = []byte("POST http://example.com/post HTTP/1.0\r\nContent-Length: 7\r\nHost: www.w3.org\r\n\r\na=1&b=2") - payload_after = []byte("POST http://new.com/post HTTP/1.0\r\nContent-Length: 7\r\nHost: www.w3.org\r\n\r\na=1&b=2") + payloadAfter = []byte("POST http://new.com/post HTTP/1.0\r\nContent-Length: 7\r\nHost: www.w3.org\r\n\r\na=1&b=2") - if payload = SetHost(payload, []byte("http://new.com"), []byte("new.com")); !bytes.Equal(payload, payload_after) { + if payload = SetHost(payload, []byte("http://new.com"), []byte("new.com")); !bytes.Equal(payload, payloadAfter) { t.Error("Should replace host", string(payload)) } } diff --git a/raw_socket_listener/listener.go b/raw_socket_listener/listener.go index dce1dc0..9fd0c65 100644 --- a/raw_socket_listener/listener.go +++ b/raw_socket_listener/listener.go @@ -69,34 +69,34 @@ func (t *Listener) listen() { } } func (t *Listener) readRAWSocket() { - conn, e := net.ListenPacket("ip4:tcp", t.addr) + conn, e := net.ListenPacket("ip4:tcp", t.addr) - if e != nil { - log.Fatal(e) - } + if e != nil { + log.Fatal(e) + } - defer conn.Close() + defer conn.Close() - for { - buf := make([]byte, 64*1024) // 64kb - // Note: ReadFrom receive messages without IP header - n, addr, err := conn.ReadFrom(buf) + for { + buf := make([]byte, 64*1024) // 64kb + // Note: ReadFrom receive messages without IP header + n, addr, err := conn.ReadFrom(buf) - if err != nil { - log.Println("Error:", err) - continue - } + if err != nil { + log.Println("Error:", err) + continue + } - if n > 0 { - go t.parsePacket(addr, buf[:n]) - } - } + if n > 0 { + go t.parsePacket(addr, buf[:n]) + } + } } func (t *Listener) parsePacket(addr net.Addr, buf []byte) { - if t.isIncomingDataPacket(buf) { - t.c_packets <- ParseTCPPacket(addr, buf) - } + if t.isIncomingDataPacket(buf) { + t.c_packets <- ParseTCPPacket(addr, buf) + } } func (t *Listener) isIncomingDataPacket(buf []byte) bool { diff --git a/raw_socket_listener/tcp_packet.go b/raw_socket_listener/tcp_packet.go index b07109b..4d3f4e6 100644 --- a/raw_socket_listener/tcp_packet.go +++ b/raw_socket_listener/tcp_packet.go @@ -59,7 +59,7 @@ func (t *TCPPacket) Parse() { // ParseBasic set of fields func (t *TCPPacket) ParseBasic() { - t.SrcPort = binary.BigEndian.Uint16(t.Data[0:2]) + t.SrcPort = binary.BigEndian.Uint16(t.Data[0:2]) t.Seq = binary.BigEndian.Uint32(t.Data[4:8]) t.Ack = binary.BigEndian.Uint32(t.Data[8:12]) t.DataOffset = (t.Data[12] & 0xF0) >> 4