refactor: refactor project structure to use golang best practices

internal/ssh/client.go

@@ -0,0 +1,55 @@
+package ssh
+
+import (
+	"fmt"
+	"log/slog"
+	"os"
+	"path/filepath"
+	"time"
+
+	"golang.org/x/crypto/ssh"
+	"golang.org/x/crypto/ssh/knownhosts"
+)
+
+type Connection struct {
+	Client *ssh.Client
+}
+
+func NewConnection(user, host string, port int, auth ssh.AuthMethod) (*Connection, error) {
+	home, err := os.UserHomeDir()
+	if err != nil {
+		return nil, fmt.Errorf("failed to get home dir: %w", err)
+	}
+	knownHostsPath := filepath.Join(home, ".ssh", "known_hosts")
+
+	hkCallback, err := knownhosts.New(knownHostsPath)
+	if err != nil {
+		slog.Warn("Could not load known_hosts, ensure you connected manually once.", "path", knownHostsPath)
+		return nil, fmt.Errorf("known_hosts error: %w", err)
+	}
+
+	cfg := &ssh.ClientConfig{
+		User:            user,
+		Auth:            []ssh.AuthMethod{auth},
+		HostKeyCallback: hkCallback,
+		Timeout:         10 * time.Second,
+	}
+
+	addr := fmt.Sprintf("%s:%d", host, port)
+	slog.Debug("Dialing SSH", "target", addr)
+
+	client, err := ssh.Dial("tcp", addr, cfg)
+	if err != nil {
+		return nil, fmt.Errorf("ssh dial failed: %w", err)
+	}
+	slog.Debug("SSH connection established", "target", addr)
+
+	return &Connection{Client: client}, nil
+}
+
+func (c *Connection) Close() error {
+	if c.Client != nil {
+		return c.Client.Close()
+	}
+	return nil
+}

internal/ssh/forwarder.go

@@ -0,0 +1,103 @@
+package ssh
+
+import (
+	"context"
+	"fmt"
+	"io"
+	"log/slog"
+	"net"
+	"sync"
+	"time"
+
+	"golang.org/x/crypto/ssh"
+)
+
+type Forwarder struct {
+	sshClient  *ssh.Client
+	localPort  string
+	remotePort string
+	remoteHost string
+}
+
+func NewForwarder(client *ssh.Client, localPort, remotePort, remoteHost string) *Forwarder {
+	return &Forwarder{
+		sshClient:  client,
+		localPort:  localPort,
+		remotePort: remotePort,
+		remoteHost: remoteHost,
+	}
+}
+
+func (f *Forwarder) Start(ctx context.Context) error {
+	localAddr := "127.0.0.1:" + f.localPort
+	remoteAddr := net.JoinHostPort(f.remoteHost, f.remotePort)
+
+	listener, err := net.Listen("tcp", localAddr)
+	if err != nil {
+		return fmt.Errorf("failed to listen on %s: %w", localAddr, err)
+	}
+
+	go func() {
+		<-ctx.Done()
+		listener.Close()
+	}()
+
+	slog.Info("Port forwarder active", "local", localAddr, "remote", remoteAddr)
+
+	for {
+		localConn, err := listener.Accept()
+		if err != nil {
+			select {
+			case <-ctx.Done():
+				return nil
+			default:
+			}
+			slog.Error("Accept failed", "error", err)
+			time.Sleep(100 * time.Millisecond)
+			continue
+		}
+
+		go f.handleConnection(ctx, localConn, remoteAddr)
+	}
+}
+
+func (f *Forwarder) handleConnection(ctx context.Context, localConn net.Conn, remoteAddr string) {
+	defer localConn.Close()
+
+	_, cancel := context.WithTimeout(ctx, 10*time.Second)
+	defer cancel()
+
+	remoteConn, err := f.sshClient.Dial("tcp", remoteAddr)
+	if err != nil {
+		slog.Error("Failed to dial remote via SSH", "target", remoteAddr, "error", err)
+		return
+	}
+	defer remoteConn.Close()
+
+	var wg sync.WaitGroup
+	wg.Add(2)
+
+	go func() {
+		defer wg.Done()
+		_, _ = io.Copy(localConn, remoteConn)
+		// localConn.SetWriteDeadline(time.Now())
+		localConn.Close()
+	}()
+
+	go func() {
+		defer wg.Done()
+		_, _ = io.Copy(remoteConn, localConn)
+		remoteConn.Close()
+	}()
+
+	done := make(chan struct{})
+	go func() {
+		wg.Wait()
+		close(done)
+	}()
+
+	select {
+	case <-done:
+	case <-ctx.Done():
+	}
+}