fix: reject stale sessions without retaining past responses (#1282)

Author: KazuCocoa

WebDriverAgentLib/Routing/FBHTTPServer.m

@@ -13,6 +13,7 @@
 #import "FBLogger.h"
 #import "FBResponsePayload.h"
 #import "FBTCPSocket.h"
+#import "FBSession.h"
 
 static NSData *FBCRLFCRLFData(void)
 {
@@ -131,10 +132,6 @@ @interface FBHTTPServer () <FBTCPSocketDelegate>
 // standalone or not (except DELETE /session itself - see -dispatchMethod:). See
 // -abandonPendingRequestsForSessionID:. Guarded by @synchronized(self.pendingSessionRequests).
 @property (nonatomic, strong) NSMutableDictionary<NSString *, NSMutableSet<FBPendingRequest *> *> *pendingSessionRequests;
-// Already-abandoned sessions mapped to the response they were abandoned with, so a request parsed
-// after that point is answered at once instead of queueing for a session that is gone. Kept for
-// the server's lifetime; ids are UUIDs. Guarded by @synchronized(self.pendingSessionRequests).
-@property (nonatomic, strong) NSMutableDictionary<NSString *, RouteResponse *> *abandonedSessionResponses;
 // When each connection started waiting for its current request. The reaper closes connections
 // whose entry outlives FBIncompleteRequestTimeout; idle keep-alive connections have no entry and
 // are exempt. Guarded by @synchronized(self.connectionBuffers).
@@ -163,7 +160,6 @@ - (instancetype)init
     _connectionsAwaitingResponse = [NSMutableSet set];
     _standaloneWaiters = [NSMutableDictionary dictionary];
     _pendingSessionRequests = [NSMutableDictionary dictionary];
-    _abandonedSessionResponses = [NSMutableDictionary dictionary];
     _incompleteRequestStarts = [NSMapTable mapTableWithKeyOptions:(NSPointerFunctionsOptions)(NSMapTableObjectPointerPersonality | NSMapTableStrongMemory)
                                                      valueOptions:(NSPointerFunctionsOptions)NSMapTableStrongMemory];
   }
@@ -739,9 +735,15 @@ - (void)dispatchMethod:(NSString *)method pathAndQuery:(NSString *)pathAndQuery
 - (nullable RouteResponse *)trackPendingRequest:(FBPendingRequest *)pendingRequest forSessionID:(NSString *)sessionID
 {
   @synchronized (self.pendingSessionRequests) {
-    RouteResponse *abandonedResponse = self.abandonedSessionResponses[sessionID];
-    if (nil != abandonedResponse) {
-      return abandonedResponse;
+    // Checking the current session under this lock closes the admission/teardown
+    // race without retaining a tombstone and a full response for every past UUID.
+    // -kill clears the active session before posting its abandonment notification.
+    if (nil == [FBSession sessionWithIdentifier:sessionID]) {
+      RouteResponse *response = [RouteResponse new];
+      [FBResponseWithStatus([FBCommandStatus noSuchDriverErrorWithMessage:@"Session does not exist"
+                                                               traceback:nil]) dispatchWithResponse:response];
+      [self applyDefaultHeadersToResponse:response];
+      return response;
     }
     NSMutableSet<FBPendingRequest *> *pendingRequests = self.pendingSessionRequests[sessionID];
     if (nil == pendingRequests) {
@@ -776,8 +778,6 @@ - (void)abandonPendingRequestsForSessionID:(NSString *)sessionID withResponse:(R
   @synchronized (self.pendingSessionRequests) {
     pendingRequests = [self.pendingSessionRequests[sessionID] copy];
     [self.pendingSessionRequests removeObjectForKey:sessionID];
-    // Recorded before the lock is dropped, so requests admitted from here on are rejected.
-    self.abandonedSessionResponses[sessionID] = response;
   }
   for (FBPendingRequest *pendingRequest in pendingRequests) {
     [self writeResponse:response toClient:pendingRequest.client];

WebDriverAgentTests/UnitTests/FBHTTPServerSessionTests.m

@@ -13,6 +13,7 @@
 #import <sys/socket.h>
 
 #import "FBHTTPServer.h"
+#import "FBSession-Private.h"
 
 static atomic_int gSessionProbeHits;
 
@@ -23,10 +24,26 @@ @interface FBHTTPServerSessionTests : XCTestCase
 
 @implementation FBHTTPServerSessionTests
 
+- (void)testAbandonmentDoesNotRetainResponsesAfterPendingRequestsDrain
+{
+  __weak RouteResponse *weakResponse;
+  @autoreleasepool {
+    RouteResponse *response = [RouteResponse new];
+    [response respondWithString:@"deleted"];
+    weakResponse = response;
+    for (NSUInteger index = 0; index < 1000; index++) {
+      [self.server abandonPendingRequestsForSessionID:[NSString stringWithFormat:@"old-%lu", (unsigned long)index]
+                                         withResponse:response];
+    }
+  }
+  XCTAssertNil(weakResponse);
+}
+
 - (void)setUp
 {
   [super setUp];
   atomic_store(&gSessionProbeHits, 0);
+  [FBSession initWithApplication:nil].identifier = @"live-session";
   self.server = [FBHTTPServer new];
   [self.server get:@"/session/:sessionID/probe" withBlock:^(RouteRequest *request, RouteResponse *response) {
     atomic_fetch_add(&gSessionProbeHits, 1);
@@ -42,6 +59,7 @@ - (void)tearDown
 {
   [self.server stop:NO];
   self.server = nil;
+  [FBSession.activeSession kill];
   [super tearDown];
 }
 
@@ -93,14 +111,14 @@ - (void)testRequestForAlreadyAbandonedSessionIsRejectedImmediately
 
   NSString *response = [self responseForRawPayload:(NSData * _Nonnull)[@"GET /session/dead-session/probe HTTP/1.1\r\n\r\n" dataUsingEncoding:NSUTF8StringEncoding]
                                             timeout:5.0];
-  XCTAssertTrue([response containsString:@"session-was-deleted"], @"%@", response);
+  XCTAssertTrue([response containsString:@"invalid session id"], @"%@", response);
   XCTAssertEqual(atomic_load(&gSessionProbeHits), 0, @"the route must not run for a deleted session");
 }
 
 - (void)testAbandonedSessionIsRememberedAfterManyLaterAbandonments
 {
-  // Abandoned ids are kept for the server's lifetime; evicting them would let a stale request
-  // queue on a possibly wedged route queue again, which is the hang this rejection prevents.
+  // A stale UUID must still be rejected after arbitrarily many other sessions,
+  // without needing to remember every deleted UUID.
   RouteResponse *abandonedResponse = [RouteResponse new];
   [abandonedResponse respondWithString:@"session-was-deleted"];
   [self.server abandonPendingRequestsForSessionID:@"dead-session" withResponse:abandonedResponse];
@@ -113,7 +131,7 @@ - (void)testAbandonedSessionIsRememberedAfterManyLaterAbandonments
 
   NSString *response = [self responseForRawPayload:(NSData * _Nonnull)[@"GET /session/dead-session/probe HTTP/1.1\r\n\r\n" dataUsingEncoding:NSUTF8StringEncoding]
                                             timeout:5.0];
-  XCTAssertTrue([response containsString:@"session-was-deleted"], @"%@", response);
+  XCTAssertTrue([response containsString:@"invalid session id"], @"%@", response);
   XCTAssertEqual(atomic_load(&gSessionProbeHits), 0, @"the route must not run for a deleted session");
 }