--- truecrypt/boot/windows/bootdiskio.cpp	2018/04/24 16:49:04	1.1
+++ truecrypt/boot/windows/bootdiskio.cpp	2018/04/24 16:52:09	1.1.1.2
@@ -9,9 +9,32 @@
 #include "Bios.h"
 #include "BootConsoleIo.h"
 #include "BootDebug.h"
+#include "BootDefs.h"
 #include "BootDiskIo.h"
 
 
+byte SectorBuffer[TC_LB_SIZE];
+
+#ifdef TC_BOOT_DEBUG_ENABLED
+static bool SectorBufferInUse = false;
+
+void AcquireSectorBuffer ()
+{
+	if (SectorBufferInUse)
+		TC_THROW_FATAL_EXCEPTION;
+
+	SectorBufferInUse = true;
+}
+
+
+void ReleaseSectorBuffer ()
+{
+	SectorBufferInUse = false;
+}
+
+#endif
+
+
 bool IsLbaSupported (byte drive)
 {
 	uint16 result = 0;
@@ -65,26 +88,28 @@ void Print (const ChsAddress &chs)
 }
 
 
-BiosResult ReadWriteSectors (bool write, byte *buffer, byte drive, const ChsAddress &chs, byte sectorCount, bool silent)
+BiosResult ReadWriteSectors (bool write, uint16 bufferSegment, uint16 bufferOffset, byte drive, const ChsAddress &chs, byte sectorCount, bool silent)
 {
+	CheckStack();
+
 	byte cylinderLow = (byte) chs.Cylinder;
 	byte sector = chs.Sector;
 	sector |= byte (chs.Cylinder >> 2) & 0xc0;
 	byte function = write ? 0x03 : 0x02;
-	
+
 	BiosResult result;
 	__asm
 	{
 		push es
-		mov ax, ds
+		mov ax, bufferSegment
 		mov	es, ax
-		mov	bx, ss:buffer
-		mov dl, ss:drive
-		mov ch, ss:cylinderLow
+		mov	bx, bufferOffset
+		mov dl, drive
+		mov ch, cylinderLow
 		mov si, chs
-		mov dh, ss:[si].Head
-		mov cl, ss:sector
-		mov	al, ss:sectorCount
+		mov dh, [si].Head
+		mov cl, sector
+		mov	al, sectorCount
 		mov	ah, function
 		int	0x13
 		mov	result, ah
@@ -101,6 +126,14 @@ BiosResult ReadWriteSectors (bool write,
 }
 
 
+BiosResult ReadWriteSectors (bool write, byte *buffer, byte drive, const ChsAddress &chs, byte sectorCount, bool silent)
+{
+	uint16 codeSeg;
+	__asm mov codeSeg, cs
+	return ReadWriteSectors (false, codeSeg, (uint16) buffer, drive, chs, sectorCount, silent);
+}
+
+
 BiosResult ReadSectors (byte *buffer, byte drive, const ChsAddress &chs, byte sectorCount, bool silent)
 {
 	return ReadWriteSectors (false, buffer, drive, chs, sectorCount, silent);
@@ -115,6 +148,8 @@ BiosResult WriteSectors (byte *buffer, b
 
 static BiosResult ReadWriteSectors (bool write, BiosLbaPacket &dapPacket, byte drive, const uint64 &sector, uint16 sectorCount, bool silent)
 {
+	CheckStack();
+
 	dapPacket.Size = sizeof (dapPacket);
 	dapPacket.Reserved = 0;
 	dapPacket.SectorCount = sectorCount;
@@ -160,10 +195,10 @@ BiosResult ReadWriteSectors (bool write,
 }
 
 
-BiosResult ReadSectors (byte *buffer, byte drive, const uint64 &sector, uint16 sectorCount, bool silent)
+BiosResult ReadSectors (uint16 bufferSegment, uint16 bufferOffset, byte drive, const uint64 &sector, uint16 sectorCount, bool silent)
 {
 	if (IsLbaSupported (drive))
-		return ReadWriteSectors (false, buffer, drive, sector, sectorCount, silent);
+		return ReadWriteSectors (false, bufferSegment, bufferOffset, drive, sector, sectorCount, silent);
 
 	DriveGeometry geometry;
 
@@ -173,7 +208,15 @@ BiosResult ReadSectors (byte *buffer, by
 
 	ChsAddress chs;
 	LbaToChs (geometry, sector, chs);
-	return ReadWriteSectors (false, buffer, drive, chs, sectorCount, silent);
+	return ReadWriteSectors (false, bufferSegment, bufferOffset, drive, chs, sectorCount, silent);
+}
+
+
+BiosResult ReadSectors (byte *buffer, byte drive, const uint64 &sector, uint16 sectorCount, bool silent)
+{
+	uint16 codeSeg;
+	__asm mov codeSeg, cs
+	return ReadSectors (codeSeg, (uint16) buffer, drive, sector, sectorCount, silent);
 }
 
 
@@ -185,6 +228,8 @@ BiosResult WriteSectors (byte *buffer, b
 
 BiosResult GetDriveGeometry (byte drive, DriveGeometry &geometry, bool silent)
 {
+	CheckStack();
+
 	byte maxCylinderLow, maxHead, maxSector;
 	BiosResult result;
 	__asm
@@ -236,17 +281,6 @@ void LbaToChs (const DriveGeometry &geom
 }
 
 
-BiosResult ReadMBR (byte drive, MBR &mbr, bool silent)
-{
-	ChsAddress chs;
-	chs.Cylinder = 0;
-	chs.Head = 0;
-	chs.Sector = 1;
-
-	return ReadSectors ((byte *) &mbr, drive, chs, 1, silent);
-}
-
-
 void PartitionEntryMBRToPartition (const PartitionEntryMBR &partEntry, Partition &partition)
 {
 	partition.Active = partEntry.BootIndicator == 0x80;
@@ -256,31 +290,36 @@ void PartitionEntryMBRToPartition (const
 	partition.SectorCount.LowPart = partEntry.SectorCountLBA;
 	partition.StartSector.HighPart = 0;
 	partition.StartSector.LowPart = partEntry.StartLBA;
-
-	partition.StartChsAddress.Cylinder = partEntry.StartCylinder | (uint16 (partEntry.StartCylSector & 0xc0) << 2);
-	partition.StartChsAddress.Head = partEntry.StartHead;
-	partition.StartChsAddress.Sector = partEntry.StartCylSector & ~0xc0;
-
 	partition.Type = partEntry.Type;
 }
 
 
 BiosResult GetDrivePartitions (byte drive, Partition *partitionArray, size_t partitionArrayCapacity, size_t &partitionCount, bool activeOnly, bool silent)
 {
-	MBR mbr;
-	BiosResult result = ReadMBR (drive, mbr, silent);
-	partitionCount = 0;
+	ChsAddress chs;
+	chs.Cylinder = 0;
+	chs.Head = 0;
+	chs.Sector = 1;
+
+	AcquireSectorBuffer();
+	BiosResult result = ReadSectors (SectorBuffer, drive, chs, 1, silent);
 	
-	if (result != BiosResultSuccess || mbr.Signature != 0xaa55)
+	ReleaseSectorBuffer();
+	partitionCount = 0;
+
+	MBR *mbr = (MBR *) SectorBuffer;
+	if (result != BiosResultSuccess || mbr->Signature != 0xaa55)
 		return result;
 
+	PartitionEntryMBR mbrPartitions[4];
+	memcpy (mbrPartitions, mbr->Partitions, sizeof (mbrPartitions));
 	size_t partitionArrayPos = 0, partitionNumber;
 	
 	for (partitionNumber = 0;
-		partitionNumber < array_capacity (mbr.Partitions) && partitionArrayPos < partitionArrayCapacity;
+		partitionNumber < array_capacity (mbrPartitions) && partitionArrayPos < partitionArrayCapacity;
 		++partitionNumber)
 	{
-		const PartitionEntryMBR &partEntry = mbr.Partitions[partitionNumber];
+		const PartitionEntryMBR &partEntry = mbrPartitions[partitionNumber];
 		
 		if (partEntry.SectorCountLBA > 0)
 		{
@@ -300,16 +339,16 @@ BiosResult GetDrivePartitions (byte driv
 					// Find all extended partitions
 					uint64 firstExtStartLBA = partition.StartSector;
 					uint64 extStartLBA = partition.StartSector;
-					MBR extMbr;
+					MBR *extMbr = (MBR *) SectorBuffer;
 
 					while (partitionArrayPos < partitionArrayCapacity &&
-						(result = ReadSectors ((byte *) &extMbr, drive, extStartLBA, 1, silent)) == BiosResultSuccess
-						&& extMbr.Signature == 0xaa55)
+						(result = ReadSectors ((byte *) extMbr, drive, extStartLBA, 1, silent)) == BiosResultSuccess
+						&& extMbr->Signature == 0xaa55)
 					{
-						if (extMbr.Partitions[0].SectorCountLBA > 0)
+						if (extMbr->Partitions[0].SectorCountLBA > 0)
 						{
 							Partition &logPart = partitionArray[partitionArrayPos++];
-							PartitionEntryMBRToPartition (extMbr.Partitions[0], logPart);
+							PartitionEntryMBRToPartition (extMbr->Partitions[0], logPart);
 							logPart.Drive = drive;
 
 							logPart.Number = partitionNumber++;
@@ -320,11 +359,11 @@ BiosResult GetDrivePartitions (byte driv
 						}
 
 						// Secondary extended
-						if (extMbr.Partitions[1].Type != 0x5 && extMbr.Partitions[1].Type == 0xf
-							|| extMbr.Partitions[1].SectorCountLBA == 0)
+						if (extMbr->Partitions[1].Type != 0x5 && extMbr->Partitions[1].Type == 0xf
+							|| extMbr->Partitions[1].SectorCountLBA == 0)
 							break;
 
-						extStartLBA.LowPart = extMbr.Partitions[1].StartLBA + firstExtStartLBA.LowPart;
+						extStartLBA.LowPart = extMbr->Partitions[1].StartLBA + firstExtStartLBA.LowPart;
 					}
 				}
 			}