Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Reduce memory usage during decryption #1318

Closed
wants to merge 2 commits into from
Closed

Reduce memory usage during decryption #1318

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
c3a5287
Reduce memory usage during decryption
ilyfairy Dec 28, 2023
b6e785b
RandomNumberGenerator uses static methods
ilyfairy Dec 30, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
102 changes: 40 additions & 62 deletions SteamKit2/SteamKit2/Util/CryptoHelper.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -202,41 +202,40 @@ public static byte[] SymmetricEncryptWithHMACIV( byte[] input, byte[] key, byte[
var random = GenerateRandomBlock( 3 );
Array.Copy( random, 0, iv, iv.Length - random.Length, random.Length );

using ( var hmac = new HMACSHA1( hmacSecret ) )
using ( var ms = new MemoryStream() )
using ( var ms = new MemoryStream( random.Length + input.Length ) )
Copy link
Member

@xPaw xPaw Dec 28, 2023 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I feel like ArrayPool.Shared should be used in these methods.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I found that if it is a small array, ArrayPool will consume more performance, but stackalloc is worth trying

Copy link
Member

@xPaw xPaw Dec 29, 2023 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Are you sure that will hold true for long running processes? Also take a look at #683

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

image
I think it's fine, the memory will always be collect in time

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can you benchmark the old against the new please, including memory stats? I’d expect the change in allocations to reflect a change in CPU time, but it would be good to have actual figures.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

image

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I ran a couple benchmarks on my laptop too, just to confirm.

This PR drops execution time about 10-30%. If we use ArrayPool<byte>.Shared instead of new MemoryStream then we can get a further 15-25% savings.

However, considering that on my laptop these calls are all measured in microseconds, it's hardly worth the effort to argue over it.

image

I'm happy to merge this if you are @xPaw.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

However, its memory usage is much reduced compared to the original, especially in multi-threading. Multiple threads download chunkdata and decrypt it at the same time, which can reduce the memory by more than 200MB (500MB->300MB).

Copy link
Member

@xPaw xPaw Feb 5, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Okay I looked at this code again, a temporary array is not even required.

Make an instance of HMACSHA1, and then use TransformBlock / TransformFinalBlock if I remember correctly.

But this code also allocates other random arrays like iv and GenerateRandomBlock, this could also be refactored (and I would use ArrayPool).

If the entire encryption flow was refactored, I'm almost positive it's possible to get it down to zero allocations (besides arraypool where it really is needed). But this requires changing code not to return byte arrays, but rather take an existing output array is a parameter.

{
ms.Write( random, 0, random.Length );
ms.Write( input, 0, input.Length );
ms.Seek( 0, SeekOrigin.Begin );

var hash = hmac.ComputeHash( ms );
Array.Copy( hash, iv, iv.Length - random.Length );
var hash = HMACSHA1.HashData( hmacSecret, ms );
Array.Copy( hash, iv, iv.Length - random.Length );
}

return SymmetricEncryptWithIV( input, key, iv );
}

/// <summary>
/// Decrypts using AES/CBC/PKCS7 with an input byte array and key, using the random IV prepended using AES/ECB/None
/// </summary>
public static byte[] SymmetricDecrypt( byte[] input, byte[] key )
{
ArgumentNullException.ThrowIfNull( input );

ArgumentNullException.ThrowIfNull( key );

return SymmetricDecrypt( input, key, out _ );
}

/// <summary>
/// Decrypts using AES/CBC/PKCS7 with an input byte array and key, using the IV (comprised of random bytes and the HMAC-SHA1 of the random bytes and plaintext) prepended using AES/ECB/None
/// </summary>
public static byte[] SymmetricDecryptHMACIV( byte[] input, byte[] key, byte[] hmacSecret )
{
ArgumentNullException.ThrowIfNull( input );

ArgumentNullException.ThrowIfNull( key );

return SymmetricEncryptWithIV( input, key, iv );
}

/// <summary>
/// Decrypts using AES/CBC/PKCS7 with an input byte array and key, using the random IV prepended using AES/ECB/None
/// </summary>
public static byte[] SymmetricDecrypt( byte[] input, byte[] key )
{
ArgumentNullException.ThrowIfNull( input );

ArgumentNullException.ThrowIfNull( key );

return SymmetricDecrypt( input, key, out _ );
}

/// <summary>
/// Decrypts using AES/CBC/PKCS7 with an input byte array and key, using the IV (comprised of random bytes and the HMAC-SHA1 of the random bytes and plaintext) prepended using AES/ECB/None
/// </summary>
public static byte[] SymmetricDecryptHMACIV( byte[] input, byte[] key, byte[] hmacSecret )
{
ArgumentNullException.ThrowIfNull( input );

ArgumentNullException.ThrowIfNull( key );

ArgumentNullException.ThrowIfNull( hmacSecret );

DebugLog.Assert( key.Length >= 16, "CryptoHelper", "SymmetricDecryptHMACIV used with a key smaller than 16 bytes." );
Expand All @@ -247,19 +246,18 @@ public static byte[] SymmetricDecryptHMACIV( byte[] input, byte[] key, byte[] hm

// validate HMAC
byte[] hmacBytes;
using ( var hmac = new HMACSHA1( hmacSecret ) )
using ( var ms = new MemoryStream() )
using ( var ms = new MemoryStream( plaintextData.Length + 3 ) )
{
ms.Write( iv, iv.Length - 3, 3 );
ms.Write( plaintextData, 0, plaintextData.Length );
ms.Seek( 0, SeekOrigin.Begin );

hmacBytes = hmac.ComputeHash( ms );
hmacBytes = HMACSHA1.HashData( hmacSecret, ms );
}

if ( !hmacBytes.Take( iv.Length - 3 ).SequenceEqual( iv.Take( iv.Length - 3 ) ) )
if ( !hmacBytes.AsSpan( 0, iv.Length - 3 ).SequenceEqual( iv.AsSpan( 0, iv.Length - 3 ) ) )
{
throw new CryptographicException( string.Format( CultureInfo.InvariantCulture, "{0} was unable to decrypt packet: HMAC from server did not match computed HMAC.", nameof(NetFilterEncryption) ) );
throw new CryptographicException( string.Format( CultureInfo.InvariantCulture, "{0} was unable to decrypt packet: HMAC from server did not match computed HMAC.", nameof( NetFilterEncryption ) ) );
}

return plaintextData;
Expand All @@ -282,17 +280,15 @@ static byte[] SymmetricDecrypt( byte[] input, byte[] key, out byte[] iv )

// first 16 bytes of input is the ECB encrypted IV
byte[] cryptedIv = new byte[ 16 ];
iv = new byte[ cryptedIv.Length ];
Array.Copy( input, 0, cryptedIv, 0, cryptedIv.Length );

// the rest is ciphertext
byte[] cipherText = new byte[ input.Length - cryptedIv.Length ];
Array.Copy( input, cryptedIv.Length, cipherText, 0, cipherText.Length );
// ciphertext length
int cipherTextLength = input.Length - cryptedIv.Length;

// decrypt the IV using ECB
aes.Mode = CipherMode.ECB;
aes.Padding = PaddingMode.None;

using ( var aesTransform = aes.CreateDecryptor( key, null ) )
{
iv = aesTransform.TransformFinalBlock( cryptedIv, 0, cryptedIv.Length );
Expand All @@ -301,21 +297,11 @@ static byte[] SymmetricDecrypt( byte[] input, byte[] key, out byte[] iv )
// decrypt the remaining ciphertext in cbc with the decrypted IV
aes.Mode = CipherMode.CBC;
aes.Padding = PaddingMode.PKCS7;
aes.Key = key;

using ( var aesTransform = aes.CreateDecryptor( key, iv ) )
using ( var ms = new MemoryStream( cipherText ) )
using ( var cs = new CryptoStream( ms, aesTransform, CryptoStreamMode.Read ) )
{
// plaintext is never longer than ciphertext
byte[] plaintext = new byte[ cipherText.Length ];

int len = cs.ReadAll( plaintext );

byte[] output = new byte[ len ];
Array.Copy( plaintext, 0, output, 0, len );
var output = aes.DecryptCbc( input.AsSpan( start: cryptedIv.Length ), iv, PaddingMode.PKCS7 );

return output;
}
return output;
}

/// <summary>
Expand All @@ -331,10 +317,7 @@ static byte[] SymmetricDecrypt( byte[] input, byte[] key, out byte[] iv )
byte[] password_bytes = Encoding.UTF8.GetBytes( password );
key = SHA256.HashData( password_bytes );

using ( HMACSHA1 hmac = new HMACSHA1( key ) )
{
hash = hmac.ComputeHash( input, 0, 32 );
}
hash = HMACSHA1.HashData( key, input.AsSpan( 0, 32 ) );

for ( int i = 32; i < input.Length; i++ )
if ( input[ i ] != hash[ i % 32 ] )
Expand Down Expand Up @@ -375,7 +358,7 @@ public static byte[] SymmetricDecryptECB( byte[] input, byte[] key )
public static byte[] AdlerHash( byte[] input )
{
ArgumentNullException.ThrowIfNull( input );

uint a = 0, b = 0;
for ( int i = 0 ; i < input.Length ; i++ )
{
Expand All @@ -390,12 +373,7 @@ public static byte[] AdlerHash( byte[] input )
/// </summary>
public static byte[] GenerateRandomBlock( int size )
{
using var rng = RandomNumberGenerator.Create();
var block = new byte[ size ];

rng.GetBytes( block );

return block;
return RandomNumberGenerator.GetBytes(size);
}
}
}
23 changes: 18 additions & 5 deletions SteamKit2/SteamKit2/Util/VZipUtil.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -33,7 +33,10 @@ public static byte[] Decompress(byte[] buffer)
/* uint creationTimestampOrSecondaryCRC = */ reader.ReadUInt32();

byte[] properties = reader.ReadBytes( 5 );
byte[] compressedBuffer = reader.ReadBytes( ( int )ms.Length - HeaderLength - FooterLength - 5 );
var compressedPosition = ms.Position;
var compressedLength = ( int )ms.Length - HeaderLength - FooterLength - 5;
//byte[] compressedBuffer = reader.ReadBytes( ( int )ms.Length - HeaderLength - FooterLength - 5 );
ms.Position += compressedLength;

uint outputCRC = reader.ReadUInt32();
uint sizeDecompressed = reader.ReadUInt32();
Expand All @@ -46,11 +49,21 @@ public static byte[] Decompress(byte[] buffer)
SevenZip.Compression.LZMA.Decoder decoder = new SevenZip.Compression.LZMA.Decoder();
decoder.SetDecoderProperties( properties );

using MemoryStream inputStream = new MemoryStream( compressedBuffer );
//using MemoryStream inputStream = new MemoryStream( compressedBuffer );
using MemoryStream outStream = new MemoryStream( ( int )sizeDecompressed );
decoder.Code( inputStream, outStream, compressedBuffer.Length, sizeDecompressed, null );

var outData = outStream.ToArray();
ms.Position = compressedPosition; // Redirect the location of compressed data, decoder.Code does not read the last 10 bytes
decoder.Code( ms, outStream, compressedLength, sizeDecompressed, null );

byte[] outData;
if ( sizeDecompressed == outStream.Position && sizeDecompressed == outStream.Capacity && sizeDecompressed == outStream.Length)
{
outData = outStream.GetBuffer(); // After specifying sizeDecompressed, MemoryStream will not be expanded. Use GetBuffer to reduce copying.
// At this time, sizeDecompressed == Position == Length == Capacity
}
else
{
outData = outStream.ToArray();
}
if ( Crc32.HashToUInt32( outData ) != outputCRC )
{
throw new InvalidDataException( "CRC does not match decompressed data. VZip data may be corrupted." );
Expand Down
Loading