Add SSH Port Forwarding (#677)

Co-authored-by: giulio93 <pilotto.giulio@gmail.com>
Co-authored-by: Luca Ronca <luca.ronca07@gmail.com>
Co-authored-by: lucarin91 <lucarin@protonmail.com>

Author: 3 people

pkg/board/remote/adb/adb.go

@@ -8,8 +8,6 @@ import (
 	"io"
 	"io/fs"
 	"log/slog"
-	"math/rand/v2"
-	"net"
 	"os"
 	"os/user"
 	"path"
@@ -20,6 +18,7 @@ import (
 	"github.com/arduino/go-paths-helper"
 
 	"github.com/bcmi-labs/orchestrator/pkg/board/remote"
+	"github.com/bcmi-labs/orchestrator/pkg/board/remote/common"
 )
 
 const username = "arduino"
@@ -58,7 +57,7 @@ func FromHost(host string, adbPath string) (*ADBConnection, error) {
 }
 
 func (a *ADBConnection) Forward(ctx context.Context, remotePort int) (int, error) {
-	hostAvailablePort, err := getAvailablePort()
+	hostAvailablePort, err := common.GetAvailablePort()
 	if err != nil {
 		return 0, fmt.Errorf("failed to find an available port: %w", err)
 	}
@@ -311,35 +310,3 @@ func FindAdbPath() string {
 
 	return adbPath
 }
-
-func isPortAvailable(port int) bool {
-	listener, err := net.Listen("tcp", fmt.Sprintf("127.0.0.1:%d", port))
-	if err != nil {
-		return false
-	}
-	listener.Close()
-	return true
-}
-
-func getRandomPort() int {
-	port := 1000 + rand.IntN(9000) // nolint:gosec
-	return port
-}
-
-const forwardPortAttempts = 10
-
-func getAvailablePort() (int, error) {
-	tried := make(map[int]any, forwardPortAttempts)
-	for len(tried) < forwardPortAttempts {
-		port := getRandomPort()
-		if _, seen := tried[port]; seen {
-			continue
-		}
-		tried[port] = struct{}{}
-
-		if isPortAvailable(port) {
-			return port, nil
-		}
-	}
-	return 0, fmt.Errorf("no available port found in range 1000-9999 after %d attempts", forwardPortAttempts)
-}

pkg/board/remote/common/common.go

@@ -0,0 +1,39 @@
+package common
+
+import (
+	"fmt"
+	"math/rand/v2"
+	"net"
+)
+
+func isPortAvailable(port int) bool {
+	listener, err := net.Listen("tcp", fmt.Sprintf("127.0.0.1:%d", port))
+	if err != nil {
+		return false
+	}
+	listener.Close()
+	return true
+}
+
+func getRandomPort() int {
+	port := 1000 + rand.IntN(9000) // nolint:gosec
+	return port
+}
+
+const forwardPortAttempts = 10
+
+func GetAvailablePort() (int, error) {
+	tried := make(map[int]any, forwardPortAttempts)
+	for len(tried) < forwardPortAttempts {
+		port := getRandomPort()
+		if _, seen := tried[port]; seen {
+			continue
+		}
+		tried[port] = struct{}{}
+
+		if isPortAvailable(port) {
+			return port, nil
+		}
+	}
+	return 0, fmt.Errorf("no available port found in range 1000-9999 after %d attempts", forwardPortAttempts)
+}

pkg/board/remote/remote_test.go

@@ -1,13 +1,19 @@
 package remote_test
 
 import (
+	"context"
+	"fmt"
+
 	"io"
+	"os/exec"
+	"strconv"
 	"strings"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
+	"github.com/bcmi-labs/orchestrator/cmd/feedback"
 	"github.com/bcmi-labs/orchestrator/pkg/board/remote"
 	"github.com/bcmi-labs/orchestrator/pkg/board/remote/adb"
 	"github.com/bcmi-labs/orchestrator/pkg/board/remote/local"
@@ -162,4 +168,80 @@ func TestSSHShell(t *testing.T) {
 			})
 		}
 	}
+
+}
+
+func TestSSHForwarder(t *testing.T) {
+	name, _, sshPort := testtools.StartAdbDContainer(t)
+	t.Cleanup(func() { testtools.StopAdbDContainer(t, name) })
+
+	conn, err := ssh.FromHost("arduino", "arduino", fmt.Sprintf("%s:%s", "localhost", sshPort))
+	require.NoError(t, err)
+
+	t.Run("Forward ADB", func(t *testing.T) {
+		ctx, cancel := context.WithCancel(t.Context())
+		defer cancel()
+
+		forwardPort, err := conn.Forward(ctx, 5555)
+		if err != nil {
+			t.Errorf("Forward failed: %v", err)
+		}
+		if forwardPort <= 0 || forwardPort > 65535 {
+			t.Fatalf("invalid port: %d", forwardPort)
+		}
+		adb_forwarded_endpoint := fmt.Sprintf("localhost:%s", strconv.Itoa(forwardPort))
+
+		out, err := exec.Command("adb", "connect", adb_forwarded_endpoint).CombinedOutput()
+		require.NoError(t, err, "adb connect output: %q", out)
+
+		cmd := exec.Command("adb", "-s", adb_forwarded_endpoint, "shell", "echo", "Hello, World!")
+		out, err = cmd.CombinedOutput()
+		require.NoError(t, err, "command output: %q", out)
+		feedback.Printf("Command output:\n%s\n", string(out))
+		require.NotNil(t, string(out))
+	})
+}
+
+func TestSSHKillForwarder(t *testing.T) {
+	name, _, sshPort := testtools.StartAdbDContainer(t)
+	t.Cleanup(func() { testtools.StopAdbDContainer(t, name) })
+
+	conn, err := ssh.FromHost("arduino", "arduino", fmt.Sprintf("%s:%s", "localhost", sshPort))
+	require.NoError(t, err)
+
+	t.Run("KillAllForwards", func(t *testing.T) {
+		ctx, cancel := context.WithCancel(t.Context())
+		defer cancel()
+
+		forwardPort, err := conn.Forward(ctx, 5555)
+		if err != nil {
+			t.Errorf("Forward failed: %v", err)
+		}
+		if forwardPort <= 0 || forwardPort > 65535 {
+			t.Fatalf("invalid port: %d", forwardPort)
+		}
+		adb_forwarded_endpoint := fmt.Sprintf("localhost:%s", strconv.Itoa(forwardPort))
+
+		out, err := exec.Command("adb", "connect", adb_forwarded_endpoint).CombinedOutput()
+		require.NoError(t, err, "adb connect output: %q", out)
+
+		cmd := exec.Command("adb", "-s", adb_forwarded_endpoint, "shell", "echo", "Hello, World!")
+		out, err = cmd.CombinedOutput()
+		require.NoError(t, err, "command output: %q", out)
+		feedback.Printf("Command output:\n%s\n", string(out))
+		require.NotNil(t, string(out))
+
+		err = conn.ForwardKillAll(t.Context())
+		require.NoError(t, err)
+		out, err = exec.Command("adb", "disconnect", adb_forwarded_endpoint).CombinedOutput()
+		require.NoError(t, err, "adb disconnect output: %q", out)
+
+		out, err = exec.Command("adb", "connect", adb_forwarded_endpoint).CombinedOutput()
+		require.NoError(t, err, "adb connect output: %q", out)
+
+		cmd = exec.Command("adb", "-s", adb_forwarded_endpoint, "shell", "echo", "Hello, World!")
+		out, err = cmd.CombinedOutput()
+		require.Error(t, err, "command output: %q", out)
+		feedback.Printf("Command output:\n%s\n", string(out))
+	})
 }

