kra-new/internal/server/gin.go

package server

import (
	"log/slog"
	"net/http"
	"strings"
	"time"

	"kra/internal/config"
	websocket "kra/internal/integration/websocket"
	"kra/internal/server/httpx"
	"kra/internal/server/middleware"
	"kra/internal/server/staticfiles"
	systemservice "kra/internal/service/system"
	platformmodule "kra/pkg/module"

	"github.com/gin-gonic/gin"
	kratoshttp "github.com/go-kratos/kratos/v3/transport/http"
	melody "github.com/olahol/melody"
)

func NewGinEngineWithRuntime(runtime *config.Store, access *systemservice.AccessControlService, auth middleware.TokenAuthenticator, security *systemservice.SecurityService, audit *systemservice.AuditRecorder, logger *slog.Logger, version string, routes platformmodule.RouteRegistrar, ws *websocket.Server) *gin.Engine {
	if runtime == nil {
		runtime = config.NewStore(nil)
	}
	gin.SetMode(gin.ReleaseMode)
	engine := gin.New()
	if err := engine.SetTrustedProxies(nil); err != nil && logger != nil {
		logger.Error("disable trusted proxies failed", "error", err)
	}
	engine.Use(middleware.RequestMeta(), middleware.Recovery(logger), websocketHandshakeLogger(logger), middleware.AccessLog(runtime, logger, version), middleware.CORS(runtime), middleware.ErrorAudit(logger), middleware.SecurityRateLimit(security))

	prefix := ""
	snapshot := runtime.Snapshot()
	if snapshot != nil && snapshot.Admin != nil {
		prefix = strings.TrimSuffix(snapshot.Admin.RouterPrefix, "/")
	}
	public := engine.Group(prefix)
	public.GET("/health", func(c *gin.Context) { c.JSON(http.StatusOK, "ok") })

	private := engine.Group(prefix)
	// The reference administration behavior installs operation recording after JWT,
	// password and Casbin/data-scope middleware. Keeping the recorder after the
	// equivalent Kra guards avoids persisting rejected unauthenticated or
	// unauthorized requests as successful business operations.
	private.Use(middleware.Auth(auth), middleware.MustChangePassword(), middleware.AccessControl(runtime, access), middleware.OperationAudit(runtime, audit))
	if routes != nil {
		routes.RegisterRoutes(public, private, engine)
	}
	if ws != nil && logger != nil {
		ws.OnConnect(func(session *melody.Session) {
			request := session.Request
			logger.InfoContext(request.Context(), "websocket connection established",
				"mod", "websocket", "path", request.URL.Path, "remote_addr", request.RemoteAddr,
				"origin", request.Header.Get("Origin"), "user_id", session.Keys["user_id"],
				"request_id", session.Keys["request_id"])
		})
		ws.OnDisconnect(func(session *melody.Session) {
			request := session.Request
			logger.InfoContext(request.Context(), "websocket connection closed",
				"mod", "websocket", "path", request.URL.Path, "remote_addr", request.RemoteAddr,
				"user_id", session.Keys["user_id"], "request_id", session.Keys["request_id"])
		})
	}
	handleWebSocket := func(c *gin.Context) {
		wsLogger := logger
		hasToken := middleware.RequestToken(c, true) != ""
		expectedPath := websocketRoutePath(prefix, ws)
		if ws == nil || !ws.Enabled() || c.Request.URL.Path != expectedPath {
			if wsLogger != nil {
				wsLogger.WarnContext(c.Request.Context(), "websocket handshake rejected: endpoint unavailable",
					"mod", "websocket", "path", c.Request.URL.Path, "expected_path", expectedPath,
					"enabled", ws != nil && ws.Enabled(), "has_token", hasToken)
			}
			c.Status(http.StatusNotFound)
			return
		}
		if !middleware.AuthenticateWebSocket(c, auth) {
			if wsLogger != nil {
				wsLogger.WarnContext(c.Request.Context(), "websocket handshake rejected: authentication failed",
					"mod", "websocket", "path", c.Request.URL.Path, "status", c.Writer.Status(),
					"has_token", hasToken, "origin", c.GetHeader("Origin"))
			}
			return
		}
		claims := middleware.Claims(c)
		keys := map[string]any{
			"request_id": contextString(c, "request_id"),
		}
		if claims != nil {
			keys["user_id"] = claims.ID
		}
		if wsLogger != nil {
			wsLogger.InfoContext(c.Request.Context(), "websocket authentication accepted; upgrading connection",
				"mod", "websocket", "path", c.Request.URL.Path, "origin", c.GetHeader("Origin"),
				"remote_addr", c.Request.RemoteAddr, "user_id", keys["user_id"],
				"request_id", keys["request_id"])
		}
		if err := ws.HandleRequestWithKeys(c.Writer, c.Request, keys); err != nil && wsLogger != nil {
			wsLogger.WarnContext(c.Request.Context(), "websocket upgrade failed", "mod", "websocket",
				"path", c.Request.URL.Path, "origin", c.GetHeader("Origin"), "error", err,
				"request_id", keys["request_id"])
		}
	}
	if ws != nil && ws.Enabled() {
		engine.GET(ws.Path(), handleWebSocket)
	}
	if logger != nil {
		logger.Info("websocket endpoint initialized", "mod", "websocket",
			"enabled", ws != nil && ws.Enabled(), "path", websocketPath(ws))
	}
	env := ""
	if snapshot != nil && snapshot.Admin != nil && snapshot.Admin.App != nil {
		env = strings.ToLower(strings.TrimSpace(snapshot.Admin.App.Env))
	}
	if env == "" || env == "development" || env == "dev" || env == "test" || env == "local" {
		registerSwagger(engine, prefix, version, logger)
	}
	staticfiles.Register(engine, runtime)
	if logger != nil {
		for _, route := range engine.Routes() {
			logger.Info("router registered", "method", route.Method, "path", route.Path)
		}
		logger.Info("router register success", "route_count", len(engine.Routes()))
	}

	engine.NoRoute(func(c *gin.Context) {
		if ws != nil && ws.Enabled() && c.Request.Method == http.MethodGet && c.Request.URL.Path == websocketRoutePath(prefix, ws) {
			handleWebSocket(c)
			return
		}
		if staticfiles.Serve(c, runtime) {
			return
		}
		httpx.Fail(c, "请求的接口不存在")
	})
	return engine
}

