Files
NixMsg/cmd/nixmsg/serve.go

339 lines
9.3 KiB
Go

package main
import (
"context"
"errors"
"fmt"
"io/fs"
"log/slog"
"net"
"net/http"
"os"
"os/signal"
"path/filepath"
"strings"
"syscall"
"time"
"git.asio.asia/nixevol/NixMsg/internal/admin"
"git.asio.asia/nixevol/NixMsg/internal/app/group"
"git.asio.asia/nixevol/NixMsg/internal/app/identity"
"git.asio.asia/nixevol/NixMsg/internal/app/message"
"git.asio.asia/nixevol/NixMsg/internal/app/port"
"git.asio.asia/nixevol/NixMsg/internal/app/presence"
"git.asio.asia/nixevol/NixMsg/internal/auth"
"git.asio.asia/nixevol/NixMsg/internal/broker"
"git.asio.asia/nixevol/NixMsg/internal/config"
"git.asio.asia/nixevol/NixMsg/internal/httpx"
"git.asio.asia/nixevol/NixMsg/internal/listener"
"git.asio.asia/nixevol/NixMsg/internal/metrics"
"git.asio.asia/nixevol/NixMsg/internal/store"
"git.asio.asia/nixevol/NixMsg/web"
)
func cmdServe(_ []string) error {
cfgPath := config.PathFromEnv()
cfg, err := config.Load(cfgPath)
if err != nil {
return err
}
if err := cfg.Validate(); err != nil {
return err
}
setupJSONLogger(cfg.Log)
ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
defer stop()
return runServe(ctx, cfg)
}
func runServe(ctx context.Context, cfg config.Config) error {
if err := os.MkdirAll(cfg.DataDir, 0o755); err != nil {
return fmt.Errorf("mkdir data_dir: %w", err)
}
db, err := store.Open(cfg.DataDir, cfg.SQLiteSynchronous)
if err != nil {
return err
}
defer func() { _ = db.Close() }()
ok, err := store.HasAdminPassword(ctx, db.Write)
if err != nil {
return err
}
if !ok {
return errors.New("admin password not initialized; run: nixmsg admin init")
}
hashPool := auth.NewPool()
sessionTokens := auth.NewSessionTokens()
apiTokens := auth.NewAPITokens()
loginLocks := auth.NewLoginLocks()
memConns := message.NewMemoryConns()
msgLim := message.LimitsFromFullConfig(cfg)
login := broker.NewLogin(broker.LoginOptions{
DB: db,
Pool: hashPool,
Tokens: sessionTokens,
Locks: loginLocks,
IdleDays: cfg.SessionIdleDays,
})
// 先用空下行建 message,拿到 App 指针;broker 建好后再 WithDownlink 重建并回填 uplink.msg。
msgApp := message.New(db, msgLim, hashPool,
message.WithLocks(loginLocks),
message.WithConnRegistry(memConns),
)
uplink := &appUplink{
msg: msgApp,
conns: memConns,
log: slog.Default(),
}
sess := broker.NewSession(broker.SessionOptions{
Login: login,
Inner: uplink,
Limits: broker.HelloLimits{
MaxBodyBytes: cfg.Limits.MaxBodyBytes,
MaxMetaBytes: cfg.Limits.MaxMetaBytes,
MaxFrameBytes: cfg.Limits.MaxFrameBytes,
MaxTTLSeconds: int64(cfg.Limits.MaxTTLSeconds),
MaxScheduleSeconds: int64(cfg.Limits.MaxScheduleSeconds),
AckTimeoutSeconds: int64(cfg.Limits.AckTimeoutSeconds),
ServerVersion: Version,
},
Logger: slog.Default(),
})
brk, err := broker.New(broker.Options{
Authenticator: login,
Uplink: sess,
Logger: slog.Default(),
OnPublishDropped: func(dropCtx context.Context, endpointID string, connID port.ConnID, payload []byte) {
if dropErr := msgApp.OnPublishDropped(dropCtx, endpointID, connID, payload); dropErr != nil {
slog.Error("on publish dropped", "endpoint", endpointID, "err", dropErr)
}
},
})
if err != nil {
return fmt.Errorf("broker: %w", err)
}
defer func() { _ = brk.Close() }()
sess.Attach(brk)
msgApp = message.New(db, msgLim, hashPool,
message.WithLocks(loginLocks),
message.WithConnRegistry(memConns),
message.WithDownlink(brk),
)
uplink.msg = msgApp
uplink.down = brk
presApp := presence.New(presence.Config{
DB: db,
Downlink: brk,
Conns: &presenceConnTable{conns: memConns},
})
sess.SetPresence(presApp)
uplink.presence = presApp
idApp := identity.New(identity.Config{
DB: db,
Hash: hashPool,
Locks: loginLocks,
Sessions: sessionTokens,
MaxScheduleSeconds: int64(cfg.Limits.MaxScheduleSeconds),
Logger: slog.Default(),
ConnControl: brk,
})
uplink.identity = idApp
groupApp := group.New(group.Config{
DB: db,
Talk: idApp,
Online: presApp,
Downlink: brk,
MaxGroupMembers: cfg.Limits.MaxGroupMembers,
})
uplink.groups = groupApp
if recoverErr := msgApp.RecoverOnStart(ctx); recoverErr != nil {
return fmt.Errorf("message recover: %w", recoverErr)
}
trustedNets := httpx.ParseCIDRs(cfg.TrustedProxies)
adminHandler := admin.New(admin.Deps{
DB: db,
Hash: hashPool,
Tokens: apiTokens,
Locks: loginLocks,
Logger: slog.Default(),
TrustedProxies: trustedNets,
Identity: idApp,
Groups: groupApp,
Config: cfg,
Version: Version,
KickEndpoint: func(kickCtx context.Context, endpointID string) (bool, error) {
if _, found := brk.ConnInfoOf(endpointID); !found {
return false, nil
}
if kickErr := sess.Kick(kickCtx, endpointID); kickErr != nil {
return false, kickErr
}
return true, nil
},
})
metricsReg := metrics.New()
buildHandlers := func(proxies *listener.ProxySet) listener.Handlers {
return listener.Handlers{
MQTT: brk.WSHandler(proxies),
ClientAPI: idApp.Handler(),
AdminAPI: adminHandler,
Metrics: metricsReg.Handler(),
MetricsToken: cfg.Metrics.Token,
Static: staticFileHandler(),
Healthz: http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
w.WriteHeader(http.StatusOK)
_, _ = w.Write([]byte("ok"))
}),
Readyz: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if readyErr := db.Ready(r.Context()); readyErr != nil {
slog.Error("readyz failed", "err", readyErr)
http.Error(w, "not ready", http.StatusServiceUnavailable)
return
}
w.WriteHeader(http.StatusOK)
_, _ = w.Write([]byte("ok"))
}),
}
}
handlers := buildHandlers(nil)
shared := cfg.AdminListen == ""
var clientHandler, adminHTTP http.Handler
if shared {
clientHandler = listener.NewMux(listener.RoleShared, handlers)
} else {
clientHandler = listener.NewMux(listener.RoleClient, handlers)
adminHTTP = listener.NewMux(listener.RoleAdmin, handlers)
}
lnOpts := listener.Options{
Listen: cfg.Listen,
AdminListen: cfg.AdminListen,
DataDir: cfg.DataDir,
CertFile: cfg.TLS.CertFile,
KeyFile: cfg.TLS.KeyFile,
AllowPlaintext: cfg.TLS.AllowPlaintext,
TrustedProxies: cfg.TrustedProxies,
ClientHandler: clientHandler,
AdminHandler: adminHTTP,
OnMQTT: func(conn net.Conn) {
_ = brk.AttachTCP(conn)
},
Logger: slog.Default(),
}
lnSrv, err := listener.New(lnOpts)
if err != nil {
return fmt.Errorf("listener: %w", err)
}
handlers = buildHandlers(lnSrv.Proxies())
if shared {
lnOpts.ClientHandler = listener.NewMux(listener.RoleShared, handlers)
} else {
lnOpts.ClientHandler = listener.NewMux(listener.RoleClient, handlers)
lnOpts.AdminHandler = listener.NewMux(listener.RoleAdmin, handlers)
}
lnSrv, err = listener.New(lnOpts)
if err != nil {
return fmt.Errorf("listener: %w", err)
}
if err := lnSrv.Start(ctx); err != nil {
return err
}
defer func() { _ = lnSrv.Close() }()
if err := writeListenAddr(cfg.DataDir, lnSrv.ListenAddr()); err != nil {
return err
}
if cfg.AdminListen != "" && lnSrv.AdminAddr() != "" {
if err := os.WriteFile(filepath.Join(cfg.DataDir, "admin.addr"), []byte(lnSrv.AdminAddr()+"\n"), 0o644); err != nil {
return err
}
}
loopCtx, loopCancel := context.WithCancel(ctx)
defer loopCancel()
go messageLoops(loopCtx, msgApp, memConns)
<-ctx.Done()
loopCancel()
_ = lnSrv.Close()
drainCtx, drainCancel := context.WithTimeout(context.Background(), 10*time.Second)
defer drainCancel()
if drainErr := db.Queue.Drain(drainCtx); drainErr != nil && !errors.Is(drainErr, context.DeadlineExceeded) {
slog.Error("write queue drain", "err", drainErr)
}
return nil
}
func messageLoops(ctx context.Context, msgApp *message.App, conns *message.MemoryConns) {
t := time.NewTicker(time.Second)
defer t.Stop()
for {
select {
case <-ctx.Done():
return
case <-t.C:
nowMs := time.Now().UnixMilli()
if _, err := msgApp.DispatchDue(ctx, nowMs, 100); err != nil {
slog.Error("dispatch due", "err", err)
}
for ep, live := range conns.Snapshot() {
if err := msgApp.PushPending(ctx, ep, live.ConnID); err != nil {
slog.Debug("push pending", "endpoint", ep, "err", err)
}
}
if err := msgApp.CleanupOnce(ctx, nowMs); err != nil {
slog.Error("cleanup once", "err", err)
}
}
}
}
func staticFileHandler() http.Handler {
root := web.Dist()
for _, prefix := range []string{"dist", "stub"} {
sub, err := fs.Sub(root, prefix)
if err != nil {
continue
}
if f, err := sub.Open("index.html"); err == nil {
_ = f.Close()
return http.FileServer(http.FS(sub))
}
}
return http.FileServer(http.FS(root))
}
func writeListenAddr(dataDir, addr string) error {
path := filepath.Join(dataDir, "listen.addr")
return os.WriteFile(path, []byte(addr+"\n"), 0o644)
}
func setupJSONLogger(cfg config.LogConfig) {
level := slog.LevelInfo
switch strings.ToLower(strings.TrimSpace(cfg.Level)) {
case "debug":
level = slog.LevelDebug
case "warn", "warning":
level = slog.LevelWarn
case "error":
level = slog.LevelError
}
h := slog.NewJSONHandler(os.Stdout, &slog.HandlerOptions{Level: level})
slog.SetDefault(slog.New(h))
}