Skip to content

Commit

Permalink
feat(pkg): add support for downloading GeoIP databases from a mirror
Browse files Browse the repository at this point in the history
- Added `SHELLHUB_MAXMIND_MIRROR` environment variable to allow
  specifying an alternate mirror for downloading MaxMind databases. When
  `SHELLHUB_MAXMIND_MIRROR` is set, it takes precedence over
  `SHELLHUB_MAXMIND_LICENSE`, directing the API to use the mirror for
  database downloads.
  • Loading branch information
heiytor committed Oct 31, 2024
1 parent 3ab7b3f commit 65dfadf
Show file tree
Hide file tree
Showing 4 changed files with 61 additions and 6 deletions.
11 changes: 10 additions & 1 deletion .env
Original file line number Diff line number Diff line change
Expand Up @@ -47,8 +47,17 @@ SHELLHUB_NETWORK=shellhub_network
# Values: A valid domain name
SHELLHUB_PUBLIC_URL_DOMAIN=


# Specifies an alternative mirror URL for downloading the GeoIP databases. This
# field takes precedence over SHELLHUB_MAXMIND_LICENSE; when both are
# configured, SHELLHUB_MAXMIND_MIRROR will be used as the primary source for
# database downloads. Leave both blank to disable the feature.
SHELLHUB_MAXMIND_MIRROR=

# Specifies a MaxMind license key used to authenticate requests for downloading
# the GeoIP database directly from MaxMind. Leave it blank to disable the feature.
# the GeoIP database directly from MaxMind. If SHELLHUB_MAXMIND_MIRROR is not
# set, this license key will be used as the fallback method for fetching the
# database. Leave both blank to disable the feature.
SHELLHUB_MAXMIND_LICENSE=

# The schedule for worker tasks.
Expand Down
25 changes: 20 additions & 5 deletions api/server.go
Original file line number Diff line number Diff line change
Expand Up @@ -106,8 +106,14 @@ type config struct {
// Check [https://github.com/hibiken/asynq/wiki/Task-aggregation] for more information.
AsynqGroupMaxSize int `env:"ASYNQ_GROUP_MAX_SIZE,default=1000"`

// GeoipMaxmindLicense is the Maxmind license key used to authenticate requests for
// downloading the GeoIP database directly from MaxMind.
// GeoipMirror specifies an alternative mirror URL for downloading the GeoIP databases.
// This field takes precedence over [GeoipMaxmindLicense]; when both are configured,
// GeoipMirror will be used as the primary source for database downloads.
GeoipMirror string `env:"MAXMIND_MIRROR,default="`

// GeoipMaxmindLicense is the MaxMind license key used to authenticate requests for
// downloading the GeoIP database directly from MaxMind. If [GeoipMirror] is not set,
// this license key will be used as the fallback method for fetching the database.
GeoipMaxmindLicense string `env:"MAXMIND_LICENSE,default="`
}

Expand Down Expand Up @@ -148,15 +154,24 @@ func startServer(ctx context.Context, cfg *config, store store.Store, cache stor

servicesOptions := []services.Option{}

if cfg.GeoipMaxmindLicense != "" {
if cfg.GeoipMirror != "" {
log.Info("GeoIP feature is enable")

locator, err := geolite2.NewLocator(ctx, geolite2.FetchFromLicenseKey(cfg.GeoipMaxmindLicense))
locator, err := geolite2.NewLocator(ctx, geolite2.FetchFromMirror()(cfg.GeoipMirror))
if err != nil {
log.WithError(err).Fatal("Failed to init GeoIP")
} else {
servicesOptions = append(servicesOptions, services.WithLocator(locator))
}
} else if cfg.GeoipMaxmindLicense != "" {
log.Info("GeoIP feature is enable")

servicesOptions = append(servicesOptions, services.WithLocator(locator))
locator, err := geolite2.NewLocator(ctx, geolite2.FetchFromLicenseKey(cfg.GeoipMaxmindLicense))
if err != nil {
log.WithError(err).Fatal("Failed to init GeoIP")
} else {
servicesOptions = append(servicesOptions, services.WithLocator(locator))
}
}

service := services.NewService(store, nil, nil, cache, apiClient, servicesOptions...)
Expand Down
1 change: 1 addition & 0 deletions docker-compose.yml
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,7 @@ services:
- SHELLHUB_ENTERPRISE=${SHELLHUB_ENTERPRISE}
- SHELLHUB_BILLING=${SHELLHUB_BILLING}
- SHELLHUB_CLOUD=${SHELLHUB_CLOUD}
- MAXMIND_MIRROR=${SHELLHUB_MAXMIND_MIRROR}
- MAXMIND_LICENSE=${SHELLHUB_MAXMIND_LICENSE}
- RECORD_RETENTION=${SHELLHUB_RECORD_RETENTION}
- TELEMETRY=${SHELLHUB_TELEMETRY:-}
Expand Down
30 changes: 30 additions & 0 deletions pkg/geoip/geolite2/fetcher.go
Original file line number Diff line number Diff line change
Expand Up @@ -41,3 +41,33 @@ func FetchFromLicenseKey(licenseKey string) GeoliteFetcher {
return nil
}
}

func FetchFromMirror(mirror string) GeoliteFetcher {
return func(ctx context.Context) error {
urls := []string{}
for _, id := range []string{dbCountryID, dbCityID} {
_, err := os.Stat(filepath.Join(dbPath, id+dbExtension))
switch {
case errors.Is(err, fs.ErrNotExist):
u, err := url.Parse(mirror)
if err != nil {
return err
}

u.Path = "/" + id + ".tar.gz"

urls = append(urls, u.String())
default:
return err
}
}

if len(urls) > 0 {
if err := fetchDBs(ctx, urls); err != nil {
return err
}
}

return nil
}
}

0 comments on commit 65dfadf

Please sign in to comment.