panel/agent/cmd/agent/main.go

package main

import (
	"context"
	"crypto/tls"
	"crypto/x509"
	"errors"
	"flag"
	"fmt"
	"io"
	"log/slog"
	"math/rand/v2"
	"os"
	"os/signal"
	"path/filepath"
	"runtime"
	"strings"
	"syscall"
	"time"

	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials"
	"google.golang.org/grpc/credentials/insecure"
	"google.golang.org/protobuf/types/known/timestamppb"

	"github.com/dbledeez/panel/agent/internal/dispatch"
	agentruntime "github.com/dbledeez/panel/agent/internal/runtime"
	modulepkg "github.com/dbledeez/panel/pkg/module"
	"github.com/dbledeez/panel/pkg/version"
	panelv1 "github.com/dbledeez/panel/proto/panel/v1"
)

// loadClientTLS returns gRPC credentials configured for mTLS against
// the controller's CA, or nil if cert files aren't present (caller then
// decides whether to refuse or fall back to plaintext). The returned CN
// is the agent identity the controller issued during pairing — callers
// should use it as the agent_id so identity stays cryptographically bound
// to the cert instead of drifting to whatever hostname the process sees.
func loadClientTLS(certDir, controllerAddr string) (credentials.TransportCredentials, string, error) {
	caPath := filepath.Join(certDir, "ca.crt")
	certPath := filepath.Join(certDir, "agent.crt")
	keyPath := filepath.Join(certDir, "agent.key")

	for _, p := range []string{caPath, certPath, keyPath} {
		if _, err := os.Stat(p); err != nil {
			return nil, "", fmt.Errorf("missing %s: %w", p, err)
		}
	}
	caPEM, err := os.ReadFile(caPath)
	if err != nil {
		return nil, "", err
	}
	pool := x509.NewCertPool()
	if !pool.AppendCertsFromPEM(caPEM) {
		return nil, "", fmt.Errorf("bad CA PEM")
	}
	cert, err := tls.LoadX509KeyPair(certPath, keyPath)
	if err != nil {
		return nil, "", fmt.Errorf("load cert+key: %w", err)
	}
	var cn string
	if len(cert.Certificate) > 0 {
		if leaf, perr := x509.ParseCertificate(cert.Certificate[0]); perr == nil {
			cn = leaf.Subject.CommonName
		}
	}
	// ServerName defaults to the host part of controllerAddr. If it's
	// an IP we still set ServerName so Go's TLS stack uses SNI; the
	// controller cert includes 127.0.0.1 as an IP SAN so this works.
	serverName := controllerAddr
	if i := strings.LastIndex(serverName, ":"); i > 0 {
		serverName = serverName[:i]
	}
	tlsCfg := &tls.Config{
		Certificates: []tls.Certificate{cert},
		RootCAs:      pool,
		ServerName:   serverName,
		MinVersion:   tls.VersionTLS12,
	}
	return credentials.NewTLS(tlsCfg), cn, nil
}

// agentVersion mirrors the shared build version (ldflags-injected via
// github.com/dbledeez/panel/pkg/version.Version; "dev" in ad-hoc builds).
var agentVersion = version.Version

func hostname() string {
	h, err := os.Hostname()
	if err != nil || h == "" {
		return "unknown"
	}
	return h
}

