From c7d8592969c4ce8723148f9f73d1773339222f4e Mon Sep 17 00:00:00 2001 From: fionera Date: Fri, 10 Jul 2026 02:32:38 +0000 Subject: [PATCH] fix(spop): support configurable frame sizes --- spop/README.md | 23 +- spop/agent.go | 29 +- spop/agent_test.go | 36 +++ spop/error.go | 26 ++ spop/frame.go | 62 +++- spop/frame_test.go | 38 ++- spop/frames.go | 24 +- spop/protocol.go | 157 +++++++-- spop/protocol_test.go | 708 +++++++++++++++++++++++++++++++++++++++++ spop/scheduler.go | 2 + spop/scheduler_test.go | 25 ++ spop/server_test.go | 81 +++-- 12 files changed, 1141 insertions(+), 70 deletions(-) create mode 100644 spop/agent_test.go create mode 100644 spop/protocol_test.go create mode 100644 spop/scheduler_test.go diff --git a/spop/README.md b/spop/README.md index f9d2059..3536148 100644 --- a/spop/README.md +++ b/spop/README.md @@ -1,9 +1,28 @@ # A HAProxy SPOE implementation in Go +## Frame size -# References +Agents accept SPOP frames up to 65,535 bytes by default. Set `Agent.MaxFrameSize` +when HAProxy is configured to advertise larger frames: + +```go +agent := spop.Agent{ + Addr: ":9000", + Handler: handler, + MaxFrameSize: 256*1024 - 4, +} +log.Fatal(agent.ListenAndServe()) +``` + +The configured value excludes the four-byte frame-length prefix. It is an +upper bound: each connection negotiates the lower of this value and HAProxy's +advertised `max-frame-size`. Values below 256 are invalid. Each active pooled +frame buffer is allocated at the configured size, so use the smallest ceiling +that accommodates the HAProxy configuration. + +## References https://www.haproxy.org/download/2.0/doc/SPOE.txt -# Alternative implementations +## Alternative implementations https://github.com/criteo/haproxy-spoe-go https://github.com/negasus/haproxy-spoe-go diff --git a/spop/agent.go b/spop/agent.go index 545944b..cbb1ffe 100644 --- a/spop/agent.go +++ b/spop/agent.go @@ -14,6 +14,11 @@ type Agent struct { Handler Handler BaseContext context.Context Addr string + + // MaxFrameSize is the maximum SPOP frame size, excluding the four-byte + // length prefix. A zero value uses DefaultMaxFrameSize; values below 256 + // are invalid. Each pooled frame buffer is sized to this limit. + MaxFrameSize uint32 } func ListenAndServe(addr string, handler Handler) error { @@ -37,6 +42,12 @@ func (a *Agent) Serve(l net.Listener) error { a.BaseContext = context.Background() } + maxFrameSize, err := a.configuredMaxFrameSize() + if err != nil { + return err + } + framePool := newFramePool(maxFrameSize) + go func() { <-a.BaseContext.Done() l.Close() @@ -60,7 +71,7 @@ func (a *Agent) Serve(l net.Listener) error { } } - p := newProtocolClient(a.BaseContext, nc, as, a.Handler) + p := newProtocolClient(a.BaseContext, nc, as, framePool, a.Handler) go func() { defer nc.Close() defer p.Close() @@ -73,6 +84,22 @@ func (a *Agent) Serve(l net.Listener) error { } } +func (a *Agent) configuredMaxFrameSize() (uint32, error) { + maxFrameSize := a.MaxFrameSize + if maxFrameSize == 0 { + maxFrameSize = DefaultMaxFrameSize + } + + if maxFrameSize < minFrameSize { + return 0, fmt.Errorf("max frame size must be at least %d: %d", minFrameSize, maxFrameSize) + } + if uint64(maxFrameSize) > uint64(^uint(0)>>1) { + return 0, fmt.Errorf("max frame size exceeds platform limit: %d", maxFrameSize) + } + + return maxFrameSize, nil +} + func wrapPanic(fn func() error) (err error) { didPanic := true defer func() { diff --git a/spop/agent_test.go b/spop/agent_test.go new file mode 100644 index 0000000..6b8a46a --- /dev/null +++ b/spop/agent_test.go @@ -0,0 +1,36 @@ +package spop + +import "testing" + +func TestAgentConfiguredMaxFrameSize(t *testing.T) { + tests := []struct { + name string + value uint32 + want uint32 + wantErr bool + }{ + {name: "default", want: DefaultMaxFrameSize}, + {name: "protocol minimum", value: minFrameSize, want: minFrameSize}, + {name: "large frame", value: 262140, want: 262140}, + {name: "below protocol minimum", value: minFrameSize - 1, wantErr: true}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + a := Agent{MaxFrameSize: tt.value} + got, err := a.configuredMaxFrameSize() + if tt.wantErr { + if err == nil { + t.Fatal("expected an error") + } + return + } + if err != nil { + t.Fatalf("configured max frame size: %v", err) + } + if got != tt.want { + t.Fatalf("expected %d, got %d", tt.want, got) + } + }) + } +} diff --git a/spop/error.go b/spop/error.go index b7c49e8..956c652 100644 --- a/spop/error.go +++ b/spop/error.go @@ -1,9 +1,35 @@ package spop import ( + "errors" "fmt" ) +type protocolError struct { + err error + code errorCode +} + +func newProtocolError(code errorCode, format string, args ...any) error { + return &protocolError{code: code, err: fmt.Errorf(format, args...)} +} + +func (e *protocolError) Error() string { + return e.err.Error() +} + +func (e *protocolError) Unwrap() error { + return e.err +} + +func protocolErrorCode(err error) errorCode { + var protocolErr *protocolError + if errors.As(err, &protocolErr) { + return protocolErr.code + } + return ErrorUnknown +} + type errorCode int func (e errorCode) String() string { diff --git a/spop/frame.go b/spop/frame.go index d476e11..939cc21 100644 --- a/spop/frame.go +++ b/spop/frame.go @@ -13,17 +13,31 @@ import ( const uint32Len = 4 -var framePool = sync.Pool{ - New: func() any { +type framePool struct { + pool sync.Pool + maxFrameSize uint32 +} + +func newFramePool(maxFrameSize uint32) *framePool { + p := &framePool{maxFrameSize: maxFrameSize} + p.pool.New = func() any { return &frame{ + buf: buffer.NewSliceBuffer(int(maxFrameSize)), length: make([]byte, uint32Len), - buf: buffer.NewSliceBuffer(maxFrameSize), + pool: p, } - }, + } + return p +} + +var defaultFramePool = newFramePool(DefaultMaxFrameSize) + +func (p *framePool) acquire() *frame { + return p.pool.Get().(*frame) } func acquireFrame() *frame { - return framePool.Get().(*frame) + return defaultFramePool.acquire() } func releaseFrame(f *frame) { @@ -31,7 +45,7 @@ func releaseFrame(f *frame) { f.frameType = 0 f.meta = frameMetadata{} - framePool.Put(f) + f.pool.pool.Put(f) } type frameMetadata struct { @@ -41,7 +55,8 @@ type frameMetadata struct { } type frame struct { - buf *buffer.SliceBuffer + buf *buffer.SliceBuffer + pool *framePool length []byte meta frameMetadata @@ -50,13 +65,25 @@ type frame struct { } func (f *frame) ReadFrom(r io.Reader) (int64, error) { + return f.readFrom(r, f.pool.maxFrameSize) +} + +func (f *frame) readFrom(r io.Reader, maxFrameSize uint32) (int64, error) { if _, err := io.ReadFull(r, f.length); err != nil { return 0, fmt.Errorf("reading frame length: %w", err) } frameLen := binary.BigEndian.Uint32(f.length) + if maxFrameSize > f.pool.maxFrameSize { + maxFrameSize = f.pool.maxFrameSize + } if frameLen > maxFrameSize { - return int64(len(f.length)), fmt.Errorf("frame length %d exceeds maximum %d", frameLen, maxFrameSize) + return int64(len(f.length)), newProtocolError( + ErrorTooBig, + "frame length %d exceeds maximum %d", + frameLen, + maxFrameSize, + ) } f.buf.Reset() @@ -72,7 +99,24 @@ func (f *frame) ReadFrom(r io.Reader) (int64, error) { } func (f *frame) WriteTo(w io.Writer) (int64, error) { - binary.BigEndian.PutUint32(f.length, uint32(f.buf.Len())) + return f.writeTo(w, f.pool.maxFrameSize) +} + +func (f *frame) writeTo(w io.Writer, maxFrameSize uint32) (int64, error) { + if maxFrameSize > f.pool.maxFrameSize { + maxFrameSize = f.pool.maxFrameSize + } + frameLen := uint32(f.buf.Len()) + if frameLen > maxFrameSize { + return 0, newProtocolError( + ErrorTooBig, + "frame length %d exceeds maximum %d", + frameLen, + maxFrameSize, + ) + } + + binary.BigEndian.PutUint32(f.length, frameLen) if n, err := w.Write(f.length); err != nil { return int64(n), err diff --git a/spop/frame_test.go b/spop/frame_test.go index 4cc6984..47e0e46 100644 --- a/spop/frame_test.go +++ b/spop/frame_test.go @@ -37,9 +37,9 @@ func TestFrame_ReadFrom_ExceedsMaxFrameSize(t *testing.T) { }) t.Run("accepts frame length at maxFrameSize boundary", func(t *testing.T) { - buf := writeFrameLength(maxFrameSize) + buf := writeFrameLength(DefaultMaxFrameSize) - frameData := make([]byte, maxFrameSize) + frameData := make([]byte, DefaultMaxFrameSize) frameData[0] = byte(frameTypeIDHaproxyHello) binary.BigEndian.PutUint32(frameData[1:5], 0) frameData[5] = 0 // streamID varint @@ -59,7 +59,7 @@ func TestFrame_ReadFrom_ExceedsMaxFrameSize(t *testing.T) { }) t.Run("rejects frame length one byte over maxFrameSize", func(t *testing.T) { - buf := writeFrameLength(maxFrameSize + 1) + buf := writeFrameLength(DefaultMaxFrameSize + 1) f := acquireFrame() defer releaseFrame(f) @@ -67,4 +67,34 @@ func TestFrame_ReadFrom_ExceedsMaxFrameSize(t *testing.T) { _, err := f.ReadFrom(buf) assertError(t, err) }) -} \ No newline at end of file + + t.Run("accepts configured frame length above default", func(t *testing.T) { + const configuredMaxFrameSize uint32 = 262140 + buf := writeFrameLength(configuredMaxFrameSize) + + frameData := make([]byte, configuredMaxFrameSize) + frameData[0] = byte(frameTypeIDHaproxyHello) + binary.BigEndian.PutUint32(frameData[1:5], 0) + frameData[5] = 0 + frameData[6] = 0 + buf.Write(frameData) + + f := newFramePool(configuredMaxFrameSize).acquire() + defer releaseFrame(f) + + if _, err := f.ReadFrom(buf); err != nil { + t.Fatalf("unexpected error: %v", err) + } + }) + + t.Run("rejects frame above configured maximum", func(t *testing.T) { + const configuredMaxFrameSize uint32 = 262140 + buf := writeFrameLength(configuredMaxFrameSize + 1) + + f := newFramePool(configuredMaxFrameSize).acquire() + defer releaseFrame(f) + + _, err := f.ReadFrom(buf) + assertError(t, err) + }) +} diff --git a/spop/frames.go b/spop/frames.go index c3fb7e9..45f42ae 100644 --- a/spop/frames.go +++ b/spop/frames.go @@ -43,7 +43,11 @@ type AgentDisconnectFrame struct { } func (a *AgentDisconnectFrame) WriteTo(w io.Writer) (int64, error) { - f := acquireFrame() + return a.writeTo(w, defaultFramePool, DefaultMaxFrameSize) +} + +func (a *AgentDisconnectFrame) writeTo(w io.Writer, pool *framePool, maxFrameSize uint32) (int64, error) { + f := pool.acquire() defer releaseFrame(f) f.frameType = frameTypeIDAgentDisconnect @@ -66,7 +70,7 @@ func (a *AgentDisconnectFrame) WriteTo(w io.Writer) (int64, error) { f.buf.AdvanceW(kvw.Off()) - return f.WriteTo(w) + return f.writeTo(w, maxFrameSize) } const ( @@ -89,7 +93,11 @@ type AgentHelloFrame struct { } func (a *AgentHelloFrame) WriteTo(w io.Writer) (int64, error) { - f := acquireFrame() + return a.writeTo(w, defaultFramePool, DefaultMaxFrameSize) +} + +func (a *AgentHelloFrame) writeTo(w io.Writer, pool *framePool, maxFrameSize uint32) (int64, error) { + f := pool.acquire() defer releaseFrame(f) f.frameType = frameTypeIDAgentHello @@ -116,7 +124,7 @@ func (a *AgentHelloFrame) WriteTo(w io.Writer) (int64, error) { } f.buf.AdvanceW(kvw.Off()) - return f.WriteTo(w) + return f.writeTo(w, maxFrameSize) } type AckFrame struct { @@ -126,7 +134,11 @@ type AckFrame struct { } func (a *AckFrame) WriteTo(w io.Writer) (int64, error) { - f := acquireFrame() + return a.writeTo(w, defaultFramePool, DefaultMaxFrameSize) +} + +func (a *AckFrame) writeTo(w io.Writer, pool *framePool, maxFrameSize uint32) (int64, error) { + f := pool.acquire() defer releaseFrame(f) f.frameType = frameTypeIDAck @@ -148,5 +160,5 @@ func (a *AckFrame) WriteTo(w io.Writer) (int64, error) { f.buf.AdvanceW(aw.Off()) - return f.WriteTo(w) + return f.writeTo(w, maxFrameSize) } diff --git a/spop/protocol.go b/spop/protocol.go index c38189c..fdf982d 100644 --- a/spop/protocol.go +++ b/spop/protocol.go @@ -2,20 +2,32 @@ package spop import ( "context" + "encoding/binary" "errors" "fmt" "io" + "net" + "sync" "syscall" + "time" "github.com/dropmorepackets/haproxy-go/pkg/encoding" ) -func newProtocolClient(ctx context.Context, rw io.ReadWriter, as *asyncScheduler, handler Handler) *protocolClient { +func newProtocolClient( + ctx context.Context, + rw io.ReadWriter, + as *asyncScheduler, + framePool *framePool, + handler Handler, +) *protocolClient { var c protocolClient c.rw = rw c.handler = handler c.ctx, c.ctxCancel = context.WithCancelCause(ctx) c.as = as + c.framePool = framePool + c.maxFrameSize = framePool.maxFrameSize return &c } @@ -26,8 +38,11 @@ type protocolClient struct { ctxCancel context.CancelCauseFunc as *asyncScheduler + framePool *framePool + engineID string + + closeOnce sync.Once - engineID string maxFrameSize uint32 gotHello bool @@ -38,17 +53,36 @@ func (c *protocolClient) Close() error { return c.ctx.Err() } - // We ignore any error since the disconnect frame is delivered on - // best effort anyway. - _, _ = (&AgentDisconnectFrame{ - ErrCode: ErrorUnknown, - }).WriteTo(c.rw) - - c.ctxCancel(fmt.Errorf("closing client")) + c.terminate(fmt.Errorf("closing client")) return nil } +func (c *protocolClient) terminate(err error) { + c.terminateWithCode(err, protocolErrorCode(err)) +} + +func (c *protocolClient) terminateWithCode(err error, code errorCode) { + c.closeOnce.Do(func() { + c.ctxCancel(err) + + if conn, ok := c.rw.(net.Conn); ok { + if deadlineErr := conn.SetWriteDeadline(time.Now().Add(disconnectWriteTimeout)); deadlineErr == nil { + _, _ = (&AgentDisconnectFrame{ + ErrCode: code, + }).writeTo(conn, c.framePool, c.maxFrameSize) + } + _ = conn.Close() + return + } + + // A generic writer cannot provide a bounded best-effort write. + if closer, ok := c.rw.(io.Closer); ok { + _ = closer.Close() + } + }) +} + func (c *protocolClient) frameHandler(f *frame) error { defer releaseFrame(f) @@ -65,17 +99,36 @@ func (c *protocolClient) frameHandler(f *frame) error { } func (c *protocolClient) Serve() error { - for { - f := acquireFrame() - if _, err := f.ReadFrom(c.rw); err != nil { - if c.ctx.Err() != nil { - return context.Cause(c.ctx) - } + f, err := c.readFrame() + if err != nil { + err = c.readError(err) + if err != nil { + c.terminate(err) + } + return err + } + if f.frameType != frameTypeIDHaproxyHello { + firstFrameType := f.frameType + releaseFrame(f) + err := newProtocolError(ErrorInvalid, "first frame must be HAPROXY-HELLO, got type %d", firstFrameType) + c.terminate(err) + return err + } + if err := c.frameHandler(f); err != nil { + c.terminate(err) + return err + } + if c.ctx.Err() != nil { + return context.Cause(c.ctx) + } - if errors.Is(err, io.EOF) || errors.Is(err, syscall.ECONNRESET) { - return nil + for { + f, err := c.readFrame() + if err != nil { + err = c.readError(err) + if err != nil { + c.terminate(err) } - return err } @@ -83,12 +136,34 @@ func (c *protocolClient) Serve() error { } } +func (c *protocolClient) readFrame() (*frame, error) { + f := c.framePool.acquire() + if _, err := f.readFrom(c.rw, c.maxFrameSize); err != nil { + releaseFrame(f) + return nil, err + } + return f, nil +} + +func (c *protocolClient) readError(err error) error { + if c.ctx.Err() != nil { + return context.Cause(c.ctx) + } + if errors.Is(err, io.EOF) || errors.Is(err, syscall.ECONNRESET) { + return nil + } + return err +} + const ( version = "2.0" - // maxFrameSize represents the maximum frame size allowed by this library - // it also represents the maximum slice size that is allowed on stack - maxFrameSize = 64<<10 - 1 + minFrameSize = 256 + disconnectWriteTimeout = 100 * time.Millisecond + + // DefaultMaxFrameSize is the maximum SPOP frame size used by an Agent when + // MaxFrameSize is not configured. + DefaultMaxFrameSize uint32 = 64<<10 - 1 ) func (c *protocolClient) onHAProxyHello(f *frame) error { @@ -102,12 +177,33 @@ func (c *protocolClient) onHAProxyHello(f *frame) error { k := encoding.AcquireKVEntry() defer encoding.ReleaseKVEntry(k) + gotMaxFrameSize := false for s.Next(k) { switch { case k.NameEquals(helloKeyMaxFrameSize): - c.maxFrameSize = uint32(k.ValueInt()) - if c.maxFrameSize > maxFrameSize { - return fmt.Errorf("maxFrameSize bigger than maximum allowed size: %d < %d", maxFrameSize, c.maxFrameSize) + if k.Type() != encoding.DataTypeUInt32 { + return newProtocolError(ErrorBadFrameSize, "peer max frame size must be uint32") + } + gotMaxFrameSize = true + peerMaxFrameSize := uint32(k.ValueInt()) + if peerMaxFrameSize < minFrameSize { + return newProtocolError( + ErrorBadFrameSize, + "peer max frame size must be at least %d: %d", + minFrameSize, + peerMaxFrameSize, + ) + } + if binary.BigEndian.Uint32(f.length) > peerMaxFrameSize { + return newProtocolError( + ErrorBadFrameSize, + "HAPROXY-HELLO frame length %d exceeds peer maximum %d", + binary.BigEndian.Uint32(f.length), + peerMaxFrameSize, + ) + } + if peerMaxFrameSize < c.maxFrameSize { + c.maxFrameSize = peerMaxFrameSize } case k.NameEquals(helloKeyEngineID): @@ -126,12 +222,15 @@ func (c *protocolClient) onHAProxyHello(f *frame) error { if err := s.Error(); err != nil { return err } + if !gotMaxFrameSize { + return newProtocolError(ErrorNoFrameSize, "HAPROXY-HELLO missing %q", helloKeyMaxFrameSize) + } _, err := (&AgentHelloFrame{ Version: version, MaxFrameSize: c.maxFrameSize, Capabilities: []string{}, - }).WriteTo(c.rw) + }).writeTo(c.rw, c.framePool, c.maxFrameSize) return err } @@ -164,7 +263,7 @@ func (c *protocolClient) onNotify(f *frame) error { FrameID: f.meta.FrameID, StreamID: f.meta.StreamID, ActionWriterCallback: fn, - }).WriteTo(c.rw) + }).writeTo(c.rw, c.framePool, c.maxFrameSize) return err } @@ -200,10 +299,14 @@ func (c *protocolClient) onHAProxyDisconnect(f *frame) error { // HAProxy returns an IO error when it doesn't require a connection // anymore. case ErrorIO, ErrorTimeout, ErrorNone: + err = context.Canceled default: err = fmt.Errorf("disconnect frame with code %d: %s", code, code) } - c.ctxCancel(err) + c.terminateWithCode(err, code) + if code == ErrorIO || code == ErrorTimeout || code == ErrorNone { + return nil + } return err } diff --git a/spop/protocol_test.go b/spop/protocol_test.go new file mode 100644 index 0000000..910a1b9 --- /dev/null +++ b/spop/protocol_test.go @@ -0,0 +1,708 @@ +package spop + +import ( + "bytes" + "context" + "encoding/binary" + "errors" + "fmt" + "io" + "net" + "strings" + "testing" + "time" + + "github.com/dropmorepackets/haproxy-go/pkg/encoding" +) + +func TestProtocolNegotiatesMaxFrameSize(t *testing.T) { + tests := []struct { + name string + agentMax uint32 + peerMax uint32 + want uint32 + }{ + {name: "peer offer is larger", agentMax: DefaultMaxFrameSize, peerMax: 262140, want: DefaultMaxFrameSize}, + {name: "agent maximum is larger", agentMax: 262140, peerMax: DefaultMaxFrameSize, want: DefaultMaxFrameSize}, + {name: "limits are equal", agentMax: 262140, peerMax: 262140, want: 262140}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + pool := newFramePool(tt.agentMax) + var rw bytes.Buffer + client := newProtocolClient(context.Background(), &rw, nil, pool, HandlerFunc(func( + context.Context, + *encoding.ActionWriter, + *encoding.Message, + ) { + })) + + frame := testHelloFrame(t, pool, &tt.peerMax) + if err := client.onHAProxyHello(frame); err != nil { + releaseFrame(frame) + t.Fatalf("handle HAPROXY-HELLO: %v", err) + } + releaseFrame(frame) + + if client.maxFrameSize != tt.want { + t.Fatalf("expected negotiated maximum %d, got %d", tt.want, client.maxFrameSize) + } + response := readTestWireFrame(t, &rw) + if response.frameType != frameTypeIDAgentHello { + t.Fatalf("expected AGENT-HELLO, got frame type %d", response.frameType) + } + if got := testHelloMaxFrameSize(t, response.payload); got != tt.want { + t.Fatalf("expected advertised maximum %d, got %d", tt.want, got) + } + }) + } +} + +func TestProtocolRejectsInvalidMaxFrameSize(t *testing.T) { + tests := []struct { + name string + peerMax *uint32 + want errorCode + }{ + {name: "missing", want: ErrorNoFrameSize}, + {name: "below protocol minimum", peerMax: uint32Ptr(minFrameSize - 1), want: ErrorBadFrameSize}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + pool := newFramePool(DefaultMaxFrameSize) + client := newProtocolClient(context.Background(), &bytes.Buffer{}, nil, pool, HandlerFunc(func( + context.Context, + *encoding.ActionWriter, + *encoding.Message, + ) { + })) + frame := testHelloFrame(t, pool, tt.peerMax) + defer releaseFrame(frame) + + err := client.onHAProxyHello(frame) + if err == nil { + t.Fatal("expected HAPROXY-HELLO to be rejected") + } + if got := protocolErrorCode(err); got != tt.want { + t.Fatalf("expected protocol error code %d, got %d", tt.want, got) + } + }) + } +} + +func TestProtocolRejectsHelloAboveAdvertisedMaxFrameSize(t *testing.T) { + const peerMax uint32 = minFrameSize + serverConn, clientConn := net.Pipe() + defer serverConn.Close() + defer clientConn.Close() + if err := clientConn.SetDeadline(time.Now().Add(5 * time.Second)); err != nil { + t.Fatal(err) + } + + client := newProtocolClient( + context.Background(), + serverConn, + nil, + newFramePool(DefaultMaxFrameSize), + HandlerFunc(func(context.Context, *encoding.ActionWriter, *encoding.Message) {}), + ) + serveDone := make(chan error, 1) + go func() { + serveDone <- client.Serve() + }() + + payload := append(testHelloPayload(t, peerMax), bytes.Repeat([]byte{0}, int(peerMax))...) + writeTestWireFrame(t, clientConn, frameTypeIDHaproxyHello, 0, 0, payload) + disconnect := readTestWireFrame(t, clientConn) + if disconnect.frameType != frameTypeIDAgentDisconnect { + t.Fatalf("expected AGENT-DISCONNECT, got frame type %d", disconnect.frameType) + } + if code := testDisconnectCode(t, disconnect.payload); code != ErrorBadFrameSize { + t.Fatalf("expected disconnect code %d, got %d", ErrorBadFrameSize, code) + } + + select { + case err := <-serveDone: + if err == nil || !strings.Contains(err.Error(), "exceeds peer maximum 256") { + t.Fatalf("expected peer maximum error, got %v", err) + } + case <-time.After(5 * time.Second): + t.Fatal("protocol did not reject the oversized HAPROXY-HELLO") + } +} + +func TestAgentLargeFrameRoundTrip(t *testing.T) { + const maxFrameSize uint32 = 262140 + requestBody := bytes.Repeat([]byte{0xa5}, 90<<10) + responseBody := bytes.Repeat([]byte{0x5a}, 90<<10) + handlerErr := make(chan error, 1) + + handler := HandlerFunc(func(_ context.Context, writer *encoding.ActionWriter, message *encoding.Message) { + if string(message.NameBytes()) != "large" { + handlerErr <- fmt.Errorf("expected message %q, got %q", "large", message.NameBytes()) + return + } + + entry := encoding.AcquireKVEntry() + defer encoding.ReleaseKVEntry(entry) + if !message.KV.Next(entry) { + handlerErr <- fmt.Errorf("missing payload entry: %v", message.KV.Error()) + return + } + if !entry.NameEquals("payload") || !bytes.Equal(entry.ValueBytes(), requestBody) { + handlerErr <- fmt.Errorf("received an unexpected payload") + return + } + if err := writer.SetBinary(encoding.VarScopeTransaction, "result", responseBody); err != nil { + handlerErr <- fmt.Errorf("write response action: %w", err) + return + } + handlerErr <- nil + }) + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + listener, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + defer listener.Close() + + agentDone := make(chan error, 1) + agent := Agent{ + BaseContext: ctx, + Handler: handler, + MaxFrameSize: maxFrameSize, + } + go func() { + agentDone <- agent.Serve(listener) + }() + + clientConn, err := net.DialTimeout("tcp", listener.Addr().String(), 5*time.Second) + if err != nil { + t.Fatal(err) + } + defer clientConn.Close() + if err := clientConn.SetDeadline(time.Now().Add(5 * time.Second)); err != nil { + t.Fatal(err) + } + + writeTestWireFrame(t, clientConn, frameTypeIDHaproxyHello, 0, 0, testHelloPayload(t, maxFrameSize)) + hello := readTestWireFrame(t, clientConn) + if hello.frameType != frameTypeIDAgentHello { + t.Fatalf("expected AGENT-HELLO, got frame type %d", hello.frameType) + } + if got := testHelloMaxFrameSize(t, hello.payload); got != maxFrameSize { + t.Fatalf("expected negotiated maximum %d, got %d", maxFrameSize, got) + } + + const streamID, frameID = 17, 23 + writeTestWireFrame( + t, + clientConn, + frameTypeIDNotify, + streamID, + frameID, + testNotifyPayload(t, "large", "payload", requestBody), + ) + ack := readTestWireFrame(t, clientConn) + if ack.frameType != frameTypeIDAck { + t.Fatalf("expected ACK, got frame type %d", ack.frameType) + } + if ack.streamID != streamID || ack.frameID != frameID { + t.Fatalf("expected stream/frame IDs %d/%d, got %d/%d", streamID, frameID, ack.streamID, ack.frameID) + } + if ack.length <= DefaultMaxFrameSize { + t.Fatalf("expected ACK larger than the default maximum, got %d bytes", ack.length) + } + expectedAction := make([]byte, len(responseBody)+256) + actionWriter := encoding.NewActionWriter(expectedAction, 0) + if err := actionWriter.SetBinary(encoding.VarScopeTransaction, "result", responseBody); err != nil { + t.Fatal(err) + } + if !bytes.Equal(ack.payload, actionWriter.Bytes()) { + t.Fatal("ACK contained an unexpected action payload") + } + if err := <-handlerErr; err != nil { + t.Fatal(err) + } + + if err := clientConn.Close(); err != nil { + t.Fatal(err) + } + cancel() + select { + case <-agentDone: + case <-time.After(5 * time.Second): + t.Fatal("agent did not stop after cancellation") + } +} + +func TestProtocolEnforcesNegotiatedMaxFrameSize(t *testing.T) { + const ( + agentMax uint32 = 262140 + peerMax uint32 = 4096 + ) + serverConn, clientConn := net.Pipe() + defer serverConn.Close() + defer clientConn.Close() + if err := clientConn.SetDeadline(time.Now().Add(5 * time.Second)); err != nil { + t.Fatal(err) + } + + client := newProtocolClient( + context.Background(), + serverConn, + nil, + newFramePool(agentMax), + HandlerFunc(func(context.Context, *encoding.ActionWriter, *encoding.Message) {}), + ) + serveDone := make(chan error, 1) + go func() { + serveDone <- client.Serve() + }() + + writeTestWireFrame(t, clientConn, frameTypeIDHaproxyHello, 0, 0, testHelloPayload(t, peerMax)) + hello := readTestWireFrame(t, clientConn) + if got := testHelloMaxFrameSize(t, hello.payload); got != peerMax { + t.Fatalf("expected negotiated maximum %d, got %d", peerMax, got) + } + + length := make([]byte, uint32Len) + binary.BigEndian.PutUint32(length, peerMax+1) + if _, err := clientConn.Write(length); err != nil { + t.Fatalf("write oversized frame length: %v", err) + } + disconnect := readTestWireFrame(t, clientConn) + if disconnect.frameType != frameTypeIDAgentDisconnect { + t.Fatalf("expected AGENT-DISCONNECT, got frame type %d", disconnect.frameType) + } + if code := testDisconnectCode(t, disconnect.payload); code != ErrorTooBig { + t.Fatalf("expected disconnect code %d, got %d", ErrorTooBig, code) + } + + select { + case err := <-serveDone: + if err == nil || !strings.Contains(err.Error(), "exceeds maximum 4096") { + t.Fatalf("expected negotiated maximum error, got %v", err) + } + case <-time.After(5 * time.Second): + t.Fatal("protocol did not reject the oversized frame") + } +} + +func TestProtocolClosesOnAckAboveNegotiatedMaxFrameSize(t *testing.T) { + const ( + agentMax uint32 = 262140 + peerMax uint32 = DefaultMaxFrameSize + ) + responseBody := bytes.Repeat([]byte{0x5a}, 90<<10) + handlerErr := make(chan error, 1) + handler := HandlerFunc(func(_ context.Context, writer *encoding.ActionWriter, _ *encoding.Message) { + handlerErr <- writer.SetBinary(encoding.VarScopeTransaction, "result", responseBody) + }) + + serverConn, clientConn := net.Pipe() + defer serverConn.Close() + defer clientConn.Close() + if err := clientConn.SetDeadline(time.Now().Add(5 * time.Second)); err != nil { + t.Fatal(err) + } + + client := newProtocolClient( + context.Background(), + serverConn, + newTestAsyncScheduler(), + newFramePool(agentMax), + handler, + ) + serveDone := make(chan error, 1) + go func() { + serveDone <- client.Serve() + }() + + writeTestWireFrame(t, clientConn, frameTypeIDHaproxyHello, 0, 0, testHelloPayload(t, peerMax)) + hello := readTestWireFrame(t, clientConn) + if got := testHelloMaxFrameSize(t, hello.payload); got != peerMax { + t.Fatalf("expected negotiated maximum %d, got %d", peerMax, got) + } + writeTestWireFrame( + t, + clientConn, + frameTypeIDNotify, + 1, + 1, + testNotifyPayload(t, "large", "payload", []byte("request")), + ) + + disconnect := readTestWireFrame(t, clientConn) + if disconnect.frameType != frameTypeIDAgentDisconnect { + t.Fatalf("expected AGENT-DISCONNECT, got frame type %d", disconnect.frameType) + } + if code := testDisconnectCode(t, disconnect.payload); code != ErrorTooBig { + t.Fatalf("expected disconnect code %d, got %d", ErrorTooBig, code) + } + if err := <-handlerErr; err != nil { + t.Fatalf("write action: %v", err) + } + + select { + case err := <-serveDone: + if err == nil || !strings.Contains(err.Error(), "exceeds maximum 65535") { + t.Fatalf("expected negotiated maximum error, got %v", err) + } + case <-time.After(5 * time.Second): + t.Fatal("protocol did not close after rejecting the ACK") + } +} + +func TestProtocolRepliesToHAProxyDisconnect(t *testing.T) { + serverConn, clientConn := net.Pipe() + defer serverConn.Close() + defer clientConn.Close() + if err := clientConn.SetDeadline(time.Now().Add(5 * time.Second)); err != nil { + t.Fatal(err) + } + + client := newProtocolClient( + context.Background(), + serverConn, + newTestAsyncScheduler(), + newFramePool(DefaultMaxFrameSize), + HandlerFunc(func(context.Context, *encoding.ActionWriter, *encoding.Message) {}), + ) + serveDone := make(chan error, 1) + go func() { + serveDone <- client.Serve() + }() + + writeTestWireFrame( + t, + clientConn, + frameTypeIDHaproxyHello, + 0, + 0, + testHelloPayload(t, DefaultMaxFrameSize), + ) + _ = readTestWireFrame(t, clientConn) + writeTestWireFrame( + t, + clientConn, + frameTypeIDHaproxyDisconnect, + 0, + 0, + testDisconnectPayload(t, ErrorNone), + ) + + disconnect := readTestWireFrame(t, clientConn) + if disconnect.frameType != frameTypeIDAgentDisconnect { + t.Fatalf("expected AGENT-DISCONNECT, got frame type %d", disconnect.frameType) + } + if code := testDisconnectCode(t, disconnect.payload); code != ErrorNone { + t.Fatalf("expected disconnect code %d, got %d", ErrorNone, code) + } + + select { + case err := <-serveDone: + if !errors.Is(err, context.Canceled) { + t.Fatalf("expected canceled protocol result, got %v", err) + } + case <-time.After(5 * time.Second): + t.Fatal("protocol did not close after HAPROXY-DISCONNECT") + } +} + +func TestProtocolTerminationIsBounded(t *testing.T) { + serverConn, clientConn := net.Pipe() + defer serverConn.Close() + defer clientConn.Close() + client := newProtocolClient( + context.Background(), + serverConn, + nil, + newFramePool(DefaultMaxFrameSize), + HandlerFunc(func(context.Context, *encoding.ActionWriter, *encoding.Message) {}), + ) + + done := make(chan struct{}) + go func() { + client.terminate(newProtocolError(ErrorTooBig, "test termination")) + close(done) + }() + + select { + case <-done: + case <-time.After(time.Second): + t.Fatal("protocol termination blocked on the disconnect write") + } +} + +func TestProtocolTerminationSkipsWriteWithoutDeadline(t *testing.T) { + serverConn, clientConn := net.Pipe() + defer serverConn.Close() + defer clientConn.Close() + client := newProtocolClient( + context.Background(), + &writeDeadlineErrorConn{Conn: serverConn}, + nil, + newFramePool(DefaultMaxFrameSize), + HandlerFunc(func(context.Context, *encoding.ActionWriter, *encoding.Message) {}), + ) + + done := make(chan struct{}) + go func() { + client.terminate(newProtocolError(ErrorTooBig, "test termination")) + close(done) + }() + + select { + case <-done: + case <-time.After(time.Second): + t.Fatal("protocol termination attempted an unbounded disconnect write") + } +} + +func TestFramePoolsKeepLimitsIsolated(t *testing.T) { + const smallMax uint32 = 4096 + const largeMax uint32 = 262140 + smallPool := newFramePool(smallMax) + largePool := newFramePool(largeMax) + + for i := 0; i < 2; i++ { + smallFrame := smallPool.acquire() + if _, err := smallFrame.ReadFrom(testFrameBytes(t, smallMax+1)); err == nil { + releaseFrame(smallFrame) + t.Fatal("expected the small pool to reject an oversized frame") + } + releaseFrame(smallFrame) + + largeFrame := largePool.acquire() + if _, err := largeFrame.ReadFrom(testFrameBytes(t, DefaultMaxFrameSize+1)); err != nil { + releaseFrame(largeFrame) + t.Fatalf("large pool rejected a valid frame: %v", err) + } + releaseFrame(largeFrame) + } +} + +type testWireFrame struct { + payload []byte + length uint32 + streamID uint64 + frameID uint64 + frameType frameType +} + +type writeDeadlineErrorConn struct { + net.Conn +} + +func (c *writeDeadlineErrorConn) SetWriteDeadline(time.Time) error { + return errors.ErrUnsupported +} + +func testHelloFrame(t *testing.T, pool *framePool, maxFrameSize *uint32) *frame { + t.Helper() + f := pool.acquire() + f.frameType = frameTypeIDHaproxyHello + f.meta.Flags = frameFlagFin + if err := f.encodeHeader(); err != nil { + releaseFrame(f) + t.Fatal(err) + } + headerLen := f.buf.Len() + + writer := encoding.NewKVWriter(f.buf.WriteBytes(), 0) + if err := writer.SetString(helloKeySupportedVersions, version); err != nil { + releaseFrame(f) + t.Fatal(err) + } + if maxFrameSize != nil { + if err := writer.SetUInt32(helloKeyMaxFrameSize, *maxFrameSize); err != nil { + releaseFrame(f) + t.Fatal(err) + } + } + if err := writer.SetString(helloKeyCapabilities, ""); err != nil { + releaseFrame(f) + t.Fatal(err) + } + f.buf.AdvanceW(writer.Off()) + binary.BigEndian.PutUint32(f.length, uint32(f.buf.Len())) + f.buf.AdvanceR(headerLen) + return f +} + +func testHelloPayload(t *testing.T, maxFrameSize uint32) []byte { + t.Helper() + buf := make([]byte, 128) + writer := encoding.NewKVWriter(buf, 0) + if err := writer.SetString(helloKeySupportedVersions, version); err != nil { + t.Fatal(err) + } + if err := writer.SetUInt32(helloKeyMaxFrameSize, maxFrameSize); err != nil { + t.Fatal(err) + } + if err := writer.SetString(helloKeyCapabilities, ""); err != nil { + t.Fatal(err) + } + return writer.Bytes() +} + +func testNotifyPayload(t *testing.T, messageName, entryName string, value []byte) []byte { + t.Helper() + buf := make([]byte, len(value)+256) + off, err := encoding.PutBytes(buf, []byte(messageName)) + if err != nil { + t.Fatal(err) + } + buf[off] = 1 + off++ + + writer := encoding.NewKVWriter(buf[off:], 0) + if err := writer.SetBinary(entryName, value); err != nil { + t.Fatal(err) + } + off += writer.Off() + return buf[:off] +} + +func testDisconnectPayload(t *testing.T, code errorCode) []byte { + t.Helper() + buf := make([]byte, 128) + writer := encoding.NewKVWriter(buf, 0) + if err := writer.SetUInt32("status-code", uint32(code)); err != nil { + t.Fatal(err) + } + if err := writer.SetString("message", code.String()); err != nil { + t.Fatal(err) + } + return writer.Bytes() +} + +func writeTestWireFrame( + t *testing.T, + w io.Writer, + frameType frameType, + streamID uint64, + frameID uint64, + payload []byte, +) { + t.Helper() + body := make([]byte, 1+uint32Len+20+len(payload)) + body[0] = byte(frameType) + binary.BigEndian.PutUint32(body[1:1+uint32Len], uint32(frameFlagFin)) + off := 1 + uint32Len + n, err := encoding.PutVarint(body[off:], streamID) + if err != nil { + t.Fatal(err) + } + off += n + n, err = encoding.PutVarint(body[off:], frameID) + if err != nil { + t.Fatal(err) + } + off += n + copy(body[off:], payload) + body = body[:off+len(payload)] + + wire := make([]byte, uint32Len+len(body)) + binary.BigEndian.PutUint32(wire, uint32(len(body))) + copy(wire[uint32Len:], body) + if _, err := w.Write(wire); err != nil { + t.Fatalf("write frame: %v", err) + } +} + +func readTestWireFrame(t *testing.T, r io.Reader) testWireFrame { + t.Helper() + lengthBytes := make([]byte, uint32Len) + if _, err := io.ReadFull(r, lengthBytes); err != nil { + t.Fatalf("read frame length: %v", err) + } + length := binary.BigEndian.Uint32(lengthBytes) + body := make([]byte, length) + if _, err := io.ReadFull(r, body); err != nil { + t.Fatalf("read frame body: %v", err) + } + if len(body) < 1+uint32Len { + t.Fatalf("frame body is too short: %d", len(body)) + } + + off := 1 + uint32Len + streamID, n, err := encoding.Varint(body[off:]) + if err != nil { + t.Fatalf("read stream ID: %v", err) + } + off += n + frameID, n, err := encoding.Varint(body[off:]) + if err != nil { + t.Fatalf("read frame ID: %v", err) + } + off += n + + return testWireFrame{ + payload: body[off:], + length: length, + streamID: streamID, + frameID: frameID, + frameType: frameType(body[0]), + } +} + +func testHelloMaxFrameSize(t *testing.T, payload []byte) uint32 { + t.Helper() + scanner := encoding.NewKVScanner(payload, -1) + entry := encoding.AcquireKVEntry() + defer encoding.ReleaseKVEntry(entry) + for scanner.Next(entry) { + if entry.NameEquals(helloKeyMaxFrameSize) { + if entry.Type() != encoding.DataTypeUInt32 { + t.Fatalf("%q has type %d, expected uint32", helloKeyMaxFrameSize, entry.Type()) + } + return uint32(entry.ValueInt()) + } + } + if err := scanner.Error(); err != nil { + t.Fatalf("scan HELLO payload: %v", err) + } + t.Fatalf("HELLO payload is missing %q", helloKeyMaxFrameSize) + return 0 +} + +func testDisconnectCode(t *testing.T, payload []byte) errorCode { + t.Helper() + scanner := encoding.NewKVScanner(payload, -1) + entry := encoding.AcquireKVEntry() + defer encoding.ReleaseKVEntry(entry) + for scanner.Next(entry) { + if entry.NameEquals("status-code") { + if entry.Type() != encoding.DataTypeUInt32 { + t.Fatalf("status-code has type %d, expected uint32", entry.Type()) + } + return errorCode(entry.ValueInt()) + } + } + if err := scanner.Error(); err != nil { + t.Fatalf("scan disconnect payload: %v", err) + } + t.Fatal("disconnect payload is missing status-code") + return ErrorUnknown +} + +func testFrameBytes(t *testing.T, length uint32) *bytes.Reader { + t.Helper() + wire := make([]byte, uint32Len+length) + binary.BigEndian.PutUint32(wire, length) + if length >= 7 { + wire[uint32Len] = byte(frameTypeIDNotify) + binary.BigEndian.PutUint32(wire[uint32Len+1:uint32Len+5], uint32(frameFlagFin)) + } + return bytes.NewReader(wire) +} + +func uint32Ptr(value uint32) *uint32 { + return &value +} diff --git a/spop/scheduler.go b/spop/scheduler.go index 8d5e531..e63cef8 100644 --- a/spop/scheduler.go +++ b/spop/scheduler.go @@ -66,6 +66,7 @@ func (bq *queue) Get() queueElem { } item := bq.elems[bq.head] + bq.elems[bq.head] = queueElem{} bq.head = (bq.head + 1) % len(bq.elems) bq.size-- @@ -96,6 +97,7 @@ func (a *asyncScheduler) queueWorker() { return qe.pc.frameHandler(qe.f) }) if err != nil { + qe.pc.terminate(err) log.Println(err) continue } diff --git a/spop/scheduler_test.go b/spop/scheduler_test.go new file mode 100644 index 0000000..79bb90c --- /dev/null +++ b/spop/scheduler_test.go @@ -0,0 +1,25 @@ +package spop + +import "testing" + +func newTestAsyncScheduler() *asyncScheduler { + scheduler := &asyncScheduler{q: newQueue(2)} + go scheduler.queueWorker() + return scheduler +} + +func TestQueueGetClearsConsumedSlot(t *testing.T) { + q := newQueue(1) + f := acquireFrame() + defer releaseFrame(f) + client := &protocolClient{} + + q.Put(f, client) + got := q.Get() + if got.f != f || got.pc != client { + t.Fatal("queue returned an unexpected element") + } + if q.elems[0].f != nil || q.elems[0].pc != nil { + t.Fatal("queue retained the consumed element") + } +} diff --git a/spop/server_test.go b/spop/server_test.go index 2940e79..5e5bd09 100644 --- a/spop/server_test.go +++ b/spop/server_test.go @@ -3,6 +3,7 @@ package spop import ( "context" "encoding/binary" + "fmt" "io" "log" "math/rand" @@ -15,44 +16,76 @@ import ( func TestFakeCon(t *testing.T) { log.SetFlags(log.LstdFlags | log.Lshortfile) - ctx, cancel := context.WithCancel(context.Background()) + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() pipe, pipeConn := testutil.PipeConn() + defer pipe.Close() + defer pipeConn.Close() + peerDone := make(chan error, 1) go func() { - defer cancel() - if err := newHelloFrame(pipe); err != nil { - t.Error(err) + peerDone <- err + return + } + if err := readExpectedFrame(pipe, frameTypeIDAgentHello); err != nil { + peerDone <- err return } if err := newNotifyFrame(pipe); err != nil { - t.Error(err) + peerDone <- err return } - }() - - go func() { - defer cancel() - - select { - case <-ctx.Done(): - case <-time.After(5 * time.Second): - t.Error("timeout") + if err := readExpectedFrame(pipe, frameTypeIDAck); err != nil { + peerDone <- err + return } + peerDone <- nil }() handler := HandlerFunc(func(_ context.Context, _ *encoding.ActionWriter, m *encoding.Message) { log.Println(m.NameBytes()) - cancel() }) - pc := newProtocolClient(context.Background(), pipeConn, newAsyncScheduler(), handler) - defer pc.Close() - defer pipe.Close() - go pc.Serve() + pc := newProtocolClient(ctx, pipeConn, newTestAsyncScheduler(), defaultFramePool, handler) + serveDone := make(chan error, 1) + go func() { + serveDone <- pc.Serve() + }() + + select { + case err := <-peerDone: + if err != nil { + t.Fatal(err) + } + case <-ctx.Done(): + t.Fatal("peer exchange timed out") + } + if err := pipe.Close(); err != nil { + t.Fatal(err) + } + select { + case err := <-serveDone: + if err != nil { + t.Fatalf("serve protocol: %v", err) + } + case <-ctx.Done(): + t.Fatal("protocol shutdown timed out") + } +} - <-ctx.Done() +func readExpectedFrame(r io.Reader, expected frameType) error { + f := acquireFrame() + defer releaseFrame(f) + + if _, err := f.ReadFrom(r); err != nil { + return err + } + if f.frameType != expected { + return fmt.Errorf("expected frame type %d, got %d", expected, f.frameType) + } + return nil } func newNotifyFrame(wr io.Writer) error { @@ -104,7 +137,13 @@ func newHelloFrame(wr io.Writer) error { w := encoding.AcquireKVWriter(f.buf.WriteBytes(), 0) defer encoding.ReleaseKVWriter(w) - if err := w.SetUInt32(helloKeyMaxFrameSize, maxFrameSize); err != nil { + if err := w.SetString(helloKeySupportedVersions, version); err != nil { + return err + } + if err := w.SetUInt32(helloKeyMaxFrameSize, DefaultMaxFrameSize); err != nil { + return err + } + if err := w.SetString(helloKeyCapabilities, ""); err != nil { return err }