Skip to content

Commit fa77edf

Browse files
psviderskipsviderski
committed
fix(proxy): don't shutdown 'uc proxy' when a client connection aborts
1 parent b99abb7 commit fa77edf

4 files changed

Lines changed: 276 additions & 66 deletions

File tree

‎cmd/uc/proxy.go‎

Lines changed: 15 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -132,35 +132,34 @@ func runProxy(ctx context.Context, uncli *cli.CLI, opts proxyOptions) error {
132132
// endpoint and shuffles the data, *it* will actually experience errors.
133133
remoteAddr := net.JoinHostPort(ip.String(), strconv.Itoa(opts.remotePort))
134134

135-
ctx, cancel := context.WithCancel(ctx)
136-
defer cancel()
137-
138135
p := &proxy.Proxy{
139136
Listener: listener,
140137
RemoteAddr: remoteAddr,
141138
DialContext: dialer.DialContext,
142139
OnError: func(err error) {
143-
fmt.Printf("Failed to proxy to '%s': %v\n", remoteAddr, err)
144-
cancel()
140+
if proxy.IsConnectionClosedError(err) {
141+
return
142+
}
143+
// A more actionable error instead of the cryptic [ssh -W] command error.
144+
if strings.Contains(err.Error(), "Session open refused by peer") {
145+
fmt.Printf("Could not connect to '%s': connection refused. "+
146+
"Check that the service is running and listening on port %d inside the container.\n",
147+
remoteAddr, opts.remotePort)
148+
return
149+
}
150+
151+
fmt.Printf("Failed to proxy a connection to '%s': %v\n", remoteAddr, err)
145152
},
146153
}
147154

148-
// Run the proxy in the background and signal when it has fully shut down.
149-
done := make(chan struct{})
150-
go func() {
151-
p.Run(ctx)
152-
close(done)
153-
}()
154-
155155
// Prefix the local address with the scheme for common HTTP ports so it becomes control-clickable in most
156156
// terminals. We assume plain HTTP since TLS is typically terminated by Caddy in front of the service.
157157
fmt.Printf("%s%s → %s (%s%s%s)\n", schemeForPort(opts.remotePort), p.Listener.Addr().String(),
158158
remoteAddr, opts.service, tui.Faint.Render("/"), containerID)
159159

160-
<-ctx.Done()
161-
// Wait for the proxy to drain in-flight connections and shut down gracefully.
162-
<-done
163-
160+
if err = p.Run(ctx); err != nil {
161+
return fmt.Errorf("run proxy to '%s': %w", remoteAddr, err)
162+
}
164163
return nil
165164
}
166165

‎internal/proxy/proxy.go‎

Lines changed: 61 additions & 48 deletions
Original file line numberDiff line numberDiff line change
@@ -2,11 +2,12 @@ package proxy
22

33
import (
44
"context"
5+
"errors"
56
"fmt"
67
"io"
78
"net"
8-
"os"
99
"sync"
10+
"syscall"
1011
"time"
1112
)
1213

@@ -15,58 +16,53 @@ type Proxy struct {
1516
Listener net.Listener
1617
RemoteAddr string
1718
DialContext func(ctx context.Context, network, address string) (net.Conn, error)
19+
// OnError is called for errors that occur during proxying individual connections. It may be called concurrently
20+
// for different connections.
1821
OnError func(error)
1922
activeConns sync.WaitGroup
2023
}
2124

22-
// deadliner is an interface for listeners that support setting deadlines.
23-
type deadliner interface {
24-
SetDeadline(t time.Time) error
25-
}
26-
2725
// halfCloser is an interface for connections that support half-close.
2826
type halfCloser interface {
2927
CloseWrite() error
3028
}
3129

