feat: Adding ssl support for registry server. (#4718)

lokeshrangineni · web-flow · commit ccf7a55e1116 · 2024-11-01T13:57:47.000-04:00
diff --git a/docs/reference/feature-servers/offline-feature-server.md b/docs/reference/feature-servers/offline-feature-server.md
@@ -23,7 +23,7 @@ helm install feast-offline-server feast-charts/feast-feature-server --set feast_
 
 ## Server Example
 
-The complete example can be find under [remote-offline-store-example](../../../examples/remote-offline-store)
+The complete example can be found under [remote-offline-store-example](../../../examples/remote-offline-store)
 
 ## How to configure the client
 
diff --git a/docs/reference/feature-servers/python-feature-server.md b/docs/reference/feature-servers/python-feature-server.md
@@ -200,12 +200,12 @@ requests.post(
     data=json.dumps(push_data))
 ```
 
-## Starting the feature server in SSL mode
+## Starting the feature server in TLS(SSL) mode
 
-Enabling SSL mode ensures that data between the Feast client and server is transmitted securely. For an ideal production environment, it is recommended to start the feature server in SSL mode.
+Enabling TLS mode ensures that data between the Feast client and server is transmitted securely. For an ideal production environment, it is recommended to start the feature server in TLS mode.
 
-### Obtaining a self-signed SSL certificate and key
-In development mode we can generate a self-signed certificate for testing. In an actual production environment it is always recommended to get it from a trusted SSL certificate provider.
+### Obtaining a self-signed TLS certificate and key
+In development mode we can generate a self-signed certificate for testing. In an actual production environment it is always recommended to get it from a trusted TLS certificate provider.
 
 ```shell
 openssl req -x509 -newkey rsa:2048 -keyout key.pem -out cert.pem -days 365 -nodes
@@ -215,11 +215,11 @@ The above command will generate two files
 * `key.pem` : certificate private key
 * `cert.pem`: certificate public key
 
-### Starting the Online Server in SSL Mode
-To start the feature server in SSL mode, you need to provide the private and public keys using the `--ssl-key-path` and `--ssl-cert-path` arguments with the `feast serve` command.
+### Starting the Online Server in TLS(SSL) Mode
+To start the feature server in TLS mode, you need to provide the private and public keys using the `--key` and `--cert` arguments with the `feast serve` command.
 
 ```shell
-feast serve --ssl-key-path key.pem --ssl-cert-path cert.pem
+feast serve --key /path/to/key.pem --cert /path/to/cert.pem
 ```
 
 # Online Feature Server Permissions and Access Control
diff --git a/docs/reference/online-stores/remote.md b/docs/reference/online-stores/remote.md
@@ -16,14 +16,14 @@ provider: local
 online_store:
   path: http://localhost:6566
   type: remote
-  ssl_cert_path: /path/to/cert.pem
+  cert: /path/to/cert.pem
 entity_key_serialization_version: 2
 auth:
   type: no_auth
 ```
 {% endcode %}
 
-`ssl_cert_path` is an optional configuration to the public certificate path when the online server starts in SSL mode. This may be needed if the online server is started with a self-signed certificate, typically this file ends with `*.crt`, `*.cer`, or `*.pem`.
+`cert` is an optional configuration to the public certificate path when the online server starts in TLS(SSL) mode. This may be needed if the online server is started with a self-signed certificate, typically this file ends with `*.crt`, `*.cer`, or `*.pem`.
 
 ## How to configure Authentication and Authorization
 Please refer the [page](./../../../docs/getting-started/concepts/permission.md) for more details on how to configure authentication and authorization.
diff --git a/sdk/python/feast/cli.py b/sdk/python/feast/cli.py
@@ -912,20 +912,22 @@ def init_command(project_directory, minimal: bool, template: str):
     show_default=True,
 )
 @click.option(
-    "--ssl-key-path",
+    "--key",
     "-k",
+    "tls_key_path",
     type=click.STRING,
     default="",
     show_default=False,
-    help="path to SSL certificate private key. You need to pass ssl-cert-path as well to start server in SSL mode",
+    help="path to TLS certificate private key. You need to pass --cert as well to start server in TLS mode",
 )
 @click.option(
-    "--ssl-cert-path",
+    "--cert",
     "-c",
+    "tls_cert_path",
     type=click.STRING,
     default="",
     show_default=False,
-    help="path to SSL certificate public key. You need to pass ssl-key-path as well to start server in SSL mode",
+    help="path to TLS certificate public key. You need to pass --key as well to start server in TLS mode",
 )
 @click.option(
     "--metrics",
@@ -944,14 +946,14 @@ def serve_command(
     workers: int,
     metrics: bool,
     keep_alive_timeout: int,
-    ssl_key_path: str,
-    ssl_cert_path: str,
+    tls_key_path: str,
+    tls_cert_path: str,
     registry_ttl_sec: int = 5,
 ):
     """Start a feature server locally on a given port."""
-    if (ssl_key_path and not ssl_cert_path) or (not ssl_key_path and ssl_cert_path):
+    if (tls_key_path and not tls_cert_path) or (not tls_key_path and tls_cert_path):
         raise click.BadParameter(
-            "Please configure ssl-cert-path and ssl-key-path args to start the feature server in SSL mode."
+            "Please pass --cert and --key args to start the feature server in TLS mode."
         )
 
     store = create_feature_store(ctx)
@@ -964,8 +966,8 @@ def serve_command(
         workers=workers,
         metrics=metrics,
         keep_alive_timeout=keep_alive_timeout,
-        ssl_key_path=ssl_key_path,
-        ssl_cert_path=ssl_cert_path,
+        tls_key_path=tls_key_path,
+        tls_cert_path=tls_cert_path,
         registry_ttl_sec=registry_ttl_sec,
     )
 
@@ -1035,12 +1037,39 @@ def serve_transformations_command(ctx: click.Context, port: int):
     default=DEFAULT_REGISTRY_SERVER_PORT,
     help="Specify a port for the server",
 )
+@click.option(
+    "--key",
+    "-k",
+    "tls_key_path",
+    type=click.STRING,
+    default="",
+    show_default=False,
+    help="path to TLS certificate private key. You need to pass --cert as well to start server in TLS mode",
+)
+@click.option(
+    "--cert",
+    "-c",
+    "tls_cert_path",
+    type=click.STRING,
+    default="",
+    show_default=False,
+    help="path to TLS certificate public key. You need to pass --key as well to start server in TLS mode",
+)
 @click.pass_context
-def serve_registry_command(ctx: click.Context, port: int):
+def serve_registry_command(
+    ctx: click.Context,
+    port: int,
+    tls_key_path: str,
+    tls_cert_path: str,
+):
     """Start a registry server locally on a given port."""
+    if (tls_key_path and not tls_cert_path) or (not tls_key_path and tls_cert_path):
+        raise click.BadParameter(
+            "Please pass --cert and --key args to start the registry server in TLS mode."
+        )
     store = create_feature_store(ctx)
 
-    store.serve_registry(port)
+    store.serve_registry(port, tls_key_path, tls_cert_path)
 
 
 @cli.command("serve_offline")
diff --git a/sdk/python/feast/feature_server.py b/sdk/python/feast/feature_server.py
@@ -339,10 +339,14 @@ def start_server(
     workers: int,
     keep_alive_timeout: int,
     registry_ttl_sec: int,
-    ssl_key_path: str,
-    ssl_cert_path: str,
+    tls_key_path: str,
+    tls_cert_path: str,
     metrics: bool,
 ):
+    if (tls_key_path and not tls_cert_path) or (not tls_key_path and tls_cert_path):
+        raise ValueError(
+            "Both key and cert file paths are required to start server in TLS mode."
+        )
     if metrics:
         logger.info("Starting Prometheus Server")
         start_http_server(8000)
@@ -375,22 +379,22 @@ def start_server(
         }
 
         # Add SSL options if the paths exist
-        if ssl_key_path and ssl_cert_path:
-            options["keyfile"] = ssl_key_path
-            options["certfile"] = ssl_cert_path
+        if tls_key_path and tls_cert_path:
+            options["keyfile"] = tls_key_path
+            options["certfile"] = tls_cert_path
         FeastServeApplication(store=store, **options).run()
     else:
         import uvicorn
 
         app = get_app(store, registry_ttl_sec)
-        if ssl_key_path and ssl_cert_path:
+        if tls_key_path and tls_cert_path:
             uvicorn.run(
                 app,
                 host=host,
                 port=port,
                 access_log=(not no_access_log),
-                ssl_keyfile=ssl_key_path,
-                ssl_certfile=ssl_cert_path,
+                ssl_keyfile=tls_key_path,
+                ssl_certfile=tls_cert_path,
             )
         else:
             uvicorn.run(app, host=host, port=port, access_log=(not no_access_log))
diff --git a/sdk/python/feast/feature_store.py b/sdk/python/feast/feature_store.py
@@ -1896,8 +1896,8 @@ def serve(
         workers: int = 1,
         metrics: bool = False,
         keep_alive_timeout: int = 30,
-        ssl_key_path: str = "",
-        ssl_cert_path: str = "",
+        tls_key_path: str = "",
+        tls_cert_path: str = "",
         registry_ttl_sec: int = 2,
     ) -> None:
         """Start the feature consumption server locally on a given port."""
@@ -1915,8 +1915,8 @@ def serve(
             workers=workers,
             metrics=metrics,
             keep_alive_timeout=keep_alive_timeout,
-            ssl_key_path=ssl_key_path,
-            ssl_cert_path=ssl_cert_path,
+            tls_key_path=tls_key_path,
+            tls_cert_path=tls_cert_path,
             registry_ttl_sec=registry_ttl_sec,
         )
 
@@ -1949,11 +1949,15 @@ def serve_ui(
             root_path=root_path,
         )
 
-    def serve_registry(self, port: int) -> None:
+    def serve_registry(
+        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)
+        registry_server.start_server(
+            self, port=port, tls_key_path=tls_key_path, tls_cert_path=tls_cert_path
+        )
 
     def serve_offline(self, host: str, port: int) -> None:
         """Start offline server locally on a given port."""
diff --git a/sdk/python/feast/infra/online_stores/remote.py b/sdk/python/feast/infra/online_stores/remote.py
@@ -41,9 +41,9 @@ class RemoteOnlineStoreConfig(FeastConfigBaseModel):
     """ str: Path to metadata store.
     If type is 'remote', then this is a URL for registry server """
 
-    ssl_cert_path: StrictStr = ""
-    """ str: Path to the public certificate when the online server starts in SSL mode. This may be needed if the online server started with a self-signed certificate, typically this file ends with `*.crt`, `*.cer`, or `*.pem`.
-    If type is 'remote', then this configuration is needed to connect to remote online server in SSL mode. """
+    cert: StrictStr = ""
+    """ str: Path to the public certificate when the online server starts in TLS(SSL) mode. This may be needed if the online server started with a self-signed certificate, typically this file ends with `*.crt`, `*.cer`, or `*.pem`.
+    If type is 'remote', then this configuration is needed to connect to remote online server in TLS mode. """
 
 
 class RemoteOnlineStore(OnlineStore):
@@ -174,11 +174,11 @@ def teardown(
 def get_remote_online_features(
     session: requests.Session, config: RepoConfig, req_body: str
 ) -> requests.Response:
-    if config.online_store.ssl_cert_path:
+    if config.online_store.cert:
         return session.post(
             f"{config.online_store.path}/get-online-features",
             data=req_body,
-            verify=config.online_store.ssl_cert_path,
+            verify=config.online_store.cert,
         )
     else:
         return session.post(
diff --git a/sdk/python/feast/infra/registry/remote.py b/sdk/python/feast/infra/registry/remote.py
@@ -55,6 +55,10 @@ class RemoteRegistryConfig(RegistryConfig):
     """ str: Path to metadata store.
     If registry_type is 'remote', then this is a URL for registry server """
 
+    cert: StrictStr = ""
+    """ 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."""
+
 
 class RemoteRegistry(BaseRegistry):
     def __init__(
@@ -65,7 +69,17 @@ def __init__(
         auth_config: AuthConfig = NoAuthConfig(),
     ):
         self.auth_config = auth_config
-        self.channel = grpc.insecure_channel(registry_config.path)
+        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)
+
         auth_header_interceptor = GrpcClientAuthHeaderInterceptor(auth_config)
         self.channel = grpc.intercept_channel(self.channel, auth_header_interceptor)
         self.stub = RegistryServer_pb2_grpc.RegistryServerStub(self.channel)
diff --git a/sdk/python/feast/registry_server.py b/sdk/python/feast/registry_server.py
@@ -1,3 +1,4 @@
+import logging
 from concurrent import futures
 from datetime import datetime, timezone
 from typing import Optional, Union, cast
@@ -38,6 +39,9 @@
 from feast.saved_dataset import SavedDataset, ValidationReference
 from feast.stream_feature_view import StreamFeatureView
 
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.INFO)
+
 
 def _build_any_feature_view_proto(feature_view: BaseFeatureView):
     if isinstance(feature_view, StreamFeatureView):
@@ -753,7 +757,13 @@ def Proto(self, request, context):
         return self.proxied_registry.proto()
 
 
-def start_server(store: FeatureStore, port: int, wait_for_termination: bool = True):
+def start_server(
+    store: FeatureStore,
+    port: int,
+    wait_for_termination: bool = True,
+    tls_key_path: str = "",
+    tls_cert_path: str = "",
+):
     auth_manager_type = str_to_auth_manager_type(store.config.auth_config.type)
     init_security_manager(auth_type=auth_manager_type, fs=store)
     init_auth_manager(
@@ -781,9 +791,25 @@ def start_server(store: FeatureStore, port: int, wait_for_termination: bool = Tr
     )
     reflection.enable_server_reflection(service_names_available_for_reflection, server)
 
-    server.add_insecure_port(f"[::]:{port}")
+    if tls_cert_path and tls_key_path:
+        with open(tls_cert_path, "rb") as cert_file, open(
+            tls_key_path, "rb"
+        ) as key_file:
+            certificate_chain = cert_file.read()
+            private_key = key_file.read()
+            server_credentials = grpc.ssl_server_credentials(
+                ((private_key, certificate_chain),)
+            )
+        logger.info("Starting grpc registry server in TLS(SSL) mode")
+        server.add_secure_port(f"[::]:{port}", server_credentials)
+    else:
+        logger.info("Starting grpc registry server in non-TLS(SSL) mode")
+        server.add_insecure_port(f"[::]:{port}")
     server.start()
     if wait_for_termination:
+        logger.info(
+            f"Grpc server started at {'https' if tls_cert_path and tls_key_path else 'http'}://localhost:{port}"
+        )
         server.wait_for_termination()
     else:
         return server
diff --git a/sdk/python/tests/conftest.py b/sdk/python/tests/conftest.py
@@ -57,6 +57,7 @@
     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
 
 logger = logging.getLogger(__name__)
@@ -511,3 +512,19 @@ def auth_config(request, is_integration_test):
             return auth_configuration.replace("KEYCLOAK_URL_PLACE_HOLDER", keycloak_url)
 
     return auth_configuration
+
+
+@pytest.fixture(params=[True, False], scope="module")
+def tls_mode(request):
+    is_tls_mode = request.param
+
+    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)
+    else:
+        tls_key_path = ""
+        tls_cert_path = ""
+
+    return is_tls_mode, tls_key_path, tls_cert_path
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
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
diff --git a/sdk/python/tests/utils/auth_permissions_util.py b/sdk/python/tests/utils/auth_permissions_util.py
