ProtocolServer: Stream the downloaded data if possible

This patchset makes ProtocolServer stream the downloads to its client
(LibProtocol), and as such changes the download API; a possible
download lifecycle could be as such:
notation = client->server:'>', server->client:'<', pipe activity:'*'
```
> StartDownload(GET, url, headers, {})
< Response(0, fd 8)
* {data, 1024b}
< HeadersBecameAvailable(0, response_headers, 200)
< DownloadProgress(0, 4K, 1024)
* {data, 1024b}
* {data, 1024b}
< DownloadProgress(0, 4K, 2048)
* {data, 1024b}
< DownloadProgress(0, 4K, 1024)
< DownloadFinished(0, true, 4K)
```

Since managing the received file descriptor is a pain, LibProtocol
implements `Download::stream_into(OutputStream)`, which can be used to
stream the download into any given output stream (be it a file, or
memory, or writing stuff with a delay, etc.).
Also, as some of the users of this API require all the downloaded data
upfront, LibProtocol also implements `set_should_buffer_all_input()`,
which causes the download instance to buffer all the data until the
download is complete, and to call the `on_buffered_download_finish`
hook.

Libraries/LibProtocol/Client.cpp

@@ -24,6 +24,7 @@
  * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
  */
 
+#include <AK/FileStream.h>
 #include <AK/SharedBuffer.h>
 #include <LibProtocol/Client.h>
 #include <LibProtocol/Download.h>
@@ -47,16 +48,20 @@ bool Client::is_supported_protocol(const String& protocol)
     return send_sync<Messages::ProtocolServer::IsSupportedProtocol>(protocol)->supported();
 }
 
-RefPtr<Download> Client::start_download(const String& method, const String& url, const HashMap<String, String>& request_headers, const ByteBuffer& request_body)
+template<typename RequestHashMapTraits>
+RefPtr<Download> Client::start_download(const String& method, const String& url, const HashMap<String, String, RequestHashMapTraits>& request_headers, ReadonlyBytes request_body)
 {
     IPC::Dictionary header_dictionary;
     for (auto& it : request_headers)
         header_dictionary.add(it.key, it.value);
 
-    i32 download_id = send_sync<Messages::ProtocolServer::StartDownload>(method, url, header_dictionary, String::copy(request_body))->download_id();
-    if (download_id < 0)
+    auto response = send_sync<Messages::ProtocolServer::StartDownload>(method, url, header_dictionary, ByteBuffer::copy(request_body));
+    auto download_id = response->download_id();
+    auto response_fd = response->response_fd().fd();
+    if (download_id < 0 || response_fd < 0)
         return nullptr;
     auto download = Download::create_from_id({}, *this, download_id);
+    download->set_download_fd({}, response_fd);
     m_downloads.set(download_id, download);
     return download;
 }
@@ -79,9 +84,8 @@ void Client::handle(const Messages::ProtocolClient::DownloadFinished& message)
 {
     RefPtr<Download> download;
     if ((download = m_downloads.get(message.download_id()).value_or(nullptr))) {
-        download->did_finish({}, message.success(), message.status_code(), message.total_size(), message.shbuf_id(), message.response_headers());
+        download->did_finish({}, message.success(), message.total_size());
     }
-    send_sync<Messages::ProtocolServer::DisownSharedBuffer>(message.shbuf_id());
     m_downloads.remove(message.download_id());
 }
 
@@ -92,6 +96,15 @@ void Client::handle(const Messages::ProtocolClient::DownloadProgress& message)
     }
 }
 
