Return X-Ray trace ID in HTTP header

Author: juffalow

src/routes/upload.ts

@@ -62,7 +62,9 @@ router.post('/', optionalAuth, uploadMiddleware, validateRequest(postSchema), as
       data: {
         file,
       }
-    }).end();
+    });
+    
+     next();
   } catch (err) {
     next(err);
   }

src/services/index.ts

@@ -67,7 +67,7 @@ const container = {
 
   get ServiceDiscovery() {
     return new ServiceDiscovery();
-  }
+  },
 };
 
 export default container;

src/services/trace/AWSXRay.ts

@@ -1,14 +1,23 @@
 import XRay from 'aws-xray-sdk';
 import http from 'http';
 import https from 'https';
+import type { Request, Response, NextFunction } from 'express';
 
 class AWSXRay implements Services.Trace {
   constructor (plugins: string) {
     this.setPlugins(plugins);
   }
 
   public openSegment(defaultName: string) {
-    return XRay.express.openSegment(defaultName);
+    const traceMiddleware = (req: Request, res: Response, next: NextFunction) => {
+      const traceId = this.getTraceId();
+
+      res.setHeader('x-request-id', traceId);
+
+      next();
+    };
+
+    return [ XRay.express.openSegment(defaultName), traceMiddleware ];
   }
 
   public closeSegment() {

src/services/trace/CLSHooked.ts

@@ -9,11 +9,11 @@ class CLSHooked implements Services.Trace {
     const namespace = this.getNamespace();
 
     return (req: Request, res: Response, next: NextFunction) => {
-      const traceId = req.query.traceId ? req.query.traceId as string : uuidv7();
+      const traceId = req.headers['x-request-id'] as string || req.query.traceId as string || uuidv7();
 
       namespace.run(() => {
         namespace.set('traceId', traceId);
-        res.header('x-request-id', traceId);
+        res.setHeader('x-request-id', traceId);
 
         next();
       });