0adcb52efc
Add error checking to all three write calls in bridge.go. A write failure now terminates that goroutine's copy loop and triggers cancel, which lets the other goroutines clean up naturally. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
86 lines
1.7 KiB
Go
86 lines
1.7 KiB
Go
package sshrelay
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"io"
|
|
"log/slog"
|
|
"sync/atomic"
|
|
"time"
|
|
|
|
"golang.org/x/crypto/ssh"
|
|
"nhooyr.io/websocket"
|
|
)
|
|
|
|
type ControlMsg struct {
|
|
Type string `json:"type"`
|
|
Cols int `json:"cols"`
|
|
Rows int `json:"rows"`
|
|
}
|
|
|
|
func bridge(ctx context.Context, cancel context.CancelFunc, ws *websocket.Conn,
|
|
sshSess *ssh.Session, stdin io.WriteCloser, stdout, stderr io.Reader, lastActive *int64) {
|
|
|
|
// WebSocket → SSH stdin
|
|
go func() {
|
|
defer cancel()
|
|
for {
|
|
typ, data, err := ws.Read(ctx)
|
|
if err != nil {
|
|
return
|
|
}
|
|
atomic.StoreInt64(lastActive, time.Now().UnixNano())
|
|
|
|
if typ == websocket.MessageText {
|
|
var ctrl ControlMsg
|
|
if json.Unmarshal(data, &ctrl) != nil {
|
|
continue
|
|
}
|
|
if ctrl.Type == "resize" && ctrl.Cols > 0 && ctrl.Cols <= 500 && ctrl.Rows > 0 && ctrl.Rows <= 200 {
|
|
sshSess.WindowChange(ctrl.Rows, ctrl.Cols)
|
|
}
|
|
continue
|
|
}
|
|
if _, err := stdin.Write(data); err != nil {
|
|
slog.Debug("SSH stdin write failed", "error", err)
|
|
return
|
|
}
|
|
}
|
|
}()
|
|
|
|
// SSH stdout → WebSocket
|
|
go func() {
|
|
defer cancel()
|
|
buf := make([]byte, 4096)
|
|
for {
|
|
n, err := stdout.Read(buf)
|
|
if err != nil {
|
|
return
|
|
}
|
|
atomic.StoreInt64(lastActive, time.Now().UnixNano())
|
|
if err := ws.Write(ctx, websocket.MessageBinary, buf[:n]); err != nil {
|
|
slog.Debug("SSH→WebSocket write failed", "error", err)
|
|
return
|
|
}
|
|
}
|
|
}()
|
|
|
|
// SSH stderr → WebSocket
|
|
go func() {
|
|
defer cancel()
|
|
buf := make([]byte, 4096)
|
|
for {
|
|
n, err := stderr.Read(buf)
|
|
if err != nil {
|
|
return
|
|
}
|
|
if err := ws.Write(ctx, websocket.MessageBinary, buf[:n]); err != nil {
|
|
slog.Debug("SSH→WebSocket write failed", "error", err)
|
|
return
|
|
}
|
|
}
|
|
}()
|
|
|
|
<-ctx.Done()
|
|
}
|