7e12327ab1
remove unnecessary float64 convert and clean code Signed-off-by: xuweiwei <xuweiwei_yewu@cmss.chinamobile.com>
95 lines
2.6 KiB
Go
95 lines
2.6 KiB
Go
// Package geoip implements a max mind database plugin.
|
|
package geoip
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"net"
|
|
"path/filepath"
|
|
|
|
"github.com/coredns/coredns/plugin"
|
|
clog "github.com/coredns/coredns/plugin/pkg/log"
|
|
"github.com/coredns/coredns/request"
|
|
|
|
"github.com/miekg/dns"
|
|
"github.com/oschwald/geoip2-golang"
|
|
)
|
|
|
|
var log = clog.NewWithPlugin(pluginName)
|
|
|
|
// GeoIP is a plugin that add geo location data to the request context by looking up a maxmind
|
|
// geoIP2 database, and which data can be later consumed by other middlewares.
|
|
type GeoIP struct {
|
|
Next plugin.Handler
|
|
db db
|
|
}
|
|
|
|
type db struct {
|
|
*geoip2.Reader
|
|
// provides defines the schemas that can be obtained by querying this database, by using
|
|
// bitwise operations.
|
|
provides int
|
|
}
|
|
|
|
const (
|
|
city = 1 << iota
|
|
)
|
|
|
|
var probingIP = net.ParseIP("127.0.0.1")
|
|
|
|
func newGeoIP(dbPath string) (*GeoIP, error) {
|
|
reader, err := geoip2.Open(dbPath)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("failed to open database file: %v", err)
|
|
}
|
|
db := db{Reader: reader}
|
|
schemas := []struct {
|
|
provides int
|
|
name string
|
|
validate func() error
|
|
}{
|
|
{name: "city", provides: city, validate: func() error { _, err := reader.City(probingIP); return err }},
|
|
}
|
|
// Query the database to figure out the database type.
|
|
for _, schema := range schemas {
|
|
if err := schema.validate(); err != nil {
|
|
// If we get an InvalidMethodError then we know this database does not provide that schema.
|
|
if _, ok := err.(geoip2.InvalidMethodError); !ok {
|
|
return nil, fmt.Errorf("unexpected failure looking up database %q schema %q: %v", filepath.Base(dbPath), schema.name, err)
|
|
}
|
|
} else {
|
|
db.provides |= schema.provides
|
|
}
|
|
}
|
|
|
|
if db.provides&city == 0 {
|
|
return nil, fmt.Errorf("database does not provide city schema")
|
|
}
|
|
|
|
return &GeoIP{db: db}, nil
|
|
}
|
|
|
|
// ServeDNS implements the plugin.Handler interface.
|
|
func (g GeoIP) ServeDNS(ctx context.Context, w dns.ResponseWriter, r *dns.Msg) (int, error) {
|
|
return plugin.NextOrFailure(pluginName, g.Next, ctx, w, r)
|
|
}
|
|
|
|
// Metadata implements the metadata.Provider Interface in the metadata plugin, and is used to store
|
|
// the data associated with the source IP of every request.
|
|
func (g GeoIP) Metadata(ctx context.Context, state request.Request) context.Context {
|
|
srcIP := net.ParseIP(state.IP())
|
|
|
|
switch {
|
|
case g.db.provides&city == city:
|
|
data, err := g.db.City(srcIP)
|
|
if err != nil {
|
|
log.Debugf("Setting up metadata failed due to database lookup error: %v", err)
|
|
return ctx
|
|
}
|
|
g.setCityMetadata(ctx, data)
|
|
}
|
|
return ctx
|
|
}
|
|
|
|
// Name implements the Handler interface.
|
|
func (g GeoIP) Name() string { return pluginName }
|