LibCore: Use length-checking stream reads and writes for SOCKS5

Userland/Libraries/LibCore/SOCKSProxyClient.cpp

@@ -39,27 +39,52 @@ struct [[gnu::packed]] Socks5VersionIdentifierAndMethodSelectionMessage {
     u8 methods[1];
 };
 
+template<>
+struct AK::Traits<Socks5VersionIdentifierAndMethodSelectionMessage> : public AK::GenericTraits<Socks5VersionIdentifierAndMethodSelectionMessage> {
+    static constexpr bool is_trivially_serializable() { return true; }
+};
+
 struct [[gnu::packed]] Socks5InitialResponse {
     u8 version_identifier;
     u8 method;
 };
 
+template<>
+struct AK::Traits<Socks5InitialResponse> : public AK::GenericTraits<Socks5InitialResponse> {
+    static constexpr bool is_trivially_serializable() { return true; }
+};
+
 struct [[gnu::packed]] Socks5ConnectRequestHeader {
     u8 version_identifier;
     u8 command;
     u8 reserved;
 };
 
+template<>
+struct AK::Traits<Socks5ConnectRequestHeader> : public AK::GenericTraits<Socks5ConnectRequestHeader> {
+    static constexpr bool is_trivially_serializable() { return true; }
+};
+
 struct [[gnu::packed]] Socks5ConnectRequestTrailer {
     u16 port;
 };
 
+template<>
+struct AK::Traits<Socks5ConnectRequestTrailer> : public AK::GenericTraits<Socks5ConnectRequestTrailer> {
+    static constexpr bool is_trivially_serializable() { return true; }
+};
+
 struct [[gnu::packed]] Socks5ConnectResponseHeader {
     u8 version_identifier;
     u8 status;
     u8 reserved;
 };
 
+template<>
+struct AK::Traits<Socks5ConnectResponseHeader> : public AK::GenericTraits<Socks5ConnectResponseHeader> {
+    static constexpr bool is_trivially_serializable() { return true; }
+};
+
 struct [[gnu::packed]] Socks5ConnectResponseTrailer {
     u8 bind_port;
 };
@@ -69,6 +94,11 @@ struct [[gnu::packed]] Socks5UsernamePasswordResponse {
     u8 status;
 };
 
