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.

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

Improve transfer speed (sending) #178

Merged
merged 11 commits into from
Jul 30, 2024
6 changes: 5 additions & 1 deletion Directory.Build.targets
Original file line number Diff line number Diff line change
@@ -1,11 +1,15 @@
<Project>
<ItemGroup Label="Global Dependencies">
<PackageReference Include="ShortDev.IO" Version="0.1.3" />
<PackageReference Include="ShortDev.IO" Version="0.1.4" />
</ItemGroup>

<ItemGroup Label="Global Usings">
<Using Include="ShortDev.IO" />
<Using Include="ShortDev.IO.Input" />
<Using Include="ShortDev.IO.Output" />
</ItemGroup>

<ItemGroup Label="Tests">
<InternalsVisibleTo Include="ShortDev.Microsoft.ConnectedDevices.Test" />
</ItemGroup>
</Project>
Original file line number Diff line number Diff line change
@@ -1,5 +1,4 @@
using System.IO;
using System.Text;
using System.Text;

namespace ShortDev.Microsoft.ConnectedDevices.NearShare;

Expand Down Expand Up @@ -43,12 +42,10 @@ public static CdpFileProvider FromStream(string fileName, Stream stream)
public ulong FileSize
=> (ulong)_buffer.Length;

public ReadOnlySpan<byte> ReadBlob(ulong start, uint length)
public void ReadBlob(ulong start, Span<byte> buffer)
{
Span<byte> buffer = new byte[length];
_buffer.Position = (long)start;
_buffer.Read(buffer);
return buffer;
}

public void Dispose()
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,47 @@
using ShortDev.Microsoft.ConnectedDevices.Serialization;

namespace ShortDev.Microsoft.ConnectedDevices.NearShare.Messages;
internal static class FetchDataResponse
{
public static void Write(EndianWriter writer, uint contentId, ulong start, int length, out Span<byte> blob)
{
CompactBinaryBondWriter bondWriter = new(writer.Buffer);

bondWriter.WriteFieldBegin(Bond.BondDataType.BT_MAP, 1);
bondWriter.WriteContainerBegin(count: 4, Bond.BondDataType.BT_WSTRING, Bond.BondDataType.BT_STRUCT);

WritePropertyBegin(ref bondWriter, "ControlMessage", PropertyType.PropertyType_UInt32);
bondWriter.WriteFieldBegin(Bond.BondDataType.BT_UINT32, 104);
bondWriter.WriteUInt32((uint)NearShareControlMsgType.FetchDataResponse);
bondWriter.WriteStructEnd();

WritePropertyBegin(ref bondWriter, "ContentId", PropertyType.PropertyType_UInt32);
bondWriter.WriteFieldBegin(Bond.BondDataType.BT_UINT32, 104);
bondWriter.WriteUInt32(contentId);
bondWriter.WriteStructEnd();

WritePropertyBegin(ref bondWriter, "BlobPosition", PropertyType.PropertyType_UInt64);
bondWriter.WriteFieldBegin(Bond.BondDataType.BT_UINT64, 106);
bondWriter.WriteUInt64(start);
bondWriter.WriteStructEnd();

WritePropertyBegin(ref bondWriter, "DataBlob", PropertyType.PropertyType_UInt8Array);
bondWriter.WriteFieldBegin(Bond.BondDataType.BT_LIST, 200);
bondWriter.WriteContainerBegin(length, Bond.BondDataType.BT_UINT8);

blob = writer.Buffer.GetSpan(length)[..length];
writer.Buffer.Advance(length);

bondWriter.WriteStructEnd();

bondWriter.WriteStructEnd();
}

static void WritePropertyBegin(ref CompactBinaryBondWriter writer, string name, PropertyType type)
{
writer.WriteWString(name);

writer.WriteFieldBegin(Bond.BondDataType.BT_INT32, 0);
writer.WriteInt32((int)type);
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,8 @@
using ShortDev.Microsoft.ConnectedDevices.NearShare.Apps;
using ShortDev.Microsoft.ConnectedDevices.NearShare.Messages;
using ShortDev.Microsoft.ConnectedDevices.Serialization;
using System.Buffers;
using System.Diagnostics;

namespace ShortDev.Microsoft.ConnectedDevices.NearShare;

Expand Down Expand Up @@ -176,21 +178,20 @@ void HandleDataRequest(BinaryMsgHeader header, ValueSet payload)
var length = payload.Get<uint>("BlobSize");

var fileProvider = _files?[(int)contentId] ?? throw new NullReferenceException("Could not access files to transfer");
var blob = fileProvider.ReadBlob(start, length);
Channel.SendBinaryMessage(writer =>
{
FetchDataResponse.Write(writer, contentId, start, (int)length, out var blob);
Debug.Assert(blob.Length == length);

fileProvider.ReadBlob(start, blob);
}, header.MessageId);

_fileProgress?.Report(new()
{
TransferedBytes = Interlocked.Add(ref _bytesSent, length),
TotalBytes = _bytesToSend,
TotalFiles = (uint)_files.Count
});

ValueSet response = new();
response.Add("ControlMessage", (uint)NearShareControlMsgType.FetchDataResponse);
response.Add("ContentId", contentId);
response.Add("BlobPosition", start);
response.Add("DataBlob", blob.ToArray().ToList()); // ToDo: Remove allocation
SendValueSet(response, header.MessageId);
}
}
}
73 changes: 45 additions & 28 deletions lib/ShortDev.Microsoft.ConnectedDevices/Encryption/CdpCryptor.cs
Original file line number Diff line number Diff line change
@@ -1,5 +1,8 @@
using ShortDev.Microsoft.ConnectedDevices.Exceptions;
using ShortDev.Microsoft.ConnectedDevices.Messages;
using ShortDev.Microsoft.ConnectedDevices.Transports;
using System.Buffers;
using System.Buffers.Binary;
using System.Diagnostics;
using System.Security.Cryptography;

