lost_partition_finder.cpp

Go to the documentation of this file.

#include "lost_partition_finder.hpp"
#include "i18n.hpp"
#include "xfs_format.hpp"
#include "hfs_format.hpp"
 
#include <algorithm>
#include <array>
#include <string>
#include <vector>
 
namespace {
constexpr uint64_t MiB = 1024ull * 1024ull;
constexpr uint64_t GiB = 1024ull * 1024ull * 1024ull;
constexpr uint64_t SECTOR_SIZE = 512ull;
 
// Bug #13 fix: modern partitions are 1 MiB aligned (Windows Vista+, Linux ≥2010).
// A coarser stride misses partitions that happen to live on a 1 MiB boundary
// not divisible by the chosen stride. The cost is a few extra reads, which are
// negligible compared to producing a usable result.
constexpr uint64_t SCAN_STRIDE = 1ull * MiB;
 
// Bug #10 fix: the previous 16 GiB cap made the lost-partition finder useless
// on any modern disk. We now scan up to 4 TiB by default (covers most consumer
// drives) and let larger drives be scanned in full if the user has time.
constexpr uint64_t DEFAULT_MAX_SCAN = 4ull * 1024ull * GiB;   // 4 TiB
constexpr uint64_t HARD_MAX_SCAN    = 64ull * 1024ull * GiB;  // 64 TiB safety cap
 
struct Range {
    uint64_t start;
    uint64_t end;
};
 
bool overlaps(const std::vector<Range>& ranges, uint64_t start, uint64_t end) {
    for (const auto& r : ranges) {
        if (start < r.end && end > r.start) {
            return true;
        }
    }
    return false;
}
 
bool is_ntfs(const std::array<unsigned char, 512>& b) {
    return std::string(reinterpret_cast<const char*>(b.data() + 3), 8) == "NTFS    " &&
           b[510] == 0x55 && b[511] == 0xAA;
}
 
bool is_fat32(const std::array<unsigned char, 512>& b) {
    return std::string(reinterpret_cast<const char*>(b.data() + 82), 8) == "FAT32   " &&
           b[510] == 0x55 && b[511] == 0xAA;
}
 
bool is_exfat(const std::array<unsigned char, 512>& b) {
    return std::string(reinterpret_cast<const char*>(b.data() + 3), 8) == "EXFAT   " &&
           b[510] == 0x55 && b[511] == 0xAA;
}
 
bool looks_ext4(DiskReader& reader, uint64_t offset) {
    std::array<unsigned char, 2> sig{};
    return reader.read_at(offset + 1024 + 56, sig.data(), sig.size()) && sig[0] == 0x53 && sig[1] == 0xEF;
}
 
bool looks_btrfs(DiskReader& reader, uint64_t offset) {
    std::array<char, 8> sig{};
    return reader.read_at(offset + 0x10040, sig.data(), sig.size()) && std::string(sig.data(), sig.size()) == "_BHRfS_M";
}
 
// XFS quick check (magic only). For full validation we use parse_superblock
// in xfs_size_bytes() below.
bool looks_xfs(DiskReader& reader, uint64_t offset) {
    std::array<unsigned char, 4> sig{};
    if (!reader.read_at(offset, sig.data(), sig.size())) return false;
    return sig[0] == 'X' && sig[1] == 'F' && sig[2] == 'S' && sig[3] == 'B';
}
 
// Returns 0 if the sector at `offset` does not actually contain a valid
// XFS superblock (e.g. magic was a coincidence), otherwise volume size.
uint64_t xfs_size_bytes(DiskReader& reader, uint64_t offset) {
    std::array<unsigned char, 512> sector{};
    if (!reader.read_at(offset, sector.data(), sector.size())) return 0;
    xfs::Superblock sb;
    if (!xfs::parse_superblock(sector.data(), sb)) return 0;
    return sb.volume_size_bytes();
}
 
// HFS+ / HFSX: volume header at +1024 from volume start. Magic in the
// first 2 bytes is "H+" or "HX".
bool looks_hfs(DiskReader& reader, uint64_t offset) {
    std::array<unsigned char, 4> sig{};
    if (!reader.read_at(offset + hfs::HFSPLUS_VOLUME_HEADER_OFFSET,
                         sig.data(), sig.size())) {
        return false;
    }
    const uint16_t s = static_cast<uint16_t>((sig[0] << 8) | sig[1]);
    return s == hfs::HFSPLUS_SIGNATURE || s == hfs::HFSX_SIGNATURE;
}
 
uint64_t hfs_size_bytes(DiskReader& reader, uint64_t offset) {
    std::array<unsigned char, 512> sector{};
    if (!reader.read_at(offset + hfs::HFSPLUS_VOLUME_HEADER_OFFSET,
                         sector.data(), sector.size())) {
        return 0;
    }
    hfs::VolumeHeader vh;
    if (!hfs::parse_volume_header(sector.data(), vh)) return 0;
    return vh.volume_size_bytes();
}
 
uint64_t le64(const unsigned char* p) {
    return static_cast<uint64_t>(p[0]) |
           (static_cast<uint64_t>(p[1]) << 8) |
           (static_cast<uint64_t>(p[2]) << 16) |
           (static_cast<uint64_t>(p[3]) << 24) |
           (static_cast<uint64_t>(p[4]) << 32) |
           (static_cast<uint64_t>(p[5]) << 40) |
           (static_cast<uint64_t>(p[6]) << 48) |
           (static_cast<uint64_t>(p[7]) << 56);
}
 
uint32_t le32(const unsigned char* p) {
    return static_cast<uint32_t>(p[0]) |
           (static_cast<uint32_t>(p[1]) << 8) |
           (static_cast<uint32_t>(p[2]) << 16) |
           (static_cast<uint32_t>(p[3]) << 24);
}
 
uint64_t ntfs_size_bytes(const std::array<unsigned char, 512>& b) {
    return le64(b.data() + 40) * SECTOR_SIZE;
}
 
uint64_t fat32_size_bytes(const std::array<unsigned char, 512>& b) {
    const uint32_t sectors32 = le32(b.data() + 32);
    return static_cast<uint64_t>(sectors32) * SECTOR_SIZE;
}
 
uint64_t exfat_size_bytes(const std::array<unsigned char, 512>& b) {
    return le64(b.data() + 72) * SECTOR_SIZE;
}
 
// Bug #20 fix: read the ext4 superblock and compute real volume size.
uint64_t ext4_size_bytes(DiskReader& reader, uint64_t offset) {
    std::array<unsigned char, 1024> sb{};
    if (!reader.read_at(offset + 1024, sb.data(), sb.size())) return 0;
    const uint32_t blocks_lo = le32(sb.data() + 0x04);
    const uint32_t log_block_size = le32(sb.data() + 0x18);
    const uint64_t block_size = 1024ull << log_block_size;
    if (block_size < 1024 || block_size > 65536) return 0;
    return static_cast<uint64_t>(blocks_lo) * block_size;
}
 
// Bug #20 fix: read the btrfs superblock total_bytes
uint64_t btrfs_size_bytes(DiskReader& reader, uint64_t offset) {
    std::array<unsigned char, 8> total{};
    // btrfs superblock at +0x10000, total_bytes is at offset 0x70 within sb
    if (!reader.read_at(offset + 0x10000 + 0x70, total.data(), total.size())) return 0;
    return le64(total.data());
}
}
 
