/*
* PROJECT: Alcyone  Kernel
* LICENSE: BSD Clause 3 
* PURPOSE: Cache Controller:: MDL Handler   
* NT KERNEL: 5.11.9360 
* COPYRIGHT: 2023-2029 Dibymartanda Samanta <>
*             
*/
#include <ntoskrnl.h>
#define NTDEBUG
#include <debug.h>
#include "ccinternal.hpp"
#include "ccmdl.hpp"


extern "C"

/*Internal Function*/

VOID
NTAPI
CcMdlReadComplete2 (
    IN PFILE_OBJECT FileObject,
    IN PMDL MdlChain
)
{
	 /* Check if MDL Chain is Valid */
	if (MdlChain)
	{    
        /*Iterate Through the MDL CHain*/
        for( auto it = begin(MdlChain); it != end(MdlChain);++it)
		{    
	        /*Unlock the memory pages associated with the MDL*/
			MmUnlockPages(*it);
		    /*Free The MDL */
           IoFreeMdl(*it); 
		}
	}
}
CcMdlWriteComplete2(IN PFILE_OBJECT FileObject,
                    IN PLARGE_INTEGER FileOffset,
                    IN PMDL MdlChain)
{

PLARGE_INTEGER FOffset = nullptr;
NTSTATUS Exception = {0};
PIO_STATUS_BLOCK IoStatus = nullptr;


for( auto it = begin(MdlChain); it != end(MdlChain);++it)
{

    if (!(MdlChain->MdlFlags & 2) == 0) )
    {
      MmUnlockPages(*it);
     
    }
    if ( (FileObject->Flags & 16) != 0 )
    {
      MmFlushSection((*it)->ByteCount, SharedCacheMap, &IoStatus, 1);
      if ( !NTSTATUS(IoStatus.Status) )
        Exception = IoStatus.Status;
    }
    else
    {
      CcSetDirtyInMask(SharedCacheMap, &FOffset, (*it)->ByteCount);
    }
    FOffset.QuadPart += (*it)->ByteCount;
}
  
  if (!(MdlChain->MdlFlags & 2) == 0))
  {
    SpinLockGuard lockGuard(LockQueueMasterLock);
    CcDecrementOpenCount(SharedCacheMap);
    
  }
  if (!NT_SUCCESS(Exception))                          
  {
    RtlRaiseStatus(FsRtlNormalizeNtstatus(Exception, WRITE_EXCEPTION));
  }
  /* Write Operation is COmplete, Now Free the MDL Chain*/
for( auto it = begin(MdlChain); it != end(MdlChain);++it)
  {    
		    /*Free The MDL */
           IoFreeMdl(*it); 
  }
}

