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

package server

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

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

// LabelManager manages Traefik dynamic configuration by writing YAML config
// files to the Traefik host via SSH. This uses the Traefik file provider
// (--providers.file.directory + --providers.file.watch=true) and avoids
// docker service update which restarts the server container.
type LabelManager struct {
	mu           sync.Mutex
	remoteHost   string // Traefik host, e.g. "ingress.nixc.us:65522"
	remoteUser   string // SSH user
	signer       ssh.Signer
	serviceName  string // Swarm service name (used for backend URL)
	entrypoint   string // e.g. "websecure"
	certResolver string // e.g. "letsencryptresolver"
	configDir    string // remote dir for Traefik dynamic configs
	configs      map[string]bool
}

// NewLabelManager creates a manager that writes Traefik file-provider configs
// to the remote host via SSH.
func NewLabelManager(
	remoteHost, remoteUser string,
	signer ssh.Signer,
	serviceName, entrypoint, certResolver string,
) (*LabelManager, error) {
	configDir := "/root/traefik/dynamic"

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

	// Ensure the config directory exists.
	cmd := fmt.Sprintf("mkdir -p %s", configDir)
	if err := lm.runRemote(cmd); err != nil {
		log.Printf("WARN: could not ensure config dir %s: %v", configDir, err)
	}

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

	return lm, nil
}

// Add writes a Traefik dynamic config YAML to the remote host for a tunnel.
// If authUser and authPass are non-empty, a basicauth middleware is included.
func (lm *LabelManager) Add(tunKey, domain string, port int, authUser, authPass string) error {
	lm.mu.Lock()
	defer lm.mu.Unlock()

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

	// Build optional middleware section.
	var middlewareYAML string
	var routerMiddleware string

	if authUser != "" && authPass != "" {
		middlewareName := fmt.Sprintf("tunnel-%s-auth", tunKey)
		htpasswd, err := generateHTPasswd(authUser, authPass)
		if err != nil {
			return fmt.Errorf("generate htpasswd for %s: %w", domain, err)
		}
		middlewareYAML = fmt.Sprintf(
			"  middlewares:\n    %s:\n      basicAuth:\n        users:\n          - %q\n",
			middlewareName, htpasswd)
		routerMiddleware = fmt.Sprintf(
			"\n      middlewares:\n        - %s", middlewareName)
		log.Printf("BasicAuth middleware %s configured for %s", middlewareName, domain)
	}

	yaml := fmt.Sprintf(
		"http:\n"+
			"  routers:\n"+
			"    %s:\n"+
			"      rule: \"Host(`%s`)\"\n"+
			"      entryPoints:\n"+
			"        - %s\n"+
			"      tls:\n"+
			"        certResolver: %s\n"+
			"      service: %s%s\n"+
			"  services:\n"+
			"    %s:\n"+
			"      loadBalancer:\n"+
			"        servers:\n"+
			"          - url: \"http://%s:%d\"\n"+
			"%s",
		routerName,
		domain,
		lm.entrypoint,
		lm.certResolver,
		svcName,
		routerMiddleware,
		svcName,
		lm.serviceName,
		port,
		middlewareYAML,
	)

	filePath := fmt.Sprintf("%s/tunnel-%s.yml", lm.configDir, tunKey)

	// Write via heredoc so we don't need to escape anything.
	cmd := fmt.Sprintf("cat > '%s' << 'TRAEFIKEOF'\n%sTRAEFIKEOF", filePath, yaml)

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

	lm.configs[tunKey] = true
	log.Printf("Wrote Traefik config: %s -> %s:%d (%s)", domain, lm.serviceName, port, filePath)
	return nil
}

// Remove deletes the Traefik dynamic config file for a tunnel.
func (lm *LabelManager) Remove(tunKey string) error {
	lm.mu.Lock()
	defer lm.mu.Unlock()

	if !lm.configs[tunKey] {
		return nil
	}

	filePath := fmt.Sprintf("%s/tunnel-%s.yml", lm.configDir, tunKey)
	cmd := fmt.Sprintf("rm -f '%s'", filePath)

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

	delete(lm.configs, tunKey)
	log.Printf("Removed Traefik config: %s", filePath)
	return nil
}

// generateHTPasswd creates a bcrypt-hashed htpasswd entry for Traefik basicauth.
func generateHTPasswd(user, pass string) (string, error) {
	hash, err := bcrypt.GenerateFromPassword([]byte(pass), bcrypt.DefaultCost)
	if err != nil {
		return "", fmt.Errorf("bcrypt hash: %w", err)
	}
	return fmt.Sprintf("%s:%s", user, string(hash)), nil
}

// runRemote executes a command on the Traefik host 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
}