From cb9859c762894c25ccfc26c78cfcc716393621e0 Mon Sep 17 00:00:00 2001 From: Federico Kamelhar Date: Wed, 23 Sep 2026 23:43:43 -0400 Subject: [PATCH] fix(oci): authenticate with profiles that use a security token `get_client_config` returns the profile as the SDK config and every OCI client is then built from it without a signer, so the SDK builds its default API key signer. A profile created by `oci session authenticate` carries a short-lived security token and an ephemeral key pair rather than a registered API key, and has no `user` entry, so that signer cannot authenticate with it: configuring the backend fails with `creds_valid()` returning False even though the same profile works with the OCI CLI. Build an `oci.auth.signers.SecurityTokenSigner` when the config has `security_token_file`, and pass it to every client the backend constructs. API key profiles keep using the SDK default signer. Verified against a session token profile: `creds_valid` now returns True, the backend configures, and offers are returned. --- .../_internal/core/backends/oci/auth.py | 32 +++++++++- .../_internal/core/backends/oci/region.py | 15 +++-- .../_internal/core/backends/oci/test_auth.py | 64 +++++++++++++++++++ 3 files changed, 103 insertions(+), 8 deletions(-) create mode 100644 src/tests/_internal/core/backends/oci/test_auth.py diff --git a/src/dstack/_internal/core/backends/oci/auth.py b/src/dstack/_internal/core/backends/oci/auth.py index 9681ef8eec..586b46c6c5 100644 --- a/src/dstack/_internal/core/backends/oci/auth.py +++ b/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 diff --git a/src/dstack/_internal/core/backends/oci/region.py b/src/dstack/_internal/core/backends/oci/region.py index 806ed76fd1..388f6e0193 100644 --- a/src/dstack/_internal/core/backends/oci/region.py +++ b/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,6 +15,7 @@ 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: @@ -22,27 +23,27 @@ def name(self) -> str: @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]: diff --git a/src/tests/_internal/core/backends/oci/test_auth.py b/src/tests/_internal/core/backends/oci/test_auth.py new file mode 100644 index 0000000000..ec3e980f7b --- /dev/null +++ b/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"