#include "precomp.h"
#include <ntddvol.h>

// 만약 헤더가 없어서 컴파일 에러가 난다면 아래 주석을 해제하여 사용하세요.
/*
DEFINE_DEVPROPKEY(DEVPKEY_Device_InstanceId,
	0x78c34fc8, 0x104a, 0x4aca, 0x9e, 0xa4, 0x52, 0x4d, 0x52, 0x99, 0x6e, 0x57, 256);
*/


NTSTATUS
GetPhysicalDiskDeviceObject(
	__in PDEVICE_OBJECT VolumeDeviceObject,
	__out PDEVICE_OBJECT* PhysicalDiskDeviceObject
)
{
	NTSTATUS status;
	PIRP Irp;
	KEVENT Event;
	IO_STATUS_BLOCK Iosb;
	PVOLUME_DISK_EXTENTS pDiskExtents = NULL;
	// 넉넉하게 4개 정도의 Extent 공간 확보
	ULONG extentsSize = sizeof(VOLUME_DISK_EXTENTS) + (sizeof(DISK_EXTENT) * 4);
	UNICODE_STRING physicalDiskName;
	WCHAR nameBuffer[64] = { 0 };
	PFILE_OBJECT pFileObject = NULL;
	PDEVICE_OBJECT pDevObj = NULL;

	*PhysicalDiskDeviceObject = NULL;

	// 1. 메모리 할당
	pDiskExtents = (PVOLUME_DISK_EXTENTS)ExAllocatePoolWithTag(NonPagedPool, extentsSize, 'dExt');
	if (!pDiskExtents) return STATUS_INSUFFICIENT_RESOURCES;

	// 메모리 초기화 (중요)
	RtlZeroMemory(pDiskExtents, extentsSize);

	KeInitializeEvent(&Event, NotificationEvent, FALSE);
	Irp = IoBuildDeviceIoControlRequest(
		IOCTL_VOLUME_GET_VOLUME_DISK_EXTENTS,
		VolumeDeviceObject,
		NULL, 0,
		pDiskExtents, extentsSize,
		FALSE, &Event, &Iosb
	);

	if (!Irp) {
		ExFreePoolWithTag(pDiskExtents, 'dExt');
		return STATUS_INSUFFICIENT_RESOURCES;
	}

	status = IoCallDriver(VolumeDeviceObject, Irp);
	if (status == STATUS_PENDING) {
		KeWaitForSingleObject(&Event, Executive, KernelMode, FALSE, NULL);
		status = Iosb.Status;
	}

	if (!NT_SUCCESS(status)) {
		ExFreePoolWithTag(pDiskExtents, 'dExt');
		return status;
	}

	// [BSOD 방지 핵심] Extents 개수 확인!
	if (pDiskExtents->NumberOfDiskExtents == 0)
	{
		KLogEx(DEBUG_TRACE_INFO, "No Disk Extents found.\n");
		ExFreePoolWithTag(pDiskExtents, 'dExt');
		return STATUS_UNSUCCESSFUL;
	}

	__try
	{
		// 2. 물리 디스크 이름 생성
		ULONG diskNumber = pDiskExtents->Extents[0].DiskNumber;

		// 안전한 문자열 함수 사용
		if (!NT_SUCCESS(RtlStringCbPrintfW(nameBuffer, sizeof(nameBuffer), L"\\Device\\Harddisk%u\\DR0", diskNumber)))
		{
			status = STATUS_BUFFER_OVERFLOW;
			__leave;
		}

		RtlInitUnicodeString(&physicalDiskName, nameBuffer);

		KLogEx(DEBUG_TRACE_INFO, "Trying to open: %S\n", nameBuffer);

		// 3. 디바이스 포인터 획득 (가장 위험한 구간)
		// InstanceSetup에서 이 호출이 위험할 수 있으므로 실패 시 즉시 빠져나옴
		status = IoGetDeviceObjectPointer(
			&physicalDiskName,
			FILE_READ_ATTRIBUTES, // 권한 최소화
			&pFileObject,
			&pDevObj
		);

		if (NT_SUCCESS(status))
		{
			ObReferenceObject(pDevObj); // 디바이스 객체 참조
			ObDereferenceObject(pFileObject); // 파일 객체 닫기
			*PhysicalDiskDeviceObject = pDevObj;
		}
		else
		{
			KLogEx(DEBUG_TRACE_ERROR, "IoGetDeviceObjectPointer Failed: 0x%x\n", status);
		}
	}
	__except (EXCEPTION_EXECUTE_HANDLER)
	{
		KLogEx(DEBUG_TRACE_ERROR, "Exception in GetPhysicalDiskDeviceObject\n");
		status = STATUS_UNHANDLED_EXCEPTION;
	}

	ExFreePoolWithTag(pDiskExtents, 'dExt');
	return status;
}

NTSTATUS
QueryDeviceProperty(
	__in PDEVICE_OBJECT pDevice,
	__in DEVICE_REGISTRY_PROPERTY DevProperty,
	__out PVOID* ppBuffer,
	__out PULONG pResultLength
)
{

	NTSTATUS status;
	PVOID pBuffer = NULL;
	ULONG BufferSize = 0;
	PDEVICE_OBJECT pdo = NULL;
	ASSERT(ARGUMENT_PRESENT(pDevice));

	// sskim - 171122 Added.
	// FltGetDiskDeviceObject로 구한 pDevice는 PDO여서 Verifier에서 터짐 (FDO로 가져오게 추가)
	pdo = IoGetDeviceAttachmentBaseRef(pDevice);
	

	status = IoGetDeviceProperty(
		pdo,
		DevProperty,
		BufferSize,
		NULL,
		&BufferSize
	);

	if (NT_SUCCESS(status))
	{
		ObDereferenceObject(pdo);
		return STATUS_NOT_SUPPORTED;
	}

	while (status == STATUS_BUFFER_TOO_SMALL)
	{
		pBuffer = ExAllocatePoolWithTag(PagedPool, BufferSize, _ALLOC_TAG);
		if (!pBuffer)
		{
			ObDereferenceObject(pdo);
			return STATUS_NO_MEMORY;
		}

		status = IoGetDeviceProperty(
			pdo,
			DevProperty,
			BufferSize,
			pBuffer,
			&BufferSize
		);

		if (NT_SUCCESS(status))
		{
			*ppBuffer = pBuffer;
			*pResultLength = BufferSize;
			ObDereferenceObject(pdo);
			return status;
		}

		ExFreePool(pBuffer);
		pBuffer = NULL;
	}

	ASSERT(!pBuffer);

	ObDereferenceObject(pdo);
	return status;
}

