1#[cfg(feature = "credential-vendor-aws")]
63pub mod aws;
64
65#[cfg(feature = "credential-vendor-azure")]
66pub mod azure;
67
68#[cfg(feature = "credential-vendor-gcp")]
69pub mod gcp;
70
71#[cfg(any(
74 feature = "credential-vendor-aws",
75 feature = "credential-vendor-azure",
76 feature = "credential-vendor-gcp"
77))]
78pub mod cache;
79
80use std::collections::HashMap;
81use std::str::FromStr;
82
83use async_trait::async_trait;
84use lance_core::Result;
85use lance_io::object_store::uri_to_url;
86use lance_namespace::models::Identity;
87
88pub const DEFAULT_CREDENTIAL_DURATION_MILLIS: u64 = 3600 * 1000;
90
91pub fn redact_credential(credential: &str) -> String {
104 const SHOW_START: usize = 8;
105 const SHOW_END: usize = 4;
106 const MIN_LENGTH_FOR_BOTH_ENDS: usize = SHOW_START + SHOW_END + 4; if credential.is_empty() {
109 return "[empty]".to_string();
110 }
111
112 if credential.len() < MIN_LENGTH_FOR_BOTH_ENDS {
113 let show = credential.len().min(SHOW_START);
115 format!("{}***", &credential[..show])
116 } else {
117 format!(
119 "{}***{}",
120 &credential[..SHOW_START],
121 &credential[credential.len() - SHOW_END..]
122 )
123 }
124}
125
126#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
139pub enum VendedPermission {
140 #[default]
142 Read,
143 Write,
149 Admin,
151}
152
153impl VendedPermission {
154 pub fn can_write(&self) -> bool {
156 matches!(self, Self::Write | Self::Admin)
157 }
158
159 pub fn can_delete(&self) -> bool {
161 matches!(self, Self::Admin)
162 }
163}
164
165impl FromStr for VendedPermission {
166 type Err = String;
167
168 fn from_str(s: &str) -> std::result::Result<Self, Self::Err> {
169 match s.to_lowercase().as_str() {
170 "read" => Ok(Self::Read),
171 "write" => Ok(Self::Write),
172 "admin" => Ok(Self::Admin),
173 _ => Err(format!(
174 "Invalid permission '{}'. Must be one of: read, write, admin",
175 s
176 )),
177 }
178 }
179}
180
181impl std::fmt::Display for VendedPermission {
182 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
183 match self {
184 Self::Read => write!(f, "read"),
185 Self::Write => write!(f, "write"),
186 Self::Admin => write!(f, "admin"),
187 }
188 }
189}
190
191pub const PROPERTY_PREFIX: &str = "credential_vendor.";
194
195pub const ENABLED: &str = "enabled";
197
198pub const PERMISSION: &str = "permission";
200
201pub const CACHE_ENABLED: &str = "cache_enabled";
204
205pub const API_KEY_SALT: &str = "api_key_salt";
208
209pub const API_KEY_HASH_PREFIX: &str = "api_key_hash.";
212
213#[cfg(feature = "credential-vendor-aws")]
215pub mod aws_props {
216 pub const ROLE_ARN: &str = "aws_role_arn";
217 pub const EXTERNAL_ID: &str = "aws_external_id";
218 pub const REGION: &str = "aws_region";
219 pub const ROLE_SESSION_NAME: &str = "aws_role_session_name";
220 pub const DURATION_MILLIS: &str = "aws_duration_millis";
223}
224
225#[cfg(feature = "credential-vendor-gcp")]
227pub mod gcp_props {
228 pub const SERVICE_ACCOUNT: &str = "gcp_service_account";
229
230 pub const WORKLOAD_IDENTITY_PROVIDER: &str = "gcp_workload_identity_provider";
233
234 pub const IMPERSONATION_SERVICE_ACCOUNT: &str = "gcp_impersonation_service_account";
237}
238
239#[cfg(feature = "credential-vendor-azure")]
241pub mod azure_props {
242 pub const TENANT_ID: &str = "azure_tenant_id";
243 pub const ACCOUNT_NAME: &str = "azure_account_name";
245 pub const DURATION_MILLIS: &str = "azure_duration_millis";
248
249 pub const FEDERATED_CLIENT_ID: &str = "azure_federated_client_id";
252}
253
254#[derive(Clone)]
256pub struct VendedCredentials {
257 pub storage_options: HashMap<String, String>,
262
263 pub expires_at_millis: u64,
265}
266
267impl std::fmt::Debug for VendedCredentials {
268 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
269 f.debug_struct("VendedCredentials")
270 .field(
271 "storage_options",
272 &format!("[{} keys redacted]", self.storage_options.len()),
273 )
274 .field("expires_at_millis", &self.expires_at_millis)
275 .finish()
276 }
277}
278
279impl VendedCredentials {
280 pub fn new(storage_options: HashMap<String, String>, expires_at_millis: u64) -> Self {
282 Self {
283 storage_options,
284 expires_at_millis,
285 }
286 }
287
288 pub fn is_expired(&self) -> bool {
290 let now_millis = std::time::SystemTime::now()
291 .duration_since(std::time::UNIX_EPOCH)
292 .expect("time went backwards")
293 .as_millis() as u64;
294 now_millis >= self.expires_at_millis
295 }
296}
297
298#[async_trait]
304pub trait CredentialVendor: Send + Sync + std::fmt::Debug {
305 async fn vend_credentials(
328 &self,
329 table_location: &str,
330 identity: Option<&Identity>,
331 ) -> Result<VendedCredentials>;
332
333 fn provider_name(&self) -> &'static str;
335
336 fn permission(&self) -> VendedPermission;
338}
339
340pub fn detect_provider_from_uri(uri: &str) -> &'static str {
349 let Ok(url) = uri_to_url(uri) else {
350 return "unknown";
351 };
352
353 match url.scheme() {
354 "s3" => "aws",
355 "gs" => "gcp",
356 "az" => "azure",
357 _ => "unknown",
358 }
359}
360
361pub fn has_credential_vendor_config(properties: &HashMap<String, String>) -> bool {
366 properties
367 .get(ENABLED)
368 .map(|v| v.eq_ignore_ascii_case("true"))
369 .unwrap_or(false)
370}
371
372#[allow(unused_variables)]
394pub async fn create_credential_vendor_for_location(
395 table_location: &str,
396 properties: &HashMap<String, String>,
397) -> Result<Option<Box<dyn CredentialVendor>>> {
398 let provider = detect_provider_from_uri(table_location);
399
400 let vendor: Option<Box<dyn CredentialVendor>> = match provider {
401 #[cfg(feature = "credential-vendor-aws")]
402 "aws" => create_aws_vendor(properties).await?,
403
404 #[cfg(feature = "credential-vendor-gcp")]
405 "gcp" => create_gcp_vendor(properties).await?,
406
407 #[cfg(feature = "credential-vendor-azure")]
408 "azure" => create_azure_vendor(properties)?,
409
410 _ => None,
411 };
412
413 #[cfg(any(
415 feature = "credential-vendor-aws",
416 feature = "credential-vendor-azure",
417 feature = "credential-vendor-gcp"
418 ))]
419 if let Some(v) = vendor {
420 let cache_enabled = properties
421 .get(CACHE_ENABLED)
422 .map(|s| !s.eq_ignore_ascii_case("false"))
423 .unwrap_or(true);
424
425 if cache_enabled {
426 return Ok(Some(Box::new(cache::CachingCredentialVendor::new(v))));
427 } else {
428 return Ok(Some(v));
429 }
430 }
431
432 #[cfg(not(any(
433 feature = "credential-vendor-aws",
434 feature = "credential-vendor-azure",
435 feature = "credential-vendor-gcp"
436 )))]
437 let _ = vendor;
438
439 Ok(None)
440}
441
442#[allow(dead_code)]
444fn parse_permission(properties: &HashMap<String, String>) -> VendedPermission {
445 properties
446 .get(PERMISSION)
447 .and_then(|s| s.parse().ok())
448 .unwrap_or_default()
449}
450
451#[allow(dead_code)]
453fn parse_duration_millis(properties: &HashMap<String, String>, key: &str) -> u64 {
454 properties
455 .get(key)
456 .and_then(|s| s.parse::<u64>().ok())
457 .unwrap_or(DEFAULT_CREDENTIAL_DURATION_MILLIS)
458}
459
460#[cfg(feature = "credential-vendor-aws")]
461async fn create_aws_vendor(
462 properties: &HashMap<String, String>,
463) -> Result<Option<Box<dyn CredentialVendor>>> {
464 use aws::{AwsCredentialVendor, AwsCredentialVendorConfig};
465 use lance_core::Error;
466
467 let role_arn = properties
469 .get(aws_props::ROLE_ARN)
470 .ok_or_else(|| Error::InvalidInput {
471 source: "AWS credential vending requires 'credential_vendor.aws_role_arn' to be set"
472 .into(),
473 location: snafu::location!(),
474 })?;
475
476 let duration_millis = parse_duration_millis(properties, aws_props::DURATION_MILLIS);
477
478 let permission = parse_permission(properties);
479
480 let mut config = AwsCredentialVendorConfig::new(role_arn)
481 .with_duration_millis(duration_millis)
482 .with_permission(permission);
483
484 if let Some(external_id) = properties.get(aws_props::EXTERNAL_ID) {
485 config = config.with_external_id(external_id);
486 }
487 if let Some(region) = properties.get(aws_props::REGION) {
488 config = config.with_region(region);
489 }
490 if let Some(session_name) = properties.get(aws_props::ROLE_SESSION_NAME) {
491 config = config.with_role_session_name(session_name);
492 }
493
494 let vendor = AwsCredentialVendor::new(config).await?;
495 Ok(Some(Box::new(vendor)))
496}
497
498#[cfg(feature = "credential-vendor-gcp")]
499async fn create_gcp_vendor(
500 properties: &HashMap<String, String>,
501) -> Result<Option<Box<dyn CredentialVendor>>> {
502 use gcp::{GcpCredentialVendor, GcpCredentialVendorConfig};
503
504 let permission = parse_permission(properties);
505
506 let mut config = GcpCredentialVendorConfig::new().with_permission(permission);
507
508 if let Some(sa) = properties.get(gcp_props::SERVICE_ACCOUNT) {
509 config = config.with_service_account(sa);
510 }
511
512 let vendor = GcpCredentialVendor::new(config).await?;
513 Ok(Some(Box::new(vendor)))
514}
515
516#[cfg(feature = "credential-vendor-azure")]
517fn create_azure_vendor(
518 properties: &HashMap<String, String>,
519) -> Result<Option<Box<dyn CredentialVendor>>> {
520 use azure::{AzureCredentialVendor, AzureCredentialVendorConfig};
521 use lance_core::Error;
522
523 let account_name =
525 properties
526 .get(azure_props::ACCOUNT_NAME)
527 .ok_or_else(|| {
528 Error::InvalidInput {
529 source:
530 "Azure credential vending requires 'credential_vendor.azure_account_name' to be set"
531 .into(),
532 location: snafu::location!(),
533 }
534 })?;
535
536 let duration_millis = parse_duration_millis(properties, azure_props::DURATION_MILLIS);
537 let permission = parse_permission(properties);
538
539 let mut config = AzureCredentialVendorConfig::new()
540 .with_account_name(account_name)
541 .with_duration_millis(duration_millis)
542 .with_permission(permission);
543
544 if let Some(tenant_id) = properties.get(azure_props::TENANT_ID) {
545 config = config.with_tenant_id(tenant_id);
546 }
547
548 let vendor = AzureCredentialVendor::new(config);
549 Ok(Some(Box::new(vendor)))
550}
551
552#[cfg(test)]
553mod tests {
554 use super::*;
555
556 #[test]
557 fn test_detect_provider_from_uri() {
558 assert_eq!(detect_provider_from_uri("s3://bucket/path"), "aws");
560 assert_eq!(detect_provider_from_uri("S3://bucket/path"), "aws");
561
562 assert_eq!(detect_provider_from_uri("gs://bucket/path"), "gcp");
564 assert_eq!(detect_provider_from_uri("GS://bucket/path"), "gcp");
565
566 assert_eq!(detect_provider_from_uri("az://container/path"), "azure");
568
569 assert_eq!(detect_provider_from_uri("/local/path"), "unknown");
571 assert_eq!(detect_provider_from_uri("file:///local/path"), "unknown");
572 assert_eq!(detect_provider_from_uri("memory://test"), "unknown");
573 assert_eq!(detect_provider_from_uri("s3a://bucket/path"), "unknown");
575 assert_eq!(
576 detect_provider_from_uri("abfss://container@account.dfs.core.windows.net/path"),
577 "unknown"
578 );
579 assert_eq!(
580 detect_provider_from_uri("wasbs://container@account.blob.core.windows.net/path"),
581 "unknown"
582 );
583 }
584
585 #[test]
586 fn test_vended_permission_from_str() {
587 assert_eq!(
589 "read".parse::<VendedPermission>().unwrap(),
590 VendedPermission::Read
591 );
592 assert_eq!(
593 "READ".parse::<VendedPermission>().unwrap(),
594 VendedPermission::Read
595 );
596 assert_eq!(
597 "write".parse::<VendedPermission>().unwrap(),
598 VendedPermission::Write
599 );
600 assert_eq!(
601 "WRITE".parse::<VendedPermission>().unwrap(),
602 VendedPermission::Write
603 );
604 assert_eq!(
605 "admin".parse::<VendedPermission>().unwrap(),
606 VendedPermission::Admin
607 );
608 assert_eq!(
609 "Admin".parse::<VendedPermission>().unwrap(),
610 VendedPermission::Admin
611 );
612
613 let err = "invalid".parse::<VendedPermission>().unwrap_err();
615 assert!(err.contains("Invalid permission"));
616 assert!(err.contains("invalid"));
617
618 let err = "".parse::<VendedPermission>().unwrap_err();
619 assert!(err.contains("Invalid permission"));
620
621 let err = "readwrite".parse::<VendedPermission>().unwrap_err();
622 assert!(err.contains("Invalid permission"));
623 }
624
625 #[test]
626 fn test_vended_permission_display() {
627 assert_eq!(VendedPermission::Read.to_string(), "read");
628 assert_eq!(VendedPermission::Write.to_string(), "write");
629 assert_eq!(VendedPermission::Admin.to_string(), "admin");
630 }
631
632 #[test]
633 fn test_parse_permission_with_invalid_values() {
634 let mut props = HashMap::new();
636 props.insert(PERMISSION.to_string(), "invalid".to_string());
637 assert_eq!(parse_permission(&props), VendedPermission::Read);
638
639 props.insert(PERMISSION.to_string(), "".to_string());
641 assert_eq!(parse_permission(&props), VendedPermission::Read);
642
643 let empty_props: HashMap<String, String> = HashMap::new();
645 assert_eq!(parse_permission(&empty_props), VendedPermission::Read);
646 }
647
648 #[test]
649 fn test_parse_duration_millis_with_invalid_values() {
650 const TEST_KEY: &str = "test_duration_millis";
651
652 let mut props = HashMap::new();
654 props.insert(TEST_KEY.to_string(), "not_a_number".to_string());
655 assert_eq!(
656 parse_duration_millis(&props, TEST_KEY),
657 DEFAULT_CREDENTIAL_DURATION_MILLIS
658 );
659
660 props.insert(TEST_KEY.to_string(), "-1000".to_string());
662 assert_eq!(
663 parse_duration_millis(&props, TEST_KEY),
664 DEFAULT_CREDENTIAL_DURATION_MILLIS
665 );
666
667 props.insert(TEST_KEY.to_string(), "".to_string());
669 assert_eq!(
670 parse_duration_millis(&props, TEST_KEY),
671 DEFAULT_CREDENTIAL_DURATION_MILLIS
672 );
673
674 let empty_props: HashMap<String, String> = HashMap::new();
676 assert_eq!(
677 parse_duration_millis(&empty_props, TEST_KEY),
678 DEFAULT_CREDENTIAL_DURATION_MILLIS
679 );
680
681 props.insert(TEST_KEY.to_string(), "7200000".to_string());
683 assert_eq!(parse_duration_millis(&props, TEST_KEY), 7200000);
684 }
685
686 #[test]
687 fn test_has_credential_vendor_config() {
688 let mut props = HashMap::new();
690 props.insert(ENABLED.to_string(), "true".to_string());
691 assert!(has_credential_vendor_config(&props));
692
693 props.insert(ENABLED.to_string(), "TRUE".to_string());
695 assert!(has_credential_vendor_config(&props));
696
697 props.insert(ENABLED.to_string(), "false".to_string());
699 assert!(!has_credential_vendor_config(&props));
700
701 props.insert(ENABLED.to_string(), "yes".to_string());
703 assert!(!has_credential_vendor_config(&props));
704
705 let empty_props: HashMap<String, String> = HashMap::new();
707 assert!(!has_credential_vendor_config(&empty_props));
708 }
709
710 #[test]
711 fn test_vended_credentials_debug_redacts_secrets() {
712 let mut storage_options = HashMap::new();
713 storage_options.insert(
714 "aws_access_key_id".to_string(),
715 "AKIAIOSFODNN7EXAMPLE".to_string(),
716 );
717 storage_options.insert(
718 "aws_secret_access_key".to_string(),
719 "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY".to_string(),
720 );
721 storage_options.insert(
722 "aws_session_token".to_string(),
723 "FwoGZXIvYXdzE...".to_string(),
724 );
725
726 let creds = VendedCredentials::new(storage_options, 1234567890);
727 let debug_output = format!("{:?}", creds);
728
729 assert!(!debug_output.contains("AKIAIOSFODNN7EXAMPLE"));
731 assert!(!debug_output.contains("wJalrXUtnFEMI"));
732 assert!(!debug_output.contains("FwoGZXIvYXdzE"));
733
734 assert!(debug_output.contains("redacted"));
736 assert!(debug_output.contains("3 keys"));
737
738 assert!(debug_output.contains("1234567890"));
740 }
741
742 #[test]
743 fn test_vended_credentials_is_expired() {
744 let past_millis = std::time::SystemTime::now()
746 .duration_since(std::time::UNIX_EPOCH)
747 .unwrap()
748 .as_millis() as u64
749 - 1000; let expired_creds = VendedCredentials::new(HashMap::new(), past_millis);
752 assert!(expired_creds.is_expired());
753
754 let future_millis = std::time::SystemTime::now()
756 .duration_since(std::time::UNIX_EPOCH)
757 .unwrap()
758 .as_millis() as u64
759 + 3600000; let valid_creds = VendedCredentials::new(HashMap::new(), future_millis);
762 assert!(!valid_creds.is_expired());
763 }
764
765 #[test]
766 fn test_redact_credential() {
767 assert_eq!(redact_credential("AKIAIOSFODNN7EXAMPLE"), "AKIAIOSF***MPLE");
769
770 assert_eq!(redact_credential("1234567890123456"), "12345678***3456");
772
773 assert_eq!(redact_credential("short1234567"), "short123***");
775 assert_eq!(redact_credential("short123"), "short123***");
776 assert_eq!(redact_credential("tiny"), "tiny***");
777 assert_eq!(redact_credential("ab"), "ab***");
778 assert_eq!(redact_credential("a"), "a***");
779
780 assert_eq!(redact_credential(""), "[empty]");
782
783 assert_eq!(redact_credential("AKIAIOSFODNN7EXAMPLE"), "AKIAIOSF***MPLE");
786
787 let long_token = "ya29.a0AfH6SMBx1234567890abcdefghijklmnopqrstuvwxyz";
789 assert_eq!(redact_credential(long_token), "ya29.a0A***wxyz");
790
791 let sas_token = "sv=2021-06-08&ss=b&srt=sco&sp=rwdlacuiytfx&se=2024-12-31";
793 assert_eq!(redact_credential(sas_token), "sv=2021-***2-31");
794 }
795}