Repository: caikit/caikit
Pull request #565 (2023-11-03 21:58:42)

Health probes no localhost
**What this PR does / why we need it**:

This PR fixes the `caikit-health-probe` to support probing a (m)TLS server over `gRPC` when the server's cert does not contain `localhost` as a Subject Alternate Name.

**Special notes for your reviewer**:

This PR does introduce a dependency on `cryptography` which is a widely used third party library in the python ecosystem. The need for `cryptography` stems from the fact that `grpc` does not allow TLS hostname verification to be fully disabled on the client side. As such, we need to parse the correct `hostname` override from the server's cert in order to set the value correctly in the client. There are numerous libraries to do this, and [it is technically possible using the builtin `ssl`](https://stackoverflow.com/a/50072461), but that requires accessing implementation details that are not guaranteed to be API stable.


Patch:
diff --git a/caikit_health_probe/__main__.py b/caikit_health_probe/__main__.py
index 53aa33c9..b73cbc42 100644
--- a/caikit_health_probe/__main__.py
+++ b/caikit_health_probe/__main__.py
@@ -23,6 +23,10 @@ import sys
 import tempfile
 import warnings
 
+# Third Party
+from cryptography import x509
+from cryptography.hazmat.backends import default_backend
+
 # First Party
 import alog
 
@@ -217,7 +221,19 @@ def _grpc_health_probe(
             credentials = grpc.ssl_channel_credentials(
                 root_certificates=tls_server_cert,
             )
-        channel = grpc.secure_channel(hostname, credentials=credentials)
+
+        # 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),),
+        )
     else:
         log.debug("Probing INSECURE gRPC server")
         channel = grpc.insecure_channel(hostname)
@@ -268,6 +284,23 @@ 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 49983c31..4fd0cb8c 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -45,6 +45,9 @@ 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 04ce61e2..bab997ce 100644
--- a/tests/runtime/conftest.py
+++ b/tests/runtime/conftest.py
@@ -438,12 +438,10 @@ def _check_server_readiness(server):
 
 def _check_http_server_readiness(server, config_overrides: Dict[str, Dict]):
     mode = "http"
-    verify = None
     cert = None
     # tls
     if config_overrides:
         mode = "https"
