zero/internal/daemon/protocol_test.go

package daemon

import (
	"bytes"
	"encoding/binary"
	"errors"
	"io"
	"strings"
	"testing"
)

func TestFrameRoundTrip(t *testing.T) {
	var buf bytes.Buffer
	payloads := [][]byte{[]byte("hello"), {}, []byte(strings.Repeat("x", 4096))}
	for _, p := range payloads {
		if err := WriteFrame(&buf, KindCtrl, p); err != nil {
			t.Fatalf("WriteFrame: %v", err)
		}
	}
	for i, want := range payloads {
		kind, got, err := ReadFrame(&buf)
		if err != nil {
			t.Fatalf("ReadFrame[%d]: %v", i, err)
		}
		if kind != KindCtrl {
			t.Fatalf("frame[%d] kind = %d, want KindCtrl", i, kind)
		}
		if !bytes.Equal(got, want) {
			t.Fatalf("frame[%d] payload = %q, want %q", i, got, want)
		}
	}
}

func TestControlRoundTrip(t *testing.T) {
	var buf bytes.Buffer
	msg := Ctrl{Type: CtrlRun, Session: "s1", Cwd: "/tmp", Args: []string{"--prompt", "hi"}}
	if err := WriteControl(&buf, msg); err != nil {
		t.Fatalf("WriteControl: %v", err)
	}
	got, err := ReadControl(&buf)
	if err != nil {
		t.Fatalf("ReadControl: %v", err)
	}
	if got.Type != CtrlRun || got.Session != "s1" || got.Cwd != "/tmp" || len(got.Args) != 2 {
		t.Fatalf("round-trip mismatch: %+v", got)
	}
}

func TestWriteFrameRejectsOversize(t *testing.T) {
	var buf bytes.Buffer
	err := WriteFrame(&buf, KindCtrl, make([]byte, MaxFrameSize+1))
	if !errors.Is(err, ErrFrameTooLarge) {
		t.Fatalf("WriteFrame oversize err = %v, want ErrFrameTooLarge", err)
	}
	if buf.Len() != 0 {
		t.Fatalf("oversize WriteFrame must not emit any bytes, wrote %d", buf.Len())
	}
}

func TestReadFrameRejectsOversizeBeforeAlloc(t *testing.T) {
	// A header advertising a 4 GiB payload must be rejected without reading or
	// allocating that payload.
	header := make([]byte, frameHeaderSize)
	binary.BigEndian.PutUint32(header[:4], MaxFrameSize+1)
	header[4] = byte(KindCtrl)
	_, _, err := ReadFrame(bytes.NewReader(header))
	if !errors.Is(err, ErrFrameTooLarge) {
		t.Fatalf("ReadFrame oversize err = %v, want ErrFrameTooLarge", err)
	}
}

func TestReadFrameUnknownKind(t *testing.T) {
	var buf bytes.Buffer
	// Hand-write a frame with an unknown kind byte (7) and a small payload.
	header := make([]byte, frameHeaderSize)
	binary.BigEndian.PutUint32(header[:4], 2)
	header[4] = 7
	buf.Write(header)
	buf.Write([]byte("ab"))
	_, _, err := ReadFrame(&buf)
	if !errors.Is(err, ErrUnknownFrameKind) {
		t.Fatalf("unknown kind err = %v, want ErrUnknownFrameKind", err)
	}
}

func TestReadFrameTruncated(t *testing.T) {
	// Header promises 10 bytes but only 3 follow → unexpected EOF.
	var buf bytes.Buffer
	header := make([]byte, frameHeaderSize)
	binary.BigEndian.PutUint32(header[:4], 10)
	header[4] = byte(KindCtrl)
	buf.Write(header)
	buf.Write([]byte("abc"))
	_, _, err := ReadFrame(&buf)
	if err == nil || errors.Is(err, ErrFrameTooLarge) {
		t.Fatalf("truncated frame err = %v, want a read error", err)
	}
	if !errors.Is(err, io.ErrUnexpectedEOF) {
		t.Fatalf("truncated frame err = %v, want io.ErrUnexpectedEOF", err)
	}
}

func TestNegotiateVersion(t *testing.T) {
	cases := []struct {
		client int
		want   int
		wantOK bool
	}{
		{client: 1, want: 1, wantOK: true},
		{client: 99, want: ProtoVersion, wantOK: true}, // clamp down to ours
		{client: 0, want: 0, wantOK: false},
		{client: -3, want: 0, wantOK: false},
	}
	for _, tc := range cases {
		got, ok := NegotiateVersion(tc.client)
		if got != tc.want || ok != tc.wantOK {
			t.Fatalf("NegotiateVersion(%d) = (%d,%v), want (%d,%v)", tc.client, got, ok, tc.want, tc.wantOK)
		}
	}
}