pangolin/server/cmd/server/main.go

package main

import (
	"context"
	"crypto/tls"
	"encoding/json"
	"flag"
	"log"
	"log/slog"
	"net"
	"net/http"
	"os"

	"github.com/go-chi/chi/v5"
	"github.com/go-chi/chi/v5/middleware"
	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials"

	"github.com/wangjia/pangolin/server/api/gen"
	"github.com/wangjia/pangolin/server/internal/admin"
	"github.com/wangjia/pangolin/server/internal/db"
	"github.com/wangjia/pangolin/server/internal/httpapi"
	"github.com/wangjia/pangolin/server/internal/mtls"
	"github.com/wangjia/pangolin/server/internal/nodes"
	agentv1 "github.com/wangjia/pangolin/server/internal/pb/agentv1"
	"github.com/wangjia/pangolin/server/internal/redisutil"
	"github.com/wangjia/pangolin/server/internal/scheduler/probe"
)

func main() {
	addr := flag.String("addr", "", "HTTP listen address (default :8080, overridden by ADDR env)")
	flag.Parse()

	if *addr == "" {
		if v := os.Getenv("ADDR"); v != "" {
			*addr = v
		} else {
			*addr = ":8080"
		}
	}

	// Admin backend (separate internal listener). Started only when explicitly
	// configured so the public API can run on its own. See internal/admin.
	startAdminIfConfigured()

	r := chi.NewRouter()
	r.Use(middleware.Logger)
	r.Use(middleware.Recoverer)

	// Health check — outside the versioned API.
	r.Get("/healthz", func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(http.StatusOK)
		_ = json.NewEncoder(w).Encode(map[string]string{"status": "ok"})
	})

	// --- Probe ingest route (optional) ---
	// Requires:
	//   PROBE_SECRETS  – JSON object mapping probeId → HMAC secret, e.g.
	//                    '{"probe-sg-01":"s3cr3t1","probe-jp-01":"s3cr3t2"}'
	//   REDIS_ADDR     – Redis address, default 127.0.0.1:6379
	//   REDIS_PASSWORD – Redis password (optional)
	//
	// If PROBE_SECRETS is empty the route is not registered and the server
	// starts normally without the probe ingest endpoint.
	probeSecretsJSON := os.Getenv("PROBE_SECRETS")
	if probeSecretsJSON != "" {
		redisAddr := getenvDefault("REDIS_ADDR", "127.0.0.1:6379")
		rdb, err := redisutil.New(redisAddr, os.Getenv("REDIS_PASSWORD"), 0)
		if err != nil {
			log.Printf("probe: redis connect failed (%v) – probe route disabled", err)
		} else {
			var secretMap map[string]string
			if err := json.Unmarshal([]byte(probeSecretsJSON), &secretMap); err != nil {
				log.Fatalf("probe: invalid PROBE_SECRETS JSON: %v", err)
			}
			reg := probe.NewMapRegistry(secretMap)
			st := probe.NewStore(rdb)
			h := probe.NewIngestHandler(reg, st)
			r.Post("/probe/report", h.ServeHTTP)
			log.Printf("probe ingest route registered (%d probe(s))", len(secretMap))
		}
	}
	// Mount all /v1/... routes.  HandlerFromMuxWithBaseURL registers every route
	// from the OpenAPI spec onto the provided chi router with the given prefix,
	// so the router itself is the http.Handler we serve.
	gen.HandlerFromMuxWithBaseURL(&httpapi.UnimplementedServer{}, r, "/v1")

	// ─── gRPC agent server (optional) ────────────────────────────────────────
	// All five env vars must be set to activate the gRPC listener.

	grpcAddr := os.Getenv("GRPC_ADDR")
	caKeyPath := os.Getenv("CA_KEY_PATH")
	caCertPath := os.Getenv("CA_CERT_PATH")
	grpcCertPath := os.Getenv("GRPC_CERT_PATH")
	grpcKeyPath := os.Getenv("GRPC_KEY_PATH")
	dbDSN := os.Getenv("DB_DSN")

	if grpcAddr != "" && caKeyPath != "" && caCertPath != "" &&
		grpcCertPath != "" && grpcKeyPath != "" && dbDSN != "" {
		redisAddr := getenvDefault("REDIS_ADDR", "127.0.0.1:6379")
		go startGRPC(grpcAddr, caKeyPath, caCertPath, grpcCertPath, grpcKeyPath,
			redisAddr, os.Getenv("REDIS_PASSWORD"), dbDSN)
	} else if grpcAddr != "" {
		slog.Warn("grpc server disabled: one or more required env vars are missing",
			"GRPC_ADDR", grpcAddr,
			"CA_KEY_PATH_set", caKeyPath != "",
			"CA_CERT_PATH_set", caCertPath != "",
			"GRPC_CERT_PATH_set", grpcCertPath != "",
			"GRPC_KEY_PATH_set", grpcKeyPath != "",
			"DB_DSN_set", dbDSN != "")
	}

	// ─── Start HTTP server (blocking) ────────────────────────────────────────

	log.Printf("pangolin server listening on %s", *addr)
	if err := http.ListenAndServe(*addr, r); err != nil {
		log.Fatalf("server error: %v", err)
	}
}

