binhive.c

来自「一个类似windows」· C语言 代码 · 共 1,660 行 · 第 1/3 页

		}
	      RegistryHive->FreeListSize--;

	      return TRUE;
	    }
	  else if (BlockOffset + BlockSize == FreeOffset)
	    {
	      DbgPrint((DPRINT_REGISTRY, "Merge current block with previous block\n"));

	      RegistryHive->FreeList[i]->CellSize += FreeBlock->CellSize;
	      FreeBlock->CellSize = 0;

	      return TRUE;
	    }
	  else if (FreeOffset + FreeBlock->CellSize == BlockOffset)
	    {
	      DbgPrint((DPRINT_REGISTRY, "Merge current block with next block\n"));

	      FreeBlock->CellSize += RegistryHive->FreeList[i]->CellSize;
	      RegistryHive->FreeList[i]->CellSize = 0;
	      RegistryHive->FreeList[i] = FreeBlock;
	      RegistryHive->FreeListOffset[i] = FreeOffset;

	      return TRUE;
	    }
	}
    }

  return FALSE;
}


static BOOL
CmiAddFree(PREGISTRY_HIVE RegistryHive,
	   PCELL_HEADER FreeBlock,
	   BLOCK_OFFSET FreeOffset,
	   BOOL MergeFreeBlocks)
{
  PCELL_HEADER *tmpList;
  BLOCK_OFFSET *tmpListOffset;
  LONG minInd;
  LONG maxInd;
  LONG medInd;

  assert(RegistryHive);
  assert(FreeBlock);

  DbgPrint((DPRINT_REGISTRY, "FreeBlock %.08lx  FreeOffset %.08lx\n",
	   FreeBlock, FreeOffset));

  /* Merge free blocks */
  if (MergeFreeBlocks == TRUE)
    {
      if (CmiMergeFree(RegistryHive, FreeBlock, FreeOffset))
	return TRUE;
    }

  if ((RegistryHive->FreeListSize + 1) > RegistryHive->FreeListMax)
    {
      tmpList = MmAllocateMemory (sizeof(PCELL_HEADER) * (RegistryHive->FreeListMax + 32));
      if (tmpList == NULL)
	{
	  return FALSE;
	}

      tmpListOffset = MmAllocateMemory (sizeof(BLOCK_OFFSET) * (RegistryHive->FreeListMax + 32));
      if (tmpListOffset == NULL)
	{
	  MmFreeMemory (tmpList);
	  return FALSE;
	}

      if (RegistryHive->FreeListMax)
	{
	  memmove (tmpList,
		   RegistryHive->FreeList,
		   sizeof(PCELL_HEADER) * (RegistryHive->FreeListMax));
	  memmove (tmpListOffset,
		   RegistryHive->FreeListOffset,
		   sizeof(BLOCK_OFFSET) * (RegistryHive->FreeListMax));
	  MmFreeMemory (RegistryHive->FreeList);
	  MmFreeMemory (RegistryHive->FreeListOffset);
	}
      RegistryHive->FreeList = tmpList;
      RegistryHive->FreeListOffset = tmpListOffset;
      RegistryHive->FreeListMax += 32;
    }

  /* Add new offset to free list, maintaining list in ascending order */
  if ((RegistryHive->FreeListSize == 0)
     || (RegistryHive->FreeListOffset[RegistryHive->FreeListSize-1] < FreeOffset))
    {
      /* Add to end of list */
      RegistryHive->FreeList[RegistryHive->FreeListSize] = FreeBlock;
      RegistryHive->FreeListOffset[RegistryHive->FreeListSize++] = FreeOffset;
    }
  else if (RegistryHive->FreeListOffset[0] > FreeOffset)
    {
      /* Add to begin of list */
      memmove (&RegistryHive->FreeList[1],
	       &RegistryHive->FreeList[0],
	       sizeof(RegistryHive->FreeList[0]) * RegistryHive->FreeListSize);
      memmove (&RegistryHive->FreeListOffset[1],
	       &RegistryHive->FreeListOffset[0],
	       sizeof(RegistryHive->FreeListOffset[0]) * RegistryHive->FreeListSize);
      RegistryHive->FreeList[0] = FreeBlock;
      RegistryHive->FreeListOffset[0] = FreeOffset;
      RegistryHive->FreeListSize++;
    }
  else
    {
      /* Search where to insert */
      minInd = 0;
      maxInd = RegistryHive->FreeListSize - 1;
      while ((maxInd - minInd) > 1)
	{
	  medInd = (minInd + maxInd) / 2;
	  if (RegistryHive->FreeListOffset[medInd] > FreeOffset)
	    maxInd = medInd;
	  else
	    minInd = medInd;
	}

      /* Insert before maxInd */
      memmove (&RegistryHive->FreeList[maxInd+1],
	       &RegistryHive->FreeList[maxInd],
	       sizeof(RegistryHive->FreeList[0]) * (RegistryHive->FreeListSize - minInd));
      memmove (&RegistryHive->FreeListOffset[maxInd + 1],
	       &RegistryHive->FreeListOffset[maxInd],
	       sizeof(RegistryHive->FreeListOffset[0]) * (RegistryHive->FreeListSize-minInd));
      RegistryHive->FreeList[maxInd] = FreeBlock;
      RegistryHive->FreeListOffset[maxInd] = FreeOffset;
      RegistryHive->FreeListSize++;
    }

  return TRUE;
}


