navigator/main.go

package main

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

	"git.eve.moe/jackyyf/navigator/ipgeo"
	"git.eve.moe/jackyyf/navigator/mapping"
	"github.com/ipipdotnet/ipdb-go"
)

const (
	errIPv4Only      = "Navigator works for valid IPv4 only :)"
	remoteAddrHeader = "X-NAV-REMOTE-IP"
)

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

var (
	listen_spec = flag.String("bind", "127.0.0.1:8086", "http server bind spec")
)

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")
	resp.WriteHeader(statusCode)
	encoder := json.NewEncoder(resp)
	err := encoder.Encode(&errorMessage{
		Error: message,
	})
	if err != nil {
		// This should never happen
		panic("json marshal failed, check code")
	}
}

func getRemoteIP(req *http.Request) string {
	if addr := req.Header.Get(remoteAddrHeader); addr != "" {
		if net.ParseIP(addr).To4() == nil {
			return ""
		}
		return addr
	}
	host, _, err := net.SplitHostPort(req.RemoteAddr)
	if err != nil {
		return ""
	}
	if net.ParseIP(host).To4() == nil {
		return ""
	}
	return host
}

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() {
	flag.Parse()
	ipgeo.Initialize()
	mapping.Initialize()

	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) {
		var host string
		if argIp := req.FormValue("ip"); argIp != "" {
			host = argIp
			if net.ParseIP(host).To4() == nil {
				responseWithError(resp, http.StatusPreconditionFailed, errIPv4Only)
				return
			}
		} else {
			host = getRemoteIP(req)
		}
		db := ipgeo.Get()

		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)
		server := mapping.Get(host)
		fmt.Fprintln(resp, "您的IP:", host)
		fmt.Fprintln(resp, "数据库中IP所属位置:", buildLocation(info_cn))
		if info_cn.IspDomain != "" {
			fmt.Fprintln(resp, "数据库中IP所属运营商:", info_cn.IspDomain)
		}
		fmt.Fprintln(resp, "您被分配的CDN节点为:", server)
		fmt.Fprintln(resp)
		fmt.Fprintln(resp, strings.Repeat("=", 72))

		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)
		}
		fmt.Fprintln(resp, "Allocated CDN node for you:", server)
		fmt.Fprintln(resp)
	})

	http.HandleFunc("/mapping", func(resp http.ResponseWriter, req *http.Request) {
		host := getRemoteIP(req)
		resp.Header().Set("Content-Type", "text/plain")
		resp.WriteHeader(http.StatusOK)
		server := mapping.Get(host)
		fmt.Fprint(resp, server)
	})
	http.HandleFunc("/getMapping", func(resp http.ResponseWriter, req *http.Request) {
		ip := req.FormValue("ip")
		if net.ParseIP(ip).To4() == nil {
			responseWithError(resp, http.StatusBadRequest, errIPv4Only)
			return
		}
		resp.Header().Set("Content-Type", "text/plain")
		resp.WriteHeader(http.StatusOK)
		server := mapping.Get(ip)
		log.Printf("%s => %s\n", ip, server)
		fmt.Fprint(resp, server)
	})
	http.HandleFunc("/getNodes", func(resp http.ResponseWriter, req *http.Request) {
		ip := req.FormValue("ip")
		if net.ParseIP(ip).To4() == nil {
			responseWithError(resp, http.StatusBadRequest, errIPv4Only)
			return
		}
		nodes := mapping.GetNodes()
		if nodes == nil {
			responseWithJsonError(resp, http.StatusInternalServerError, "Unable to get nodes")
			return
		}
		suffix := mapping.GetSuffix(ip)
		resp.Header().Set("Content-Type", "application/json")
		resp.WriteHeader(http.StatusOK)
		jsonEncoder := json.NewEncoder(resp)
		ret := make([]string, 0, len(nodes))
		for _, node := range nodes {
			ret = append(ret, node+suffix)
		}
		jsonEncoder.Encode(ret)
	})
	clientApi := http.NewServeMux()
	http.Handle("/client/", http.StripPrefix("/client", clientApi))
	clientV1Api := http.NewServeMux()
	clientApi.Handle("/v1/", http.StripPrefix("/v1", clientV1Api))
	clientV1Api.HandleFunc("/getNodes", func(resp http.ResponseWriter, req *http.Request) {
		host := getRemoteIP(req)
		nodes := mapping.GetNodes()
		if nodes == nil {
			responseWithJsonError(resp, http.StatusInternalServerError, "Unable to get nodes")
			return
		}
		suffix := mapping.GetSuffix(host)
		resp.Header().Set("Content-Type", "application/json")
		resp.WriteHeader(http.StatusOK)
		jsonEncoder := json.NewEncoder(resp)
		jsonEncoder.Encode(map[string]interface{}{
			"nodes":  nodes,
			"suffix": suffix,
		})
	})
	log.Println("HTTP server is running on", *listen_spec)
	http.ListenAndServe(*listen_spec, nil)
}