فهرست منبع

Add replacers

Andrey Blinov 7 سال پیش
والد
کامیت
76eeba9090
4فایلهای تغییر یافته به همراه155 افزوده شده و 140 حذف شده
  1. 0 89
      handler.go
  2. 0 50
      handler_test.go
  3. 84 1
      setup.go
  4. 71 0
      setup_test.go

+ 0 - 89
handler.go

@@ -1,89 +0,0 @@
-package geoip
-
-import (
-	"errors"
-	"log"
-	"net"
-	"net/http"
-	"strings"
-
-	"github.com/kodnaplakal/caddy/caddyhttp/httpserver"
-)
-
-func (gip GeoIP) ServeHTTP(w http.ResponseWriter, r *http.Request) (int, error) {
-	gip.addPlaceholders(w, r)
-	return gip.Next.ServeHTTP(w, r)
-}
-
-var record struct {
-	Country struct {
-		ISOCode           string            `maxminddb:"iso_code"`
-		IsInEuropeanUnion bool              `maxminddb:"is_in_european_union"`
-		Names             map[string]string `maxminddb:"names"`
-	} `maxminddb:"country"`
-
-	City struct {
-		Names map[string]string `maxminddb:"names"`
-	} `maxminddb:"city"`
-
-	Location struct {
-		Latitude  float64 `maxminddb:"latitude"`
-		Longitude float64 `maxminddb:"longitude"`
-		TimeZone  string  `maxminddb:"time_zone"`
-	} `maxminddb:"location"`
-}
-
-func (gip GeoIP) addPlaceholders(w http.ResponseWriter, r *http.Request) {
-	clientIP, _ := getClientIP(r, true)
-
-	err := gip.DBHandler.Lookup(clientIP, &record)
-	if err != nil {
-		log.Println(err)
-	}
-
-	if rr, ok := w.(*httpserver.ResponseRecorder); ok && rr.Replacer != nil {
-		rr.Replacer.Set("geoip_country_code", record.Country.ISOCode)
-	}
-
-	// repl.Set("sd", record.Country.ISOCode)
-
-	r.Header.Set("ssdsd", record.Country.ISOCode)
-	// r.Header.Set(gip.Config.HeaderNameCountryIsEU, strconv.FormatBool(record.Country.IsInEuropeanUnion))
-	// r.Header.Set(gip.Config.HeaderNameCountryName, record.Country.Names["en"])
-
-	// r.Header.Set(gip.Config.HeaderNameCityName, record.City.Names["en"])
-
-	// 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) {
-	var ip string
-
-	// Use the client ip from the 'X-Forwarded-For' header, if available.
-	if fwdFor := r.Header.Get("X-Forwarded-For"); fwdFor != "" && !strict {
-		ips := strings.Split(fwdFor, ", ")
-		ip = ips[0]
-	} else {
-		// Otherwise, get the client ip from the request remote address.
-		var err error
-		ip, _, err = net.SplitHostPort(r.RemoteAddr)
-		if err != nil {
-			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
-			}
-		}
-	}
-
-	// Parse the ip address string into a net.IP.
-	parsedIP := net.ParseIP(ip)
-	if parsedIP == nil {
-		return nil, errors.New("unable to parse address")
-	}
-
-	return parsedIP, nil
-}

+ 0 - 50
handler_test.go

@@ -1,50 +0,0 @@
-package geoip
-
-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("./test-data/GeoLite2-City.mmdb")
-	if err != nil {
-		t.Errorf("geoip: Can't open database: GeoLite2-City.mmdb")
-	}
-
-	config := Config{}
-
-	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)
-	}
-}

+ 84 - 1
setup.go

@@ -1,18 +1,43 @@
 package geoip
 
 import (
+	"errors"
+	"log"
+	"net"
+	"net/http"
+	"strconv"
+	"strings"
+
 	"github.com/mholt/caddy"
 	"github.com/mholt/caddy/caddyhttp/httpserver"
 	maxminddb "github.com/oschwald/maxminddb-golang"
 )
 
-// GeoIP Comments me
+// GeoIP represents a middleware instance
 type GeoIP struct {
 	Next      httpserver.Handler
 	DBHandler *maxminddb.Reader
 	Config    Config
 }
 
