From 3a44acdb8c3ca80d78892372056c7df3a0d122d7 Mon Sep 17 00:00:00 2001
From: Chan Chiem Jeffery Saeteurn <chiemsaeteurn@gmail.com>
Date: Fri, 18 Jan 2019 15:31:49 -0800
Subject: [PATCH] Fix httplib invalid scheme detection for HTTPS (#121). (#122)

* Libraries utilizing urllib3 now properly get matched as https when an https request is made.
        + botocore and requests utilize urllib3, so any underlying https request now properly
          identifies as https.
---
 aws_xray_sdk/ext/httplib/patch.py | 12 +++++++++++-
 tests/ext/httplib/test_httplib.py | 29 +++++++++++++++++++++++++++--
 2 files changed, 38 insertions(+), 3 deletions(-)

diff --git a/aws_xray_sdk/ext/httplib/patch.py b/aws_xray_sdk/ext/httplib/patch.py
index 269d8fcc..a2db41d0 100644
--- a/aws_xray_sdk/ext/httplib/patch.py
+++ b/aws_xray_sdk/ext/httplib/patch.py
@@ -2,6 +2,8 @@
 import sys
 import wrapt
 
+import urllib3.connection
+
 from aws_xray_sdk.core import xray_recorder
 from aws_xray_sdk.core.models import http
 from aws_xray_sdk.core.exceptions.exceptions import SegmentNotFoundException
@@ -92,7 +94,15 @@ def decompose_args(method, url, body, headers, encode_chunked=False):
         if subsegment:
             inject_trace_header(headers, subsegment)
 
-        ssl_cxt = getattr(instance, '_context', None)
+        if issubclass(instance.__class__, urllib3.connection.HTTPSConnection):
+            ssl_cxt = getattr(instance, 'ssl_context', None)
+        elif issubclass(instance.__class__, httplib.HTTPSConnection):
+            ssl_cxt = getattr(instance, '_context', None)
+        else:
+            # In this case, the patcher can't determine which module the connection instance is from.
+            # We default to it to check ssl_context but may be None so that the default scheme would be
+            # (and may falsely be) http.
+            ssl_cxt = getattr(instance, 'ssl_context', None)
         scheme = 'https' if ssl_cxt and type(ssl_cxt).__name__ == 'SSLContext' else 'http'
         xray_url = '{}://{}{}'.format(scheme, instance.host, url)
         xray_data = _XRay_Data(method, instance.host, xray_url)
diff --git a/tests/ext/httplib/test_httplib.py b/tests/ext/httplib/test_httplib.py
index 80de9d0a..ec528648 100644
--- a/tests/ext/httplib/test_httplib.py
+++ b/tests/ext/httplib/test_httplib.py
@@ -37,12 +37,15 @@ def construct_ctx():
     unpatch()
 
 
-def _do_req(url, method='GET'):
+def _do_req(url, method='GET', use_https=True):
     parts = urlparse(url)
     host, _, port = parts.netloc.partition(':')
     if port == '':
         port = None
-    conn = httplib.HTTPSConnection(parts.netloc, port)
+    if use_https:
+        conn = httplib.HTTPSConnection(parts.netloc, port)
+    else:
+        conn = httplib.HTTPConnection(parts.netloc, port)
 
     path = '{}?{}'.format(parts.path, parts.query) if parts.query else parts.path
     conn.request(method, path)
@@ -116,3 +119,25 @@ def test_invalid_url():
 
     exception = subsegment.cause['exceptions'][0]
     assert exception.type == 'gaierror'
+
+
+def test_correct_identify_http():
+    status_code = 200
+    url = 'http://{}/status/{}?foo=bar&baz=foo'.format(BASE_URL, status_code)
+    _do_req(url, use_https=False)
+    subsegment = xray_recorder.current_segment().subsegments[0]
+    assert subsegment.name == strip_url(url)
+
+    http_meta = subsegment.http
+    assert http_meta['request']['url'].split(":")[0] == 'http'
+
+
+def test_correct_identify_https():
+    status_code = 200
+    url = 'https://{}/status/{}?foo=bar&baz=foo'.format(BASE_URL, status_code)
+    _do_req(url, use_https=True)
+    subsegment = xray_recorder.current_segment().subsegments[0]
+    assert subsegment.name == strip_url(url)
+
+    https_meta = subsegment.http
+    assert https_meta['request']['url'].split(":")[0] == 'https'