qpack.go

package http3

import (
	"bytes"
	"io"
	"sync/atomic"

	"github.com/quic-go/qpack"
	"gnalloy.org/gnalloy/buffer"
	"gnalloy.org/gnalloy/channel"
)

// HeaderField 是 HTTP/3 QPACK 解压后的 header 字段。
type HeaderField struct {
	Name  string
	Value string
}

// HeadersBlock 表示解压后的 HTTP/3 HEADERS 语义消息。
type HeadersBlock struct {
	Fields []HeaderField
}

// Release 保持 HeadersBlock 可进入统一的 pipeline 释放路径。
func (HeadersBlock) Release() {}

// PushPromiseBlock 表示解压后的 HTTP/3 PUSH_PROMISE 语义消息。
type PushPromiseBlock struct {
	PushID uint64
	Fields []HeaderField
}

// Release 保持 PushPromiseBlock 可进入统一的 pipeline 释放路径。
func (PushPromiseBlock) Release() {}

// HeaderCodecConfig 描述 HTTP/3 QPACK 编解码边界。
type HeaderCodecConfig struct {
	// MaxHeaderListSize 按 RFC 语义统计 name/value 加 32 字节开销，0 表示不额外限制。
	MaxHeaderListSize uint64
}

// HeaderDecoder 把 HEADERS/PUSH_PROMISE 的 QPACK block 解码为字段。
type HeaderDecoder struct {
	decoder           *qpack.Decoder
	maxHeaderListSize uint64
}

// HeaderEncoder 把字段编码为 QPACK block，再包装为 HTTP/3 frame。
type HeaderEncoder struct {
	encoder *qpack.Encoder
	buf     bytes.Buffer
}

// NewHeaderDecoder 创建 HTTP/3 QPACK header decoder。
func NewHeaderDecoder(cfg HeaderCodecConfig) *HeaderDecoder {
	return &HeaderDecoder{decoder: qpack.NewDecoder(), maxHeaderListSize: cfg.MaxHeaderListSize}
}

// NewHeaderEncoder 创建 HTTP/3 QPACK header encoder。
func NewHeaderEncoder() *HeaderEncoder {
	e := &HeaderEncoder{}
	e.encoder = qpack.NewEncoder(&e.buf)
	return e
}

// ChannelRead 解码 HEADERS 与 PUSH_PROMISE，其它消息继续透传。
func (d *HeaderDecoder) ChannelRead(ctx *channel.HandlerContext, msg any) {
	switch frame := msg.(type) {
	case HeadersFrame:
		fields, err := d.decodeFields(frame.HeaderBlock)
		frame.Release()
		if err != nil {
			ctx.FireExceptionCaught(err)
			return
		}
		ctx.FireChannelRead(HeadersBlock{Fields: fields})
	case PushPromiseFrame:
		fields, err := d.decodeFields(frame.HeaderBlock)
		frame.Release()
		if err != nil {
			ctx.FireExceptionCaught(err)
			return
		}
		ctx.FireChannelRead(PushPromiseBlock{PushID: frame.PushID, Fields: fields})
	default:
		ctx.FireChannelRead(msg)
	}
}

// Write 编码 HTTP/3 header 语义消息，其它消息继续透传。
func (e *HeaderEncoder) Write(ctx *channel.HandlerContext, msg any) error {
	switch block := msg.(type) {
	case HeadersBlock:
		headerBlock, err := e.encodeFields(ctx, block.Fields)
		if err != nil {
			return err
		}
		return writeHeaderFrame(ctx, HeadersFrame{HeaderBlock: headerBlock})
	case PushPromiseBlock:
		headerBlock, err := e.encodeFields(ctx, block.Fields)
		if err != nil {
			return err
		}
		return writeHeaderFrame(ctx, PushPromiseFrame{PushID: block.PushID, HeaderBlock: headerBlock})
	default:
		return ctx.Write(msg)
	}
}

func (d *HeaderDecoder) decodeFields(block buffer.ByteBuf) ([]HeaderField, error) {
	if block == nil {
		return nil, nil
	}
	decode := d.decoder.Decode(headerBlockBytes(block))
	fields := make([]HeaderField, 0, 8)
	var size uint64
	for {
		field, err := decode()
		if err == io.EOF {
			return fields, nil
		}
		if err != nil {
			return nil, err
		}
		size += uint64(len(field.Name) + len(field.Value) + 32)
		if d.maxHeaderListSize > 0 && size > d.maxHeaderListSize {
			return nil, ErrHeaderListTooLarge
		}
		fields = append(fields, HeaderField{Name: field.Name, Value: field.Value})
	}
}

