private static byte[] DecompressForChained(Smb2CompressedPacket packet, Smb2CompressionInfo compressionInfo, Smb2Role role)
{
var p = packet as Smb2ChainedCompressedPacket;
var result = p.Payloads.SelectMany(tuple =>
{
var payloadHeader = tuple.Item1;
switch (payloadHeader.CompressionAlgorithm)
{
case CompressionAlgorithm.NONE:
{
return(tuple.Item2 as byte[]);
}
break;
case CompressionAlgorithm.LZNT1:
case CompressionAlgorithm.LZ77:
case CompressionAlgorithm.LZ77Huffman:
{
var decompressor = GetDecompressor(payloadHeader.CompressionAlgorithm);
var decompressedData = decompressor.Decompress(tuple.Item2 as byte[]);
if (decompressedData.Length != payloadHeader.OriginalPayloadSize)
{
throw new InvalidOperationException("The length decompressed chained payload is inconsistent with OriginalPayloadSize of payload header.");
}
return(decompressedData);
}
break;
case CompressionAlgorithm.Pattern_V1:
{
var pattern = (SMB2_COMPRESSION_PATTERN_PAYLOAD_V1)tuple.Item2;
return(Enumerable.Repeat(pattern.Pattern, (int)pattern.Repetitions).ToArray());
}
break;
default:
{
throw new InvalidOperationException("Unexpected compression algorithm!");
}
break;
}
});
return(result.ToArray());
}
/// <summary>
/// Compress SMB2 packet.
/// </summary>
/// <param name="packet">The SMB2 packet.</param>
/// <param name="compressionInfo">Compression info.</param>
/// <param name="role">SMB2 role.</param>
/// <param name="offset">The offset where compression start, default zero.</param>
/// <returns></returns>
public static Smb2Packet Compress(Smb2CompressiblePacket packet, Smb2CompressionInfo compressionInfo, Smb2Role role, uint offset = 0)
{
var compressionAlgorithm = GetCompressionAlgorithm(packet, compressionInfo, role);
if (compressionAlgorithm == CompressionAlgorithm.NONE)
{
return(packet);
}
var packetBytes = packet.ToBytes();
var compressor = GetCompressor(compressionAlgorithm);
var compressedPacket = new Smb2CompressedPacket();
compressedPacket.Header.ProtocolId = Smb2Consts.ProtocolIdInCompressionTransformHeader;
compressedPacket.Header.OriginalCompressedSegmentSize = (uint)packetBytes.Length;
compressedPacket.Header.CompressionAlgorithm = compressionAlgorithm;
compressedPacket.Header.Reserved = 0;
compressedPacket.Header.Offset = offset;
compressedPacket.UncompressedData = packetBytes.Take((int)offset).ToArray();
compressedPacket.CompressedData = compressor.Compress(packetBytes.Skip((int)offset).ToArray());
// HACK: fake size
if (((Smb2SinglePacket)packet).Header.Command == Smb2Command.WRITE)
{
((Smb2WriteRequestPacket)packet).PayLoad.Length += 0x1000;
compressedPacket.Header.OriginalCompressedSegmentSize += 0x1000;
}
// HACK: force compressed packet to be sent
return(compressedPacket);
// var compressedPackectBytes = compressedPacket.ToBytes();
// Check whether compression shrinks the on-wire packet size
// if (compressedPackectBytes.Length < packetBytes.Length)
// {
// compressedPacket.OriginalPacket = packet;
// return compressedPacket;
// }
// else
// {
// return packet;
// }
}
/// <summary>
/// Decompress the Smb2CompressedPacket.
/// </summary>
/// <param name="packet">The compressed packet.</param>
/// <param name="compressionInfo">Compression info.</param>
/// <param name="role">SMB2 role.</param>
/// <returns>Byte array containing the decompressed packet.</returns>
public static byte[] Decompress(Smb2CompressedPacket packet, Smb2CompressionInfo compressionInfo, Smb2Role role)
{
if (packet.Header.CompressionAlgorithm == CompressionAlgorithm.NONE)
{
throw new InvalidOperationException("Invalid CompressionAlgorithm in header!");
}
if (!compressionInfo.CompressionIds.Any(compressionAlgorithm => compressionAlgorithm == packet.Header.CompressionAlgorithm))
{
throw new InvalidOperationException("The CompressionAlgorithm is not supported!");
}
var decompressor = GetDecompressor(packet.Header.CompressionAlgorithm);
var decompressedBytes = decompressor.Decompress(packet.CompressedData);
var originalPacketBytes = packet.UncompressedData.Concat(decompressedBytes).ToArray();
return(originalPacketBytes);
}
/// <summary>
/// Decompress the Smb2CompressedPacket.
/// </summary>
/// <param name="packet">The compressed packet.</param>
/// <param name="compressionInfo">Compression info.</param>
/// <param name="role">SMB2 role.</param>
/// <returns>Byte array containing the decompressed packet.</returns>
public static byte[] Decompress(Smb2CompressedPacket packet, Smb2CompressionInfo compressionInfo, Smb2Role role)
{
bool isChained = packet.Header.Flags.HasFlag(Compression_Transform_Header_Flags.SMB2_COMPRESSION_FLAG_CHAINED);
byte[] decompressedData;
if (isChained)
{
decompressedData = DecompressForChained(packet, compressionInfo, role);
}
else
{
decompressedData = DecompressForNonChained(packet, compressionInfo, role);
}
if (decompressedData.Length != packet.Header.OriginalCompressedSegmentSize)
{
throw new InvalidOperationException($"The length of decompressed data (0x{decompressedData.Length:X08}) is inconsistent with compression header (0x{packet.Header.OriginalCompressedSegmentSize:X08}).");
}
return(decompressedData);
}
/// <summary>
/// Run compression test against given input parameters.
/// 1. Write testData to test file given by treeId and fileId, and compress message based on compressWriteRequest using compressionAlgorithmForTest.
/// 2. Read out the data just written, and request compressing READ response message based on compressReadRequest.
/// 3. Check whether the READ response based on readResponseShouldBeCompressed.
/// 4. Check whether data read out is equal to test data.
/// </summary>
/// <param name="client">SMB2 functional client to use.</param>
/// <param name="treeId">TreeId to use.</param>
/// <param name="fileId">FileId to use.</param>
public void Run(Smb2FunctionalClient client, uint treeId, FILEID fileId)
{
if (compressionAlgorithmForTest != CompressionAlgorithm.NONE)
{
if (!client.Smb2Client.CompressionInfo.CompressionIds.Any(compressionAlgorithmSupported => compressionAlgorithmSupported == compressionAlgorithmForTest))
{
// The specified compression algorithm is not supported by SUT.
return;
}
}
BaseTestSite.Log.Add(
LogEntryKind.TestStep,
"Test will trigger WRITE request with CompressWrite: {0} and preferred compression algorithm: {1}.",
compressWriteRequest,
compressionAlgorithmForTest
);
// Specify the compression algorithm for write request.
client.Smb2Client.CompressionInfo.PreferredCompressionAlgorithm = compressionAlgorithmForTest;
if (compressWriteRequestBufferOnly)
{
client.Smb2Client.CompressionInfo.CompressBufferOnly = true;
}
client.Write(treeId, fileId, testData, compressWrite: compressWriteRequest);
if (compressWriteRequestBufferOnly)
{
client.Smb2Client.CompressionInfo.CompressBufferOnly = false;
}
byte[] readOutData = null;
bool readResponseIsCompressed = false;
bool readResponseIsChained = false;
Smb2ReadResponsePacket readResponsePacket = null;
Smb2CompressedPacket compressedPacket = null;
Action <Smb2Packet> Smb2Client_PacketReceived = (Smb2Packet obj) =>
{
if (obj is Smb2ReadResponsePacket)
{
readResponsePacket = obj as Smb2ReadResponsePacket;
readResponseIsCompressed = readResponsePacket.Compressed;
compressedPacket = readResponsePacket.CompressedPacket;
if (compressedPacket is Smb2ChainedCompressedPacket)
{
readResponseIsChained = true;
}
}
};
BaseTestSite.Log.Add(
LogEntryKind.TestStep,
"Test will trigger READ request with CompressRead: {0} and check whether READ response is compressed: {1}.",
compressReadRequest,
readResponseShouldBeCompressed
);
client.Smb2Client.PacketReceived += Smb2Client_PacketReceived;
client.Read(treeId, fileId, 0, (uint)testData.Length, out readOutData, compressRead: compressReadRequest);
client.Smb2Client.PacketReceived -= Smb2Client_PacketReceived;
if (compressReadRequest)
{
if (readResponseShouldBeCompressed)
{
BaseTestSite.Assert.IsTrue(readResponseIsCompressed && compressedPacket != null, "[MS-SMB2] section 3.3.5.12: When SMB2_READFLAG_REQUEST_COMPRESSED is specified in read request, the server MUST compress the message if compression will shrink the message size.");
BaseTestSite.Log.Add(LogEntryKind.Debug, "Read response is compressed using {0}.", compressedPacket.Header.CompressionAlgorithm);
if (readResponseShouldBeChained)
{
BaseTestSite.Assert.IsTrue(readResponseIsChained, "The read response should be chained.");
}
else
{
BaseTestSite.Assert.IsFalse(readResponseIsChained, "The read response should not be chained.");
}
}
else
{
BaseTestSite.Assert.IsTrue(!readResponseIsCompressed && compressedPacket == null, "[MS-SMB2] section 3.3.5.12: When SMB2_READFLAG_REQUEST_COMPRESSED is specified in read request, the server MUST NOT compress the message if compression will not shrink the message size.");
}
}
else
{
BaseTestSite.Log.Add(LogEntryKind.Debug, "SMB2_READFLAG_REQUEST_COMPRESSED is not specified in read request, and read response is compressed: {0}.", readResponseIsCompressed);
if (readResponseIsCompressed)
{
BaseTestSite.Assert.IsTrue(compressedPacket != null, "Compressed packet is received.");
BaseTestSite.Log.Add(LogEntryKind.Debug, "Read response is compressed using {0}.", compressedPacket.Header.CompressionAlgorithm);
}
}
BaseTestSite.Assert.IsTrue(Enumerable.SequenceEqual(testData, readOutData), "The read out content MUST be the same with that is written.");
}
/// <summary>
/// Run compression test against given input parameters.
/// 1. Write testData to test file given by treeId and fileId, and compress message based on compressWriteRequest using compressionAlgorithmForTest.
/// 2. Read out the data just written, and request compressing READ response message based on compressReadRequest.
/// 3. Check whether the READ response based on readResponseShouldBeCompressed.
/// 4. Check whether data read out is equal to test data.
/// </summary>
/// <param name="client">SMB2 functional client to use.</param>
/// <param name="variant">Compression test variant.</param>
/// <param name="treeId">TreeId to use.</param>
/// <param name="fileId">FileId to use.</param>
/// <param name="isLargeFile">Whether is large file.</param>
public void Run(Smb2FunctionalClient client, CompressionTestVariant variant, uint treeId, FILEID fileId, bool isLargeFile = false)
{
if (compressionAlgorithmForTest != CompressionAlgorithm.NONE)
{
if (!client.Smb2Client.CompressionInfo.CompressionIds.Any(compressionAlgorithmSupported => compressionAlgorithmSupported == compressionAlgorithmForTest))
{
// The specified compression algorithm is not supported by SUT.
return;
}
}
BaseTestSite.Log.Add(
LogEntryKind.TestStep,
"Test will trigger WRITE request with CompressWrite: {0} and preferred compression algorithm: {1}.",
compressWriteRequest,
compressionAlgorithmForTest
);
// Specify the compression algorithm for write request.
client.Smb2Client.CompressionInfo.PreferredCompressionAlgorithm = compressionAlgorithmForTest;
if (compressWriteRequestBufferOnly)
{
client.Smb2Client.CompressionInfo.CompressBufferOnly = true;
}
if (isLargeFile)
{
// Write several times for existing testData for large data write and read tests
// Base testData:
// Length of LZ77 and LZ77 Huffman testData is 300 bytes.
// Length of LZNT1 testData is 142 bytes.
// Length of Pattern_V1 testData is 256 bytes.
// Length of Compressible data testData is 2048 bytes.
// We will change testData to 1 MB, and write 100 times to generate a 100 MB file
int writeRequestCount = 100;
ulong offset = 0;
// According to real experience with copying large file to shared folder, the client will request 1 MB packet per request.
// Change testData to 512 bytes * 2048 = 1 megabytes, so we can test large file
if (compressionAlgorithmForTest == CompressionAlgorithm.LZ77Huffman)
{
// LZ77Huffman match length needs to be less than 65538, so we use 512*128=65536.
var test512Bytes = Enumerable.Repeat(testData, 2).SelectMany(a => a).Take(512).ToArray();
testData = Enumerable.Repeat(test512Bytes, 128).SelectMany(b => b).ToArray();
}
else if (variant == CompressionTestVariant.ChainedCompressibleWritePatternV1AtFront)
{
// For ChainedCompressibleWritePatternV1AtFront, we need to prepend some data which can be compressed with PatternV1.
var test256Bytes = commonCompressibleData.Take(256).ToArray();
var newCommonCompressibleData = Enumerable.Repeat(test256Bytes, 4 * 1024 - 1).SelectMany(b => b).ToArray();
testData = GenerateByteArray(exampleTestData[CompressionAlgorithm.Pattern_V1], newCommonCompressibleData);
}
else if (variant == CompressionTestVariant.ChainedCompressibleRead || variant == CompressionTestVariant.ChainedCompressibleWritePatternV1AtEnd)
{
// For ChainedCompressibleRead or ChainedCompressibleWritePatternV1AtEnd, we need to append some data which can be compressed with PatternV1.
var test256Bytes = commonCompressibleData.Take(256).ToArray();
var newCommonCompressibleData = Enumerable.Repeat(test256Bytes, 4 * 1024 - 1).SelectMany(b => b).ToArray();
testData = GenerateByteArray(newCommonCompressibleData, exampleTestData[CompressionAlgorithm.Pattern_V1]);
}
else if (variant == CompressionTestVariant.ChainedCompressibleWritePatternV1AtFrontAndEnd)
{
// For ChainedCompressibleWritePatternV1AtFrontAndEnd, we need to add start and end with PatternV1.
var test256Bytes = commonCompressibleData.Take(256).ToArray();
var newCommonCompressibleData = Enumerable.Repeat(test256Bytes, 4 * 1024 - 2).SelectMany(b => b).ToArray();
testData = GenerateByteArray(exampleTestData[CompressionAlgorithm.Pattern_V1], newCommonCompressibleData, exampleTestData[CompressionAlgorithm.Pattern_V1]);
}
else
{
// For other cases, testData with a length of 1 MB will be used.
var test512Bytes = Enumerable.Repeat(testData, 4).SelectMany(a => a).Take(512).ToArray();
testData = Enumerable.Repeat(test512Bytes, 2048).SelectMany(b => b).ToArray();
}
int requestBytes = testData.Length;
for (int time = 0; time < writeRequestCount; time++)
{
client.Write(treeId, fileId, testData, offset, compressWrite: compressWriteRequest);
offset += (uint)requestBytes;
}
}
else
{
client.Write(treeId, fileId, testData, compressWrite: compressWriteRequest);
}
if (compressWriteRequestBufferOnly)
{
client.Smb2Client.CompressionInfo.CompressBufferOnly = false;
}
byte[] readOutData = null;
bool readResponseIsCompressed = false;
bool readResponseIsChained = false;
Smb2ReadResponsePacket readResponsePacket = null;
Smb2CompressedPacket compressedPacket = null;
Action <Smb2Packet> Smb2Client_PacketReceived = (Smb2Packet obj) =>
{
if (obj is Smb2ReadResponsePacket)
{
readResponsePacket = obj as Smb2ReadResponsePacket;
readResponseIsCompressed = readResponsePacket.Compressed;
compressedPacket = readResponsePacket.CompressedPacket;
if (compressedPacket is Smb2ChainedCompressedPacket)
{
readResponseIsChained = true;
}
}
};
BaseTestSite.Log.Add(
LogEntryKind.TestStep,
"Test will trigger READ request with CompressRead: {0} and check whether READ response is compressed: {1}.",
compressReadRequest,
readResponseShouldBeCompressed
);
client.Smb2Client.PacketReceived += Smb2Client_PacketReceived;
if (isLargeFile)
{
// Read several times for exist testData for large data write and read tests
int requestBytes = testData.Length;
int readRequestCount = 100;
ulong offset = 0;
for (int time = 0; time < readRequestCount; time++)
{
client.Read(treeId, fileId, (uint)offset, (uint)requestBytes, out readOutData, compressRead: compressReadRequest);
BaseTestSite.Assert.IsTrue(Enumerable.SequenceEqual(testData, readOutData), $"Request times:{time + 1}, packet offset: {offset}, byteSize:{requestBytes}, the read out content MUST be the same with that is written.");
offset += (uint)requestBytes;
}
}
else
{
client.Read(treeId, fileId, 0, (uint)testData.Length, out readOutData, compressRead: compressReadRequest);
BaseTestSite.Assert.IsTrue(Enumerable.SequenceEqual(testData, readOutData), "The read out content MUST be the same with that is written.");
}
client.Smb2Client.PacketReceived -= Smb2Client_PacketReceived;
if (compressReadRequest)
{
if (readResponseShouldBeCompressed)
{
BaseTestSite.Assert.IsTrue(readResponseIsCompressed && compressedPacket != null, "[MS-SMB2] section 3.3.5.12: When SMB2_READFLAG_REQUEST_COMPRESSED is specified in read request, the server MUST compress the message if compression will shrink the message size.");
BaseTestSite.Log.Add(LogEntryKind.Debug, "Read response is compressed using {0}.", compressedPacket.Header.CompressionAlgorithm);
if (readResponseShouldBeChained)
{
BaseTestSite.Assert.IsTrue(readResponseIsChained, "The read response should be chained.");
}
else
{
BaseTestSite.Assert.IsFalse(readResponseIsChained, "The read response should not be chained.");
}
}
else
{
BaseTestSite.Assert.IsTrue(!readResponseIsCompressed && compressedPacket == null, "[MS-SMB2] section 3.3.5.12: When SMB2_READFLAG_REQUEST_COMPRESSED is specified in read request, the server MUST NOT compress the message if compression will not shrink the message size.");
}
}
else
{
BaseTestSite.Log.Add(LogEntryKind.Debug, "SMB2_READFLAG_REQUEST_COMPRESSED is not specified in read request, and read response is compressed: {0}.", readResponseIsCompressed);
if (readResponseIsCompressed)
{
BaseTestSite.Assert.IsTrue(compressedPacket != null, "Compressed packet is received.");
BaseTestSite.Log.Add(LogEntryKind.Debug, "Read response is compressed using {0}.", compressedPacket.Header.CompressionAlgorithm);
}
}
}
