Repository: caikit/caikit
Pull request #567 (2023-11-06 21:11:39)

No crypto health probes
**What this PR does / why we need it**:

This PR is one more attempt to build a robust `caikit-health-probe`, but this time without relying on the `cryptography` package. The decision to not add `cryptography` as a dependency is based on the significant security surface opened up by having this tool in the dependency chain.

With this PR, the health probe can fail when `localhost` is not in the CN/SANs for he server's cert AND the server is not running with the insecure unix socket open. This is an acceptable limitation to trade for not adding a crypto-sensitive library in the dependency list.


Patch:
diff --git a/caikit_health_probe/__main__.py b/caikit_health_probe/__main__.py
index b73cbc42..86c37a5f 100644
--- a/caikit_health_probe/__main__.py
+++ b/caikit_health_probe/__main__.py
@@ -23,10 +23,6 @@ import sys
 import tempfile
 import warnings
 
-# Third Party
-from cryptography import x509
-from cryptography.hazmat.backends import default_backend
-
 # First Party
 import alog
 
@@ -205,8 +201,17 @@ def _grpc_health_probe(
         )
         import grpc  # pylint: disable=import-outside-toplevel
 
+    # Server hostname to use unless using socket mode
     hostname = f"localhost:{port}"
-    if tls_key and tls_cert:
+    socket_file = get_config().runtime.grpc.unix_socket_path
+
+    # If available, use a unix socket
+    if socket_file and os.path.exists(os.path.dirname(socket_file)):
+        socket_address = f"unix://{socket_file}"
+        log.debug("Probing gRPC server over unix socket: %s", socket_file)
+        channel = grpc.insecure_channel(socket_address)
+
+    elif tls_key and tls_cert:
         tls_server_key = bytes(_load_secret(tls_key), "utf-8")
         tls_server_cert = bytes(_load_secret(tls_cert), "utf-8")
         if client_ca:
@@ -222,18 +227,14 @@ def _grpc_health_probe(
                 root_certificates=tls_server_cert,
             )
 
-        # Since we are not guaranteed to know the hostname that the server is
-        # using and the server may not have 'localhost' in the SANs for the cert
-        # it is using, we need to parse out the correct hostname here so that it
-        # can be passed to grpc hostname verification. There is no option to
-        # fully disable grpc hostname verification as of the writing of this
-        # code.
-        host_name = _get_host_name(tls_server_cert)
-        channel = grpc.secure_channel(
-            hostname,
-            credentials=credentials,
-            options=(("grpc.ssl_target_name_override", host_name),),
-        )
+        # NOTE: If the server's certificate does not have 'localhost' in it,
+        #   this will cause certificate validation errors and fail. The original
+        #   workaround for this was to parse the cert's SANs and use hostname
+        #   overrides, but that requires a full cryptographic PEM parser which
+        #   is a security-sensitive dependency to pull that we want to avoid.
+        #   Instead, the workaround is to use the unix socket server option
+        #   above.
+        channel = grpc.secure_channel(hostname, credentials=credentials)
     else:
         log.debug("Probing INSECURE gRPC server")
         channel = grpc.insecure_channel(hostname)
@@ -284,23 +285,6 @@ def _load_secret(secret: str) -> str:
     return secret
 
 
-def _get_host_name(cert_bytes: bytes) -> str:
-    """Get the Common Name from the given certificate in order to provide
-    hostname verification for the gRPC health check
-    """
-    cert = x509.load_pem_x509_certificate(cert_bytes, default_backend())
-    if sans := cert.subject.get_attributes_for_oid(x509.OID_SUBJECT_ALTERNATIVE_NAME):
-        return sans[0].value
-    if ext_san := cert.extensions.get_extension_for_oid(
-        x509.OID_SUBJECT_ALTERNATIVE_NAME
-    ):
-        if hn_sans := ext_san.value.get_values_for_type(x509.DNSName):
-            return hn_sans[0]
-        if ip_sans := ext_san.value.get_values_for_type(x509.IPAddress):
-            return ip_sans[0]
-    return cert.subject.get_attributes_for_oid(x509.OID_COMMON_NAME)[0].value
-
-
 ## Main ########################################################################
 def main():
     caikit_config = get_config()
