Files
goproxy/internal/rewriter/rewriter.go
2025-03-13 15:56:33 +08:00

99 lines
2.1 KiB
Go

package rewriter
import (
"net/http"
"regexp"
"strings"
)
// Rewriter URL重写器
// 用于在反向代理中重写请求URL
type Rewriter struct {
// 重写规则列表
rules []*RewriteRule
}
// RewriteRule 重写规则
type RewriteRule struct {
// 匹配模式
Pattern string
// 替换模式
Replacement string
// 是否使用正则表达式
UseRegex bool
// 编译后的正则表达式
regex *regexp.Regexp
}
// NewRewriter 创建URL重写器
func NewRewriter() *Rewriter {
return &Rewriter{
rules: make([]*RewriteRule, 0),
}
}
// AddRule 添加重写规则
func (r *Rewriter) AddRule(pattern, replacement string, useRegex bool) error {
rule := &RewriteRule{
Pattern: pattern,
Replacement: replacement,
UseRegex: useRegex,
}
if useRegex {
regex, err := regexp.Compile(pattern)
if err != nil {
return err
}
rule.regex = regex
}
r.rules = append(r.rules, rule)
return nil
}
// Rewrite 重写URL
func (r *Rewriter) Rewrite(req *http.Request) {
path := req.URL.Path
for _, rule := range r.rules {
if rule.UseRegex {
if rule.regex.MatchString(path) {
req.URL.Path = rule.regex.ReplaceAllString(path, rule.Replacement)
break
}
} else {
if strings.HasPrefix(path, rule.Pattern) {
req.URL.Path = strings.Replace(path, rule.Pattern, rule.Replacement, 1)
break
}
}
}
}
// RewriteResponse 重写响应
// 主要用于处理响应中的Location头和内容中的URL
func (r *Rewriter) RewriteResponse(resp *http.Response, originHost string) {
// 处理重定向头
location := resp.Header.Get("Location")
if location != "" {
// 将后端服务器的域名替换成代理服务器的域名
for _, rule := range r.rules {
if rule.UseRegex && rule.regex != nil {
if rule.regex.MatchString(location) {
newLocation := rule.regex.ReplaceAllString(location, rule.Replacement)
resp.Header.Set("Location", newLocation)
break
}
}
}
}
}
// LoadRulesFromFile 从文件加载重写规则
func (r *Rewriter) LoadRulesFromFile(filename string) error {
// 实现从配置文件加载规则的逻辑
// 这里省略实现细节
return nil
}