freegeoip web server using httpmux

freegeoip-httpmux.go

package main

import (
	"bytes"
	"encoding/csv"
	"encoding/json"
	"encoding/xml"
	"fmt"
	"log"
	"math/rand"
	"net"
	"net/http"
	"os"
	"strconv"
	"time"

	"github.com/fiorix/freegeoip"
	"golang.org/x/net/context"

	"github.com/go-web/httplog"
	"github.com/go-web/httpmux"
)

func init() {
	rand.Seed(time.Now().UnixNano())
}

var maxmindDB = "http://geolite.maxmind.com/download/geoip/database/GeoLite2-City.mmdb.gz"

func main() {
	db, err := freegeoip.OpenURL(maxmindDB, 24*time.Hour, time.Hour)
	if err != nil {
		log.Fatal(err)
	}
	f := &handler{db: db}
	mux := httpmux.New()
	l := log.New(os.Stderr, "", 0)
	mux.Use(httplog.ApacheCombinedFormat(l))
	mux.Use(f.QueryIP)
	mux.GET("/json/*host", f.JSON)
	mux.GET("/xml/*host", f.XML)
	mux.GET("/csv/*host", f.CSV)
	err = http.ListenAndServe(":8080", mux)
	if err != nil {
		log.Fatal(err)
	}
}

type handler struct {
	db *freegeoip.DB
}

func (f *handler) QueryIP(next http.HandlerFunc) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		host := httpmux.Params(r).ByName("host")
		switch len(host) {
		case 1:
			host, _, _ = net.SplitHostPort(r.RemoteAddr)
			if host == "" {
				host = r.RemoteAddr
			}
		default:
			host = host[1:]
		}
		ips, err := net.LookupIP(host)
		if err != nil || len(ips) == 0 {
			http.NotFound(w, r)
			return
		}
		q := &lookupQuery{}
		ip := ips[rand.Intn(len(ips))]
		err = f.db.Lookup(ip, q)
		if err != nil {
			http.Error(w, "Try again later.", http.StatusServiceUnavailable)
			return
		}
		resp := q.Record(ip, r.Header.Get("Accept-Language"))
		ctx := context.WithValue(httpmux.Context(r), "resp", resp)
		httpmux.SetContext(ctx, r)
		next(w, r)
	}
}

func (f *handler) JSON(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "application/json")
	resp := httpmux.Context(r).Value("resp")
	json.NewEncoder(w).Encode(resp)
}

func (f *handler) XML(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "application/xml")
	resp := httpmux.Context(r).Value("resp")
	x := xml.NewEncoder(w)
	x.Indent("", "\t")
	x.Encode(resp)
	w.Write([]byte{'\n'})
}

func (f *handler) CSV(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "text/csv")
	resp := httpmux.Context(r).Value("resp")
	fmt.Fprint(w, resp)
}

type lookupQuery struct {
	Country struct {
		ISOCode string            `maxminddb:"iso_code"`
		Names   map[string]string `maxminddb:"names"`
	} `maxminddb:"country"`
	Region []struct {
		ISOCode string            `maxminddb:"iso_code"`
		Names   map[string]string `maxminddb:"names"`
	} `maxminddb:"subdivisions"`
	City struct {
		Names map[string]string `maxminddb:"names"`
	} `maxminddb:"city"`
	Location struct {
		Latitude  float64 `maxminddb:"latitude"`
		Longitude float64 `maxminddb:"longitude"`
		MetroCode uint    `maxminddb:"metro_code"`
		TimeZone  string  `maxminddb:"time_zone"`
	} `maxminddb:"location"`
	Postal struct {
		Code string `maxminddb:"code"`
	} `maxminddb:"postal"`
}

func (lq *lookupQuery) Record(ip net.IP, lang string) *responseRecord {
	// TODO: parse accept-language value from lang.
	if lq.Country.Names[lang] == "" {
		lang = "en"
	}
	r := &responseRecord{
		IP:          ip.String(),
		CountryCode: lq.Country.ISOCode,
		CountryName: lq.Country.Names[lang],
		City:        lq.City.Names[lang],
		ZipCode:     lq.Postal.Code,
		TimeZone:    lq.Location.TimeZone,
		Latitude:    lq.Location.Latitude,
		Longitude:   lq.Location.Longitude,
		MetroCode:   lq.Location.MetroCode,
	}
	if len(lq.Region) > 0 {
		r.RegionCode = lq.Region[0].ISOCode
		r.RegionName = lq.Region[0].Names[lang]
	}
	return r
}

type responseRecord struct {
	XMLName     xml.Name `xml:"Response" json:"-"`
	IP          string   `json:"ip"`
	CountryCode string   `json:"country_code"`
	CountryName string   `json:"country_name"`
	RegionCode  string   `json:"region_code"`
	RegionName  string   `json:"region_name"`
	City        string   `json:"city"`
	ZipCode     string   `json:"zip_code"`
	TimeZone    string   `json:"time_zone"`
	Latitude    float64  `json:"latitude"`
	Longitude   float64  `json:"longitude"`
	MetroCode   uint     `json:"metro_code"`
}

func (rr *responseRecord) String() string {
	b := &bytes.Buffer{}
	w := csv.NewWriter(b)
	w.UseCRLF = true
	w.Write([]string{
		rr.IP,
		rr.CountryCode,
		rr.CountryName,
		rr.RegionCode,
		rr.RegionName,
		rr.City,
		rr.ZipCode,
		rr.TimeZone,
		strconv.FormatFloat(rr.Latitude, 'f', 2, 64),
		strconv.FormatFloat(rr.Longitude, 'f', 2, 64),
		strconv.Itoa(int(rr.MetroCode)),
	})
	w.Flush()
	return b.String()
}