CSHARP-5618: CSOT: Race condition in ExclusiveConnectionPool could lead to ArgumentOutOfRangeException (#1712)

Author: sanych-sun

src/MongoDB.Driver/OperationContext.cs

@@ -17,9 +17,7 @@
 using System.Diagnostics;
 using System.Threading;
 using System.Threading.Tasks;
-#if !NET6_0_OR_GREATER
 using MongoDB.Driver.Core.Misc;
-#endif
 
 namespace MongoDB.Driver
 {
@@ -36,7 +34,7 @@ public OperationContext(TimeSpan timeout, CancellationToken cancellationToken)
         internal OperationContext(Stopwatch stopwatch, TimeSpan timeout, CancellationToken cancellationToken)
         {
             Stopwatch = stopwatch;
-            Timeout = timeout;
+            Timeout = Ensure.IsInfiniteOrGreaterThanOrEqualToZero(timeout, nameof(timeout));
             CancellationToken = cancellationToken;
         }
 
@@ -53,7 +51,13 @@ public TimeSpan RemainingTimeout
                     return System.Threading.Timeout.InfiniteTimeSpan;
                 }
 
-                return Timeout - Stopwatch.Elapsed;
+                var result = Timeout - Stopwatch.Elapsed;
+                if (result < TimeSpan.Zero)
+                {
+                    result = TimeSpan.Zero;
+                }
+
+                return result;
             }
         }
 
@@ -69,7 +73,7 @@ public bool IsTimedOut()
                 return false;
             }
 
-            return remainingTimeout < TimeSpan.Zero;
+            return remainingTimeout == TimeSpan.Zero;
         }
 
         public void ThrowIfTimedOutOrCanceled()
@@ -141,6 +145,8 @@ public async Task WaitTaskAsync(Task task)
 
         public OperationContext WithTimeout(TimeSpan timeout)
         {
+            Ensure.IsInfiniteOrGreaterThanOrEqualToZero(timeout, nameof(timeout));
+
             var remainingTimeout = RemainingTimeout;
             if (timeout == System.Threading.Timeout.InfiniteTimeSpan)
             {

tests/MongoDB.Bson.TestHelpers/Threading/ThreadingUtilities.cs

@@ -29,7 +29,7 @@ public static void ExecuteOnNewThreads(int threadsCount, Action<int> action, int
 
             if (exceptions.Any())
             {
-                throw exceptions.First();
+                throw new AggregateException(exceptions);
             }
         }
 

tests/MongoDB.Driver.Tests/OperationContextTests.cs

@@ -42,6 +42,14 @@ public void Constructor_should_initialize_properties()
             operationContext.ParentContext.Should().BeNull();
         }
 
+        [Fact]
+        public void Constructor_throws_on_negative_timeout()
+        {
+            var exception = Record.Exception(() => new OperationContext(TimeSpan.FromSeconds(-5), CancellationToken.None));
+
+            exception.Should().BeOfType<ArgumentOutOfRangeException>();
+        }
+
         [Fact]
         public void RemainingTimeout_should_calculate()
         {
@@ -68,16 +76,12 @@ public void RemainingTimeout_should_return_infinite_for_infinite_timeout()
         }
 
         [Fact]
-        public void RemainingTimeout_could_be_negative()
+        public void RemainingTimeout_should_return_zero_for_timeout_context()
         {
-            var timeout = TimeSpan.FromMilliseconds(5);
-            var stopwatch = Stopwatch.StartNew();
+            var operationContext = new OperationContext(TimeSpan.FromMilliseconds(5), CancellationToken.None);
             Thread.Sleep(10);
-            stopwatch.Stop();
 
-            var operationContext = new OperationContext(stopwatch, timeout, CancellationToken.None);
-
-            operationContext.RemainingTimeout.Should().Be(timeout - stopwatch.Elapsed);
+            operationContext.RemainingTimeout.Should().Be(TimeSpan.Zero);
         }
 
         [Theory]
@@ -276,6 +280,29 @@ public void WithTimeout_should_set_ParentContext()
 
             resultContext.ParentContext.Should().Be(operationContext);
         }
+
+        [Fact]
+        public void WithTimeout_should_create_timed_out_context_on_timed_out_context()
+        {
+            var operationContext = new OperationContext(TimeSpan.FromMilliseconds(5), CancellationToken.None);
+            Thread.Sleep(10);
+            operationContext.IsTimedOut().Should().BeTrue();
+
+            var resultContext = operationContext.WithTimeout(TimeSpan.FromSeconds(10));
+
+            resultContext.IsTimedOut().Should().BeTrue();
+        }
+
+        [Fact]
+        public void WithTimeout_throws_on_negative_timeout()
+        {
+            var operationContext = new OperationContext(Timeout.InfiniteTimeSpan, CancellationToken.None);
+
+            var exception = Record.Exception(() => operationContext.WithTimeout(TimeSpan.FromSeconds(-5)));
+
+            exception.Should().BeOfType<ArgumentOutOfRangeException>()
+                .Subject.ParamName.Should().Be("timeout");
+        }
     }
 }