//#define QUERY_BUFFER_SIZE 0x2000
//NTSTATUS SafeSetVolumContextParam(PUCHAR pucDest, DWORD dwSize, PUCHAR pucSrc, DWORD dwOffset)
//{
//	int size = 0;
//	PCHAR pszName = NULL;
//
//	if (dwOffset == 0)
//		return STATUS_UNSUCCESSFUL;
//
//	if (dwOffset >= dwSize)
//		return STATUS_UNSUCCESSFUL;
//
//	if (dwOffset >= (QUERY_BUFFER_SIZE - 10))
//		return STATUS_UNSUCCESSFUL;
//
//	pszName = (char*)GetPtr(pucSrc, dwOffset);
//	size = (int)strlen(pszName);
//
//	KLogEx(DEBUG_TRACE_INFO, "SafeSetVolumContextParam, size(%d)\n", size);
//
//	if (size > VOLUME_DESCRIPTION_LENGTH)
//	{
//		size = VOLUME_DESCRIPTION_LENGTH - 1;
//	}
//
//	RtlCopyMemory(pucDest, pszName, size);
//	return STATUS_SUCCESS;
//}


// 매크로 정의 (없다면 추가 필요)
#ifndef QUERY_BUFFER_SIZE
#define QUERY_BUFFER_SIZE 0x2000
#endif

// 안전한 포인터 계산 및 복사 함수 (대문자 변환 포함)
NTSTATUS SafeSetVolumContextParam(
	__in PUCHAR pucDest,        // 대상 버퍼 (Context 내 변수)
	__in ULONG  ulSrcTotalSize, // 원본 전체 버퍼 크기 (pDesc->Size)
	__in PUCHAR pucSrcBase,     // 원본 버퍼 시작 주소 (pDesc)
	__in ULONG  ulOffset        // 문자열 오프셋
)
{
	ULONG i = 0;
	ULONG destIndex = 0;
	PCHAR pSrcStr = NULL;
	BOOLEAN bLeadingSpace = TRUE; // 앞 공백 여부 플래그
	ULONG lastNonSpaceIndex = 0;  // 마지막으로 공백이 아닌 문자의 인덱스

	// 1. 오프셋 유효성 검사
	if (ulOffset == 0 || ulOffset >= ulSrcTotalSize)
	{
		if (pucDest) pucDest[0] = 0;
		return STATUS_SUCCESS;
	}

	// 2. 문자열 시작 포인터 계산
	pSrcStr = (PCHAR)(pucSrcBase + ulOffset);

	// 3. 복사 루프 (대문자 변환 + 앞 공백 제거)
	for (i = 0; i < VOLUME_DESCRIPTION_LENGTH - 1; i++)
	{
		// A. 원본 버퍼 범위 체크 (Overflow 방지)
		if ((ULONG_PTR)(pSrcStr + i) >= (ULONG_PTR)(pucSrcBase + ulSrcTotalSize))
			break;

		// B. 원본 문자열 끝(NULL) 체크
		if (pSrcStr[i] == 0)
			break;

		// 문자 가져오기
		CHAR c = pSrcStr[i];

		// [Trim Start] 앞쪽 공백 제거 로직
		if (bLeadingSpace)
		{
			if (c == ' ' || c == '\t' || c == '\r' || c == '\n')
			{
				continue; // 앞쪽 공백은 건너뜀
			}
			else
			{
				bLeadingSpace = FALSE; // 공백이 아닌 문자를 처음 만남 -> 이제부터 복사 시작
			}
		}

		// [대문자 변환] 및 복사
		pucDest[destIndex] = RtlUpperChar(c);

		// [Trim End 준비] 공백이 아닌 문자의 위치를 기억
		if (pucDest[destIndex] != ' ' && pucDest[destIndex] != '\t' &&
			pucDest[destIndex] != '\r' && pucDest[destIndex] != '\n')
		{
			lastNonSpaceIndex = destIndex;
		}

		destIndex++;
	}

	// 4. [Trim End] 뒤쪽 공백 제거 및 NULL Termination
	// 복사된 내용이 하나라도 있다면, 마지막 공백이 아닌 문자 다음을 NULL로 만듦
	if (destIndex > 0 && !bLeadingSpace)
	{
		pucDest[lastNonSpaceIndex + 1] = 0;
	}
	else
	{
		// 공백만 있었거나 빈 문자열인 경우
		pucDest[0] = 0;
	}

	KLogEx(DEBUG_TRACE_INFO, "Copied Length(%d), Content(%s)\n", lastNonSpaceIndex + 1, pucDest);

	return STATUS_SUCCESS;
}

typedef struct _STORAGE_DEVICE_ID_DESCRIPTOR {
	ULONG Version;
	ULONG Size;
	ULONG NumberOfIdentifiers;
	UCHAR Identifiers[1]; // 가변 길이 문자열 배열
} STORAGE_DEVICE_ID_DESCRIPTOR, * PSTORAGE_DEVICE_ID_DESCRIPTOR;

