Просмотр исходного кода

Kernel: Make UDPSocket::create() API OOM safe

Brian Gianforcaro 4 лет назад
Родитель
Сommit
2e34714ba1
3 измененных файлов с 13 добавлено и 5 удалено
  1. 6 2
      Kernel/Net/IPv4Socket.cpp
  2. 5 2
      Kernel/Net/UDPSocket.cpp
  3. 2 1
      Kernel/Net/UDPSocket.h

+ 6 - 2
Kernel/Net/IPv4Socket.cpp

@@ -42,8 +42,12 @@ KResultOr<NonnullRefPtr<Socket>> IPv4Socket::create(int type, int protocol)
             return tcp_socket.error();
         return tcp_socket.release_value();
     }
-    if (type == SOCK_DGRAM)
-        return UDPSocket::create(protocol);
+    if (type == SOCK_DGRAM) {
+        auto udp_socket = UDPSocket::create(protocol);
+        if (udp_socket.is_error())
+            return udp_socket.error();
+        return udp_socket.release_value();
+    }
     if (type == SOCK_RAW) {
         auto raw_socket = adopt_ref_if_nonnull(new IPv4Socket(type, protocol));
         if (raw_socket)

+ 5 - 2
Kernel/Net/UDPSocket.cpp

@@ -54,9 +54,12 @@ UDPSocket::~UDPSocket()
     sockets_by_port().resource().remove(local_port());
 }
 
-NonnullRefPtr<UDPSocket> UDPSocket::create(int protocol)
+KResultOr<NonnullRefPtr<UDPSocket>> UDPSocket::create(int protocol)
 {
-    return adopt_ref(*new UDPSocket(protocol));
+    auto socket = adopt_ref_if_nonnull(new UDPSocket(protocol));
+    if (socket)
+        return socket.release_nonnull();
+    return ENOMEM;
 }
 
 KResultOr<size_t> UDPSocket::protocol_receive(ReadonlyBytes raw_ipv4_packet, UserOrKernelBuffer& buffer, size_t buffer_size, [[maybe_unused]] int flags)

+ 2 - 1
Kernel/Net/UDPSocket.h

@@ -6,13 +6,14 @@
 
 #pragma once
 
+#include <Kernel/KResult.h>
 #include <Kernel/Net/IPv4Socket.h>
 
 namespace Kernel {
 
 class UDPSocket final : public IPv4Socket {
 public:
-    static NonnullRefPtr<UDPSocket> create(int protocol);
+    static KResultOr<NonnullRefPtr<UDPSocket>> create(int protocol);
     virtual ~UDPSocket() override;
 
     static SocketHandle<UDPSocket> from_port(u16);