From 6a8a72ac18c733180475fac010a0a2ba9690dfcc Mon Sep 17 00:00:00 2001
From: tedli <lizhen@outlook.jp>
Date: Sat, 14 Dec 2019 16:45:35 +0800
Subject: [PATCH] add extra checking of header buffer, to support multi line
 header value (#123) (#688)

---
 header.go      | 86 ++++++++++++++++++++++++++++++++++++++++++++++++--
 header_test.go | 38 ++++++++++++++++++++++
 2 files changed, 122 insertions(+), 2 deletions(-)

diff --git a/header.go b/header.go
index c0fd5f61f2..a04f4f90f5 100644
--- a/header.go
+++ b/header.go
@@ -2018,9 +2018,24 @@ type headerScanner struct {
 	hLen int
 
 	disableNormalizing bool
+
+	// by checking whether the next line contains a colon or not to tell
+	// it's a header entry or a multi line value of current header entry.
+	// the side effect of this operation is that we know the index of the
+	// next colon and new line, so this can be used during next iteration,
+	// instead of find them again.
+	nextColon   int
+	nextNewLine int
+
+	initialized bool
 }
 
 func (s *headerScanner) next() bool {
+	if !s.initialized {
+		s.nextColon = -1
+		s.nextNewLine = -1
+		s.initialized = true
+	}
 	bLen := len(s.b)
 	if bLen >= 2 && s.b[0] == '\r' && s.b[1] == '\n' {
 		s.b = s.b[2:]
@@ -2032,7 +2047,13 @@ func (s *headerScanner) next() bool {
 		s.hLen++
 		return false
 	}
-	n := bytes.IndexByte(s.b, ':')
+	var n int
+	if s.nextColon >= 0 {
+		n = s.nextColon
+		s.nextColon = -1
+	} else {
+		n = bytes.IndexByte(s.b, ':')
+	}
 	if n < 0 {
 		s.err = errNeedMore
 		return false
@@ -2042,14 +2063,48 @@ func (s *headerScanner) next() bool {
 	n++
 	for len(s.b) > n && s.b[n] == ' ' {
 		n++
+		// the newline index is a relative index, and lines below trimed `s.b` by `n`,
+		// so the relative newline index also shifted forward. it's safe to decrease
+		// to a minus value, it means it's invalid, and will find the newline again.
+		s.nextNewLine--
 	}
 	s.hLen += n
 	s.b = s.b[n:]
-	n = bytes.IndexByte(s.b, '\n')
+	if s.nextNewLine >= 0 {
+		n = s.nextNewLine
+		s.nextNewLine = -1
+	} else {
+		n = bytes.IndexByte(s.b, '\n')
+	}
 	if n < 0 {
 		s.err = errNeedMore
 		return false
 	}
+	isMultiLineValue := false
+	for {
+		if n+1 >= len(s.b) {
+			break
+		}
+		d := bytes.IndexByte(s.b[n+1:], '\n')
+		if d <= 0 {
+			break
+		} else if d == 1 && s.b[n+1] == '\r' {
+			break
+		}
+		e := n + d + 1
+		if c := bytes.IndexByte(s.b[n+1:e], ':'); c >= 0 {
+			s.nextColon = c
+			s.nextNewLine = d - c - 1
+			break
+		}
+		isMultiLineValue = true
+		n = e
+	}
+	if n >= len(s.b) {
+		s.err = errNeedMore
+		return false
+	}
+	oldB := s.b
 	s.value = s.b[:n]
 	s.hLen += n + 1
 	s.b = s.b[n+1:]
@@ -2061,6 +2116,9 @@ func (s *headerScanner) next() bool {
 		n--
 	}
 	s.value = s.value[:n]
+	if isMultiLineValue {
+		s.value, s.b, s.hLen = normalizeHeaderValue(s.value, oldB, s.hLen)
+	}
 	return true
 }
 
@@ -2129,6 +2187,30 @@ func getHeaderKeyBytes(kv *argsKV, key string, disableNormalizing bool) []byte {
 	return kv.key
 }
 
+func normalizeHeaderValue(ov, ob []byte, headerLength int) (nv, nb []byte, nhl int) {
+	nv = ov
+	length := len(ov)
+	if length <= 0 {
+		return
+	}
+	write := 0
+	shrunk := 0
+	for read := 0; read < length; read++ {
+		c := ov[read]
+		if c == '\r' || c == '\n' {
+			shrunk++
+			continue
+		}
+		nv[write] = c
+		write++
+	}
+	nv = nv[:write]
+	copy(ob[write:], ob[write+shrunk:])
+	nb = ob[write+2 : len(ob)-shrunk]
+	nhl = headerLength - shrunk
+	return
+}
+
 func normalizeHeaderKey(b []byte, disableNormalizing bool) {
 	if disableNormalizing {
 		return
diff --git a/header_test.go b/header_test.go
index 3debd976a0..6f3aa0fe1a 100644
--- a/header_test.go
+++ b/header_test.go
@@ -6,11 +6,49 @@ import (
 	"fmt"
 	"io"
 	"io/ioutil"
+	"net/http"
 	"reflect"
 	"strings"
 	"testing"
 )
 
+func TestResponseHeaderMultiLineValue(t *testing.T) {
+	s := "HTTP/1.1 200 OK\r\n" +
+		"EmptyValue1:\r\n" +
+		"Content-Type: foo/bar;\r\n\tnewline;\r\n another/newline\r\n" + // the '\t' will be kept, won't be removed
+		"Foo: Bar\r\n" +
+		"Multi-Line: one;\r\n two\r\n" +
+		"Values: v1;\r\n v2;\r\n v3; v4\r\n" +
+		"\r\n"
+	expectContentType := "foo/bar;\tnewline; another/newline"
+	// net/http not only remove "\r\n" but also replace \t to space
+	expectNetHttpContentType := "foo/bar; newline; another/newline"
+	expectMultiLine := "one; two"
+	header := new(ResponseHeader)
+	_, err := header.parse([]byte(s))
+	if err != nil {
+		t.Fatalf("parse headers with multi-line values failed, %s", err)
+	}
+	gotContentType := header.Peek("Content-Type")
+	if string(gotContentType) != expectContentType {
+		t.Fatalf("unexpected content-type: %q. Expecting %q", gotContentType, expectContentType)
+	}
+	gotMultiLine := header.Peek("Multi-Line")
+	if string(gotMultiLine) != expectMultiLine {
+		t.Fatalf("unexpected multi-line: %q. Expecting %q", gotMultiLine, expectMultiLine)
+	}
+	// ensure behave same as net/http
+	response, err := http.ReadResponse(bufio.NewReader(strings.NewReader(s)), nil)
+	if err != nil {
+		t.Fatalf("parse response using net/http failed, %s", err)
+	}
+	gotNetHttpContentType := response.Header.Get("Content-Type")
+	if gotNetHttpContentType != expectNetHttpContentType {
+		t.Fatalf("unexpected content-type (net/http): %q. Expecting %q",
+			gotNetHttpContentType, expectNetHttpContentType)
+	}
+}
+
 func TestResponseHeaderEmptyValueFromHeader(t *testing.T) {
 	t.Parallel()