Skip to content

Commit

Permalink
Update / Sync
Browse files Browse the repository at this point in the history
  • Loading branch information
@gus33000
gus33000 committed Jan 29, 2025
1 parent 2d9b607 commit 0cd9cc8
Show file tree
Hide file tree
Showing 4 changed files with 108 additions and 75 deletions.
6 changes: 3 additions & 3 deletions Img2Ffu.Library/Writer/FFUFactory.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -127,7 +127,7 @@ public static void GenerateFFU(
List<GPT.Partition> partitions,
Memory<byte> StoreHeaderBuffer,
Memory<byte> WriteDescriptorBuffer,
KeyValuePair<ByteArrayKey, BlockPayload>[] BlockPayloads,
List<KeyValuePair<ByteArrayKey, BlockPayload>> BlockPayloads,
VirtualDisk? InputDisk
)> StoreGenerationParameters = [];

Expand All @@ -148,7 +148,7 @@ public static void GenerateFFU(
List<GPT.Partition> partitions,
Memory<byte> StoreHeaderBuffer,
Memory<byte> WriteDescriptorBuffer,
KeyValuePair<ByteArrayKey, BlockPayload>[] BlockPayloads,
List<KeyValuePair<ByteArrayKey, BlockPayload>> BlockPayloads,
VirtualDisk? InputDisk
) GeneratedStoreParameters = StoreFactory.GenerateStore(
inputForStore,
Expand Down Expand Up @@ -227,7 +227,7 @@ public static void GenerateFFU(
List<GPT.Partition> _,
Memory<byte> StoreHeaderBuffer,
Memory<byte> WriteDescriptorBuffer,
KeyValuePair<ByteArrayKey, BlockPayload>[] _,
List<KeyValuePair<ByteArrayKey, BlockPayload>> _,
VirtualDisk? _
) in StoreGenerationParameters)
{
Expand Down
152 changes: 91 additions & 61 deletions Img2Ffu.Library/Writer/Flashing/BlockPayloadsGenerator.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,7 @@ namespace Img2Ffu.Writer.Flashing
{
internal static class BlockPayloadsGenerator
{
internal static KeyValuePair<ByteArrayKey, BlockPayload>[] GetGPTPayloads(KeyValuePair<ByteArrayKey, BlockPayload>[] blockPayloads, Stream stream, uint BlockSize, bool IsFixedDiskLength)
internal static List<KeyValuePair<ByteArrayKey, BlockPayload>> GetGPTPayloads(List<KeyValuePair<ByteArrayKey, BlockPayload>> blockPayloads, Stream stream, uint BlockSize, bool IsFixedDiskLength)
{
List<KeyValuePair<ByteArrayKey, BlockPayload>> blockPayloadsList = [.. blockPayloads];

Expand Down Expand Up @@ -148,7 +148,22 @@ internal static KeyValuePair<ByteArrayKey, BlockPayload>[] GetGPTPayloads(KeyVal
ulong endGPTChunkStartLocation = (ulong)stream.Length - BlockSize;
byte[] SecondaryGPTBuffer = new byte[(int)BlockSize];
_ = stream.Seek((long)endGPTChunkStartLocation, SeekOrigin.Begin);
_ = stream.Read(SecondaryGPTBuffer);

try
{
_ = stream.Read(SecondaryGPTBuffer);
}
catch (EndOfStreamException)
{
// For some reason there is a bug with DiscUtils itself that can cause this exception to get thrown
// Simply ignore it and replace the payload with a whole lot of nothing.
SecondaryGPTBuffer = new byte[BlockSize];
}
catch
{
throw;
}


MemoryStream secondaryGPTStream = new(SecondaryGPTBuffer);

Expand All @@ -175,20 +190,21 @@ internal static KeyValuePair<ByteArrayKey, BlockPayload>[] GetGPTPayloads(KeyVal
)));
}

return [.. blockPayloadsList];
return blockPayloadsList;
}

internal static KeyValuePair<ByteArrayKey, BlockPayload>[] GetOptimizedPayloads(FlashPart[] flashParts, uint BlockSize, uint BlankSectorBufferSize, ILogging Logging)
internal static List<KeyValuePair<ByteArrayKey, BlockPayload>> GetOptimizedPayloads(FlashPart[] flashParts, uint BlockSize, uint BlankSectorBufferSize, ILogging Logging)
{
List<KeyValuePair<ByteArrayKey, BlockPayload>> hashedBlocks = [];

if (flashParts == null)
{
return [.. hashedBlocks];
return hashedBlocks;
}

ulong CurrentBlockCount = 0;
ulong TotalBlockCount = 0;

foreach (FlashPart flashPart in flashParts)
{
TotalBlockCount += (ulong)flashPart.Stream.Length / BlockSize;
Expand All @@ -207,79 +223,93 @@ internal static KeyValuePair<ByteArrayKey, BlockPayload>[] GetOptimizedPayloads(
List<KeyValuePair<ByteArrayKey, BlockPayload>> blankBlocks = [];

Memory<byte> blockBuffer = new byte[BlockSize];
Span<byte> span = blockBuffer.Span;
Span<byte> FFUBlockPayload = blockBuffer.Span;

foreach (FlashPart flashPart in flashParts)
{
_ = flashPart.Stream.Seek(0, SeekOrigin.Begin);
long totalBlockCount = flashPart.Stream.Length / BlockSize;

for (uint blockIndex = 0; blockIndex < totalBlockCount; blockIndex++)
ulong streamLength = (ulong)flashPart.Stream.Length;
ulong totalBlockCount = streamLength / BlockSize;

for (ulong blockIndex = 0; blockIndex < totalBlockCount; blockIndex++)
{
long streamPosition = flashPart.Stream.Position;
_ = flashPart.Stream.Read(span);
byte[] blockHash = SHA256.HashData(span);
ulong streamPosition = (ulong)flashPart.Stream.Position;

if (!StructuralComparisons.StructuralEqualityComparer.Equals(EMPTY_BLOCK_HASH, blockHash))
try
{
hashedBlocks.Add(new KeyValuePair<ByteArrayKey, BlockPayload>(new ByteArrayKey(blockHash), new BlockPayload(
new WriteDescriptor()
{
BlockDataEntry = new BlockDataEntry()
_ = flashPart.Stream.Read(FFUBlockPayload);
}
catch (EndOfStreamException)
{
// For some reason there is a bug with DiscUtils itself that can cause this exception to get thrown
// Simply ignore it and replace the payload with a whole lot of nothing.
blockBuffer = new byte[BlockSize];
FFUBlockPayload = blockBuffer.Span;
}
catch
{
throw;
}

byte[] FFUBlockHash = SHA256.HashData(FFUBlockPayload);

if (!StructuralComparisons.StructuralEqualityComparer.Equals(EMPTY_BLOCK_HASH, FFUBlockHash) ||
blankPayloadCount < BlankSectorBufferSize)
{
ulong FFUBlockIndex = (flashPart.StartLocation / BlockSize) + blockIndex;

if (FFUBlockIndex > uint.MaxValue)
{
throw new NotSupportedException("The image requires more block than the FFU format can support.");
}

KeyValuePair<ByteArrayKey, BlockPayload> blockDataKeyPair = new(
new ByteArrayKey(FFUBlockHash),
new BlockPayload(
new WriteDescriptor()
{
BlockCount = 1,
LocationCount = 1
},
DiskLocations =
[
new DiskLocation()
BlockDataEntry = new BlockDataEntry()
{
BlockCount = 1,
LocationCount = 1
},
DiskLocations =
[
new DiskLocation()
{
BlockIndex = (uint)((flashPart.StartLocation / BlockSize) + blockIndex),
BlockIndex = (uint)FFUBlockIndex,
DiskAccessMethod = 0
}
]
},
flashPart.Stream,
(ulong)streamPosition
)));
]
},
flashPart.Stream,
streamPosition
));

if (blankPayloadPhase && blankPayloadCount < BlankSectorBufferSize)
if (!StructuralComparisons.StructuralEqualityComparer.Equals(EMPTY_BLOCK_HASH, FFUBlockHash))
{
foreach (KeyValuePair<ByteArrayKey, BlockPayload> blankPayload in blankBlocks)
hashedBlocks.Add(blockDataKeyPair);

if (blankPayloadPhase && blankPayloadCount < BlankSectorBufferSize)
{
hashedBlocks.Add(blankPayload);
foreach (KeyValuePair<ByteArrayKey, BlockPayload> blankPayload in blankBlocks)
{
hashedBlocks.Add(blankPayload);
}
}
}

blankPayloadPhase = false;
blankPayloadCount = 0;
blankBlocks.Clear();
}
else if (blankPayloadCount < BlankSectorBufferSize)
{
blankPayloadPhase = true;
blankPayloadCount++;
blankPayloadPhase = false;
blankPayloadCount = 0;
blankBlocks.Clear();
}
else if (blankPayloadCount < BlankSectorBufferSize)
{
blankPayloadPhase = true;
blankPayloadCount++;

blankBlocks.Add(new KeyValuePair<ByteArrayKey, BlockPayload>(new ByteArrayKey(blockHash), new BlockPayload(
new WriteDescriptor()
{
BlockDataEntry = new BlockDataEntry()
{
BlockCount = 1,
LocationCount = 1
},
DiskLocations =
[
new DiskLocation()
{
BlockIndex = (uint)((flashPart.StartLocation / BlockSize) + blockIndex),
DiskAccessMethod = 0
}
]
},
flashPart.Stream,
(ulong)streamPosition
)));
blankBlocks.Add(blockDataKeyPair);
}
}
else if (blankPayloadCount >= BlankSectorBufferSize && blankBlocks.Count > 0)
{
Expand All @@ -299,7 +329,7 @@ internal static KeyValuePair<ByteArrayKey, BlockPayload>[] GetOptimizedPayloads(
Logging.Log("");
Logging.Log($"FFU Block Count: {hashedBlocks.Count} - {hashedBlocks.Count * BlockSize / (1024 * 1024 * 1024)}GB");

return [.. hashedBlocks];
return hashedBlocks;
}
}
}
11 changes: 6 additions & 5 deletions Img2Ffu.Library/Writer/ImageSplitter.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,6 @@
* OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
* SOFTWARE.
*/
using static Img2Ffu.GPT;

namespace Img2Ffu.Writer
{
Expand All @@ -30,7 +29,7 @@ private static GPT GetGPT(Stream stream, uint BlockSize, uint sectorSize, ILoggi
byte[] GPTBuffer = new byte[BlockSize];
_ = stream.Read(GPTBuffer);

uint requiredGPTBufferSize = GetGPTSize(GPTBuffer, sectorSize);
uint requiredGPTBufferSize = Img2Ffu.GPT.GetGPTSize(GPTBuffer, sectorSize);
if (BlockSize < requiredGPTBufferSize)
{
string errorMessage = $"The Block size is too small to contain the GPT, the GPT is {requiredGPTBufferSize} bytes long, the Block size is {BlockSize} bytes long";
Expand All @@ -42,9 +41,11 @@ private static GPT GetGPT(Stream stream, uint BlockSize, uint sectorSize, ILoggi

GPT GPT = new(GPTBuffer, sectorSize);

if (BlockSize > requiredGPTBufferSize && GPT.Partitions.OrderBy(x => x.FirstSector).Any(x => x.FirstSector < sectorsInABlock))
IOrderedEnumerable<GPT.Partition> orderedGPTPartitions = GPT.Partitions.OrderBy(x => x.FirstSector);

if (BlockSize > requiredGPTBufferSize && orderedGPTPartitions.Any(x => x.FirstSector < sectorsInABlock))
{
Partition conflictingPartition = GPT.Partitions.OrderBy(x => x.FirstSector).First(x => x.FirstSector < sectorsInABlock);
GPT.Partition conflictingPartition = orderedGPTPartitions.First(x => x.FirstSector < sectorsInABlock);

string errorMessage = $"The Block size is too big to contain only the GPT, the GPT is {requiredGPTBufferSize} bytes long, the Block size is {BlockSize} bytes long. The overlapping partition is {conflictingPartition.Name} at {conflictingPartition.FirstSector * sectorSize}";
Logging.Log(errorMessage, ILoggingLevel.Error);
Expand All @@ -54,7 +55,7 @@ private static GPT GetGPT(Stream stream, uint BlockSize, uint sectorSize, ILoggi
return GPT;
}

internal static (FlashPart[], List<Partition> partitions) GetImageSlices(Stream stream, uint BlockSize, string[] ExcludedPartitionNames, uint sectorSize, ILogging Logging)
internal static (FlashPart[], List<GPT.Partition> partitions) GetImageSlices(Stream stream, uint BlockSize, string[] ExcludedPartitionNames, uint sectorSize, ILogging Logging)
{
GPT GPT = GetGPT(stream, BlockSize, sectorSize, Logging);
uint sectorsInABlock = BlockSize / sectorSize;
Expand Down
14 changes: 8 additions & 6 deletions Img2Ffu.Library/Writer/StoreFactory.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,7 @@ namespace Img2Ffu.Writer
{
internal static class StoreFactory
{
private static Memory<byte> GetWriteDescriptorsBuffer(KeyValuePair<ByteArrayKey, BlockPayload>[] payloads, FlashUpdateVersion storeHeaderVersion)
private static Memory<byte> GetWriteDescriptorsBuffer(List<KeyValuePair<ByteArrayKey, BlockPayload>> payloads, FlashUpdateVersion storeHeaderVersion)
{
using MemoryStream WriteDescriptorsStream = new();
using BinaryWriter binaryWriter = new(WriteDescriptorsStream);
Expand All @@ -51,7 +51,7 @@ private static (Stream InputStream, VirtualDisk? InputDisk) OpenInput(string Inp
Stream InputStream;
VirtualDisk? InputDisk = null;

if (InputFile.Contains(@"\\.\physicaldrive", StringComparison.CurrentCultureIgnoreCase))
if (InputFile.Contains(@"\\.\PhysicalDrive", StringComparison.CurrentCultureIgnoreCase))
{
InputStream = new DeviceStream(InputFile, FileAccess.Read);
}
Expand Down Expand Up @@ -87,7 +87,7 @@ internal static (
List<GPT.Partition> partitions,
Memory<byte> StoreHeaderBuffer,
Memory<byte> WriteDescriptorBuffer,
KeyValuePair<ByteArrayKey, BlockPayload>[] BlockPayloads,
List<KeyValuePair<ByteArrayKey, BlockPayload>> BlockPayloads,
VirtualDisk? InputDisk
) GenerateStore(
InputForStore InputForStore,
Expand All @@ -102,6 +102,8 @@ internal static (
Logging.Log("Opening input file...");
(Stream InputStream, VirtualDisk? InputDisk) = OpenInput(InputForStore.InputFile, Logging);

Logging.Log($"Store Length: {InputStream.Length}");

Logging.Log("Generating Image Slices...");
(FlashPart[] flashParts, List<GPT.Partition> partitions) = ImageSplitter.GetImageSlices(InputStream,
BlockSize,
Expand All @@ -110,7 +112,7 @@ internal static (
Logging);

Logging.Log("Generating Block Payloads...");
KeyValuePair<ByteArrayKey, BlockPayload>[] BlockPayloads = BlockPayloadsGenerator.GetOptimizedPayloads(flashParts,
List<KeyValuePair<ByteArrayKey, BlockPayload>> BlockPayloads = BlockPayloadsGenerator.GetOptimizedPayloads(flashParts,
BlockSize,
InputForStore.MaximumNumberOfBlankBlocksAllowed,
Logging);
Expand All @@ -123,14 +125,14 @@ internal static (
Logging.Log("Generating store header...");
StoreHeader store = new()
{
WriteDescriptorCount = (uint)BlockPayloads.LongLength,
WriteDescriptorCount = (uint)BlockPayloads.LongCount(),
WriteDescriptorLength = (uint)WriteDescriptorBuffer.Length,
PlatformIds = PlatformIDs.ToArray(),
BlockSize = BlockSize,
NumberOfStores = NumberOfStores,
StoreIndex = StoreIndex,
DevicePath = InputForStore.DevicePath,
StorePayloadSize = (ulong)BlockPayloads.LongLength * BlockSize
StorePayloadSize = (ulong)BlockPayloads.LongCount() * BlockSize
};

Memory<byte> StoreHeaderBuffer = store.GetResultingBuffer(FlashUpdateVersion, FlashUpdateType.Full, CompressionAlgorithm.None);
Expand Down

0 comments on commit 0cd9cc8

Please sign in to comment.