feat: add tenant-id to invoke context and structured log messages (#128)

* feat: add tenant id to invoke context

* test: add unit tests for next invocation through rapid client

* feat: add tenant-id to structured log message

---------

Co-authored-by: Ramisa Alam <[email protected]>

Author: ramisa2108

deps/aws-lambda-cpp-0.2.8.tar.gz

@@ -0,0 +1,40 @@
+diff --git a/include/aws/lambda-runtime/runtime.h b/include/aws/lambda-runtime/runtime.h
+index f7cb8ef..833b69d 100644
+--- a/include/aws/lambda-runtime/runtime.h
++++ b/include/aws/lambda-runtime/runtime.h
+@@ -56,6 +56,11 @@ struct invocation_request {
+      */
+     std::string function_arn;
+ 
++    /**
++     * The Tenant ID of the current invocation.
++     */
++    std::string tenant_id;
++
+     /**
+      * Function execution deadline counted in milliseconds since the Unix epoch.
+      */
+diff --git a/src/runtime.cpp b/src/runtime.cpp
+index d1e655f..bcc217d 100644
+--- a/src/runtime.cpp
++++ b/src/runtime.cpp
+@@ -40,6 +40,7 @@ static constexpr auto CLIENT_CONTEXT_HEADER = "lambda-runtime-client-context";
+ static constexpr auto COGNITO_IDENTITY_HEADER = "lambda-runtime-cognito-identity";
+ static constexpr auto DEADLINE_MS_HEADER = "lambda-runtime-deadline-ms";
+ static constexpr auto FUNCTION_ARN_HEADER = "lambda-runtime-invoked-function-arn";
++static constexpr auto TENANT_ID_HEADER = "lambda-runtime-aws-tenant-id";
+ 
+ enum Endpoints {
+     INIT,
+@@ -296,6 +297,11 @@ runtime::next_outcome runtime::get_next()
+         req.function_arn = std::move(out).get_result();
+     }
+ 
++    out = resp.get_header(TENANT_ID_HEADER);
++    if (out.is_success()) {
++        req.tenant_id = std::move(out).get_result();
++    }
++
+     out = resp.get_header(DEADLINE_MS_HEADER);
+     if (out.is_success()) {
+         auto const& deadline_string = std::move(out).get_result();

deps/patches/aws-lambda-cpp-add-tenant-id.patch

@@ -29,7 +29,8 @@ wget -c https://github.com/awslabs/aws-lambda-cpp/archive/refs/tags/v$AWS_LAMBDA
 # Apply patches
 (
   cd aws-lambda-cpp-$AWS_LAMBDA_CPP_RELEASE && \
-    patch -p1 < ../patches/aws-lambda-cpp-add-xray-response.patch
+    patch -p1 < ../patches/aws-lambda-cpp-add-xray-response.patch && \
+    patch -p1 < ../patches/aws-lambda-cpp-add-tenant-id.patch
 )
 
 # Pack again and remove the folder

scripts/update_dependencies.sh

@@ -9,7 +9,7 @@
 'use strict';
 
 const assert = require('assert').strict;
-let { setCurrentRequestId } = require('./LogPatch');
+let { setCurrentRequestId, setCurrentTenantId } = require('./LogPatch');
 
 const INVOKE_HEADER = {
   ClientContext: 'lambda-runtime-client-context',
@@ -18,6 +18,7 @@ const INVOKE_HEADER = {
   AWSRequestId: 'lambda-runtime-aws-request-id',
   DeadlineMs: 'lambda-runtime-deadline-ms',
   XRayTrace: 'lambda-runtime-trace-id',
+  TenantId: 'lambda-runtime-aws-tenant-id',
 };
 
 module.exports = class InvokeContext {
@@ -34,11 +35,19 @@ module.exports = class InvokeContext {
     return id;
   }
 
+  /**
+   * The tenantId for this request.
+   */
+  get tenantId() {
+    return this.headers[INVOKE_HEADER.TenantId];
+  }
+
   /**
    * Push relevant invoke data into the logging context.
    */
   updateLoggingContext() {
     setCurrentRequestId(this.invokeId);
+    setCurrentTenantId(this.tenantId);
   }
 
   /**
@@ -91,6 +100,7 @@ module.exports = class InvokeContext {
       ),
       invokedFunctionArn: this.headers[INVOKE_HEADER.ARN],
       awsRequestId: this.headers[INVOKE_HEADER.AWSRequestId],
+      tenantId: this.headers[INVOKE_HEADER.TenantId],
       getRemainingTimeInMillis: function () {
         return deadline - Date.now();
       },

src/InvokeContext.js

@@ -28,13 +28,23 @@ const jsonErrorReplacer = (_, value) => {
   return value;
 };
 
-function formatJsonMessage(requestId, timestamp, level, ...messageParams) {
+function formatJsonMessage(
+  requestId,
+  timestamp,
+  level,
+  tenantId,
+  ...messageParams
+) {
   let result = {
     timestamp: timestamp,
     level: level.name,
     requestId: requestId,
   };
 
+  if (tenantId != undefined && tenantId != null) {
+    result.tenantId = tenantId;
+  }
+
   if (messageParams.length === 1) {
     result.message = messageParams[0];
     try {
@@ -65,6 +75,13 @@ let _currentRequestId = {
   set: (id) => (global[REQUEST_ID_SYMBOL] = id),
 };
 
+/* Use a unique symbol to provide global access without risk of name clashes. */
+const TENANT_ID_SYMBOL = Symbol.for('aws.lambda.runtime.tenantId');
+let _currentTenantId = {
+  get: () => global[TENANT_ID_SYMBOL],
+  set: (id) => (global[TENANT_ID_SYMBOL] = id),
+};
+
 /**
  * Write logs to stdout.
  */
@@ -82,7 +99,15 @@ let logTextToStdout = (level, message, ...params) => {
 let logJsonToStdout = (level, message, ...params) => {
   let time = new Date().toISOString();
   let requestId = _currentRequestId.get();
-  let line = formatJsonMessage(requestId, time, level, message, ...params);
+  let tenantId = _currentTenantId.get();
+  let line = formatJsonMessage(
+    requestId,
+    time,
+    level,
+    tenantId,
+    message,
+    ...params,
+  );
   line = line.replace(/\n/g, '\r');
   process.stdout.write(line + '\n');
 };
@@ -125,10 +150,12 @@ let logJsonToFd = function (logTarget) {
     let date = new Date();
     let time = date.toISOString();
     let requestId = _currentRequestId.get();
+    let tenantId = _currentTenantId.get();
     let enrichedMessage = formatJsonMessage(
       requestId,
       time,
       level,
+      tenantId,
       message,
       ...params,
     );
@@ -239,6 +266,7 @@ let _patchConsole = () => {
 
 module.exports = {
   setCurrentRequestId: _currentRequestId.set,
+  setCurrentTenantId: _currentTenantId.set,
   patchConsole: _patchConsole,
   structuredConsole: structuredConsole,
 };

src/LogPatch.js

@@ -82,6 +82,11 @@ class RuntimeApiNextPromiseWorker : public Napi::AsyncWorker {
                     Napi::String::New(env, "lambda-runtime-cognito-identity"),
                     Napi::String::New(env, response.cognito_identity.c_str()));
         }
+        if (response.tenant_id != "") {
+            headers.Set(
+                    Napi::String::New(env, "lambda-runtime-aws-tenant-id"),
+                    Napi::String::New(env, response.tenant_id.c_str()));
+        }
 
         auto ret = Napi::Object::New(env);
         ret.Set(Napi::String::New(env, "bodyJson"), response_data);

src/rapid-client.cc

@@ -35,3 +35,32 @@ describe('Getting remaining invoke time', () => {
     remainingTime.should.lessThanOrEqual(1000);
   });
 });
+
+describe('Verifying tenant id', () => {
+  it('should return undefined if tenant id is not set', () => {
+    let ctx = new InvokeContext({});
+
+    (ctx._headerData().tenantId === undefined).should.be.true();
+  });
+  it('should return undefined if tenant id is set to undefined', () => {
+    let ctx = new InvokeContext({
+      'lambda-runtime-aws-tenant-id': undefined,
+    });
+
+    (ctx._headerData().tenantId === undefined).should.be.true();
+  });
+  it('should return empty if tenant id is set to empty string', () => {
+    let ctx = new InvokeContext({
+      'lambda-runtime-aws-tenant-id': '',
+    });
+
+    (ctx._headerData().tenantId === '').should.be.true();
+  });
+  it('should return the same id if a valid tenant id is set', () => {
+    let ctx = new InvokeContext({
+      'lambda-runtime-aws-tenant-id': 'blue',
+    });
+
+    ctx._headerData().tenantId.should.equal('blue');
+  });
+});

test/unit/InvokeContextTest.js

@@ -415,9 +415,38 @@ describe('The multiline log patch', () => {
         );
         receivedMessage.should.have.property('level', logFunctions[fIdx][1]);
         receivedMessage.should.have.property('requestId', EXPECTED_ID);
+        receivedMessage.should.not.have.property('tenantId');
       }
     });
 
+    it('should format messages with tenant id as json correctly', () => {
+      const EXPECTED_TENANT_ID = 'tenantId';
+      LogPatch.setCurrentTenantId(EXPECTED_TENANT_ID);
+
+      for (let fIdx = 0; fIdx < logFunctions.length; fIdx++) {
+        logFunctions[fIdx][0]('structured logging with tenant id');
+        let receivedMessage = telemetryTarget.readLine(
+          logFunctions[fIdx][1],
+          'JSON',
+        );
+        receivedMessage = JSON.parse(receivedMessage);
+
+        receivedMessage.should.have.property('timestamp');
+        let receivedTime = new Date(receivedMessage.timestamp);
+        let now = new Date();
+        assert(now >= receivedTime && now - receivedTime <= 1000);
+
+        receivedMessage.should.have.property(
+          'message',
+          'structured logging with tenant id',
+        );
+        receivedMessage.should.have.property('level', logFunctions[fIdx][1]);
+        receivedMessage.should.have.property('requestId', EXPECTED_ID);
+        receivedMessage.should.have.property('tenantId', EXPECTED_TENANT_ID);
+      }
+      LogPatch.setCurrentTenantId(undefined);
+    });
+
     it('should filter messages correctly', () => {
       const loglevelSettings = [
         undefined,

test/unit/LogPatchTest.js

@@ -47,6 +47,23 @@ class NoOpNativeHttp {
   }
 }
 
+class MockNativeClient {
+  constructor(response) {
+    this.response = response;
+    this.called = false;
+    this.shouldThrowError = false;
+  }
+
+  next() {
+    this.called = true;
+    if (this.shouldThrowError) {
+      return Promise.reject(new Error('Failed to get next invocation'));
+    } else {
+      return Promise.resolve(this.response);
+    }
+  }
+}
+
 class EvilError extends Error {
   get name() {
     throw 'gotcha';
@@ -115,3 +132,82 @@ describe('invalid request id works', () => {
     });
   });
 });
+
+describe('next invocation with native client works', () => {
+  it('should call the native client next() method', async () => {
+    const mockNative = new MockNativeClient({
+      bodyJson: '',
+      headers: {
+        'lambda-runtime-aws-request-id': 'test-request-id',
+      },
+    });
+    const client = new RAPIDClient('notUsed:1337', undefined, mockNative);
+    client.useAlternativeClient = false;
+
+    await client.nextInvocation();
+    // verify native client was called
+    mockNative.called.should.be.true();
+  });
+  it('should parse all required headers', async () => {
+    const mockResponse = {
+      bodyJson: '{"message":"Hello from Lambda!"}',
+      headers: {
+        'lambda-runtime-aws-request-id': 'test-request-id',
+        'lambda-runtime-deadline-ms': 1619712000000,
+        'lambda-runtime-trace-id': 'test-trace-id',
+        'lambda-runtime-invoked-function-arn': 'test-function-arn',
+        'lambda-runtime-client-context': '{"client":{"app_title":"MyApp"}}',
+        'lambda-runtime-cognito-identity':
+          '{"identityId":"id123","identityPoolId":"pool123"}',
+        'lambda-runtime-aws-tenant-id': 'test-tenant-id',
+      },
+    };
+
+    const mockNative = new MockNativeClient(mockResponse);
+    const client = new RAPIDClient('notUsed:1337', undefined, mockNative);
+
+    client.useAlternativeClient = false;
+    const response = await client.nextInvocation();
+
+    // Verify all headers are present
+    response.headers.should.have.property(
+      'lambda-runtime-aws-request-id',
+      'test-request-id',
+    );
+    response.headers.should.have.property(
+      'lambda-runtime-deadline-ms',
+      1619712000000,
+    );
+    response.headers.should.have.property(
+      'lambda-runtime-trace-id',
+      'test-trace-id',
+    );
+    response.headers.should.have.property(
+      'lambda-runtime-invoked-function-arn',
+      'test-function-arn',
+    );
+    response.headers.should.have.property(
+      'lambda-runtime-client-context',
+      '{"client":{"app_title":"MyApp"}}',
+    );
+    response.headers.should.have.property(
+      'lambda-runtime-cognito-identity',
+      '{"identityId":"id123","identityPoolId":"pool123"}',
+    );
+    response.headers.should.have.property(
+      'lambda-runtime-aws-tenant-id',
+      'test-tenant-id',
+    );
+    // Verify body is correctly passed through
+    response.bodyJson.should.equal('{"message":"Hello from Lambda!"}');
+  });
+  it('should handle native client errors', async () => {
+    const nativeClient = new MockNativeClient({});
+    nativeClient.shouldThrowError = true;
+
+    const client = new RAPIDClient('localhost:8080', null, nativeClient);
+    client.useAlternativeClient = false;
+
+    await client.nextInvocation().should.be.rejectedWith('Failed to get next invocation');
+  });
+});