lolly/internal/middleware/rewrite/rewrite.go

// Package rewrite 提供 URL 重写中间件，支持正则表达式匹配和多种重写标志。
package rewrite

import (
	"regexp"
	"strings"

	"github.com/valyala/fasthttp"
	"rua.plus/lolly/internal/config"
)

// Flag 重写标志类型。
type Flag int

const (
	// FlagLast 继续匹配其他规则。
	FlagLast Flag = iota
	// FlagRedirect 返回 302 临时重定向。
	FlagRedirect
	// FlagPermanent 返回 301 永久重定向。
	FlagPermanent
	// FlagBreak 停止匹配规则。
	FlagBreak
)

// parseFlag 解析配置中的标志字符串。
func parseFlag(s string) Flag {
	switch strings.ToLower(s) {
	case "redirect":
		return FlagRedirect
	case "permanent":
		return FlagPermanent
	case "break":
		return FlagBreak
	default:
		return FlagLast
	}
}

// Rule 编译后的重写规则。
type Rule struct {
	pattern     *regexp.Regexp
	replacement string
	flag        Flag
}

// RewriteMiddleware URL 重写中间件。
type RewriteMiddleware struct {
	rules []Rule
}

// New 创建重写中间件。
func New(rules []config.RewriteRule) (*RewriteMiddleware, error) {
	compiled := make([]Rule, 0, len(rules))
	for _, r := range rules {
		re, err := regexp.Compile(r.Pattern)
		if err != nil {
			return nil, err
		}
		compiled = append(compiled, Rule{
			pattern:     re,
			replacement: r.Replacement,
			flag:        parseFlag(r.Flag),
		})
	}
	return &RewriteMiddleware{rules: compiled}, nil
}

// Name 返回中间件名称。
func (m *RewriteMiddleware) Name() string {
	return "rewrite"
}

// Process 应用重写规则。
func (m *RewriteMiddleware) Process(next fasthttp.RequestHandler) fasthttp.RequestHandler {
	return func(ctx *fasthttp.RequestCtx) {
		path := string(ctx.Path())
		originalPath := path

		for _, rule := range m.rules {
			if rule.pattern.MatchString(path) {
				// 执行正则替换
				newPath := rule.pattern.ReplaceAllString(path, rule.replacement)

				switch rule.flag {
				case FlagRedirect:
					ctx.Redirect(newPath, fasthttp.StatusFound)
					return
				case FlagPermanent:
					ctx.Redirect(newPath, fasthttp.StatusMovedPermanently)
					return
				case FlagBreak:
					// 修改路径后停止匹配
					ctx.Request.SetRequestURI(newPath)
					next(ctx)
					return
				case FlagLast:
					// 修改路径，继续匹配其他规则
					path = newPath
					ctx.Request.SetRequestURI(path)
				}
			}
		}

		// 如果路径被修改过，需要重新设置
		if path != originalPath {
			ctx.Request.SetRequestURI(path)
		}

		next(ctx)
	}
}

// Rules 返回编译后的规则列表（用于调试）。
func (m *RewriteMiddleware) Rules() []Rule {
	return m.rules
}