[URL Session] Avoid Closing Sockets Before Cancelling Event Sources

When you create an event source for a handle in dispatch, dispatch
assumes it owns it, however, when curl finishes with the socket, it
closes it, which leaves dispatch with an invalid handle.

This sets the curl close callback to do nothing which allows us to close
the socket after cancelling the dispatch sources.

Author: gmittert

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()
         }