From 17d3592cabffe3bc413df45cb4b166b55d600f44 Mon Sep 17 00:00:00 2001
From: sin-ack <sin-ack@users.noreply.github.com>
Date: Sun, 6 Feb 2022 16:40:12 +0000
Subject: [PATCH] LibCore: Add a timeout option to UDPSocket::connect

This allows us to set a timeout during connection and during receive and
send operations. I didn't add this to the other connect calls as it's
not used anywhere else for the time being.
---
 Userland/Libraries/LibCore/Stream.cpp | 15 ++++++++++++---
 Userland/Libraries/LibCore/Stream.h   |  6 ++++--
 2 files changed, 16 insertions(+), 5 deletions(-)

diff --git a/Userland/Libraries/LibCore/Stream.cpp b/Userland/Libraries/LibCore/Stream.cpp
index 21988042b3..0498ab3f46 100644
--- a/Userland/Libraries/LibCore/Stream.cpp
+++ b/Userland/Libraries/LibCore/Stream.cpp
@@ -390,6 +390,12 @@ ErrorOr<void> PosixSocketHelper::set_close_on_exec(bool enabled)
     return {};
 }
 
+ErrorOr<void> PosixSocketHelper::set_receive_timeout(Time timeout)
+{
+    auto timeout_spec = timeout.to_timespec();
+    return System::setsockopt(m_fd, SOL_SOCKET, SO_RCVTIMEO, &timeout_spec, sizeof(timeout_spec));
+}
+
 void PosixSocketHelper::setup_notifier()
 {
     if (!m_notifier)
@@ -438,18 +444,21 @@ ErrorOr<size_t> PosixSocketHelper::pending_bytes() const
     return static_cast<size_t>(value);
 }
 
-ErrorOr<NonnullOwnPtr<UDPSocket>> UDPSocket::connect(String const& host, u16 port)
+ErrorOr<NonnullOwnPtr<UDPSocket>> UDPSocket::connect(String const& host, u16 port, Optional<Time> timeout)
 {
     auto ip_address = TRY(resolve_host(host, SocketType::Datagram));
-    return connect(SocketAddress { ip_address, port });
+    return connect(SocketAddress { ip_address, port }, timeout);
 }
 
-ErrorOr<NonnullOwnPtr<UDPSocket>> UDPSocket::connect(SocketAddress const& address)
+ErrorOr<NonnullOwnPtr<UDPSocket>> UDPSocket::connect(SocketAddress const& address, Optional<Time> timeout)
 {
     auto socket = TRY(adopt_nonnull_own_or_enomem(new (nothrow) UDPSocket()));
 
     auto fd = TRY(create_fd(SocketDomain::Inet, SocketType::Datagram));
     socket->m_helper.set_fd(fd);
+    if (timeout.has_value()) {
+        TRY(socket->m_helper.set_receive_timeout(timeout.value()));
+    }
 
     TRY(connect_inet(fd, address));
 
diff --git a/Userland/Libraries/LibCore/Stream.h b/Userland/Libraries/LibCore/Stream.h
index d8bcb07dd0..3757115425 100644
--- a/Userland/Libraries/LibCore/Stream.h
+++ b/Userland/Libraries/LibCore/Stream.h
@@ -14,6 +14,7 @@
 #include <AK/Result.h>
 #include <AK/Span.h>
 #include <AK/String.h>
+#include <AK/Time.h>
 #include <AK/Variant.h>
 #include <LibCore/Notifier.h>
 #include <LibCore/SocketAddress.h>
@@ -247,6 +248,7 @@ public:
 
     ErrorOr<void> set_blocking(bool enabled);
     ErrorOr<void> set_close_on_exec(bool enabled);
+    ErrorOr<void> set_receive_timeout(Time timeout);
 
     void setup_notifier();
     RefPtr<Core::Notifier> notifier() { return m_notifier; }
@@ -321,8 +323,8 @@ private:
 
 class UDPSocket final : public Socket {
 public:
-    static ErrorOr<NonnullOwnPtr<UDPSocket>> connect(String const& host, u16 port);
-    static ErrorOr<NonnullOwnPtr<UDPSocket>> connect(SocketAddress const& address);
+    static ErrorOr<NonnullOwnPtr<UDPSocket>> connect(String const& host, u16 port, Optional<Time> timeout = {});
+    static ErrorOr<NonnullOwnPtr<UDPSocket>> connect(SocketAddress const& address, Optional<Time> timeout = {});
 
     UDPSocket(UDPSocket&& other)
         : Socket(static_cast<Socket&&>(other))