fix: potential memory leak in string conversions (#119)

pykeio · Nov 21, 2023 · b4b59ee · b4b59ee
2 parents fd84626 + 449d9b3
commit b4b59ee
Show file tree

Hide file tree

Showing 3 changed files with 44 additions and 10 deletions.
diff --git a/src/execution_providers/mod.rs b/src/execution_providers/mod.rs
@@ -61,7 +61,13 @@ pub trait ExecutionProvider {
 		}
 
 		for i in 0..num_providers {
-			let avail = char_p_to_string(unsafe { *providers.offset(i as isize) })?;
+			let avail = match char_p_to_string(unsafe { *providers.offset(i as isize) }) {
+				Ok(avail) => avail,
+				Err(e) => {
+					let _ = ortsys![unsafe ReleaseAvailableProviders(providers, num_providers)];
+					return Err(e);
+				}
+			};
 			if self.as_str() == avail {
 				let _ = ortsys![unsafe ReleaseAvailableProviders(providers, num_providers)];
 				return Ok(true);

diff --git a/src/metadata.rs b/src/metadata.rs
@@ -18,7 +18,13 @@ impl ModelMetadata {
 		let mut str_bytes: *mut c_char = std::ptr::null_mut();
 		ortsys![unsafe ModelMetadataGetDescription(self.metadata_ptr, self.allocator_ptr, &mut str_bytes) -> Error::GetModelMetadata; nonNull(str_bytes)];
 
-		let value = char_p_to_string(str_bytes)?;
+		let value = match char_p_to_string(str_bytes) {
+			Ok(value) => value,
+			Err(e) => {
+				ortfree!(unsafe self.allocator_ptr, str_bytes);
+				return Err(e);
+			}
+		};
 		ortfree!(unsafe self.allocator_ptr, str_bytes);
 		Ok(value)
 	}
@@ -28,7 +34,13 @@ impl ModelMetadata {
 		let mut str_bytes: *mut c_char = std::ptr::null_mut();
 		ortsys![unsafe ModelMetadataGetProducerName(self.metadata_ptr, self.allocator_ptr, &mut str_bytes) -> Error::GetModelMetadata; nonNull(str_bytes)];
 
-		let value = char_p_to_string(str_bytes)?;
+		let value = match char_p_to_string(str_bytes) {
+			Ok(value) => value,
+			Err(e) => {
+				ortfree!(unsafe self.allocator_ptr, str_bytes);
+				return Err(e);
+			}
+		};
 		ortfree!(unsafe self.allocator_ptr, str_bytes);
 		Ok(value)
 	}
@@ -38,7 +50,13 @@ impl ModelMetadata {
 		let mut str_bytes: *mut c_char = std::ptr::null_mut();
 		ortsys![unsafe ModelMetadataGetGraphName(self.metadata_ptr, self.allocator_ptr, &mut str_bytes) -> Error::GetModelMetadata; nonNull(str_bytes)];
 
-		let value = char_p_to_string(str_bytes)?;
+		let value = match char_p_to_string(str_bytes) {
+			Ok(value) => value,
+			Err(e) => {
+				ortfree!(unsafe self.allocator_ptr, str_bytes);
+				return Err(e);
+			}
+		};
 		ortfree!(unsafe self.allocator_ptr, str_bytes);
 		Ok(value)
 	}
@@ -56,11 +74,15 @@ impl ModelMetadata {
 		let key_str = CString::new(key)?;
 		ortsys![unsafe ModelMetadataLookupCustomMetadataMap(self.metadata_ptr, self.allocator_ptr, key_str.as_ptr(), &mut str_bytes) -> Error::GetModelMetadata];
 		if !str_bytes.is_null() {
-			unsafe {
-				let value = char_p_to_string(str_bytes)?;
-				ortfree!(self.allocator_ptr, str_bytes);
-				Ok(Some(value))
-			}
+			let value = match char_p_to_string(str_bytes) {
+				Ok(value) => value,
+				Err(e) => {
+					ortfree!(unsafe self.allocator_ptr, str_bytes);
+					return Err(e);
+				}
+			};
+			ortfree!(unsafe self.allocator_ptr, str_bytes);
+			Ok(Some(value))
 		} else {
 			Ok(None)
 		}

diff --git a/src/session/mod.rs b/src/session/mod.rs
@@ -760,7 +760,13 @@ mod dangerous {
 	}
 
 	pub(crate) fn raw_pointer_to_string(allocator_ptr: *mut ort_sys::OrtAllocator, c_str: *mut c_char) -> Result<String> {
-		let name = char_p_to_string(c_str)?;
+		let name = match char_p_to_string(c_str) {
+			Ok(name) => name,
+			Err(e) => {
+				ortfree!(unsafe allocator_ptr, c_str);
+				return Err(e);
+			}
+		};
 		ortfree!(unsafe allocator_ptr, c_str);
 		Ok(name)
 	}