func main() {
	controllerAddr := flag.String("controller", "localhost:8443", "controller gRPC address host:port")
	agentID := flag.String("agent-id", "", "agent identifier (default: mTLS cert CN, or hostname when --insecure)")
	modulesDir := flag.String("modules-dir", "./modules", "path to modules directory")
	dataRoot := flag.String("data-root", "./data/instances", "root directory for instance data")
	metaDir := flag.String("meta-dir", "./data/instance-meta", "directory for per-instance sidecar metadata (survives agent restart)")
	backupDir := flag.String("backup-dir", "./data/backups", "directory where instance backup archives are written")
	certDir := flag.String("cert-dir", "./data/certs", "directory holding agent.crt + agent.key + ca.crt after pairing")
	heartbeatInterval := flag.Duration("heartbeat", 10*time.Second, "heartbeat interval")
	insecureDial := flag.Bool("insecure", false, "force plaintext (no TLS). Use --insecure=true only for dev/local setups.")
	pairToken := flag.String("pair-token", "", "if set, agent runs in pair mode: fetches CA + signed cert from --pair-url and exits")
	pairURL := flag.String("pair-url", "http://localhost:8080", "controller HTTP URL for pair/CA/issue endpoints (used with --pair-token)")
	logLevel := flag.String("log-level", "info", "log level: debug, info, warn, error")
	flag.Parse()

	var lvl slog.Level
	if err := lvl.UnmarshalText([]byte(*logLevel)); err != nil {
		fmt.Fprintf(os.Stderr, "invalid log level %q: %v\n", *logLevel, err)
		os.Exit(2)
	}
	// Pair mode: fetch CA + sign a client cert, then exit. Pair runs before
	// we touch the local cert so the agent-id has to come from either --agent-id
	// or hostname — there's no cert CN to read yet.
	if *pairToken != "" {
		if *agentID == "" {
			*agentID = hostname()
		}
		pairLogger := slog.New(slog.NewTextHandler(os.Stdout, &slog.HandlerOptions{Level: lvl})).
			With("agent_id", *agentID)
		if err := runPair(pairConfig{
			AgentID: *agentID,
			PairURL: *pairURL,
			Token:   *pairToken,
			CertDir: *certDir,
			Logger:  pairLogger,
		}); err != nil {
			pairLogger.Error("pair failed", "err", err)
			os.Exit(1)
		}
		return
	}

	// TLS creds: try loading cert files; fall back if --insecure.
	// The cert's CN is the canonical agent identity. If --agent-id wasn't
	// set explicitly, adopt the CN; if it was set to something different,
	// refuse — a mismatched identity just silently orphans instances in the
	// DB (the exact bug that motivated this change).
	var tlsCreds credentials.TransportCredentials
	var certCN string
	creds, cn, tlsErr := loadClientTLS(*certDir, *controllerAddr)
	switch {
	case tlsErr == nil:
		tlsCreds = creds
		certCN = cn
	case *insecureDial:
		// Plaintext path — cert CN unavailable, fall through to hostname.
	default:
		fmt.Fprintf(os.Stderr, "no TLS certs found and --insecure not set. Run pairing first:\n"+
			"  agent --pair-token <token> --pair-url http://controller:8080 --cert-dir %s\n"+
			"detail: %v\n", *certDir, tlsErr)
		os.Exit(2)
	}
	switch {
	case *agentID == "" && certCN != "":
		*agentID = certCN
	case *agentID == "":
		*agentID = hostname()
	case certCN != "" && *agentID != certCN:
		fmt.Fprintf(os.Stderr, "--agent-id=%q does not match mTLS cert CN=%q; refusing to start.\n"+
			"Re-pair the agent with the intended id or drop --agent-id to adopt the cert CN.\n",
			*agentID, certCN)
		os.Exit(2)
	}
	logger := slog.New(slog.NewTextHandler(os.Stdout, &slog.HandlerOptions{Level: lvl})).
		With("agent_id", *agentID)
	if tlsCreds != nil {
		logger.Info("mTLS enabled", "cert_dir", *certDir, "cert_cn", certCN)
	} else {
		logger.Warn("running WITHOUT TLS (--insecure) — agent traffic is plaintext")
	}

	absDataRoot, err := filepath.Abs(*dataRoot)
	if err != nil {
		logger.Error("resolve data-root", "err", err)
		os.Exit(1)
	}
	if err := os.MkdirAll(absDataRoot, 0o755); err != nil {
		logger.Error("mkdir data-root", "err", err)
		os.Exit(1)
	}
	absMetaDir, err := filepath.Abs(*metaDir)
	if err != nil {
		logger.Error("resolve meta-dir", "err", err)
		os.Exit(1)
	}
	if err := os.MkdirAll(absMetaDir, 0o755); err != nil {
		logger.Error("mkdir meta-dir", "err", err)
		os.Exit(1)
	}
	absBackupDir, err := filepath.Abs(*backupDir)
	if err != nil {
		logger.Error("resolve backup-dir", "err", err)
		os.Exit(1)
	}
	if err := os.MkdirAll(absBackupDir, 0o755); err != nil {
		logger.Error("mkdir backup-dir", "err", err)
		os.Exit(1)
	}

	registry := modulepkg.NewRegistry()
	if err := registry.LoadDir(*modulesDir); err != nil {
		logger.Warn("module registry load failed", "dir", *modulesDir, "err", err)
	}
	logger.Info("modules loaded", "count", registry.Len(), "dir", *modulesDir)
	for _, m := range registry.List() {
		logger.Info("module registered", "id", m.ID, "name", m.Name, "version", m.Version)
	}

	dockerRT, err := agentruntime.NewDocker()
	if err != nil {
		logger.Error("docker runtime init", "err", err)
		os.Exit(1)
	}
	defer dockerRT.Close()

	ctx, cancel := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
	defer cancel()

	if err := dockerRT.Ping(ctx); err != nil {
		logger.Warn("docker daemon not reachable; instance commands will fail", "err", err)
	} else {
		logger.Info("docker runtime ready")
	}

	disp := dispatch.New(logger, registry, dockerRT, absDataRoot, absMetaDir, absBackupDir)

	if err := disp.Rehydrate(ctx); err != nil {
		logger.Warn("rehydrate failed (agent will start fresh)", "err", err)
	}

	// Reconnect loop: keep dialing the controller until ctx is cancelled.
	// A session ends either because ctx was cancelled (operator signal)
	// or because the stream died (controller blip, network, etc.).
	backoff := 2 * time.Second
	const maxBackoff = 30 * time.Second
	for {
		if ctx.Err() != nil {
			break
		}
		err := runSession(ctx, logger, *controllerAddr, *agentID, *heartbeatInterval, disp, tlsCreds)
		if ctx.Err() != nil {
			break
		}
		logger.Warn("session ended, will reconnect", "err", err, "in", backoff)
		select {
		case <-ctx.Done():
		case <-time.After(backoff + jitter(backoff)):
		}
		if backoff < maxBackoff {
			backoff *= 2
			if backoff > maxBackoff {
				backoff = maxBackoff
			}
		}
	}

	logger.Info("operator signal — gracefully stopping running instances")
	disp.Shutdown(context.Background())
}

