From 5119094dabd5f86b3636b1df614900591d0d8aa6 Mon Sep 17 00:00:00 2001
From: xiongzhaozhi <xiongzhaozhi@bilibili.com>
Date: Wed, 12 Apr 2023 00:23:22 +0800
Subject: [PATCH] fix crash

---
 src/audio-capture-helper.cpp | 27 ++++++++++++++++++---------
 src/audio-capture.cpp        | 13 ++++++++-----
 src/session-monitor.hpp      |  1 +
 3 files changed, 27 insertions(+), 14 deletions(-)

diff --git a/src/audio-capture-helper.cpp b/src/audio-capture-helper.cpp
index 7b3e546..3c25397 100644
--- a/src/audio-capture-helper.cpp
+++ b/src/audio-capture-helper.cpp
@@ -46,16 +46,17 @@ void AudioCaptureHelper::InitClient()
 	auto propvariant = GetPropvariant(&params);
 
 	wil::com_ptr<IActivateAudioInterfaceAsyncOperation> async_op;
-	CompletionHandler completion_handler;
+
+	Microsoft::WRL::ComPtr<CompletionHandler> completion_handler = Microsoft::WRL::Make<CompletionHandler>();
 
 	THROW_IF_FAILED(ActivateAudioInterfaceAsync(VIRTUAL_AUDIO_DEVICE_PROCESS_LOOPBACK,
 						    __uuidof(IAudioClient), &propvariant,
-						    &completion_handler, &async_op));
+						    completion_handler.Get(), &async_op));
 
-	completion_handler.event_finished.wait();
-	THROW_IF_FAILED(completion_handler.activate_hr);
+	completion_handler->event_finished.wait();
+	THROW_IF_FAILED(completion_handler->activate_hr);
 
-	client = completion_handler.client;
+	client = completion_handler->client;
 
 	THROW_IF_FAILED(
 		client->Initialize(AUDCLNT_SHAREMODE_SHARED,
@@ -102,12 +103,20 @@ void AudioCaptureHelper::ForwardPacket()
 	THROW_IF_FAILED(capture_client->GetNextPacketSize(&num_frames));
 
 	while (num_frames > 0) {
-		BYTE *new_data;
+		BYTE *new_data = nullptr;
 		DWORD flags;
 		UINT64 qpc_position;
 
-		THROW_IF_FAILED(capture_client->GetBuffer(&new_data, &num_frames, &flags, NULL,
-							  &qpc_position));
+		HRESULT hr = capture_client->GetBuffer(&new_data, &num_frames, &flags, NULL,
+			&qpc_position);
+
+		THROW_IF_FAILED(hr);
+		if (hr == AUDCLNT_S_BUFFER_EMPTY || new_data == nullptr)
+		{
+			THROW_IF_FAILED(capture_client->ReleaseBuffer(num_frames));
+			THROW_IF_FAILED(capture_client->GetNextPacketSize(&num_frames));
+			continue;
+		}
 
 		if (!(flags & AUDCLNT_BUFFERFLAGS_SILENT))
 			ForwardToMixers(qpc_position, new_data, num_frames);
@@ -164,7 +173,7 @@ void AudioCaptureHelper::CaptureSafe()
 AudioCaptureHelper::AudioCaptureHelper(Mixer *mixer, WAVEFORMATEX format, DWORD pid)
 	: mixers{mixer}, format{format}, pid{pid}
 {
-	for (auto &event : events)
+ 	for (auto &event : events)
 		event.create();
 
 	capture_thread = std::thread(&AudioCaptureHelper::CaptureSafe, this);
diff --git a/src/audio-capture.cpp b/src/audio-capture.cpp
index f800002..0d24b1d 100644
--- a/src/audio-capture.cpp
+++ b/src/audio-capture.cpp
@@ -87,12 +87,15 @@ AudioCapture::DeDuplicateCaptureList(const std::set<DWORD> &pids,
 		for (auto pid : explicitly_captured_pids)
 			uncaptured_pids.erase(pid);
 
-		for (auto pid : uncaptured_pids) {
-			if (!explicitly_captured_pids.contains(parents[pid]))
+		auto iter = uncaptured_pids.begin();
+		while (iter != uncaptured_pids.end()) {
+			if (!explicitly_captured_pids.contains(parents[*iter]))
+			{
+				++iter;
 				continue;
-
-			implicitly_captured_pids.insert(pid);
-			uncaptured_pids.erase(pid);
+			}
+			implicitly_captured_pids.insert(*iter);
+			iter = uncaptured_pids.erase(iter);
 		}
 	}
 
diff --git a/src/session-monitor.hpp b/src/session-monitor.hpp
index 6340afe..7491844 100644
--- a/src/session-monitor.hpp
+++ b/src/session-monitor.hpp
@@ -15,6 +15,7 @@
 #include <wrl/implements.h>
 
 #include "common.hpp"
+#include <optional>
 
 using namespace Microsoft::WRL;