diff --git a/cmd/ssh-proxy/main.go b/cmd/ssh-proxy/main.go index bed23c1..b2baa19 100644 --- a/cmd/ssh-proxy/main.go +++ b/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) } } diff --git a/go.mod b/go.mod index 1c29af6..7c2ba1c 100644 --- a/go.mod +++ b/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 diff --git a/go.sum b/go.sum index 8ab7361..194b8a1 100644 --- a/go.sum +++ b/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= diff --git a/proxy.go b/proxy.go index 9ab203e..791c3df 100644 --- a/proxy.go +++ b/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 } -} \ No newline at end of file +} diff --git a/websocket.go b/websocket.go index a859eea..1fc04af 100644 --- a/websocket.go +++ b/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 } \ No newline at end of file