usb-server/internal/client/direct_test.go

package client

import (
	"bytes"
	"net"
	"strconv"
	"testing"
	"time"

	"github.com/duffy/usb-server/internal/crypto"
	"github.com/duffy/usb-server/internal/protocol"
)

const (
	testTok1 = "111111111111111111111111111111111111111111="
	testTok2 = "222222222222222222222222222222222222222222="
	testTok3 = "333333333333333333333333333333333333333333="
)

func testSecret(t *testing.T) *crypto.TunnelSecret {
	t.Helper()
	s, err := crypto.DeriveTunnelSecret(testTok1, testTok2, testTok3)
	if err != nil {
		t.Fatalf("DeriveTunnelSecret: %v", err)
	}
	return s
}

const testTunnelID = "0123456789abcdef" // exactly TunnelHeaderSize

func TestHandshakeRoundTrip(t *testing.T) {
	s := testSecret(t)
	token := s.PeerToken(testTunnelID)

	greeting, err := buildHandshake(testTunnelID, token)
	if err != nil {
		t.Fatalf("buildHandshake: %v", err)
	}
	if len(greeting) != handshakeSize {
		t.Fatalf("greeting is %d bytes, want %d", len(greeting), handshakeSize)
	}

	gotID, gotToken, err := parseHandshake(greeting)
	if err != nil {
		t.Fatalf("parseHandshake: %v", err)
	}
	if gotID != testTunnelID {
		t.Errorf("tunnel ID = %q, want %q", gotID, testTunnelID)
	}
	if gotToken != token {
		t.Errorf("token = %q, want %q", gotToken, token)
	}
}

func TestParseHandshakeRejectsMalformed(t *testing.T) {
	s := testSecret(t)
	valid, _ := buildHandshake(testTunnelID, s.PeerToken(testTunnelID))

	tests := []struct {
		name string
		data []byte
	}{
		{"empty", nil},
		{"truncated", valid[:handshakeSize-1]},
		{"too long", append(append([]byte{}, valid...), 0x00)},
		{"bad magic", func() []byte {
			b := append([]byte{}, valid...)
			b[0] = 'X'
			return b
		}()},
		{"unsupported version", func() []byte {
			b := append([]byte{}, valid...)
			b[4] = 99
			return b
		}()},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			if _, _, err := parseHandshake(tt.data); err == nil {
				t.Error("malformed handshake was accepted")
			}
		})
	}
}

func TestHandshakeReplyStatuses(t *testing.T) {
	if err := parseHandshakeReply(buildHandshakeReply(directAccepted)); err != nil {
		t.Errorf("accepted reply reported an error: %v", err)
	}
	for _, status := range []byte{directUnknownTun, directBadToken, directWrongVerson, 99} {
		if err := parseHandshakeReply(buildHandshakeReply(status)); err == nil {
			t.Errorf("status %d was treated as success", status)
		}
	}
}

func TestBuildHandshakeRejectsBadInput(t *testing.T) {
	s := testSecret(t)
	good := s.PeerToken(testTunnelID)

	if _, err := buildHandshake("short", good); err == nil {
		t.Error("a wrong-length tunnel ID was accepted")
	}
	if _, err := buildHandshake(testTunnelID, "not-hex"); err == nil {
		t.Error("a non-hex token was accepted")
	}
	if _, err := buildHandshake(testTunnelID, "abcd"); err == nil {
		t.Error("a short token was accepted")
	}
}

// A listener must only hand over connections whose peer proves group
// membership. The relay knows tunnel IDs, so the token is what stops it — or
// anyone else who reaches the port — from taking a device over.
func TestListenerAcceptsOnlyValidToken(t *testing.T) {
	s := testSecret(t)

	dl, err := newDirectListener(0, s)
	if err != nil {
		t.Fatalf("newDirectListener: %v", err)
	}
	defer dl.Close()

	accepted, err := dl.Expect(testTunnelID)
	if err != nil {
		t.Fatalf("Expect: %v", err)
	}

	addr := net.JoinHostPort("127.0.0.1", strconv.Itoa(dl.Port()))

	t.Run("wrong token is rejected", func(t *testing.T) {
		other, _ := crypto.DeriveTunnelSecret(testTok1, testTok2, "different")
		greeting, _ := buildHandshake(testTunnelID, other.PeerToken(testTunnelID))

		conn, err := net.DialTimeout("tcp", addr, time.Second)
		if err != nil {
			t.Fatalf("dial: %v", err)
		}
		defer conn.Close()

		conn.Write(greeting)
		reply := make([]byte, handshakeReplySize)
		conn.SetReadDeadline(time.Now().Add(2 * time.Second))
		if _, err := readFull(conn, reply); err != nil {
			t.Fatalf("reading reply: %v", err)
		}
		if err := parseHandshakeReply(reply); err == nil {
			t.Fatal("listener accepted a connection with the wrong token")
		}
	})

	t.Run("unknown tunnel is rejected", func(t *testing.T) {
		greeting, _ := buildHandshake("fedcba9876543210", s.PeerToken("fedcba9876543210"))

		conn, err := net.DialTimeout("tcp", addr, time.Second)
		if err != nil {
			t.Fatalf("dial: %v", err)
		}
		defer conn.Close()

		conn.Write(greeting)
		reply := make([]byte, handshakeReplySize)
		conn.SetReadDeadline(time.Now().Add(2 * time.Second))
		readFull(conn, reply)
		if err := parseHandshakeReply(reply); err == nil {
			t.Fatal("listener accepted a connection for an unregistered tunnel")
		}
	})

	t.Run("valid token is accepted", func(t *testing.T) {
		conn, _, err := dialDirect([]string{addr}, testTunnelID, s.PeerToken(testTunnelID))
		if err != nil {
			t.Fatalf("dialDirect: %v", err)
		}
		defer conn.Close()

		select {
		case got := <-accepted:
			if got == nil {
				t.Fatal("listener delivered a nil connection")
			}
			got.Close()
		case <-time.After(2 * time.Second):
			t.Fatal("listener never delivered the accepted connection")
		}
	})
}