32-
// Run starts the proxy and runs until the context is canceled.
33-
func (p *Proxy) Run(ctx context.Context) {
30+
// IsConnectionClosedError reports whether err indicates that a connection was closed or aborted by either peer.
31+
// Callers can use it to ignore routine connection shutdown or broken pipe errors reported to Proxy.OnError.
32+
func IsConnectionClosedError(err error) bool {
33+
return errors.Is(err, net.ErrClosed) || errors.Is(err, io.ErrClosedPipe) ||
34+
errors.Is(err, syscall.EPIPE) || errors.Is(err, syscall.ECONNRESET)
35+
}
36+
37+
// Run starts the proxy and runs until the context is canceled or the listener fails. It returns nil when the context
38+
// is canceled. Errors handling individual connections are reported to OnError and do not stop the proxy.
39+
func (p *Proxy) Run(ctx context.Context) error {
3440
if p.DialContext == nil {
3541
p.DialContext = (&net.Dialer{}).DialContext
3642
}
3743

44+
ctx, cancel := context.WithCancel(ctx)
45+
defer cancel()
3846
defer p.Listener.Close()
3947

40-
// Handle incoming connections until context is canceled.
41-
Loop:
42-
for {
43-
select {
44-
case <-ctx.Done():
45-
break Loop
46-
default:
47-
}
48-
49-
// Set a deadline on the listener if supported to check context periodically.
50-
if dl, ok := p.Listener.(deadliner); ok {
51-
dl.SetDeadline(time.Now().Add(1 * time.Second))
52-
}
48+
// Closing the listener unblocks Accept when the context is canceled. This works for both TCP and Unix listeners
49+
// and avoids polling with listener deadlines.
50+
stopClose := context.AfterFunc(ctx, func() {
51+
p.Listener.Close()
52+
})
53+
defer stopClose()
5354

55+
var runErr error
56+
for {
5457
conn, err := p.Listener.Accept()
5558
if err != nil {
56-
if os.IsTimeout(err) {
57-
// Just a timeout, continue to check context and accept again.
58-
continue
59+
if ctx.Err() != nil {
60+
break
5961
}
6062

61-
select {
62-
case <-ctx.Done():
63-
break Loop
64-
default:
65-
if p.OnError != nil {
66-
p.OnError(fmt.Errorf("accept local connection: %w", err))
67-
}
68-
continue
69-
}
63+
runErr = fmt.Errorf("accept local connection: %w", err)
64+
cancel()
65+
break
7066
}
7167

7268
p.activeConns.Add(1)
@@ -75,6 +71,7 @@ Loop:
7571

7672
// Wait for all connections to finish.
7773
p.activeConns.Wait()
74+
return runErr
7875
}
7976

8077
func (p *Proxy) handleConnection(ctx context.Context, localConn net.Conn) {
@@ -87,46 +84,62 @@ func (p *Proxy) handleConnection(ctx context.Context, localConn net.Conn) {
8784

8885
remoteConn, err := p.DialContext(dialCtx, "tcp", p.RemoteAddr)
8986
if err != nil {
90-
if p.OnError != nil {
87+
if ctx.Err() == nil && p.OnError != nil {
9188
p.OnError(fmt.Errorf("connect remote address '%s': %w", p.RemoteAddr, err))
9289
}
9390
return
9491
}
9592
defer remoteConn.Close()
9693

97-
// Bidirectional copy with proper half-close handling.
94+
// Closing both connections aborts both copies after cancellation or a copy error. A clean EOF still uses
95+
// half-close so the other direction can finish sending any remaining data.
96+
closeConnections := func() {
97+
localConn.Close()
98+
remoteConn.Close()
99+
}
100+
stopClose := context.AfterFunc(ctx, closeConnections)
101+
defer stopClose()
102+
98103
done := make(chan error, 2)
99104

100105
go func() {
101106
_, err := io.Copy(remoteConn, localConn)
107+
if err != nil {
108+
done <- err
109+
closeConnections()
110+
return
111+
}
102112
// Close write half of remote connection if supported.
103113
if hc, ok := remoteConn.(halfCloser); ok {
104114
hc.CloseWrite()
105115
}
106-
done <- err
116+
done <- nil
107117
}()
108118

109119
go func() {
110120
_, err := io.Copy(localConn, remoteConn)
121+
if err != nil {
122+
done <- err
123+
closeConnections()
124+
return
125+
}
111126
// Close write half of local connection if supported.
112127
if hc, ok := localConn.(halfCloser); ok {
113128
hc.CloseWrite()
114129
}
115-
done <- err
130+
done <- nil
116131
}()
117132

118-
// Wait for both copies to complete or context cancel.
133+
// Wait for both copies to complete. The first error is the original failure because a copy reports it before
134+
// closing the connections to unblock the other copy.
135+
var copyErr error
119136
for range 2 {
120-
select {
121-
case <-ctx.Done():
122-
// Close connections to abort ongoing copies.
123-
localConn.Close()
124-
remoteConn.Close()
125-
return
126-
case err = <-done:
127-
if err != nil && p.OnError != nil {
128-
p.OnError(fmt.Errorf("data copy: %w", err))
129-
}
137+
if err = <-done; err != nil && copyErr == nil {
138+
copyErr = err
130139
}
131140
}
141+
142+
if copyErr != nil && ctx.Err() == nil && p.OnError != nil {
143+
p.OnError(fmt.Errorf("data copy: %w", copyErr))
144+
}
132145
}

0 commit comments

Comments
 (0)