From 9a9be34bb127bdb1f926dada193885a4039d7202 Mon Sep 17 00:00:00 2001
From: Matt Robenolt <matt@ydekproductions.com>
Date: Fri, 26 Oct 2018 14:12:39 -0700
Subject: [PATCH] feat: Support unix sockets for urllib3

---
 src/sentry/net/http.py              | 35 +++++++++++++++++++++++++++++
 src/sentry/nodestore/riak/client.py | 25 ++++++++++++---------
 2 files changed, 50 insertions(+), 10 deletions(-)

diff --git a/src/sentry/net/http.py b/src/sentry/net/http.py
index 9ea7d1ae974e6..46fe4619bdbab 100644
--- a/src/sentry/net/http.py
+++ b/src/sentry/net/http.py
@@ -1,5 +1,6 @@
 from __future__ import absolute_import
 
+import socket
 from socket import error as SocketError, timeout as SocketTimeout
 
 from requests import Session as _Session
@@ -8,6 +9,7 @@
 from urllib3.connection import HTTPConnection, HTTPSConnection
 from urllib3.exceptions import NewConnectionError, ConnectTimeoutError
 from urllib3.poolmanager import PoolManager
+from urllib3.util.connection import _set_socket_options
 
 from sentry import VERSION as SENTRY_VERSION
 from sentry.net.socket import safe_create_connection
@@ -149,3 +151,36 @@ def __init__(self):
         adapter = BlacklistAdapter()
         self.mount('https://', adapter)
         self.mount('http://', adapter)
+
+
+class UnixHTTPConnection(HTTPConnection):
+    default_socket_options = []
+
+    def __init__(self, host, **kwargs):
+        # We're using the `host` as the socket path, but
+        # urllib3 uses this host as the Host header by default.
+        # If we send along the socket path as a Host header, this is
+        # never what you want and would typically be malformed value.
+        # So we fake this by sending along `localhost` by default as
+        # other libraries do.
+        self.socket_path = host
+        super(UnixHTTPConnection, self).__init__(host='localhost', **kwargs)
+
+    def _new_conn(self):
+        sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
+
+        # If provided, set socket level options before connecting.
+        _set_socket_options(sock, self.socket_options)
+
+        if self.timeout is not socket._GLOBAL_DEFAULT_TIMEOUT:
+            sock.settimeout(self.timeout)
+        sock.connect(self.socket_path)
+        return sock
+
+
+class UnixHTTPConnectionPool(HTTPConnectionPool):
+    ConnectionCls = UnixHTTPConnection
+
+    def __str__(self):
+        return '%s(host=%r)' % (type(self).__name__,
+                                self.host)
diff --git a/src/sentry/nodestore/riak/client.py b/src/sentry/nodestore/riak/client.py
index 077425940ce94..ca16897e07546 100644
--- a/src/sentry/nodestore/riak/client.py
+++ b/src/sentry/nodestore/riak/client.py
@@ -23,9 +23,11 @@
 from requests.certs import where as ca_certs
 from six.moves.urllib.parse import urlencode, quote_plus
 from urllib3 import HTTPConnectionPool, HTTPSConnectionPool
-from urllib3.connection import HTTPConnection
 from urllib3.exceptions import HTTPError
 
+from sentry.net.http import UnixHTTPConnectionPool
+
+
 DEFAULT_NODES = ({'host': '127.0.0.1', 'port': 8098}, )
 
 
@@ -243,11 +245,6 @@ def create_pool(self, host):
         if 'basic_auth' in host:
             options['headers']['authorization'] = encode_basic_auth(host['basic_auth'])
 
-        if self.tcp_keepalive:
-            options['socket_options'] = HTTPConnection.default_socket_options + [
-                (socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1),
-            ]
-
         # Support backwards compatibility with `http_port`
         if 'http_port' in host:
             import warnings
@@ -257,10 +254,12 @@ def create_pool(self, host):
         addr = host.get('host', '127.0.0.1')
         port = int(host.get('port', 8098))
         secure = host.get('secure', False)
-        if not secure:
-            connection_cls = HTTPConnectionPool
+        if addr[:1] == '/':
+            pool_cls = UnixHTTPConnectionPool
+        elif not secure:
+            pool_cls = HTTPConnectionPool
         else:
-            connection_cls = HTTPSConnectionPool
+            pool_cls = HTTPSConnectionPool
             verify_ssl = host.get('verify_ssl', False)
             if verify_ssl:
                 options.extend(
@@ -269,7 +268,13 @@ def create_pool(self, host):
                         'ca_certs': host.get('ca_certs', ca_certs())
                     }
                 )
-        return connection_cls(addr, port, **options)
+
+        if self.tcp_keepalive:
+            options['socket_options'] = pool_cls.ConnectionCls.default_socket_options + [
+                (socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1),
+            ]
+
+        return pool_cls(addr, port, **options)
 
     def urlopen(self, method, path, headers=None, **kwargs):
         """