luci-app-mosdns/mosdns/patches/122-add-new-string_exp-matcher.patch

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 <https://www.gnu.org/licenses/>.
+ */
+
+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 <https://www.gnu.org/licenses/>.
+ */
+
+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