actually add tests

Signed-off-by: Jason Hall <jason@chainguard.dev>

proxy_test.go

@@ -0,0 +1,286 @@
+package sshproxy
+
+import (
+	"context"
+	"net"
+	"net/http"
+	"net/http/httptest"
+	"net/url"
+	"strings"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/google/go-cmp/cmp"
+	"github.com/gorilla/websocket"
+	"golang.org/x/oauth2"
+)
+
+var ts = oauth2.StaticTokenSource(&oauth2.Token{AccessToken: "mock-token"})
+
+func TestProxySSHToWebSocket_BidirectionalTraffic(t *testing.T) {
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+
+	var wsConn *websocket.Conn
+	var wsMutex sync.Mutex
+	receivedFromSSH := make([][]byte, 0)
+
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		// Check auth header
+		if auth := r.Header.Get("Authorization"); auth != "Bearer mock-token" {
+			t.Errorf("Expected auth header 'Bearer mock-token', got %q", auth)
+		}
+
+		conn, err := upgrader.Upgrade(w, r, nil)
+		if err != nil {
+			t.Errorf("Failed to upgrade connection: %v", err)
+			return
+		}
+
+		wsMutex.Lock()
+		wsConn = conn
+		wsMutex.Unlock()
+
+		// Read messages from SSH connection via WebSocket
+		for {
+			if _, message, err := conn.ReadMessage(); err != nil {
+				return
+			} else {
+				receivedFromSSH = append(receivedFromSSH, message)
+			}
+		}
+	}))
+	defer server.Close()
+
+	// Convert to WebSocket URL
+	wsURL, _ := url.Parse("ws" + strings.TrimPrefix(server.URL, "http"))
+
+	// Create SSH connection pair
+	sshClient, sshServer := net.Pipe()
+	defer sshClient.Close()
+	defer sshServer.Close()
+
+	// Track data received on SSH side
+	var sshMutex sync.Mutex
+	receivedFromWS := make([][]byte, 0)
+
+	// Read from SSH server side in background
+	go func() {
+		buf := make([]byte, 1024)
+		for {
+			n, err := sshServer.Read(buf)
+			if err != nil {
+				return
+			}
+			sshMutex.Lock()
+			receivedFromWS = append(receivedFromWS, append([]byte(nil), buf[:n]...))
+			sshMutex.Unlock()
+		}
+	}()
+
+	// Start the proxy
+	go func() {
+		if err := ProxySSHToWebSocket(ctx, sshClient, wsURL, ts); err != nil && ctx.Err() == nil {
+			t.Errorf("ProxySSHToWebSocket failed: %v", err)
+		}
+	}()
+
+	// Wait for WebSocket connection to be established
+	time.Sleep(100 * time.Millisecond)
+
+	// Test SSH -> WebSocket direction
+	testDataSSHToWS := []byte("Hello from SSH to WebSocket!")
+	if _, err := sshServer.Write(testDataSSHToWS); err != nil {
+		t.Fatalf("Failed to write to SSH server: %v", err)
+	}
+
+	// Wait for message to propagate
+	time.Sleep(100 * time.Millisecond)
+
+	if len(receivedFromSSH) == 0 {
+		t.Error("No data received from SSH connection via WebSocket")
+	} else if string(receivedFromSSH[0]) != string(testDataSSHToWS) {
+		t.Errorf("SSH->WebSocket: expected %q, got %q", testDataSSHToWS, receivedFromSSH[0])
+	}
+
+	// Test WebSocket -> SSH direction
+	wsMutex.Lock()
+	if wsConn != nil {
+		testDataWSToSSH := []byte("Hello from WebSocket to SSH!")
+		if err := wsConn.WriteMessage(websocket.BinaryMessage, testDataWSToSSH); err != nil {
+			wsMutex.Unlock()
+			t.Fatalf("Failed to write to WebSocket: %v", err)
+		}
+		wsMutex.Unlock()
+
+		// Wait for message to propagate
+		time.Sleep(100 * time.Millisecond)
+
+		sshMutex.Lock()
+		if len(receivedFromWS) == 0 {
+			t.Error("No data received from WebSocket via SSH connection")
+		} else if string(receivedFromWS[0]) != string(testDataWSToSSH) {
+			t.Errorf("WebSocket->SSH: expected %q, got %q", testDataWSToSSH, receivedFromWS[0])
+		}
+		sshMutex.Unlock()
+	} else {
+		wsMutex.Unlock()
+		t.Error("WebSocket connection was not established")
+	}
+}
+
+func TestProxySSHToWebSocket_MultipleMessages(t *testing.T) {
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		conn, err := upgrader.Upgrade(w, r, nil)
+		if err != nil {
+			t.Errorf("Failed to upgrade connection: %v", err)
+			return
+		}
+		defer conn.Close()
+
+		// Echo messages back
+		for {
+			if _, message, err := conn.ReadMessage(); err != nil {
+				return
+			} else {
+				// Echo back with prefix
+				response := append([]byte("echo: "), message...)
+				if err := conn.WriteMessage(websocket.BinaryMessage, response); err != nil {
+					return
+				}
+			}
+		}
+	}))
+	defer server.Close()
+
+	wsURL, _ := url.Parse("ws" + strings.TrimPrefix(server.URL, "http"))
+
+	// Create SSH connection pair
+	sshClient, sshServer := net.Pipe()
+	defer sshClient.Close()
+	defer sshServer.Close()
+
+	// Track responses
+	responses := make(chan []byte, 10)
+
+	go func() {
+		buf := make([]byte, 1024)
+		for {
+			n, err := sshServer.Read(buf)
+			if err != nil {
+				close(responses)
+				return
+			}
+			responses <- append([]byte(nil), buf[:n]...)
+		}
+	}()
+
+	// Start the proxy
+	go func() {
+
+		if err := ProxySSHToWebSocket(ctx, sshClient, wsURL, ts); err != nil && ctx.Err() == nil {
+			t.Errorf("ProxySSHToWebSocket failed: %v", err)
+		}
+	}()
+
+	// Wait for connection
+	time.Sleep(100 * time.Millisecond)
+
+	// Send multiple messages
+	testMessages := []string{"msg1", "msg2", "msg3"}
+	for _, msg := range testMessages {
+		_, err := sshServer.Write([]byte(msg))
+		if err != nil {
+			t.Fatalf("Failed to write message %q: %v", msg, err)
+		}
+	}
+
+	// Verify echo responses
+	for i, expectedMsg := range testMessages {
+		select {
+		case response := <-responses:
+			expected := "echo: " + expectedMsg
+			if string(response) != expected {
+				t.Errorf("Message %d: expected %q, got %q", i, expected, string(response))
+			}
+		case <-time.After(1 * time.Second):
+			t.Errorf("Timeout waiting for response to message %d", i)
+		}
+	}
+}
+
+func TestProxySSHToWebSocket_LargeData(t *testing.T) {
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+
+	// Create mock WebSocket server
+	upgrader := websocket.Upgrader{
+		CheckOrigin: func(r *http.Request) bool { return true },
+	}
+
+	var receivedData []byte
+	var dataMutex sync.Mutex
+
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		conn, err := upgrader.Upgrade(w, r, nil)
+		if err != nil {
+			t.Errorf("Failed to upgrade connection: %v", err)
+			return
+		}
+		defer conn.Close()
+
+		for {
+			_, message, err := conn.ReadMessage()
+			if err != nil {
+				return
+			}
+			dataMutex.Lock()
+			receivedData = append(receivedData, message...)
+			dataMutex.Unlock()
+		}
+	}))
+	defer server.Close()
+
+	wsURL, _ := url.Parse("ws" + strings.TrimPrefix(server.URL, "http"))
+
+	// Create SSH connection pair
+	sshClient, sshServer := net.Pipe()
+	defer sshClient.Close()
+	defer sshServer.Close()
+
+	// Start the proxy
+	go func() {
+		if err := ProxySSHToWebSocket(ctx, sshClient, wsURL, ts); err != nil && ctx.Err() == nil {
+			t.Errorf("ProxySSHToWebSocket failed: %v", err)
+		}
+	}()
+
+	// Wait for connection
+	time.Sleep(100 * time.Millisecond)
+
+	// Send large data (50KB)
+	largeData := make([]byte, 50*1024)
+	for i := range largeData {
+		largeData[i] = byte(i % 256)
+	}
+
+	if _, err := sshServer.Write(largeData); err != nil {
+		t.Fatalf("Failed to write large data: %v", err)
+	}
+
+	// Wait for all data to be received
+	time.Sleep(500 * time.Millisecond)
+
+	dataMutex.Lock()
+	if len(receivedData) != len(largeData) {
+		t.Errorf("Large data: expected %d bytes, got %d bytes", len(largeData), len(receivedData))
+	}
+	if d := cmp.Diff(largeData, receivedData); d != "" {
+		t.Errorf("Data mismatch (-want +got):\n%s", d)
+	}
+	dataMutex.Unlock()
+}

