From 44c8cc690f3d4e2b1e7c69201d4cf617bfaaf25a Mon Sep 17 00:00:00 2001 From: alecthw Date: Sun, 17 Nov 2024 09:29:43 +0800 Subject: [PATCH] add mmdb support for rest ip match --- go.mod | 2 + go.sum | 4 + plugin/data_provider/iface.go | 5 + plugin/data_provider/mmdb/mmdb.go | 61 +++++++++++ plugin/enabled_plugins.go | 2 + plugin/matcher/resp_ip_mmdb/resp_ip_mmdb.go | 107 ++++++++++++++++++++ 6 files changed, 181 insertions(+) create mode 100644 plugin/data_provider/mmdb/mmdb.go create mode 100644 plugin/matcher/resp_ip_mmdb/resp_ip_mmdb.go diff --git a/go.mod b/go.mod index c733c06f9..79fcdb3b0 100644 --- a/go.mod +++ b/go.mod @@ -45,6 +45,8 @@ require ( github.com/mdlayher/socket v0.5.1 // indirect github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 // indirect github.com/onsi/ginkgo/v2 v2.20.0 // indirect + github.com/oschwald/geoip2-golang v1.11.0 // indirect + github.com/oschwald/maxminddb-golang v1.13.1 // indirect github.com/pelletier/go-toml/v2 v2.2.2 // indirect github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 // indirect github.com/prometheus/client_model v0.6.1 // indirect diff --git a/go.sum b/go.sum index 7d2025af9..5562d30f0 100644 --- a/go.sum +++ b/go.sum @@ -58,6 +58,10 @@ github.com/onsi/ginkgo/v2 v2.20.0 h1:PE84V2mHqoT1sglvHc8ZdQtPcwmvvt29WLEEO3xmdZw github.com/onsi/ginkgo/v2 v2.20.0/go.mod h1:lG9ey2Z29hR41WMVthyJBGUBcBhGOtoPF2VFMvBXFCI= github.com/onsi/gomega v1.34.1 h1:EUMJIKUjM8sKjYbtxQI9A4z2o+rruxnzNvpknOXie6k= github.com/onsi/gomega v1.34.1/go.mod h1:kU1QgUvBDLXBJq618Xvm2LUX6rSAfRaFRTcdOeDLwwY= +github.com/oschwald/geoip2-golang v1.11.0 h1:hNENhCn1Uyzhf9PTmquXENiWS6AlxAEnBII6r8krA3w= +github.com/oschwald/geoip2-golang v1.11.0/go.mod h1:P9zG+54KPEFOliZ29i7SeYZ/GM6tfEL+rgSn03hYuUo= +github.com/oschwald/maxminddb-golang v1.13.1 h1:G3wwjdN9JmIK2o/ermkHM+98oX5fS+k5MbwsmL4MRQE= +github.com/oschwald/maxminddb-golang v1.13.1/go.mod h1:K4pgV9N/GcK694KSTmVSDTODk4IsCNThNdTmnaBZ/F8= github.com/pelletier/go-toml/v2 v2.2.2 h1:aYUidT7k73Pcl9nb2gScu7NSrKCSHIDE89b3+6Wq+LM= github.com/pelletier/go-toml/v2 v2.2.2/go.mod h1:1t835xjRzz80PqgE6HHgN2JOsmgYu/h4qDAS4n929Rs= github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= diff --git a/plugin/data_provider/iface.go b/plugin/data_provider/iface.go index 7ec08d205..a4563844e 100644 --- a/plugin/data_provider/iface.go +++ b/plugin/data_provider/iface.go @@ -22,6 +22,7 @@ package data_provider import ( "github.com/IrineSistiana/mosdns/v5/pkg/matcher/domain" "github.com/IrineSistiana/mosdns/v5/pkg/matcher/netlist" + "github.com/oschwald/geoip2-golang" ) type DomainMatcherProvider interface { @@ -31,3 +32,7 @@ type DomainMatcherProvider interface { type IPMatcherProvider interface { GetIPMatcher() netlist.Matcher } + +type MmdbMatcherProvider interface { + GetMmdbMatcher() *geoip2.Reader +} diff --git a/plugin/data_provider/mmdb/mmdb.go b/plugin/data_provider/mmdb/mmdb.go new file mode 100644 index 000000000..542f69d3e --- /dev/null +++ b/plugin/data_provider/mmdb/mmdb.go @@ -0,0 +1,61 @@ +/* + * 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 mmdb + +import ( + "github.com/IrineSistiana/mosdns/v5/coremain" + "github.com/IrineSistiana/mosdns/v5/plugin/data_provider" + "github.com/oschwald/geoip2-golang" +) + +const PluginType = "mmdb" + +func init() { + coremain.RegNewPluginFunc(PluginType, Init, func() any { return new(Args) }) +} + +func Init(bp *coremain.BP, args any) (any, error) { + return NewMmdb(bp, args.(*Args)) +} + +type Args struct { + File string `yaml:"file"` +} + +var _ data_provider.MmdbMatcherProvider = (*Mmdb)(nil) + +type Mmdb struct { + mmdb *geoip2.Reader +} + +func (m *Mmdb) GetMmdbMatcher() *geoip2.Reader { + return m.mmdb +} + +func NewMmdb(bp *coremain.BP, args *Args) (*Mmdb, error) { + m := &Mmdb{} + + db, err := geoip2.Open(args.File) + if err == nil { + m.mmdb = db + } + + return m, nil +} diff --git a/plugin/enabled_plugins.go b/plugin/enabled_plugins.go index 764811fc7..6e810e6f9 100644 --- a/plugin/enabled_plugins.go +++ b/plugin/enabled_plugins.go @@ -24,6 +24,7 @@ import ( // data provider _ "github.com/IrineSistiana/mosdns/v5/plugin/data_provider/domain_set" _ "github.com/IrineSistiana/mosdns/v5/plugin/data_provider/ip_set" + _ "github.com/IrineSistiana/mosdns/v5/plugin/data_provider/mmdb" // matcher _ "github.com/IrineSistiana/mosdns/v5/plugin/matcher/client_ip" @@ -38,6 +39,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/resp_ip_mmdb" _ "github.com/IrineSistiana/mosdns/v5/plugin/matcher/string_exp" // executable diff --git a/plugin/matcher/resp_ip_mmdb/resp_ip_mmdb.go b/plugin/matcher/resp_ip_mmdb/resp_ip_mmdb.go new file mode 100644 index 000000000..33d93be0f --- /dev/null +++ b/plugin/matcher/resp_ip_mmdb/resp_ip_mmdb.go @@ -0,0 +1,107 @@ +/* + * 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 resp_ip_mmdb + +import ( + "context" + "errors" + "fmt" + "github.com/IrineSistiana/mosdns/v5/pkg/query_context" + "github.com/IrineSistiana/mosdns/v5/plugin/data_provider" + "github.com/IrineSistiana/mosdns/v5/plugin/executable/sequence" + "github.com/miekg/dns" + "github.com/oschwald/geoip2-golang" + "net" + "strings" +) + +const PluginType = "resp_ip_mmdb" + +func init() { + sequence.MustRegMatchQuickSetup(PluginType, QuickSetup) +} + +func QuickSetup(bq sequence.BQ, s string) (sequence.Matcher, error) { + if len(s) == 0 { + return nil, errors.New("a iso code probability is required") + } + + args := strings.Fields(s) + if len(args) != 2 { + return nil, errors.New("probability error, must like: resp_ip_mmdb $plugin_name CN") + } + + mmdbName, _ := cutPrefix(args[0], "$") + + p := bq.M().GetPlugin(mmdbName) + provider, _ := p.(data_provider.MmdbMatcherProvider) + if provider == nil { + return nil, fmt.Errorf("cannot find mmdb %s", mmdbName) + } + m := provider.GetMmdbMatcher() + + return &Matcher{args[1], m}, nil +} + +type Matcher struct { + isoCode string + mmdb *geoip2.Reader +} + +func (m *Matcher) Match(_ context.Context, qCtx *query_context.Context) (bool, error) { + r := qCtx.R() + if r == nil { + return false, nil + } + + if m.mmdb == nil { + return false, nil + } + + for _, rr := range r.Answer { + var ip net.IP + switch rr := rr.(type) { + case *dns.A: + ip = rr.A + case *dns.AAAA: + ip = rr.AAAA + default: + continue + } + + record, err := m.mmdb.Country(ip) + if err != nil { + continue + } + + if record.Country.IsoCode == m.isoCode { + return true, nil + } + } + + return false, nil +} + +func cutPrefix(s string, p string) (string, bool) { + if strings.HasPrefix(s, p) { + return strings.TrimPrefix(s, p), true + } + return s, false +}