Avoid leaking stale filesystem store temp files

Async filesystem writes may be awaited out of order, making older writes
stale after their temporary data has already been created. Ensure those
stale or failed write attempts do not leave historical plaintext data in
*.tmp files.

Co-Authored-By: HAL 9000

This finding was discovered by Project Loupe

Author: tnull

lightning-persister/src/fs_store/common.rs

@@ -42,6 +42,14 @@ fn path_to_windows_str<T: AsRef<OsStr>>(path: &T) -> Vec<u16> {
 	path.as_ref().encode_wide().chain(Some(0)).collect()
 }
 
+fn remove_file_if_exists(path: &Path) -> lightning::io::Result<()> {
+	match fs::remove_file(path) {
+		Ok(()) => Ok(()),
+		Err(e) if e.kind() == ErrorKind::NotFound => Ok(()),
+		Err(e) => Err(e.into()),
+	}
+}
+
 // The number of times we retry listing keys in `FilesystemStore::list` before we give up reaching
 // a consistent view and error out.
 const LIST_DIR_CONSISTENCY_RETRIES: usize = 10;
@@ -277,64 +285,81 @@ impl FilesystemStoreInner {
 		let tmp_file_ext = format!("{}.tmp", self.tmp_file_counter.fetch_add(1, Ordering::AcqRel));
 		tmp_file_path.set_extension(tmp_file_ext);
 
-		{
-			let mut tmp_file = fs::File::create(&tmp_file_path)?;
-			tmp_file.write_all(&buf)?;
+		let tmp_file_res = match fs::File::create(&tmp_file_path) {
+			Ok(mut tmp_file) => (|| -> lightning::io::Result<()> {
+				tmp_file.write_all(&buf)?;
 
-			// If we need to preserve the original mtime (for updates), set it before fsync.
-			if let Some(mtime) = mtime {
-				let times = fs::FileTimes::new().set_modified(mtime);
-				tmp_file.set_times(times)?;
-			}
+				// If we need to preserve the original mtime (for updates), set it before fsync.
+				if let Some(mtime) = mtime {
+					let times = fs::FileTimes::new().set_modified(mtime);
+					tmp_file.set_times(times)?;
+				}
 
-			tmp_file.sync_all()?;
+				tmp_file.sync_all()?;
+				Ok(())
+			})(),
+			Err(e) => return Err(e.into()),
+		};
+		if let Err(e) = tmp_file_res {
+			let _ = remove_file_if_exists(&tmp_file_path);
+			return Err(e);
 		}
 
-		self.execute_locked_write(inner_lock_ref, dest_file_path.clone(), version, || {
-			#[cfg(not(target_os = "windows"))]
-			{
-				fs::rename(&tmp_file_path, &dest_file_path)?;
-				let dir_file = fs::OpenOptions::new().read(true).open(&parent_directory)?;
-				dir_file.sync_all()?;
-				Ok(())
-			}
+		let mut tmp_file_needs_cleanup = true;
+		let write_res =
+			self.execute_locked_write(inner_lock_ref, dest_file_path.clone(), version, || {
+				#[cfg(not(target_os = "windows"))]
+				{
+					fs::rename(&tmp_file_path, &dest_file_path)?;
+					tmp_file_needs_cleanup = false;
+					let dir_file = fs::OpenOptions::new().read(true).open(&parent_directory)?;
+					dir_file.sync_all()?;
+					Ok(())
+				}
 
-			#[cfg(target_os = "windows")]
-			{
-				let res = if dest_file_path.exists() {
-					call!(unsafe {
-						windows_sys::Win32::Storage::FileSystem::ReplaceFileW(
+				#[cfg(target_os = "windows")]
+				{
+					let res = if dest_file_path.exists() {
+						call!(unsafe {
+							windows_sys::Win32::Storage::FileSystem::ReplaceFileW(
 							path_to_windows_str(&dest_file_path).as_ptr(),
 							path_to_windows_str(&tmp_file_path).as_ptr(),
 							std::ptr::null(),
 							windows_sys::Win32::Storage::FileSystem::REPLACEFILE_IGNORE_MERGE_ERRORS,
 							std::ptr::null_mut() as *const core::ffi::c_void,
 							std::ptr::null_mut() as *const core::ffi::c_void,
 							)
-					})
-				} else {
-					call!(unsafe {
-						windows_sys::Win32::Storage::FileSystem::MoveFileExW(
+						})
+					} else {
+						call!(unsafe {
+							windows_sys::Win32::Storage::FileSystem::MoveFileExW(
 							path_to_windows_str(&tmp_file_path).as_ptr(),
 							path_to_windows_str(&dest_file_path).as_ptr(),
 							windows_sys::Win32::Storage::FileSystem::MOVEFILE_WRITE_THROUGH
 							| windows_sys::Win32::Storage::FileSystem::MOVEFILE_REPLACE_EXISTING,
 							)
-					})
-				};
-
-				match res {
-					Ok(()) => {
-						// We fsync the dest file in hopes this will also flush the metadata to disk.
-						let dest_file =
-							fs::OpenOptions::new().read(true).write(true).open(&dest_file_path)?;
-						dest_file.sync_all()?;
-						Ok(())
-					},
-					Err(e) => Err(e.into()),
+						})
+					};
+
+					match res {
+						Ok(()) => {
+							tmp_file_needs_cleanup = false;
+							// We fsync the dest file in hopes this will also flush the metadata to disk.
+							let dest_file = fs::OpenOptions::new()
+								.read(true)
+								.write(true)
+								.open(&dest_file_path)?;
+							dest_file.sync_all()?;
+							Ok(())
+						},
+						Err(e) => Err(e.into()),
+					}
 				}
-			}
-		})
+			});
+		if tmp_file_needs_cleanup {
+			let _ = remove_file_if_exists(&tmp_file_path);
+		}
+		write_res
 	}
 
 	fn remove_version(

lightning-persister/src/fs_store/v2.rs

@@ -444,6 +444,39 @@ mod tests {
 		assert_eq!(listed_keys.len(), 0);
 	}
 
+	#[cfg(feature = "tokio")]
+	#[tokio::test]
+	async fn stale_write_does_not_leak_tmp_file() {
+		use lightning::util::persist::KVStore;
+
+		let mut temp_path = std::env::temp_dir();
+		temp_path.push("test_stale_write_does_not_leak_tmp_file_v2");
+		let _ = fs::remove_dir_all(&temp_path);
+		let fs_store = FilesystemStoreV2::new(temp_path.clone()).unwrap();
+
+		let data1 = vec![1u8; 32];
+		let data2 = vec![2u8; 32];
+
+		let primary = "testspace";
+		let secondary = "testsubspace";
+		let key = "testkey";
+
+		let fut1 = KVStore::write(&fs_store, primary, secondary, key, data1);
+		let fut2 = KVStore::write(&fs_store, primary, secondary, key, data2);
+
+		fut2.await.unwrap();
+		fut1.await.unwrap();
+
+		let dir = temp_path.join(primary).join(secondary);
+		let tmp_files: Vec<_> = fs::read_dir(&dir)
+			.unwrap()
+			.filter_map(|e| e.ok())
+			.map(|e| e.path())
+			.filter(|p| p.extension().map_or(false, |ext| ext == "tmp"))
+			.collect();
+		assert!(tmp_files.is_empty(), "Found leaked tmp files: {:?}", tmp_files);
+	}
+
 	#[test]
 	fn test_data_migration() {
 		let mut source_temp_path = std::env::temp_dir();