Index: src/OFZIPArchive.h
==================================================================
--- src/OFZIPArchive.h
+++ src/OFZIPArchive.h
@@ -31,17 +31,24 @@
  */
 OF_SUBCLASSING_RESTRICTED
 @interface OFZIPArchive: OFObject
 {
 	OFStream *_stream;
+#ifdef OF_ZIP_ARCHIVE_M
+@public
+#endif
 	int64_t _offset;
+@protected
 	uint_least8_t _mode;
 	uint32_t _diskNumber, _centralDirectoryDisk;
 	uint64_t _centralDirectoryEntriesInDisk, _centralDirectoryEntries;
 	uint64_t _centralDirectorySize;
 	int64_t _centralDirectoryOffset;
 	OFString *_Nullable _archiveComment;
+#ifdef OF_ZIP_ARCHIVE_M
+@public
+#endif
 	OFMutableArray OF_GENERIC(OFZIPArchiveEntry *) *_entries;
 	OFMutableDictionary OF_GENERIC(OFString *, OFZIPArchiveEntry *)
 	    *_pathToEntryMap;
 	OFStream *_Nullable _lastReturnedStream;
 }

Index: src/OFZIPArchive.m
==================================================================
--- src/OFZIPArchive.m
+++ src/OFZIPArchive.m
@@ -11,10 +11,12 @@
  * Public License, either version 2 or 3, which can be found in the file
  * LICENSE.GPLv2 or LICENSE.GPLv3 respectively included in the packaging of this
  * file.
  */
 
+#define OF_ZIP_ARCHIVE_M
+
 #include "config.h"
 
 #include <errno.h>
 
 #import "OFZIPArchive.h"
@@ -58,11 +60,10 @@
 
 OF_DIRECT_MEMBERS
 @interface OFZIPArchive ()
 - (void)of_readZIPInfo;
 - (void)of_readEntries;
-- (void)of_closeLastReturnedStream;
 - (void)of_writeCentralDirectory;
 @end
 
 OF_DIRECT_MEMBERS
 @interface OFZIPArchiveLocalFileHeader: OFObject
@@ -81,33 +82,37 @@
 @end
 
 OF_DIRECT_MEMBERS
 @interface OFZIPArchiveFileReadStream: OFStream
 {
+	OFZIPArchive *_archive;
 	OFStream *_stream, *_decompressedStream;
 	OFZIPArchiveEntry *_entry;
 	unsigned long long _toRead;
 	uint32_t _CRC32;
 	bool _atEndOfStream;
 }
 
-- (instancetype)of_initWithStream: (OFStream *)stream
-			    entry: (OFZIPArchiveEntry *)entry;
+- (instancetype)of_initWithArchive: (OFZIPArchive *)archive
+			    stream: (OFStream *)stream
+			     entry: (OFZIPArchiveEntry *)entry;
 @end
 
 OF_DIRECT_MEMBERS
 @interface OFZIPArchiveFileWriteStream: OFStream
 {
+	OFZIPArchive *_archive;
 	OFStream *_stream;
 	uint32_t _CRC32;
 @public
 	unsigned long long _bytesWritten;
 	OFMutableZIPArchiveEntry *_entry;
 }
 
-- (instancetype)initWithStream: (OFStream *)stream
-			 entry: (OFMutableZIPArchiveEntry *)entry;
+- (instancetype)of_initWithArchive: (OFZIPArchive *)archive
+			    stream: (OFStream *)stream
+			     entry: (OFMutableZIPArchiveEntry *)entry;
 @end
 
 uint32_t
 OFZIPArchiveReadField32(const uint8_t **data, uint16_t *size)
 {
@@ -255,11 +260,10 @@
 
 	[_stream release];
 	[_archiveComment release];
 	[_entries release];
 	[_pathToEntryMap release];
-	[_lastReturnedStream release];
 
 	[super dealloc];
 }
 
 - (void)of_readZIPInfo
@@ -401,40 +405,10 @@
 	[old release];
 
 	objc_autoreleasePoolPop(pool);
 }
 
-- (void)of_closeLastReturnedStream
-{
-	@try {
-		[_lastReturnedStream close];
-	} @catch (OFNotOpenException *e) {
-		/* Might have already been closed by the user - that's fine. */
-	}
-
-	if ((_mode == modeWrite || _mode == modeAppend) &&
-	    [_lastReturnedStream isKindOfClass:
-	    [OFZIPArchiveFileWriteStream class]]) {
-		OFZIPArchiveFileWriteStream *stream =
-		    (OFZIPArchiveFileWriteStream *)_lastReturnedStream;
-
-		if (ULLONG_MAX - _offset < stream->_bytesWritten)
-			@throw [OFOutOfRangeException exception];
-
-		_offset += stream->_bytesWritten;
-
-		if (stream->_entry != nil) {
-			[_entries addObject: stream->_entry];
-			[_pathToEntryMap setObject: stream->_entry
-					    forKey: [stream->_entry fileName]];
-		}
-	}
-
-	[_lastReturnedStream release];
-	_lastReturnedStream = nil;
-}
-
 - (OFStream *)streamForReadingFile: (OFString *)path
 {
 	void *pool = objc_autoreleasePoolPush();
 	OFZIPArchiveEntry *entry;
 	OFZIPArchiveLocalFileHeader *localFileHeader;
@@ -449,11 +423,16 @@
 	if ((entry = [_pathToEntryMap objectForKey: path]) == nil)
 		@throw [OFOpenItemFailedException exceptionWithPath: path
 							       mode: @"r"
 							      errNo: ENOENT];
 
-	[self of_closeLastReturnedStream];
+	@try {
+		[_lastReturnedStream close];
+	} @catch (OFNotOpenException *e) {
+		/* Might have already been closed by the user - that's fine. */
+	}
+	_lastReturnedStream = nil;
 
 	offset64 = entry.of_localFileHeaderOffset;
 	if (offset64 < 0 || (OFStreamOffset)offset64 != offset64)
 		@throw [OFOutOfRangeException exception];
 
@@ -472,17 +451,18 @@
 
 		@throw [OFUnsupportedVersionException
 		    exceptionWithVersion: version];
 	}
 
