Add recovery utility to open /usr directory in disks with corrupted wfs header

CMakeLists.txt

@@ -37,6 +37,7 @@ add_library(${PROJECT_NAME}
     src/metadata_block.cpp
     src/structs.cpp
     src/sub_block_allocator.cpp
+    src/recovery.cpp
     src/wfs.cpp
     src/wfs_item.cpp
 )

include/wfslib/blocks_device.h

@@ -13,6 +13,8 @@
 #include <span>
 #include <unordered_map>
 
+#include "device_encryption.h"
+
 class Block;
 
 class Device;

include/wfslib/directory.h

@@ -41,6 +41,7 @@ class Directory : public WfsItem, public std::enable_shared_from_this<Directory>
 
  private:
   friend DirectoryItemsIterator;
+  friend class Recovery;
 
   // TODO: We may have cyclic reference here if we do cache in area.
   std::shared_ptr<Area> area_;

include/wfslib/errors.h

@@ -21,6 +21,7 @@ enum WfsError {
   kFileDataCorrupted,
   kFileMetadataCorrupted,
   kTransactionsAreaCorrupted,
+  kInvalidWfsVersion,
 };
 
 class WfsException : public std::exception {

include/wfslib/file_device.h

@@ -27,11 +27,10 @@ class FileDevice : public Device {
   uint32_t Log2SectorSize() const override { return log2_sector_size_; }
   bool IsReadOnly() const override { return read_only_; }
 
- private:
-  friend Wfs;
   void SetSectorsCount(uint32_t sectors_count) { sectors_count_ = sectors_count; }
   void SetLog2SectorSize(uint32_t log2_sector_size) { log2_sector_size_ = log2_sector_size; }
 
+ private:
   std::unique_ptr<std::iostream> file_;
   std::mutex io_lock_;
 

include/wfslib/recovery.h

@@ -0,0 +1,37 @@
+/*
+ * Copyright (C) 2024 koolkdev
+ *
+ * This software may be modified and distributed under the terms
+ * of the MIT license.  See the LICENSE file for details.
+ */
+
+#pragma once
+
+#include <expected>
+#include <memory>
+#include <optional>
+#include <vector>
+
+#include "errors.h"
+
+class FileDevice;
+class Wfs;
+class Directory;
+
+class Recovery {
+ public:
+  static bool CheckWfsKey(std::shared_ptr<FileDevice> device, std::optional<std::vector<std::byte>> key);
+
+  static std::optional<WfsError> DetectDeviceParams(std::shared_ptr<FileDevice> device,
+                                                    std::optional<std::vector<std::byte>> key);
+
+  // Open a WFS device without knowning the exact device parameters
+  static std::expected<std::shared_ptr<Wfs>, WfsError> OpenWfsWithoutDeviceParams(
+      std::shared_ptr<FileDevice> device,
+      std::optional<std::vector<std::byte>> key);
+
+  // Open a WFS device with /usr as root directory while skipping the wfs header
+  static std::expected<std::shared_ptr<Wfs>, WfsError> OpenUsrDirectoryWithoutWfsHeader(
+      std::shared_ptr<FileDevice> device,
+      std::optional<std::vector<std::byte>> key);
+};

include/wfslib/wfs.h

@@ -34,9 +34,6 @@ class Wfs {
   std::shared_ptr<File> GetFile(const std::string& filename);
   std::shared_ptr<Directory> GetDirectory(const std::string& filename);
 
-  static void DetectDeviceSectorSizeAndCount(std::shared_ptr<FileDevice> device,
-                                             std::optional<std::vector<std::byte>> key = std::nullopt);
-
   std::shared_ptr<Area> GetRootArea() { return root_area_; }
 
   void Flush();

include/wfslib/wfslib.h

@@ -8,7 +8,9 @@
 #pragma once
 
 #include "directory.h"
+#include "errors.h"
 #include "file.h"
 #include "file_device.h"
 #include "key_file.h"
+#include "recovery.h"
 #include "wfs.h"

src/area.h

@@ -70,6 +70,7 @@ class Area : public std::enable_shared_from_this<Area> {
 
  private:
   friend class Wfs;
+  friend class Recovery;
 
   static constexpr uint32_t FreeBlocksAllocatorBlockNumber = 1;
 

src/errors.cpp

@@ -29,6 +29,8 @@ char const* WfsException::what() const noexcept {
       return "File metadata corrupted";
     case WfsError::kTransactionsAreaCorrupted:
       return "Transactions area corrupted";
+    case WfsError::kInvalidWfsVersion:
+      return "Invalid WFS version";
   }
   return "";
 }

src/recovery.cpp

@@ -0,0 +1,267 @@
+/*
+ * Copyright (C) 2024 koolkdev
+ *
+ * This software may be modified and distributed under the terms
+ * of the MIT license.  See the LICENSE file for details.
+ */
+
+#include "recovery.h"
+
+#include <array>
+
+#include "area.h"
+#include "blocks_device.h"
+#include "device_encryption.h"
+#include "directory.h"
+#include "file_device.h"
+#include "metadata_block.h"
+#include "structs.h"
+#include "wfs.h"
+
+namespace {
+
+bool RestoreMetadataBlockIVParameters(std::shared_ptr<FileDevice> device,
+                                      std::shared_ptr<BlocksDevice> blocks_device,
+                                      uint32_t block_number,
+                                      uint32_t area_start_block_number,
+                                      Block::BlockSize block_size,
+                                      uint32_t& area_xor_wfs_iv) {
+  // The IV of the encryption of block looks like that:
+  // DWORD[0]: block offset in disk in bytes
+  // DWORD[1]: block area ^ device wfs 32 bit IVs ^ block number (the first one is in the area header, the second one in
+  // the wfs header, the block number is in basic size which is 0x1000 bytes per block)
+  // DWORD[2]: disk sectors count (LBAs available for WFS)
+  // DWORD[3]: disk sector size in bytes (LBA size)
+  // Without prior information DWORD[0] is known, DWORD[1] is unknown. DWORD[2]/[3] can be determined from the device
+  // itself, but using the following methodolgy, we can restore the three of them anyway.
+  // The encryption is CBC mode, so IV is xored with the data AFTER decryption of the first block, so if we can predict
+  // the data of the first block, we can restore the IV. The reason it is possible is that all of those dwords fall on
+  // bytes 4-16 while in metadata blocks the hash will in bytes 4-24. So if we decrypt the block with the right key, we
+  // can get the right hash, and than we can xor the expected value with the value that we got after decryption, and we
+  // will get the xor of our IV with the real.
+  // The sector size will usually be 512 bytes anyway, so lets start with that.
+  device->SetLog2SectorSize(9);
+  // Set the sectors count to be enough to read our block
+  device->SetSectorsCount((block_number << (Block::BlockSize::Basic - 9)) + (1 << (block_size - 9)));
+  // This is the initial iv that the encryption will use
+  std::array<uint32_t, 4> iv = {block_number << Block::BlockSize::Basic, 0, device->SectorsCount(),
+                                device->SectorSize()};
+  std::shared_ptr<const MetadataBlock> block =
+      *MetadataBlock::LoadBlock(blocks_device, block_number, block_size, 0, false);
+  // The two last dwords of the IV is the sectors count and sector size, right now it is xored with our fake sector size
+  // and sector count, and with the hash
+  std::vector<std::byte> data{block->data().begin(), block->data().end()};
+  auto first_4_dwords = reinterpret_cast<uint32_be_t*>(data.data());
+  iv[0] ^= first_4_dwords[0].value();
+  iv[1] ^= first_4_dwords[1].value();
+  iv[2] ^= first_4_dwords[2].value();
+  iv[3] ^= first_4_dwords[3].value();
+  // Lets calculate the hash of the block
+  DeviceEncryption::CalculateHash(data,
+                                  {data.data() + offsetof(MetadataBlockHeader, hash), DeviceEncryption::DIGEST_SIZE});
+  iv[0] ^= first_4_dwords[0].value();
+  iv[1] ^= first_4_dwords[1].value();
+  iv[2] ^= first_4_dwords[2].value();
+  iv[3] ^= first_4_dwords[3].value();
+  auto iv32 = iv[1];
+  auto sectors_count = iv[2];
+  auto sector_size = iv[3];
+  if (std::popcount(sector_size) != 1) {
+    // Not pow of 2
+    return false;
+  }
+  auto log2_sector_size = std::bit_width(sector_size) - 1;
+
+  device->SetLog2SectorSize(log2_sector_size);
+  device->SetSectorsCount(sectors_count);
+  block.reset();
+
+  // Now try to fetch block again, this time check the hash.
+  if (!MetadataBlock::LoadBlock(blocks_device, block_number, block_size, iv32, true).has_value())
+    return false;
+
+  area_xor_wfs_iv = iv32 - ((block_number - area_start_block_number) << (Block::BlockSize::Basic - log2_sector_size));
+  return true;
+}
+
+}  // namespace
+
+bool Recovery::CheckWfsKey(std::shared_ptr<FileDevice> device, std::optional<std::vector<std::byte>> key) {
+  auto blocks_device = std::make_shared<BlocksDevice>(device, key);
+  std::shared_ptr<const MetadataBlock> block =
+      *MetadataBlock::LoadBlock(blocks_device, 0, Block::BlockSize::Basic, 0, false);
+  auto wfs_header = reinterpret_cast<const WfsHeader*>(&block->data()[sizeof(MetadataBlockHeader)]);
+  // Check wfs header
+  return wfs_header->version.value() == WFS_VERSION;
+}
+
+std::optional<WfsError> Recovery::DetectDeviceParams(std::shared_ptr<FileDevice> device,
+                                                     std::optional<std::vector<std::byte>> key) {
+  if (!CheckWfsKey(device, key))
+    return WfsError::kInvalidWfsVersion;
+  auto blocks_device = std::make_shared<BlocksDevice>(device, key);
+  std::shared_ptr<const MetadataBlock> block =
+      *MetadataBlock::LoadBlock(blocks_device, 0, Block::BlockSize::Basic, 0, false);
+  auto wfs_header = block->get_object<WfsHeader>(sizeof(MetadataBlockHeader));
+  auto block_size = Block::BlockSize::Basic;
+  if (!(wfs_header->root_area_attributes.flags.value() & Attributes::Flags::AREA_SIZE_BASIC) &&
+      (wfs_header->root_area_attributes.flags.value() & Attributes::Flags::AREA_SIZE_REGULAR))
+    block_size = Block::BlockSize::Regular;
+  block.reset();
+  uint32_t area_xor_wfs_iv;
+  if (!RestoreMetadataBlockIVParameters(device, blocks_device, 0, 0, block_size, area_xor_wfs_iv)) {
+    return WfsError::kAreaHeaderCorrupted;
+  }
+  return std::nullopt;
+}
+
+std::expected<std::shared_ptr<Wfs>, WfsError> Recovery::OpenWfsWithoutDeviceParams(
+    std::shared_ptr<FileDevice> device,
+    std::optional<std::vector<std::byte>> key) {
+  auto res = DetectDeviceParams(device, key);
+  if (res.has_value())
+    return std::unexpected(*res);
+  return std::make_shared<Wfs>(device, key);
+}
+
+namespace {
+class FakeWfsHeaderBlocksDevice final : public BlocksDevice {
+ public:
+  FakeWfsHeaderBlocksDevice(std::shared_ptr<Device> device,
+                            std::optional<std::vector<std::byte>> key,
+                            std::vector<std::byte> wfs_header_block)
+      : BlocksDevice(std::move(device), std::move(key)), wfs_header_block_(std::move(wfs_header_block)) {}
+  ~FakeWfsHeaderBlocksDevice() final override = default;
+
+  bool ReadBlock(uint32_t block_number,
+                 uint32_t size_in_blocks,
+                 const std::span<std::byte>& data,
+                 const std::span<const std::byte>& hash,
+                 uint32_t iv,
+                 bool encrypt,
+                 bool check_hash) override {
+    if (block_number == 0) {
+      if (data.size() != wfs_header_block_.size())
+        return false;
+      std::copy(wfs_header_block_.begin(), wfs_header_block_.end(), data.begin());
+      return true;
+    } else if (block_number < 0x1000) {
+      // The first 0x1000 blocks not supperted and not supposed to be reached in this mode.
+      return false;
+    }
+    return BlocksDevice::ReadBlock(block_number, size_in_blocks, data, hash, iv, encrypt, check_hash);
+  }
+
+ private:
+  std::vector<std::byte> wfs_header_block_;
+};
+
+}  // namespace
+
+std::expected<std::shared_ptr<Wfs>, WfsError> Recovery::OpenUsrDirectoryWithoutWfsHeader(
+    std::shared_ptr<FileDevice> device,
+    std::optional<std::vector<std::byte>> key) {
+  // First step: Read the /usr directory and restore root area ^ wfs ivs
+  const uint32_t kUsrDirectoryBlockNumber = 0x1000;
+  // Set the sectors count to be enough to read our block
+  device->SetLog2SectorSize(9);
+  device->SetSectorsCount((kUsrDirectoryBlockNumber + 2) << 3);
+  auto blocks_device = std::make_shared<BlocksDevice>(device, key);
+  std::shared_ptr<const MetadataBlock> block =
+      *MetadataBlock::LoadBlock(blocks_device, kUsrDirectoryBlockNumber, Block::BlockSize::Basic, 0, false);
+  auto metadata_header = block->get_object<MetadataBlockHeader>(0);
+  // Check wfs header
+  if ((metadata_header->block_flags.value() >> 20) != 0xe00) {
+    return std::unexpected(WfsError::kInvalidWfsVersion);
+  }
+  block.reset();
+
+  uint32_t root_area_xor_wfs_iv;
+  // Assume regular sized area
+  if (!RestoreMetadataBlockIVParameters(device, blocks_device, kUsrDirectoryBlockNumber, 0, Block::BlockSize::Regular,
+                                        root_area_xor_wfs_iv)) {
+    return std::unexpected(WfsError::kAreaHeaderCorrupted);
+  }
+
+  // Create an initial fake first block, fill the important fields only for parsing
+  MetadataBlockHeader header;
+  std::memset(&header, 0, sizeof(header));
+  header.block_flags = MetadataBlockHeader::Flags::AREA | MetadataBlockHeader::Flags::ROOT_AREA;
+  WfsHeader wfs_header;
+  std::memset(&wfs_header, 0, sizeof(wfs_header));
+  wfs_header.iv = 0;
+  wfs_header.version = 0x01010800;
+  wfs_header.root_area_attributes.flags = static_cast<uint32_t>(
+      Attributes::Flags::AREA_SIZE_REGULAR | Attributes::Flags::QUOTA | Attributes::Flags::DIRECTORY);
+  WfsAreaHeader area_header;
+  std::memset(&area_header, 0, sizeof(area_header));
+  // This is the initial iv until we correct it
+  area_header.iv = root_area_xor_wfs_iv;
+  // Set the root directory to be /usr
+  area_header.root_directory_block_number =
+      kUsrDirectoryBlockNumber >> (Block::BlockSize::Regular - Block::BlockSize::Basic);
+  area_header.log2_block_size = Block::BlockSize::Regular;
+
+  // Now create the block
+  std::vector<std::byte> first_wfs_block;
+  std::copy(reinterpret_cast<std::byte*>(&header), reinterpret_cast<std::byte*>(&header + 1),
+            std::back_inserter(first_wfs_block));
+  std::copy(reinterpret_cast<std::byte*>(&wfs_header), reinterpret_cast<std::byte*>(&wfs_header + 1),
+            std::back_inserter(first_wfs_block));
+  std::copy(reinterpret_cast<std::byte*>(&area_header), reinterpret_cast<std::byte*>(&area_header + 1),
+            std::back_inserter(first_wfs_block));
+  std::fill_n(std::back_inserter(first_wfs_block), (1 << Block::BlockSize::Regular) - first_wfs_block.size(),
+              std::byte{0});
+  auto first_fixed_device = std::make_shared<FakeWfsHeaderBlocksDevice>(device, key, std::move(first_wfs_block));
+
+  // Now go to the the /usr/system/save directory, there must be a sub-area there
+  auto wfs = std::make_unique<Wfs>(first_fixed_device);
+  auto system_save_dir = wfs->GetDirectory("/save/system");
+  if (!system_save_dir) {
+    throw std::logic_error("Failed to find /usr/save/system");
+  }
+  std::shared_ptr<const Area> sub_area;
+  for (auto [name, item_or_error] : *system_save_dir) {
+    if (item_or_error.has_value())
+      continue;
+    // this is probably a corrupted quota, let's check
+    const auto attributes = system_save_dir->GetObjectAttributes(system_save_dir->block_, name);
+    if (!attributes.has_value() || !attributes->Attributes()->IsQuota())
+      continue;
+    // ok this is quota
+    auto new_area = system_save_dir->area()->GetArea(attributes->Attributes()->directory_block_number.value(), name,
+                                                     *attributes, Block::BlockSize::Regular);
+    if (!new_area.has_value())
+      return std::unexpected(new_area.error());
+    sub_area = std::move(*new_area);
+    break;
+  }
+  if (!sub_area) {
+    throw std::logic_error("Failed to find sub-quota under /usr/save/system");
+  }
+  uint32_t sub_area_xor_wfs_iv;
+  if (!RestoreMetadataBlockIVParameters(
+          device, blocks_device, sub_area->AbsoluteBlockNumber(sub_area->header()->root_directory_block_number.value()),
+          sub_area->AbsoluteBlockNumber(0), Block::BlockSize::Regular, sub_area_xor_wfs_iv)) {
+    return std::unexpected(WfsError::kAreaHeaderCorrupted);
+  }
+  uint32_t wfs_iv = sub_area_xor_wfs_iv ^ sub_area->header()->iv.value();
+  uint32_t root_area_iv = root_area_xor_wfs_iv ^ wfs_iv;
+
+  // Now fix the header
+  wfs_header.iv = wfs_iv;
+  area_header.iv = root_area_iv;
+
+  // Recreate the first block
+  first_wfs_block.clear();
+  std::copy(reinterpret_cast<std::byte*>(&header), reinterpret_cast<std::byte*>(&header + 1),
+            std::back_inserter(first_wfs_block));
+  std::copy(reinterpret_cast<std::byte*>(&wfs_header), reinterpret_cast<std::byte*>(&wfs_header + 1),
+            std::back_inserter(first_wfs_block));
+  std::copy(reinterpret_cast<std::byte*>(&area_header), reinterpret_cast<std::byte*>(&area_header + 1),
+            std::back_inserter(first_wfs_block));
+  std::fill_n(std::back_inserter(first_wfs_block), (1 << Block::BlockSize::Regular) - first_wfs_block.size(),
+              std::byte{0});
+  auto final_device = std::make_shared<FakeWfsHeaderBlocksDevice>(device, key, std::move(first_wfs_block));
+  return std::make_shared<Wfs>(final_device);
+}