[fix] issues related to tcp-keepalive and heartbeat

Author: Musixal

internal/client/client.go

@@ -50,8 +50,8 @@ func (c *Client) Start() {
 		tcpConfig := &transport.TcpConfig{
 			RemoteAddr:    c.config.RemoteAddr,
 			Nodelay:       c.config.Nodelay,
-			KeepAlive:     time.Duration(c.config.Keepalive),
-			RetryInterval: time.Duration(c.config.RetryInterval),
+			KeepAlive:     time.Duration(c.config.Keepalive) * time.Second,
+			RetryInterval: time.Duration(c.config.RetryInterval) * time.Second,
 			Token:         c.config.Token,
 			Forwarder:     c.forwarderReader(c.config.Forwarder),
 		}
@@ -62,8 +62,8 @@ func (c *Client) Start() {
 		tcpMuxConfig := &transport.TcpMuxConfig{
 			RemoteAddr:    c.config.RemoteAddr,
 			Nodelay:       c.config.Nodelay,
-			KeepAlive:     time.Duration(c.config.Keepalive),
-			RetryInterval: time.Duration(c.config.RetryInterval),
+			KeepAlive:     time.Duration(c.config.Keepalive) * time.Second,
+			RetryInterval: time.Duration(c.config.RetryInterval) * time.Second,
 			Token:         c.config.Token,
 			MuxSession:    c.config.MuxSession,
 
@@ -76,8 +76,8 @@ func (c *Client) Start() {
 		WsConfig := &transport.WsConfig{
 			RemoteAddr:    c.config.RemoteAddr,
 			Nodelay:       c.config.Nodelay,
-			KeepAlive:     time.Duration(c.config.Keepalive),
-			RetryInterval: time.Duration(c.config.RetryInterval),
+			KeepAlive:     time.Duration(c.config.Keepalive) * time.Second,
+			RetryInterval: time.Duration(c.config.RetryInterval) * time.Second,
 			Token:         c.config.Token,
 			Forwarder:     c.forwarderReader(c.config.Forwarder),
 		}

internal/client/transport/tcp.go

@@ -87,7 +87,7 @@ func (c *TcpTransport) ChannelDialer() {
 			tunnelTCPConn, err := c.tcpDialer(c.config.RemoteAddr, c.config.Nodelay)
 			if err != nil {
 				c.logger.Error(err)
-				time.Sleep(c.config.RetryInterval * time.Second)
+				time.Sleep(c.config.RetryInterval)
 				continue
 			}
 
@@ -111,7 +111,7 @@ func (c *TcpTransport) ChannelDialer() {
 					c.logger.Error("unable to establish new control channel")
 				}
 				tunnelTCPConn.Close() // Close connection on error or timeout
-				time.Sleep(c.config.RetryInterval* time.Second)
+				time.Sleep(c.config.RetryInterval)
 				continue
 			}
 
@@ -125,7 +125,7 @@ func (c *TcpTransport) ChannelDialer() {
 			} else {
 				c.logger.Error("invalid token received, retrying...")
 				tunnelTCPConn.Close() // Close connection if the token is invalid
-				time.Sleep(c.config.RetryInterval * time.Second)
+				time.Sleep(c.config.RetryInterval)
 				continue
 			}
 		}

internal/client/transport/tcpmux.go

@@ -44,6 +44,7 @@ func NewMuxClient(parentCtx context.Context, config *TcpMuxConfig, logger *logru
 		cancel:      cancel,
 		logger:      logger,
 		smuxSession: make([]*smux.Session, config.MuxSession),
+		timeout:     5 * time.Second, // Default timeout
 	}
 
 	return client
@@ -87,7 +88,7 @@ func (c *TcpMuxTransport) MuxDialer() {
 				tunnelTCPConn, err := c.tcpDialer(c.config.RemoteAddr, c.config.Nodelay)
 				if err != nil {
 					c.logger.Error("failed to dial tunnel server: ", err)
-					time.Sleep(c.config.RetryInterval * time.Second)
+					time.Sleep(c.config.RetryInterval)
 					continue
 				}
 

internal/client/transport/ws.go

@@ -91,7 +91,7 @@ func (c *WsTransport) ChannelDialer() {
 			tunnelWSConn, err := c.wsDialer(c.config.RemoteAddr, "/channel")
 			if err != nil {
 				c.logger.Error(err)
-				time.Sleep(c.config.RetryInterval * time.Second)
+				time.Sleep(c.config.RetryInterval)
 				continue
 			}
 			c.controlChannel = tunnelWSConn

internal/server/server.go

@@ -43,7 +43,7 @@ func (s *Server) Start() {
 		tcpConfig := &transport.TcpConfig{
 			BindAddr:       s.config.BindAddr,
 			Nodelay:        s.config.Nodelay,
-			KeepAlive:      time.Duration(s.config.Keepalive),
+			KeepAlive:      time.Duration(s.config.Keepalive) * time.Second,
 			ConnectionPool: s.config.ConnectionPool,
 			Token:          s.config.Token,
 			ChannelSize:    s.config.ChannelSize,
@@ -57,7 +57,7 @@ func (s *Server) Start() {
 		tcpMuxConfig := &transport.TcpMuxConfig{
 			BindAddr:    s.config.BindAddr,
 			Nodelay:     s.config.Nodelay,
-			KeepAlive:   time.Duration(s.config.Keepalive),
+			KeepAlive:   time.Duration(s.config.Keepalive) * time.Second,
 			Token:       s.config.Token,
 			MuxSession:  s.config.MuxSession,
 			ChannelSize: s.config.ChannelSize,
@@ -71,7 +71,7 @@ func (s *Server) Start() {
 		wsConfig := &transport.WsConfig{
 			BindAddr:       s.config.BindAddr,
 			Nodelay:        s.config.Nodelay,
-			KeepAlive:      time.Duration(s.config.Keepalive),
+			KeepAlive:      time.Duration(s.config.Keepalive) * time.Second,
 			ConnectionPool: s.config.ConnectionPool,
 			Token:          s.config.Token,
 			ChannelSize:    s.config.ChannelSize,

internal/server/transport/tcp.go

@@ -227,12 +227,16 @@ func (s *TcpTransport) channelListener() {
 }
 
 func (s *TcpTransport) heartbeat() {
+	ticker := time.NewTicker(s.heartbeatDuration)
+	defer ticker.Stop()
+
 	for {
 		select {
 		case <-s.ctx.Done():
 			return
-		default:
+		case <-ticker.C:
 			if s.controlChannel == nil {
+				go s.Restart()
 				return
 			}
 			err := utils.SendBinaryString(s.controlChannel, s.heartbeatSig)
@@ -241,12 +245,11 @@ func (s *TcpTransport) heartbeat() {
 				go s.Restart()
 				return
 			}
-			s.logger.Debug("heartbeat sended successfully")
-			time.Sleep(s.heartbeatDuration * time.Second)
+			s.logger.Debug("heartbeat sent successfully")
 		}
 	}
-
 }
+
 func (s *TcpTransport) poolChecker() {
 	for {
 		select {

internal/server/transport/ws.go

@@ -116,24 +116,25 @@ func (s *WsTransport) portConfigReader() {
 }
 
 func (s *WsTransport) heartbeat() {
+	ticker := time.NewTicker(s.heartbeatDuration)
+	defer ticker.Stop()
+
 	for {
 		select {
 		case <-s.ctx.Done():
 			return
-		default:
+		case <-ticker.C:
 			if s.controlChannel == nil {
+				go s.Restart()
 				return
 			}
-			s.mu.Lock()
 			err := s.controlChannel.WriteMessage(websocket.TextMessage, []byte(s.heartbeatSig))
-			s.mu.Unlock()
 			if err != nil {
 				s.logger.Error("unable to send heartbeat. restarting server...")
 				go s.Restart()
 				return
 			}
 			s.logger.Debug("heartbeat sent successfully")
-			time.Sleep(s.heartbeatDuration * time.Second)
 		}
 	}
 }