Fix OoB checks in ZSO + Updated Python ZSO script (#1209)

* fixed ziso_total_block calculation

* adjust oob reads

* fix python script

* clang format

* cleanup ziso script

modules/isofs/zso.c

@@ -21,8 +21,9 @@ void ziso_init(ZISO_header *header, u32 first_block)
     // read header information
     ziso_align = header->align;
     ziso_idx_start_block = -1;
-    // calculate number of blocks without using uncompressed_size (avoid 64bit division)
-    ziso_total_block = ((((first_block & 0x7FFFFFFF) << ziso_align) - sizeof(ZISO_header)) / 4) - 1;
+    // calculate number of blocks without using 64bit division library
+    u32 *total_bytes_p = (u32 *)&(header->total_bytes);
+    ziso_total_block = (total_bytes_p[0] >> 11) | ((total_bytes_p[1] & 0x7ff) << 21);
     // allocate memory
     if (ziso_tmp_buf == NULL) {
         ziso_tmp_buf = ziso_alloc(2048 + sizeof(u32) * ZISO_IDX_MAX_ENTRIES + 64);
@@ -49,6 +50,10 @@ int ziso_read_sector(u8 *addr, u32 lsn, unsigned int count)
         return 0; // can't seek beyond file
     }
 
+    if (lsn + count > ziso_total_block) {
+        count = ziso_total_block - lsn; // adjust oob reads
+    }
+
     // refresh index table if needed
     if (ziso_idx_start_block < 0 || lsn < ziso_idx_start_block || lsn + count >= ziso_idx_start_block + ZISO_IDX_MAX_ENTRIES - 1) {
         read_raw_data((u8 *)ziso_idx_cache, ZISO_IDX_MAX_ENTRIES * sizeof(u32), lsn * 4 + sizeof(ZISO_header), 0);

modules/isofs/zso.h

@@ -8,7 +8,7 @@
 
 #define ZSO_MAGIC 0x4F53495A // ZISO
 
-// no game should request more than 256 sectors per read
+// no game should request more than 256 sectors per read (512KB of data)
 // should allow us to decompress all data with only 2 IO calls at most.
 #define ZISO_IDX_MAX_ENTRIES 257
 

pc/ziso.py

@@ -33,7 +33,8 @@
 
 ZISO_MAGIC = 0x4F53495A
 DEFAULT_ALIGN = 0
-COMPRESS_THREHOLD = 100
+DEFAULT_BLOCK_SIZE = 0x800
+COMPRESS_THREHOLD = 95
 DEFAULT_PADDING = br'X'
 
 MP = False
@@ -47,11 +48,15 @@ def hexdump(data):
 
 
 def lz4_compress(plain, level=9):
-    return lz4.block.compress(plain, store_size=False)
+    mode = "high_compression" if level > 1 else "default"
+    return lz4.block.compress(plain, mode=mode, compression=level, store_size=False)
 
 
 def lz4_compress_mp(i):
-    return lz4.block.compress(i[0], store_size=False)
+    plain = i[0]
+    level = i[1]
+    mode = "high_compression" if level > 1 else "default"
+    return lz4.block.compress(plain, mode=mode, compression=level, store_size=False)
 
 
 def lz4_decompress(compressed, block_size):
@@ -68,8 +73,9 @@ def lz4_decompress(compressed, block_size):
 
 def usage():
     print("Usage: ziso [-c level] [-m] [-t percent] [-h] infile outfile")
-    print("  -c level: 1-9 compress ISO to ZSO, use any non-zero number it has no effect")
+    print("  -c level: 1-12 compress ISO to ZSO, 1 for standard compression, >1 for high compression")
     print("              0 decompress ZSO to ISO")
+    print("  -b size:  2048-8192, specify block size (2048 by default)")
     print("  -m Use multiprocessing acceleration for compressing")
     print("  -t percent Compression Threshold (1-100)")
     print("  -a align Padding alignment 0=small/slow 6=fast/large")
@@ -112,12 +118,13 @@ def generate_zso_header(magic, header_size, total_bytes, block_size, ver, align)
     return data
 
 
-def show_zso_info(fname_in, fname_out, total_bytes, block_size, total_block, align):
+def show_zso_info(fname_in, fname_out, total_bytes, block_size, total_block, ver, align):
     print("Decompress '%s' to '%s'" % (fname_in, fname_out))
     print("Total File Size %ld bytes" % (total_bytes))
     print("block size      %d  bytes" % (block_size))
     print("total blocks    %d  blocks" % (total_block))
     print("index align     %d" % (align))
+    print("version         %d" % (ver))
 
 
 def decompress_zso(fname_in, fname_out):
@@ -136,7 +143,7 @@ def decompress_zso(fname_in, fname_out):
         index_buf.append(unpack('I', fin.read(4))[0])
 
     show_zso_info(fname_in, fname_out, total_bytes,
-                  block_size, total_block, align)
+                  block_size, total_block, ver, align)
 
     block = 0
     percent_period = total_block/100
@@ -189,12 +196,13 @@ def decompress_zso(fname_in, fname_out):
     print("ziso decompress completed")
 
 
-def show_comp_info(fname_in, fname_out, total_bytes, block_size, align, level):
+def show_comp_info(fname_in, fname_out, total_bytes, block_size, ver, align, level):
     print("Compress '%s' to '%s'" % (fname_in, fname_out))
     print("Total File Size %ld bytes" % (total_bytes))
     print("block size      %d  bytes" % (block_size))
     print("index align     %d" % (1 << align))
     print("compress level  %d" % (level))
+    print("version         %d" % (ver))
     if MP:
         print("multiprocessing %s" % (MP))
 
@@ -208,13 +216,13 @@ def set_align(fout, write_pos, align):
     return write_pos
 
 
-def compress_zso(fname_in, fname_out, level):
+def compress_zso(fname_in, fname_out, level, bsize):
     fin, fout = open_input_output(fname_in, fname_out)
     fin.seek(0, os.SEEK_END)
     total_bytes = fin.tell()
     fin.seek(0)
 
-    magic, header_size, block_size, ver, align = ZISO_MAGIC, 0x18, 0x800, 1, DEFAULT_ALIGN
+    magic, header_size, block_size, ver, align = ZISO_MAGIC, 0x18, bsize, 1, DEFAULT_ALIGN
 
     # We have to use alignment on any ZSO files which > 2GB, for MSB bit of index as the plain indicator
     # If we don't then the index can be larger than 2GB, which its plain indicator was improperly set
@@ -228,7 +236,7 @@ def compress_zso(fname_in, fname_out, level):
     index_buf = [0 for i in range(total_block + 1)]
 
     fout.write(b"\x00\x00\x00\x00" * len(index_buf))
-    show_comp_info(fname_in, fname_out, total_bytes, block_size, align, level)
+    show_comp_info(fname_in, fname_out, total_bytes, block_size, ver, align, level)
 
     write_pos = fout.tell()
     percent_period = total_block/100
@@ -252,7 +260,7 @@ def compress_zso(fname_in, fname_out, level):
                     block / percent_period, 0), file=sys.stderr, end='\r')
             else:
                 print("compress %3d%% avarage rate %3d%%\r" % (
-                    block / percent_period, 100*write_pos/(block*0x800)), file=sys.stderr, end='\r')
+                    block / percent_period, 100*write_pos/(block*block_size)), file=sys.stderr, end='\r')
 
         if MP:
             iso_data = [(fin.read(block_size), level)
@@ -323,17 +331,20 @@ def parse_args():
         sys.exit(-1)
 
     try:
-        optlist, args = gnu_getopt(sys.argv, "c:mt:a:p:h")
+        optlist, args = gnu_getopt(sys.argv, "c:b:mt:a:p:h")
     except GetoptError as err:
         print(str(err))
         usage()
         sys.exit(-1)
 
     level = None
+    bsize = DEFAULT_BLOCK_SIZE
 
     for o, a in optlist:
         if o == '-c':
             level = int(a)
+        elif o == '-b':
+            bsize = int(a)
         elif o == '-m':
             MP = True
         elif o == '-t':
@@ -352,7 +363,11 @@ def parse_args():
         print("You have to specify input/output filename: %s", err)
         sys.exit(-1)
 
-    return level, fname_in, fname_out
+    if bsize%2048 != 0:
+        print("Error, invalid block size. Must be multiple of 2048.")
+        sys.exit(-1)
+
+    return level, bsize, fname_in, fname_out
 
 
 def load_sector_table(sector_table_fn, total_block, default_level=9):
@@ -391,12 +406,12 @@ def load_sector_table(sector_table_fn, total_block, default_level=9):
 
 def main():
     print("ziso-python %s by %s" % (__version__, __author__))
-    level, fname_in, fname_out = parse_args()
+    level, bsize, fname_in, fname_out = parse_args()
 
     if level == 0:
         decompress_zso(fname_in, fname_out)
     else:
-        compress_zso(fname_in, fname_out, level)
+        compress_zso(fname_in, fname_out, level, bsize)
 
 
 PROFILE = False