+void Client::handle(const Messages::ProtocolClient::HeadersBecameAvailable& message)
+{
+    if (auto download = const_cast<Download*>(m_downloads.get(message.download_id()).value_or(nullptr))) {
+        HashMap<String, String, CaseInsensitiveStringTraits> headers;
+        message.response_headers().for_each_entry([&](auto& name, auto& value) { headers.set(name, value); });
+        download->did_receive_headers({}, headers, message.status_code());
+    }
+}
+
 OwnPtr<Messages::ProtocolClient::CertificateRequestedResponse> Client::handle(const Messages::ProtocolClient::CertificateRequested& message)
 {
     if (auto download = const_cast<Download*>(m_downloads.get(message.download_id()).value_or(nullptr))) {
@@ -102,3 +115,6 @@ OwnPtr<Messages::ProtocolClient::CertificateRequestedResponse> Client::handle(co
 }
 
 }
+
+template RefPtr<Protocol::Download> Protocol::Client::start_download(const String& method, const String& url, const HashMap<String, String>& request_headers, ReadonlyBytes request_body);
+template RefPtr<Protocol::Download> Protocol::Client::start_download(const String& method, const String& url, const HashMap<String, String, CaseInsensitiveStringTraits>& request_headers, ReadonlyBytes request_body);

Libraries/LibProtocol/Client.h

@@ -44,7 +44,8 @@ public:
     virtual void handshake() override;
 
     bool is_supported_protocol(const String&);
-    RefPtr<Download> start_download(const String& method, const String& url, const HashMap<String, String>& request_headers = {}, const ByteBuffer& request_body = {});
+    template<typename RequestHashMapTraits = Traits<String>>
+    RefPtr<Download> start_download(const String& method, const String& url, const HashMap<String, String, RequestHashMapTraits>& request_headers = {}, ReadonlyBytes request_body = {});
 
     bool stop_download(Badge<Download>, Download&);
     bool set_certificate(Badge<Download>, Download&, String, String);
@@ -55,6 +56,7 @@ private:
     virtual void handle(const Messages::ProtocolClient::DownloadProgress&) override;
     virtual void handle(const Messages::ProtocolClient::DownloadFinished&) override;
     virtual OwnPtr<Messages::ProtocolClient::CertificateRequestedResponse> handle(const Messages::ProtocolClient::CertificateRequested&) override;
+    virtual void handle(const Messages::ProtocolClient::HeadersBecameAvailable&) override;
 
     HashMap<i32, RefPtr<Download>> m_downloads;
 };

Libraries/LibProtocol/Download.cpp

@@ -41,25 +41,81 @@ bool Download::stop()
     return m_client->stop_download({}, *this);
 }
 
-void Download::did_finish(Badge<Client>, bool success, Optional<u32> status_code, u32 total_size, i32 shbuf_id, const IPC::Dictionary& response_headers)
+void Download::stream_into(OutputStream& stream)
+{
+    ASSERT(!m_internal_stream_data);
+
+    auto notifier = Core::Notifier::construct(fd(), Core::Notifier::Read);
+
+    m_internal_stream_data = make<InternalStreamData>(fd());
+    m_internal_stream_data->read_notifier = notifier;
+
+    auto user_on_finish = move(on_finish);
+    on_finish = [this](auto success, auto total_size) {
+        m_internal_stream_data->success = success;
+        m_internal_stream_data->total_size = total_size;
+        m_internal_stream_data->download_done = true;
+    };
+
+    notifier->on_ready_to_read = [this, &stream, user_on_finish = move(user_on_finish)] {
+        constexpr size_t buffer_size = 1 * KiB;
+        static char buf[buffer_size];
+        auto nread = m_internal_stream_data->read_stream.read({ buf, buffer_size });
+        if (!stream.write_or_error({ buf, nread })) {
+            // FIXME: What do we do here?
+            TODO();
+        }
+
+        if (m_internal_stream_data->read_stream.eof() || (m_internal_stream_data->download_done && !m_internal_stream_data->success)) {
+            m_internal_stream_data->read_notifier->close();
+            user_on_finish(m_internal_stream_data->success, m_internal_stream_data->total_size);
+        } else {
+            m_internal_stream_data->read_stream.handle_any_error();
+        }
+    };
+}
+
+void Download::set_should_buffer_all_input(bool value)
+{
+    if (m_should_buffer_all_input == value)
+        return;
+
+    if (m_internal_buffered_data && !value) {
+        m_internal_buffered_data = nullptr;
+        m_should_buffer_all_input = false;
+        return;
+    }
+
+    ASSERT(!m_internal_stream_data);
+    ASSERT(!m_internal_buffered_data);
+    ASSERT(on_buffered_download_finish); // Not having this set makes no sense.
+    m_internal_buffered_data = make<InternalBufferedData>(fd());
+    m_should_buffer_all_input = true;
+
+    on_headers_received = [this](auto& headers, auto response_code) {
+        m_internal_buffered_data->response_headers = headers;
+        m_internal_buffered_data->response_code = move(response_code);
+    };
+
+    on_finish = [this](auto success, u32 total_size) {
+        auto output_buffer = m_internal_buffered_data->payload_stream.copy_into_contiguous_buffer();
+        on_buffered_download_finish(
+            success,
+            total_size,
+            m_internal_buffered_data->response_headers,
+            m_internal_buffered_data->response_code,
+            output_buffer);
+    };
+
+    stream_into(m_internal_buffered_data->payload_stream);
+}
+
+void Download::did_finish(Badge<Client>, bool success, u32 total_size)
 {
     if (!on_finish)
         return;
 
-    ReadonlyBytes payload;
-    RefPtr<SharedBuffer> shared_buffer;
-    if (success && shbuf_id != -1) {
-        shared_buffer = SharedBuffer::create_from_shbuf_id(shbuf_id);
-        payload = { shared_buffer->data<void>(), total_size };
-    }
-
-    // FIXME: It's a bit silly that we copy the response headers here just so we can move them into a HashMap with different traits.
-    HashMap<String, String, CaseInsensitiveStringTraits> caseless_response_headers;
-    response_headers.for_each_entry([&](auto& name, auto& value) {
-        caseless_response_headers.set(name, value);
-    });
-
-    on_finish(success, payload, move(shared_buffer), caseless_response_headers, status_code);
+    on_finish(success, total_size);
 }
 
 void Download::did_progress(Badge<Client>, Optional<u32> total_size, u32 downloaded_size)
