Merge pull request swiftlang#2424 from gmittert/DontStealTheSources

[URL Session] Avoid Closing Sockets Before Cancelling Event Sources

Author: millenomi

CoreFoundation/URL.subproj/CFURLSessionInterface.c

@@ -110,6 +110,10 @@ CFURLSessionEasyCode CFURLSession_easy_setopt_tc(CFURLSessionEasyHandle _Nonnull
     return MakeEasyCode(curl_easy_setopt(curl,  option.value, a));
 }
 
+CFURLSessionEasyCode CFURLSession_easy_setopt_csf(CFURLSessionEasyHandle _Nonnull curl, CFURLSessionOption option, int (*_Nonnull a)(void *_Nullable clientp, CFURLSession_socket_t)) {
+     return MakeEasyCode(curl_easy_setopt(curl, option.value, a));
+}
+
 CFURLSessionEasyCode CFURLSession_easy_getinfo_long(CFURLSessionEasyHandle _Nonnull curl, CFURLSessionInfo info, long *_Nonnull a) {
     return MakeEasyCode(curl_easy_getinfo(curl, info.value, a));
 }

CoreFoundation/URL.subproj/CFURLSessionInterface.h

@@ -601,7 +601,7 @@ CF_EXPORT CFURLSessionMultiCode CFURLSession_multi_setopt_ptr(CFURLSessionMultiH
 CF_EXPORT CFURLSessionMultiCode CFURLSession_multi_setopt_l(CFURLSessionMultiHandle _Nonnull multi_handle, CFURLSessionMultiOption option, long a);
 CF_EXPORT CFURLSessionMultiCode CFURLSession_multi_setopt_sf(CFURLSessionMultiHandle _Nonnull multi_handle, CFURLSessionMultiOption option, int (*_Nonnull a)(CFURLSessionEasyHandle _Nonnull, CFURLSession_socket_t, int, void *_Nullable, void *_Nullable));
 CF_EXPORT CFURLSessionMultiCode CFURLSession_multi_setopt_tf(CFURLSessionMultiHandle _Nonnull multi_handle, CFURLSessionMultiOption option, int (*_Nonnull a)(CFURLSessionMultiHandle _Nonnull, long, void *_Nullable));
-
+CF_EXPORT CFURLSessionEasyCode CFURLSession_easy_setopt_csf(CFURLSessionEasyHandle _Nonnull curl, CFURLSessionOption option, int (*_Nonnull a)(void* _Nullable clientp, CFURLSession_socket_t));
 CF_EXPORT CFURLSessionEasyCode CFURLSessionInit(void);
 
 

Foundation/URLSession/libcurl/EasyHandle.swift

@@ -486,6 +486,14 @@ fileprivate extension _EasyHandle {
                 return 1
             }
         }.asError()
+
+        try! CFURLSession_easy_setopt_csf(rawHandle, CFURLSessionOptionCLOSESOCKETFUNCTION) { (clientp: UnsafeMutableRawPointer?, socket: CFURLSession_socket_t) -> Int32 in
+            // Don't let CURL close the socket here because the
+            // dispatch sources are associated with it and we need to
+            // cancel them before closing the file descriptor.
+            return 0
+        }.asError()
+
         // seeking in input stream
         try! CFURLSession_easy_setopt_ptr(rawHandle, CFURLSessionOptionSEEKDATA, UnsafeMutableRawPointer(Unmanaged.passUnretained(self).toOpaque())).asError()
         try! CFURLSession_easy_setopt_seek(rawHandle, CFURLSessionOptionSEEKFUNCTION, { (userdata, offset, origin) -> Int32 in

Foundation/URLSession/libcurl/MultiHandle.swift

@@ -122,7 +122,9 @@ fileprivate extension URLSession._MultiHandle {
         } else if socketSources != nil && action == .unregister {
             // We need to release the stored pointer:
             if let opaque = socketSourcePtr {
-                Unmanaged<_SocketSources>.fromOpaque(opaque).release()
+                let s: Unmanaged<_SocketSources> = Unmanaged<_SocketSources>.fromOpaque(opaque)
+                s.takeUnretainedValue().tearDown(socket: socket, queue: queue)
+                s.release()
             }
             socketSources = nil
         }
@@ -407,6 +409,7 @@ fileprivate extension URLSession._MultiHandle._Timeout {
 fileprivate class _SocketSources {
     var readSource: DispatchSource?
     var writeSource: DispatchSource?
+    let activeSockets = DispatchGroup()
 
     func createReadSource(socket: CFURLSession_socket_t, queue: DispatchQueue, handler: DispatchWorkItem) {
         guard readSource == nil else { return }
@@ -415,7 +418,13 @@ fileprivate class _SocketSources {
 #else
         let s = DispatchSource.makeReadSource(fileDescriptor: socket, queue: queue)
 #endif
+        activeSockets.enter()
         s.setEventHandler(handler: handler)
+        s.setCancelHandler(handler: DispatchWorkItem { [weak self] in
+            guard let self = self else { return }
+            self.activeSockets.leave()
+        })
+
         readSource = s as? DispatchSource
         s.resume()
     }
@@ -427,12 +436,26 @@ fileprivate class _SocketSources {
 #else
         let s = DispatchSource.makeWriteSource(fileDescriptor: socket, queue: queue)
 #endif
+        activeSockets.enter()
+        s.setCancelHandler(handler: DispatchWorkItem { [weak self] in
+            guard let self = self else { return }
+            self.activeSockets.leave()
+        })
         s.setEventHandler(handler: handler)
         writeSource = s as? DispatchSource
         s.resume()
     }
 
-    func tearDown() {
+    func tearDown(socket: CFURLSession_socket_t, queue: DispatchQueue) {
+        activeSockets.notify(queue: queue) {
+            withExtendedLifetime(self) {
+#if os(Windows)
+                closesocket(socket)
+#else
+                close(socket)
+#endif
+            }
+        }
         if let s = readSource {
             s.cancel()
         }