feat(rotor): support loading MaxMind databases from a local path

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 <edition>.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();
+});

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) {

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<GeoResolver> {
-  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<Buffer> {
+  // path may point to a directory with `<Edition>.mmdb` or `<Edition>.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<Buffer> {
   try {
     const command = new GetObjectCommand({ Bucket: bucket, Key: edition + ".tar.gz" });

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(),