VerseVDI-Data-Plane/gateway/gateway_test.go

package gateway

import (
	"context"
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/tls"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/hex"
	"errors"
	"math/big"
	"net"
	"os"
	"strings"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	protocol "git.sechmachine.io.vn/sechmachine/VerseVDI-Protocol/gen/go/protocol"
)

func TestFrameValidationAndFragmentation(t *testing.T) {
	frames, err := FragmentPayload(ChannelVideo, 7, 11, make([]byte, 1180))
	if err != nil || len(frames) != 2 || len(frames[0].Payload) != 1179 || len(frames[1].Payload) != 1 {
		t.Fatalf("fragmentation = %#v, err = %v", frames, err)
	}
	encoded, err := EncodeFrame(frames[0])
	if err != nil {
		t.Fatal(err)
	}
	decoded, err := DecodeFrame(encoded)
	if err != nil || string(decoded.Payload) != string(frames[0].Payload) {
		t.Fatalf("decoded = %#v, err = %v", decoded, err)
	}
	for _, raw := range [][]byte{
		{0x56, 0x44},
		append([]byte(nil), encoded[:len(encoded)-1]...),
		append(append([]byte(nil), encoded...), 0),
	} {
		if err := ValidateFrame(raw); err == nil {
			t.Fatalf("accepted malformed frame %x", raw)
		}
	}
}

func FuzzDecodeFrame(f *testing.F) {
	seed, _ := hex.DecodeString("5644010a0000000000000000000000000000010000")
	f.Add(seed)
	f.Add([]byte("not-a-frame"))
	f.Fuzz(func(t *testing.T, data []byte) {
		_, _ = DecodeFrame(data)
	})
}

func FuzzDecodeControlPacket(f *testing.F) {
	seed, _ := EncodeControlPacket(ControlPacket{Kind: 1, Sequence: 2, Payload: []byte("fixture")})
	f.Add(seed)
	f.Add([]byte("APC1"))
	f.Fuzz(func(t *testing.T, data []byte) {
		_, _ = DecodeControlPacket(data)
	})
}

func FuzzDecodeInputEvent(f *testing.F) {
	seed, _ := EncodeInputEvent(InputEvent{Sequence: 1, Device: "keyboard", Code: 7, Pressed: true})
	f.Add(seed)
	f.Add([]byte("INP1"))
	f.Fuzz(func(t *testing.T, data []byte) {
		_, _ = DecodeInputEvent(data)
	})
}

func TestCapabilityIntersectionAndBoundedQueue(t *testing.T) {
	capabilities := DefaultCapabilities()
	if _, err := IntersectCapabilities(capabilities, capabilities); err != nil {
		t.Fatal(err)
	}
	other := capabilities
	other.Audio = "different"
	if !errors.Is(func() error { _, err := IntersectCapabilities(capabilities, other); return err }(), ErrNoCapabilityOverlap) {
		t.Fatal("capability mismatch was accepted")
	}
	queue := NewBoundedQueue[int](2)
	_ = queue.PushLatest(1)
	_ = queue.PushLatest(2)
	_ = queue.PushLatest(3)
	if queue.Dropped() != 1 || queue.Len() != 2 {
		t.Fatalf("queue length=%d dropped=%d", queue.Len(), queue.Dropped())
	}
	ctx := context.Background()
	first, _ := queue.Pop(ctx)
	second, _ := queue.Pop(ctx)
	if first != 2 || second != 3 {
		t.Fatalf("queue values=%d,%d", first, second)
	}
}

func TestSyntheticImpairmentPacingAndResourceBounds(t *testing.T) {
	payload := make([]byte, 1179*16+1)
	if _, err := FragmentPayload(ChannelVideo, 1, 0, payload); !errors.Is(err, ErrFrameFragmentedLimit) {
		t.Fatalf("oversized media payload accepted: %v", err)
	}
	frames, err := FragmentPayload(ChannelVideo, 1, 0, bytesRepeat(0x5a, 1179*4))
	if err != nil {
		t.Fatal(err)
	}
	var delivered int
	var deliveredBytes int
	for index, frame := range frames {
		if (index+1)%3 == 0 { // deterministic synthetic loss profile: every third frame.
			continue
		}
		delivered++
		encoded, encodeErr := EncodeFrame(frame)
		if encodeErr != nil {
			t.Fatal(encodeErr)
		}
		decoded, decodeErr := DecodeFrame(encoded)
		if decodeErr != nil {
			t.Fatal(decodeErr)
		}
		deliveredBytes += len(decoded.Payload)
	}
	if delivered != 3 || deliveredBytes != 1179*3 {
		t.Fatalf("synthetic impairment delivered=%d bytes=%d", delivered, deliveredBytes)
	}
	pacer := NewPacer(1)
	ctx, cancel := context.WithTimeout(context.Background(), time.Millisecond)
	defer cancel()
	if err := pacer.Wait(ctx, 100); !errors.Is(err, context.DeadlineExceeded) {
		t.Fatalf("pacer ignored bounded context: %v", err)
	}
}

