improve PRS efficiency further

src/Compression.cc

@@ -58,15 +58,22 @@ void PRSCompressor::advance() {
 
   for (auto it = start_offsets.begin(); (it != start_offsets.end()) && (best_match_size < 0x100); it++) {
     size_t match_offset = *it;
+    if (match_offset == this->compression_offset - 0x2000) {
+      continue;
+    }
+
     size_t match_size = 0;
+    size_t match_loop_bytes = this->compression_offset - match_offset;
     while ((match_size < 0x100) &&
-           (match_offset + match_size < this->compression_offset) &&
            (this->compression_offset + match_size < this->input_bytes) &&
-           (this->reverse_log[(match_offset + match_size) & 0x1FFF] == this->forward_log[(this->compression_offset + match_size) & 0xFF])) {
+           (this->reverse_log[(match_offset + (match_size % match_loop_bytes)) & 0x1FFF] == this->forward_log[(this->compression_offset + match_size) & 0xFF])) {
       match_size++;
     }
 
-    if (match_size > best_match_size) {
+    // If there are multiple matches of the longest length, use the latest one,
+    // since it's more likely that it can be expressed as a short copy instead
+    // of a long copy.
+    if (match_size >= best_match_size) {
       best_match_offset = match_offset;
       best_match_size = match_size;
     }
@@ -76,7 +83,7 @@ void PRSCompressor::advance() {
   bool should_write_literal = false;
   size_t advance_bytes = 0;
   ssize_t backreference_offset = best_match_offset - this->compression_offset;
-  if (best_match_size < 2 || backreference_offset == -0x2000) {
+  if (best_match_size < 2) {
     should_write_literal = true;
 
   } else {
@@ -142,9 +149,11 @@ void PRSCompressor::advance() {
     uint8_t existing_v = this->reverse_log[reverse_log_offset];
     uint8_t new_v = this->forward_log[this->compression_offset & 0xFF];
 
-    this->reverse_log_index[existing_v].erase(this->compression_offset - this->reverse_log.size());
+    if (this->compression_offset & (~0x1FFF)) {
+      this->reverse_log_index[existing_v].pop_front();
+    }
     this->reverse_log[reverse_log_offset] = new_v;
-    this->reverse_log_index[new_v].emplace(this->compression_offset);
+    this->reverse_log_index[new_v].emplace_back(this->compression_offset);
     this->compression_offset++;
   }
 }
@@ -381,6 +390,57 @@ size_t prs_decompress_size(const string& data, size_t max_output_size) {
 
 
 
+void prs_disassemble(FILE* stream, const void* data, size_t size) {
+  size_t output_bytes = 0;
+  StringReader r(data, size);
+  ControlStreamReader cr(r);
+
+  while (!r.eof()) {
+    size_t r_offset = r.where();
+    if (cr.read()) {
+      fprintf(stream, "[%zX => %zX] literal %02hhX\n", r_offset, output_bytes, r.get_u8());
+      output_bytes++;
+
+    } else {
+      ssize_t offset;
+      size_t count;
+
+      bool is_long_copy = cr.read();
+      if (is_long_copy) {
+        uint16_t a = r.get_u8();
+        a |= (r.get_u8() << 8);
+        offset = (a >> 3) | (~0x1FFF);
+        if (offset == ~0x1FFF) {
+          fprintf(stream, "[%zX => %zX] end\n", r_offset, output_bytes);
+          break;
+        }
+        count = (a & 7) ? ((a & 7) + 2) : (r.get_u8() + 1);
+
+      } else {
+        count = cr.read() << 1;
+        count = (count | cr.read()) + 2;
+        offset = r.get_u8() | (~0xFF);
+      }
+
+      size_t read_offset = output_bytes + offset;
+      fprintf(stream, "[%zX => %zX] %s copy -%zX (from %zX) %zX\n",
+          r_offset, output_bytes, is_long_copy ? "long" : "short",
+          -offset, read_offset, count);
+
+      if (read_offset >= output_bytes) {
+        throw runtime_error("backreference offset beyond beginning of output");
+      }
+      output_bytes += count;
+    }
+  }
+}
+
+void prs_disassemble(FILE* stream, const std::string& data) {
+  return prs_disassemble(stream, data.data(), data.size());
+}
+
+
+
 // BC0 is a compression algorithm fairly similar to PRS, but with a simpler set
 // of commands. Like PRS, there is a control stream, indicating when to copy a
 // literal byte from the input and when to copy from a backreference; unlike
@@ -394,13 +454,15 @@ string bc0_compress(
 
   parray<uint8_t, 0x1000> memo;
   uint16_t memo_offset = 0x0FEE;
-  vector<set<size_t>> memo_index(0x100);
+  vector<deque<size_t>> memo_index(0x100);
   auto write_memo = [&](uint8_t new_v) -> void {
     uint8_t existing_v = memo[memo_offset];
     if (existing_v != new_v) {
-      memo_index[existing_v].erase(memo_offset);
+      if (!memo_index[existing_v].empty()) {
+        memo_index[existing_v].pop_front();
+      }
       memo[memo_offset] = new_v;
-      memo_index[new_v].emplace(memo_offset);
+      memo_index[new_v].emplace_back(memo_offset);
     }
     memo_offset = (memo_offset + 1) & 0xFFF;
   };
@@ -427,6 +489,9 @@ string bc0_compress(
         // Forbid matches that span the memo boundary - during decompression,
         // the client will be overwriting its memo while reading from it and
         // will likely generate incorrect data
+        // TODO: We can actually support this (and it will improve compression),
+        // but we have to set a loop boundary like we have in prs_compress and
+        // I'm lazy.
         size_t start_memo_offset = offset;
         size_t end_memo_offset = (offset + match_size) & 0xFFF;
         if (end_memo_offset < start_memo_offset) {
@@ -512,8 +577,8 @@ string bc0_compress(
 // The BC0 decompression implementation in PSO GC is vulnerable to overflow
 // attacks - there is no bounds checking on the output buffer. It is unlikely
 // that this can be usefully exploited (e.g. for RCE) because the output pointer
-// is checked before every byte is written, so we cannot change the output
-// pointer to any arbitrary address.
+// is loaded from memory before every byte is written, so we cannot change the
+// output pointer to any arbitrary address.
 
 string bc0_decompress(const string& data) {
   StringReader r(data);

src/Compression.hh

@@ -4,7 +4,7 @@
 
 #include <string>
 #include <functional>
-#include <set>
+#include <deque>
 
 #include "Text.hh"
 
@@ -51,7 +51,7 @@ private:
   parray<uint8_t, 0x100> forward_log;
   size_t compression_offset;
   parray<uint8_t, 0x2000> reverse_log;
-  std::vector<std::set<size_t>> reverse_log_index;
+  std::vector<std::deque<size_t>> reverse_log_index;
 
   StringWriter output;
 };
@@ -71,6 +71,10 @@ std::string prs_decompress(const std::string& data, size_t max_output_size = 0);
 size_t prs_decompress_size(const void* data, size_t size, size_t max_output_size = 0);
 size_t prs_decompress_size(const std::string& data, size_t max_output_size = 0);
 
+// Prints the command stream from a PRS-compressed buffer.
+void prs_disassemble(FILE* stream, const void* data, size_t size);
+void prs_disassemble(FILE* stream, const std::string& data);
+
 // Compresses and decompresses data using the BC0 algorithm.
 std::string bc0_compress(const std::string& data, std::function<void(size_t, size_t)> progress_fn = nullptr);
 std::string bc0_decompress(const std::string& data);

src/Main.cc

@@ -380,6 +380,7 @@ enum class Behavior {
   COMPRESS_BC0,
   DECOMPRESS_BC0,
   PRS_SIZE,
+  PRS_DISASSEMBLE,
   ENCRYPT_DATA,
   DECRYPT_DATA,
   DECRYPT_TRIVIAL_DATA,
@@ -399,6 +400,7 @@ static bool behavior_takes_input_filename(Behavior b) {
          (b == Behavior::COMPRESS_BC0) ||
          (b == Behavior::DECOMPRESS_BC0) ||
          (b == Behavior::PRS_SIZE) ||
+         (b == Behavior::PRS_DISASSEMBLE) ||
          (b == Behavior::ENCRYPT_DATA) ||
          (b == Behavior::DECRYPT_DATA) ||
          (b == Behavior::DECRYPT_TRIVIAL_DATA) ||
@@ -504,6 +506,8 @@ int main(int argc, char** argv) {
           behavior = Behavior::DECOMPRESS_BC0;
         } else if (!strcmp(argv[x], "prs-size")) {
           behavior = Behavior::PRS_SIZE;
+        } else if (!strcmp(argv[x], "disassemble-prs")) {
+          behavior = Behavior::PRS_DISASSEMBLE;
         } else if (!strcmp(argv[x], "encrypt-data")) {
           behavior = Behavior::ENCRYPT_DATA;
         } else if (!strcmp(argv[x], "decrypt-data")) {
@@ -597,6 +601,7 @@ int main(int argc, char** argv) {
             input_progress, input_bytes, progress, output_progress, size_ratio);
       };
 
+      uint64_t start = now();
       if (behavior == Behavior::COMPRESS_PRS) {
         data = prs_compress(data, progress_fn);
       } else if (behavior == Behavior::DECOMPRESS_PRS) {
@@ -608,9 +613,12 @@ int main(int argc, char** argv) {
       } else {
         throw logic_error("invalid behavior");
       }
+      uint64_t end = now();
+      string time_str = format_duration(end - start);
+
       float size_ratio = static_cast<float>(data.size() * 100) / input_bytes;
-      log_info("%zu (0x%zX) bytes input => %zu (0x%zX) bytes output (%g%%)",
-          input_bytes, input_bytes, data.size(), data.size(), size_ratio);
+      log_info("%zu (0x%zX) bytes input => %zu (0x%zX) bytes output (%g%%) in %s",
+          input_bytes, input_bytes, data.size(), data.size(), size_ratio, time_str.c_str());
 
       write_output_data(data.data(), data.size());
       break;
@@ -625,6 +633,11 @@ int main(int argc, char** argv) {
       break;
     }
 
+    case Behavior::PRS_DISASSEMBLE: {
+      prs_disassemble(stdout, read_input_data());
+      break;
+    }
+
     case Behavior::DECRYPT_DATA:
     case Behavior::ENCRYPT_DATA: {
       shared_ptr<PSOEncryption> crypt;