func (e *HeaderEncoder) encodeFields(ctx *channel.HandlerContext, fields []HeaderField) (buffer.ByteBuf, error) {
	e.buf.Reset()
	if len(fields) == 0 {
		return http3BufferFromBytes(ctx, []byte{0, 0})
	}
	defer func() { _ = e.encoder.Close() }()
	for _, field := range fields {
		if err := e.encoder.WriteField(qpack.HeaderField{Name: field.Name, Value: field.Value}); err != nil {
			return nil, err
		}
	}
	return http3BufferFromBytes(ctx, e.buf.Bytes())
}

func writeHeaderFrame(ctx *channel.HandlerContext, frame any) error {
	if err := ctx.Write(frame); err != nil {
		releaseMessage(frame)
		return err
	}
	return nil
}

func http3BufferFromBytes(ctx *channel.HandlerContext, data []byte) (buffer.ByteBuf, error) {
	out, err := ctx.Channel().Allocator().Acquire(len(data))
	if err != nil {
		return nil, err
	}
	if _, err := out.WriteBytes(data); err != nil {
		out.Release()
		return nil, err
	}
	return out, nil
}

// FrameDirection 描述 HTTP/3 frame 在 Pipeline 中的方向。
type FrameDirection uint8

const (
	FrameDirectionInbound FrameDirection = iota + 1
	FrameDirectionOutbound
)

// StatsSnapshot 是 HTTP/3 frame 统计的无锁快照。
type StatsSnapshot struct {
	InboundFrames       uint64
	OutboundFrames      uint64
	InboundBytes        uint64
	OutboundBytes       uint64
	InboundDataBytes    uint64
	OutboundDataBytes   uint64
	InboundHeaderBytes  uint64
	OutboundHeaderBytes uint64
	SettingsFrames      uint64
	GoAwayFrames        uint64
	UnknownFrames       uint64
}

// StatsRecorder 是 HTTP/3 frame 统计接入点。
type StatsRecorder interface {
	RecordHTTP3Frame(direction FrameDirection, frameType FrameType, payloadBytes int)
}

// AtomicStatsRecorder 是低基数、并发安全的 HTTP/3 frame 统计器。
type AtomicStatsRecorder struct {
	inboundFrames       atomic.Uint64
	outboundFrames      atomic.Uint64
	inboundBytes        atomic.Uint64
	outboundBytes       atomic.Uint64
	inboundDataBytes    atomic.Uint64
	outboundDataBytes   atomic.Uint64
	inboundHeaderBytes  atomic.Uint64
	outboundHeaderBytes atomic.Uint64
	settingsFrames      atomic.Uint64
	goAwayFrames        atomic.Uint64
	unknownFrames       atomic.Uint64
}

// NewAtomicStatsRecorder 创建 HTTP/3 原子统计器。
func NewAtomicStatsRecorder() *AtomicStatsRecorder {
	return &AtomicStatsRecorder{}
}

// RecordHTTP3Frame 记录单个 HTTP/3 frame，不持有消息引用。
func (r *AtomicStatsRecorder) RecordHTTP3Frame(direction FrameDirection, frameType FrameType, payloadBytes int) {
	if r == nil {
		return
	}
	if payloadBytes < 0 {
		payloadBytes = 0
	}
	switch direction {
	case FrameDirectionInbound:
		r.inboundFrames.Add(1)
		r.inboundBytes.Add(uint64(payloadBytes))
		if frameType == FrameData {
			r.inboundDataBytes.Add(uint64(payloadBytes))
		}
		if frameType == FrameHeaders {
			r.inboundHeaderBytes.Add(uint64(payloadBytes))
		}
	case FrameDirectionOutbound:
		r.outboundFrames.Add(1)
		r.outboundBytes.Add(uint64(payloadBytes))
		if frameType == FrameData {
			r.outboundDataBytes.Add(uint64(payloadBytes))
		}
		if frameType == FrameHeaders {
			r.outboundHeaderBytes.Add(uint64(payloadBytes))
		}
	default:
		return
	}
	switch frameType {
	case FrameSettings:
		r.settingsFrames.Add(1)
	case FrameGoAway:
		r.goAwayFrames.Add(1)
	default:
		if frameType != FrameData &&
			frameType != FrameHeaders &&
			frameType != FrameCancelPush &&
			frameType != FramePushPromise &&
			frameType != FrameMaxPushID &&
			frameType != FramePriorityUpdateStream &&
			frameType != FramePriorityUpdatePush {
			r.unknownFrames.Add(1)
		}
	}
}

