1
- // Geneva Config Client with TLS (PKCS#12) and TODO: Managed Identity support
1
+ // Geneva Config Client with TLS (PKCS#12) and Azure Workload Identity support TODO: Azure Arc support
2
2
3
3
use base64:: { engine:: general_purpose, Engine as _} ;
4
4
use reqwest:: {
@@ -18,15 +18,20 @@ use std::fs;
18
18
use std:: path:: PathBuf ;
19
19
use std:: sync:: RwLock ;
20
20
21
- // Azure Identity imports for MSI authentication
21
+ // Azure Identity imports for MSI and Workload Identity authentication
22
22
use azure_core:: credentials:: TokenCredential ;
23
- use azure_identity:: { ManagedIdentityCredential , ManagedIdentityCredentialOptions , UserAssignedId } ;
23
+ use azure_identity:: {
24
+ ManagedIdentityCredential , ManagedIdentityCredentialOptions , UserAssignedId ,
25
+ WorkloadIdentityCredential ,
26
+ } ;
24
27
25
28
/// Authentication methods for the Geneva Config Client.
26
29
///
27
- /// The client supports two authentication methods:
28
- /// - Certificate-based authentication using PKCS#12 (.p12) files
29
- /// - Managed Identity (Azure) - planned for future implementation
30
+ /// The client supports the following authentication methods:
31
+ /// - Certificate-based authentication (mTLS) using PKCS#12 (.p12) files
32
+ /// - Azure Managed Identity (System-assigned or User-assigned)
33
+ /// - Azure Workload Identity (Federated Identity for Kubernetes)
34
+ /// - Mock authentication for testing (feature-gated)
30
35
///
31
36
/// # Certificate Format
32
37
/// Certificates should be in PKCS#12 (.p12) format for client TLS authentication.
@@ -65,6 +70,19 @@ pub enum AuthMethod {
65
70
UserManagedIdentityByObjectId { object_id : String } ,
66
71
/// User-assigned managed identity by resource ID
67
72
UserManagedIdentityByResourceId { resource_id : String } ,
73
+ /// Azure Workload Identity authentication (Federated Identity for Kubernetes)
74
+ ///
75
+ /// The following environment variables must be set in the pod spec:
76
+ /// * `AZURE_CLIENT_ID` - Azure AD Application (client) ID (set explicitly in pod env)
77
+ /// * `AZURE_TENANT_ID` - Azure AD Tenant ID (set explicitly in pod env)
78
+ /// * `AZURE_FEDERATED_TOKEN_FILE` - Path to service account token file (auto-injected by workload identity webhook)
79
+ ///
80
+ /// These variables are automatically read by the Azure Identity SDK at runtime.
81
+ ///
82
+ /// # Arguments
83
+ /// * `resource` - Azure AD resource URI for token acquisition
84
+ /// (e.g., <https://monitor.azure.com> for Azure Public Cloud)
85
+ WorkloadIdentity { resource : String } ,
68
86
#[ cfg( feature = "mock_auth" ) ]
69
87
MockAuth , // No authentication, used for testing purposes
70
88
}
@@ -78,6 +96,8 @@ pub(crate) enum GenevaConfigClientError {
78
96
JwtTokenError ( String ) ,
79
97
#[ error( "Certificate error: {0}" ) ]
80
98
Certificate ( String ) ,
99
+ #[ error( "Workload Identity authentication error: {0}" ) ]
100
+ WorkloadIdentityAuth ( String ) ,
81
101
#[ error( "MSI authentication error: {0}" ) ]
82
102
MsiAuth ( String ) ,
83
103
@@ -257,6 +277,10 @@ impl GenevaConfigClient {
257
277
. map_err ( |e| GenevaConfigClientError :: Certificate ( e. to_string ( ) ) ) ?;
258
278
client_builder = client_builder. use_preconfigured_tls ( tls_connector) ;
259
279
}
280
+ AuthMethod :: WorkloadIdentity { .. } => {
281
+ // No special HTTP client configuration needed for Workload Identity
282
+ // Authentication is done via Bearer token in request headers
283
+ }
260
284
AuthMethod :: SystemManagedIdentity
261
285
| AuthMethod :: UserManagedIdentity { .. }
262
286
| AuthMethod :: UserManagedIdentityByObjectId { .. }
@@ -276,13 +300,14 @@ impl GenevaConfigClient {
276
300
let version_str = format ! ( "Ver{0}v0" , config. config_major_version) ;
277
301
278
302
// Use different API endpoints based on authentication method
279
- // Certificate auth uses "api", MSI auth uses "userapi"
303
+ // Certificate auth uses "api", MSI auth and Workload Identity use "userapi"
280
304
let api_path = match & config. auth_method {
281
305
AuthMethod :: Certificate { .. } => "api" ,
282
306
AuthMethod :: SystemManagedIdentity
283
307
| AuthMethod :: UserManagedIdentity { .. }
284
308
| AuthMethod :: UserManagedIdentityByObjectId { .. }
285
- | AuthMethod :: UserManagedIdentityByResourceId { .. } => "userapi" ,
309
+ | AuthMethod :: UserManagedIdentityByResourceId { .. }
310
+ | AuthMethod :: WorkloadIdentity { .. } => "userapi" ,
286
311
#[ cfg( feature = "mock_auth" ) ]
287
312
AuthMethod :: MockAuth => "api" , // treat mock like certificate path for URL shape
288
313
} ;
@@ -329,6 +354,55 @@ impl GenevaConfigClient {
329
354
headers
330
355
}
331
356
357
+ /// Get Azure AD token using Workload Identity (Federated Identity)
358
+ ///
359
+ /// Reads AZURE_CLIENT_ID, AZURE_TENANT_ID, and AZURE_FEDERATED_TOKEN_FILE from environment variables.
360
+ /// In Kubernetes:
361
+ /// - AZURE_CLIENT_ID and AZURE_TENANT_ID must be set explicitly in the pod spec
362
+ /// - AZURE_FEDERATED_TOKEN_FILE is auto-injected by the workload identity webhook
363
+ async fn get_workload_identity_token ( & self ) -> Result < String > {
364
+ let resource =
365
+ match & self . config . auth_method {
366
+ AuthMethod :: WorkloadIdentity { resource } => resource,
367
+ _ => return Err ( GenevaConfigClientError :: WorkloadIdentityAuth (
368
+ "get_workload_identity_token called but auth method is not WorkloadIdentity"
369
+ . to_string ( ) ,
370
+ ) ) ,
371
+ } ;
372
+
373
+ // TODO: Extract scope generation logic into helper function shared with get_msi_token()
374
+ let base = resource. trim_end_matches ( "/.default" ) . trim_end_matches ( '/' ) ;
375
+ let mut scope_candidates: Vec < String > = vec ! [ format!( "{base}/.default" ) , base. to_string( ) ] ;
376
+ // TODO - below check is not required, as we alread trim "/"
377
+ if !base. ends_with ( '/' ) {
378
+ scope_candidates. push ( format ! ( "{base}/" ) ) ;
379
+ }
380
+
381
+ // TODO: Consider caching WorkloadIdentityCredential if profiling shows credential creation overhead
382
+ // Pass None to let azure_identity crate read AZURE_CLIENT_ID, AZURE_TENANT_ID,
383
+ // and AZURE_FEDERATED_TOKEN_FILE from environment variables automatically
384
+ let credential = WorkloadIdentityCredential :: new ( None ) . map_err ( |e| {
385
+ GenevaConfigClientError :: WorkloadIdentityAuth ( format ! (
386
+ "Failed to create WorkloadIdentityCredential. Ensure AZURE_CLIENT_ID, AZURE_TENANT_ID, and AZURE_FEDERATED_TOKEN_FILE environment variables are set: {e}"
387
+ ) )
388
+ } ) ?;
389
+
390
+ let mut last_err: Option < String > = None ;
391
+ for scope in & scope_candidates {
392
+ //TODO - It looks like the get_token API accepts a slice of &str
393
+ match credential. get_token ( & [ scope. as_str ( ) ] , None ) . await {
394
+ Ok ( token) => return Ok ( token. token . secret ( ) . to_string ( ) ) ,
395
+ Err ( e) => last_err = Some ( e. to_string ( ) ) ,
396
+ }
397
+ }
398
+
399
+ let detail = last_err. unwrap_or_else ( || "no error detail" . into ( ) ) ;
400
+ Err ( GenevaConfigClientError :: WorkloadIdentityAuth ( format ! (
401
+ "Workload Identity token acquisition failed. Scopes tried: {scopes}. Last error: {detail}" ,
402
+ scopes = scope_candidates. join( ", " )
403
+ ) ) )
404
+ }
405
+
332
406
/// Get MSI token for GCS authentication
333
407
async fn get_msi_token ( & self ) -> Result < String > {
334
408
let resource = self . config . msi_resource . as_ref ( ) . ok_or_else ( || {
@@ -337,17 +411,14 @@ impl GenevaConfigClient {
337
411
)
338
412
} ) ?;
339
413
340
- // Normalize resource (strip trailing "/.default" if provided by user )
414
+ // TODO: Extract scope generation logic into helper function shared with get_workload_identity_token( )
341
415
let base = resource. trim_end_matches ( "/.default" ) . trim_end_matches ( '/' ) ;
342
-
343
- // Candidate scopes tried with Azure Identity
344
416
let mut scope_candidates: Vec < String > = vec ! [ format!( "{base}/.default" ) , base. to_string( ) ] ;
345
- // Add variant with trailing slash if not already present
417
+ // TODO - below check is not required, as we alread trim "/"
346
418
if !base. ends_with ( '/' ) {
347
419
scope_candidates. push ( format ! ( "{base}/" ) ) ;
348
420
}
349
421
350
- // Build credential based on selector
351
422
let user_assigned_id = match & self . config . auth_method {
352
423
AuthMethod :: SystemManagedIdentity => None ,
353
424
AuthMethod :: UserManagedIdentity { client_id } => {
@@ -367,6 +438,7 @@ impl GenevaConfigClient {
367
438
}
368
439
} ;
369
440
441
+ // TODO: Consider caching ManagedIdentityCredential if profiling shows credential creation overhead
370
442
let options = ManagedIdentityCredentialOptions {
371
443
user_assigned_id,
372
444
..Default :: default ( )
@@ -382,6 +454,7 @@ impl GenevaConfigClient {
382
454
Err ( e) => last_err = Some ( e. to_string ( ) ) ,
383
455
}
384
456
}
457
+
385
458
let detail = last_err. unwrap_or_else ( || "no error detail" . into ( ) ) ;
386
459
Err ( GenevaConfigClientError :: MsiAuth ( format ! (
387
460
"Managed Identity token acquisition failed. Scopes tried: {scopes}. Last error: {detail}. IMDS fallback intentionally disabled." ,
@@ -506,8 +579,8 @@ impl GenevaConfigClient {
506
579
507
580
/// Internal method that actually fetches data from Geneva Config Service
508
581
async fn fetch_ingestion_info ( & self ) -> Result < ( IngestionGatewayInfo , MonikerInfo ) > {
509
- let tag_id = Uuid :: new_v4 ( ) . to_string ( ) ; //TODO - uuid is costly, check if counter is enough?
510
- let mut url = String :: with_capacity ( self . precomputed_url_prefix . len ( ) + 50 ) ; // Pre-allocate with reasonable capacity
582
+ let tag_id = Uuid :: new_v4 ( ) . to_string ( ) ; // TODO: consider cheaper counter if perf-critical
583
+ let mut url = String :: with_capacity ( self . precomputed_url_prefix . len ( ) + 50 ) ;
511
584
write ! ( & mut url, "{}&TagId={tag_id}" , self . precomputed_url_prefix) . map_err ( |e| {
512
585
GenevaConfigClientError :: InternalError ( format ! ( "Failed to write URL: {e}" ) )
513
586
} ) ?;
@@ -518,48 +591,44 @@ impl GenevaConfigClient {
518
591
519
592
request = request. header ( "x-ms-client-request-id" , req_id) ;
520
593
521
- // Add MSI authentication for managed identity auth method
594
+ // Add appropriate authentication header
522
595
match & self . config . auth_method {
596
+ AuthMethod :: WorkloadIdentity { .. } => {
597
+ let token = self . get_workload_identity_token ( ) . await ?;
598
+ request = request. header ( AUTHORIZATION , format ! ( "Bearer {}" , token) ) ;
599
+ }
523
600
AuthMethod :: SystemManagedIdentity
524
601
| AuthMethod :: UserManagedIdentity { .. }
525
602
| AuthMethod :: UserManagedIdentityByObjectId { .. }
526
603
| AuthMethod :: UserManagedIdentityByResourceId { .. } => {
527
- let msi_token = self . get_msi_token ( ) . await ?;
528
- request = request. header ( AUTHORIZATION , format ! ( "Bearer {}" , msi_token ) ) ;
604
+ let token = self . get_msi_token ( ) . await ?;
605
+ request = request. header ( AUTHORIZATION , format ! ( "Bearer {}" , token ) ) ;
529
606
}
530
607
AuthMethod :: Certificate { .. } => { /* mTLS only */ }
531
608
#[ cfg( feature = "mock_auth" ) ]
532
609
AuthMethod :: MockAuth => { /* no auth header */ }
533
610
}
534
611
535
- // Log the request details for debugging
612
+ // Send HTTP request
536
613
let response = match request. send ( ) . await {
537
- Ok ( response) => response,
538
- Err ( e) => {
539
- return Err ( GenevaConfigClientError :: Http ( e) ) ;
540
- }
614
+ Ok ( resp) => resp,
615
+ Err ( e) => return Err ( GenevaConfigClientError :: Http ( e) ) ,
541
616
} ;
542
617
543
- // Check if the response is successful
544
618
let status = response. status ( ) ;
545
619
let body = response. text ( ) . await ?;
620
+
546
621
if status. is_success ( ) {
547
- let parsed = match serde_json:: from_str :: < GenevaResponse > ( & body) {
548
- Ok ( response) => response,
549
- Err ( e) => {
550
- return Err ( GenevaConfigClientError :: AuthInfoNotFound ( format ! (
551
- "Failed to parse response: {e}"
552
- ) ) ) ;
553
- }
554
- } ;
622
+ let parsed = serde_json:: from_str :: < GenevaResponse > ( & body) . map_err ( |e| {
623
+ GenevaConfigClientError :: AuthInfoNotFound ( format ! ( "Failed to parse response: {e}" ) )
624
+ } ) ?;
555
625
556
626
for account in parsed. storage_account_keys {
557
627
if account. is_primary_moniker && account. account_moniker_name . contains ( "diag" ) {
558
628
let moniker_info = MonikerInfo {
559
629
name : account. account_moniker_name ,
560
630
account_group : account. account_group_name ,
561
631
} ;
562
-
563
632
return Ok ( ( parsed. ingestion_gateway_info , moniker_info) ) ;
564
633
}
565
634
}
@@ -610,16 +679,21 @@ fn extract_endpoint_from_token(token: &str) -> Result<String> {
610
679
_ => payload. to_string ( ) ,
611
680
} ;
612
681
613
- // Decode the Base64-encoded payload into raw bytes with a more tolerant approach.
682
+ // Decode the Base64-encoded payload into raw bytes.
683
+ // Try URL-safe (with and without padding), then fall back to standard Base64.
614
684
let decoded = match general_purpose:: URL_SAFE_NO_PAD . decode ( & payload) {
615
685
Ok ( b) => b,
616
- Err ( e_url ) => match general_purpose:: STANDARD . decode ( & payload) {
686
+ Err ( e_url_no_pad ) => match general_purpose:: URL_SAFE . decode ( & payload) {
617
687
Ok ( b) => b,
618
- Err ( e_std) => {
619
- return Err ( GenevaConfigClientError :: JwtTokenError ( format ! (
620
- "Failed to decode JWT (url_safe and standard): url_err={e_url}; std_err={e_std}"
621
- ) ) )
622
- }
688
+ Err ( e_url_pad) => match general_purpose:: STANDARD . decode ( & payload) {
689
+ Ok ( b) => b,
690
+ Err ( e_std) => {
691
+ return Err ( GenevaConfigClientError :: JwtTokenError ( format ! (
692
+ "Failed to decode JWT (URL_SAFE_NO_PAD, URL_SAFE, and STANDARD): \
693
+ no_pad_err={e_url_no_pad}; pad_err={e_url_pad}; std_err={e_std}"
694
+ ) ) )
695
+ }
696
+ } ,
623
697
} ,
624
698
} ;
625
699
0 commit comments