make SocketAddress useable after receive timeout (dotnet#102142)

* make SocketAddress useable after receive timeout * Update src/libraries/System.Net.Sockets/tests/FunctionalTests/SendReceive/SendReceiveMisc.cs Co-authored-by: Stephen Toub <stoub@microsoft.com> --------- Co-authored-by: Stephen Toub <stoub@microsoft.com>
sbomer · May 16, 2024 · ffa1bea · ffa1bea
1 parent ffad857
commit ffa1bea
Show file tree

Hide file tree

Showing 2 changed files with 37 additions and 2 deletions.
diff --git a/src/libraries/System.Net.Sockets/src/System/Net/Sockets/Socket.cs b/src/libraries/System.Net.Sockets/src/System/Net/Sockets/Socket.cs
@@ -1909,8 +1909,6 @@ public int ReceiveFrom(Span<byte> buffer, SocketFlags socketFlags, SocketAddress
 
             int bytesTransferred;
             SocketError errorCode = SocketPal.ReceiveFrom(_handle, buffer, socketFlags, receivedAddress.Buffer, out int socketAddressSize, out bytesTransferred);
-            receivedAddress.Size = socketAddressSize;
-
             UpdateReceiveSocketErrorForDisposed(ref errorCode, bytesTransferred);
             // If the native call fails we'll throw a SocketException.
             if (errorCode != SocketError.Success)
@@ -1927,6 +1925,7 @@ public int ReceiveFrom(Span<byte> buffer, SocketFlags socketFlags, SocketAddress
                 if (SocketType == SocketType.Dgram) SocketsTelemetry.Log.DatagramReceived();
             }
 
+            receivedAddress.Size = socketAddressSize;
             return bytesTransferred;
         }
 

diff --git a/src/libraries/System.Net.Sockets/tests/FunctionalTests/SendReceive/SendReceiveMisc.cs b/src/libraries/System.Net.Sockets/tests/FunctionalTests/SendReceive/SendReceiveMisc.cs
@@ -283,5 +283,41 @@ public async Task Socket_ReceiveFlags_Success()
                 Assert.Equal(0, receiver.Available);
             }
         }
+
+        [Theory]
+        [InlineData(true)]
+        [InlineData(false)]
+        public void ReceiveFrom_MultipleRounds_Success(bool async)
+        {
+            using (var socket = new Socket(AddressFamily.InterNetwork, SocketType.Dgram, ProtocolType.Udp))
+            {
+                socket.ReceiveTimeout = 100;
+                socket.BindToAnonymousPort(IPAddress.Loopback);
+
+                var address = new SocketAddress(AddressFamily.InterNetwork);
+                var buffer = new byte[100];
+                int receivedLength;
+
+                for (int i = 0; i < 5; i++)
+                {
+                    try
+                    {
+                        if (async)
+                        {
+                            using var cts = new CancellationTokenSource();
+                            cts.CancelAfter(100);
+                            receivedLength = socket.ReceiveFromAsync(buffer, SocketFlags.None, address, cts.Token).AsTask().GetAwaiter().GetResult();
+                        }
+                        else
+                        {
+                            receivedLength = socket.ReceiveFrom(buffer, SocketFlags.None, address);
+                        }
+                        Assert.Equal(0, receivedLength);
+                    }
+                    catch (SocketException ex) when (ex.SocketErrorCode == SocketError.TimedOut) { }
+                    catch (OperationCanceledException) { }
+                }
+            }
+        }
     }
 }