fix: big oopsie

This commit is contained in:
skidoodle 2024-07-11 10:14:43 +02:00
parent 0fb3779ee0
commit 7dd9039b92

90
main.go
View file

@ -17,12 +17,11 @@ import (
"github.com/oschwald/maxminddb-golang" "github.com/oschwald/maxminddb-golang"
) )
// GeoIP database readers var cityDB *maxminddb.Reader
var ( var asnDB *maxminddb.Reader
cityDB *maxminddb.Reader var dbMtx = new(sync.RWMutex)
asnDB *maxminddb.Reader
dbMtx sync.RWMutex
var (
currCityFilename = time.Now().Format("2006-01") + "-city.mmdb" currCityFilename = time.Now().Format("2006-01") + "-city.mmdb"
currASNFilename = time.Now().Format("2006-01") + "-asn.mmdb" currASNFilename = time.Now().Format("2006-01") + "-asn.mmdb"
) )
@ -33,66 +32,60 @@ const (
) )
func main() { func main() {
initDBs() initDatabases()
go runUpdater() go startUpdater()
runServer() startServer()
} }
// Initialize the databases func initDatabases() {
func initDBs() {
var err error var err error
cityDB, err = openDB(currCityFilename)
if err != nil {
log.Printf("Error opening city database: %v", err)
updateDB(cityDBURL, &cityDB, &currCityFilename)
}
asnDB, err = openDB(currASNFilename)
if err != nil {
log.Printf("Error opening ASN database: %v", err)
updateDB(asnDBURL, &asnDB, &currASNFilename)
}
}
// Open a MaxMind DB file cityDB, err = maxminddb.Open(currCityFilename)
func openDB(filename string) (*maxminddb.Reader, error) { if err != nil {
db, err := maxminddb.Open(filename) if os.IsNotExist(err) {
if err != nil && !os.IsNotExist(err) { currCityFilename = ""
return nil, err
}
return db, nil
}
// Download and set the database
func updateDB(urlTemplate string, db **maxminddb.Reader, currFilename *string) {
*currFilename = ""
doUpdate() doUpdate()
if *db == nil { if cityDB == nil {
log.Fatalf("Failed to initialize database from %s", urlTemplate) log.Fatalf("Failed to initialize city database: %v", err)
}
} else {
log.Fatalf("Error opening city database: %v", err)
}
}
asnDB, err = maxminddb.Open(currASNFilename)
if err != nil {
if os.IsNotExist(err) {
currASNFilename = ""
doUpdate()
if asnDB == nil {
log.Fatalf("Failed to initialize ASN database: %v", err)
}
} else {
log.Fatalf("Error opening ASN database: %v", err)
}
} }
} }
// Periodically update the databases func startUpdater() {
func runUpdater() { for range time.Tick(time.Hour * 24 * 7) {
for range time.Tick(24 * time.Hour * 7) {
doUpdate() doUpdate()
} }
} }
// Start the HTTP server func startServer() {
func runServer() {
log.Println("Server listening on :3000") log.Println("Server listening on :3000")
http.HandleFunc("/", handler) http.HandleFunc("/", handler)
log.Fatal(http.ListenAndServe(":3000", nil)) log.Fatal(http.ListenAndServe(":3000", nil))
} }
// Fetch and update the GeoIP databases
func doUpdate() { func doUpdate() {
log.Println("Fetching updates...") log.Println("Fetching updates...")
currMonth := time.Now().Format("2006-01") currMonth := time.Now().Format("2006-01")
newCityFilename := currMonth + "-city.mmdb" newCityFilename := currMonth + "-city.mmdb"
newASNFilename := currMonth + "-asn.mmdb" newASNFilename := currMonth + "-asn.mmdb"
updateDatabase(cityDBURL, newCityFilename, func(newDB *maxminddb.Reader) { updateDatabase(fmt.Sprintf(cityDBURL, currMonth), newCityFilename, func(newDB *maxminddb.Reader) {
dbMtx.Lock() dbMtx.Lock()
defer dbMtx.Unlock() defer dbMtx.Unlock()
if cityDB != nil { if cityDB != nil {
@ -103,7 +96,7 @@ func doUpdate() {
log.Printf("City GeoIP database updated to %s\n", currMonth) log.Printf("City GeoIP database updated to %s\n", currMonth)
}) })
updateDatabase(asnDBURL, newASNFilename, func(newDB *maxminddb.Reader) { updateDatabase(fmt.Sprintf(asnDBURL, currMonth), newASNFilename, func(newDB *maxminddb.Reader) {
dbMtx.Lock() dbMtx.Lock()
defer dbMtx.Unlock() defer dbMtx.Unlock()
if asnDB != nil { if asnDB != nil {
@ -115,22 +108,21 @@ func doUpdate() {
}) })
} }
// Update the specified database func updateDatabase(url, dstFilename string, updateFunc func(*maxminddb.Reader)) {
func updateDatabase(urlTemplate, dstFilename string, updateFunc func(*maxminddb.Reader)) { resp, err := http.Get(url)
resp, err := http.Get(fmt.Sprintf(urlTemplate, time.Now().Format("2006-01")))
if err != nil { if err != nil {
log.Printf("Error fetching the updated DB: %v\n", err) log.Printf("Error fetching the updated DB from %s: %v\n", url, err)
return return
} }
defer resp.Body.Close() defer resp.Body.Close()
if resp.StatusCode != http.StatusOK { if resp.StatusCode != http.StatusOK {
log.Printf("Non-200 status code (%d), retry later...\n", resp.StatusCode) log.Printf("Non-200 status code (%d) from %s, retry later...\n", resp.StatusCode, url)
return return
} }
dst, err := os.Create(dstFilename) dst, err := os.Create(dstFilename)
if err != nil { if err != nil {
log.Printf("Error creating file: %v\n", err) log.Printf("Error creating file %s: %v\n", dstFilename, err)
return return
} }
defer dst.Close() defer dst.Close()