Make wxSocket::Peek() work with UDP too

Fix wxSocketImpl::RecvDgram() when caller requests fewer bytes than
are present in datagram, as we must always read all the available bytes
when using UDP -- they are not buffered by the OS and so won't be
provided by the next Read() call.

See #23594, #23604.

(cherry picked from commit b9d0541f9ac4d4c356b6ab03772c1fe10a2fd1bc)

docs/changes.txt

@@ -239,6 +239,7 @@ All:
 
 - Add move ctor and assignment operator to wxString (Pavel Tyunin, #23224).
 - Enable large file support in Unix CMake builds (Maarten Bent, #22750).
+- Make wxSocket::Peek() work with UDP too (Brian Nixon, #23594, #23604).
 - Update Georgian translations (NorwayFun, #22673).
 - Don't use invalid iterator in wxString in UTF-8 build (Ian Day, #23305).
 - Fix wrong description for some languages (Ulrich Telle, #23419).

src/common/socket.cpp

@@ -689,7 +689,15 @@ int wxSocketImpl::RecvDgram(void *buffer, int size)
                                   0, &from.addr, &fromlen) );
 
     if ( ret == SOCKET_ERROR )
-        return SOCKET_ERROR;
+    {
+#ifdef __WINDOWS__
+        if ( WSAGetLastError() == WSAEMSGSIZE )
+            ret = size;
+        else
+#endif // __WINDOWS__
+            return SOCKET_ERROR;
+    }
+
 
     m_peer = wxSockAddressImpl(from.addr, fromlen);
     if ( !m_peer.IsOk() )
@@ -1144,9 +1152,35 @@ wxSocketBase& wxSocketBase::Peek(void* buffer, wxUint32 nbytes)
     // Peek() should never block
     wxSocketWaitModeChanger changeFlags(this, wxSOCKET_NOWAIT);
 
-    m_lcount = DoRead(buffer, nbytes);
+    // Guard against data loss when reading fewer bytes
+    // than are present in a received datagram
+    void* readbuf;
+    wxUint32 readbytes;
+    const wxUint32 DGRAM_MIN_READ = 65536;  // 64K is enough for UDP
+    std::vector<unsigned char> peekbuf;
+    bool usePeekbuf = !m_impl->m_stream && nbytes < DGRAM_MIN_READ;
+    if ( usePeekbuf )
+    {
+        // Allocate our own buffer
+        peekbuf.resize(DGRAM_MIN_READ);
+        readbuf = &peekbuf[0];
+        readbytes = DGRAM_MIN_READ;
+    }
+    else
+    {
+        // Use the caller-supplied buffer directly
+        readbuf = buffer;
+        readbytes = nbytes;
+    }
 
-    Pushback(buffer, m_lcount);
+    wxUint32 lcount = DoRead(readbuf, readbytes);
+
+    Pushback(readbuf, lcount);
+
+    if ( usePeekbuf )
+        lcount = GetPushback(buffer, nbytes, true);
+
+    m_lcount = lcount;
 
     return *this;
 }

tests/net/socket.cpp

@@ -301,4 +301,60 @@ void SocketTestCase::UrlTest()
     CPPUNIT_ASSERT_EQUAL( wxSTREAM_EOF, in->Read(out).GetLastError() );
 }
 
+TEST_CASE("wxDatagramSocket::ShortRead", "[socket][dgram]")
+{
+    // Check that reading fewer bytes than are present in a
+    // datagram does not leave the socket in an error state
+
+    wxIPV4address addr;
+    addr.LocalHost();
+    addr.Service(19898);// Arbitrary port number
+    wxDatagramSocket sock(addr);
+
+    // Send ourselves a datagram
+    unsigned int sendbuf[4] = {1, 2, 3, 4};
+    sock.SendTo(addr, sendbuf, sizeof(sendbuf));
+
+    // Read less than we know we sent
+    unsigned int recvbuf[1] = {0};
+    sock.Read(recvbuf, sizeof(recvbuf));
+    CHECK(!sock.Error());
+    CHECK(sock.LastReadCount() == sizeof(recvbuf));
+    CHECK(recvbuf[0] == sendbuf[0]);
+}
+
+TEST_CASE("wxDatagramSocket::ShortPeek", "[socket][dgram]")
+{
+    // Check that peeking fewer bytes than are present in a datagram
+    // does not lose the rest of the data in that datagram (#23594)
+
+    wxIPV4address addr;
+    addr.LocalHost();
+    addr.Service(27384);// Arbitrary port number
+    wxDatagramSocket sock(addr);
+
+    // Send ourselves 2 datagrams
+    unsigned int sendbuf1[2] = {1, 2};
+    sock.SendTo(addr, sendbuf1, sizeof(sendbuf1));
+    unsigned int sendbuf2[2] = {3, 4};
+    sock.SendTo(addr, sendbuf2, sizeof(sendbuf2));
+
+    long timeout_s = 1;
+    if ( !sock.WaitForRead(timeout_s) )
+        return;
+
+    // Peek the first word
+    unsigned int peekbuf[1] = {0};
+    sock.Peek(peekbuf, sizeof(peekbuf));
+    CHECK(sock.LastCount() == sizeof(peekbuf));
+    CHECK(peekbuf[0] == sendbuf1[0]);
+
+    // Read the whole of the first datagram
+    unsigned int recvbuf[2] = {0};
+    sock.Read(recvbuf, sizeof(recvbuf));
+    CHECK(sock.LastReadCount() == sizeof(recvbuf));
+    CHECK(recvbuf[0] == sendbuf1[0]);
+    CHECK(recvbuf[1] == sendbuf1[1]);
+}
+
 #endif // wxUSE_SOCKETS