From 71145e797f3748b2b608d7f2e0319339fbd41f5b Mon Sep 17 00:00:00 2001
From: Irine Sistiana <49315432+IrineSistiana@users.noreply.github.com>
Date: Fri, 22 Sep 2023 17:35:01 +0800
Subject: [PATCH 6/9] add new string_exp matcher
---
plugin/enabled_plugins.go | 1 +
plugin/matcher/string_exp/string_exp.go | 184 +++++++++++++++++++
plugin/matcher/string_exp/string_exp_test.go | 67 +++++++
3 files changed, 252 insertions(+)
create mode 100644 plugin/matcher/string_exp/string_exp.go
create mode 100644 plugin/matcher/string_exp/string_exp_test.go
diff --git a/plugin/enabled_plugins.go b/plugin/enabled_plugins.go
index 0f7531b..dfb311b 100644
--- a/plugin/enabled_plugins.go
+++ b/plugin/enabled_plugins.go
@@ -38,6 +38,7 @@ import (
_ "github.com/IrineSistiana/mosdns/v5/plugin/matcher/random"
_ "github.com/IrineSistiana/mosdns/v5/plugin/matcher/rcode"
_ "github.com/IrineSistiana/mosdns/v5/plugin/matcher/resp_ip"
+ _ "github.com/IrineSistiana/mosdns/v5/plugin/matcher/string_exp"
// executable
_ "github.com/IrineSistiana/mosdns/v5/plugin/executable/arbitrary"
diff --git a/plugin/matcher/string_exp/string_exp.go b/plugin/matcher/string_exp/string_exp.go
new file mode 100644
index 0000000..692f4e3
--- /dev/null
+++ b/plugin/matcher/string_exp/string_exp.go
@@ -0,0 +1,184 @@
+/*
+ * Copyright (C) 2020-2022, IrineSistiana
+ *
+ * This file is part of mosdns.
+ *
+ * mosdns is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * mosdns is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program. If not, see .
+ */
+
+package string_exp
+
+import (
+ "context"
+ "errors"
+ "fmt"
+ "os"
+ "regexp"
+ "strings"
+
+ "github.com/IrineSistiana/mosdns/v5/pkg/query_context"
+ "github.com/IrineSistiana/mosdns/v5/plugin/executable/sequence"
+)
+
+const PluginType = "string_exp"
+
+func init() {
+ sequence.MustRegMatchQuickSetup(PluginType, QuickSetup)
+}
+
+var _ sequence.Matcher = (*Matcher)(nil)
+
+type Matcher struct {
+ getStr GetStrFunc
+ m StringMatcher
+}
+
+type StringMatcher interface {
+ MatchStr(s string) bool
+}
+
+type GetStrFunc func(qCtx *query_context.Context) string
+
+func (m *Matcher) Match(_ context.Context, qCtx *query_context.Context) (bool, error) {
+ return m.match(qCtx), nil
+}
+
+func (m *Matcher) match(qCtx *query_context.Context) bool {
+ return m.m.MatchStr(m.getStr(qCtx))
+}
+
+func NewMatcher(f GetStrFunc, sm StringMatcher) *Matcher {
+ m := &Matcher{
+ getStr: f,
+ m: sm,
+ }
+ return m
+}
+
+// Format: "scr_string_name op [string]..."
+// scr_string_name = {url_path|server_name|$env_key}
+// op = {zl|eq|prefix|suffix|contains|regexp}
+func QuickSetupFromStr(s string) (sequence.Matcher, error) {
+ sf := strings.Fields(s)
+ if len(sf) < 2 {
+ return nil, errors.New("not enough args")
+ }
+ srcStrName := sf[0]
+ op := sf[1]
+ args := sf[2:]
+
+ var sm StringMatcher
+ switch op {
+ case "zl":
+ sm = opZl{}
+ case "eq":
+ m := make(map[string]struct{})
+ for _, s := range args {
+ m[s] = struct{}{}
+ }
+ sm = &opEq{m: m}
+ case "regexp":
+ var exps []*regexp.Regexp
+ for _, s := range args {
+ exp, err := regexp.Compile(s)
+ if err != nil {
+ return nil, fmt.Errorf("invalid reg expression, %w", err)
+ }
+ exps = append(exps, exp)
+ }
+ sm = &opRegExp{exp: exps}
+ case "prefix":
+ sm = &opF{s: args, f: strings.HasPrefix}
+ case "suffix":
+ sm = &opF{s: args, f: strings.HasSuffix}
+ case "contains":
+ sm = &opF{s: args, f: strings.Contains}
+ default:
+ return nil, fmt.Errorf("invalid operator %s", op)
+ }
+
+ var gf GetStrFunc
+ if strings.HasPrefix(srcStrName, "$") {
+ // Env
+ envKey := strings.TrimPrefix(srcStrName, "$")
+ gf = func(_ *query_context.Context) string {
+ return os.Getenv(envKey)
+ }
+ } else {
+ switch srcStrName {
+ case "url_path":
+ gf = getUrlPath
+ case "server_name":
+ gf = getServerName
+ default:
+ return nil, fmt.Errorf("invalid src string name %s", srcStrName)
+ }
+ }
+ return NewMatcher(gf, sm), nil
+}
+
+// QuickSetup returns a sequence.ExecQuickSetupFunc.
+func QuickSetup(_ sequence.BQ, s string) (sequence.Matcher, error) {
+ return QuickSetupFromStr(s)
+}
+
+type opZl struct{}
+
+func (op opZl) MatchStr(s string) bool {
+ return len(s) == 0
+}
+
+type opEq struct {
+ m map[string]struct{}
+}
+
+func (op *opEq) MatchStr(s string) bool {
+ _, ok := op.m[s]
+ return ok
+}
+
+type opF struct {
+ s []string
+ f func(s, arg string) bool
+}
+
+func (op *opF) MatchStr(s string) bool {
+ for _, sub := range op.s {
+ if op.f(s, sub) {
+ return true
+ }
+ }
+ return false
+}
+
+type opRegExp struct {
+ exp []*regexp.Regexp
+}
+
+func (op *opRegExp) MatchStr(s string) bool {
+ for _, exp := range op.exp {
+ if exp.MatchString(s) {
+ return true
+ }
+ }
+ return false
+}
+
+func getUrlPath(qCtx *query_context.Context) string {
+ return qCtx.QueryMeta().UrlPath
+}
+
+func getServerName(qCtx *query_context.Context) string {
+ return qCtx.QueryMeta().ServerName
+}
diff --git a/plugin/matcher/string_exp/string_exp_test.go b/plugin/matcher/string_exp/string_exp_test.go
new file mode 100644
index 0000000..9140191
--- /dev/null
+++ b/plugin/matcher/string_exp/string_exp_test.go
@@ -0,0 +1,67 @@
+/*
+ * Copyright (C) 2020-2022, IrineSistiana
+ *
+ * This file is part of mosdns.
+ *
+ * mosdns is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * mosdns is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program. If not, see .
+ */
+
+package string_exp
+
+import (
+ "context"
+ "os"
+ "testing"
+
+ "github.com/IrineSistiana/mosdns/v5/pkg/query_context"
+ "github.com/miekg/dns"
+ "github.com/stretchr/testify/require"
+)
+
+func TestMatcher_Match(t *testing.T) {
+ r := require.New(t)
+ q := new(dns.Msg)
+ qc := query_context.NewContext(q, query_context.QueryMeta{UrlPath: "/dns-query", ServerName: "a.b.c"})
+ os.Setenv("STRING_EXP_TEST", "abc")
+
+ doTest := func(arg string, want bool) {
+ t.Helper()
+ urlMatcher, err := QuickSetupFromStr(arg)
+ r.NoError(err)
+ got, err := urlMatcher.Match(context.Background(), qc)
+ r.NoError(err)
+ r.Equal(want, got)
+ }
+
+ doTest("url_path zl", false)
+ doTest("url_path eq /dns-query", true)
+ doTest("url_path eq /123 /dns-query /abc", true)
+ doTest("url_path eq /123 /abc", false)
+ doTest("url_path contains abc dns def", true)
+ doTest("url_path contains abc def", false)
+ doTest("url_path prefix abc /dns def", true)
+ doTest("url_path prefix abc def", false)
+ doTest("url_path suffix abc query def", true)
+ doTest("url_path suffix abc def", false)
+ doTest("url_path regexp ^/dns-query$", true)
+ doTest("url_path regexp ^abc", false)
+
+ doTest("server_name eq abc a.b.c def", true)
+ doTest("server_name eq abc def", false)
+
+ doTest("$STRING_EXP_TEST eq 123 abc def", true)
+ doTest("$STRING_EXP_TEST eq 123 def", false)
+ doTest("$STRING_EXP_TEST_NOT_EXIST eq 123 abc def", false)
+ doTest("$STRING_EXP_TEST_NOT_EXIST zl", true)
+}
--
2.34.8