* Adding the option to pass public certificate to pass as part of the feature_store.yaml so that we can test self-signed certificate.

lokeshrangineni · lokeshrangineni · commit 5f99a8d7b349 · 2024-10-24T16:50:05.000-04:00
* Adding the integration test to run the remote online server in SSL and non SSL mode.

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/sdk/python/feast/infra/online_stores/remote.py b/sdk/python/feast/infra/online_stores/remote.py
@@ -41,6 +41,10 @@ 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 path in case if the online server starts in SSL mode. This may be needed especially if it is 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. """
+
 
 class RemoteOnlineStore(OnlineStore):
     """
@@ -170,6 +174,13 @@ def teardown(
 def get_remote_online_features(
     session: requests.Session, config: RepoConfig, req_body: str
 ) -> requests.Response:
-    return session.post(
-        f"{config.online_store.path}/get-online-features", data=req_body
-    )
+    if config.online_store.ssl_cert_path:
+        return session.post(
+            f"{config.online_store.path}/get-online-features",
+            data=req_body,
+            verify=config.online_store.ssl_cert_path,
+        )
+    else:
+        return session.post(
+            f"{config.online_store.path}/get-online-features", data=req_body
+        )
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
@@ -15,12 +15,13 @@
     start_feature_server,
 )
 from tests.utils.cli_repo_creator import CliRunner
+from tests.utils.generate_self_signed_certifcate_util import generate_self_signed_cert
 from tests.utils.http_server import free_port
 
 
-@pytest.mark.parametrize("run_ssl", [True, False])
+@pytest.mark.parametrize("ssl_mode", [True, False])
 @pytest.mark.integration
-def test_remote_online_store_read(auth_config, run_ssl):
+def test_remote_online_store_read(auth_config, ssl_mode):
     with tempfile.TemporaryDirectory() as remote_server_tmp_dir, tempfile.TemporaryDirectory() as remote_client_tmp_dir:
         permissions_list = [
             Permission(
@@ -42,12 +43,12 @@ def test_remote_online_store_read(auth_config, run_ssl):
                 actions=[AuthzedAction.READ_ONLINE],
             ),
         ]