// End-to-end over a real socket pair: the two ends must agree on framing and
// on which direction each encrypts in.
func TestDirectTunnelCarriesTrafficBothWays(t *testing.T) {
	s := testSecret(t)

	dl, err := newDirectListener(0, s)
	if err != nil {
		t.Fatalf("newDirectListener: %v", err)
	}
	defer dl.Close()

	accepted, err := dl.Expect(testTunnelID)
	if err != nil {
		t.Fatalf("Expect: %v", err)
	}

	addr := net.JoinHostPort("127.0.0.1", strconv.Itoa(dl.Port()))
	useConn, _, err := dialDirect([]string{addr}, testTunnelID, s.PeerToken(testTunnelID))
	if err != nil {
		t.Fatalf("dialDirect: %v", err)
	}
	defer useConn.Close()

	var shareConn *directConn
	select {
	case shareConn = <-accepted:
	case <-time.After(2 * time.Second):
		t.Fatal("listener never delivered the connection")
	}
	defer shareConn.Close()

	shareCodec, err := newTunnelCodec(s, testTunnelID, crypto.DirShareToUse)
	if err != nil {
		t.Fatalf("share codec: %v", err)
	}
	useCodec, err := newTunnelCodec(s, testTunnelID, crypto.DirUseToShare)
	if err != nil {
		t.Fatalf("use codec: %v", err)
	}

	// use -> share
	want := []byte("USBIP CMD_SUBMIT payload")
	if err := send(useCodec, directSender(useConn), want); err != nil {
		t.Fatalf("sending use->share: %v", err)
	}
	frame, err := shareConn.ReadFrame()
	if err != nil {
		t.Fatalf("share reading frame: %v", err)
	}
	got, err := shareCodec.decode(frame)
	if err != nil {
		t.Fatalf("share decoding frame: %v", err)
	}
	if !bytes.Equal(got, want) {
		t.Errorf("share received %q, want %q", got, want)
	}

	// share -> use
	want2 := []byte("USBIP RET_SUBMIT payload")
	if err := send(shareCodec, directSender(shareConn), want2); err != nil {
		t.Fatalf("sending share->use: %v", err)
	}
	frame2, err := useConn.ReadFrame()
	if err != nil {
		t.Fatalf("use reading frame: %v", err)
	}
	got2, err := useCodec.decode(frame2)
	if err != nil {
		t.Fatalf("use decoding frame: %v", err)
	}
	if !bytes.Equal(got2, want2) {
		t.Errorf("use received %q, want %q", got2, want2)
	}
}

func TestDirectConnFramingPreservesBoundaries(t *testing.T) {
	client, server := net.Pipe()
	defer client.Close()
	defer server.Close()

	sender := newDirectConn(client)
	receiver := newDirectConn(server)

	payloads := [][]byte{
		[]byte("a"),
		bytes.Repeat([]byte("x"), 1000),
		[]byte("last one"),
	}

	go func() {
		for _, p := range payloads {
			if err := sender.WriteFrame(p); err != nil {
				return
			}
		}
	}()

	for i, want := range payloads {
		got, err := receiver.ReadFrame()
		if err != nil {
			t.Fatalf("frame %d: %v", i, err)
		}
		if !bytes.Equal(got, want) {
			t.Errorf("frame %d is %d bytes, want %d", i, len(got), len(want))
		}
	}
}

func TestDirectConnRejectsOversizedLength(t *testing.T) {
	client, server := net.Pipe()
	defer client.Close()
	defer server.Close()

	go func() {
		// A length prefix claiming far more than the cap must be refused
		// before anything is allocated.
		client.Write([]byte{0xFF, 0xFF, 0xFF, 0xFF})
	}()

	receiver := newDirectConn(server)
	if _, err := receiver.ReadFrame(); err == nil {
		t.Error("an oversized frame length was accepted")
	}
}