NTSTATUS
GetStorageDeviceId(
	__in PDEVICE_OBJECT DeviceObject,
	__out PWCHAR* OutInstanceId
)
{
	NTSTATUS status = STATUS_SUCCESS;
	PIRP Irp;
	KEVENT Event;
	IO_STATUS_BLOCK Iosb;
	STORAGE_PROPERTY_QUERY PropQuery = { 0 };
	PSTORAGE_DEVICE_ID_DESCRIPTOR pDesc = NULL;
	PVOID QueryBuffer = NULL;
	ULONG QuerySize = 1024; // 넉넉하게 잡음
	PDEVICE_OBJECT pTopDevice = NULL;
	*OutInstanceId = NULL;
	pTopDevice = IoGetAttachedDeviceReference(DeviceObject);
	// 1. 메모리 할당
	QueryBuffer = ExAllocatePoolWithTag(PagedPool, QuerySize, 'StoI');
	if (!QueryBuffer) return STATUS_INSUFFICIENT_RESOURCES;

	RtlZeroMemory(QueryBuffer, QuerySize);

	// 2. 쿼리 설정: StorageDeviceIdProperty (중요!)
	PropQuery.PropertyId = StorageDeviceIdProperty; // enum 값: 2
	PropQuery.QueryType = PropertyStandardQuery;

	KeInitializeEvent(&Event, NotificationEvent, FALSE);

	// 3. IRP 생성 및 전송
	Irp = IoBuildDeviceIoControlRequest(
		IOCTL_STORAGE_QUERY_PROPERTY,
		pTopDevice,
		&PropQuery,
		sizeof(PropQuery),
		QueryBuffer,
		QuerySize,
		FALSE,
		&Event,
		&Iosb
	);

	if (!Irp)
	{
		ExFreePoolWithTag(QueryBuffer, 'StoI');
		return STATUS_INSUFFICIENT_RESOURCES;
	}

	status = IoCallDriver(pTopDevice, Irp);

	if (status == STATUS_PENDING)
	{
		KeWaitForSingleObject(&Event, Executive, KernelMode, FALSE, NULL);
		status = Iosb.Status;
	}

	// 4. 결과 파싱
	if (NT_SUCCESS(status) && Iosb.Information > sizeof(STORAGE_DEVICE_ID_DESCRIPTOR))
	{
		pDesc = (PSTORAGE_DEVICE_ID_DESCRIPTOR)QueryBuffer;

		if (pDesc->NumberOfIdentifiers > 0)
		{
			// Identifiers는 Null로 구분된 ANSI 문자열들의 배열입니다.
			// 첫 번째 문자열이 보통 가장 구체적인 PnP ID (Hardware ID)입니다.
			PCHAR pAnsiId = (PCHAR)pDesc->Identifiers;

			// ANSI String -> Unicode String 변환 필요
			ANSI_STRING ansiStr;
			UNICODE_STRING uniStr;

			RtlInitAnsiString(&ansiStr, pAnsiId);

			status = RtlAnsiStringToUnicodeString(&uniStr, &ansiStr, TRUE); // TRUE: Allocate Memory

			if (NT_SUCCESS(status))
			{
				// 호출자에게 전달 (호출자가 나중에 ExFreePool 해야 함)
				// RtlAnsiStringToUnicodeString이 할당한 버퍼는 PagedPool입니다.
				// 여기서는 호출자가 편하게 쓰도록 새 버퍼에 복사하거나, uniStr.Buffer를 그대로 줍니다.

				// 안전하게 복사본을 만들어 리턴 (NonPagedPool 권장 시)
				ULONG uniSize = uniStr.Length + sizeof(WCHAR);
				PWCHAR pResult = ExAllocatePoolWithTag(NonPagedPool, uniSize, 'IDst');

				if (pResult)
				{
					RtlZeroMemory(pResult, uniSize);
					RtlCopyMemory(pResult, uniStr.Buffer, uniStr.Length);
					*OutInstanceId = pResult;
				}

				RtlFreeUnicodeString(&uniStr); // 임시 변환 버퍼 해제
			}
		}
	}

	ExFreePoolWithTag(QueryBuffer, 'StoI');
	return status;
}


NTSTATUS
GetStorageProperty(
	__in PDEVICE_OBJECT device_object,
	__in PVOLUME_CONTEXT pVolumeContext
)
{
	PIRP Irp;
	KEVENT Event;
	NTSTATUS status = STATUS_SUCCESS;
	IO_STATUS_BLOCK Iosb;
	STORAGE_PROPERTY_QUERY PropQuery = { 0, };
	PVOID QueryBuffer = NULL;
	ULONG QuerySize = QUERY_BUFFER_SIZE;
	PSTORAGE_DEVICE_DESCRIPTOR pDesc;
	//PDEVICE_OBJECT pTopDevice = NULL;
	//ULONG size = 0;

	__try
	{
		//pTopDevice = IoGetAttachedDeviceReference(device_object);

		QueryBuffer = ExAllocatePoolWithTag(PagedPool, QuerySize, _ALLOC_TAG);
		if (!QueryBuffer)
		{
			status = STATUS_INSUFFICIENT_RESOURCES;
			__leave;
		}

		RtlZeroMemory(&PropQuery, sizeof(PropQuery));
		RtlZeroMemory(QueryBuffer, QuerySize);
		PropQuery.PropertyId = StorageDeviceProperty;
		PropQuery.QueryType = PropertyStandardQuery;

		KeInitializeEvent(&Event, NotificationEvent, FALSE);

		Irp = IoBuildDeviceIoControlRequest(
			IOCTL_STORAGE_QUERY_PROPERTY,
			device_object,
			&PropQuery,
			sizeof(PropQuery),
			QueryBuffer,
			QuerySize,
			FALSE,
			&Event,
			&Iosb
		);

		if (!Irp)
		{
			status = STATUS_INSUFFICIENT_RESOURCES;
			__leave;
		}

		status = IoCallDriver(device_object, Irp);

		if (status == STATUS_PENDING)
		{
			KeWaitForSingleObject(
				&Event,
				Executive,
				KernelMode,
				FALSE,
				(PLARGE_INTEGER)NULL
			);

			status = Iosb.Status;
		}

		if (!NT_SUCCESS(status))
		{
			__leave;
		}

		if (Iosb.Information < sizeof(STORAGE_DEVICE_DESCRIPTOR))
		{
			status = STATUS_UNSUCCESSFUL;
			__leave;
		}

		pDesc = (PSTORAGE_DEVICE_DESCRIPTOR)QueryBuffer;
		if (!pDesc)
			__leave;

		if (pDesc->Size > Iosb.Information)
		{
			if (pDesc->Size > QuerySize) pDesc->Size = (ULONG)Iosb.Information;
		}

		KLogEx(DEBUG_TRACE_INFO, "Version(%x),Size(%x),DeviceType(%x),DeviceTypeModifier(%x)\n", pDesc->Version, pDesc->Size, pDesc->DeviceType, pDesc->DeviceTypeModifier);
		KLogEx(DEBUG_TRACE_INFO, "bustype(%d),Removal(%d),vidoffset(%x),pidoffset(%x)\n", pDesc->BusType, pDesc->RemovableMedia, pDesc->VendorIdOffset, pDesc->ProductIdOffset);
		KLogEx(DEBUG_TRACE_INFO, "ProductRevisionOffset(%d),SerialNumberOffset(%d)\n", pDesc->ProductRevisionOffset, pDesc->SerialNumberOffset);

		pVolumeContext->bustype = pDesc->BusType;
		SafeSetVolumContextParam(pVolumeContext->vendorid, pDesc->Size, (PUCHAR)pDesc, pDesc->VendorIdOffset);
		SafeSetVolumContextParam(pVolumeContext->productid, pDesc->Size, (PUCHAR)pDesc, pDesc->ProductIdOffset);
		SafeSetVolumContextParam(pVolumeContext->productrevisionlevel, pDesc->Size, (PUCHAR)pDesc, pDesc->ProductRevisionOffset);
		SafeSetVolumContextParam(pVolumeContext->vendorspecific, pDesc->Size, (PUCHAR)pDesc, pDesc->SerialNumberOffset);
	}
	__finally
	{
		if (QueryBuffer)
		{
			ExFreePoolWithTag(QueryBuffer, _ALLOC_TAG);
			QueryBuffer = NULL;
		}
	}

	return status;
}