func websocketHandshakeLogger(logger *slog.Logger) gin.HandlerFunc {
	return func(c *gin.Context) {
		if logger != nil && strings.EqualFold(strings.TrimSpace(c.GetHeader("Upgrade")), "websocket") {
			logger.InfoContext(c.Request.Context(), "websocket handshake packet received",
				"mod", "websocket", "path", c.Request.URL.Path, "host", c.Request.Host,
				"remote_addr", c.Request.RemoteAddr, "origin", c.GetHeader("Origin"),
				"has_token", middleware.RequestToken(c, true) != "",
				"connection", c.GetHeader("Connection"), "user_agent", c.Request.UserAgent(),
				"request_id", contextString(c, "request_id"))
		}
		c.Next()
	}
}

func contextString(c *gin.Context, key string) string {
	value, _ := c.Get(key)
	text, _ := value.(string)
	return text
}

func websocketPath(ws *websocket.Server) string {
	if ws == nil {
		return ""
	}
	return ws.Path()
}

func websocketRoutePath(prefix string, ws *websocket.Server) string {
	path := websocketPath(ws)
	if path == "" {
		return ""
	}
	if prefix == "" {
		return path
	}
	return strings.TrimSuffix(prefix, "/") + "/" + strings.TrimPrefix(path, "/")
}

func NewGinServer(c *config.Server, engine *gin.Engine) *kratoshttp.Server {
	network, address := "tcp", ":8000"
	if c != nil && c.HTTP != nil {
		if c.HTTP.Network != "" {
			network = c.HTTP.Network
		}
		if c.HTTP.Addr != "" {
			address = c.HTTP.Addr
		}
	}
	timeout := 10 * time.Minute
	if c != nil && c.HTTP != nil && c.HTTP.Timeout > 0 {
		timeout = c.HTTP.Timeout
	}
	server := kratoshttp.NewServer(kratoshttp.Network(network), kratoshttp.Address(address), kratoshttp.Timeout(timeout))
	server.HandlePrefix("/", engine)
	server.Server.ReadHeaderTimeout = 10 * time.Second
	server.Server.ReadTimeout = timeout
	server.Server.WriteTimeout = timeout
	server.Server.IdleTimeout = 2 * time.Minute
	server.Server.MaxHeaderBytes = 1 << 20
	return server
}