Andrey Blinov 7 years ago
parent
commit
0b1d748f17
5 changed files with 97 additions and 42 deletions
  1. 2 15
      config.go
  2. 9 7
      config_test.go
  3. 16 11
      handler.go
  4. 58 4
      handler_test.go
  5. 12 5
      setup.go

+ 2 - 15
config.go

@@ -2,12 +2,11 @@ package geoip
 
 import (
 	"github.com/mholt/caddy"
-	maxminddb "github.com/oschwald/maxminddb-golang"
 )
 
 // Config specifies configuration parsed for Caddyfile
 type Config struct {
-	DBHandler *maxminddb.Reader // Database's handler if it gets opened.
+	DatabasePath string
 
 	// Yout can set returned header names in config
 	// Country
@@ -51,19 +50,7 @@ func parseConfig(c *caddy.Controller) (Config, error) {
 				if !c.NextArg() {
 					continue
 				}
-				// Check if a database has already been opened
-				if config.DBHandler != nil {
-					continue
-				}
-
-				database := c.Val()
-
-				// Open the database.
-				var err error
-				config.DBHandler, err = maxminddb.Open(database)
-				if err != nil {
-					return config, c.Err("geoip: Can't open database: " + database)
-				}
+				config.DatabasePath = c.Val()
 			case "set_header_country_code":
 				if !c.NextArg() {
 					continue

+ 9 - 7
config_test.go

@@ -11,13 +11,14 @@ func TestParseConfig(t *testing.T) {
 	controller := caddy.NewTestController("http", `
 		localhost:8080
 		geoip {
-			set_header_country_code "Code"
-			set_header_country_name "CountryName"
-			set_header_country_eu "Eu"
-			set_header_city_name "CityName"
-			set_header_location_lat "Lat"
-			set_header_location_lon "Lon"
-			set_header_location_tz "TZ"
+			database path/to/maxmind/db
+			set_header_country_code Code
+			set_header_country_name CountryName
+			set_header_country_eu Eu
+			set_header_city_name CityName
+			set_header_location_lat Lat
+			set_header_location_lon Lon
+			set_header_location_tz TZ
 		}
 	`)
 	actual, err := parseConfig(controller)
@@ -25,6 +26,7 @@ func TestParseConfig(t *testing.T) {
 		t.Errorf("parseConfig return err: %v", err)
 	}
 	expected := Config{
+		DatabasePath:               "path/to/maxmind/db",
 		HeaderNameCountryCode:      "Code",
 		HeaderNameCountryName:      "CountryName",
 		HeaderNameCountryIsEU:      "Eu",

+ 16 - 11
handler.go

@@ -35,20 +35,20 @@ var record struct {
 func (gip GeoIP) addGeoIPHeaders(w http.ResponseWriter, r *http.Request) {
 	clientIP, _ := getClientIP(r, true)
 
-	err := gip.Config.DBHandler.Lookup(clientIP, &record)
+	err := gip.DBHandler.Lookup(clientIP, &record)
 	if err != nil {
-		log.Fatal(err)
+		log.Println(err)
 	}
 
-	w.Header().Add(gip.Config.HeaderNameCountryCode, record.Country.ISOCode)
-	w.Header().Add(gip.Config.HeaderNameCountryIsEU, strconv.FormatBool(record.Country.IsInEuropeanUnion))
-	w.Header().Add(gip.Config.HeaderNameCountryName, record.Country.Names["en"])
+	r.Header.Set(gip.Config.HeaderNameCountryCode, record.Country.ISOCode)
+	r.Header.Set(gip.Config.HeaderNameCountryIsEU, strconv.FormatBool(record.Country.IsInEuropeanUnion))
+	r.Header.Set(gip.Config.HeaderNameCountryName, record.Country.Names["en"])
 
-	w.Header().Add(gip.Config.HeaderNameCityName, record.City.Names["en"])
+	r.Header.Set(gip.Config.HeaderNameCityName, record.City.Names["en"])
 
-	w.Header().Add(gip.Config.HeaderNameLocationLat, strconv.FormatFloat(record.Location.Latitude, 'f', 6, 64))
-	w.Header().Add(gip.Config.HeaderNameLocationLon, strconv.FormatFloat(record.Location.Longitude, 'f', 6, 64))
-	w.Header().Add(gip.Config.HeaderNameLocationTimeZone, record.Location.TimeZone)
+	r.Header.Set(gip.Config.HeaderNameLocationLat, strconv.FormatFloat(record.Location.Latitude, 'f', 6, 64))
+	r.Header.Set(gip.Config.HeaderNameLocationLon, strconv.FormatFloat(record.Location.Longitude, 'f', 6, 64))
+	r.Header.Set(gip.Config.HeaderNameLocationTimeZone, record.Location.TimeZone)
 }
 
 func getClientIP(r *http.Request, strict bool) (net.IP, error) {
@@ -63,10 +63,15 @@ func getClientIP(r *http.Request, strict bool) (net.IP, error) {
 		var err error
 		ip, _, err = net.SplitHostPort(r.RemoteAddr)
 		if err != nil {
-			return nil, err
+			if serr, ok := err.(*net.AddrError); ok && serr.Err == "missing port in address" { // It's not critical try parse
+				ip = r.RemoteAddr
+			} else {
+				log.Printf("Error when SplitHostPort: %v", serr.Err)
+				return nil, err
+			}
 		}
 	}
-	//	ip = "212.50.99.193"
+
 	// Parse the ip address string into a net.IP.
 	parsedIP := net.ParseIP(ip)
 	if parsedIP == nil {

+ 58 - 4
handler_test.go

@@ -1,6 +1,60 @@
 package geoip
 
-// func TestInvokeOK(t *testing.T) {
-// 	//gip := GeoIP{}
-// 	//caddy.NewTestController("http", test.input)
-// }
+import (
+	"net/http"
+	"net/http/httptest"
+	"reflect"
+	"strings"
+	"testing"
+
+	"github.com/mholt/caddy/caddyhttp/httpserver"
+	maxminddb "github.com/oschwald/maxminddb-golang"
+)
+
+func TestToResolveGeoip(t *testing.T) {
+	dbhandler, err := maxminddb.Open("./db/GeoLite2-City.mmdb")
+	if err != nil {
+		t.Errorf("geoip: Can't open database: GeoLite2-City.mmdb")
+	}
+
+	config := Config{}
+
+	config.HeaderNameCountryCode = "X-Geoip-Country-Code"
+	config.HeaderNameCountryIsEU = "X-Geoip-Country-Eu"
+	config.HeaderNameCountryName = "X-Geoip-Country-Name"
+
+	config.HeaderNameCityName = "X-Geoip-City-Name"
+
+	config.HeaderNameLocationLat = "X-Geoip-Location-Lat"
+	config.HeaderNameLocationLon = "X-Geoip-Location-Lon"
+	config.HeaderNameLocationTimeZone = "X-Geoip-Location-Tz"
+
+	var (
+		gotHeaders      http.Header
+		expectedHeaders = http.Header{
+			"X-Geoip-Country-Code": []string{"CY"},
+			"X-Geoip-Location-Lat": []string{"34.684100"},
+			"X-Geoip-Location-Lon": []string{"33.037900"},
+			"X-Geoip-Location-Tz":  []string{"Asia/Nicosia"},
+			"X-Geoip-Country-Eu":   []string{"false"},
+			"X-Geoip-Country-Name": []string{"Cyprus"},
+			"X-Geoip-City-Name":    []string{"Limassol"},
+		}
+	)
+	l := GeoIP{
+		Next: httpserver.HandlerFunc(func(w http.ResponseWriter, r *http.Request) (int, error) {
+			gotHeaders = r.Header
+			return 0, nil
+		}),
+		DBHandler: dbhandler,
+		Config:    config,
+	}
+
+	r := httptest.NewRequest("GET", "/", strings.NewReader(""))
+	r.RemoteAddr = "212.50.99.193"
+	l.ServeHTTP(httptest.NewRecorder(), r)
+
+	if !reflect.DeepEqual(expectedHeaders, gotHeaders) {
+		t.Errorf("Expected %v actual %v", expectedHeaders, gotHeaders)
+	}
+}

+ 12 - 5
setup.go

@@ -3,12 +3,14 @@ package geoip
 import (
 	"github.com/mholt/caddy"
 	"github.com/mholt/caddy/caddyhttp/httpserver"
+	maxminddb "github.com/oschwald/maxminddb-golang"
 )
 
 // GeoIP Comments me
 type GeoIP struct {
-	Next   httpserver.Handler
-	Config Config
+	Next      httpserver.Handler
+	DBHandler *maxminddb.Reader
+	Config    Config
 }
 
 // Init initializes the plugin
@@ -20,16 +22,21 @@ func init() {
 }
 
 func setup(c *caddy.Controller) error {
-	ifconfig, err := parseConfig(c)
+	config, err := parseConfig(c)
 	if err != nil {
 		return err
 	}
 
+	dbhandler, err := maxminddb.Open(config.DatabasePath)
+	if err != nil {
+		return c.Err("geoip: Can't open database: " + config.DatabasePath)
+	}
 	// Create new middleware
 	newMiddleWare := func(next httpserver.Handler) httpserver.Handler {
 		return &GeoIP{
-			Next:   next,
-			Config: ifconfig,
+			Next:      next,
+			DBHandler: dbhandler,
+			Config:    config,
 		}
 	}
 	// Add middleware