typedef struct _GET_ID_WORK_CONTEXT {
	PFLT_FILTER Filter;
	PFLT_INSTANCE Instance; // 인스턴스 핸들 (컨텍스트 접근용)
	PDEVICE_OBJECT VolumeDeviceObject; // 볼륨 디바이스
} GET_ID_WORK_CONTEXT, * PGET_ID_WORK_CONTEXT;

// [작업자 스레드 함수] 여기서 안전하게 Hardware ID를 구합니다.
VOID
Bs1FltGetHardwareIdWorker(
	__in PFLT_GENERIC_WORKITEM FltWorkItem,
	__in PVOID ConnectionCookie,
	__in PVOID Context_
)
{
	PGET_ID_WORK_CONTEXT pWorkCtx = (PGET_ID_WORK_CONTEXT)Context_;
	PVOLUME_CONTEXT pVolumeCtx = NULL;
	PDEVICE_OBJECT pPhysicalDisk = NULL;
	PDEVICE_OBJECT pPhysicalPDO = NULL;
	NTSTATUS status;
	PWCHAR pHardwareId = NULL;
	ULONG PropertySize = 0;

	UNREFERENCED_PARAMETER(ConnectionCookie);

	KLogEx(DEBUG_TRACE_INFO, "[Worker] Starting to query Hardware ID async...\n");

	// 1. 인스턴스 컨텍스트 가져오기
	status = FltGetInstanceContext(pWorkCtx->Instance, &pVolumeCtx);
	if (NT_SUCCESS(status))
	{
		// 2. 물리 디스크 객체 구하기 (이제 안전함!)
		// (이전에 만든 GetPhysicalDiskDeviceObject 함수 사용)
		status = GetPhysicalDiskDeviceObject(pWorkCtx->VolumeDeviceObject, &pPhysicalDisk);

		if (NT_SUCCESS(status) && pPhysicalDisk != NULL)
		{
			// 3. PDO 찾기
			pPhysicalPDO = IoGetDeviceAttachmentBaseRef(pPhysicalDisk);
			if (pPhysicalPDO)
			{
				// 4. Hardware ID 쿼리
				status = QueryDeviceProperty(pPhysicalPDO, DevicePropertyHardwareID, &pHardwareId, &PropertySize);
				if (NT_SUCCESS(status) && pHardwareId)
				{
					KLogEx(DEBUG_TRACE_INFO, "[Worker] Found Real Hardware ID: %S\n", pHardwareId);

					// TODO: pVolumeCtx에 ID 저장 (필요시 파싱)
					// 예: RtlStringCbCopyW(pVolumeCtx->HardwareId, ...);

					ExFreePool(pHardwareId);
				}
				ObDereferenceObject(pPhysicalPDO);
			}
			ObDereferenceObject(pPhysicalDisk);
		}
		else
		{
			KLogEx(DEBUG_TRACE_INFO, "[Worker] Failed to get Physical Disk: 0x%x\n", status);
		}

		// 컨텍스트 해제
		FltReleaseContext(pVolumeCtx);
	}

	// 3. 정리
	ObDereferenceObject(pWorkCtx->VolumeDeviceObject);
	FltObjectDereference(pWorkCtx->Instance); // 인스턴스 참조 해제
	FltFreeGenericWorkItem(FltWorkItem);
	ExFreePool(pWorkCtx);
}

