mybeacon-daemon/cmd/beacon-daemon/ssh_tunnel.go

package main

import (
	"fmt"
	"log"
	"os"
	"os/exec"
	"sync"
	"time"
)

// TunnelConfig contains configuration for a specific tunnel
type TunnelConfig struct {
	Enabled           bool
	Server            string
	Port              int
	User              string
	RemotePort        int
	KeepaliveInterval int
	ReconnectDelay    int
}

// SSHTunnel manages reverse SSH tunnel to server
type SSHTunnel struct {
	name      string // "ssh" or "dashboard"
	localPort int    // Local port to forward (22 for SSH, 8080 for dashboard)

	cfg       *TunnelConfig
	client    *APIClient
	cmd       *exec.Cmd
	stopChan  chan struct{}
	mu        sync.Mutex
}

// NewSSHTunnel creates a new SSH tunnel manager
func NewSSHTunnel(name string, localPort int, cfg *TunnelConfig, client *APIClient) *SSHTunnel {
	return &SSHTunnel{
		name:      name,
		localPort: localPort,
		cfg:       cfg,
		client:    client,
		stopChan:  make(chan struct{}),
	}
}

// Start initiates the SSH tunnel
func (t *SSHTunnel) Start() error {
	if !t.cfg.Enabled {
		log.Printf("[%s-tunnel] Tunnel disabled", t.name)
		return nil
	}

	if t.cfg.RemotePort == 0 {
		log.Printf("[%s-tunnel] Remote port not allocated yet, waiting...", t.name)
		return nil
	}

	keyPath := "/opt/mybeacon/etc/ssh_tunnel_ed25519"

	// Verify key exists
	if _, err := os.Stat(keyPath); os.IsNotExist(err) {
		return fmt.Errorf("SSH key not found: %s", keyPath)
	}

	// Build reverse tunnel string: remote_port:localhost:local_port
	reverseSpec := fmt.Sprintf("%d:localhost:%d", t.cfg.RemotePort, t.localPort)

	args := []string{
		"-N", // No command execution
		"-R", reverseSpec, // Reverse tunnel with fixed port
		"-o", fmt.Sprintf("ServerAliveInterval=%d", t.cfg.KeepaliveInterval),
		"-o", "ServerAliveCountMax=3",
		"-o", "ExitOnForwardFailure=yes",
		"-o", "StrictHostKeyChecking=accept-new",
		"-i", keyPath,
		"-p", fmt.Sprintf("%d", t.cfg.Port),
		fmt.Sprintf("%s@%s", t.cfg.User, t.cfg.Server),
	}

	t.cmd = exec.Command("ssh", args...)

	if err := t.cmd.Start(); err != nil {
		return err
	}

	log.Printf("[%s-tunnel] Started: %s:%d -> localhost:%d (remote_port=%d)",
		t.name, t.cfg.Server, t.cfg.Port, t.localPort, t.cfg.RemotePort)

	// Report tunnel connection to server
	go func() {
		// Wait a bit for SSH to establish connection
		time.Sleep(2 * time.Second)
		if err := t.client.ReportTunnelPort(t.name, t.cfg.RemotePort, "connected"); err != nil {
			log.Printf("[%s-tunnel] Failed to report connection to server: %v", t.name, err)
		} else {
			log.Printf("[%s-tunnel] Reported connection to server (port %d)", t.name, t.cfg.RemotePort)
		}
	}()

	// Monitor process
	go t.monitor()

	return nil
}

// monitor watches SSH process and handles reconnection
func (t *SSHTunnel) monitor() {
	err := t.cmd.Wait()

	if err != nil {
		log.Printf("[%s-tunnel] Tunnel exited with error: %v", t.name, err)
	} else {
		log.Printf("[%s-tunnel] Tunnel exited", t.name)
	}

	// Auto-reconnect after delay
	select {
	case <-t.stopChan:
		log.Printf("[%s-tunnel] Stopped, not reconnecting", t.name)
		return
	case <-time.After(time.Duration(t.cfg.ReconnectDelay) * time.Second):
		log.Printf("[%s-tunnel] Reconnecting in %ds...", t.name, t.cfg.ReconnectDelay)
		if err := t.Start(); err != nil {
			log.Printf("[%s-tunnel] Reconnect failed: %v", t.name, err)
			// Will retry after another delay via monitor()
		}
	}
}

// Stop terminates the SSH tunnel
func (t *SSHTunnel) Stop() {
	log.Printf("[%s-tunnel] Stopping...", t.name)

	t.mu.Lock()
	defer t.mu.Unlock()

	select {
	case <-t.stopChan:
		// Already closed
		return
	default:
		close(t.stopChan)
	}

	if t.cmd != nil && t.cmd.Process != nil {
		t.cmd.Process.Kill()

		// Report disconnection to server
		if t.cfg.RemotePort != 0 {
			if err := t.client.ReportTunnelPort(t.name, 0, "disconnected"); err != nil {
				log.Printf("[%s-tunnel] Failed to report disconnection to server: %v", t.name, err)
			}
		}
	}
}

// Restart restarts the tunnel (useful when config changes)
func (t *SSHTunnel) Restart() error {
	t.Stop()
	time.Sleep(2 * time.Second)

	// Recreate stop channel
	t.mu.Lock()
	t.stopChan = make(chan struct{})
	t.mu.Unlock()

	return t.Start()
}

// UpdateConfig updates the tunnel configuration
func (t *SSHTunnel) UpdateConfig(cfg *TunnelConfig) {
	t.mu.Lock()
	defer t.mu.Unlock()
	t.cfg = cfg
}