/* External Function */ 
VOID
NTAPI
CcMdlRead (
    IN PFILE_OBJECT FileObject,
    IN PLARGE_INTEGER FileOffset,
    IN ULONG Length,
    OUT PMDL * MdlChain,
    OUT PIO_STATUS_BLOCK IoStatus)
{
		
	ULONG LENGTH_FINAL = Length;
	PVOID VirtualAddress = nullptr;
	PVACB VACB  = nullptr;
	ULONG ReceivedLength = {0};
	ULONG Information = {0};
	PMDL Allocated = nullptr;
	LARGE_INTEGER FOffset = {0};
	PSHARED_CACHE_MAP sharedcachmap = FileObject->SectionObjectPointer->SharedCacheMap; 
	PRIVATE_CACHE_MAP PrivateCacheMap = FileObject->PrivateCacheMap;
	PVOID VirtualAddress = nullptr;
	LARGE_INTEGER Offset_Lenghth = {0};


_try
{

/* Iterate until it is final length is null*/ 	
	while(Length != NULL)
	{
		
		VirtualAddress = static_cast<void*>(CcGetVirtualAddress(sharedcachmap, &VACB, FOffset,(LARGE_INTEGER)&ReceivedLength, 0));
		INT PagePriority = static_cast<int>(((&sharedcachmap->NodeTypeCode >> 18) & 7));
		
		if ((sharedcachmap->Flags & 8)== null )
		{
			/* Our Memory Manager doesn't support prefetching yet, will be added in 2025*/
		//CcFetchDataForRead(FileObject,&FOffset,Length,Wait_Type_MDL_Read,&FaultedData,&VACB,PagePriority,MDL_READTYPE_DiskIoAttribution); // Return
		}
		
		if ( ReceivedLength > Length )
		{   
			ReceivedLength = Length;
		}
	
	  
	  Offset_Lenghth =  FOffset.QuadPart +  static_cast<LONGLONG>(ReceivedLength);
		
	
		/* Allocate the MDL ,it is not associated with Specific IRP*/
		Allocated = IoAllocateMdl(VirtualAddress,ReceivedLength , false, false, optional_irp);
		
		/* Assign the MDL to MmProbeAndLockPages to prevent Page Fault/ Data Corruption */
		MmProbeAndLockPages(Allocated, KernelMode, IoReadAccess);
		/* Free the VACB */
		CcFreeVirtualAddress(VACB);
		
		/* Now Update the MDL Chain with allocated one */

/* Check if Valid MDL Chain already exist, if so update it with allocated one else assign MDL Chain with one we allocated  */
   if ( *MdlChain != NULL)
    {   for( auto it = begin(MdlChain); it != end(MdlChain);++it)
		{
			*it = Allocated;
		}
    }
	else
    {
      *MdlChain = Allocated;
    }
		
	
	
	FOffset.QuadPart = Offset_Lenghth.QuadPart;
	
	/* Update the IOSTATUS Information FLag */
	Information += ReceivedLength;
	/* Adjust the length , else you will be in mess of a infinite nonsensical loop fighting with Exception Handler */ 
	Length -= ReceivedLength;
}
	
}

_finally{
	
 if ( (FileObject->Flags & 0x100000) = = 0 ||(&PrivateCacheMap->NodeType_UlongFlags & 0x20000 == 0) && FaultedData )
 {
    CcScheduleReadAheadEx(FileObject, FileOffset,Length,NULL);
	 _InterlockedOr(&PrivateCacheMap->NodeType_UlongFlags, FLAG_UPDATE_READAHEAD);
   
 }

  
/* No mater What , Catch History must be updated even if Aliens Invasion Happens that day*/
/* UPDATE CATCH HISTORY */
InterlockedExchange64(&PrivateCachemap->FileOffset1,PrivateCachemap->FileOffset2);
InterlockedExchange64(&PrivateCachemap->BeyondLastByte1,PrivateCachemap->BeyondLastByte2);
InterlockedExchange64(&PrivateCachemap->FileOffset2,*FileOffset);
LONG LONG FinalLength =FileOffset->QuadPart+ static_cast<LONGLONG>(Length);
InterlockedExchange64(&PrivateCachemap->BeyondLastByte2.Quadpart,FinalLength);

/* Update Shared Cachemap Flags */
if ( (SharedCacheMap->Flags & 8) != 0 )
{
	  SpinLockGuard lockGuard(LockQueueMasterLock);
    SharedCacheMap->Flags &= ~8;
    
}

/* Update IRP Bytes */ 
IoStatus->Status = Status_Sucess;
IoStatus->Information = Information;
}
}

VOID
NTAPI
CcMdlReadComplete(IN PFILE_OBJECT FileObject,
                  IN PMDL MdlChain)
{
    PDEVICE_OBJECT RelatedDeviceObject = nullptr;
    PFAST_IO_DISPATCH FastIoDispatch = nullptr;
    IO_MDLREAD_COMPLETE_ROUTINE MdlReadComplete;

    RelatedDeviceObject = IoGetRelatedDeviceObject(FileObject);
    FastIoDispatch = RelatedDeviceObject->DriverObject->FastIoDispatch;
    if(MdlChain == nullptr)
      {   DBGPRINT("Invalid MDLChain \n");
           return;
      }
    if (FastIoDispatch &&
        FastIoDispatch->SizeOfFastIoDispatch >= sizeof(FAST_IO_DISPATCH) &&
        (MdlReadComplete = FastIoDispatch->MdlReadComplete) != NULL &&
        MdlReadComplete(FileObject, MdlChain, RelatedDeviceObject))
    {
        /* Faster Fast I/O MDL path successful, if it fails, we will go for our slower implementation*/
        return;
    }

    CcMdlReadComplete2(FileObject, MdlChain);
}

