async_hooks: use correct resource for AsyncLocalStorage

Fixes: nodejs#40693

Signed-off-by: Darshan Sen <darshan.sen@postman.com>

lib/async_hooks.js

@@ -252,10 +252,29 @@ class AsyncResource {
   }
 }
 
+function lookupResourceWithStorage() {
+  // When a TCP/UDP object is created, its owner_symbol is uninitialized, so
+  // calling executionAsyncResource() would return the TCP/UDP object itself.
+  // So the store is initially saved on the TCP/UDP object. Later on, when the
+  // Socket/TLSSocket object is created, the TCP/UDP object is wrapped by the
+  // Socket/TLSSocket object in the _handle property and the owner_symbol of the
+  // TCP/UDP object is set to the corresponding Socket/TLSSocket object. So
+  // calling executionAsyncResource() now would return the Socket/TLSSocket
+  // object, which does not contain the storage that was saved initially. Hence,
+  // in the case of a Socket/TLSSocket object, we must use the underlying
+  // TCP/UDP object, if available, for storage.
+  let resource = executionAsyncResource();
+  if ((resource.constructor.name === 'Socket' ||
+       resource.constructor.name === 'TLSSocket') &&
+      resource._handle != null)
+    resource = resource._handle;
+  return resource;
+}
+
 const storageList = [];
 const storageHook = createHook({
   init(asyncId, type, triggerAsyncId, resource) {
-    const currentResource = executionAsyncResource();
+    const currentResource = lookupResourceWithStorage();
     // Value of currentResource is always a non null object
     for (let i = 0; i < storageList.length; ++i) {
       storageList[i]._propagate(resource, currentResource);
@@ -299,7 +318,7 @@ class AsyncLocalStorage {
 
   enterWith(store) {
     this._enable();
-    const resource = executionAsyncResource();
+    const resource = lookupResourceWithStorage();
     resource[this.kResourceStore] = store;
   }
 
@@ -311,7 +330,7 @@ class AsyncLocalStorage {
 
     this._enable();
 
-    const resource = executionAsyncResource();
+    const resource = lookupResourceWithStorage();
     const oldStore = resource[this.kResourceStore];
 
     resource[this.kResourceStore] = store;
@@ -337,7 +356,7 @@ class AsyncLocalStorage {
 
   getStore() {
     if (this.enabled) {
-      const resource = executionAsyncResource();
+      const resource = lookupResourceWithStorage();
       return resource[this.kResourceStore];
     }
   }

test/async-hooks/test-async-local-storage-dgram.js

@@ -0,0 +1,24 @@
+'use strict';
+
+require('../common');
+
+const assert = require('assert');
+const dgram = require('dgram');
+const { AsyncLocalStorage } = require('async_hooks');
+
+dgram.createSocket('udp4')
+  .on('message', function(msg, rinfo) { this.send(msg, rinfo.port); })
+  .on('listening', function() {
+    const asyncLocalStorage = new AsyncLocalStorage();
+    const store = { val: 'abcd' };
+    asyncLocalStorage.run(store, () => {
+      const client = dgram.createSocket('udp4');
+      client.on('message', (msg, rinfo) => {
+        assert.deepStrictEqual(asyncLocalStorage.getStore(), store);
+        client.close();
+        this.close();
+      });
+      client.send('Hello, world!', this.address().port);
+    });
+  })
+  .bind(0);

test/async-hooks/test-async-local-storage-socket.js

@@ -0,0 +1,25 @@
+'use strict';
+
+require('../common');
+
+const assert = require('assert');
+const net = require('net');
+const { AsyncLocalStorage } = require('async_hooks');
+
+net
+  .createServer((socket) => {
+    socket.write('Hello, world!');
+    socket.pipe(socket);
+  })
+  .listen(0, function() {
+    const asyncLocalStorage = new AsyncLocalStorage();
+    const store = { val: 'abcd' };
+    asyncLocalStorage.run({ val: 'abcd' }, () => {
+      const client = net.connect({ port: this.address().port });
+      client.on('data', () => {
+        assert.deepStrictEqual(asyncLocalStorage.getStore(), store);
+        client.end();
+        this.close();
+      });
+    });
+  });

test/async-hooks/test-async-local-storage-tlssocket.js

@@ -0,0 +1,34 @@
+'use strict';
+
+const common = require('../common');
+if (!common.hasCrypto)
+  common.skip('missing crypto');
+
+const assert = require('assert');
+const fixtures = require('../common/fixtures');
+const tls = require('tls');
+const { AsyncLocalStorage } = require('async_hooks');
+
+const options = {
+  cert: fixtures.readKey('rsa_cert.crt'),
+  key: fixtures.readKey('rsa_private.pem'),
+  rejectUnauthorized: false
+};
+
+tls
+  .createServer(options, (socket) => {
+    socket.write('Hello, world!');
+    socket.pipe(socket);
+  })
+  .listen(0, function() {
+    const asyncLocalStorage = new AsyncLocalStorage();
+    const store = { val: 'abcd' };
+    asyncLocalStorage.run({ val: 'abcd' }, () => {
+      const client = tls.connect({ port: this.address().port, ...options });
+      client.on('data', () => {
+        assert.deepStrictEqual(asyncLocalStorage.getStore(), store);
+        client.end();
+        this.close();
+      });
+    });
+  });