| 1 | import { defineConfig } from './config' |
| 2 | import { CFG, DAY, httpStream, isIpLan, isLocalHost, statWithTimeout, unzip } from './misc' |
| 3 | import { rename } from 'node:fs/promises' |
| 4 | import { IP2Location } from 'ip2location-nodejs' |
| 5 | import _ from 'lodash' |
| 6 | import { Middleware } from 'koa' |
| 7 | import { disconnect, updateConnection } from './connections' |
| 8 | |
| 9 | const ip2location = new IP2Location() |
| 10 | const enabled = defineConfig(CFG.geo_enable, false) |
| 11 | const allow = defineConfig<boolean | null>(CFG.geo_allow, null) |
| 12 | const list = defineConfig(CFG.geo_list, [] as string[]) |
| 13 | const allowUnknown = defineConfig(CFG.geo_allow_unknown, false) |
| 14 | enabled.sub(checkFiles) |
| 15 | setInterval(checkFiles, DAY) // keep updated at run-time |
| 16 | |
| 17 | // benchmark: memoize can make this 44x faster |
| 18 | export const ip2country = _.memoize((ip: string) => ip2location.getCountryShortAsync(ip).then(v => v === '-' ? '' : v, () => '')) |
| 19 | |
| 20 | export const geoFilter: Middleware = async (ctx, next) => { |
| 21 | if (enabled.get() && !isLocalHost(ctx) && !isIpLan(ctx.ip)) { |
| 22 | const { connection } = ctx.state |
| 23 | const country = connection.country ??= await ip2country(ctx.ip) |
| 24 | if (country) |
| 25 | updateConnection(connection, { country }) |
| 26 | if (!ctx.state.skipFilters && allow.get() !== null) |
| 27 | if (country ? list.get().includes(country) !== allow.get() : !allowUnknown.get()) |
| 28 | return disconnect(ctx, 'geo-filter') |
| 29 | } |
| 30 | await next() |
| 31 | } |
| 32 | |
| 33 | function isOpen() { |
| 34 | return Boolean(ip2location.getPackageVersion()) |
| 35 | } |
| 36 | |
| 37 | async function checkFiles() { |
| 38 | if (!enabled.get()) return |
| 39 | const BIN_FILE = 'IP2LOCATION-LITE-DB1.IPV6.BIN' |
| 40 | const URL = `https://download.ip2location.com/lite/${BIN_FILE}.ZIP` |
| 41 | const LOCAL_FILE = 'geo_ip.bin' |
| 42 | const TEMP = LOCAL_FILE + '.downloading' |
| 43 | const { mtime=0 } = await statWithTimeout(LOCAL_FILE).catch(() => ({ mtime: 0 })) |
| 44 | const name = 'geo-ip db' |
| 45 | const now = Date.now() |
| 46 | if (+mtime < now - 31 * DAY) // month-old or non-existing |
| 47 | try { |
| 48 | const req = await httpStream(URL) |
| 49 | console.log(`Downloading ${name}`) |
| 50 | await unzip(req, path => path.toUpperCase().endsWith(BIN_FILE) && TEMP) // give a temp name |
| 51 | const s = await statWithTimeout(TEMP) // check existence |
| 52 | if (s.size < 1E6) |
| 53 | throw `Bad size for geo_ip: ${s.size}` |
| 54 | if (isOpen()) |
| 55 | ip2location.close() |
| 56 | await rename(TEMP, LOCAL_FILE) |
| 57 | ip2country.cache.clear?.() |
| 58 | console.log(`${name} download completed`) |
| 59 | } |
| 60 | catch (e: any) { |
| 61 | console.error(`Failed to download ${name}${mtime ? ", falling back on old data" : ''}:`, e?.message || String(e)) |
| 62 | } |
| 63 | else if (isOpen()) return |
| 64 | console.debug(`Loading ${name}`) |
| 65 | ip2location.open(LOCAL_FILE) // using openAsync causes a DEP0137 error within 10 seconds |
| 66 | } |