MHR-16: NewBoot - BDiskFormatter and QR.

Signed-off-by: Amlal El Mahrouss <amlal.elmahrouss@icloud.com>

4 files changed, 25 insertions, 12 deletions

diff --git a/Private/NewBoot/BootKit/BootKit.hxx b/Private/NewBoot/BootKit/BootKit.hxx
index b3a81160..76bfbfb0 100644
--- a/Private/NewBoot/BootKit/BootKit.hxx
+++ b/Private/NewBoot/BootKit/BootKit.hxx
@@ -213,6 +213,7 @@ public:
     /// @retval False failed to format.
     Boolean Format(const char* partName, BFileDescriptor* fileBlobs, SizeT blobCount);
 
+    /// @brief check if partition is good.
     operator bool() noexcept {
         fDiskDev.Leak().mBase = (kNewFSAddressAsLba);
         fDiskDev.Leak().mSize = BootDev::kSectorSize;
@@ -228,9 +229,17 @@ public:
                 return false;
         }
 
+        if (blockPart->DiskSize != this->fDiskDev.GetDiskSize()) {
+            EFI::ThrowError(L"Invalid-Disk-Geometry", L"Invalid disk geometry.");
+        }
+
+        if (blockPart->DiskSize < 1) {
+                    EFI::ThrowError(L"Invalid-Disk-Geometry", L"Invalid disk geometry.");
+        }
+
         BTextWriter writer;
 
-        writer.Write(L"Disk Partition: ").Write(blockPart->PartitionName).Write(L" is okay.\r\n");
+        writer.Write(L"Disk Partition: ").Write(blockPart->PartitionName).Write(L" is healthy.\r\n");
 
         return true;
     }
@@ -275,7 +284,6 @@ private:
             forkKind->DataOffset = (startLba + sizeof(NewCatalog) + sizeof(NewFork));
             forkKind->DataSize = blob->fBlobSz;
 
-            Lba lbaStart = forkKind->DataOffset;
             SizeT cur = 0UL;
 
             writer.Write((catalogKind->Kind == kNewFSCatalogKindFile) ? L"New Boot: Write-File: " :
@@ -288,16 +296,14 @@ private:
 
             fDiskDev.Write((Char*)forkKind, sizeof(NewFork));
 
-            while (cur < forkKind->DataSize) {
+            do {
                 this->fDiskDev.Leak().mSize = BootDev::kSectorSize;
                 this->fDiskDev.Leak().mBase = (forkKind->DataOffset + cur);
 
                 this->fDiskDev.Write((Char*)(blob->fBlob) + cur, BootDev::kSectorSize);
 
                 cur += BootDev::kSectorSize;
-                lbaStart += BootDev::kSectorSize;
-            }
-
+            } while (cur < forkKind->DataSize);
 
             /// Fork is done.
 
diff --git a/Private/NewBoot/Source/HEL/AMD64/BootATA.cxx b/Private/NewBoot/Source/HEL/AMD64/BootATA.cxx
index 41ea04ac..4bf7bc96 100644
--- a/Private/NewBoot/Source/HEL/AMD64/BootATA.cxx
+++ b/Private/NewBoot/Source/HEL/AMD64/BootATA.cxx
@@ -237,15 +237,20 @@ SizeT BootDeviceATA::GetDiskSize() noexcept {
 
     boot_ata_wait_io(this->Leak().mBus);
 
-    SizeT result = 0;
+    UInt64 result = 0;
 
-    result += In8(this->Leak().mBus + ATA_CYL_LOW);
+    result = In8(this->Leak().mBus + ATA_CYL_LOW);
     result += In8(this->Leak().mBus + ATA_CYL_MID) << 8;
     result += In8(this->Leak().mBus + ATA_CYL_HIGH) << 16;
-    result += In8(this->Leak().mBus + ATA_CYL_HIGH) << 24;
+
+    Out8(this->Leak().mBus + ATA_REG_CONTROL, 0x80);
+
+    result += In8(this->Leak().mBus + ATA_CYL_LOW) << 24;
+    result += In8(this->Leak().mBus + ATA_CYL_MID) << 32;
+    result += In8(this->Leak().mBus + ATA_CYL_HIGH) << 40;
 
     BTextWriter writer;
-    writer.Write(L"Disk-Size: ").Write(result).Write(L"\r\n");
+    writer.Write(L"Device-Size: ").Write(result).Write(L"\r\n");
 
     return result;
 }
diff --git a/Private/NewBoot/Source/HEL/AMD64/BootFileReader.cxx b/Private/NewBoot/Source/HEL/AMD64/BootFileReader.cxx
index 0e126ada..b2d728ae 100644
--- a/Private/NewBoot/Source/HEL/AMD64/BootFileReader.cxx
+++ b/Private/NewBoot/Source/HEL/AMD64/BootFileReader.cxx
@@ -110,7 +110,7 @@ Void BFileReader::ReadAll(SizeT until, SizeT chunk) {
     if (auto err = BS->AllocatePool(EfiLoaderCode, until, (VoidPtr*)&mBlob) !=
                    kEfiOk) {
       mWriter.Write(L"*** EFI-Code: ").Write(err).Write(L" ***\r\n");
-      EFI::ThrowError(L"NewBoot_PageError", L"Allocation error.");
+      EFI::ThrowError(L"OutOfMemory", L"Allocation error.");
     }
   }
 
diff --git a/Private/NewBoot/Source/HEL/AMD64/BootMain.cxx b/Private/NewBoot/Source/HEL/AMD64/BootMain.cxx
index a928a41a..1c969b28 100644
--- a/Private/NewBoot/Source/HEL/AMD64/BootMain.cxx
+++ b/Private/NewBoot/Source/HEL/AMD64/BootMain.cxx
@@ -181,9 +181,11 @@ EFI_EXTERN_C EFI_API Int Main(EfiHandlePtr ImageHandle,
   /// format the disk.
   //
 
+  constexpr auto binarySize = KIB(512);
+
   /// need this as well, to invoke BExecutableLoader.
   BFileReader readerKernel(L"NewKernel.exe", ImageHandle);
-  readerKernel.ReadAll(MIB(1), 4096);
+  readerKernel.ReadAll(binarySize, BootDeviceATA::kSectorSize);
 
   BDiskFormatFactory<BootDeviceATA> diskFormatter;