From 32e4e92284e00ce6b272fef6f05ba88d804f4dde Mon Sep 17 00:00:00 2001
From: Trent Mick <trent.mick@elastic.co>
Date: Wed, 1 Sep 2021 14:40:59 -0700
Subject: [PATCH] fix: make 'elasticsearch' instrumentation host/port parsing
 more resilient (#2312)

In particular support the `host` or `hosts` config options being arrays.

Refs: #2306
---
 CHANGELOG.asciidoc                            |  6 +-
 lib/instrumentation/modules/elasticsearch.js  | 56 ++++++++++++++++++-
 .../modules/elasticsearch.test.js             | 37 ++++++++++++
 3 files changed, 95 insertions(+), 4 deletions(-)

diff --git a/CHANGELOG.asciidoc b/CHANGELOG.asciidoc
index 1d8bec88c0f..c42484379c4 100644
--- a/CHANGELOG.asciidoc
+++ b/CHANGELOG.asciidoc
@@ -44,6 +44,10 @@ Notes:
 [float]
 ===== Bug fixes
 
+* Fix a crash in instrumentation of the old Elasticsearch client
+  (`elasticsearch`) for some rarer cases of client options -- for example
+  passing multiple hosts. ({pull}2312[#2312])
+
 * Ensure the internal HTTP(S) client requests made by the APM agent to APM
   server are not themselves traced. ({issues}1168[#1168], {issues}1136[#1136])
 
@@ -90,7 +94,7 @@ Notes:
 
 * Add instrumentation of the AWS SNS publish method when using the
   https://www.npmjs.com/package/aws-sdk[JavaScript AWS SDK v2] (`aws-sdk`). ({pull}2157[#2157])
-  
+
 [float]
 ===== Bug fixes
 
diff --git a/lib/instrumentation/modules/elasticsearch.js b/lib/instrumentation/modules/elasticsearch.js
index b09793fd292..d71c57fe736 100644
--- a/lib/instrumentation/modules/elasticsearch.js
+++ b/lib/instrumentation/modules/elasticsearch.js
@@ -1,9 +1,58 @@
 'use strict'
 
+const URL = require('url').URL
+
 var shimmer = require('../shimmer')
 var { getDBDestination } = require('../context')
 const { setElasticsearchDbContext } = require('../elasticsearch-shared')
 
+const startsWithProtocolRE = /^([a-z]+:)?\/\//i
+const DEFAULT_PORT = 9200
+const DEFAULT_PORT_FROM_PROTO = {
+  'http:': 80,
+  'https:': 443
+}
+
+// This is an imperfect equivalent of the handling in the `Transport`
+// constructor and internal `Host` parsing function in the ES client.
+function getHostAndPortFromTransportConfig (config) {
+  const transportHosts = config ? config.host || config.hosts : null
+  if (!transportHosts) {
+    return null
+  }
+
+  let firstTransportHost = Array.isArray(transportHosts)
+    ? transportHosts[0] : transportHosts
+  if (!firstTransportHost) {
+    return null
+  }
+
+  if (typeof firstTransportHost === 'string') {
+    // "example.com:42" or "someprotocol://example.com:42" or
+    // "someprotocol://example.com".
+    if (!startsWithProtocolRE.test(firstTransportHost)) {
+      firstTransportHost = 'http://' + firstTransportHost
+    }
+    let u
+    try {
+      u = new URL(firstTransportHost)
+    } catch (_err) {
+      return null
+    }
+    if (!u.port) {
+      u.port = DEFAULT_PORT_FROM_PROTO[u.protocol]
+    }
+    return [u.hostname, u.port]
+  } else if (typeof firstTransportHost === 'object') {
+    return [
+      firstTransportHost.hostname || firstTransportHost.host,
+      firstTransportHost.port || DEFAULT_PORT
+    ]
+  }
+
+  return null
+}
+
 module.exports = function (elasticsearch, agent, { enabled }) {
   if (!enabled) return elasticsearch
 
@@ -28,10 +77,11 @@ module.exports = function (elasticsearch, agent, { enabled }) {
           params && params.body)
 
         // Get the remote host information from elasticsearch Transport options.
-        const transportConfig = this._config
         let host, port
-        if (typeof transportConfig === 'object' && transportConfig.host) {
-          [host, port] = transportConfig.host.split(':')
+        const hostAndPort = getHostAndPortFromTransportConfig(this._config)
+        if (hostAndPort) {
+          host = hostAndPort[0]
+          port = hostAndPort[1]
         }
         span.setDestinationContext(getDBDestination(span, host, port))
 
diff --git a/test/instrumentation/modules/elasticsearch.test.js b/test/instrumentation/modules/elasticsearch.test.js
index f4a3219286e..82ec69ff9a2 100644
--- a/test/instrumentation/modules/elasticsearch.test.js
+++ b/test/instrumentation/modules/elasticsearch.test.js
@@ -186,6 +186,43 @@ test('client.count with callback', function userLandCode (t) {
   })
 })
 
+test('client with host=<array of host:port>', function userLandCode (t) {
+  resetAgent(done(t, 'HEAD', '/'))
+  agent.startTransaction('foo')
+  var client = new elasticsearch.Client({ host: [host] })
+  client.ping(function (err) {
+    t.error(err)
+    agent.endTransaction()
+    agent.flush()
+  })
+})
+
+test('client with hosts=<array of host:port>', function userLandCode (t) {
+  resetAgent(done(t, 'HEAD', '/'))
+  agent.startTransaction('foo')
+  var client = new elasticsearch.Client({ hosts: [host, host] })
+  client.ping(function (err) {
+    t.error(err)
+    agent.endTransaction()
+    agent.flush()
+  })
+})
+
+test('client with hosts="http://host:port"', function userLandCode (t) {
+  resetAgent(done(t, 'HEAD', '/'))
+  agent.startTransaction('foo')
+  let hostWithProto = host
+  if (!hostWithProto.startsWith('http')) {
+    hostWithProto = 'http://' + host
+  }
+  var client = new elasticsearch.Client({ hosts: hostWithProto })
+  client.ping(function (err) {
+    t.error(err)
+    agent.endTransaction()
+    agent.flush()
+  })
+})
+
 function done (t, method, path, query, abort = false) {
   return function (data, cb) {
     t.strictEqual(data.transactions.length, 1, 'should have 1 transaction')