cs: implement TakeScreenShot command

libraries/libstratosphere/include/stratosphere/capsrv/capsrv_screen_shot_control_api.hpp

@@ -25,6 +25,10 @@ namespace ams::capsrv {
     Result InitializeScreenShotControl();
     void   FinalizeScreenShotControl();
 
+    Result OpenRawScreenShotReadStreamForDevelop(size_t *out_data_size, s32 *out_width, s32 *out_height, vi::LayerStack layer_stack, TimeSpan timeout);
+    Result ReadRawScreenShotReadStreamForDevelop(size_t *out_read_size, void *dst, size_t dst_size, std::ptrdiff_t offset);
+    void CloseRawScreenShotReadStreamForDevelop();
+
     Result CaptureJpegScreenshot(u64 *out_size, void *dst, size_t dst_size, vi::LayerStack layer_stack, TimeSpan timeout);
 
 }

libraries/libstratosphere/include/stratosphere/cs/cs_command_processor.hpp

@@ -16,6 +16,7 @@
 #pragma once
 #include <vapours.hpp>
 #include <stratosphere/scs/scs_command_processor.hpp>
+#include <stratosphere/vi/vi_layer_stack.hpp>
 
 namespace ams::cs {
 
@@ -25,6 +26,8 @@ namespace ams::cs {
     class CommandProcessor : public scs::CommandProcessor {
         public:
             virtual bool ProcessCommand(const CommandHeader &header, const u8 *body, s32 socket) override;
+        private:
+            void TakeScreenShot(const CommandHeader &header, s32 socket, vi::LayerStack layer_stack);
         private:
             static void SendFirmwareVersion(s32 socket, const CommandHeader &header);
     };

libraries/libstratosphere/source/capsrv/capsrv_screen_shot_control_api.cpp

@@ -29,4 +29,27 @@ namespace ams::capsrv {
         return ::capsscCaptureJpegScreenShot(out_size, dst, dst_size, static_cast<::ViLayerStack>(layer_stack), timeout.GetNanoSeconds());
     }
 
+    Result OpenRawScreenShotReadStreamForDevelop(size_t *out_data_size, s32 *out_width, s32 *out_height, vi::LayerStack layer_stack, TimeSpan timeout) {
+        u64 data_size, width, height;
+        R_TRY(::capsscOpenRawScreenShotReadStream(std::addressof(data_size), std::addressof(width), std::addressof(height), static_cast<::ViLayerStack>(layer_stack), timeout.GetNanoSeconds()));
+
+        *out_data_size = static_cast<size_t>(data_size);
+        *out_width     = static_cast<s32>(width);
+        *out_height    = static_cast<s32>(height);
+
+        return ResultSuccess();
+    }
+
+    Result ReadRawScreenShotReadStreamForDevelop(size_t *out_read_size, void *dst, size_t dst_size, std::ptrdiff_t offset) {
+        u64 read_size;
+        R_TRY(::capsscReadRawScreenShotReadStream(std::addressof(read_size), dst, dst_size, static_cast<u64>(offset)));
+
+        *out_read_size = static_cast<size_t>(read_size);
+        return ResultSuccess();
+    }
+
+    void CloseRawScreenShotReadStreamForDevelop() {
+        ::capsscCloseRawScreenShotReadStream();
+    }
+
 }

libraries/libstratosphere/source/cs/cs_command_impl.cpp

@@ -18,9 +18,65 @@
 
 namespace ams::cs {
 
+    namespace {
+
+        void SendEmptyData(const CommandDataTakeScreenShot &params, size_t remaining_size) {
+            /* Clear the data buffer. */
+            std::memset(params.buffer, 0, params.buffer_size);
+
+            /* Send data until the end. */
+            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);
+
+                /* Advance. */
+                remaining_size -= cur_size;
+            }
+        }
+
+    }
+
     Result DoGetFirmwareVersionCommand(settings::system::FirmwareVersion *out) {
         settings::system::GetFirmwareVersion(out);
         return ResultSuccess();
     }
 
