152 lines
3.0 KiB
Go
152 lines
3.0 KiB
Go
package main
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"log"
|
|
"net/url"
|
|
"os"
|
|
"time"
|
|
|
|
"github.com/gofiber/fiber/v2"
|
|
|
|
"geoip-rest/internal/geo"
|
|
"geoip-rest/internal/schedule"
|
|
)
|
|
|
|
const (
|
|
defaultPort = "8080"
|
|
defaultDBPath = "/initial_data/GeoLite2-City.mmdb"
|
|
defaultCron = ""
|
|
defaultScript = "./scripts/dump_and_import.sh"
|
|
)
|
|
|
|
func main() {
|
|
backend := geo.Backend(env("GEOIP_BACKEND", string(geo.BackendMMDB)))
|
|
dbPath := env("GEOIP_DB_PATH", defaultDBPath)
|
|
dbURL := os.Getenv("DATABASE_URL")
|
|
lookupQuery := os.Getenv("GEOIP_LOOKUP_QUERY")
|
|
port := env("PORT", defaultPort)
|
|
|
|
resolver, err := geo.NewResolver(geo.Config{
|
|
Backend: backend,
|
|
MMDBPath: dbPath,
|
|
DatabaseURL: dbURL,
|
|
LookupQuery: lookupQuery,
|
|
})
|
|
if err != nil {
|
|
log.Fatalf("failed to initialize resolver: %v", err)
|
|
}
|
|
defer resolver.Close()
|
|
|
|
app := fiber.New(fiber.Config{
|
|
DisableStartupMessage: true,
|
|
})
|
|
|
|
app.Get("/", func(c *fiber.Ctx) error {
|
|
return c.JSON(fiber.Map{
|
|
"service": "geoip-rest",
|
|
"endpoints": []string{
|
|
"/health",
|
|
"/lookup?ip=<IPv4|IPv6>",
|
|
},
|
|
})
|
|
})
|
|
|
|
app.Get("/health", func(c *fiber.Ctx) error {
|
|
return c.JSON(fiber.Map{"status": "ok"})
|
|
})
|
|
|
|
app.Get("/lookup", func(c *fiber.Ctx) error {
|
|
ip := c.Query("ip")
|
|
if ip == "" {
|
|
ip = c.IP()
|
|
}
|
|
|
|
location, err := resolver.Lookup(ip)
|
|
if err != nil {
|
|
switch {
|
|
case errors.Is(err, geo.ErrInvalidIP):
|
|
return c.Status(fiber.StatusBadRequest).JSON(fiber.Map{
|
|
"error": "invalid ip address",
|
|
})
|
|
case errors.Is(err, geo.ErrNotFound):
|
|
return c.Status(fiber.StatusNotFound).JSON(fiber.Map{
|
|
"error": "location not found",
|
|
})
|
|
default:
|
|
return c.Status(fiber.StatusInternalServerError).JSON(fiber.Map{
|
|
"error": "lookup failed",
|
|
})
|
|
}
|
|
}
|
|
|
|
return c.JSON(location)
|
|
})
|
|
|
|
log.Printf("starting GeoIP API on :%s backend=%s", port, backend)
|
|
switch backend {
|
|
case geo.BackendPostgres:
|
|
log.Printf("using postgres DSN %s", sanitizeDBURL(dbURL))
|
|
default:
|
|
log.Printf("using mmdb path %s", dbPath)
|
|
}
|
|
|
|
stopScheduler := maybeStartScheduler()
|
|
defer func() {
|
|
if stopScheduler != nil {
|
|
ctx := stopScheduler()
|
|
<-ctx.Done()
|
|
}
|
|
}()
|
|
|
|
if err := app.Listen(":" + port); err != nil {
|
|
log.Fatalf("server stopped: %v", err)
|
|
}
|
|
}
|
|
|
|
func env(key, fallback string) string {
|
|
if val := os.Getenv(key); val != "" {
|
|
return val
|
|
}
|
|
return fallback
|
|
}
|
|
|
|
func sanitizeDBURL(raw string) string {
|
|
u, err := url.Parse(raw)
|
|
if err != nil {
|
|
return "postgres"
|
|
}
|
|
return u.Redacted()
|
|
}
|
|
|
|
func maybeStartScheduler() func() context.Context {
|
|
cronExpr := env("USER_PROGRAM_CRON", defaultCron)
|
|
if cronExpr == "" {
|
|
return nil
|
|
}
|
|
script := env("USER_PROGRAM_SCRIPT", defaultScript)
|
|
|
|
sched, err := schedule.Start(schedule.Config{
|
|
CronExpr: cronExpr,
|
|
ScriptPath: script,
|
|
Logger: log.Default(),
|
|
})
|
|
if err != nil {
|
|
log.Printf("scheduler not started (error=%v)", err)
|
|
return nil
|
|
}
|
|
|
|
return func() context.Context {
|
|
ctx := sched.Stop()
|
|
timer := time.NewTimer(2 * time.Second)
|
|
select {
|
|
case <-ctx.Done():
|
|
timer.Stop()
|
|
return ctx
|
|
case <-timer.C:
|
|
return ctx
|
|
}
|
|
}
|
|
}
|