OCI: support profiles that authenticate with a security token

src/dstack/_internal/core/backends/oci/auth.py

@@ -1,3 +1,6 @@
+from pathlib import Path
+from typing import Dict, Optional
+
 import oci
 from typing_extensions import Any, Mapping
 
@@ -11,10 +14,37 @@ def get_client_config(creds: AnyOCICreds) -> Mapping[str, Any]:
     return creds.model_dump(exclude={"type"})
 
 
+def get_signer(config: Mapping[str, Any]) -> Optional[oci.signer.AbstractBaseSigner]:
+    """
+    Build a signer for a config that authenticates with a security token.
+
+    `oci session authenticate` writes a profile holding a short-lived security token and
+    an ephemeral key pair instead of a registered API key, so the API key signer the SDK
+    builds by default cannot be used with such a profile. Returns None for API key
+    profiles, letting the SDK build its default signer.
+    """
+    token_file = config.get("security_token_file")
+    if not token_file:
+        return None
+    token = Path(token_file).expanduser().read_text().strip()
+    private_key = oci.signer.load_private_key_from_file(
+        Path(config["key_file"]).expanduser(), config.get("pass_phrase")
+    )
+    return oci.auth.signers.SecurityTokenSigner(token, private_key)
+
+
+def make_client_kwargs(config: Mapping[str, Any]) -> Dict[str, Any]:
+    """
+    Extra keyword arguments every OCI client must be constructed with for `config`.
+    """
+    signer = get_signer(config)
+    return {"signer": signer} if signer is not None else {}
+
+
 def creds_valid(creds: AnyOCICreds) -> bool:
     try:
         config = get_client_config(creds)
-        client = oci.identity.IdentityClient(config)
+        client = oci.identity.IdentityClient(config, **make_client_kwargs(config))
         client.get_tenancy(config["tenancy"])
     except any_oci_exception:
         return False

src/dstack/_internal/core/backends/oci/region.py

@@ -4,7 +4,7 @@
 
 import oci
 
-from dstack._internal.core.backends.oci.auth import get_client_config
+from dstack._internal.core.backends.oci.auth import get_client_config, make_client_kwargs
 from dstack._internal.core.backends.oci.models import AnyOCICreds
 
 
@@ -15,34 +15,35 @@ class OCIRegionClient:
 
     def __init__(self, client_config: Mapping[str, Any]):
         self.client_config = client_config
+        self.client_kwargs = make_client_kwargs(client_config)
 
     @property
     def name(self) -> str:
         return self.client_config["region"]
 
     @cached_property
     def compute_client(self) -> oci.core.ComputeClient:
-        return oci.core.ComputeClient(self.client_config)
+        return oci.core.ComputeClient(self.client_config, **self.client_kwargs)
 
     @cached_property
     def identity_client(self) -> oci.identity.IdentityClient:
-        return oci.identity.IdentityClient(self.client_config)
+        return oci.identity.IdentityClient(self.client_config, **self.client_kwargs)
 
     @cached_property
     def marketplace_client(self) -> oci.marketplace.MarketplaceClient:
-        return oci.marketplace.MarketplaceClient(self.client_config)
+        return oci.marketplace.MarketplaceClient(self.client_config, **self.client_kwargs)
 
     @cached_property
     def object_storage_client(self) -> oci.object_storage.ObjectStorageClient:
-        return oci.object_storage.ObjectStorageClient(self.client_config)
+        return oci.object_storage.ObjectStorageClient(self.client_config, **self.client_kwargs)
 
     @cached_property
     def virtual_network_client(self) -> oci.core.VirtualNetworkClient:
-        return oci.core.VirtualNetworkClient(self.client_config)
+        return oci.core.VirtualNetworkClient(self.client_config, **self.client_kwargs)
 
     @cached_property
     def work_request_client(self) -> oci.work_requests.WorkRequestClient:
-        return oci.work_requests.WorkRequestClient(self.client_config)
+        return oci.work_requests.WorkRequestClient(self.client_config, **self.client_kwargs)
 
     @cached_property
     def availability_domains(self) -> List[oci.identity.models.AvailabilityDomain]:

src/tests/_internal/core/backends/oci/test_auth.py

@@ -0,0 +1,64 @@
+import oci
+import pytest
+from cryptography.hazmat.primitives import serialization
+from cryptography.hazmat.primitives.asymmetric import rsa
+
+from dstack._internal.core.backends.oci.auth import get_signer, make_client_kwargs
+
+
+@pytest.fixture
+def key_file(tmp_path) -> str:
+    key = rsa.generate_private_key(public_exponent=65537, key_size=2048)
+    path = tmp_path / "oci_api_key.pem"
+    path.write_bytes(
+        key.private_bytes(
+            encoding=serialization.Encoding.PEM,
+            format=serialization.PrivateFormat.TraditionalOpenSSL,
+            encryption_algorithm=serialization.NoEncryption(),
+        )
+    )
+    return str(path)
+
+
+@pytest.fixture
+def token_file(tmp_path) -> str:
+    path = tmp_path / "token"
+    path.write_text("a.security.token\n")
+    return str(path)
+
+
+class TestGetSigner:
+    def test_api_key_config_uses_the_sdk_default_signer(self, key_file: str):
+        config = {
+            "user": "ocid1.user.oc1..aaaa",
+            "tenancy": "ocid1.tenancy.oc1..aaaa",
+            "fingerprint": "00:11:22",
+            "key_file": key_file,
+            "region": "us-ashburn-1",
+        }
+        assert get_signer(config) is None
+        assert make_client_kwargs(config) == {}
+
+    def test_security_token_config_gets_a_security_token_signer(
+        self, key_file: str, token_file: str
+    ):
+        config = {
+            "tenancy": "ocid1.tenancy.oc1..aaaa",
+            "fingerprint": "00:11:22",
+            "key_file": key_file,
+            "security_token_file": token_file,
+            "region": "us-ashburn-1",
+        }
+        assert isinstance(get_signer(config), oci.auth.signers.SecurityTokenSigner)
+        kwargs = make_client_kwargs(config)
+        assert list(kwargs) == ["signer"]
+        assert isinstance(kwargs["signer"], oci.auth.signers.SecurityTokenSigner)
+
+    def test_security_token_is_stripped(self, key_file: str, token_file: str):
+        config = {
+            "tenancy": "ocid1.tenancy.oc1..aaaa",
+            "key_file": key_file,
+            "security_token_file": token_file,
+            "region": "us-ashburn-1",
+        }
+        assert get_signer(config).api_key == "ST$a.security.token"