Use smart pointers in NetworkSessionCocoa::SessionWrapper

https://bugs.webkit.org/show_bug.cgi?id=259396
rdar://112666063

Reviewed by Chris Dumez.

* Source/WebKit/NetworkProcess/WebSocketTask.h:
* Source/WebKit/NetworkProcess/cocoa/NetworkDataTaskCocoa.mm:
(WebKit::NetworkDataTaskCocoa::~NetworkDataTaskCocoa):
* Source/WebKit/NetworkProcess/cocoa/NetworkSessionCocoa.h:
* Source/WebKit/NetworkProcess/cocoa/NetworkSessionCocoa.mm:
(-[WKNetworkSessionDelegate existingTask:]):
(-[WKNetworkSessionDelegate URLSession:task:didSendBodyData:totalBytesSent:totalBytesExpectedToSend:]):
(-[WKNetworkSessionDelegate URLSession:task:needNewBodyStream:]):
(-[WKNetworkSessionDelegate URLSession:task:willPerformHTTPRedirection:newRequest:completionHandler:]):
(-[WKNetworkSessionDelegate URLSession:task:_schemeUpgraded:completionHandler:]):
(-[WKNetworkSessionDelegate sessionFromTask:]):
(-[WKNetworkSessionDelegate URLSession:task:didReceiveChallenge:completionHandler:]):
(-[WKNetworkSessionDelegate URLSession:task:didCompleteWithError:]):
(-[WKNetworkSessionDelegate URLSession:task:didFinishCollectingMetrics:]):
(-[WKNetworkSessionDelegate URLSession:task:didReceiveInformationalResponse:]):
(-[WKNetworkSessionDelegate URLSession:dataTask:didReceiveResponse:completionHandler:]):
(-[WKNetworkSessionDelegate URLSession:dataTask:didReceiveData:]):
(-[WKNetworkSessionDelegate URLSession:dataTask:didBecomeDownloadTask:]):
(-[WKNetworkSessionDelegate existingWebSocketTask:]):
(WebKit::NetworkSessionCocoa::continueDidReceiveChallenge):

Canonical link: https://commits.webkit.org/266219@main

Source/WebKit/NetworkProcess/WebSocketTask.h

