diff --git a/internal/config/config.go b/internal/config/config.go index 3da0713..a1ed60d 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -118,7 +118,7 @@ func New() Config { Upstream: ConfigUpstream{ Mode: SingleUpstream, Upstream: []string{"ip-api.com"}, - RotatedInterval: upstreamRotatedInterval(time.Duration.Hours(24)), + RotatedInterval: upstreamRotatedInterval(24 * time.Hour), }, Dev: ConfigDev{ Debug: false, diff --git a/internal/upstream/selector.go b/internal/upstream/selector.go index a574129..d879083 100644 --- a/internal/upstream/selector.go +++ b/internal/upstream/selector.go @@ -1,7 +1,10 @@ package upstream import ( + "errors" + "fmt" "math/rand/v2" + "sync" "time" "github.com/SourLemonJuice/ipapi-agent/internal/config" @@ -9,43 +12,119 @@ import ( ) var ( + rotationMu sync.RWMutex rotatedName string - rotatedCycleEnded bool = false + rotatedCycleEnded bool + rotationTimer *time.Timer ) -func InitSelector(conf config.ConfigUpstream) { +func InitSelector(conf config.ConfigUpstream) error { + if err := validateUpstreamConfig(conf); err != nil { + return err + } + switch conf.Mode { case config.RotatedUpstream: - rotatedName = randomCodename(conf.Upstream) - newRotatedCycle(conf) + if err := newRotatedCycle(conf); err != nil { + return err + } } + + return nil } -func SelectAPI(conf config.ConfigUpstream) API { +func SelectAPI(conf config.ConfigUpstream) (API, error) { + if len(conf.Upstream) == 0 { + return nil, errors.New("upstream pool is empty") + } + switch conf.Mode { case config.SingleUpstream: return New(conf.Upstream[0]) case config.RandomUpstream: - return New(randomCodename(conf.Upstream)) + codename, err := randomCodename(conf.Upstream) + if err != nil { + return nil, err + } + return New(codename) case config.RotatedUpstream: - if rotatedCycleEnded { - rotatedCycleEnded = false - newRotatedCycle(conf) + rotationMu.RLock() + name := rotatedName + cycleEnded := rotatedCycleEnded + rotationMu.RUnlock() + + if cycleEnded { + if err := newRotatedCycle(conf); err != nil { + return nil, err + } + + rotationMu.RLock() + name = rotatedName + rotationMu.RUnlock() + } + + if name == "" { + return nil, errors.New("rotated upstream not initialized") } - return New(rotatedName) + + return New(name) } - return nil + return nil, errors.New("unsupported upstream mode") } -func randomCodename(list []string) string { - return list[rand.IntN(len(list))] +func randomCodename(list []string) (string, error) { + if len(list) == 0 { + return "", errors.New("upstream pool is empty") + } + return list[rand.IntN(len(list))], nil } -func newRotatedCycle(conf config.ConfigUpstream) { - debug.Logger.Println("new rotation cycle") - time.AfterFunc(time.Duration(conf.RotatedInterval), func() { +func newRotatedCycle(conf config.ConfigUpstream) error { + name, err := randomCodename(conf.Upstream) + if err != nil { + return err + } + + interval := time.Duration(conf.RotatedInterval) + if interval <= 0 { + return errors.New("rotated interval must be greater than zero") + } + + rotationMu.Lock() + if rotationTimer != nil { + rotationTimer.Stop() + } + rotatedName = name + rotatedCycleEnded = false + rotationTimer = time.AfterFunc(interval, func() { + rotationMu.Lock() rotatedCycleEnded = true - rotatedName = randomCodename(conf.Upstream) + rotationTimer = nil + rotationMu.Unlock() }) + rotationMu.Unlock() + + debug.Logger.Println("new rotation cycle") + return nil +} + +func validateUpstreamConfig(conf config.ConfigUpstream) error { + if len(conf.Upstream) == 0 { + return errors.New("upstream pool is empty") + } + + for _, codename := range conf.Upstream { + if _, err := New(codename); err != nil { + return fmt.Errorf("unsupported upstream %q: %w", codename, err) + } + } + + if conf.Mode == config.RotatedUpstream { + if time.Duration(conf.RotatedInterval) <= 0 { + return errors.New("rotated interval must be greater than zero") + } + } + + return nil } diff --git a/internal/upstream/upstream.go b/internal/upstream/upstream.go index c582249..fa2e93b 100644 --- a/internal/upstream/upstream.go +++ b/internal/upstream/upstream.go @@ -18,17 +18,17 @@ type API interface { Fill(resp *response.Query) error } -func New(codename string) API { +func New(codename string) (API, error) { switch codename { case "ip-api.com": - return &ipApiCom{} + return &ipApiCom{}, nil case "ipinfo-free": - return &ipinfoFree{} + return &ipinfoFree{}, nil case "ipapi.co": - return &ipapiCo{} + return &ipapiCo{}, nil } - return nil + return nil, fmt.Errorf("unknown upstream codename: %s", codename) } func getJSON(url string, data API) error { diff --git a/main.go b/main.go index 710e429..45c4fac 100644 --- a/main.go +++ b/main.go @@ -60,7 +60,9 @@ func main() { gin.SetMode(gin.ReleaseMode) } - upstream.InitSelector(conf.Upstream) + if err := upstream.InitSelector(conf.Upstream); err != nil { + log.Fatalf("failed to initialize upstream selector: %v", err) + } router := gin.New() router.RedirectTrailingSlash = true @@ -70,7 +72,9 @@ func main() { router.Use(gin.Logger()) } - router.SetTrustedProxies(conf.TrustedProxies) + if err := router.SetTrustedProxies(conf.TrustedProxies); err != nil { + log.Fatalf("invalid trusted proxies configuration: %v", err) + } router.GET("/", getRoot) router.GET("/query", getQuery) @@ -155,7 +159,13 @@ func getRoot(c *gin.Context) { // let struct cache compatible with getQuery() resp.Status = "success" - apidata := upstream.SelectAPI(conf.Upstream) + apidata, err := upstream.SelectAPI(conf.Upstream) + if err != nil { + log.Printf("Upstream selection error: %v", err) + c.Abort() + c.String(http.StatusInternalServerError, respTXTFailure(colorful, "Internal Server Error")) + return + } err = apidata.Request(addrStr) if err != nil { log.Printf("Data source error: %v", err) @@ -265,7 +275,15 @@ func getQuery(c *gin.Context) { resp.Status = "success" - apidata := upstream.SelectAPI(conf.Upstream) + apidata, err := upstream.SelectAPI(conf.Upstream) + if err != nil { + log.Printf("Upstream selection error: %v", err) + c.AbortWithStatusJSON(http.StatusInternalServerError, gin.H{ + "status": "failure", + "message": "Internal Server Error", + }) + return + } err = apidata.Request(addrStr) if err != nil { log.Printf("Data source error: %v", err)