+var record struct {
+	Country struct {
+		ISOCode           string            `maxminddb:"iso_code"`
+		IsInEuropeanUnion bool              `maxminddb:"is_in_european_union"`
+		Names             map[string]string `maxminddb:"names"`
+	} `maxminddb:"country"`
+
+	City struct {
+		Names map[string]string `maxminddb:"names"`
+	} `maxminddb:"city"`
+
+	Location struct {
+		Latitude  float64 `maxminddb:"latitude"`
+		Longitude float64 `maxminddb:"longitude"`
+		TimeZone  string  `maxminddb:"time_zone"`
+	} `maxminddb:"location"`
+}
+
 // Init initializes the plugin
 func init() {
 	caddy.RegisterPlugin("geoip", caddy.Plugin{
@@ -45,3 +70,61 @@ func setup(c *caddy.Controller) error {
 
 	return nil
 }
+
+func (gip GeoIP) ServeHTTP(w http.ResponseWriter, r *http.Request) (int, error) {
+	gip.lookupLocation(w, r)
+	return gip.Next.ServeHTTP(w, r)
+}
+
+func (gip GeoIP) lookupLocation(w http.ResponseWriter, r *http.Request) {
+	clientIP, _ := getClientIP(r, true)
+	replacer := newReplacer(r)
+
+	err := gip.DBHandler.Lookup(clientIP, &record)
+	if err != nil {
+		log.Println(err)
+	}
+	rr, ok := w.(*httpserver.ResponseRecorder)
+	log.Printf("%v %v", rr, ok)
+
+	replacer.Set("geoip_country_code", record.Country.ISOCode)
+	replacer.Set("geoip_country_name", record.Country.Names["en"])
+	replacer.Set("geoip_country_eu", strconv.FormatBool(record.Country.IsInEuropeanUnion))
+	replacer.Set("geoip_city_name", record.City.Names["en"])
+	replacer.Set("geoip_latitude", strconv.FormatFloat(record.Location.Latitude, 'f', 6, 64))
+	replacer.Set("geoip_longitude", strconv.FormatFloat(record.Location.Longitude, 'f', 6, 64))
+	replacer.Set("geoip_time_zone", record.Location.TimeZone)
+}
+
+func getClientIP(r *http.Request, strict bool) (net.IP, error) {
+	var ip string
+
+	// Use the client ip from the 'X-Forwarded-For' header, if available.
+	if fwdFor := r.Header.Get("X-Forwarded-For"); fwdFor != "" && !strict {
+		ips := strings.Split(fwdFor, ", ")
+		ip = ips[0]
+	} else {
+		// Otherwise, get the client ip from the request remote address.
+		var err error
+		ip, _, err = net.SplitHostPort(r.RemoteAddr)
+		if err != nil {
+			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
+			}
+		}
+	}
+	// Parse the ip address string into a net.IP.
+	parsedIP := net.ParseIP(ip)
+	if parsedIP == nil {
+		return nil, errors.New("unable to parse address")
+	}
+
+	return parsedIP, nil
+}
+
+func newReplacer(r *http.Request) httpserver.Replacer {
+	return httpserver.NewReplacer(r, nil, "")
+}

+ 71 - 0
setup_test.go

@@ -0,0 +1,71 @@
+package geoip
+
+import (
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"testing"
+
+	"github.com/mholt/caddy/caddyhttp/httpserver"
+	maxminddb "github.com/oschwald/maxminddb-golang"
+)
+
+type testResponseRecorder struct {
+	*httpserver.ResponseWriterWrapper
+}
+
+func (testResponseRecorder) CloseNotify() <-chan bool { return nil }
+
+func TestReplacers(t *testing.T) {
+	dbhandler, err := maxminddb.Open("./test-data/GeoLite2-City.mmdb")
+	if err != nil {
+		t.Errorf("geoip: Can't open database: GeoLite2-City.mmdb")
+	}
+
+	config := Config{}
+
+	l := GeoIP{
+		Next: httpserver.HandlerFunc(func(w http.ResponseWriter, r *http.Request) (int, error) {
+			return 0, nil
+		}),
+		DBHandler: dbhandler,
+		Config:    config,
+	}
+
+	r := httptest.NewRequest("GET", "/", strings.NewReader(""))
+	r.RemoteAddr = "212.50.99.193"
+	rr := httpserver.NewResponseRecorder(testResponseRecorder{
+		ResponseWriterWrapper: &httpserver.ResponseWriterWrapper{ResponseWriter: httptest.NewRecorder()},
+	})
+	rr.Replacer = httpserver.NewReplacer(r, rr, "-")
+
+	l.ServeHTTP(rr, r)
+
+	if got, want := rr.Replacer.Replace("{geoip_country_code}"), "CY"; got != want {
+		t.Errorf("Expected custom placeholder {geoip_country_code} to be set (%s), but it wasn't; got: %s", want, got)
+	}
+
+	if got, want := rr.Replacer.Replace("{geoip_country_name}"), "Cyprus"; got != want {
+		t.Errorf("Expected custom placeholder {geoip_country_name} to be set (%s), but it wasn't; got: %s", want, got)
+	}
+
+	if got, want := rr.Replacer.Replace("{geoip_country_eu}"), "false"; got != want {
+		t.Errorf("Expected custom placeholder {geoip_country_eu} to be set (%s), but it wasn't; got: %s", want, got)
+	}
+
+	if got, want := rr.Replacer.Replace("{geoip_city_name}"), "Limassol"; got != want {
+		t.Errorf("Expected custom placeholder {geoip_city_name} to be set (%s), but it wasn't; got: %s", want, got)
+	}
+
+	if got, want := rr.Replacer.Replace("{geoip_latitude}"), "34.684100"; got != want {
+		t.Errorf("Expected custom placeholder {geoip_latitude} to be set (%s), but it wasn't; got: %s", want, got)
+	}
+
+	if got, want := rr.Replacer.Replace("{geoip_longitude}"), "33.037900"; got != want {
+		t.Errorf("Expected custom placeholder {geoip_longitude} to be set (%s), but it wasn't; got: %s", want, got)
+	}
+
+	if got, want := rr.Replacer.Replace("{geoip_time_zone}"), "Asia/Nicosia"; got != want {
+		t.Errorf("Expected custom placeholder {geoip_time_zone} to be set (%s), but it wasn't; got: %s", want, got)
+	}
+}