Files
pangolin/server/cmd/server/main.go
T
wangjia 636a3bbf2f feat(server/stats): stats-overhaul Phase2 — 每设备归因 + GB 综合配额
服务端记账从「账户」细到「每设备」(每设备独立 dp_uuid),配额单位分钟→GB
按账户综合卡控。000015_per_device_usage 迁移(mysql+sqlite 双份):devices.dp_uuid
+ usage_device_daily 表 + plans.daily_mb。handler_grpc 按 dp_uuid 回映射
(user_id,device_id) 双写账户+每设备;usage 服务/handler 暴露 /v1/usage(/devices)。
含 sqlite_per_device / usage handler 测试。

注:此迁移 prod 已 migrate up 运行、部署二进制已内嵌;本次补提交使 git 与线上一致。

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
2026-06-28 18:12:00 +08:00

502 lines
19 KiB
Go
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
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)
deviceUsageHandler := usage.NewDeviceUsageHandler(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/redeemapp 仍用 /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.Get("/usage/devices", deviceUsageHandler.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
}