Skip to content

Commit

Permalink
add mmdb support for rest ip match
Browse files Browse the repository at this point in the history
  • Loading branch information
alecthw committed Nov 17, 2024
1 parent 025823c commit 44c8cc6
Show file tree
Hide file tree
Showing 6 changed files with 181 additions and 0 deletions.
2 changes: 2 additions & 0 deletions go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
4 changes: 4 additions & 0 deletions go.sum
Original file line number Diff line number Diff line change
Expand Up @@ -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=
Expand Down
5 changes: 5 additions & 0 deletions plugin/data_provider/iface.go
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand All @@ -31,3 +32,7 @@ type DomainMatcherProvider interface {
type IPMatcherProvider interface {
GetIPMatcher() netlist.Matcher
}

type MmdbMatcherProvider interface {
GetMmdbMatcher() *geoip2.Reader
}
61 changes: 61 additions & 0 deletions plugin/data_provider/mmdb/mmdb.go
Original file line number Diff line number Diff line change
@@ -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 <https://www.gnu.org/licenses/>.
*/

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
}
2 changes: 2 additions & 0 deletions plugin/enabled_plugins.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand All @@ -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
Expand Down
107 changes: 107 additions & 0 deletions plugin/matcher/resp_ip_mmdb/resp_ip_mmdb.go
Original file line number Diff line number Diff line change
@@ -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 <https://www.gnu.org/licenses/>.
*/

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
}

0 comments on commit 44c8cc6

Please sign in to comment.