// runSession opens one gRPC connection, handshakes, (re)announces current
// instance state, then runs the recv/send loops until the stream ends.
// Returns nil if the operator signalled shutdown, or an error if the
// connection broke.
func runSession(ctx context.Context, logger *slog.Logger, addr, agentID string, hb time.Duration, disp *dispatch.Dispatcher, tlsCreds credentials.TransportCredentials) error {
	var dialOpt grpc.DialOption
	if tlsCreds != nil {
		dialOpt = grpc.WithTransportCredentials(tlsCreds)
	} else {
		dialOpt = grpc.WithTransportCredentials(insecure.NewCredentials())
	}
	// Lift the default 4 MiB per-message cap so big file uploads
	// (FsWrite carrying mod zips, scenario archives, save backups)
	// flow through end-to-end. Matches the controller-side cap.
	const grpcMaxMsgSize = 5 * 1024 * 1024 * 1024
	conn, err := grpc.NewClient(addr, dialOpt,
		grpc.WithDefaultCallOptions(
			grpc.MaxCallRecvMsgSize(grpcMaxMsgSize),
			grpc.MaxCallSendMsgSize(grpcMaxMsgSize),
		),
	)
	if err != nil {
		return fmt.Errorf("grpc client: %w", err)
	}
	defer conn.Close()

	client := panelv1.NewAgentClient(conn)

	logger.Info("connecting to controller", "addr", addr)
	stream, err := client.Connect(ctx)
	if err != nil {
		return fmt.Errorf("open stream: %w", err)
	}

	hello := &panelv1.AgentEnvelope{
		CorrelationId: "hello",
		SentAt:        timestamppb.Now(),
		Payload: &panelv1.AgentEnvelope_Hello{
			Hello: &panelv1.AgentHello{
				AgentId:           agentID,
				AgentVersion:      agentVersion,
				HostOs:            runtime.GOOS,
				HostArch:          runtime.GOARCH,
				Hostname:          hostname(),
				SupportedRuntimes: []string{"docker"},
				Modules:           summarizeModules(disp),
			},
		},
	}
	if err := stream.Send(hello); err != nil {
		return fmt.Errorf("send hello: %w", err)
	}

	resp, err := stream.Recv()
	if err != nil {
		return fmt.Errorf("recv hello: %w", err)
	}
	ch := resp.GetHello()
	if ch == nil {
		return fmt.Errorf("unexpected first message: %T", resp.Payload)
	}
	logger.Info("handshake complete", "controller_version", ch.ControllerVersion)

	// Install the current stream as the upward sender and re-announce.
	disp.SetSender(stream)
	disp.Announce(ctx)

	// Drop the sender before we return so no one sends into a dead stream.
	defer disp.SetSender(nil)

	recvDone := make(chan struct{})
	recvErr := make(chan error, 1)
	go func() {
		defer close(recvDone)
		for {
			env, err := stream.Recv()
			if err != nil {
				if errors.Is(err, io.EOF) {
					logger.Info("controller closed stream")
					recvErr <- nil
				} else if ctx.Err() == nil {
					recvErr <- err
				} else {
					recvErr <- nil
				}
				return
			}
			// Dispatch each envelope in its OWN goroutine so a slow/blocked
			// handler can't freeze the recv loop (and thus the whole agent).
			// This was the Valheim+stdio hang: an RCON write to a wedged stdin
			// pipe blocked here forever, so the subsequent stop never got read.
			// Safe to parallelize: handlers guard shared state with d.mu, and the
			// controller serializes dependent ops (create→start) via its own
			// correlation-id awaits, so it never sends a dependent op before the
			// prior reply arrives.
			go disp.Handle(ctx, env)
		}
	}()

	ticker := time.NewTicker(hb)
	defer ticker.Stop()

	for {
		select {
		case <-ctx.Done():
			_ = stream.CloseSend()
			<-recvDone
			return nil
		case <-recvDone:
			return <-recvErr
		case t := <-ticker.C:
			// Heartbeats go through the dispatcher's send queue (droppable)
			// so sendLoop stays the sole owner of stream.Send — sending here
			// directly raced the drainer goroutine on the same gRPC stream.
			// A dead stream is detected by the recv goroutine (stream.Recv
			// errors → recvDone), not by a failed heartbeat send.
			disp.SendHeartbeat(t)
		}
	}
}

