text: /** @file
    Copyright (C) 2019, vit9696. All rights reserved.
  
    All rights reserved.
  
    This program and the accompanying materials
    are licensed and made available under the terms and conditions of the BSD License
    which accompanies this distribution.  The full text of the license may be found at
    http://opensource.org/licenses/bsd-license.php
  
    THE PROGRAM IS DISTRIBUTED UNDER THE BSD LICENSE ON AN "AS IS" BASIS,
    WITHOUT WARRANTIES OR REPRESENTATIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED.
  **/
  
  #include <Uefi.h>
  
  #include <Library/BaseLib.h>
  #include <Library/BaseMemoryLib.h>
  #include <Library/DebugLib.h>
  #include <Library/MemoryAllocationLib.h>
  #include <Library/UefiBootServicesTableLib.h>
  #include <Library/OcGuardLib.h>
  #include <Library/OcVirtualFsLib.h>
  
  #include <Guid/FileInfo.h>
  
  #include "VirtualFsInternal.h"
  
  STATIC EFI_HANDLE_PROTOCOL mOriginalHandleProtocol;
  STATIC EFI_LOCATE_PROTOCOL mOriginalLocateProtocol;
  STATIC EFI_FILE_OPEN       mOpenCallback;
  STATIC UINT32              mEntranceCount;
  
  STATIC
  VOID
  VirtualFsWrapProtocol (
    IN  EFI_GUID          *Protocol,
    OUT VOID              **Interface
    )
  {
    EFI_STATUS                       Status;
    EFI_SIMPLE_FILE_SYSTEM_PROTOCOL  *FileSystem;
  
    if (!CompareGuid (Protocol, &gEfiSimpleFileSystemProtocolGuid)) {
      return;
    }
  
    Status = CreateVirtualFs (*Interface, mOpenCallback, &FileSystem);
    if (!EFI_ERROR (Status)) {
      *Interface = FileSystem;
    }
  }
  
  STATIC
  EFI_STATUS
  EFIAPI
  VirtualFsHandleProtocol (
    IN  EFI_HANDLE        Handle,
    IN  EFI_GUID          *Protocol,
    OUT VOID              **Interface
    )
  {
    EFI_STATUS  Status;
  
    Status = mOriginalHandleProtocol (Handle, Protocol, Interface);
  
    if (!EFI_ERROR (Status) && Interface != NULL && mEntranceCount == 0) {
      ++mEntranceCount;
      VirtualFsWrapProtocol (Protocol, Interface);
      --mEntranceCount;
    }
  
    return Status;
  }
  
  STATIC
  EFI_STATUS
  EFIAPI
  VirtualFsLocateProtocol (
    IN  EFI_GUID  *Protocol,
    IN  VOID      *Registration OPTIONAL,
    OUT VOID      **Interface
    )
  {
    EFI_STATUS  Status;
  
    Status = mOriginalLocateProtocol (Protocol, Registration, Interface);
  
    if (!EFI_ERROR (Status) && Interface != NULL && mEntranceCount == 0) {
      ++mEntranceCount;
      VirtualFsWrapProtocol (Protocol, Interface);
      --mEntranceCount;
    }
  
    return Status;
  }
  
  EFI_STATUS
  EnableVirtualFs (
    IN OUT EFI_BOOT_SERVICES       *BootServices,
    IN     EFI_FILE_OPEN           OpenCallback
    )
  {
    if (mOriginalHandleProtocol != NULL
      || BootServices->HandleProtocol == VirtualFsHandleProtocol) {
      return EFI_ALREADY_STARTED;
    }
  
    //
    // Note, we should normally also hook OpenProtocol, but it results in freezes on APTIO IV and V.
    //
    mOpenCallback                = OpenCallback;
    mOriginalHandleProtocol      = BootServices->HandleProtocol;
    mOriginalLocateProtocol      = BootServices->LocateProtocol;
    BootServices->HandleProtocol = VirtualFsHandleProtocol;
    BootServices->LocateProtocol = VirtualFsLocateProtocol;
    BootServices->Hdr.CRC32      = 0;
    BootServices->CalculateCrc32 (BootServices, BootServices->Hdr.HeaderSize, &BootServices->Hdr.CRC32);
  
    return EFI_SUCCESS;
  }
  
  EFI_STATUS
  DisableVirtualFs (
    IN OUT EFI_BOOT_SERVICES       *BootServices
    )
  {
    if (mOriginalHandleProtocol == NULL
      || BootServices->HandleProtocol != VirtualFsHandleProtocol) {
      return EFI_ALREADY_STARTED;
    }
  
    BootServices->HandleProtocol = mOriginalHandleProtocol;
    BootServices->LocateProtocol = mOriginalLocateProtocol;
    mOpenCallback                = NULL;
    mOriginalHandleProtocol      = NULL;
    mOriginalLocateProtocol      = NULL;
    BootServices->Hdr.CRC32      = 0;
    BootServices->CalculateCrc32 (BootServices, BootServices->Hdr.HeaderSize, &BootServices->Hdr.CRC32);
  
    return EFI_SUCCESS;
  }
  
  