From 4f364c691dc49826f96a4c19b4e596fe0eea9391 Mon Sep 17 00:00:00 2001 From: jiaqizho Date: Mon, 27 Apr 2026 15:47:25 +0800 Subject: [PATCH] feat: add Azure cross-tenant support for external tables This adds the C++ side for Azure cross-tenant access. The new fs.azure_client_id and fs.azure_tenant_id properties are parsed into the filesystem config, and AzureFileSystem can now use our Managed Identity to exchange for a customer-tenant storage bearer without holding the customer's account key. For Iceberg, ABFS/ABFSS reads can now go through a custom ADLS storage factory when cross-tenant options are present. The bridge handles the Managed Identity to AAD token exchange, injects the bearer into ADLS requests, and keeps the normal Azure path unchanged when cross-tenant mode is not enabled. For Lance, az:// reads now get a per-call custom object store session with a bearer-token credential provider. This lets Lance read customer Azure storage through the same cross-tenant flow while filter ing out account keys, SAS tokens, and other credentials that could bypass the intended path. Signed-off-by: jiaqizho --- .../azure/azure_cross_tenant_credential.h | 96 ++++ .../milvus-storage/filesystem/azure/azurefs.h | 8 + cpp/include/milvus-storage/filesystem/fs.h | 9 + cpp/include/milvus-storage/properties.h | 2 + .../azure/azure_cross_tenant_credential.cpp | 237 +++++++++ .../filesystem/azure/azure_fs_producer.cpp | 10 +- cpp/src/filesystem/azure/azurefs.cc | 17 + cpp/src/filesystem/fs.cpp | 4 + cpp/src/format/bridge/rust/Cargo.lock | 1 + cpp/src/format/bridge/rust/Cargo.toml | 6 +- .../bridge/rust/src/azure_adls_provider.rs | 473 ++++++++++++++++++ .../rust/src/azure_cross_tenant_provider.rs | 207 ++++++++ .../bridge/rust/src/azure_federation.rs | 357 +++++++++++++ .../bridge/rust/src/iceberg_bridgeimpl.rs | 15 +- .../bridge/rust/src/lance_bridgeimpl.rs | 75 ++- cpp/src/format/bridge/rust/src/lib.rs | 3 + cpp/src/format/iceberg/iceberg_common.cpp | 21 +- cpp/src/format/lance/lance_common.cpp | 24 +- cpp/src/properties.cpp | 13 + cpp/test/format/external_table_arn_test.cpp | 288 ++++++++++- 20 files changed, 1849 insertions(+), 17 deletions(-) create mode 100644 cpp/include/milvus-storage/filesystem/azure/azure_cross_tenant_credential.h create mode 100644 cpp/src/filesystem/azure/azure_cross_tenant_credential.cpp create mode 100644 cpp/src/format/bridge/rust/src/azure_adls_provider.rs create mode 100644 cpp/src/format/bridge/rust/src/azure_cross_tenant_provider.rs create mode 100644 cpp/src/format/bridge/rust/src/azure_federation.rs diff --git a/cpp/include/milvus-storage/filesystem/azure/azure_cross_tenant_credential.h b/cpp/include/milvus-storage/filesystem/azure/azure_cross_tenant_credential.h new file mode 100644 index 000000000..5a02bac6f --- /dev/null +++ b/cpp/include/milvus-storage/filesystem/azure/azure_cross_tenant_credential.h @@ -0,0 +1,96 @@ +// Copyright 2025 Zilliz +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + +#pragma once + +#include +#include +#include +#include +#include + +#include +#include +#include + +namespace milvus_storage::fs { + +/// \brief Azure TokenCredential that performs the IMDS → AAD two-hop OAuth2 +/// federated client_assertion exchange to mint a customer-tenant Bearer. +/// +/// The two-hop flow: +/// 1. GET 169.254.169.254/metadata/identity/oauth2/token +/// ?api-version=2018-02-01 +/// &resource=api://AzureADTokenExchange +/// → JWT signed by *our* tenant, audience=api://AzureADTokenExchange. +/// 2. POST https://login.microsoftonline.com/{customer_tenant}/oauth2/v2.0/token +/// Body: client_id={customer_app}, scope=https://storage.azure.com/.default, +/// grant_type=client_credentials, client_assertion={Step-1 JWT}, +/// client_assertion_type=urn:ietf:params:oauth:client-assertion-type:jwt-bearer +/// → Customer-tenant Bearer (~3600s lifetime). +/// +/// This credential class is the C++ side mirror of `azure_federation.rs` in +/// the Rust bridge: same protocol, same audience constants. We need both +/// because the bridge handles Lance/Iceberg `plan_files` (Rust path) while +/// this class handles Iceberg parquet data reads through the C++ Arrow +/// Azure FS (since `IcebergFormatReader` calls `parquet::ParquetFormatReader` +/// with an `arrow::fs::FileSystem`). +/// +/// `ClientAssertionCredential` from Azure Identity 1.10+ would let us do +/// this with just an assertion callback, but the version vendored via Conan +/// (1.7.0-beta.3) doesn't have it yet. Subclassing `TokenCredential` +/// directly is the portable path that works against any 1.x. +/// +/// The cached bearer is refreshed when within `kRefreshOffset` seconds of +/// expiry; concurrent `GetToken` calls that miss the cache may both fetch +/// (no global mutex held during HTTP), but the cache stores whichever +/// completes last — credential thrashing is bounded by the small number +/// of in-flight FS operations rather than queue depth. +class AzureCrossTenantCredential final : public Azure::Core::Credentials::TokenCredential { + public: + /// \param tenant_id Customer's Entra ID tenant ID (the AAD authority). + /// \param client_id Customer's App Registration client_id with a + /// Federated Identity Credential trusting our MI. + AzureCrossTenantCredential(std::string tenant_id, std::string client_id); + + ~AzureCrossTenantCredential() override = default; + + Azure::Core::Credentials::AccessToken GetToken( + Azure::Core::Credentials::TokenRequestContext const& tokenRequestContext, + Azure::Core::Context const& context) const override; + + private: + /// Refresh the cached bearer when the remaining lifetime drops below this. + /// Mirrors `REFRESH_OFFSET_SECS` in `azure_federation.rs`. + static constexpr std::chrono::seconds kRefreshOffset{300}; + + /// Fetch step-1 MI assertion from IMDS. Returns the JWT string. + /// `transport` is reused across the two hops to share keep-alive. + std::string FetchManagedIdentityAssertion(Azure::Core::Http::HttpTransport& transport, + Azure::Core::Context const& context) const; + + /// Fetch step-2 customer-tenant bearer. `mi_assertion` is the JWT from + /// step 1. Returns (bearer, expires_at). + std::pair ExchangeForStorageBearer(Azure::Core::Http::HttpTransport& transport, + std::string const& mi_assertion, + Azure::Core::Context const& context) const; + + std::string tenant_id_; + std::string client_id_; + + // Mutable state for cache. Mutex protects only cache entry, NOT the HTTP + // calls — those run unlocked so concurrent misses don't serialize. + mutable std::mutex cache_mu_; + mutable std::optional cached_; +}; + +} // namespace milvus_storage::fs diff --git a/cpp/include/milvus-storage/filesystem/azure/azurefs.h b/cpp/include/milvus-storage/filesystem/azure/azurefs.h index dff1f3141..a76d0d0b7 100644 --- a/cpp/include/milvus-storage/filesystem/azure/azurefs.h +++ b/cpp/include/milvus-storage/filesystem/azure/azurefs.h @@ -215,6 +215,14 @@ struct ARROW_EXPORT AzureOptions { Status ConfigureCLICredential(); Status ConfigureWorkloadIdentityCredential(); Status ConfigureEnvironmentCredential(); + /// \brief Cross-tenant Managed-Identity → AAD federated client_assertion. + /// + /// Uses our local IMDS-attached MI to request an `api://AzureADTokenExchange` + /// audience JWT, then exchanges it at the customer's tenant + /// (`{tenant_id}/oauth2/v2.0/token`) for a Bearer scoped to + /// `https://storage.azure.com/.default`. Customer-side prerequisite: + /// App Registration with a Federated Identity Credential trusting our MI. + Status ConfigureCrossTenantCredential(const std::string& tenant_id, const std::string& client_id); bool Equals(const AzureOptions& other) const; diff --git a/cpp/include/milvus-storage/filesystem/fs.h b/cpp/include/milvus-storage/filesystem/fs.h index 11e92957a..2e666d571 100644 --- a/cpp/include/milvus-storage/filesystem/fs.h +++ b/cpp/include/milvus-storage/filesystem/fs.h @@ -199,6 +199,15 @@ struct ArrowFileSystemConfig { // Target service account email to impersonate (e.g., "sa@project.iam.gserviceaccount.com") std::string gcp_target_service_account = ""; + // Azure cross-tenant access (Managed Identity → customer-tenant bearer via + // OAuth2 federated client_assertion). When both are set, format readers + // emit bridge-private storage_options keys consumed by custom Lance/Iceberg + // providers; account_key/SAS are NOT used. + // - azure_client_id: customer's App Registration client_id + // - azure_tenant_id: customer's Entra ID tenant_id + std::string azure_client_id = ""; + std::string azure_tenant_id = ""; + // Lifetime requested for cross-tenant temporary credentials, in seconds. // Shared across providers that mint short-lived tokens: // - AWS STS AssumeRole: STS session length (valid range [900, 43200], diff --git a/cpp/include/milvus-storage/properties.h b/cpp/include/milvus-storage/properties.h index 8e1a4153d..a1becee69 100644 --- a/cpp/include/milvus-storage/properties.h +++ b/cpp/include/milvus-storage/properties.h @@ -101,6 +101,8 @@ struct PropertyInfo { // Cross-tenant access properties #define PROPERTY_FS_GCP_TARGET_SERVICE_ACCOUNT "fs.gcp_target_service_account" +#define PROPERTY_FS_AZURE_CLIENT_ID "fs.azure_client_id" +#define PROPERTY_FS_AZURE_TENANT_ID "fs.azure_tenant_id" // --- External Filesystem Properties --- // External filesystems are configured with properties following the pattern: diff --git a/cpp/src/filesystem/azure/azure_cross_tenant_credential.cpp b/cpp/src/filesystem/azure/azure_cross_tenant_credential.cpp new file mode 100644 index 000000000..51917e37e --- /dev/null +++ b/cpp/src/filesystem/azure/azure_cross_tenant_credential.cpp @@ -0,0 +1,237 @@ +// Copyright 2025 Zilliz +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + +#include "milvus-storage/filesystem/azure/azure_cross_tenant_credential.h" + +#include +#include +#include + +#include +#include +#include +#include +#include +#include + +#include + +#include "milvus-storage/common/log.h" + +namespace milvus_storage::fs { + +namespace { + +// IMDS endpoint and audience are fixed by the Azure platform; not +// configurable per VM. The audience `api://AzureADTokenExchange` is the +// well-known audience that AAD's `oauth2/v2.0/token` accepts as a +// `client_assertion` for federated identity credentials. +constexpr char kImdsTokenUrl[] = "http://169.254.169.254/metadata/identity/oauth2/token"; +constexpr char kImdsApiVersion[] = "2018-02-01"; +constexpr char kMiAudience[] = "api://AzureADTokenExchange"; + +// AAD authority host. Public Azure cloud only; sovereign clouds use +// different hosts but cross-tenant FIC across sovereign boundaries is not a +// supported scenario in this MVP. +constexpr char kAadAuthority[] = "https://login.microsoftonline.com"; + +// Scope for the customer-tenant Bearer. `.default` asks for whatever +// permissions the customer's App Registration was granted on +// `https://storage.azure.com/`. +constexpr char kStorageScope[] = "https://storage.azure.com/.default"; + +constexpr char kClientAssertionType[] = "urn:ietf:params:oauth:client-assertion-type:jwt-bearer"; + +/// Percent-encode a value for an `application/x-www-form-urlencoded` body. +/// Keep alnum + `- _ . ~` literal; everything else percent-escaped. Equivalent +/// to RFC 3986 unreserved. +std::string UrlEncode(const std::string& s) { + std::string out; + out.reserve(s.size() * 11 / 10 + 2); + for (unsigned char c : s) { + if (std::isalnum(c) || c == '-' || c == '_' || c == '.' || c == '~') { + out.push_back(static_cast(c)); + } else { + char buf[4]; + std::snprintf(buf, sizeof(buf), "%%%02X", c); + out.append(buf, 3); + } + } + return out; +} + +/// Read full response body. Use `ExtractBodyStream` rather than `GetBody` +/// because `CurlTransport` may return responses with the body still +/// streaming and `GetBody` would be empty. +std::string ReadResponseBodyAsString(Azure::Core::Http::RawResponse& response, Azure::Core::Context const& context) { + auto stream = response.ExtractBodyStream(); + std::string body; + if (stream) { + auto bytes = stream->ReadToEnd(context); + body.assign(bytes.begin(), bytes.end()); + } + // Some transport configurations populate `m_body` directly instead of + // setting a stream; fall back to that. + if (body.empty()) { + auto const& raw = response.GetBody(); + body.assign(raw.begin(), raw.end()); + } + return body; +} + +} // namespace + +AzureCrossTenantCredential::AzureCrossTenantCredential(std::string tenant_id, std::string client_id) + : Azure::Core::Credentials::TokenCredential("AzureCrossTenantCredential"), + tenant_id_(std::move(tenant_id)), + client_id_(std::move(client_id)) {} + +Azure::Core::Credentials::AccessToken AzureCrossTenantCredential::GetToken( + Azure::Core::Credentials::TokenRequestContext const& /*tokenRequestContext*/, + Azure::Core::Context const& context) const { + // Fast path: cached token still has comfortable lifetime. + { + std::lock_guard g(cache_mu_); + if (cached_) { + // Both sides as Azure::DateTime so the time_point template params line + // up; DateTime inherits from system_clock::time_point but operator- + // doesn't pick up the unrelated raw time_point on the RHS. + auto remaining = cached_->ExpiresOn - Azure::DateTime(std::chrono::system_clock::now()); + if (remaining > kRefreshOffset) { + return *cached_; + } + } + } + + // Slow path: fetch fresh. We deliberately don't hold cache_mu_ across the + // HTTP calls — concurrent misses may both fetch, but only one bearer + // wins the cache slot. That is bounded waste; holding the mutex would + // serialize unrelated FS operations behind a 200-500ms AAD round-trip. + Azure::Core::Http::CurlTransport transport; + std::string mi_assertion; + std::pair exchanged; + try { + mi_assertion = FetchManagedIdentityAssertion(transport, context); + exchanged = ExchangeForStorageBearer(transport, mi_assertion, context); + } catch (const std::exception& e) { + throw Azure::Core::Credentials::AuthenticationException(std::string("AzureCrossTenantCredential failed: ") + + e.what()); + } + + Azure::Core::Credentials::AccessToken fresh; + fresh.Token = std::move(exchanged.first); + fresh.ExpiresOn = exchanged.second; + + { + std::lock_guard g(cache_mu_); + cached_ = fresh; + } + return fresh; +} + +std::string AzureCrossTenantCredential::FetchManagedIdentityAssertion(Azure::Core::Http::HttpTransport& transport, + Azure::Core::Context const& context) const { + Azure::Core::Url url(kImdsTokenUrl); + url.AppendQueryParameter("api-version", kImdsApiVersion); + url.AppendQueryParameter("resource", kMiAudience); + + Azure::Core::Http::Request request(Azure::Core::Http::HttpMethod::Get, url); + request.SetHeader("Metadata", "true"); + + auto response = transport.Send(request, context); + if (!response) { + throw std::runtime_error("IMDS token request returned no response"); + } + auto status = response->GetStatusCode(); + if (status != Azure::Core::Http::HttpStatusCode::Ok) { + auto body = ReadResponseBodyAsString(*response, context); + throw std::runtime_error("IMDS token request failed status=" + std::to_string(static_cast(status)) + + " body=" + body); + } + + auto body = ReadResponseBodyAsString(*response, context); + folly::dynamic parsed; + try { + parsed = folly::parseJson(body); + } catch (const std::exception& e) { + throw std::runtime_error(std::string("IMDS response was not valid JSON: ") + e.what()); + } + auto access_token = parsed.getDefault("access_token", "").asString(); + if (access_token.empty()) { + throw std::runtime_error("IMDS response missing access_token"); + } + return access_token; +} + +std::pair AzureCrossTenantCredential::ExchangeForStorageBearer( + Azure::Core::Http::HttpTransport& transport, + std::string const& mi_assertion, + Azure::Core::Context const& context) const { + // POST to https://login.microsoftonline.com/{tenant_id}/oauth2/v2.0/token + // with form-urlencoded body. AAD requires: + // client_id={customer_app} + // scope=https://storage.azure.com/.default + // grant_type=client_credentials + // client_assertion={MI assertion JWT} + // client_assertion_type=urn:ietf:params:oauth:client-assertion-type:jwt-bearer + std::string token_url = std::string(kAadAuthority) + "/" + tenant_id_ + "/oauth2/v2.0/token"; + Azure::Core::Url url(token_url); + + std::string body_str; + body_str.reserve(2048); + body_str.append("client_id=").append(UrlEncode(client_id_)); + body_str.append("&scope=").append(UrlEncode(kStorageScope)); + body_str.append("&grant_type=client_credentials"); + body_str.append("&client_assertion_type=").append(UrlEncode(kClientAssertionType)); + body_str.append("&client_assertion=").append(UrlEncode(mi_assertion)); + + std::vector body_bytes(body_str.begin(), body_str.end()); + Azure::Core::IO::MemoryBodyStream body_stream(body_bytes.data(), body_bytes.size()); + + Azure::Core::Http::Request request(Azure::Core::Http::HttpMethod::Post, url, &body_stream); + request.SetHeader("Content-Type", "application/x-www-form-urlencoded"); + request.SetHeader("Content-Length", std::to_string(body_bytes.size())); + + auto response = transport.Send(request, context); + if (!response) { + throw std::runtime_error("AAD token exchange returned no response"); + } + auto status = response->GetStatusCode(); + if (status != Azure::Core::Http::HttpStatusCode::Ok) { + auto err_body = ReadResponseBodyAsString(*response, context); + throw std::runtime_error( + "AAD token exchange failed (the customer's App Registration likely has no " + "Federated Identity Credential trusting our MI, or the audience/issuer/subject " + "in the FIC don't match) tenant=" + + tenant_id_ + " client_id=" + client_id_ + " status=" + std::to_string(static_cast(status)) + + " body=" + err_body); + } + + auto resp_body = ReadResponseBodyAsString(*response, context); + folly::dynamic parsed; + try { + parsed = folly::parseJson(resp_body); + } catch (const std::exception& e) { + throw std::runtime_error(std::string("AAD response was not valid JSON: ") + e.what()); + } + auto access_token = parsed.getDefault("access_token", "").asString(); + if (access_token.empty()) { + throw std::runtime_error("AAD response missing access_token"); + } + // expires_in is seconds-from-now. AAD always sends an int here for + // client_credentials grants. + int64_t expires_in = parsed.getDefault("expires_in", 3600).asInt(); + auto expires_at = std::chrono::system_clock::now() + std::chrono::seconds(expires_in); + return {std::move(access_token), Azure::DateTime(expires_at)}; +} + +} // namespace milvus_storage::fs diff --git a/cpp/src/filesystem/azure/azure_fs_producer.cpp b/cpp/src/filesystem/azure/azure_fs_producer.cpp index 558a56abc..64e9544d8 100644 --- a/cpp/src/filesystem/azure/azure_fs_producer.cpp +++ b/cpp/src/filesystem/azure/azure_fs_producer.cpp @@ -51,7 +51,15 @@ arrow::Result AzureFileSystemProducer::Make() { } options.background_writes = config_.background_writes; - if (config_.use_iam) { + if (!config_.azure_client_id.empty() && !config_.azure_tenant_id.empty()) { + // Cross-tenant via Managed Identity. The customer's storage account is + // in their own tenant; we have no shared key. Our local IMDS-attached + // MI is exchanged at the customer's AAD authority for a Storage Bearer. + // See `AzureCrossTenantCredential` for the full two-hop protocol. + LOG_STORAGE_DEBUG_ << "Azure cross-tenant: tenant=" << config_.azure_tenant_id + << " client=" << config_.azure_client_id; + ARROW_RETURN_NOT_OK(options.ConfigureCrossTenantCredential(config_.azure_tenant_id, config_.azure_client_id)); + } else if (config_.use_iam) { const char* federated_token = getenv("AZURE_FEDERATED_TOKEN_FILE"); if (federated_token != nullptr && strlen(federated_token) > 0) { // Workload Identity diff --git a/cpp/src/filesystem/azure/azurefs.cc b/cpp/src/filesystem/azure/azurefs.cc index eff4c686e..7a1b7e02d 100644 --- a/cpp/src/filesystem/azure/azurefs.cc +++ b/cpp/src/filesystem/azure/azurefs.cc @@ -23,6 +23,7 @@ #include "milvus-storage/filesystem/azure/azurefs.h" #include "milvus-storage/filesystem/azure/azurefs_internal.h" +#include "milvus-storage/filesystem/azure/azure_cross_tenant_credential.h" #include "milvus-storage/common/extend_status.h" #include "arrow/io/memory.h" @@ -440,6 +441,22 @@ Status AzureOptions::ConfigureWorkloadIdentityCredential() { return Status::OK(); } +Status AzureOptions::ConfigureCrossTenantCredential(const std::string& tenant_id, + const std::string& client_id) { + if (tenant_id.empty() || client_id.empty()) { + return Status::Invalid( + "ConfigureCrossTenantCredential requires non-empty tenant_id and client_id"); + } + // Re-use kManagedIdentity as the dispatch case; the only thing it controls + // in MakeBlobServiceClient / MakeDataLakeServiceClient is "use + // token_credential_" — which is exactly what our custom TokenCredential + // wants. Adding a new enum value would touch every switch statement in + // this vendored file for no behavioural difference. + credential_kind_ = CredentialKind::kManagedIdentity; + token_credential_ = std::make_shared(tenant_id, client_id); + return Status::OK(); +} + Status AzureOptions::ConfigureEnvironmentCredential() { credential_kind_ = CredentialKind::kEnvironment; token_credential_ = std::make_shared(); diff --git a/cpp/src/filesystem/fs.cpp b/cpp/src/filesystem/fs.cpp index 7ce89b577..e8d5eb99d 100644 --- a/cpp/src/filesystem/fs.cpp +++ b/cpp/src/filesystem/fs.cpp @@ -148,6 +148,10 @@ arrow::Status ArrowFileSystemConfig::create_file_system_config(const milvus_stor api::GetValue(properties_map, PROPERTY_FS_USE_CRC32C_CHECKSUM)); ARROW_ASSIGN_OR_RAISE(result.gcp_target_service_account, api::GetValue(properties_map, PROPERTY_FS_GCP_TARGET_SERVICE_ACCOUNT)); + ARROW_ASSIGN_OR_RAISE(result.azure_client_id, + api::GetValue(properties_map, PROPERTY_FS_AZURE_CLIENT_ID)); + ARROW_ASSIGN_OR_RAISE(result.azure_tenant_id, + api::GetValue(properties_map, PROPERTY_FS_AZURE_TENANT_ID)); return arrow::Status::OK(); } diff --git a/cpp/src/format/bridge/rust/Cargo.lock b/cpp/src/format/bridge/rust/Cargo.lock index ec49f411d..aaba6504c 100644 --- a/cpp/src/format/bridge/rust/Cargo.lock +++ b/cpp/src/format/bridge/rust/Cargo.lock @@ -6311,6 +6311,7 @@ dependencies = [ "cxx-build", "futures", "hmac", + "http 1.4.0", "iceberg", "iceberg-storage-opendal", "lance", diff --git a/cpp/src/format/bridge/rust/Cargo.toml b/cpp/src/format/bridge/rust/Cargo.toml index 2c0f16745..0d99c250b 100644 --- a/cpp/src/format/bridge/rust/Cargo.toml +++ b/cpp/src/format/bridge/rust/Cargo.toml @@ -69,7 +69,7 @@ uuid = { version = "1", features = ["v7"] } # AWS SDK (for Lance AssumeRole credential provider) aws-config = "1" aws-credential-types = "1" -object_store = { version = "0.12", features = ["aws", "gcp"] } +object_store = { version = "0.12", features = ["aws", "gcp", "azure"] } # HTTP client for the GCP Service Account Impersonation flow in # gcp_impersonation.rs (calls metadata.google.internal + iamcredentials.googleapis.com). @@ -80,6 +80,10 @@ reqwest = { version = "0.12", default-features = false, features = ["json", "rus # `url::Url` is named in lance-io's ObjectStoreProvider::new_store signature # we implement; keep as a direct dep rather than leaning on transitive re-export. url = "2" +# Named in opendal's HttpFetch trait signature (`http::Request` → +# `http::Response`); already in the transitive tree via reqwest / +# opendal but listed directly so azure_adls_provider's HttpFetch impl resolves. +http = "1" # Crypto + encoding for Aliyun POP v1 signing in aliyun_oss_provider.rs # (ECS IMDS → sts:AssumeRole fallback when OIDC machine identity is absent). diff --git a/cpp/src/format/bridge/rust/src/azure_adls_provider.rs b/cpp/src/format/bridge/rust/src/azure_adls_provider.rs new file mode 100644 index 000000000..2b747c188 --- /dev/null +++ b/cpp/src/format/bridge/rust/src/azure_adls_provider.rs @@ -0,0 +1,473 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright Zilliz + +//! Azure ADLS Gen2 cross-tenant `Storage` for the iceberg `FileIO` pipeline. +//! +//! `iceberg-storage-opendal::azdls_config_parse` only forwards a fixed set of +//! `adls.*` keys to opendal's `AzdlsConfig`, and `AzdlsConfig` itself only +//! supports `account_key` / SAS / `client_secret` credentials — there is no +//! bearer-token field, no IMDS-with-explicit-audience, and no two-hop +//! Federated Identity flow. +//! +//! For Azure cross-tenant Managed Identity access, this module: +//! +//! 1. At factory build time, extracts the bridge-private keys +//! (`adls.cross-tenant-client-id`, `adls.cross-tenant-tenant-id`, +//! `adls.cross-tenant-refresh-secs`) plus the standard `adls.account-name` +//! and `adls.endpoint-suffix`, and stashes them on the [`AzdlsCrossTenant +//! Storage`] instance. +//! 2. On each `Storage::*` call, builds an opendal `Azdls` operator pointed +//! at the customer's account and attaches a custom [`HttpFetch`] +//! implementation that injects `Authorization: Bearer ` into +//! every outbound request. +//! 3. The bearer comes from a process-wide [`CrossTenantBearerCache`] +//! (`azure_federation`) that performs the IMDS → AAD two-hop on demand +//! and refreshes ahead of expiry. +//! +//! # Why a placeholder `account_key` is needed +//! +//! opendal's `AzdlsBackend::sign` calls `loader.load_credential().await?` +//! before issuing any HTTP request and aborts with `"no valid credential +//! found"` if the loader returns `None`. The loader has four sources +//! (`config` AK/SAS, `client_secret`, `workload_identity`, `imds`) and none +//! of them returns a credential we can use for cross-tenant. So we feed +//! `account_key` a syntactically valid base64 placeholder — the SharedKey +//! signing path runs, writes `Authorization: SharedKey ...`, and our +//! `HttpFetch` overwrites that header with the real Bearer before the +//! request leaves the process. The placeholder is never transmitted. +//! +//! Verified end-to-end in `azure_bearer_spike` — see commit history for the +//! spike code if revisiting this design. + +use std::sync::Arc; +use std::time::Duration; + +use opendal::raw::{HttpBody, HttpClient, HttpFetch}; +use opendal::services::Azdls; +use opendal::{Buffer, Operator}; +use serde::{Deserialize, Serialize}; +use tokio::sync::OnceCell; + +use iceberg::io::{ + FileMetadata, FileRead, FileWrite, InputFile, OutputFile, Storage as IcebergStorage, + StorageConfig, StorageFactory, +}; +use iceberg::{Error as IcebergError, ErrorKind as IcebergErrorKind, Result as IcebergResult}; + +use crate::azure_federation::{CrossTenantBearerCache, REFRESH_OFFSET_SECS}; + +/// Property names consumed by this factory. All three are required when +/// the factory is selected; absence triggers an error at build time. +const PROP_ACCOUNT_NAME: &str = "adls.account-name"; +const PROP_ENDPOINT_SUFFIX: &str = "adls.endpoint-suffix"; +const PROP_CLIENT_ID: &str = "adls.cross-tenant-client-id"; +const PROP_TENANT_ID: &str = "adls.cross-tenant-tenant-id"; +const PROP_REFRESH_SECS: &str = "adls.cross-tenant-refresh-secs"; + +/// Marker key used by the bridge to detect "cross-tenant mode" without +/// looking at every individual key. A caller's presence of `PROP_CLIENT_ID` +/// is enough. +pub const CROSS_TENANT_MARKER_KEY: &str = PROP_CLIENT_ID; + +/// Floor / ceiling on refresh_offset (seconds before expiry to refresh). +/// AAD-issued bearers last ~3600s; refresh_offset >= 1800 means we'd +/// refresh on every call (bearer always within window). Cap to keep behavior +/// reasonable regardless of what `load_frequency` value made it down here. +const MIN_REFRESH_OFFSET_SECS: u64 = 60; +const MAX_REFRESH_OFFSET_SECS: u64 = 1800; + +/// Placeholder Azure Storage Shared Key. Any valid base64 string works — +/// 64 zero bytes gives the canonical 88-char form. Never transmitted: our +/// `HttpFetch` overwrites the SharedKey Authorization header before send. +const PLACEHOLDER_ACCOUNT_KEY: &str = + "AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA"; + +// ============================================================================ +// StorageFactory + Storage (typetag wiring for iceberg's FileIO) +// ============================================================================ + +fn from_opendal_error(e: opendal::Error) -> IcebergError { + IcebergError::new(IcebergErrorKind::Unexpected, "Failure in doing io operation").with_source(e) +} + +#[derive(Clone, Debug, Default, Serialize, Deserialize)] +pub struct AzdlsCrossTenantStorageFactory; + +#[typetag::serde(name = "AzdlsCrossTenantStorageFactory")] +impl StorageFactory for AzdlsCrossTenantStorageFactory { + fn build(&self, config: &StorageConfig) -> IcebergResult> { + let props = config.props(); + + let account_name = props + .get(PROP_ACCOUNT_NAME) + .cloned() + .filter(|s| !s.is_empty()) + .ok_or_else(|| { + IcebergError::new( + IcebergErrorKind::DataInvalid, + format!("Azure cross-tenant: missing {PROP_ACCOUNT_NAME}"), + ) + })?; + let client_id = props + .get(PROP_CLIENT_ID) + .cloned() + .filter(|s| !s.is_empty()) + .ok_or_else(|| { + IcebergError::new( + IcebergErrorKind::DataInvalid, + format!("Azure cross-tenant: missing {PROP_CLIENT_ID}"), + ) + })?; + let tenant_id = props + .get(PROP_TENANT_ID) + .cloned() + .filter(|s| !s.is_empty()) + .ok_or_else(|| { + IcebergError::new( + IcebergErrorKind::DataInvalid, + format!("Azure cross-tenant: missing {PROP_TENANT_ID}"), + ) + })?; + let endpoint_suffix = props.get(PROP_ENDPOINT_SUFFIX).cloned().unwrap_or_default(); + + // Honor caller's refresh_secs but clamp to a sensible window. Default + // when missing/zero/unparsable is the shared REFRESH_OFFSET_SECS. + let refresh_secs = props + .get(PROP_REFRESH_SECS) + .and_then(|s| s.parse::().ok()) + .filter(|n| *n > 0) + .unwrap_or(REFRESH_OFFSET_SECS) + .clamp(MIN_REFRESH_OFFSET_SECS, MAX_REFRESH_OFFSET_SECS); + + Ok(Arc::new(AzdlsCrossTenantStorage { + account_name, + endpoint_suffix, + client_id, + tenant_id, + refresh_secs, + cache: Arc::new(OnceCell::new()), + })) + } +} + +#[derive(Clone, Debug, Default, Serialize, Deserialize)] +pub struct AzdlsCrossTenantStorage { + account_name: String, + /// Used to reconstruct fully-qualified endpoint when paths don't carry + /// `container@account.dfs.suffix` form. Falls back to "core.windows.net" + /// if empty. + endpoint_suffix: String, + client_id: String, + tenant_id: String, + refresh_secs: u64, + /// Lazily initialized on first I/O call so the struct stays Default + + /// Serialize-able for typetag's dyn-trait routing. The cache itself + /// holds a reqwest::Client and a tokio RwLock — neither of which is + /// Serialize, hence the `#[serde(skip)]`. + #[serde(skip)] + cache: Arc>>, +} + +impl AzdlsCrossTenantStorage { + async fn cache(&self) -> &Arc { + self.cache + .get_or_init(|| async { + Arc::new(CrossTenantBearerCache::new( + self.tenant_id.clone(), + self.client_id.clone(), + Duration::from_secs(self.refresh_secs), + )) + }) + .await + } + + /// Parse `abfss://container@account.dfs.suffix/path` → (container, + /// endpoint_url, relative_path). Also accepts the simpler form + /// `abfss://container/path`, in which case `account_name` and + /// `endpoint_suffix` from the storage config fill in the host. + fn parse_abfss_uri(&self, uri: &str) -> IcebergResult<(String, String, String)> { + let scheme_end = uri.find("://").ok_or_else(|| { + IcebergError::new( + IcebergErrorKind::DataInvalid, + format!("expected abfs[s]:// URI, got: {uri}"), + ) + })?; + let scheme = &uri[..scheme_end]; + if scheme != "abfs" && scheme != "abfss" { + return Err(IcebergError::new( + IcebergErrorKind::DataInvalid, + format!("expected abfs/abfss scheme, got: {scheme}"), + )); + } + let http_scheme = if scheme == "abfss" { "https" } else { "http" }; + let rest = &uri[scheme_end + 3..]; + let first_slash = rest.find('/').unwrap_or(rest.len()); + let authority = &rest[..first_slash]; + let path = if first_slash < rest.len() { + &rest[first_slash + 1..] + } else { + "" + }; + + let (container, host) = match authority.find('@') { + Some(at) => (&authority[..at], authority[at + 1..].to_string()), + None => { + // No `@` → authority is just the container, fill host from + // configured account/suffix. + let suffix = if self.endpoint_suffix.is_empty() { + "core.windows.net".to_string() + } else { + self.endpoint_suffix.clone() + }; + ( + authority, + format!("{}.dfs.{}", self.account_name, suffix), + ) + } + }; + if container.is_empty() { + return Err(IcebergError::new( + IcebergErrorKind::DataInvalid, + format!("abfss URI missing container: {uri}"), + )); + } + let endpoint = format!("{http_scheme}://{host}"); + Ok((container.to_string(), endpoint, path.to_string())) + } + + async fn create_operator(&self, abfss_path: &str) -> IcebergResult<(Operator, String)> { + let (container, endpoint, relative) = self.parse_abfss_uri(abfss_path)?; + let cache = self.cache().await.clone(); + let fetcher = BearerInjectingFetcher { + inner: reqwest::Client::new(), + cache, + }; + let http_client = HttpClient::with(fetcher); + + // Spike A confirmed: the placeholder account_key satisfies opendal's + // "credential required before we'll send any HTTP" check; the + // SharedKey Authorization header signer writes is then overwritten + // by `BearerInjectingFetcher::fetch`. See module docs for why. + #[allow(deprecated)] + let builder = Azdls::default() + .account_name(&self.account_name) + .endpoint(&endpoint) + .filesystem(&container) + .account_key(PLACEHOLDER_ACCOUNT_KEY) + .http_client(http_client); + let op = Operator::new(builder) + .map_err(|e| { + IcebergError::new( + IcebergErrorKind::Unexpected, + format!("Failed to build Azdls operator: {e}"), + ) + })? + .finish(); + Ok((op, relative)) + } +} + +#[typetag::serde(name = "AzdlsCrossTenantStorage")] +#[async_trait::async_trait] +impl IcebergStorage for AzdlsCrossTenantStorage { + async fn exists(&self, path: &str) -> IcebergResult { + let (op, rel) = self.create_operator(path).await?; + op.exists(&rel).await.map_err(from_opendal_error) + } + + async fn metadata(&self, path: &str) -> IcebergResult { + let (op, rel) = self.create_operator(path).await?; + let meta = op.stat(&rel).await.map_err(from_opendal_error)?; + Ok(FileMetadata { + size: meta.content_length(), + }) + } + + async fn read(&self, path: &str) -> IcebergResult { + let (op, rel) = self.create_operator(path).await?; + Ok(op.read(&rel).await.map_err(from_opendal_error)?.to_bytes()) + } + + async fn reader(&self, path: &str) -> IcebergResult> { + let (op, rel) = self.create_operator(path).await?; + Ok(Box::new(OpenDalReader( + op.reader(&rel).await.map_err(from_opendal_error)?, + ))) + } + + async fn write(&self, path: &str, bs: bytes::Bytes) -> IcebergResult<()> { + let (op, rel) = self.create_operator(path).await?; + op.write(&rel, bs).await.map_err(from_opendal_error)?; + Ok(()) + } + + async fn writer(&self, path: &str) -> IcebergResult> { + let (op, rel) = self.create_operator(path).await?; + Ok(Box::new(OpenDalWriter( + op.writer(&rel).await.map_err(from_opendal_error)?, + ))) + } + + async fn delete(&self, path: &str) -> IcebergResult<()> { + let (op, rel) = self.create_operator(path).await?; + op.delete(&rel).await.map_err(from_opendal_error) + } + + async fn delete_prefix(&self, path: &str) -> IcebergResult<()> { + let (op, rel) = self.create_operator(path).await?; + let prefixed = if rel.ends_with('/') { + rel.clone() + } else { + format!("{rel}/") + }; + op.remove_all(&prefixed).await.map_err(from_opendal_error) + } + + fn new_input(&self, path: &str) -> IcebergResult { + Ok(InputFile::new(Arc::new(self.clone()), path.to_string())) + } + + fn new_output(&self, path: &str) -> IcebergResult { + Ok(OutputFile::new(Arc::new(self.clone()), path.to_string())) + } +} + +// ============================================================================ +// HTTP fetch wrapper that overwrites Authorization with our cross-tenant Bearer +// ============================================================================ + +struct BearerInjectingFetcher { + inner: reqwest::Client, + cache: Arc, +} + +impl HttpFetch for BearerInjectingFetcher { + async fn fetch( + &self, + mut req: http::Request, + ) -> opendal::Result> { + let bearer = self.cache.current().await.map_err(|e| { + opendal::Error::new( + opendal::ErrorKind::Unexpected, + format!("cross-tenant bearer fetch failed: {e}"), + ) + })?; + let header_value = http::HeaderValue::from_str(&format!("Bearer {}", *bearer)) + .map_err(|e| { + opendal::Error::new( + opendal::ErrorKind::Unexpected, + format!("bearer is not a valid header value: {e}"), + ) + })?; + req.headers_mut() + .insert(http::header::AUTHORIZATION, header_value); + // Delegate to opendal's built-in `impl HttpFetch for reqwest::Client`. + self.inner.fetch(req).await + } +} + +// ============================================================================ +// Iceberg-side opendal Reader/Writer wrappers (mirrors aliyun_oss_provider) +// ============================================================================ + +/// `FileRead` over an opendal reader. Same shape as the (`pub(crate)`) wrapper +/// in `iceberg-storage-opendal` and `aliyun_oss_provider`; duplicated here to +/// keep this module self-contained. +struct OpenDalReader(opendal::Reader); + +#[async_trait::async_trait] +impl FileRead for OpenDalReader { + async fn read(&self, range: std::ops::Range) -> IcebergResult { + Ok(opendal::Reader::read(&self.0, range) + .await + .map_err(from_opendal_error)? + .to_bytes()) + } +} + +/// `FileWrite` over an opendal writer. +struct OpenDalWriter(opendal::Writer); + +#[async_trait::async_trait] +impl FileWrite for OpenDalWriter { + async fn write(&mut self, bs: bytes::Bytes) -> IcebergResult<()> { + Ok(opendal::Writer::write(&mut self.0, bs) + .await + .map_err(from_opendal_error)?) + } + + async fn close(&mut self) -> IcebergResult<()> { + let _ = opendal::Writer::close(&mut self.0) + .await + .map_err(from_opendal_error)?; + Ok(()) + } +} + +// ============================================================================ +// Tests +// ============================================================================ + +#[cfg(test)] +mod tests { + use super::*; + + fn make_storage(account: &str, suffix: &str) -> AzdlsCrossTenantStorage { + AzdlsCrossTenantStorage { + account_name: account.into(), + endpoint_suffix: suffix.into(), + client_id: "client".into(), + tenant_id: "tenant".into(), + refresh_secs: REFRESH_OFFSET_SECS, + cache: Arc::new(OnceCell::new()), + } + } + + #[test] + fn parse_abfss_with_full_authority() { + let s = make_storage("ignored", "ignored"); + let (container, endpoint, rel) = s + .parse_abfss_uri("abfss://data@acme.dfs.core.windows.net/path/to/file.parquet") + .unwrap(); + assert_eq!(container, "data"); + assert_eq!(endpoint, "https://acme.dfs.core.windows.net"); + assert_eq!(rel, "path/to/file.parquet"); + } + + #[test] + fn parse_abfss_with_short_authority_uses_config() { + let s = make_storage("acme", "core.windows.net"); + let (container, endpoint, rel) = + s.parse_abfss_uri("abfss://data/path/to/file.parquet").unwrap(); + assert_eq!(container, "data"); + assert_eq!(endpoint, "https://acme.dfs.core.windows.net"); + assert_eq!(rel, "path/to/file.parquet"); + } + + #[test] + fn parse_abfss_short_authority_default_suffix() { + // Empty endpoint_suffix falls back to core.windows.net. + let s = make_storage("acme", ""); + let (_, endpoint, _) = + s.parse_abfss_uri("abfss://data/file.parquet").unwrap(); + assert_eq!(endpoint, "https://acme.dfs.core.windows.net"); + } + + #[test] + fn parse_abfss_rejects_other_schemes() { + let s = make_storage("acme", "core.windows.net"); + assert!(s.parse_abfss_uri("s3://bucket/key").is_err()); + assert!(s.parse_abfss_uri("not-a-url").is_err()); + } + + #[test] + fn factory_requires_client_and_tenant() { + // Build a minimal StorageConfig stand-in by going through the public + // `with_prop` builder. iceberg 0.9 exposes StorageConfig only via + // FileIOBuilder.with_prop().build() chain — replicating that here + // would pull half of FileIO into a unit test. Skip: integration + // coverage in the e2e test exercises the full factory path. + let factory = AzdlsCrossTenantStorageFactory::default(); + let _ = factory; // keep symbol used; full coverage is e2e + } +} diff --git a/cpp/src/format/bridge/rust/src/azure_cross_tenant_provider.rs b/cpp/src/format/bridge/rust/src/azure_cross_tenant_provider.rs new file mode 100644 index 000000000..7ac872440 --- /dev/null +++ b/cpp/src/format/bridge/rust/src/azure_cross_tenant_provider.rs @@ -0,0 +1,207 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright Zilliz + +//! Azure cross-tenant Managed Identity for `lance-io`'s `az` scheme. +//! +//! `object_store::azure::MicrosoftAzureBuilder` has a clean +//! `with_credentials(Arc>)` +//! hook. We wire up an [`AzureCredential::BearerToken`] sourced from the +//! shared [`CrossTenantBearerCache`] (`azure_federation`), which performs +//! the IMDS → AAD two-hop exchange and refreshes ahead of expiry. +//! +//! Why a custom provider rather than relying on object_store's built-in +//! Azure auth paths: +//! +//! * **`WorkloadIdentityOAuthProvider`** wants a federated token *file*, +//! which AKS workload-identity sets up via `AZURE_FEDERATED_TOKEN_FILE`. +//! Our deployment is a plain Azure VM with a system-assigned Managed +//! Identity — no token file exists. +//! * **`ImdsManagedIdentityProvider`** asks IMDS for a +//! `https://storage.azure.com/` audience token in *our* tenant. The +//! customer's storage account lives in a *different* tenant and rejects +//! that token's issuer. +//! * **`ClientSecretOAuthProvider`** would need the customer to share a +//! long-lived secret, which defeats the cross-tenant-without-secret point. +//! +//! Lance registers `az` to its built-in `AzureBlobStoreProvider` which uses +//! `with_url + with_config(k, v)` and never plumbs custom credentials. We +//! override `az` in a per-call `Session`'s `ObjectStoreRegistry` (see +//! `lance_bridgeimpl::pick_custom_session`) so only opens that opt in via +//! cross-tenant storage_options pick this provider up; everything else +//! continues to go through stock. + +use std::collections::HashMap; +use std::sync::Arc; +use std::time::Duration; + +use async_trait::async_trait; +use object_store::azure::{AzureCredential, MicrosoftAzureBuilder}; +use object_store::{ + CredentialProvider, ObjectStore as OSObjectStore, RetryConfig, Result as ObjectStoreResult, +}; +use snafu::location; +use url::Url; + +use lance::{Error as LanceError, Result as LanceResult}; +use lance_io::object_store::{ + ObjectStore, ObjectStoreParams, ObjectStoreProvider, StorageOptions, + DEFAULT_CLOUD_IO_PARALLELISM, +}; + +use crate::azure_federation::{into_object_store_err, CrossTenantBearerCache}; + +/// lance-io's `DEFAULT_CLOUD_BLOCK_SIZE` is crate-private; mirror its 64 KiB +/// value so opens through this provider behave the same as the stock Azure one. +const AZURE_DEFAULT_BLOCK_SIZE: usize = 64 * 1024; + +/// `object_store::CredentialProvider` returning a customer-tenant Bearer +/// minted via the shared cache. `get_credential` is a hot path +/// (object_store calls it on every outbound request); cache hits return in a +/// single read-lock acquisition, so the steady-state cost is negligible. +#[derive(Debug)] +pub struct CrossTenantAzureCredentialProvider { + cache: Arc, +} + +impl CrossTenantAzureCredentialProvider { + pub fn new(cache: Arc) -> Self { + Self { cache } + } +} + +#[async_trait] +impl CredentialProvider for CrossTenantAzureCredentialProvider { + type Credential = AzureCredential; + + async fn get_credential(&self) -> ObjectStoreResult> { + let bearer = self + .cache + .current() + .await + .map_err(into_object_store_err)?; + // `(*bearer).clone()` materializes the inner `String` for the + // BearerToken variant — `Arc` itself can't be moved into + // `BearerToken(String)` without dereferencing. + Ok(Arc::new(AzureCredential::BearerToken((*bearer).clone()))) + } +} + +/// Lance `ObjectStoreProvider` for `az://` opens that should use cross-tenant +/// MI. Wires a [`CrossTenantAzureCredentialProvider`] into the standard +/// `MicrosoftAzureBuilder`, and forwards every non-credential storage option +/// (account name, retry knobs, etc.) so end users can still tune behaviour. +#[derive(Debug)] +pub struct CrossTenantAzureStoreProvider { + cache: Arc, +} + +impl CrossTenantAzureStoreProvider { + pub fn new(cache: Arc) -> Self { + Self { cache } + } +} + +/// Storage-option keys that would compete with our credential provider if +/// passed through to `MicrosoftAzureBuilder` (and silently win, since the +/// builder's auth selection runs independently of `with_credentials`). +/// We strip them defensively even though the C++ side +/// (`lance_common.cpp` Azure cross-tenant branch) already declines to emit +/// them — a property file or env var sweep could still leak one in. +const CONFLICTING_AZURE_KEYS: &[&str] = &[ + "azure_storage_account_key", + "azure_account_key", + "account_key", + "azure_storage_sas_token", + "azure_storage_sas_key", + "sas_token", + "sas_key", + "azure_client_secret", + "client_secret", + "azure_storage_token", + "azure_bearer_token", + "bearer_token", +]; + +#[async_trait] +impl ObjectStoreProvider for CrossTenantAzureStoreProvider { + async fn new_store( + &self, + base_path: Url, + params: &ObjectStoreParams, + ) -> LanceResult { + let block_size = params.block_size.unwrap_or(AZURE_DEFAULT_BLOCK_SIZE); + + // Pre-filter storage_options before lance-io's `StorageOptions` env + // sweep runs, so a stray account_key in the process env can't shadow + // our credentials. + let raw_options = params.storage_options.clone().unwrap_or_default(); + let filtered: HashMap = raw_options + .into_iter() + .filter(|(k, _)| { + let lower = k.to_ascii_lowercase(); + !CONFLICTING_AZURE_KEYS.contains(&lower.as_str()) + }) + .collect(); + let mut storage_options = StorageOptions(filtered); + storage_options.with_env_azure(); + let download_retry_count = storage_options.download_retry_count(); + let max_retries = storage_options.client_max_retries(); + let retry_timeout = storage_options.client_retry_timeout(); + + let retry_config = RetryConfig { + backoff: Default::default(), + max_retries, + retry_timeout: Duration::from_secs(retry_timeout), + }; + + let mut builder = MicrosoftAzureBuilder::new() + .with_url(base_path.as_ref()) + .with_retry(retry_config); + // Forward Azure-recognized config keys (account name, endpoint, etc.). + // `as_azure_options()` filters via `AzureConfigKey::from_str`, so it + // naturally drops our bridge-private `azure_cross_tenant_*` keys + // (those don't parse as known Azure config keys). + for (key, value) in storage_options.as_azure_options() { + // Defense-in-depth: even though we filtered raw keys above, + // AzureConfigKey enumerates more than CONFLICTING_AZURE_KEYS + // covers (FabricToken*, etc.). Skip the credential-bearing keys + // a second time at the AzureConfigKey enum level. + if matches!( + key, + object_store::azure::AzureConfigKey::AccessKey + | object_store::azure::AzureConfigKey::SasKey + | object_store::azure::AzureConfigKey::ClientSecret + | object_store::azure::AzureConfigKey::Token + | object_store::azure::AzureConfigKey::FederatedTokenFile + ) { + continue; + } + builder = builder.with_config(key, value); + } + + // Plug our credential provider in last so it wins over anything + // `with_url` / `with_config` may have inferred. + let credential_provider: Arc> = + Arc::new(CrossTenantAzureCredentialProvider::new(self.cache.clone())); + builder = builder.with_credentials(credential_provider); + + let store = builder.build().map_err(|e| LanceError::IO { + source: Box::new(e), + location: location!(), + })?; + let inner = Arc::new(store) as Arc; + + Ok(ObjectStore::new( + inner, + base_path, + Some(block_size), + params.object_store_wrapper.clone(), + params.use_constant_size_upload_parts, + // Azure list is lexically ordered (matches stock AzureBlobStoreProvider). + true, + DEFAULT_CLOUD_IO_PARALLELISM, + download_retry_count, + params.storage_options.as_ref(), + )) + } +} diff --git a/cpp/src/format/bridge/rust/src/azure_federation.rs b/cpp/src/format/bridge/rust/src/azure_federation.rs new file mode 100644 index 000000000..6b4c0fd07 --- /dev/null +++ b/cpp/src/format/bridge/rust/src/azure_federation.rs @@ -0,0 +1,357 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright Zilliz + +//! Azure cross-tenant access via Managed Identity → customer-tenant bearer. +//! +//! Neither `object_store` nor `opendal`/`reqsign` natively supports this flow: +//! +//! * **`object_store::azure`**: has `WorkloadIdentityOAuthProvider` that takes +//! a federated token *file*, but on plain Azure VMs we have no such file — +//! only an IMDS-attached Managed Identity. Its built-in `ImdsManagedIdentity +//! Provider` requests `https://storage.azure.com/` audience tokens against +//! *our* tenant, which the customer's storage account rejects. +//! * **`reqsign 0.16`** (under `opendal 0.55`): same set of paths, same gap. +//! Its `load_via_imds` is single-tenant only. +//! +//! This module hand-rolls the two-hop OAuth2 exchange: +//! +//! 1. `GET 169.254.169.254/metadata/identity/oauth2/token` +//! `?resource=api://AzureADTokenExchange` — fetch a JWT signed by *our* +//! tenant's STS that names the MI as subject. This JWT is intended to be +//! used as a `client_assertion` (audience `api://AzureADTokenExchange`). +//! 2. `POST https://login.microsoftonline.com/{customer_tenant}/oauth2/v2.0/token` +//! with that JWT as `client_assertion`, `grant_type=client_credentials`, +//! `scope=https://storage.azure.com/.default`. AAD validates the customer +//! App Registration's Federated Identity Credential trusts our MI and +//! issues a customer-tenant bearer. +//! +//! The result is a short-lived bearer (~1h, AAD-decided) that authenticates +//! against the customer's storage account, without our process ever holding +//! the customer's secret/key. +//! +//! # How callers use this +//! +//! [`CrossTenantBearerCache`] is the cached, refresh-on-demand entry point — +//! both bridges share it: +//! +//! * **Iceberg** (`azure_adls_provider.rs`): the cache is wrapped in a custom +//! `iceberg::io::Storage`. opendal `AzdlsConfig` has no bearer field, so +//! we build the `Operator` with a placeholder `account_key` (any valid +//! base64) to satisfy reqsign's "credential present" check, then attach +//! an `HttpFetch` wrapper that overwrites `Authorization: Bearer ...` with +//! our cache's current value before each outbound request. +//! * **Lance** (`azure_cross_tenant_provider.rs`): the cache is wrapped in +//! an `object_store::CredentialProvider` returning +//! `AzureCredential::BearerToken`, plugged into `MicrosoftAzureBuilder:: +//! with_credentials`. Clean — no placeholder hack needed. +//! +//! # Why the endpoint URLs are inlined +//! +//! Same reasoning as `gcp_impersonation.rs`: no Rust crate in our tree +//! exposes IMDS or AAD as a library call we could reuse. Pulling in +//! `azure_identity` would drag the full Azure SDK transitive tree for two +//! JSON HTTP calls, which isn't worth it. The URLs are stable documented +//! Azure endpoints. + +use std::sync::Arc; +use std::time::{Duration, SystemTime, UNIX_EPOCH}; + +use serde::Deserialize; +use tokio::sync::RwLock; + +/// IMDS endpoint for Azure Managed Identity. Fixed by the platform, not +/// configurable per VM. The hop 1 audience `api://AzureADTokenExchange` is +/// the well-known audience that AAD's `oauth2/v2.0/token` accepts as a +/// `client_assertion` for federated identity credentials. +const IMDS_TOKEN_URL: &str = "http://169.254.169.254/metadata/identity/oauth2/token"; +const IMDS_API_VERSION: &str = "2018-02-01"; +const MI_AUDIENCE: &str = "api://AzureADTokenExchange"; + +/// AAD authority host. Fixed for the public Azure cloud; sovereign clouds +/// (Azure China, Azure Government) use different hosts but cross-tenant FIC +/// across sovereign boundaries is not a typical scenario, and adding the +/// configurability would require plumbing through the C++ side. Revisit if +/// needed. +const AAD_AUTHORITY: &str = "https://login.microsoftonline.com"; + +/// Scope passed to AAD when requesting the storage bearer. The `.default` +/// suffix asks for whatever permissions the customer's App Registration has +/// been granted on `https://storage.azure.com/` — typically Storage Blob +/// Data Reader/Contributor RBAC roles. +const STORAGE_SCOPE: &str = "https://storage.azure.com/.default"; + +const CLIENT_ASSERTION_TYPE: &str = + "urn:ietf:params:oauth:client-assertion-type:jwt-bearer"; + +/// How long before a cached bearer's expiry we trigger a refresh. Mirrors +/// the AWS / GCP paths so callers see consistent refresh behavior across +/// providers. +pub const REFRESH_OFFSET_SECS: u64 = 300; + +const HTTP_CONNECT_TIMEOUT_SECS: u64 = 10; +const HTTP_REQUEST_TIMEOUT_SECS: u64 = 30; + +/// Neutral store name used in error wrapping; matches the per-provider +/// pattern in `gcp_impersonation.rs`. +const STORE_NAME: &str = "azure_cross_tenant"; + +fn build_http_client() -> reqwest::Client { + reqwest::Client::builder() + .connect_timeout(Duration::from_secs(HTTP_CONNECT_TIMEOUT_SECS)) + .timeout(Duration::from_secs(HTTP_REQUEST_TIMEOUT_SECS)) + .build() + .expect("reqwest client builder: valid config") +} + +/// IMDS token endpoint response. We only care about the access_token; the +/// other fields are not used (we re-derive expiry from AAD's response, since +/// the second hop's bearer is what actually goes on the wire). +#[derive(Deserialize)] +struct ImdsTokenResponse { + access_token: String, +} + +/// AAD `oauth2/v2.0/token` response. `expires_in` is seconds-from-now for +/// successful client_credentials grants. +#[derive(Deserialize)] +struct AadTokenResponse { + access_token: String, + expires_in: u64, +} + +/// Hop 1: IMDS → MI assertion JWT (audience=api://AzureADTokenExchange). +async fn fetch_mi_assertion(http: &reqwest::Client) -> Result { + let resp = http + .get(IMDS_TOKEN_URL) + .header("Metadata", "true") + .query(&[ + ("api-version", IMDS_API_VERSION), + ("resource", MI_AUDIENCE), + ]) + .send() + .await + .and_then(|r| r.error_for_status()) + .map_err(|e| { + format!( + "IMDS token request failed (this code path requires running on an Azure \ + VM/host with a system- or user-assigned Managed Identity attached): {e}" + ) + })?; + let body: ImdsTokenResponse = resp + .json() + .await + .map_err(|e| format!("IMDS response was not valid JSON: {e}"))?; + Ok(body.access_token) +} + +/// Hop 2: MI assertion → customer-tenant storage bearer. +async fn exchange_for_storage_bearer( + http: &reqwest::Client, + tenant_id: &str, + client_id: &str, + mi_assertion: &str, +) -> Result<(String, SystemTime), String> { + let url = format!("{}/{}/oauth2/v2.0/token", AAD_AUTHORITY, tenant_id); + let form = [ + ("client_id", client_id), + ("scope", STORAGE_SCOPE), + ("grant_type", "client_credentials"), + ("client_assertion_type", CLIENT_ASSERTION_TYPE), + ("client_assertion", mi_assertion), + ]; + let resp = http + .post(&url) + .form(&form) + .send() + .await + .and_then(|r| r.error_for_status()) + .map_err(|e| { + format!( + "AAD token exchange failed for tenant={tenant_id}, client_id={client_id} \ + (the customer's App Registration likely has no Federated Identity \ + Credential trusting our MI, or the audience/issuer/subject in the FIC \ + don't match): {e}" + ) + })?; + let body: AadTokenResponse = resp + .json() + .await + .map_err(|e| format!("AAD token response was not valid JSON: {e}"))?; + // expires_in is seconds-from-now; convert to absolute SystemTime. + let expires_at = SystemTime::now() + Duration::from_secs(body.expires_in); + Ok((body.access_token, expires_at)) +} + +/// One full IMDS → AAD round-trip, no caching. Useful for tests and one-shot +/// callers; long-lived flows should use [`CrossTenantBearerCache`]. +pub async fn fetch_cross_tenant_bearer( + tenant_id: &str, + client_id: &str, +) -> Result<(String, SystemTime), String> { + let http = build_http_client(); + let assertion = fetch_mi_assertion(&http).await?; + exchange_for_storage_bearer(&http, tenant_id, client_id, &assertion).await +} + +#[derive(Clone)] +struct CachedBearer { + bearer: Arc, + /// Epoch milliseconds when the cached bearer expires. + expires_at_ms: u64, +} + +/// Cached customer-tenant bearer for a single (tenant_id, client_id) target. +/// +/// `current()` returns the cached bearer when fresh, or fetches+caches a +/// new one otherwise. Refresh is triggered when the current bearer is +/// within `refresh_offset` of expiry. A double-checked write-lock pattern +/// keeps concurrent `current()` calls from stampeding the AAD endpoint +/// (mirrors the GCP / AWS providers). +pub struct CrossTenantBearerCache { + tenant_id: String, + client_id: String, + refresh_offset: Duration, + http: reqwest::Client, + cache: Arc>>, +} + +impl std::fmt::Debug for CrossTenantBearerCache { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("CrossTenantBearerCache") + .field("tenant_id", &self.tenant_id) + .field("client_id", &self.client_id) + .field("refresh_offset", &self.refresh_offset) + .finish_non_exhaustive() + } +} + +impl CrossTenantBearerCache { + pub fn new(tenant_id: String, client_id: String, refresh_offset: Duration) -> Self { + Self { + tenant_id, + client_id, + refresh_offset, + http: build_http_client(), + cache: Arc::new(RwLock::new(None)), + } + } + + fn now_ms() -> u64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or(Duration::from_secs(0)) + .as_millis() as u64 + } + + fn needs_refresh(&self, cached: &Option) -> bool { + match cached { + None => true, + Some(c) => Self::now_ms() + self.refresh_offset.as_millis() as u64 >= c.expires_at_ms, + } + } + + /// Fast path with read lock; on miss, escalate to write lock and refresh. + /// Returns `Ok(None)` when the write lock is contended so the outer + /// `current` can back off briefly and retry — this matches the pattern + /// used by GCP / AWS providers. + async fn try_get(&self) -> Result>, String> { + { + let cached = self.cache.read().await; + if !self.needs_refresh(&cached) { + if let Some(c) = &*cached { + return Ok(Some(c.bearer.clone())); + } + } + } + let Ok(mut cache) = self.cache.try_write() else { + return Ok(None); + }; + // Double-check after acquiring write lock — another task may have + // just refreshed. + if !self.needs_refresh(&cache) { + if let Some(c) = &*cache { + return Ok(Some(c.bearer.clone())); + } + } + let assertion = fetch_mi_assertion(&self.http).await?; + let (bearer, expires_at) = + exchange_for_storage_bearer(&self.http, &self.tenant_id, &self.client_id, &assertion) + .await?; + let expires_at_ms = expires_at + .duration_since(UNIX_EPOCH) + .map_err(|e| format!("expires_at before UNIX epoch: {e}"))? + .as_millis() as u64; + let entry = CachedBearer { + bearer: Arc::new(bearer), + expires_at_ms, + }; + *cache = Some(entry.clone()); + Ok(Some(entry.bearer)) + } + + /// Return the current customer-tenant bearer, refreshing on demand. + pub async fn current(&self) -> Result, String> { + loop { + if let Some(b) = self.try_get().await? { + return Ok(b); + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + } +} + +/// Convenience: wrap a `String` error as `object_store::Error::Generic` for +/// callers in the Lance bridge that need an `object_store::Result`. +pub fn into_object_store_err(msg: String) -> object_store::Error { + object_store::Error::Generic { + store: STORE_NAME, + source: msg.into(), + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn needs_refresh_when_empty() { + let cache = CrossTenantBearerCache::new( + "tenant".into(), + "client".into(), + Duration::from_secs(REFRESH_OFFSET_SECS), + ); + assert!(cache.needs_refresh(&None)); + } + + #[test] + fn needs_refresh_within_offset() { + let cache = CrossTenantBearerCache::new( + "tenant".into(), + "client".into(), + Duration::from_secs(REFRESH_OFFSET_SECS), + ); + // Token expires in 100s, refresh offset is 300s → must refresh. + let expires_soon = CrossTenantBearerCache::now_ms() + 100_000; + let cached = Some(CachedBearer { + bearer: Arc::new("x".into()), + expires_at_ms: expires_soon, + }); + assert!(cache.needs_refresh(&cached)); + } + + #[test] + fn no_refresh_when_fresh() { + let cache = CrossTenantBearerCache::new( + "tenant".into(), + "client".into(), + Duration::from_secs(REFRESH_OFFSET_SECS), + ); + // Token expires in 1h, refresh offset is 300s → fresh. + let expires_far = CrossTenantBearerCache::now_ms() + 3_600_000; + let cached = Some(CachedBearer { + bearer: Arc::new("x".into()), + expires_at_ms: expires_far, + }); + assert!(!cache.needs_refresh(&cached)); + } +} diff --git a/cpp/src/format/bridge/rust/src/iceberg_bridgeimpl.rs b/cpp/src/format/bridge/rust/src/iceberg_bridgeimpl.rs index ac473f553..5bc48ba7a 100644 --- a/cpp/src/format/bridge/rust/src/iceberg_bridgeimpl.rs +++ b/cpp/src/format/bridge/rust/src/iceberg_bridgeimpl.rs @@ -26,6 +26,7 @@ use iceberg::table::StaticTable; use iceberg_storage_opendal::OpenDalStorageFactory; use crate::aliyun_oss_provider::AliyunOssStorageFactory; +use crate::azure_adls_provider::{AzdlsCrossTenantStorageFactory, CROSS_TENANT_MARKER_KEY}; use crate::gcp_impersonation::{DEFAULT_TOKEN_LIFETIME_SECS, fetch_impersonated_bearer}; use crate::iceberg_ffi::IcebergFileInfo; @@ -48,11 +49,19 @@ pub(crate) fn vec_to_hashmap(keys: Vec, values: Vec) -> HashMap< /// Intercepts `oss://` so per-tenant `oss.role-arn` can reach opendal — /// upstream `OpenDalStorageFactory::Oss` only carries endpoint/AK/SK. -/// Every other scheme is a pure pass-through to upstream. -fn storage_factory_for_scheme(scheme: &str) -> anyhow::Result> { +/// Intercepts `abfs[s]://` when cross-tenant Managed-Identity props are set +/// — upstream `AzdlsConfig` has no bearer/MI-cross-tenant path. +/// Every other scheme/case is a pure pass-through to upstream. +fn storage_factory_for_scheme( + scheme: &str, + props: &HashMap, +) -> anyhow::Result> { if scheme == "oss" { return Ok(Arc::new(AliyunOssStorageFactory::default())); } + if matches!(scheme, "abfs" | "abfss") && props.contains_key(CROSS_TENANT_MARKER_KEY) { + return Ok(Arc::new(AzdlsCrossTenantStorageFactory::default())); + } upstream_opendal_factory(scheme) } @@ -91,7 +100,7 @@ pub(crate) fn build_file_io( scheme: &str, props: &HashMap, ) -> anyhow::Result { - let factory = storage_factory_for_scheme(scheme)?; + let factory = storage_factory_for_scheme(scheme, props)?; let mut builder = FileIOBuilder::new(factory); for (k, v) in props { builder = builder.with_prop(k, v); diff --git a/cpp/src/format/bridge/rust/src/lance_bridgeimpl.rs b/cpp/src/format/bridge/rust/src/lance_bridgeimpl.rs index 2b960a128..e4b35e8c7 100644 --- a/cpp/src/format/bridge/rust/src/lance_bridgeimpl.rs +++ b/cpp/src/format/bridge/rust/src/lance_bridgeimpl.rs @@ -39,8 +39,12 @@ use lance_table::utils::stream::ReadBatchFutStream; use lance::io::ObjectStoreParams; use lance::session::Session; -use lance_io::object_store::{ObjectStoreRegistry, StorageOptionsProvider}; +use lance_io::object_store::ObjectStoreRegistry; +use crate::azure_cross_tenant_provider::CrossTenantAzureStoreProvider; +use crate::azure_federation::{ + CrossTenantBearerCache, REFRESH_OFFSET_SECS as AZURE_REFRESH_OFFSET_SECS, +}; use crate::gcp_impersonation::{ImpersonatingGcsStoreProvider, REFRESH_OFFSET_SECS}; #[derive(Clone)] @@ -483,17 +487,78 @@ fn build_gcp_impersonation_session(config: &GcpImpersonationConfig) -> Arc) -> Result> { + let Some(client_id) = storage_options.remove("azure_cross_tenant_client_id") else { + return Ok(None); + }; + let Some(tenant_id) = storage_options.remove("azure_cross_tenant_tenant_id") else { + return Ok(None); + }; + if client_id.is_empty() || tenant_id.is_empty() { + return Ok(None); + } + let refresh_offset_secs: u64 = storage_options + .remove("azure_cross_tenant_refresh_secs") + .and_then(|s| s.parse::().ok()) + .filter(|n| *n > 0) + .unwrap_or(AZURE_REFRESH_OFFSET_SECS) + .clamp(60, 1800); + Ok(Some(Self { + client_id, + tenant_id, + refresh_offset_secs, + })) + } +} + +/// Build a `Session` whose `ObjectStoreRegistry` overrides the `az` scheme +/// with a `CrossTenantAzureStoreProvider`. Stock `az` keeps working for +/// non-cross-tenant opens (which won't reach this code path). +fn build_azure_cross_tenant_session(config: &AzureCrossTenantConfig) -> Arc { + let cache = Arc::new(CrossTenantBearerCache::new( + config.tenant_id.clone(), + config.client_id.clone(), + std::time::Duration::from_secs(config.refresh_offset_secs), + )); + let registry = ObjectStoreRegistry::default(); + registry.insert("az", Arc::new(CrossTenantAzureStoreProvider::new(cache))); + Arc::new(Session::new(0, 0, Arc::new(registry))) +} + /// Pick a per-call Session if any cross-tenant credential feature is active. -/// The two supported features are mutually exclusive at the URI level (a URI -/// is either `gs://` or `oss://`), so at most one override is installed per -/// call. Returns `None` when no override is needed, so lance falls back to -/// its default session. +/// The supported features are mutually exclusive at the URI level (a URI is +/// either `gs://`, `az://`, or `oss://`), so at most one override is +/// installed per call. Returns `None` when no override is needed, so lance +/// falls back to its default session. fn pick_custom_session( storage_options: &mut HashMap, ) -> Result>> { if let Some(cfg) = GcpImpersonationConfig::extract(storage_options)? { return Ok(Some(build_gcp_impersonation_session(&cfg))); } + if let Some(cfg) = AzureCrossTenantConfig::extract(storage_options)? { + return Ok(Some(build_azure_cross_tenant_session(&cfg))); + } if storage_options.contains_key("oss_role_arn") { return Ok(Some(crate::aliyun_oss_provider::build_aliyun_oss_session())); } diff --git a/cpp/src/format/bridge/rust/src/lib.rs b/cpp/src/format/bridge/rust/src/lib.rs index 00e7fb4b5..e0e145ccc 100644 --- a/cpp/src/format/bridge/rust/src/lib.rs +++ b/cpp/src/format/bridge/rust/src/lib.rs @@ -13,6 +13,9 @@ // limitations under the License. mod aliyun_oss_provider; +mod azure_adls_provider; +mod azure_cross_tenant_provider; +mod azure_federation; mod gcp_impersonation; mod lance_bridgeimpl; mod vortex_bridgeimpl; diff --git a/cpp/src/format/iceberg/iceberg_common.cpp b/cpp/src/format/iceberg/iceberg_common.cpp index 82b83cbd3..3799b778f 100644 --- a/cpp/src/format/iceberg/iceberg_common.cpp +++ b/cpp/src/format/iceberg/iceberg_common.cpp @@ -70,11 +70,12 @@ std::unordered_map ToStorageOptions(const ArrowFileSys const auto& provider = config.cloud_provider; LOG_STORAGE_DEBUG_ << fmt::format( "provider={}, endpoint={}, use_ssl={}, use_iam={}, has_aksk={}, role_arn={}, external_id_set={}, " - "gcp_target_sa={}", + "gcp_target_sa={}, azure_cross_tenant={}", provider, config.address, config.use_ssl, config.use_iam, !config.access_key_id.empty() && !config.access_key_value.empty(), config.role_arn.empty() ? "(empty)" : config.role_arn, !config.external_id.empty(), - config.gcp_target_service_account.empty() ? "(empty)" : config.gcp_target_service_account); + config.gcp_target_service_account.empty() ? "(empty)" : config.gcp_target_service_account, + (!config.azure_client_id.empty() && !config.azure_tenant_id.empty()) ? "yes" : "no"); if (provider == kCloudProviderAWS) { if (!config.role_arn.empty()) { // AssumeRole: set ARN fields + region/endpoint; do NOT set AKSK so opendal @@ -99,7 +100,21 @@ std::unordered_map ToStorageOptions(const ArrowFileSys // Pass the endpoint suffix so the Rust bridge can reconstruct the full // Azure DFS endpoint (account.dfs.suffix) from scheme://container/path URIs. set("adls.endpoint-suffix", config.address); - if (config.use_iam) { + if (!config.azure_client_id.empty() && !config.azure_tenant_id.empty()) { + // Cross-tenant via Managed Identity. Bridge-private keys consumed by + // AzdlsCrossTenantStorageFactory in iceberg_bridgeimpl.rs. opendal 0.55 + // AzdlsConfig has no bearer field and reqsign's IMDS path requests an + // `https://storage.azure.com/` audience in *our* tenant (wrong audience + // for cross-tenant), so we route through a custom Storage that does + // the IMDS → AAD exchange and injects Authorization: Bearer ... at + // request time via opendal's HttpFetch hook. + set("adls.cross-tenant-client-id", config.azure_client_id); + set("adls.cross-tenant-tenant-id", config.azure_tenant_id); + if (config.load_frequency > 0) { + options["adls.cross-tenant-refresh-secs"] = std::to_string(config.load_frequency); + } + // Do NOT set adls.account-key / adls.client-secret on this branch. + } else if (config.use_iam) { auto* client_id = std::getenv("AZURE_CLIENT_ID"); if (client_id) set("adls.client-id", client_id); diff --git a/cpp/src/format/lance/lance_common.cpp b/cpp/src/format/lance/lance_common.cpp index 09e466ed2..4b8b111aa 100644 --- a/cpp/src/format/lance/lance_common.cpp +++ b/cpp/src/format/lance/lance_common.cpp @@ -45,11 +45,12 @@ StorageOptions ToStorageOptions(const ArrowFileSystemConfig& config) { const auto& provider = config.cloud_provider; LOG_STORAGE_DEBUG_ << fmt::format( "provider={}, endpoint={}, use_ssl={}, use_iam={}, has_aksk={}, role_arn={}, external_id_set={}, " - "gcp_target_sa={}", + "gcp_target_sa={}, azure_cross_tenant={}", provider, config.address, config.use_ssl, config.use_iam, !config.access_key_id.empty() && !config.access_key_value.empty(), config.role_arn.empty() ? "(empty)" : config.role_arn, !config.external_id.empty(), - config.gcp_target_service_account.empty() ? "(empty)" : config.gcp_target_service_account); + config.gcp_target_service_account.empty() ? "(empty)" : config.gcp_target_service_account, + (!config.azure_client_id.empty() && !config.azure_tenant_id.empty()) ? "yes" : "no"); if (provider == kCloudProviderAWS) { if (!config.role_arn.empty()) { // AssumeRole: set region/endpoint + ARN fields; do NOT set AKSK so the @@ -74,7 +75,24 @@ StorageOptions ToStorageOptions(const ArrowFileSystemConfig& config) { } } else if (provider == kCloudProviderAzure) { set("azure_storage_account_name", config.access_key_id); - if (!config.use_iam) { + if (!config.azure_client_id.empty() && !config.azure_tenant_id.empty()) { + // Cross-tenant via Managed Identity. Bridge-private keys consumed by + // CrossTenantAzureStoreProvider (see lance_bridgeimpl.rs::pick_custom_session). + // object_store has no native MI-cross-tenant path: WorkloadIdentityOAuth + // wants a federated_token_file we don't have on plain VMs, and the + // built-in IMDS provider asks for `https://storage.azure.com/` audience + // in *our* tenant — wrong tenant for the customer's storage account. + // Hence a custom CredentialProvider doing the IMDS → AAD two-hop + // exchange and returning AzureCredential::BearerToken. + set("azure_cross_tenant_client_id", config.azure_client_id); + set("azure_cross_tenant_tenant_id", config.azure_tenant_id); + if (config.load_frequency > 0) { + // Refresh-ahead interval. AAD-issued bearer is fixed at ~1h; we don't + // request a lifetime, only schedule when to refresh. + options["azure_cross_tenant_refresh_secs"] = std::to_string(config.load_frequency); + } + // Do NOT set azure_storage_account_key on this branch. + } else if (!config.use_iam) { set("azure_storage_account_key", config.access_key_value); } if (!config.address.empty()) { diff --git a/cpp/src/properties.cpp b/cpp/src/properties.cpp index 7cfde2eff..f341b9dc6 100644 --- a/cpp/src/properties.cpp +++ b/cpp/src/properties.cpp @@ -416,6 +416,19 @@ static std::unordered_map property_infos = { "The target GCP service account email for cross-project impersonation.", "", std::nullopt), + REGISTER_PROPERTY(PROPERTY_FS_AZURE_CLIENT_ID, + PropertyType::STRING, + "The customer's App Registration client_id for Azure cross-tenant access. " + "When set together with fs.azure_tenant_id, our process exchanges its " + "Managed Identity for a customer-tenant bearer via OAuth2 federated credentials.", + "", + std::nullopt), + REGISTER_PROPERTY(PROPERTY_FS_AZURE_TENANT_ID, + PropertyType::STRING, + "The customer's Entra ID tenant_id for Azure cross-tenant access. " + "Used as the authority in the OAuth2 token exchange. See PROPERTY_FS_AZURE_CLIENT_ID.", + "", + std::nullopt), // --- writer properties define --- REGISTER_PROPERTY(PROPERTY_WRITER_POLICY, PropertyType::STRING, diff --git a/cpp/test/format/external_table_arn_test.cpp b/cpp/test/format/external_table_arn_test.cpp index a3e73e1c8..4ecc1fe5b 100644 --- a/cpp/test/format/external_table_arn_test.cpp +++ b/cpp/test/format/external_table_arn_test.cpp @@ -1654,4 +1654,290 @@ TEST_F(ExternalTableAliyunOIDCArnTest, ReadTwoParquetFilesWithOIDCChain) { loon_properties_free(&loon_props); } -} // namespace milvus_storage +// =========================================================================== +// Integration tests for reading external Azure tables via cross-tenant +// Managed Identity. +// +// These tests verify that the storage layer can use {azure_client_id, +// azure_tenant_id} — pointing at a customer's App Registration with a +// Federated Identity Credential trusting our MI — to access data in a +// customer's Azure storage account, without us holding any of the customer's +// secrets. +// +// The flow exercised end-to-end: +// * Lance — bridge `CrossTenantAzureStoreProvider` is registered against +// `az` in a per-call Session; its `CredentialProvider` returns +// `AzureCredential::BearerToken(...)` minted via the IMDS → AAD two-hop +// in `azure_federation::CrossTenantBearerCache`. +// * Iceberg — bridge `AzdlsCrossTenantStorageFactory` is selected for +// `abfs[s]://` schemes when the bridge-private cross-tenant keys are +// present; uses opendal `Azdls` with a placeholder `account_key` and an +// `HttpFetch` wrapper that rewrites Authorization with the bearer. +// +// Test data is written using Azure account_key (same-tenant write, an +// admin-owned bypass for fixture seeding only — production reads never carry +// account_key on the cross-tenant path), then read back using only +// {azure_client_id, azure_tenant_id}. +// +// Required environment variables (all must be set; test is skipped otherwise): +// =========================================================================== +#define AZURE_CT_ENV_ADDRESS "AZURE_CT_TEST_ENV_ADDRESS" // Endpoint suffix (e.g. "core.windows.net") +#define AZURE_CT_ENV_ACCOUNT "AZURE_CT_TEST_ENV_ACCOUNT_NAME" // Customer storage account +#define AZURE_CT_ENV_CONTAINER "AZURE_CT_TEST_ENV_CONTAINER" // Customer container/filesystem +#define AZURE_CT_ENV_ACCOUNT_KEY "AZURE_CT_TEST_ENV_ACCOUNT_KEY" // Customer account key (write only) +#define AZURE_CT_ENV_CLIENT_ID "AZURE_CT_TEST_ENV_CLIENT_ID" // Customer App Registration client_id +#define AZURE_CT_ENV_TENANT_ID "AZURE_CT_TEST_ENV_TENANT_ID" // Customer Entra ID tenant_id + +struct AzureCtWriteResult { + api::ColumnGroupFile cgfile; + std::shared_ptr schema; // nullptr for Iceberg + uint64_t num_rows; + std::string explore_dir; // Full URI with address for loon_exttable_explore + int64_t iceberg_snapshot_id; // Only used for iceberg +}; + +class ExternalTableAzureCrossTenantTest : public ::testing::TestWithParam { + protected: + void SetUp() override { + // Our-side bucket (IAM-based, for writing manifest) + our_address_ = GetEnvVar(OUR_ENV_ADDRESS).ValueOr(""); + our_bucket_ = GetEnvVar(OUR_ENV_BUCKET).ValueOr(""); + our_region_ = GetEnvVar(OUR_ENV_REGION).ValueOr(""); + our_cloud_provider_ = GetEnvVar(OUR_ENV_CLOUD_PROVIDER).ValueOr(""); + + // Customer-side Azure storage + address_ = GetEnvVar(AZURE_CT_ENV_ADDRESS).ValueOr(""); + account_name_ = GetEnvVar(AZURE_CT_ENV_ACCOUNT).ValueOr(""); + container_ = GetEnvVar(AZURE_CT_ENV_CONTAINER).ValueOr(""); + account_key_ = GetEnvVar(AZURE_CT_ENV_ACCOUNT_KEY).ValueOr(""); + client_id_ = GetEnvVar(AZURE_CT_ENV_CLIENT_ID).ValueOr(""); + tenant_id_ = GetEnvVar(AZURE_CT_ENV_TENANT_ID).ValueOr(""); + + if (our_address_.empty() || our_bucket_.empty() || our_cloud_provider_.empty() || address_.empty() || + account_name_.empty() || container_.empty() || account_key_.empty() || client_id_.empty() || + tenant_id_.empty()) { + GTEST_SKIP() << "Azure cross-tenant tests require all env vars: " << OUR_ENV_ADDRESS << ", " << OUR_ENV_BUCKET + << ", " << OUR_ENV_REGION << ", " << OUR_ENV_CLOUD_PROVIDER << ", " << AZURE_CT_ENV_ADDRESS << ", " + << AZURE_CT_ENV_ACCOUNT << ", " << AZURE_CT_ENV_CONTAINER << ", " << AZURE_CT_ENV_ACCOUNT_KEY << ", " + << AZURE_CT_ENV_CLIENT_ID << ", " << AZURE_CT_ENV_TENANT_ID; + } + + // --- Write properties: Azure account_key (same-tenant, fixture-seeding) --- + // PROPERTY_FS_ACCESS_KEY_ID maps to account_name on Azure (account-level), + // PROPERTY_FS_ACCESS_KEY_VALUE to account_key. PROPERTY_FS_BUCKET_NAME is + // the container/filesystem name. + api::SetValue(write_props_, PROPERTY_FS_STORAGE_TYPE, "remote"); + api::SetValue(write_props_, PROPERTY_FS_CLOUD_PROVIDER, "azure"); + api::SetValue(write_props_, PROPERTY_FS_ADDRESS, address_.c_str()); + api::SetValue(write_props_, PROPERTY_FS_BUCKET_NAME, container_.c_str()); + api::SetValue(write_props_, PROPERTY_FS_ACCESS_KEY_ID, account_name_.c_str()); + api::SetValue(write_props_, PROPERTY_FS_ACCESS_KEY_VALUE, account_key_.c_str()); + api::SetValue(write_props_, PROPERTY_FS_USE_SSL, "true"); + + // --- Read properties: extfs.azct.* with cross-tenant client_id + tenant_id --- + // No account_key on this branch — the bridge's cross-tenant providers + // mint a customer-tenant Bearer via the IMDS → AAD two-hop. + api::SetValue(read_props_, "extfs.azct.storage_type", "remote"); + api::SetValue(read_props_, "extfs.azct.cloud_provider", "azure"); + api::SetValue(read_props_, "extfs.azct.address", address_.c_str()); + api::SetValue(read_props_, "extfs.azct.bucket_name", container_.c_str()); + api::SetValue(read_props_, "extfs.azct.access_key_id", account_name_.c_str()); + api::SetValue(read_props_, "extfs.azct.use_ssl", "true"); + api::SetValue(read_props_, "extfs.azct.azure_client_id", client_id_.c_str()); + api::SetValue(read_props_, "extfs.azct.azure_tenant_id", tenant_id_.c_str()); + + FilesystemCache::getInstance().clean(); + + auto ts = std::chrono::steady_clock::now().time_since_epoch().count(); + test_base_ = "zc/azct-test-" + std::to_string(ts); + } + + // No TearDown cleanup: we intentionally don't call GetFileSystem() with + // cloud_provider=azure here — the AzureFileSystemProducer's + // ConfigureAccountKeyCredential path leaves global Azure SDK state that + // can collide with subsequent tests. Test data is left in the customer + // bucket; rely on container lifecycle rules. + void TearDown() override { FilesystemCache::getInstance().clean(); } + + arrow::Result CreateTestTable(const std::string& format, uint64_t num_rows) { + if (format == LOON_FORMAT_LANCE_TABLE) { + return CreateLanceTable(num_rows); + } else if (format == LOON_FORMAT_ICEBERG_TABLE) { + return CreateIcebergTable(num_rows); + } + return arrow::Status::Invalid("Unknown format: " + format); + } + + // Our-side + std::string our_address_; + std::string our_bucket_; + std::string our_region_; + std::string our_cloud_provider_; + // Customer-side + std::string address_; + std::string account_name_; + std::string container_; + std::string account_key_; + std::string client_id_; + std::string tenant_id_; + + api::Properties write_props_; + api::Properties read_props_; + std::string test_base_; + + private: + arrow::Result CreateLanceTable(uint64_t num_rows) { + ARROW_ASSIGN_OR_RAISE(auto schema, CreateTestSchema({true, true, true, false})); + ARROW_ASSIGN_OR_RAISE(auto batch, CreateTestData(schema, 0, false, num_rows, 4, 50, {true, true, true, false})); + auto path = test_base_ + "/lance"; + lance::LanceTableWriter writer(path, schema, write_props_); + ARROW_RETURN_NOT_OK(writer.Write(batch)); + ARROW_ASSIGN_OR_RAISE(auto cgfile, writer.Close()); + std::cout << "[Azure CT Test] Lance cgfile: " << cgfile.ToString() << std::endl; + // explore_dir mirrors `BuildLanceBaseUri` shape but with address inserted + // for extfs matching: az://
// + auto explore_dir = "az://" + address_ + "/" + container_ + "/" + path; + return AzureCtWriteResult{std::move(cgfile), schema, num_rows, explore_dir, 0}; + } + + arrow::Result CreateIcebergTable(uint64_t num_rows) { + auto path = test_base_ + "/iceberg"; + // iceberg_create_test_table writes via opendal Azdls (account-key path + // works same-tenant). Use abfss URI in container@account.dfs.suffix form + // — this is the canonical fully-qualified shape that opendal accepts + // without needing extra config to reconstruct the endpoint. + auto table_uri = "abfss://" + container_ + "@" + account_name_ + ".dfs." + address_ + "/" + path; + + ArrowFileSystemConfig write_config; + ARROW_RETURN_NOT_OK(ArrowFileSystemConfig::create_file_system_config(write_props_, write_config)); + auto storage_options = iceberg::ToStorageOptions(write_config); + + auto table_info = iceberg::CreateTestTable(table_uri, num_rows, false, {}, storage_options); + auto file_infos = iceberg::PlanFiles(table_info.metadata_location, table_info.snapshot_id, storage_options); + if (file_infos.empty()) { + return arrow::Status::Invalid("PlanFiles returned no files"); + } + + auto milvus_path = iceberg::ToMilvusUri(file_infos[0].data_file_path, address_); + api::ColumnGroupFile cg_file{milvus_path, 0, static_cast(file_infos[0].record_count), {}}; + std::cout << "[Azure CT Test] Iceberg cgfile: " << cg_file.ToString() << std::endl; + auto explore_dir = iceberg::ToMilvusUri(table_info.metadata_location, address_); + return AzureCtWriteResult{std::move(cg_file), nullptr, num_rows, explore_dir, table_info.snapshot_id}; + } +}; + +TEST_P(ExternalTableAzureCrossTenantTest, ReadWithCrossTenantMI) { + const auto& format = GetParam(); + const uint64_t num_rows = 100; + + // Step 1: Write test data with account_key (same-tenant, fixture seeding) + ASSERT_AND_ASSIGN(auto result, CreateTestTable(format, num_rows)); + + std::cout << "[Azure CT Test] Format: " << format << std::endl; + std::cout << "[Azure CT Test] Written to: " << result.cgfile.path << std::endl; + std::cout << "[Azure CT Test] Explore dir: " << result.explore_dir << std::endl; + std::cout << "[Azure CT Test] Customer client_id: " << client_id_ << std::endl; + std::cout << "[Azure CT Test] Customer tenant_id: " << tenant_id_ << std::endl; + + // Step 2: Build properties for loon_exttable_explore + // - fs.*: our-side bucket with IAM for manifest storage + // - extfs.azct.*: customer-side cross-tenant for external data access + auto manifest_base = test_base_ + "/manifest"; + + std::vector> props = { + {PROPERTY_FS_STORAGE_TYPE, "remote"}, + {PROPERTY_FS_CLOUD_PROVIDER, our_cloud_provider_}, + {PROPERTY_FS_ADDRESS, our_address_}, + {PROPERTY_FS_BUCKET_NAME, our_bucket_}, + {PROPERTY_FS_REGION, our_region_}, + {PROPERTY_FS_USE_SSL, "true"}, + {PROPERTY_FS_USE_IAM, "true"}, + // extfs.azct: cross-tenant Azure for external data access + {"extfs.azct.storage_type", "remote"}, + {"extfs.azct.cloud_provider", "azure"}, + {"extfs.azct.address", address_}, + {"extfs.azct.bucket_name", container_}, + {"extfs.azct.access_key_id", account_name_}, + {"extfs.azct.use_ssl", "true"}, + {"extfs.azct.azure_client_id", client_id_}, + {"extfs.azct.azure_tenant_id", tenant_id_}, + }; + if (format == LOON_FORMAT_ICEBERG_TABLE) { + props.emplace_back(PROPERTY_ICEBERG_SNAPSHOT_ID, std::to_string(result.iceberg_snapshot_id)); + } + + std::vector c_keys, c_values; + c_keys.reserve(props.size()); + c_values.reserve(props.size()); + for (const auto& [k, v] : props) { + c_keys.push_back(k.c_str()); + c_values.push_back(v.c_str()); + } + + LoonProperties loon_props = {}; + auto rc = loon_properties_create(c_keys.data(), c_values.data(), c_keys.size(), &loon_props); + ASSERT_TRUE(loon_ffi_is_success(&rc)) << loon_ffi_get_errmsg(&rc); + + // Step 3: loon_exttable_explore — discovers files via cross-tenant MI bearer + const char* columns_arr[] = {"id", "name", "value"}; + uint64_t out_num_files = 0; + char* out_manifest_path = nullptr; + + rc = loon_exttable_explore(columns_arr, 3, format.c_str(), manifest_base.c_str(), result.explore_dir.c_str(), + &loon_props, &out_num_files, &out_manifest_path); + ASSERT_TRUE(loon_ffi_is_success(&rc)) << loon_ffi_get_errmsg(&rc); + ASSERT_GT(out_num_files, 0u); + ASSERT_NE(out_manifest_path, nullptr); + + std::cout << "[Azure CT Test] loon_exttable_explore: found " << out_num_files + << " files, manifest=" << out_manifest_path << std::endl; + + // Step 4: Read manifest via FFI to get ColumnGroupFiles + LoonManifest* out_manifest = nullptr; + rc = loon_exttable_read_manifest(out_manifest_path, &loon_props, &out_manifest); + ASSERT_TRUE(loon_ffi_is_success(&rc)) << loon_ffi_get_errmsg(&rc); + ASSERT_NE(out_manifest, nullptr); + ASSERT_EQ(out_manifest->column_groups.num_of_column_groups, 1u); + + auto* cg = &out_manifest->column_groups.column_group_array[0]; + ASSERT_EQ(cg->num_of_files, out_num_files); + + std::cout << "[Azure CT Test] manifest has " << cg->num_of_files << " files" << std::endl; + + // Step 5: Read data using FormatReader with cross-tenant MI + std::vector columns = {"id", "name", "value"}; + + int64_t total_rows = 0; + for (uint64_t f = 0; f < cg->num_of_files; ++f) { + auto& loon_file = cg->files[f]; + api::ColumnGroupFile cgfile; + cgfile.path = loon_file.path; + cgfile.start_index = loon_file.start_index; + cgfile.end_index = loon_file.end_index; + if (loon_file.property_keys != nullptr) { + for (uint32_t p = 0; p < loon_file.num_properties; ++p) { + cgfile.properties[loon_file.property_keys[p]] = loon_file.property_values[p]; + } + } + + ASSERT_AND_ASSIGN(auto reader, FormatReader::create(result.schema, format, cgfile, read_props_, columns, nullptr)); + ASSERT_AND_ASSIGN(auto rg_infos, reader->get_row_group_infos()); + + for (size_t i = 0; i < rg_infos.size(); ++i) { + ASSERT_AND_ASSIGN(auto batch, reader->get_chunk(i)); + total_rows += batch->num_rows(); + } + } + ASSERT_EQ(total_rows, static_cast(num_rows)); + std::cout << "[Azure CT Test] FormatReader read " << total_rows << " rows via cross-tenant MI OK" << std::endl; + + loon_manifest_destroy(out_manifest); + free(out_manifest_path); + loon_properties_free(&loon_props); +} + +INSTANTIATE_TEST_SUITE_P(AzureCrossTenantFormats, + ExternalTableAzureCrossTenantTest, + ::testing::Values(LOON_FORMAT_LANCE_TABLE, LOON_FORMAT_ICEBERG_TABLE)); + +} // namespace milvus_storage \ No newline at end of file