From 691553417aca1aa82c31a59383501c4d47bc89d8 Mon Sep 17 00:00:00 2001 From: alecthw Date: Fri, 27 Oct 2023 14:56:01 +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 bd524b20b..fc4541bd7 100644 --- a/go.mod +++ b/go.mod @@ -48,6 +48,8 @@ require ( github.com/mdlayher/netlink v1.7.2 // indirect github.com/mdlayher/socket v0.5.0 // indirect github.com/onsi/ginkgo/v2 v2.13.0 // indirect + github.com/oschwald/geoip2-golang v1.9.0 // indirect + github.com/oschwald/maxminddb-golang v1.11.0 // indirect github.com/pelletier/go-toml/v2 v2.1.0 // indirect github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 // indirect github.com/prometheus/client_model v0.5.0 // indirect diff --git a/go.sum b/go.sum index 93ab36f2d..0a1af2339 100644 --- a/go.sum +++ b/go.sum @@ -183,6 +183,10 @@ github.com/onsi/ginkgo/v2 v2.13.0 h1:0jY9lJquiL8fcf3M4LAXN5aMlS/b2BV86HFFPCPMgE4 github.com/onsi/ginkgo/v2 v2.13.0/go.mod h1:TE309ZR8s5FsKKpuB1YAQYBzCaAfUgatB/xlT/ETL/o= github.com/onsi/gomega v1.27.10 h1:naR28SdDFlqrG6kScpT8VWpu1xWY5nJRCF3XaYyBjhI= github.com/onsi/gomega v1.27.10/go.mod h1:RsS8tutOdbdgzbPtzzATp12yT7kM5I5aElG3evPbQ0M= +github.com/oschwald/geoip2-golang v1.9.0 h1:uvD3O6fXAXs+usU+UGExshpdP13GAqp4GBrzN7IgKZc= +github.com/oschwald/geoip2-golang v1.9.0/go.mod h1:BHK6TvDyATVQhKNbQBdrj9eAvuwOMi2zSFXizL3K81Y= +github.com/oschwald/maxminddb-golang v1.11.0 h1:aSXMqYR/EPNjGE8epgqwDay+P30hCBZIveY0WZbAWh0= +github.com/oschwald/maxminddb-golang v1.11.0/go.mod h1:YmVI+H0zh3ySFR3w+oz8PCfglAFj3PuCmui13+P9zDg= github.com/pelletier/go-toml/v2 v2.1.0 h1:FnwAJ4oYMvbT/34k9zzHuZNrhlz48GB3/s6at6/MHO4= github.com/pelletier/go-toml/v2 v2.1.0/go.mod h1:tJU2Z3ZkXwnxa4DPO899bsyIoywizdUvyaeZurnPPDc= github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0= 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 01f9f6c68..4028ca004 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 +}