fix: Crash issues on waitgroups

Author: dogukanoksuz

app/handlers/command.go

@@ -54,7 +54,7 @@ func OutsideCommandRunner(c *fiber.Ctx) error {
 		}
 	}
 
-	session, err := bridge.Connections.GetRaw(
+	session, err := bridge.Sessions.GetRaw(
 		c.Locals("user_id").(string),
 		c.FormValue("remote_host"),
 		c.FormValue("username"),
@@ -70,7 +70,7 @@ func OutsideCommandRunner(c *fiber.Ctx) error {
 		)
 
 		if isCreated {
-			bridge.Connections.SetRaw(
+			bridge.Sessions.SetRaw(
 				c.Locals("user_id").(string),
 				c.FormValue("remote_host"),
 				c.FormValue("username"),
@@ -88,7 +88,7 @@ func OutsideCommandRunner(c *fiber.Ctx) error {
 
 	if c.FormValue("disconnect") == "1" {
 		session.CloseAllConnections()
-		bridge.Connections.Delete(c.Locals("user_id").(string) + c.FormValue("remote_host") + c.FormValue("username"))
+		bridge.Sessions.Delete(c.Locals("user_id").(string) + c.FormValue("remote_host") + c.FormValue("username"))
 	}
 
 	return c.SendString(output)

internal/bridge/clean.go

@@ -11,7 +11,7 @@ import (
 // Clean clears long standing sessions from memory
 func Clean() {
 	now := time.Now()
-	for key, session := range Connections {
+	for key, session := range Sessions.Connections {
 		if now.Sub(session.LastConnection).Seconds() > 266 {
 			closeSession(session, key)
 			continue
@@ -65,9 +65,9 @@ func Clean() {
 		}
 	}
 
-	for key, tunnel := range Tunnels {
+	for _, tunnel := range Tunnels.Connections {
 		if now.Sub(tunnel.LastConnection).Seconds() > 266 {
-			closeTunnel(tunnel, key)
+			closeTunnel(tunnel)
 			continue
 		}
 
@@ -81,7 +81,7 @@ func Clean() {
 			10*time.Second,
 		)
 		if err != nil {
-			closeTunnel(tunnel, key)
+			closeTunnel(tunnel)
 			continue
 		}
 
@@ -94,7 +94,7 @@ func Clean() {
 			default:
 				_, _, err := tunnel.SshClient.SendRequest("[email protected]", true, nil)
 				if err != nil {
-					closeTunnel(tunnel, key)
+					closeTunnel(tunnel)
 					logger.Sugar().Warnw("error when sending request")
 				}
 
@@ -106,7 +106,7 @@ func Clean() {
 		case <-ch:
 			continue
 		case <-time.After(10 * time.Second):
-			closeTunnel(tunnel, key)
+			closeTunnel(tunnel)
 			continue
 		}
 	}
@@ -120,13 +120,13 @@ func closeSession(s *Session, key string) {
 	s.CloseAllConnections()
 	s.Mutex.Unlock()
 
-	Connections.Delete(key)
+	Sessions.Delete(key)
 }
 
-func closeTunnel(t *Tunnel, key string) {
+func closeTunnel(t *Tunnel) {
 	t.Mutex.Lock()
 	t.Stop()
 	t.Mutex.Unlock()
 
-	Tunnels.Delete(key)
+	t.errHandler()
 }

internal/bridge/connection.go

@@ -7,18 +7,23 @@ import (
 	"time"
 )
 
-type Pool map[string]*Session
-type TunnelPool map[string]*Tunnel
+type Pool struct {
+	Connections map[string]*Session
+	sync.Mutex
+}
+type TunnelPool struct {
+	Connections map[string]*Tunnel
+	sync.Mutex
+}
 
 var (
-	Connections Pool       = make(Pool)
-	Tunnels     TunnelPool = make(TunnelPool)
-	mutex       sync.Mutex
+	Sessions Pool       = Pool{Connections: make(map[string]*Session)}
+	Tunnels  TunnelPool = TunnelPool{Connections: make(map[string]*Tunnel)}
 )
 
 // Get connection from pool
 func (p *Pool) Get(userID, serverID string) (*Session, error) {
-	if conn, ok := Connections[userID+serverID]; ok {
+	if conn, ok := p.Connections[userID+serverID]; ok {
 		return conn, nil
 	} else {
 		return nil, errors.New("connection does not exist")
@@ -27,12 +32,14 @@ func (p *Pool) Get(userID, serverID string) (*Session, error) {
 
 // Set connection on pool
 func (p *Pool) Set(userID, serverID string, session *Session) {
-	Connections[userID+serverID] = session
+	p.Lock()
+	defer p.Unlock()
+	p.Connections[userID+serverID] = session
 }
 
 // GetRaw connection object from pool
 func (p *Pool) GetRaw(userID, remoteHost, username string) (*Session, error) {
-	if conn, ok := Connections[userID+remoteHost+username]; ok {
+	if conn, ok := p.Connections[userID+remoteHost+username]; ok {
 		conn.LastConnection = time.Now()
 		return conn, nil
 	} else {
@@ -42,19 +49,21 @@ func (p *Pool) GetRaw(userID, remoteHost, username string) (*Session, error) {
 
 // SetRaw connection object on pool
 func (p *Pool) SetRaw(userID, remoteHost, username string, session *Session) {
-	Connections[userID+remoteHost+username] = session
+	p.Lock()
+	defer p.Unlock()
+	p.Connections[userID+remoteHost+username] = session
 }
 
 // Delete connection object from pool
 func (p *Pool) Delete(key string) {
-	mutex.Lock()
-	defer mutex.Unlock()
-	delete(Connections, key)
+	p.Lock()
+	defer p.Unlock()
+	delete(p.Connections, key)
 }
 
 // Get Tunnel connection from pool
 func (t *TunnelPool) Get(remoteHost, remotePort, username string) (*Tunnel, error) {
-	if tunnel, ok := Tunnels[remoteHost+":"+remotePort+":"+username]; ok {
+	if tunnel, ok := Tunnels.Connections[remoteHost+":"+remotePort+":"+username]; ok {
 		tunnel.LastConnection = time.Now()
 		return tunnel, nil
 	} else {
@@ -64,16 +73,16 @@ func (t *TunnelPool) Get(remoteHost, remotePort, username string) (*Tunnel, erro
 
 // Set Tunnel connection to pool
 func (t *TunnelPool) Set(remoteHost, remotePort, username string, tunnel *Tunnel) {
-	mutex.Lock()
-	defer mutex.Unlock()
-	Tunnels[remoteHost+":"+remotePort+":"+username] = tunnel
+	t.Lock()
+	defer t.Unlock()
+	Tunnels.Connections[remoteHost+":"+remotePort+":"+username] = tunnel
 }
 
 // Delete Tunnel connection from pool
 func (t *TunnelPool) Delete(key string) {
-	mutex.Lock()
-	defer mutex.Unlock()
-	delete(Tunnels, key)
+	t.Lock()
+	defer t.Unlock()
+	delete(Tunnels.Connections, key)
 }
 
 // VerifyAuth verifies key when connecting

internal/bridge/initialize.go

@@ -12,14 +12,14 @@ import (
 
 // GetSession gets session from memory and if does not exists creates a new one
 func GetSession(userID, serverID, host string) (*Session, error) {
-	session, err := Connections.Get(userID, serverID)
+	session, err := Sessions.Get(userID, serverID)
 	if err != nil {
 		session = &Session{}
 		isConnected := session.CreateShell(userID, serverID, host)
 		if !isConnected {
 			return nil, logger.FiberError(fiber.StatusForbidden, "cannot connect to server")
 		}
-		Connections.Set(userID, serverID, session)
+		Sessions.Set(userID, serverID, session)
 	}
 
 	session.LastConnection = time.Now()

internal/bridge/session.go

@@ -25,14 +25,15 @@ type Session struct {
 	SFTP           *sftp.Client
 	SMB            *smb2.Session
 	WinRM          *winrm.Client
-	Mutex          sync.Mutex
 	LastConnection time.Time
 	WindowsLetter  string
 	WindowsPath    string
 	Username       string
 	IpAddr         string
 	Port           string
 	password       string
+
+	sync.Mutex
 }
 
 // CloseAllConnections closes all connections on session