pangolin/server/cmd/server/main.go

package main

import (
	"context"
	"crypto/tls"
	"database/sql"
	"encoding/hex"
	"encoding/json"
	"flag"
	"log"
	"log/slog"
	"net"
	"net/http"
	"os"
	"os/signal"
	"strconv"
	"strings"
	"syscall"
	"time"

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

	"github.com/wangjia/pangolin/server/internal/admin"
	"github.com/wangjia/pangolin/server/internal/apierr"
	"github.com/wangjia/pangolin/server/internal/auth"
	"github.com/wangjia/pangolin/server/internal/codes"
	"github.com/wangjia/pangolin/server/internal/db"
	"github.com/wangjia/pangolin/server/internal/devices"
	"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/provision"
	"github.com/wangjia/pangolin/server/internal/provision/providers"
	"github.com/wangjia/pangolin/server/internal/redisutil"
	"github.com/wangjia/pangolin/server/internal/scheduler"
	"github.com/wangjia/pangolin/server/internal/scheduler/probe"
	"github.com/wangjia/pangolin/server/internal/usage"
)

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

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

	// Admin backend (separate internal listener, optional).
	startAdminIfConfigured()

	// ─── Shared: Redis ────────────────────────────────────────────────────────
	redisAddr := getenvDefault("REDIS_ADDR", "127.0.0.1:6379")
	rdb, err := redisutil.New(redisAddr, os.Getenv("REDIS_PASSWORD"), 0)
	if err != nil {
		log.Fatalf("redis connect: %v", err)
	}

	// ─── Shared: DB ───────────────────────────────────────────────────────────
	dbDSN := os.Getenv("DB_DSN")
	var sqlDB *sql.DB
	if dbDSN != "" {
		sqlDB, err = db.Open(dbDSN)
		if err != nil {
			log.Fatalf("db open: %v", err)
		}
	}

	// ─── Shared: nodes.Service (Hub shared by HTTP connect + gRPC agent) ──────
	// Constructed here so both sides use the same Hub instance.
	var nodeSvc *nodes.Service
	if sqlDB != nil {
		nodeStore := nodes.NewSQLNodeStore(sqlDB)

		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")

		if grpcAddr != "" && caKeyPath != "" && caCertPath != "" &&
			grpcCertPath != "" && grpcKeyPath != "" {

			ca, err := mtls.NewCA(mtls.CAConfig{KeyPath: caKeyPath, CertPath: caCertPath})
			if err != nil {
				log.Fatalf("load node CA: %v", err)
			}
			tokens := mtls.NewBootstrapTokenManager(rdb)
			crl := mtls.NewCRL(rdb, nil)
			nodeSvc = nodes.NewService(ca, tokens, rdb, nodeStore)
			nodeSvc.Hub().Start(context.Background())

			go startGRPCWithService(nodeSvc, grpcAddr, grpcCertPath, grpcKeyPath, ca, crl)

		} else {
			if grpcAddr != "" {
				slog.Warn("grpc agent server disabled: one or more required env vars missing",
					"CA_KEY_PATH_set", caKeyPath != "",
					"CA_CERT_PATH_set", caCertPath != "",
					"GRPC_CERT_PATH_set", grpcCertPath != "",
					"GRPC_KEY_PATH_set", grpcKeyPath != "")
			}
			// Even without gRPC, build a Hub-only nodeSvc so the HTTP /v1/nodes
			// routes (ListNodes/Connect) work locally. nil CA/tokens means
			// agent Enroll would panic if called — acceptable when gRPC is off.
			nodeSvc = nodes.NewService(nil, nil, rdb, nodeStore)
			nodeSvc.Hub().Start(context.Background())
			log.Printf("nodes.Service: gRPC not configured; Hub active for command queueing (HTTP /v1/nodes enabled)")
		}
	}

	// ─── HTTP router ──────────────────────────────────────────────────────────
	r := chi.NewRouter()
	r.Use(chimw.Logger)
	r.Use(chimw.Recoverer)
	r.Use(apierr.Middleware)

	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"})
	})

	// Optional probe ingest route. sharedProbeStore is reused by the scheduler
	// (below) when both are enabled, so they share one Redis-backed store.
	var sharedProbeStore *probe.Store
	if probeJSON := os.Getenv("PROBE_SECRETS"); probeJSON != "" {
		pRDB, 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(probeJSON), &secretMap); err != nil {
				log.Fatalf("probe: invalid PROBE_SECRETS JSON: %v", err)
			}
			reg := probe.NewMapRegistry(secretMap)
			st := probe.NewStore(pRDB)
			sharedProbeStore = st
			r.Post("/probe/report", probe.NewIngestHandler(reg, st).ServeHTTP)
			log.Printf("probe ingest route registered (%d probe(s))", len(secretMap))
		}
	}

	// ─── Scheduler (optional) ─────────────────────────────────────────────────
	//
	// Set SCHED_ENABLED=true to start the three leader-elected scheduler loops
	// (DetectLoop / OrchestrateLoop / CapacityLoop) in this process alongside the
	// HTTP and gRPC servers. Rollback: unset SCHED_ENABLED and restart; all other
	// server functionality is unaffected.
	if os.Getenv("SCHED_ENABLED") == "true" {
		schedRDB, err := redisutil.New(redisAddr, os.Getenv("REDIS_PASSWORD"), 0)
		if err != nil {
			log.Printf("scheduler: redis connect failed (%v) — scheduler disabled", err)
		} else {
			// Reuse the probe store if the probe route is wired up; otherwise
			// create a standalone store on the scheduler Redis client.
			ps := sharedProbeStore
			if ps == nil {
				ps = probe.NewStore(schedRDB)
			}
			// Prefer real lifecycle (SQL) + provision (#14) wiring when a DB is
			// available; fall back to no-op stubs otherwise. Without vendor
			// credentials the provision service's CreateNode fails gracefully
			// (pending replacements stay pending) while lifecycle reads/transitions
			// still operate on real node rows.
			var cfg scheduler.Config
			if sqlDB != nil {
				provStore := provision.NewMySQLStore(sqlDB)
				provSvc, perr := provision.NewService(provision.Config{
					Store:    provStore,
					Adapters: providers.NewRegistry(),
				})
				if perr != nil {
					log.Printf("scheduler: provision service init failed (%v) — using stub config", perr)
					cfg = scheduler.BuildStubConfig(schedRDB, ps)
				} else {
					cfg = scheduler.BuildRealConfig(schedRDB, ps, sqlDB,
						nodes.NewLoadCache(schedRDB), provSvc, provStore)
					log.Printf("scheduler: real lifecycle + provision wired")
				}
			} else {
				cfg = scheduler.BuildStubConfig(schedRDB, ps)
			}
			sched := scheduler.New(cfg)

			// sigCtx cancels on SIGINT/SIGTERM, giving the scheduler up to
			// TickTimeout to finish any in-flight tick before process exit.
			sigCtx, stopSig := signal.NotifyContext(context.Background(),
				os.Interrupt, syscall.SIGTERM)
			go func() {
				defer stopSig()
				if err := sched.Start(sigCtx); err != nil {
					slog.Error("scheduler: Start error", "error", err)
				}
			}()
			log.Printf("scheduler started (SCHED_ENABLED=true)")
		}
	}

	// ─── /v1 routes (requires DB) ────────────────────────────────────────────
	if sqlDB != nil {
		mountV1(r, sqlDB, rdb, nodeSvc)
	} else {
		log.Printf("DB_DSN not set; /v1 routes disabled")
	}

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

