Skip to content

Commit

Permalink
Fixed SquashFs fragment table issue
Browse files Browse the repository at this point in the history
  • Loading branch information
LTRData committed Apr 4, 2024
1 parent 8337290 commit cc31cce
Show file tree
Hide file tree
Showing 6 changed files with 119 additions and 126 deletions.
2 changes: 2 additions & 0 deletions Library/DiscUtils.SquashFs/FragmentWriter.cs
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 +48,8 @@ public FragmentWriter(BuilderContext context)

public int FragmentCount { get; private set; }

public int FragmentBlocksCount => _fragmentBlocks.Count;

public void Flush()
{
if (_currentOffset != 0)
Expand Down
29 changes: 20 additions & 9 deletions Library/DiscUtils.SquashFs/MetablockWriter.cs
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
//
// Copyright (c) 2008-2011, Kenneth Bell
// Copyright (c) 2008-2024, Kenneth Bell, Olof Lagerkvist
//
// Permission is hereby granted, free of charge, to any person obtaining a
// copy of this software and associated documentation files (the "Software"),
Expand All @@ -26,6 +26,7 @@
using DiscUtils.Compression;
using DiscUtils.Streams;
using DiscUtils.Streams.Compatibility;
using LTRData.Extensions.Buffers;

namespace DiscUtils.SquashFs;

Expand Down Expand Up @@ -58,13 +59,19 @@ public void Dispose()
}

public void Write(byte[] buffer, int offset, int count)
=> Write(buffer.AsSpan(offset, count));

public void Write(ReadOnlySpan<byte> buffer)
{
var totalStored = 0;

while (totalStored < count)
while (totalStored < buffer.Length)
{
var toCopy = Math.Min(_currentBlock.Length - _currentOffset, count - totalStored);
System.Buffer.BlockCopy(buffer, offset + totalStored, _currentBlock, _currentOffset, toCopy);
var toCopy = Math.Min(_currentBlock.Length - _currentOffset, buffer.Length - totalStored);
var sourceSlice = buffer.Slice(totalStored, toCopy);
var destinationSlice = new Span<byte>(_currentBlock, _currentOffset, toCopy);
sourceSlice.CopyTo(destinationSlice);

_currentOffset += toCopy;
totalStored += toCopy;

Expand All @@ -83,7 +90,7 @@ internal void Persist(Stream output)
NextBlock();
}

output.Write(_buffer.ToArray(), 0, (int)_buffer.Length);
output.Write(_buffer.AsSpan());
}

internal long DistanceFrom(MetadataRef startPos)
Expand All @@ -94,29 +101,33 @@ internal long DistanceFrom(MetadataRef startPos)

private void NextBlock()
{
const int SQUASHFS_COMPRESSED_BIT = 1 << 15;

var compressed = new MemoryStream();
using (var compStream = new ZlibStream(compressed, CompressionMode.Compress, true))
{
compStream.Write(_currentBlock, 0, _currentOffset);
}

byte[] writeData;
Span<byte> writeData;
ushort writeLen;

if (compressed.Length < _currentOffset)
{
writeData = compressed.ToArray();
var compressedData = compressed.AsSpan();
writeData = compressedData;
writeLen = (ushort)compressed.Length;
}
else
{
writeData = _currentBlock;
writeLen = (ushort)(_currentOffset | 0x8000);
writeLen = (ushort)(_currentOffset | SQUASHFS_COMPRESSED_BIT);
}

Span<byte> header = stackalloc byte[2];
EndianUtilities.WriteBytesLittleEndian(writeLen, header);
_buffer.Write(header);
_buffer.Write(writeData, 0, writeLen & 0x7FFF);
_buffer.Write(writeData.Slice(0, writeLen & 0x7FFF));

++_currentBlockNum;
}
Expand Down
139 changes: 32 additions & 107 deletions Library/DiscUtils.SquashFs/SquashFileSystemBuilder.cs
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
//
// Copyright (c) 2008-2011, Kenneth Bell
// Copyright (c) 2008-2024, Kenneth Bell, Olof Lagerkvist
//
// Permission is hereby granted, free of charge, to any person obtaining a
// copy of this software and associated documentation files (the "Software"),
Expand Down Expand Up @@ -30,6 +30,7 @@
using DiscUtils.Internal;
using DiscUtils.Streams;
using DiscUtils.Streams.Compatibility;
using LTRData.Extensions.Buffers;

