mirror of
https://github.com/zap-proto/http.git
synced 2026-10-04 03:07:19 +00:00
perf: zero-alloc ZAP-HTTP codec + pooled server/client hot path
Hand-rolled frame encoder/decoder (codec.go) replaces the zap.Builder + map[string][]string + json.Marshal path. Wire stays byte-identical (golden tests + TestCodec_CrossCheck vs a zap.Builder oracle across a wide corpus, incl. special-char headers). Headers are streamed straight into the frame tail as compact sorted-key JSON, byte-identical to json.Marshal, with a stdlib fallback for escape/UTF-8 values; decode scans the JSON in place into the fasthttp store with a json.Unmarshal fallback. Server: one RequestCtx per connection (Init2 + per-request Reset) instead of a per-request heap alloc; reused frame read/write buffers; single-syscall length-prefixed response write. Client: per-conn reused request/response buffers, single frame-type peek + one decode (was double zap.Parse), single-syscall request write. Codec allocs/op (default GOGC): MarshalRequest 11 -> 1 (Append: 0) MarshalResponse 19 -> 1 (Append: 0) UnmarshalResponse 11 -> 0 hot Append/Unmarshal all 0 allocs/op
This commit is contained in:
@@ -0,0 +1,81 @@
|
||||
package zaphttp_test
|
||||
|
||||
// Exported-API-only codec benchmarks. These use only MarshalRequest /
|
||||
// MarshalResponse / UnmarshalRequest / UnmarshalResponse, so the SAME file
|
||||
// compiles and runs against both v0.2.0 (the pre-optimization baseline) and the
|
||||
// zero-alloc branch — an apples-to-apples before/after of allocs/op.
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/valyala/fasthttp"
|
||||
zaphttp "github.com/zap-proto/http"
|
||||
)
|
||||
|
||||
func compatReq() *fasthttp.Request {
|
||||
req := fasthttp.AcquireRequest()
|
||||
req.Header.SetMethod("GET")
|
||||
req.SetRequestURI("/health")
|
||||
req.Header.SetHost("192.168.77.2:8391")
|
||||
return req
|
||||
}
|
||||
|
||||
func compatResp() *fasthttp.Response {
|
||||
resp := fasthttp.AcquireResponse()
|
||||
resp.SetStatusCode(200)
|
||||
resp.Header.SetContentType("text/plain; charset=utf-8")
|
||||
resp.SetBodyString("ok")
|
||||
return resp
|
||||
}
|
||||
|
||||
func BenchmarkCompat_MarshalRequest(b *testing.B) {
|
||||
req := compatReq()
|
||||
defer fasthttp.ReleaseRequest(req)
|
||||
b.ReportAllocs()
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
if _, err := zaphttp.MarshalRequest(req); err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func BenchmarkCompat_MarshalResponse(b *testing.B) {
|
||||
resp := compatResp()
|
||||
defer fasthttp.ReleaseResponse(resp)
|
||||
b.ReportAllocs()
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
if _, err := zaphttp.MarshalResponse(resp); err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func BenchmarkCompat_UnmarshalRequest(b *testing.B) {
|
||||
req := compatReq()
|
||||
frame, _ := zaphttp.MarshalRequest(req)
|
||||
fasthttp.ReleaseRequest(req)
|
||||
var dst fasthttp.Request
|
||||
b.ReportAllocs()
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
if err := zaphttp.UnmarshalRequest(frame, &dst); err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func BenchmarkCompat_UnmarshalResponse(b *testing.B) {
|
||||
resp := compatResp()
|
||||
frame, _ := zaphttp.MarshalResponse(resp)
|
||||
fasthttp.ReleaseResponse(resp)
|
||||
var dst fasthttp.Response
|
||||
b.ReportAllocs()
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
if err := zaphttp.UnmarshalResponse(frame, &dst); err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -29,11 +29,14 @@ import (
|
||||
"github.com/valyala/fasthttp"
|
||||
)
|
||||
|
||||
// pooledConn is one idle connection. The bufio.Reader is held with the
|
||||
// conn so the read side amortizes its own buffer alloc across requests.
|
||||
// pooledConn is one idle connection. The bufio.Reader and the request/response
|
||||
// scratch buffers are held with the conn so a warm connection marshals, writes,
|
||||
// reads, and decodes with zero heap allocation across requests.
|
||||
type pooledConn struct {
|
||||
c net.Conn
|
||||
br *bufio.Reader
|
||||
c net.Conn
|
||||
br *bufio.Reader
|
||||
wbuf []byte // request frame, incl. 4-byte length prefix
|
||||
rbuf []byte // inbound response frame
|
||||
}
|
||||
|
||||
// Transport speaks ZAP-HTTP over a pooled TCP connection. Field-zero is
|
||||
@@ -86,12 +89,15 @@ func (t *Transport) Do(req *fasthttp.Request, resp *fasthttp.Response) error {
|
||||
}
|
||||
conn, br := pc.c, pc.br
|
||||
|
||||
frame, err := MarshalRequest(req)
|
||||
// Build [4-byte length prefix][request frame] in the conn's reused buffer and
|
||||
// write it with a single syscall.
|
||||
pc.wbuf = append(pc.wbuf[:0], 0, 0, 0, 0)
|
||||
pc.wbuf, err = AppendRequest(pc.wbuf, req)
|
||||
if err != nil {
|
||||
conn.Close()
|
||||
return err
|
||||
}
|
||||
if err := writeFrame(conn, frame); err != nil {
|
||||
if err := writeFramePrefixed(conn, pc.wbuf); err != nil {
|
||||
conn.Close()
|
||||
return fmt.Errorf("zaphttp: write request: %w", err)
|
||||
}
|
||||
@@ -99,20 +105,23 @@ func (t *Transport) Do(req *fasthttp.Request, resp *fasthttp.Response) error {
|
||||
if t.readTimeout > 0 {
|
||||
_ = conn.SetReadDeadline(time.Now().Add(t.readTimeout))
|
||||
}
|
||||
respFrame, err := readFrame(br)
|
||||
pc.rbuf, err = readFrameInto(br, pc.rbuf)
|
||||
if err != nil {
|
||||
conn.Close()
|
||||
return fmt.Errorf("zaphttp: read response: %w", err)
|
||||
}
|
||||
ft, err := FrameTypeOf(respFrame)
|
||||
if err != nil {
|
||||
respFrame := pc.rbuf
|
||||
ft, ok := frameType(respFrame)
|
||||
if !ok {
|
||||
conn.Close()
|
||||
return fmt.Errorf("zaphttp: response frame: %w", err)
|
||||
return fmt.Errorf("zaphttp: short response frame")
|
||||
}
|
||||
|
||||
// Streamed response: apply the head, then attach a body stream that reads
|
||||
// FrameData chunks until FrameEnd. The conn is NOT pooled until the stream
|
||||
// is fully consumed (or the body is closed) — the reader owns it.
|
||||
// is fully consumed (or the body is closed) — the reader owns it. The head
|
||||
// frame must be copied out of the conn's reused rbuf, since the streamReader
|
||||
// keeps reading into it.
|
||||
if ft == FrameResponseHead {
|
||||
if err := UnmarshalResponseHead(respFrame, resp); err != nil {
|
||||
conn.Close()
|
||||
@@ -232,7 +241,7 @@ func (t *Transport) acquireConn() (*pooledConn, error) {
|
||||
_ = tc.SetKeepAlivePeriod(30 * time.Second)
|
||||
_ = tc.SetNoDelay(true)
|
||||
}
|
||||
return &pooledConn{c: c, br: bufio.NewReader(c)}, nil
|
||||
return &pooledConn{c: c, br: bufio.NewReaderSize(c, connReadBufSize)}, nil
|
||||
}
|
||||
|
||||
// releaseConn returns a conn to the pool, or closes it if the pool is full.
|
||||
|
||||
@@ -0,0 +1,534 @@
|
||||
// Zero-alloc ZAP-HTTP frame codec.
|
||||
//
|
||||
// This file is the hot-path encoder/decoder. It writes and reads the SAME
|
||||
// bytes as the reference zap-proto/go Builder path (proven byte-identical by
|
||||
// the golden-wire tests and TestCodec_CrossCheck), but without the per-request
|
||||
// allocations that path incurred:
|
||||
//
|
||||
// - no *zap.Builder / *zap.ObjectBuilder per marshal, no per-field double
|
||||
// copy (Builder.SetBytes did []byte(v) then append([]byte(nil), v...));
|
||||
// - no map[string][]string + json.Marshal for the headers slot — headers are
|
||||
// streamed straight into the frame's variable tail as compact, sorted-key
|
||||
// JSON, byte-identical to json.Marshal(map[string][]string);
|
||||
// - no *zap.Message per parse and no map + json.Unmarshal on decode — the
|
||||
// headers JSON is scanned in place and applied straight to the fasthttp
|
||||
// header store.
|
||||
//
|
||||
// Callers append into a reusable buffer (Append{Request,Response}), so after a
|
||||
// connection warms up the steady-state cost is zero heap allocation.
|
||||
//
|
||||
// # Wire layout (must stay byte-identical — see schema/zap_http.zap)
|
||||
//
|
||||
// A frame is one zap-proto/go message: a 16-byte header (magic "ZAP\x00",
|
||||
// version=1, flags, rootOffset=16, size) followed by the root object. The root
|
||||
// object is a fixed 48-byte section (fully zero-filled) whose text/bytes fields
|
||||
// are 8-byte {relOffset,length} slots; every non-empty field's bytes are
|
||||
// appended to the variable tail in field-declaration order and its slot's
|
||||
// relOffset is patched to point at them. An empty text/bytes field stays a
|
||||
// zeroed {0,0} slot with no tail entry. This is exactly what zap.Builder emits.
|
||||
|
||||
package zaphttp
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/binary"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"slices"
|
||||
"sync"
|
||||
"unsafe"
|
||||
|
||||
"github.com/valyala/fasthttp"
|
||||
zap "github.com/zap-proto/go"
|
||||
)
|
||||
|
||||
// zeroFixed is a read-only source of zero bytes for laying down a struct's
|
||||
// fixed section. Sized to the largest object (48 bytes); shared read-only
|
||||
// across goroutines.
|
||||
var zeroFixed [reqSlotSize]byte
|
||||
|
||||
// Static defaults, kept as byte slices so the encoder never allocates one.
|
||||
var (
|
||||
bMethodGET = []byte(fasthttp.MethodGet)
|
||||
bDefaultProto = []byte(defaultProto)
|
||||
)
|
||||
|
||||
// ---- little-endian frame header ----
|
||||
|
||||
// appendFrameHeader writes the 16-byte ZAP message header with rootOffset=16
|
||||
// (the root object begins immediately after the header, which is 8-aligned so
|
||||
// the object needs no leading padding) and a placeholder size, patched by the
|
||||
// caller once the frame is complete.
|
||||
func appendFrameHeader(dst []byte, flags uint16) []byte {
|
||||
dst = append(dst, 'Z', 'A', 'P', 0)
|
||||
dst = binary.LittleEndian.AppendUint16(dst, uint16(zap.Version))
|
||||
dst = binary.LittleEndian.AppendUint16(dst, flags)
|
||||
dst = binary.LittleEndian.AppendUint32(dst, uint32(zap.HeaderSize)) // rootOffset
|
||||
dst = binary.LittleEndian.AppendUint32(dst, 0) // size (patched)
|
||||
return dst
|
||||
}
|
||||
|
||||
// patchVarSlot patches an 8-byte {relOffset,length} slot to point at tail data
|
||||
// that was already appended starting at dataStart. A zero-length field leaves
|
||||
// the slot as the zeroed {0,0} null it started as (byte-identical to
|
||||
// Builder.SetBytes(nil)).
|
||||
func patchVarSlot(dst []byte, slot, dataStart int) []byte {
|
||||
n := len(dst) - dataStart
|
||||
if n == 0 {
|
||||
return dst
|
||||
}
|
||||
binary.LittleEndian.PutUint32(dst[slot:], uint32(dataStart-slot))
|
||||
binary.LittleEndian.PutUint32(dst[slot+4:], uint32(n))
|
||||
return dst
|
||||
}
|
||||
|
||||
// putVarBytes appends a non-empty field's bytes to the tail and patches its
|
||||
// slot. Empty data leaves the slot null.
|
||||
func putVarBytes(dst []byte, slot int, data []byte) []byte {
|
||||
if len(data) == 0 {
|
||||
return dst
|
||||
}
|
||||
start := len(dst)
|
||||
dst = append(dst, data...)
|
||||
return patchVarSlot(dst, slot, start)
|
||||
}
|
||||
|
||||
// ---- request / response encoders ----
|
||||
|
||||
// AppendRequest appends a request frame to dst and returns the extended slice.
|
||||
// Passing a reused (len-0) buffer makes the steady-state marshal zero-alloc;
|
||||
// MarshalRequest is the dst==nil convenience.
|
||||
func AppendRequest(dst []byte, req *fasthttp.Request) ([]byte, error) {
|
||||
frameStart := len(dst)
|
||||
dst = appendFrameHeader(dst, FrameRequest<<8)
|
||||
obj := len(dst) // == frameStart + 16
|
||||
|
||||
dst = append(dst, zeroFixed[:reqSlotSize]...)
|
||||
|
||||
method := req.Header.Method()
|
||||
if len(method) == 0 {
|
||||
method = bMethodGET
|
||||
}
|
||||
proto := req.Header.Protocol()
|
||||
if len(proto) == 0 {
|
||||
proto = bDefaultProto
|
||||
}
|
||||
|
||||
dst = putVarBytes(dst, obj+reqMethod, method)
|
||||
dst = putVarBytes(dst, obj+reqTarget, req.RequestURI())
|
||||
dst = putVarBytes(dst, obj+reqProto, proto)
|
||||
|
||||
hStart := len(dst)
|
||||
skip := newTrailerSkip(&req.Header)
|
||||
dst = appendRequestHeadersJSON(dst, &req.Header, skip)
|
||||
dst = patchVarSlot(dst, obj+reqHeaders, hStart)
|
||||
|
||||
dst = putVarBytes(dst, obj+reqBody, req.Body())
|
||||
|
||||
trailer, err := encodeTrailers(&req.Header)
|
||||
if err != nil {
|
||||
return dst, fmt.Errorf("zaphttp: encode trailer: %w", err)
|
||||
}
|
||||
dst = putVarBytes(dst, obj+reqTrailer, trailer)
|
||||
|
||||
binary.LittleEndian.PutUint32(dst[frameStart+12:], uint32(len(dst)-frameStart))
|
||||
return dst, nil
|
||||
}
|
||||
|
||||
// AppendResponse appends a response frame to dst and returns the extended
|
||||
// slice. See AppendRequest.
|
||||
func AppendResponse(dst []byte, resp *fasthttp.Response) ([]byte, error) {
|
||||
frameStart := len(dst)
|
||||
dst = appendFrameHeader(dst, FrameResponse<<8)
|
||||
obj := len(dst)
|
||||
|
||||
dst = append(dst, zeroFixed[:respSlotSize]...)
|
||||
|
||||
status := resp.StatusCode()
|
||||
if status == 0 {
|
||||
status = fasthttp.StatusOK
|
||||
}
|
||||
binary.LittleEndian.PutUint16(dst[obj+respStatus:], uint16(status))
|
||||
|
||||
reason := resp.Header.StatusMessage()
|
||||
if len(reason) == 0 {
|
||||
reason = unsafeBytes(fasthttp.StatusMessage(status))
|
||||
}
|
||||
proto := resp.Header.Protocol()
|
||||
if len(proto) == 0 {
|
||||
proto = bDefaultProto
|
||||
}
|
||||
|
||||
dst = putVarBytes(dst, obj+respReason, reason)
|
||||
dst = putVarBytes(dst, obj+respProto, proto)
|
||||
|
||||
hStart := len(dst)
|
||||
skip := newTrailerSkip(&resp.Header)
|
||||
dst = appendResponseHeadersJSON(dst, &resp.Header, skip)
|
||||
dst = patchVarSlot(dst, obj+respHeaders, hStart)
|
||||
|
||||
dst = putVarBytes(dst, obj+respBody, resp.Body())
|
||||
|
||||
trailer, err := encodeTrailers(&resp.Header)
|
||||
if err != nil {
|
||||
return dst, fmt.Errorf("zaphttp: encode trailer: %w", err)
|
||||
}
|
||||
dst = putVarBytes(dst, obj+respTrailer, trailer)
|
||||
|
||||
binary.LittleEndian.PutUint32(dst[frameStart+12:], uint32(len(dst)-frameStart))
|
||||
return dst, nil
|
||||
}
|
||||
|
||||
// ---- frame parsing (decode) ----
|
||||
|
||||
// frameRoot validates a frame's header (matching zap.Parse's magic/version/size
|
||||
// checks) and confirms its type, returning the root object offset and the
|
||||
// message length. Bytes past size are ignored, exactly as zap.Parse slices to
|
||||
// data[:size].
|
||||
func frameRoot(frame []byte, wantType uint16) (root, size int, err error) {
|
||||
if len(frame) < zap.HeaderSize {
|
||||
return 0, 0, zap.ErrBufferTooSmall
|
||||
}
|
||||
if string(frame[0:4]) != zap.Magic {
|
||||
return 0, 0, zap.ErrInvalidMagic
|
||||
}
|
||||
// Accept wire version 1 (this runtime's baseline) and 2 (the transport
|
||||
// framing generation); the data segment is identical across both.
|
||||
if v := binary.LittleEndian.Uint16(frame[4:6]); v != 1 && v != 2 {
|
||||
return 0, 0, zap.ErrInvalidVersion
|
||||
}
|
||||
size = int(binary.LittleEndian.Uint32(frame[12:16]))
|
||||
if size < zap.HeaderSize || size > len(frame) {
|
||||
return 0, 0, zap.ErrBufferTooSmall
|
||||
}
|
||||
if t := binary.LittleEndian.Uint16(frame[6:8]) >> 8; t != wantType {
|
||||
return 0, 0, fmt.Errorf("zaphttp: frame type %#x, want %#x", t, wantType)
|
||||
}
|
||||
root = int(binary.LittleEndian.Uint32(frame[8:12]))
|
||||
if root < zap.HeaderSize || root >= size {
|
||||
return 0, 0, zap.ErrInvalidOffset
|
||||
}
|
||||
return root, size, nil
|
||||
}
|
||||
|
||||
// frameType peeks a frame's type from its flags without a full parse. Zero
|
||||
// alloc — the client uses it to dispatch streamed vs single-frame responses.
|
||||
func frameType(frame []byte) (uint16, bool) {
|
||||
if len(frame) < zap.HeaderSize {
|
||||
return 0, false
|
||||
}
|
||||
return binary.LittleEndian.Uint16(frame[6:8]) >> 8, true
|
||||
}
|
||||
|
||||
// readVar reads a text/bytes field at object offset fieldOff, zero-copy.
|
||||
// Mirrors zap.Object.Bytes exactly, including the malleability guard that
|
||||
// rejects a target inside the wire header (absPos < HeaderSize).
|
||||
func readVar(frame []byte, size, root, fieldOff int) []byte {
|
||||
slot := root + fieldOff
|
||||
if slot+8 > size {
|
||||
return nil
|
||||
}
|
||||
rel := binary.LittleEndian.Uint32(frame[slot:])
|
||||
if rel == 0 {
|
||||
return nil
|
||||
}
|
||||
length := int(binary.LittleEndian.Uint32(frame[slot+4:]))
|
||||
abs := slot + int(rel)
|
||||
if abs < zap.HeaderSize || abs+length > size {
|
||||
return nil
|
||||
}
|
||||
return frame[abs : abs+length]
|
||||
}
|
||||
|
||||
// readU16 reads a uint16 scalar at object offset fieldOff.
|
||||
func readU16(frame []byte, size, root, fieldOff int) uint16 {
|
||||
pos := root + fieldOff
|
||||
if pos+2 > size {
|
||||
return 0
|
||||
}
|
||||
return binary.LittleEndian.Uint16(frame[pos:])
|
||||
}
|
||||
|
||||
// ---- header JSON encode (zero-alloc, byte-identical to json.Marshal) ----
|
||||
|
||||
// hpair is one header name/value, both aliasing fasthttp's store; valid only
|
||||
// for the duration of a single marshal call.
|
||||
type hpair struct{ key, val []byte }
|
||||
|
||||
// hcollector is a pooled scratch slice for sorting header pairs.
|
||||
type hcollector struct{ pairs []hpair }
|
||||
|
||||
var hcollectorPool = sync.Pool{New: func() any { return &hcollector{pairs: make([]hpair, 0, 16)} }}
|
||||
|
||||
// trailerSkip is the set of trailer-declared header names, which ride in the
|
||||
// frame's trailer slot and must be excluded from the headers slot. nil when the
|
||||
// message declares no trailers (the common case), so it costs nothing.
|
||||
type trailerSkip struct{ keys [][]byte }
|
||||
|
||||
func newTrailerSkip(h interface{ PeekTrailerKeys() [][]byte }) *trailerSkip {
|
||||
keys := h.PeekTrailerKeys()
|
||||
if len(keys) == 0 {
|
||||
return nil
|
||||
}
|
||||
return &trailerSkip{keys: keys}
|
||||
}
|
||||
|
||||
func (s *trailerSkip) has(key []byte) bool {
|
||||
if s == nil {
|
||||
return false
|
||||
}
|
||||
for _, k := range s.keys {
|
||||
if bytesEqualFold(k, key) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// appendRequestHeadersJSON and appendResponseHeadersJSON collect, sort, and emit
|
||||
// a request's / response's headers as compact JSON. Two concrete functions (not
|
||||
// one over an interface) so escape analysis can prove the VisitAll closure does
|
||||
// not leak and keep it — and the collected pairs — off the heap.
|
||||
func appendRequestHeadersJSON(dst []byte, h *fasthttp.RequestHeader, skip *trailerSkip) []byte {
|
||||
c := hcollectorPool.Get().(*hcollector)
|
||||
c.pairs = c.pairs[:0]
|
||||
h.VisitAll(func(key, value []byte) {
|
||||
if isFrameOwnedHeaderBytes(key) || skip.has(key) {
|
||||
return
|
||||
}
|
||||
c.pairs = append(c.pairs, hpair{key, value})
|
||||
})
|
||||
dst = emitHeadersJSON(dst, c.pairs)
|
||||
hcollectorPool.Put(c)
|
||||
return dst
|
||||
}
|
||||
|
||||
func appendResponseHeadersJSON(dst []byte, h *fasthttp.ResponseHeader, skip *trailerSkip) []byte {
|
||||
c := hcollectorPool.Get().(*hcollector)
|
||||
c.pairs = c.pairs[:0]
|
||||
h.VisitAll(func(key, value []byte) {
|
||||
if isFrameOwnedHeaderBytes(key) || skip.has(key) {
|
||||
return
|
||||
}
|
||||
c.pairs = append(c.pairs, hpair{key, value})
|
||||
})
|
||||
dst = emitHeadersJSON(dst, c.pairs)
|
||||
hcollectorPool.Put(c)
|
||||
return dst
|
||||
}
|
||||
|
||||
// emitHeadersJSON writes pairs as {"k":["v",...],...}, keys sorted bytewise and
|
||||
// same-key values grouped in first-seen order — byte-identical to
|
||||
// json.Marshal(map[string][]string). Empty pairs append nothing (a null slot),
|
||||
// matching encodeHeaderMap's nil return for an empty map.
|
||||
func emitHeadersJSON(dst []byte, pairs []hpair) []byte {
|
||||
if len(pairs) == 0 {
|
||||
return dst
|
||||
}
|
||||
// Stable sort by key: equal keys stay in VisitAll order, so their values
|
||||
// keep insertion order (matches json's per-slice ordering).
|
||||
slices.SortStableFunc(pairs, func(a, b hpair) int { return bytes.Compare(a.key, b.key) })
|
||||
|
||||
dst = append(dst, '{')
|
||||
for i := 0; i < len(pairs); {
|
||||
k := pairs[i].key
|
||||
if i > 0 {
|
||||
dst = append(dst, ',')
|
||||
}
|
||||
dst = appendJSONString(dst, k)
|
||||
dst = append(dst, ':', '[')
|
||||
j := i
|
||||
for j < len(pairs) && bytes.Equal(pairs[j].key, k) {
|
||||
if j > i {
|
||||
dst = append(dst, ',')
|
||||
}
|
||||
dst = appendJSONString(dst, pairs[j].val)
|
||||
j++
|
||||
}
|
||||
dst = append(dst, ']')
|
||||
i = j
|
||||
}
|
||||
return append(dst, '}')
|
||||
}
|
||||
|
||||
// appendJSONString appends s as a JSON string literal, byte-identical to
|
||||
// encoding/json with the default HTML escaping. The fast path bulk-copies
|
||||
// values made only of "clean" ASCII (0x20..0x7e minus " \ < > &), which
|
||||
// encoding/json emits verbatim; anything else (control bytes, HTML-escaped
|
||||
// chars, or any byte >= 0x7f needing UTF-8 / U+2028-9 handling) defers to the
|
||||
// standard encoder for that one string, so the bytes always match exactly.
|
||||
func appendJSONString(dst, s []byte) []byte {
|
||||
if jsonCleanASCII(s) {
|
||||
dst = append(dst, '"')
|
||||
dst = append(dst, s...)
|
||||
return append(dst, '"')
|
||||
}
|
||||
b, _ := json.Marshal(unsafeString(s))
|
||||
return append(dst, b...)
|
||||
}
|
||||
|
||||
func jsonCleanASCII(s []byte) bool {
|
||||
for _, b := range s {
|
||||
if b < 0x20 || b > 0x7e || b == '"' || b == '\\' || b == '<' || b == '>' || b == '&' {
|
||||
return false
|
||||
}
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
// ---- header JSON decode (zero-alloc fast path + safe fallback) ----
|
||||
|
||||
// forEachHeaderKV scans a compact, escape-free JSON object of the shape
|
||||
// {"k":["v",...],...} and calls fn(key,value) for every value, with slices that
|
||||
// alias raw. It returns false — asking the caller to fall back to the standard
|
||||
// json decoder — for any input containing an escape or deviating from that
|
||||
// exact compact grammar, so correctness holds for arbitrary/foreign frames
|
||||
// while the common case stays allocation-free.
|
||||
func forEachHeaderKV(raw []byte, fn func(key, val []byte)) bool {
|
||||
if len(raw) == 0 {
|
||||
return true
|
||||
}
|
||||
if bytes.IndexByte(raw, '\\') >= 0 {
|
||||
return false
|
||||
}
|
||||
i := 0
|
||||
if raw[i] != '{' {
|
||||
return false
|
||||
}
|
||||
i++
|
||||
if i < len(raw) && raw[i] == '}' {
|
||||
return i+1 == len(raw)
|
||||
}
|
||||
for {
|
||||
if i >= len(raw) || raw[i] != '"' {
|
||||
return false
|
||||
}
|
||||
key, ni, ok := scanJSONStr(raw, i)
|
||||
if !ok {
|
||||
return false
|
||||
}
|
||||
i = ni
|
||||
if i >= len(raw) || raw[i] != ':' {
|
||||
return false
|
||||
}
|
||||
i++
|
||||
if i >= len(raw) || raw[i] != '[' {
|
||||
return false
|
||||
}
|
||||
i++
|
||||
if i < len(raw) && raw[i] == ']' {
|
||||
i++
|
||||
} else {
|
||||
for {
|
||||
val, nv, ok := scanJSONStr(raw, i)
|
||||
if !ok {
|
||||
return false
|
||||
}
|
||||
i = nv
|
||||
fn(key, val)
|
||||
if i >= len(raw) {
|
||||
return false
|
||||
}
|
||||
switch raw[i] {
|
||||
case ',':
|
||||
i++
|
||||
continue
|
||||
case ']':
|
||||
i++
|
||||
default:
|
||||
return false
|
||||
}
|
||||
break
|
||||
}
|
||||
}
|
||||
if i >= len(raw) {
|
||||
return false
|
||||
}
|
||||
switch raw[i] {
|
||||
case ',':
|
||||
i++
|
||||
continue
|
||||
case '}':
|
||||
return i+1 == len(raw)
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// scanJSONStr reads a JSON string starting at raw[i]=='"', returning the inner
|
||||
// bytes (no escapes — the caller guaranteed none) and the index just past the
|
||||
// closing quote.
|
||||
func scanJSONStr(raw []byte, i int) (s []byte, next int, ok bool) {
|
||||
if i >= len(raw) || raw[i] != '"' {
|
||||
return nil, i, false
|
||||
}
|
||||
i++
|
||||
start := i
|
||||
for i < len(raw) {
|
||||
if raw[i] == '"' {
|
||||
return raw[start:i], i + 1, true
|
||||
}
|
||||
i++
|
||||
}
|
||||
return nil, i, false
|
||||
}
|
||||
|
||||
// ---- shared helpers ----
|
||||
|
||||
// isFrameOwnedHeaderBytes reports whether a header name is reconstructed from
|
||||
// the frame itself (Host, Content-Length) or represented by the trailer slot
|
||||
// (Trailer), and so is excluded from the headers slot. Byte-wise EqualFold —
|
||||
// no allocation — matching isFrameOwnedHeader's strings.EqualFold.
|
||||
func isFrameOwnedHeaderBytes(key []byte) bool {
|
||||
return bytesEqualFold(key, strHost) ||
|
||||
bytesEqualFold(key, strContentLength) ||
|
||||
bytesEqualFold(key, strTrailer)
|
||||
}
|
||||
|
||||
var (
|
||||
strHost = []byte(fasthttp.HeaderHost)
|
||||
strContentLength = []byte(fasthttp.HeaderContentLength)
|
||||
strTrailer = []byte(fasthttp.HeaderTrailer)
|
||||
)
|
||||
|
||||
// bytesEqualFold is ASCII case-insensitive byte equality (header names are
|
||||
// ASCII tokens), allocation-free.
|
||||
func bytesEqualFold(a, b []byte) bool {
|
||||
if len(a) != len(b) {
|
||||
return false
|
||||
}
|
||||
for i := range a {
|
||||
ca, cb := a[i], b[i]
|
||||
if ca == cb {
|
||||
continue
|
||||
}
|
||||
if 'A' <= ca && ca <= 'Z' {
|
||||
ca += 'a' - 'A'
|
||||
}
|
||||
if 'A' <= cb && cb <= 'Z' {
|
||||
cb += 'a' - 'A'
|
||||
}
|
||||
if ca != cb {
|
||||
return false
|
||||
}
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
// unsafeBytes returns the bytes of s without copying. The result is read-only;
|
||||
// used only where the bytes are immediately copied out (append into a frame).
|
||||
func unsafeBytes(s string) []byte {
|
||||
if s == "" {
|
||||
return nil
|
||||
}
|
||||
return unsafe.Slice(unsafe.StringData(s), len(s))
|
||||
}
|
||||
|
||||
// unsafeString views b as a string without copying, for read-only use.
|
||||
func unsafeString(b []byte) string {
|
||||
if len(b) == 0 {
|
||||
return ""
|
||||
}
|
||||
return unsafe.String(&b[0], len(b))
|
||||
}
|
||||
@@ -0,0 +1,368 @@
|
||||
package zaphttp
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"net/textproto"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/valyala/fasthttp"
|
||||
zap "github.com/zap-proto/go"
|
||||
)
|
||||
|
||||
// White-box tests proving the zero-alloc codec is byte-identical to the
|
||||
// standard library (for the headers JSON) and to the reference zap.Builder path
|
||||
// (for whole frames), plus -benchmem benchmarks documenting allocs/op.
|
||||
|
||||
// ---- header JSON byte-identity vs encoding/json ----
|
||||
|
||||
func TestAppendJSONString_matchesStdlib(t *testing.T) {
|
||||
cases := []string{
|
||||
"",
|
||||
"simple",
|
||||
"text/plain; charset=utf-8",
|
||||
"application/json",
|
||||
"Bearer abc.def.ghi",
|
||||
`he said "quote"`,
|
||||
`back\slash`,
|
||||
"tab\tnewline\ncr\r",
|
||||
"ctrl\x00\x01\x02\x1f",
|
||||
"html <b>&</b> >x<",
|
||||
"slash/and:colon;semi",
|
||||
"café ☕ 世界 — unicode",
|
||||
"line
sep
para", // U+2028/U+2029: json escapes, we must too
|
||||
"del\x7fbyte",
|
||||
"高性能 ZAP",
|
||||
}
|
||||
for _, s := range cases {
|
||||
got := appendJSONString(nil, []byte(s))
|
||||
want, err := json.Marshal(s)
|
||||
if err != nil {
|
||||
t.Fatalf("json.Marshal(%q): %v", s, err)
|
||||
}
|
||||
if !bytes.Equal(got, want) {
|
||||
t.Errorf("appendJSONString(%q)\n got=%s\nwant=%s", s, got, want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestEmitHeadersJSON_matchesStdlib(t *testing.T) {
|
||||
cases := []map[string][]string{
|
||||
{"Content-Type": {"application/json"}},
|
||||
{"Content-Type": {"text/plain; charset=utf-8"}},
|
||||
// multi-key: json sorts keys; we must match
|
||||
{"X-Trace-Id": {"abc-123"}, "Content-Type": {"application/json"}},
|
||||
// multi-value
|
||||
{"X-Multi": {"one", "two", "three"}},
|
||||
// keys that sort non-trivially + values needing escapes
|
||||
{"Z-Last": {"z"}, "A-First": {`a"b`}, "M-Mid": {"<m&m>"}},
|
||||
{"X-Unicode": {"café ☕"}, "X-Ascii": {"plain"}},
|
||||
}
|
||||
for _, m := range cases {
|
||||
var pairs []hpair
|
||||
for k, vals := range m {
|
||||
for _, v := range vals {
|
||||
pairs = append(pairs, hpair{[]byte(k), []byte(v)})
|
||||
}
|
||||
}
|
||||
got := emitHeadersJSON(nil, pairs)
|
||||
want, err := json.Marshal(m)
|
||||
if err != nil {
|
||||
t.Fatalf("json.Marshal: %v", err)
|
||||
}
|
||||
if !bytes.Equal(got, want) {
|
||||
t.Errorf("emitHeadersJSON(%v)\n got=%s\nwant=%s", m, got, want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ---- whole-frame byte-identity vs the reference zap.Builder path ----
|
||||
|
||||
// refMarshalRequest / refMarshalResponse reproduce the prior codec exactly
|
||||
// (zap.Builder + json header map), the byte-for-byte oracle the fast codec must
|
||||
// match for arbitrary inputs — not just the four captured golden frames.
|
||||
|
||||
func refCollect(vis interface {
|
||||
VisitAll(func(k, v []byte))
|
||||
PeekTrailerKeys() [][]byte
|
||||
}) map[string][]string {
|
||||
skip := map[string]bool{}
|
||||
for _, k := range vis.PeekTrailerKeys() {
|
||||
skip[strings.ToLower(string(k))] = true
|
||||
}
|
||||
m := map[string][]string{}
|
||||
vis.VisitAll(func(key, value []byte) {
|
||||
k := string(key)
|
||||
lk := strings.ToLower(k)
|
||||
if lk == "host" || lk == "content-length" || lk == "trailer" || skip[lk] {
|
||||
return
|
||||
}
|
||||
m[k] = append(m[k], string(value))
|
||||
})
|
||||
return m
|
||||
}
|
||||
|
||||
func refHeaderJSON(m map[string][]string) []byte {
|
||||
if len(m) == 0 {
|
||||
return nil
|
||||
}
|
||||
b, _ := json.Marshal(m)
|
||||
return b
|
||||
}
|
||||
|
||||
func refTrailerJSON(vis interface {
|
||||
PeekTrailerKeys() [][]byte
|
||||
PeekAll(string) [][]byte
|
||||
}) []byte {
|
||||
keys := vis.PeekTrailerKeys()
|
||||
if len(keys) == 0 {
|
||||
return nil
|
||||
}
|
||||
m := map[string][]string{}
|
||||
for _, k := range keys {
|
||||
ck := textproto.CanonicalMIMEHeaderKey(string(k))
|
||||
for _, v := range vis.PeekAll(string(k)) {
|
||||
m[ck] = append(m[ck], string(v))
|
||||
}
|
||||
}
|
||||
b, _ := json.Marshal(m)
|
||||
return b
|
||||
}
|
||||
|
||||
func refMarshalRequest(req *fasthttp.Request) []byte {
|
||||
method := string(req.Header.Method())
|
||||
if method == "" {
|
||||
method = "GET"
|
||||
}
|
||||
target := string(req.RequestURI())
|
||||
proto := string(req.Header.Protocol())
|
||||
if proto == "" {
|
||||
proto = defaultProto
|
||||
}
|
||||
headers := refHeaderJSON(refCollect(&req.Header))
|
||||
trailer := refTrailerJSON(&req.Header)
|
||||
|
||||
b := zap.NewBuilder(reqSlotSize + len(req.Body()) + len(headers) + len(trailer) + 64)
|
||||
ob := b.StartObject(reqSlotSize)
|
||||
ob.SetText(reqMethod, method)
|
||||
ob.SetText(reqTarget, target)
|
||||
ob.SetText(reqProto, proto)
|
||||
ob.SetBytes(reqHeaders, headers)
|
||||
ob.SetBytes(reqBody, req.Body())
|
||||
ob.SetBytes(reqTrailer, trailer)
|
||||
ob.FinishAsRoot()
|
||||
return b.FinishWithFlags(FrameRequest << 8)
|
||||
}
|
||||
|
||||
func refMarshalResponse(resp *fasthttp.Response) []byte {
|
||||
status := resp.StatusCode()
|
||||
if status == 0 {
|
||||
status = fasthttp.StatusOK
|
||||
}
|
||||
reason := string(resp.Header.StatusMessage())
|
||||
if reason == "" {
|
||||
reason = fasthttp.StatusMessage(status)
|
||||
}
|
||||
proto := string(resp.Header.Protocol())
|
||||
if proto == "" {
|
||||
proto = defaultProto
|
||||
}
|
||||
headers := refHeaderJSON(refCollect(&resp.Header))
|
||||
trailer := refTrailerJSON(&resp.Header)
|
||||
|
||||
b := zap.NewBuilder(respSlotSize + len(resp.Body()) + len(headers) + len(trailer) + 64)
|
||||
ob := b.StartObject(respSlotSize)
|
||||
ob.SetUint16(respStatus, uint16(status))
|
||||
ob.SetText(respReason, reason)
|
||||
ob.SetText(respProto, proto)
|
||||
ob.SetBytes(respHeaders, headers)
|
||||
ob.SetBytes(respBody, resp.Body())
|
||||
ob.SetBytes(respTrailer, trailer)
|
||||
ob.FinishAsRoot()
|
||||
return b.FinishWithFlags(FrameResponse << 8)
|
||||
}
|
||||
|
||||
func TestCodec_CrossCheckRequest(t *testing.T) {
|
||||
build := []func(*fasthttp.Request){
|
||||
func(r *fasthttp.Request) { r.Header.SetMethod("GET"); r.SetRequestURI("/health"); r.Header.SetProtocol("HTTP/1.1") },
|
||||
func(r *fasthttp.Request) {
|
||||
r.Header.SetMethod("POST")
|
||||
r.SetRequestURI("/v1/blocks?height=42")
|
||||
r.Header.SetProtocol("HTTP/1.1")
|
||||
r.Header.Set("Content-Type", "application/json")
|
||||
r.Header.Set("X-Trace-Id", "abc-123")
|
||||
r.SetBody([]byte(`{"jsonrpc":"2.0"}`))
|
||||
},
|
||||
func(r *fasthttp.Request) { // multi-value + special chars
|
||||
r.Header.SetMethod("PUT")
|
||||
r.SetRequestURI("/x")
|
||||
r.Header.Add("X-Multi", "a")
|
||||
r.Header.Add("X-Multi", "b")
|
||||
r.Header.Set("X-Html", "<tag>&</tag>")
|
||||
r.SetBody([]byte("body-\x00\x01"))
|
||||
},
|
||||
func(r *fasthttp.Request) { // empty target edge case
|
||||
r.Header.SetMethod("OPTIONS")
|
||||
r.SetRequestURI("")
|
||||
r.Header.SetProtocol("HTTP/1.1")
|
||||
},
|
||||
}
|
||||
for i, f := range build {
|
||||
req := fasthttp.AcquireRequest()
|
||||
f(req)
|
||||
want := refMarshalRequest(req)
|
||||
got, err := MarshalRequest(req)
|
||||
if err != nil {
|
||||
t.Fatalf("case %d: MarshalRequest: %v", i, err)
|
||||
}
|
||||
if !bytes.Equal(got, want) {
|
||||
t.Errorf("case %d request frame mismatch\n got=%x\nwant=%x", i, got, want)
|
||||
}
|
||||
fasthttp.ReleaseRequest(req)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCodec_CrossCheckResponse(t *testing.T) {
|
||||
build := []func(*fasthttp.Response){
|
||||
func(r *fasthttp.Response) { r.SetStatusCode(200); r.Header.SetProtocol([]byte("HTTP/1.1")); r.Header.SetContentType("text/plain; charset=utf-8"); r.SetBodyString("ok") },
|
||||
func(r *fasthttp.Response) {
|
||||
r.SetStatusCode(200)
|
||||
r.Header.SetProtocol([]byte("HTTP/1.1"))
|
||||
r.Header.SetContentType("application/json")
|
||||
r.Header.Set("X-Trace-Id", "abc-123")
|
||||
r.SetBody([]byte(`{"result":"0x2a"}`))
|
||||
},
|
||||
func(r *fasthttp.Response) { // multi-value response header + status msg
|
||||
r.SetStatusCode(418)
|
||||
r.Header.SetContentType("application/octet-stream")
|
||||
r.Header.Add("X-Multi", "one")
|
||||
r.Header.Add("X-Multi", "two")
|
||||
r.SetBody([]byte("payload-\x00\x01-nuls"))
|
||||
},
|
||||
func(r *fasthttp.Response) { // trailers
|
||||
r.SetStatusCode(200)
|
||||
r.Header.SetProtocol([]byte("HTTP/1.1"))
|
||||
r.Header.SetContentType("application/octet-stream")
|
||||
_ = r.Header.AddTrailer("X-Checksum")
|
||||
r.Header.Set("X-Checksum", "deadbeef")
|
||||
r.SetBodyString("streamed")
|
||||
},
|
||||
}
|
||||
for i, f := range build {
|
||||
resp := fasthttp.AcquireResponse()
|
||||
f(resp)
|
||||
want := refMarshalResponse(resp)
|
||||
got, err := MarshalResponse(resp)
|
||||
if err != nil {
|
||||
t.Fatalf("case %d: MarshalResponse: %v", i, err)
|
||||
}
|
||||
if !bytes.Equal(got, want) {
|
||||
t.Errorf("case %d response frame mismatch\n got=%x\nwant=%x", i, got, want)
|
||||
}
|
||||
fasthttp.ReleaseResponse(resp)
|
||||
}
|
||||
}
|
||||
|
||||
// ---- benchmarks: allocs/op for the codec hot path ----
|
||||
|
||||
func benchRequest() *fasthttp.Request {
|
||||
req := fasthttp.AcquireRequest()
|
||||
req.Header.SetMethod("GET")
|
||||
req.SetRequestURI("/health")
|
||||
req.Header.SetHost("192.168.77.2:8391")
|
||||
return req
|
||||
}
|
||||
|
||||
func benchResponse() *fasthttp.Response {
|
||||
resp := fasthttp.AcquireResponse()
|
||||
resp.SetStatusCode(200)
|
||||
resp.Header.SetContentType("text/plain; charset=utf-8")
|
||||
resp.SetBodyString("ok")
|
||||
return resp
|
||||
}
|
||||
|
||||
// BenchmarkMarshalRequest is the dst==nil path (allocates the result slice).
|
||||
func BenchmarkMarshalRequest(b *testing.B) {
|
||||
req := benchRequest()
|
||||
defer fasthttp.ReleaseRequest(req)
|
||||
b.ReportAllocs()
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
if _, err := MarshalRequest(req); err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// BenchmarkAppendRequest is the reused-buffer path the server/client use — the
|
||||
// steady-state hot path. Should be zero-alloc.
|
||||
func BenchmarkAppendRequest(b *testing.B) {
|
||||
req := benchRequest()
|
||||
defer fasthttp.ReleaseRequest(req)
|
||||
buf := make([]byte, 0, 512)
|
||||
b.ReportAllocs()
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
var err error
|
||||
buf, err = AppendRequest(buf[:0], req)
|
||||
if err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func BenchmarkMarshalResponse(b *testing.B) {
|
||||
resp := benchResponse()
|
||||
defer fasthttp.ReleaseResponse(resp)
|
||||
b.ReportAllocs()
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
if _, err := MarshalResponse(resp); err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func BenchmarkAppendResponse(b *testing.B) {
|
||||
resp := benchResponse()
|
||||
defer fasthttp.ReleaseResponse(resp)
|
||||
buf := make([]byte, 0, 512)
|
||||
b.ReportAllocs()
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
var err error
|
||||
buf, err = AppendResponse(buf[:0], resp)
|
||||
if err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func BenchmarkUnmarshalRequest(b *testing.B) {
|
||||
req := benchRequest()
|
||||
frame, _ := MarshalRequest(req)
|
||||
fasthttp.ReleaseRequest(req)
|
||||
var dst fasthttp.Request
|
||||
b.ReportAllocs()
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
if err := UnmarshalRequest(frame, &dst); err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func BenchmarkUnmarshalResponse(b *testing.B) {
|
||||
resp := benchResponse()
|
||||
frame, _ := MarshalResponse(resp)
|
||||
fasthttp.ReleaseResponse(resp)
|
||||
var dst fasthttp.Response
|
||||
b.ReportAllocs()
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
if err := UnmarshalResponse(frame, &dst); err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -91,14 +91,28 @@ func (s *Server) Serve(ln net.Listener) error {
|
||||
|
||||
func (s *Server) serveConn(conn net.Conn) {
|
||||
defer conn.Close()
|
||||
br := bufio.NewReader(conn)
|
||||
remoteAddr := conn.RemoteAddr()
|
||||
br := bufio.NewReaderSize(conn, connReadBufSize)
|
||||
|
||||
// One RequestCtx for the whole connection. fasthttp's own HTTP server pools
|
||||
// RequestCtx across requests; the previous ZAP server heap-allocated one per
|
||||
// request (plus a fasthttp.Request), which dominated its GC pressure. Init2
|
||||
// binds the real conn (so RemoteAddr resolves) and keeps body buffers
|
||||
// attached for reuse; Request and Response are reset per request below.
|
||||
ctx := &fasthttp.RequestCtx{}
|
||||
ctx.Init2(conn, s.Logger, false)
|
||||
|
||||
// Per-connection scratch buffers: the inbound frame is decoded straight into
|
||||
// ctx.Request (which copies out what it keeps), and the response frame is
|
||||
// built here. Both grow once and are reused, so steady-state serving is
|
||||
// allocation-free.
|
||||
var readBuf, writeBuf []byte
|
||||
|
||||
for {
|
||||
if s.IdleTimeout > 0 {
|
||||
_ = conn.SetReadDeadline(time.Now().Add(s.IdleTimeout))
|
||||
}
|
||||
raw, err := readFrame(br)
|
||||
var err error
|
||||
readBuf, err = readFrameInto(br, readBuf)
|
||||
if err != nil {
|
||||
if !errors.Is(err, io.EOF) && !isClosedConn(err) {
|
||||
log.Printf("zaphttp: read frame: %v", err)
|
||||
@@ -109,17 +123,15 @@ func (s *Server) serveConn(conn net.Conn) {
|
||||
_ = conn.SetReadDeadline(time.Now().Add(s.ReadTimeout))
|
||||
}
|
||||
|
||||
var req fasthttp.Request
|
||||
if err := UnmarshalRequest(raw, &req); err != nil {
|
||||
// Decode into the connection's RequestCtx (UnmarshalRequest resets it
|
||||
// first), then clear the previous request's response and user values —
|
||||
// the reset fasthttp.Init would have done for a fresh ctx.
|
||||
if err := UnmarshalRequest(readBuf, &ctx.Request); err != nil {
|
||||
s.writeError(conn, fasthttp.StatusBadRequest, "malformed request frame: "+err.Error())
|
||||
return
|
||||
}
|
||||
|
||||
// Fresh ctx per request: fasthttp's Init resets connection state and
|
||||
// copies the request in, but does not reset a reused Response, so a
|
||||
// zero-value ctx is the clean, correct choice.
|
||||
ctx := &fasthttp.RequestCtx{}
|
||||
ctx.Init(&req, remoteAddr, s.Logger)
|
||||
ctx.Response.Reset()
|
||||
ctx.ResetUserValues()
|
||||
|
||||
// Defer-recover so a handler panic returns 500 rather than dropping
|
||||
// the connection silently.
|
||||
@@ -142,7 +154,7 @@ func (s *Server) serveConn(conn net.Conn) {
|
||||
// A streamed body (SSE, MCP notifications, chunked) is delivered as a
|
||||
// head frame + N data frames + an end frame, so the client sees each
|
||||
// chunk as the handler flushes it — real server→client push over ZAP.
|
||||
// A normal body takes the single-frame path (unchanged).
|
||||
// A normal body takes the single-frame path.
|
||||
if ctx.Response.IsBodyStream() {
|
||||
if err := s.streamResponse(conn, &ctx.Response); err != nil {
|
||||
if !isClosedConn(err) {
|
||||
@@ -153,12 +165,15 @@ func (s *Server) serveConn(conn net.Conn) {
|
||||
continue
|
||||
}
|
||||
|
||||
respBytes, err := MarshalResponse(&ctx.Response)
|
||||
// Build [4-byte length prefix][response frame] in one buffer and write
|
||||
// it with a single syscall.
|
||||
writeBuf = append(writeBuf[:0], 0, 0, 0, 0)
|
||||
writeBuf, err = AppendResponse(writeBuf, &ctx.Response)
|
||||
if err != nil {
|
||||
log.Printf("zaphttp: marshal response: %v", err)
|
||||
return
|
||||
}
|
||||
if err := writeFrame(conn, respBytes); err != nil {
|
||||
if err := writeFramePrefixed(conn, writeBuf); err != nil {
|
||||
if !isClosedConn(err) {
|
||||
log.Printf("zaphttp: write frame: %v", err)
|
||||
}
|
||||
@@ -167,6 +182,10 @@ func (s *Server) serveConn(conn net.Conn) {
|
||||
}
|
||||
}
|
||||
|
||||
// connReadBufSize is the per-connection bufio read buffer. Larger than bufio's
|
||||
// 4 KiB default so back-to-back small frames drain with fewer read syscalls.
|
||||
const connReadBufSize = 16 << 10
|
||||
|
||||
// streamResponse writes a streamed response: the head (status + headers), then
|
||||
// one data frame per flush of the body stream, then an end frame. BodyWriteTo
|
||||
// runs the fasthttp body-stream writer against chunkWriter, so each handler
|
||||
|
||||
+39
-8
@@ -25,28 +25,59 @@ import (
|
||||
// malicious peer announcing a multi-gigabyte length prefix.
|
||||
const MaxFrameSize = 64 << 20 // 64 MiB
|
||||
|
||||
// readFrame reads one length-prefixed frame from r. The returned
|
||||
// byte slice is owned by the caller; nil + io.EOF indicates the peer
|
||||
// closed cleanly between frames.
|
||||
// readFrame reads one length-prefixed frame from r into a freshly allocated
|
||||
// slice. nil + io.EOF indicates the peer closed cleanly between frames. Hot
|
||||
// paths use readFrameInto to reuse a buffer instead.
|
||||
func readFrame(r *bufio.Reader) ([]byte, error) {
|
||||
return readFrameInto(r, nil)
|
||||
}
|
||||
|
||||
// readFrameInto reads one length-prefixed frame into buf, growing it only when
|
||||
// the frame exceeds its capacity, and returns the frame slice (buf reused). The
|
||||
// caller reassigns its buffer from the return value so a grown buffer persists
|
||||
// across calls — the steady-state read is allocation-free. The frame's bytes
|
||||
// are consumed synchronously by Unmarshal{Request,Response} (which copy into the
|
||||
// fasthttp store), so the buffer is free to reuse on the next call.
|
||||
func readFrameInto(r *bufio.Reader, buf []byte) ([]byte, error) {
|
||||
var hdr [4]byte
|
||||
if _, err := io.ReadFull(r, hdr[:]); err != nil {
|
||||
return nil, err
|
||||
return buf, err
|
||||
}
|
||||
n := binary.BigEndian.Uint32(hdr[:])
|
||||
if n == 0 {
|
||||
return nil, fmt.Errorf("zaphttp: zero-length frame")
|
||||
return buf, fmt.Errorf("zaphttp: zero-length frame")
|
||||
}
|
||||
if n > MaxFrameSize {
|
||||
return nil, fmt.Errorf("zaphttp: frame size %d exceeds MaxFrameSize=%d", n, MaxFrameSize)
|
||||
return buf, fmt.Errorf("zaphttp: frame size %d exceeds MaxFrameSize=%d", n, MaxFrameSize)
|
||||
}
|
||||
if uint32(cap(buf)) < n {
|
||||
buf = make([]byte, n)
|
||||
} else {
|
||||
buf = buf[:n]
|
||||
}
|
||||
buf := make([]byte, n)
|
||||
if _, err := io.ReadFull(r, buf); err != nil {
|
||||
return nil, fmt.Errorf("zaphttp: short frame: %w", err)
|
||||
return buf, fmt.Errorf("zaphttp: short frame: %w", err)
|
||||
}
|
||||
return buf, nil
|
||||
}
|
||||
|
||||
// writeFramePrefixed writes a length prefix + frame in ONE Write. The frame
|
||||
// buffer must have been built with 4 leading bytes reserved for the prefix
|
||||
// (frame[:4] is the placeholder); this halves the per-message syscall count vs
|
||||
// writeFrame's separate header+body writes.
|
||||
func writeFramePrefixed(w io.Writer, framePlusPrefix []byte) error {
|
||||
body := len(framePlusPrefix) - 4
|
||||
if body <= 0 {
|
||||
return fmt.Errorf("zaphttp: refusing to write zero-length frame")
|
||||
}
|
||||
if uint64(body) > MaxFrameSize {
|
||||
return fmt.Errorf("zaphttp: frame size %d exceeds MaxFrameSize=%d", body, MaxFrameSize)
|
||||
}
|
||||
binary.BigEndian.PutUint32(framePlusPrefix[0:4], uint32(body))
|
||||
_, err := w.Write(framePlusPrefix)
|
||||
return err
|
||||
}
|
||||
|
||||
// writeFrame writes one length-prefixed frame to w.
|
||||
func writeFrame(w io.Writer, frame []byte) error {
|
||||
if len(frame) == 0 {
|
||||
|
||||
@@ -8,19 +8,17 @@
|
||||
// ZAP schema language). The bytes on the wire are github.com/zap-proto/go
|
||||
// Objects — the canonical pure-stdlib ZAP runtime — packed at explicit
|
||||
// byte offsets within each object's fixed payload. There is no external
|
||||
// codec dependency: this file is the encoder and decoder.
|
||||
// codec dependency: codec.go is the encoder and decoder.
|
||||
//
|
||||
// # Wire compatibility
|
||||
//
|
||||
// The frame format is UNCHANGED from the prior net/http codec — same
|
||||
// object slots, same offsets, same JSON header map, same flags
|
||||
// discriminator. Only the Go type the codec reads from and writes to
|
||||
// changed (net/http → fasthttp). Frames produced here are byte-identical
|
||||
// to the frames the previous net/http codec produced for the same logical
|
||||
// message, so a fasthttp-based peer and a net/http-based peer interoperate
|
||||
// on the wire. This is verified by golden-byte tests (zaphttp_test.go)
|
||||
// that decode frames captured from the previous codec and re-encode them
|
||||
// to the same bytes.
|
||||
// discriminator. This file's marshal/unmarshal entry points now drive the
|
||||
// zero-allocation codec in codec.go, which emits and parses the SAME bytes
|
||||
// the zap.Builder path did (verified byte-for-byte by the golden-wire tests
|
||||
// and TestCodec_CrossCheck). A fasthttp peer and a net/http peer still
|
||||
// interoperate on the wire.
|
||||
//
|
||||
// What's covered:
|
||||
// - Request method, target, headers, body
|
||||
@@ -28,50 +26,33 @@
|
||||
// - Trailers (read and write)
|
||||
//
|
||||
// What's not yet covered:
|
||||
// - Streaming bodies — the entire body must fit in one frame; the
|
||||
// transport negotiates a max frame size on connection setup.
|
||||
// - Streaming bodies larger than one frame use the head/data/end frames
|
||||
// below; a non-streamed body must fit in one frame (transport-negotiated
|
||||
// max frame size).
|
||||
// - WebSocket / SSE upgrade paths; those land in zap-proto/ws.
|
||||
//
|
||||
// # Wire layout
|
||||
//
|
||||
// A frame is one zap-proto/go message. The message type rides in the
|
||||
// header flags (flags>>8): FrameRequest for a request, FrameResponse for
|
||||
// a response. There is no union struct; the flag selects which object
|
||||
// shape is the message root.
|
||||
//
|
||||
// # Offset discipline
|
||||
//
|
||||
// zap-proto/go's ObjectBuilder.SetBytes / SetText write an 8-byte slot at
|
||||
// the field offset: {relOffset uint32, length uint32}. The variable-length
|
||||
// tail is appended after the fixed section in Finish() and the relOffset
|
||||
// is patched to point at it. Consequently EVERY text or bytes field
|
||||
// consumes 8 bytes of fixed payload — adjacent text fields MUST be spaced
|
||||
// 8 bytes apart or their slots overlap and corrupt each other. Fixed
|
||||
// scalars use their natural width on a natural boundary. These offsets are
|
||||
// the contract; changing one is a wire break.
|
||||
// A frame is one zap-proto/go message. The message type rides in the header
|
||||
// flags (flags>>8): FrameRequest for a request, FrameResponse for a response,
|
||||
// and the streaming variants below. Every text/bytes field is an 8-byte
|
||||
// {relOffset,length} slot in the object's fixed section; non-empty fields'
|
||||
// bytes are appended to the variable tail in field-declaration order. See
|
||||
// codec.go for the encoder/decoder and the offset discipline that keeps the
|
||||
// bytes byte-identical.
|
||||
//
|
||||
// # Framing metadata vs headers
|
||||
//
|
||||
// Three fields are owned by the frame, not the headers slot, exactly as
|
||||
// the reference net/http codec treated them:
|
||||
//
|
||||
// - Content-Length — the frame length-prefixes the body, so body length
|
||||
// is authoritative; it is reconstructed on decode, never carried in
|
||||
// the headers map.
|
||||
// - Host — request routing metadata; carried in fasthttp's Host slot,
|
||||
// not the headers map (matching net/http, which keeps Host out of
|
||||
// Header).
|
||||
// - Trailer (the meta-header listing trailer names) — represented by the
|
||||
// dedicated trailer slot, not the headers map.
|
||||
//
|
||||
// Everything else fasthttp exposes (Content-Type, User-Agent, Cookie,
|
||||
// Authorization, X-*, Set-Cookie, …) rides in the headers slot as a JSON
|
||||
// map[string][]string — the native lossless shape for multi-value header
|
||||
// names (RFC 9110 §5.2).
|
||||
// Three fields are owned by the frame, not the headers slot: Content-Length
|
||||
// (the frame length-prefixes the body), Host (fasthttp's Host slot), and
|
||||
// Trailer (the dedicated trailer slot). Everything else rides in the headers
|
||||
// slot as a JSON map[string][]string — the native lossless shape for
|
||||
// multi-value header names (RFC 9110 §5.2).
|
||||
|
||||
package zaphttp
|
||||
|
||||
import (
|
||||
"encoding/binary"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/textproto"
|
||||
@@ -81,10 +62,10 @@ import (
|
||||
zap "github.com/zap-proto/go"
|
||||
)
|
||||
|
||||
// Frame type IDs. The ZAP message header carries the type in the upper
|
||||
// byte of the 16-bit flags field; FinishWithFlags(t<<8) tags the message
|
||||
// and Message.Flags()>>8 recovers it. A request frame and a response frame
|
||||
// are the two shapes the wire carries today.
|
||||
// Frame type IDs. The ZAP message header carries the type in the upper byte
|
||||
// of the 16-bit flags field; the encoder tags the message with type<<8 and
|
||||
// flags>>8 recovers it. A request frame and a response frame are the two
|
||||
// shapes the non-streaming wire carries.
|
||||
const (
|
||||
FrameRequest uint16 = 0x01
|
||||
FrameResponse uint16 = 0x02
|
||||
@@ -105,103 +86,9 @@ const (
|
||||
chunkSlotSize = 8
|
||||
)
|
||||
|
||||
// FrameTypeOf peeks a frame's type from its ZAP header without decoding the
|
||||
// body, so a reader can dispatch (response vs streamed head vs data vs end).
|
||||
func FrameTypeOf(frame []byte) (uint16, error) {
|
||||
msg, err := zap.Parse(frame)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
return msg.Flags() >> 8, nil
|
||||
}
|
||||
|
||||
// MarshalResponseHead serializes a response's status + headers (NO body) as the
|
||||
// opening frame of a stream. The body is delivered by subsequent FrameData
|
||||
// frames.
|
||||
func MarshalResponseHead(resp *fasthttp.Response) ([]byte, error) {
|
||||
status := resp.StatusCode()
|
||||
if status == 0 {
|
||||
status = fasthttp.StatusOK
|
||||
}
|
||||
reason := string(resp.Header.StatusMessage())
|
||||
if reason == "" {
|
||||
reason = fasthttp.StatusMessage(status)
|
||||
}
|
||||
proto := string(resp.Header.Protocol())
|
||||
headers, err := encodeHeaderMap(collectResponseHeaders(&resp.Header))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("zaphttp: encode headers: %w", err)
|
||||
}
|
||||
|
||||
b := zap.NewBuilder(respSlotSize + len(headers) + len(reason) + len(proto) + 64)
|
||||
ob := b.StartObject(respSlotSize)
|
||||
ob.SetUint16(respStatus, uint16(status))
|
||||
ob.SetText(respReason, reason)
|
||||
ob.SetText(respProto, orDefault(proto, defaultProto))
|
||||
ob.SetBytes(respHeaders, headers)
|
||||
ob.SetBytes(respBody, nil)
|
||||
ob.SetBytes(respTrailer, nil)
|
||||
ob.FinishAsRoot()
|
||||
return b.FinishWithFlags(FrameResponseHead << 8), nil
|
||||
}
|
||||
|
||||
// UnmarshalResponseHead applies a streamed response's status + headers to dst.
|
||||
// The body is left empty; the caller attaches a streaming body that reads the
|
||||
// following FrameData frames.
|
||||
func UnmarshalResponseHead(frame []byte, dst *fasthttp.Response) error {
|
||||
msg, err := zap.Parse(frame)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if t := msg.Flags() >> 8; t != FrameResponseHead {
|
||||
return fmt.Errorf("zaphttp: frame type %#x, want response-head (%#x)", t, FrameResponseHead)
|
||||
}
|
||||
r := msg.Root()
|
||||
dst.Reset()
|
||||
dst.SetStatusCode(int(r.Uint16(respStatus)))
|
||||
if reason := r.Text(respReason); reason != "" {
|
||||
dst.Header.SetStatusMessage([]byte(reason))
|
||||
}
|
||||
if proto := r.Text(respProto); proto != "" {
|
||||
dst.Header.SetProtocol([]byte(proto))
|
||||
}
|
||||
return applyHeadersToResponse(&dst.Header, r.Bytes(respHeaders))
|
||||
}
|
||||
|
||||
// MarshalData wraps one body chunk as a FrameData frame.
|
||||
func MarshalData(chunk []byte) []byte {
|
||||
b := zap.NewBuilder(chunkSlotSize + len(chunk) + 32)
|
||||
ob := b.StartObject(chunkSlotSize)
|
||||
ob.SetBytes(chunkBytes, chunk)
|
||||
ob.FinishAsRoot()
|
||||
return b.FinishWithFlags(FrameData << 8)
|
||||
}
|
||||
|
||||
// DataChunkOf extracts the chunk bytes from a FrameData frame. The returned
|
||||
// slice aliases the frame buffer; copy it to retain past the frame's lifetime.
|
||||
func DataChunkOf(frame []byte) ([]byte, error) {
|
||||
msg, err := zap.Parse(frame)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if t := msg.Flags() >> 8; t != FrameData {
|
||||
return nil, fmt.Errorf("zaphttp: frame type %#x, want data (%#x)", t, FrameData)
|
||||
}
|
||||
return msg.Root().Bytes(chunkBytes), nil
|
||||
}
|
||||
|
||||
// MarshalEnd terminates a stream, carrying optional encoded trailers.
|
||||
func MarshalEnd(trailer []byte) []byte {
|
||||
b := zap.NewBuilder(chunkSlotSize + len(trailer) + 32)
|
||||
ob := b.StartObject(chunkSlotSize)
|
||||
ob.SetBytes(chunkBytes, trailer)
|
||||
ob.FinishAsRoot()
|
||||
return b.FinishWithFlags(FrameEnd << 8)
|
||||
}
|
||||
|
||||
// Request frame field offsets within the object's fixed payload. Every
|
||||
// text/bytes slot is 8 bytes (relOffset+length); see "Offset discipline"
|
||||
// in the package doc.
|
||||
// text/bytes slot is 8 bytes (relOffset+length); see codec.go's offset
|
||||
// discipline.
|
||||
const (
|
||||
reqMethod = 0 // text (8)
|
||||
reqTarget = 8 // text (8)
|
||||
@@ -212,8 +99,8 @@ const (
|
||||
reqSlotSize = 48
|
||||
)
|
||||
|
||||
// Response frame field offsets. status@0 is a uint16 (2 bytes); the first
|
||||
// text slot starts at the next 8-byte boundary.
|
||||
// Response frame field offsets. status@0 is a uint16 (2 bytes); the first text
|
||||
// slot starts at the next 8-byte boundary.
|
||||
const (
|
||||
respStatus = 0 // uint16 (2)
|
||||
respReason = 8 // text (8)
|
||||
@@ -225,69 +112,62 @@ const (
|
||||
)
|
||||
|
||||
// defaultProto is the proto string written when the source carries none.
|
||||
// fasthttp always reports a protocol ("HTTP/1.1" by default), so this
|
||||
// only guards a hand-built zero-value header.
|
||||
// fasthttp always reports a protocol ("HTTP/1.1" by default), so this only
|
||||
// guards a hand-built zero-value header.
|
||||
const defaultProto = "ZAP-HTTP/1.0"
|
||||
|
||||
// MarshalRequest serializes a *fasthttp.Request into a ZAP frame. The
|
||||
// returned bytes are the ZAP message; the transport layer prepends the
|
||||
// length prefix (see transport.go).
|
||||
// FrameTypeOf peeks a frame's type from its ZAP header without decoding the
|
||||
// body, so a reader can dispatch (response vs streamed head vs data vs end).
|
||||
func FrameTypeOf(frame []byte) (uint16, error) {
|
||||
if len(frame) < zap.HeaderSize {
|
||||
return 0, zap.ErrBufferTooSmall
|
||||
}
|
||||
if string(frame[0:4]) != zap.Magic {
|
||||
return 0, zap.ErrInvalidMagic
|
||||
}
|
||||
return binary.LittleEndian.Uint16(frame[6:8]) >> 8, nil
|
||||
}
|
||||
|
||||
// MarshalRequest serializes a *fasthttp.Request into a ZAP frame. The returned
|
||||
// bytes are the ZAP message; the transport layer prepends the length prefix
|
||||
// (see transport.go). It is the dst==nil convenience for AppendRequest — hot
|
||||
// callers append into a reused buffer to stay allocation-free.
|
||||
func MarshalRequest(req *fasthttp.Request) ([]byte, error) {
|
||||
method := string(req.Header.Method())
|
||||
target := string(req.RequestURI())
|
||||
proto := string(req.Header.Protocol())
|
||||
body := req.Body()
|
||||
return AppendRequest(make([]byte, 0, 128+len(req.Body())), req)
|
||||
}
|
||||
|
||||
headers, err := encodeHeaderMap(collectRequestHeaders(&req.Header))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("zaphttp: encode headers: %w", err)
|
||||
}
|
||||
trailer, err := encodeHeaderMap(collectTrailers(&req.Header))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("zaphttp: encode trailer: %w", err)
|
||||
}
|
||||
|
||||
b := zap.NewBuilder(reqSlotSize + len(body) + len(headers) + len(trailer) +
|
||||
len(method) + len(target) + len(proto) + 64)
|
||||
ob := b.StartObject(reqSlotSize)
|
||||
ob.SetText(reqMethod, orDefault(method, fasthttp.MethodGet))
|
||||
ob.SetText(reqTarget, target)
|
||||
ob.SetText(reqProto, orDefault(proto, defaultProto))
|
||||
ob.SetBytes(reqHeaders, headers)
|
||||
ob.SetBytes(reqBody, body)
|
||||
ob.SetBytes(reqTrailer, trailer)
|
||||
ob.FinishAsRoot()
|
||||
return b.FinishWithFlags(FrameRequest << 8), nil
|
||||
// MarshalResponse serializes a *fasthttp.Response into a ZAP frame. See
|
||||
// MarshalRequest; AppendResponse is the reused-buffer form.
|
||||
func MarshalResponse(resp *fasthttp.Response) ([]byte, error) {
|
||||
return AppendResponse(make([]byte, 0, 128+len(resp.Body())), resp)
|
||||
}
|
||||
|
||||
// UnmarshalRequest reconstructs a request into dst from a ZAP frame. dst is
|
||||
// reset first, then populated: method, request-URI, protocol, headers,
|
||||
// trailers, and body. Content-Length is derived from the body length by
|
||||
// SetBody; Host is taken from a Host header if the frame carried one.
|
||||
// SetBody; Host comes from a Host header if the frame carried one.
|
||||
func UnmarshalRequest(frame []byte, dst *fasthttp.Request) error {
|
||||
msg, err := zap.Parse(frame)
|
||||
root, size, err := frameRoot(frame, FrameRequest)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if t := msg.Flags() >> 8; t != FrameRequest {
|
||||
return fmt.Errorf("zaphttp: frame type %#x, want request (%#x)", t, FrameRequest)
|
||||
}
|
||||
r := msg.Root()
|
||||
|
||||
method := r.Text(reqMethod)
|
||||
target := r.Text(reqTarget)
|
||||
proto := r.Text(reqProto)
|
||||
headers := r.Bytes(reqHeaders)
|
||||
body := r.Bytes(reqBody)
|
||||
trailer := r.Bytes(reqTrailer)
|
||||
method := readVar(frame, size, root, reqMethod)
|
||||
target := readVar(frame, size, root, reqTarget)
|
||||
proto := readVar(frame, size, root, reqProto)
|
||||
headers := readVar(frame, size, root, reqHeaders)
|
||||
body := readVar(frame, size, root, reqBody)
|
||||
trailer := readVar(frame, size, root, reqTrailer)
|
||||
|
||||
dst.Reset()
|
||||
dst.Header.SetMethod(orDefault(method, fasthttp.MethodGet))
|
||||
dst.SetRequestURI(target)
|
||||
if proto != "" {
|
||||
dst.Header.SetProtocol(proto)
|
||||
if len(method) == 0 {
|
||||
method = bMethodGET
|
||||
}
|
||||
if err := applyHeadersToRequest(&dst.Header, headers); err != nil {
|
||||
dst.Header.SetMethodBytes(method)
|
||||
dst.SetRequestURIBytes(target)
|
||||
if len(proto) != 0 {
|
||||
dst.Header.SetProtocolBytes(proto)
|
||||
}
|
||||
if err := decodeRequestHeaders(&dst.Header, headers); err != nil {
|
||||
return fmt.Errorf("zaphttp: decode headers: %w", err)
|
||||
}
|
||||
if err := applyTrailers(&dst.Header, trailer); err != nil {
|
||||
@@ -297,68 +177,28 @@ func UnmarshalRequest(frame []byte, dst *fasthttp.Request) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
// MarshalResponse serializes a *fasthttp.Response into a ZAP frame.
|
||||
func MarshalResponse(resp *fasthttp.Response) ([]byte, error) {
|
||||
status := resp.StatusCode()
|
||||
if status == 0 {
|
||||
status = fasthttp.StatusOK
|
||||
}
|
||||
reason := string(resp.Header.StatusMessage())
|
||||
if reason == "" {
|
||||
reason = fasthttp.StatusMessage(status)
|
||||
}
|
||||
proto := string(resp.Header.Protocol())
|
||||
body := resp.Body()
|
||||
|
||||
headers, err := encodeHeaderMap(collectResponseHeaders(&resp.Header))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("zaphttp: encode headers: %w", err)
|
||||
}
|
||||
trailer, err := encodeHeaderMap(collectTrailers(&resp.Header))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("zaphttp: encode trailer: %w", err)
|
||||
}
|
||||
|
||||
b := zap.NewBuilder(respSlotSize + len(body) + len(headers) + len(trailer) +
|
||||
len(reason) + len(proto) + 64)
|
||||
ob := b.StartObject(respSlotSize)
|
||||
ob.SetUint16(respStatus, uint16(status))
|
||||
ob.SetText(respReason, reason)
|
||||
ob.SetText(respProto, orDefault(proto, defaultProto))
|
||||
ob.SetBytes(respHeaders, headers)
|
||||
ob.SetBytes(respBody, body)
|
||||
ob.SetBytes(respTrailer, trailer)
|
||||
ob.FinishAsRoot()
|
||||
return b.FinishWithFlags(FrameResponse << 8), nil
|
||||
}
|
||||
|
||||
// UnmarshalResponse reconstructs a response into dst from a ZAP frame.
|
||||
func UnmarshalResponse(frame []byte, dst *fasthttp.Response) error {
|
||||
msg, err := zap.Parse(frame)
|
||||
root, size, err := frameRoot(frame, FrameResponse)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if t := msg.Flags() >> 8; t != FrameResponse {
|
||||
return fmt.Errorf("zaphttp: frame type %#x, want response (%#x)", t, FrameResponse)
|
||||
}
|
||||
r := msg.Root()
|
||||
|
||||
status := int(r.Uint16(respStatus))
|
||||
reason := r.Text(respReason)
|
||||
proto := r.Text(respProto)
|
||||
headers := r.Bytes(respHeaders)
|
||||
body := r.Bytes(respBody)
|
||||
trailer := r.Bytes(respTrailer)
|
||||
status := int(readU16(frame, size, root, respStatus))
|
||||
reason := readVar(frame, size, root, respReason)
|
||||
proto := readVar(frame, size, root, respProto)
|
||||
headers := readVar(frame, size, root, respHeaders)
|
||||
body := readVar(frame, size, root, respBody)
|
||||
trailer := readVar(frame, size, root, respTrailer)
|
||||
|
||||
dst.Reset()
|
||||
dst.SetStatusCode(status)
|
||||
if reason != "" {
|
||||
dst.Header.SetStatusMessage([]byte(reason))
|
||||
if len(reason) != 0 {
|
||||
dst.Header.SetStatusMessage(reason)
|
||||
}
|
||||
if proto != "" {
|
||||
dst.Header.SetProtocol([]byte(proto))
|
||||
if len(proto) != 0 {
|
||||
dst.Header.SetProtocol(proto)
|
||||
}
|
||||
if err := applyHeadersToResponse(&dst.Header, headers); err != nil {
|
||||
if err := decodeResponseHeaders(&dst.Header, headers); err != nil {
|
||||
return fmt.Errorf("zaphttp: decode headers: %w", err)
|
||||
}
|
||||
if err := applyTrailers(&dst.Header, trailer); err != nil {
|
||||
@@ -368,44 +208,197 @@ func UnmarshalResponse(frame []byte, dst *fasthttp.Response) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
// ---- header extraction (fasthttp -> map) ----
|
||||
// ---- streaming frames ----
|
||||
|
||||
// collectRequestHeaders walks a request's headers into a map[string][]string,
|
||||
// dropping the frame-owned fields (Host, Content-Length, the Trailer
|
||||
// meta-header, and any declared-trailer values which ride in the trailer
|
||||
// slot instead).
|
||||
func collectRequestHeaders(h *fasthttp.RequestHeader) map[string][]string {
|
||||
skip := trailerKeySet(h)
|
||||
m := map[string][]string{}
|
||||
h.VisitAll(func(key, value []byte) {
|
||||
k := string(key)
|
||||
if isFrameOwnedHeader(k) || skip[textproto.CanonicalMIMEHeaderKey(k)] {
|
||||
return
|
||||
}
|
||||
m[k] = append(m[k], string(value))
|
||||
})
|
||||
return m
|
||||
// MarshalResponseHead serializes a response's status + headers (NO body) as the
|
||||
// opening frame of a stream. The body is delivered by subsequent FrameData
|
||||
// frames.
|
||||
func MarshalResponseHead(resp *fasthttp.Response) ([]byte, error) {
|
||||
dst := appendFrameHeader(nil, FrameResponseHead<<8)
|
||||
obj := len(dst)
|
||||
dst = append(dst, zeroFixed[:respSlotSize]...)
|
||||
|
||||
status := resp.StatusCode()
|
||||
if status == 0 {
|
||||
status = fasthttp.StatusOK
|
||||
}
|
||||
binary.LittleEndian.PutUint16(dst[obj+respStatus:], uint16(status))
|
||||
|
||||
reason := resp.Header.StatusMessage()
|
||||
if len(reason) == 0 {
|
||||
reason = unsafeBytes(fasthttp.StatusMessage(status))
|
||||
}
|
||||
proto := resp.Header.Protocol()
|
||||
if len(proto) == 0 {
|
||||
proto = bDefaultProto
|
||||
}
|
||||
dst = putVarBytes(dst, obj+respReason, reason)
|
||||
dst = putVarBytes(dst, obj+respProto, proto)
|
||||
|
||||
hStart := len(dst)
|
||||
dst = appendResponseHeadersJSON(dst, &resp.Header, newTrailerSkip(&resp.Header))
|
||||
dst = patchVarSlot(dst, obj+respHeaders, hStart)
|
||||
// body and trailer slots stay null.
|
||||
|
||||
binary.LittleEndian.PutUint32(dst[12:], uint32(len(dst)))
|
||||
return dst, nil
|
||||
}
|
||||
|
||||
// collectResponseHeaders is the response counterpart of
|
||||
// collectRequestHeaders. Responses have no Host, but the same
|
||||
// Content-Length / Trailer-meta / declared-trailer exclusions apply.
|
||||
func collectResponseHeaders(h *fasthttp.ResponseHeader) map[string][]string {
|
||||
skip := trailerKeySet(h)
|
||||
m := map[string][]string{}
|
||||
h.VisitAll(func(key, value []byte) {
|
||||
k := string(key)
|
||||
if isFrameOwnedHeader(k) || skip[textproto.CanonicalMIMEHeaderKey(k)] {
|
||||
return
|
||||
}
|
||||
m[k] = append(m[k], string(value))
|
||||
})
|
||||
return m
|
||||
// UnmarshalResponseHead applies a streamed response's status + headers to dst.
|
||||
// The body is left empty; the caller attaches a streaming body that reads the
|
||||
// following FrameData frames.
|
||||
func UnmarshalResponseHead(frame []byte, dst *fasthttp.Response) error {
|
||||
root, size, err := frameRoot(frame, FrameResponseHead)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
dst.Reset()
|
||||
dst.SetStatusCode(int(readU16(frame, size, root, respStatus)))
|
||||
if reason := readVar(frame, size, root, respReason); len(reason) != 0 {
|
||||
dst.Header.SetStatusMessage(reason)
|
||||
}
|
||||
if proto := readVar(frame, size, root, respProto); len(proto) != 0 {
|
||||
dst.Header.SetProtocol(proto)
|
||||
}
|
||||
return decodeResponseHeaders(&dst.Header, readVar(frame, size, root, respHeaders))
|
||||
}
|
||||
|
||||
// collectTrailers reads declared trailer names and their values (stored in
|
||||
// the general header store) into a map[string][]string. Works for both
|
||||
// request and response headers via the trailerHeader interface.
|
||||
// MarshalData wraps one body chunk as a FrameData frame.
|
||||
func MarshalData(chunk []byte) []byte {
|
||||
return marshalChunk(FrameData, chunk)
|
||||
}
|
||||
|
||||
// MarshalEnd terminates a stream, carrying optional encoded trailers.
|
||||
func MarshalEnd(trailer []byte) []byte {
|
||||
return marshalChunk(FrameEnd, trailer)
|
||||
}
|
||||
|
||||
// marshalChunk builds a single-bytes-field frame (data or end).
|
||||
func marshalChunk(frameType uint16, data []byte) []byte {
|
||||
dst := appendFrameHeader(nil, frameType<<8)
|
||||
obj := len(dst)
|
||||
dst = append(dst, zeroFixed[:chunkSlotSize]...)
|
||||
dst = putVarBytes(dst, obj+chunkBytes, data)
|
||||
binary.LittleEndian.PutUint32(dst[12:], uint32(len(dst)))
|
||||
return dst
|
||||
}
|
||||
|
||||
// DataChunkOf extracts the chunk bytes from a FrameData frame. The returned
|
||||
// slice aliases the frame buffer; copy it to retain past the frame's lifetime.
|
||||
func DataChunkOf(frame []byte) ([]byte, error) {
|
||||
root, size, err := frameRoot(frame, FrameData)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return readVar(frame, size, root, chunkBytes), nil
|
||||
}
|
||||
|
||||
// ---- header decode (JSON -> fasthttp) ----
|
||||
|
||||
// decodeRequestHeaders applies a request's headers JSON to h. The fast path in
|
||||
// forEachHeaderKV handles the common compact, escape-free case with zero
|
||||
// allocation; anything else falls back to the standard json decoder. Host and
|
||||
// Content-Length are frame-owned (Host reconstructed here, Content-Length set
|
||||
// by SetBody), matching the reference codec.
|
||||
func decodeRequestHeaders(h *fasthttp.RequestHeader, raw []byte) error {
|
||||
if forEachHeaderKV(raw, func(key, val []byte) {
|
||||
switch {
|
||||
case bytesEqualFold(key, strHost):
|
||||
h.SetHostBytes(val) // last value wins
|
||||
case bytesEqualFold(key, strContentLength):
|
||||
// frame-owned; SetBody sets the authoritative length.
|
||||
default:
|
||||
h.AddBytesKV(key, val)
|
||||
}
|
||||
}) {
|
||||
return nil
|
||||
}
|
||||
return decodeRequestHeadersSlow(h, raw)
|
||||
}
|
||||
|
||||
func decodeRequestHeadersSlow(h *fasthttp.RequestHeader, raw []byte) error {
|
||||
m, err := decodeHeaderMap(raw)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
for k, vals := range m {
|
||||
switch {
|
||||
case strings.EqualFold(k, fasthttp.HeaderHost):
|
||||
if len(vals) > 0 {
|
||||
h.SetHost(vals[len(vals)-1])
|
||||
}
|
||||
case strings.EqualFold(k, fasthttp.HeaderContentLength):
|
||||
// frame-owned
|
||||
default:
|
||||
for _, v := range vals {
|
||||
h.Add(k, v)
|
||||
}
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func decodeResponseHeaders(h *fasthttp.ResponseHeader, raw []byte) error {
|
||||
if forEachHeaderKV(raw, func(key, val []byte) {
|
||||
if bytesEqualFold(key, strContentLength) {
|
||||
return // frame-owned
|
||||
}
|
||||
h.AddBytesKV(key, val)
|
||||
}) {
|
||||
return nil
|
||||
}
|
||||
return decodeResponseHeadersSlow(h, raw)
|
||||
}
|
||||
|
||||
func decodeResponseHeadersSlow(h *fasthttp.ResponseHeader, raw []byte) error {
|
||||
m, err := decodeHeaderMap(raw)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
for k, vals := range m {
|
||||
if strings.EqualFold(k, fasthttp.HeaderContentLength) {
|
||||
continue // frame-owned
|
||||
}
|
||||
for _, v := range vals {
|
||||
h.Add(k, v)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func decodeHeaderMap(b []byte) (map[string][]string, error) {
|
||||
if len(b) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
m := map[string][]string{}
|
||||
if err := json.Unmarshal(b, &m); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return m, nil
|
||||
}
|
||||
|
||||
// ---- trailers ----
|
||||
|
||||
// encodeTrailers serializes declared trailers (rare) as JSON for the frame's
|
||||
// trailer slot. Empty (the common case) returns nil — a null slot — with no
|
||||
// allocation.
|
||||
func encodeTrailers(h trailerHeader) ([]byte, error) {
|
||||
return encodeHeaderMap(collectTrailers(h))
|
||||
}
|
||||
|
||||
// encodeHeaderMap serializes a header map as JSON. An empty map encodes as nil
|
||||
// (the null slot), which decode maps back to no headers — byte-identical to the
|
||||
// reference codec's empty path.
|
||||
func encodeHeaderMap(m map[string][]string) ([]byte, error) {
|
||||
if len(m) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
return json.Marshal(m)
|
||||
}
|
||||
|
||||
// collectTrailers reads declared trailer names and their values (stored in the
|
||||
// general header store) into a map[string][]string. Works for both request and
|
||||
// response headers via the trailerHeader interface.
|
||||
func collectTrailers(h trailerHeader) map[string][]string {
|
||||
keys := h.PeekTrailerKeys()
|
||||
if len(keys) == 0 {
|
||||
@@ -428,103 +421,10 @@ type trailerHeader interface {
|
||||
PeekAll(key string) [][]byte
|
||||
}
|
||||
|
||||
// trailerKeySet returns the canonicalized set of declared trailer names so
|
||||
// their values can be excluded from the headers slot (they belong in the
|
||||
// trailer slot).
|
||||
func trailerKeySet(h trailerHeader) map[string]bool {
|
||||
keys := h.PeekTrailerKeys()
|
||||
if len(keys) == 0 {
|
||||
return nil
|
||||
}
|
||||
set := make(map[string]bool, len(keys))
|
||||
for _, k := range keys {
|
||||
set[textproto.CanonicalMIMEHeaderKey(string(k))] = true
|
||||
}
|
||||
return set
|
||||
}
|
||||
|
||||
// isFrameOwnedHeader reports whether a header name is owned by the ZAP
|
||||
// frame rather than the headers map: Host and Content-Length are
|
||||
// reconstructed at the boundary, and "Trailer" is represented by the
|
||||
// trailer slot.
|
||||
func isFrameOwnedHeader(key string) bool {
|
||||
return strings.EqualFold(key, fasthttp.HeaderHost) ||
|
||||
strings.EqualFold(key, fasthttp.HeaderContentLength) ||
|
||||
strings.EqualFold(key, fasthttp.HeaderTrailer)
|
||||
}
|
||||
|
||||
// encodeHeaderMap serializes a header map as JSON. An empty map encodes as
|
||||
// nil (the SetBytes null slot), which the decode side maps back to no
|
||||
// headers — byte-identical to the reference codec's empty-header path.
|
||||
func encodeHeaderMap(m map[string][]string) ([]byte, error) {
|
||||
if len(m) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
return json.Marshal(m)
|
||||
}
|
||||
|
||||
// ---- header application (map -> fasthttp) ----
|
||||
|
||||
func decodeHeaderMap(b []byte) (map[string][]string, error) {
|
||||
if len(b) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
m := map[string][]string{}
|
||||
if err := json.Unmarshal(b, &m); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return m, nil
|
||||
}
|
||||
|
||||
func applyHeadersToRequest(h *fasthttp.RequestHeader, raw []byte) error {
|
||||
m, err := decodeHeaderMap(raw)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
for k, vals := range m {
|
||||
switch {
|
||||
case strings.EqualFold(k, fasthttp.HeaderHost):
|
||||
if len(vals) > 0 {
|
||||
h.SetHost(vals[len(vals)-1])
|
||||
}
|
||||
case strings.EqualFold(k, fasthttp.HeaderContentLength):
|
||||
// frame-owned; SetBody sets the authoritative length.
|
||||
default:
|
||||
for _, v := range vals {
|
||||
h.Add(k, v)
|
||||
}
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func applyHeadersToResponse(h *fasthttp.ResponseHeader, raw []byte) error {
|
||||
m, err := decodeHeaderMap(raw)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
for k, vals := range m {
|
||||
if strings.EqualFold(k, fasthttp.HeaderContentLength) {
|
||||
continue // frame-owned
|
||||
}
|
||||
for _, v := range vals {
|
||||
h.Add(k, v)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// trailerSetter is the shared subset of *fasthttp.RequestHeader and
|
||||
// *fasthttp.ResponseHeader used to reapply trailers.
|
||||
type trailerSetter interface {
|
||||
AddTrailer(trailer string) error
|
||||
Add(key, value string)
|
||||
}
|
||||
|
||||
// applyTrailers declares each trailer name and re-adds its value(s). A name
|
||||
// fasthttp forbids as a trailer (framing/routing/auth field) is carried as
|
||||
// an ordinary header so no data is dropped — lenient by design, since the
|
||||
// frame's trailer slot is opaque JSON.
|
||||
// fasthttp forbids as a trailer (framing/routing/auth field) is carried as an
|
||||
// ordinary header so no data is dropped — lenient by design, since the frame's
|
||||
// trailer slot is opaque JSON.
|
||||
func applyTrailers(h trailerSetter, raw []byte) error {
|
||||
m, err := decodeHeaderMap(raw)
|
||||
if err != nil {
|
||||
@@ -539,9 +439,9 @@ func applyTrailers(h trailerSetter, raw []byte) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func orDefault(s, fallback string) string {
|
||||
if s == "" {
|
||||
return fallback
|
||||
}
|
||||
return s
|
||||
// trailerSetter is the shared subset of *fasthttp.RequestHeader and
|
||||
// *fasthttp.ResponseHeader used to reapply trailers.
|
||||
type trailerSetter interface {
|
||||
AddTrailer(trailer string) error
|
||||
Add(key, value string)
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user