pkg/board/remote/ssh/ssh.go

@@ -3,20 +3,31 @@ package ssh
 import (
 	"bytes"
 	"context"
+	"errors"
 	"fmt"
 	"io"
 	"io/fs"
 	"log"
+	"log/slog"
+	"net"
 	"path"
 	"strings"
+	"sync"
 
 	"golang.org/x/crypto/ssh"
 
 	"github.com/bcmi-labs/orchestrator/pkg/board/remote"
+	"github.com/bcmi-labs/orchestrator/pkg/board/remote/common"
 )
 
+var ErrAuthFailed = errors.New("ssh authentication failed")
+
 type SSHConnection struct {
 	client *ssh.Client
+	wg     sync.WaitGroup
+
+	mu        sync.Mutex
+	Listeners []net.Listener
 }
 
 // Ensures SSHConnection implements the RemoteConn interface at compile time.
@@ -32,7 +43,13 @@ func FromHost(user, password, address string) (*SSHConnection, error) {
 		HostKeyCallback: ssh.InsecureIgnoreHostKey(), // nolint:gosec
 	})
 	if err != nil {
-		log.Fatalf("Failed to dial: %s", err)
+		msg := err.Error()
+		if strings.Contains(msg, "unable to authenticate") ||
+			strings.Contains(msg, "no supported methods remain") ||
+			strings.Contains(msg, "permission denied") {
+			return nil, ErrAuthFailed
+		}
+		return nil, fmt.Errorf("failed to dial SSH: %w", err)
 	}
 
 	return &SSHConnection{
@@ -41,11 +58,82 @@ func FromHost(user, password, address string) (*SSHConnection, error) {
 }
 
 func (a *SSHConnection) Forward(ctx context.Context, remotePort int) (int, error) {
-	panic("`Forward` is not implemented for SSHConnection")
+
+	// Get a random available port for remote connection
+	randomPort, err := common.GetAvailablePort()
+	if err != nil {
+		log.Printf("failed to get available port: %v", err)
+		return 0, err
+	}
+	// Listen locally on remote port
+	listener, err := net.Listen("tcp", fmt.Sprintf("%s:%d", "localhost", randomPort))
+	if err != nil {
+		return 0, err
+	}
+
+	a.mu.Lock()
+	a.Listeners = append(a.Listeners, listener)
+	a.mu.Unlock()
+
+	a.wg.Add(1)
+	go func() {
+		defer listener.Close()
+		defer a.wg.Done()
+
+		for {
+			localConn, err := listener.Accept()
+			if err != nil {
+				if !errors.Is(err, net.ErrClosed) {
+					slog.Warn("failed to accept local connection:", slog.Any("error", err))
+				}
+				return
+			}
+
+			go func(localConn net.Conn, remotePort int) {
+				defer localConn.Close()
+
+				// TODO: the kill operation should forcefully terminate the connection that was already estabish
+
+				// Open remote connection through SSH
+				remoteConn, err := a.client.Dial("tcp", fmt.Sprintf("localhost:%d", remotePort))
+				if err != nil {
+					slog.Warn("failed to dial remote host:", slog.Any("error", err))
+					return
+
+				}
+				defer remoteConn.Close()
+
+				// Bidirectional copy
+				var wg sync.WaitGroup
+				wg.Go(func() { copyAndLog(remoteConn, localConn) })
+				wg.Go(func() { copyAndLog(localConn, remoteConn) })
+				wg.Wait()
+			}(localConn, remotePort)
+		}
+	}()
+
+	return randomPort, nil
+
+}
+
+func copyAndLog(dst io.Writer, src io.Reader) {
+	_, err := io.Copy(dst, src)
+	if err != nil {
+		slog.Warn("failed to copy connection", slog.Any("error", err))
+	}
 }
 
 func (a *SSHConnection) ForwardKillAll(ctx context.Context) error {
-	panic("`ForwardKillAll` is not implemented for SSHConnection")
+	a.mu.Lock()
+	defer a.mu.Unlock()
+	for _, listener := range a.Listeners {
+		if err := listener.Close(); err != nil {
+			return err
+		}
+	}
+	a.wg.Wait()
+	a.Listeners = make([]net.Listener, 0)
+	return nil
 }
 
 func (a *SSHConnection) List(path string) ([]remote.FileInfo, error) {