Merged PR 7444: Remove unwrap() from message, handler, and security modules

saurabh500 · saurabh500 · commit a66638fd53db · 2026-05-02T23:15:03.000Z
Replace .unwrap() calls in the message, handler, and security modules with proper error handling in preparation for enabling the unwrap_used clippy lint.

## Changes

- handler_factory.rs: Replace unwrap() with proper error propagation
- bulk_load.rs: Replace unwrap() with error handling for bulk load operations
- fedauth.rs: Replace unwrap() with safe alternative
- login.rs: Replace unwrap() with error propagation for login message construction
- security/mock.rs: Replace unwrap() with .expect() since this is test-only infrastructure

----
#### AI description  (iteration 1)
#### PR Classification
Code cleanup to replace panic-prone `unwrap()` calls with explicit error handling in security, message, and handler modules to support stricter Clippy lints.

#### PR Summary
This pull request eliminates `unwrap()` calls across security, message, and handler modules by replacing them with `expect()` calls containing descriptive error messages or proper error handling logic.

- `security/mock.rs`: Replaced all `unwrap()` calls with `expect()` containing context-specific error messages in tests and documentation examples
- `message/login.rs`: Added proper error handling for routing information to return `ProtocolError` instead of panicking when redirection target is missing
- `message/bulk_load.rs`: Replaced `unwrap()` with `ok_or_else()` to return `ImplementationError` when first row column count is unexpectedly missing
- `handler/handler_factory.rs`: Refactored to use `unwrap_or(false)` for database instance validation and `if let Some()` pattern for TDS error handling
- `message/features/fedauth.rs`: Changed `unwrap()` to `unwrap_or(i32::MAX)` for length conversion fallback
&lt;!-- GitOpsUserAgent=GitOps.Apps.Server.pullrequestcopilot --&gt;


Related work items: #44826
diff --git a/mssql-tds/src/handler/handler_factory.rs b/mssql-tds/src/handler/handler_factory.rs
@@ -349,14 +349,13 @@ impl PreloginHandler<'_> {
         let response = PreloginResponse {};
 
         let response_model = &response.deserialize(reader_writer).await?;
