SpiceAgent: Handle file transfer requests properly :^)

Now, we write the data recieved to a file when the user drags a file
onto the Spice Viewer window. Once complete, the FileExplorer will open
with the copied file highlighted.

6 files changed, 215 insertions, 2 deletions

diff --git a/Userland/Services/SpiceAgent/CMakeLists.txt b/Userland/Services/SpiceAgent/CMakeLists.txt
index f2eb8c3a2b..9702ef01d4 100644
--- a/Userland/Services/SpiceAgent/CMakeLists.txt
+++ b/Userland/Services/SpiceAgent/CMakeLists.txt
@@ -5,10 +5,11 @@ serenity_component(
 
 set(SOURCES
     main.cpp
+    FileTransferOperation.cpp
     Message.cpp
     SpiceAgent.cpp
 )
 
 serenity_bin(SpiceAgent)
-target_link_libraries(SpiceAgent PRIVATE LibCore LibGfx LibGUI LibMain)
+target_link_libraries(SpiceAgent PRIVATE LibCore LibDesktop LibFileSystem LibGfx LibGUI LibMain)
 add_dependencies(SpiceAgent Clipboard)
diff --git a/Userland/Services/SpiceAgent/FileTransferOperation.cpp b/Userland/Services/SpiceAgent/FileTransferOperation.cpp
new file mode 100644
index 0000000000..44ba04f94b
--- /dev/null
+++ b/Userland/Services/SpiceAgent/FileTransferOperation.cpp
@@ -0,0 +1,92 @@
+/*
+ * Copyright (c) 2023, Caoimhe Byrne <caoimhebyrne06@gmail.com>
+ *
+ * SPDX-License-Identifier: BSD-2-Clause
+ */
+
+#include "FileTransferOperation.h"
+#include "SpiceAgent.h"
+#include <AK/URL.h>
+#include <LibCore/StandardPaths.h>
+#include <LibDesktop/Launcher.h>
+#include <LibFileSystem/FileSystem.h>
+
+namespace SpiceAgent {
+
+ErrorOr<NonnullRefPtr<FileTransferOperation>> FileTransferOperation::create(FileTransferStartMessage& message)
+{
+    // Attempt to construct a path.
+    StringBuilder destination_builder;
+    TRY(destination_builder.try_append(Core::StandardPaths::downloads_directory()));
+    TRY(destination_builder.try_append('/'));
+    TRY(destination_builder.try_append(message.metadata().name));
+
+    auto destination_path = TRY(destination_builder.to_string());
+
+    // Ensure that the file doesn't already exist, and if it does, remove it.
+    if (FileSystem::exists(destination_path)) {
+        // If that "file" is a directory, we should stop doing anything else.
+        if (FileSystem::is_directory(destination_path)) {
+            return Error::from_string_literal("The name of the file being transferred is already taken by a directory!");
+        }
+
+        TRY(FileSystem::remove(destination_path, FileSystem::RecursionMode::Disallowed));
+    }
+
+    auto file = TRY(Core::File::open(destination_path, Core::File::OpenMode::ReadWrite));
+    return TRY(AK::adopt_nonnull_ref_or_enomem(new (nothrow) FileTransferOperation(message.id(), message.metadata(), move(file))));
+}
+
+ErrorOr<void> FileTransferOperation::begin_transfer(SpiceAgent& agent)
+{
+    // Ensure that we are in the `Pending` status.
+    if (m_status != Status::Pending) {
+        return Error::from_string_literal("Attempt to start a file transfer which has already been started!");
+    }
+
+    // Send the CanSendData status to the server.
+    auto status_message = FileTransferStatusMessage(m_id, FileTransferStatus::CanSendData);
+    TRY(agent.send_message(status_message));
+
+    // We are now in the transferring stage!
+    set_status(Status::Transferring);
+
+    return {};
+}
+
+ErrorOr<void> FileTransferOperation::complete_transfer(SpiceAgent& agent)
+{
+    // Ensure that we are in the `Transferring` status.
+    if (m_status != Status::Transferring) {
+        return Error::from_string_literal("Attempt to call `on_data_received` on a file transfer which has already been completed!");
+    }
+
+    // We are now in the complete stage :^)
+    set_status(Status::Complete);
+
+    // Send the Success status to the server, since we have received the data, and handled it correctly
+    auto status_message = FileTransferStatusMessage(m_id, FileTransferStatus::Success);
+    TRY(agent.send_message(status_message));
+
+    // Open the file manager for the user :^)
+    // FIXME: This currently opens a new window for each successful file transfer...
+    //        Is there a way/can we make a way for it to highlight a new file in an already-open window?
+    Desktop::Launcher::open(URL::create_with_file_scheme(Core::StandardPaths::downloads_directory(), m_metadata.name.to_deprecated_string()));
+
+    return {};
+}
+
+ErrorOr<void> FileTransferOperation::on_data_received(FileTransferDataMessage& message)
+{
+    // Ensure that we are in the `Transferring` status.
+    if (m_status != Status::Transferring) {
+        return Error::from_string_literal("Attempt to call `on_data_received` on a file transfer which has already been completed!");
+    }
+
+    // Attempt to write more data to the file.
+    TRY(m_destination->write_until_depleted(message.contents()));
+
+    return {};
+}
+
+}
diff --git a/Userland/Services/SpiceAgent/FileTransferOperation.h b/Userland/Services/SpiceAgent/FileTransferOperation.h
new file mode 100644
index 0000000000..2e0d7b1079
--- /dev/null
+++ b/Userland/Services/SpiceAgent/FileTransferOperation.h
@@ -0,0 +1,64 @@
+/*
+ * Copyright (c) 2023, Caoimhe Byrne <caoimhebyrne06@gmail.com>
+ *
+ * SPDX-License-Identifier: BSD-2-Clause
+ */
+
+#pragma once
+
+#include "Message.h"
+#include <AK/Types.h>
+#include <LibCore/File.h>
+
+namespace SpiceAgent {
+
+// Forward declaration
+class SpiceAgent;
+
+class FileTransferOperation : public RefCounted<FileTransferOperation> {
+public:
+    enum class Status {
+        // If we haven't accepted the transfer yet.
+        Pending,
+
+        // If we are awaiting data from the server.
+        Transferring,
+
+        // If we've received all the data.
+        Complete
+    };
+
+    static ErrorOr<NonnullRefPtr<FileTransferOperation>> create(FileTransferStartMessage& message);
+
+    // Fired by the SpiceAgent when it wants the data transfer to begin.
+    ErrorOr<void> begin_transfer(SpiceAgent& agent);
+
+    // Fired by SpiceAgent when we have received all of the data needed for this transfer.
+    ErrorOr<void> complete_transfer(SpiceAgent& agent);
+
+    // Fired by the SpiceAgent when it recieves data related to this transfer.
+    ErrorOr<void> on_data_received(FileTransferDataMessage& message);
+
+private:
+    // All file transfers start off as Pending.
+    FileTransferOperation(u32 id, FileTransferStartMessage::Metadata metadata, NonnullOwnPtr<Core::File> destination)
+        : m_destination(move(destination))
+        , m_metadata(move(metadata))
+        , m_id(id)
+        , m_status(Status::Pending)
+    {
+    }
+
+    void set_status(Status const& value)
+    {
+        m_status = value;
+    }
+
+    NonnullOwnPtr<Core::File> m_destination;
+    FileTransferStartMessage::Metadata m_metadata;
+
+    u32 m_id { 0 };
+    Status m_status { Status::Pending };
+};
+
+}
diff --git a/Userland/Services/SpiceAgent/SpiceAgent.cpp b/Userland/Services/SpiceAgent/SpiceAgent.cpp
index 075f4b6c53..04d100a455 100644
--- a/Userland/Services/SpiceAgent/SpiceAgent.cpp
+++ b/Userland/Services/SpiceAgent/SpiceAgent.cpp
@@ -158,12 +158,56 @@ ErrorOr<void> SpiceAgent::on_message_received()
         auto message = TRY(FileTransferStatusMessage::read_from_stream(stream));
         dbgln("File transfer {} has been cancelled: {}", message.id(), message.status());
 
+        m_file_transfer_operations.remove(message.id());
+
         break;
     }
 
