/** @file
  Copyright (C) 2019, Goldfish64. 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/BaseOverflowLib.h>
#include <Library/DebugLib.h>
#include <Library/MemoryAllocationLib.h>
#include <Library/OcAppleDiskImageLib.h>
#include <Library/OcXmlLib.h>

#include "OcAppleDiskImageLibInternal.h"

STATIC
BOOLEAN
InternalFindPlistDictChild (
  IN  XML_NODE  *Node,
  IN  CHAR8     *KeyName,
  OUT XML_NODE  **Key,
  OUT XML_NODE  **Value
  )
{
  UINT32       ChildCount;
  XML_NODE     *ChildValue;
  XML_NODE     *ChildKey;
  CONST CHAR8  *ChildKeyName;
  UINT32       Index;

  ASSERT (Node != NULL);
  ASSERT (KeyName != NULL);
  ASSERT (Key != NULL);
  ASSERT (Value != NULL);

  ChildCount = PlistDictChildren (Node);
  for (Index = 0; Index < ChildCount; ++Index) {
    ChildKey = PlistDictChild (Node, Index, &ChildValue);
    if (ChildKey == NULL) {
      break;
    }

    ChildKeyName = PlistKeyValue (ChildKey);
    if (ChildKeyName == NULL) {
      break;
    }

    if (AsciiStrCmp (ChildKeyName, KeyName) == 0) {
      *Key   = ChildKey;
      *Value = ChildValue;
      return TRUE;
    }
  }

  return FALSE;
}

STATIC
BOOLEAN
InternalSwapBlockData (
  IN OUT APPLE_DISK_IMAGE_BLOCK_DATA  *BlockData,
  IN     UINT32                       MaxSize,
  IN     UINTN                        SectorCount,
  IN     UINTN                        DataForkOffset,
  IN     UINTN                        DataForkSize
  )
{
  UINT32                  ChunksSize;
  UINTN                   MaxOffset;
  BOOLEAN                 Result;
  APPLE_DISK_IMAGE_CHUNK  *Chunk;
  UINT32                  Index;
  UINT64                  BlockSectorTop;
  UINT64                  ChunkSectorTop;
  UINT64                  OffsetTop;

  ASSERT (MaxSize >= sizeof (*BlockData));

  BlockData->ChunkCount = SwapBytes32 (BlockData->ChunkCount);

  Result = BaseOverflowMulU32 (
             BlockData->ChunkCount,
             sizeof (*BlockData->Chunks),
             &ChunksSize
             );
  if (Result || (ChunksSize > (MaxSize - sizeof (*BlockData)))) {
    return FALSE;
  }

  MaxOffset = (DataForkOffset + DataForkSize);

  BlockData->Version          = SwapBytes32 (BlockData->Version);
  BlockData->SectorNumber     = SwapBytes64 (BlockData->SectorNumber);
  BlockData->SectorCount      = SwapBytes64 (BlockData->SectorCount);
  BlockData->DataOffset       = SwapBytes64 (BlockData->DataOffset);
  BlockData->BuffersNeeded    = SwapBytes32 (BlockData->BuffersNeeded);
  BlockData->BlockDescriptors = SwapBytes32 (BlockData->BlockDescriptors);
  BlockData->Checksum.Type    = SwapBytes32 (BlockData->Checksum.Type);
  BlockData->Checksum.Size    = SwapBytes32 (BlockData->Checksum.Size);

  if (  (BlockData->DataOffset > DataForkOffset)
     || (BlockData->Checksum.Size > (sizeof (BlockData->Checksum.Data) * 8))
     || (BlockData->DataOffset > MaxOffset))
  {
    return FALSE;
  }

  Result = BaseOverflowAddU64 (
             BlockData->SectorNumber,
             BlockData->SectorCount,
             &BlockSectorTop
             );
  if (Result || (BlockSectorTop > SectorCount)) {
    DEBUG ((DEBUG_ERROR, "OCDI: Block sectors exceed DMG sectors %lu %lu\n", BlockSectorTop, SectorCount));
    return FALSE;
  }

  for (Index = 0; Index < APPLE_DISK_IMAGE_CHECKSUM_SIZE; ++Index) {
    BlockData->Checksum.Data[Index] = SwapBytes32 (
                                        BlockData->Checksum.Data[Index]
                                        );
  }

  for (Index = 0; Index < BlockData->ChunkCount; ++Index) {
    Chunk = &BlockData->Chunks[Index];

    Chunk->Type             = SwapBytes32 (Chunk->Type);
    Chunk->Comment          = SwapBytes32 (Chunk->Comment);
    Chunk->SectorNumber     = SwapBytes64 (Chunk->SectorNumber);
    Chunk->SectorCount      = SwapBytes64 (Chunk->SectorCount);
    Chunk->CompressedOffset = SwapBytes64 (Chunk->CompressedOffset);
    Chunk->CompressedLength = SwapBytes64 (Chunk->CompressedLength);

    Result = BaseOverflowAddU64 (
               Chunk->SectorNumber,
               Chunk->SectorCount,
               &ChunkSectorTop
               );
    if (Result || (ChunkSectorTop > BlockSectorTop)) {
      return FALSE;
    }

    Result = BaseOverflowAddU64 (
               Chunk->CompressedOffset,
               Chunk->CompressedLength,
               &OffsetTop
               );
    if (Result || (OffsetTop > MaxOffset)) {
      return FALSE;
    }
  }

  return TRUE;
}

BOOLEAN
InternalParsePlist (
  IN  CHAR8                        *Plist,
  IN  UINT32                       PlistSize,
  IN  UINTN                        SectorCount,
  IN  UINTN                        DataForkOffset,
  IN  UINTN                        DataForkSize,
  OUT UINT32                       *BlockCount,
  OUT APPLE_DISK_IMAGE_BLOCK_DATA  ***Blocks
  )
{
  BOOLEAN  Result;

  XML_DOCUMENT  *XmlPlistDoc;
  XML_NODE      *NodeRoot;
  XML_NODE      *NodeResourceForkKey;
  XML_NODE      *NodeResourceForkValue;
  XML_NODE      *NodeBlockListKey;
  XML_NODE      *NodeBlockListValue;

  XML_NODE  *NodeBlockDict;
  XML_NODE  *BlockDictChildKey;
  XML_NODE  *BlockDictChildValue;
  UINT32    BlockDictChildDataSize;

  UINT32                       NumDmgBlocks;
  UINT32                       DmgBlocksSize;
  APPLE_DISK_IMAGE_BLOCK_DATA  **DmgBlocks;
  APPLE_DISK_IMAGE_BLOCK_DATA  *Block;

  UINT32  Index;

  ASSERT (Plist != NULL);
  ASSERT (PlistSize > 0);
  ASSERT (BlockCount != NULL);
  ASSERT (Blocks != NULL);

  DmgBlocks = NULL;

  XmlPlistDoc = NULL;

  XmlPlistDoc = XmlDocumentParse (Plist, PlistSize, FALSE);
  if (XmlPlistDoc == NULL) {
    Result = FALSE;
    goto DONE_ERROR;
  }

  NodeRoot = PlistDocumentRoot (XmlPlistDoc);
  if (NodeRoot == NULL) {
    Result = FALSE;
    goto DONE_ERROR;
  }

  Result = InternalFindPlistDictChild (
             NodeRoot,
             DMG_PLIST_RESOURCE_FORK_KEY,
             &NodeResourceForkKey,
             &NodeResourceForkValue
             );
  if (!Result) {
    goto DONE_ERROR;
  }

  Result = InternalFindPlistDictChild (
             NodeResourceForkValue,
             DMG_PLIST_BLOCK_LIST_KEY,
             &NodeBlockListKey,
             &NodeBlockListValue
             );
  if (!Result) {
    goto DONE_ERROR;
  }

  NumDmgBlocks = XmlNodeChildren (NodeBlockListValue);
  if (NumDmgBlocks == 0) {
    Result = FALSE;
    goto DONE_ERROR;
  }

  Result = !BaseOverflowMulU32 (NumDmgBlocks, sizeof (*DmgBlocks), &DmgBlocksSize);
  if (!Result) {
    ///< Result must be FALSE on error, it's checked at DONE_ERROR
    goto DONE_ERROR;
  }

  DmgBlocks = AllocatePool (DmgBlocksSize);
  if (DmgBlocks == NULL) {
    Result = FALSE;
    goto DONE_ERROR;
  }

  for (Index = 0; Index < NumDmgBlocks; ++Index) {
    NodeBlockDict = XmlNodeChild (NodeBlockListValue, Index);

    Result = InternalFindPlistDictChild (
               NodeBlockDict,
               DMG_PLIST_DATA,
               &BlockDictChildKey,
               &BlockDictChildValue
               );
    if (!Result) {
      goto DONE_ERROR;
    }

    Result = PlistDataSize (BlockDictChildValue, &BlockDictChildDataSize);
    if (!Result || (BlockDictChildDataSize < sizeof (*Block))) {
      Result = FALSE;
      goto DONE_ERROR;
    }

    Block = AllocatePool (BlockDictChildDataSize);
    if (Block == NULL) {
      Result = FALSE;
      goto DONE_ERROR;
    }

    DmgBlocks[Index] = Block;

    Result = PlistDataValue (
               BlockDictChildValue,
               (UINT8 *)Block,
               &BlockDictChildDataSize
               );
    if (!Result) {
      FreePool (Block);
      goto DONE_ERROR;
    }

    Result = InternalSwapBlockData (
               Block,
               BlockDictChildDataSize,
               SectorCount,
               DataForkOffset,
               DataForkSize
               );
    if (!Result) {
      FreePool (Block);
      goto DONE_ERROR;
    }
  }

  *BlockCount = NumDmgBlocks;
  *Blocks     = DmgBlocks;
  Result      = TRUE;

DONE_ERROR:
  if (!Result && (DmgBlocks != NULL)) {
    while ((Index--) != 0) {
      FreePool (DmgBlocks[Index]);
    }

    FreePool (DmgBlocks);
  }

  if (XmlPlistDoc != NULL) {
    XmlDocumentFree (XmlPlistDoc);
  }

  return Result;
}

BOOLEAN
InternalGetBlockChunk (
  IN  OC_APPLE_DISK_IMAGE_CONTEXT  *Context,
  IN  UINTN                        Lba,
  OUT APPLE_DISK_IMAGE_BLOCK_DATA  **Data,
  OUT APPLE_DISK_IMAGE_CHUNK       **Chunk
  )
{
  UINT32                       BlockIndex;
  UINT32                       ChunkIndex;
  APPLE_DISK_IMAGE_BLOCK_DATA  *BlockData;
  APPLE_DISK_IMAGE_CHUNK       *BlockChunk;

  for (BlockIndex = 0; BlockIndex < Context->BlockCount; ++BlockIndex) {
    BlockData = Context->Blocks[BlockIndex];

    if (  (Lba >= BlockData->SectorNumber)
       && (Lba < (BlockData->SectorNumber + BlockData->SectorCount)))
    {
      for (ChunkIndex = 0; ChunkIndex < BlockData->ChunkCount; ++ChunkIndex) {
        BlockChunk = &BlockData->Chunks[ChunkIndex];

        if (  (Lba >= DMG_SECTOR_START_ABS (BlockData, BlockChunk))
           && (Lba < (DMG_SECTOR_START_ABS (BlockData, BlockChunk) + BlockChunk->SectorCount)))
        {
          *Data  = BlockData;
          *Chunk = BlockChunk;
          return TRUE;
        }
      }
    }
  }

  return FALSE;
}
