From 41a4bf53f6fcb9abb38175feb4f96e50dd5a779c Mon Sep 17 00:00:00 2001
From: Michael Scire <SciresM@gmail.com>
Date: Thu, 7 Oct 2021 19:48:40 -0700
Subject: [PATCH] cs: prevent exceptions from manifesting

---
 .../source/cs/cs_command_impl.hpp             |  5 ++-
 ...s_command_impl.cpp => cs_command_impl.inc} | 16 ++++-----
 .../source/cs/cs_command_processor.cpp        | 33 ++++++++++---------
 3 files changed, 27 insertions(+), 27 deletions(-)
 rename libraries/libstratosphere/source/cs/{cs_command_impl.cpp => cs_command_impl.inc} (86%)

diff --git a/libraries/libstratosphere/source/cs/cs_command_impl.hpp b/libraries/libstratosphere/source/cs/cs_command_impl.hpp
index 5f399b581..b8613088e 100644
--- a/libraries/libstratosphere/source/cs/cs_command_impl.hpp
+++ b/libraries/libstratosphere/source/cs/cs_command_impl.hpp
@@ -20,14 +20,13 @@ namespace ams::cs {
 
     struct CommandDataTakeScreenShot {
         vi::LayerStack layer_stack;
-        std::function<void (s32, s32, s32)> send_header;
-        std::function<void (u8 *, size_t)> send_data;
         u8 *buffer;
         size_t buffer_size;
     };
 
     Result DoGetFirmwareVersionCommand(settings::system::FirmwareVersion *out);
 
-    Result DoTakeScreenShotCommand(const CommandDataTakeScreenShot &params);
+    template<typename SendHeader, typename SendData>
+    Result DoTakeScreenShotCommand(const CommandDataTakeScreenShot &params, SendHeader send_header, SendData send_data);
 
 }
diff --git a/libraries/libstratosphere/source/cs/cs_command_impl.cpp b/libraries/libstratosphere/source/cs/cs_command_impl.inc
similarity index 86%
rename from libraries/libstratosphere/source/cs/cs_command_impl.cpp
rename to libraries/libstratosphere/source/cs/cs_command_impl.inc
index 0aeb9ae8e..5e8affc3b 100644
--- a/libraries/libstratosphere/source/cs/cs_command_impl.cpp
+++ b/libraries/libstratosphere/source/cs/cs_command_impl.inc
@@ -13,14 +13,13 @@
  * You should have received a copy of the GNU General Public License
  * along with this program.  If not, see <http://www.gnu.org/licenses/>.
  */
-#include <stratosphere.hpp>
-#include "cs_command_impl.hpp"
 
 namespace ams::cs {
 
     namespace {
 
-        void SendEmptyData(const CommandDataTakeScreenShot &params, size_t remaining_size) {
+        template<typename SendData>
+        void SendEmptyData(const CommandDataTakeScreenShot &params, size_t remaining_size, SendData send_data) {
             /* Clear the data buffer. */
             std::memset(params.buffer, 0, params.buffer_size);
 
@@ -28,7 +27,7 @@ namespace ams::cs {
             while (remaining_size > 0) {
                 /* Send as much as we can. */
                 const auto cur_size = std::min(remaining_size, params.buffer_size);
-                params.send_data(params.buffer, cur_size);
+                send_data(params.buffer, cur_size);
 
                 /* Advance. */
                 remaining_size -= cur_size;
@@ -42,7 +41,8 @@ namespace ams::cs {
         return ResultSuccess();
     }
 
-    Result DoTakeScreenShotCommand(const CommandDataTakeScreenShot &params) {
+    template<typename SendHeader, typename SendData>
+    Result DoTakeScreenShotCommand(const CommandDataTakeScreenShot &params, SendHeader send_header, SendData send_data) {
         /* Initialize screenshot control. */
         R_TRY(capsrv::InitializeScreenShotControl());
 
@@ -58,18 +58,18 @@ namespace ams::cs {
         ON_SCOPE_EXIT { capsrv::CloseRawScreenShotReadStreamForDevelop(); };
 
         /* Send the header. */
-        params.send_header(static_cast<s32>(data_size), width, height);
+        send_header(static_cast<s32>(data_size), width, height);
 
         /* Read and send data. */
         size_t total_read_size = 0;
-        auto data_guard = SCOPE_GUARD { SendEmptyData(params, data_size - total_read_size); };
+        auto data_guard = SCOPE_GUARD { SendEmptyData(params, data_size - total_read_size, send_data); };
         while (total_read_size < data_size) {
             /* Read data from the stream. */
             size_t read_size;
             R_TRY(capsrv::ReadRawScreenShotReadStreamForDevelop(std::addressof(read_size), params.buffer, params.buffer_size, total_read_size));
 
             /* Send the data that was read. */
-            params.send_data(params.buffer, read_size);
+            send_data(params.buffer, read_size);
 
             /* Advance. */
             total_read_size += read_size;
diff --git a/libraries/libstratosphere/source/cs/cs_command_processor.cpp b/libraries/libstratosphere/source/cs/cs_command_processor.cpp
index 7892d215f..6b1924aea 100644
--- a/libraries/libstratosphere/source/cs/cs_command_processor.cpp
+++ b/libraries/libstratosphere/source/cs/cs_command_processor.cpp
@@ -16,6 +16,9 @@
 #include <stratosphere.hpp>
 #include "cs_command_impl.hpp"
 
+/* Include command implementations. */
+#include "cs_command_impl.inc"
+
 namespace ams::cs {
 
     namespace {
@@ -82,7 +85,18 @@ namespace ams::cs {
         /* Create the command data. */
         const CommandDataTakeScreenShot params = {
             .layer_stack = layer_stack,
-            .send_header = [&](s32 data_size, s32 width, s32 height) {
+            .buffer      = g_data,
+            .buffer_size = sizeof(g_data),
+        };
+
+        /* Take the screenshot. */
+        Result result;
+        {
+            /* Acquire the send lock. */
+            auto lk = MakeSendGuardBlock();
+
+            /* Perform the command. */
+            result = DoTakeScreenShotCommand(params, [&](s32 data_size, s32 width, s32 height) {
                 /* Use global buffer for response. */
                 ResponseTakeScreenShot *response = reinterpret_cast<ResponseTakeScreenShot *>(g_data);
 
@@ -100,23 +114,10 @@ namespace ams::cs {
 
                 /* Send data. */
                 Send(socket, response, sizeof(*response));
-            },
-            .send_data   = [&](u8 *data, size_t data_size) {
+            }, [&](u8 *data, size_t data_size) {
                 /* Send data. */
                 Send(socket, data, data_size);
-            },
-            .buffer      = g_data,
-            .buffer_size = sizeof(g_data),
-        };
-
-        /* Take the screenshot. */
-        Result result;
-        {
-            /* Acquire the send lock. */
-            auto lk = MakeSendGuardBlock();
-
-            /* Perform the command. */
-            result = DoTakeScreenShotCommand(params);
+            });
         }
 
         /* Handle the error case. */