From 13471bd4a664be13b17fcedd0abeda785acc3453 Mon Sep 17 00:00:00 2001 From: Samuel K Date: Fri, 11 Sep 2026 01:08:50 -0600 Subject: [PATCH 1/5] fix(ssh): harden tunnel forwarding reliability --- cmd/workspace/ssh.go | 53 +++++++++++++- go.mod | 20 ++--- go.sum | 40 +++++----- pkg/devcontainer/sshtunnel/sshtunnel.go | 13 ++-- pkg/ssh/connection_counter.go | 65 ++++++++++++----- pkg/ssh/connection_counter_test.go | 16 ++++ pkg/ssh/forward.go | 97 ++++++++++++++----------- pkg/ssh/forward_test.go | 2 +- pkg/ssh/helper.go | 25 +++++++ pkg/tunnel/services.go | 21 ++---- 10 files changed, 239 insertions(+), 113 deletions(-) diff --git a/cmd/workspace/ssh.go b/cmd/workspace/ssh.go index b12b7a032..59d367a47 100644 --- a/cmd/workspace/ssh.go +++ b/cmd/workspace/ssh.go @@ -29,7 +29,9 @@ import ( ) const ( - DisableSSHKeepAlive time.Duration = 0 * time.Second + DisableSSHKeepAlive time.Duration = 0 * time.Second + sshKeepAliveProbeTimeout = 10 * time.Second + sshKeepAliveMaxFailures = 3 ) // SSHCmd holds the ssh cmd flags. @@ -685,27 +687,70 @@ func (cmd *SSHCmd) startServices( } } +type sshKeepAliveClient interface { + SendRequest(string, bool, []byte) (bool, []byte, error) + Close() error +} + func startSSHKeepAlive( ctx context.Context, - client *ssh.Client, + client sshKeepAliveClient, interval time.Duration, +) { + startSSHKeepAliveWithOptions(ctx, client, interval, sshKeepAliveProbeTimeout, sshKeepAliveMaxFailures) +} + +func startSSHKeepAliveWithOptions( + ctx context.Context, + client sshKeepAliveClient, + interval, probeTimeout time.Duration, + maxFailures int, ) { ticker := time.NewTicker(interval) defer ticker.Stop() + failures := 0 for { select { case <-ctx.Done(): return case <-ticker.C: - ok, _, err := client.SendRequest("keepalive@openssh.com", true, nil) - if err := checkKeepAliveResponse(ok, err); err != nil { + if err := sendBoundedKeepAlive(ctx, client, probeTimeout); err != nil { + if ctx.Err() != nil { + return + } log.Errorf("failed to send keepalive: %v", err) + failures++ + if failures >= maxFailures { + log.Errorf("SSH keepalive failed %d consecutive times; closing client", failures) + _ = client.Close() + return + } + continue } + failures = 0 } } } +func sendBoundedKeepAlive(ctx context.Context, client sshKeepAliveClient, timeout time.Duration) error { + result := make(chan error, 1) + go func() { + ok, _, err := client.SendRequest("keepalive@openssh.com", true, nil) + result <- checkKeepAliveResponse(ok, err) + }() + timer := time.NewTimer(timeout) + defer timer.Stop() + select { + case err := <-result: + return err + case <-ctx.Done(): + return ctx.Err() + case <-timer.C: + return fmt.Errorf("keepalive probe timed out after %s", timeout) + } +} + // A positive or negative SSH reply proves that the peer is alive. Only a // transport error means that the keepalive request failed. func checkKeepAliveResponse(_ bool, err error) error { diff --git a/go.mod b/go.mod index e267767be..45f5dd0b3 100644 --- a/go.mod +++ b/go.mod @@ -23,7 +23,7 @@ require ( github.com/devsy-org/agentapi v1.0.1 github.com/devsy-org/api v1.1.0 github.com/devsy-org/apiserver v1.5.3 - github.com/devsy-org/ssh v1.2.5 + github.com/devsy-org/ssh v1.2.9 github.com/distribution/reference v0.6.0 github.com/docker/cli v29.7.2+incompatible github.com/docker/docker v28.5.2+incompatible @@ -64,11 +64,11 @@ require ( go.uber.org/atomic v1.11.0 go.uber.org/goleak v1.3.0 go.uber.org/zap v1.28.0 - golang.org/x/crypto v0.55.0 - golang.org/x/mod v0.38.0 - golang.org/x/sync v0.22.0 - golang.org/x/sys v0.47.0 - golang.org/x/term v0.45.0 + golang.org/x/crypto v0.57.0 + golang.org/x/mod v0.41.0 + golang.org/x/sync v0.23.0 + golang.org/x/sys v0.48.0 + golang.org/x/term v0.46.0 google.golang.org/grpc v1.83.0 google.golang.org/protobuf v1.36.12-0.20260120151049-f2248ac996af gopkg.in/yaml.v3 v3.0.1 @@ -506,12 +506,12 @@ require ( go4.org/netipx v0.0.0-20231129151722-fdeea329fbba // indirect gocloud.dev v0.46.0 // indirect golang.org/x/exp v0.0.0-20260603202125-055de637280b // indirect - golang.org/x/net v0.57.1-0.20260729233039-99c3b0a8f463 // indirect + golang.org/x/net v0.58.0 // indirect golang.org/x/oauth2 v0.36.0 // indirect - golang.org/x/telemetry v0.0.0-20260708182218-49f421fb7959 // indirect - golang.org/x/text v0.41.0 // indirect + golang.org/x/telemetry v0.0.0-20260811182544-a038080d80e5 // indirect + golang.org/x/text v0.42.0 // indirect golang.org/x/time v0.15.0 // indirect - golang.org/x/tools v0.48.0 // indirect + golang.org/x/tools v0.49.0 // indirect golang.org/x/xerrors v0.0.0-20240903120638-7835f813f4da // indirect golang.zx2c4.com/wintun v0.0.0-20230126152724-0fa3db229ce2 // indirect golang.zx2c4.com/wireguard/windows v0.5.3 // indirect diff --git a/go.sum b/go.sum index 12d81e5f2..6c5253940 100644 --- a/go.sum +++ b/go.sum @@ -414,8 +414,8 @@ github.com/devsy-org/api v1.1.0 h1:l7T9k7RVwatwN4lxeDTF3iN6EYmfGZgR3ZTJMDGha1M= github.com/devsy-org/api v1.1.0/go.mod h1:mAZklKdnywJYiXDReBLte/H+3m69z6G7RHB3n1lI53Q= github.com/devsy-org/apiserver v1.5.3 h1:tFKMgPxxfvojJ+C+wo0oqSmbE9PTJ2ur1E2yFj0eISI= github.com/devsy-org/apiserver v1.5.3/go.mod h1:m7gpbrh++Hp8iEM5jaP7vjjODPbs9Bmh3l4+iPFg1jE= -github.com/devsy-org/ssh v1.2.5 h1:Z7gTanYs2ZslT1swTw4leoVVuDEmuNNhQi48W2kNqMU= -github.com/devsy-org/ssh v1.2.5/go.mod h1:6r5tZ+H9JFoMl6NrxXRgltg9HhEryTd69avY831Lio8= +github.com/devsy-org/ssh v1.2.9 h1:KHqX1xAplGFanm0FMSAojtiC9nV/UFxUeP/5jU5quak= +github.com/devsy-org/ssh v1.2.9/go.mod h1:Uff10+cSSDZk3bG07u5D9+eQ8GMGqsgE70WCUJWcHv4= github.com/devsy-org/tailscale v1.102.2 h1:9SB6htvO+HmG8alal8WGCshHapfHR7dUFRSMYiFIzIM= github.com/devsy-org/tailscale v1.102.2/go.mod h1:kQUA0lYb/bqCJJZzShx+gqLSxggEFgXB4VDVZDzxWc8= github.com/dghubble/go-twitter v0.0.0-20211115160449-93a8679adecb h1:7ENzkH+O3juL+yj2undESLTaAeRllHwCs/b8z6aWSfc= @@ -1459,8 +1459,8 @@ golang.org/x/crypto v0.0.0-20210921155107-089bfa567519/go.mod h1:GvvjBRRGRdwPK5y golang.org/x/crypto v0.0.0-20220622213112-05595931fe9d/go.mod h1:IxCIyHEi3zRg3s0A5j5BB6A9Jmi73HwBIUl50j+osU4= golang.org/x/crypto v0.0.0-20220722155217-630584e8d5aa/go.mod h1:IxCIyHEi3zRg3s0A5j5BB6A9Jmi73HwBIUl50j+osU4= golang.org/x/crypto v0.17.0/go.mod h1:gCAAfMLgwOJRpTjQ2zCCt2OcSfYMTeZVSRtQlPC7Nq4= -golang.org/x/crypto v0.55.0 h1:+KWHjbgOaAQ66dh/YlkZKHlz9ZUlq61AFirAR9ntP8M= -golang.org/x/crypto v0.55.0/go.mod h1:uq0V9dE/fzQuJtbnL+2EhWOE63vo164FY8xqEnV9xis= +golang.org/x/crypto v0.57.0 h1:3ZVCjf8Ggz7zneR/EHRVx68Ctf+2pmIMP2UFhh9cC6M= +golang.org/x/crypto v0.57.0/go.mod h1:Fdz0i5U6CoizGwLda9DttjSk6qlZo25zYNtR+ycvuZA= golang.org/x/exp v0.0.0-20190121172915-509febef88a4/go.mod h1:CJ0aWSM057203Lf6IL+f9T1iT9GByDxfZKAQTCR3kQA= golang.org/x/exp v0.0.0-20260603202125-055de637280b h1:v1uXiEBHo8QA0LiGCo7UgHMzHT4Kdfpl2zmtH5vaP1Q= golang.org/x/exp v0.0.0-20260603202125-055de637280b/go.mod h1:d2fgXJLVs4dYDHUk5lwMIfzRzSrWCfGZb0ZqeLa/Vcw= @@ -1478,8 +1478,8 @@ golang.org/x/mod v0.3.0/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA= golang.org/x/mod v0.4.2/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA= golang.org/x/mod v0.6.0-dev.0.20220419223038-86c51ed26bb4/go.mod h1:jJ57K6gSWd91VN4djpZkiMVwK6gcyfeH4XE8wZrZaV4= golang.org/x/mod v0.8.0/go.mod h1:iBbtSCu2XBx23ZKBPSOrRkjjQPZFPuis4dIYUhu/chs= -golang.org/x/mod v0.38.0 h1:MECBjubtXD7yj4HrhIUcywNaGeNVUdfVnxmPajOk4yk= -golang.org/x/mod v0.38.0/go.mod h1:V6Xz0pq8TQ3dGqVQ1FVHuelZpAL0uNhSkk9ogYP3c40= +golang.org/x/mod v0.41.0 h1:qJmnOUb4YB+FsEuM3HcWucdZASCPGhsX6uljO6pog0c= +golang.org/x/mod v0.41.0/go.mod h1:Ek9pY8RKWXwsWvd3rQiHYtMqkjSUV+s1Rj7j4H5Ur6o= golang.org/x/net v0.0.0-20180724234803-3673e40ba225/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4= golang.org/x/net v0.0.0-20180826012351-8a410e7b638d/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4= golang.org/x/net v0.0.0-20190213061140-3a22650c66bd/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4= @@ -1495,8 +1495,8 @@ golang.org/x/net v0.0.0-20211112202133-69e39bad7dc2/go.mod h1:9nx3DQGgdP8bBQD5qx golang.org/x/net v0.0.0-20220722155237-a158d28d115b/go.mod h1:XRhObCWvk6IyKnWLug+ECip1KBveYUHfp+8e9klMJ9c= golang.org/x/net v0.6.0/go.mod h1:2Tu9+aMcznHK/AK1HMvgo6xiTLG5rD5rZLDS+rp2Bjs= golang.org/x/net v0.10.0/go.mod h1:0qNGK6F8kojg2nk9dLZ2mShWaEBan6FAoqfSigmmuDg= -golang.org/x/net v0.57.1-0.20260729233039-99c3b0a8f463 h1:56+vNxUDjzdwAJjglL/DMnp31nd8UeytlPRLbZfz55Q= -golang.org/x/net v0.57.1-0.20260729233039-99c3b0a8f463/go.mod h1:KpXc8iv+r3XplLAG/f7Jsf9RPszJzdR0f58q9vGOuEU= +golang.org/x/net v0.58.0 h1:ynWG7rqYi4ccpTEuPZ2QGWHktVEM9DMCj9yzDE0Q7To= +golang.org/x/net v0.58.0/go.mod h1:YwCddHnFlT7eLQqVprV19OnhLGtc5xOKgE0RyqgfWAU= golang.org/x/oauth2 v0.0.0-20180821212333-d2e6202438be/go.mod h1:N/0e6XlmueqKjAGxoOufVs8QHGRruUQn6yWY3a++T0U= golang.org/x/oauth2 v0.36.0 h1:peZ/1z27fi9hUOFCAZaHyrpWG5lwe0RJEEEeH0ThlIs= golang.org/x/oauth2 v0.36.0/go.mod h1:YDBUJMTkDnJS+A4BP4eZBjCqtokkg1hODuPjwiGPO7Q= @@ -1508,8 +1508,8 @@ golang.org/x/sync v0.0.0-20201020160332-67f06af15bc9/go.mod h1:RxMgew5VJxzue5/jJ golang.org/x/sync v0.0.0-20210220032951-036812b2e83c/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= golang.org/x/sync v0.0.0-20220722155255-886fb9371eb4/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= golang.org/x/sync v0.1.0/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= -golang.org/x/sync v0.22.0 h1:SZjpbeLmrCk4xhRSZFNZW5gFUeCeFgjekvI/+gfScek= -golang.org/x/sync v0.22.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0= +golang.org/x/sync v0.23.0 h1:KameEIfc1IkluZyXWLn39Wd4tURc6GbCiISGiZm2bQk= +golang.org/x/sync v0.23.0/go.mod h1:sUUOizhqBxiL6pEWpqNLUiaJn1ShEbZ6BBqskPbjZm0= golang.org/x/sys v0.0.0-20180830151530-49385e6e1522/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= golang.org/x/sys v0.0.0-20190222072716-a9d3bda3a223/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= @@ -1537,17 +1537,17 @@ golang.org/x/sys v0.5.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.8.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.15.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA= -golang.org/x/sys v0.47.0 h1:o7XGOvZQCADBQQ4Y7VNq2dRWQR7JmOUW8Kxx4ZsNgWs= -golang.org/x/sys v0.47.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw= -golang.org/x/telemetry v0.0.0-20260708182218-49f421fb7959 h1:RJhm5l6Fo4rmEIcndxDllNhhf/fAx8qIm4t6A7vpm2A= -golang.org/x/telemetry v0.0.0-20260708182218-49f421fb7959/go.mod h1:LV7u5Oco+Z/g6XI7PqN+EUUUGGkEcmB1uj2ceI0fOVg= +golang.org/x/sys v0.48.0 h1:bbX/i/6MgT9BVLM9RT1thmxL04yeTAhbEz4SyadbXoo= +golang.org/x/sys v0.48.0/go.mod h1:hNLxWAXmnKAxqDtdwIYC4bM9oQPEecfsnNMuSxOs3og= +golang.org/x/telemetry v0.0.0-20260811182544-a038080d80e5 h1:ZUSxONxc981v7AW7QUg+I9WwZzSTTJ019ENBYr5pV/Q= +golang.org/x/telemetry v0.0.0-20260811182544-a038080d80e5/go.mod h1:LVehoXe41cL5SCVQilsV7Gg6BNG+Js6P9PhSbYTIUkQ= golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo= golang.org/x/term v0.0.0-20210927222741-03fcf44c2211/go.mod h1:jbD1KX2456YbFQfuXm/mYQcufACuNUgVhRMnK/tPxf8= golang.org/x/term v0.5.0/go.mod h1:jMB1sMXY+tzblOD4FWmEbocvup2/aLOaQEp7JmGp78k= golang.org/x/term v0.8.0/go.mod h1:xPskH00ivmX89bAKVGSKKtLOWNx2+17Eiy94tnKShWo= golang.org/x/term v0.15.0/go.mod h1:BDl952bC7+uMoWR75FIrCDx79TPU9oHkTZ9yRbYOrX0= -golang.org/x/term v0.45.0 h1:NwWyBmoJCbfTHpxrWoZ9C6/VxOf7ic219I8xZZFdrf0= -golang.org/x/term v0.45.0/go.mod h1:9aqxs0blBcrm/n0L9QW0aRVD+ktan8ssZromtqJC43w= +golang.org/x/term v0.46.0 h1:3+OXuTbaKDgwk8jTi3aSLHRlmWqHEUDUtxnbFigO4YE= +golang.org/x/term v0.46.0/go.mod h1:+K02xbkittuwc0Am4abfA3Fc+XRGXkvBXNO88NCXPoc= golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ= golang.org/x/text v0.3.2/go.mod h1:bEr9sfX3Q8Zfm5fL9x+3itogRgK3+ptLWKqgva+5dAk= golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ= @@ -1558,8 +1558,8 @@ golang.org/x/text v0.4.0/go.mod h1:mrYo+phRRbMaCq/xk9113O4dZlRixOauAjOtrjsXDZ8= golang.org/x/text v0.7.0/go.mod h1:mrYo+phRRbMaCq/xk9113O4dZlRixOauAjOtrjsXDZ8= golang.org/x/text v0.9.0/go.mod h1:e1OnstbJyHTd6l/uOt8jFFHp6TRDWZR/bV3emEE/zU8= golang.org/x/text v0.14.0/go.mod h1:18ZOQIKpY8NJVqYksKHtTdi31H5itFRjB5/qKTNYzSU= -golang.org/x/text v0.41.0 h1:vz/seA0lnX87Othu2f/0L24RcgrXD9/YFTSuGjj3rH8= -golang.org/x/text v0.41.0/go.mod h1:jvf1O8ajNzZqhSrQBPbutR/EB83Cc0CFrezNQIwbb5M= +golang.org/x/text v0.42.0 h1:JbOZXgfeCPU9gacVtYliJqOhD+zhrEqK4LfdpmlUZqI= +golang.org/x/text v0.42.0/go.mod h1:ojzP1Z+2QtioaF8DTtO8K5q7JWVVYwZKenzujK0Zd0E= golang.org/x/time v0.15.0 h1:bbrp8t3bGUeFOx08pvsMYRTCVSMk89u4tKbNOZbp88U= golang.org/x/time v0.15.0/go.mod h1:Y4YMaQmXwGQZoFaVFk4YpCt4FLQMYKZe9oeV/f4MSno= golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ= @@ -1577,8 +1577,8 @@ golang.org/x/tools v0.0.0-20210106214847-113979e3529a/go.mod h1:emZCQorbCU4vsT4f golang.org/x/tools v0.1.5/go.mod h1:o0xws9oXOQQZyjljx8fwUC0k7L1pTE6eaCbjGeHmOkk= golang.org/x/tools v0.1.12/go.mod h1:hNGJHUnrk76NpqgfD5Aqm5Crs+Hm0VOH/i9J2+nxYbc= golang.org/x/tools v0.6.0/go.mod h1:Xwgl3UAJ/d3gWutnCtw505GrjyAbvKui8lOU390QaIU= -golang.org/x/tools v0.48.0 h1:3+hClM1aLL5mjMKm5ovokw9epgRXPuu2tILgismM6RE= -golang.org/x/tools v0.48.0/go.mod h1:08xX0orndb/F7jJxGDicx061tyd5pcMto75YMAXr6lk= +golang.org/x/tools v0.49.0 h1:3NI7VXzL9+1WZD52Dx2ttoPwD5DWrFGpl9mFZDlmisI= +golang.org/x/tools v0.49.0/go.mod h1:SJNXV9DBKT0UbdttsQjbfJlAE/q+y36++zo3uL3N0Oo= golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= golang.org/x/xerrors v0.0.0-20191011141410-1b5146add898/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= diff --git a/pkg/devcontainer/sshtunnel/sshtunnel.go b/pkg/devcontainer/sshtunnel/sshtunnel.go index 8c100c2ba..ab9841871 100644 --- a/pkg/devcontainer/sshtunnel/sshtunnel.go +++ b/pkg/devcontainer/sshtunnel/sshtunnel.go @@ -153,8 +153,8 @@ func runSSHTunnel(ctx context.Context, p sshTunnelParams) (*config2.Result, erro err = p.grpcBridge.RunPair(ctx, func(ctx context.Context, stdin, stdout *os.File) error { return runCommandInSSHTunnel(ctx, sshCommandParams{ - sshClient: sshClient, - command: p.opts.Command, + session: sess, + command: p.opts.Command, }, stdin, stdout) }, func(ctx context.Context, stdout, stdin *os.File) error { @@ -231,8 +231,8 @@ func setupSSHAgentForwarding(sshClient *ssh.Client, sess *ssh.Session) { } type sshCommandParams struct { - sshClient *ssh.Client - command string + session *ssh.Session + command string } func runCommandInSSHTunnel(ctx context.Context, p sshCommandParams, stdin, stdout *os.File) error { @@ -240,8 +240,7 @@ func runCommandInSSHTunnel(ctx context.Context, p sshCommandParams, stdin, stdou defer func() { _ = streamer.Close() }() log.Debugf("running agent command in SSH tunnel: %q", p.command) - err := devssh.Run(ctx, devssh.RunOptions{ - Client: p.sshClient, + err := devssh.RunSession(ctx, p.session, devssh.RunSessionOptions{ Command: p.command, Stdin: stdin, Stdout: stdout, @@ -257,7 +256,7 @@ func runCommandInSSHTunnel(ctx context.Context, p sshCommandParams, stdin, stdou return nil } -const maxLogLines = 1 +const maxLogLines = 25 func newSSHTunnelJSONLogStreamer() *log.JSONLogStreamer { return log.NewJSONLogStreamer(log.StreamerOptions{ diff --git a/pkg/ssh/connection_counter.go b/pkg/ssh/connection_counter.go index c0b6ede66..8344cf598 100644 --- a/pkg/ssh/connection_counter.go +++ b/pkg/ssh/connection_counter.go @@ -14,12 +14,16 @@ func newConnectionCounter( onTimeout func(), address string, ) *connectionCounter { - return &connectionCounter{ + c := &connectionCounter{ ctx: ctx, address: address, timeout: timeout, onTimeout: onTimeout, } + c.m.Lock() + c.armTimeoutLocked() + c.m.Unlock() + return c } type connectionCounter struct { @@ -31,13 +35,43 @@ type connectionCounter struct { m sync.Mutex connections int - generation int + timer *time.Timer + closed bool +} + +func (c *connectionCounter) armTimeoutLocked() { + if c.closed || c.connections != 0 || c.timeout <= 0 || c.ctx.Err() != nil { + return + } + if c.timer != nil { + c.timer.Stop() + } + var timer *time.Timer + timer = time.AfterFunc(c.timeout, func() { + c.m.Lock() + if c.timer != timer || c.closed || c.connections != 0 || c.ctx.Err() != nil { + c.m.Unlock() + return + } + c.timer = nil + onTimeout := c.onTimeout + c.m.Unlock() + onTimeout() + }) + c.timer = timer } func (c *connectionCounter) Add() { c.m.Lock() defer c.m.Unlock() + if c.closed { + return + } + if c.timer != nil { + c.timer.Stop() + c.timer = nil + } c.connections++ log.Debugf("New connection on %s (Total: %d)", c.address, c.connections) } @@ -46,27 +80,24 @@ func (c *connectionCounter) Dec() { c.m.Lock() defer c.m.Unlock() + if c.closed { + return + } if c.connections <= 0 { c.connections = 0 } else { c.connections-- } log.Debugf("Closed connection on %s (Total: %d)", c.address, c.connections) - if c.connections <= 0 && c.timeout > 0 { - c.generation++ - - go func(generation int) { - select { - case <-c.ctx.Done(): - return - case <-time.After(c.timeout): - c.m.Lock() - defer c.m.Unlock() + c.armTimeoutLocked() +} - if c.generation == generation && c.connections <= 0 { - c.onTimeout() - } - } - }(c.generation) +func (c *connectionCounter) Close() { + c.m.Lock() + defer c.m.Unlock() + c.closed = true + if c.timer != nil { + c.timer.Stop() + c.timer = nil } } diff --git a/pkg/ssh/connection_counter_test.go b/pkg/ssh/connection_counter_test.go index 8a4df9693..c541aaf58 100644 --- a/pkg/ssh/connection_counter_test.go +++ b/pkg/ssh/connection_counter_test.go @@ -43,6 +43,22 @@ func TestConnectionCounter_IdleTimeoutFires(t *testing.T) { assert.Equal(t, 0, c.connections, "count must remain zero when the timeout fires") } +func TestConnectionCounter_InitialZeroConnectionsStartsTimeout(t *testing.T) { + c, calls := newRecordingCounter(t, 10*time.Millisecond) + defer c.Close() + + require.Eventually(t, func() bool { return calls.Load() == 1 }, + time.Second, time.Millisecond, "an initially idle counter should time out") +} + +func TestConnectionCounter_CloseStopsPendingTimeout(t *testing.T) { + c, calls := newRecordingCounter(t, 10*time.Millisecond) + c.Close() + + time.Sleep(30 * time.Millisecond) + assert.Zero(t, calls.Load()) +} + func TestConnectionCounter_NewConnectionBeforeTimeoutCancelsIt(t *testing.T) { c, calls := newRecordingCounter(t, 50*time.Millisecond) diff --git a/pkg/ssh/forward.go b/pkg/ssh/forward.go index 9fd4b9a34..e82161a24 100644 --- a/pkg/ssh/forward.go +++ b/pkg/ssh/forward.go @@ -3,9 +3,9 @@ package ssh import ( "context" "errors" + "fmt" "io" "net" - "sync" "time" "github.com/devsy-org/devsy/pkg/log" @@ -21,6 +21,7 @@ var ErrIdleTimeout = errors.New("port forward idle timeout") var ErrTransportClosed = errors.New("ssh transport closed") type ForwardingFunction func( + context.Context, net.Conn, *ssh.Client, string, @@ -151,6 +152,7 @@ func portForwarding( ) cancel(ErrIdleTimeout) }, srcAddr) + defer counter.Close() for { // waiting for a new connection connection, err := listener.Accept() @@ -171,7 +173,7 @@ func portForwarding( go func() { defer counter.Dec() - forwardFn(connection, client, dstNetwork, dstAddr) + forwardFn(fwdCtx, connection, client, dstNetwork, dstAddr) }() } } @@ -202,10 +204,12 @@ func watchTransportClosed( } func forward( + ctx context.Context, localConn net.Conn, client *ssh.Client, remoteNetwork, remoteAddr string, ) { + defer func() { _ = localConn.Close() }() // Setup sshConn (type net.Conn) sshConn, err := client.Dial(remoteNetwork, remoteAddr) if err != nil { @@ -213,35 +217,18 @@ func forward( return } defer func() { _ = sshConn.Close() }() - - // Copy localConn.Reader to sshConn.Writer - waitGroup := sync.WaitGroup{} - waitGroup.Go(func() { - defer func() { _ = sshConn.Close() }() - - _, err = io.Copy(sshConn, localConn) - if err != nil { - log.Debugf("error copying to remote: %v", err) - } - }) - - // Copy sshConn.Reader to localConn.Writer - waitGroup.Go(func() { - defer func() { _ = localConn.Close() }() - - _, err = io.Copy(localConn, sshConn) - if err != nil { - log.Debugf("error copying to local: %v", err) - } - }) - waitGroup.Wait() + if err := relayDuplex(ctx, localConn, sshConn); err != nil { + log.Debugf("error forwarding connection: %v", err) + } } func reverseForward( + ctx context.Context, remoteConn net.Conn, client *ssh.Client, localNetwork, localAddr string, ) { + defer func() { _ = remoteConn.Close() }() // Setup localConn (type net.Conn) localConn, err := net.Dial(localNetwork, localAddr) if err != nil { @@ -249,26 +236,54 @@ func reverseForward( return } defer func() { _ = localConn.Close() }() + if err := relayDuplex(ctx, remoteConn, localConn); err != nil { + log.Debugf("error forwarding reverse connection: %v", err) + } +} - // Copy localConn.Reader to sshConn.Writer - waitGroup := sync.WaitGroup{} - waitGroup.Go(func() { - defer func() { _ = localConn.Close() }() +type closeWriter interface{ CloseWrite() error } - _, err = io.Copy(localConn, remoteConn) - if err != nil { - log.Debugf("error copying to local: %v", err) - } - }) +type relayResult struct { + direction string + err error +} - // Copy sshConn.Reader to localConn.Writer - waitGroup.Go(func() { - defer func() { _ = remoteConn.Close() }() +func relayOneWay(dst net.Conn, src net.Conn, direction string, results chan<- relayResult) { + _, err := io.Copy(dst, src) + if err == nil { + cw, ok := dst.(closeWriter) + if !ok { + err = errors.New("destination does not support CloseWrite") + } else if closeErr := cw.CloseWrite(); closeErr != nil && !errors.Is(closeErr, io.EOF) { + err = closeErr + } + } + results <- relayResult{direction: direction, err: err} +} - _, err = io.Copy(remoteConn, localConn) - if err != nil { - log.Debugf("error copying to remote: %v", err) +func relayDuplex(ctx context.Context, left, right net.Conn) error { + results := make(chan relayResult, 2) + stop := make(chan struct{}) + defer close(stop) + go func() { + select { + case <-ctx.Done(): + _ = left.Close() + _ = right.Close() + case <-stop: } - }) - waitGroup.Wait() + }() + go relayOneWay(right, left, "local-to-remote", results) + go relayOneWay(left, right, "remote-to-local", results) + + var firstErr error + for i := 0; i < 2; i++ { + result := <-results + if result.err != nil && firstErr == nil { + firstErr = fmt.Errorf("%s relay: %w", result.direction, result.err) + _ = left.Close() + _ = right.Close() + } + } + return firstErr } diff --git a/pkg/ssh/forward_test.go b/pkg/ssh/forward_test.go index 88f2d8831..fd1dd4a2e 100644 --- a/pkg/ssh/forward_test.go +++ b/pkg/ssh/forward_test.go @@ -14,7 +14,7 @@ import ( // noopForward is a ForwardingFunction that immediately closes the local // connection. This lets portForwarding exercise its idle/cancellation paths // without needing a real *ssh.Client. -func noopForward(localConn net.Conn, _ *ssh.Client, _, _ string) { +func noopForward(_ context.Context, localConn net.Conn, _ *ssh.Client, _, _ string) { _ = localConn.Close() } diff --git a/pkg/ssh/helper.go b/pkg/ssh/helper.go index 1fe7ce2b6..6293e6391 100644 --- a/pkg/ssh/helper.go +++ b/pkg/ssh/helper.go @@ -136,6 +136,14 @@ type RunOptions struct { EnvVars map[string]string } +type RunSessionOptions struct { + Command string + Stdin io.Reader + Stdout io.Writer + Stderr io.Writer + EnvVars map[string]string +} + // ExitError wraps an SSH exit error with the exit code. type ExitError struct { ExitCode int @@ -173,7 +181,24 @@ func Run(ctx context.Context, opts RunOptions) error { return fmt.Errorf("failed to create SSH session: %w", err) } defer func() { _ = sess.Close() }() + return RunSession(ctx, sess, RunSessionOptions{ + Command: opts.Command, + Stdin: opts.Stdin, + Stdout: opts.Stdout, + Stderr: opts.Stderr, + EnvVars: opts.EnvVars, + }) +} +// RunSession executes a command on a caller-owned SSH session. It never closes +// the session; callers retain ownership and are responsible for cleanup. +func RunSession(ctx context.Context, sess *ssh.Session, opts RunSessionOptions) error { + if sess == nil { + return fmt.Errorf("SSH session is required") + } + if opts.Command == "" { + return fmt.Errorf("command is required") + } // Set environment variables (best effort - SSH servers may reject env vars or not support them) for k, v := range opts.EnvVars { _ = sess.Setenv(k, v) // Ignore errors - command should work without env vars diff --git a/pkg/tunnel/services.go b/pkg/tunnel/services.go index 7c6223a57..b26ff6274 100644 --- a/pkg/tunnel/services.go +++ b/pkg/tunnel/services.go @@ -248,12 +248,13 @@ func RunServices(ctx context.Context, opts RunServicesOptions) error { extraListeners: opts.ExtraListeners, } - forwardedPorts, err := forwardDevContainerPorts(ctx, fp) + result, err := getContainerResult(ctx, fp) if err != nil { - return fmt.Errorf("forward ports: %w", err) + return fmt.Errorf("retrieve container result: %w", err) } + forwardedPorts := forwardDevContainerPorts(ctx, fp, result) - resolver := buildPortAttributeResolver(ctx, fp) + resolver := buildPortAttributeResolver(result) return retry.OnError(wait.Backoff{ Steps: maxRetrySteps, @@ -268,9 +269,8 @@ func RunServices(ctx context.Context, opts RunServicesOptions) error { } // buildPortAttributeResolver loads port attributes from the container result. -func buildPortAttributeResolver(ctx context.Context, p portForwardParams) PortAttributeResolver { - result, err := getContainerResult(ctx, p) - if err != nil || result == nil || result.MergedConfig == nil { +func buildPortAttributeResolver(result *config2.Result) PortAttributeResolver { + if result == nil || result.MergedConfig == nil { return nil } mc := result.MergedConfig @@ -289,18 +289,13 @@ type portForwardParams struct { } // forwardDevContainerPorts forwards all the ports defined in the devcontainer.json. -func forwardDevContainerPorts(ctx context.Context, p portForwardParams) ([]string, error) { - result, err := getContainerResult(ctx, p) - if err != nil { - return nil, err - } - +func forwardDevContainerPorts(ctx context.Context, p portForwardParams, result *config2.Result) []string { forwardedPorts := []string{} forwardedPorts = append(forwardedPorts, forwardExtraPorts(ctx, p)...) forwardedPorts = append(forwardedPorts, forwardAppPorts(ctx, p, result)...) forwardedPorts = append(forwardedPorts, forwardConfigPorts(ctx, p, result)...) - return forwardedPorts, nil + return forwardedPorts } // getContainerResult retrieves and parses the container result. From b7cabe19fb1caade55b8b7bda30c433171cdfc6d Mon Sep 17 00:00:00 2001 From: Samuel K Date: Fri, 11 Sep 2026 01:28:02 -0600 Subject: [PATCH 2/5] fix(lint): resolve CI lint findings --- cmd/workspace/ssh.go | 32 ++++++++++++----- pkg/ssh/connection_counter.go | 67 ++++++++++++++++++++++++----------- pkg/ssh/forward.go | 28 ++++++++------- pkg/ssh/forward_test.go | 2 +- pkg/tunnel/services.go | 6 +++- 5 files changed, 91 insertions(+), 44 deletions(-) diff --git a/cmd/workspace/ssh.go b/cmd/workspace/ssh.go index 59d367a47..d350ef432 100644 --- a/cmd/workspace/ssh.go +++ b/cmd/workspace/ssh.go @@ -692,21 +692,30 @@ type sshKeepAliveClient interface { Close() error } +type sshKeepAliveOptions struct { + interval time.Duration + probeTimeout time.Duration + maxFailures int +} + func startSSHKeepAlive( ctx context.Context, client sshKeepAliveClient, interval time.Duration, ) { - startSSHKeepAliveWithOptions(ctx, client, interval, sshKeepAliveProbeTimeout, sshKeepAliveMaxFailures) + startSSHKeepAliveWithOptions(ctx, client, sshKeepAliveOptions{ + interval: interval, + probeTimeout: sshKeepAliveProbeTimeout, + maxFailures: sshKeepAliveMaxFailures, + }) } func startSSHKeepAliveWithOptions( ctx context.Context, client sshKeepAliveClient, - interval, probeTimeout time.Duration, - maxFailures int, + opts sshKeepAliveOptions, ) { - ticker := time.NewTicker(interval) + ticker := time.NewTicker(opts.interval) defer ticker.Stop() failures := 0 @@ -715,14 +724,17 @@ func startSSHKeepAliveWithOptions( case <-ctx.Done(): return case <-ticker.C: - if err := sendBoundedKeepAlive(ctx, client, probeTimeout); err != nil { + if err := sendBoundedKeepAlive(ctx, client, opts.probeTimeout); err != nil { if ctx.Err() != nil { return } log.Errorf("failed to send keepalive: %v", err) failures++ - if failures >= maxFailures { - log.Errorf("SSH keepalive failed %d consecutive times; closing client", failures) + if failures >= opts.maxFailures { + log.Errorf( + "SSH keepalive failed %d consecutive times; closing client", + failures, + ) _ = client.Close() return } @@ -733,7 +745,11 @@ func startSSHKeepAliveWithOptions( } } -func sendBoundedKeepAlive(ctx context.Context, client sshKeepAliveClient, timeout time.Duration) error { +func sendBoundedKeepAlive( + ctx context.Context, + client sshKeepAliveClient, + timeout time.Duration, +) error { result := make(chan error, 1) go func() { ok, _, err := client.SendRequest("keepalive@openssh.com", true, nil) diff --git a/pkg/ssh/connection_counter.go b/pkg/ssh/connection_counter.go index 8344cf598..dbe3a7896 100644 --- a/pkg/ssh/connection_counter.go +++ b/pkg/ssh/connection_counter.go @@ -36,30 +36,11 @@ type connectionCounter struct { m sync.Mutex connections int timer *time.Timer + timerToken *connectionTimerToken closed bool } -func (c *connectionCounter) armTimeoutLocked() { - if c.closed || c.connections != 0 || c.timeout <= 0 || c.ctx.Err() != nil { - return - } - if c.timer != nil { - c.timer.Stop() - } - var timer *time.Timer - timer = time.AfterFunc(c.timeout, func() { - c.m.Lock() - if c.timer != timer || c.closed || c.connections != 0 || c.ctx.Err() != nil { - c.m.Unlock() - return - } - c.timer = nil - onTimeout := c.onTimeout - c.m.Unlock() - onTimeout() - }) - c.timer = timer -} +type connectionTimerToken struct{} func (c *connectionCounter) Add() { c.m.Lock() @@ -72,6 +53,7 @@ func (c *connectionCounter) Add() { c.timer.Stop() c.timer = nil } + c.timerToken = nil c.connections++ log.Debugf("New connection on %s (Total: %d)", c.address, c.connections) } @@ -100,4 +82,47 @@ func (c *connectionCounter) Close() { c.timer.Stop() c.timer = nil } + c.timerToken = nil +} + +func (c *connectionCounter) armTimeoutLocked() { + if !c.canArmTimeoutLocked() { + return + } + if c.timer != nil { + c.timer.Stop() + c.timer = nil + } + token := new(connectionTimerToken) + c.timerToken = token + c.timer = time.AfterFunc(c.timeout, func() { + c.handleTimeout(token) + }) +} + +func (c *connectionCounter) canArmTimeoutLocked() bool { + if c.closed || c.connections != 0 { + return false + } + if c.timeout <= 0 { + return false + } + return c.ctx.Err() == nil +} + +func (c *connectionCounter) handleTimeout(token *connectionTimerToken) { + c.m.Lock() + if c.timerToken != token { + c.m.Unlock() + return + } + if !c.canArmTimeoutLocked() { + c.m.Unlock() + return + } + c.timer = nil + c.timerToken = nil + onTimeout := c.onTimeout + c.m.Unlock() + onTimeout() } diff --git a/pkg/ssh/forward.go b/pkg/ssh/forward.go index e82161a24..3970a8076 100644 --- a/pkg/ssh/forward.go +++ b/pkg/ssh/forward.go @@ -20,13 +20,12 @@ var ErrIdleTimeout = errors.New("port forward idle timeout") var ErrTransportClosed = errors.New("ssh transport closed") -type ForwardingFunction func( - context.Context, - net.Conn, - *ssh.Client, - string, - string, -) +type forwardTarget struct { + network string + address string +} + +type ForwardingFunction func(context.Context, net.Conn, *ssh.Client, forwardTarget) func PortForward( ctx context.Context, @@ -173,7 +172,10 @@ func portForwarding( go func() { defer counter.Dec() - forwardFn(fwdCtx, connection, client, dstNetwork, dstAddr) + forwardFn(fwdCtx, connection, client, forwardTarget{ + network: dstNetwork, + address: dstAddr, + }) }() } } @@ -207,11 +209,11 @@ func forward( ctx context.Context, localConn net.Conn, client *ssh.Client, - remoteNetwork, remoteAddr string, + target forwardTarget, ) { defer func() { _ = localConn.Close() }() // Setup sshConn (type net.Conn) - sshConn, err := client.Dial(remoteNetwork, remoteAddr) + sshConn, err := client.Dial(target.network, target.address) if err != nil { log.Debugf("error dialing remote: %v", err) return @@ -226,11 +228,11 @@ func reverseForward( ctx context.Context, remoteConn net.Conn, client *ssh.Client, - localNetwork, localAddr string, + target forwardTarget, ) { defer func() { _ = remoteConn.Close() }() // Setup localConn (type net.Conn) - localConn, err := net.Dial(localNetwork, localAddr) + localConn, err := net.Dial(target.network, target.address) if err != nil { log.Debugf("error dialing remote: %v", err) return @@ -277,7 +279,7 @@ func relayDuplex(ctx context.Context, left, right net.Conn) error { go relayOneWay(left, right, "remote-to-local", results) var firstErr error - for i := 0; i < 2; i++ { + for range 2 { result := <-results if result.err != nil && firstErr == nil { firstErr = fmt.Errorf("%s relay: %w", result.direction, result.err) diff --git a/pkg/ssh/forward_test.go b/pkg/ssh/forward_test.go index fd1dd4a2e..4271ea792 100644 --- a/pkg/ssh/forward_test.go +++ b/pkg/ssh/forward_test.go @@ -14,7 +14,7 @@ import ( // noopForward is a ForwardingFunction that immediately closes the local // connection. This lets portForwarding exercise its idle/cancellation paths // without needing a real *ssh.Client. -func noopForward(_ context.Context, localConn net.Conn, _ *ssh.Client, _, _ string) { +func noopForward(_ context.Context, localConn net.Conn, _ *ssh.Client, _ forwardTarget) { _ = localConn.Close() } diff --git a/pkg/tunnel/services.go b/pkg/tunnel/services.go index b26ff6274..b0a84862f 100644 --- a/pkg/tunnel/services.go +++ b/pkg/tunnel/services.go @@ -289,7 +289,11 @@ type portForwardParams struct { } // forwardDevContainerPorts forwards all the ports defined in the devcontainer.json. -func forwardDevContainerPorts(ctx context.Context, p portForwardParams, result *config2.Result) []string { +func forwardDevContainerPorts( + ctx context.Context, + p portForwardParams, + result *config2.Result, +) []string { forwardedPorts := []string{} forwardedPorts = append(forwardedPorts, forwardExtraPorts(ctx, p)...) forwardedPorts = append(forwardedPorts, forwardAppPorts(ctx, p, result)...) From 7a91f057a241559d484f25692a04e163640ff5f0 Mon Sep 17 00:00:00 2001 From: Samuel K Date: Fri, 11 Sep 2026 08:50:17 -0600 Subject: [PATCH 3/5] fix(ssh): address review feedback --- pkg/ssh/connection_counter_test.go | 2 +- pkg/tunnel/services.go | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/pkg/ssh/connection_counter_test.go b/pkg/ssh/connection_counter_test.go index c541aaf58..97fcab255 100644 --- a/pkg/ssh/connection_counter_test.go +++ b/pkg/ssh/connection_counter_test.go @@ -52,7 +52,7 @@ func TestConnectionCounter_InitialZeroConnectionsStartsTimeout(t *testing.T) { } func TestConnectionCounter_CloseStopsPendingTimeout(t *testing.T) { - c, calls := newRecordingCounter(t, 10*time.Millisecond) + c, calls := newRecordingCounter(t, time.Second) c.Close() time.Sleep(30 * time.Millisecond) diff --git a/pkg/tunnel/services.go b/pkg/tunnel/services.go index b0a84862f..f61181530 100644 --- a/pkg/tunnel/services.go +++ b/pkg/tunnel/services.go @@ -250,7 +250,7 @@ func RunServices(ctx context.Context, opts RunServicesOptions) error { result, err := getContainerResult(ctx, fp) if err != nil { - return fmt.Errorf("retrieve container result: %w", err) + return err } forwardedPorts := forwardDevContainerPorts(ctx, fp, result) From 4b99be354af87b4abccbc50b9fab3588a9317e1e Mon Sep 17 00:00:00 2001 From: Samuel K Date: Fri, 11 Sep 2026 09:20:21 -0600 Subject: [PATCH 4/5] fix(ssh): guard timeout dispatch and bound diagnostics --- pkg/devcontainer/sshtunnel/sshtunnel.go | 6 +++- pkg/log/streamer.go | 44 +++++++++++++++++++++---- pkg/log/streamer_test.go | 24 ++++++++++++++ pkg/ssh/connection_counter.go | 9 +++-- pkg/ssh/connection_counter_test.go | 28 +++++++++++++--- pkg/ssh/forward.go | 5 ++- 6 files changed, 100 insertions(+), 16 deletions(-) diff --git a/pkg/devcontainer/sshtunnel/sshtunnel.go b/pkg/devcontainer/sshtunnel/sshtunnel.go index ab9841871..4d179709c 100644 --- a/pkg/devcontainer/sshtunnel/sshtunnel.go +++ b/pkg/devcontainer/sshtunnel/sshtunnel.go @@ -256,12 +256,16 @@ func runCommandInSSHTunnel(ctx context.Context, p sshCommandParams, stdin, stdou return nil } -const maxLogLines = 25 +const ( + maxLogLines = 25 + maxLogBytes = 256 * 1024 +) func newSSHTunnelJSONLogStreamer() *log.JSONLogStreamer { return log.NewJSONLogStreamer(log.StreamerOptions{ FallbackLevel: log.LevelDebug, CaptureLines: maxLogLines, + CaptureBytes: maxLogBytes, DetectLevelPrefixes: true, TreatUnknownJSONAsDebug: true, }) diff --git a/pkg/log/streamer.go b/pkg/log/streamer.go index 32f605609..648e2abd3 100644 --- a/pkg/log/streamer.go +++ b/pkg/log/streamer.go @@ -16,6 +16,9 @@ type StreamerOptions struct { FallbackLevel int // CaptureLines retains this many raw lines for ErrorOutput. CaptureLines int + // CaptureBytes bounds the total raw bytes retained for ErrorOutput, + // including newline separators. A zero value means no byte limit. + CaptureBytes int // DetectLevelPrefixes preserves a level from timestamp-prefixed plain text. DetectLevelPrefixes bool // TreatUnknownJSONAsDebug preserves the historical tunnel behavior for @@ -34,11 +37,13 @@ type JSONLogStreamer struct { detectLevelPrefixes bool treatUnknownJSONAsDebug bool captureLines int + captureBytes int - mu sync.Mutex - lastLines []string - closeOnce sync.Once - closeErr error + mu sync.Mutex + lastLines []string + capturedBytes int + closeOnce sync.Once + closeErr error } // NewJSONLogStreamer returns a writer that decodes Devsy JSON log lines while @@ -52,6 +57,7 @@ func NewJSONLogStreamer(options StreamerOptions) *JSONLogStreamer { detectLevelPrefixes: options.DetectLevelPrefixes, treatUnknownJSONAsDebug: options.TreatUnknownJSONAsDebug, captureLines: options.CaptureLines, + captureBytes: options.CaptureBytes, } if options.CaptureLines > 0 { streamer.lastLines = make([]string, 0, options.CaptureLines) @@ -126,15 +132,39 @@ func (s *JSONLogStreamer) process(reader io.Reader) { } func (s *JSONLogStreamer) capture(line string) { - if s.captureLines <= 0 { + if s.captureLines <= 0 && s.captureBytes <= 0 { return } + if s.captureBytes > 0 && len(line) > s.captureBytes { + line = line[len(line)-s.captureBytes:] + } s.mu.Lock() defer s.mu.Unlock() - if len(s.lastLines) >= s.captureLines { - s.lastLines = s.lastLines[1:] + for (s.captureLines > 0 && len(s.lastLines) >= s.captureLines) || + (s.captureBytes > 0 && len(s.lastLines) > 0 && + s.capturedBytes+1+len(line) > s.captureBytes) { + s.dropOldestLine() + } + if s.captureBytes > 0 && len(s.lastLines) > 0 && + s.capturedBytes+1+len(line) > s.captureBytes { + line = line[len(line)-(s.captureBytes-s.capturedBytes-1):] + } + if len(s.lastLines) > 0 { + s.capturedBytes++ } s.lastLines = append(s.lastLines, line) + s.capturedBytes += len(line) +} + +func (s *JSONLogStreamer) dropOldestLine() { + if len(s.lastLines) == 0 { + return + } + s.capturedBytes -= len(s.lastLines[0]) + if len(s.lastLines) > 1 { + s.capturedBytes-- + } + s.lastLines = s.lastLines[1:] } type jsonLine struct { diff --git a/pkg/log/streamer_test.go b/pkg/log/streamer_test.go index dfaf268a6..4464af2b0 100644 --- a/pkg/log/streamer_test.go +++ b/pkg/log/streamer_test.go @@ -165,6 +165,30 @@ func TestJSONLogStreamerCaptureLinesIsBounded(t *testing.T) { assert.Len(t, logs.All(), 2) } +func TestJSONLogStreamerCaptureBytesIsBounded(t *testing.T) { + const maxBytes = 16 + streamer := NewJSONLogStreamer(StreamerOptions{ + CaptureLines: 25, + CaptureBytes: maxBytes, + }) + _, err := streamer.Write([]byte("first line\nsecond line\n")) + require.NoError(t, err) + require.NoError(t, streamer.Close()) + + assert.LessOrEqual(t, len(streamer.ErrorOutput()), maxBytes) + assert.Equal(t, "second line", streamer.ErrorOutput()) +} + +func TestJSONLogStreamerCaptureBytesTruncatesOversizedLine(t *testing.T) { + const maxBytes = 8 + streamer := NewJSONLogStreamer(StreamerOptions{CaptureBytes: maxBytes}) + _, err := streamer.Write([]byte("prefix-0123456789\n")) + require.NoError(t, err) + require.NoError(t, streamer.Close()) + + assert.Equal(t, "23456789", streamer.ErrorOutput()) +} + func TestJSONLogStreamerFormattedOutputHasNoNestedEnvelope(t *testing.T) { var output bytes.Buffer previous := sugar.Load() diff --git a/pkg/ssh/connection_counter.go b/pkg/ssh/connection_counter.go index dbe3a7896..44a9947e3 100644 --- a/pkg/ssh/connection_counter.go +++ b/pkg/ssh/connection_counter.go @@ -38,16 +38,17 @@ type connectionCounter struct { timer *time.Timer timerToken *connectionTimerToken closed bool + timingOut bool } type connectionTimerToken struct{} -func (c *connectionCounter) Add() { +func (c *connectionCounter) Add() bool { c.m.Lock() defer c.m.Unlock() - if c.closed { - return + if c.closed || c.timingOut { + return false } if c.timer != nil { c.timer.Stop() @@ -56,6 +57,7 @@ func (c *connectionCounter) Add() { c.timerToken = nil c.connections++ log.Debugf("New connection on %s (Total: %d)", c.address, c.connections) + return true } func (c *connectionCounter) Dec() { @@ -120,6 +122,7 @@ func (c *connectionCounter) handleTimeout(token *connectionTimerToken) { c.m.Unlock() return } + c.timingOut = true c.timer = nil c.timerToken = nil onTimeout := c.onTimeout diff --git a/pkg/ssh/connection_counter_test.go b/pkg/ssh/connection_counter_test.go index 97fcab255..93063aafd 100644 --- a/pkg/ssh/connection_counter_test.go +++ b/pkg/ssh/connection_counter_test.go @@ -51,6 +51,25 @@ func TestConnectionCounter_InitialZeroConnectionsStartsTimeout(t *testing.T) { time.Second, time.Millisecond, "an initially idle counter should time out") } +func TestConnectionCounter_TimeoutDispatchRejectsNewConnection(t *testing.T) { + started := make(chan struct{}) + release := make(chan struct{}) + c := newConnectionCounter(context.Background(), 10*time.Millisecond, func() { + close(started) + <-release + }, "test") + + select { + case <-started: + case <-time.After(time.Second): + t.Fatal("timed out waiting for timeout dispatch") + } + assert.False(t, c.Add(), "connections must be rejected after timeout dispatch starts") + assert.Zero(t, c.connections) + close(release) + c.Close() +} + func TestConnectionCounter_CloseStopsPendingTimeout(t *testing.T) { c, calls := newRecordingCounter(t, time.Second) c.Close() @@ -86,13 +105,14 @@ func TestConnectionCounter_SpuriousDecClampsAtZero(t *testing.T) { c.Dec() c.Dec() // spurious Dec with no matching Add + // A subsequent Add must be tracked against a zero baseline, not a negative one. + assert.True(t, c.Add()) + assert.Equal(t, 1, c.connections, "Add after a spurious Dec must account correctly") + c.Dec() + require.Eventually(t, func() bool { return calls.Load() == 1 }, time.Second, time.Millisecond, "onTimeout should fire exactly once") assert.Equal(t, 0, c.connections, "count must never go negative after a spurious Dec") - - // A subsequent Add must be tracked against a zero baseline, not a negative one. - c.Add() - assert.Equal(t, 1, c.connections, "Add after a spurious Dec must account correctly") } func TestConnectionCounter_CancelledContextDoesNotFireTimeout(t *testing.T) { diff --git a/pkg/ssh/forward.go b/pkg/ssh/forward.go index 3970a8076..65efb0516 100644 --- a/pkg/ssh/forward.go +++ b/pkg/ssh/forward.go @@ -166,7 +166,10 @@ func portForwarding( } // tell the counter there is a connection - counter.Add() + if !counter.Add() { + _ = connection.Close() + continue + } // forward connection go func() { From 83be05777266291f00e76e3e12c804faa89e6c2c Mon Sep 17 00:00:00 2001 From: Samuel K Date: Fri, 11 Sep 2026 09:23:02 -0600 Subject: [PATCH 5/5] fix(ssh): stop timeout timer rearming during dispatch --- pkg/ssh/connection_counter.go | 2 +- pkg/ssh/connection_counter_test.go | 6 +++--- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/pkg/ssh/connection_counter.go b/pkg/ssh/connection_counter.go index 44a9947e3..3caf3f6a4 100644 --- a/pkg/ssh/connection_counter.go +++ b/pkg/ssh/connection_counter.go @@ -103,7 +103,7 @@ func (c *connectionCounter) armTimeoutLocked() { } func (c *connectionCounter) canArmTimeoutLocked() bool { - if c.closed || c.connections != 0 { + if c.closed || c.timingOut || c.connections != 0 { return false } if c.timeout <= 0 { diff --git a/pkg/ssh/connection_counter_test.go b/pkg/ssh/connection_counter_test.go index 93063aafd..9f3f22df7 100644 --- a/pkg/ssh/connection_counter_test.go +++ b/pkg/ssh/connection_counter_test.go @@ -71,11 +71,11 @@ func TestConnectionCounter_TimeoutDispatchRejectsNewConnection(t *testing.T) { } func TestConnectionCounter_CloseStopsPendingTimeout(t *testing.T) { - c, calls := newRecordingCounter(t, time.Second) + c, calls := newRecordingCounter(t, 10*time.Millisecond) c.Close() - time.Sleep(30 * time.Millisecond) - assert.Zero(t, calls.Load()) + require.Never(t, func() bool { return calls.Load() > 0 }, + 50*time.Millisecond, time.Millisecond) } func TestConnectionCounter_NewConnectionBeforeTimeoutCancelsIt(t *testing.T) {