static BOOL
CmiAddBin(PREGISTRY_HIVE RegistryHive,
	  ULONG BlockCount,
	  PVOID *NewBlock,
	  PBLOCK_OFFSET NewBlockOffset)
{
  PCELL_HEADER tmpBlock;
  PHBIN *BlockList;
  PHBIN tmpBin;
  ULONG BinSize;
  ULONG i;

  BinSize = BlockCount * REG_BLOCK_SIZE;
  tmpBin = AllocateMbMemory (BinSize);
  if (tmpBin == NULL)
    {
      return FALSE;
    }
  memset (tmpBin, 0, BinSize);

  tmpBin->HeaderId = REG_BIN_ID;
  tmpBin->BinOffset = RegistryHive->FileSize - REG_BLOCK_SIZE;
  RegistryHive->FileSize += BinSize;
  tmpBin->BinSize = BinSize;
  tmpBin->DateModified = 0ULL;
  tmpBin->MemAlloc = 0;

  /* Increase size of list of blocks */
  BlockList = MmAllocateMemory (sizeof(PHBIN) * (RegistryHive->BlockListSize + BlockCount));
  if (BlockList == NULL)
    {
      return FALSE;
    }

  if (RegistryHive->BlockListSize > 0)
    {
      memcpy (BlockList,
	      RegistryHive->BlockList,
	      sizeof(PHBIN) * RegistryHive->BlockListSize);
      MmFreeMemory (RegistryHive->BlockList);
    }

  RegistryHive->BlockList = BlockList;
  for (i = 0; i < BlockCount; i++)
    RegistryHive->BlockList[RegistryHive->BlockListSize + i] = tmpBin;
  RegistryHive->BlockListSize += BlockCount;

  /* Initialize a free block in this heap : */
  tmpBlock = (PCELL_HEADER)((ULONG) tmpBin + REG_HBIN_DATA_OFFSET);
  tmpBlock->CellSize = (REG_BLOCK_SIZE - REG_HBIN_DATA_OFFSET);

  *NewBlock = (PVOID) tmpBlock;

  if (NewBlockOffset)
    *NewBlockOffset = tmpBin->BinOffset + REG_HBIN_DATA_OFFSET;

  return TRUE;
}


static BOOL
CmiAllocateCell (PREGISTRY_HIVE RegistryHive,
		 LONG CellSize,
		 PVOID *Block,
		 PBLOCK_OFFSET pBlockOffset)
{
  PCELL_HEADER NewBlock;
  ULONG i;

  *Block = NULL;

  /* Round to 16 bytes multiple */
  CellSize = ROUND_UP(CellSize, 16);

  /* first search in free blocks */
  NewBlock = NULL;
  for (i = 0; i < RegistryHive->FreeListSize; i++)
    {
      if (RegistryHive->FreeList[i]->CellSize >= CellSize)
	{
	  NewBlock = RegistryHive->FreeList[i];
	  if (pBlockOffset)
	    *pBlockOffset = RegistryHive->FreeListOffset[i];

	  if ((i + 1) < RegistryHive->FreeListSize)
	    {
	      memmove (&RegistryHive->FreeList[i],
		       &RegistryHive->FreeList[i + 1],
		       sizeof(RegistryHive->FreeList[0])
		         * (RegistryHive->FreeListSize - i - 1));
	      memmove (&RegistryHive->FreeListOffset[i],
		       &RegistryHive->FreeListOffset[i + 1],
		       sizeof(RegistryHive->FreeListOffset[0])
		         * (RegistryHive->FreeListSize - i - 1));
	    }
	  RegistryHive->FreeListSize--;
	  break;
	}
    }

  /* Need to extend hive file */
  if (NewBlock == NULL)
    {
      /* Add a new block */
      if (!CmiAddBin(RegistryHive,
		     ((sizeof(HBIN) + CellSize - 1) / REG_BLOCK_SIZE) + 1,
		     (PVOID *)(PVOID)&NewBlock,
		     pBlockOffset))
	return FALSE;
    }

  *Block = NewBlock;

  /* Split the block in two parts */
  if (NewBlock->CellSize > CellSize)
    {
      NewBlock = (PCELL_HEADER) ((ULONG)NewBlock + CellSize);
      NewBlock->CellSize = ((PCELL_HEADER) (*Block))->CellSize - CellSize;
      CmiAddFree (RegistryHive,
		  NewBlock,
		  *pBlockOffset + CellSize,
		  TRUE);
    }
  else if (NewBlock->CellSize < CellSize)
    {
      return FALSE;
    }

  memset(*Block, 0, CellSize);
  ((PCELL_HEADER)(*Block))->CellSize = -CellSize;

  return TRUE;
}


