t5de-proxy/internal/imq/frame_test.go

package imq

import (
	"bytes"
	"errors"
	"testing"
)

// frame encodes a frame as the client does: field 1 the type, field 2 the data.
func frame(t uint64, data []byte) []byte {
	b := append([]byte{0x08}, uvarint(t)...)
	b = append(b, 0x12)
	b = append(b, uvarint(uint64(len(data)))...)
	return append(b, data...)
}

func uvarint(v uint64) []byte {
	var b []byte
	for v >= 0x80 {
		b = append(b, byte(v)|0x80)
		v >>= 7
	}
	return append(b, byte(v))
}

func TestReaderSplitsFrames(t *testing.T) {
	stream := append(frame(1302, []byte("hello")), frame(3103, bytes.Repeat([]byte{1}, 300))...)
	stream = append(stream, frame(1301, nil)...)
	// Fed a byte at a time, the frames still come out whole and in order.
	var r Reader
	var got []Frame
	for i := range stream {
		frames, err := r.Write(stream[i : i+1])
		if err != nil {
			t.Fatal(err)
		}
		got = append(got, frames...)
	}
	want := []Frame{{1302, 5, []byte("hello")}, {3103, 300, bytes.Repeat([]byte{1}, 300)}, {1301, 0, nil}}
	if len(got) != len(want) {
		t.Fatalf("got %v, want %v", got, want)
	}
	for i := range want {
		if got[i].Type != want[i].Type || got[i].Size != want[i].Size || !bytes.Equal(got[i].Data, want[i].Data) {
			t.Fatalf("got %v, want %v", got, want)
		}
	}
}

func TestReaderRejectsOtherStreams(t *testing.T) {
	var r Reader
	// A TLS ClientHello isn't IMQ.
	if _, err := r.Write([]byte{0x16, 0x03, 0x01, 0x00, 0xa5}); !errors.Is(err, ErrMalformed) {
		t.Fatalf("got %v, want ErrMalformed", err)
	}
	if _, err := r.Write(frame(1301, nil)); !errors.Is(err, ErrMalformed) {
		t.Fatal("a reader keeps failing after a malformed frame")
	}
}

func TestTypeName(t *testing.T) {
	if TypeName(3104) != "G2cSendMessage" || TypeName(7) != "type 7" {
		t.Fatal(TypeName(3104), TypeName(7))
	}
}

func TestSendMessageRoundTrip(t *testing.T) {
	m := SendMessage{OpID: 300, Queue: "/chat/1", Mount: "messages", Message: []byte(`{"message": "*use 80 1"}`)}
	data := m.Marshal()
	got, err := ParseSendMessage(data)
	if err != nil || got.OpID != 300 || got.Queue != "/chat/1" || got.Mount != "messages" || string(got.Message) != string(m.Message) {
		t.Fatalf("%+v %v", got, err)
	}
	// A frame of it reads back as one frame with the same data.
	var r Reader
	frames, err := r.Write(EncodeFrame(C2gSendMessage, data))
	if err != nil || len(frames) != 1 || frames[0].Type != C2gSendMessage || !bytes.Equal(frames[0].Data, data) {
		t.Fatalf("%v %v", frames, err)
	}
	if _, err := ParseSendMessage(Result(1, 0)); err == nil {
		t.Fatal("a G2cResult parsed as a C2gSendMessage")
	}
}

func TestSubscriptionOps(t *testing.T) {
	// A subscription with a plain queue, and one with an op ID and a property (field 3), which survives re-encoding.
	withProps := appendBytesField(appendVarintField(appendBytesField(nil, 1, []byte("/chat/9")), 2, 7), 3, []byte("x"))
	data := appendBytesField(appendBytesField(nil, 1, []byte("/user/1")), 2, withProps)
	s, err := ParseSubscriptionOps(data)
	if err != nil || len(s.Queues) != 1 || len(s.WithResults) != 1 || s.WithResults[0].Queue != "/chat/9" || s.WithResults[0].OpID != 7 {
		t.Fatalf("%+v %v", s, err)
	}
	if string(s.Marshal()) != string(data) {
		t.Fatal("re-encoded differently")
	}
	if all := s.All(); len(all) != 2 || all[1] != "/chat/9" {
		t.Fatalf("all %v", all)
	}
}

func TestFloodgates(t *testing.T) {
	state := func(q string, seq, op uint64) []byte {
		return appendVarintField(appendVarintField(appendBytesField(nil, 1, []byte(q)), 2, seq), 3, op)
	}
	data := appendBytesField(appendBytesField(nil, 1, state("/user/1", 38, 19)), 1, state("/chat/2100000001", 29, 20))
	states, err := ParseFloodgates(data)
	if err != nil || len(states) != 2 || states[1].Queue != "/chat/2100000001" || states[1].OpID != 20 {
		t.Fatalf("%+v %v", states, err)
	}
	if string(MarshalFloodgates(states)) != string(data) || string(MarshalFloodgates(states[:1])) != string(appendBytesField(nil, 1, state("/user/1", 38, 19))) {
		t.Fatal("re-encoded differently")
	}
}