namespace DiscUtils.SquashFs;

Expand Down Expand Up @@ -395,7 +396,7 @@ public override void Build(Stream output)
Magic = SuperBlock.SquashFsMagic,
CreationTime = DateTime.Now,
BlockSize = (uint)_context.DataBlockSize,
Compression = 1 // DEFLATE
Compression = SuperBlock.CompressionType.ZLib
};
superBlock.BlockSizeLog2 = (ushort)MathUtilities.Log2(superBlock.BlockSize);
superBlock.MajorVersion = 4;
Expand All @@ -408,8 +409,9 @@ public override void Build(Stream output)
fragWriter.Flush();
superBlock.RootInode = GetRoot().InodeRef;
superBlock.InodesCount = _nextInode - 1;
superBlock.FragmentsCount = (uint)fragWriter.FragmentCount;
superBlock.FragmentsCount = (uint)fragWriter.FragmentBlocksCount;
superBlock.UidGidCount = (ushort)idWriter.IdCount;
superBlock.Flags = SuperBlock.SuperBlockFlags.NoXAttrs | SuperBlock.SuperBlockFlags.FragmentsAlwaysGenerated;

superBlock.InodeTableStart = output.Position;
inodeWriter.Persist(output);
Expand All @@ -421,115 +423,33 @@ public override void Build(Stream output)
superBlock.LookupTableStart = -1;
superBlock.UidGidTableStart = idWriter.Persist();
superBlock.ExtendedAttrsTableStart = -1;

superBlock.BytesUsed = output.Position;

// Pad to 4KB
var end = MathUtilities.RoundUp(output.Position, 4 * Sizes.OneKiB);
if (end != output.Position)
{
var padding = new byte[(int)(end - output.Position)];
output.Write(padding, 0, padding.Length);
Span<byte> padding = stackalloc byte[(int)(end - output.Position)];
padding.Clear();
output.Write(padding);
}

// Go back and write the superblock
output.Position = 0;
var buffer = new byte[superBlock.Size];
Span<byte> buffer = stackalloc byte[superBlock.Size];
superBlock.WriteTo(buffer);
output.Write(buffer, 0, buffer.Length);
output.Write(buffer);
output.Position = end;
}