Expand Down Expand Up @@ -27,16 +30,13 @@ void GenerateIV(CommonHeader header, Span<byte> destination)
{
Debug.Assert(destination.Length == Constants.IVSize);

var aes = _ivAes;
Span<byte> raw = stackalloc byte[Constants.IVSize];
BinaryPrimitives.WriteUInt64BigEndian(raw[..8], header.SessionId);
BinaryPrimitives.WriteUInt32BigEndian(raw[8..12], header.SequenceNumber);
BinaryPrimitives.WriteUInt16BigEndian(raw[12..14], header.FragmentIndex);
BinaryPrimitives.WriteUInt16BigEndian(raw[14..16], header.FragmentCount);

EndianWriter writer = new(Endianness.BigEndian, Constants.IVSize);

writer.Write(header.SessionId);
writer.Write(header.SequenceNumber);
writer.Write(header.FragmentIndex);
writer.Write(header.FragmentCount);

int bytesWritten = aes.EncryptCbc(writer.Buffer.AsSpan(), _ivData, destination, PaddingMode.None);
int bytesWritten = _ivAes.EncryptCbc(raw, _ivData, destination, PaddingMode.None);
Debug.Assert(bytesWritten == destination.Length);
}

Expand Down Expand Up @@ -86,36 +86,53 @@ void VerifyHMac(CommonHeader header, ReadOnlySpan<byte> payload, ReadOnlySpan<by
throw new CdpSecurityException("Invalid hmac!");
}