func TestApolloFixturesAndLifecycle(t *testing.T) {
	management, err := os.ReadFile("testdata/apollo-management.xml")
	if err != nil {
		t.Fatal(err)
	}
	info, err := ParseManagementXML(management)
	if err != nil || info.Identity.UniqueID != "apollo-fixture-1" {
		t.Fatalf("management = %#v, err = %v", info, err)
	}
	rtspText, err := os.ReadFile("testdata/rtsp-setup-response.txt")
	if err != nil {
		t.Fatal(err)
	}
	rtsp, err := ParseRTSPResponse([]byte(strings.ReplaceAll(string(rtspText), `\r\n`, "\r\n")))
	if err != nil || rtsp.StatusCode != 200 {
		t.Fatalf("RTSP = %#v, err = %v", rtsp, err)
	}
	video, _ := hex.DecodeString(strings.TrimSpace(string(mustRead(t, "testdata/encoded-video.hex"))))
	audio, _ := hex.DecodeString(strings.TrimSpace(string(mustRead(t, "testdata/encoded-audio.hex"))))
	now := time.Date(2026, time.January, 1, 0, 0, 0, 0, time.UTC)
	identity := ProviderIdentity{UniqueID: "apollo-fixture-1", Fingerprint: "sha256:fixture-apollo-1"}
	fake := NewFakeApollo(FakeApolloConfig{Identity: identity, Now: now, Video: [][]byte{video}, Audio: [][]byte{audio}})
	session, err := fake.Start(context.Background(), LaunchRequest{SessionID: "session-1", ProviderProfile: ProviderProfileApollo, ProviderIdentity: identity.Key(), Capabilities: DefaultCapabilities()})
	if err != nil {
		t.Fatal(err)
	}
	if got := <-session.Video(); string(got) != string(video) {
		t.Fatalf("video changed: %x", got)
	}
	if got := <-session.Audio(); string(got) != string(audio) {
		t.Fatalf("audio changed: %x", got)
	}
	if err := session.Input(context.Background(), InputEvent{Sequence: 1, Device: "keyboard", Code: 7, Pressed: true}); err != nil {
		t.Fatal(err)
	}
	if err := session.ReleaseAll(context.Background()); err != nil {
		t.Fatal(err)
	}
	if err := session.Terminate(context.Background()); err != nil {
		t.Fatal(err)
	}
	if state := session.State(); state.State != ProviderStateTerminated || state.CleanupPending {
		t.Fatalf("state = %#v", state)
	}

	identityFailure := NewFakeApollo(FakeApolloConfig{Identity: identity, Now: now, Failure: FakeFailureIdentity})
	if _, err := identityFailure.Start(context.Background(), LaunchRequest{SessionID: "session-2", ProviderProfile: ProviderProfileApollo, ProviderIdentity: identity.Key(), Capabilities: DefaultCapabilities()}); !errors.Is(err, ErrProviderIdentity) {
		t.Fatalf("identity failure = %v", err)
	}
}

func TestProviderTimeoutAndBoundedInput(t *testing.T) {
	now := time.Date(2026, time.January, 1, 0, 0, 0, 0, time.UTC)
	fake := NewFakeApollo(FakeApolloConfig{Now: now, Failure: FakeFailureReadinessTimeout})
	ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond)
	defer cancel()
	started := time.Now()
	_, err := fake.Start(ctx, LaunchRequest{SessionID: "session-timeout", ProviderProfile: ProviderProfileApollo, ProviderIdentity: fake.config.Identity.Key(), Capabilities: DefaultCapabilities()})
	if !errors.Is(err, ErrProviderTimeout) || time.Since(started) > time.Second {
		t.Fatalf("readiness timeout = %v after %s", err, time.Since(started))
	}
	if _, err := EncodeInputEvent(InputEvent{Device: strings.Repeat("d", 65)}); !errors.Is(err, ErrInputMalformed) {
		t.Fatalf("oversized input accepted: %v", err)
	}
	if _, err := DecodeControlPacket([]byte("APC1")); !errors.Is(err, ErrProviderMalformed) {
		t.Fatalf("truncated control accepted: %v", err)
	}
}

