Improve transfer speed (sending) (#178)

nearby-sharing · Jul 30, 2024 · eb8135d · eb8135d
1 parent eafdf2d
commit eb8135d
Show file tree

Hide file tree

Showing 14 changed files with 751 additions and 58 deletions.
diff --git a/Directory.Build.targets b/Directory.Build.targets
@@ -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>
diff --git a/lib/ShortDev.Microsoft.ConnectedDevices.NearShare/CdpFileProvider.cs b/lib/ShortDev.Microsoft.ConnectedDevices.NearShare/CdpFileProvider.cs
@@ -1,5 +1,4 @@
-using System.IO;
-using System.Text;
+using System.Text;
 
 namespace ShortDev.Microsoft.ConnectedDevices.NearShare;
 
@@ -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()

diff --git a/lib/ShortDev.Microsoft.ConnectedDevices.NearShare/Messages/FetchDataResponse.cs b/lib/ShortDev.Microsoft.ConnectedDevices.NearShare/Messages/FetchDataResponse.cs
@@ -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);
+    }
+}
diff --git a/lib/ShortDev.Microsoft.ConnectedDevices.NearShare/NearShareSender.cs b/lib/ShortDev.Microsoft.ConnectedDevices.NearShare/NearShareSender.cs
@@ -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;
 
@@ -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);
         }
     }
 }
diff --git a/lib/ShortDev.Microsoft.ConnectedDevices/Encryption/CdpCryptor.cs b/lib/ShortDev.Microsoft.ConnectedDevices/Encryption/CdpCryptor.cs
@@ -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;
 
@@ -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);
     }
 
@@ -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)

diff --git a/lib/ShortDev.Microsoft.ConnectedDevices/Extensions.cs b/lib/ShortDev.Microsoft.ConnectedDevices/Extensions.cs
@@ -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;
 

diff --git a/lib/ShortDev.Microsoft.ConnectedDevices/Messages/Session/BinaryMsgHeader.cs b/lib/ShortDev.Microsoft.ConnectedDevices/Messages/Session/BinaryMsgHeader.cs
@@ -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()