Add proto.DelHeader

Update proto.go to handle headers with whitespace after

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)

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