static PVOID
CmiGetCell (PREGISTRY_HIVE Hive,
	    BLOCK_OFFSET BlockOffset)
{
  PHBIN Bin;
  ULONG BlockIndex;

  if (BlockOffset == (ULONG) -1)
    return NULL;

  BlockIndex = BlockOffset / REG_BLOCK_SIZE;
  if (BlockIndex >= Hive->BlockListSize)
    return NULL;

  Bin = Hive->BlockList[BlockIndex];
  if (Bin == NULL)
    return NULL;

  return (PVOID)((ULONG)Bin + (BlockOffset - Bin->BinOffset));
}


static BOOL
CmiAllocateHashTableCell (PREGISTRY_HIVE Hive,
			  PBLOCK_OFFSET HBOffset,
			  ULONG SubKeyCount)
{
  PHASH_TABLE_CELL HashCell;
  ULONG NewHashSize;
  BOOL Status;

  NewHashSize = sizeof(HASH_TABLE_CELL) +
		(SubKeyCount * sizeof(HASH_RECORD));
  Status = CmiAllocateCell (Hive,
			    NewHashSize,
			    (PVOID*)(PVOID)&HashCell,
			    HBOffset);
  if ((HashCell == NULL) || (Status == FALSE))
    {
      return FALSE;
    }

  HashCell->Id = REG_HASH_TABLE_BLOCK_ID;
  HashCell->HashTableSize = SubKeyCount;

  return TRUE;
}


static BOOL
CmiAddKeyToParentHashTable (PREGISTRY_HIVE Hive,
			    BLOCK_OFFSET ParentKeyOffset,
			    PKEY_CELL NewKeyCell,
			    BLOCK_OFFSET NKBOffset)
{
  PHASH_TABLE_CELL HashBlock;
  PKEY_CELL ParentKeyCell;
  ULONG i;

  ParentKeyCell = CmiGetCell (Hive, ParentKeyOffset);
  if (ParentKeyCell == NULL)
    {
      DbgPrint((DPRINT_REGISTRY, "CmiGetCell() failed\n"));
      return FALSE;
    }

  HashBlock =CmiGetCell (Hive, ParentKeyCell->HashTableOffset);
  if (HashBlock == NULL)
    {
      DbgPrint((DPRINT_REGISTRY, "CmiGetCell() failed\n"));
      return FALSE;
    }

  for (i = 0; i < HashBlock->HashTableSize; i++)
    {
      if (HashBlock->Table[i].KeyOffset == 0)
	{
	  HashBlock->Table[i].KeyOffset = NKBOffset;
	  memcpy (&HashBlock->Table[i].HashValue,
		  NewKeyCell->Name,
		  min(NewKeyCell->NameSize, sizeof(ULONG)));
	  ParentKeyCell->NumberOfSubKeys++;
	  return TRUE;
	}
    }

  return FALSE;
}


static BOOL
CmiAllocateValueListCell (PREGISTRY_HIVE Hive,
			  PBLOCK_OFFSET ValueListOffset,
			  ULONG ValueCount)
{
  PVALUE_LIST_CELL ValueListCell;
  ULONG ValueListSize;
  BOOL Status;

  ValueListSize = sizeof(VALUE_LIST_CELL) +
		  (ValueCount * sizeof(BLOCK_OFFSET));
  Status = CmiAllocateCell (Hive,
			    ValueListSize,
			    (PVOID)&ValueListCell,
			    ValueListOffset);
  if ((ValueListCell == NULL) || (Status == FALSE))
    {
      DbgPrint((DPRINT_REGISTRY, "CmiAllocateCell() failed\n"));
      return FALSE;
    }

  return TRUE;
}


