viosock: Reference also the internal file objects

The Socket driver stores socket-related data in memory associated with file objects representing the sockets. It uses WdfObjectReference/WdfObjectDereference to track ownership and protect sockets being in use from deletion. However, these routines do not affect reference count of the underlying FILE_OBJECT structures (served by Object Manager) -- they work only with the reference count of the WDF object wrapped around the FILE_OBJECT structure. Thus, when the underlying FILE_OBJECT is deleted, WdfObjectReference does not protect the socket context from deallocation. This commit fixes this issue by introducing VioSockReference and VioSockDereference which also work with reference count of the wrapped FILE_OBJECT structures. Also, VioSockClose is replaced with VioSockCleanup -- the latter one is invoked when the last handle for a file object (FILE_OBJECT, not WDFFILEOBJECT) gets closed (the former one is invoked when the last reference is gone -- which does not happen if VioSock(De)reference are used). For usermode, this makes no difference since user applications work only with handles, not pointer references to file objects. Signed-off-by:¨Martin Drab <[email protected]>
virtio-win · Jan 28, 2024 · e175594 · e175594
1 parent 0a6c5fc
commit e175594
Show file tree

Hide file tree

Showing 4 changed files with 39 additions and 13 deletions.
diff --git a/viosock/sys/Device.c b/viosock/sys/Device.c
@@ -282,8 +282,8 @@ VIOSockEvtDeviceAdd(
     WDF_FILEOBJECT_CONFIG_INIT(
         &fileConfig,
         VIOSockCreateStub,
-        VIOSockClose,
-        WDF_NO_EVENT_CALLBACK // Cleanup
+        WDF_NO_EVENT_CALLBACK, // Close
+        VIOSockCleanup
     );
     fileConfig.FileObjectClass = WdfFileObjectWdfCanUseFsContext;
 
@@ -778,7 +778,7 @@ VIOSockSelectCleanupFds(
         ASSERT(pHandleSet[i].Socket);
 
         InterlockedDecrement(&GetSocketContext(pHandleSet[i].Socket)->SelectRefs[iFdSet]); //dereference socket
-        WdfObjectDereference(pHandleSet[i].Socket);
+        VioSockDereference(pHandleSet[i].Socket);
     }
 
     pPkt->FdCount[iFdSet] = 0;

diff --git a/viosock/sys/Loopback.c b/viosock/sys/Loopback.c
@@ -77,7 +77,7 @@ VIOSockLoopbackAcceptEnqueue(
 
             //link accepted socket to connecting one
             pAcceptSocket->LoopbackSocket = pConnectSocket->ThisSocket;
-            WdfObjectReference(pAcceptSocket->LoopbackSocket);
+            VioSockReference(pAcceptSocket->LoopbackSocket);
             VIOSockSetFlag(pAcceptSocket, SOCK_LOOPBACK);
 
             VIOSockAcceptInitSocket(pAcceptSocket, pListenSocket);
@@ -95,7 +95,7 @@ VIOSockLoopbackAcceptEnqueue(
 
         pAccept->Memory = Memory;
         pAccept->ConnectSocket = pConnectSocket->ThisSocket;
-        WdfObjectReference(pAccept->ConnectSocket);
+        VioSockReference(pAccept->ConnectSocket);
 
         pAccept->dst_cid = pContext->Config.guest_cid;
         pAccept->dst_port = pConnectSocket->src_port;
@@ -141,7 +141,7 @@ VIOSockLoopbackAcceptDequeue(
     {
         ASSERT(FALSE);
         //skip accept entry
-        WdfObjectDereference(pAcceptEntry->ConnectSocket);
+        VioSockDereference(pAcceptEntry->ConnectSocket);
         bRes = FALSE;
     }
 
@@ -238,7 +238,7 @@ VIOSockLoopbackHandleConnecting(
                 ASSERT(pSrcSocket);
                 //link connecting socket to accepted one
                 pDestSocket->LoopbackSocket = pSrcSocket->ThisSocket;
-                WdfObjectReference(pDestSocket->LoopbackSocket);
+                VioSockReference(pDestSocket->LoopbackSocket);
 
                 WdfSpinLockAcquire(pDestSocket->StateLock);
                 VIOSockStateSet(pDestSocket, VIOSOCK_STATE_CONNECTED);

diff --git a/viosock/sys/Socket.c b/viosock/sys/Socket.c
@@ -122,7 +122,7 @@ EVT_WDF_TIMER          VIOSockPendedTimerFunc;
 #pragma alloc_text (PAGE, VIOSockConnectedListInit)
 
 #pragma alloc_text (PAGE, VIOSockCreateStub)
-#pragma alloc_text (PAGE, VIOSockClose)
+#pragma alloc_text (PAGE, VIOSockCleanup)
 
 #pragma alloc_text (PAGE, VIOSockBind)
 #pragma alloc_text (PAGE, VIOSockConnect)
@@ -988,7 +988,7 @@ VIOSockAcceptCleanup(
 
         if (pAccept->ConnectSocket != WDF_NO_HANDLE)
         {
-            WdfObjectDereference(pAccept->ConnectSocket);
+            VioSockDereference(pAccept->ConnectSocket);
         }
 
         WdfObjectDelete(pAccept->Memory);
@@ -1253,7 +1253,7 @@ VIOSockDoClose(
 }
 
 VOID
-VIOSockClose(
+VIOSockCleanup(
     IN WDFFILEOBJECT FileObject
 )
 {
@@ -1321,7 +1321,7 @@ VIOSockClose(
     }
 
     if (pSocket->LoopbackSocket != WDF_NO_HANDLE)
-        WdfObjectDereference(pSocket->LoopbackSocket);
+        VioSockDereference(pSocket->LoopbackSocket);
 
     if (PrevState == VIOSOCK_STATE_LISTEN)
         VIOSockAcceptCleanup(pSocket);
@@ -2710,7 +2710,7 @@ VIOSockGetSocketFromHandle(
 
         if (pSocket)
         {
-            WdfObjectReference(pSocket->ThisSocket);
+            VioSockReference(pSocket->ThisSocket);
             return pSocket->ThisSocket;
         }
     }

diff --git a/viosock/sys/viosock.h b/viosock/sys/viosock.h
@@ -299,7 +299,7 @@ VIOSockInterruptInit(
 //////////////////////////////////////////////////////////////////////////
 //Socket functions
 EVT_WDF_DEVICE_FILE_CREATE  VIOSockCreateStub;
-EVT_WDF_FILE_CLOSE          VIOSockClose;
+EVT_WDF_FILE_CLEANUP          VIOSockCleanup;
 
 NTSTATUS
 VIOSockDeviceControl(
@@ -917,4 +917,30 @@ VIOSockTimerResume(
 
 //////////////////////////////////////////////////////////////////////////
 
+__inline
+VOID
+VioSockReference(
+    IN WDFFILEOBJECT Socket
+)
+{
+    ObReferenceObject(WdfFileObjectWdmGetFileObject(Socket));
+    WdfObjectReference(Socket);
+
+    return;
+}
+
+
+__inline
+VOID
+VioSockDereference(
+    IN WDFFILEOBJECT Socket
+)
+{
+    WdfObjectDereference(Socket);
+    ObDereferenceObject(WdfFileObjectWdmGetFileObject(Socket));
+
+    return;
+}
+
+
 #endif /* VIOSOCK_H */