Also try to login to Azure via SDK provided methods (#5573)

## Problem We ideally use the Azure SDK's way of obtaining authorization, as pointed out in https://github.com/neondatabase/neon/pull/5546#discussion_r1360619178 . ## Summary of changes This PR adds support for Azure SDK based authentication, using [DefaultAzureCredential](https://docs.rs/azure_identity/0.16.1/azure_identity/struct.DefaultAzureCredential.html), which tries the following credentials: * [EnvironmentCredential](https://docs.rs/azure_identity/0.16.1/azure_identity/struct.EnvironmentCredential.html), reading from various env vars * [ImdsManagedIdentityCredential](https://docs.rs/azure_identity/0.16.1/azure_identity/struct.ImdsManagedIdentityCredential.html), using managed identity * [AzureCliCredential](https://docs.rs/azure_identity/0.16.1/azure_identity/struct.AzureCliCredential.html), using Azure CLI closes #5566.
2026-01-05 20:42:54 +00:00 · 2023-10-17 12:59:57 +02:00
parent 9256788273
commit 00c71bb93a
5 changed files with 81 additions and 9 deletions
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -853,6 +853,27 @@ dependencies = [
 "uuid",
 ]

+[[package]]
+name = "azure_identity"
+version = "0.16.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5b67b337346da8739e91ea1e9400a6ebc9bc54e0b2af1d23c9bcd565950588f9"
+dependencies = [
+ "async-lock",
+ "async-trait",
+ "azure_core",
+ "futures",
+ "log",
+ "oauth2",
+ "pin-project",
+ "serde",
+ "serde_json",
+ "time 0.3.21",
+ "tz-rs",
+ "url",
+ "uuid",
+]
+
 [[package]]
 name = "azure_storage"
 version = "0.16.0"
@@ -1131,9 +1152,11 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "4e3c5919066adf22df73762e50cffcde3a758f2a848b113b586d1f86728b673b"
 dependencies = [
 "iana-time-zone",
+ "js-sys",
 "num-integer",
 "num-traits",
 "serde",
+ "wasm-bindgen",
 "winapi",
 ]

@@ -2949,6 +2972,34 @@ dependencies = [
 "libc",
 ]

+[[package]]
+name = "num_threads"
+version = "0.1.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "2819ce041d2ee131036f4fc9d6ae7ae125a3a40e97ba64d04fe799ad9dabbb44"
+dependencies = [
+ "libc",
+]
+
+[[package]]
+name = "oauth2"
+version = "4.4.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c38841cdd844847e3e7c8d29cef9dcfed8877f8f56f9071f77843ecf3baf937f"
+dependencies = [
+ "base64 0.13.1",
+ "chrono",
+ "getrandom 0.2.9",
+ "http",
+ "rand 0.8.5",
+ "serde",
+ "serde_json",
+ "serde_path_to_error",
+ "sha2 0.10.6",
+ "thiserror",
+ "url",
+]
+
 [[package]]
 name = "object"
 version = "0.30.3"
@@ -4021,6 +4072,7 @@ dependencies = [
 "aws-smithy-http",
 "aws-types",
 "azure_core",
+ "azure_identity",
 "azure_storage",
 "azure_storage_blobs",
 "bytes",
@@ -5219,6 +5271,8 @@ checksum = "8f3403384eaacbca9923fa06940178ac13e4edb725486d70e8e15881d0c836cc"
 dependencies = [
 "itoa",
 "js-sys",
+ "libc",
+ "num_threads",
 "serde",
 "time-core",
 "time-macros 0.2.9",
@@ -5770,6 +5824,15 @@ version = "1.16.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "497961ef93d974e23eb6f433eb5fe1b7930b659f06d12dec6fc44a8f554c0bba"

+[[package]]
+name = "tz-rs"
+version = "0.6.14"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "33851b15c848fad2cf4b105c6bb66eb9512b6f6c44a4b13f57c53c73c707e2b4"
+dependencies = [
+ "const_fn",
+]
+
 [[package]]
 name = "uname"
 version = "0.1.1"
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -36,9 +36,10 @@ license = "Apache-2.0"
 [workspace.dependencies]
 anyhow = { version = "1.0", features = ["backtrace"] }
 async-compression = { version = "0.4.0", features = ["tokio", "gzip"] }
-azure_core = "0.16.0"
+azure_core = "0.16"
+azure_identity = "0.16"
 azure_storage = "0.16"
-azure_storage_blobs = "0.16.0"
+azure_storage_blobs = "0.16"
 flate2 = "1.0.26"
 async-stream = "0.3"
 async-trait = "0.1"
--- a/libs/remote_storage/Cargo.toml
+++ b/libs/remote_storage/Cargo.toml
@@ -28,6 +28,7 @@ utils.workspace = true
 pin-project-lite.workspace = true
 workspace_hack.workspace = true
 azure_core.workspace = true
+azure_identity.workspace = true
 azure_storage.workspace = true
 azure_storage_blobs.workspace = true
 futures-util.workspace = true
--- a/libs/remote_storage/src/azure_blob.rs
+++ b/libs/remote_storage/src/azure_blob.rs
@@ -1,12 +1,15 @@
 //! Azure Blob Storage wrapper

+use std::env;
 use std::num::NonZeroU32;
+use std::sync::Arc;
 use std::{borrow::Cow, collections::HashMap, io::Cursor};

 use super::REMOTE_STORAGE_PREFIX_SEPARATOR;
 use anyhow::Result;
 use azure_core::request_options::{MaxResults, Metadata, Range};
 use azure_core::Header;
+use azure_identity::DefaultAzureCredential;
 use azure_storage::StorageCredentials;
 use azure_storage_blobs::prelude::ClientBuilder;
 use azure_storage_blobs::{
@@ -38,12 +41,16 @@ impl AzureBlobStorage {
            azure_config.container_name
        );

-        let account =
-            std::env::var("AZURE_STORAGE_ACCOUNT").expect("missing AZURE_STORAGE_ACCOUNT");
-        let access_key =
-            std::env::var("AZURE_STORAGE_ACCESS_KEY").expect("missing AZURE_STORAGE_ACCESS_KEY");
+        let account = env::var("AZURE_STORAGE_ACCOUNT").expect("missing AZURE_STORAGE_ACCOUNT");

-        let credentials = StorageCredentials::access_key(account.clone(), access_key);
+        // If the `AZURE_STORAGE_ACCESS_KEY` env var has an access key, use that,
+        // otherwise try the token based credentials.
+        let credentials = if let Ok(access_key) = env::var("AZURE_STORAGE_ACCESS_KEY") {
+            StorageCredentials::access_key(account.clone(), access_key)
+        } else {
+            let token_credential = DefaultAzureCredential::default();
+            StorageCredentials::token_credential(Arc::new(token_credential))
+        };

        let builder = ClientBuilder::new(account, credentials);

--- a/workspace_hack/Cargo.toml
+++ b/workspace_hack/Cargo.toml
@@ -21,7 +21,7 @@ aws-smithy-http = { version = "0.56", default-features = false, features = ["eve
 axum = { version = "0.6", features = ["ws"] }
 base64 = { version = "0.21", features = ["alloc"] }
 bytes = { version = "1", features = ["serde"] }
-chrono = { version = "0.4", default-features = false, features = ["clock", "serde"] }
+chrono = { version = "0.4", default-features = false, features = ["clock", "serde", "wasmbind"] }
 clap = { version = "4", features = ["derive", "string"] }
 clap_builder = { version = "4", default-features = false, features = ["color", "help", "std", "string", "suggestions", "usage"] }
 crossbeam-utils = { version = "0.8" }
@@ -57,7 +57,7 @@ serde_json = { version = "1", features = ["raw_value"] }
 smallvec = { version = "1", default-features = false, features = ["write"] }
 socket2 = { version = "0.4", default-features = false, features = ["all"] }
 standback = { version = "0.2", default-features = false, features = ["std"] }
-time = { version = "0.3", features = ["macros", "serde-well-known"] }
+time = { version = "0.3", features = ["local-offset", "macros", "serde-well-known"] }
 tokio = { version = "1", features = ["fs", "io-std", "io-util", "macros", "net", "process", "rt-multi-thread", "signal", "test-util"] }
 tokio-rustls = { version = "0.24" }
 tokio-util = { version = "0.7", features = ["codec", "io"] }