rdp-proxy/agents/rap-node-agent/internal/fabricproto/session_frames_test.go

package fabricproto

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

func TestHandleFrameOpensStreamAndReceivesData(t *testing.T) {
	session := NewSession(SessionConfig{})

	event, responses, err := session.HandleFrame(Frame{
		Type:         FrameOpenStream,
		TrafficClass: TrafficClassInteractive,
		StreamID:     7,
	})
	if err != nil {
		t.Fatalf("handle open stream: %v", err)
	}
	if event.Type != SessionEventStreamOpened || event.StreamID != 7 || len(responses) != 0 {
		t.Fatalf("open event/responses = %+v / %+v", event, responses)
	}

	event, responses, err = session.HandleFrame(Frame{
		Type:         FrameData,
		TrafficClass: TrafficClassInteractive,
		StreamID:     7,
		Sequence:     11,
		Payload:      []byte("rdp-input"),
	})
	if err != nil {
		t.Fatalf("handle data: %v", err)
	}
	if event.Type != SessionEventData || string(event.Payload) != "rdp-input" {
		t.Fatalf("data event = %+v", event)
	}
	if len(responses) != 1 || responses[0].Type != FrameAck || responses[0].StreamID != 7 || responses[0].Sequence != 11 {
		t.Fatalf("responses = %+v, want ack for stream 7 seq 11", responses)
	}
	if !bytes.Equal(responses[0].Payload, DataAckPayload([]byte("rdp-input"))) {
		t.Fatalf("ack checksum = %x, want sha256 payload checksum", responses[0].Payload)
	}
	snapshot := session.Snapshot()
	if snapshot.FramesReceived != 1 || snapshot.Streams[7].Received != 1 {
		t.Fatalf("received metrics = %+v stream=%+v", snapshot, snapshot.Streams[7])
	}
}

func TestHandleFramePingReturnsPong(t *testing.T) {
	session := NewSession(SessionConfig{})

	event, responses, err := session.HandleFrame(Frame{
		Type:     FramePing,
		Sequence: 99,
		Payload:  []byte("probe"),
	})
	if err != nil {
		t.Fatalf("handle ping: %v", err)
	}
	if event.Type != SessionEventPing {
		t.Fatalf("event = %+v, want ping", event)
	}
	if len(responses) != 1 || responses[0].Type != FramePong || responses[0].Sequence != 99 || string(responses[0].Payload) != "probe" {
		t.Fatalf("responses = %+v, want pong echo", responses)
	}
}

func TestHandleFrameCreditUnblocksOutgoingStream(t *testing.T) {
	session := NewSession(SessionConfig{InitialStreamCredit: 1})
	mustOpenStream(t, session, 1, TrafficClassReliable)
	mustEnqueue(t, session, 1, "one")
	if _, err := session.EnqueueData(1, []byte("blocked")); !errors.Is(err, ErrStreamCreditExhausted) {
		t.Fatalf("credit error = %v, want %v", err, ErrStreamCreditExhausted)
	}

	event, responses, err := session.HandleFrame(Frame{
		Type:     FrameStreamCredit,
		StreamID: 1,
		Sequence: 2,
	})
	if err != nil {
		t.Fatalf("handle credit: %v", err)
	}
	if event.Type != SessionEventCredit || len(responses) != 0 {
		t.Fatalf("credit event/responses = %+v / %+v", event, responses)
	}
	mustEnqueue(t, session, 1, "two")
}

func TestHandleFrameAckUpdatesSession(t *testing.T) {
	session := NewSession(SessionConfig{})
	mustOpenStream(t, session, 1, TrafficClassReliable)
	mustEnqueue(t, session, 1, "one")
	mustEnqueue(t, session, 1, "two")

	event, responses, err := session.HandleFrame(Frame{
		Type:     FrameAck,
		StreamID: 1,
		Sequence: 2,
	})
	if err != nil {
		t.Fatalf("handle ack: %v", err)
	}
	if event.Type != SessionEventAck || len(responses) != 0 {
		t.Fatalf("ack event/responses = %+v / %+v", event, responses)
	}
	snapshot := session.Snapshot()
	if snapshot.Streams[1].Acked != 2 || snapshot.FramesAcked != 2 {
		t.Fatalf("ack metrics = %+v stream=%+v", snapshot, snapshot.Streams[1])
	}
}

func TestHandleFrameResetDoesNotAffectOtherStreams(t *testing.T) {
	session := NewSession(SessionConfig{})
	mustOpenStream(t, session, 1, TrafficClassBulk)
	mustOpenStream(t, session, 2, TrafficClassControl)
	mustEnqueue(t, session, 1, "bulk")
	mustEnqueue(t, session, 2, "control")

	event, _, err := session.HandleFrame(Frame{Type: FrameResetStream, StreamID: 1})
	if err != nil {
		t.Fatalf("handle reset: %v", err)
	}
	if event.Type != SessionEventStreamReset {
		t.Fatalf("event = %+v, want reset", event)
	}
	frame, ok := session.DequeueNext()
	if !ok || frame.StreamID != 2 {
		t.Fatalf("dequeued = %+v ok=%v, want control stream 2", frame, ok)
	}
}

func TestHandleFrameGoAwayClosesSession(t *testing.T) {
	session := NewSession(SessionConfig{})
	mustOpenStream(t, session, 1, TrafficClassReliable)

	event, _, err := session.HandleFrame(Frame{Type: FrameGoAway})
	if err != nil {
		t.Fatalf("handle goaway: %v", err)
	}
	if event.Type != SessionEventGoAway {
		t.Fatalf("event = %+v, want goaway", event)
	}
	if _, err := session.EnqueueData(1, []byte("after-goaway")); !errors.Is(err, ErrSessionClosed) {
		t.Fatalf("enqueue err = %v, want %v", err, ErrSessionClosed)
	}
}

func TestHandleFrameRejectsDataForMissingStream(t *testing.T) {
	session := NewSession(SessionConfig{})

	_, _, err := session.HandleFrame(Frame{
		Type:         FrameData,
		TrafficClass: TrafficClassReliable,
		StreamID:     44,
		Sequence:     1,
	})
	if !errors.Is(err, ErrStreamNotFound) {
		t.Fatalf("error = %v, want %v", err, ErrStreamNotFound)
	}
}

func TestHandleFrameRejectsTrafficClassMismatch(t *testing.T) {
	session := NewSession(SessionConfig{})
	mustOpenStream(t, session, 3, TrafficClassInteractive)

	_, _, err := session.HandleFrame(Frame{
		Type:         FrameData,
		TrafficClass: TrafficClassBulk,
		StreamID:     3,
		Sequence:     1,
	})
	if !errors.Is(err, ErrTrafficClassMismatch) {
		t.Fatalf("error = %v, want %v", err, ErrTrafficClassMismatch)
	}
}