subconverter-go/internal/middleware/security.go

package middleware

import (
	"strings"

	"github.com/gofiber/fiber/v2"
	"github.com/subconverter-go/internal/logging"
)

// SetupSecurity 设置安全中间件
// 添加各种安全相关的HTTP头
func SetupSecurity() fiber.Handler {
	return func(c *fiber.Ctx) error {
		// 设置安全相关的HTTP头
		c.Set("X-Content-Type-Options", "nosniff")
		c.Set("X-Frame-Options", "DENY")
		c.Set("X-XSS-Protection", "1; mode=block")
		c.Set("Referrer-Policy", "strict-origin-when-cross-origin")
		c.Set("Content-Security-Policy", "default-src 'self'")

		// 移除服务器信息
		c.Set("Server", "")

		return c.Next()
	}
}

// SetupAuthentication 设置认证中间件
// 验证访问令牌，保护API端点
func SetupAuthentication(logger *logging.Logger, validTokens []string) fiber.Handler {
	allowed := make(map[string]struct{}, len(validTokens))
	for _, token := range validTokens {
		if token == "" {
			continue
		}
		allowed[strings.TrimSpace(token)] = struct{}{}
	}

	return func(c *fiber.Ctx) error {
		// 如果没有配置令牌，跳过认证
		if len(allowed) == 0 {
			return c.Next()
		}

		token := ""

		authHeader := c.Get("Authorization")
		if strings.HasPrefix(authHeader, "Bearer ") {
			token = strings.TrimSpace(strings.TrimPrefix(authHeader, "Bearer "))
		}

		if token == "" {
			token = strings.TrimSpace(c.Query("token"))
		}

		if token == "" {
			logger.Warn("Missing authorization token")
			return forbiddenResponse(c)
		}

		if _, ok := allowed[token]; !ok {
			logger.WithField("token", token).Warn("Invalid authorization token")
			return forbiddenResponse(c)
		}

		c.Locals("authenticated", true)
		c.Locals("token", token)

		return c.Next()
	}
}

func forbiddenResponse(c *fiber.Ctx) error {
	c.Set("Content-Type", "text/plain")
	return c.Status(fiber.StatusForbidden).SendString("Forbidden\n")
}

// SetupCompression 设置压缩中间件
// 启用响应压缩，减少传输数据量
func SetupCompression() fiber.Handler {
	return func(c *fiber.Ctx) error {
		// 设置Accept-Encoding头
		c.Set("Accept-Encoding", "gzip, deflate")

		return c.Next()
	}
}