From 54c4a050622727bfefc75fea63245e5db45aa5f9 Mon Sep 17 00:00:00 2001
From: RoosterDragon <RoosterDragon@outlook.com>
Date: Sun, 18 Oct 2020 15:28:33 +0100
Subject: [PATCH] Classes derived from Stream override ReadByte.

The Stream.ReadByte method is implemented by allocating a 1 byte buffer and calling into Read(byte[], int, int). Override ReadByte in derived classes to avoid needing to allocate this small temp buffer.

Also, fix some bugs in the stream implementations. Remove Write capability from MergedStream that didn't make sense. Add guards into SegmentStream to ensure reads and writes belonged to the segment - otherwise a reader or writer could access regions of the base stream that were outside the intended segment.
---
 OpenRA.Game/Primitives/MergedStream.cs        | 31 ++++++++++++++-----
 .../Primitives/ReadOnlyAdapterStream.cs       | 19 +++++++++++-
 OpenRA.Game/Primitives/SegmentStream.cs       | 31 +++++++++++++++++--
 3 files changed, 71 insertions(+), 10 deletions(-)

diff --git a/OpenRA.Game/Primitives/MergedStream.cs b/OpenRA.Game/Primitives/MergedStream.cs
index af0e7ac7a0..e4dd99d6ed 100644
--- a/OpenRA.Game/Primitives/MergedStream.cs
+++ b/OpenRA.Game/Primitives/MergedStream.cs
@@ -9,6 +9,7 @@
  */
 #endregion
 
+using System;
 using System.IO;
 
 namespace OpenRA.Primitives
@@ -18,7 +19,7 @@ namespace OpenRA.Primitives
 		public Stream Stream1 { get; set; }
 		public Stream Stream2 { get; set; }
 
-		long VirtualLength { get; set; }
+		long VirtualLength { get; }
 		long position;
 
 		public MergedStream(Stream stream1, Stream stream2)
@@ -61,7 +62,21 @@ namespace OpenRA.Primitives
 
 		public override void SetLength(long value)
 		{
-			VirtualLength = value;
+			throw new NotSupportedException();
+		}
+
+		public override int ReadByte()
+		{
+			int value;
+
+			if (position >= Stream1.Length)
+				value = Stream2.ReadByte();
+			else
+				value = Stream1.ReadByte();
+
+			position++;
+
+			return value;
 		}
 
 		public override int Read(byte[] buffer, int offset, int count)
@@ -83,12 +98,14 @@ namespace OpenRA.Primitives
 			return bytesRead;
 		}
 
+		public override void WriteByte(byte value)
+		{
+			throw new NotSupportedException();
+		}
+
 		public override void Write(byte[] buffer, int offset, int count)
 		{
-			if (position >= Stream1.Length)
-				Stream2.Write(buffer, offset - (int)Stream1.Length, count);
-			else
-				Stream1.Write(buffer, offset, count);
+			throw new NotSupportedException();
 		}
 
 		public override bool CanRead
@@ -103,7 +120,7 @@ namespace OpenRA.Primitives
 
 		public override bool CanWrite
 		{
-			get { return Stream1.CanWrite && Stream2.CanWrite; }
+			get { return false; }
 		}
 
 		public override long Length
diff --git a/OpenRA.Game/Primitives/ReadOnlyAdapterStream.cs b/OpenRA.Game/Primitives/ReadOnlyAdapterStream.cs
index 8b53435b8a..1001332b0c 100644
--- a/OpenRA.Game/Primitives/ReadOnlyAdapterStream.cs
+++ b/OpenRA.Game/Primitives/ReadOnlyAdapterStream.cs
@@ -48,9 +48,25 @@ namespace OpenRA.Primitives
 
 		public sealed override long Seek(long offset, SeekOrigin origin) { throw new NotSupportedException(); }
 		public sealed override void SetLength(long value) { throw new NotSupportedException(); }
+		public sealed override void WriteByte(byte value) { throw new NotSupportedException(); }
 		public sealed override void Write(byte[] buffer, int offset, int count) { throw new NotSupportedException(); }
 		public sealed override void Flush() { throw new NotSupportedException(); }
 
+		public sealed override int ReadByte()
+		{
+			if (data.Count > 0)
+				return data.Dequeue();
+
+			while (!baseStreamEmpty)
+			{
+				baseStreamEmpty = BufferData(baseStream, data);
+				if (data.Count > 0)
+					return data.Dequeue();
+			}
+
+			return -1;
+		}
+
 		public sealed override int Read(byte[] buffer, int offset, int count)
 		{
 			var copied = 0;
@@ -66,7 +82,8 @@ namespace OpenRA.Primitives
 		}
 
 		/// <summary>
-		/// Reads data into a buffer, which will be used to satisfy <see cref="Read(byte[], int, int)"/> calls.
+		/// Reads data into a buffer, which will be used to satisfy <see cref="ReadByte()"/> and
+		/// <see cref="Read(byte[], int, int)"/> calls.
 		/// </summary>
 		/// <param name="baseStream">The source stream from which bytes should be read.</param>
 		/// <param name="data">The queue where bytes should be enqueued. Do not dequeue from this buffer.</param>
diff --git a/OpenRA.Game/Primitives/SegmentStream.cs b/OpenRA.Game/Primitives/SegmentStream.cs
index 66080af01a..7b24bdba7b 100644
--- a/OpenRA.Game/Primitives/SegmentStream.cs
+++ b/OpenRA.Game/Primitives/SegmentStream.cs
@@ -59,8 +59,35 @@ namespace OpenRA.Primitives
 			set { BaseStream.Position = BaseOffset + value; }
 		}
 
-		public override int Read(byte[] buffer, int offset, int count) { return BaseStream.Read(buffer, offset, count); }
-		public override void Write(byte[] buffer, int offset, int count) { BaseStream.Write(buffer, offset, count); }
+		public override int ReadByte()
+		{
+			if (Position < Length)
+				return BaseStream.ReadByte();
+			return -1;
+		}
+
+		public override int Read(byte[] buffer, int offset, int count)
+		{
+			var remaining = Length - Position;
+			if (remaining <= 0)
+				return 0;
+			return BaseStream.Read(buffer, offset, (int)Math.Min(remaining, count));
+		}
+
+		public override void WriteByte(byte value)
+		{
+			if (Position < Length)
+				BaseStream.WriteByte(value);
+			throw new IOException("Attempted to write past the end of the stream.");
+		}
+
+		public override void Write(byte[] buffer, int offset, int count)
+		{
+			if (Position + count >= Length)
+				throw new IOException("Attempted to write past the end of the stream.");
+			BaseStream.Write(buffer, offset, count);
+		}
+
 		public override void Flush() { BaseStream.Flush(); }
 		public override long Seek(long offset, SeekOrigin origin)
 		{