// startAdminIfConfigured launches the admin listener in a background goroutine
// when ADMIN_SECRET_KEY and DB_DSN are present. The admin port binds an
// internal address only (enforced by admin.FromEnv) and is fronted by an IP
// allowlist + two-factor login.
func startAdminIfConfigured() {
	if os.Getenv("ADMIN_SECRET_KEY") == "" || os.Getenv("DB_DSN") == "" {
		log.Printf("admin backend disabled (set DB_DSN and ADMIN_SECRET_KEY to enable)")
		return
	}
	cfg, err := admin.FromEnv()
	if err != nil {
		log.Fatalf("admin config: %v", err)
	}
	database, err := db.Open(os.Getenv("DB_DSN"))
	if err != nil {
		log.Fatalf("admin db: %v", err)
	}
	rdb, err := redisutil.New(
		getenvDefault("REDIS_ADDR", "127.0.0.1:6379"),
		os.Getenv("REDIS_PASSWORD"), 0)
	if err != nil {
		log.Fatalf("admin redis: %v", err)
	}

	svc := admin.BuildServices(database, rdb, 5, cfg.LoginLockDuration)
	handler, err := admin.NewHandler(cfg, database, rdb, svc, log.Default())
	if err != nil {
		log.Fatalf("admin handler: %v", err)
	}

	go func() {
		log.Printf("admin backend listening on %s (internal only)", cfg.Listen)
		if err := http.ListenAndServe(cfg.Listen, handler); err != nil {
			log.Fatalf("admin server error: %v", err)
		}
	}()
}

// startGRPC initialises and starts the mTLS gRPC agent server.
// Must be called in a goroutine. Calls log.Fatalf on unrecoverable errors.
func startGRPC(
	addr, caKeyPath, caCertPath, grpcCertPath, grpcKeyPath,
	redisAddr, redisPassword, dbDSN string,
) {
	// Pangolin Node CA (signs + verifies node client certs).
	ca, err := mtls.NewCA(mtls.CAConfig{
		KeyPath:  caKeyPath,
		CertPath: caCertPath,
	})
	if err != nil {
		log.Fatalf("grpc: load CA: %v", err)
	}

	// Redis (hub command queue, load cache, CRL, bootstrap tokens).
	rdb, err := redisutil.New(redisAddr, redisPassword, 0)
	if err != nil {
		log.Fatalf("grpc: redis connect: %v", err)
	}

	// CRL + one-time bootstrap token manager.
	crl := mtls.NewCRL(rdb, nil)
	tokens := mtls.NewBootstrapTokenManager(rdb)

	// MySQL for node catalogue and usage accumulation.
	sqlDB, err := db.Open(dbDSN)
	if err != nil {
		log.Fatalf("grpc: db open: %v", err)
	}

	// Nodes service (wires Hub, LoadCache, Handler).
	store := nodes.NewSQLNodeStore(sqlDB)
	svc := nodes.NewService(ca, tokens, rdb, store)
	svc.Hub().Start(context.Background()) // pub/sub goroutine runs for process lifetime

	// gRPC transport TLS (server's own cert, typically from Let's Encrypt).
	serverCert, err := tls.LoadX509KeyPair(grpcCertPath, grpcKeyPath)
	if err != nil {
		log.Fatalf("grpc: load transport TLS cert: %v", err)
	}
	tlsCfg := mtls.NewServerTLSConfig(ca, crl)
	tlsCfg.Certificates = []tls.Certificate{serverCert}

	// Build gRPC server with mTLS transport + identity interceptors.
	srv := grpc.NewServer(
		grpc.Creds(credentials.NewTLS(tlsCfg)),
		grpc.ChainUnaryInterceptor(mtls.UnaryServerInterceptor()),
		grpc.ChainStreamInterceptor(mtls.StreamServerInterceptor()),
	)
	agentv1.RegisterAgentServiceServer(srv, svc.Handler())

	lis, err := net.Listen("tcp", addr)
	if err != nil {
		log.Fatalf("grpc: listen %s: %v", addr, err)
	}
	slog.Info("grpc agent server listening", "addr", addr)
	if err := srv.Serve(lis); err != nil {
		log.Fatalf("grpc: serve: %v", err)
	}
}

func getenvDefault(key, def string) string {
	if v := os.Getenv(key); v != "" {
		return v
	}
	return def
}