diff --git a/pyproject.toml b/pyproject.toml
index 99bf3ca2..266dbf10 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -45,9 +45,6 @@ runtime-grpc = [
     "grpcio-reflection>=1.35.0,<2.0",
     "prometheus_client>=0.12.0,<1.0",
     "py-grpc-prometheus>=0.7.0,<0.8",
-    # This is needed to support health probes for TLS gRPC servers since we
-    # cannot fully disable hostname verification.
-    "cryptography>=40.0.1,<41",
 ]
 
 runtime-http = [


Test patch:
diff --git a/tests/runtime/conftest.py b/tests/runtime/conftest.py
index bab997ce..1390ab6a 100644
--- a/tests/runtime/conftest.py
+++ b/tests/runtime/conftest.py
@@ -14,6 +14,7 @@ import sys
 import tempfile
 import threading
 import time
+import warnings
 
 # Third Party
 from grpc_health.v1 import health_pb2, health_pb2_grpc
@@ -151,9 +152,7 @@ def runtime_grpc_test_server(open_port, *args, **kwargs):
 
 
 @pytest.fixture(scope="session")
-def runtime_grpc_server(
-    session_scoped_open_port,
-) -> RuntimeGRPCServer:
+def runtime_grpc_server(session_scoped_open_port) -> RuntimeGRPCServer:
     with runtime_grpc_test_server(
         session_scoped_open_port,
     ) as server:
@@ -183,12 +182,14 @@ def runtime_http_test_server(open_port, *args, **kwargs):
             # IFF the configs contain actual TLS (indicated by the presence of
             # the special "use_in_test" element).
             config_overrides = kwargs.pop("tls_config_override", {})
-            if "use_in_test" not in config_overrides:
-                config_overrides = {}
+            if tls_config_override := config_overrides.get("runtime", {}).get("tls"):
+                kwargs["tls_config_override"] = tls_config_override
             else:
-                kwargs["tls_config_override"] = config_overrides["runtime"]["tls"]
+                config_overrides = {}
+            check_readiness = kwargs.pop("check_readiness", True)
             with http_server.RuntimeHTTPServer(*args, **kwargs) as server:
-                _check_http_server_readiness(server, config_overrides)
+                if check_readiness:
+                    _check_http_server_readiness(server, config_overrides)
                 # Give tests access to the workdir
                 server.workdir = workdir
                 yield server
@@ -451,11 +452,13 @@ def _check_http_server_readiness(server, config_overrides: Dict[str, Dict]):
     done = False
     while not done:
         try:
-            response = requests.get(
-                f"{mode}://localhost:{server.port}{http_server.HEALTH_ENDPOINT}",
-                verify=False,
-                cert=cert,
-            )
+            with warnings.catch_warnings():
+                warnings.filterwarnings("ignore", module="urllib3")
+                response = requests.get(
+                    f"{mode}://localhost:{server.port}{http_server.HEALTH_ENDPOINT}",
+                    verify=False,
+                    cert=cert,
+                )
             assert response.status_code == 200
             assert response.text == "OK"
             done = True
diff --git a/tests/runtime/http_server/test_http_server.py b/tests/runtime/http_server/test_http_server.py
index d46e830f..0d4ab553 100644
--- a/tests/runtime/http_server/test_http_server.py
+++ b/tests/runtime/http_server/test_http_server.py
@@ -99,6 +99,8 @@ def generate_tls_configs(
         config_overrides = {}
         client_keyfile, client_certfile = None, None
         ca_cert, server_cert, server_key = None, None, None
+        use_in_test = config_overrides.setdefault("use_in_test", {})
+        use_in_test["workdir"] = workdir
         if mtls or tls:
             ca_key = tls_test_tools.generate_key()[0]
             ca_cert = tls_test_tools.generate_ca_cert(ca_key)
@@ -124,11 +126,9 @@ def generate_tls_configs(
             # need to save this ca_certfile in config_overrides so the tls
             # tests below can access it from client side
             ca_certfile, _ = save_key_cert_pair("ca", workdir, cert=ca_cert)
-            config_overrides["use_in_test"] = {
-                "ca_cert": ca_certfile,
-                "server_key": server_keyfile,
-                "server_cert": server_certfile,
-            }
+            use_in_test["ca_cert"] = ca_certfile
+            use_in_test["server_key"] = server_keyfile
+            use_in_test["server_cert"] = server_certfile
 
             # also saving a bad ca_certfile for a failure test case
             bad_ca_file = os.path.join(workdir, "bad_ca_cert.crt")
@@ -137,7 +137,7 @@ def generate_tls_configs(
                     "-----BEGIN CERTIFICATE-----\nfoobar\n-----END CERTIFICATE-----"
                 )
                 handle.write(bad_cert)
-            config_overrides["use_in_test"]["bad_ca_cert"] = bad_ca_file
+            use_in_test["bad_ca_cert"] = bad_ca_file
 
             if mtls:
                 if separate_client_ca:
@@ -171,8 +171,8 @@ def generate_tls_configs(
                     ),
                 )
                 # need to save the client cert and key in config_overrides so the mtls test below can access it
-                config_overrides["use_in_test"]["client_cert"] = client_certfile
-                config_overrides["use_in_test"]["client_key"] = client_keyfile
+                use_in_test["client_cert"] = client_certfile
+                use_in_test["client_key"] = client_keyfile
 
             config_overrides["runtime"] = {"tls": tls_config.to_dict()}
         config_overrides.setdefault("runtime", {})["http"] = {
diff --git a/tests/runtime/test_caikit_health_probe.py b/tests/runtime/test_caikit_health_probe.py
index cb6e1dba..a36769e3 100644
--- a/tests/runtime/test_caikit_health_probe.py
+++ b/tests/runtime/test_caikit_health_probe.py
@@ -21,9 +21,10 @@ early causing some of the core tests to fail!
 """
 # Standard
 from contextlib import contextmanager
+from dataclasses import dataclass
 from enum import Enum
-from typing import Tuple
 from unittest import mock
+import os
 
 # Third Party
 import pytest
@@ -44,21 +45,6 @@ from tests.runtime.http_server.test_http_server import generate_tls_configs
 log = alog.use_channel("TEST")
 
 
-class TlsMode(Enum):
-    INSECURE = 0
-    TLS = 1
-    TLS_NO_LOCALHOST = 2
-    MTLS_COMMON_CLIENT_CA = 3
-    MTLS_SEPARATE_CLIENT_CA = 4
-    MTLS_NO_LOCALHOST = 5
-
-
-class ServerMode(Enum):
-    HTTP = 0
-    GRPC = 1
-    BOTH = 2
-
-
 @contextmanager
 def maybe_runtime_grpc_test_server(*args, **kwargs):
     if get_config().runtime.grpc.enabled:
@@ -85,54 +71,101 @@ def temp_probe_config(*args, **kwargs):
             yield
 
 
+class TlsMode(Enum):
+    INSECURE = 0
+    TLS = 1
+    MTLS = 3
+
+
+class ServerMode(Enum):
+    HTTP = 0
+    GRPC = 1
+    BOTH = 2
+
+
+@dataclass
+class ProbeTestConfig:
+    tls_mode: TlsMode
+    server_mode: ServerMode
+    # TLS blobs passed as inline strings instead of files
+    inline: bool = False
+    # Run the unix socket grpc server
+    unix_socket: bool = True
+    # Put "localhost" in the SAN list for the server's cert
+    localhost_in_cert: bool = True
+    # Use a common CA for client and server certs (mTLS only)
+    common_client_ca: bool = True
+    # Whether the test should eventually become healthy
+    should_become_healthy: bool = True
+
+
 ## Tests #######################################################################
 
 
 @pytest.mark.parametrize(
     "test_config",
     [
-        (TlsMode.INSECURE, True),
-        (TlsMode.TLS, True),
-        (TlsMode.TLS, False),
-        (TlsMode.MTLS_COMMON_CLIENT_CA, True),
-        (TlsMode.MTLS_COMMON_CLIENT_CA, False),
-        (TlsMode.MTLS_SEPARATE_CLIENT_CA, True),
-        (TlsMode.MTLS_SEPARATE_CLIENT_CA, False),
-        # Tests for no "localhost" in the SAN list
-        (TlsMode.TLS_NO_LOCALHOST, False),
-        (TlsMode.MTLS_NO_LOCALHOST, False),
+        # Insecure
+        ProbeTestConfig(TlsMode.INSECURE, ServerMode.HTTP),
+        ProbeTestConfig(TlsMode.INSECURE, ServerMode.GRPC),
+        ProbeTestConfig(TlsMode.INSECURE, ServerMode.BOTH),
+        # TLS
+        ProbeTestConfig(TlsMode.TLS, ServerMode.HTTP),
+        ProbeTestConfig(TlsMode.TLS, ServerMode.GRPC),
+        ProbeTestConfig(TlsMode.TLS, ServerMode.BOTH),
+        ProbeTestConfig(TlsMode.TLS, ServerMode.BOTH, inline=True),
+        ProbeTestConfig(TlsMode.TLS, ServerMode.BOTH, localhost_in_cert=False),
+        # mTLS
+        ProbeTestConfig(TlsMode.MTLS, ServerMode.HTTP),
+        ProbeTestConfig(TlsMode.MTLS, ServerMode.GRPC),
+        ProbeTestConfig(TlsMode.MTLS, ServerMode.BOTH),
+        ProbeTestConfig(TlsMode.MTLS, ServerMode.BOTH, inline=True),
+        ProbeTestConfig(TlsMode.MTLS, ServerMode.BOTH, localhost_in_cert=False),
+        ProbeTestConfig(TlsMode.MTLS, ServerMode.BOTH, common_client_ca=False),
+        # Invalid configs that never pass
+        ProbeTestConfig(
+            TlsMode.TLS,
+            ServerMode.GRPC,
+            localhost_in_cert=False,
+            unix_socket=False,
+            should_become_healthy=False,
+        ),
+        ProbeTestConfig(
+            TlsMode.TLS,
+            ServerMode.BOTH,
+            localhost_in_cert=False,
+            unix_socket=False,
+            should_become_healthy=False,
+        ),
+        ProbeTestConfig(
+            TlsMode.MTLS,
+            ServerMode.GRPC,
+            localhost_in_cert=False,
+            unix_socket=False,
+            should_become_healthy=False,
+        ),
     ],
 )
-@pytest.mark.parametrize("server_mode", ServerMode.__members__.values())
-def test_health_probe(test_config: Tuple[TlsMode, bool], server_mode: ServerMode):
+def test_health_probe(test_config: ProbeTestConfig):
     """Test all of the different ways that the servers could be running"""
-    tls_mode, inline = test_config
-    with alog.ContextLog(
-        log.info,
-        "---LOG CONFIG: tls: %s, inline: %s, server: %s---",
-        tls_mode.name,
-        inline,
-        server_mode.name,
-    ):
+    with alog.ContextLog(log.info, "---LOG CONFIG: %s---", test_config):
         # Get ports for both servers
         http_port = tls_test_tools.open_port()
         grpc_port = tls_test_tools.open_port()
 
         # Set up SAN lists if not putting "localhost" in
         server_sans, client_sans = None, None
-        if "NO_LOCALHOST" in tls_mode.name:
+        if not test_config.localhost_in_cert:
             server_sans = ["foo.bar"]
             client_sans = ["baz.bat"]
 
         # Set up tls values if needed
-        tls = tls_mode.name.startswith("TLS")
-        mtls = tls_mode.name.startswith("MTLS")
         with generate_tls_configs(
             port=http_port,
-            tls=tls,
-            mtls=mtls,
-            inline=inline,
-            separate_client_ca=tls_mode == TlsMode.MTLS_SEPARATE_CLIENT_CA,
+            tls=test_config.tls_mode == TlsMode.TLS,
+            mtls=test_config.tls_mode == TlsMode.MTLS,
+            inline=test_config.inline,
+            separate_client_ca=not test_config.common_client_ca,
             server_sans=server_sans,
             client_sans=client_sans,
         ) as config_overrides:
@@ -141,11 +174,17 @@ def test_health_probe(test_config: Tuple[TlsMode, bool], server_mode: ServerMode
                     "runtime": {
                         "grpc": {
                             "port": grpc_port,
-                            "enabled": server_mode
+                            "enabled": test_config.server_mode
                             in [ServerMode.GRPC, ServerMode.BOTH],
+                            "unix_socket_path": os.path.join(
+                                config_overrides["use_in_test"]["workdir"],
+                                "grpc.sock",
+                            )
+                            if test_config.unix_socket
+                            else None,
                         },
                         "http": {
-                            "enabled": server_mode
+                            "enabled": test_config.server_mode
                             in [ServerMode.HTTP, ServerMode.BOTH],
                         },
                     }
@@ -158,11 +197,17 @@ def test_health_probe(test_config: Tuple[TlsMode, bool], server_mode: ServerMode
                 with maybe_runtime_grpc_test_server(grpc_port):
                     # If only running gRPC, health probe should pass
                     assert caikit_health_probe.health_probe() == (
-                        server_mode == ServerMode.GRPC
+                        test_config.should_become_healthy
+                        and test_config.server_mode == ServerMode.GRPC
                     )
                     # If booting the HTTP server, do so
                     with maybe_runtime_http_test_server(
-                        http_port, tls_config_override=config_overrides
+                        http_port,
+                        tls_config_override=config_overrides,
+                        check_readiness=test_config.should_become_healthy,
                     ):
                         # Probe should always pass with both possible servers up
-                        assert caikit_health_probe.health_probe()
+                        assert (
+                            caikit_health_probe.health_probe()
+                            == test_config.should_become_healthy
+                        )