pizzasql/pkg/pgserver/protocol_test.go

package pgserver

import (
	"bytes"
	"encoding/binary"
	"testing"
)

func TestReadMessageRejectsOversizedMessage(t *testing.T) {
	var buf bytes.Buffer
	buf.WriteByte('Q')
	if err := binary.Write(&buf, binary.BigEndian, uint32(MaxMessageSize+1)); err != nil {
		t.Fatal(err)
	}
	if _, err := ReadMessage(&buf); err == nil {
		t.Fatal("expected error for message larger than MaxMessageSize")
	}
}

func TestReadMessageAcceptsNormalMessage(t *testing.T) {
	payload := []byte("SELECT 1\x00")
	var buf bytes.Buffer
	buf.WriteByte(MsgQuery)
	if err := binary.Write(&buf, binary.BigEndian, uint32(len(payload)+4)); err != nil {
		t.Fatal(err)
	}
	buf.Write(payload)

	msg, err := ReadMessage(&buf)
	if err != nil {
		t.Fatalf("ReadMessage: %v", err)
	}
	if msg.Type != MsgQuery {
		t.Fatalf("type = %c, want %c", msg.Type, MsgQuery)
	}
	if !bytes.Equal(msg.Data, payload) {
		t.Fatalf("data = %q, want %q", msg.Data, payload)
	}
}

func TestReadStartupMessageRejectsOversizedMessage(t *testing.T) {
	var buf bytes.Buffer
	if err := binary.Write(&buf, binary.BigEndian, uint32(MaxStartupMessageSize+1)); err != nil {
		t.Fatal(err)
	}
	if _, err := ReadStartupMessage(&buf); err == nil {
		t.Fatal("expected error for startup message larger than MaxStartupMessageSize")
	}
}

func TestReadStartupMessageAcceptsNormalMessage(t *testing.T) {
	// Build a minimal startup payload: protocol version + user param.
	var payload bytes.Buffer
	if err := binary.Write(&payload, binary.BigEndian, uint32(196608)); err != nil { // 3.0
		t.Fatal(err)
	}
	payload.WriteString("user\x00tester\x00\x00")

	var buf bytes.Buffer
	if err := binary.Write(&buf, binary.BigEndian, uint32(payload.Len()+4)); err != nil {
		t.Fatal(err)
	}
	buf.Write(payload.Bytes())

	params, err := ReadStartupMessage(&buf)
	if err != nil {
		t.Fatalf("ReadStartupMessage: %v", err)
	}
	if params["user"] != "tester" {
		t.Fatalf("user = %q, want %q", params["user"], "tester")
	}
	if params["protocol_version"] != "196608" {
		t.Fatalf("protocol_version = %q, want %q", params["protocol_version"], "196608")
	}
}