navigator/main.go

package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"log"
	"net"
	"net/http"
	_ "net/http/pprof"
	"strconv"
	"strings"

	"github.com/ipipdotnet/ipdb-go"
)

var (
	requiredFields = []string{
		"country_name",
		"region_name",
		"city_name",
		"isp_domain",
		"country_code",
		"continent_code",
	}
)

const (
	errIPv4Only = "Navigator works for IPv4 only :)"
)

type errorMessage struct {
	Error string `json:error`
}

func responseWithError(resp http.ResponseWriter, statusCode int, message string) {
	resp.Header().Set("Content-Type", "text/plain")
	resp.WriteHeader(statusCode)
	resp.Write([]byte("error: " + message))
}

func responseWithJsonError(resp http.ResponseWriter, statusCode int, message string) {
	resp.Header().Set("Content-Type", "application/json")
	body, err := json.Marshal(&errorMessage{
		Error: message,
	})
	resp.Header().Set("Content-Length", strconv.Itoa(len(body)))
	if err != nil {
		// This should never happen
		panic("json marshal failed, check code")
	}
	resp.WriteHeader(statusCode)
	resp.Write(body)
}

func buildLocation(info *ipdb.CityInfo) string {
	ret := ""
	if info.CountryName != "" {
		ret += info.CountryName + " "
	}
	if info.RegionName != "" {
		ret += info.RegionName + " "
	}
	if info.CityName != "" {
		ret += info.CityName + " "
	}
	return strings.TrimSpace(ret)
}

func main() {
	ipipdb := flag.String("ipdb-database", "ipip.db", "path to ipip database")
	flag.Parse()

	db, err := ipdb.NewCity(*ipipdb)
	if err != nil {
		log.Fatalln("Unable to open ipdb:", err.Error())
	}
	if !db.IsIPv4() {
		log.Fatalln("This IPIP.net database has no IPv4 information!")
	}
	{
		fields := db.Fields()
		for _, requiredField := range requiredFields {
			ok := false
			for _, field := range fields {
				if field == requiredField {
					ok = true
					break
				}
			}
			if !ok {
				log.Fatalln("This IPIP.net database has no required field", requiredField)
			}
		}
	}

	http.HandleFunc("/healthz", func(resp http.ResponseWriter, req *http.Request) {
		resp.WriteHeader(200)
		resp.Write([]byte("ok"))
	})

	http.HandleFunc("/info", func(resp http.ResponseWriter, req *http.Request) {
		host, _, err := net.SplitHostPort(req.RemoteAddr)
		if err != nil {
			responseWithError(resp, http.StatusPreconditionFailed, errIPv4Only)
			return
		}
		if net.ParseIP(host).To4() == nil {
			responseWithError(resp, http.StatusPreconditionFailed, errIPv4Only)
			return
		}

		info_cn, err := db.FindInfo(host, "CN")
		if err != nil {
			fmt.Fprintf(resp, "IP %s not found in the database.", host)
			return
		}
		info_en, err := db.FindInfo(host, "EN")
		if err != nil {
			fmt.Fprintf(resp, "IP %s not found in the database.", host)
			return
		}
		resp.Header().Set("Content-Type", "text/plain")
		resp.WriteHeader(http.StatusOK)
		fmt.Fprintln(resp, "您的IP:", host)
		fmt.Fprintln(resp, "数据库中IP所属位置:", buildLocation(info_cn))
		if info_cn.IspDomain != "" {
			fmt.Fprintln(resp, "数据库中IP所属运营商:", info_cn.IspDomain)
		}
		// TODO(jackyyf): add mapping info here

		fmt.Fprintln(resp, "Your IP:", host)
		fmt.Fprintln(resp, "Location for your IP according our database:", buildLocation(info_en))
		if info_en.IspDomain != "" {
			fmt.Fprintln(resp, "ISP for your IP according our database:", info_en.IspDomain)
		}
		// TODO(jackyyf): add mapping info here

	})
	http.ListenAndServe("0.0.0.0:8086", nil)
}