From a4e608b9633c85ed2a208030d49a1a5e0719cbde Mon Sep 17 00:00:00 2001
From: Blayne Chard <bchard@linz.govt.nz>
Date: Tue, 26 Jul 2022 12:29:15 +1200
Subject: [PATCH] fix: ensure logging the request is the last thing that is
 done before sending a response

---
 src/function.ts                       | 27 ++++++++++++++-------------
 src/http/__test__/router.hook.test.ts | 26 ++++++++++++++++++++++++--
 src/http/router.ts                    |  6 +++---
 3 files changed, 41 insertions(+), 18 deletions(-)

diff --git a/src/function.ts b/src/function.ts
index df4a85b..dc34cf1 100644
--- a/src/function.ts
+++ b/src/function.ts
@@ -37,6 +37,8 @@ export async function runFunction<T extends LambdaRequest, K>(
   req: T,
   fn: (req: T) => K | Promise<K>,
 ): Promise<K | LambdaHttpResponse> {
+  if (!req.timer.timers.has('lambda')) req.timer.start('lambda');
+
   try {
     return await fn(req);
   } catch (error) {
@@ -49,14 +51,7 @@ export async function runFunction<T extends LambdaRequest, K>(
   }
 }
 
-export async function execute<T extends LambdaRequest, K>(
-  req: T,
-  fn: (req: T) => K | Promise<K>,
-): Promise<K | LambdaHttpResponse> {
-  req.timer.start('lambda');
-
-  const res = await runFunction(req, fn);
-
+function after<T extends LambdaRequest, K>(req: T, res: K): K | LambdaHttpResponse {
   let status = 200;
   if (LambdaHttpResponse.is(res)) {
     status = res.status;
@@ -70,7 +65,7 @@ export async function execute<T extends LambdaRequest, K>(
 
   if (versionInfo.hash) req.set('package', versionInfo);
 
-  const duration = req.timer.end('lambda');
+  const duration = req.timer.timers.has('lambda') ? req.timer.end('lambda') : 0;
   req.set('unfinished', req.timer.unfinished);
   req.set('duration', duration);
 
@@ -154,10 +149,14 @@ export class lf {
       const lambdaId = context.awsRequestId;
       req.set('aws', { lambdaId });
 
-      execute(req, fn).then((res) => {
-        if (opts.rejectOnError && LambdaHttpResponse.is(res)) {
-          if (req.logContext['err']) return callback(req.logContext['err'] as Error);
-          if (res.status > 399) return callback(req.toResponse(res) as unknown as string);
+      runFunction(req, fn).then((res) => {
+        after(req, res);
+
+        if (LambdaHttpResponse.is(res)) {
+          if (opts.rejectOnError) {
+            if (req.logContext['err']) return callback(req.logContext['err'] as Error);
+            if (res.status > 399) return callback(req.toResponse(res) as unknown as string);
+          }
         }
         return callback(null, req.toResponse(res));
       });
@@ -202,6 +201,8 @@ export class lf {
 
         if (lf.ServerName) res.header(HttpHeader.Server, `${lf.ServerName}-${version}`);
 
+        after(req, res);
+
         callback(null, req.toResponse(res));
       });
     }
diff --git a/src/http/__test__/router.hook.test.ts b/src/http/__test__/router.hook.test.ts
index 22d2c89..3ca953f 100644
--- a/src/http/__test__/router.hook.test.ts
+++ b/src/http/__test__/router.hook.test.ts
@@ -1,7 +1,8 @@
-import { Context } from 'aws-lambda';
+import { ALBResult, Context } from 'aws-lambda';
 import o from 'ospec';
 import sinon from 'sinon';
-import { UrlExample } from '../../__test__/examples.js';
+import { lf } from '../../function.js';
+import { AlbExample, UrlExample } from '../../__test__/examples.js';
 import { fakeLog } from '../../__test__/log.js';
 import { LambdaUrlRequest } from '../request.url.js';
 import { LambdaHttpResponse } from '../response.http.js';
@@ -94,5 +95,26 @@ o.spec('RouterHook', () => {
       o(res.status).equals(500);
       o(res.statusDescription).equals('Internal Server Error');
     });
+
+    o('should log after the response hook', async () => {
+      fakeLog.logs = [];
+      const http = lf.http(fakeLog);
+
+      http.router.hook('response', (req, res) => {
+        if (res.status !== 404) throw new Error('status should be 404');
+        res.status = 200; // Convert the response to a 200!
+        req.set('logParam', 'response'); // Add a new log parameter to be logged
+      });
+
+      const res = await new Promise<ALBResult>((r) => http(AlbExample, fakeContext, (err, res) => r(res as ALBResult)));
+
+      o(fakeLog.logs.length).equals(1);
+      const [firstLog] = fakeLog.logs;
+      o(firstLog.logParam).equals('response');
+      o(firstLog['@type']).equals('report');
+      o(firstLog['status']).equals(200);
+
+      o(res.statusCode).equals(200);
+    });
   });
 });
diff --git a/src/http/router.ts b/src/http/router.ts
index bfc5ff7..a103e4b 100644
--- a/src/http/router.ts
+++ b/src/http/router.ts
@@ -1,7 +1,7 @@
-import { execute, runFunction } from '../function.js';
+import FindMyWay from 'find-my-way';
+import { runFunction } from '../function.js';
 import { LambdaHttpRequest, RequestTypes } from './request.http.js';
 import { LambdaHttpResponse } from './response.http.js';
-import FindMyWay from 'find-my-way';
 
 export type Route<T extends RequestTypes = RequestTypes> = (
   req: LambdaHttpRequest<T>,
@@ -41,7 +41,7 @@ export class Router {
     this.router.on(method, path, async (req: unknown, res, params) => {
       if (!(req instanceof LambdaHttpRequest)) return new LambdaHttpResponse(500, 'Internal server error');
       req.params = params;
-      return execute(req, fn);
+      return runFunction(req, fn);
     });
   }