Fix attachment uploading

Author: Trolldemorted

libsignal-service-dotnet/crypto/AttachmentCipherOutputStream.cs

@@ -37,34 +37,45 @@ internal static long GetCiphertextLength(long plaintextLength)
 
         public override void Write(byte[] buffer, int offset, int count)
         {
-            Cipher.Write(buffer, offset, count);
+            WriteToCipherStream(buffer, offset, count);
             byte[] cipherBuffer = new byte[Aes.BlockSize];
             int read = TmpStream.Read(cipherBuffer, 0, cipherBuffer.Length);
             while (read > 0)
             {
-                Mac.AppendData(buffer);
+                Mac.AppendData(cipherBuffer, 0, read);
                 base.Write(cipherBuffer, 0, read);
+                read = TmpStream.Read(cipherBuffer, 0, cipherBuffer.Length);
             }
         }
 
         public override void Flush()
         {
-            Cipher.FlushFinalBlock();
+            FlushFinalBlock();
             byte[] cipherBuffer = new byte[Aes.BlockSize];
             int read = TmpStream.Read(cipherBuffer, 0, cipherBuffer.Length);
             while (read > 0)
             {
                 Mac.AppendData(cipherBuffer, 0, read);
                 base.Write(cipherBuffer, 0, read);
+                read = TmpStream.Read(cipherBuffer, 0, cipherBuffer.Length);
             }
             byte[] auth = Mac.GetHashAndReset();
             base.Write(auth, 0, auth.Length);
             base.Flush();
         }
 
-        public static void ReadIntoBuffer(Stream s)
+        private void WriteToCipherStream(byte[] buffer, int offset, int count)
         {
+            var oldPos = TmpStream.Position;
+            Cipher.Write(buffer, offset, count);
+            TmpStream.Position = oldPos;
+        }
 
+        private void FlushFinalBlock()
+        {
+            var oldPos = TmpStream.Position;
+            Cipher.FlushFinalBlock();
+            TmpStream.Position = oldPos;
         }
     }
 }

libsignal-service-dotnet/crypto/DigestingOutputStream.cs

