diff --git a/lib/bb/fetch2/__init__.py b/lib/bb/fetch2/__init__.py index a5772f8a..ea007f36 100644 --- a/lib/bb/fetch2/__init__.py +++ b/lib/bb/fetch2/__init__.py @@ -2127,26 +2127,28 @@ class FetchConnectionCache(object): def __init__(self): self.cache = {} - def get_connection_name(self, host, port): - return host + ':' + str(port) + def get_connection_name(self, host, port, connection_id=None): + if connection_id is None: + return host + ':' + str(port) + return (host, port, connection_id) - def add_connection(self, host, port, connection): - cn = self.get_connection_name(host, port) + def add_connection(self, host, port, connection, connection_id=None): + cn = self.get_connection_name(host, port, connection_id) if cn not in self.cache: self.cache[cn] = connection - def get_connection(self, host, port): + def get_connection(self, host, port, connection_id=None): connection = None - cn = self.get_connection_name(host, port) + cn = self.get_connection_name(host, port, connection_id) if cn in self.cache: connection = self.cache[cn] return connection - def remove_connection(self, host, port): - cn = self.get_connection_name(host, port) + def remove_connection(self, host, port, connection_id=None): + cn = self.get_connection_name(host, port, connection_id) if cn in self.cache: self.cache[cn].close() del self.cache[cn] diff --git a/lib/bb/fetch2/wget.py b/lib/bb/fetch2/wget.py index 141c2d06..ad98d967 100644 --- a/lib/bb/fetch2/wget.py +++ b/lib/bb/fetch2/wget.py @@ -151,27 +151,54 @@ class Wget(FetchMethod): return True def checkstatus(self, fetch, ud, d, try_again=True): + check_certs = self.check_certs(d) + newenv = bb.fetch2.get_fetcher_environment(d) + class HTTPConnectionCache(http.client.HTTPConnection): + def cache_id(self): + return None + if fetch.connection_cache: def connect(self): """Connect to the host and port specified in __init__.""" - sock = fetch.connection_cache.get_connection(self.host, self.port) + sock = fetch.connection_cache.get_connection( + self.host, self.port, self.cache_id()) if sock: self.sock = sock else: self.sock = socket.create_connection((self.host, self.port), self.timeout, self.source_address) - fetch.connection_cache.add_connection(self.host, self.port, self.sock) + fetch.connection_cache.add_connection( + self.host, self.port, self.sock, self.cache_id()) if self._tunnel_host: self._tunnel() + class HTTPSConnectionCache(http.client.HTTPSConnection): + def cache_id(self): + return ("https", check_certs, + newenv.get("SSL_CERT_FILE"), + self._tunnel_host, self._tunnel_port) + + if fetch.connection_cache: + def connect(self): + """Reuse an established TLS connection when available.""" + + sock = fetch.connection_cache.get_connection( + self.host, self.port, self.cache_id()) + if sock: + self.sock = sock + else: + super().connect() + fetch.connection_cache.add_connection( + self.host, self.port, self.sock, self.cache_id()) + class CacheHTTPHandler(urllib.request.HTTPHandler): def http_open(self, req): return self.do_open(HTTPConnectionCache, req) - def do_open(self, http_class, req): + def do_open(self, http_class, req, **http_conn_args): """Return an addinfourl object for the request, using http_class. http_class must implement the HTTPConnection API from httplib. @@ -185,7 +212,7 @@ class Wget(FetchMethod): if not host: raise urllib.error.URLError('no host given') - h = http_class(host, timeout=req.timeout) # will parse host:port + h = http_class(host, timeout=req.timeout, **http_conn_args) # will parse host:port h.set_debuglevel(self._debuglevel) headers = dict(req.unredirected_hdrs) @@ -231,7 +258,8 @@ class Wget(FetchMethod): # If it still fails, we give up, which can happen for bad # HTTP proxy settings. if fetch.connection_cache: - fetch.connection_cache.remove_connection(h.host, h.port) + fetch.connection_cache.remove_connection( + h.host, h.port, h.cache_id()) h.close() raise @@ -265,10 +293,21 @@ class Wget(FetchMethod): # Close connection when server request it. if fetch.connection_cache is not None: if 'Connection' in r.msg and r.msg['Connection'] == 'close': - fetch.connection_cache.remove_connection(h.host, h.port) + fetch.connection_cache.remove_connection( + h.host, h.port, h.cache_id()) return resp + class CacheHTTPSHandler(CacheHTTPHandler, urllib.request.HTTPSHandler): + def __init__(self, debuglevel=0, context=None, check_hostname=None): + urllib.request.HTTPSHandler.__init__(self, debuglevel, context, + check_hostname) + + def https_open(self, req): + return self.do_open(HTTPSConnectionCache, req, + context=self._context, + check_hostname=self._check_hostname) + class HTTPMethodFallback(urllib.request.BaseHandler): """ Fallback to GET if HEAD is not allowed (405 HTTP error) @@ -370,12 +409,10 @@ class Wget(FetchMethod): # Avoid tramping the environment too much by using bb.utils.environment # to scope the changes to the build_opener request, which is when the # environment lookups happen. - newenv = bb.fetch2.get_fetcher_environment(d) - with bb.utils.environment(**newenv): import ssl - if self.check_certs(d): + if check_certs: context = ssl.create_default_context() else: context = ssl._create_unverified_context() @@ -384,7 +421,7 @@ class Wget(FetchMethod): HTTPMethodFallback, urllib.request.ProxyHandler(), CacheHTTPHandler(), - urllib.request.HTTPSHandler(context=context)] + CacheHTTPSHandler(context=context)] opener = urllib.request.build_opener(*handlers) try: diff --git a/lib/bb/tests/fetch.py b/lib/bb/tests/fetch.py index cd50c37a..fd064743 100644 --- a/lib/bb/tests/fetch.py +++ b/lib/bb/tests/fetch.py @@ -1781,6 +1781,72 @@ class FetchCheckStatusTest(FetcherTest): connection_cache.close_connections() + @unittest.skipUnless(shutil.which("openssl"), "openssl not installed") + def test_wget_checkstatus_https_connection_cache(self): + import ssl + from socketserver import ThreadingMixIn + from bb.fetch2 import FetchConnectionCache + + class HTTPSRequestHandler(http.server.BaseHTTPRequestHandler): + protocol_version = "HTTP/1.1" + + def do_HEAD(self): + self.send_response(200) + self.send_header("Content-Length", "0") + self.end_headers() + + def log_message(self, format_str, *args): + pass + + class HTTPSServer(ThreadingMixIn, http.server.HTTPServer): + daemon_threads = True + + def __init__(self, *args, **kwargs): + self.connection_count = 0 + super().__init__(*args, **kwargs) + + def get_request(self): + request, client_address = super().get_request() + self.connection_count += 1 + return request, client_address + + certificate = os.path.join(self.tempdir, "certificate.pem") + private_key = os.path.join(self.tempdir, "private-key.pem") + subprocess.check_call( + ["openssl", "req", "-x509", "-newkey", "rsa:2048", "-nodes", + "-keyout", private_key, "-out", certificate, "-days", "1", + "-subj", "/CN=127.0.0.1"], + stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL) + + server = HTTPSServer(("127.0.0.1", 0), HTTPSRequestHandler) + context = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER) + context.load_cert_chain(certificate, private_key) + server.socket = context.wrap_socket(server.socket, server_side=True) + server_thread = threading.Thread(target=server.serve_forever) + server_thread.daemon = True + server_thread.start() + + connection_cache = FetchConnectionCache() + try: + url = "https://127.0.0.1:%s/test" % server.server_port + self.d.setVar("BB_CHECK_SSL_CERTS", "0") + fetch = bb.fetch2.Fetch([url], self.d, + connection_cache=connection_cache) + ud = fetch.ud[url] + self.assertTrue(ud.method.checkstatus(fetch, ud, self.d)) + self.assertTrue(ud.method.checkstatus(fetch, ud, self.d)) + self.assertEqual(server.connection_count, 1) + + # A connection established without certificate checks must not be + # reused after certificate checking is enabled. + self.d.setVar("BB_CHECK_SSL_CERTS", "1") + self.assertFalse(ud.method.checkstatus(fetch, ud, self.d)) + finally: + connection_cache.close_connections() + server.shutdown() + server_thread.join() + server.server_close() + def test_wget_checkstatus_same_origin_redirect_keeps_auth(self): server = self._start_checkstatus_server() server.redirect_url = "http://127.0.0.1:%s/b" % server.server_port