diff --git a/proto/proto.go b/proto/proto.go index 0e74cc8..4e15a5a 100644 --- a/proto/proto.go +++ b/proto/proto.go @@ -123,7 +123,7 @@ func headerIndex(payload []byte, name []byte) 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) { +func header(payload []byte, name []byte) (value []byte, headerStart, headerEnd, valueStart, valueEnd int) { headerStart = headerIndex(payload, name) if headerStart == -1 { @@ -131,24 +131,37 @@ 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++ - } - headerEnd = valueStart + bytes.IndexByte(payload[valueStart:], '\n') - if payload[headerEnd-1] == '\r' { - headerEnd -= 1 + for valueStart < headerEnd { // Ignore empty space after ':' + if payload[valueStart] == ' ' { + valueStart++ + } else { + break + } } - value = payload[valueStart:headerEnd] + valueEnd = valueStart + bytes.IndexByte(payload[valueStart:], '\n') + + if payload[headerEnd-1] == '\r' { + valueEnd -= 1 + } + + for valueStart < valueEnd { // ignore empty space at end of header value + if payload[valueEnd-1] == ' ' { + valueEnd -= 1 + } else { + break + } + } + value = payload[valueStart:valueEnd] return } // Header returns header value, if header not found, value will be blank func Header(payload, name []byte) []byte { - val, _, _, _ := header(payload, name) + val, _, _, _, _ := header(payload, name) return val } @@ -156,11 +169,11 @@ func Header(payload, name []byte) []byte { // 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, ve := header(payload, name) if hs != -1 { - // If header found we just repace its value - return byteutils.Replace(payload, vs, he, value) + // If header found we just replace its value + return byteutils.Replace(payload, vs, ve, value) } return AddHeader(payload, name, value) @@ -180,6 +193,19 @@ func AddHeader(payload, name, value []byte) []byte { return byteutils.Insert(payload, mimeStart, header) } +// DelHeader takes http payload and removes header name from headers section +// Returns modified request payload +func DelHeader(payload, name[]byte) []byte { + _, hs, he, _, _ := header(payload, name) + if hs != -1 { + newHeader := make([]byte, len(payload) - (he - hs) - 1) + copy(newHeader[:hs], payload[:hs]) + copy(newHeader[hs:], payload[he + 1:]) + return newHeader + } + return payload +} + // Body returns request/response body func Body(payload []byte) []byte { // 4 -> len(EMPTY_LINE) diff --git a/proto/proto_test.go b/proto/proto_test.go index a8d9179..2510237 100644 --- a/proto/proto_test.go +++ b/proto/proto_test.go @@ -16,6 +16,13 @@ func TestHeader(t *testing.T) { t.Error("Should find header value") } + // Value with space at end + 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 7") + } + // Value without space at start payload = []byte("POST /post HTTP/1.1\r\nContent-Length:7\r\nHost: www.w3.org\r\n\r\na=1&b=2") @@ -38,7 +45,7 @@ func TestHeader(t *testing.T) { } // Header not found - if _, headerStart, _, _ = header(payload, []byte("Not-Found")); headerStart != -1 { + if _, headerStart, _, _, _ = header(payload, []byte("Not-Found")); headerStart != -1 { t.Error("Should not found header") } @@ -97,6 +104,25 @@ func TestSetHeader(t *testing.T) { } } +func TestDelHeader(t *testing.T) { + var payload, payloadAfter []byte + + payload = []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\nContent-Length: 7\r\nHost: www.w3.org\r\n\r\na=1&b=2") + + if payload = DelHeader(payload, []byte("User-Agent")); !bytes.Equal(payload, payloadAfter) { + t.Error("Should delete header if found", string(payload), string(payloadAfter)) + } + + //Whitespace at end of User-Agent + payload = []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\nContent-Length: 7\r\nHost: www.w3.org\r\n\r\na=1&b=2") + + if payload = DelHeader(payload, []byte("User-Agent")); !bytes.Equal(payload, payloadAfter) { + t.Error("Should delete header if found", string(payload), string(payloadAfter)) + } +} + func TestPath(t *testing.T) { var path, payload []byte