// jitter returns up to ±25% of d as a random offset, so N agents don't
// synchronize their reconnects against a restarting controller.
func jitter(d time.Duration) time.Duration {
	if d <= 0 {
		return 0
	}
	half := int64(d / 4)
	return time.Duration(rand.Int64N(2*half+1) - half)
}

// summarizeModules reduces the agent's loaded manifests to the small
// ModuleSummary record the controller caches for UI picking.
func summarizeModules(disp *dispatch.Dispatcher) []*panelv1.ModuleSummary {
	mods := disp.Modules().List()
	out := make([]*panelv1.ModuleSummary, 0, len(mods))
	for _, m := range mods {
		providers := make([]*panelv1.ModuleUpdateProviderSummary, 0, len(m.UpdateProviders))
		for _, p := range m.UpdateProviders {
			providers = append(providers, &panelv1.ModuleUpdateProviderSummary{
				Id:                 p.ID,
				Kind:               p.Kind,
				RequiresSteamLogin: p.RequiresSteamLogin,
			})
		}
		out = append(out, &panelv1.ModuleSummary{
			Id:              m.ID,
			Name:            m.Name,
			Version:         m.Version,
			SupportedModes:  m.SupportedModes,
			UpdateProviders: providers,
			HasRcon:         m.RCON != nil,
			Authors:         m.Authors,
		})
	}
	return out
}