feat: Loading the CA trusted store certificate into Feast to verify the public certificate. (#4852)

lokeshrangineni · web-flow · commit 132ce2a6c9e3 · 2024-12-18T12:55:45.000-05:00
* Initial Draft version to load the CA trusted store code.

Signed-off-by: lrangine &lt;19699092+lokeshrangineni@users.noreply.github.com&gt;

* Initial Draft version to load the CA trusted store code.

Signed-off-by: lrangine &lt;19699092+lokeshrangineni@users.noreply.github.com&gt;

* Fixing the lint error.

Signed-off-by: lrangine &lt;19699092+lokeshrangineni@users.noreply.github.com&gt;

* Trying to fix the online store test cases.

Signed-off-by: lrangine &lt;19699092+lokeshrangineni@users.noreply.github.com&gt;

* Formatted the python to fix lint errors.

Signed-off-by: lrangine &lt;19699092+lokeshrangineni@users.noreply.github.com&gt;

* Fixing the unit test cases.

Signed-off-by: lrangine &lt;19699092+lokeshrangineni@users.noreply.github.com&gt;

* Fixing the unit test cases.

Signed-off-by: lrangine &lt;19699092+lokeshrangineni@users.noreply.github.com&gt;

* removing unnecessary cli args.

Signed-off-by: lrangine &lt;19699092+lokeshrangineni@users.noreply.github.com&gt;

* Now configuring the SSL ca store configurations on the feast client side rather than on the server side. And also fixing the integration tests.

Signed-off-by: lrangine &lt;19699092+lokeshrangineni@users.noreply.github.com&gt;

* Renamed the remote registry is_tls_mode variable to is_tls.
Changed the offline store TLS setting decision from cert to scheme.

Signed-off-by: lrangine &lt;19699092+lokeshrangineni@users.noreply.github.com&gt;

* Adding the existing trust store certificates to the newly created trust store.

Signed-off-by: lrangine &lt;19699092+lokeshrangineni@users.noreply.github.com&gt;

* Clearing the existing trust store configuration to see if it fixes the PR integration failures.

Signed-off-by: lrangine &lt;19699092+lokeshrangineni@users.noreply.github.com&gt;

* Clearing the existing trust store configuration to see if it fixes the PR integration failures.

Signed-off-by: lrangine &lt;19699092+lokeshrangineni@users.noreply.github.com&gt;

* Clearing the existing trust store configuration to see if it fixes the PR integration failures.

Signed-off-by: lrangine &lt;19699092+lokeshrangineni@users.noreply.github.com&gt;

* combining the default system ca store with the custom one to fix the integration tests.

Signed-off-by: lrangine &lt;19699092+lokeshrangineni@users.noreply.github.com&gt;

* Final clean up and adding documentation.

Signed-off-by: lrangine &lt;19699092+lokeshrangineni@users.noreply.github.com&gt;

* Incorporating the code review comments from Francisco.

Signed-off-by: lrangine &lt;19699092+lokeshrangineni@users.noreply.github.com&gt;

---------

Signed-off-by: lrangine &lt;19699092+lokeshrangineni@users.noreply.github.com&gt;
diff --git a/docs/how-to-guides/starting-feast-servers-tls-mode.md b/docs/how-to-guides/starting-feast-servers-tls-mode.md
@@ -189,3 +189,8 @@ INFO:     Waiting for application startup.
 INFO:     Application startup complete.
 INFO:     Uvicorn running on https://0.0.0.0:8888 (Press CTRL+C to quit)
 ```
+
+
+## Adding public key to CA trust store and configuring the feast to use the trust store.
+You can pass the public key for SSL verification using the `cert` parameter, however, it is sometimes difficult to maintain individual certificates and pass them individually.
+The alternative recommendation is to add the public certificate to CA trust store and set the path as an environment variable (e.g., `FEAST_CA_CERT_FILE_PATH`). Feast will use the trust store path in the  `FEAST_CA_CERT_FILE_PATH` environment variable.
diff --git a/sdk/python/feast/cli.py b/sdk/python/feast/cli.py
@@ -982,7 +982,6 @@ def serve_command(
         raise click.BadParameter(
             "Please pass --cert and --key args to start the feature server in TLS mode."
         )
-
     store = create_feature_store(ctx)
 
     store.serve(
diff --git a/sdk/python/feast/feature_store.py b/sdk/python/feast/feature_store.py
@@ -86,6 +86,7 @@
 from feast.repo_config import RepoConfig, load_repo_config
 from feast.repo_contents import RepoContents
 from feast.saved_dataset import SavedDataset, SavedDatasetStorage, ValidationReference
+from feast.ssl_ca_trust_store_setup import configure_ca_trust_store_env_variables
 from feast.stream_feature_view import StreamFeatureView
 from feast.utils import _utc_now
 
@@ -129,6 +130,8 @@ def __init__(
         if fs_yaml_file is not None and config is not None:
             raise ValueError("You cannot specify both fs_yaml_file and config.")
 
+        configure_ca_trust_store_env_variables()
+
         if repo_path:
             self.repo_path = Path(repo_path)
         else:
@@ -1949,13 +1952,19 @@ def serve_ui(
         )
 
     def serve_registry(
-        self, port: int, tls_key_path: str = "", tls_cert_path: str = ""
+        self,
+        port: int,
+        tls_key_path: str = "",
+        tls_cert_path: str = "",
     ) -> None:
         """Start registry server locally on a given port."""
         from feast import registry_server
 
         registry_server.start_server(
-            self, port=port, tls_key_path=tls_key_path, tls_cert_path=tls_cert_path
+            self,
+            port=port,
+            tls_key_path=tls_key_path,
+            tls_cert_path=tls_cert_path,
         )
 
     def serve_offline(
diff --git a/sdk/python/feast/infra/offline_stores/remote.py b/sdk/python/feast/infra/offline_stores/remote.py
@@ -74,7 +74,7 @@ def build_arrow_flight_client(
     scheme: str, host: str, port, auth_config: AuthConfig, cert: str = ""
 ):
     arrow_scheme = "grpc+tcp"
-    if cert:
+    if scheme == "https":
         logger.info(
             "Scheme is https so going to connect offline server in SSL(TLS) mode."
         )
diff --git a/sdk/python/feast/infra/registry/remote.py b/sdk/python/feast/infra/registry/remote.py
@@ -1,3 +1,4 @@
+import os
 from datetime import datetime
 from pathlib import Path
 from typing import List, Optional, Union
@@ -59,6 +60,12 @@ class RemoteRegistryConfig(RegistryConfig):
     """ str: Path to the public certificate when the registry server starts in TLS(SSL) mode. This may be needed if the registry server started with a self-signed certificate, typically this file ends with `*.crt`, `*.cer`, or `*.pem`.
     If registry_type is 'remote', then this configuration is needed to connect to remote registry server in TLS mode. If the remote registry started in non-tls mode then this configuration is not needed."""
 
+    is_tls: bool = False
+    """     bool: Set to `True` if you plan to connect to a registry server running in TLS (SSL) mode.
+    If you intend to add the public certificate to the trust store instead of passing it via the `cert` parameter, this field must be set to `True`.
+    If you are planning to add the public certificate as part of the trust store instead of passing it as a `cert` parameters then setting this field to `true` is mandatory.
+    """
+
 
 class RemoteRegistry(BaseRegistry):
     def __init__(
@@ -70,20 +77,32 @@ def __init__(
     ):
         self.auth_config = auth_config
         assert isinstance(registry_config, RemoteRegistryConfig)
-        if registry_config.cert:
-            with open(registry_config.cert, "rb") as cert_file:
-                trusted_certs = cert_file.read()
-                tls_credentials = grpc.ssl_channel_credentials(
-                    root_certificates=trusted_certs
-                )
-            self.channel = grpc.secure_channel(registry_config.path, tls_credentials)
-        else:
-            self.channel = grpc.insecure_channel(registry_config.path)
+        self.channel = self._create_grpc_channel(registry_config)
 
         auth_header_interceptor = GrpcClientAuthHeaderInterceptor(auth_config)
         self.channel = grpc.intercept_channel(self.channel, auth_header_interceptor)
         self.stub = RegistryServer_pb2_grpc.RegistryServerStub(self.channel)
 
+    def _create_grpc_channel(self, registry_config):
+        assert isinstance(registry_config, RemoteRegistryConfig)
+        if registry_config.cert or registry_config.is_tls:
+            cafile = os.getenv("SSL_CERT_FILE") or os.getenv("REQUESTS_CA_BUNDLE")
+            if not cafile and not registry_config.cert:
+                raise EnvironmentError(
+                    "SSL_CERT_FILE or REQUESTS_CA_BUNDLE environment variable must be set to use secure TLS or set the cert parameter in feature_Store.yaml file under remote registry configuration."
+                )
+            with open(
+                registry_config.cert if registry_config.cert else cafile, "rb"
+            ) as cert_file:
+                trusted_certs = cert_file.read()
+            tls_credentials = grpc.ssl_channel_credentials(
+                root_certificates=trusted_certs
+            )
+            return grpc.secure_channel(registry_config.path, tls_credentials)
+        else:
+            # Create an insecure gRPC channel
+            return grpc.insecure_channel(registry_config.path)
+
     def close(self):
         if self.channel:
             self.channel.close()
diff --git a/sdk/python/feast/ssl_ca_trust_store_setup.py b/sdk/python/feast/ssl_ca_trust_store_setup.py
@@ -0,0 +1,22 @@
+import logging
+import os
+
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.INFO)
+
+
+def configure_ca_trust_store_env_variables():
+    """
+    configures the environment variable so that other libraries or servers refer to the TLS ca file path.
+    :param ca_file_path:
+    :return:
+    """
+    if (
+        "FEAST_CA_CERT_FILE_PATH" in os.environ
+        and os.environ["FEAST_CA_CERT_FILE_PATH"]
+    ):
+        logger.info(
+            f"Feast CA Cert file path found in environment variable FEAST_CA_CERT_FILE_PATH={os.environ['FEAST_CA_CERT_FILE_PATH']}. Going to refer this path."
+        )
+        os.environ["SSL_CERT_FILE"] = os.environ["FEAST_CA_CERT_FILE_PATH"]
+        os.environ["REQUESTS_CA_BUNDLE"] = os.environ["FEAST_CA_CERT_FILE_PATH"]
diff --git a/sdk/python/tests/conftest.py b/sdk/python/tests/conftest.py
@@ -57,8 +57,12 @@
     location,
 )
 from tests.utils.auth_permissions_util import default_store
-from tests.utils.generate_self_signed_certifcate_util import generate_self_signed_cert
 from tests.utils.http_server import check_port_open, free_port  # noqa: E402
+from tests.utils.ssl_certifcates_util import (
+    combine_trust_stores,
+    create_ca_trust_store,
+    generate_self_signed_cert,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -514,17 +518,36 @@ def auth_config(request, is_integration_test):
     return auth_configuration
 
 
-@pytest.fixture(params=[True, False], scope="module")
+@pytest.fixture(scope="module")
 def tls_mode(request):
-    is_tls_mode = request.param
+    is_tls_mode = request.param[0]
+    output_combined_truststore_path = ""
 
     if is_tls_mode:
         certificates_path = tempfile.mkdtemp()
         tls_key_path = os.path.join(certificates_path, "key.pem")
         tls_cert_path = os.path.join(certificates_path, "cert.pem")
+
         generate_self_signed_cert(cert_path=tls_cert_path, key_path=tls_key_path)
+        is_ca_trust_store_set = request.param[1]
+        if is_ca_trust_store_set:
+            # Paths
+            feast_ca_trust_store_path = os.path.join(
+                certificates_path, "feast_trust_store.pem"
+            )
+            create_ca_trust_store(
+                public_key_path=tls_cert_path,
+                private_key_path=tls_key_path,
+                output_trust_store_path=feast_ca_trust_store_path,
+            )
+
+            # Combine trust stores
+            output_combined_path = os.path.join(
+                certificates_path, "combined_trust_store.pem"
+            )
+            combine_trust_stores(feast_ca_trust_store_path, output_combined_path)
     else:
         tls_key_path = ""
         tls_cert_path = ""
 
-    return is_tls_mode, tls_key_path, tls_cert_path
+    return is_tls_mode, tls_key_path, tls_cert_path, output_combined_truststore_path
diff --git a/sdk/python/tests/integration/feature_repos/universal/data_sources/file.py b/sdk/python/tests/integration/feature_repos/universal/data_sources/file.py
@@ -34,8 +34,8 @@
     DataSourceCreator,
 )
 from tests.utils.auth_permissions_util import include_auth_config
-from tests.utils.generate_self_signed_certifcate_util import generate_self_signed_cert
 from tests.utils.http_server import check_port_open, free_port  # noqa: E402
+from tests.utils.ssl_certifcates_util import generate_self_signed_cert
 
 logger = logging.getLogger(__name__)
 
diff --git a/sdk/python/tests/integration/online_store/test_remote_online_store.py b/sdk/python/tests/integration/online_store/test_remote_online_store.py
@@ -22,6 +22,9 @@
 
 
 @pytest.mark.integration
+@pytest.mark.parametrize(
+    "tls_mode", [("True", "True"), ("True", "False"), ("False", "")], indirect=True
+)
 def test_remote_online_store_read(auth_config, tls_mode):
     with (
         tempfile.TemporaryDirectory() as remote_server_tmp_dir,
@@ -56,13 +59,13 @@ def test_remote_online_store_read(auth_config, tls_mode):
             )
         )
         assert None not in (server_store, server_url, registry_path)
-        _, _, tls_cert_path = tls_mode
+
         client_store = _create_remote_client_feature_store(
             temp_dir=remote_client_tmp_dir,
             server_registry_path=str(registry_path),
             feature_server_url=server_url,
             auth_config=auth_config,
-            tls_cert_path=tls_cert_path,
+            tls_mode=tls_mode,
         )
         assert client_store is not None
         _assert_non_existing_entity_feature_views_entity(
@@ -172,14 +175,15 @@ def _create_server_store_spin_feature_server(
 ):
     store = default_store(str(temp_dir), auth_config, permissions_list)
     feast_server_port = free_port()
-    is_tls_mode, tls_key_path, tls_cert_path = tls_mode
+    is_tls_mode, tls_key_path, tls_cert_path, ca_trust_store_path = tls_mode
 
     server_url = next(
         start_feature_server(
             repo_path=str(store.repo_path),
             server_port=feast_server_port,
             tls_key_path=tls_key_path,
             tls_cert_path=tls_cert_path,
+            ca_trust_store_path=ca_trust_store_path,
         )
     )
     if is_tls_mode:
@@ -203,20 +207,33 @@ def _create_remote_client_feature_store(
     server_registry_path: str,
     feature_server_url: str,
     auth_config: str,
-    tls_cert_path: str = "",
+    tls_mode,
 ) -> FeatureStore:
     project_name = "REMOTE_ONLINE_CLIENT_PROJECT"
     runner = CliRunner()
     result = runner.run(["init", project_name], cwd=temp_dir)
     assert result.returncode == 0
     repo_path = os.path.join(temp_dir, project_name, "feature_repo")
-    _overwrite_remote_client_feature_store_yaml(
-        repo_path=str(repo_path),
-        registry_path=server_registry_path,
-        feature_server_url=feature_server_url,
-        auth_config=auth_config,
-        tls_cert_path=tls_cert_path,
-    )
+    is_tls_mode, _, tls_cert_path, ca_trust_store_path = tls_mode
+    if is_tls_mode and not ca_trust_store_path:
+        _overwrite_remote_client_feature_store_yaml(
+            repo_path=str(repo_path),
+            registry_path=server_registry_path,
+            feature_server_url=feature_server_url,
+            auth_config=auth_config,
+            tls_cert_path=tls_cert_path,
+        )
+    else:
+        _overwrite_remote_client_feature_store_yaml(
+            repo_path=str(repo_path),
+            registry_path=server_registry_path,
+            feature_server_url=feature_server_url,
+            auth_config=auth_config,
+        )
+
+    if is_tls_mode and ca_trust_store_path:
+        # configure trust store path only when is_tls_mode and ca_trust_store_path exists.
+        os.environ["FEAST_CA_CERT_FILE_PATH"] = ca_trust_store_path
 
     return FeatureStore(repo_path=repo_path)
 
diff --git a/sdk/python/tests/unit/permissions/auth/server/test_auth_registry_server.py b/sdk/python/tests/unit/permissions/auth/server/test_auth_registry_server.py
@@ -44,7 +44,7 @@ def start_registry_server(
 
     assertpy.assert_that(server_port).is_not_equal_to(0)
 
-    is_tls_mode, tls_key_path, tls_cert_path = tls_mode
+    is_tls_mode, tls_key_path, tls_cert_path, tls_ca_file_path = tls_mode
     if is_tls_mode:
         print(f"Starting Registry in TLS mode at {server_port}")
         server = start_server(
@@ -74,6 +74,9 @@ def start_registry_server(
     server.stop(grace=None)  # Teardown server
 
 
+@pytest.mark.parametrize(
+    "tls_mode", [("True", "True"), ("True", "False"), ("False", "")], indirect=True
+)
 def test_registry_apis(
     auth_config,
     tls_mode,
diff --git a/sdk/python/tests/utils/auth_permissions_util.py b/sdk/python/tests/utils/auth_permissions_util.py
@@ -60,6 +60,7 @@ def start_feature_server(
     metrics: bool = False,
     tls_key_path: str = "",
     tls_cert_path: str = "",
+    ca_trust_store_path: str = "",
 ):
     host = "0.0.0.0"
     cmd = [
@@ -127,18 +128,30 @@ def start_feature_server(
 
 
 def get_remote_registry_store(server_port, feature_store, tls_mode):
-    is_tls_mode, _, tls_cert_path = tls_mode
+    is_tls_mode, _, tls_cert_path, ca_trust_store_path = tls_mode
     if is_tls_mode:
-        registry_config = RemoteRegistryConfig(
-            registry_type="remote",
-            path=f"localhost:{server_port}",
-            cert=tls_cert_path,
-        )
+        if ca_trust_store_path:
+            registry_config = RemoteRegistryConfig(
+                registry_type="remote",
+                path=f"localhost:{server_port}",
+                is_tls=True,
+            )
+        else:
+            registry_config = RemoteRegistryConfig(
+                registry_type="remote",
+                path=f"localhost:{server_port}",
+                is_tls=True,
+                cert=tls_cert_path,
+            )
     else:
         registry_config = RemoteRegistryConfig(
             registry_type="remote", path=f"localhost:{server_port}"
         )
 
+    if is_tls_mode and ca_trust_store_path:
+        # configure trust store path only when is_tls_mode and ca_trust_store_path exists.
+        os.environ["FEAST_CA_CERT_FILE_PATH"] = ca_trust_store_path
+
     store = FeatureStore(
         config=RepoConfig(
             project=PROJECT_NAME,
diff --git a/sdk/python/tests/utils/generate_self_signed_certifcate_util.py b/sdk/python/tests/utils/generate_self_signed_certifcate_util.py
diff --git a/sdk/python/tests/utils/ssl_certifcates_util.py b/sdk/python/tests/utils/ssl_certifcates_util.py

Original file line number	Diff line number	Diff line change
`@@ -982,7 +982,6 @@ def serve_command(`
`982`	`982`	`raise click.BadParameter(`
`983`	`983`	`"Please pass --cert and --key args to start the feature server in TLS mode."`
`984`	`984`	`)`
`985`		`-`
`986`	`985`	`store = create_feature_store(ctx)`
`987`	`986`
`988`	`987`	`store.serve(`
Original file line number	Diff line number	Diff line change
`@@ -74,7 +74,7 @@ def build_arrow_flight_client(`
`74`	`74`	`scheme: str, host: str, port, auth_config: AuthConfig, cert: str = ""`
`75`	`75`	`):`
`76`	`76`	`arrow_scheme = "grpc+tcp"`
`77`		`- if cert:`
	`77`	`+ if scheme == "https":`
`78`	`78`	`logger.info(`
`79`	`79`	`"Scheme is https so going to connect offline server in SSL(TLS) mode."`
`80`	`80`	`)`
Original file line number	Diff line number	Diff line change
`@@ -34,8 +34,8 @@`
`34`	`34`	`DataSourceCreator,`
`35`	`35`	`)`
`36`	`36`	`from tests.utils.auth_permissions_util import include_auth_config`
`37`		`-from tests.utils.generate_self_signed_certifcate_util import generate_self_signed_cert`
`38`	`37`	`from tests.utils.http_server import check_port_open, free_port # noqa: E402`
	`38`	`+from tests.utils.ssl_certifcates_util import generate_self_signed_cert`
`39`	`39`
`40`	`40`	`logger = logging.getLogger(__name__)`
`41`	`41`