NTSTATUS
GetVolumeProperties(
	__in PCFLT_RELATED_OBJECTS FltObjects,
	__in PVOLUME_CONTEXT pVolumeContext
)
{

	NTSTATUS status = STATUS_SUCCESS;
	PDEVICE_OBJECT pDevice = NULL;
	PVOID pBuffer = NULL;
	ULONG PropertySize = 0;
	//PWCHAR pDeviceId = NULL;
	//PDEVICE_OBJECT physicalDeviceObject = NULL;

	ASSERT(ARGUMENT_PRESENT(pVolumeContext));

	__try
	{
		status = FltGetDiskDeviceObject(FltObjects->Volume, &pDevice);
		if (!NT_SUCCESS(status))
		{
			pDevice = NULL;
			__leave;
		}

		//status = GetPhysicalDiskDeviceObject((PDEVICE_OBJECT)FltObjects->Volume, &pDevice);

		status = GetStorageProperty(pDevice, pVolumeContext);
		status = QueryDeviceProperty(pDevice, DevicePropertyRemovalPolicy, &pBuffer, &PropertySize);
		
		if (NT_SUCCESS(status))
		{
			PDEVICE_REMOVAL_POLICY pRemovalPolicy = (PDEVICE_REMOVAL_POLICY)pBuffer;
			pVolumeContext->removalpolicy = *pRemovalPolicy;
			//pVolumeContext->device_object = pDevice;
			KLogEx(DEBUG_TRACE_INFO, " REMOVAL_POLICY pDevice(%p),removalpolicy(%d)\n", pDevice, pVolumeContext->removalpolicy);
			ExFreePoolWithTag(pBuffer, _ALLOC_TAG);
			pBuffer = NULL;
		}
		else
		{
			KLogEx(DEBUG_TRACE_INFO, "Query DevicePropertyRemovalPolicy Failed 0x%x\n", status);
		}

		status = QueryDeviceProperty(pDevice, DevicePropertyFriendlyName, &pBuffer, &PropertySize);
		if (NT_SUCCESS(status))
		{
			KLogEx(DEBUG_TRACE_INFO, "DevicePropertyFriendlyName(%S)\n", (wchar_t*)pBuffer);
			ExFreePoolWithTag(pBuffer, _ALLOC_TAG);
			pBuffer = NULL;
		}
		else
		{
			KLogEx(DEBUG_TRACE_INFO, "Query DevicePropertyRemovalPolicy Failed 0x%x\n", status);
		}

		status= QueryDeviceProperty(pDevice, DevicePropertyDeviceDescription, &pBuffer, &PropertySize);
		if (NT_SUCCESS(status))
		{
			KLogEx(DEBUG_TRACE_INFO, "DevicePropertyDeviceDescription(%S)\n", (wchar_t*)pBuffer);
			ExFreePoolWithTag(pBuffer, _ALLOC_TAG);
			pBuffer = NULL;
		}
		else
		{
			KLogEx(DEBUG_TRACE_INFO, "Query DevicePropertyRemovalPolicy Failed 0x%x\n", status);
		}

		status = QueryDeviceProperty(pDevice, DevicePropertyHardwareID, &pBuffer, &PropertySize);
		if (NT_SUCCESS(status))
		{
			KLogEx(DEBUG_TRACE_INFO, "DevicePropertyHardwareID(%S)\n", (wchar_t*)pBuffer);
			ExFreePoolWithTag(pBuffer, _ALLOC_TAG);
			pBuffer = NULL;
		}
		else
		{
			KLogEx(DEBUG_TRACE_INFO, "Query DevicePropertyRemovalPolicy Failed 0x%x\n", status);
		}

		status = QueryDeviceProperty(pDevice, DevicePropertyEnumeratorName, &pBuffer, &PropertySize);
		if (NT_SUCCESS(status))
		{
			KLogEx(DEBUG_TRACE_INFO, "DevicePropertyEnumeratorName(%S)\n", (wchar_t*)pBuffer);
			ExFreePoolWithTag(pBuffer, _ALLOC_TAG);
			pBuffer = NULL;
		}
		else
		{
			KLogEx(DEBUG_TRACE_INFO, "Query DevicePropertyRemovalPolicy Failed 0x%x\n", status);
		}


		//status = GetStorageDeviceId(pDevice, &pDeviceId);
		//if (NT_SUCCESS(status) && pDeviceId != NULL)
		//{
		//	KLogEx(DEBUG_TRACE_INFO, "Found Device ID: %S\n", pDeviceId);

		//	// 결과 예시: "USB\VID_346D&PID_5678&REV_0100"
		//	// 여기서 문자열 파싱하여 정책 적용

		//	ExFreePoolWithTag(pDeviceId, 'IDst');
		//}
		//else
		//{
		//	KLogEx(DEBUG_TRACE_INFO, "[IOCTL] Failed 0x%x. Trying PnP Query...\n", status);

		//	// =================================================================
		//	// 시도 2: PnP 방식 (PDO에게 요청) - HardwareID / InstanceID
		//	// =================================================================
		//	// IOCTL을 처리하는 FDO의 가장 밑바닥(PDO)을 구합니다.
		//	physicalDeviceObject = IoGetDeviceAttachmentBaseRef(pDevice);

		//	if (physicalDeviceObject)
		//	{
		//		// PnP 쿼리는 PDO에게 해야 합니다.
		//		// DevicePropertyHardwareID (리스트) 또는 DevicePropertyInstanceID (단일) 시도
		//		PWCHAR pHardwareId = NULL;

		//		// 보통 HardwareID가 가장 확실하게 정보(VID/PID)를 줍니다.
		//		// HardwareID는 MULTI_SZ 형식이므로 첫 번째 문자열만 쓰면 됩니다.
		//		status = QueryDeviceProperty(physicalDeviceObject, DevicePropertyHardwareID, &pHardwareId, &PropertySize);

		//		if (NT_SUCCESS(status) && pHardwareId)
		//		{
		//			KLogEx(DEBUG_TRACE_INFO, "[PnP] Hardware ID: %S\n", pHardwareId);
		//			ExFreePoolWithTag(pHardwareId, _ALLOC_TAG);
		//		}
		//		else
		//		{
		//			KLogEx(DEBUG_TRACE_INFO, "[PnP] Query DevicePropertyHardwareID Failed 0x%x\n", status);
		//		}

		//		ObDereferenceObject(physicalDeviceObject);
		//	}
		//}

		status = STATUS_SUCCESS;
	}
	__finally
	{
		if (pDevice)
			ObDereferenceObject(pDevice);
	}

	return status;
}


// -----------------------------------------------------------------------------
// Helper: PnP 속성(Registry Property) 조회 함수
// 주의: 반드시 PDO(Physical Device Object)를 넘겨줘야 정확한 값이 나옵니다.
// -----------------------------------------------------------------------------
NTSTATUS
QueryDeviceRegistryProperty(
	__in PDEVICE_OBJECT pPdo,
	__in DEVICE_REGISTRY_PROPERTY DeviceProperty,
	__out PWCHAR pBuffer,
	__in ULONG BufferLength
)
{
	NTSTATUS status;
	ULONG ResultLength = 0;

	// 버퍼 초기화
	RtlZeroMemory(pBuffer, BufferLength);

	status = IoGetDeviceProperty(
		pPdo,
		DeviceProperty,
		BufferLength,
		pBuffer,
		&ResultLength
	);

	return status;
}

