Skip to content

Commit

Permalink
Add Locations endpoints (#1516)
Browse files Browse the repository at this point in the history 
* add locations endpoints

* Add sqlite3 check and database generation in geolite script

* Add SQLite storage for geolocation data

* Refactor file existence check into a separate function

* Integrate geolocation services into management application

* Refactoring

* Refactor city retrieval to include Geonames ID

* Add signature verification for GeoLite2 database download

* Change to in-memory database for geolocation store

* Merge manager to geolocation

* Update GetAllCountries to return Country name and iso code

* fix tests

* Add reload to SqliteStore

* Add geoname indexes

* move db file check to connectDB

* Add concurrency safety to SQL queries and database reloading

The commit adds mutex locks to the GetAllCountries and GetCitiesByCountry functions to ensure thread-safety during database queries. Additionally, it introduces a mechanism to safely close the old database connection before a new connection is established upon reloading, which improves the reliability of database operations. Lastly, it moves the checking of database file existence to the connectDB function.

* Add sha256 sum check to geolocation store before reload

* Use read lock

* Check SHA256 twice when reload geonames db

---------

Co-authored-by: Yury Gargay <yury.gargay@gmail.com>
  • Loading branch information
@bcmmbaga @surik
bcmmbaga and surik authored Feb 7, 2024
1 parent 4bcee77 commit 6b11bf0
Show file tree
Hide file tree
Showing 9 changed files with 604 additions and 46 deletions.
124 changes: 89 additions & 35 deletions infrastructure_files/download-geolite2.sh
View file
Original file line number Diff line number Diff line change
Expand Up @@ -22,41 +22,95 @@ then
exit 1
fi

DATABASE_URL="https://download.maxmind.com/geoip/databases/GeoLite2-City/download?suffix=tar.gz"
SIGNATURE_URL="https://download.maxmind.com/geoip/databases/GeoLite2-City/download?suffix=tar.gz.sha256"

# Download the database and signature files
echo "Downloading database file..."
DATABASE_FILE=$(curl -s -u "$MM_ACCOUNT_ID":"$MM_LICENSE_KEY" -L -O -J "$DATABASE_URL" -w "%{filename_effective}")
echo "Downloading signature file..."
SIGNATURE_FILE=$(curl -s -u "$MM_ACCOUNT_ID":"$MM_LICENSE_KEY" -L -O -J "$SIGNATURE_URL" -w "%{filename_effective}")

# Verify the signature
echo "Verifying signature..."
if sha256sum -c --status "$SIGNATURE_FILE"; then
echo "Signature is valid."
else
echo "Signature is invalid. Aborting."
if ! command -v sqlite3 &> /dev/null
then
echo "sqlite3 is not installed or not in PATH, please install with your package manager. e.g. sudo apt install sqlite3" > /dev/stderr
exit 1
fi

# Unpack the database file
EXTRACTION_DIR=$(basename "$DATABASE_FILE" .tar.gz)
echo "Unpacking $DATABASE_FILE..."
mkdir -p "$EXTRACTION_DIR"
tar -xzvf "$DATABASE_FILE"

# Create a SHA256 signature file
MMDB_FILE="GeoLite2-City.mmdb"
cd "$EXTRACTION_DIR"
sha256sum "$MMDB_FILE" > "$MMDB_FILE.sha256"
echo "SHA256 signature created for $MMDB_FILE."
cd -

# Remove downloaded files
rm "$DATABASE_FILE" "$SIGNATURE_FILE"

# Done. Print next steps
echo "Process completed successfully."
echo "Now you can place $EXTRACTION_DIR/$MMDB_FILE to 'datadir' of management service."
echo -e "Example:\n\tdocker compose cp $EXTRACTION_DIR/$MMDB_FILE management:/var/lib/netbird/"
download_geolite_mmdb() {
DATABASE_URL="https://download.maxmind.com/geoip/databases/GeoLite2-City/download?suffix=tar.gz"
SIGNATURE_URL="https://download.maxmind.com/geoip/databases/GeoLite2-City/download?suffix=tar.gz.sha256"

# Download the database and signature files
echo "Downloading database file..."
DATABASE_FILE=$(curl -s -u "$MM_ACCOUNT_ID":"$MM_LICENSE_KEY" -L -O -J "$DATABASE_URL" -w "%{filename_effective}")
echo "Downloading signature file..."
SIGNATURE_FILE=$(curl -s -u "$MM_ACCOUNT_ID":"$MM_LICENSE_KEY" -L -O -J "$SIGNATURE_URL" -w "%{filename_effective}")

# Verify the signature
echo "Verifying signature..."
if sha256sum -c --status "$SIGNATURE_FILE"; then
echo "Signature is valid."
else
echo "Signature is invalid. Aborting."
exit 1
fi

# Unpack the database file
EXTRACTION_DIR=$(basename "$DATABASE_FILE" .tar.gz)
echo "Unpacking $DATABASE_FILE..."
mkdir -p "$EXTRACTION_DIR"
tar -xzvf "$DATABASE_FILE"

# Create a SHA256 signature file
MMDB_FILE="GeoLite2-City.mmdb"
cd "$EXTRACTION_DIR"
sha256sum "$MMDB_FILE" > "$MMDB_FILE.sha256"
echo "SHA256 signature created for $MMDB_FILE."
cd -

# Remove downloaded files
rm "$DATABASE_FILE" "$SIGNATURE_FILE"

# Done. Print next steps
echo "Process completed successfully."
echo "Now you can place $EXTRACTION_DIR/$MMDB_FILE to 'datadir' of management service."
echo -e "Example:\n\tdocker compose cp $EXTRACTION_DIR/$MMDB_FILE management:/var/lib/netbird/"
}


download_geolite_csv_and_create_sqlite_db() {
DATABASE_URL="https://download.maxmind.com/geoip/databases/GeoLite2-City-CSV/download?suffix=zip"
SIGNATURE_URL="https://download.maxmind.com/geoip/databases/GeoLite2-City-CSV/download?suffix=zip.sha256"


# Download the database file
echo "Downloading database file..."
DATABASE_FILE=$(curl -s -u "$MM_ACCOUNT_ID":"$MM_LICENSE_KEY" -L -O -J "$DATABASE_URL" -w "%{filename_effective}")
echo "Downloading signature file..."
SIGNATURE_FILE=$(curl -s -u "$MM_ACCOUNT_ID":"$MM_LICENSE_KEY" -L -O -J "$SIGNATURE_URL" -w "%{filename_effective}")

# Verify the signature
echo "Verifying signature..."
if sha256sum -c --status "$SIGNATURE_FILE"; then
echo "Signature is valid."
else
echo "Signature is invalid. Aborting."
exit 1
fi

# Unpack the database file
EXTRACTION_DIR=$(basename "$DATABASE_FILE" .zip)
DB_NAME="geonames.db"

unzip "$DATABASE_FILE"

# Create SQLite database and import data from CSV
sqlite3 "$DB_NAME" <<EOF
.mode csv
.import "$EXTRACTION_DIR/GeoLite2-City-Locations-en.csv" geonames
EOF


# Remove downloaded and extracted files
rm -r -r "$EXTRACTION_DIR"
rm "$DATABASE_FILE" "$SIGNATURE_FILE"

echo "SQLite database '$DB_NAME' created successfully."
echo "Now you can place $DB_NAME to 'datadir' of management service."
echo -e "Example:\n\tdocker compose cp $DB_NAME management:/var/lib/netbird/"
}

download_geolite_mmdb
download_geolite_csv_and_create_sqlite_db
3 changes: 2 additions & 1 deletion management/cmd/management.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,7 @@ import (

"github.com/grpc-ecosystem/go-grpc-middleware/v2/interceptors/realip"
"github.com/netbirdio/management-integrations/integrations"

"github.com/netbirdio/netbird/encryption"
mgmtProto "github.com/netbirdio/netbird/management/proto"
"github.com/netbirdio/netbird/management/server"
Expand Down Expand Up @@ -230,7 +231,7 @@ var (
UserIDClaim: config.HttpConfig.AuthUserIDClaim,
KeysLocation: config.HttpConfig.AuthKeysLocation,
}
httpAPIHandler, err := httpapi.APIHandler(accountManager, *jwtValidator, appMetrics, httpAPIAuthCfg)
httpAPIHandler, err := httpapi.APIHandler(accountManager, geo, *jwtValidator, appMetrics, httpAPIAuthCfg)
if err != nil {
return fmt.Errorf("failed creating HTTP API handler: %v", err)
}
Expand Down
66 changes: 65 additions & 1 deletion management/server/geolocation/geolocation.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@ type Geolocation struct {
mux *sync.RWMutex
sha256sum []byte
db *maxminddb.Reader
locationDB *SqliteStore
stopCh chan struct{}
reloadCheckInterval time.Duration
}
Expand All @@ -43,6 +44,16 @@ type Record struct {
} `maxminddb:"country"`
}

type City struct {
GeoNameID int `gorm:"column:geoname_id"`
CityName string
}

type Country struct {
CountryISOCode string `gorm:"column:country_iso_code"`
CountryName string
}

func NewGeolocation(datadir string) (*Geolocation, error) {
mmdbPath := path.Join(datadir, mmdbFileName)

Expand All @@ -56,11 +67,17 @@ func NewGeolocation(datadir string) (*Geolocation, error) {
return nil, err
}

locationDB, err := NewSqliteStore(datadir)
if err != nil {
return nil, err
}

geo := &Geolocation{
mmdbPath: mmdbPath,
mux: &sync.RWMutex{},
sha256sum: sha256sum,
db: db,
locationDB: locationDB,
reloadCheckInterval: 60 * time.Second, // TODO: make configurable
stopCh: make(chan struct{}),
}
Expand Down Expand Up @@ -115,6 +132,38 @@ func (gl *Geolocation) Lookup(ip net.IP) (*Record, error) {
return &record, nil
}

// GetAllCountries retrieves a list of all countries.
func (gl *Geolocation) GetAllCountries() ([]Country, error) {
allCountries, err := gl.locationDB.GetAllCountries()
if err != nil {
return nil, err
}

countries := make([]Country, 0)
for _, country := range allCountries {
if country.CountryName != "" {
countries = append(countries, country)
}
}
return countries, nil
}

// GetCitiesByCountry retrieves a list of cities in a specific country based on the country's ISO code.
func (gl *Geolocation) GetCitiesByCountry(countryISOCode string) ([]City, error) {
allCities, err := gl.locationDB.GetCitiesByCountry(countryISOCode)
if err != nil {
return nil, err
}

cities := make([]City, 0)
for _, city := range allCities {
if city.CityName != "" {
cities = append(cities, city)
}
}
return cities, nil
}

func (gl *Geolocation) Stop() error {
close(gl.stopCh)
if gl.db != nil {
Expand All @@ -129,6 +178,10 @@ func (gl *Geolocation) reloader() {
case <-gl.stopCh:
return
case <-time.After(gl.reloadCheckInterval):
if err := gl.locationDB.reload(); err != nil {
log.Errorf("geonames db reload failed: %s", err)
}

newSha256sum1, err := getSha256sum(gl.mmdbPath)
if err != nil {
log.Errorf("failed to calculate sha256 sum for '%s': %s", gl.mmdbPath, err)
Expand All @@ -149,7 +202,7 @@ func (gl *Geolocation) reloader() {
}
err = gl.reload(newSha256sum2)
if err != nil {
log.Errorf("reload failed: %s", err)
log.Errorf("mmdb reload failed: %s", err)
}
} else {
log.Debugf("No changes in '%s', no need to reload. Next check is in %.0f seconds.",
Expand Down Expand Up @@ -182,3 +235,14 @@ func (gl *Geolocation) reload(newSha256sum []byte) error {

return nil
}

func fileExists(filePath string) (bool, error) {
_, err := os.Stat(filePath)
if err == nil {
return true, nil
}
if os.IsNotExist(err) {
return false, fmt.Errorf("%v does not exist", filePath)
}
return false, err
}
12 changes: 10 additions & 2 deletions management/server/geolocation/geolocation_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,10 +4,12 @@ import (
"net"
"os"
"path"
"sync"
"testing"

"github.com/netbirdio/netbird/util"
"github.com/stretchr/testify/assert"

"github.com/netbirdio/netbird/util"
)

// from https://github.com/maxmind/MaxMind-DB/blob/main/test-data/GeoLite2-City-Test.mmdb
Expand All @@ -25,8 +27,14 @@ func TestGeoLite_Lookup(t *testing.T) {
}
}()

geo, err := NewGeolocation(tempDir)
db, err := openDB(mmdbPath)
assert.NoError(t, err)

geo := &Geolocation{
mux: &sync.RWMutex{},
db: db,
stopCh: make(chan struct{}),
}
assert.NotNil(t, geo)
defer func() {
err = geo.Stop()
Expand Down
Loading

0 comments on commit 6b11bf0

Please sign in to comment.