Fix fibers api.

Add fibers test.
ntoskrnl7 · May 23, 2022 · bd6cb30 · bd6cb30
1 parent 0c327a5
commit bd6cb30
Show file tree

Hide file tree

Showing 11 changed files with 288 additions and 108 deletions.
diff --git a/README.md b/README.md
@@ -101,7 +101,7 @@ project(MyProject)
 
 # add dependencies
 include(cmake/CPM.cmake)
-CPMAddPackage("gh:ntoskrnl7/ldk@0.6.3")
+CPMAddPackage("gh:ntoskrnl7/ldk@0.7.0")
 
 # add dependencies
 CPMAddPackage("gh:ntoskrnl7/FindWDK#master")
@@ -237,9 +237,5 @@ DriverUnload (
  아직 Dll을 로드하여 사용하기에는 API가 조금밖에 구현되지 않아서 간단한 Dll 밖에는 사용할 수 없습니다.
 
 - [ ] 빠른 구현을 위해서 ReactOS 코드를 일부 사용하였으며, 추후 자체 구현해야합니다.
-- [x] Kernel32
-  - [x] LoadLibrary 시 환경 변수 등 특정 경로의 모듈을 로드할수있도록 개선해야함
-- [x] Ntdll
-  - [x] LdrLoadDll를 정상적으로 구현해야함.
 - [ ] 문서화
 - [ ] 기타
diff --git a/include/Ldk/internal/version.h b/include/Ldk/internal/version.h
@@ -25,7 +25,7 @@ Module Name:
 #define _LDK_VERSION_H_
 
 #define LDK_VERSION_MAJOR 0
-#define LDK_VERSION_MINOR 6
-#define LDK_VERSION_PATCH 3
+#define LDK_VERSION_MINOR 7
+#define LDK_VERSION_PATCH 0
 
 #endif // _LDK_VERSION_H_
diff --git a/src/internal.h b/src/internal.h
@@ -2,6 +2,11 @@
 
 
 
+#define LDK_TLS_SLOTS_SIZE			1024
+#define LDK_FLS_SLOTS_SIZE			1024
+
+
+
 #ifndef Add2Ptr
 #define Add2Ptr(P,I) ((PVOID)((PUCHAR)(P) + (I)))
 #endif

diff --git a/src/kernel32/fibersapi.c b/src/kernel32/fibersapi.c
@@ -17,30 +17,29 @@ FlsAlloc (
     _In_opt_ PFLS_CALLBACK_FUNCTION lpCallback
     )
 {
-	DWORD index = FLS_OUT_OF_INDEXES;
+	DWORD Index = FLS_OUT_OF_INDEXES;
 
 	PAGED_CODE();
 
 	RtlAcquirePebLock();
 	try {
-		index = RtlFindClearBitsAndSet( &NtCurrentPeb()->FlsBitmap,
+		Index = RtlFindClearBitsAndSet( &NtCurrentPeb()->FlsBitmap,
 										1,
 										0 );
-		if (index == 0xFFFFFFFF) {
+		if (Index == 0xFFFFFFFF) {
 			LdkSetLastNTError( STATUS_NO_MEMORY );
 		} else {
-			PLDK_FLS_SLOT slot = &NtCurrentTeb()->FlsSlots[index];
 #pragma warning(disable:4054 4055)
-			InterlockedExchangePointer( (PVOID *)&slot->Callback,
+			InterlockedExchangePointer( (PVOID *)&NtCurrentPeb()->FlsCallbacks[Index],
 										(PVOID)lpCallback );
 #pragma warning(default:4054 4055)
-			InterlockedExchangePointer( &slot->Data,
+			InterlockedExchangePointer( &NtCurrentTeb()->FlsSlots[Index],
 										NULL );
 		}
 	} finally {
 		RtlReleasePebLock();
 	}
-	return index;
+	return Index;
 }
 
 WINBASEAPI
@@ -51,7 +50,7 @@ FlsGetValue (
     )
 {
 	if (dwFlsIndex < LDK_FLS_SLOTS_SIZE) {
-		return NtCurrentTeb()->FlsSlots[dwFlsIndex].Data;
+		return NtCurrentTeb()->FlsSlots[dwFlsIndex];
 	}
 	LdkSetLastNTError( STATUS_INVALID_PARAMETER );
 	return NULL;
@@ -66,23 +65,27 @@ FlsSetValue (
     )
 {
 	if (dwFlsIndex < LDK_FLS_SLOTS_SIZE) {
-		InterlockedExchangePointer( &NtCurrentTeb()->FlsSlots[dwFlsIndex].Data,
+		InterlockedExchangePointer( &NtCurrentTeb()->FlsSlots[dwFlsIndex],
 									lpFlsData );
 		return TRUE;
 	}
-	LdkSetLastNTError(STATUS_INVALID_PARAMETER);
+	LdkSetLastNTError( STATUS_INVALID_PARAMETER );
 	return FALSE;
 }
 
+PTEB
+LdkGetNextTebRundownProtection (
+	_In_ PTEB Teb
+	);
+
 WINBASEAPI
 BOOL
 WINAPI
 FlsFree (
     _In_ DWORD dwFlsIndex
     )
 {
-	BOOLEAN index = FALSE;
-	PRTL_BITMAP bitmap;
+	BOOLEAN Found = FALSE;
 
 	PAGED_CODE();
 
@@ -91,33 +94,39 @@ FlsFree (
 		if (dwFlsIndex >= LDK_FLS_SLOTS_SIZE) {
 			leave;
 		}
-		bitmap = &NtCurrentPeb()->FlsBitmap;
-		index = RtlAreBitsSet( bitmap,
+		PRTL_BITMAP Bitmap = &NtCurrentPeb()->FlsBitmap;
+		Found = RtlAreBitsSet( Bitmap,
 							   dwFlsIndex,
 							   1 );
-		if (index) {
-			PLDK_TEB Teb = NtCurrentTeb();
-			if (ExAcquireRundownProtection( &Teb->RundownProtect )) {
-				PVOID data;
-				PFLS_CALLBACK_FUNCTION callback;
-				PLDK_FLS_SLOT slot = &Teb->FlsSlots[dwFlsIndex];
+		if (Found) {
+			PFLS_CALLBACK_FUNCTION Callback;
 #pragma warning(disable:4055)
-				callback = (PFLS_CALLBACK_FUNCTION)InterlockedExchangePointer( (PVOID *)&slot->Callback,
-																			   NULL );
+			Callback = (PFLS_CALLBACK_FUNCTION)InterlockedExchangePointer( (PVOID *)&NtCurrentPeb()->FlsCallbacks[dwFlsIndex],
+																		   NULL );
 #pragma warning(default:4055)
-				data = InterlockedExchangePointer( &slot->Data,
-												   NULL );
-				if (callback) {
-					callback( data );
+			if (Callback) {
+				PTEB InitialTeb = NtCurrentTeb();
+				if (! ExAcquireRundownProtection( &InitialTeb->RundownProtect )) {
+					return FALSE;
 				}
-				ExReleaseRundownProtection(&Teb->RundownProtect);
+				PTEB Teb = InitialTeb;
+				PVOID Data;
+				do {
+					Data = InterlockedExchangePointer( &Teb->FlsSlots[dwFlsIndex],
+													   NULL );
+					if (Data) {
+						Callback( Data );
+					}
+					Teb = LdkGetNextTebRundownProtection( Teb );
+				} while (Teb != InitialTeb);
+				ExReleaseRundownProtection( &Teb->RundownProtect );
 			}
-			RtlClearBits( bitmap,
+			RtlClearBits( Bitmap,
 						  dwFlsIndex,
 						  1 );
 		}
 	} finally {
 		RtlReleasePebLock();
 	}
-	return index;
+	return Found;
 }
diff --git a/src/kernel32/processthreadsapi.c b/src/kernel32/processthreadsapi.c
@@ -45,6 +45,14 @@ typedef struct _LDK_THREAD_CONTEXT {
 PAGED_LOOKASIDE_LIST LdkpThreadContextLookaside;
 
 
+
+VOID
+LdkpInvokeFlsCallback (
+	_Inout_ PLDK_TEB Teb
+	);
+
+
+
 NTSTATUS
 LdkpInitializeThreadContexts (
 	VOID
@@ -82,10 +90,14 @@ LdkpThreadStartExpandStackAndCallout (
 {
 	PAGED_CODE();
 
-	Context->ThreadStartRoutine( Context->lpThreadParameter );
+	LPVOID lpThreadParameter = Context->lpThreadParameter;
 
 	ExFreeToPagedLookasideList( &LdkpThreadContextLookaside,
 								Context );
+
+	Context->ThreadStartRoutine( lpThreadParameter );
+
+	LdkpInvokeFlsCallback( NtCurrentTeb() );
 }
 
 _IRQL_requires_same_
@@ -110,10 +122,14 @@ LdkpThreadStartRoutine (
 		}
 	}
 
-	Context->ThreadStartRoutine( Context->lpThreadParameter );
+	LPVOID lpThreadParameter = Context->lpThreadParameter;
 
 	ExFreeToPagedLookasideList( &LdkpThreadContextLookaside,
 								Context );
+
+	Context->ThreadStartRoutine( lpThreadParameter );
+
+	LdkpInvokeFlsCallback( NtCurrentTeb() );
 }
 
 WINBASEAPI
@@ -343,6 +359,9 @@ ExitThread (
     )
 {
 	PAGED_CODE();
+
+	LdkpInvokeFlsCallback( NtCurrentTeb() );
+
 	PsTerminateSystemThread( (NTSTATUS)dwExitCode );
 }
 

diff --git a/src/ntdll/eballoc.c b/src/ntdll/eballoc.c
@@ -16,8 +16,11 @@ RtlAcquirePebLock (
     VOID
     )
 {
+    PAGED_CODE();
+
     KeEnterCriticalRegion();
-    ExAcquireResourceExclusiveLite(&LdkCurrentPeb()->ModuleListResource, TRUE);
+    ExAcquireResourceExclusiveLite( &LdkCurrentPeb()->ModuleListResource,
+                                    TRUE );
 }
 
 VOID
@@ -26,6 +29,8 @@ RtlReleasePebLock (
     VOID
     )
 {
-	ExReleaseResourceLite(&LdkCurrentPeb()->ModuleListResource);
+    PAGED_CODE();
+
+	ExReleaseResourceLite( &LdkCurrentPeb()->ModuleListResource );
 	KeLeaveCriticalRegion();
 }
diff --git a/src/peb.h b/src/peb.h
@@ -71,6 +71,7 @@ typedef struct _LDK_PEB {
 
 	RTL_BITMAP FlsBitmap;
 	ULONG FlsBitmapBits[32];
+     PFLS_CALLBACK_FUNCTION FlsCallbacks[LDK_FLS_SLOTS_SIZE];
 
 	ERESOURCE ModuleListResource;
 	LIST_ENTRY ModuleListHead;