GeoIP project to support lookups for mirror

geoip.go

@@ -0,0 +1,94 @@
+package geoip
+
+import (
+	"fmt"
+	"net"
+	"sync"
+	"time"
+
+	"github.com/IncSW/geoip2"
+)
+
+// GeoIPHandler will keep itself up-to-date with the latest GeoIP database by updating once every 24 hours
+// Provides methods to lookup IP addresses and return the associated latitude and longitude
+// The structure should be created with it's maxmind license key
+type GeoIPHandler struct {
+	sync.RWMutex
+
+	// The stop channel is used to end the goroutine that updates the database
+	stop chan struct{}
+	// Maxmind license key can be created at https://www.maxmind.com
+	licenseKey string
+	// underlying database object
+	db         *geoip2.CityReader
+}
+
+// NewGeoIPHandler creates a new GeoIPHandler with the given license key
+func NewGeoIPHandler(licenseKey string) (*GeoIPHandler, error) {
+	// Download the database
+	bytes, err := downloadAndCheckHash(licenseKey)
+	if err != nil {
+		return nil, err
+	}
+
+	// Create the database
+	db, err := geoip2.NewCityReader(bytes)
+	if err != nil {
+		return nil, err
+	}
+
+	// Create the handler
+	handler := &GeoIPHandler{
+		stop: 	 make(chan struct{}),
+		licenseKey: licenseKey,
+		db:         db,
+	}
+
+	// update the database every 24 hours
+	go handler.update(handler.stop)
+
+	return handler, nil
+}
+
+func (g *GeoIPHandler) update(stop chan struct{}) {
+	// update the database every 24 hours
+	for {
+		select {
+		case <-stop:
+			return
+		case <-time.After(24 * time.Hour):
+			// Lock the database
+			g.Lock()
+
+			// Download the database
+			bytes, err := downloadAndCheckHash(g.licenseKey)
+			if err != nil {
+				fmt.Println(err)
+				g.Unlock()
+				continue
+			}
+
+			// Create the database
+			db, err := geoip2.NewCityReader(bytes)
+
+			if err != nil {
+				fmt.Println(err)
+				g.Unlock()
+				continue
+			}
+
+			// Create the handler
+			g.db = db
+
+			g.Unlock()
+		}
+	}
+}
+
+func (g *GeoIPHandler) Close() {
+	g.stop <- struct{}{}
+}
+
+func (g *GeoIPHandler) Lookup(ip net.IP) (*geoip2.CityResult, error) {
+	return g.db.Lookup(ip)
+}

go.mod

@@ -0,0 +1,5 @@
+module github.com/COSI-Lab/Mirror/geoip
+
+go 1.18
+
+require github.com/IncSW/geoip2 v0.1.2

go.sum

@@ -0,0 +1,2 @@
+github.com/IncSW/geoip2 v0.1.2 h1:v7iAyDiNZjHES45P1JPM3SMvkw0VNeJtz0XSVxkRwOY=
+github.com/IncSW/geoip2 v0.1.2/go.mod h1:adcasR40vXiUBjtzdaTTKL/6wSf+fgO4M8Gve/XzPUk=

update.go

