Extract iputil package

cmd/ipd/main.go

@@ -6,6 +6,8 @@ import (
 	"os"
 
 	"github.com/mpolden/ipd/http"
+	"github.com/mpolden/ipd/iputil"
+	"github.com/mpolden/ipd/iputil/db"
 	"github.com/sirupsen/logrus"
 )
 
@@ -32,32 +34,28 @@ func main() {
 	}
 	log.Level = level
 
-	oracle := http.NewOracle()
+	ipDb := db.Empty()
+	if opts.CountryDBPath != "" || opts.CityDBPath != "" {
+		ipDb, err = db.Open(opts.CountryDBPath, opts.CityDBPath)
+		if err != nil {
+			log.Fatal(err)
+		}
+	}
+	var lookupAddr http.LookupAddr
+	var lookupPort http.LookupPort
 	if opts.ReverseLookup {
 		log.Println("Enabling reverse lookup")
-		oracle.EnableLookupAddr()
+		lookupAddr = iputil.LookupAddr
 	}
 	if opts.PortLookup {
 		log.Println("Enabling port lookup")
-		oracle.EnableLookupPort()
+		lookupPort = iputil.LookupPort
-	}
-	if opts.CountryDBPath != "" {
-		log.Printf("Enabling country lookup (using database: %s)", opts.CountryDBPath)
-		if err := oracle.EnableLookupCountry(opts.CountryDBPath); err != nil {
-			log.Fatal(err)
-		}
-	}
-	if opts.CityDBPath != "" {
-		log.Printf("Enabling city lookup (using database: %s)", opts.CityDBPath)
-		if err := oracle.EnableLookupCity(opts.CityDBPath); err != nil {
-			log.Fatal(err)
-		}
 	}
 	if opts.IPHeader != "" {
 		log.Printf("Trusting header %s to contain correct remote IP", opts.IPHeader)
 	}
 
-	server := http.New(oracle, log)
+	server := http.New(ipDb, lookupAddr, lookupPort, log)
 	server.Template = opts.Template
 	server.IPHeader = opts.IPHeader
 

http/http.go

