simplify and provide ProxySSHToWebSocket as a func too

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

cmd/ssh-proxy/main.go

@@ -21,10 +21,7 @@ var env = envconfig.MustProcess(context.Background(), &struct {
 func main() {
 	ctx, cancel := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
 	defer cancel()
-	if err := sshproxy.NewProxy(sshproxy.ProxyConfig{
-		WebsocketURL: env.WebsocketURL,
-		SSHAddr:      env.SSHAddr,
-	}).Start(ctx); err != nil {
+	if err := sshproxy.NewProxy(env.WebsocketURL, env.SSHAddr).Start(ctx); err != nil {
 		clog.FatalContext(ctx, "failed to start SSH proxy", "error", err)
 	}
 }

go.mod

@@ -6,7 +6,6 @@ require (
 	github.com/chainguard-dev/clog v1.7.0
 	github.com/gorilla/websocket v1.5.3
 	github.com/sethvargo/go-envconfig v1.3.0
-	golang.org/x/crypto v0.41.0
 	google.golang.org/api v0.246.0
 )
 
@@ -25,6 +24,7 @@ require (
 	go.opentelemetry.io/otel v1.36.0 // indirect
 	go.opentelemetry.io/otel/metric v1.36.0 // indirect
 	go.opentelemetry.io/otel/trace v1.36.0 // indirect
+	golang.org/x/crypto v0.41.0 // indirect
 	golang.org/x/net v0.42.0 // indirect
 	golang.org/x/oauth2 v0.30.0 // indirect
 	golang.org/x/sys v0.35.0 // indirect

go.sum

@@ -61,8 +61,6 @@ golang.org/x/sync v0.16.0 h1:ycBJEhp9p4vXvUZNszeOq0kGTPghopOL8q0fq3vstxw=
 golang.org/x/sync v0.16.0/go.mod h1:1dzgHSNfp02xaA81J2MS99Qcpr2w7fw1gpm99rleRqA=
 golang.org/x/sys v0.35.0 h1:vz1N37gP5bs89s7He8XuIYXpyY0+QlsKmzipCbUtyxI=
 golang.org/x/sys v0.35.0/go.mod h1:BJP2sWEmIv4KK5OTEluFJCKSidICx8ciO85XgH3Ak8k=
-golang.org/x/term v0.34.0 h1:O/2T7POpk0ZZ7MAzMeWFSg6S5IpWd/RXDlM9hgM3DR4=
-golang.org/x/term v0.34.0/go.mod h1:5jC53AEywhIVebHgPVeg0mj8OD3VO9OzclacVrqpaAw=
 golang.org/x/text v0.28.0 h1:rhazDwis8INMIwQ4tpjLDzUhx6RlXqZNPEM0huQojng=
 golang.org/x/text v0.28.0/go.mod h1:U8nCwOR8jO/marOQ0QbDiOngZVEBB7MAiitBuMjXiNU=
 golang.org/x/time v0.12.0 h1:ScB/8o8olJvc+CQPWrK3fPZNfh7qgwCrY0zJmoEQLSE=

proxy.go

@@ -2,7 +2,6 @@ package sshproxy
 
 import (
 	"context"
-	"encoding/base64"
 	"errors"
 	"fmt"
 	"io"
@@ -15,39 +14,28 @@ import (
 	"google.golang.org/api/idtoken"
 )
 
-// ProxyConfig holds configuration for the SSH proxy
-type ProxyConfig struct {
+// Proxy handles SSH-to-WebSocket proxying
+type Proxy struct {
 	WebsocketURL string
 	SSHAddr      string
 }
 
-// Proxy handles SSH-to-WebSocket proxying
-type Proxy struct {
-	config ProxyConfig
-}
-
 // NewProxy creates a new SSH proxy
-func NewProxy(config ProxyConfig) *Proxy {
-	return &Proxy{config: config}
+func NewProxy(websocketURL, sshAddr string) *Proxy {
+	return &Proxy{WebsocketURL: websocketURL, SSHAddr: sshAddr}
 }
 
 // Start starts the SSH proxy server
 func (p *Proxy) Start(ctx context.Context) error {
 	log := clog.FromContext(ctx)
 
-	// Parse the URL to validate it
-	u, err := url.Parse(p.config.WebsocketURL)
+	listener, err := net.Listen("tcp", p.SSHAddr)
 	if err != nil {
-		return fmt.Errorf("invalid WebSocket URL: %w", err)
-	}
-
-	listener, err := net.Listen("tcp", p.config.SSHAddr)
-	if err != nil {
-		return fmt.Errorf("failed to listen on %s: %w", p.config.SSHAddr, err)
+		return fmt.Errorf("failed to listen on %s: %w", p.SSHAddr, err)
 	}
 	defer listener.Close()
 
-	log.Info("SSH proxy listening", "address", p.config.SSHAddr, "forwarding", p.config.WebsocketURL)
+	log.Info("SSH proxy listening", "address", p.SSHAddr, "forwarding", p.WebsocketURL)
 
 	// Accept SSH connections and proxy them to WebSocket
 	for {
@@ -63,128 +51,39 @@ func (p *Proxy) Start(ctx context.Context) error {
 			continue
 		}
 
-		go p.handleSSHConnection(ctx, conn, u)
+		go ProxySSHToWebSocket(ctx, conn, p.WebsocketURL)
 	}
 }
 
-func (p *Proxy) handleSSHConnection(ctx context.Context, sshConn net.Conn, wsURL *url.URL) {
-	ctx = clog.WithValues(ctx, "remote_addr", sshConn.RemoteAddr())
-	log := clog.FromContext(ctx)
-	defer sshConn.Close()
-
-	// Get identity token for authentication to Cloud Run
-	audience := fmt.Sprintf("https://%s", wsURL.Host)
-
-	// Get ID token
-	log.Info("Getting identity token for WebSocket authentication", "audience", audience)
-	tokenSource, err := idtoken.NewTokenSource(ctx, audience)
-	if err != nil {
-		log.Info("Failed to create token source", "error", err)
-		return
-	}
-
-	token, err := tokenSource.Token()
-	if err != nil {
-		log.Info("Failed to get identity token", "error", err)
-		return
-	}
-	log.Info("Got identity token, expires at", "expiry", token.Expiry)
-
-	// Create WebSocket dialer
-	dialer := websocket.DefaultDialer
-
-	// Create request headers with authentication
-	header := http.Header{}
-	header.Set("Authorization", fmt.Sprintf("Bearer %s", token.AccessToken))
-	log.Info("Connecting to WebSocket URL with auth header", "wsURL", wsURL.String())
-
-	// Connect to WebSocket with authentication
-	wsConn, resp, err := dialer.Dial(wsURL.String(), header)
-	if err != nil {
-		log.Info("Failed to connect to WebSocket", "error", err)
-		if resp != nil {
-			log.Info("Response status", "status", fmt.Sprintf("%d %s", resp.StatusCode, resp.Status))
-			_ = resp.Body.Close()
-		}
-		return
-	}
-	if resp != nil {
-		_ = resp.Body.Close()
-	}
-	defer wsConn.Close()
-
-	log.Info("Proxying SSH connection")
-
-	// Create error channel for goroutines
-	errCh := make(chan error, 2)
-
-	// SSH -> WebSocket
-	go func() {
-		buf := make([]byte, 32*1024)
-		for {
-			n, err := sshConn.Read(buf)
-			if err != nil {
-				errCh <- err
-				return
-			}
-
-			// Send as binary message
-			err = wsConn.WriteMessage(websocket.BinaryMessage, buf[:n])
-			if err != nil {
-				errCh <- err
-				return
-			}
-		}
-	}()
-
-	// WebSocket -> SSH
-	go func() {
-		for {
-			messageType, message, err := wsConn.ReadMessage()
-			if err != nil {
-				errCh <- err
-				return
-			}
-
-			var data []byte
-			switch messageType {
-			case websocket.BinaryMessage:
-				data = message
-			case websocket.TextMessage:
-				// Decode base64 if text message
-				decoded, err := base64.StdEncoding.DecodeString(string(message))
-				if err != nil {
-					log.Info("Failed to decode base64", "error", err)
-					continue
-				}
-				data = decoded
-			default:
-				continue
-			}
-
-			_, err = sshConn.Write(data)
-			if err != nil {
-				errCh <- err
-				return
-			}
-		}
-	}()
-
-	// Wait for error
-	<-errCh
-	log.Info("Connection closed for SSH")
-}
-
 // ProxySSHToWebSocket creates a bidirectional proxy between SSH and WebSocket
-func ProxySSHToWebSocket(ctx context.Context, sshConn net.Conn, wsURL string) error {
-	// Connect to WebSocket
-	dialer := websocket.DefaultDialer
-	wsConn, resp, err := dialer.Dial(wsURL, nil)
+func ProxySSHToWebSocket(ctx context.Context, sshConn net.Conn, wsAddr string) error {
+	log := clog.FromContext(ctx)
+
+	wsURL, err := url.Parse(wsAddr)
 	if err != nil {
-		return fmt.Errorf("connecting to WebSocket: %w", err)
+		return fmt.Errorf("failed to parse WebSocket URL: %w", err)
 	}
+
+	audience := fmt.Sprintf("https://%s", wsURL.Host)
+	log.Info("Getting identity token for WebSocket authentication", "audience", audience)
+	tokenSource, err := idtoken.NewTokenSource(ctx, audience)
+	if err != nil {
+		return fmt.Errorf("failed to create token source: %w", err)
+	}
+	token, err := tokenSource.Token()
+	if err != nil {
+		return fmt.Errorf("failed to get identity token: %w", err)
+	}
+
+	// Connect to WebSocket with authentication
+	header := http.Header{}
+	header.Set("Authorization", fmt.Sprintf("Bearer %s", token.AccessToken))
+	wsConn, resp, err := websocket.DefaultDialer.Dial(wsURL.String(), header)
 	if resp != nil {
-		_ = resp.Body.Close()
+		defer resp.Body.Close()
+	}
+	if err != nil {
+		return fmt.Errorf("failed to connect to WebSocket: %w", err)
 	}
 	defer wsConn.Close()
 
@@ -203,8 +102,7 @@ func ProxySSHToWebSocket(ctx context.Context, sshConn net.Conn, wsURL string) er
 				return
 			}
 
-			err = wsConn.WriteMessage(websocket.BinaryMessage, buf[:n])
-			if err != nil {
+			if err := wsConn.WriteMessage(websocket.BinaryMessage, buf[:n]); err != nil {
 				errCh <- err
 				return
 			}
@@ -222,8 +120,7 @@ func ProxySSHToWebSocket(ctx context.Context, sshConn net.Conn, wsURL string) er
 				return
 			}
 
-			_, err = sshConn.Write(message)
-			if err != nil {
+			if _, err := sshConn.Write(message); err != nil {
 				errCh <- err
 				return
 			}
@@ -237,4 +134,4 @@ func ProxySSHToWebSocket(ctx context.Context, sshConn net.Conn, wsURL string) er
 	case err := <-errCh:
 		return err
 	}
-}
+}

websocket.go

@@ -2,175 +2,90 @@
 package sshproxy
 
 import (
-	"encoding/base64"
+	"errors"
 	"io"
 	"net"
 	"net/http"
-	"time"
 
 	"github.com/chainguard-dev/clog"
 	"github.com/gorilla/websocket"
-	"golang.org/x/crypto/ssh"
 )
 
-var upgrader = websocket.Upgrader{
-	CheckOrigin: func(r *http.Request) bool {
-		// Allow connections from any origin for now
-		// In production, you'd want to restrict this
-		// TODO: Implement proper origin checking
-		return true
-	},
-}
+func ProxyWebSocketToSSH(tcpAddr string, upgrader websocket.Upgrader) http.HandlerFunc {
+	return func(w http.ResponseWriter, r *http.Request) {
+		log := clog.FromContext(r.Context())
 
-// WebSocketServer handles WebSocket connections that tunnel SSH
-type WebSocketServer struct {
-	sshConfig *ssh.ServerConfig
-	handler   func(ssh.Channel, <-chan *ssh.Request, *ssh.Permissions)
-}
-
-// NewWebSocketServer creates a new SSH-over-WebSocket server
-func NewWebSocketServer(sshConfig *ssh.ServerConfig, handler func(ssh.Channel, <-chan *ssh.Request, *ssh.Permissions)) *WebSocketServer {
-	return &WebSocketServer{
-		sshConfig: sshConfig,
-		handler:   handler,
-	}
-}
-
-// ServeHTTP handles WebSocket upgrade and SSH tunneling
-func (s *WebSocketServer) ServeHTTP(w http.ResponseWriter, r *http.Request) {
-	ctx := r.Context()
-	log := clog.FromContext(ctx)
-
-	log.Info("WebSocket connection request", "remote", r.RemoteAddr)
-
-	conn, err := upgrader.Upgrade(w, r, nil)
-	if err != nil {
-		log.Error("WebSocket upgrade failed", "error", err)
-		return
-	}
-	defer conn.Close()
-
-	log.Info("WebSocket connection established", "remote", r.RemoteAddr)
-
-	// Create bidirectional pipes for SSH connection
-	clientReader, clientWriter := io.Pipe()
-	serverReader, serverWriter := io.Pipe()
-
-	// Handle SSH connection in a goroutine
-	go func() {
-		sshConn, chans, reqs, err := ssh.NewServerConn(&pipeConn{
-			Reader: clientReader,
-			Writer: serverWriter,
-		}, s.sshConfig)
+		// Upgrade HTTP connection to WebSocket
+		conn, err := upgrader.Upgrade(w, r, nil)
 		if err != nil {
-			log.Error("SSH handshake failed", "error", err)
-			_ = clientWriter.Close()
-			_ = serverWriter.Close()
+			log.Error("Failed to upgrade WebSocket connection", "error", err)
 			return
 		}
-		defer sshConn.Close()
+		defer conn.Close()
 
-		log.Info("SSH connection established", "user", sshConn.User())
-
-		// Handle SSH requests
-		go ssh.DiscardRequests(reqs)
-
-		// Handle SSH channels
-		for newChannel := range chans {
-			if newChannel.ChannelType() != "session" {
-				_ = newChannel.Reject(ssh.UnknownChannelType, "unknown channel type")
-				continue
-			}
-
-			channel, requests, err := newChannel.Accept()
-			if err != nil {
-				log.Error("Failed to accept channel", "error", err)
-				continue
-			}
-
-			go s.handler(channel, requests, sshConn.Permissions)
+		// Connect to local SSH server
+		tcpConn, err := net.Dial("tcp", tcpAddr)
+		if err != nil {
+			log.Error("Failed to connect to SSH server", "error", err, "address", tcpAddr)
+			return
 		}
-	}()
+		defer tcpConn.Close()
 
-	// Bridge WebSocket and SSH connections
-	errCh := make(chan error, 2)
+		log.Info("WebSocket connection established", "remote_addr", r.RemoteAddr)
 
-	// WebSocket -> SSH
-	go func() {
-		defer clientWriter.Close()
-		for {
-			messageType, message, err := conn.ReadMessage()
-			if err != nil {
-				log.Debug("WebSocket read error", "error", err)
-				errCh <- err
-				return
-			}
+		// Create bidirectional proxy
 
-			log.Debug("Received WebSocket message", "type", messageType, "size", len(message))
-			if messageType == websocket.BinaryMessage {
-				_, err = clientWriter.Write(message)
+		// WebSocket -> TCP
+		done := make(chan string)
+		go func() {
+			defer func() {
+				select {
+				case done <- "ws_to_tcp_closed":
+				default: // other goroutine already closed, don't block
+				}
+			}()
+			for {
+				messageType, data, err := conn.ReadMessage()
 				if err != nil {
-					log.Debug("Pipe write error", "error", err)
-					errCh <- err
+					if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
+						log.Error("Unexpected WebSocket closure", "error", err)
+					}
+					// TODO: what other errors need to be handled here?
 					return
 				}
-			} else if messageType == websocket.TextMessage {
-				// Handle base64-encoded binary data for browser compatibility
-				decoded, err := base64.StdEncoding.DecodeString(string(message))
-				if err != nil {
-					log.Warn("Failed to decode base64 message", "error", err)
-					continue
+				if messageType == websocket.BinaryMessage {
+					if _, err := tcpConn.Write(data); err != nil {
+						log.Error("TCP write error", "error", err)
+						return
+					}
 				}
-				_, err = clientWriter.Write(decoded)
-				if err != nil {
-					errCh <- err
+			}
+		}()
+
+		// TCP -> WebSocket
+		go func() {
+			defer func() {
+				select {
+				case done <- "tcp_to_ws_closed":
+				default: // other goroutine already closed, don't block
+				}
+			}()
+			buf := make([]byte, 4096)
+			for {
+				n, err := tcpConn.Read(buf)
+				if err != nil && !errors.Is(err, io.EOF) {
+					log.Error("TCP read error", "error", err)
+					return
+				}
+				if err = conn.WriteMessage(websocket.BinaryMessage, buf[:n]); err != nil {
+					log.Error("WebSocket write error", "error", err)
 					return
 				}
 			}
-		}
-	}()
+		}()
 
-	// SSH -> WebSocket
-	go func() {
-		buf := make([]byte, 32*1024)
-		for {
-			n, err := serverReader.Read(buf)
-			if err != nil {
-				errCh <- err
-				return
-			}
-
-			err = conn.WriteMessage(websocket.BinaryMessage, buf[:n])
-			if err != nil {
-				errCh <- err
-				return
-			}
-		}
-	}()
-
-	// Wait for error or connection close
-	<-errCh
-	log.Info("WebSocket connection closed", "remote", r.RemoteAddr)
-}
-
-// pipeConn wraps io.Reader and io.Writer to implement net.Conn
-type pipeConn struct {
-	io.Reader
-	io.Writer
-}
-
-func (c *pipeConn) Close() error {
-	if closer, ok := c.Writer.(io.Closer); ok {
-		_ = closer.Close()
+		// Wait for either direction to close
+		msg := <-done
+		log.Info("WebSocket proxy connection closed", "reason", msg)
 	}
-	if closer, ok := c.Reader.(io.Closer); ok {
-		_ = closer.Close()
-	}
-	return nil
 }
-
-func (c *pipeConn) LocalAddr() net.Addr                { return &net.TCPAddr{} }
-func (c *pipeConn) RemoteAddr() net.Addr               { return &net.TCPAddr{} }
-func (c *pipeConn) SetDeadline(t time.Time) error      { return nil }
-func (c *pipeConn) SetReadDeadline(t time.Time) error  { return nil }
-func (c *pipeConn) SetWriteDeadline(t time.Time) error { return nil }