VOID
NTAPI
CcMdlWriteAbort(IN PFILE_OBJECT FileObject,
                IN PMDL MdlChain)
{
PSHARED_CACHE_MAP SharedCacheMap = FileObject->SectionObjectPointer->SharedCacheMap;


for( auto it = begin(MdlChain); it != end(MdlChain);++it)
{
    if ((MdlChain->MdlFlags & 2) != 0 )
   {
     MmUnlockPages(*it);
   }
    IoFreeMdl(*it);
}
if ( (MdlChain->MdlFlags & 2) != 0 )
{
    SpinLockGuard lockGuard(LockQueueMasterLock);
    CcDecrementOpenCount(SharedCacheMap);
    
}
}

VOID
NTAPI
CcMdlWriteComplete(IN PFILE_OBJECT FileObject,
                   IN PLARGE_INTEGER FileOffset,
                   IN PMDL MdlChain)
{
	
	  PDEVICE_OBJECT RelatedDeviceObject = nullptr;
    PFAST_IO_DISPATCH FastIoDispatch = nullptr;
    IO_MDLWRITE_COMPLETE_ROUTINE MdlWriteComplete;
    if(MdlChain == nullptr)
    {   DBGPRINT("Invalid MDLChain \n");
    return;
    }
    RelatedDeviceObject = IoGetRelatedDeviceObject(FileObject);
    FastIoDispatch = RelatedDeviceObject->DriverObject->FastIoDispatch;

    if (FastIoDispatch &&
        FastIoDispatch->SizeOfFastIoDispatch >= sizeof(FAST_IO_DISPATCH) &&
        (MdlWriteComplete = FastIoDispatch->MdlWriteComplete) != NULL &&
        MdlWriteComplete(FileObject, FileOffset, MdlChain, RelatedDeviceObject))
    {
        /* Faster Fast I/O MDL path successful, if it fails, we will go for our slower implementation*/
        return;
    }

    CcMdlWriteComplete2(FileObject, FileOffset, MdlChain);
}