@@ -5,6 +5,8 @@ import (
 	"fmt"
 	"html/template"
 
+	"github.com/mpolden/ipd/iputil"
+	"github.com/mpolden/ipd/iputil/db"
 	"github.com/mpolden/ipd/useragent"
 	"github.com/sirupsen/logrus"
 
@@ -23,11 +25,16 @@ const (
 	textMediaType = "text/plain"
 )
 
+type LookupAddr func(net.IP) ([]string, error)
+type LookupPort func(net.IP, uint64) error
+
 type Server struct {
-	Template string
+	Template   string
-	IPHeader string
+	IPHeader   string
-	oracle   Oracle
+	lookupAddr LookupAddr
-	log      *logrus.Logger
+	lookupPort LookupPort
+	db         db.Database
+	log        *logrus.Logger
 }
 
 type Response struct {
@@ -45,18 +52,8 @@ type PortResponse struct {
 	Reachable bool   `json:"reachable"`
 }
 
-func New(oracle Oracle, logger *logrus.Logger) *Server {
+func New(db db.Database, lookupAddr LookupAddr, lookupPort LookupPort, logger *logrus.Logger) *Server {
-	return &Server{oracle: oracle, log: logger}
+	return &Server{lookupAddr: lookupAddr, lookupPort: lookupPort, db: db, log: logger}
-}
-
-func ipToDecimal(ip net.IP) *big.Int {
-	i := big.NewInt(0)
-	if to4 := ip.To4(); to4 != nil {
-		i.SetBytes(to4)
-	} else {
-		i.SetBytes(ip)
-	}
-	return i
 }
 
 func ipFromRequest(header string, r *http.Request) (net.IP, error) {
@@ -80,28 +77,24 @@ func (s *Server) newResponse(r *http.Request) (Response, error) {
 	if err != nil {
 		return Response{}, err
 	}
-	ipDecimal := ipToDecimal(ip)
+	ipDecimal := iputil.ToDecimal(ip)
-	country, err := s.oracle.LookupCountry(ip)
+	country, err := s.db.Country(ip)
 	if err != nil {
 		s.log.Debug(err)
 	}
-	countryISO, err := s.oracle.LookupCountryISO(ip)
+	city, err := s.db.City(ip)
 	if err != nil {
 		s.log.Debug(err)
 	}
-	city, err := s.oracle.LookupCity(ip)
+	hostnames, err := s.lookupAddr(ip)
-	if err != nil {
-		s.log.Debug(err)
-	}
-	hostnames, err := s.oracle.LookupAddr(ip)
 	if err != nil {
 		s.log.Debug(err)
 	}
 	return Response{
 		IP:         ip,
 		IPDecimal:  ipDecimal,
-		Country:    country,
+		Country:    country.Name,
-		CountryISO: countryISO,
+		CountryISO: country.ISO,
 		City:       city,
 		Hostname:   strings.Join(hostnames, " "),
 	}, nil
@@ -120,7 +113,7 @@ func (s *Server) newPortResponse(r *http.Request) (PortResponse, error) {
 	if err != nil {
 		return PortResponse{Port: port}, err
 	}
-	err = s.oracle.LookupPort(ip, port)
+	err = s.lookupPort(ip, port)
 	return PortResponse{
 		IP:        ip,
 		Port:      port,
@@ -201,11 +194,23 @@ func (s *Server) DefaultHandler(w http.ResponseWriter, r *http.Request) *appErro
 	if err != nil {
 		return internalServerError(err)
 	}
+	json, err := json.MarshalIndent(response, "", "  ")
+	if err != nil {
+		return internalServerError(err)
+	}
 	var data = struct {
-		Host string
 		Response
-		Oracle
+		Host string
-	}{r.Host, response, s.oracle}
+		JSON string
+		Port bool
+		Map  bool
+	}{
+		response,
+		r.Host,
+		string(json),
+		s.lookupPort != nil,
+		response.Country != "" && response.City != "",
+	}
 	if err := t.Execute(w, &data); err != nil {
 		return internalServerError(err)
 	}

http/http_test.go

@@ -3,27 +3,25 @@ package http
 import (
 	"io/ioutil"
 	"log"
-	"math/big"
 	"net"
 	"net/http"
 	"net/http/httptest"
 	"testing"
+
+	"github.com/mpolden/ipd/iputil/db"
 )
 
-type mockOracle struct{}
+type database struct{}
 
-func (r *mockOracle) LookupAddr(net.IP) ([]string, error)     { return []string{"localhost"}, nil }
+func lookupAddr(net.IP) ([]string, error) { return []string{"localhost"}, nil }
-func (r *mockOracle) LookupCountry(net.IP) (string, error)    { return "Elbonia", nil }
+func lookupPort(net.IP, uint64) error     { return nil }
-func (r *mockOracle) LookupCountryISO(net.IP) (string, error) { return "EB", nil }
+func (d *database) Country(net.IP) (db.Country, error) {
-func (r *mockOracle) LookupCity(net.IP) (string, error)       { return "Bornyasherk", nil }
+	return db.Country{Name: "Elbonia", ISO: "EB"}, nil
-func (r *mockOracle) LookupPort(net.IP, uint64) error         { return nil }
+}
-func (r *mockOracle) IsLookupAddrEnabled() bool               { return true }
+func (d *database) City(net.IP) (string, error) { return "Bornyasherk", nil }
-func (r *mockOracle) IsLookupCountryEnabled() bool            { return true }
-func (r *mockOracle) IsLookupCityEnabled() bool               { return true }
-func (r *mockOracle) IsLookupPortEnabled() bool               { return true }
 
 func newTestAPI() *Server {
-	return &Server{oracle: &mockOracle{}}
+	return &Server{db: &database{}, lookupAddr: lookupAddr, lookupPort: lookupPort}
 }
 
 func httpGet(url string, acceptMediaType string, userAgent string) (string, int, error) {
@@ -168,19 +166,3 @@ func TestCLIMatcher(t *testing.T) {
 		}
 	}
 }
-
-func TestIPToDecimal(t *testing.T) {
-	var tests = []struct {
-		in  string
-		out *big.Int
-	}{
-		{"127.0.0.1", big.NewInt(2130706433)},
-		{"::1", big.NewInt(1)},
-	}
-	for _, tt := range tests {
-		i := ipToDecimal(net.ParseIP(tt.in))
-		if i.Cmp(tt.out) != 0 {
-			t.Errorf("Expected %d, got %d for IP %s", tt.out, i, tt.in)
-		}
-	}
-}

http/oracle.go

@@ -1,163 +0,0 @@
-package http
-
-import (
-	"fmt"
-	"net"
-	"strings"
-	"time"
-
-	"github.com/oschwald/geoip2-golang"
-)
-
-type Oracle interface {
-	LookupAddr(net.IP) ([]string, error)
-	LookupCountry(net.IP) (string, error)
-	LookupCountryISO(net.IP) (string, error)
-	LookupCity(net.IP) (string, error)
-	LookupPort(net.IP, uint64) error
-	IsLookupAddrEnabled() bool
-	IsLookupCountryEnabled() bool
-	IsLookupCityEnabled() bool
-	IsLookupPortEnabled() bool
-}
-
-type DefaultOracle struct {
-	lookupAddr           func(net.IP) ([]string, error)
-	lookupCountry        func(net.IP) (string, error)
-	lookupCountryISO     func(net.IP) (string, error)
-	lookupCity           func(net.IP) (string, error)
-	lookupPort           func(net.IP, uint64) error
-	lookupAddrEnabled    bool
-	lookupCountryEnabled bool
-	lookupCityEnabled    bool
-	lookupPortEnabled    bool
-}
-
-func NewOracle() *DefaultOracle {
-	return &DefaultOracle{
-		lookupAddr:       func(net.IP) ([]string, error) { return nil, nil },
-		lookupCountry:    func(net.IP) (string, error) { return "", nil },
-		lookupCountryISO: func(net.IP) (string, error) { return "", nil },
-		lookupCity:       func(net.IP) (string, error) { return "", nil },
-		lookupPort:       func(net.IP, uint64) error { return nil },
-	}
-}
-
-func (r *DefaultOracle) LookupAddr(ip net.IP) ([]string, error) {
-	return r.lookupAddr(ip)
-}
-
-func (r *DefaultOracle) LookupCountry(ip net.IP) (string, error) {
-	return r.lookupCountry(ip)
-}
-
-func (r *DefaultOracle) LookupCountryISO(ip net.IP) (string, error) {
-	return r.lookupCountryISO(ip)
-}
-
-func (r *DefaultOracle) LookupCity(ip net.IP) (string, error) {
-	return r.lookupCity(ip)
-}
-
-func (r *DefaultOracle) LookupPort(ip net.IP, port uint64) error {
-	return r.lookupPort(ip, port)
-}
-
-func (r *DefaultOracle) EnableLookupAddr() {
-	r.lookupAddr = lookupAddr
-	r.lookupAddrEnabled = true
-}
-
-func (r *DefaultOracle) EnableLookupCountry(filepath string) error {
-	db, err := geoip2.Open(filepath)
-	if err != nil {
-		return err
-	}
-	r.lookupCountry = func(ip net.IP) (string, error) {
-		return lookupCountry(db, ip)
-	}
-	r.lookupCountryISO = func(ip net.IP) (string, error) {
-		return lookupCountryISO(db, ip)
-	}
-	r.lookupCountryEnabled = true
-	return nil
-}
-
-func (r *DefaultOracle) EnableLookupCity(filepath string) error {
-	db, err := geoip2.Open(filepath)
-	if err != nil {
-		return err
-	}
-	r.lookupCity = func(ip net.IP) (string, error) {
-		return lookupCity(db, ip)
-	}
-	r.lookupCityEnabled = true
-	return nil
-}
-
-func (r *DefaultOracle) EnableLookupPort() {
-	r.lookupPort = lookupPort
-	r.lookupPortEnabled = true
-}
-
-func (r *DefaultOracle) IsLookupAddrEnabled() bool    { return r.lookupAddrEnabled }
-func (r *DefaultOracle) IsLookupCountryEnabled() bool { return r.lookupCountryEnabled }
-func (r *DefaultOracle) IsLookupCityEnabled() bool    { return r.lookupCityEnabled }
-func (r *DefaultOracle) IsLookupPortEnabled() bool    { return r.lookupPortEnabled }
-
-func lookupAddr(ip net.IP) ([]string, error) {
-	names, err := net.LookupAddr(ip.String())
-	for i, _ := range names {
-		names[i] = strings.TrimRight(names[i], ".") // Always return unrooted name
-	}
-	return names, err
-}
-
-func lookupPort(ip net.IP, port uint64) error {
-	address := fmt.Sprintf("[%s]:%d", ip, port)
-	conn, err := net.DialTimeout("tcp", address, 2*time.Second)
-	if err != nil {
-		return err
-	}
-	defer conn.Close()
-	return nil
-}
-
-func lookupCountry(db *geoip2.Reader, ip net.IP) (string, error) {
-	record, err := db.Country(ip)
-	if err != nil {
-		return "", err
-	}
-	if country, exists := record.Country.Names["en"]; exists {
-		return country, nil
-	}
-	if country, exists := record.RegisteredCountry.Names["en"]; exists {
-		return country, nil
-	}
-	return "Unknown", fmt.Errorf("could not determine country for IP: %s", ip)
-}
-
-func lookupCountryISO(db *geoip2.Reader, ip net.IP) (string, error) {
-	record, err := db.City(ip)
-	if err != nil {
-		return "", err
-	}
-	if record.Country.IsoCode != "" {
-		return record.Country.IsoCode, nil
-	}
-	if record.RegisteredCountry.IsoCode != "" {
-		return record.RegisteredCountry.IsoCode, nil
-	}
-	return "Unknown", fmt.Errorf("could not determine country ISO Code for IP: %s", ip)
-}
-
-func lookupCity(db *geoip2.Reader, ip net.IP) (string, error) {
-	record, err := db.City(ip)
-	if err != nil {
-		return "", err
-	}
-	if city, exists := record.City.Names["en"]; exists {
-		return city, nil
-	}
-	return "Unknown", fmt.Errorf("could not determine city for IP: %s", ip)
-}

index.html

@@ -5,9 +5,9 @@
     <title>What is my IP address? &mdash; {{ .Host }}</title>
     <meta name="viewport" content="width=device-width, initial-scale=1.0">
     <meta name="description" content="What is my IP address?">
-    <link href="//fonts.googleapis.com/css?family=Open+Sans" rel="stylesheet">
+    <link href="https://fonts.googleapis.com/css?family=Open+Sans" rel="stylesheet">
-    <link rel="stylesheet" href="//cdnjs.cloudflare.com/ajax/libs/pure/0.6.2/pure-min.css">
+    <link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/pure/0.6.2/pure-min.css">
-    <link rel="stylesheet" href="//cdnjs.cloudflare.com/ajax/libs/pure/0.6.2/grids-responsive-min.css">
+    <link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/pure/0.6.2/grids-responsive-min.css">
     <style>
       html, .pure-g [class *= "pure-u"] {
         font-family: "Open Sans", sans-serif;
@@ -58,7 +58,7 @@ $ fetch -qo- https://{{ .Host }}
 
 $ bat -print=b {{ .Host }}/ip
 {{ .IP }}</pre>
-{{ if .IsLookupCountryEnabled }}
+{{ if .Country }}
         <h2>Country lookup</h2>
         <pre>
 $ http {{ .Host }}/country
@@ -67,7 +67,7 @@ $ http {{ .Host }}/country
 $ http {{ .Host }}/country-iso
 {{ .CountryISO }}</pre>
 {{ end }}
-{{ if .IsLookupCityEnabled }}
+{{ if .City }}
         <h2>City lookup</h2>
         <pre>
 $ http {{ .Host }}/city
@@ -78,22 +78,14 @@ $ http {{ .Host }}/city
         <h2>JSON output</h2>
         <pre>
 $ http {{ .Host }}/json
-{ {{ if .IsLookupCityEnabled }}
+{{ .JSON }}</pre>
-  "city": "{{ .City }}",
-  "country": "{{ .Country }}",
-  "country_iso": "{{ .CountryISO }}",{{ end }}{{ if .IsLookupAddrEnabled }}
-  "hostname": "{{ .Hostname }}",{{ end }}
-  "ip": "{{ .IP }}",
-  "ip_decimal": {{ .IPDecimal }}
-}</pre>
-
         <p>Setting the Accept header to application/json also works.</p>
         <h2>Plain output</h2>
         <p>Always returns the IP address including a trailing newline, regardless of user agent.</p>
         <pre>
 $ http {{ .Host }}/ip
 {{ .IP }}</pre>
-{{ if .IsLookupPortEnabled }}
+{{ if .Port }}
         <h2>Port testing</h2>
         <pre>
 $ http {{ .Host }}/port/8080
@@ -104,13 +96,13 @@ $ http {{ .Host }}/port/8080
 }</pre>
 {{ end }}
       </div>
-{{ if .IsLookupCountryEnabled }}{{ if ne .Country "Unknown"}}{{ if .IsLookupCityEnabled }}{{ if ne .City "Unknown"}}
+{{ if .Map }}
       <div class="pure-u-1 pure-u-md-1-2">
         <h2>Map</h2>
         <p><img src="https://maps.googleapis.com/maps/api/staticmap?size=500x166&amp;scale=2&amp;markers=color%3Aorange%7Clabel%3AS%7c{{ .City }},{{ .Country }}" width="500" height="166"/>
         </p>
-{{ end }}{{ end }}{{ end }}{{ end }}
       </div>
+{{ end }}
       <div class="pure-u-1 pure-u-md-1-2">
         <h2>FAQ</h2>
         <h3>How do I force IPv4 or IPv6 lookup?</h3>

iputil/db/db.go

@@ -0,0 +1,89 @@
+package db
+
+import (
+	"net"
+
+	geoip2 "github.com/oschwald/geoip2-golang"
+)
+
+type Database interface {
+	Country(net.IP) (Country, error)
+	City(net.IP) (string, error)
+}
+
+type Country struct {
+	Name string
+	ISO  string
+}
+
+type geoip struct {
+	country *geoip2.Reader
+	city    *geoip2.Reader
+}
+
+type empty struct{}
+
+func (d *empty) Country(ip net.IP) (Country, error) { return Country{}, nil }
+func (d *empty) City(ip net.IP) (string, error)     { return "", nil }
+
+func Empty() Database { return &empty{} }
+
+func Open(countryDB, cityDB string) (Database, error) {
+	var (
+		country *geoip2.Reader
+		city    *geoip2.Reader
+	)
+	if countryDB != "" {
+		r, err := geoip2.Open(countryDB)
+		if err != nil {
+			return nil, err
+		}
+		country = r
+	}
+	if cityDB != "" {
+		r, err := geoip2.Open(cityDB)
+		if err != nil {
+			return nil, err
+		}
+		city = r
+	}
+	return &geoip{country: country, city: city}, nil
+}
+
+func (g *geoip) Country(ip net.IP) (Country, error) {
+	country := Country{}
+	if g.country == nil {
+		return country, nil
+	}
+	record, err := g.country.Country(ip)
+	if err != nil {
+		return country, err
+	}
+	if c, exists := record.Country.Names["en"]; exists {
+		country.Name = c
+	}
+	if c, exists := record.RegisteredCountry.Names["en"]; exists && country.Name == "" {
+		country.Name = c
+	}
+	if record.Country.IsoCode != "" {
+		country.ISO = record.Country.IsoCode
+	}
+	if record.RegisteredCountry.IsoCode != "" && country.ISO == "" {
+		country.ISO = record.RegisteredCountry.IsoCode
+	}
+	return country, nil
+}
+
+func (g *geoip) City(ip net.IP) (string, error) {
+	if g.city == nil {
+		return "", nil
+	}
+	record, err := g.city.City(ip)
+	if err != nil {
+		return "", err
+	}
+	if city, exists := record.City.Names["en"]; exists {
+		return city, nil
+	}
+	return "", nil
+}

iputil/iputil.go

@@ -0,0 +1,37 @@
+package iputil
+
+import (
+	"fmt"
+	"math/big"
+	"net"
+	"strings"
+	"time"
+)
+
+func LookupAddr(ip net.IP) ([]string, error) {
+	names, err := net.LookupAddr(ip.String())
+	for i, _ := range names {
+		names[i] = strings.TrimRight(names[i], ".") // Always return unrooted name
+	}
+	return names, err
+}
+
+func LookupPort(ip net.IP, port uint64) error {
+	address := fmt.Sprintf("[%s]:%d", ip, port)
+	conn, err := net.DialTimeout("tcp", address, 2*time.Second)
+	if err != nil {
+		return err
+	}
+	defer conn.Close()
+	return nil
+}
+
+func ToDecimal(ip net.IP) *big.Int {
+	i := big.NewInt(0)
+	if to4 := ip.To4(); to4 != nil {
+		i.SetBytes(to4)
+	} else {
+		i.SetBytes(ip)
+	}
+	return i
+}

iputil/iputil_test.go

@@ -0,0 +1,23 @@
+package iputil
+
+import (
+	"math/big"
+	"net"
+	"testing"
+)
+
+func TestToDecimal(t *testing.T) {
+	var tests = []struct {
+		in  string
+		out *big.Int
+	}{
+		{"127.0.0.1", big.NewInt(2130706433)},
+		{"::1", big.NewInt(1)},
+	}
+	for _, tt := range tests {
+		i := ToDecimal(net.ParseIP(tt.in))
+		if i.Cmp(tt.out) != 0 {
+			t.Errorf("Expected %d, got %d for IP %s", tt.out, i, tt.in)
+		}
+	}
+}