+    Result DoTakeScreenShotCommand(const CommandDataTakeScreenShot &params) {
+        /* Initialize screenshot control. */
+        R_TRY(capsrv::InitializeScreenShotControl());
+
+        /* Finalize screenshot control when we're done. */
+        ON_SCOPE_EXIT { capsrv::FinalizeScreenShotControl(); };
+
+        /* Open screenshot read stream. */
+        size_t data_size;
+        s32 width, height;
+        R_TRY(capsrv::OpenRawScreenShotReadStreamForDevelop(std::addressof(data_size), std::addressof(width), std::addressof(height), params.layer_stack, TimeSpan::FromSeconds(10)));
+
+        /* Close the screenshot stream when we're done. */
+        ON_SCOPE_EXIT { capsrv::CloseRawScreenShotReadStreamForDevelop(); };
+
+        /* Send the header. */
+        params.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); };
+        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);
+
+            /* Advance. */
+            total_read_size += read_size;
+        }
+        data_guard.Cancel();
+
+        return ResultSuccess();
+    }
+
 }

libraries/libstratosphere/source/cs/cs_command_impl.hpp

@@ -18,6 +18,16 @@
 
 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);
+
 }

libraries/libstratosphere/source/cs/cs_command_processor.cpp

@@ -25,6 +25,15 @@ namespace ams::cs {
             settings::system::FirmwareVersion firmware_version;
         };
 
+        struct ResponseTakeScreenShot {
+            ResponseHeader header;
+            s32 data_size;
+            s32 width;
+            s32 height;
+        };
+
+        constinit u8 g_data[0x1000];
+
     }
 
     bool CommandProcessor::ProcessCommand(const CommandHeader &header, const u8 *body, s32 socket) {
@@ -32,7 +41,12 @@ namespace ams::cs {
             case Command_GetFirmwareVersion:
                 SendFirmwareVersion(socket, header);
                 break;
-            /* TODO: Command support. */
+            case Command_TakeScreenShot:
+                this->TakeScreenShot(header, socket, vi::LayerStack_ApplicationForDebug);
+                break;
+            case Command_TakeForegroundScreenShot:
+                this->TakeScreenShot(header, socket, vi::LayerStack_LastFrame);
+                break;
             /* TODO: Command support. */
             default:
                 scs::CommandProcessor::ProcessCommand(header, body, socket);
@@ -64,4 +78,45 @@ namespace ams::cs {
         }
     }
 
+    void CommandProcessor::TakeScreenShot(const CommandHeader &header, s32 socket, vi::LayerStack layer_stack) {
+        /* Create the command data. */
+        const CommandDataTakeScreenShot params = {
+            .layer_stack = layer_stack,
+            .send_header = [&](s32 data_size, s32 width, s32 height) {
+                /* Use global buffer for response. */
+                ResponseTakeScreenShot *response = reinterpret_cast<ResponseTakeScreenShot *>(g_data);
+
+                /* Set response header. */
+                *response = {
+                    .header    = {
+                        .id = header.id,
+                        .response  = Response_ScreenShot,
+                        .body_size = static_cast<u32>(sizeof(data_size) + sizeof(width) + sizeof(height) + data_size),
+                    },
+                    .data_size = data_size,
+                    .width     = width,
+                    .height    = height,
+                };
+
+                /* Send data. */
+                Send(socket, response, sizeof(*response));
+            },
+            .send_data   = [&](u8 *data, size_t data_size) {
+                /* Send data. */
+                Send(socket, data, data_size);
+            },
+            .buffer      = g_data,
+            .buffer_size = sizeof(g_data),
+        };
+
+        /* Acquire the send lock. */
+        auto lk = MakeSendGuardBlock();
+
+        /* Take the screenshot. */
+        const Result result = DoTakeScreenShotCommand(params);
+        if (R_FAILED(result)) {
+            SendErrorResult(socket, header, result);
+        }
+    }
+
 }