Source/Core/DiscIO/Src/CompressedBlob.cpp

   1 // Copyright (C) 2003 Dolphin Project.
   2
   3 // This program is free software: you can redistribute it and/or modify
   4 // it under the terms of the GNU General Public License as published by
   5 // the Free Software Foundation, version 2.0.
   6
   7 // This program is distributed in the hope that it will be useful,
   8 // but WITHOUT ANY WARRANTY; without even the implied warranty of
   9 // MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  10 // GNU General Public License 2.0 for more details.
  11
  12 // A copy of the GPL 2.0 should have been included with the program.
  13 // If not, see http://www.gnu.org/licenses/
  14
  15 // Official SVN repository and contact information can be found at
  16 // http://code.google.com/p/dolphin-emu/
  17
  18 #include "stdafx.h"
  19
  20 #ifdef _WIN32
  21 #include <io.h>
  22 #include <windows.h>
  23 #else
  24 #include <unistd.h>
  25 #endif
  26
  27 #include "Common.h"
  28 #include "CompressedBlob.h"
  29 #include "DiscScrubber.h"
  30 #include "FileUtil.h"
  31 #include "Hash.h"
  32
  33 #include "zlib.h"
  34
  35 namespace DiscIO
  36 {
  37
  38 CompressedBlobReader::CompressedBlobReader(const char *filename)
  39 {
  40         file_name = filename;
  41         file = fopen(filename, "rb");
  42         fseek(file, 0, SEEK_END);
  43         file_size = ftell(file);
  44         fseek(file, 0, SEEK_SET);
  45         fread(&header, sizeof(CompressedBlobHeader), 1, file);
  46
  47         SetSectorSize(header.block_size);
  48
  49         // cache block pointers and hashes
  50         block_pointers = new u64[header.num_blocks];
  51         fread(block_pointers, sizeof(u64), header.num_blocks, file);
  52         hashes = new u32[header.num_blocks];
  53         fread(hashes, sizeof(u32), header.num_blocks, file);
  54
  55         data_offset = (sizeof(CompressedBlobHeader))
  56                 + (sizeof(u64)) * header.num_blocks   // skip block pointers
  57                 + (sizeof(u32)) * header.num_blocks;  // skip hashes
  58
  59         // A compressed block is never ever longer than a decompressed block, so just header.block_size should be fine.
  60         // I still add some safety margin.
  61         zlib_buffer_size = header.block_size + 64;
  62         zlib_buffer = new u8[zlib_buffer_size];
  63         memset(zlib_buffer, 0, zlib_buffer_size);
  64 }
  65
  66 CompressedBlobReader* CompressedBlobReader::Create(const char* filename)
  67 {
  68         if (IsCompressedBlob(filename))
  69                 return new CompressedBlobReader(filename);
  70         else
  71                 return 0;
  72 }
  73
  74 CompressedBlobReader::~CompressedBlobReader()
  75 {
  76         delete [] zlib_buffer;
  77         delete [] block_pointers;
  78         delete [] hashes;
  79         fclose(file);
  80         file = 0;
  81 }
  82
  83 // IMPORTANT: Calling this function invalidates all earlier pointers gotten from this function.
  84 u64 CompressedBlobReader::GetBlockCompressedSize(u64 block_num) const
  85 {
  86         u64 start = block_pointers[block_num];
  87         if (block_num < header.num_blocks - 1)
  88                 return block_pointers[block_num + 1] - start;
  89         else if (block_num == header.num_blocks - 1)
  90                 return header.compressed_data_size - start;
  91         else
  92                 PanicAlert("GetBlockCompressedSize - illegal block number %i", (int)block_num);
  93         return 0;
  94 }
  95
  96 void CompressedBlobReader::GetBlock(u64 block_num, u8 *out_ptr)
  97 {
  98         bool uncompressed = false;
  99         u32 comp_block_size = (u32)GetBlockCompressedSize(block_num);
 100         u64 offset = block_pointers[block_num] + data_offset;
 101
 102         if (offset & (1ULL << 63))
 103         {
 104                 if (comp_block_size != header.block_size)
 105                         PanicAlert("Uncompressed block with wrong size");
 106                 uncompressed = true;
 107                 offset &= ~(1ULL << 63);
 108         }
 109
 110         // clear unused part of zlib buffer. maybe this can be deleted when it works fully.
 111         memset(zlib_buffer + comp_block_size, 0, zlib_buffer_size - comp_block_size);
 112
 113         fseek(file, offset, SEEK_SET);
 114         fread(zlib_buffer, 1, comp_block_size, file);
 115
 116         u8* source = zlib_buffer;
 117         u8* dest = out_ptr;
 118
 119         // First, check hash.
 120         u32 block_hash = HashAdler32(source, comp_block_size);
 121         if (block_hash != hashes[block_num])
 122                 PanicAlert("Hash of block %i is %08x instead of %08x.\n"
 123                            "Your ISO, %s, is corrupt.",
 124                            block_num, block_hash, hashes[block_num],
 125                                    file_name.c_str());
 126
 127         if (uncompressed)
 128         {
 129                 memcpy(dest, source, comp_block_size);
 130         }
 131         else
 132         {
 133                 z_stream z;
 134                 memset(&z, 0, sizeof(z));
 135                 z.next_in  = source;
 136                 z.avail_in = comp_block_size;
 137                 if (z.avail_in > header.block_size)
 138                 {
 139                         PanicAlert("We have a problem");
 140                 }
 141                 z.next_out  = dest;
 142                 z.avail_out = header.block_size;
 143                 inflateInit(&z);
 144                 int status = inflate(&z, Z_FULL_FLUSH);
 145                 u32 uncomp_size = header.block_size - z.avail_out;
 146                 if (status != Z_STREAM_END)
 147                 {
 148                         // this seem to fire wrongly from time to time
 149                         // to be sure, don't use compressed isos :P
 150                         PanicAlert("Failure reading block %i - out of data and not at end.", block_num);
 151                 }
 152                 inflateEnd(&z);
 153                 if (uncomp_size != header.block_size)
 154                         PanicAlert("Wrong block size");
 155         }
 156 }
 157
 158 bool CompressFileToBlob(const char* infile, const char* outfile, u32 sub_type,
 159                                                 int block_size, CompressCB callback, void* arg)
 160 {
 161         bool scrubbing = false;
 162
 163         if (IsCompressedBlob(infile))
 164         {
 165                 PanicAlert("%s is already compressed! Cannot compress it further.", infile);
 166                 return false;
 167         }
 168
 169         if (sub_type == 1)
 170         {
 171                 if (!DiscScrubber::SetupScrub(infile, block_size))
 172                 {
 173                         PanicAlert("%s failed to be scrubbed. Probably the image is corrupt.", infile);
 174                         return false;
 175                 }
 176
 177                 scrubbing = true;
 178         }
 179
 180         FILE* inf = fopen(infile, "rb");
 181         if (!inf)
 182                 return false;
 183
 184         FILE* f = fopen(outfile, "wb");
 185         if (!f)
 186         {
 187                 fclose(inf);
 188                 return false;
 189         }
 190
 191         callback("Files opened, ready to compress.", 0, arg);
 192
 193         fseek(inf, 0, SEEK_END);
 194         s64 insize = ftell(inf);
 195         fseek(inf, 0, SEEK_SET);
 196         CompressedBlobHeader header;
 197         header.magic_cookie = kBlobCookie;
 198         header.sub_type   = sub_type;
 199         header.block_size = block_size;
 200         header.data_size  = insize;
 201
 202         // round upwards!
 203         header.num_blocks = (u32)((header.data_size + (block_size - 1)) / block_size);
 204
 205         u64* offsets = new u64[header.num_blocks];
 206         u32* hashes = new u32[header.num_blocks];
 207         u8* out_buf = new u8[block_size];
 208         u8* in_buf = new u8[block_size];
 209
 210         // seek past the header (we will write it at the end)
 211         fseek(f, sizeof(CompressedBlobHeader), SEEK_CUR);
 212         // seek past the offset and hash tables (we will write them at the end)
 213         fseek(f, (sizeof(u64) + sizeof(u32)) * header.num_blocks, SEEK_CUR);
 214
 215         // Now we are ready to write compressed data!
 216         u64 position = 0;
 217         int num_compressed = 0;
 218         int num_stored = 0;
 219         int progress_monitor = max<int>(1, header.num_blocks / 1000);
 220
 221         for (u32 i = 0; i < header.num_blocks; i++)
 222         {
 223                 if (i % progress_monitor == 0)
 224                 {
 225                         u64 inpos = ftell(inf);
 226                         int ratio = 0;
 227                         if (inpos != 0)
 228                                 ratio = (int)(100 * position / inpos);
 229                         char temp[512];
 230                         sprintf(temp, "%i of %i blocks. compression ratio %i%%", i, header.num_blocks, ratio);
 231                         callback(temp, (float)i / (float)header.num_blocks, arg);
 232                 }
 233
 234                 offsets[i] = position;
 235                 // u64 start = i * header.block_size;
 236                 // u64 size = header.block_size;
 237                 std::fill(in_buf, in_buf + header.block_size, 0);
 238                 if (scrubbing)
 239                         DiscScrubber::GetNextBlock(inf, in_buf);
 240                 else
 241                         fread(in_buf, header.block_size, 1, inf);
 242                 z_stream z;
 243                 memset(&z, 0, sizeof(z));
 244                 z.zalloc = Z_NULL;
 245                 z.zfree  = Z_NULL;
 246                 z.opaque = Z_NULL;
 247                 z.next_in   = in_buf;
 248                 z.avail_in  = header.block_size;
 249                 z.next_out  = out_buf;
 250                 z.avail_out = block_size;
 251                 int retval = deflateInit(&z, 9);
 252
 253                 if (retval != Z_OK)
 254                 {
 255                         ERROR_LOG(DISCIO, "Deflate failed");
 256                         goto cleanup;
 257                 }
 258
 259                 int status = deflate(&z, Z_FINISH);
 260                 int comp_size = block_size - z.avail_out;
 261                 if ((status != Z_STREAM_END) || (z.avail_out < 10))
 262                 {
 263                         //PanicAlert("%i %i Store %i", i*block_size, position, comp_size);
 264                         // let's store uncompressed
 265                         offsets[i] |= 0x8000000000000000ULL;
 266                         fwrite(in_buf, block_size, 1, f);
 267                         hashes[i] = HashAdler32(in_buf, block_size);
 268                         position += block_size;
 269                         num_stored++;
 270                 }
 271                 else
 272                 {
 273                         // let's store compressed
 274                         //PanicAlert("Comp %i to %i", block_size, comp_size);
 275                         fwrite(out_buf, comp_size, 1, f);
 276                         hashes[i] = HashAdler32(out_buf, comp_size);
 277                         position += comp_size;
 278                         num_compressed++;
 279                 }
 280
 281                 deflateEnd(&z);
 282         }
 283
 284         header.compressed_data_size = position;
 285
 286         // Okay, go back and fill in headers
 287         fseek(f, 0, SEEK_SET);
 288         fwrite(&header, sizeof(header), 1, f);
 289         fwrite(offsets, sizeof(u64), header.num_blocks, f);
 290         fwrite(hashes, sizeof(u32), header.num_blocks, f);
 291
 292 cleanup:
 293         // Cleanup
 294         delete[] in_buf;
 295         delete[] out_buf;
 296         delete[] offsets;
 297         delete[] hashes;
 298         fclose(f);
 299         fclose(inf);
 300         DiscScrubber::Cleanup();
 301         callback("Done compressing disc image.", 1.0f, arg);
 302         return true;
 303 }
 304
 305 bool DecompressBlobToFile(const char* infile, const char* outfile, CompressCB callback, void* arg)
 306 {
 307         if (!IsCompressedBlob(infile))
 308         {
 309                 PanicAlert("File not compressed");
 310                 return false;
 311         }
 312
 313         CompressedBlobReader* reader = CompressedBlobReader::Create(infile);
 314         if (!reader) return false;
 315
 316         FILE* f = fopen(outfile, "wb");
 317         const CompressedBlobHeader &header = reader->GetHeader();
 318         u8* buffer = new u8[header.block_size];
 319         int progress_monitor = max<int>(1, header.num_blocks / 100);
 320
 321         for (u64 i = 0; i < header.num_blocks; i++)
 322         {
 323                 if (i % progress_monitor == 0)
 324                 {
 325                         callback("Unpacking", (float)i / (float)header.num_blocks, arg);
 326                 }
 327                 reader->Read(i * header.block_size, header.block_size, buffer);
 328                 fwrite(buffer, header.block_size, 1, f);
 329         }
 330
 331         delete[] buffer;
 332
 333 #ifdef _WIN32
 334         // ector: _chsize sucks, not 64-bit safe
 335         // F|RES: changed to _chsize_s. i think it is 64-bit safe
 336         _chsize_s(_fileno(f), header.data_size);
 337 #else
 338         ftruncate(fileno(f), header.data_size);
 339 #endif
 340         fclose(f);
 341
 342         delete reader;
 343
 344         return true;
 345 }
 346
 347 bool IsCompressedBlob(const char* filename)
 348 {
 349         FILE* f = fopen(filename, "rb");
 350
 351         if (!f)
 352                 return false;
 353
 354         CompressedBlobHeader header;
 355         fread(&header, sizeof(header), 1, f);
 356         fclose(f);
 357         return header.magic_cookie == kBlobCookie;
 358 }
 359
 360 }  // namespace