kra-new/internal/server/swagger.go

package server

import (
	"encoding/json"
	"log/slog"
	"regexp"
	"sort"
	"strings"
	"sync"

	"kra/internal/service"

	"github.com/gin-gonic/gin"
	swaggerFiles "github.com/swaggo/files"
	ginSwagger "github.com/swaggo/gin-swagger"
	"github.com/swaggo/swag"
)

const swaggerInstanceName = "kra-admin"

var (
	swaggerPathParameter = regexp.MustCompile(`:([A-Za-z0-9_]+)`)
	swaggerRegistration  sync.Once
	swaggerDocument      runtimeSwaggerDocument
)

type runtimeSwaggerDocument struct {
	mu  sync.RWMutex
	doc string
}

func (d *runtimeSwaggerDocument) ReadDoc() string {
	d.mu.RLock()
	defer d.mu.RUnlock()
	return d.doc
}

func (d *runtimeSwaggerDocument) replace(doc string) {
	d.mu.Lock()
	d.doc = doc
	d.mu.Unlock()
}

func registerSwagger(engine *gin.Engine, prefix, version string, logger *slog.Logger) {
	swaggerRegistration.Do(func() { swag.Register(swaggerInstanceName, &swaggerDocument) })
	swaggerDocument.replace(buildSwaggerDocument(engine.Routes(), prefix, version))
	path := strings.TrimSuffix(prefix, "/") + "/swagger/*any"
	engine.GET(path, ginSwagger.WrapHandler(
		swaggerFiles.Handler,
		ginSwagger.InstanceName(swaggerInstanceName),
		ginSwagger.URL("doc.json"),
		ginSwagger.PersistAuthorization(true),
	))
	if logger != nil {
		logger.Info("register swagger handler", "mod", "system", "path", path)
	}
}

func buildSwaggerDocument(routes []gin.RouteInfo, prefix, version string) string {
	basePath := strings.TrimSuffix(prefix, "/")
	if basePath == "" {
		basePath = "/"
	}
	sort.Slice(routes, func(i, j int) bool {
		if routes[i].Path == routes[j].Path {
			return routes[i].Method < routes[j].Method
		}
		return routes[i].Path < routes[j].Path
	})
	paths := make(map[string]map[string]any, len(routes))
	for _, route := range routes {
		method := strings.ToLower(route.Method)
		switch method {
		case "get", "post", "put", "delete", "patch":
		default:
			continue
		}
		apiPath := route.Path
		if prefix != "" {
			apiPath = strings.TrimPrefix(apiPath, strings.TrimSuffix(prefix, "/"))
		}
		if apiPath == "" {
			apiPath = "/"
		}
		documentPath := swaggerPathParameter.ReplaceAllString(apiPath, `{$1}`)
		group, description := service.RouteMetadata(route.Method, apiPath)
		if description == "" {
			description = route.Method + " " + apiPath
		}
		operation := map[string]any{
			"tags":        []string{group},
			"summary":     description,
			"operationId": swaggerOperationID(route.Method, apiPath),
			"produces":    []string{"application/json"},
			"responses": map[string]any{
				"200": map[string]any{"description": "OK", "schema": map[string]any{"$ref": "#/definitions/Response"}},
			},
		}
		if parameters := swaggerPathParameters(apiPath); len(parameters) > 0 {
			operation["parameters"] = parameters
		}
		if !swaggerPublicPath(apiPath) {
			operation["security"] = []map[string][]string{{"ApiKeyAuth": {}}}
		}
		if paths[documentPath] == nil {
			paths[documentPath] = map[string]any{}
		}
		paths[documentPath][method] = operation
	}
	document := map[string]any{
		"swagger":  "2.0",
		"info":     map[string]any{"title": "Kra Administration API", "version": version},
		"basePath": basePath,
		"schemes":  []string{"http", "https"},
		"consumes": []string{"application/json"},
		"produces": []string{"application/json"},
		"securityDefinitions": map[string]any{
			"ApiKeyAuth": map[string]any{"type": "apiKey", "name": "x-token", "in": "header"},
		},
		"paths": paths,
		"definitions": map[string]any{
			"Response": map[string]any{
				"type": "object",
				"properties": map[string]any{
					"code": map[string]any{"type": "integer"},
					"data": map[string]any{"type": "object"},
					"msg":  map[string]any{"type": "string"},
				},
			},
		},
	}
	raw, _ := json.Marshal(document)
	return string(raw)
}

func swaggerOperationID(method, path string) string {
	value := strings.ToLower(method) + "_" + strings.Trim(path, "/")
	value = swaggerPathParameter.ReplaceAllString(value, "$1")
	return strings.Map(func(r rune) rune {
		if r >= 'a' && r <= 'z' || r >= '0' && r <= '9' || r == '_' {
			return r
		}
		return '_'
	}, value)
}

func swaggerPathParameters(path string) []map[string]any {
	matches := swaggerPathParameter.FindAllStringSubmatch(path, -1)
	parameters := make([]map[string]any, 0, len(matches))
	for _, match := range matches {
		parameters = append(parameters, map[string]any{"name": match[1], "in": "path", "required": true, "type": "string"})
	}
	return parameters
}

func swaggerPublicPath(path string) bool {
	for _, marker := range []string{
		"/health", "/base/login", "/base/captcha", "/init/checkdb", "/init/initdb",
		"/api/freshCasbin", "/sysExportTemplate/exportExcelByToken", "/sysExportTemplate/exportTemplateByToken",
		"/sysError/createSysError", "/info/getInfoDataSource", "/info/getInfoPublic",
	} {
		if path == marker {
			return true
		}
	}
	return false
}