func TestDialDirectFailsWithoutReachableAddress(t *testing.T) {
	s := testSecret(t)

	// Port 1 on loopback refuses immediately, so this stays fast.
	_, _, err := dialDirect([]string{"127.0.0.1:1"}, testTunnelID, s.PeerToken(testTunnelID))
	if err == nil {
		t.Fatal("dialDirect succeeded against a closed port")
	}

	if _, _, err := dialDirect(nil, testTunnelID, s.PeerToken(testTunnelID)); err == nil {
		t.Error("dialDirect succeeded with no candidate addresses")
	}
}

// The dialer races candidates; an unreachable one alongside a good one must
// not stop the good one from winning.
func TestDialDirectPicksTheReachableAddress(t *testing.T) {
	s := testSecret(t)

	dl, err := newDirectListener(0, s)
	if err != nil {
		t.Fatalf("newDirectListener: %v", err)
	}
	defer dl.Close()

	accepted, _ := dl.Expect(testTunnelID)
	good := net.JoinHostPort("127.0.0.1", strconv.Itoa(dl.Port()))

	conn, addr, err := dialDirect(
		[]string{"127.0.0.1:1", good, "127.0.0.1:2"},
		testTunnelID, s.PeerToken(testTunnelID))
	if err != nil {
		t.Fatalf("dialDirect: %v", err)
	}
	defer conn.Close()

	if addr != good {
		t.Errorf("connected to %s, want %s", addr, good)
	}
	select {
	case c := <-accepted:
		c.Close()
	case <-time.After(2 * time.Second):
		t.Error("listener never saw the connection")
	}
}

func TestLocalEndpointsExcludeLoopback(t *testing.T) {
	if got := localEndpoints(0); got != nil {
		t.Errorf("localEndpoints(0) = %v, want nil — port 0 means no listener", got)
	}

	for _, ep := range localEndpoints(9000) {
		host, port, err := net.SplitHostPort(ep)
		if err != nil {
			t.Errorf("endpoint %q is not host:port: %v", ep, err)
			continue
		}
		if port != "9000" {
			t.Errorf("endpoint %q has port %q, want 9000", ep, port)
		}
		ip := net.ParseIP(host)
		if ip == nil {
			t.Errorf("endpoint %q has an unparseable host", ep)
			continue
		}
		if ip.IsLoopback() {
			t.Errorf("endpoint %q is loopback; a peer cannot reach that", ep)
		}
	}
}

func TestTunnelCodecNilPassesThrough(t *testing.T) {
	var codec *tunnelCodec

	if codec.encrypted() {
		t.Error("a nil codec reported itself as encrypted")
	}

	payload := []byte("cleartext")
	encoded, err := codec.encode(payload)
	if err != nil {
		t.Fatalf("encode: %v", err)
	}
	if !bytes.Equal(encoded, payload) {
		t.Error("a nil codec altered the payload")
	}

	decoded, err := codec.decode(encoded)
	if err != nil {
		t.Fatalf("decode: %v", err)
	}
	if !bytes.Equal(decoded, payload) {
		t.Error("round trip through a nil codec changed the payload")
	}
}

func TestTunnelCodecEncryptsWhenSecretPresent(t *testing.T) {
	s := testSecret(t)

	codec, err := newTunnelCodec(s, testTunnelID, crypto.DirShareToUse)
	if err != nil {
		t.Fatalf("newTunnelCodec: %v", err)
	}
	if !codec.encrypted() {
		t.Fatal("codec with a secret reported itself as unencrypted")
	}

	payload := []byte("this must not appear on the wire")
	encoded, err := codec.encode(payload)
	if err != nil {
		t.Fatalf("encode: %v", err)
	}
	if bytes.Contains(encoded, payload) {
		t.Error("the encoded frame contains its plaintext")
	}

	peer, _ := newTunnelCodec(s, testTunnelID, crypto.DirUseToShare)
	decoded, err := peer.decode(encoded)
	if err != nil {
		t.Fatalf("peer decode: %v", err)
	}
	if !bytes.Equal(decoded, payload) {
		t.Errorf("peer decoded %q, want %q", decoded, payload)
	}
}

func TestConstantTimeEqual(t *testing.T) {
	if !constantTimeEqual("abc", "abc") {
		t.Error("equal strings compared unequal")
	}
	if constantTimeEqual("abc", "abd") {
		t.Error("different strings compared equal")
	}
	if constantTimeEqual("abc", "abcd") {
		t.Error("strings of different length compared equal")
	}
	if !constantTimeEqual("", "") {
		t.Error("empty strings compared unequal")
	}
}

// Guards the assumption baked into the wire format.
func TestTunnelIDFitsHandshake(t *testing.T) {
	if protocol.TunnelHeaderSize != 16 {
		t.Fatalf("TunnelHeaderSize is %d; the handshake layout assumes 16", protocol.TunnelHeaderSize)
	}
	if len(testTunnelID) != protocol.TunnelHeaderSize {
		t.Fatalf("test tunnel ID is %d bytes, want %d", len(testTunnelID), protocol.TunnelHeaderSize)
	}
}

// --- helpers ---

func readFull(conn net.Conn, buf []byte) (int, error) {
	total := 0
	for total < len(buf) {
		n, err := conn.Read(buf[total:])
		total += n
		if err != nil {
			return total, err
		}
	}
	return total, nil
}