diff --git a/services/rotor/__tests__/maxmind-path.test.ts b/services/rotor/__tests__/maxmind-path.test.ts new file mode 100644 index 000000000..8897cc6f4 --- /dev/null +++ b/services/rotor/__tests__/maxmind-path.test.ts @@ -0,0 +1,39 @@ +import * as fs from "fs"; +import * as os from "os"; +import * as path from "path"; +import * as tar from "tar"; +import { expect, test } from "vitest"; +import { loadFromPath } from "../src/lib/maxmind"; + +const cityDb = Buffer.from("fake-city-db"); + +const tmpDir = () => fs.mkdtempSync(path.join(os.tmpdir(), "maxmind-test-")); + +test("loadFromPath loads .mmdb from a directory", async () => { + const dir = tmpDir(); + fs.writeFileSync(path.join(dir, "GeoLite2-City.mmdb"), cityDb); + expect(await loadFromPath(dir, "GeoLite2-City")).toEqual(cityDb); + await expect(loadFromPath(dir, "GeoLite2-Country")).rejects.toThrow( + "neither GeoLite2-Country.mmdb nor GeoLite2-Country.tar.gz found" + ); +}); + +test("loadFromPath loads a single .mmdb file named after its edition", async () => { + const dir = tmpDir(); + const file = path.join(dir, "GeoIP2-City.mmdb"); + fs.writeFileSync(file, cityDb); + expect(await loadFromPath(file, "GeoIP2-City")).toEqual(cityDb); + await expect(loadFromPath(file, "GeoIP2-Country")).rejects.toThrow("doesn't contain GeoIP2-Country edition"); +}); + +test("loadFromPath extracts .mmdb from a tar.gz archive", async () => { + const dir = tmpDir(); + fs.writeFileSync(path.join(dir, "GeoLite2-City.mmdb"), cityDb); + await tar.c({ gzip: true, file: path.join(dir, "GeoLite2-City.tar.gz"), cwd: dir }, ["GeoLite2-City.mmdb"]); + fs.rmSync(path.join(dir, "GeoLite2-City.mmdb")); + expect(await loadFromPath(dir, "GeoLite2-City")).toEqual(cityDb); +}); + +test("loadFromPath fails on a missing path", async () => { + await expect(loadFromPath(path.join(tmpDir(), "nope"), "GeoLite2-City")).rejects.toThrow(); +}); diff --git a/services/rotor/src/index.ts b/services/rotor/src/index.ts index b1f75ec7c..5622dba32 100644 --- a/services/rotor/src/index.ts +++ b/services/rotor/src/index.ts @@ -101,6 +101,7 @@ async function main() { licenseKey: serverEnv.MAXMIND_LICENSE_KEY, url: serverEnv.MAXMIND_URL, s3Bucket: serverEnv.MAXMIND_S3_BUCKET, + path: serverEnv.MAXMIND_PATH, }); metricsServer = initMetricsServer(); } catch (e) { diff --git a/services/rotor/src/lib/maxmind.ts b/services/rotor/src/lib/maxmind.ts index c380273b2..d3b827324 100644 --- a/services/rotor/src/lib/maxmind.ts +++ b/services/rotor/src/lib/maxmind.ts @@ -1,4 +1,6 @@ import { Reader, ReaderModel, City, Isp, Names } from "@maxmind/geoip2-node"; +import * as fs from "fs"; +import * as nodePath from "path"; import * as zlib from "zlib"; import * as tar from "tar"; import { Geo } from "@jitsu/protocols/analytics"; @@ -72,15 +74,18 @@ export async function initMaxMindClient(opts: { licenseKey?: string; url?: string; s3Bucket?: string; + path?: string; }): Promise { - const { licenseKey, s3Bucket, url } = opts; - if (!licenseKey && !url && !s3Bucket) { - log.atWarn().log("licenseKey, url or s3Bucket must be provided. GeoIP resolution will not work."); + const { licenseKey, s3Bucket, url, path } = opts; + if (!licenseKey && !url && !s3Bucket && !path) { + log.atWarn().log("licenseKey, url, s3Bucket or path must be provided. GeoIP resolution will not work."); return DummyResolver; } let loadFunc: LoadFunction; let s3client: S3Client = undefined as any as S3Client; - if (s3Bucket) { + if (path) { + loadFunc = (edition: Edition) => loadFromPath(path, edition); + } else if (s3Bucket) { s3client = createS3Client(); loadFunc = (edition: Edition) => loadFromS3(s3client, s3Bucket, edition); } else { @@ -254,6 +259,27 @@ export async function initMaxMindClient(opts: { } } +export async function loadFromPath(localPath: string, edition: Edition): Promise { + // path may point to a directory with `.mmdb` or `.tar.gz` files + // (the layout produced by MaxMind's geoipupdate tool) or to a single such file + let filePath = localPath; + if ((await fs.promises.stat(localPath)).isDirectory()) { + const fileName = [`${edition}.mmdb`, `${edition}.tar.gz`].find(name => + fs.existsSync(nodePath.join(localPath, name)) + ); + if (!fileName) { + throw new Error(`neither ${edition}.mmdb nor ${edition}.tar.gz found in ${localPath}`); + } + filePath = nodePath.join(localPath, fileName); + } else if (nodePath.basename(localPath).replace(/\.(mmdb|tar\.gz|tgz)$/, "") !== edition) { + throw new Error( + `${localPath} doesn't contain ${edition} edition. Rename the file to ${edition}.mmdb or ${edition}.tar.gz` + ); + } + const buffer = await fs.promises.readFile(filePath); + return filePath.endsWith(".gz") || filePath.endsWith(".tgz") ? await untar(buffer) : buffer; +} + async function loadFromS3(client: S3Client, bucket: string, edition: Edition): Promise { try { const command = new GetObjectCommand({ Bucket: bucket, Key: edition + ".tar.gz" }); diff --git a/services/rotor/src/serverEnv.ts b/services/rotor/src/serverEnv.ts index e5708c7f7..d27e60e38 100644 --- a/services/rotor/src/serverEnv.ts +++ b/services/rotor/src/serverEnv.ts @@ -31,6 +31,7 @@ const ServerEnvSchema = z.object({ // MaxMind Configuration MAXMIND_LICENSE_KEY: z.string().optional(), MAXMIND_URL: z.string().optional(), + MAXMIND_PATH: z.string().optional(), MAXMIND_S3_BUCKET: z.string().optional(), MAXMIND_S3_REGION: z.string().optional(), MAXMIND_S3_ACCESS_KEY_ID: z.string().optional(),