+    // Received when the user drags a file onto the virtual machine.
     case Message::Type::FileTransferStart: {
         auto message = TRY(FileTransferStartMessage::read_from_stream(stream));
-        dbgln("File transfer request received: {}", TRY(message.debug_description()));
+        auto operation = TRY(FileTransferOperation::create(message));
+
+        // Tell the operation to start the file transfer.
+        TRY(operation->begin_transfer(*this));
+        m_file_transfer_operations.set(message.id(), operation);
+
+        break;
+    }
+
+    // Received when the server has data related to a file transfer for us.
+    case Message::Type::FileTransferData: {
+        auto message = TRY(FileTransferDataMessage::read_from_stream(stream));
+        auto optional_operation = m_file_transfer_operations.get(message.id());
+        if (!optional_operation.has_value()) {
+            return Error::from_string_literal("Attempt to supply data to a file transfer operation which doesn't exist!");
+        }
+
+        // Inform the operation that we have received new data.
+        auto* operation = optional_operation.release_value();
+        auto result = operation->on_data_received(message);
+        if (result.is_error()) {
+            // We can also discard of this transfer operation, since it will be cancelled by the server after our status message.
+            m_file_transfer_operations.remove(message.id());
+
+            // Inform the server that the operation has failed
+            auto status_message = FileTransferStatusMessage(message.id(), FileTransferStatus::Error);
+            TRY(this->send_message(status_message));
+
+            return result.release_error();
+        }
+
+        // The maximum amount of data that a FileTransferData message can hold is 65536.
+        // If it's less than 65536, this is the only (or last) message in relation to this transfer.
+        // Otherwise, we must wait for more data to be received.
+        auto transfer_is_complete = message.contents().size() < file_transfer_buffer_threshold;
+        if (!transfer_is_complete) {
+            return {};
+        }
+
+        // The transfer is now complete, let's write the data to the file!
+        TRY(operation->complete_transfer(*this));
+        m_file_transfer_operations.remove(message.id());
 
         break;
     }