// mountV1 wires all /v1 routes onto r.
func mountV1(r chi.Router, sqlDB *sql.DB, rdb *redis.Client, nodeSvc *nodes.Service) {
	// ── JWT TokenManager ──────────────────────────────────────────────────────
	jwtPrivPath := os.Getenv("JWT_PRIVATE_KEY_PATH")
	jwtKID := os.Getenv("JWT_KEY_ID")
	var tm *auth.TokenManager
	if jwtPrivPath != "" && jwtKID != "" {
		tc, err := auth.LoadTokenConfig(jwtPrivPath, jwtKID, parseKeyMap(os.Getenv("JWT_PUBLIC_KEYS")))
		if err != nil {
			log.Fatalf("JWT keys: %v", err)
		}
		tm, err = auth.NewTokenManager(rdb, tc)
		if err != nil {
			log.Fatalf("JWT token manager: %v", err)
		}
	} else {
		log.Printf("JWT not configured — /v1 protected routes will be unavailable")
	}

	// ── Auth ──────────────────────────────────────────────────────────────────
	var authHandler *auth.Handler
	if tm != nil {
		var mailer auth.Mailer
		if smtpHost := os.Getenv("SMTP_HOST"); smtpHost != "" {
			mailer = auth.NewSMTPMailer(auth.SMTPConfig{
				Host:     smtpHost,
				Port:     intEnvDefault("SMTP_PORT", 587),
				Username: os.Getenv("SMTP_USERNAME"),
				Password: os.Getenv("SMTP_PASSWORD"),
				From:     getenvDefault("SMTP_FROM", "no-reply@pangolin.app"),
			})
		} else {
			mailer = auth.NewLogMailer(nil) // dev: code printed to log
		}
		rl := auth.NewRateLimiter(rdb, nil)
		authStore := auth.NewSQLStore(sqlDB)
		authSvc := auth.NewService(authStore, rdb, rl, tm, mailer, auth.ServiceConfig{}, nil)
		authHandler = auth.NewHandler(authSvc)
	}

	// ── User TOTP (web 用户中心 2FA) ─────────────────────────────────────────────
	// Only enabled when a valid 32-byte key is configured (USER_TOTP_ENC_KEY,
	// raw 32 bytes or 64 hex chars); secrets are encrypted at rest under it.
	var totpHandler *auth.TOTPHandler
	if tm != nil {
		if key := parseTOTPKey(os.Getenv("USER_TOTP_ENC_KEY")); key != nil {
			totpHandler = auth.NewTOTPHandler(sqlDB, key, tm, rdb)
		} else if os.Getenv("USER_TOTP_ENC_KEY") != "" {
			log.Printf("USER_TOTP_ENC_KEY invalid (need 32 bytes or 64 hex chars); user TOTP disabled")
		}
	}

	// ── Codes ────────────────────────────────────────────────────────────────
	codesStore := codes.NewStore(sqlDB)
	codesSvc := codes.NewService(codesStore, rdb, 5, time.Hour)
	redeemHandler := codes.NewRedeemHandler(codesSvc)
	webhookHandler := codes.NewWebhookHandler(codesStore, rdb,
		os.Getenv("WEBHOOK_SECRET"), 5*time.Minute, 15*time.Minute)

	// ── Devices ───────────────────────────────────────────────────────────────
	devicesStore := devices.NewStore(sqlDB)
	devicesSvc := devices.NewService(devicesStore, nil) // NoopRevoker for MVP
	devicesHandler := devices.NewHandler(devicesSvc)

	// ── Usage ─────────────────────────────────────────────────────────────────
	usageStore := usage.NewStore(sqlDB)
	usageSvc := usage.NewService(usageStore, rdb, nil, time.Hour)
	usageHandler := usage.NewUsageHandler(usageSvc)
	adsHandler := usage.NewAdsUnlockHandler(usageSvc)

	// ── Account / Plans / Notices ─────────────────────────────────────────────
	accountAPI := httpapi.NewAccountAPI(sqlDB)

	// ── Nodes + Connect ───────────────────────────────────────────────────────
	var nodeAPI *httpapi.NodeAPI
	var nodeStore nodes.NodeStore
	if nodeSvc != nil {
		nodeStore = nodeSvc.Store()
		publicURL := os.Getenv("PANGOLIN_PUBLIC_URL")
		if publicURL == "" {
			// 没设 PANGOLIN_PUBLIC_URL → BuildClientConfig 会静默跳过国内分流
			// (rule_set 没有公网基址可下载),客户端 split_cn=1 也无效、全量走隧道。
			// 这是个隐蔽坑(国内流量绕道出海、变慢),启动期显式告警。
			slog.Warn("PANGOLIN_PUBLIC_URL 未设置:国内分流(split_cn)将被静默跳过,客户端全量走隧道。" +
				"如需国内直连,设为控制面对外公网基址(如 http://<公网IP>:8080)")
		}
		nodeAPI = httpapi.NewNodeAPI(nodeStore, nodeSvc.Hub(), os.Getenv("NODE_DERIVE_KEY"), publicURL)
	}

	// 国内分流(#5)的 rule-set 静态服务:GET /v1/rules/{name}.srs(自托管,
	// sing-box 直连下载;PANGOLIN_RULES_DIR 默认 /var/lib/pangolin/rules)。
	rulesHandler := httpapi.NewRulesHandler(os.Getenv("PANGOLIN_RULES_DIR"))

	// ── Subscription URL (web user-center 订阅导入) ───────────────────────────
	subAPI := httpapi.NewSubscriptionAPI(sqlDB, nodeStore, os.Getenv("NODE_DERIVE_KEY"), os.Getenv("SUB_BASE"))

	// Public subscription endpoint (no JWT): clients fetch their sing-box config
	// from the per-user URL. Only available when a node store is configured.
	if nodeStore != nil {
		r.Get("/sub/{token}", subAPI.ServeSub)
	}

	// ─── Mount under /v1 ─────────────────────────────────────────────────────
	r.Route("/v1", func(v1 chi.Router) {
		// Public (no auth): 国内分流 rule-set 下载(sing-box 不带 token)。
		v1.Get("/rules/{name}", rulesHandler.Serve)

		// Public (no auth): 诊断端点,量节点→目标出海段耗时(白名单限定,防 SSRF),
		// 供白盒拆「接入段 vs 出海段」。只返回耗时数字,不传业务数据。
		v1.Get("/diag/egress", httpapi.NewDiagHandler().EgressTiming)

		// Public (no auth): auth endpoints.
		if authHandler != nil {
			v1.Post("/auth/code", authHandler.SendCode)
			v1.Post("/auth/register", authHandler.Register)
			v1.Post("/auth/login", authHandler.Login)
			v1.Post("/auth/refresh", authHandler.Refresh)
			v1.Post("/auth/logout", authHandler.Logout)
			if totpHandler != nil {
				v1.Post("/auth/login/totp", totpHandler.LoginTOTP)
			}
		}

		// Webhook: HMAC-authenticated, no JWT.
		v1.Post("/webhook/store/codes", webhookHandler.ServeHTTP)

		// Protected: all routes that require a valid Bearer JWT.
		if tm != nil {
			v1.Group(func(protected chi.Router) {
				protected.Use(auth.RequireAuth(tm))

				protected.Route("/me", func(me chi.Router) {
					me.Get("/", accountAPI.GetMe) // 子路由根，避免与 Route("/me") 冲突致 404
					devicesHandler.RegisterRoutes(me)
					// Web 用户中心调用 /v1/me/redeem；app 仍用 /v1/redeem。两者同处理器。
					me.Post("/redeem", redeemHandler.ServeHTTP)
					me.Get("/subscription", subAPI.GetSubscription)
					me.Post("/subscription/reset", subAPI.ResetSubscription)
					if totpHandler != nil {
						me.Post("/totp/setup", totpHandler.Setup)
						me.Post("/totp/verify", totpHandler.Verify)
						me.Post("/totp/disable", totpHandler.Disable)
					}
				})
				protected.Post("/redeem", redeemHandler.ServeHTTP)
				protected.Get("/usage", usageHandler.ServeHTTP)
				protected.Post("/ads/unlock", adsHandler.ServeHTTP)
				protected.Get("/plans", accountAPI.ListPlans)
				protected.Get("/notices", accountAPI.ListNotices)

				if nodeAPI != nil {
					protected.Get("/nodes", nodeAPI.ListNodes)
					protected.Post("/nodes/{id}/connect", nodeAPI.ConnectNode)
					protected.Post("/nodes/{id}/disconnect", nodeAPI.DisconnectNode)
				}
			})
		}
	})
}

