github.com/yaling888/clash@v1.53.0/component/script/matcher.go (about)

     1  package script
     2  
     3  import (
     4  	"errors"
     5  	"fmt"
     6  	T "time"
     7  
     8  	"go.starlark.net/lib/time"
     9  	"go.starlark.net/starlark"
    10  	"go.starlark.net/syntax"
    11  
    12  	C "github.com/yaling888/clash/constant"
    13  )
    14  
    15  const metadataLocalKey = "local.metadata_key"
    16  
    17  var allowKeywords = map[string]bool{
    18  	"_metadata":     true,
    19  	"now":           true,
    20  	"type":          true,
    21  	"network":       true,
    22  	"host":          true,
    23  	"process_name":  true,
    24  	"process_path":  true,
    25  	"src_ip":        true,
    26  	"src_port":      true,
    27  	"dst_ip":        true,
    28  	"dst_port":      true,
    29  	"user_agent":    true,
    30  	"special_proxy": true,
    31  	"inbound_port":  true,
    32  }
    33  
    34  var parseOption = syntax.LegacyFileOptions()
    35  
    36  var nowErrFunc = func() (T.Time, error) {
    37  	return T.Now(), nil
    38  }
    39  
    40  var _ C.Matcher = (*Matcher)(nil)
    41  
    42  type Matcher struct {
    43  	name string
    44  	key  string
    45  
    46  	program *starlark.Program
    47  }
    48  
    49  func (m *Matcher) Name() string {
    50  	return m.name
    51  }
    52  
    53  func (m *Matcher) Eval(metadata *C.Metadata) (string, error) {
    54  	metadataDict, err := metadataToDict(metadata)
    55  	if err != nil {
    56  		return "", err
    57  	}
    58  
    59  	predefined := make(starlark.StringDict)
    60  	predefined["_metadata"] = metadataDict
    61  
    62  	thread := &starlark.Thread{
    63  		Print: func(_ *starlark.Thread, _ string) {},
    64  	}
    65  
    66  	thread.SetLocal(metadataLocalKey, metadata)
    67  
    68  	time.SetNow(thread, nowErrFunc)
    69  
    70  	results, err := m.program.Init(thread, predefined)
    71  	if err != nil {
    72  		return "", err
    73  	}
    74  
    75  	evalResult := results[m.key]
    76  	if v, ok := evalResult.(starlark.String); ok {
    77  		return v.GoString(), nil
    78  	}
    79  
    80  	if evalResult == nil {
    81  		return "", errors.New("invalid return type, got <nil>, want string")
    82  	}
    83  
    84  	return "", fmt.Errorf("invalid return type, got %s, want string", evalResult.Type())
    85  }
    86  
    87  func (m *Matcher) Match(metadata *C.Metadata) (bool, error) {
    88  	predefined, err := metadataToStringDict(metadata, nil)
    89  	if err != nil {
    90  		return false, err
    91  	}
    92  
    93  	predefined["now"] = time.Time(T.Now())
    94  
    95  	thread := &starlark.Thread{
    96  		Print: func(_ *starlark.Thread, _ string) {},
    97  	}
    98  
    99  	thread.SetLocal(metadataLocalKey, metadata)
   100  
   101  	time.SetNow(thread, nowErrFunc)
   102  
   103  	results, err := m.program.Init(thread, predefined)
   104  	if err != nil {
   105  		return false, err
   106  	}
   107  
   108  	evalResult := results[m.key]
   109  	if v, ok := evalResult.(starlark.Bool); ok {
   110  		return bool(v), nil
   111  	}
   112  
   113  	if evalResult == nil {
   114  		return false, errors.New("invalid return type, got <nil>, want bool")
   115  	}
   116  
   117  	return false, fmt.Errorf("invalid return type, got %s, want bool", evalResult.Type())
   118  }
   119  
   120  func NewMatcher(name, filename, code string) (_ *Matcher, err error) {
   121  	defer func() {
   122  		if r := recover(); r != nil {
   123  			err = fmt.Errorf("parse script code panic: %v", r)
   124  		}
   125  	}()
   126  
   127  	if filename == "" {
   128  		filename = name + ".star"
   129  	}
   130  
   131  	key := fmt.Sprintf("_%s_eval", name)
   132  	if name == "main" {
   133  		code = fmt.Sprintf("%s\n\n%s=main(_clash_ctx, _metadata)\n", code, key)
   134  	} else {
   135  		code = fmt.Sprintf("%s=(%s)", key, code)
   136  	}
   137  
   138  	starFile, err := parseOption.Parse(filename, code, 0)
   139  	if err != nil {
   140  		return nil, fmt.Errorf("parse script code error: %w", err)
   141  	}
   142  
   143  	program, err := starlark.FileProgram(starFile, func(s string) bool {
   144  		rs, ok := allowKeywords[s]
   145  		return ok && rs
   146  	})
   147  	if err != nil {
   148  		return nil, fmt.Errorf("compile script code error: %w", err)
   149  	}
   150  
   151  	return &Matcher{
   152  		name:    name,
   153  		key:     key,
   154  		program: program,
   155  	}, nil
   156  }