diff --git a/Userland/Services/SpiceAgent/SpiceAgent.h b/Userland/Services/SpiceAgent/SpiceAgent.h
index 2ea1893f6e..98994228fb 100644
--- a/Userland/Services/SpiceAgent/SpiceAgent.h
+++ b/Userland/Services/SpiceAgent/SpiceAgent.h
@@ -8,6 +8,7 @@
 #pragma once
 
 #include "ChunkHeader.h"
+#include "FileTransferOperation.h"
 #include "Message.h"
 #include "MessageHeader.h"
 #include <AK/MemoryStream.h>
@@ -22,6 +23,9 @@ namespace SpiceAgent {
 // If the buffer's length is equal to this, then the next data recieved will be more data from the same buffer.
 constexpr u32 message_buffer_threshold = 2048;
 
+// The maximum amount of data that can be received in one file transfer message
+constexpr u32 file_transfer_buffer_threshold = 65536;
+
 class SpiceAgent {
 public:
     static ErrorOr<NonnullOwnPtr<SpiceAgent>> create(StringView device_path);
@@ -61,6 +65,7 @@ public:
 private:
     NonnullOwnPtr<Core::File> m_spice_device;
     Vector<Capability> m_capabilities;
+    HashMap<u32, NonnullRefPtr<FileTransferOperation>> m_file_transfer_operations;
 
     RefPtr<Core::Notifier> m_notifier;
 
diff --git a/Userland/Services/SpiceAgent/main.cpp b/Userland/Services/SpiceAgent/main.cpp
index ff2caafa05..ca3f62c4f6 100644
--- a/Userland/Services/SpiceAgent/main.cpp
+++ b/Userland/Services/SpiceAgent/main.cpp
@@ -6,7 +6,10 @@
  */
 
 #include "SpiceAgent.h"
+#include <AK/URL.h>
+#include <LibCore/StandardPaths.h>
 #include <LibCore/System.h>
+#include <LibDesktop/Launcher.h>
 #include <LibGUI/Application.h>
 #include <LibGUI/Clipboard.h>
 #include <LibIPC/ConnectionToServer.h>
@@ -20,11 +23,15 @@ ErrorOr<int> serenity_main(Main::Arguments arguments)
     // We use the application to be able to easily write to the user's clipboard.
     auto app = TRY(GUI::Application::create(arguments));
 
+    TRY(Desktop::Launcher::add_allowed_url(URL::create_with_file_scheme(Core::StandardPaths::downloads_directory())));
+    TRY(Desktop::Launcher::seal_allowlist());
+
     // FIXME: Make Core::File support reading and writing, but without creating:
     //        By default, Core::File opens the file descriptor with O_CREAT when using OpenMode::Write (and subsequently, OpenMode::ReadWrite).
     //        To minimise confusion for people that have already used Core::File, we can probably just do `OpenMode::ReadWrite | OpenMode::DontCreate`.
     TRY(Core::System::pledge("unix rpath wpath stdio sendfd recvfd cpath"));
     TRY(Core::System::unveil(SPICE_DEVICE, "rwc"sv));
+    TRY(Core::System::unveil(Core::StandardPaths::downloads_directory(), "rwc"sv));
     TRY(Core::System::unveil(nullptr, nullptr));
 
     auto agent = TRY(SpiceAgent::SpiceAgent::create(SPICE_DEVICE));