VOID
NTAPI
CcPrepareMdlWrite(IN PFILE_OBJECT FileObject,
                  IN PLARGE_INTEGER FileOffset,
                  IN ULONG Length,
                  OUT PMDL *MdlChain,
                  OUT PIO_STATUS_BLOCK IoStatus)
{
  PMDL Mdl = nullptr;
  PVOID VirtualAddress = nullptr;
  PSHARED_CACHE_MAP SharedCacheMap = FileObject->SectionObjectPointer->SharedCacheMap;
  PLARGE_INTEGER FOffset = nullptr;
  PLARGE_INTEGER ReadAheadLenghth = nullptr;
  LARGE_INTEGER Final_Length = {0};
  LARGE_INTEGER ReadAhead_Length =  {0};
  LARGE_INTEGER Information =  {0};
  ULONG ReceivedLength =  {0};
  ULONG SavedState =  {0};
  ULONG ZeroOpsFlags =  {0};
  PVACB Vacb = nullptr;
  BOOL MapFlag = false;
  PKTHREAD currentThread = nullptr;
  UCHAR  SavedState ={0};
  if ((FileObject->Flags & 0x10) == 0)
    {
        if (IoIsFileOriginRemote(FileObject) && !CcCanIWriteStream(FileObject, Length, 0, 0))
            RtlRaiseStatus(STATUS_INSUFFICIENT_RESOURCES );
    }
 _try{

    while (Length != NULL)
    {
        VirtualAddress = CcGetVirtualAddress(SharedCacheMap, &ReceivedLength, &Vacb, 0, FOffset, ReadAheadLenghth);
		
	     /* Should not exceed PageLimit */
        if (ReceivedLength > Length)
            ReceivedLength = Length;
        Final_Length.QuadPart =  FileOffset->QuadPart;
        ReadAhead_Length = (ReadAheadLenghth->QuadPart) + (FileOffset->QuadPart);
        ZeroOpsFlags = 2;
        /* Set the flags for CcMapAndRead */
        if ((FOffset->Lowpart & 0xFFF) == 0 && ReceivedLength >= 0x1000)
        {
            ZeroOpsFlags = 3;
        }
        if ((Final_Length & 0xFFF) == 0)
        {
            ZeroOpsFlags |= 4;
        }
        KeAcquireInStackQueuedSpinLock(&SharedCacheMap->BcbSpinLock, &LockHandle);
        LONGLONG ReadAhead_State = SharedCacheMap->ValidDataGoal.QuadPart - (ReadAheadLenghth << 12 | (FOffset->Lowpart & 0xFFFFF000));
        ReadAheadLenghth = static_cast<ULONG>(ReadAhead_State >> 32);
        KeReleaseInStackQueuedSpinLock(&LockHandle);
        if (ReadAheadLenghth > 0)
        {
            SavedState = static_cast<ULONG>(ReadAhead_State);
        }
        else if ((ReadAheadLenghth & 0x80000000) != 0 || (SavedState = static_cast<ULONG>(ReadAhead_State)) == 0)
        {
            ZeroOpsFlags |= 7;
        }
        if (!ReadAheadLenghth && SavedState <= 0x1000)
        {
            ZeroOpsFlags |= 6;
        }
        MapFlag = CcMapAndRead(ReceivedLength, ZeroOpsFlags, 1, VirtualAddress);
		    if(!MapFlag)
		    {
			      DBGPRINT(" CcPrepareMdlWrite:: Page Maping Failed \n");
		    }
        Mdl = IoAllocateMdl(VirtualAddress, ReceivedLength, FALSE, FALSE, NULL);
        if (!Mdl)
            RtlRaiseStatus(STATUS_INSUFFICIENT_RESOURCES );
      
        currentThread = KeGetCurrentThread();
        SavedState = currentThread->UserIdealProcessor ;
        savedDisablePageFaultClustering = currentThread->DisablePageFaultClustering;
        currentThread->DisablePageFaultClustering = 1;
        MmProbeAndLockPages(Mdl, KernelMode, IoWriteAccess);
        currentThread->UserIdealProcessor = SavedState;
        currentThread->DisablePageFaultClustering = savedDisablePageFaultClustering;
        /* Update the Read Aheas Stats*/
        KeAcquireInStackQueuedSpinLock(&SharedCacheMap->BcbSpinLock, &LockHandle);
        if ( ReadAhead_Length.QuadPart > SharedCacheMap->ValidDataGoal.QuadPart )
        {
          SharedCacheMap->ValidDataGoal.QuadPart = ReadAhead_Length.QuadPart;
        }
        KeReleaseInStackQueuedSpinLock(&LockHandle);
        /*Free The VACB and Dereferance it*/
        CcFreeVirtualAddress(Vacb);
        Vacb = nullptr;
        /* Check if Valid MDL Chain already exist, if so update it with allocated one else assign MDL Chain with one we allocated  */
        if (*MdlChain)
        {
          /*MDL Iterator in Action, While Loopers take your C++ Coolaid now */
          for(auto it = begin(MdlChain); it != end(MdlChain);++it)
          {
            *it = Allocated;
          }
        }
        else
        {
          *MdlChain = Mdl;
        }
        Mdl = nullptr;
        FOffset.QuadPart = Final_Length;
        ReadAheadLenghth = ReadAhead_Length;
        /* Update the IOSTATUS Information FLag with amount of Byte Counts */
        Information += ReceivedLength;
        /* Adjust the length as per bytes written, else you will be in mess of a infinite nonsensical loop, Might result in Kaboom aka BSOD */
        Length -= ReceivedLength;
    }
 }
 _finally{
    IoStatus->Status = STATUS_SUCESS;
    IoStatus->Information = Information;
    KIRQL LAST_IRQL = KeAcquireQueuedSpinLock(LockQueueMasterLock);
    SharedCacheMap->OpenCount++;
    KeReleaseQueuedSpinLock(LockQueueMasterLock, LAST_IRQL);
}