From efad0db67edfd550ee2ae881274d218de203696e Mon Sep 17 00:00:00 2001
From: dviejokfs <dviejo@kungfusoftware.es>
Date: Fri, 24 Jun 2022 12:37:53 +0200
Subject: [PATCH] Close connections on error

---
 cmd/server/server.go | 8 ++++++++
 pkg/tunnel/tunnel.go | 3 ++-
 2 files changed, 10 insertions(+), 1 deletion(-)

diff --git a/cmd/server/server.go b/cmd/server/server.go
index a650f2e..cec5c86 100644
--- a/cmd/server/server.go
+++ b/cmd/server/server.go
@@ -112,6 +112,8 @@ func (c *serverCmd) handleTunnelRequest(mux *vhost.TLSMuxer, conn net.Conn) erro
 	sni := msg.GetTls().GetSni()
 	s := c.sessionRegistry.find(sni)
 	if s != nil {
+		defer conn.Close()
+		defer sess.Close()
 		log.Trace().Msgf("Session already exists in the registry for %s", sni)
 		err = c.returnResponse(initialConn, messages.TunnelStatus_ALREADY_EXISTS)
 		if err != nil {
@@ -127,6 +129,8 @@ func (c *serverCmd) handleTunnelRequest(mux *vhost.TLSMuxer, conn net.Conn) erro
 		if muxListener != nil {
 			muxListener.Close()
 		}
+		defer conn.Close()
+		defer sess.Close()
 		if strings.Contains(strings.ToLower(err.Error()), "already bound") {
 			err = c.returnResponse(initialConn, messages.TunnelStatus_ALREADY_EXISTS)
 			if err != nil {
@@ -302,6 +306,10 @@ func (c *serverCmd) run() error {
 		r.GET("/tunnels", func(c1 *gin.Context) {
 			c1.JSON(200, c.sessionRegistry.sessions)
 		})
+		r.DELETE("/tunnels/:sni", func(c1 *gin.Context) {
+			c.sessionRegistry.delete(c1.Param("sni"))
+			c1.JSON(200, c.sessionRegistry.sessions)
+		})
 		log.Info().Msgf("admin server listening on %s", c.adminAddr)
 		err := r.Run(c.adminAddr)
 		if err != nil {
diff --git a/pkg/tunnel/tunnel.go b/pkg/tunnel/tunnel.go
index f4fe6bc..761131f 100644
--- a/pkg/tunnel/tunnel.go
+++ b/pkg/tunnel/tunnel.go
@@ -28,7 +28,8 @@ func (c *tunnelClient) StartTlsTunnel(sni string, remoteAddress string) error {
 		log.Trace().Msgf("Failed to connect to tunnel: %v", err)
 		return err
 	}
-	session, err := yamux.Client(conn, nil)
+	cfg := yamux.DefaultConfig()
+	session, err := yamux.Client(conn, cfg)
 	if err != nil {
 		log.Trace().Msgf("Failed to create yamux session: %v", err)
 		return err