feat(agent,ssh,pkg): add connection versioning to support for v1 and v2 tunnels

This change refactors the agent to support multiple connection protocol
versions, introducing a new configuration option ConnectionVersion to
select between them. Version 1 uses an Echo-based tunnel with HTTP-like
handlers, while version 2 leverages a Yamux-based tunnel integrated with
the client. The agent’s Listen method now delegates to version-specific
listeners, and related handlers were updated accordingly. The tunnel
abstraction was also revised to adopt generics, renaming Handler to
HandlerFunc and splitting out implementations for v1 and v2.
Additionally, reverse listener creation was modernized in the API client
with explicit NewReverseListenerV1 and NewReverseListenerV2 methods,
each with their own configuration management. These changes simplify the
agent’s architecture while allowing incremental migration to the new
tunnel protocol.

Author: henrybarreto

pkg/agent/agent.go

@@ -42,20 +42,16 @@ import (
 	"context"
 	"crypto/rsa"
 	"fmt"
-	"io"
 	"math/rand"
 	"net"
-	"net/netip"
 	"net/url"
 	"os"
 	"runtime"
 	"strings"
-	"sync"
 	"sync/atomic"
 	"time"
 
 	"github.com/Masterminds/semver"
-	dockerclient "github.com/docker/docker/client"
 	"github.com/pkg/errors"
 	"github.com/shellhub-io/shellhub/pkg/agent/pkg/keygen"
 	"github.com/shellhub-io/shellhub/pkg/agent/pkg/sysinfo"
@@ -125,6 +121,10 @@ type Config struct {
 	// MaxRetryConnectionTimeout specifies the maximum time, in seconds, that an agent will wait
 	// before attempting to reconnect to the ShellHub server. Default is 60 seconds.
 	MaxRetryConnectionTimeout int `env:"MAX_RETRY_CONNECTION_TIMEOUT,default=60" validate:"min=10,max=120"`
+
+	// ConnectionVersion specifies the version of the connection protocol to use.
+	// Supported values are 1 and 2. Default is 1.
+	ConnectionVersion int `env:"CONNECTION_VERSION,default=1"`
 }
 
 func LoadConfigFromEnv() (*Config, map[string]interface{}, error) {
@@ -174,12 +174,11 @@ type Agent struct {
 	cli        client.Client
 	serverInfo *models.Info
 	server     *server.Server
-	tunnel     *tunnel.Tunnel
 	listening  chan bool
 	closed     atomic.Bool
 	mode       Mode
-	// conn is the current connection to the server.
-	conn net.Conn
+	// listener is the current connection to the server.
+	listener atomic.Pointer[net.Listener]
 }
 
 // NewAgent creates a new agent instance, requiring the ShellHub server's address to connect to, the namespace's tenant
@@ -385,237 +384,109 @@ func (a *Agent) isClosed() bool {
 func (a *Agent) Close() error {
 	a.closed.Store(true)
 
-	return a.conn.Close()
-}
-
-// httpProxyHandler handlers proxy connections to the required address.
-func httpProxyHandler(agent *Agent) tunnel.Handler {
-	const ProxyHandlerNetwork = "tcp"
+	l := a.listener.Load()
+	if l == nil {
+		return nil
+	}
 
-	return func(ctx tunnel.Context, rwc io.ReadWriteCloser) error {
-		headers, err := ctx.Headers()
-		if err != nil {
-			log.WithError(err).Error("failed to get the headers from the connection")
+	return (*l).Close()
+}
 
-			return err
-		}
+func (a *Agent) Listen(ctx context.Context) error {
+	a.mode.Serve(a)
 
-		id := headers["id"]
-		host := headers["host"]
-		port := headers["port"]
+	switch a.config.ConnectionVersion {
+	case 1:
+		return a.listenV1(ctx)
+	case 2:
+		return a.listenV2(ctx)
+	default:
+		return fmt.Errorf("unsupported connection version: %d", a.config.ConnectionVersion)
+	}
+}
 
-		logger := log.WithFields(log.Fields{
-			"id":   id,
-			"host": host,
-			"port": port,
-		})
+func (a *Agent) listenV1(ctx context.Context) error {
+	tun := tunnel.NewTunnelV1()
 
-		if _, ok := agent.mode.(*ConnectorMode); ok {
-			cli, err := dockerclient.NewClientWithOpts(dockerclient.FromEnv, dockerclient.WithAPIVersionNegotiation())
-			if err != nil {
-				log.WithError(err).Error("failed to create the Docker client")
+	tun.Handle(HandleSSHOpenV1, sshHandlerV1(a))
+	tun.Handle(HandleSSHCloseV1, sshCloseHandlerV1(a))
+	tun.Handle(HandleHTTPProxyV1, httpProxyHandlerV1(a))
 
-				return ctx.Error(errors.New("failed to connect to the Docker Engine"))
-			}
-
-			container, err := cli.ContainerInspect(context.Background(), agent.server.ContainerID)
-			if err != nil {
-				log.WithError(err).Error("failed to inspect the container")
+	go a.ping(ctx, AgentPingDefaultInterval) //nolint:errcheck
 
-				return ctx.Error(errors.New("failed to inspect the container"))
-			}
+	logger := log.WithFields(log.Fields{
+		"version":            AgentVersion,
+		"tenant_id":          a.authData.Namespace,
+		"server_address":     a.config.ServerAddress,
+		"ssh_endpoint":       a.serverInfo.Endpoints.SSH,
+		"api_endpoint":       a.serverInfo.Endpoints.API,
+		"connection_version": a.config.ConnectionVersion,
+		"sshid":              fmt.Sprintf("%s.%s@%s", a.authData.Namespace, a.authData.Name, strings.Split(a.serverInfo.Endpoints.SSH, ":")[0]),
+	})
 
-			var target string
+	ctx, cancel := context.WithCancel(ctx)
+	go func() {
+		for {
+			if a.isClosed() {
+				logger.Info("Stopped listening for connections")
 
-			addr, err := netip.ParseAddr(host)
-			if err != nil {
-				log.WithError(err).Error("failed to parse the address on proxy")
+				cancel()
 
-				return ctx.Error(errors.New("failed to parse the address on proxy"))
+				return
 			}
 
-			if addr.IsLoopback() {
-				log.Trace("host is a loopback address, using the container IP address")
-
-				for _, network := range container.NetworkSettings.Networks {
-					target = network.IPAddress
-
-					break
-				}
-			} else {
-				for _, network := range container.NetworkSettings.Networks {
-					subnet, err := netip.ParsePrefix(fmt.Sprintf("%s/%d", network.Gateway, network.IPPrefixLen))
-					if err != nil {
-						logger.WithError(err).Error("failed to parse the gateway on proxy")
-
-						continue
-					}
-
-					ip, err := netip.ParseAddr(host)
-					if err != nil {
-						logger.WithError(err).Error("failed to parse the address on proxy")
+			ShellHubConnectV1Path := "/ssh/connection"
 
-						continue
-					}
+			logger.Debug("Using tunnel version 1")
 
-					if subnet.Contains(ip) {
-						target = ip.String()
+			listener, err := a.cli.NewReverseListenerV1(
+				ctx,
+				a.authData.Token,
+				ShellHubConnectV1Path,
+			)
+			if err != nil {
+				logger.Error("Failed to connect to server through reverse tunnel. Retry in 10 seconds")
 
-						break
-					}
-				}
-			}
+				time.Sleep(time.Second * 10)
 
-			if target == "" {
-				return ctx.Error(errors.New("address not found on the device"))
+				continue
 			}
+			a.listener.Store(&listener)
 
-			host = target
-		}
-
-		ErrFailedDialToAddressAndPort := errors.New("failed to dial to the address and port")
-
-		logger.Trace("proxy handler connecting to the address")
-
-		in, err := net.Dial(ProxyHandlerNetwork, net.JoinHostPort(host, port))
-		if err != nil {
-			logger.WithError(err).Error("proxy handler failed to dial to the address")
-
-			return ctx.Error(ErrFailedDialToAddressAndPort)
-		}
-
-		defer in.Close()
-
-		logger.Trace("proxy handler dialed to the address")
-
-		// TODO: Add consts for status values.
-		if err := ctx.Status("ok"); err != nil {
-			logger.WithError(err).Error("proxy handler failed to send status response")
-
-			return err
-		}
-
-		wg := new(sync.WaitGroup)
-		done := sync.OnceFunc(func() {
-			defer in.Close()
-			defer rwc.Close()
-
-			logger.Trace("close called on in and out connections")
-		})
-
-		wg.Add(1)
-		go func() {
-			defer done()
-			defer wg.Done()
-
-			if _, err := io.Copy(in, rwc); err != nil && err != io.EOF {
-				logger.WithError(err).Error("proxy handler copy from rwc to in failed")
-			}
-		}()
+			logger.Info("Server connection established")
 
-		wg.Add(1)
-		go func() {
-			defer done()
-			defer wg.Done()
+			a.listening <- true
 
-			if _, err := io.Copy(rwc, in); err != nil && err != io.EOF {
-				logger.WithError(err).Error("proxy handler copy from in to rwc failed")
+			if err := tun.Listen(ctx, listener); err != nil {
+				logger.WithError(err).Error("Tunnel listener exited with error")
 			}
-		}()
-
-		logger.WithError(err).Info("proxy handler waiting for data pipe")
-
-		wg.Wait()
-
-		logger.WithError(err).Info("proxy handler done")
-
-		return nil
-	}
-}
-
-func sshHandler(agent *Agent) tunnel.Handler {
-	return func(ctx tunnel.Context, rwc io.ReadWriteCloser) error {
-		defer rwc.Close()
-
-		headers, err := ctx.Headers()
-		if err != nil {
-			log.WithError(err).Error("failed to get the headers from the connection")
-
-			return err
-		}
-
-		id := headers["id"]
-
-		conn, ok := rwc.(net.Conn)
-		if !ok {
-			log.Error("failed to cast the ReadWriteCloser to net.Conn")
-
-			return errors.New("failed to cast the ReadWriteCloser to net.Conn")
-		}
 
-		agent.server.Sessions.Store(id, conn)
-		agent.server.HandleConn(conn)
-
-		return nil
-	}
-}
-
-func sshCloseHandler(agent *Agent) tunnel.Handler {
-	return func(ctx tunnel.Context, rwc io.ReadWriteCloser) error {
-		defer rwc.Close()
-
-		headers, err := ctx.Headers()
-		if err != nil {
-			log.WithError(err).Error("failed to get the headers from the connection")
-
-			return err
+			a.listening <- false
 		}
+	}()
 
-		id := headers["id"]
-
-		agent.server.CloseSession(id)
-
-		log.WithFields(
-			log.Fields{
-				"id":             id,
-				"version":        AgentVersion,
-				"tenant_id":      agent.authData.Namespace,
-				"server_address": agent.config.ServerAddress,
-			},
-		).Info("A tunnel connection was closed")
+	<-ctx.Done()
 
-		return nil
-	}
+	return a.Close()
 }
 
-const (
-	// HandleSSHOpen is the protocol used to open a new SSH connection.
-	HandleSSHOpen = "/ssh/open/1.0.0"
-	// HandleSSHClose is the protocol used to close an existing SSH connection.
-	HandleSSHClose = "/ssh/close/1.0.0"
-	// HandleHTTPProxy is the protocol used to open a new HTTP proxy connection.
-	HandleHTTPProxy = "/http/proxy/1.0.0"
-)
-
-// Listen creates the SSH server and listening for connections.
-func (a *Agent) Listen(ctx context.Context) error {
-	a.mode.Serve(a)
-
-	a.tunnel = tunnel.NewTunnel()
+func (a *Agent) listenV2(ctx context.Context) error {
+	tun := tunnel.NewTunnelV2(a.cli)
 
-	a.tunnel.Handle(HandleSSHOpen, sshHandler(a))
-	a.tunnel.Handle(HandleSSHClose, sshCloseHandler(a))
-	a.tunnel.Handle(HandleHTTPProxy, httpProxyHandler(a))
+	tun.Handle(HandleSSHOpenV2, sshHandlerV2(a))
+	tun.Handle(HandleSSHCloseV2, sshCloseHandlerV2(a))
+	tun.Handle(HandleHTTPProxyV2, httpProxyHandlerV2(a))
 
 	go a.ping(ctx, AgentPingDefaultInterval) //nolint:errcheck
 
 	logger := log.WithFields(log.Fields{
-		"version":        AgentVersion,
-		"tenant_id":      a.authData.Namespace,
-		"server_address": a.config.ServerAddress,
-		"ssh_endpoint":   a.serverInfo.Endpoints.SSH,
-		"api_endpoint":   a.serverInfo.Endpoints.API,
-		"sshid":          fmt.Sprintf("%s.%s@%s", a.authData.Namespace, a.authData.Name, strings.Split(a.serverInfo.Endpoints.SSH, ":")[0]),
+		"version":            AgentVersion,
+		"tenant_id":          a.authData.Namespace,
+		"server_address":     a.config.ServerAddress,
+		"ssh_endpoint":       a.serverInfo.Endpoints.SSH,
+		"api_endpoint":       a.serverInfo.Endpoints.API,
+		"connection_version": a.config.ConnectionVersion,
+		"sshid":              fmt.Sprintf("%s.%s@%s", a.authData.Namespace, a.authData.Name, strings.Split(a.serverInfo.Endpoints.SSH, ":")[0]),
 	})
 
 	ctx, cancel := context.WithCancel(ctx)
@@ -629,24 +500,30 @@ func (a *Agent) Listen(ctx context.Context) error {
 				return
 			}
 
-			DefaultShellHubConnectPath := "/connection"
+			ShellHubConnectV2Path := "/connection"
+
+			logger.Debug("Using tunnel version 2")
 
-			conn, err := a.cli.Connect(ctx, a.authData.Token, DefaultShellHubConnectPath)
+			listener, err := a.cli.NewReverseListenerV2(
+				ctx,
+				a.authData.Token,
+				ShellHubConnectV2Path,
+				client.NewReverseV2ConfigFromMap(a.authData.Config),
+			)
 			if err != nil {
 				logger.Error("Failed to connect to server through reverse tunnel. Retry in 10 seconds")
 
 				time.Sleep(time.Second * 10)
 
 				continue
 			}
-
-			a.conn = conn
+			a.listener.Store(&listener)
 
 			logger.Info("Server connection established")
 
 			a.listening <- true
 
-			if err := a.tunnel.Listen(conn, tunnel.NewConfigFromMap(a.authData.Config)); err != nil {
+			if err := tun.Listen(ctx, listener); err != nil {
 				logger.WithError(err).Error("Tunnel listener exited with error")
 			}