func TestNativeApolloEncodedRelay(t *testing.T) {
	provider, peer := net.Pipe()
	session := newNativeApolloSession(provider, "session-native")
	go session.readMedia()
	go func() {
		_, _ = peer.Write([]byte{'$', 0, 0, 3, 1, 2, 3})
		_, _ = peer.Write([]byte{'$', 1, 0, 2, 4, 5})
	}()
	select {
	case payload := <-session.Video():
		if string(payload) != string([]byte{1, 2, 3}) {
			t.Fatalf("video payload changed: %x", payload)
		}
	case <-time.After(time.Second):
		t.Fatal("video payload not relayed")
	}
	select {
	case payload := <-session.Audio():
		if string(payload) != string([]byte{4, 5}) {
			t.Fatalf("audio payload changed: %x", payload)
		}
	case <-time.After(time.Second):
		t.Fatal("audio payload not relayed")
	}
	terminateCtx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond)
	defer cancel()
	_ = session.Terminate(terminateCtx)
	_ = peer.Close()
}

func TestAdmissionQUICMTLSRelayAndCleanup(t *testing.T) {
	serverTLS, clientTLS := testTLS(t)
	fake := NewFakeApollo(FakeApolloConfig{Now: time.Date(2026, time.January, 1, 0, 0, 0, 0, time.UTC)})
	authority := protocol.SessionAuthority{Version: "1", SessionID: "session-1", GatewayID: "gateway-1", Audience: "versevdi-gateway", ReconnectSequence: 0, ExpiresAt: time.Now().Add(5 * time.Second).UTC().Format(time.RFC3339Nano), Capabilities: DefaultCapabilities(), ProviderProfile: ProviderProfileApollo, ProviderIdentity: fake.config.Identity.Key()}
	admission := &oneTimeAdmission{authority: authority, released: make(chan struct{})}
	reporter := &recordingProviderStateReporter{}
	server, err := NewServer(ServerConfig{ListenAddress: "127.0.0.1:0", TLSConfig: serverTLS, GatewayID: "gateway-1", Capabilities: DefaultCapabilities(), ProviderCapabilities: DefaultCapabilities(), Admission: admission, ProviderStateReporter: reporter, Provider: fake})
	if err != nil {
		t.Fatal(err)
	}
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	serveDone := make(chan error, 1)
	go func() { serveDone <- server.Serve(ctx) }()
	request := protocol.TunnelAdmissionRequest{Version: "1", SessionID: "session-1", GatewayID: "gateway-1", Audience: "versevdi-gateway", Grant: strings.Repeat("g", 64), ReconnectSequence: 0, ClientNonce: "nonce-0000000001", DeviceSignature: strings.Repeat("s", 86), Capabilities: DefaultCapabilities()}
	client, err := Dial(context.Background(), server.Addr().String(), clientTLS, request)
	if err != nil {
		_ = server.Close()
		t.Fatal(err)
	}
	for i := 0; i < 2; i++ {
		frame, receiveErr := client.ReceiveFrame(context.Background())
		if receiveErr != nil {
			t.Fatal(receiveErr)
		}
		if frame.Channel != ChannelVideo && frame.Channel != ChannelAudio {
			t.Fatalf("unexpected media channel %d", frame.Channel)
		}
	}
	if err := client.SendInput(InputEvent{Sequence: 1, Device: "keyboard", Code: 7, Pressed: true}); err != nil {
		t.Fatal(err)
	}
	_ = client.Close()
	deadline := time.NewTimer(2 * time.Second)
	select {
	case <-admission.released:
	case <-deadline.C:
		t.Fatal("gateway did not release admission")
	}
	deadline.Stop()
	if fake.LastSession().State().State != ProviderStateTerminated {
		t.Fatalf("provider state = %#v", fake.LastSession().State())
	}
	if _, err := Dial(context.Background(), server.Addr().String(), clientTLS, request); err == nil {
		t.Fatal("replayed grant was accepted")
	}
	_ = server.Close()
	if err := <-serveDone; err != nil {
		t.Fatal(err)
	}
	states := reporter.States()
	if len(states) != 3 || states[0].State != ProviderStateStarting || states[1].State != ProviderStateReady || states[2].State != ProviderStateTerminated {
		t.Fatalf("provider states = %#v", states)
	}
}

