diff --git a/client.go b/client.go index a0c0533..8cff536 100644 --- a/client.go +++ b/client.go @@ -1,73 +1,15 @@ package main -// echo -n "test out the server" | nc localhost 3333 - import ( - "encoding/json" - "io/ioutil" "math/rand" "net" - "net/http" - "strings" "sync/atomic" "time" "github.com/golang/glog" - "github.com/pierrre/geohash" "github.com/prometheus/client_golang/prometheus" ) -type geoIP struct { - Ip string `json:""` - 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:"zipcode"` - Latitude float64 `json:"latitude"` - Longitude float64 `json:"longitude"` - MetroCode int `json:"metro_code"` - AreaCode int `json:"area_code"` -} - -func getGeohashAndLocation(address string) (string, string, string, error) { - var geo geoIP - response, err := http.Get("https://freegeoip.live/json/" + address) - if err != nil { - return "s000", "Unknown", "Unknown", err - } - defer response.Body.Close() - - body, err := ioutil.ReadAll(response.Body) - if err != nil { - return "s000", "Unknown", "Unknown", err - } - - err = json.Unmarshal(body, &geo) - if err != nil { - return "s000", "Unknown", "Unknown", err - } - - var locations []string - for _, s := range []string{geo.CountryName, geo.RegionName, geo.City} { - if strings.TrimSpace(s) != "" { - locations = append(locations, s) - } - } - location := strings.Join(locations, ", ") - if location == "" { - location = "Unknown" - } - country := geo.CountryName - if country == "" { - country = "Unknown" - } - gh := geohash.EncodeAuto(geo.Latitude, geo.Longitude) - - return gh, country, location, nil -} - var letterBytes = []byte(" abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ01234567890!@#$%^&*()-=_+[]{}|;:',./<>?") func randStringBytes(n int64) []byte { @@ -80,22 +22,23 @@ func randStringBytes(n int64) []byte { } type client struct { - conn net.Conn - last time.Time - next time.Time - start time.Time - interval time.Duration - geohash string - country string - location string - bytes_sent int + conn net.Conn + last time.Time + next time.Time + start time.Time + interval time.Duration + geoipSupplier string + geohash string + country string + location string + bytes_sent int } -func NewClient(conn net.Conn, interval time.Duration, maxClient int64) *client { +func NewClient(conn net.Conn, interval time.Duration, maxClient int64, geoipSupplier string) *client { addr := conn.RemoteAddr().(*net.TCPAddr) atomic.AddInt64(&numCurrentClients, 1) atomic.AddInt64(&numTotalClients, 1) - geohash, country, location, err := getGeohashAndLocation(addr.IP.String()) + geohash, country, location, err := geohashAndLocation(addr.IP.String(), geoipSupplier) if err != nil { glog.Warningf("Failed to obatin the geohash of %v: %v.", addr.IP, err) } diff --git a/geoip.go b/geoip.go new file mode 100644 index 0000000..db5e561 --- /dev/null +++ b/geoip.go @@ -0,0 +1,128 @@ +package main + +import ( + "encoding/json" + "fmt" + "io/ioutil" + "net/http" + "strings" + + "github.com/pierrre/geohash" +) + +type freegeoip struct { + 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:"zipcode"` + Latitude float64 `json:"latitude"` + Longitude float64 `json:"longitude"` + MetroCode int `json:"metro_code"` + AreaCode int `json:"area_code"` +} + +func geohashAndLocationFromFreegeoip(address string) (string, string, string, error) { + var geo freegeoip + response, err := http.Get("https://freegeoip.live/json/" + address) + if err != nil { + return "s000", "Unknown", "Unknown", err + } + defer response.Body.Close() + + body, err := ioutil.ReadAll(response.Body) + if err != nil { + return "s000", "Unknown", "Unknown", err + } + + err = json.Unmarshal(body, &geo) + if err != nil { + return "s000", "Unknown", "Unknown", err + } + + var locations []string + for _, s := range []string{geo.CountryName, geo.RegionName, geo.City} { + if strings.TrimSpace(s) != "" { + locations = append(locations, s) + } + } + location := strings.Join(locations, ", ") + if location == "" { + location = "Unknown" + } + country := geo.CountryName + if country == "" { + country = "Unknown" + } + gh := geohash.EncodeAuto(geo.Latitude, geo.Longitude) + + return gh, country, location, nil +} + +type ipapi struct { + Status string `json:"status"` + Message string `json:"message"` + Ip string `json:"query"` + CountryCode string `json:"countryCode"` + CountryName string `json:"country"` + RegionCode string `json:"region"` + RegionName string `json:"regionName"` + City string `json:"city"` + Zipcode string `json:"zip"` + Latitude float64 `json:"lat"` + Longitude float64 `json:"lon"` +} + +func geohashAndLocationFromIpapi(address string) (string, string, string, error) { + var geo ipapi + response, err := http.Get("http://ip-api.com/json/" + address) + if err != nil { + return "s000", "Unknown", "Unknown", err + } + defer response.Body.Close() + + body, err := ioutil.ReadAll(response.Body) + if err != nil { + return "s000", "Unknown", "Unknown", err + } + + err = json.Unmarshal(body, &geo) + if err != nil { + return "s000", "Unknown", "Unknown", err + } + + if geo.Status != "success" { + return "s000", "Unknown", "Unknown", fmt.Errorf("failed to query %v via ip-api: status: %v, message: %v", address, geo.Status, geo.Message) + } + + var locations []string + for _, s := range []string{geo.CountryName, geo.RegionName, geo.City} { + if strings.TrimSpace(s) != "" { + locations = append(locations, s) + } + } + location := strings.Join(locations, ", ") + if location == "" { + location = "Unknown" + } + country := geo.CountryName + if country == "" { + country = "Unknown" + } + gh := geohash.EncodeAuto(geo.Latitude, geo.Longitude) + + return gh, country, location, nil +} + +func geohashAndLocation(address string, geoipSupplier string) (string, string, string, error) { + switch geoipSupplier { + case "ip-api": + return geohashAndLocationFromIpapi(address) + case "freegeoip": + return geohashAndLocationFromFreegeoip(address) + default: + return "s000", "Unknown", "Unknown", fmt.Errorf("unknown geoipSupplier %v.", geoipSupplier) + } +} diff --git a/main.go b/main.go index 2a40b2c..d1a6126 100644 --- a/main.go +++ b/main.go @@ -1,7 +1,5 @@ package main -// echo -n "test out the server" | nc localhost 3333 - import ( "flag" "fmt" @@ -100,6 +98,7 @@ func main() { enablePrometheus := flag.Bool("enable_prometheus", false, "Enable prometheus") prometheusPort := flag.String("prometheus_port", "2112", "The port for prometheus") prometheusEntry := flag.String("prometheus_entry", "metrics", "Entry point for prometheus") + geoipSupplier := flag.String("geoip_supplier", "ip-api", "Supplier to obtain Geohash of IPs. Possible values are \"ip-api\", \"freegeoip\"") flag.Usage = func() { fmt.Fprintf(flag.CommandLine.Output(), "Usage of %v \n", os.Args[0]) @@ -127,7 +126,7 @@ func main() { glog.Infof("Listening on %v:%v", *connHost, *connPort) clients := make(chan *client, *maxClients) - go func(clients chan *client, interval time.Duration, bannerMaxLength int64) { + go func() { for { c, more := <-clients if !more { @@ -136,15 +135,15 @@ func main() { if time.Now().Before(c.next) { time.Sleep(c.next.Sub(time.Now())) } - err := c.Send(bannerMaxLength) + err := c.Send(*bannerMaxLength) if err != nil { c.Close() continue } go func() { clients <- c }() } - }(clients, interval, *bannerMaxLength) - listener := func(clients chan *client, interval time.Duration, maxClients int64) { + }() + listener := func() { for { // Listen for an incoming connection. conn, err := l.Accept() @@ -153,11 +152,11 @@ func main() { os.Exit(1) } // Handle connections in a new goroutine. - for numCurrentClients >= maxClients { + for numCurrentClients >= *maxClients { time.Sleep(interval) } - clients <- NewClient(conn, interval, maxClients) + clients <- NewClient(conn, interval, *maxClients, *geoipSupplier) } } - listener(clients, interval, *maxClients) + listener() }