-	_lastReturnedStream = [[OFZIPArchiveFileReadStream alloc]
-	     of_initWithStream: _stream
-			 entry: entry];
-
 	objc_autoreleasePoolPop(pool);
 
-	return [[_lastReturnedStream retain] autorelease];
+	_lastReturnedStream = [[[OFZIPArchiveFileReadStream alloc]
+	    of_initWithArchive: self
+			stream: _stream
+			 entry: entry] autorelease];
+
+	return _lastReturnedStream;
 }
 
 - (OFStream *)streamForWritingEntry: (OFZIPArchiveEntry *)entry_
 {
 	/* TODO: Avoid data descriptor when _stream is an OFSeekableStream */
@@ -510,11 +490,16 @@
 
 	if (entry.compressionMethod != OFZIPArchiveEntryCompressionMethodNone)
 		@throw [OFNotImplementedException exceptionWithSelector: _cmd
 								 object: self];
 
-	[self of_closeLastReturnedStream];
+	@try {
+		[_lastReturnedStream close];
+	} @catch (OFNotOpenException *e) {
+		/* Might have already been closed by the user - that's fine. */
+	}
+	_lastReturnedStream = nil;
 
 	fileName = entry.fileName;
 	fileNameLength = fileName.UTF8StringLength;
 	extraField = entry.extraField;
 	extraFieldLength = extraField.count;
@@ -563,16 +548,17 @@
 		@throw [OFOutOfRangeException exception];
 
 	_offset += offsetAdd;
 
 	_lastReturnedStream = [[OFZIPArchiveFileWriteStream alloc]
-	     initWithStream: _stream
-		      entry: entry];
+	    of_initWithArchive: self
+			stream: _stream
+			 entry: entry];
 
 	objc_autoreleasePoolPop(pool);
 
-	return [[_lastReturnedStream retain] autorelease];
+	return [_lastReturnedStream autorelease];
 }
 
 - (void)of_writeCentralDirectory
 {
 	void *pool = objc_autoreleasePoolPush();
@@ -625,11 +611,16 @@
 - (void)close
 {
 	if (_stream == nil)
 		@throw [OFNotOpenException exceptionWithObject: self];
 
-	[self of_closeLastReturnedStream];
+	@try {
+		[_lastReturnedStream close];
+	} @catch (OFNotOpenException *e) {
+		/* Might have already been closed by the user - that's fine. */
+	}
+	_lastReturnedStream = nil;
 
 	if (_mode == modeWrite || _mode == modeAppend)
 		[self of_writeCentralDirectory];
 
 	[_stream release];
@@ -735,16 +726,18 @@
 	return true;
 }
 @end
 
 @implementation OFZIPArchiveFileReadStream
-- (instancetype)of_initWithStream: (OFStream *)stream
-			    entry: (OFZIPArchiveEntry *)entry
+- (instancetype)of_initWithArchive: (OFZIPArchive *)archive
+			    stream: (OFStream *)stream
+			     entry: (OFZIPArchiveEntry *)entry
 {
 	self = [super init];
 
 	@try {
+		_archive = [archive retain];
 		_stream = [stream retain];
 
 		switch (entry.compressionMethod) {
 		case OFZIPArchiveEntryCompressionMethodNone:
 			_decompressedStream = [stream retain];
@@ -778,10 +771,15 @@
 {
 	if (_stream != nil || _decompressedStream != nil)
 		[self close];
 
 	[_entry release];
+
+	if (_archive->_lastReturnedStream == self)
+		_archive->_lastReturnedStream = nil;
+
+	[_archive release];
 
 	[super dealloc];
 }
 
 - (bool)lowlevelIsAtEndOfStream
@@ -862,15 +860,17 @@
 	[super close];
 }
 @end
 
 @implementation OFZIPArchiveFileWriteStream
-- (instancetype)initWithStream: (OFStream *)stream
-			 entry: (OFMutableZIPArchiveEntry *)entry
+- (instancetype)of_initWithArchive: (OFZIPArchive *)archive
+			    stream: (OFStream *)stream
+			     entry: (OFMutableZIPArchiveEntry *)entry
 {
 	self = [super init];
 
+	_archive = [archive retain];
 	_stream = [stream retain];
 	_entry = [entry retain];
 	_CRC32 = ~0;
 
 	return self;
@@ -880,10 +880,15 @@
 {
 	if (_stream != nil)
 		[self close];
 
 	[_entry release];
+
+	if (_archive->_lastReturnedStream == self)
+		_archive->_lastReturnedStream = nil;
+
+	[_archive release];
 
 	[super dealloc];
 }
 
 - (size_t)lowlevelWriteBuffer: (const void *)buffer length: (size_t)length
@@ -936,9 +941,17 @@
 	_entry.compressedSize = _bytesWritten;
 	_entry.uncompressedSize = _bytesWritten;
 	[_entry makeImmutable];
 
 	_bytesWritten += (2 * 4 + 2 * 8);
+
+	[_archive->_entries addObject: _entry];
+	[_archive->_pathToEntryMap setObject: _entry forKey: _entry.fileName];
+
+	if (ULLONG_MAX - _archive->_offset < _bytesWritten)
+		@throw [OFOutOfRangeException exception];
+
+	_archive->_offset += _bytesWritten;
 
 	[super close];
 }
 @end