@@ -11,7 +11,7 @@ public class DigestingOutputStream : Stream
     {
         private IncrementalHash RunningDigest = IncrementalHash.CreateHash(HashAlgorithmName.SHA256);
         private Stream OutputStream { get; }
-        public override bool CanRead => throw new NotImplementedException();
+        public override bool CanRead => false;
         public override bool CanSeek => false;
         public override bool CanWrite => true;
         public override long Length => throw new NotImplementedException();

libsignal-service-dotnet/crypto/PaddingInputStream.cs

@@ -11,9 +11,9 @@ internal class PaddingInputStream : Stream
         private readonly Stream InputStream;
         private long PaddingRemaining;
 
-        public override bool CanRead => throw new NotImplementedException();
-        public override bool CanSeek => throw new NotImplementedException();
-        public override bool CanWrite => throw new NotImplementedException();
+        public override bool CanRead => true;
+        public override bool CanSeek => false;
+        public override bool CanWrite => false;
         public override long Length { get => InputStream.Length + Util.ToIntExact(PaddingRemaining); }
         public override long Position { get => throw new NotImplementedException(); set => throw new NotImplementedException(); }
 

libsignal-service-dotnet/messages/SignalServiceAttachment.cs

@@ -1,81 +1,86 @@
-using System;
-using System.IO;
-
-namespace libsignalservice.messages
+using System;
+using System.IO;
+
+namespace libsignalservice.messages
 {
 #pragma warning disable CS1591 // Missing XML comment for publicly visible type or member
     public abstract class SignalServiceAttachment
-    {
-        public String ContentType { get; }
-
-        internal SignalServiceAttachment(String contentType)
-        {
-            this.ContentType = contentType;
-        }
-
-        public abstract bool IsStream();
-
-        public abstract bool IsPointer();
-
-        public SignalServiceAttachmentStream AsStream()
-        {
-            return (SignalServiceAttachmentStream)this;
-        }
-
-        public SignalServiceAttachmentPointer AsPointer()
-        {
-            return (SignalServiceAttachmentPointer)this;
-        }
-
-        public class Builder
-        {
-            private Stream InputStream;
-            private string ContentType;
-            private string FileName;
-            private long Length;
-            private IProgressListener Listener;
+    {
+        public String ContentType { get; }
+
+        internal SignalServiceAttachment(String contentType)
+        {
+            this.ContentType = contentType;
+        }
+
+        public abstract bool IsStream();
+
+        public abstract bool IsPointer();
+
+        public SignalServiceAttachmentStream AsStream()
+        {
+            return (SignalServiceAttachmentStream)this;
+        }
+
+        public SignalServiceAttachmentPointer AsPointer()
+        {
+            return (SignalServiceAttachmentPointer)this;
+        }
+
+        public static Builder NewStreamBuilder()
+        {
+            return new Builder();
+        }
+
+        public class Builder
+        {
+            private Stream InputStream;
+            private string ContentType;
+            private string FileName;
+            private long Length;
+            private IProgressListener Listener;
             private bool VoiceNote;
             private int Width;
-            private int Height;
-
-            Builder()
-            {
-            }
-
-            public Builder WithStream(Stream inputStream)
-            {
-                InputStream = inputStream;
-                return this;
-            }
-
-            public Builder WithContentType(string contentType)
-            {
-                ContentType = contentType;
-                return this;
-            }
-
-            public Builder WithLength(long length)
-            {
-                Length = length;
-                return this;
-            }
-
-            public Builder WithFileName(string fileName)
-            {
-                FileName = fileName;
-                return this;
-            }
-
-            public Builder WithListener(IProgressListener listener)
-            {
-                Listener = listener;
-                return this;
-            }
-
-            public Builder WithVoiceNote(bool voiceNote)
-            {
-                VoiceNote = voiceNote;
-                return this;
+            private int Height;
+
+            internal Builder()
+            {
+            }
+
+            public Builder WithStream(Stream inputStream)
+            {
+                InputStream = inputStream;
+                return this;
+            }
+
+            public Builder WithContentType(string contentType)
+            {
+                ContentType = contentType;
+                return this;
+            }
+
+            public Builder WithLength(long length)
+            {
+                Length = length;
+                return this;
+            }
+
+            public Builder WithFileName(string fileName)
+            {
+                FileName = fileName;
+                return this;
+            }
+
+            public Builder WithListener(IProgressListener listener)
+            {
+                Listener = listener;
+                return this;
+            }
+
+            public Builder WithVoiceNote(bool voiceNote)
+            {
+                VoiceNote = voiceNote;
+                return this;
             }
 
             public Builder WithWidth(int width)
@@ -88,36 +93,36 @@ public Builder WithHeight(int height)
             {
                 Height = height;
                 return this;
-            }
-
-            public SignalServiceAttachmentStream Build()
-            {
-                if (InputStream == null)
-                {
-                    throw new ArgumentException("Must specify stream!");
-                }
-                if (ContentType == null)
-                {
-                    throw new ArgumentException("No content type specified!");
-                }
-                if (Length == 0)
-                {
-                    throw new ArgumentException("No length specified!");
-                }
-
-                return new SignalServiceAttachmentStream(InputStream, ContentType, (uint)Length, FileName, VoiceNote, null, Width, Height, Listener);
-            }
-        }
-
-        public interface IProgressListener
-        {
-            /// <summary>
-            /// Called on a progress change event.
-            /// </summary>
-            /// <param name="total">The total amount of transmit/receive in bytes.</param>
-            /// <param name="progress">The amount that has been transmitted/received in bytes thus far</param>
-            void OnAttachmentProgress(long total, long progress);
-        }
+            }
+
+            public SignalServiceAttachmentStream Build()
+            {
+                if (InputStream == null)
+                {
+                    throw new ArgumentException("Must specify stream!");
+                }
+                if (ContentType == null)
+                {
+                    throw new ArgumentException("No content type specified!");
+                }
+                if (Length == 0)
+                {
+                    throw new ArgumentException("No length specified!");
+                }
+
+                return new SignalServiceAttachmentStream(InputStream, ContentType, (uint)Length, FileName, VoiceNote, null, Width, Height, Listener);
+            }
+        }
+
+        public interface IProgressListener
+        {
+            /// <summary>
+            /// Called on a progress change event.
+            /// </summary>
+            /// <param name="total">The total amount of transmit/receive in bytes.</param>
+            /// <param name="progress">The amount that has been transmitted/received in bytes thus far</param>
+            void OnAttachmentProgress(long total, long progress);
+        }
     }
-#pragma warning restore CS1591 // Missing XML comment for publicly visible type or member
-}
+#pragma warning restore CS1591 // Missing XML comment for publicly visible type or member
+}

libsignal-service-dotnet/push/PushServiceSocket.cs

@@ -10,6 +10,7 @@
 using libsignalservice.profiles;
 using libsignalservice.push.exceptions;
 using libsignalservice.util;
+using Microsoft.Extensions.Logging;
 using Newtonsoft.Json;
 
 using System;
@@ -58,6 +59,7 @@ internal class PushServiceSocket
 
         private static readonly string PROFILE_PATH = "/v1/profile/%s";
 
+        private readonly ILogger Logger = LibsignalLogging.CreateLogger<PushServiceSocket>();
         private readonly SignalServiceConfiguration SignalConnectionInformation;
         private readonly ICredentialsProvider CredentialsProvider;
         private readonly string UserAgent;
@@ -337,7 +339,7 @@ public async Task<bool> SetCurrentSignedPreKey(CancellationToken token, SignedPr
         {
             var (id, location) = await RetrieveAttachmentUploadUrl(token);
 
-            byte[] digest = await UploadAttachment("PUT", location, attachment.Data,
+            byte[] digest = await UploadAttachment(token, "PUT", location, attachment.Data,
                 attachment.DataSize, attachment.OutputFactory, attachment.Listener);
 
             return (id, digest);
@@ -364,7 +366,7 @@ public async Task<bool> SetCurrentSignedPreKey(CancellationToken token, SignedPr
         public async Task RetrieveAttachment(CancellationToken token, string relay, ulong attachmentId, Stream tmpDestination, int maxSizeBytes)
         {
             string attachmentUrlLocation = await RetrieveAttachmentDownloadUrl(token, relay, attachmentId);
-            await DownloadAttachment(attachmentUrlLocation, tmpDestination);
+            await DownloadAttachment(token, attachmentUrlLocation, tmpDestination);
         }
 
         /// <summary>
@@ -490,7 +492,7 @@ public void CancelInFlightRequests()
             throw new NotImplementedException();
         }
 
-        private async Task DownloadAttachment(string url, Stream localDestination)
+        private async Task DownloadAttachment(CancellationToken token, string url, Stream localDestination)
         {
             try
             {
@@ -500,7 +502,7 @@ private async Task DownloadAttachment(string url, Stream localDestination)
                 HttpRequestMessage req = new HttpRequestMessage(HttpMethod.Get, url);
                 req.Content = new StringContent("");
                 req.Content.Headers.ContentType = new System.Net.Http.Headers.MediaTypeHeaderValue("application/octet-stream");
-                using (var resp = await connection.SendAsync(req))
+                using (var resp = await connection.SendAsync(req, token))
                 {
                     Stream input = await resp.Content.ReadAsStreamAsync();
                     byte[] buffer = new byte[4096];
@@ -526,19 +528,26 @@ private async Task DownloadAttachment(string url, Stream localDestination)
             }
         }
 
-        private async Task<byte[]> UploadAttachment(string method, string url, Stream data, long dataSize,
+        private async Task<byte[]> UploadAttachment(CancellationToken token, string method, string url, Stream data, long dataSize,
             OutputStreamFactory outputStreamFactory, IProgressListener listener)
         {
+            // buffer payload in memory...
             MemoryStream tmpStream = new MemoryStream();
             DigestingOutputStream outputStream = outputStreamFactory.CreateFor(tmpStream);
             StreamContent streamContent = new StreamContent(tmpStream);
-            var request = new HttpRequestMessage(HttpMethod.Put, url);
-            request.Content = streamContent;
+            data.CopyTo(outputStream);
+            outputStream.Flush();
+            tmpStream.Position = 0;
+
+            // ... and upload it!
+            var request = new HttpRequestMessage(HttpMethod.Put, url)
+            {
+                Content = new StreamContent(tmpStream)
+            };
             request.Content.Headers.ContentType = new MediaTypeHeaderValue("application/octet-stream");
             request.Headers.ConnectionClose = true;
-
             HttpClient client = new HttpClient();
-            HttpResponseMessage response = await client.SendAsync(request);
+            HttpResponseMessage response = await client.SendAsync(request, token);
             if (response.StatusCode != HttpStatusCode.OK)
             {
                 throw new IOException($"Bad response: {response.StatusCode} {await response.Content.ReadAsStringAsync()}");