// -----------------------------------------------------------------------------
// Helper: 스토리지 속성(IOCTL) 조회 함수
// 주의: IOCTL을 처리할 수 있는 스택 상단(Top Device) 혹은 FDO에 보내야 합니다.
// -----------------------------------------------------------------------------
NTSTATUS
GetStorageHardwareProperty(
	__in PDEVICE_OBJECT pDeviceObject,
	__inout PFULL_DISK_INFO pInfo
)
{
	NTSTATUS status = STATUS_SUCCESS;
	PIRP Irp;
	KEVENT Event;
	IO_STATUS_BLOCK Iosb;
	STORAGE_PROPERTY_QUERY Query = { 0 };
	PSTORAGE_DEVICE_DESCRIPTOR pDesc = NULL;
	PVOID pBuffer = NULL;
	ULONG BufferSize = 1024; // 넉넉하게 잡음

	// 버퍼 할당
	pBuffer = ExAllocatePoolWithTag(PagedPool, BufferSize, _ALLOC_TAG);
	if (!pBuffer) return STATUS_INSUFFICIENT_RESOURCES;

	RtlZeroMemory(pBuffer, BufferSize);

	// 쿼리 설정
	Query.PropertyId = StorageDeviceProperty;
	Query.QueryType = PropertyStandardQuery;

	KeInitializeEvent(&Event, NotificationEvent, FALSE);

	// IOCTL 생성
	Irp = IoBuildDeviceIoControlRequest(
		IOCTL_STORAGE_QUERY_PROPERTY,
		pDeviceObject,
		&Query,
		sizeof(Query),
		pBuffer,
		BufferSize,
		FALSE,
		&Event,
		&Iosb
	);

	if (!Irp)
	{
		ExFreePoolWithTag(pBuffer, _ALLOC_TAG);
		return STATUS_INSUFFICIENT_RESOURCES;
	}

	// 드라이버 호출
	status = IoCallDriver(pDeviceObject, Irp);

	if (status == STATUS_PENDING)
	{
		KeWaitForSingleObject(&Event, Executive, KernelMode, FALSE, NULL);
		status = Iosb.Status;
	}

	if (NT_SUCCESS(status))
	{
		pDesc = (PSTORAGE_DEVICE_DESCRIPTOR)pBuffer;

		pInfo->BusType = pDesc->BusType;
		pInfo->RemovableMedia = pDesc->RemovableMedia;

		SafeSetVolumContextParam((PUCHAR)pInfo->VendorId, pDesc->Size, (PUCHAR)pDesc, pDesc->VendorIdOffset);
		SafeSetVolumContextParam((PUCHAR)pInfo->ProductId, pDesc->Size, (PUCHAR)pDesc, pDesc->ProductIdOffset);
		SafeSetVolumContextParam((PUCHAR)pInfo->Productrevisionlevel, pDesc->Size, (PUCHAR)pDesc, pDesc->ProductRevisionOffset);
		SafeSetVolumContextParam((PUCHAR)pInfo->SerialNumber, pDesc->Size, (PUCHAR)pDesc, pDesc->SerialNumberOffset);
	}

	ExFreePoolWithTag(pBuffer, _ALLOC_TAG);
	return status;
}

NTSTATUS
GetDeviceInstanceId(
	__in PDEVICE_OBJECT pDeviceObject,
	__in const DEVPROPKEY* pKey,
	__out PWCHAR pBuffer,
	__in ULONG BufferLength
)
{
	NTSTATUS status;
	//PDEVICE_OBJECT pPdo = NULL;
	ULONG returnLength = 0;
	DEVPROPTYPE propType;

	//// 1. PDO 구하기 (필수)
	//pPdo = IoGetDeviceAttachmentBaseRef(pDeviceObject);
	//if (pPdo == NULL)
	//{
	//	return STATUS_NO_SUCH_DEVICE;
	//}

	// 2. IoGetDevicePropertyData 호출 (Vista 이상 지원)
	// DEVPKEY_Device_InstanceId 키를 사용하여 인스턴스 경로를 요청합니다.
	status = IoGetDevicePropertyData(
		pDeviceObject,
		pKey,
		LOCALE_NEUTRAL,
		0,
		BufferLength,
		pBuffer,
		&returnLength,
		&propType
	);

	// 3. 참조 해제
	//ObDereferenceObject(pPdo);

	return status;
}

NTSTATUS
GetUsbIdFromVolume(
	__in PDEVICE_OBJECT pVolumeDeviceObject,
	__out PWCHAR pBuffer,
	__in ULONG BufferLength
)
{
	NTSTATUS status;
	PIRP pIrp;
	KEVENT event;
	IO_STATUS_BLOCK ioStatus;
	PVOLUME_DISK_EXTENTS pExtents = NULL;
	ULONG extentsSize = sizeof(VOLUME_DISK_EXTENTS) + sizeof(DISK_EXTENT) * 4; // 넉넉하게 잡음

	UNICODE_STRING diskName;
	WCHAR nameBuffer[64];
	PFILE_OBJECT pDiskFileObject = NULL;
	PDEVICE_OBJECT pDiskDeviceObject = NULL;
	PDEVICE_OBJECT pDiskPdo = NULL;

	// ---------------------------------------------------------
	// 1단계: 볼륨이 위치한 디스크 번호 알아내기 (IOCTL)
	// ---------------------------------------------------------
	pExtents = (PVOLUME_DISK_EXTENTS)ExAllocatePoolWithTag(NonPagedPool, extentsSize, 'Tag1');
	if (!pExtents) return STATUS_INSUFFICIENT_RESOURCES;

	KeInitializeEvent(&event, NotificationEvent, FALSE);

	// 볼륨 드라이버 스택의 최상단 장치 구하기 (IOCTL 전송용)
	PDEVICE_OBJECT pTopDevice = IoGetAttachedDeviceReference(pVolumeDeviceObject);

	pIrp = IoBuildDeviceIoControlRequest(
		IOCTL_VOLUME_GET_VOLUME_DISK_EXTENTS,
		pTopDevice,
		NULL, 0,
		pExtents, extentsSize,
		FALSE, &event, &ioStatus
	);

	if (pIrp == NULL) {
		ObDereferenceObject(pTopDevice);
		ExFreePool(pExtents);
		return STATUS_INSUFFICIENT_RESOURCES;
	}

	status = IoCallDriver(pTopDevice, pIrp);
	if (status == STATUS_PENDING) {
		KeWaitForSingleObject(&event, Executive, KernelMode, FALSE, NULL);
		status = ioStatus.Status;
	}

	ObDereferenceObject(pTopDevice); // 참조 해제

	if (!NT_SUCCESS(status) || pExtents->NumberOfDiskExtents == 0) {
		ExFreePool(pExtents);
		return status; // 디스크 정보를 얻을 수 없음
	}

	// ---------------------------------------------------------
	// 2단계: 디스크 번호로 물리 디스크 장치 객체(Pointer) 얻기
	// ---------------------------------------------------------
	// 보통 첫 번째 Extent가 해당 디스크입니다.
	ULONG diskNumber = pExtents->Extents[0].DiskNumber;
	ExFreePool(pExtents); // 더 이상 필요 없음

	// 디스크 장치 이름 생성 (예: \Device\Harddisk2\DR2)
	RtlStringCchPrintfW(nameBuffer, 64, L"\\Device\\Harddisk%d\\DR%d", diskNumber, diskNumber);
	RtlInitUnicodeString(&diskName, nameBuffer);

	status = IoGetDeviceObjectPointer(&diskName, FILE_READ_ATTRIBUTES, &pDiskFileObject, &pDiskDeviceObject);
	if (!NT_SUCCESS(status)) {
		return status;
	}

	// ---------------------------------------------------------
	// 3단계: 물리 디스크의 PDO 구하기 (핵심)
	// ---------------------------------------------------------
	pDiskPdo = IoGetDeviceAttachmentBaseRef(pDiskDeviceObject);

	// IoGetDeviceObjectPointer로 얻은 FileObject는 해제해야 함
	ObDereferenceObject(pDiskFileObject);

	if (!pDiskPdo) {
		return STATUS_NO_SUCH_DEVICE;
	}

	// ---------------------------------------------------------
	// 4단계: 드디어! USB\VID... 진짜 ID 쿼리
	// ---------------------------------------------------------
	// 기존에 만드신 함수나 직접 호출 사용
	// 여기서는 직접 호출 예시:
	ULONG returnLength = 0;
	DEVPROPTYPE propType;

	status = IoGetDevicePropertyData(
		pDiskPdo,
		&DEVPKEY_Device_InstanceId, // USB\VID_...가 들어있는 키
		LOCALE_NEUTRAL,
		0,
		BufferLength,
		pBuffer,
		&returnLength,
		&propType
	);

	// PDO 참조 해제
	ObDereferenceObject(pDiskPdo);

	return status;
}