@@ -39,7 +39,7 @@ namespace WebKit {
 
 struct SessionSet;
 
-class WebSocketTask {
+class WebSocketTask : public CanMakeWeakPtr<WebSocketTask> {
     WTF_MAKE_FAST_ALLOCATED;
 public:
     typedef uint64_t TaskIdentifier;

Source/WebKit/NetworkProcess/cocoa/NetworkDataTaskCocoa.mm

@@ -336,8 +336,11 @@ static float toNSURLSessionTaskPriority(WebCore::ResourceLoadPriority priority)
         WTFEndSignpost(m_task.get(), "DataTask");
 
     if (m_task && m_sessionWrapper) {
-        auto dataTask = m_sessionWrapper->dataTaskMap.take([m_task taskIdentifier]);
-        RELEASE_ASSERT(dataTask == this);
+        auto& map = m_sessionWrapper->dataTaskMap;
+        auto iterator = map.find([m_task taskIdentifier]);
+        RELEASE_ASSERT(iterator != map.end());
+        ASSERT(!iterator->value.get());
+        map.remove(iterator);
     }
 }
 

Source/WebKit/NetworkProcess/cocoa/NetworkSessionCocoa.h

@@ -64,10 +64,10 @@ struct SessionWrapper : public CanMakeWeakPtr<SessionWrapper> {
 
     RetainPtr<NSURLSession> session;
     RetainPtr<WKNetworkSessionDelegate> delegate;
-    HashMap<NetworkDataTaskCocoa::TaskIdentifier, NetworkDataTaskCocoa*> dataTaskMap;
+    HashMap<NetworkDataTaskCocoa::TaskIdentifier, ThreadSafeWeakPtr<NetworkDataTaskCocoa>> dataTaskMap;
     HashMap<NetworkDataTaskCocoa::TaskIdentifier, DownloadID> downloadMap;
 #if HAVE(NSURLSESSION_WEBSOCKET)
-    HashMap<NetworkDataTaskCocoa::TaskIdentifier, WebSocketTask*> webSocketDataTaskMap;
+    HashMap<NetworkDataTaskCocoa::TaskIdentifier, WeakPtr<WebSocketTask>> webSocketDataTaskMap;
 #endif
 };
 
@@ -119,7 +119,7 @@ class NetworkSessionCocoa final : public NetworkSession {
     static bool allowsSpecificHTTPSCertificateForHost(const WebCore::AuthenticationChallenge&);
     void setClientAuditToken(const WebCore::AuthenticationChallenge&);
 
-    void continueDidReceiveChallenge(SessionWrapper&, const WebCore::AuthenticationChallenge&, NegotiatedLegacyTLS, NetworkDataTaskCocoa::TaskIdentifier, NetworkDataTaskCocoa*, CompletionHandler<void(WebKit::AuthenticationChallengeDisposition, const WebCore::Credential&)>&&);
+    void continueDidReceiveChallenge(SessionWrapper&, const WebCore::AuthenticationChallenge&, NegotiatedLegacyTLS, NetworkDataTaskCocoa::TaskIdentifier, RefPtr<NetworkDataTaskCocoa>, CompletionHandler<void(WebKit::AuthenticationChallengeDisposition, const WebCore::Credential&)>&&);
 
     SessionWrapper& sessionWrapperForDownloadResume() { return m_defaultSessionSet->sessionWithCredentialStorage; }
 

Source/WebKit/NetworkProcess/cocoa/NetworkSessionCocoa.mm

@@ -514,15 +514,15 @@ - (void)sessionInvalidated
     _sessionWrapper = nullptr;
 }
 
-- (NetworkDataTaskCocoa*)existingTask:(NSURLSessionTask *)task
+- (RefPtr<NetworkDataTaskCocoa>)existingTask:(NSURLSessionTask *)task
 {
     if (!_sessionWrapper)
         return nullptr;
 
     if (!task)
         return nullptr;
 
-    return _sessionWrapper->dataTaskMap.get(task.taskIdentifier);
+    return _sessionWrapper->dataTaskMap.get(task.taskIdentifier).get();
 }
 
 - (void)URLSession:(NSURLSession *)session didBecomeInvalidWithError:(NSError *)error
@@ -532,13 +532,13 @@ - (void)URLSession:(NSURLSession *)session didBecomeInvalidWithError:(NSError *)
 
 - (void)URLSession:(NSURLSession *)session task:(NSURLSessionTask *)task didSendBodyData:(int64_t)bytesSent totalBytesSent:(int64_t)totalBytesSent totalBytesExpectedToSend:(int64_t)totalBytesExpectedToSend
 {
-    if (auto* networkDataTask = [self existingTask:task])
+    if (auto networkDataTask = [self existingTask:task])
         networkDataTask->didSendData(totalBytesSent, totalBytesExpectedToSend);
 }
 
 - (void)URLSession:(NSURLSession *)session task:(NSURLSessionTask *)task needNewBodyStream:(void (^)(NSInputStream *bodyStream))completionHandler
 {
-    auto* networkDataTask = [self existingTask:task];
+    auto networkDataTask = [self existingTask:task];
     if (!networkDataTask) {
         completionHandler(nil);
         return;
@@ -615,7 +615,7 @@ - (void)URLSession:(NSURLSession *)session task:(NSURLSessionTask *)task willPer
     auto taskIdentifier = task.taskIdentifier;
     LOG(NetworkSession, "%llu willPerformHTTPRedirection from %s to %s", taskIdentifier, response.URL.absoluteString.UTF8String, request.URL.absoluteString.UTF8String);
 
-    if (auto* networkDataTask = [self existingTask:task]) {
+    if (auto networkDataTask = [self existingTask:task]) {
         bool shouldIgnoreHSTS = false;
 #if ENABLE(TRACKING_PREVENTION)
         if (auto* sessionCocoa = networkDataTask->networkSession()) {
@@ -669,7 +669,7 @@ - (void)URLSession:(NSURLSession *)session task:(NSURLSessionTask*)task _schemeU
     auto taskIdentifier = task.taskIdentifier;
     LOG(NetworkSession, "%llu _schemeUpgraded %s", taskIdentifier, request.URL.absoluteString.UTF8String);
 
-    if (auto* networkDataTask = [self existingTask:task]) {
+    if (auto networkDataTask = [self existingTask:task]) {
         bool shouldIgnoreHSTS = false;
 #if ENABLE(TRACKING_PREVENTION)
         if (auto* sessionCocoa = networkDataTask->networkSession()) {
@@ -717,7 +717,7 @@ static inline void processServerTrustEvaluation(NetworkSessionCocoa& session, Se
 }
 
 - (NetworkSessionCocoa*)sessionFromTask:(NSURLSessionTask *)task {
-    if (auto* networkDataTask = [self existingTask:task])
+    if (auto networkDataTask = [self existingTask:task])
         return static_cast<NetworkSessionCocoa*>(networkDataTask->networkSession());
 
     if (!_sessionWrapper)
@@ -730,7 +730,7 @@ - (NetworkSessionCocoa*)sessionFromTask:(NSURLSessionTask *)task {
     }
 
 #if HAVE(NSURLSESSION_WEBSOCKET)
-    if (auto* webSocketTask = _sessionWrapper->webSocketDataTaskMap.get(task.taskIdentifier))
+    if (auto webSocketTask = _sessionWrapper->webSocketDataTaskMap.get(task.taskIdentifier).get())
         return webSocketTask->networkSession();
 #endif
 
@@ -789,7 +789,7 @@ - (void)URLSession:(NSURLSession *)session task:(NSURLSessionTask *)task didRece
 
         // Handle server trust evaluation at platform-level if requested, for performance reasons and to use ATS defaults.
         if (sessionCocoa->fastServerTrustEvaluationEnabled() && negotiatedLegacyTLS == NegotiatedLegacyTLS::No) {
-            auto* networkDataTask = [self existingTask:task];
+            auto networkDataTask = [self existingTask:task];
             if (networkDataTask) {
                 NSURLProtectionSpace *protectionSpace = challenge.protectionSpace;
                 networkDataTask->didNegotiateModernTLS(URL { makeString(String(protectionSpace.protocol), "://", String(protectionSpace.host), ':', protectionSpace.port) });
@@ -889,7 +889,7 @@ - (void)URLSession:(NSURLSession *)session task:(NSURLSessionTask *)task didComp
         NSDictionary *oldUserInfo = [error userInfo];
         NSMutableDictionary *newUserInfo = oldUserInfo ? [NSMutableDictionary dictionaryWithDictionary:oldUserInfo] : [NSMutableDictionary dictionary];
         newUserInfo[@"networkTaskDescription"] = [task description];
-        if (RefPtr networkDataTask = [self existingTask:task]) {
+        if (auto networkDataTask = [self existingTask:task]) {
 #if HAVE(NETWORK_CONNECTION_PRIVACY_STANCE)
             newUserInfo[@"networkTaskMetricsPrivacyStance"] = privacyStanceToString(networkDataTask->networkLoadMetrics().privacyStance);
 #endif
@@ -905,7 +905,7 @@ - (void)URLSession:(NSURLSession *)session task:(NSURLSessionTask *)task didComp
         error = [NSError errorWithDomain:[error domain] code:[error code] userInfo:newUserInfo];
     }
 
-    if (auto* networkDataTask = [self existingTask:task])
+    if (auto networkDataTask = [self existingTask:task])
         networkDataTask->didCompleteWithError(error, networkDataTask->networkLoadMetrics());
     else if (error) {
         if (!_sessionWrapper)
@@ -939,7 +939,7 @@ - (void)URLSession:(NSURLSession *)session task:(NSURLSessionTask *)task didComp
 - (void)URLSession:(NSURLSession *)session task:(NSURLSessionTask *)task didFinishCollectingMetrics:(NSURLSessionTaskMetrics *)metrics
 {
     LOG(NetworkSession, "%llu didFinishCollectingMetrics", task.taskIdentifier);
-    if (auto* networkDataTask = [self existingTask:task]) {
+    if (auto networkDataTask = [self existingTask:task]) {
         NSArray<NSURLSessionTaskTransactionMetrics *> *transactionMetrics = metrics.transactionMetrics;
         NSURLSessionTaskTransactionMetrics *m = transactionMetrics.lastObject;
 
@@ -1050,7 +1050,7 @@ - (void)URLSession:(NSURLSession *)session task:(NSURLSessionTask *)task didFini
 
 - (void)URLSession:(NSURLSession *)session task:(NSURLSessionTask *)task didReceiveInformationalResponse:(NSHTTPURLResponse *)response
 {
-    if (auto* networkDataTask = [self existingTask:task]) {
+    if (auto networkDataTask = [self existingTask:task]) {
         ASSERT(RunLoop::isMain());
 
         WebCore::ResourceResponse resourceResponse(response);
@@ -1069,7 +1069,7 @@ - (void)URLSession:(NSURLSession *)session dataTask:(NSURLSessionDataTask *)data
 {
     auto taskIdentifier = dataTask.taskIdentifier;
     LOG(NetworkSession, "%llu didReceiveResponse", taskIdentifier);
-    if (auto* networkDataTask = [self existingTask:dataTask]) {
+    if (auto networkDataTask = [self existingTask:dataTask]) {
         ASSERT(RunLoop::isMain());
 
         NegotiatedLegacyTLS negotiatedLegacyTLS = NegotiatedLegacyTLS::No;
@@ -1123,7 +1123,7 @@ - (void)URLSession:(NSURLSession *)session dataTask:(NSURLSessionDataTask *)data
 
 - (void)URLSession:(NSURLSession *)session dataTask:(NSURLSessionDataTask *)dataTask didReceiveData:(NSData *)data
 {
-    if (auto* networkDataTask = [self existingTask:dataTask])
+    if (auto networkDataTask = [self existingTask:dataTask])
         networkDataTask->didReceiveData(WebCore::SharedBuffer::create(data));
 }
 
@@ -1161,7 +1161,7 @@ - (void)URLSession:(NSURLSession *)session downloadTask:(NSURLSessionDownloadTas
 
 - (void)URLSession:(NSURLSession *)session dataTask:(NSURLSessionDataTask *)dataTask didBecomeDownloadTask:(NSURLSessionDownloadTask *)downloadTask
 {
-    auto* networkDataTask = [self existingTask:dataTask];
+    auto networkDataTask = [self existingTask:dataTask];
     if (!networkDataTask)
         return;
     auto* sessionCocoa = networkDataTask->networkSession();
@@ -1190,7 +1190,7 @@ - (WebSocketTask*)existingWebSocketTask:(NSURLSessionTask *)task
     if (!task)
         return nullptr;
 
-    return _sessionWrapper->webSocketDataTaskMap.get(task.taskIdentifier);
+    return _sessionWrapper->webSocketDataTaskMap.get(task.taskIdentifier).get();
 }
 
 
@@ -1851,11 +1851,11 @@ static void activateSessionCleanup(NetworkSessionCocoa& session, const NetworkSe
     };
 }
 
-void NetworkSessionCocoa::continueDidReceiveChallenge(SessionWrapper& sessionWrapper, const WebCore::AuthenticationChallenge& challenge, NegotiatedLegacyTLS negotiatedLegacyTLS, NetworkDataTaskCocoa::TaskIdentifier taskIdentifier, NetworkDataTaskCocoa* networkDataTask, CompletionHandler<void(WebKit::AuthenticationChallengeDisposition, const WebCore::Credential&)>&& completionHandler)
+void NetworkSessionCocoa::continueDidReceiveChallenge(SessionWrapper& sessionWrapper, const WebCore::AuthenticationChallenge& challenge, NegotiatedLegacyTLS negotiatedLegacyTLS, NetworkDataTaskCocoa::TaskIdentifier taskIdentifier, RefPtr<NetworkDataTaskCocoa> networkDataTask, CompletionHandler<void(WebKit::AuthenticationChallengeDisposition, const WebCore::Credential&)>&& completionHandler)
 {
     if (!networkDataTask) {
 #if HAVE(NSURLSESSION_WEBSOCKET)
-        if (auto* webSocketTask = sessionWrapper.webSocketDataTaskMap.get(taskIdentifier)) {
+        if (auto webSocketTask = sessionWrapper.webSocketDataTaskMap.get(taskIdentifier).get()) {
             auto challengeCompletionHandler = createChallengeCompletionHandler(networkProcess(), sessionID(), challenge, webSocketTask->partition(), 0, WTFMove(completionHandler));
             networkProcess().authenticationManager().didReceiveAuthenticationChallenge(sessionID(), webSocketTask->webProxyPageID(), !webSocketTask->topOrigin().isNull() ? &webSocketTask->topOrigin() : nullptr, challenge, negotiatedLegacyTLS, WTFMove(challengeCompletionHandler));
             return;