+template<>
+struct AK::Traits<Socks5UsernamePasswordResponse> : public AK::GenericTraits<Socks5UsernamePasswordResponse> {
+    static constexpr bool is_trivially_serializable() { return true; }
+};
+
 namespace {
 StringView reply_response_name(Reply reply)
 {
@@ -102,16 +132,9 @@ ErrorOr<void> send_version_identifier_and_method_selection_message(Core::Socket&
         .method_count = 1,
         .methods = { to_underlying(method) },
     };
-    // FIXME: This should write the entire span.
+    TRY(socket.write_value(message));
-    auto size = TRY(socket.write_some({ &message, sizeof(message) }));
-    if (size != sizeof(message))
-        return Error::from_string_literal("SOCKS negotiation failed: Failed to send version identifier and method selection message");
 
-    Socks5InitialResponse response;
+    auto response = TRY(socket.read_value<Socks5InitialResponse>());
-    // FIXME: This should read the entire span.
-    size = TRY(socket.read_some({ &response, sizeof(response) })).size();
-    if (size != sizeof(response))
-        return Error::from_string_literal("SOCKS negotiation failed: Failed to receive initial response");
 
     if (response.version_identifier != to_underlying(version))
         return Error::from_string_literal("SOCKS negotiation failed: Invalid version identifier");
@@ -135,22 +158,15 @@ ErrorOr<Reply> send_connect_request_message(Core::Socket& socket, Core::SOCKSPro
         .port = htons(port),
     };
 
-    // FIXME: This should write the entire span.
+    TRY(stream.write_value(header));
-    auto size = TRY(stream.write_some({ &header, sizeof(header) }));
-    if (size != sizeof(header))
-        return Error::from_string_literal("SOCKS negotiation failed: Failed to send connect request header");
 
     TRY(target.visit(
         [&](DeprecatedString const& hostname) -> ErrorOr<void> {
             u8 address_data[2];
             address_data[0] = to_underlying(AddressType::DomainName);
             address_data[1] = hostname.length();
-            // FIXME: This should write the entire span.
+            TRY(stream.write_until_depleted({ address_data, sizeof(address_data) }));
-            auto size = TRY(stream.write_some({ address_data, sizeof(address_data) }));
+            TRY(stream.write_until_depleted({ hostname.characters(), hostname.length() }));
-            if (size != array_size(address_data))
-                return Error::from_string_literal("SOCKS negotiation failed: Failed to send connect request address data");
-            // FIXME: This should write the entire span.
-            TRY(stream.write_some({ hostname.characters(), hostname.length() }));
             return {};
         },
         [&](u32 ipv4) -> ErrorOr<void> {
@@ -158,62 +174,33 @@ ErrorOr<Reply> send_connect_request_message(Core::Socket& socket, Core::SOCKSPro
             address_data[0] = to_underlying(AddressType::IPV4);
             u32 network_ordered_ipv4 = NetworkOrdered<u32>(ipv4);
             memcpy(address_data + 1, &network_ordered_ipv4, sizeof(network_ordered_ipv4));
-            // FIXME: This should write the entire span.
+            TRY(stream.write_until_depleted({ address_data, sizeof(address_data) }));
-            auto size = TRY(stream.write_some({ address_data, sizeof(address_data) }));
-            if (size != array_size(address_data))
-                return Error::from_string_literal("SOCKS negotiation failed: Failed to send connect request address data");
             return {};
         }));
 
-    // FIXME: This should write the entire span.
+    TRY(stream.write_value(trailer));
-    size = TRY(stream.write_some({ &trailer, sizeof(trailer) }));
-    if (size != sizeof(trailer))
-        return Error::from_string_literal("SOCKS negotiation failed: Failed to send connect request trailer");
 
     auto buffer = TRY(ByteBuffer::create_uninitialized(stream.used_buffer_size()));
     TRY(stream.read_until_filled(buffer.bytes()));
+    TRY(socket.write_until_depleted(buffer));
 
-    // FIXME: This should write the entire span.
+    auto response_header = TRY(socket.read_value<Socks5ConnectResponseHeader>());
-    size = TRY(socket.write_some({ buffer.data(), buffer.size() }));
-    if (size != buffer.size())
-        return Error::from_string_literal("SOCKS negotiation failed: Failed to send connect request");
-
-    Socks5ConnectResponseHeader response_header;
-    // FIXME: This should read the entire span.
-    size = TRY(socket.read_some({ &response_header, sizeof(response_header) })).size();
-    if (size != sizeof(response_header))
-        return Error::from_string_literal("SOCKS negotiation failed: Failed to receive connect response header");
 
     if (response_header.version_identifier != to_underlying(version))
         return Error::from_string_literal("SOCKS negotiation failed: Invalid version identifier");
 
-    u8 response_address_type;
+    auto response_address_type = TRY(socket.read_value<u8>());
-    // FIXME: This should read the entire span.
-    size = TRY(socket.read_some({ &response_address_type, sizeof(response_address_type) })).size();
-    if (size != sizeof(response_address_type))
-        return Error::from_string_literal("SOCKS negotiation failed: Failed to receive connect response address type");
 
     switch (AddressType(response_address_type)) {
     case AddressType::IPV4: {
         u8 response_address_data[4];
-        // FIXME: This should read the entire span.
+        TRY(socket.read_until_filled({ response_address_data, sizeof(response_address_data) }));
-        size = TRY(socket.read_some({ response_address_data, sizeof(response_address_data) })).size();
-        if (size != sizeof(response_address_data))
-            return Error::from_string_literal("SOCKS negotiation failed: Failed to receive connect response address data");
         break;
     }
     case AddressType::DomainName: {
-        u8 response_address_length;
+        auto response_address_length = TRY(socket.read_value<u8>());
-        // FIXME: This should read the entire span.
+        auto buffer = TRY(ByteBuffer::create_uninitialized(response_address_length));
-        size = TRY(socket.read_some({ &response_address_length, sizeof(response_address_length) })).size();
+        TRY(socket.read_until_filled(buffer));
-        if (size != sizeof(response_address_length))
-            return Error::from_string_literal("SOCKS negotiation failed: Failed to receive connect response address length");
-        ByteBuffer buffer;
-        buffer.resize(response_address_length);
-        // FIXME: This should read the entire span.
-        size = TRY(socket.read_some(buffer)).size();
-        if (size != response_address_length)
-            return Error::from_string_literal("SOCKS negotiation failed: Failed to receive connect response address data");
         break;
     }
     case AddressType::IPV6:
@@ -221,11 +208,7 @@ ErrorOr<Reply> send_connect_request_message(Core::Socket& socket, Core::SOCKSPro
         return Error::from_string_literal("SOCKS negotiation failed: Invalid connect response address type");
     }
 
-    u16 bound_port;
+    [[maybe_unused]] auto bound_port = TRY(socket.read_value<u16>());
-    // FIXME: This should read the entire span.
-    size = TRY(socket.read_some({ &bound_port, sizeof(bound_port) })).size();
-    if (size != sizeof(bound_port))
-        return Error::from_string_literal("SOCKS negotiation failed: Failed to receive connect response bound port");
 
     return Reply(response_header.status);
 }
@@ -235,46 +218,24 @@ ErrorOr<u8> send_username_password_authentication_message(Core::Socket& socket,
     AllocatingMemoryStream stream;
 
     u8 version = 0x01;
-    // FIXME: This should write the entire span.
+    TRY(stream.write_value(version));
-    auto size = TRY(stream.write_some({ &version, sizeof(version) }));
-    if (size != sizeof(version))
-        return Error::from_string_literal("SOCKS negotiation failed: Failed to send username/password authentication message");
 
     u8 username_length = auth_data.username.length();
-    // FIXME: This should write the entire span.
+    TRY(stream.write_value(username_length));
-    size = TRY(stream.write_some({ &username_length, sizeof(username_length) }));
-    if (size != sizeof(username_length))
-        return Error::from_string_literal("SOCKS negotiation failed: Failed to send username/password authentication message");
 
-    // FIXME: This should write the entire span.
+    TRY(stream.write_until_depleted({ auth_data.username.characters(), auth_data.username.length() }));
-    size = TRY(stream.write_some({ auth_data.username.characters(), auth_data.username.length() }));
-    if (size != auth_data.username.length())
-        return Error::from_string_literal("SOCKS negotiation failed: Failed to send username/password authentication message");
 
     u8 password_length = auth_data.password.length();
-    // FIXME: This should write the entire span.
+    TRY(stream.write_value(password_length));
-    size = TRY(stream.write_some({ &password_length, sizeof(password_length) }));
-    if (size != sizeof(password_length))
-        return Error::from_string_literal("SOCKS negotiation failed: Failed to send username/password authentication message");
 
-    // FIXME: This should write the entire span.
+    TRY(stream.write_until_depleted({ auth_data.password.characters(), auth_data.password.length() }));
-    size = TRY(stream.write_some({ auth_data.password.characters(), auth_data.password.length() }));
-    if (size != auth_data.password.length())
-        return Error::from_string_literal("SOCKS negotiation failed: Failed to send username/password authentication message");
 
     auto buffer = TRY(ByteBuffer::create_uninitialized(stream.used_buffer_size()));
     TRY(stream.read_until_filled(buffer.bytes()));
 
-    // FIXME: This should write the entire span.
+    TRY(socket.write_until_depleted(buffer));
-    size = TRY(socket.write_some(buffer));
-    if (size != buffer.size())
-        return Error::from_string_literal("SOCKS negotiation failed: Failed to send username/password authentication message");
 
-    Socks5UsernamePasswordResponse response;
+    auto response = TRY(socket.read_value<Socks5UsernamePasswordResponse>());
-    // FIXME: This should read the entire span.
-    size = TRY(socket.read_some({ &response, sizeof(response) })).size();
-    if (size != sizeof(response))
-        return Error::from_string_literal("SOCKS negotiation failed: Failed to receive username/password authentication response");
 
     if (response.version_identifier != version)
         return Error::from_string_literal("SOCKS negotiation failed: Invalid version identifier");