static BOOL
CmiAllocateValueCell(PREGISTRY_HIVE Hive,
		     PVALUE_CELL *ValueCell,
		     BLOCK_OFFSET *ValueCellOffset,
		     PWCHAR ValueName)
{
  PVALUE_CELL NewValueCell;
  ULONG NameSize;
  BOOL Status;
  BOOLEAN Packable = TRUE;
  ULONG i;

  NameSize = (ValueName == NULL) ? 0 : wcslen (ValueName);
  for (i = 0; i < NameSize; i++)
    {
      if (ValueName[i] & 0xFF00)
        {
          NameSize *= sizeof(WCHAR);
          Packable = FALSE;
          break;
        }
    }
  Status = CmiAllocateCell (Hive,
			    sizeof(VALUE_CELL) + NameSize,
			    (PVOID*)(PVOID)&NewValueCell,
			    ValueCellOffset);
  if ((NewValueCell == NULL) || (Status == FALSE))
    {
      DbgPrint((DPRINT_REGISTRY, "CmiAllocateCell() failed\n"));
      return FALSE;
    }

  NewValueCell->Id = REG_VALUE_CELL_ID;
  NewValueCell->NameSize = NameSize;
  NewValueCell->Flags = 0;
  if (NameSize > 0)
    {
      if (Packable)
        {
          for (i = 0; i < NameSize; i++)
            {
              ((PCHAR)NewValueCell->Name)[i] = (CHAR)ValueName[i];
            }
          NewValueCell->Flags |= REG_VALUE_NAME_PACKED;
        }
      else
        {
          memcpy (NewValueCell->Name,
	          ValueName,
	          NameSize);
        }
    }
  NewValueCell->DataType = 0;
  NewValueCell->DataSize = 0;
  NewValueCell->DataOffset = -1;

  *ValueCell = NewValueCell;

  return TRUE;
}


static BOOL
CmiAddValueToKeyValueList(PREGISTRY_HIVE Hive,
			  BLOCK_OFFSET KeyCellOffset,
			  BLOCK_OFFSET ValueCellOffset)
{
  PVALUE_LIST_CELL ValueListCell;
  PKEY_CELL KeyCell;

  KeyCell = CmiGetCell (Hive, KeyCellOffset);
  if (KeyCell == NULL)
    {
      DbgPrint((DPRINT_REGISTRY, "CmiGetCell() failed\n"));
      return FALSE;
    }

  ValueListCell = CmiGetCell (Hive, KeyCell->ValueListOffset);
  if (ValueListCell == NULL)
    {
      DbgPrint((DPRINT_REGISTRY, "CmiGetCell() failed\n"));
      return FALSE;
    }

  ValueListCell->ValueOffset[KeyCell->NumberOfValues] = ValueCellOffset;
  KeyCell->NumberOfValues++;

  return TRUE;
}

static BOOL
CmiExportValue (PREGISTRY_HIVE Hive,
		BLOCK_OFFSET KeyCellOffset,
		FRLDRHKEY Key,
		PVALUE Value)
{
  BLOCK_OFFSET ValueCellOffset;
  BLOCK_OFFSET DataCellOffset;
  PVALUE_CELL ValueCell;
  PDATA_CELL DataCell;
  ULONG DataSize;
  ULONG DataType;
  PCHAR Data;

  DbgPrint((DPRINT_REGISTRY, "CmiExportValue('%S') called\n",
	   (Value == NULL) ? "<default>" : (PCHAR)Value->Name));
  DbgPrint((DPRINT_REGISTRY, "DataSize %lu\n",
	   (Value == NULL) ? Key->DataSize : Value->DataSize));

  /* Allocate value cell */
  if (!CmiAllocateValueCell(Hive, &ValueCell, &ValueCellOffset, (Value == NULL) ? NULL : Value->Name))
    {
      return FALSE;
    }

  if (!CmiAddValueToKeyValueList(Hive, KeyCellOffset, ValueCellOffset))
    {
      return FALSE;
    }

  if (Value == NULL)
    {
      DataType = Key->DataType;
      DataSize = Key->DataSize;
      Data = Key->Data;
    }
  else
    {
      DataType = Value->DataType;
      DataSize = Value->DataSize;
      Data = Value->Data;
    }

  if (DataSize <= sizeof(BLOCK_OFFSET))
    {
      ValueCell->DataSize = DataSize | REG_DATA_IN_OFFSET;
      ValueCell->DataType = DataType;
      memcpy (&ValueCell->DataOffset,
	      &Data,
	      DataSize);
    }
  else
    {
      /* Allocate data cell */
      if (!CmiAllocateCell (Hive,
			    sizeof(CELL_HEADER) + DataSize,
			    (PVOID *)(PVOID)&DataCell,
			    &DataCellOffset))
	{
	  return FALSE;
	}

      ValueCell->DataOffset = DataCellOffset;
      ValueCell->DataSize = DataSize;
      ValueCell->DataType = DataType;

      memcpy (DataCell->Data,
	      Data,
	      DataSize);
    }

  return TRUE;
}