websocket_test.go

@@ -0,0 +1,347 @@
+package sshproxy
+
+import (
+	"io"
+	"net"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/gorilla/websocket"
+)
+
+var upgrader = websocket.Upgrader{CheckOrigin: func(r *http.Request) bool { return true }}
+
+func TestProxyWebSocketToSSH_BidirectionalTraffic(t *testing.T) {
+	// Create mock SSH server
+	sshServerAddr := "127.0.0.1:0"
+	listener, err := net.Listen("tcp", sshServerAddr)
+	if err != nil {
+		t.Fatalf("Failed to create SSH server: %v", err)
+	}
+	defer listener.Close()
+	actualAddr := listener.Addr().String()
+
+	// Track data received by SSH server
+	var sshMutex sync.Mutex
+	receivedFromWS := make([][]byte, 0)
+	var sshConn net.Conn
+
+	go func() {
+		for {
+			conn, err := listener.Accept()
+			if err != nil {
+				return
+			}
+			sshConn = conn
+
+			// Read from WebSocket via SSH connection
+			buf := make([]byte, 1024)
+			for {
+				n, err := conn.Read(buf)
+				if err != nil {
+					conn.Close()
+					return
+				}
+				sshMutex.Lock()
+				receivedFromWS = append(receivedFromWS, append([]byte(nil), buf[:n]...))
+				sshMutex.Unlock()
+			}
+		}
+	}()
+
+	// Create test server with proxy handler
+	server := httptest.NewServer(http.HandlerFunc(ProxyWebSocketToSSH(actualAddr, upgrader)))
+	defer server.Close()
+
+	// Connect WebSocket client
+	wsURL := "ws" + strings.TrimPrefix(server.URL, "http")
+	wsConn, _, err := websocket.DefaultDialer.Dial(wsURL, nil)
+	if err != nil {
+		t.Fatalf("Failed to connect WebSocket: %v", err)
+	}
+	defer wsConn.Close()
+
+	// Wait for SSH connection to be established
+	time.Sleep(100 * time.Millisecond)
+
+	// Test WebSocket -> SSH direction
+	testDataWSToSSH := []byte("Hello from WebSocket to SSH!")
+	err = wsConn.WriteMessage(websocket.BinaryMessage, testDataWSToSSH)
+	if err != nil {
+		t.Fatalf("Failed to write to WebSocket: %v", err)
+	}
+
+	// Wait for message to propagate
+	time.Sleep(100 * time.Millisecond)
+
+	sshMutex.Lock()
+	if len(receivedFromWS) == 0 {
+		t.Error("No data received from WebSocket via SSH connection")
+	} else if string(receivedFromWS[0]) != string(testDataWSToSSH) {
+		t.Errorf("WebSocket->SSH: expected %q, got %q", testDataWSToSSH, receivedFromWS[0])
+	}
+	sshMutex.Unlock()
+
+	// Test SSH -> WebSocket direction
+	if sshConn != nil {
+		testDataSSHToWS := []byte("Hello from SSH to WebSocket!")
+		_, err = sshConn.Write(testDataSSHToWS)
+		if err != nil {
+			t.Fatalf("Failed to write to SSH connection: %v", err)
+		}
+
+		// Read response from WebSocket
+		if _, message, err := wsConn.ReadMessage(); err != nil {
+			t.Fatalf("Failed to read from WebSocket: %v", err)
+		} else if string(message) != string(testDataSSHToWS) {
+			t.Errorf("SSH->WebSocket: expected %q, got %q", testDataSSHToWS, message)
+		}
+	} else {
+		t.Error("SSH connection was not established")
+	}
+}
+
+func TestProxyWebSocketToSSH_MultipleMessages(t *testing.T) {
+	// Create echo SSH server
+	listener, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		t.Fatalf("Failed to create SSH server: %v", err)
+	}
+	defer listener.Close()
+	actualAddr := listener.Addr().String()
+
+	go func() {
+		for {
+			conn, err := listener.Accept()
+			if err != nil {
+				return
+			}
+
+			// Echo server: read and write back with prefix
+			go func(c net.Conn) {
+				defer c.Close()
+				buf := make([]byte, 1024)
+				for {
+					n, err := c.Read(buf)
+					if err != nil {
+						return
+					}
+					// Echo back with prefix
+					response := append([]byte("ssh-echo: "), buf[:n]...)
+					c.Write(response)
+				}
+			}(conn)
+		}
+	}()
+
+	// Create test server
+	server := httptest.NewServer(http.HandlerFunc(ProxyWebSocketToSSH(actualAddr, upgrader)))
+	defer server.Close()
+
+	// Connect WebSocket client
+	wsURL := "ws" + strings.TrimPrefix(server.URL, "http")
+	wsConn, _, err := websocket.DefaultDialer.Dial(wsURL, nil)
+	if err != nil {
+		t.Fatalf("Failed to connect WebSocket: %v", err)
+	}
+	defer wsConn.Close()
+
+	// Wait for connection
+	time.Sleep(100 * time.Millisecond)
+
+	// Send multiple messages and verify echoes
+	testMessages := []string{"msg1", "msg2", "msg3"}
+	for i, msg := range testMessages {
+		// Send message
+		if err = wsConn.WriteMessage(websocket.BinaryMessage, []byte(msg)); err != nil {
+			t.Fatalf("Failed to write message %d: %v", i, err)
+		}
+
+		// Read echo response
+		expected := "ssh-echo: " + msg
+		if _, response, err := wsConn.ReadMessage(); err != nil {
+			t.Fatalf("Failed to read response %d: %v", i, err)
+		} else if string(response) != expected {
+			t.Errorf("Message %d: expected %q, got %q", i, expected, string(response))
+		}
+	}
+}
+
+func TestProxyWebSocketToSSH_LargeData(t *testing.T) {
+	// Create SSH server that accumulates data
+	listener, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		t.Fatalf("Failed to create SSH server: %v", err)
+	}
+	defer listener.Close()
+	actualAddr := listener.Addr().String()
+
+	var receivedData []byte
+	var dataMutex sync.Mutex
+	dataDone := make(chan struct{})
+
+	go func() {
+		conn, err := listener.Accept()
+		if err != nil {
+			return
+		}
+		defer conn.Close()
+
+		buf := make([]byte, 4096)
+		for {
+			n, err := conn.Read(buf)
+			if err != nil {
+				if err == io.EOF {
+					close(dataDone)
+				}
+				return
+			}
+			dataMutex.Lock()
+			receivedData = append(receivedData, buf[:n]...)
+			dataMutex.Unlock()
+
+			// Check if we've received all expected data
+			if len(receivedData) >= 50*1024 {
+				close(dataDone)
+				return
+			}
+		}
+	}()
+
+	// Create test server
+	server := httptest.NewServer(http.HandlerFunc(ProxyWebSocketToSSH(actualAddr, upgrader)))
+	defer server.Close()
+
+	// Connect WebSocket client
+	wsURL := "ws" + strings.TrimPrefix(server.URL, "http")
+	wsConn, _, err := websocket.DefaultDialer.Dial(wsURL, nil)
+	if err != nil {
+		t.Fatalf("Failed to connect WebSocket: %v", err)
+	}
+	defer wsConn.Close()
+
+	// Wait for connection
+	time.Sleep(100 * time.Millisecond)
+
+	// Send large data (50KB)
+	largeData := make([]byte, 50*1024)
+	for i := range largeData {
+		largeData[i] = byte(i % 256)
+	}
+
+	err = wsConn.WriteMessage(websocket.BinaryMessage, largeData)
+	if err != nil {
+		t.Fatalf("Failed to write large data: %v", err)
+	}
+
+	// Wait for all data to be received
+	select {
+	case <-dataDone:
+		// Success
+	case <-time.After(5 * time.Second):
+		t.Fatal("Timeout waiting for large data")
+	}
+
+	dataMutex.Lock()
+	defer dataMutex.Unlock()
+
+	if len(receivedData) != len(largeData) {
+		t.Errorf("Large data: expected %d bytes, got %d bytes", len(largeData), len(receivedData))
+	}
+
+	// Verify data integrity
+	for i, b := range receivedData {
+		if i >= len(largeData) {
+			break
+		}
+		if b != largeData[i] {
+			t.Errorf("Data mismatch at byte %d: expected %d, got %d", i, largeData[i], b)
+			break
+		}
+	}
+}
+
+func TestProxyWebSocketToSSH_ConnectionClosed(t *testing.T) {
+	// Create SSH server that closes immediately
+	listener, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		t.Fatalf("Failed to create SSH server: %v", err)
+	}
+	defer listener.Close()
+	actualAddr := listener.Addr().String()
+
+	go func() {
+		conn, err := listener.Accept()
+		if err != nil {
+			return
+		}
+		// Close immediately after accepting
+		conn.Close()
+	}()
+
+	// Create test server
+	server := httptest.NewServer(http.HandlerFunc(ProxyWebSocketToSSH(actualAddr, upgrader)))
+	defer server.Close()
+
+	// Connect WebSocket client
+	wsURL := "ws" + strings.TrimPrefix(server.URL, "http")
+	wsConn, _, err := websocket.DefaultDialer.Dial(wsURL, nil)
+	if err != nil {
+		t.Fatalf("Failed to connect WebSocket: %v", err)
+	}
+	defer wsConn.Close()
+
+	// Try to send data
+	if err := wsConn.WriteMessage(websocket.BinaryMessage, []byte("test")); err != nil {
+		t.Fatalf("Failed to write to WebSocket: %v", err)
+	}
+
+	// The WebSocket connection should close due to SSH connection being closed
+	// We'll wait a bit and then try to read, expecting an error
+	time.Sleep(200 * time.Millisecond)
+
+	wsConn.SetReadDeadline(time.Now().Add(1 * time.Second))
+	if _, _, err := wsConn.ReadMessage(); err != nil {
+		// Either we get an error (connection closed) or no error (connection still open)
+		// Both are valid since the proxy handles connection closure asynchronously
+		t.Logf("WebSocket closed as expected: %v", err)
+	} else {
+		t.Logf("WebSocket connection still open, which is acceptable")
+	}
+}
+
+func TestProxyWebSocketToSSH_SSHServerUnavailable(t *testing.T) {
+	// Use a non-existent address
+	unavailableAddr := "127.0.0.1:0" // This will be assigned but not listened on
+
+	// Get an available port but don't listen on it
+	tempListener, err := net.Listen("tcp", unavailableAddr)
+	if err != nil {
+		t.Fatalf("Failed to get available port: %v", err)
+	}
+	unavailableAddr = tempListener.Addr().String()
+	tempListener.Close() // Close so it becomes unavailable
+
+	// Create test server
+	server := httptest.NewServer(http.HandlerFunc(ProxyWebSocketToSSH(unavailableAddr, upgrader)))
+	defer server.Close()
+
+	// Try to connect WebSocket client
+	wsURL := "ws" + strings.TrimPrefix(server.URL, "http")
+	wsConn, _, err := websocket.DefaultDialer.Dial(wsURL, nil)
+
+	// The WebSocket upgrade might succeed, but the proxy should fail quickly
+	if err == nil {
+		defer wsConn.Close()
+		// The connection should close quickly due to SSH connection failure
+		wsConn.SetReadDeadline(time.Now().Add(2 * time.Second))
+		if _, _, err := wsConn.ReadMessage(); err == nil {
+			t.Error("Expected error when SSH server is unavailable, but connection stayed open")
+		}
+	}
+	// If WebSocket upgrade fails, that's also acceptable behavior
+}