Skip to content

Commit e46eda9

Browse files
mikewmikew
committed
dunno what openhands did here
1 parent 56a5415 commit e46eda9

2 files changed

Lines changed: 214 additions & 53 deletions

File tree

‎src/nvrh_binary_ssh/main.go‎

Lines changed: 63 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -1,16 +1,18 @@
11
package nvrh_binary_ssh
22

33
import (
4+
"context"
45
"log/slog"
56
"os"
67
"os/exec"
8+
"time"
79

8-
"nvrh/src/context"
10+
nvrhcontext "nvrh/src/context"
911
"nvrh/src/ssh_tunnel_info"
1012
)
1113

1214
type NvrhBinarySshClient struct {
13-
Ctx *context.NvrhContext
15+
Ctx *nvrhcontext.NvrhContext
1416
}
1517

1618
func (c *NvrhBinarySshClient) Close() error {
@@ -51,28 +53,66 @@ func (c *NvrhBinarySshClient) Run(command string, tunnelInfo *ssh_tunnel_info.Ss
5153
}
5254

5355
func (c *NvrhBinarySshClient) TunnelSocket(tunnelInfo *ssh_tunnel_info.SshTunnelInfo) {
54-
sshCommand := exec.Command(
55-
c.Ctx.SshPath,
56-
"-NL",
57-
tunnelInfo.BoundToIp(),
58-
c.Ctx.Endpoint.Given,
59-
)
60-
61-
slog.Info("Tunneling SSH socket", "tunnelInfo", tunnelInfo)
62-
63-
c.Ctx.CommandsToKill = append(c.Ctx.CommandsToKill, sshCommand)
64-
if c.Ctx.Debug {
65-
sshCommand.Stdout = os.Stdout
66-
sshCommand.Stderr = os.Stderr
67-
}
56+
c.TunnelSocketWithTimeout(tunnelInfo, 30*time.Second, 3)
57+
}
6858

69-
if err := sshCommand.Start(); err != nil {
70-
return
59+
// TunnelSocketWithTimeout creates an SSH tunnel with automatic cleanup after timeout or repeated errors
60+
func (c *NvrhBinarySshClient) TunnelSocketWithTimeout(tunnelInfo *ssh_tunnel_info.SshTunnelInfo, timeout time.Duration, maxErrors int) {
61+
ctx, cancel := context.WithTimeout(context.Background(), timeout)
62+
defer cancel()
63+
64+
errorCount := 0
65+
66+
for errorCount < maxErrors {
67+
sshCommand := exec.CommandContext(ctx,
68+
c.Ctx.SshPath,
69+
"-NL",
70+
tunnelInfo.BoundToIp(),
71+
c.Ctx.Endpoint.Given,
72+
)
73+
74+
slog.Info("Tunneling SSH socket", "tunnelInfo", tunnelInfo, "timeout", timeout, "attempt", errorCount+1)
75+
76+
c.Ctx.CommandsToKill = append(c.Ctx.CommandsToKill, sshCommand)
77+
if c.Ctx.Debug {
78+
sshCommand.Stdout = os.Stdout
79+
sshCommand.Stderr = os.Stderr
80+
}
81+
82+
if err := sshCommand.Start(); err != nil {
83+
slog.Error("Failed to start SSH tunnel", "error", err, "attempt", errorCount+1)
84+
errorCount++
85+
time.Sleep(1 * time.Second) // Brief delay before retry
86+
continue
87+
}
88+
89+
// Monitor for context cancellation or command completion
90+
done := make(chan error, 1)
91+
go func() {
92+
done <- sshCommand.Wait()
93+
}()
94+
95+
select {
96+
case <-ctx.Done():
97+
slog.Warn("SSH tunnel timeout reached, killing process", "timeout", timeout)
98+
if sshCommand.Process != nil {
99+
sshCommand.Process.Kill()
100+
}
101+
return
102+
case err := <-done:
103+
if err != nil {
104+
slog.Error("SSH tunnel process exited with error", "error", err, "attempt", errorCount+1)
105+
errorCount++
106+
if errorCount < maxErrors {
107+
slog.Info("Retrying SSH tunnel", "remaining_attempts", maxErrors-errorCount)
108+
time.Sleep(2 * time.Second) // Longer delay before retry on error
109+
}
110+
} else {
111+
slog.Info("SSH tunnel process completed successfully")
112+
return
113+
}
114+
}
71115
}
72116

73-
defer sshCommand.Process.Kill()
74-
75-
if err := sshCommand.Wait(); err != nil {
76-
return
77-
}
117+
slog.Error("SSH tunnel failed after maximum attempts", "max_errors", maxErrors, "timeout", timeout)
78118
}

‎src/nvrh_internal_ssh/main.go‎

Lines changed: 151 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -1,24 +1,38 @@
11
package nvrh_internal_ssh
22

33
import (
4+
"context"
45
"fmt"
56
"io"
67
"log/slog"
78
"net"
89
"os"
10+
"sync"
11+
"time"
912

1013
"golang.org/x/crypto/ssh"
1114

12-
"nvrh/src/context"
15+
nvrhcontext "nvrh/src/context"
1316
"nvrh/src/ssh_tunnel_info"
1417
)
1518

1619
type NvrhInternalSshClient struct {
17-
Ctx *context.NvrhContext
20+
Ctx *nvrhcontext.NvrhContext
1821
SshClient *ssh.Client
22+
tunnelCtx context.Context
23+
cancelTunnel context.CancelFunc
24+
tunnelMutex sync.Mutex
1925
}
2026

2127
func (c *NvrhInternalSshClient) Close() error {
28+
c.tunnelMutex.Lock()
29+
defer c.tunnelMutex.Unlock()
30+
31+
// Cancel any active tunnels
32+
if c.cancelTunnel != nil {
33+
c.cancelTunnel()
34+
}
35+
2236
if c.SshClient == nil {
2337
return fmt.Errorf("ssh client not initialized")
2438
}
@@ -58,48 +72,155 @@ func (c *NvrhInternalSshClient) Run(command string, tunnelInfo *ssh_tunnel_info.
5872
}
5973

6074
func (c *NvrhInternalSshClient) TunnelSocket(tunnelInfo *ssh_tunnel_info.SshTunnelInfo) {
61-
if c.SshClient == nil {
62-
return
63-
}
75+
c.TunnelSocketWithTimeout(tunnelInfo, 30*time.Second, 3)
76+
}
6477

65-
// Listen on the local Unix socket
66-
localListener, err := LocalListenerFromTunnelInfo(tunnelInfo)
67-
if err != nil {
68-
slog.Error("Failed to listen on local socket", "err", err)
78+
// TunnelSocketWithTimeout creates an SSH tunnel with automatic cleanup after timeout or repeated errors
79+
func (c *NvrhInternalSshClient) TunnelSocketWithTimeout(tunnelInfo *ssh_tunnel_info.SshTunnelInfo, timeout time.Duration, maxErrors int) {
80+
if c.SshClient == nil {
81+
slog.Error("SSH client not initialized")
6982
return
7083
}
7184

72-
defer localListener.Close()
85+
c.tunnelMutex.Lock()
86+
c.tunnelCtx, c.cancelTunnel = context.WithTimeout(context.Background(), timeout)
87+
ctx := c.tunnelCtx
88+
cancel := c.cancelTunnel
89+
c.tunnelMutex.Unlock()
90+
91+
defer cancel()
7392

74-
// Clean up local socket file
75-
defer func() {
76-
if tunnelInfo.Mode == "unix" {
77-
os.Remove(tunnelInfo.LocalSocket)
93+
errorCount := 0
94+
95+
for errorCount < maxErrors {
96+
select {
97+
case <-ctx.Done():
98+
slog.Warn("SSH tunnel timeout reached", "timeout", timeout)
99+
return
100+
default:
78101
}
79-
}()
80-
81-
slog.Info("Tunneling SSH socket", "tunnelInfo", tunnelInfo)
82102

83-
for {
84-
// Accept incoming connections
85-
localConn, err := localListener.Accept()
103+
// Listen on the local socket
104+
localListener, err := LocalListenerFromTunnelInfo(tunnelInfo)
86105
if err != nil {
87-
slog.Error("Failed to accept connection", "err", err)
88-
continue
106+
slog.Error("Failed to listen on local socket", "error", err, "attempt", errorCount+1)
107+
errorCount++
108+
if errorCount < maxErrors {
109+
time.Sleep(2 * time.Second)
110+
continue
111+
}
112+
break
89113
}
90114

91-
// Establish a connection to the remote socket via SSH
92-
remoteConn, err := RemoteListenerFromTunnelInfo(tunnelInfo, c.SshClient)
93-
if err != nil {
94-
slog.Error("Failed to dial remote socket", "err", err)
95-
localConn.Close()
96-
continue
115+
// Clean up local socket file
116+
defer func() {
117+
localListener.Close()
118+
if tunnelInfo.Mode == "unix" {
119+
os.Remove(tunnelInfo.LocalSocket)
120+
}
121+
}()
122+
123+
slog.Info("Tunneling SSH socket", "tunnelInfo", tunnelInfo, "timeout", timeout, "attempt", errorCount+1)
124+
125+
// Accept connections with timeout
126+
connectionErrors := 0
127+
for {
128+
select {
129+
case <-ctx.Done():
130+
slog.Warn("SSH tunnel context cancelled", "reason", ctx.Err())
131+
return
132+
default:
133+
}
134+
135+
// Set a deadline for accepting connections
136+
if tcpListener, ok := localListener.(*net.TCPListener); ok {
137+
tcpListener.SetDeadline(time.Now().Add(1 * time.Second))
138+
} else if unixListener, ok := localListener.(*net.UnixListener); ok {
139+
unixListener.SetDeadline(time.Now().Add(1 * time.Second))
140+
}
141+
142+
localConn, err := localListener.Accept()
143+
if err != nil {
144+
if netErr, ok := err.(net.Error); ok && netErr.Timeout() {
145+
// Timeout is expected, continue the loop
146+
continue
147+
}
148+
slog.Error("Failed to accept connection", "error", err)
149+
connectionErrors++
150+
if connectionErrors >= 5 {
151+
slog.Error("Too many connection errors, restarting listener")
152+
localListener.Close()
153+
errorCount++
154+
break
155+
}
156+
continue
157+
}
158+
159+
// Reset connection error count on successful accept
160+
connectionErrors = 0
161+
162+
// Establish a connection to the remote socket via SSH
163+
remoteConn, err := RemoteListenerFromTunnelInfo(tunnelInfo, c.SshClient)
164+
if err != nil {
165+
slog.Error("Failed to dial remote socket", "error", err)
166+
localConn.Close()
167+
continue
168+
}
169+
170+
// Start a goroutine to handle the connection with context
171+
go c.handleConnectionWithContext(ctx, localConn, remoteConn)
97172
}
98173

99-
// Start a goroutine to handle the connection
100-
go handleConnection(localConn, remoteConn)
174+
if errorCount >= maxErrors {
175+
break
176+
}
177+
178+
time.Sleep(2 * time.Second) // Wait before retry
101179
}
102180

181+
slog.Error("SSH tunnel failed after maximum attempts", "max_errors", maxErrors, "timeout", timeout)
182+
}
183+
184+
// handleConnectionWithContext handles a connection with context cancellation support
185+
func (c *NvrhInternalSshClient) handleConnectionWithContext(ctx context.Context, localConn net.Conn, remoteConn net.Conn) {
186+
defer localConn.Close()
187+
defer remoteConn.Close()
188+
189+
// Create a context that gets cancelled when the parent context is cancelled
190+
connCtx, cancel := context.WithCancel(ctx)
191+
defer cancel()
192+
193+
// Channel to signal when copying is done
194+
done := make(chan struct{}, 2)
195+
196+
// Copy data from local to remote
197+
go func() {
198+
defer func() { done <- struct{}{} }()
199+
io.Copy(remoteConn, localConn)
200+
}()
201+
202+
// Copy data from remote to local
203+
go func() {
204+
defer func() { done <- struct{}{} }()
205+
io.Copy(localConn, remoteConn)
206+
}()
207+
208+
// Wait for either context cancellation or connection completion
209+
select {
210+
case <-connCtx.Done():
211+
slog.Debug("Connection cancelled due to context")
212+
return
213+
case <-done:
214+
// One direction finished, wait for the other or timeout
215+
select {
216+
case <-done:
217+
slog.Debug("Connection completed normally")
218+
case <-time.After(5 * time.Second):
219+
slog.Debug("Connection cleanup timeout")
220+
case <-connCtx.Done():
221+
slog.Debug("Connection cancelled during cleanup")
222+
}
223+
}
103224
}
104225

105226
func handleConnection(localConn net.Conn, remoteConn net.Conn) {

0 commit comments

Comments
 (0)