@@ -0,0 +1,112 @@
+package geoip
+
+import (
+	"archive/tar"
+	"bytes"
+	"compress/gzip"
+	"crypto/sha256"
+	"fmt"
+	"io"
+	"io/ioutil"
+	"net/http"
+	"strings"
+)
+
+// Creates a new geoip2 reader by downloading the database from MaxMind.
+// We also preform a sha256 check to ensure the database is not corrupt.
+
+const CHECKSUM_URL string = "https://download.maxmind.com/app/geoip_download?edition_id=GeoLite2-City&suffix=tar.gz.sha256&license_key="
+const DATABASE_URL string = "https://download.maxmind.com/app/geoip_download?edition_id=GeoLite2-City&suffix=tar.gz&license_key="
+
+// Uses the MaxMind permalink to download the most recent sha256 checksum of the database.
+func downloadHash(licenseKey string) (string, error) {
+	resp, err := http.Get(CHECKSUM_URL + licenseKey)
+	if err != nil {
+		return "", err
+	}
+	defer resp.Body.Close()
+
+	// Read the response body
+	body, err := ioutil.ReadAll(resp.Body)
+	if err != nil {
+		return "", err
+	}
+
+	// Check the status code
+	if resp.StatusCode != 200 {
+		return "", fmt.Errorf("HTTP Status %d while trying to download the sha256 checksum", resp.StatusCode)
+	}
+
+	// Return the sha256 checksum
+	return string(body[:64]), nil
+}
+
+func downloadDatabase(licenseKey, checksum string) ([]byte, error) {
+	resp, err := http.Get(DATABASE_URL + licenseKey)
+	if err != nil {
+		return nil, err
+	}
+	defer resp.Body.Close()
+
+	// Check the status code
+	if resp.StatusCode != 200 {
+		return nil, fmt.Errorf("HTTP Status %d while trying to download the database", resp.StatusCode)
+	}
+
+	// Read the response body
+	body, err := ioutil.ReadAll(resp.Body)
+	if err != nil {
+		fmt.Println(err)
+		return nil, err
+	}
+
+	// Calculate the sha256 checksum of the tarball
+	calculatedHash := sha256.Sum256(body)
+	calculatedHashString := fmt.Sprintf("%x", calculatedHash)
+
+	// Check the checksum
+	if checksum != calculatedHashString {
+		return nil, fmt.Errorf("checksum mismatch. Expected %s, got %s", checksum, calculatedHashString)
+	}
+
+	// Here we have a tar.gz file. We need to extract it.
+	gzr, err := gzip.NewReader(bytes.NewReader(body))
+	if err != nil {
+		fmt.Println("Error creating gzip reader:", err)
+		return nil, err
+	}
+	defer gzr.Close()
+
+	// Read the files names of the things in the tar file
+	tarReader := tar.NewReader(gzr)
+	for {
+		header, err := tarReader.Next()
+
+		if err != nil {
+			if err == io.EOF {
+				break
+			}
+			return nil, err
+		}
+
+		// Name ends with "GeoLite2-City.mmdb"
+		if strings.HasSuffix(header.Name, "GeoLite2-City.mmdb") {
+			// We found the database file. Read it.
+			return ioutil.ReadAll(tarReader)
+		}
+	}
+
+	// Return the database
+	return nil, fmt.Errorf("database not found in the tarball")
+}
+
+func downloadAndCheckHash(licenseKey string) ([]byte, error) {
+	// Download the hash
+	hash, err := downloadHash(licenseKey)
+	if err != nil {
+		return nil, err
+	}
+
+	// Download the database
+	return downloadDatabase(licenseKey, hash)
+}

update_test.go

@@ -0,0 +1,88 @@
+package geoip
+
+import (
+	"encoding/hex"
+	"net"
+	"os"
+	"testing"
+
+	"github.com/IncSW/geoip2"
+)
+
+// GeoIP tests
+func TestGetSHA256(t *testing.T) {
+	// Get the license key through environment variables
+	licenseKey := os.Getenv("MAXMIND_LICENSE_KEY")
+	if licenseKey == "" {
+		t.Error("MAXMIND_LICENSE_KEY environment variable not set")
+		return
+	}
+
+	// Download the sha256 checksum
+	sha256, err := downloadHash(licenseKey)
+	if err != nil {
+		t.Error(err)
+	}
+
+	// the checksum should be a hex string that is 64 characters long
+	if len(sha256) != 64 {
+		t.Error("sha256 checksum is not 64 characters long")
+	}
+
+	// decode the sha256 checksum
+	_, err = hex.DecodeString(sha256)
+	if err != nil {
+		t.Error(err)
+	}
+}
+
+// Download a new database
+func TestDownloadDatabase(t *testing.T) {
+	// Get the license key through environment variables
+	licenseKey := os.Getenv("MAXMIND_LICENSE_KEY")
+	if licenseKey == "" {
+		t.Error("MAXMIND_LICENSE_KEY environment variable not set")
+		return
+	}
+	// Prepare the checksum
+	sha256, err := downloadHash(licenseKey)
+	if err != nil {
+		t.Error(err)
+	}
+
+	// Download the database
+	bytes, err := downloadDatabase(licenseKey, sha256)
+	if err != nil {
+		t.Error(err)
+	}
+
+	// Verify that the database can be opened
+	_, err = geoip2.NewCityReader(bytes)
+	if err != nil {
+		t.Error(err)
+	}
+}
+
+func TestLookups(t *testing.T) {
+	// Get the license key through environment variables
+	licenseKey := "O1MFhnxCY9aKUyio"
+
+	geoip, err := NewGeoIPHandler(licenseKey)
+	if err != nil {
+		t.Error(err)
+		return
+	}
+
+	// Lookup some IP addresses
+	ips := []string{"128.153.145.19", "2605:6480:c051:100::1"}
+	for _, ip := range ips {
+		_, err := geoip.Lookup(net.ParseIP(ip))
+		if err != nil {
+			t.Error(ip, err)
+		}
+	}
+
+	// TODO: Add a test that ensures maxmind knows where mirror is
+
+	geoip.Close()
+}