-        if request_model.mars_enabled && !response_model.mars_enabled.unwrap_or(false) {
+        if request_model.mars_enabled && !response_model.mars_enabled {
             return Err(Error::ProtocolError(
                 "Server does not support MARS (Multiple Active Result Sets)".to_string(),
             ));
         }
 
-        if !response_model.dbinstance_valid.unwrap() {
-            // Non-fatal behaviour.
+        if !response_model.dbinstance_valid {
             warn!("Database instance validation failed");
         }
 
@@ -607,8 +606,7 @@ impl LoginHandler<'_> {
         let response_model = response
             .deserialize(reader_writer, requested_features)
             .await?;
-        if response_model.tds_error.is_some() {
-            let tds_error = response_model.tds_error.unwrap();
+        if let Some(tds_error) = response_model.tds_error.as_ref() {
             Err(Error::from_sql_error(crate::error::SqlErrorInfo {
                 message: tds_error.get_message(),
                 state: tds_error.error_token.state,
diff --git a/mssql-tds/src/message/bulk_load.rs b/mssql-tds/src/message/bulk_load.rs
@@ -282,7 +282,11 @@ impl<'a> StreamingBulkLoadWriter<'a> {
             );
         } else {
             // Subsequent rows: validate against first row's column count
-            let expected_count = self.first_row_column_count.unwrap();
+            let expected_count = self.first_row_column_count.ok_or_else(|| {
+                Error::ImplementationError(
+                    "First row column count is missing after initial row write".to_string(),
+                )
+            })?;
             if column_index != expected_count {
                 return Err(Error::UsageError(format!(
                     "Row {} has {} columns, but first row had {} columns. All rows must have the same number of columns as the first row.",
diff --git a/mssql-tds/src/message/features/fedauth.rs b/mssql-tds/src/message/features/fedauth.rs
@@ -121,13 +121,17 @@ impl FedAuthFeature {
     /// The payload length is calculated based on the presence of an access token.
     /// If an access token is present, the length includes the size of the token.
     /// If no access token is present, the length is 2.
-    fn get_payload_length(&self) -> i32 {
+    fn get_payload_length(&self) -> TdsResult<i32> {
         let len = if let Some(bytes) = &self.access_token_bytes {
             1 + size_of::<i32>() + bytes.len()
         } else {
             2
         };
-        len.try_into().unwrap()
+        i32::try_from(len).map_err(|_| {
+            crate::error::Error::ProtocolError(format!(
+                "FedAuth payload length {len} exceeds i32 range"
+            ))
+        })
     }
 }
 
@@ -138,7 +142,9 @@ impl Feature for FedAuthFeature {
     }
 
     fn data_length(&self) -> i32 {
-        let data_length = self.get_payload_length();
+        // get_payload_length only fails if token > 2GB, which cannot happen in practice
+        #[allow(clippy::unwrap_used)]
+        let data_length = self.get_payload_length().unwrap();
         let base_length = size_of::<u8>() + size_of::<i32>();
         data_length + base_length as i32
     }
@@ -153,7 +159,7 @@ impl Feature for FedAuthFeature {
             .write_byte_async(self.feature_identifier().as_u8())
             .await?;
         packet_writer
-            .write_i32_async(self.get_payload_length())
+            .write_i32_async(self.get_payload_length()?)
             .await?;
 
         packet_writer.write_byte_async(self.get_options()).await?;
diff --git a/mssql-tds/src/message/login.rs b/mssql-tds/src/message/login.rs
@@ -486,9 +486,14 @@ impl LoginResponseModel {
             },
             EnvChangeContainer::RoutingType(routing_change) => {
                 self.change_properties.routing_information = routing_change.new_value().clone();
+                let routing_info = routing_change.new_value().as_ref().ok_or_else(|| {
+                    crate::error::Error::ProtocolError(
+                        "Routing env change did not include a redirection target".to_string(),
+                    )
+                })?;
                 Err(crate::error::Error::Redirection {
-                    host: routing_change.new_value().as_ref().unwrap().server.clone(),
-                    port: routing_change.new_value().as_ref().unwrap().port,
+                    host: routing_info.server.clone(),
+                    port: routing_info.port,
                 })
             }
             _ => {
diff --git a/mssql-tds/src/message/prelogin.rs b/mssql-tds/src/message/prelogin.rs
@@ -152,8 +152,8 @@ impl PreloginRequestModel {
 pub struct PreloginResponseModel {
     pub encryption: EncryptionType,
     pub federated_auth_supported: bool,
-    pub dbinstance_valid: Option<bool>,
-    pub mars_enabled: Option<bool>,
+    pub dbinstance_valid: bool,
+    pub mars_enabled: bool,
     pub server_version: Version,
     pub sql_server_version: SQLServerVersion,
 }
@@ -164,8 +164,8 @@ impl PreloginResponseModel {
             server_version: Version::new(0, 0, 0, 0),
             sql_server_version: SQLServerVersion::SqlServerNotsupported,
             encryption: EncryptionType::Off,
-            dbinstance_valid: Option::from(false),
-            mars_enabled: Option::from(false),
+            dbinstance_valid: false,
+            mars_enabled: false,
             federated_auth_supported: false,
         }
     }
@@ -242,10 +242,10 @@ impl PreloginResponse {
                     // encryption type.
                 }
                 OptionType::InstOpt => {
-                    result.dbinstance_valid = Option::from(packet_reader.read_byte().await? == 0);
+                    result.dbinstance_valid = packet_reader.read_byte().await? == 0;
                 }
                 OptionType::Mars => {
-                    result.mars_enabled = Option::from(packet_reader.read_byte().await? == 1);
+                    result.mars_enabled = packet_reader.read_byte().await? == 1;
                 }
                 OptionType::FedAuthRequired => {
                     result.federated_auth_supported = packet_reader.read_byte().await? == 1;
@@ -545,7 +545,7 @@ pub(crate) mod tests {
         let response_model = block_on(response.deserialize(&mut mocked_packet_reader)).unwrap();
 
         // Compare the guid, which is auto-generated.
-        assert_eq!(response_model.mars_enabled, Option::from(true));
+        assert!(response_model.mars_enabled);
         assert_eq!(
             response_model.sql_server_version,
             SQLServerVersion::SqlServer2019
diff --git a/mssql-tds/src/security/mock.rs b/mssql-tds/src/security/mock.rs
@@ -21,15 +21,21 @@ use crate::security::{SecurityContext, SecurityError, SspiAuthToken};
 ///
 /// // Single-round authentication
 /// let mut ctx = MockSecurityContext::single_round(vec![1, 2, 3, 4]);
-/// let token = ctx.generate_token(None).unwrap();
+/// let token = ctx
+///     .generate_token(None)
+///     .expect("single-round mock should yield a token");
 /// assert_eq!(token.data, vec![1, 2, 3, 4]);
 /// assert!(token.is_complete);
 ///
 /// // Multi-round authentication
 /// let mut ctx = MockSecurityContext::multi_round(vec![1, 2], vec![3, 4]);
-/// let initial = ctx.generate_token(None).unwrap();
+/// let initial = ctx
+///     .generate_token(None)
+///     .expect("first mock round should yield a token");
 /// assert!(!initial.is_complete);
-/// let response = ctx.generate_token(Some(&[5, 6])).unwrap();
+/// let response = ctx
+///     .generate_token(Some(&[5, 6]))
+///     .expect("second mock round should yield a token");
 /// assert!(response.is_complete);
 /// ```
 #[derive(Debug, Clone)]
@@ -236,7 +242,9 @@ mod tests {
         assert!(!ctx.is_complete());
         assert_eq!(ctx.current_round(), 0);
 
-        let token = ctx.generate_token(None).unwrap();
+        let token = ctx
+            .generate_token(None)
+            .expect("single-round mock should yield a token");
         assert_eq!(token.data, vec![1, 2, 3, 4]);
         assert!(token.is_complete);
         assert!(ctx.is_complete());
@@ -248,13 +256,17 @@ mod tests {
         let mut ctx = MockSecurityContext::multi_round(vec![1, 2], vec![3, 4]);
 
         // First round
-        let initial = ctx.generate_token(None).unwrap();
+        let initial = ctx
+            .generate_token(None)
+            .expect("first multi-round token should be generated");
         assert_eq!(initial.data, vec![1, 2]);
         assert!(!initial.is_complete);
         assert!(!ctx.is_complete());
 
         // Second round with server challenge
-        let response = ctx.generate_token(Some(&[5, 6, 7])).unwrap();
+        let response = ctx
+            .generate_token(Some(&[5, 6, 7]))
+            .expect("second multi-round token should be generated");
         assert_eq!(response.data, vec![3, 4]);
         assert!(response.is_complete);
         assert!(ctx.is_complete());
@@ -268,13 +280,19 @@ mod tests {
         );
 
         // Three rounds
-        let t1 = ctx.generate_token(None).unwrap();
+        let t1 = ctx
+            .generate_token(None)
+            .expect("first custom-round token should be generated");
         assert!(!t1.is_complete);
 
-        let t2 = ctx.generate_token(Some(&[])).unwrap();
+        let t2 = ctx
+            .generate_token(Some(&[]))
+            .expect("second custom-round token should be generated");
         assert!(!t2.is_complete);
 
-        let t3 = ctx.generate_token(Some(&[])).unwrap();
+        let t3 = ctx
+            .generate_token(Some(&[]))
+            .expect("third custom-round token should be generated");
         assert!(t3.is_complete);
     }
 
@@ -304,11 +322,15 @@ mod tests {
 
         assert_eq!(ctx.package_name(), "NTLM");
 
-        let negotiate = ctx.generate_token(None).unwrap();
+        let negotiate = ctx
+            .generate_token(None)
+            .expect("NTLM mock should generate a negotiate token");
         assert!(negotiate.data.starts_with(b"NTLMSSP"));
         assert!(!negotiate.is_complete);
 
-        let authenticate = ctx.generate_token(Some(&[0xAA, 0xBB])).unwrap();
+        let authenticate = ctx
+            .generate_token(Some(&[0xAA, 0xBB]))
+            .expect("NTLM mock should generate an authenticate token");
         assert!(authenticate.data.starts_with(b"NTLMSSP"));
         assert!(authenticate.is_complete);
     }
@@ -319,7 +341,9 @@ mod tests {
 
         assert_eq!(ctx.package_name(), "Kerberos");
 
-        let token = ctx.generate_token(None).unwrap();
+        let token = ctx
+            .generate_token(None)
+            .expect("single-round mock should yield a token");
         assert!(token.is_complete);
     }
 
@@ -329,7 +353,8 @@ mod tests {
             .with_expected_challenges(vec![vec![0xAA, 0xBB]]);
 
         // First round - no challenge expected
-        ctx.generate_token(None).unwrap();
+        ctx.generate_token(None)
+            .expect("first expected-challenge round should generate a token");
 
         // Second round - wrong challenge should fail
         let result = ctx.generate_token(Some(&[0xCC, 0xDD]));
@@ -341,7 +366,8 @@ mod tests {
         let mut ctx = MockSecurityContext::multi_round(vec![1], vec![2])
             .with_expected_challenges(vec![vec![0xAA, 0xBB]]);
 
-        ctx.generate_token(None).unwrap();
+        ctx.generate_token(None)
+            .expect("first expected-challenge round should generate a token");
 
         // Correct challenge
         let result = ctx.generate_token(Some(&[0xAA, 0xBB]));