std::vector<ScanResult> LostPartitionFinder::find(DiskReader& reader,
                                                   const std::vector<ScanResult>& known_results,
                                                   std::vector<std::string>& logs) {
    std::vector<ScanResult> out;
    std::vector<Range> claimed_ranges;
    claimed_ranges.reserve(known_results.size());
 
    for (const auto& r : known_results) {
        if (r.kind == ResultKind::Partition && r.size_bytes > 0) {
            claimed_ranges.push_back({r.offset_bytes, r.offset_bytes + r.size_bytes});
        }
    }
 
    const uint64_t disk_size = reader.size();
    const uint64_t max_scan = std::min<uint64_t>(disk_size, std::min<uint64_t>(DEFAULT_MAX_SCAN, HARD_MAX_SCAN));
 
    logs.push_back(trs("Lost partition finder: deep scan started."));
    logs.push_back(trs("Lost partition scan stride:") + std::string(" ") +
                   std::to_string(SCAN_STRIDE / MiB) + " MiB");
    logs.push_back(trs("Lost partition scan range:") + std::string(" ") +
                   std::to_string(max_scan / GiB) + " GiB");
 
    // Bug #10/#13 fix: scan with fixed 1 MiB stride from offset SCAN_STRIDE
    // up to max_scan. Partitions are 1 MiB aligned on every modern OS, and the
    // first MiB is reserved for MBR/GPT, never a partition boundary.
    for (uint64_t offset = SCAN_STRIDE; offset + 4096 < max_scan; offset += SCAN_STRIDE) {
        std::array<unsigned char, 512> sector{};
        if (!reader.read_at(offset, sector.data(), sector.size())) {
            continue;
        }
 
        if (is_ntfs(sector)) {
            const uint64_t size_bytes = ntfs_size_bytes(sector);
            const uint64_t end = size_bytes > 0 ? offset + size_bytes : offset + MiB;
            if (!overlaps(claimed_ranges, offset, end)) {
                out.push_back({ResultKind::LostPartitionCandidate,
                               offset,
                               size_bytes,
                               "Lost NTFS partition candidate",
                               trs("Lost partition finder"),
                               trs("Unallocated NTFS boot sector candidate"),
                               0.88});
                claimed_ranges.push_back({offset, end});
            }
        } else if (is_fat32(sector)) {
            const uint64_t size_bytes = fat32_size_bytes(sector);
            const uint64_t end = size_bytes > 0 ? offset + size_bytes : offset + MiB;
            if (!overlaps(claimed_ranges, offset, end)) {
                out.push_back({ResultKind::LostPartitionCandidate,
                               offset,
                               size_bytes,
                               "Lost FAT32 partition candidate",
                               trs("Lost partition finder"),
                               trs("Unallocated FAT32 boot sector candidate"),
                               0.84});
                claimed_ranges.push_back({offset, end});
            }
        } else if (is_exfat(sector)) {
            const uint64_t size_bytes = exfat_size_bytes(sector);
            const uint64_t end = size_bytes > 0 ? offset + size_bytes : offset + MiB;
            if (!overlaps(claimed_ranges, offset, end)) {
                out.push_back({ResultKind::LostPartitionCandidate,
                               offset,
                               size_bytes,
                               "Lost exFAT partition candidate",
                               trs("Lost partition finder"),
                               trs("Unallocated exFAT boot sector candidate"),
                               0.86});
                claimed_ranges.push_back({offset, end});
            }
        } else if (looks_ext4(reader, offset)) {
            // Bug #20 fix: compute real ext4 volume size instead of leaving 0
            const uint64_t size_bytes = ext4_size_bytes(reader, offset);
            const uint64_t end = size_bytes > 0 ? offset + size_bytes : offset + MiB;
            if (!overlaps(claimed_ranges, offset, end)) {
                out.push_back({ResultKind::LostPartitionCandidate,
                               offset,
                               size_bytes,
                               "Lost ext4 partition candidate",
                               trs("Lost partition finder"),
                               trs("Unallocated ext4 superblock candidate"),
                               0.80});
                claimed_ranges.push_back({offset, end});
            }
        } else if (looks_btrfs(reader, offset)) {
            // Bug #20 fix: compute real btrfs volume size instead of leaving 0
            const uint64_t size_bytes = btrfs_size_bytes(reader, offset);
            const uint64_t end = size_bytes > 0 ? offset + size_bytes : offset + MiB;
            if (!overlaps(claimed_ranges, offset, end)) {
                out.push_back({ResultKind::LostPartitionCandidate,
                               offset,
                               size_bytes,
                               trs("Lost Btrfs partition candidate"),
                               trs("Lost partition finder"),
                               trs("Unallocated Btrfs superblock candidate"),
                               0.81});
                claimed_ranges.push_back({offset, end});
            }
        } else if (looks_xfs(reader, offset)) {
            // Use parse_superblock to also validate geometry (rejects false
            // positives where 'XFSB' appears as random bytes).
            const uint64_t size_bytes = xfs_size_bytes(reader, offset);
            if (size_bytes > 0) {
                const uint64_t end = offset + size_bytes;
                if (!overlaps(claimed_ranges, offset, end)) {
                    out.push_back({ResultKind::LostPartitionCandidate,
                                   offset,
                                   size_bytes,
                                   "Lost XFS partition candidate",
                                   trs("Lost partition finder"),
                                   trs("Unallocated XFS superblock candidate"),
                                   0.86});
                    claimed_ranges.push_back({offset, end});
                }
            }
        } else if (looks_hfs(reader, offset)) {
            // HFS+/HFSX: validate geometry through parse_volume_header.
            const uint64_t size_bytes = hfs_size_bytes(reader, offset);
            if (size_bytes > 0) {
                const uint64_t end = offset + size_bytes;
                if (!overlaps(claimed_ranges, offset, end)) {
                    out.push_back({ResultKind::LostPartitionCandidate,
                                   offset,
                                   size_bytes,
                                   "Lost HFS+ partition candidate",
                                   trs("Lost partition finder"),
                                   trs("Unallocated HFS+/HFSX volume header candidate"),
                                   0.85});
                    claimed_ranges.push_back({offset, end});
                }
            }
        }
    }
 
    if (max_scan < disk_size) {
        logs.push_back(trs("Lost partition finder: scan limited to first ") +
                       std::to_string(max_scan / GiB) + " GiB " +
                       trs("(disk is larger; remaining area not scanned)."));
    }
    logs.push_back(trs("Lost partition finder: deep scan completed."));
    return out;
}