public void EncryptMessage(EndianWriter writer, CommonHeader header, ReadOnlySpan<byte> payloadBuffer)
public void EncryptMessage(IFragmentSender sender, CommonHeader header, ReadOnlySpan<byte> payloadBuffer)
{
EndianWriter msgWriter = new(Endianness.BigEndian);
// Prepend payload with length
ReadOnlySpan<byte> finalPayload;
{
EndianWriter payloadWriter = new(Endianness.BigEndian);
payloadWriter.Write((uint)payloadBuffer.Length);
payloadWriter.Write(payloadBuffer);

Span<byte> iv = stackalloc byte[Constants.IVSize];
GenerateIV(header, iv);
finalPayload = payloadWriter.Buffer.AsSpan();
}

// Encrypt
var msgWriter = Encrypt(header, finalPayload);

EndianWriter payloadWriter = new(Endianness.BigEndian);
payloadWriter.Write((uint)payloadBuffer.Length);
payloadWriter.Write(payloadBuffer);
// HMAC
{
var msgBuffer = msgWriter.Buffer.AsWriteableSpan();
Span<byte> hmac = stackalloc byte[Constants.HMacSize];
ComputeHmac(msgBuffer, hmac);
CommonHeader.ModifyMessageLength(msgBuffer, +Constants.HMacSize);
msgWriter.Write(hmac);
}

var buffer = payloadWriter.Buffer.AsSpan();
sender.SendFragment(msgWriter.Buffer.AsSpan());
}

EndianWriter Encrypt(CommonHeader header, ReadOnlySpan<byte> buffer)
{
// If payload size is an exact multiple of block length (16 bytes) no padding is applied
PaddingMode paddingMode = buffer.Length % 16 == 0 ? PaddingMode.None : PaddingMode.PKCS7;
var encryptedPayload = _aes.EncryptCbc(buffer, iv, paddingMode);
var encryptedPayloadLength = _aes.GetCiphertextLengthCbc(buffer.Length, paddingMode);

// Write header
EndianWriter writer = new(Endianness.BigEndian);
header.Flags |= MessageFlags.SessionEncrypted | MessageFlags.HasHMAC;
header.SetPayloadLength(encryptedPayload.Length);
header.Write(msgWriter);

msgWriter.Write(encryptedPayload);
header.SetPayloadLength(encryptedPayloadLength);
header.Write(writer);

var msgBuffer = msgWriter.Buffer.AsWriteableSpan();
Span<byte> iv = stackalloc byte[Constants.IVSize];
GenerateIV(header, iv);

Span<byte> hmac = stackalloc byte[Constants.HMacSize];
ComputeHmac(msgBuffer, hmac);
CommonHeader.ModifyMessageLength(msgBuffer, +Constants.HMacSize);
// Encrypt and write to msgWriter
_aes.EncryptCbc(buffer, iv, writer.Buffer.GetSpan(encryptedPayloadLength), paddingMode);
writer.Buffer.Advance(encryptedPayloadLength);

writer.Write(msgBuffer);
writer.Write(hmac);
return writer;
}

public void Read(ref EndianReader reader, CommonHeader header)
Expand Down
2 changes: 2 additions & 0 deletions lib/ShortDev.Microsoft.ConnectedDevices/Extensions.cs
Original file line number Diff line number Diff line change
Expand Up @@ -69,6 +69,8 @@ public readonly struct ArrayPoolToken<T>(ArrayPool<T> pool, int capacity) : IDis
private readonly int _capacity = capacity;
private readonly T[] _array = pool.Rent(capacity);

public T[] ArrayUnsafe => _array;

public Memory<T> Memory => _array.AsMemory()[0.._capacity];
public Span<T> Span => Memory.Span;

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -4,11 +4,11 @@
/// cdp.dll!cdp::BinaryFragmenter::GetMessageFragments <br/>
/// <see cref="AdditionalHeaderType.UserMessageRequestId"/>
/// </summary>
public sealed class BinaryMsgHeader : ICdpHeader<BinaryMsgHeader>
public readonly struct BinaryMsgHeader() : ICdpHeader<BinaryMsgHeader>
{
public uint FragmentCount { get; set; } = 1;
public uint FragmentIndex { get; set; } = 0;
public required uint MessageId { get; set; }
public uint FragmentCount { get; init; } = 1;
public uint FragmentIndex { get; init; } = 0;
public required uint MessageId { get; init; }

public static BinaryMsgHeader Parse(ref EndianReader reader)
=> new()
Expand Down
Loading
Loading