@contrast/sources 1.4.0 → 1.5.0

package/lib/index.js

@@ -23,6 +23,11 @@ const NormalizedUriMapper = require('./normalized-uri-mapper');
 const { HttpSourceInfo } = require('./source-info');
 
 const componentName = 'sources';
+const SourceType = {
+  HTTP: 'HTTP',
+  WEBSOCKET: 'WEBSOCKET',
+};
+
 
 module.exports = Core.makeComponent({
   name: componentName,
@@ -48,20 +53,40 @@ class Sources {
     const { _hooks, _normalizedUriMapper, core } = this;
 
     return function (next, data) {
-      const { args: [event, req, res] } = data;
-
-      if (event !== 'request') {
-        if (event === 'listening') {
-          // take a snapshot of Perf.all at this point. this will get logged
-          // at some point on the perf interval timer.
-          core.Perf.mark('listening');
-          core.messages.emit(Event.SERVER_LISTENING, { type: serverType, server: data.obj });
-        }
+      const { args: [event, req, resOrSocket] } = data;
+
+      if (event === 'listening') {
+        // take a snapshot of Perf.all at this point. this will get logged
+        // at some point on the perf interval timer.
+        core.Perf.mark('listening');
+        core.messages.emit(Event.SERVER_LISTENING, { type: serverType, server: data.obj });
         return next();
       }
 
-      core.Perf.requestCount += 1;
+      const isUpgrade = event === 'upgrade';
+      let protocol = serverType == 'http' ? 'http' : 'https';
+
+      if (isUpgrade) {
+        for (let i = 0; i < req.rawHeaders.length; i += 2) {
+          if (req.rawHeaders[i].toLowerCase?.() == 'upgrade') {
+            protocol = req.rawHeaders[i + 1]?.toLowerCase?.();
+            break;
+          }
+        }
+      }
+
+      // For HTTP servers we run the "request" and "upgrade" events in a source
+      // scope. We only support "websocket" upgrades currently, but this can be
+      // extended. Support for non-HTTP sources is expected and will require new
+      // instrumentation and possibly model alterations.
+      if (
+        event !== 'request' &&
+        (!isUpgrade || (isUpgrade && protocol !== 'websocket'))
+      ) {
+        return next();
+      }
 
+      // websocket sources are http
       const sourceInfo = new HttpSourceInfo({
         serverType,
         raw: req,
@@ -69,20 +94,38 @@ class Sources {
       });
       const store = { sourceInfo };
 
-      onFinished(res, (/* err, req */) => {
-        core.messages.emit(Event.RESPONSE_FINISH, store);
-      });
+      if (isUpgrade) {
+        core.patcher.patch(resOrSocket, 'emit', {
+          name: `${serverType}.socket.emit`,
+          patchType: 'sources',
+          around(next) {
+            return core.scopes.sources.run(store, next);
+          },
+        });
+      } else {
+        onFinished(resOrSocket, (/* err, req */) => {
+          core.messages.emit(Event.RESPONSE_FINISH, store);
+        });
+      }
+
+      core.Perf.requestCount += 1;
 
       return core.scopes.sources.run(store, () => {
-        if (_hooks._events.onSource) {
-          _hooks.emit('onSource', {
-            // future: non-http sources will have their own type
-            sourceType: 'HTTP',
-            store,
-            incomingMessage: req,
-            serverResponse: res,
-          });
-        }
+        const sourceType = protocol == 'websocket' ? SourceType.WEBSOCKET : SourceType.HTTP;
+        const eventArg = {
+          sourceType,
+          store,
+          incomingMessage: req,
+        };
+
+        if (sourceType == SourceType.HTTP)
+          eventArg.serverResponse = resOrSocket;
+
+        else if (sourceType == SourceType.WEBSOCKET)
+          eventArg.socket = resOrSocket;
+
+        if (_hooks._events.onSource)
+          _hooks.emit('onSource', eventArg);
 
         return next();
       });

package/lib/index.test.js

@@ -8,6 +8,38 @@ const mocks = require('@contrast/test/mocks');
 const proxyquire = require('proxyquire');
 
 describe('agentify sources', function () {
+  let core, api, ServerMock, reqMock, resMock, onFinishedMock;
+
+  beforeEach(function () {
+    ({ core } = initProtectFixture());
+    ServerMock = function ServerMock() {
+      this.e = new EventEmitter();
+    };
+    ServerMock.prototype.emit = function (...args) {
+      this.e.emit(...args);
+    };
+    ServerMock.prototype.on = function (...args) {
+      this.e.on(...args);
+    };
+    api = {
+      Server: ServerMock,
+      createServer() {
+        return new ServerMock();
+      },
+      createSecureServer() {
+        return new ServerMock();
+      }
+    };
+    reqMock = mocks.incomingMessage();
+    // resMock = new EventEmitter();
+    onFinishedMock = sinon.stub();
+
+    core.depHooks.resolve.withArgs(sinon.match({ name: 'http' })).yields(api);
+    proxyquire('.', {
+      'on-finished': onFinishedMock,
+    })(core).install();
+  });
+
   [
     {
       name: 'http',
@@ -24,41 +56,9 @@ describe('agentify sources', function () {
         protocol: 'https',
         serverType: 'https',
       },
-    }
+    },
   ].forEach(({ name, method, expected }) => {
     describe(`${name} sources using ${method || 'Server'}()`, function () {
-      let core, api, ServerMock, reqMock, resMock, onFinishedMock;
-
-      beforeEach(function () {
-        ({ core } = initProtectFixture());
-        ServerMock = function ServerMock() {
-          this.e = new EventEmitter();
-        };
-        ServerMock.prototype.emit = function (...args) {
-          this.e.emit(...args);
-        };
-        ServerMock.prototype.on = function (...args) {
-          this.e.on(...args);
-        };
-        api = {
-          Server: ServerMock,
-          createServer() {
-            return new ServerMock();
-          },
-          createSecureServer() {
-            return new ServerMock();
-          }
-        };
-        reqMock = mocks.incomingMessage();
-        // resMock = new EventEmitter();
-        onFinishedMock = sinon.stub();
-
-        core.depHooks.resolve.withArgs(sinon.match({ name: 'http' })).yields(api);
-        proxyquire('.', {
-          'on-finished': onFinishedMock,
-        })(core).install();
-      });
-
       it('"request" events run in scope with correct sourceInfo', function () {
         const server = method ? api[method]() : new ServerMock();
         let store;
@@ -88,7 +88,52 @@ describe('agentify sources', function () {
 
         server.emit('foo', reqMock, resMock);
         expect(store).to.be.undefined;
+        expect(onFinishedMock).not.to.have.been.called;
       });
     });
   });
+
+  it('"upgrade" events run in scope with correct sourceInfo', function () {
+    const server = new ServerMock();
+    const socketMock = new EventEmitter();
+    let store;
+    let socketEventStore;
+
+    reqMock.rawHeaders.push('upgrade', 'websocket');
+
+    server.on('upgrade', function () {
+      store = core.scopes.sources.getStore();
+    });
+
+    server.emit('upgrade', reqMock, socketMock);
+
+    expect(store.sourceInfo).to.deep.include({
+      port: 8080,
+      protocol: 'http',
+      serverType: 'http',
+    });
+
+    socketMock.on('some-message', function() {
+      socketEventStore = core.scopes.sources.getStore();
+    });
+    socketMock.emit('some-message');
+
+    expect(store).to.equal(socketEventStore);
+  });
+
+  it('"upgrade" events do not run in scope if not "websocket"', function () {
+    const server = new ServerMock();
+    const socketMock = new EventEmitter();
+    let store;
+
+    reqMock.rawHeaders.push('upgrade', 'h2c');
+
+    server.on('upgrade', function () {
+      store = core.scopes.sources.getStore();
+    });
+
+    server.emit('upgrade', reqMock, socketMock);
+
+    expect(store).to.be.undefined;
+  });
 });

package/lib/source-info.js

@@ -53,7 +53,7 @@ class HttpSourceInfo {
     this.httpVersion = raw.httpVersion;
     this.ip = raw.socket.remoteAddress ? StringPrototypeReplace.call(raw.socket.remoteAddress, /::ffff:/, '') : undefined;
     this.port = raw.socket.address?.()?.port || 0;
-    this.protocol = serverType == 'http' ? 'http' : 'https'; // todo
+    this.protocol = serverType == 'http' ? 'http' : 'https';
     this.serverType = serverType;
     this.time = Date.now();
     this.method = StringPrototypeToLowerCase.call(raw.method);
@@ -62,20 +62,11 @@ class HttpSourceInfo {
     for (let i = 0; i < raw.rawHeaders.length; i += 2) {
       const iNext = i + 1;
       const headerName = StringPrototypeToLowerCase.call(raw.rawHeaders[i]);
-
       headerName == 'content-type' && (this.contentType = raw.rawHeaders[iNext]);
-
       this.rawHeaders[i] = headerName;
       this.rawHeaders[iNext] = headerName == 'content-type' ?
         StringPrototypeToLowerCase.call(raw.rawHeaders[iNext]) :
         raw.rawHeaders[iNext];
-
-      if (
-        headerName == 'upgrade' &&
-        StringPrototypeToLowerCase.call(this.rawHeaders[iNext]) == 'websocket'
-      ) {
-        this.protocol = 'ws';
-      }
     }
 
     const idx = raw.url.indexOf('?');

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@contrast/sources",
-  "version": "1.4.0",
+  "version": "1.5.0",
   "description": "Instruments to have incoming messages run in async-local request scope.",
   "main": "lib/index.js",
   "scripts": {
@@ -9,8 +9,8 @@
   "author": "",
   "license": "ISC",
   "dependencies": {
-    "@contrast/common": "1.37.0",
-    "@contrast/core": "1.58.0",
+    "@contrast/common": "1.38.0",
+    "@contrast/core": "1.59.0",
     "on-finished": "^2.4.1"
   }
 }