BPackageWriter: Add Recompress()

Allows rewriting an existing package file with a different compression.

headers/os/package/hpkg/PackageReader.h

@@ -24,6 +24,7 @@ class BAbstractBufferedDataReader;
 class BErrorOutput;
 class BLowLevelPackageContentHandler;
 class BPackageContentHandler;
+class BPackageWriter;
 
 
 class BPackageReader {
@@ -43,6 +44,9 @@ public:
 			BAbstractBufferedDataReader* HeapReader() const;
 									// Only valid as long as the reader lives.
 
+private:
+			friend class BPackageWriter;
+
 private:
 			PackageReaderImpl*	fImpl;
 };

headers/os/package/hpkg/PackageWriter.h

@@ -16,6 +16,9 @@ namespace BPackageKit {
 namespace BHPKG {
 
 
+class BPackageReader;
+
+
 namespace BPrivate {
 	class PackageWriterImpl;
 }
@@ -72,6 +75,9 @@ public:
 			status_t			AddEntry(const char* fileName, int fd = -1);
 			status_t			Finish();
 
+			status_t			Recompress(BPackageReader* reader);
+									// to be called after Init(); no Finish()
+
 private:
 			PackageWriterImpl*	fImpl;
 };

headers/private/package/hpkg/PackageFileHeapAccessorBase.h

@@ -92,6 +92,10 @@ public:
 			void				SetFD(int fd)
 									{ fFD = fd; }
 
+			uint64				HeapOverhead(uint64 uncompressedSize) const;
+									// additional bytes needed when storing
+									// the given amount of data
+
 	// BAbstractBufferedDataReader
 	virtual	status_t			ReadDataToOutput(off_t offset,
 									size_t size, BDataIO* output);

headers/private/package/hpkg/PackageReaderImpl.h

@@ -22,6 +22,9 @@ class BPackageEntryAttribute;
 namespace BPrivate {
 
 
+class PackageWriterImpl;
+
+
 class PackageReaderImpl : public ReaderImplBase {
 	typedef	ReaderImplBase		inherited;
 public:
@@ -58,6 +61,8 @@ private:
 			struct EntryAttributeHandler;
 			struct RootAttributeHandler;
 
+			friend class PackageWriterImpl;
+
 private:
 			status_t			_ParseTOC(AttributeHandlerContext* context,
 									AttributeHandler* rootAttributeHandler);

headers/private/package/hpkg/PackageWriterImpl.h

@@ -35,6 +35,7 @@ class BPackageWriterParameters;
 namespace BPrivate {
 
 
+class PackageReaderImpl;
 struct hpkg_header;
 
 
@@ -53,6 +54,9 @@ public:
 			status_t			AddEntry(const char* fileName, int fd = -1);
 			status_t			Finish();
 
+			status_t			Recompress(PackageReaderImpl* reader);
+									// to be called after Init(); no Finish()
+
 private:
 			struct Attribute;
 			struct PackageContentHandler;
@@ -67,6 +71,8 @@ private:
 									const BPackageWriterParameters& parameters);
 			status_t			_Finish();
 
+			status_t			_Recompress(PackageReaderImpl* reader);
+
 			status_t			_RegisterEntry(const char* fileName, int fd);
 			Entry*				_RegisterEntry(Entry* parent,
 									const char* name, size_t nameLength, int fd,

headers/private/package/hpkg/WriterImplBase.h

@@ -151,6 +151,7 @@ protected:
 
 	inline	int					FD() const;
 	inline	uint32				Flags() const;
+	inline	const BPackageWriterParameters& Parameters() const;
 
 	inline	const PackageAttributeList&	PackageAttributes() const;
 	inline	PackageAttributeList&	PackageAttributes();
@@ -230,6 +231,13 @@ WriterImplBase::Flags() const
 }
 
 
+inline const BPackageWriterParameters&
+WriterImplBase::Parameters() const
+{
+	return fParameters;
+}
+
+
 inline const WriterImplBase::PackageAttributeList&
 WriterImplBase::PackageAttributes() const
 {

src/kits/package/hpkg/PackageFileHeapAccessorBase.cpp

@@ -144,6 +144,16 @@ PackageFileHeapAccessorBase::~PackageFileHeapAccessorBase()
 }
 
 
+uint64
+PackageFileHeapAccessorBase::HeapOverhead(uint64 uncompressedSize) const
+{
+	// Determine number of chunks and the size of the chunk size table. Note
+	// that the size of the last chunk is not saved, since its size is implied.
+	size_t chunkCount = (uncompressedSize + kChunkSize - 1) / kChunkSize;
+	return chunkCount > 1 ? (chunkCount - 1) * 2 : 0;
+}
+
+
 status_t
 PackageFileHeapAccessorBase::ReadDataToOutput(off_t offset, size_t size,
 	BDataIO* output)

src/kits/package/hpkg/PackageWriter.cpp

@@ -8,6 +8,8 @@
 
 #include <new>
 
+#include <package/hpkg/PackageReader.h>
+
 #include <package/hpkg/PackageWriterImpl.h>
 
 
@@ -128,6 +130,16 @@ BPackageWriter::Finish()
 }
 
 
+status_t
+BPackageWriter::Recompress(BPackageReader* reader)
+{
+	if (fImpl == NULL)
+		return B_NO_INIT;
+
+	return fImpl->Recompress(reader->fImpl);
+}
+
+
 }	// namespace BHPKG
 
 }	// namespace BPackageKit