NTSTATUS
GetIdFromDiskDeviceObject(
	__in PDEVICE_OBJECT pDiskDeviceObject,
	__out PWCHAR pBuffer,
	__in ULONG BufferLength
)
{
	NTSTATUS status;
	PDEVICE_OBJECT pDiskPdo = NULL;
	ULONG returnLength = 0;
	DEVPROPTYPE propType;

	// 1. 디스크 스택의 최하단 PDO (USBSTOR 등) 구하기
	pDiskPdo = IoGetDeviceAttachmentBaseRef(pDiskDeviceObject);
	if (!pDiskPdo) {
		return STATUS_NO_SUCH_DEVICE;
	}

	// 2. PDO에게 Instance ID 물어보기
	status = IoGetDevicePropertyData(
		pDiskPdo,
		&DEVPKEY_Device_InstanceId,
		LOCALE_NEUTRAL,
		0,
		BufferLength,
		pBuffer,
		&returnLength,
		&propType
	);

	// 3. 참조 해제
	ObDereferenceObject(pDiskPdo);

	return status;
}

VOID
SafeQueryIdRoutine(
	_In_ PFLT_GENERIC_WORKITEM FltWorkItem,
	_In_ PFLT_FILTER Filter,
	_In_ PVOID Context
)
{
	PFLT_VOLUME Volume = (PFLT_VOLUME)Context;
	PDEVICE_OBJECT pVolumeDeviceObject = NULL;
	NTSTATUS status;
	WCHAR buffer[200];
	PDEVICE_OBJECT pDiskPdo = NULL;
	ULONG returnLength = 0;
	DEVPROPTYPE propType;

	UNREFERENCED_PARAMETER(Filter);
	// 1. 볼륨의 DeviceObject 얻기
	status = FltGetDiskDeviceObject(Volume, &pVolumeDeviceObject);

	if (NT_SUCCESS(status)) {
		pDiskPdo = IoGetDeviceAttachmentBaseRef(pVolumeDeviceObject);
		if (pDiskPdo) {

			status = IoGetDevicePropertyData(
				pDiskPdo,
				&DEVPKEY_Device_InstanceId,
				LOCALE_NEUTRAL,
				0, sizeof(buffer), buffer, &returnLength, &propType
			);

			if (NT_SUCCESS(status)) {
				KLogEx(DEBUG_TRACE_INFO, "USB ID Found: %ws\n", buffer);
				// TODO: 볼륨 컨텍스트에 ID 저장
			}
			else
			{
				KLogEx(DEBUG_TRACE_INFO, "DEVPKEY_Device_InstanceId fail : %x\n", status);
			}
			
			status = IoGetDevicePropertyData(
				pDiskPdo,
				&DEVPKEY_Device_Parent, // 부모 키 사용
				LOCALE_NEUTRAL,
				0, sizeof(buffer), buffer, &returnLength, &propType
			);

			if (NT_SUCCESS(status)) {
				KLogEx(DEBUG_TRACE_INFO, "[Bs1Flt] Parent ID: %ws\n", buffer);
				// 결과 예시: USB\VID_346D&PID_5678\8328501217610604362
			}
			else {
				KLogEx(DEBUG_TRACE_INFO, "[Bs1Flt] DEVPKEY_Device_Parent Failed: 0x%08x\n", status);
			}

			ObDereferenceObject(pDiskPdo);

		}

		//status = GetUsbIdFromVolume(pVolumeDeviceObject, buffer, sizeof(buffer));

		//if (NT_SUCCESS(status)) {
		//	KLogEx(DEBUG_TRACE_INFO, "USB ID Found: %ws\n", buffer);
		//	// TODO: 볼륨 컨텍스트에 ID 저장
		//}
		//else
		//{
		//	KLogEx(DEBUG_TRACE_INFO, "GetUsbIdFromVolume : %x\n", status);
		//}

		//status = GetIdFromDiskDeviceObject(pVolumeDeviceObject, buffer, sizeof(buffer));

		//if (NT_SUCCESS(status)) {
		//	KLogEx(DEBUG_TRACE_INFO, "[Bs1Flt] USB ID Found: %ws\n", buffer);

		//	// TODO: 여기서 전역 리스트나 Context에 ID를 저장하고 정책 적용
		//	// CheckUsbPolicy(instanceId); 
		//}
		//else {
		//	KLogEx(DEBUG_TRACE_INFO, "[Bs1Flt] Failed to query ID from PDO: 0x%08x\n", status);
		//}

		ObDereferenceObject(pVolumeDeviceObject);
	}

	// 3. 정리
	FltObjectDereference(Volume); // 참조 해제
	FltFreeGenericWorkItem(FltWorkItem);
}