// Snapshot 返回当前统计快照。
func (r *AtomicStatsRecorder) Snapshot() StatsSnapshot {
	if r == nil {
		return StatsSnapshot{}
	}
	return StatsSnapshot{
		InboundFrames:       r.inboundFrames.Load(),
		OutboundFrames:      r.outboundFrames.Load(),
		InboundBytes:        r.inboundBytes.Load(),
		OutboundBytes:       r.outboundBytes.Load(),
		InboundDataBytes:    r.inboundDataBytes.Load(),
		OutboundDataBytes:   r.outboundDataBytes.Load(),
		InboundHeaderBytes:  r.inboundHeaderBytes.Load(),
		OutboundHeaderBytes: r.outboundHeaderBytes.Load(),
		SettingsFrames:      r.settingsFrames.Load(),
		GoAwayFrames:        r.goAwayFrames.Load(),
		UnknownFrames:       r.unknownFrames.Load(),
	}
}

// StatsHandler 在不接管消息所有权的前提下记录 HTTP/3 frame 统计。
type StatsHandler struct {
	recorder StatsRecorder
}

// NewStatsHandler 创建 HTTP/3 frame 统计 handler。
func NewStatsHandler(recorder StatsRecorder) *StatsHandler {
	return &StatsHandler{recorder: recorder}
}

// ChannelRead 记录入站 frame 后继续传播。
func (h *StatsHandler) ChannelRead(ctx *channel.HandlerContext, msg any) {
	if frameType, payloadBytes, ok := frameMetric(msg); ok && h.recorder != nil {
		h.recorder.RecordHTTP3Frame(FrameDirectionInbound, frameType, payloadBytes)
	}
	ctx.FireChannelRead(msg)
}

// Write 记录出站 frame 后继续传播。
func (h *StatsHandler) Write(ctx *channel.HandlerContext, msg any) error {
	if frameType, payloadBytes, ok := frameMetric(msg); ok && h.recorder != nil {
		h.recorder.RecordHTTP3Frame(FrameDirectionOutbound, frameType, payloadBytes)
	}
	return ctx.Write(msg)
}

func frameMetric(msg any) (FrameType, int, bool) {
	switch frame := msg.(type) {
	case DataFrame:
		return FrameData, readable(frame.Data), true
	case *DataFrame:
		if frame == nil {
			return 0, 0, false
		}
		return FrameData, readable(frame.Data), true
	case HeadersFrame:
		return FrameHeaders, readable(frame.HeaderBlock), true
	case *HeadersFrame:
		if frame == nil {
			return 0, 0, false
		}
		return FrameHeaders, readable(frame.HeaderBlock), true
	case CancelPushFrame:
		return FrameCancelPush, 0, true
	case SettingsFrame:
		return FrameSettings, len(frame.Settings) * 2, true
	case PushPromiseFrame:
		return FramePushPromise, readable(frame.HeaderBlock), true
	case *PushPromiseFrame:
		if frame == nil {
			return 0, 0, false
		}
		return FramePushPromise, readable(frame.HeaderBlock), true
	case GoAwayFrame:
		return FrameGoAway, 0, true
	case MaxPushIDFrame:
		return FrameMaxPushID, 0, true
	case PriorityUpdateFrame:
		return frame.Type, readable(frame.FieldValue), true
	case *PriorityUpdateFrame:
		if frame == nil {
			return 0, 0, false
		}
		return frame.Type, readable(frame.FieldValue), true
	case UnknownFrame:
		return frame.Type, readable(frame.Payload), true
	case *UnknownFrame:
		if frame == nil {
			return 0, 0, false
		}
		return frame.Type, readable(frame.Payload), true
	default:
		return 0, 0, false
	}
}