src/kits/package/hpkg/PackageWriterImpl.cpp

@@ -38,6 +38,7 @@
 #include <package/hpkg/HPKGDefsPrivate.h>
 
 #include <package/hpkg/DataReader.h>
+#include <package/hpkg/PackageFileHeapReader.h>
 #include <package/hpkg/PackageFileHeapWriter.h>
 #include <package/hpkg/PackageReaderImpl.h>
 #include <package/hpkg/Stacker.h>
@@ -605,6 +606,23 @@ PackageWriterImpl::Finish()
 }
 
 
+status_t
+PackageWriterImpl::Recompress(PackageReaderImpl* reader)
+{
+	if (reader == NULL)
+		return B_BAD_VALUE;
+
+	try {
+		return _Recompress(reader);
+	} catch (status_t error) {
+		return error;
+	} catch (std::bad_alloc) {
+		fListener->PrintError("Out of memory!\n");
+		return B_NO_MEMORY;
+	}
+}
+
+
 status_t
 PackageWriterImpl::_Init(const char* fileName,
 	const BPackageWriterParameters& parameters)
@@ -724,6 +742,69 @@ PackageWriterImpl::_Finish()
 }
 
 
+status_t
+PackageWriterImpl::_Recompress(PackageReaderImpl* reader)
+{
+	if (reader == NULL)
+		return B_BAD_VALUE;
+
+	// read the header
+	hpkg_header header;
+	status_t error = reader->ReadBuffer(0, &header, sizeof(header));
+	if (error != B_OK) {
+		fListener->PrintError("Failed to reader hpkg header: %s\n",
+			strerror(error));
+		return error;
+	}
+
+	// Update some header fields, assuming no compression. We'll rewrite the
+	// header later, should compression have been used. Doing it this way allows
+	// for streaming an uncompressed package.
+	uint64 uncompressedHeapSize
+		= reader->RawHeapReader()->UncompressedHeapSize();
+	uint64 compressedHeapSize = uncompressedHeapSize
+		+ fHeapWriter->HeapOverhead(uncompressedHeapSize);
+
+	off_t totalSize = fHeapWriter->HeapOffset() + (off_t)compressedHeapSize;
+
+	header.heap_compression = B_HOST_TO_BENDIAN_INT16(B_HPKG_COMPRESSION_ZLIB);
+	header.heap_chunk_size = B_HOST_TO_BENDIAN_INT32(fHeapWriter->ChunkSize());
+	header.heap_size_compressed = B_HOST_TO_BENDIAN_INT64(compressedHeapSize);
+	header.heap_size_uncompressed
+		= B_HOST_TO_BENDIAN_INT64(uncompressedHeapSize);
+	header.total_size = B_HOST_TO_BENDIAN_INT64(totalSize);
+
+	if (Parameters().CompressionLevel() == 0)
+		RawWriteBuffer(&header, sizeof(hpkg_header), 0);
+
+	// copy the heap data
+	uint64 bytesCompressed;
+	error = fHeapWriter->AddData(*reader->RawHeapReader(), uncompressedHeapSize,
+		bytesCompressed);
+	if (error != B_OK)
+		return error;
+
+	// flush the heap
+	error = fHeapWriter->Finish();
+	if (error != B_OK)
+		return error;
+
+	// If compression is enabled, update and write the header.
+	if (Parameters().CompressionLevel() != 0) {
+		compressedHeapSize = fHeapWriter->CompressedHeapSize();
+		totalSize = fHeapWriter->HeapOffset() + (off_t)compressedHeapSize;
+		header.heap_size_compressed = B_HOST_TO_BENDIAN_INT64(compressedHeapSize);
+		header.total_size = B_HOST_TO_BENDIAN_INT64(totalSize);
+
+		// write the header
+		RawWriteBuffer(&header, sizeof(hpkg_header), 0);
+	}
+
+	SetFinished(true);
+	return B_OK;
+}
+
+
 status_t
 PackageWriterImpl::_CheckLicenses()
 {