Skip to content

Commit

Permalink
chore: new rule
Browse files Browse the repository at this point in the history
  • Loading branch information
Adlyq authored and xishang0128 committed Apr 3, 2024
1 parent 0ce18bf commit 6731bc9
Show file tree
Hide file tree
Showing 7 changed files with 77 additions and 49 deletions.
6 changes: 3 additions & 3 deletions config/config.go
Original file line number Diff line number Diff line change
Expand Up @@ -293,8 +293,8 @@ type RawTuicServer struct {
}

type RawMitm struct {
Port int `yaml:"port" json:"port"`
Rules []string `yaml:"rules" json:"rules"`
Port int `yaml:"port" json:"port"`
Rules []rewrites.RawMitmRule `yaml:"rules" json:"rules"`
}

type RawConfig struct {
Expand Down Expand Up @@ -515,7 +515,7 @@ func UnmarshalRawConfig(buf []byte) (*RawConfig, error) {
},
MITM: RawMitm{
Port: 0,
Rules: []string{},
Rules: []rewrites.RawMitmRule{},
},
Profile: Profile{
StoreSelected: true,
Expand Down
2 changes: 2 additions & 0 deletions constant/metadata.go
Original file line number Diff line number Diff line change
Expand Up @@ -177,6 +177,8 @@ func (m *Metadata) SourceAddrPort() netip.AddrPort {
func (m *Metadata) SourceDetail() string {
if m.Type == INNER {
return fmt.Sprintf("%s", MihomoName)
} else if m.Type == MITM {
return fmt.Sprintf("%s-MITM", MihomoName)
}

switch {
Expand Down
38 changes: 38 additions & 0 deletions constant/rewrite.go
Original file line number Diff line number Diff line change
@@ -1,6 +1,8 @@
package constant

import (
"encoding/json"
"errors"
regexp "github.com/dlclark/regexp2"
)

Expand Down Expand Up @@ -37,6 +39,42 @@ const (

type RewriteType int

// UnmarshalYAML unserialize RewriteType with yaml
func (e *RewriteType) UnmarshalYAML(unmarshal func(any) error) error {
var tp string
if err := unmarshal(&tp); err != nil {
return err
}
mode, exist := RewriteTypeMapping[tp]
if !exist {
return errors.New("invalid MITM Action")
}
*e = mode
return nil
}

// MarshalYAML serialize RewriteType with yaml
func (e RewriteType) MarshalYAML() (any, error) {
return e.String(), nil
}

// UnmarshalJSON unserialize RewriteType with json
func (e *RewriteType) UnmarshalJSON(data []byte) error {
var tp string
json.Unmarshal(data, &tp)
mode, exist := RewriteTypeMapping[tp]
if !exist {
return errors.New("invalid MITM Action")
}
*e = mode
return nil
}

// MarshalJSON serialize RewriteType with json
func (e RewriteType) MarshalJSON() ([]byte, error) {
return json.Marshal(e.String())
}

func (rt RewriteType) String() string {
switch rt {
case MitmReject:
Expand Down
7 changes: 6 additions & 1 deletion rewrite/handler.go
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ import (
"bufio"
"bytes"
"errors"
log "github.com/sirupsen/logrus"
"io"
"io/ioutil"
"net/http"
Expand Down Expand Up @@ -31,6 +32,8 @@ func (*RewriteHandler) HandleRequest(session *mitm.Session) (*http.Request, *htt
return nil, nil
}

log.Infof("[MITM] %s <- request %s", rule.RuleType().String(), request.URL.String())

switch rule.RuleType() {
case C.MitmReject:
response = session.NewResponse(http.StatusNotFound, nil)
Expand Down Expand Up @@ -113,6 +116,8 @@ func (*RewriteHandler) HandleResponse(session *mitm.Session) *http.Response {
return nil
}

log.Infof("[MITM] %s <- response %s", rule.RuleType().String(), request.URL.String())

switch rule.RuleType() {
case C.MitmResponseHeader:
if len(response.Header) == 0 {
Expand Down Expand Up @@ -182,7 +187,7 @@ func matchRewriteRule(url string, isRequest bool) (rr C.Rewrite, sub []string, f
if isRequest {
found = rewrites.SearchInRequest(func(r C.Rewrite) bool {
sub, err := r.URLRegx().FindStringMatch(url)
if err != nil {
if err != nil || sub == nil {
return false
}

Expand Down
61 changes: 18 additions & 43 deletions rewrite/parser.go
Original file line number Diff line number Diff line change
Expand Up @@ -7,12 +7,7 @@ import (
C "github.com/metacubex/mihomo/constant"
)

func ParseRewrite(line string) (C.Rewrite, error) {
url, others, found := strings.Cut(strings.TrimSpace(line), "url")
if !found {
return nil, errInvalid
}

func ParseRewrite(line RawMitmRule) (C.Rewrite, error) {
var (
urlRegx *regexp.Regexp
ruleType *C.RewriteType
Expand All @@ -22,57 +17,37 @@ func ParseRewrite(line string) (C.Rewrite, error) {
err error
)

url := line.Url
urlRegx, err = regexp.Compile(strings.Trim(url, " "), regexp.None)
if err != nil {
return nil, err
}

others = strings.Trim(others, " ")
first := strings.Split(others, " ")[0]
for k, v := range C.RewriteTypeMapping {
if k == others {
ruleType = &v
ruleType = &line.Action
switch *ruleType {
case C.Mitm302, C.Mitm307:
{
rulePayload = line.New
break
}
case C.MitmRequestHeader, C.MitmRequestBody, C.MitmResponseHeader, C.MitmResponseBody:
{
var old string
if line.Old == nil {
old = ".*"
} else {
old = *line.Old
}

if k != first {
continue
}

rs := trimArr(strings.Split(others, k))
l := len(rs)
if l > 2 {
continue
}

if l == 1 {
ruleType = &v
rulePayload = rs[0]
break
} else {
ruleRegx, err = regexp.Compile(rs[0], regexp.None)
re, err := regexp.Compile(old, regexp.Singleline)
if err != nil {
return nil, err
}
ruleRegx = re

ruleType = &v
rulePayload = rs[1]
break
rulePayload = line.New
}
}

if ruleType == nil {
return nil, errInvalid
}

return NewRewriteRule(urlRegx, *ruleType, ruleRegx, rulePayload), nil
}

func trimArr(arr []string) (r []string) {
for _, e := range arr {
if s := strings.Trim(e, " "); s != "" {
r = append(r, s)
}
}
return
}
8 changes: 6 additions & 2 deletions rewrite/rewrite.go
Original file line number Diff line number Diff line change
@@ -1,7 +1,6 @@
package rewrites

import (
"errors"
regexp "github.com/dlclark/regexp2"
"strconv"
"strings"
Expand All @@ -11,7 +10,12 @@ import (
"github.com/gofrs/uuid"
)

var errInvalid = errors.New("invalid rewrite rule")
type RawMitmRule struct {
Url string `yaml:"url" json:"url"`
Action C.RewriteType `yaml:"action" json:"action"`
Old *string `yaml:"old" json:"old"`
New string `yaml:"new" json:"new"`
}

type RewriteRule struct {
id string
Expand Down
4 changes: 4 additions & 0 deletions tunnel/tunnel.go
Original file line number Diff line number Diff line change
Expand Up @@ -620,6 +620,10 @@ func match(metadata *C.Metadata) (C.Proxy, C.Rule, error) {
}

for _, rule := range getRules(metadata) {
if metadata.Type == C.MITM && rule.Adapter() == "MITM" {
continue
}

if !resolved && shouldResolveIP(rule, metadata) {
func() {
ctx, cancel := context.WithTimeout(context.Background(), resolver.DefaultDNSTimeout)
Expand Down

0 comments on commit 6731bc9

Please sign in to comment.