-        verify = config_overrides["use_in_test"]["ca_cert"]
         # mtls
         if "client_cert" and "client_key" in config_overrides["use_in_test"]:
             cert = (
@@ -455,7 +453,7 @@ def _check_http_server_readiness(server, config_overrides: Dict[str, Dict]):
         try:
             response = requests.get(
                 f"{mode}://localhost:{server.port}{http_server.HEALTH_ENDPOINT}",
-                verify=verify,
+                verify=False,
                 cert=cert,
             )
             assert response.status_code == 200
diff --git a/tests/runtime/http_server/test_http_server.py b/tests/runtime/http_server/test_http_server.py
index 13287c8d..b0e28bf4 100644
--- a/tests/runtime/http_server/test_http_server.py
+++ b/tests/runtime/http_server/test_http_server.py
@@ -18,7 +18,7 @@ Tests for the caikit HTTP server
 from contextlib import contextmanager
 from io import BytesIO
 from pathlib import Path
-from typing import Dict
+from typing import Dict, List, Optional
 import json
 import os
 import signal
@@ -90,6 +90,8 @@ def generate_tls_configs(
     mtls: bool = False,
     inline: bool = False,
     separate_client_ca: bool = False,
+    server_sans: Optional[List[str]] = None,
+    client_sans: Optional[List[str]] = None,
     **http_config_overrides,
 ) -> Dict[str, Dict]:
     """Helper to generate tls configs"""
@@ -101,7 +103,8 @@ def generate_tls_configs(
             ca_key = tls_test_tools.generate_key()[0]
             ca_cert = tls_test_tools.generate_ca_cert(ca_key)
             server_key, server_cert = tls_test_tools.generate_derived_key_cert_pair(
-                ca_key=ca_key
+                ca_key=ca_key,
+                san_list=server_sans,
             )
             server_certfile, server_keyfile = save_key_cert_pair(
                 "server", workdir, server_key, server_cert
@@ -163,6 +166,7 @@ def generate_tls_configs(
                     workdir,
                     *tls_test_tools.generate_derived_key_cert_pair(
                         ca_key=client_ca_key,
+                        san_list=client_sans,
                         **subject_kwargs,
                     ),
                 )
diff --git a/tests/runtime/test_caikit_health_probe.py b/tests/runtime/test_caikit_health_probe.py
index ece6ea02..cb6e1dba 100644
--- a/tests/runtime/test_caikit_health_probe.py
+++ b/tests/runtime/test_caikit_health_probe.py
@@ -31,6 +31,7 @@ import tls_test_tools
 
 # First Party
 from caikit_health_probe import __main__ as caikit_health_probe
+import alog
 
 # Local
 from caikit import get_config
@@ -40,12 +41,16 @@ from tests.runtime.http_server.test_http_server import generate_tls_configs
 
 ## Helpers #####################################################################
 
+log = alog.use_channel("TEST")
+
 
 class TlsMode(Enum):
     INSECURE = 0
     TLS = 1
-    MTLS_COMMON_CLIENT_CA = 2
-    MTLS_SEPARATE_CLIENT_CA = 3
+    TLS_NO_LOCALHOST = 2
+    MTLS_COMMON_CLIENT_CA = 3
+    MTLS_SEPARATE_CLIENT_CA = 4
+    MTLS_NO_LOCALHOST = 5
 
 
 class ServerMode(Enum):
@@ -93,52 +98,71 @@ def temp_probe_config(*args, **kwargs):
         (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),
     ],
 )
 @pytest.mark.parametrize("server_mode", ServerMode.__members__.values())
 def test_health_probe(test_config: Tuple[TlsMode, bool], server_mode: ServerMode):
     """Test all of the different ways that the servers could be running"""
     tls_mode, inline = test_config
-
-    # Get ports for both servers
-    http_port = tls_test_tools.open_port()
-    grpc_port = tls_test_tools.open_port()
-
-    # Set up tls values if needed
-    tls = tls_mode == TlsMode.TLS
-    mtls = tls_mode in [TlsMode.MTLS_COMMON_CLIENT_CA, TlsMode.MTLS_SEPARATE_CLIENT_CA]
-    with generate_tls_configs(
-        port=http_port,
-        tls=tls,
-        mtls=mtls,
-        inline=inline,
-        separate_client_ca=tls_mode == TlsMode.MTLS_SEPARATE_CLIENT_CA,
-    ) as config_overrides:
-        with temp_probe_config(
-            {
-                "runtime": {
-                    "grpc": {
-                        "port": grpc_port,
-                        "enabled": server_mode in [ServerMode.GRPC, ServerMode.BOTH],
-                    },
-                    "http": {
-                        "enabled": server_mode in [ServerMode.HTTP, ServerMode.BOTH],
-                    },
-                }
-            },
-            "merge",
-        ):
-            # Health probe fails with no servers booted
-            assert not caikit_health_probe.health_probe()
-            # If booting the gRPC server, do so
-            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
-                )
-                # If booting the HTTP server, do so
-                with maybe_runtime_http_test_server(
-                    http_port, tls_config_override=config_overrides
-                ):
-                    # Probe should always pass with both possible servers up
-                    assert caikit_health_probe.health_probe()
+    with alog.ContextLog(
+        log.info,
+        "---LOG CONFIG: tls: %s, inline: %s, server: %s---",
+        tls_mode.name,
+        inline,
+        server_mode.name,
+    ):
+        # 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:
+            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,
+            server_sans=server_sans,
+            client_sans=client_sans,
+        ) as config_overrides:
+            with temp_probe_config(
+                {
+                    "runtime": {
+                        "grpc": {
+                            "port": grpc_port,
+                            "enabled": server_mode
+                            in [ServerMode.GRPC, ServerMode.BOTH],
+                        },
+                        "http": {
+                            "enabled": server_mode
+                            in [ServerMode.HTTP, ServerMode.BOTH],
+                        },
+                    }
+                },
+                "merge",
+            ):
+                # Health probe fails with no servers booted
+                assert not caikit_health_probe.health_probe()
+                # If booting the gRPC server, do so
+                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
+                    )
+                    # If booting the HTTP server, do so
+                    with maybe_runtime_http_test_server(
+                        http_port, tls_config_override=config_overrides
+                    ):
+                        # Probe should always pass with both possible servers up
+                        assert caikit_health_probe.health_probe()