fix: ssh not proxied when workspace restarts

internal/sshproxy/proxy_connection.go

@@ -11,9 +11,12 @@ type proxyConnection struct {
 	internalPort int
 	externalPort int
 	listener     net.Listener
+
+	// closedChan is used to notify that the proxy connection is closed
+	closedChan chan<- *proxyConnection
 }
 
-func newProxyConnection(toPort int) (*proxyConnection, error) {
+func newProxyConnection(toPort int, closedChan chan *proxyConnection) (*proxyConnection, error) {
 	l, err := net.Listen("tcp", ":0")
 	if err != nil {
 		return nil, err
@@ -25,6 +28,7 @@ func newProxyConnection(toPort int) (*proxyConnection, error) {
 		internalPort: toPort,
 		externalPort: externalPort,
 		listener:     l,
+		closedChan:   closedChan,
 	}, nil
 }
 
@@ -44,6 +48,9 @@ func (c *proxyConnection) forwardConnectionToSSH(conn net.Conn) {
 		fmt.Printf("error connecting to container ssh at port %d\n", c.internalPort)
 		return
 	}
+	defer func() {
+		c.closedChan <- c
+	}()
 	defer containerConn.Close()
 
 	var wg sync.WaitGroup

internal/sshproxy/sshproxy.go

@@ -6,17 +6,24 @@ type SSHProxy struct {
 	internalPorts map[int]int
 
 	connections map[int]*proxyConnection
+
+	closedConnections chan *proxyConnection
 }
 
 func New() *SSHProxy {
-	return &SSHProxy{
-		internalPorts: map[int]int{},
-		connections:   map[int]*proxyConnection{},
+	p := &SSHProxy{
+		internalPorts:     map[int]int{},
+		connections:       map[int]*proxyConnection{},
+		closedConnections: make(chan *proxyConnection),
 	}
+
+	go p.handleClosedConnections()
+
+	return p
 }
 
 func (p *SSHProxy) NewProxyEntryTo(toPort int) error {
-	c, err := newProxyConnection(toPort)
+	c, err := newProxyConnection(toPort, p.closedConnections)
 	if err != nil {
 		return err
 	}
@@ -35,3 +42,10 @@ func (p *SSHProxy) FindExternalPort(internalPort int) int {
 	}
 	return -1
 }
+
+func (p *SSHProxy) handleClosedConnections() {
+	for c := range p.closedConnections {
+		delete(p.internalPorts, c.internalPort)
+		delete(p.connections, c.internalPort)
+	}
+}