From 35771b555d10a46aed91f7ab042c8f32bfe756fc Mon Sep 17 00:00:00 2001
From: Tino Reichardt <milky-7zip@mcmilk.de>
Date: Sun, 12 Apr 2020 23:33:35 +0200
Subject: [PATCH] Add support of lz4 and zstd compression for squashfs files

---
 CPP/7zip/Archive/SquashfsHandler.cpp | 43 +++++++++++++++++++++++++++-
 1 file changed, 42 insertions(+), 1 deletion(-)

diff --git a/CPP/7zip/Archive/SquashfsHandler.cpp b/CPP/7zip/Archive/SquashfsHandler.cpp
index f0d4fdb9..85009f01 100644
--- a/CPP/7zip/Archive/SquashfsHandler.cpp
+++ b/CPP/7zip/Archive/SquashfsHandler.cpp
@@ -6,6 +6,7 @@
 #include "../../../C/CpuArch.h"
 #include "../../../C/LzmaDec.h"
 #include "../../../C/Xz.h"
+#include "../../../C/lz4/lz4.h"
 
 #include "../../Common/ComTry.h"
 #include "../../Common/MyLinux.h"
@@ -26,6 +27,7 @@
 #include "../Compress/CopyCoder.h"
 #include "../Compress/ZlibDecoder.h"
 // #include "../Compress/LzmaDecoder.h"
+#include "../Compress/ZstdDecoder.h"
 
 namespace NArchive {
 namespace NSquashfs {
@@ -66,6 +68,8 @@ static const UInt32 kSignature32_B2 = 0x73687371;
 #define kMethod_LZMA 2
 #define kMethod_LZO  3
 #define kMethod_XZ   4
+#define kMethod_LZ4  5
+#define kMethod_ZSTD 6
 
 static const char * const k_Methods[] =
 {
@@ -74,6 +78,8 @@ static const char * const k_Methods[] =
   , "LZMA"
   , "LZO"
   , "XZ"
+  , "LZ4"
+  , "ZSTD"
 };
 
 static const UInt32 kMetadataBlockSizeLog = 13;
@@ -873,6 +879,9 @@ class CHandler:
   NCompress::NZlib::CDecoder *_zlibDecoderSpec;
   CMyComPtr<ICompressCoder> _zlibDecoder;
   
+  NCompress::NZSTD::CDecoder *_zstdDecoderSpec;
+  CMyComPtr<ICompressCoder> _zstdDecoder;
+
   CXzUnpacker _xz;
 
   CByteBuffer _inputBuffer;
@@ -1116,6 +1125,21 @@ static HRESULT LzoDecode(Byte *dest, SizeT *destLen, const Byte *src, SizeT *src
   }
 }
 
+static HRESULT Lz4Decode(Byte *dest, SizeT *destLen, const Byte *src, SizeT *srcLen)
+{
+  const char *Src = (const char *)src;
+  char *Dst = (char *)dest;
+  int compressedSize = (int)*srcLen;
+  int dstCapacity = (int)*destLen;
+  // int LZ4_decompress_safe (const char* src, char* dst, int compressedSize, int dstCapacity);
+  int rv = LZ4_decompress_safe(Src, Dst, compressedSize, dstCapacity);
+  if (rv == 0)
+    return S_FALSE;
+
+  *destLen = rv;
+  return S_OK;
+}
+
 HRESULT CHandler::Decompress(ISequentialOutStream *outStream, Byte *outBuf, bool *outBufWasWritten, UInt32 *outBufWasWrittenSize, UInt32 inSize, UInt32 outSizeMax)
 {
   if (outBuf)
@@ -1156,7 +1180,18 @@ HRESULT CHandler::Decompress(ISequentialOutStream *outStream, Byte *outBuf, bool
     if (inSize != _zlibDecoderSpec->GetInputProcessedSize())
       return S_FALSE;
   }
-  /*
+  else if (method == kMethod_ZSTD)
+  {
+    if (!_zstdDecoder)
+    {
+      _zstdDecoderSpec = new NCompress::NZSTD::CDecoder();
+      _zstdDecoder = _zstdDecoderSpec;
+    }
+    RINOK(_zstdDecoder->Code(_limitedInStream, outStream, NULL, NULL, NULL));
+    if (inSize != _zstdDecoderSpec->GetInputProcessedSize())
+      return S_FALSE;
+  }
+   /*
   else if (method == kMethod_LZMA)
   {
     if (!_lzmaDecoder)
@@ -1210,6 +1245,10 @@ HRESULT CHandler::Decompress(ISequentialOutStream *outStream, Byte *outBuf, bool
     {
       RINOK(LzoDecode(dest, &destLen, _inputBuffer, &srcLen));
     }
+    else if (method == kMethod_LZ4)
+    {
+      RINOK(Lz4Decode(dest, &destLen, _inputBuffer, &srcLen));
+    }
     else if (method == kMethod_LZMA)
     {
       Byte props[5];
@@ -1539,6 +1578,8 @@ HRESULT CHandler::Open2(IInStream *inStream)
       case kMethod_LZMA:
       case kMethod_LZO:
       case kMethod_XZ:
+      case kMethod_LZ4:
+      case kMethod_ZSTD:
         break;
       default:
         return E_NOTIMPL;