Add option for wait for governance on cold start (#43)

* add waiting for governance option
* bump version

Author: xinghengwang

.nvmrc

@@ -0,0 +1 @@
+20

README.md

@@ -945,7 +945,38 @@ On cold starts (when a new execution environment is initialized), there can be a
 
 In typical production APIs, regular traffic keeps functions warm, so the likelihood of requests slipping through this initial window is small.
 
-In isolated tests or very low-traffic scenarios, the execution environment may be recycled and the cache cleared, recreating the short window until configuration is reloaded. When testing, please send few normal request to keep the system warm first. 
+In isolated tests or very low-traffic scenarios, the execution environment may be recycled and the cache cleared, recreating the short window until configuration is reloaded. When testing, please send few normal request to keep the system warm first.
+
+### Cold Start Governance
+
+You can enforce that initial requests during a Lambda cold start wait until Moesif Config and Governance Rules are loaded.
+
+- `waitForGovernanceOnColdStart` (Boolean): When `true`, the middleware blocks the first invocations until governance rules are fetched and the app config is available. Default: `false`.
+- `governanceLoadTimeoutMs` (Number): Maximum time in milliseconds to wait for config before proceeding even if it hasn’t loaded, to avoid indefinite blocking. Default: `5000`.
+
+Example:
+
+```javascript
+const moesif = require('moesif-aws-lambda');
+
+const options = {
+  applicationId: process.env.MOESIF_APPLICATION_ID,
+  waitForGovernanceOnColdStart: true,
+  governanceLoadTimeoutMs: 8000,
+};
+
+exports.handler = moesif(options, async (event, context) => {
+  return {
+    statusCode: 200,
+    headers: { 'Content-Type': 'application/json' },
+    body: JSON.stringify({ ok: true })
+  };
+});
+```
+
+Notes:
+- Requests during cold start may take longer but ensure governance is applied when rules exist.
+- If the timeout elapses without config, the request proceeds to avoid indefinite delays.
 
 ## How to Get Help
 If you face any issues using this middleware, try the [troubheshooting guidelines](#troubleshoot). For further assistance, reach out to our [support team](mailto:support@moesif.com).

lib/index.js

@@ -70,7 +70,7 @@ module.exports = function (options, handler) {
   // config moesifapi
   var config = moesifapi.configuration;
   config.ApplicationId = options.applicationId || options.ApplicationId || process.env.MOESIF_APPLICATION_ID;
-  config.UserAgent = 'moesif-aws-lambda-nodejs/' + '2.0.7';
+  config.UserAgent = 'moesif-aws-lambda-nodejs/' + '2.1.0';
   config.BaseUri = options.baseUri || options.BaseUri || config.BaseUri;
   var moesifController = moesifapi.ApiController;
 
@@ -125,7 +125,46 @@ module.exports = function (options, handler) {
       details);
   };
   governanceRulesManager.setLogger(logGovernance);
-  governanceRulesManager.tryGetRules();
+  var getRulesPromise = governanceRulesManager.tryGetRules();
+
+  // Option: wait for governance rules + config on cold start
+  options.waitForGovernanceOnColdStart = options.waitForGovernanceOnColdStart || false;
+  options.governanceLoadTimeoutMs = typeof options.governanceLoadTimeoutMs === 'number'
+    ? options.governanceLoadTimeoutMs
+    : 2000;
+
+  var initialLoadPromise = null;
+  if (options.waitForGovernanceOnColdStart) {
+    // Kick off config load immediately
+    moesifConfigManager.tryGetConfig();
+
+    function waitForConfig(timeoutMs) {
+      return new Promise(function(resolve) {
+        var start = Date.now();
+        (function check() {
+          if (moesifConfigManager.hasConfig()) {
+            return resolve();
+          }
+          if (timeoutMs && Date.now() - start >= timeoutMs) {
+            // Timeout reached; proceed even if config isn't available
+            return resolve();
+          }
+          // Keep nudging config fetch and poll until available or timeout
+          moesifConfigManager.tryGetConfig();
+          setTimeout(check, 100);
+        })();
+      });
+    }
+
+    initialLoadPromise = Promise.all([
+      // Use the existing rules fetch promise kicked off above
+      getRulesPromise,
+      waitForConfig(options.governanceLoadTimeoutMs)
+    ]).catch(function() {
+      // Swallow errors; proceed regardless to avoid indefinite blocking
+      return;
+    });
+  }
 
   var trySaveEventLocal = function(eventData) {
     moesifConfigManager.tryGetConfig();
@@ -179,60 +218,69 @@ module.exports = function (options, handler) {
       callback(err, modifiedResultObject)
     };
 
-    const isV1 = determineIsEventVersionV1(event);
-
-    let returnPromise;
-    if (governanceRulesManager.hasRules()) {
-      var requestDataForGovernance = prepareRequestDataForGovernance(event, context, isV1);
-
-      var governedResponseHolder = governanceRulesManager.governRequest(
-        moesifConfigManager._config,
-        // this may cause identifyUser and identifyCompany to be called twice,
-        // but this should be ok, but in order to block for governance rule
-        // we have to trigger this earlier in the stream before response might be ready
-        ensureToString(options.identifyUser(event, context)),
-        ensureToString(options.identifyCompany(event, context)),
-        requestDataForGovernance.requestFields,
-        requestDataForGovernance.requestHeaders,
-        requestDataForGovernance.requestBody
-      );
+    function proceed() {
+      const isV1 = determineIsEventVersionV1(event);
+
+      let returnPromise;
+      if (governanceRulesManager.hasRules()) {
+        var requestDataForGovernance = prepareRequestDataForGovernance(event, context, isV1);
+
+        var governedResponseHolder = governanceRulesManager.governRequest(
+          moesifConfigManager._config,
+          // this may cause identifyUser and identifyCompany to be called twice,
+          // but this should be ok, but in order to block for governance rule
+          // we have to trigger this earlier in the stream before response might be ready
+          ensureToString(options.identifyUser(event, context)),
+          ensureToString(options.identifyCompany(event, context)),
+          requestDataForGovernance.requestFields,
+          requestDataForGovernance.requestHeaders,
+          requestDataForGovernance.requestBody
+        );
 
-      if (governedResponseHolder.headers) {
-        extraHeaders = { ...extraHeaders, ...governedResponseHolder.headers };
+        if (governedResponseHolder.headers) {
+          extraHeaders = { ...extraHeaders, ...governedResponseHolder.headers };
+        }
+
+        if (governedResponseHolder.blocked_by) {
+          returnPromise = Promise.resolve({
+            isBase64Encoded: false,
+            statusCode: governedResponseHolder.status,
+            headers: {
+              ...extraHeaders,
+              'X-Moesif-Blocked-By': governedResponseHolder.blocked_by
+            },
+            body: JSON.stringify(governedResponseHolder.body)
+          });
+        }
+      }
+      if (!returnPromise) {
+        returnPromise = handler(event, context, next);
       }
 
-      if (governedResponseHolder.blocked_by) {
-        returnPromise = Promise.resolve({
-          isBase64Encoded: false,
-          statusCode: governedResponseHolder.status,
-          headers: {
+      if (returnPromise instanceof Promise) {
+        return returnPromise.then((result) => {
+          result.headers = {
+            ...result.headers,
             ...extraHeaders,
-            'X-Moesif-Blocked-By': governedResponseHolder.blocked_by
-          },
-          body: JSON.stringify(governedResponseHolder.body)
+          };
+
+          return logEvent(event, context, null, result, options, moesifController).then(() => {
+            return result;
+          });
+        }).catch((err) => {
+          logEvent(event, context, err, {}, options, moesifController);
+          throw err;
         });
       }
-    }
-    if (!returnPromise) {
-      returnPromise = handler(event, context, next);
+      return returnPromise;
     }
 
-    if (returnPromise instanceof Promise) {
-      return returnPromise.then((result) => {
-        result.headers = {
-          ...result.headers,
-          ...extraHeaders,
-        };
-
-        return logEvent(event, context, null, result, options, moesifController).then(() => {
-          return result;
-        });
-      }).catch((err) => {
-        logEvent(event, context, err, {}, options, moesifController);
-        throw err;
+    if (initialLoadPromise) {
+      return initialLoadPromise.then(function () {
+        return proceed();
       });
     }
-    return returnPromise;
+    return proceed();
   };
 
   moesifMiddleware.updateUser = function(userModel, cb) {

package-lock.json

@@ -1,6 +1,6 @@
 {
   "name": "moesif-aws-lambda",
-  "version": "2.0.7",
+  "version": "2.1.0",
   "description": "API Monitoring Middleware for AWS Lambda",
   "main": "lib/index.js",
   "keywords": [

package.json

@@ -0,0 +1,12 @@
+#!/usr/bin/env zsh
+set -euo pipefail
+
+# Use Node from .nvmrc if available
+if command -v nvm >/dev/null 2>&1; then
+  nvm use >/dev/null
+fi
+
+# Set a dummy MOESIF_APPLICATION_ID if not provided
+export MOESIF_APPLICATION_ID=application_id_placeholder
+
+node --test tests/*.js

scripts/run-tests.sh

@@ -0,0 +1,56 @@
+const path = require("path");
+require("dotenv").config({ path: path.resolve(__dirname, ".env") });
+
+// Use basic assertions to avoid Node version dependency on node:test
+const assert = require("assert");
+
+const moesif = require("../lib");
+const governanceRulesManager = require("../lib/governanceRulesManager");
+const moesifConfigManager = require("../lib/moesifConfigManager");
+
+// Stub config and rules to avoid network calls
+moesifConfigManager._config = { sample_rate: 100 };
+
+// Make rules available and return a blocking response
+const BLOCK_RULE_ID = "test-rule-id";
+governanceRulesManager.hasRules = function () { return true; };
+governanceRulesManager.governRequest = function () {
+  return {
+    status: 451,
+    headers: { "X-Blocked": "true" },
+    body: { reason: "blocked by test" },
+    blocked_by: BLOCK_RULE_ID,
+  };
+};
+
+const handler = async (event, context) => {
+  return {
+    statusCode: 200,
+    headers: { "Content-Type": "application/json" },
+    body: JSON.stringify({ ok: true }),
+  };
+};
+
+const context = { getRemainingTimeInMillis: () => 1000 };
+
+async function run() {
+  const moesifWrappedHandler = moesif({
+    applicationId: process.env.MOESIF_APPLICATION_ID || "dummy",
+    waitForGovernanceOnColdStart: true,
+    governanceLoadTimeoutMs: 1000,
+  }, handler);
+
+  const event = { headers: {}, httpMethod: "GET", path: "/", requestContext: {} };
+
+  const response = await moesifWrappedHandler(event, { ...context });
+
+  assert.deepStrictEqual(response.statusCode, 451);
+  assert.ok(response.headers["X-Moesif-Transaction-Id"]);
+  assert.deepStrictEqual(response.headers["X-Moesif-Blocked-By"], BLOCK_RULE_ID);
+  assert.ok(response.body);
+}
+
+run().catch((err) => {
+  console.error(err);
+  process.exit(1);
+});