// startAdminIfConfigured launches the admin listener when ADMIN_SECRET_KEY and
// DB_DSN are both present.
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)
		}
	}()
}

// startGRPCWithService starts the mTLS gRPC agent server using the shared
// nodes.Service so the HTTP connect handler and the gRPC server share one Hub.
func startGRPCWithService(
	svc *nodes.Service,
	addr, grpcCertPath, grpcKeyPath string,
	ca *mtls.CA,
	crl *mtls.CRL,
) {
	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}

	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)
	}
}

// parseKeyMap parses "kid1:/path1,kid2:/path2" into a map. Used for JWT_PUBLIC_KEYS.
func parseKeyMap(raw string) map[string]string {
	if raw == "" {
		return nil
	}
	m := map[string]string{}
	for _, pair := range strings.Split(raw, ",") {
		pair = strings.TrimSpace(pair)
		if pair == "" {
			continue
		}
		i := strings.IndexByte(pair, ':')
		if i <= 0 || i == len(pair)-1 {
			continue
		}
		m[strings.TrimSpace(pair[:i])] = strings.TrimSpace(pair[i+1:])
	}
	return m
}

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

// parseTOTPKey returns a 32-byte AES key from USER_TOTP_ENC_KEY, accepting either
// 64 hex chars or a raw 32-byte string. Returns nil when unset/invalid.
func parseTOTPKey(raw string) []byte {
	if len(raw) == 64 {
		if b, err := hex.DecodeString(raw); err == nil && len(b) == 32 {
			return b
		}
		return nil
	}
	if len(raw) == 32 {
		return []byte(raw)
	}
	return nil
}

func intEnvDefault(key string, def int) int {
	v := os.Getenv(key)
	if v == "" {
		return def
	}
	n, err := strconv.Atoi(v)
	if err != nil || n == 0 {
		return def
	}
	return n
}