func testTLS(t *testing.T) (*tls.Config, *tls.Config) {
	t.Helper()
	caKey, _ := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	caTemplate := &x509.Certificate{SerialNumber: big.NewInt(1), Subject: pkix.Name{CommonName: "Verse Test CA"}, NotBefore: time.Now().Add(-time.Hour), NotAfter: time.Now().Add(time.Hour), IsCA: true, BasicConstraintsValid: true, KeyUsage: x509.KeyUsageCertSign | x509.KeyUsageDigitalSignature}
	caDER, err := x509.CreateCertificate(rand.Reader, caTemplate, caTemplate, &caKey.PublicKey, caKey)
	if err != nil {
		t.Fatal(err)
	}
	caCert, err := x509.ParseCertificate(caDER)
	if err != nil {
		t.Fatal(err)
	}
	makeLeaf := func(serial int64, dns string, usage x509.ExtKeyUsage) tls.Certificate {
		key, keyErr := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
		if keyErr != nil {
			t.Fatal(keyErr)
		}
		template := &x509.Certificate{SerialNumber: big.NewInt(serial), Subject: pkix.Name{CommonName: dns}, DNSNames: []string{dns}, IPAddresses: []net.IP{net.ParseIP("127.0.0.1")}, NotBefore: time.Now().Add(-time.Hour), NotAfter: time.Now().Add(time.Hour), ExtKeyUsage: []x509.ExtKeyUsage{usage}, KeyUsage: x509.KeyUsageDigitalSignature}
		der, createErr := x509.CreateCertificate(rand.Reader, template, caCert, &key.PublicKey, caKey)
		if createErr != nil {
			t.Fatal(createErr)
		}
		return tls.Certificate{Certificate: [][]byte{der, caDER}, PrivateKey: key}
	}
	serverCert := makeLeaf(2, "gateway.test", x509.ExtKeyUsageServerAuth)
	clientCert := makeLeaf(3, "client.test", x509.ExtKeyUsageClientAuth)
	pool := x509.NewCertPool()
	pool.AddCert(caCert)
	return &tls.Config{MinVersion: tls.VersionTLS13, Certificates: []tls.Certificate{serverCert}, ClientAuth: tls.RequireAndVerifyClientCert, ClientCAs: pool}, &tls.Config{MinVersion: tls.VersionTLS13, Certificates: []tls.Certificate{clientCert}, RootCAs: pool, ServerName: "gateway.test"}
}

type oneTimeAdmission struct {
	used      atomic.Bool
	authority protocol.SessionAuthority
	releases  atomic.Int64
	released  chan struct{}
}

type recordingProviderStateReporter struct {
	mu     sync.Mutex
	states []protocol.ProviderState
}

func (r *recordingProviderStateReporter) ReportProviderState(_ context.Context, state protocol.ProviderState) error {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.states = append(r.states, state)
	return nil
}

func (r *recordingProviderStateReporter) States() []protocol.ProviderState {
	r.mu.Lock()
	defer r.mu.Unlock()
	return append([]protocol.ProviderState(nil), r.states...)
}

func (a *oneTimeAdmission) Admit(context.Context, protocol.TunnelAdmissionRequest) (protocol.SessionAuthority, error) {
	if !a.used.CompareAndSwap(false, true) {
		return protocol.SessionAuthority{}, ErrAdmissionRejected
	}
	return a.authority, nil
}

func (a *oneTimeAdmission) ProviderWork(_ context.Context, authority protocol.SessionAuthority) (protocol.ProviderSessionWork, error) {
	if authority != a.authority {
		return protocol.ProviderSessionWork{}, ErrAdmissionRejected
	}
	return protocol.ProviderSessionWork{
		Version: "1", SessionID: authority.SessionID, GatewayID: authority.GatewayID,
		ReconnectSequence: authority.ReconnectSequence, ExpiresAt: authority.ExpiresAt,
		ProviderProfile: ProviderProfileApollo, ProviderIdentity: authority.ProviderIdentity,
		PolicyVersionID: "policy-1", ApplicationID: "1", ClientID: "paired-client", ManagementHost: "apollo.test", ManagementPort: 47990,
		StreamHost: "apollo.test", StreamPort: 47984, ClientCertificatePem: "certificate",
		ClientPrivateKeyPem: "private-key", ServerCertificatePem: "server-certificate",
	}, nil
}

func (a *oneTimeAdmission) Release(context.Context, protocol.SessionAuthority) error {
	if a.releases.Add(1) == 1 {
		close(a.released)
	}
	return nil
}

func mustRead(t *testing.T, path string) []byte {
	t.Helper()
	data, err := os.ReadFile(path)
	if err != nil {
		t.Fatal(err)
	}
	return data
}

func bytesRepeat(value byte, count int) []byte {
	data := make([]byte, count)
	for index := range data {
		data[index] = value
	}
	return data
}