/// <summary>
/// Writes the file system to an existing stream.
/// </summary>
/// <param name="output">The stream to write to.</param>
/// <param name="cancellationToken"></param>
/// <remarks>The <c>output</c> stream must support seeking and writing.</remarks>
public async override Task BuildAsync(Stream output, CancellationToken cancellationToken)
public override Task BuildAsync(Stream output, CancellationToken cancellationToken)
{
if (output == null)
{
throw new ArgumentNullException(nameof(output));
}

if (!output.CanWrite)
{
throw new ArgumentException("Output stream must be writable", nameof(output));
}

if (!output.CanSeek)
{
throw new ArgumentException("Output stream must support seeking", nameof(output));
}

_context = new BuilderContext
{
RawStream = output,
DataBlockSize = DefaultBlockSize,
IoBuffer = new byte[DefaultBlockSize]
};

var inodeWriter = new MetablockWriter();
var dirWriter = new MetablockWriter();
var fragWriter = new FragmentWriter(_context);
var idWriter = new IdTableWriter(_context);

_context.AllocateInode = AllocateInode;
_context.AllocateId = idWriter.AllocateId;
_context.WriteDataBlock = WriteDataBlock;
_context.WriteFragment = fragWriter.WriteFragment;
_context.InodeWriter = inodeWriter;
_context.DirectoryWriter = dirWriter;

_nextInode = 1;

var superBlock = new SuperBlock
{
Magic = SuperBlock.SquashFsMagic,
CreationTime = DateTime.Now,
BlockSize = (uint)_context.DataBlockSize,
Compression = 1 // DEFLATE
};
superBlock.BlockSizeLog2 = (ushort)MathUtilities.Log2(superBlock.BlockSize);
superBlock.MajorVersion = 4;
superBlock.MinorVersion = 0;

output.Position = superBlock.Size;

GetRoot().Reset();
GetRoot().Write(_context);
fragWriter.Flush();
superBlock.RootInode = GetRoot().InodeRef;
superBlock.InodesCount = _nextInode - 1;
superBlock.FragmentsCount = (uint)fragWriter.FragmentCount;
superBlock.UidGidCount = (ushort)idWriter.IdCount;

superBlock.InodeTableStart = output.Position;
inodeWriter.Persist(output);
cancellationToken.ThrowIfCancellationRequested();

superBlock.DirectoryTableStart = output.Position;
dirWriter.Persist(output);
Build(output);

superBlock.FragmentTableStart = fragWriter.Persist();
superBlock.LookupTableStart = -1;
superBlock.UidGidTableStart = idWriter.Persist();
superBlock.ExtendedAttrsTableStart = -1;
superBlock.BytesUsed = output.Position;

// Pad to 4KB
var end = MathUtilities.RoundUp(output.Position, 4 * Sizes.OneKiB);
if (end != output.Position)
{
var padding = new byte[(int)(end - output.Position)];
await output.WriteAsync(padding, cancellationToken).ConfigureAwait(false);
}

// Go back and write the superblock
output.Position = 0;
var buffer = new byte[superBlock.Size];
superBlock.WriteTo(buffer);
await output.WriteAsync(buffer, cancellationToken).ConfigureAwait(false);
output.Position = end;
return Task.CompletedTask;
}

/// <summary>
Expand All @@ -552,32 +472,37 @@ private uint AllocateInode()
/// a flag indicating compression (or not).
/// </returns>
private uint WriteDataBlock(byte[] buffer, int offset, int count)
=> WriteDataBlock(buffer.AsSpan(offset, count));

private uint WriteDataBlock(ReadOnlySpan<byte> buffer)
{
const int SQUASHFS_COMPRESSED_BIT = 1 << 24;

var compressed = new MemoryStream();
using (var compStream = new ZlibStream(compressed, CompressionMode.Compress, true))
{
compStream.Write(buffer, offset, count);
compStream.Write(buffer);
}

byte[] writeData;
int writeOffset;
int writeLen;
if (compressed.Length < count)
ReadOnlySpan<byte> writeData;
int returnValue;

if (compressed.Length < buffer.Length)
{
writeData = compressed.ToArray();
writeOffset = 0;
writeLen = (int)compressed.Length;
var compressedData = compressed.AsSpan();
compressedData[1] = 0xda;
writeData = compressedData;
returnValue = writeData.Length;
}
else
{
writeData = buffer;
writeOffset = offset;
writeLen = count | 0x01000000;
returnValue = writeData.Length | SQUASHFS_COMPRESSED_BIT; // Flag to indicate uncompressed buffer
}

_context.RawStream.Write(writeData, writeOffset, writeLen & 0xFFFFFF);
_context.RawStream.Write(writeData);

return (uint)writeLen;
return (uint)returnValue;
}

/// <summary>
Expand Down
41 changes: 34 additions & 7 deletions Library/DiscUtils.SquashFs/SuperBlock.cs
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
//
// Copyright (c) 2008-2011, Kenneth Bell
// Copyright (c) 2008-2024, Kenneth Bell, Olof Lagerkvist
//
// Permission is hereby granted, free of charge, to any person obtaining a
// copy of this software and associated documentation files (the "Software"),
Expand Down Expand Up @@ -27,15 +27,42 @@ namespace DiscUtils.SquashFs;