// [헬퍼 함수] 진짜 PDO를 구하는 함수 (필수 추가)
PDEVICE_OBJECT GetPdoFromDeviceObject(PDEVICE_OBJECT DeviceObject)
{
	PDEVICE_OBJECT pdo = NULL;
	PDEVICE_RELATIONS deviceRelations = NULL;
	IO_STATUS_BLOCK ioStatus;
	KEVENT event;
	PIRP irp;
	NTSTATUS status;

	if (!DeviceObject) return NULL;

	KeInitializeEvent(&event, NotificationEvent, FALSE);

	// 스택의 최상단 디바이스를 구해서 IRP를 보냅니다.
	PDEVICE_OBJECT targetDevice = IoGetAttachedDeviceReference(DeviceObject);

	irp = IoBuildSynchronousFsdRequest(IRP_MJ_PNP, targetDevice, NULL, 0, NULL, &event, &ioStatus);
	if (!irp) {
		ObDereferenceObject(targetDevice);
		return NULL;
	}

	irp->IoStatus.Status = STATUS_NOT_SUPPORTED;
	PIO_STACK_LOCATION irpSp = IoGetNextIrpStackLocation(irp);
	irpSp->MinorFunction = IRP_MN_QUERY_DEVICE_RELATIONS;
	irpSp->Parameters.QueryDeviceRelations.Type = TargetDeviceRelation;

	status = IoCallDriver(targetDevice, irp);
	if (status == STATUS_PENDING) {
		KeWaitForSingleObject(&event, Executive, KernelMode, FALSE, NULL);
		status = ioStatus.Status;
	}

	if (NT_SUCCESS(status) && ioStatus.Information) 
	{
		deviceRelations = (PDEVICE_RELATIONS)ioStatus.Information;
		if (deviceRelations->Count > 0) 
		{
			pdo = deviceRelations->Objects[0]; // 이것이 진짜 PDO
			KLogEx(DEBUG_TRACE_ERROR, "deviceRelations->Count(% d)\n", deviceRelations->Count);
			// 만약 Count가 1보다 크면 나머지는 해제 (매우 드문 케이스)
			for (ULONG i = 1; i < deviceRelations->Count; i++) {
				ObDereferenceObject(deviceRelations->Objects[i]);
			}
		}
		ExFreePool(deviceRelations);
	}

	ObDereferenceObject(targetDevice);
	return pdo;
}


// -----------------------------------------------------------------------------
// [MAIN] 디스크 정보 통합 수집 함수
// pDeviceObject: 미니필터에서 FltGetDiskDeviceObject 등으로 구한 디바이스 객체
// -----------------------------------------------------------------------------
NTSTATUS
CollectAllDiskInfo(
	__in PDEVICE_OBJECT pDeviceObject,
	__out PFULL_DISK_INFO pOutInfo
)
{
	//NTSTATUS status;
	PDEVICE_OBJECT pPdo = NULL;
	PDEVICE_OBJECT pTopDevice = NULL;
	

	RtlZeroMemory(pOutInfo, sizeof(FULL_DISK_INFO));

	// PDO 구하기 (QueryDeviceProperty용)
	// IoGetDeviceAttachmentBaseRef는 스택의 최하단(PDO)을 반환하며 레퍼런스 카운트를 증가시킵니다.
	//pPdo = IoGetDeviceAttachmentBaseRef(pDeviceObject);
	pPdo = GetPdoFromDeviceObject(pDeviceObject);
	if (!pPdo)
	{
		return STATUS_NO_SUCH_DEVICE;
	}

	// Top Device 구하기 (GetStorageProperty용)
	// IOCTL은 스택의 최상단으로 보내는 것이 가장 안전합니다.
	pTopDevice = IoGetAttachedDeviceReference(pDeviceObject);
	if (!pTopDevice)
	{
		ObDereferenceObject(pPdo);
		return STATUS_NO_SUCH_DEVICE;
	}

	__try
	{
		// =================================================================
		// PnP 레지스트리 속성 수집 (PDO 사용)
		// =================================================================
		if (pPdo != NULL)
		{
			QueryDeviceRegistryProperty(pPdo, DevicePropertyFriendlyName, pOutInfo->FriendlyName, sizeof(pOutInfo->FriendlyName));
			QueryDeviceRegistryProperty(pPdo, DevicePropertyDeviceDescription, pOutInfo->Description, sizeof(pOutInfo->Description));
			QueryDeviceRegistryProperty(pPdo, DevicePropertyHardwareID, pOutInfo->HardwareId, sizeof(pOutInfo->HardwareId));
			QueryDeviceRegistryProperty(pPdo, DevicePropertyEnumeratorName, pOutInfo->EnumeratorName, sizeof(pOutInfo->EnumeratorName));
			GetDeviceInstanceId(pPdo, &DEVPKEY_Device_InstanceId, pOutInfo->InstanceId, sizeof(pOutInfo->InstanceId));
		}
		else
		{
			// PDO를 못 구한 경우 (가상 드라이브 등) 로그 남김
			KLogEx(DEBUG_TRACE_ERROR, "CollectAllDiskInfo: No valid PDO found. Skipping PnP props.\n");
		}
		//GetDeviceInstanceId(pPdo, &DEVPKEY_Device_Parent, pOutInfo->ParentId, sizeof(pOutInfo->ParentId));
		//DEVPKEY_Device_Parent
		// DEVPKEY_Device_InstanceId
		// =================================================================
		// 하드웨어 물리 속성 수집 (Top Device 사용)
		// =================================================================
		GetStorageHardwareProperty(pTopDevice, pOutInfo);
	}
	__finally
	{
		// 사용한 객체 참조 해제 (필수)
		if (pPdo) ObDereferenceObject(pPdo);
		if (pTopDevice) ObDereferenceObject(pTopDevice);
	}

	return STATUS_SUCCESS;
}