-        server_store, server_url, registry_path = (
+        server_store, server_url, registry_path, ssl_cert_path = (
             _create_server_store_spin_feature_server(
                 temp_dir=remote_server_tmp_dir,
                 auth_config=auth_config,
                 permissions_list=permissions_list,
-                run_ssl=run_ssl,
+                ssl_mode=ssl_mode,
             )
         )
         assert None not in (server_store, server_url, registry_path)
@@ -56,6 +57,7 @@ def test_remote_online_store_read(auth_config, run_ssl):
             server_registry_path=str(registry_path),
             feature_server_url=server_url,
             auth_config=auth_config,
+            ssl_cert_path=ssl_cert_path,
         )
         assert client_store is not None
         _assert_non_existing_entity_feature_views_entity(
@@ -161,23 +163,46 @@ def _assert_client_server_online_stores_are_matching(
 
 
 def _create_server_store_spin_feature_server(
-    temp_dir, auth_config: str, permissions_list, run_ssl: bool
+    temp_dir, auth_config: str, permissions_list, ssl_mode: bool
 ):
     store = default_store(str(temp_dir), auth_config, permissions_list)
     feast_server_port = free_port()
+    if ssl_mode:
+        certificates_path = tempfile.mkdtemp()
+        ssl_key_path = os.path.join(certificates_path, "key.pem")
+        ssl_cert_path = os.path.join(certificates_path, "cert.pem")
+        generate_self_signed_cert(cert_path=ssl_cert_path, key_path=ssl_key_path)
+    else:
+        ssl_key_path = ""
+        ssl_cert_path = ""
+
     server_url = next(
         start_feature_server(
             repo_path=str(store.repo_path),
             server_port=feast_server_port,
-            run_ssl=run_ssl,
+            ssl_key_path=ssl_key_path,
+            ssl_cert_path=ssl_cert_path,
         )
     )
-    print(f"Server started successfully, {server_url}")
-    return store, server_url, os.path.join(store.repo_path, "data", "registry.db")
+    if ssl_cert_path and ssl_key_path:
+        print(f"Online Server started successfully in SSL mode, {server_url}")
+    else:
+        print(f"Server started successfully, {server_url}")
+
+    return (
+        store,
+        server_url,
+        os.path.join(store.repo_path, "data", "registry.db"),
+        ssl_cert_path,
+    )
 
 
 def _create_remote_client_feature_store(
-    temp_dir, server_registry_path: str, feature_server_url: str, auth_config: str
+    temp_dir,
+    server_registry_path: str,
+    feature_server_url: str,
+    auth_config: str,
+    ssl_cert_path: str = "",
 ) -> FeatureStore:
     project_name = "REMOTE_ONLINE_CLIENT_PROJECT"
     runner = CliRunner()
@@ -189,27 +214,50 @@ def _create_remote_client_feature_store(
         registry_path=server_registry_path,
         feature_server_url=feature_server_url,
         auth_config=auth_config,
+        ssl_cert_path=ssl_cert_path,
     )
 
     return FeatureStore(repo_path=repo_path)
 
 
 def _overwrite_remote_client_feature_store_yaml(
-    repo_path: str, registry_path: str, feature_server_url: str, auth_config: str
+    repo_path: str,
+    registry_path: str,
+    feature_server_url: str,
+    auth_config: str,
+    ssl_cert_path: str = "",
 ):
     repo_config = os.path.join(repo_path, "feature_store.yaml")
     with open(repo_config, "w") as repo_config:
-        repo_config.write(
-            dedent(
-                f"""
-            project: {PROJECT_NAME}
-            registry: {registry_path}
-            provider: local
-            online_store:
-                path: {feature_server_url}
-                type: remote
-            entity_key_serialization_version: 2
-            """
+        if ssl_cert_path:
+            repo_config.write(
+                dedent(
+                    f"""
+                project: {PROJECT_NAME}
+                registry: {registry_path}
+                provider: local
+                online_store:
+                    path: {feature_server_url}
+                    type: remote
+                    ssl_cert_path: {ssl_cert_path}
+                entity_key_serialization_version: 2
+                """
+                )
+                + auth_config
+            )
+
+        else:
+            repo_config.write(
+                dedent(
+                    f"""
+                project: {PROJECT_NAME}
+                registry: {registry_path}
+                provider: local
+                online_store:
+                    path: {feature_server_url}
+                    type: remote
+                entity_key_serialization_version: 2
+                """
+                )
+                + auth_config
             )
-            + auth_config
-        )
diff --git a/sdk/python/tests/utils/auth_permissions_util.py b/sdk/python/tests/utils/auth_permissions_util.py
@@ -55,7 +55,11 @@ def default_store(
 
 
 def start_feature_server(
-    repo_path: str, server_port: int, metrics: bool = False, run_ssl: bool = False
+    repo_path: str,
+    server_port: int,
+    metrics: bool = False,
+    ssl_key_path: str = "",
+    ssl_cert_path: str = "",
 ):
     host = "0.0.0.0"
     cmd = [
@@ -68,12 +72,11 @@ def start_feature_server(
         str(server_port),
     ]
 
-    if run_ssl:
-        # TODO: Generate these certificates during the test run.
+    if ssl_cert_path and ssl_cert_path:
         cmd.append("--ssl-key-path")
-        cmd.append("test_key.pem")
-        cmd.append("--ssl-key-path")
-        cmd.append("test_cert.pem")
+        cmd.append(ssl_key_path)
+        cmd.append("--ssl-cert-path")
+        cmd.append(ssl_cert_path)
 
     feast_server_process = subprocess.Popen(
         cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE
@@ -101,12 +104,14 @@ def start_feature_server(
             "localhost", 8000
         ), "Prometheus server is running when it should be disabled."
 
-    yield (
+    online_server_url = (
         f"https://localhost:{server_port}"
-        if run_ssl
+        if ssl_key_path and ssl_cert_path
         else f"http://localhost:{server_port}"
     )
 
+    yield (online_server_url)
+
     if feast_server_process is not None:
         feast_server_process.kill()
 
diff --git a/sdk/python/tests/utils/generate_self_signed_certifcate_util.py b/sdk/python/tests/utils/generate_self_signed_certifcate_util.py
@@ -0,0 +1,73 @@
+import logging
+from datetime import datetime, timedelta
+
+from cryptography import x509
+from cryptography.hazmat.backends import default_backend
+from cryptography.hazmat.primitives import hashes, serialization
+from cryptography.hazmat.primitives.asymmetric import rsa
+from cryptography.x509.oid import NameOID
+
+logger = logging.getLogger(__name__)
+
+
+def generate_self_signed_cert(
+    cert_path="cert.pem", key_path="key.pem", common_name="localhost"
+):
+    """
+    Generate a self-signed certificate and save it to the specified paths.
+
+    :param cert_path: Path to save the certificate (PEM format)
+    :param key_path: Path to save the private key (PEM format)
+    :param common_name: Common name (CN) for the certificate, defaults to 'localhost'
+    """
+    # Generate private key
+    key = rsa.generate_private_key(
+        public_exponent=65537, key_size=2048, backend=default_backend()
+    )
+
+    # Create a self-signed certificate
+    subject = issuer = x509.Name(
+        [
+            x509.NameAttribute(NameOID.COUNTRY_NAME, "US"),
+            x509.NameAttribute(NameOID.STATE_OR_PROVINCE_NAME, "California"),
+            x509.NameAttribute(NameOID.LOCALITY_NAME, "San Francisco"),
+            x509.NameAttribute(NameOID.ORGANIZATION_NAME, "Feast"),
+            x509.NameAttribute(NameOID.COMMON_NAME, common_name),
+        ]
+    )
+
+    certificate = (
+        x509.CertificateBuilder()
+        .subject_name(subject)
+        .issuer_name(issuer)
+        .public_key(key.public_key())
+        .serial_number(x509.random_serial_number())
+        .not_valid_before(datetime.utcnow())
+        .not_valid_after(
+            # Certificate valid for 1 year
+            datetime.utcnow() + timedelta(days=365)
+        )
+        .add_extension(
+            x509.SubjectAlternativeName([x509.DNSName(common_name)]),
+            critical=False,
+        )
+        .sign(key, hashes.SHA256(), default_backend())
+    )
+
+    # Write the private key to a file
+    with open(key_path, "wb") as f:
+        f.write(
+            key.private_bytes(
+                encoding=serialization.Encoding.PEM,
+                format=serialization.PrivateFormat.TraditionalOpenSSL,
+                encryption_algorithm=serialization.NoEncryption(),
+            )
+        )
+
+    # Write the certificate to a file
+    with open(cert_path, "wb") as f:
+        f.write(certificate.public_bytes(serialization.Encoding.PEM))
+
+    logger.info(
+        f"Self-signed certificate and private key have been generated at {cert_path} and {key_path}."
+    )
diff --git a/sdk/python/tests/utils/test_cert.pem b/sdk/python/tests/utils/test_cert.pem
diff --git a/sdk/python/tests/utils/test_key.pem b/sdk/python/tests/utils/test_key.pem
