better-argo-tunnels/internal/server/labels.go

package server

import (
	"fmt"
	"log"
	"strings"
	"sync"

	"golang.org/x/crypto/ssh"
)

// LabelManager manages Traefik routing labels on its own Swarm service
// by SSHing into the Swarm manager and running docker service update.
type LabelManager struct {
	mu           sync.Mutex
	remoteHost   string // Swarm manager, e.g. "ingress.nixc.us"
	remoteUser   string // SSH user
	signer       ssh.Signer
	serviceName  string // Swarm service name, e.g. "better-argo-tunnels_tunnel-server"
	entrypoint   string // e.g. "websecure"
	certResolver string // e.g. "letsencryptresolver"
	labels       map[string]bool // track which tunnel keys we've added
}

// NewLabelManager creates a label manager that updates Swarm service labels via SSH.
func NewLabelManager(
	remoteHost, remoteUser string,
	signer ssh.Signer,
	serviceName, entrypoint, certResolver string,
) (*LabelManager, error) {

	lm := &LabelManager{
		remoteHost:   remoteHost,
		remoteUser:   remoteUser,
		signer:       signer,
		serviceName:  serviceName,
		entrypoint:   entrypoint,
		certResolver: certResolver,
		labels:       make(map[string]bool),
	}

	// Verify we can reach the Swarm manager and the service exists.
	cmd := fmt.Sprintf("docker service inspect --format '{{.Spec.Name}}' %s", serviceName)
	if err := lm.runRemote(cmd); err != nil {
		log.Printf("WARN: could not verify service %s (may not exist yet): %v", serviceName, err)
	} else {
		log.Printf("Verified Swarm service: %s", serviceName)
	}

	log.Printf("Label manager ready (host=%s, service=%s, ep=%s, resolver=%s)",
		remoteHost, serviceName, entrypoint, certResolver)

	return lm, nil
}

// Add adds Traefik routing labels to the Swarm service for a tunnel.
func (lm *LabelManager) Add(tunKey, domain string, port int) error {
	lm.mu.Lock()
	defer lm.mu.Unlock()

	routerName := fmt.Sprintf("tunnel-%s-router", tunKey)
	serviceName := fmt.Sprintf("tunnel-%s-service", tunKey)

	// Build the label-add flags for docker service update.
	labelArgs := []string{
		labelFlag(fmt.Sprintf("traefik.http.routers.%s.rule", routerName),
			fmt.Sprintf("Host(`%s`)", domain)),
		labelFlag(fmt.Sprintf("traefik.http.routers.%s.entrypoints", routerName),
			lm.entrypoint),
		labelFlag(fmt.Sprintf("traefik.http.routers.%s.tls", routerName),
			"true"),
		labelFlag(fmt.Sprintf("traefik.http.routers.%s.tls.certresolver", routerName),
			lm.certResolver),
		labelFlag(fmt.Sprintf("traefik.http.routers.%s.service", routerName),
			serviceName),
		labelFlag(fmt.Sprintf("traefik.http.services.%s.loadbalancer.server.port", serviceName),
			fmt.Sprintf("%d", port)),
	}

	cmd := fmt.Sprintf("docker service update --label-add %s %s",
		strings.Join(labelArgs, " --label-add "), lm.serviceName)

	if err := lm.runRemote(cmd); err != nil {
		return fmt.Errorf("add labels for %s: %w", domain, err)
	}

	lm.labels[tunKey] = true
	log.Printf("Added Swarm labels: %s -> %s:%d", domain, lm.serviceName, port)
	return nil
}

// Remove removes Traefik routing labels from the Swarm service for a tunnel.
func (lm *LabelManager) Remove(tunKey string) error {
	lm.mu.Lock()
	defer lm.mu.Unlock()

	if !lm.labels[tunKey] {
		return nil // nothing to remove
	}

	routerName := fmt.Sprintf("tunnel-%s-router", tunKey)
	serviceName := fmt.Sprintf("tunnel-%s-service", tunKey)

	// Build the label-rm flags.
	rmLabels := []string{
		fmt.Sprintf("traefik.http.routers.%s.rule", routerName),
		fmt.Sprintf("traefik.http.routers.%s.entrypoints", routerName),
		fmt.Sprintf("traefik.http.routers.%s.tls", routerName),
		fmt.Sprintf("traefik.http.routers.%s.tls.certresolver", routerName),
		fmt.Sprintf("traefik.http.routers.%s.service", routerName),
		fmt.Sprintf("traefik.http.services.%s.loadbalancer.server.port", serviceName),
	}

	cmd := fmt.Sprintf("docker service update --label-rm %s %s",
		strings.Join(rmLabels, " --label-rm "), lm.serviceName)

	if err := lm.runRemote(cmd); err != nil {
		return fmt.Errorf("remove labels for %s: %w", tunKey, err)
	}

	delete(lm.labels, tunKey)
	log.Printf("Removed Swarm labels for tunnel: %s", tunKey)
	return nil
}

// labelFlag formats a --label-add value, quoting properly for shell.
func labelFlag(key, value string) string {
	return fmt.Sprintf("'%s=%s'", key, value)
}

// runRemote executes a command on the Swarm manager via SSH.
func (lm *LabelManager) runRemote(cmd string) error {
	addr := lm.remoteHost
	if !strings.Contains(addr, ":") {
		addr = addr + ":22"
	}

	config := &ssh.ClientConfig{
		User: lm.remoteUser,
		Auth: []ssh.AuthMethod{
			ssh.PublicKeys(lm.signer),
		},
		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
	}

	client, err := ssh.Dial("tcp", addr, config)
	if err != nil {
		return fmt.Errorf("SSH dial %s: %w", addr, err)
	}
	defer client.Close()

	session, err := client.NewSession()
	if err != nil {
		return fmt.Errorf("SSH session: %w", err)
	}
	defer session.Close()

	output, err := session.CombinedOutput(cmd)
	if err != nil {
		return fmt.Errorf("remote cmd failed: %w (output: %s)", err, string(output))
	}

	return nil
}

// Close is a no-op — SSH connections are opened/closed per operation.
func (lm *LabelManager) Close() error {
	return nil
}