test(spop): cover frame size handshake

spop/protocol_test.go

@@ -0,0 +1,110 @@
+package spop
+
+import (
+	"bytes"
+	"context"
+	"encoding/binary"
+	"testing"
+
+	"github.com/dropmorepackets/haproxy-go/pkg/encoding"
+)
+
+func TestProtocolMaxFrameSizeOffer(t *testing.T) {
+	tests := []struct {
+		name    string
+		offer   uint32
+		wantErr bool
+	}{
+		{name: "at current limit", offer: uint32(maxFrameSize)},
+		{name: "above current limit", offer: 262140, wantErr: true},
+		{name: "maximum uint32", offer: ^uint32(0), wantErr: true},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			var rw bytes.Buffer
+			client := newProtocolClient(
+				context.Background(),
+				&rw,
+				nil,
+				HandlerFunc(func(context.Context, *encoding.ActionWriter, *encoding.Message) {}),
+			)
+			frame := testHelloFrame(t, tt.offer)
+			defer releaseFrame(frame)
+
+			err := client.onHAProxyHello(frame)
+			if tt.wantErr {
+				if err == nil {
+					t.Fatal("expected max-frame-size offer to be rejected")
+				}
+				return
+			}
+			if err != nil {
+				t.Fatalf("handle HAPROXY-HELLO: %v", err)
+			}
+			if client.maxFrameSize != tt.offer {
+				t.Fatalf("expected negotiated size %d, got %d", tt.offer, client.maxFrameSize)
+			}
+			if got := agentHelloMaxFrameSize(t, &rw); got != tt.offer {
+				t.Fatalf("expected AGENT-HELLO size %d, got %d", tt.offer, got)
+			}
+		})
+	}
+}
+
+func agentHelloMaxFrameSize(t *testing.T, rw *bytes.Buffer) uint32 {
+	t.Helper()
+	f := acquireFrame()
+	defer releaseFrame(f)
+	if _, err := f.ReadFrom(rw); err != nil {
+		t.Fatalf("read AGENT-HELLO: %v", err)
+	}
+	if f.frameType != frameTypeIDAgentHello {
+		t.Fatalf("expected AGENT-HELLO, got frame type %d", f.frameType)
+	}
+
+	s := encoding.AcquireKVScanner(f.buf.ReadBytes(), -1)
+	defer encoding.ReleaseKVScanner(s)
+	k := encoding.AcquireKVEntry()
+	defer encoding.ReleaseKVEntry(k)
+	for s.Next(k) {
+		if k.NameEquals(helloKeyMaxFrameSize) {
+			return uint32(k.ValueInt())
+		}
+	}
+	if err := s.Error(); err != nil {
+		t.Fatalf("scan AGENT-HELLO: %v", err)
+	}
+	t.Fatal("AGENT-HELLO missing max-frame-size")
+	return 0
+}
+
+func testHelloFrame(t *testing.T, offer uint32) *frame {
+	t.Helper()
+	f := acquireFrame()
+	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 err := writer.SetUInt32(helloKeyMaxFrameSize, offer); 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
+}

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, newAsyncScheduler(), 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,9 +137,15 @@ func newHelloFrame(wr io.Writer) error {
 	w := encoding.AcquireKVWriter(f.buf.WriteBytes(), 0)
 	defer encoding.ReleaseKVWriter(w)
 
+	if err := w.SetString(helloKeySupportedVersions, version); err != nil {
+		return err
+	}
 	if err := w.SetUInt32(helloKeyMaxFrameSize, maxFrameSize); err != nil {
 		return err
 	}
+	if err := w.SetString(helloKeyCapabilities, ""); err != nil {
+		return err
+	}
 
 	// TODO
 	if err := w.SetString(helloKeyEngineID, "random engine"); err != nil {