Improve ZInputStream/ZOutputStream constructors

2 files changed, 60 insertions, 15 deletions

diff --git a/crypto/src/util/zlib/ZInputStream.cs b/crypto/src/util/zlib/ZInputStream.cs
index 4b7351555..434fe95c8 100644
--- a/crypto/src/util/zlib/ZInputStream.cs
+++ b/crypto/src/util/zlib/ZInputStream.cs
@@ -42,9 +42,16 @@ namespace Org.BouncyCastle.Utilities.Zlib
 	public class ZInputStream
 		: Stream
 	{
-		private const int BufferSize = 512;
+        private static ZStream GetDefaultZStream(bool nowrap)
+        {
+            ZStream z = new ZStream();
+            z.inflateInit(nowrap);
+            return z;
+        }
+
+        private const int BufferSize = 512;
 
-		protected ZStream z = new ZStream();
+		protected ZStream z;
 		protected int flushLevel = JZlib.Z_NO_FLUSH;
 		// TODO Allow custom buf
 		protected byte[] buf = new byte[BufferSize];
@@ -62,24 +69,46 @@ namespace Org.BouncyCastle.Utilities.Zlib
 		}
 
 		public ZInputStream(Stream input, bool nowrap)
+            : this(input, GetDefaultZStream(nowrap))
+		{
+		}
+
+        public ZInputStream(Stream input, ZStream z)
+			: base()
 		{
 			Debug.Assert(input.CanRead);
 
-			this.input = input;
-			this.z.inflateInit(nowrap);
-			this.compress = false;
-			this.z.next_in = buf;
+            if (z == null)
+            {
+                z = new ZStream();
+            }
+
+            if (z.istate == null && z.dstate == null)
+            {
+                z.inflateInit();
+            }
+
+            this.input = input;
+            this.compress = (z.istate == null);
+            this.z = z;
+            this.z.next_in = buf;
 			this.z.next_in_index = 0;
 			this.z.avail_in = 0;
 		}
 
-		public ZInputStream(Stream input, int level)
+        public ZInputStream(Stream input, int level)
+            : this(input, level, false)
+		{
+        }
+
+        public ZInputStream(Stream input, int level, bool nowrap)
 		{
 			Debug.Assert(input.CanRead);
 			
 			this.input = input;
-			this.z.deflateInit(level);
-			this.compress = true;
+            this.compress = true;
+            this.z = new ZStream();
+			this.z.deflateInit(level, nowrap);
 			this.z.next_in = buf;
 			this.z.next_in_index = 0;
 			this.z.avail_in = 0;
diff --git a/crypto/src/util/zlib/ZOutputStream.cs b/crypto/src/util/zlib/ZOutputStream.cs
index d9f005f69..1633b2d8f 100644
--- a/crypto/src/util/zlib/ZOutputStream.cs
+++ b/crypto/src/util/zlib/ZOutputStream.cs
@@ -42,7 +42,14 @@ namespace Org.BouncyCastle.Utilities.Zlib
 	public class ZOutputStream
 		: Stream
 	{
-		private const int BufferSize = 512;
+        private static ZStream GetDefaultZStream(bool nowrap)
+        {
+            ZStream z = new ZStream();
+            z.inflateInit(nowrap);
+            return z;
+        }
+
+        private const int BufferSize = 512;
 
 		protected ZStream z;
 		protected int flushLevel = JZlib.Z_NO_FLUSH;
@@ -55,9 +62,14 @@ namespace Org.BouncyCastle.Utilities.Zlib
 		protected bool closed;
 
         public ZOutputStream(Stream output)
-			: this(output, null)
-        {
-        }
+            : this(output, false)
+		{
+		}
+
+        public ZOutputStream(Stream output, bool nowrap)
+            : this(output, GetDefaultZStream(nowrap))
+		{
+		}
 
         public ZOutputStream(Stream output, ZStream z)
 			: base()
@@ -67,12 +79,16 @@ namespace Org.BouncyCastle.Utilities.Zlib
             if (z == null)
             {
                 z = new ZStream();
+            }
+
+            if (z.istate == null && z.dstate == null)
+            {
                 z.inflateInit();
             }
 
             this.output = output;
+            this.compress = (z.istate == null);
             this.z = z;
-			this.compress = false;
 		}
 
         public ZOutputStream(Stream output, int level)
@@ -86,9 +102,9 @@ namespace Org.BouncyCastle.Utilities.Zlib
 			Debug.Assert(output.CanWrite);
 
 			this.output = output;
+            this.compress = true;
             this.z = new ZStream();
 			this.z.deflateInit(level, nowrap);
-			this.compress = true;
 		}
 
 		public sealed override bool CanRead { get { return false; } }