Cosmos: RID-aware resource references and raw-path protocol (1/3)

sdk/cosmos/.cspell.json

@@ -16,6 +16,7 @@
     "backoff",
     "backpressure",
     "BADFUNC",
+    "balxjyz",
     "bindgen",
     "brazilsouth",
     "BUFFERTOOSMALL",
@@ -159,6 +160,8 @@
     "ptrs",
     "pushback",
     "pyroscope",
+    "QBAA",
+    "QBAOW",
     "qname",
     "qself",
     "RAII",

sdk/cosmos/azure_data_cosmos_driver/CHANGELOG.md

@@ -6,6 +6,8 @@
 
 ### Breaking Changes
 
+- Resource-reference accessors now return `Option` to account for RID-addressed references that have no name. `DatabaseReference::name_based_path`, `ContainerReference::database_name`, and `ContainerReference::name_based_path` return `None` when the reference is addressed by RID (previously they returned `&str`/`String` and assumed a name was always present). Use the new `ContainerReference::base_path` to obtain the addressing-appropriate path (RID-based or name-based) when building request URLs. ([#4640](https://github.com/Azure/azure-sdk-for-rust/pull/4640))
+
 ### Bugs Fixed
 
 ### Other Changes

sdk/cosmos/azure_data_cosmos_driver/src/driver/cache/container_cache.rs

@@ -28,7 +28,16 @@ impl ContainerNameKey {
     fn from_container(c: &ContainerReference) -> Self {
         Self {
             account_endpoint: c.account().endpoint().as_str().to_owned(),
-            db_name: c.database_name().to_owned(),
+            // TODO(Slice 2): `database_name()` is `None` for a RID-addressed
+            // container, so `unwrap_or_default()` collapses the key to
+            // `{endpoint, "", container_name}`. That is harmless today because
+            // nothing produces RID-addressed references yet, but once Slice 2's
+            // resolution layer starts caching them, two RID-resolved containers
+            // that share a name across different databases would alias to the
+            // same by-name key. Slice 2 must skip the by-name index for
+            // RID-addressed containers (or key it differently) before relying on
+            // this path.
+            db_name: c.database_name().unwrap_or_default().to_owned(),
             container_name: c.name().to_owned(),
         }
     }
@@ -378,8 +387,8 @@ mod tests {
             .await
             .unwrap();
 
-        assert_eq!(r1.database_name(), "db1");
-        assert_eq!(r2.database_name(), "db2");
+        assert_eq!(r1.database_name(), Some("db1"));
+        assert_eq!(r2.database_name(), Some("db2"));
     }
 
     // --- get returns none ---

sdk/cosmos/azure_data_cosmos_driver/src/driver/pipeline/operation_pipeline.rs

@@ -27,9 +27,9 @@ use crate::{
         transport::CosmosTransport,
     },
     models::{
-        cosmos_headers::QUERY_CONTENT_TYPE, request_header_names, AccountEndpoint, ActivityId,
-        CosmosOperation, CosmosResponse, Credential, DefaultConsistencyLevel,
-        EffectivePartitionKey, OperationType, SessionToken, SubStatusCode,
+        cosmos_headers::QUERY_CONTENT_TYPE, encode_path_segments, request_header_names,
+        AccountEndpoint, ActivityId, CosmosOperation, CosmosResponse, Credential,
+        DefaultConsistencyLevel, EffectivePartitionKey, OperationType, SessionToken, SubStatusCode,
     },
     options::{
         HedgeThreshold, OperationOptionsView, ReadConsistencyStrategy, Region,
@@ -1269,7 +1269,17 @@ fn build_transport_request(
         } else {
             format!("/{}", request_path)
         };
-        base.set_path(&normalized);
+        // Name-based paths are percent-encoded so the gateway reconstructs the
+        // same resource link we signed (names may contain reserved characters).
+        // RID-based paths must be sent raw: encoding the `=` padding of a base64
+        // RID makes the gateway treat the segment as a name and reject the
+        // RID-based signature. The authorization signature is derived from
+        // `paths` below (a lowercased RID for RID-addressed requests).
+        if paths.is_rid_based() {
+            base.set_path(&normalized);
+        } else {
+            base.set_path(&encode_path_segments(&normalized));
+        }
         base
     };
 
@@ -3643,6 +3653,65 @@ mod tests {
         assert_eq!(request.url.path(), "/dbs/mydb");
     }
 
+    /// Builds a transport request for `operation` with default routing/context
+    /// and returns the final `Url::path()` after `set_path` has reprocessed it.
+    /// Used to assert the raw-vs-percent-encoded seam in `build_transport_request`.
+    fn transport_request_path(operation: &CosmosOperation) -> String {
+        let routing = test_routing();
+        let activity_id = ActivityId::from_string("default-activity".to_string());
+        let ctx = TransportRequestContext {
+            routing: &routing,
+            activity_id: &activity_id,
+            execution_context: ExecutionContext::Initial,
+            deadline: None,
+            resolved_session_token: None,
+            throughput_control: None,
+        };
+        build_transport_request(operation, &OperationOverrides::default(), None, &ctx)
+            .expect("request should build")
+            .url
+            .path()
+            .to_owned()
+    }
+
+    #[test]
+    fn build_transport_request_rid_path_is_sent_raw() {
+        // A RID-addressed path must reach the gateway raw: the base64 `=` padding
+        // survives `set_path` un-encoded (not `%3D`), otherwise the gateway rejects
+        // the RID-based signature with a `401`.
+        let db = DatabaseReference::from_rid(test_account(), "qjQBAA==");
+        let operation = CosmosOperation::read_database(db);
+        assert_eq!(transport_request_path(&operation), "/dbs/qjQBAA==");
+    }
+
+    #[test]
+    fn build_transport_request_rid_path_with_plus_is_sent_raw() {
+        // RIDs use a base64 variant that can contain `+`; it must also survive raw
+        // (not `%2B`) on a RID-addressed path.
+        let db = DatabaseReference::from_rid(test_account(), "ab+cdEAA=");
+        let operation = CosmosOperation::read_database(db);
+        assert_eq!(transport_request_path(&operation), "/dbs/ab+cdEAA=");
+    }
+
+    #[test]
+    fn build_transport_request_name_path_is_percent_encoded() {
+        // Name-based paths are percent-encoded so the gateway reconstructs the
+        // same resource link we signed: a space becomes `%20` and `+` becomes
+        // `%2B`.
+        let db = DatabaseReference::from_name(test_account(), "my db+x");
+        let operation = CosmosOperation::read_database(db);
+        assert_eq!(transport_request_path(&operation), "/dbs/my%20db%2Bx");
+    }
+
+    #[test]
+    fn build_transport_request_offer_path_is_sent_raw() {
+        // Offers are signed over the lowercased RID, so the `/offers/{rid}` path is
+        // RID-addressed and must be sent raw — reserved base64 characters in the
+        // RID survive un-encoded.
+        let operation = CosmosOperation::read_offer(test_account(), "ab+QBAA==");
+        assert_eq!(transport_request_path(&operation), "/offers/ab+QBAA==");
+    }
+
     #[test]
     fn build_transport_request_uses_operation_activity_id_when_present() {
         let operation = CosmosOperation::read_all_databases(test_account())

sdk/cosmos/azure_data_cosmos_driver/src/driver/routing/session_container.rs

@@ -33,11 +33,15 @@ struct SessionContainerInner {
     name_to_rid: HashMap<String, ResourceId>,
 }
 
-/// Returns the `dbs/{db}/colls/{coll}` name path from a [`ContainerReference`],
-/// reusing the pre-computed `name_based_path` by skipping the leading `/`.
-fn name_path(container: &ContainerReference) -> &str {
-    // name_based_path() returns "/dbs/{db}/colls/{coll}"; skip the leading '/'.
-    &container.name_based_path()[1..]
+/// Returns the container path used as the session-token index key, reusing the
+/// pre-computed path by skipping the leading `/`.
+///
+/// For name-addressed containers this is `dbs/{db}/colls/{coll}`; for RID-addressed
+/// containers it is the RID-based path. Either way it is stable per container and
+/// consistent between `set` and `resolve`, which is all the name→RID fallback needs.
+fn index_path(container: &ContainerReference) -> &str {
+    // base_path() returns "/dbs/.../colls/..."; skip the leading '/'.
+    &container.base_path()[1..]
 }
 
 impl SessionContainer {
@@ -84,7 +88,7 @@ impl SessionContainer {
         }
 
         // Fall back to name → RID → token
-        let np = name_path(container);
+        let np = index_path(container);
         if let Some(resolved_rid) = guard.name_to_rid.get(np) {
             return Self::build_composite_token(&guard, resolved_rid.as_str());
         }
@@ -106,7 +110,7 @@ impl SessionContainer {
         session_token_value: &str,
     ) {
         let collection_rid = container.rid();
-        let np = name_path(container);
+        let np = index_path(container);
         let mut guard = self.inner.write().unwrap_or_else(|e| e.into_inner());
 
         let rid = ResourceId::new(collection_rid.to_owned());