@@ -68,6 +124,12 @@ void Download::did_progress(Badge<Client>, Optional<u32> total_size, u32 downloa
         on_progress(total_size, downloaded_size);
 }
 
+void Download::did_receive_headers(Badge<Client>, const HashMap<String, String, CaseInsensitiveStringTraits>& response_headers, Optional<u32> response_code)
+{
+    if (on_headers_received)
+        on_headers_received(response_headers, response_code);
+}
+
 void Download::did_request_certificates(Badge<Client>)
 {
     if (on_certificate_requested) {

Libraries/LibProtocol/Download.h

@@ -28,10 +28,13 @@
 
 #include <AK/Badge.h>
 #include <AK/ByteBuffer.h>
+#include <AK/FileStream.h>
 #include <AK/Function.h>
+#include <AK/MemoryStream.h>
 #include <AK/RefCounted.h>
 #include <AK/String.h>
 #include <AK/WeakPtr.h>
+#include <LibCore/Notifier.h>
 #include <LibIPC/Forward.h>
 
 namespace Protocol {
@@ -51,20 +54,65 @@ public:
     }
 
     int id() const { return m_download_id; }
+    int fd() const { return m_fd; }
     bool stop();
 
-    Function<void(bool success, ReadonlyBytes payload, RefPtr<SharedBuffer> payload_storage, const HashMap<String, String, CaseInsensitiveStringTraits>& response_headers, Optional<u32> status_code)> on_finish;
+    void stream_into(OutputStream&);
+
+    bool should_buffer_all_input() const { return m_should_buffer_all_input; }
+    /// Note: Will override `on_finish', and `on_headers_received', and expects `on_buffered_download_finish' to be set!
+    void set_should_buffer_all_input(bool);
+
+    /// Note: Must be set before `set_should_buffer_all_input(true)`.
+    Function<void(bool success, u32 total_size, const HashMap<String, String, CaseInsensitiveStringTraits>& response_headers, Optional<u32> response_code, ReadonlyBytes payload)> on_buffered_download_finish;
+    Function<void(bool success, u32 total_size)> on_finish;
     Function<void(Optional<u32> total_size, u32 downloaded_size)> on_progress;
+    Function<void(const HashMap<String, String, CaseInsensitiveStringTraits>& response_headers, Optional<u32> response_code)> on_headers_received;
     Function<CertificateAndKey()> on_certificate_requested;
 
-    void did_finish(Badge<Client>, bool success, Optional<u32> status_code, u32 total_size, i32 shbuf_id, const IPC::Dictionary& response_headers);
+    void did_finish(Badge<Client>, bool success, u32 total_size);
     void did_progress(Badge<Client>, Optional<u32> total_size, u32 downloaded_size);
+    void did_receive_headers(Badge<Client>, const HashMap<String, String, CaseInsensitiveStringTraits>& response_headers, Optional<u32> response_code);
     void did_request_certificates(Badge<Client>);
 
+    RefPtr<Core::Notifier>& write_notifier(Badge<Client>) { return m_write_notifier; }
+    void set_download_fd(Badge<Client>, int fd) { m_fd = fd; }
+
 private:
     explicit Download(Client&, i32 download_id);
     WeakPtr<Client> m_client;
     int m_download_id { -1 };
+    RefPtr<Core::Notifier> m_write_notifier;
+    int m_fd { -1 };
+    bool m_should_buffer_all_input { false };
+
+    struct InternalBufferedData {
+        InternalBufferedData(int fd)
+            : read_stream(fd)
+        {
+        }
+
+        InputFileStream read_stream;
+        DuplexMemoryStream payload_stream;
+        HashMap<String, String, CaseInsensitiveStringTraits> response_headers;
+        Optional<u32> response_code;
+    };
+
+    struct InternalStreamData {
+        InternalStreamData(int fd)
+            : read_stream(fd)
+        {
+        }
+
+        InputFileStream read_stream;
+        RefPtr<Core::Notifier> read_notifier;
+        bool success;
+        u32 total_size { 0 };
+        bool download_done { false };
+    };
+
+    OwnPtr<InternalBufferedData> m_internal_buffered_data;
+    OwnPtr<InternalStreamData> m_internal_stream_data;
 };
 
 }