From 3da0c072f40d1a93a295e07766355f9278941e13 Mon Sep 17 00:00:00 2001
From: sin-ack <sin-ack@users.noreply.github.com>
Date: Sat, 11 Sep 2021 19:38:05 +0000
Subject: [PATCH] Kernel: Return the correct result for FIONREAD on datagram
 sockets

Before this commit, we only checked the receive buffer on the socket,
which is unused on datagram streams. Now we return the actual size of
the datagram without the protocol headers, which required the protocol
to tell us what the size of the payload is.
---
 Kernel/Net/IPv4Socket.cpp | 10 +++++++++-
 Kernel/Net/IPv4Socket.h   |  1 +
 Kernel/Net/TCPSocket.cpp  |  7 +++++++
 Kernel/Net/TCPSocket.h    |  1 +
 Kernel/Net/UDPSocket.cpp  |  7 +++++++
 Kernel/Net/UDPSocket.h    |  1 +
 6 files changed, 26 insertions(+), 1 deletion(-)

diff --git a/Kernel/Net/IPv4Socket.cpp b/Kernel/Net/IPv4Socket.cpp
index fb0a19052a..17d1e07300 100644
--- a/Kernel/Net/IPv4Socket.cpp
+++ b/Kernel/Net/IPv4Socket.cpp
@@ -774,7 +774,15 @@ ErrorOr<void> IPv4Socket::ioctl(OpenFileDescription&, unsigned request, Userspac
         return ioctl_arp();
 
     case FIONREAD: {
-        int readable = m_receive_buffer->immediately_readable();
+        int readable = 0;
+        if (buffer_mode() == BufferMode::Bytes) {
+            readable = static_cast<int>(m_receive_buffer->immediately_readable());
+        } else {
+            if (m_receive_queue.size() != 0u) {
+                readable = static_cast<int>(TRY(protocol_size(m_receive_queue.first().data->bytes())));
+            }
+        }
+
         return copy_to_user(static_ptr_cast<int*>(arg), &readable);
     }
     }
diff --git a/Kernel/Net/IPv4Socket.h b/Kernel/Net/IPv4Socket.h
index 6f1addbc2f..b9bcbf6821 100644
--- a/Kernel/Net/IPv4Socket.h
+++ b/Kernel/Net/IPv4Socket.h
@@ -84,6 +84,7 @@ protected:
     virtual ErrorOr<size_t> protocol_send(const UserOrKernelBuffer&, size_t) { return ENOTIMPL; }
     virtual ErrorOr<void> protocol_connect(OpenFileDescription&, ShouldBlock) { return {}; }
     virtual ErrorOr<u16> protocol_allocate_local_port() { return ENOPROTOOPT; }
+    virtual ErrorOr<size_t> protocol_size(ReadonlyBytes /* raw_ipv4_packet */) { return ENOTIMPL; }
     virtual bool protocol_is_disconnected() const { return false; }
 
     virtual void shut_down_for_reading() override;
diff --git a/Kernel/Net/TCPSocket.cpp b/Kernel/Net/TCPSocket.cpp
index b3091b85b5..fa7eb8cc10 100644
--- a/Kernel/Net/TCPSocket.cpp
+++ b/Kernel/Net/TCPSocket.cpp
@@ -161,6 +161,13 @@ ErrorOr<NonnullRefPtr<TCPSocket>> TCPSocket::try_create(int protocol, NonnullOwn
     return adopt_nonnull_ref_or_enomem(new (nothrow) TCPSocket(protocol, move(receive_buffer), move(scratch_buffer)));
 }
 
+ErrorOr<size_t> TCPSocket::protocol_size(ReadonlyBytes raw_ipv4_packet)
+{
+    auto& ipv4_packet = *reinterpret_cast<const IPv4Packet*>(raw_ipv4_packet.data());
+    auto& tcp_packet = *static_cast<const TCPPacket*>(ipv4_packet.payload());
+    return raw_ipv4_packet.size() - sizeof(IPv4Packet) - tcp_packet.header_size();
+}
+
 ErrorOr<size_t> TCPSocket::protocol_receive(ReadonlyBytes raw_ipv4_packet, UserOrKernelBuffer& buffer, size_t buffer_size, [[maybe_unused]] int flags)
 {
     auto& ipv4_packet = *reinterpret_cast<const IPv4Packet*>(raw_ipv4_packet.data());
diff --git a/Kernel/Net/TCPSocket.h b/Kernel/Net/TCPSocket.h
index b8c60c3fd6..4cf0f32296 100644
--- a/Kernel/Net/TCPSocket.h
+++ b/Kernel/Net/TCPSocket.h
@@ -174,6 +174,7 @@ private:
     virtual ErrorOr<size_t> protocol_send(const UserOrKernelBuffer&, size_t) override;
     virtual ErrorOr<void> protocol_connect(OpenFileDescription&, ShouldBlock) override;
     virtual ErrorOr<u16> protocol_allocate_local_port() override;
+    virtual ErrorOr<size_t> protocol_size(ReadonlyBytes raw_ipv4_packet) override;
     virtual bool protocol_is_disconnected() const override;
     virtual ErrorOr<void> protocol_bind() override;
     virtual ErrorOr<void> protocol_listen(bool did_allocate_port) override;
diff --git a/Kernel/Net/UDPSocket.cpp b/Kernel/Net/UDPSocket.cpp
index 0a6b4eb7f8..691cc5dd49 100644
--- a/Kernel/Net/UDPSocket.cpp
+++ b/Kernel/Net/UDPSocket.cpp
@@ -59,6 +59,13 @@ ErrorOr<NonnullRefPtr<UDPSocket>> UDPSocket::try_create(int protocol, NonnullOwn
     return adopt_nonnull_ref_or_enomem(new (nothrow) UDPSocket(protocol, move(receive_buffer)));
 }
 
+ErrorOr<size_t> UDPSocket::protocol_size(ReadonlyBytes raw_ipv4_packet)
+{
+    auto& ipv4_packet = *(const IPv4Packet*)(raw_ipv4_packet.data());
+    auto& udp_packet = *static_cast<const UDPPacket*>(ipv4_packet.payload());
+    return udp_packet.length() - sizeof(UDPPacket);
+}
+
 ErrorOr<size_t> UDPSocket::protocol_receive(ReadonlyBytes raw_ipv4_packet, UserOrKernelBuffer& buffer, size_t buffer_size, [[maybe_unused]] int flags)
 {
     auto& ipv4_packet = *(const IPv4Packet*)(raw_ipv4_packet.data());
diff --git a/Kernel/Net/UDPSocket.h b/Kernel/Net/UDPSocket.h
index 74845b7a27..0678c207be 100644
--- a/Kernel/Net/UDPSocket.h
+++ b/Kernel/Net/UDPSocket.h
@@ -27,6 +27,7 @@ private:
 
     virtual ErrorOr<size_t> protocol_receive(ReadonlyBytes raw_ipv4_packet, UserOrKernelBuffer& buffer, size_t buffer_size, int flags) override;
     virtual ErrorOr<size_t> protocol_send(const UserOrKernelBuffer&, size_t) override;
+    virtual ErrorOr<size_t> protocol_size(ReadonlyBytes raw_ipv4_packet) override;
     virtual ErrorOr<void> protocol_connect(OpenFileDescription&, ShouldBlock) override;
     virtual ErrorOr<u16> protocol_allocate_local_port() override;
     virtual ErrorOr<void> protocol_bind() override;