internal class SuperBlock : IByteArraySerializable
{
public enum CompressionType : ushort
{
Unknown,
ZLib,
LZMA,
Lzo,
Xz,
Lz4,
ZStd
}

[Flags]
public enum SuperBlockFlags : ushort
{
UncompressedInodes = 0x0001,
UncompressedBlocks = 0x0002,
UncompressedFragments = 0x0008,
FragmentsNotUsed = 0x0010,
FragmentsAlwaysGenerated = 0x0020,
DeduplicatedData = 0x0040,
NFSExportTableExists = 0x0080,
UncompressedXAttrs = 0x0100,
NoXAttrs = 0x0200,
CompressorOptionsPresent = 0x0400,
UncompressedIdTable = 0x0800
}

public const uint SquashFsMagic = 0x73717368;
public uint BlockSize;
public ushort BlockSizeLog2;
public long BytesUsed;
public ushort Compression;
public CompressionType Compression;
public DateTime CreationTime;
public long DirectoryTableStart;
public long ExtendedAttrsTableStart;
public ushort Flags;
public SuperBlockFlags Flags;
public uint FragmentsCount;
public long FragmentTableStart;
public uint InodesCount;
Expand Down Expand Up @@ -63,9 +90,9 @@ public int ReadFrom(ReadOnlySpan<byte> buffer)
CreationTime = DateTimeOffset.FromUnixTimeSeconds(EndianUtilities.ToUInt32LittleEndian(buffer.Slice(8))).DateTime;
BlockSize = EndianUtilities.ToUInt32LittleEndian(buffer.Slice(12));
FragmentsCount = EndianUtilities.ToUInt32LittleEndian(buffer.Slice(16));
Compression = EndianUtilities.ToUInt16LittleEndian(buffer.Slice(20));
Compression = (CompressionType)EndianUtilities.ToUInt16LittleEndian(buffer.Slice(20));
BlockSizeLog2 = EndianUtilities.ToUInt16LittleEndian(buffer.Slice(22));
Flags = EndianUtilities.ToUInt16LittleEndian(buffer.Slice(24));
Flags = (SuperBlockFlags)EndianUtilities.ToUInt16LittleEndian(buffer.Slice(24));
UidGidCount = EndianUtilities.ToUInt16LittleEndian(buffer.Slice(26));
MajorVersion = EndianUtilities.ToUInt16LittleEndian(buffer.Slice(28));
MinorVersion = EndianUtilities.ToUInt16LittleEndian(buffer.Slice(30));
Expand All @@ -88,9 +115,9 @@ public void WriteTo(Span<byte> buffer)
EndianUtilities.WriteBytesLittleEndian(Convert.ToUInt32(new DateTimeOffset(CreationTime).ToUnixTimeSeconds()), buffer.Slice(8));
EndianUtilities.WriteBytesLittleEndian(BlockSize, buffer.Slice(12));
EndianUtilities.WriteBytesLittleEndian(FragmentsCount, buffer.Slice(16));
EndianUtilities.WriteBytesLittleEndian(Compression, buffer.Slice(20));
EndianUtilities.WriteBytesLittleEndian((ushort)Compression, buffer.Slice(20));
EndianUtilities.WriteBytesLittleEndian(BlockSizeLog2, buffer.Slice(22));
EndianUtilities.WriteBytesLittleEndian(Flags, buffer.Slice(24));
EndianUtilities.WriteBytesLittleEndian((ushort)Flags, buffer.Slice(24));
EndianUtilities.WriteBytesLittleEndian(UidGidCount, buffer.Slice(26));
EndianUtilities.WriteBytesLittleEndian(MajorVersion, buffer.Slice(28));
EndianUtilities.WriteBytesLittleEndian(MinorVersion, buffer.Slice(30));
Expand Down
2 changes: 1 addition & 1 deletion Library/DiscUtils.SquashFs/VfsSquashFileSystemReader.cs
Original file line number Diff line number Diff line change
Expand Up @@ -59,7 +59,7 @@ public VfsSquashFileSystemReader(Stream stream)
throw new IOException("Invalid SquashFS filesystem - magic mismatch");
}

if (_context.SuperBlock.Compression != 1)
if (_context.SuperBlock.Compression != SuperBlock.CompressionType.ZLib)
{
throw new IOException("Unsupported compression used");
}
Expand Down
Loading

0 comments on commit cc31cce

Please sign in to comment.