[tunnel] close the tunnel if at least one side is dropped

Author: akosyakov

components/local-app/pkg/bastion/bastion.go

@@ -26,7 +26,6 @@ import (
 	"github.com/kevinburke/ssh_config"
 	"github.com/prometheus/common/log"
 	"golang.org/x/crypto/ssh"
-	"golang.org/x/sync/errgroup"
 	"google.golang.org/grpc"
 	"google.golang.org/protobuf/proto"
 
@@ -449,7 +448,9 @@ func (b *Bastion) establishTunnel(ctx context.Context, ws *Workspace, logprefix
 				logrus.WithError(err).WithField("workspace", ws.WorkspaceID).Warn(logprefix + ": failed to accept connection")
 				continue
 			}
+			logrus.WithField("workspace", ws.WorkspaceID).Debug(logprefix + ": accepted new connection")
 			go func() {
+				defer logrus.WithField("workspace", ws.WorkspaceID).Debug(logprefix + ": connection closed")
 				defer conn.Close()
 
 				clientCh := make(chan *TunnelClient, 1)
@@ -476,16 +477,17 @@ func (b *Bastion) establishTunnel(ctx context.Context, ws *Workspace, logprefix
 				}
 				defer sshChan.Close()
 				go ssh.DiscardRequests(reqs)
-				eg, _ := errgroup.WithContext(listenerCtx)
-				eg.Go(func() error {
+
+				ctx, cancel := context.WithCancel(listenerCtx)
+				go func() {
 					_, _ = io.Copy(sshChan, conn)
-					return nil
-				})
-				eg.Go(func() error {
+					cancel()
+				}()
+				go func() {
 					_, _ = io.Copy(conn, sshChan)
-					return nil
-				})
-				eg.Wait()
+					cancel()
+				}()
+				<-ctx.Done()
 			}()
 		}
 	}()

components/supervisor/pkg/supervisor/supervisor.go

@@ -731,6 +731,8 @@ func tunnelOverSSH(ctx context.Context, tunneled *ports.TunneledPortsService, ne
 		newCh.Reject(ssh.Prohibited, err.Error())
 		return
 	}
+	log.Debug("tunnel: accepted new connection")
+	defer log.Debug("tunnel: connection closed")
 	defer tunnel.Close()
 
 	sshChan, reqs, err := newCh.Accept()
@@ -740,17 +742,16 @@ func tunnelOverSSH(ctx context.Context, tunneled *ports.TunneledPortsService, ne
 	}
 	defer sshChan.Close()
 	go ssh.DiscardRequests(reqs)
-	var wg sync.WaitGroup
-	wg.Add(2)
+	ctx, cancel := context.WithCancel(ctx)
 	go func() {
 		_, _ = io.Copy(sshChan, tunnel)
-		wg.Done()
+		cancel()
 	}()
 	go func() {
 		_, _ = io.Copy(tunnel, sshChan)
-		wg.Done()
+		cancel()
 	}()
-	wg.Wait()
+	<-ctx.Done()
 }
 
 func startSSHServer(ctx context.Context, cfg *Config, wg *sync.WaitGroup) {