core: refactor geoip.js

Author: micooz

core/src/main.js

@@ -10,7 +10,7 @@ const bsInit = require('blinksocks/bin/init');
 const { Config } = require('blinksocks');
 
 const runServer = require('./core/server');
-const { ServiceManager, logger, GeoIP } = require('./utils');
+const { ServiceManager, logger } = require('./utils');
 
 const {
   RUN_TYPE_CLIENT,
@@ -152,7 +152,6 @@ module.exports = async function main(args) {
         const ip = await getPublicIP();
         logger.info(`public ip address is: ${ip}`);
         db.set('runtime.ip', ip).write();
-        GeoIP.put(ip, { self: true });
       } catch (err) {
         logger.error('cannot get public ip address of this machine: %s', err.stack);
       }

core/src/utils/geoip.js

@@ -52,15 +52,18 @@ function isPrivateIP(ip) {
   );
 }
 
-module.exports = {
+const MAX_QUEUE_SIZE = 500;
 
-  _uniqueKeys: new Set(),
+class GeoIP {
 
-  _store: new Map(/* <ip>: <result> */),
+  constructor() {
+    this._uniqueKeys = new Set();
+    this._store = [];
+  }
 
   // lookup ip address to geo location
   async put(arg, extra = {}) {
-    if (typeof arg !== 'string') {
+    if (typeof arg !== 'string' || !arg) {
       return;
     }
 
@@ -85,50 +88,45 @@ module.exports = {
     try {
       const geoInfo = await getGeoInfo(ip);
       const { lat, lon: lng, query } = geoInfo;
-      if (!this._store.get(ip)) {
-        // merge hostname and ip by the same geo location
-        const item = this.findItemByPosition(lat, lng);
-        if (item) {
-          const oldValue = item.value;
-          const newValue = {
-            ...oldValue,
-            ips: oldValue.ips.concat([query]),
-          };
-          if (oldValue.hostname && extra.hostname) {
-            newValue.hostname = oldValue.hostname.concat(extra.hostname);
+
+      // merge hostname and ip by the same geo location
+      const item = this._store.find((v) => v.lat === lat && v.lng === lng);
+      if (item) {
+        item.ips.push(query);
+        if (item.ips.length > 20) {
+          item.ips.shift();
+        }
+        if (item.hostname && extra.hostname) {
+          item.hostname.push(extra.hostname);
+          if (item.hostname.length > 20) {
+            item.hostname.shift();
           }
-          this._store.set(item.key, newValue);
-        } else {
-          const obj = _.pick(geoInfo, ['as', 'city', 'country', 'lat', 'org', 'regionName']);
-          obj['ips'] = [query];
-          obj['lng'] = lng;
-          this._store.set(ip, Object.assign({}, obj, extra));
         }
-        return;
+      } else {
+        const obj = _.pick(geoInfo, ['as', 'city', 'country', 'lat', 'org', 'regionName']);
+        obj['ips'] = [query];
+        obj['lng'] = lng;
+        this._store.push(Object.assign({}, obj, extra));
+        if (this._store.length > MAX_QUEUE_SIZE) {
+          this._store.shift();
+        }
       }
     } catch (err) {
       logger.error(err.message);
+      this._uniqueKeys.delete(arg);
     }
-    this._uniqueKeys.delete(arg);
-  },
+  }
 
   // return all resolved ip and its geo location
   getStore() {
     return this._store;
-  },
+  }
 
   clear() {
     this._uniqueKeys.clear();
-    this._store.clear();
-  },
+    this._store = [];
+  }
 
-  findItemByPosition(lat, lng) {
-    for (const [key, value] of this._store) {
-      if (lat === value.lat && lng === value.lng) {
-        return { key, value };
-      }
-    }
-    return null;
-  },
+}
 
-};
+module.exports = GeoIP;

core/src/utils/service_manager.js

@@ -2,6 +2,7 @@ const path = require('path');
 const child_process = require('child_process');
 const _ = require('lodash');
 
+const db = require('./db');
 const logger = require('./logger');
 const GeoIP = require('./geoip');
 
@@ -20,7 +21,9 @@ function create() {
     cwd: RUNTIME_PATH,
     silent: process.env.NODE_ENV === 'production',
   });
-  const geoip = Object.create(GeoIP);
+  const geoip = new GeoIP();
+  geoip.put(db.get('runtime.ip').value(), { self: true });
+
   const messageQueue = [];
 
   subprocess.on('message', (message) => {
@@ -199,7 +202,7 @@ module.exports = {
   getServiceGeoIPs(id) {
     const sub = this._subprocesses.get(id);
     if (sub) {
-      return [...sub.geoip.getStore().values()];
+      return sub.geoip.getStore();
     } else {
       return [];
     }