omg almost done with httpx port to super smol http

httpx/header_parse.go

@@ -6,6 +6,7 @@ import (
 	"errors"
 	"fmt"
 	"io"
+	"slices"
 	"strings"
 	"unsafe"
 )
@@ -17,49 +18,37 @@ var (
 	errNonNumericChars = errors.New("non-numeric chars found")
 )
 
-type headerScanner struct {
-	b     []byte
-	key   []byte
-	value []byte
-	err   error
-
-	// hLen stores header subslice len
-	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 (hb *headerBuf) readFromBytes(b []byte) {
+	hb.buf = append(hb.buf, b...)
 }
 
-type headerValueScanner struct {
-	b     string
-	value string
+func (hb *headerBuf) free() int { return cap(hb.buf) - len(hb.buf) }
+
+func (hb *headerBuf) readFrom(r io.Reader) error {
+	buf := hb.buf
+	free := hb.free()
+	if free == 0 {
+		return errSmallBuffer
+	}
+	n, err := r.Read(buf[len(buf):cap(buf)])
+	hb.buf = buf[:len(buf)+n]
+	return err
 }
 
-func (h *header) parse(buf []byte) (int, error) {
-	m, err := h.parseFirstLine(buf)
+func (h *header) parse() (err error) {
+	hb := &h.hbuf
+	hb.off = 0 // start parsing from 0.
+	h.method, h.requestURI, h.proto, h.flags, err = hb.parseFirstLine(h.flags)
 	if err != nil {
-		return 0, err
+		return err
 	}
 
-	h.rawHeaders, _, err = readRawHeaders(h.rawHeaders[:0], b2s(buf[m:]))
+	var ss scannerState
+	err = h.parseHeaders(&ss)
 	if err != nil {
-		return 0, err
+		return err
 	}
-	var n int
-	n, err = h.parseHeaders(buf[m:])
-	if err != nil {
-		return 0, err
-	}
-	return m + n, nil
+	return nil
 }
 
 func (hb *headerBuf) offBuf() []byte {
@@ -87,13 +76,14 @@ func (hb *headerBuf) scanUntilByte(c byte) []byte {
 	return buf
 }
 
-func (hb *headerBuf) parseFirstLine() (method, uri, proto headerSlice, flags flags, err error) {
+func (hb *headerBuf) parseFirstLine(initFlags flags) (method, uri, proto headerSlice, flags flags, err error) {
 	var b []byte
 	for len(b) == 0 {
 		b = hb.scanLine()
 	}
+	flags = initFlags
 	if len(b) < 5 {
-		return method, uri, proto, 0, errors.New("too short first HTTP line")
+		return method, uri, proto, flags, errors.New("too short first HTTP line")
 	}
 
 	methodEnd := max(0, bytes.IndexByte(b, ' '))
@@ -115,6 +105,143 @@ func (hb *headerBuf) parseFirstLine() (method, uri, proto headerSlice, flags fla
 	return method, uri, proto, flags, nil
 }
 
+type scannerState struct {
+	err error
+
+	// hLen stores header subslice len
+	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 (h *header) parseHeaders(ss *scannerState) (err error) {
+	hb := &h.hbuf
+	h.contentLength = -2
+
+	for kv := hb.nextKV(ss); kv.isValid(); kv = hb.nextKV(ss) {
+		if h.flags.hasAny(disableSpecialHeader) {
+			h.hbuf.headers = append(h.hbuf.headers, kv)
+			continue
+		}
+	}
+	if ss.err != nil && err == nil {
+		err = ss.err
+	}
+	if err != nil {
+		h.flags |= connectionClose
+		return err
+	}
+
+	// if h.contentLength < 0 {
+	// 	h.contentLengthBytes = hb.noKV().value
+	// }
+	if h.flags.hasAny(noHTTP11) && !h.flags.hasAny(connectionClose) {
+		// close connection for non-http/1.1 request unless 'Connection: keep-alive' is set.
+		if !h.hasHeaderValue(strConnection, strKeepAlive) {
+			h.flags |= connectionClose
+		}
+	}
+	return nil
+}
+
+func (h *header) hasHeaderValue(key, value string) bool {
+	kv := h.peekHeader(key)
+	return kv.isValid() && b2s(h.hbuf.musttoken(kv.value)) == value
+}
+
+func (h *header) peekHeaderBytes(key string) []byte {
+	kv := h.peekHeader(key)
+	if kv.isValid() {
+		return h.hbuf.musttoken(kv.value)
+	}
+	return nil
+}
+
+// peekHeader returns header key-value for the given key.
+//
+// The returned value is valid until the request is released,
+// either though ReleaseRequest or your request handler returning.
+// Do not store references to returned value. Make copies instead.
+func (h *header) peekHeader(key string) argsKV {
+	hb := &h.hbuf
+	for i := 0; i < len(h.hbuf.headers); i++ {
+		if b2s(hb.musttoken(h.hbuf.headers[i].key)) == key {
+			return h.hbuf.headers[i]
+		}
+	}
+	return hb.noKV()
+}
+
+func (h *header) peekPtrHeader(key string) *argsKV {
+	hb := &h.hbuf
+	for i := 0; i < len(h.hbuf.headers); i++ {
+		if b2s(hb.musttoken(h.hbuf.headers[i].key)) == key {
+			return &h.hbuf.headers[i]
+		}
+	}
+	return nil
+}
+
+func (hb *headerBuf) mustAppendSlice(value string) headerSlice {
+	L := len(hb.buf)
+	copy(hb.buf[L:L+len(value)], value)
+	hb.buf = hb.buf[:L+len(value)]
+	return hb.slice(hb.buf[L : L+len(value)])
+}
+
+func (h *header) reuseOrAppend(tok headerSlice, value string) headerSlice {
+	if tok.len > tokint(len(value)) {
+		copy(h.hbuf.musttoken(tok), value)
+		tok.len = tokint(len(value))
+		return tok
+	}
+	return h.appendSlice(value)
+}
+
+func (h *header) appendSlice(value string) headerSlice {
+	free := h.hbuf.free()
+	if len(value) > free {
+		if h.flags.hasAny(flagNoBufferGrow) {
+			h.flags |= flagOOMReached
+			return headerSlice{}
+		}
+		h.hbuf.buf = slices.Grow(h.hbuf.buf, len(value))
+	}
+	return h.hbuf.mustAppendSlice(value)
+}
+
+func (h *header) appendHeader(key, value string) {
+	hb := &h.hbuf
+	free := hb.free()
+	buf := h.hbuf.buf
+
+	if len(key)+len(value) > free {
+		if h.flags.hasAny(flagNoBufferGrow) {
+			panic(errSmallBuffer)
+		}
+		slices.Grow(buf, len(key)+len(value))
+	}
+	k := hb.mustAppendSlice(key)
+	v := hb.mustAppendSlice(value)
+	if !h.flags.hasAny(disableNormalizing) {
+		// TODO
+	}
+	hb.headers = append(hb.headers, argsKV{
+		key:   k,
+		value: v,
+	})
+}
+
 func readRawHeaders(dst []byte, buf string) ([]byte, int, error) {
 	n := strings.IndexByte(buf, nChar)
 	if n < 0 {
@@ -142,217 +269,125 @@ func readRawHeaders(dst []byte, buf string) ([]byte, int, error) {
 		}
 	}
 }
+func (hb *headerBuf) noKV() argsKV { return argsKV{} }
 
-func (h *header) parseHeaders(buf []byte) (int, error) {
-	h.contentLength = -2
-	h.scanner = headerScanner{}
-	s := &h.scanner
-	s.b = buf
-	s.disableNormalizing = h.disableNormalizing
-	var err error
-	for s.next() {
-		key := b2s(s.key)
-		value := b2s(s.value)
-		if len(key) > 0 {
-			// Spaces between the header key and colon are not allowed.
-			// See RFC 7230, Section 3.2.4.
-
-			if strings.IndexByte(key, ' ') != -1 || strings.IndexByte(key, '\t') != -1 {
-				err = fmt.Errorf("invalid header key %q", s.key)
-				continue
-			}
-
-			if h.disableSpecialHeader {
-				h.h = appendArg(h.h, key, value, argsHasValue)
-				continue
-			}
-
-			switch s.key[0] | 0x20 {
-			case 'h':
-				if caseInsensitiveCompare(key, strHost) {
-					h.host = append(h.host[:0], value...)
-					continue
-				}
-			case 'u':
-				if caseInsensitiveCompare(key, strUserAgent) {
-					h.userAgent = append(h.userAgent[:0], value...)
-					continue
-				}
-			case 'c':
-				if caseInsensitiveCompare(key, strContentType) {
-					h.contentType = append(h.contentType[:0], value...)
-					continue
-				}
-				if caseInsensitiveCompare(key, strContentLength) {
-					if h.contentLength != -1 {
-						var nerr error
-						if h.contentLength, nerr = parseContentLength(b2s(s.value)); nerr != nil {
-							if err == nil {
-								err = nerr
-							}
-							h.contentLength = -2
-						} else {
-							h.contentLengthBytes = append(h.contentLengthBytes[:0], value...)
-						}
-					}
-					continue
-				}
-				if caseInsensitiveCompare(key, strConnection) {
-					if b2s(s.value) == strClose {
-						h.connectionClose = true
-					} else {
-						h.connectionClose = false
-						h.h = appendArg(h.h, key, value, argsHasValue)
-					}
-					continue
-				}
-			case 't':
-				if caseInsensitiveCompare(key, strTransferEncoding) {
-					if value != strIdentity {
-						h.contentLength = -1
-						h.h = setArg(h.h, strTransferEncoding, strChunked, argsHasValue)
-					}
-					continue
-				}
-				if caseInsensitiveCompare(key, strTrailer) {
-					if nerr := h.SetTrailer(value); nerr != nil {
-						if err == nil {
-							err = nerr
-						}
-					}
-					continue
-				}
-			}
-		}
-		h.h = appendArg(h.h, key, value, argsHasValue)
+func (hb *headerBuf) nextKV(ss *scannerState) argsKV {
+	if !ss.initialized {
+		ss.nextColon = -1
+		ss.nextNewLine = -1
+		ss.initialized = true
 	}
-	if s.err != nil && err == nil {
-		err = s.err
+	buf := hb.buf[hb.off:]
+	bLen := len(buf)
+	if bLen >= 2 && buf[0] == rChar && buf[1] == nChar {
+		hb.off += 2
+		return hb.noKV() // \r\n\r\n Ends header.
 	}
-	if err != nil {
-		h.connectionClose = true
-		return 0, err
-	}
-
-	if h.contentLength < 0 {
-		h.contentLengthBytes = h.contentLengthBytes[:0]
-	}
-	if h.noHTTP11 && !h.connectionClose {
-		// close connection for non-http/1.1 request unless 'Connection: keep-alive' is set.
-		v := peekArgStr(h.h, strConnection)
-		h.connectionClose = !hasHeaderValue(b2s(v), strKeepAlive)
-	}
-	return s.hLen, nil
-}
-
-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] == rChar && s.b[1] == nChar {
-		s.b = s.b[2:]
-		s.hLen += 2
-		return false
-	}
-	if bLen >= 1 && s.b[0] == nChar {
-		s.b = s.b[1:]
-		s.hLen++
-		return false
+	if bLen >= 1 && buf[0] == nChar {
+		hb.off++
+		return hb.noKV() // \n\n: Ends header.
 	}
 
 	var n int
-	if s.nextColon >= 0 {
-		n = s.nextColon
-		s.nextColon = -1
+	if ss.nextColon >= 0 {
+		n = ss.nextColon
+		ss.nextColon = -1
 	} else {
-		n = bytes.IndexByte(s.b, ':')
+		n = bytes.IndexByte(buf, ':')
 
 		// There can't be a \n inside the header name, check for this.
-		x := bytes.IndexByte(s.b, nChar)
+		x := bytes.IndexByte(buf, nChar)
 		if x < 0 {
 			// A header name should always at some point be followed by a \n
 			// even if it's the one that terminates the header block.
-			s.err = errNeedMore
-			return false
+			ss.err = errNeedMore
+			return hb.noKV()
 		}
 		if x < n {
 			// There was a \n before the :
-			s.err = errInvalidName
-			return false
+			ss.err = errInvalidName
+			return hb.noKV()
 		}
 	}
 	if n < 0 {
-		s.err = errNeedMore
-		return false
+		ss.err = errNeedMore
+		return hb.noKV()
 	}
-	s.key = s.b[:n]
-	normalizeHeaderKey(s.key, s.disableNormalizing)
+
+	if bytes.IndexByte(buf[:n], ' ') >= 0 || bytes.IndexByte(buf[:n], '\t') >= 0 {
+		// Spaces between the header key and colon are not allowed.
+		// See RFC 7230, Section 3.2.4.
+		ss.err = errInvalidName
+		return hb.noKV()
+	}
+
+	var resultKV argsKV
+	resultKV.key = hb.slice(buf[:n])
+	normalizeHeaderKey(buf[:n], ss.disableNormalizing)
 	n++
-	for len(s.b) > n && s.b[n] == ' ' {
+	for len(buf) > n && buf[n] == ' ' {
 		n++
 		// the newline index is a relative index, and lines below trimmed `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--
+		ss.nextNewLine--
 	}
-	s.hLen += n
-	s.b = s.b[n:]
-	if s.nextNewLine >= 0 {
-		n = s.nextNewLine
-		s.nextNewLine = -1
+	ss.hLen += n
+	buf = buf[n:]
+	if ss.nextNewLine >= 0 {
+		n = ss.nextNewLine
+		ss.nextNewLine = -1
 	} else {
-		n = bytes.IndexByte(s.b, nChar)
+		n = bytes.IndexByte(buf, nChar)
 	}
 	if n < 0 {
-		s.err = errNeedMore
-		return false
+		ss.err = errNeedMore
+		return hb.noKV()
 	}
 	isMultiLineValue := false
 	for {
-		if n+1 >= len(s.b) {
+		if n+1 >= len(buf) {
 			break
 		}
-		if s.b[n+1] != ' ' && s.b[n+1] != '\t' {
+		if buf[n+1] != ' ' && buf[n+1] != '\t' {
 			break
 		}
-		d := bytes.IndexByte(s.b[n+1:], nChar)
+		d := bytes.IndexByte(buf[n+1:], nChar)
 		if d <= 0 {
 			break
-		} else if d == 1 && s.b[n+1] == rChar {
+		} else if d == 1 && buf[n+1] == rChar {
 			break
 		}
 		e := n + d + 1
-		if c := bytes.IndexByte(s.b[n+1:e], ':'); c >= 0 {
-			s.nextColon = c
-			s.nextNewLine = d - c - 1
+		if c := bytes.IndexByte(buf[n+1:e], ':'); c >= 0 {
+			ss.nextColon = c
+			ss.nextNewLine = d - c - 1
 			break
 		}
 		isMultiLineValue = true
 		n = e
 	}
-	if n >= len(s.b) {
-		s.err = errNeedMore
-		return false
+	if n >= len(buf) {
+		ss.err = errNeedMore
+		return hb.noKV()
 	}
-	oldB := s.b
-	s.value = s.b[:n]
-	s.hLen += n + 1
-	s.b = s.b[n+1:]
+	oldB := buf
+	value := buf[:n]
+	ss.hLen += n + 1
+	buf = buf[n+1:]
 
-	if n > 0 && s.value[n-1] == rChar {
+	if n > 0 && value[n-1] == rChar {
 		n--
 	}
-	for n > 0 && s.value[n-1] == ' ' {
+	for n > 0 && value[n-1] == ' ' {
 		n--
 	}
-	s.value = s.value[:n]
+	value = value[:n]
 	if isMultiLineValue {
-		s.value, s.b, s.hLen = normalizeHeaderValue(s.value, oldB, s.hLen)
+		value, buf, ss.hLen = normalizeHeaderValue(value, oldB, ss.hLen)
 	}
-	return true
+	resultKV.value = hb.slice(value)
+	return resultKV
 }
 
 func normalizeHeaderKey(b []byte, disableNormalizing bool) {
@@ -437,33 +472,6 @@ func parseContentLength(b string) (int, error) {
 	return v, nil
 }
 
-func hasHeaderValue(s, value string) bool {
-	var vs headerValueScanner
-	vs.b = s
-	for vs.next() {
-		if caseInsensitiveCompare(vs.value, value) {
-			return true
-		}
-	}
-	return false
-}
-
-func (s *headerValueScanner) next() bool {
-	b := s.b
-	if len(b) == 0 {
-		return false
-	}
-	n := strings.IndexByte(b, ',')
-	if n < 0 {
-		s.value = stripSpace(b)
-		s.b = b[len(b):]
-		return true
-	}
-	s.value = stripSpace(b[:n])
-	s.b = b[n+1:]
-	return true
-}
-
 func nextLine(b []byte) ([]byte, []byte, error) {
 	nNext := bytes.IndexByte(b, nChar)
 	if nNext < 0 {
@@ -552,6 +560,7 @@ func (h *header) readLoop(r *bufio.Reader, waitForMore bool) error {
 func (h *header) tryRead(r *bufio.Reader, n int) error {
 	h.resetSkipNormalize()
 	b, err := r.Peek(n)
+
 	if len(b) == 0 {
 		if err == io.EOF {
 			return err
@@ -577,35 +586,28 @@ func (h *header) tryRead(r *bufio.Reader, n int) error {
 		return fmt.Errorf("error when reading request headers: %w", err)
 	}
 	b = mustPeekBuffered(r)
-	headersLen, errParse := h.parse(b)
+	errParse := h.parse()
 	if errParse != nil {
 		return headerError("request", err, errParse, b, false)
 	}
-	mustDiscard(r, headersLen)
+	// mustDiscard(r, headersLen)
 	return nil
 }
 
+func (h *headerBuf) reset() {
+	*h = headerBuf{
+		buf:     h.buf[:0],
+		headers: h.headers[:0],
+		cookies: h.cookies[:0],
+	}
+}
+
 func (h *header) resetSkipNormalize() {
-	h.noHTTP11 = false
-	h.connectionClose = false
-
-	h.contentLength = 0
-	h.contentLengthBytes = h.contentLengthBytes[:0]
-
-	h.method = h.method[:0]
-	h.proto = h.proto[:0]
-	h.requestURI = h.requestURI[:0]
-	h.host = h.host[:0]
-	h.contentType = h.contentType[:0]
-	h.userAgent = h.userAgent[:0]
-	h.trailer = h.trailer[:0]
-	h.mulHeader = h.mulHeader[:0]
-
-	h.h = h.h[:0]
-	h.cookies = h.cookies[:0]
-	h.cookiesCollected = false
-
-	h.rawHeaders = h.rawHeaders[:0]
+	h.hbuf.reset()
+	*h = header{
+		hbuf:   h.hbuf,
+		logger: h.logger,
+	}
 }
 
 func headerError(typ string, err, errParse error, b []byte, secureErrorLogMessage bool) error {
@@ -638,6 +640,7 @@ func isOnlyCRLF(b []byte) bool {
 	}
 	return true
 }
+
 func headerErrorMsg(typ string, err error, b []byte, secureErrorLogMessage bool) error {
 	return fmt.Errorf("error when reading %s headers: %w. Buffer size=%d", typ, err, len(b))
 }
@@ -671,89 +674,19 @@ func mustDiscard(r *bufio.Reader, n int) {
 	}
 }
 
-// Peek returns header value for the given key.
-//
-// The returned value is valid until the request is released,
-// either though ReleaseRequest or your request handler returning.
-// Do not store references to returned value. Make copies instead.
-func (h *header) Peek(key string) []byte {
-	k := getHeaderKeyBytes(&h.bufKV, key, h.disableNormalizing)
-	return h.peek(b2s(k))
-}
-
 // Host returns Host header value.
 func (h *header) Host() []byte {
-	if h.disableSpecialHeader {
-		return peekArg(h.h, HeaderHost)
-	}
-	return h.host
-}
-
-func getHeaderKeyBytes(kv *argsKV, key string, disableNormalizing bool) []byte {
-	kv.key = append(kv.key[:0], key...)
-	normalizeHeaderKey(kv.key, disableNormalizing)
-	return kv.key
-}
-
-func peekArg(h []argsKV, k string) []byte {
-	for i, n := 0, len(h); i < n; i++ {
-		kv := &h[i]
-		if b2s(kv.key) == k {
-			return kv.value
-		}
-	}
-	return nil
-}
-
-func (h *header) peek(key string) []byte {
-	switch key {
-	case HeaderHost:
-		return h.Host()
-	case HeaderContentType:
-		return h.ContentType()
-	case HeaderUserAgent:
-		return h.UserAgent()
-	case HeaderConnection:
-		if h.ConnectionClose() {
-			return []byte(strClose)
-		}
-		return peekArg(h.h, key)
-	case HeaderContentLength:
-		return h.contentLengthBytes
-	case HeaderCookie:
-		if h.cookiesCollected {
-			return appendRequestCookieBytes(nil, h.cookies)
-		}
-		return peekArg(h.h, key)
-	case HeaderTrailer:
-		return appendArgsKey(nil, h.trailer, strCommaSpace)
-	default:
-		return peekArg(h.h, key)
-	}
+	return h.peekHeaderBytes(HeaderHost)
 }
 
 // ConnectionClose returns true if 'Connection: close' header is set.
 func (h *header) ConnectionClose() bool {
-	return h.connectionClose
+	return h.flags.hasAny(connectionClose)
 }
 
 // UserAgent returns User-Agent header value.
 func (h *header) UserAgent() []byte {
-	if h.disableSpecialHeader {
-		return peekArg(h.h, HeaderUserAgent)
-	}
-	return h.userAgent
-}
-
-func appendArgsKey(dst []byte, args []argsKV, sep string) []byte {
-	for i, n := 0, len(args); i < n; i++ {
-		kv := &args[i]
-		dst = append(dst, kv.key...)
-		if i+1 < n {
-			dst = append(dst, sep...)
-		}
-	}
-	return dst
+	return h.peekHeaderBytes(HeaderUserAgent)
 }
 
 // b2s converts byte slice to a string without memory allocation.