99 lines
2.1 KiB
Go
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
|
|
}
|