From a426b323c0cec6224c0f8a600e4fa8d52025ffb5 Mon Sep 17 00:00:00 2001
From: sansmoraxz <31374920+sansmoraxz@users.noreply.github.com>
Date: Thu, 13 Jun 2024 00:11:38 +0530
Subject: [PATCH] Add more decompression algorithms

---
 decompress.go | 58 +++++++++++++++++++++++++++++++++++++++++++++++++++
 download.go   | 48 ++++++++++++++++++------------------------
 go.mod        |  6 +++---
 go.sum        | 14 +++++++++++--
 utils.go      |  8 ++++---
 5 files changed, 98 insertions(+), 36 deletions(-)
 create mode 100644 decompress.go

diff --git a/decompress.go b/decompress.go
new file mode 100644
index 0000000..2b9e40f
--- /dev/null
+++ b/decompress.go
@@ -0,0 +1,58 @@
+package godm
+
+import (
+	"compress/flate"
+	"compress/gzip"
+	"compress/lzw"
+	"fmt"
+	"io"
+	"os"
+
+	"github.com/klauspost/compress/zstd"
+
+	"github.com/dsnet/compress/brotli"
+)
+
+var supportedEncodings = []string{"gzip", "br", "zstd", "compress", "deflate"}
+
+func getReader(encoding string, mainFile *os.File) (io.ReadCloser, error) {
+	switch encoding {
+	case "gzip":
+		return gzip.NewReader(mainFile)
+	case "br":
+		return brotli.NewReader(mainFile, nil)
+	case "zstd":
+		if d, err := zstd.NewReader(mainFile); err != nil {
+			return nil, err
+		} else {
+			return d.IOReadCloser(), nil
+		}
+	case "compress":
+		return lzw.NewReader(mainFile, lzw.LSB, 8), nil
+	case "deflate":
+		return flate.NewReader(mainFile), nil
+	default:
+		return nil, fmt.Errorf("unknown encoding: %s", encoding)
+	}
+}
+
+func decompressFile(mainFile *os.File, filePath string, encoding string) error {
+	mainFile.Seek(0, 0)
+
+	gr, err := getReader(encoding, mainFile)
+	if err != nil {
+		return err
+	}
+	defer gr.Close()
+	log.Info("Decompressing file...")
+	fout, err := os.Create(filePath)
+	if err != nil {
+		return err
+	}
+	defer fout.Close()
+	if _, err = io.Copy(fout, gr); err != nil {
+		return err
+	}
+	log.Info("Decompressed file...")
+	return nil
+}
diff --git a/download.go b/download.go
index 5afe075..93526cf 100644
--- a/download.go
+++ b/download.go
@@ -1,11 +1,12 @@
 package godm
 
 import (
-	"compress/gzip"
+	"errors"
 	"fmt"
 	"io"
 	"net/http"
 	"os"
+	"slices"
 	"strconv"
 	"sync"
 	"time"
@@ -35,9 +36,19 @@ func DownloadFile(filePath string, url string, displayDownloadBar bool, compress
 	isAcceptRanges := headers.Get("Accept-Ranges") == "bytes"
 	length, _ := strconv.Atoi(headers.Get("Content-Length")) // it will be 0 if not present
 	etag := headers.Get("ETag")
+	encoding := headers.Get("Content-Encoding")
+
+	if encoding == "" {
+		compress = false
+	} else if slices.Contains(supportedEncodings, encoding) {
+		compress = true
+	} else {
+		return errors.New("Unknown encoding: " + encoding)
+	}
 
 	log.Info("Downloading: ", url)
 	log.Info("Content-Length: ", length)
+	log.Info("Content-Encoding: ", encoding)
 	log.Info("Accept-Ranges: ", isAcceptRanges)
 	log.Info("ETag: ", etag)
 
@@ -117,15 +128,16 @@ func DownloadFile(filePath string, url string, displayDownloadBar bool, compress
 	var reassembledFile *os.File
 	
 	if compress {
-		// intermediate gzip file if compress is true
-		if reassembledFile, err = os.Create(filePath + ".gz"); err != nil {
+		// intermediate archive file if compress is true
+		archivePath := filePath + ".archive"
+		if reassembledFile, err = os.Create(archivePath); err != nil {
 			return err
 		}
 		defer func() {
 			if err = reassembledFile.Close(); err != nil {
 				log.Error("Error: ", err)
 			}
-			if err = os.Remove(filePath + ".gz"); err != nil {
+			if err = os.Remove(archivePath); err != nil {
 				log.Error("Error: ", err)
 			}
 		}()
@@ -137,12 +149,14 @@ func DownloadFile(filePath string, url string, displayDownloadBar bool, compress
 	}
 	log.Info("Reassembling file...")
 
-	reassembleFile(reassembledFile, toDownloadTracker, partFileNameFn)
+	if err = reassembleFile(reassembledFile, toDownloadTracker, partFileNameFn); err != nil {
+		return err
+	}
 
 	log.Info("Reassembled file")
 
 	if compress {
-		if err = decompressFile(reassembledFile, filePath); err != nil {
+		if err = decompressFile(reassembledFile, filePath, encoding); err != nil {
 			return err
 		}
 	}
@@ -185,25 +199,3 @@ func reassembleFile(mainFile *os.File, chunks map[Chunk]bool, partFileNameFn fun
 	}
 	return nil
 }
-
-
-func decompressFile(mainFile *os.File, filePath string) error {
-	mainFile.Seek(0, 0)
-
-	gr, err := gzip.NewReader(mainFile)
-	if err != nil {
-		return err
-	}
-	defer gr.Close()
-	log.Info("Decompressing file...")
-	fout, err := os.Create(filePath)
-	if err != nil {
-		return err
-	}
-	defer fout.Close()
-	if _, err = io.Copy(fout, gr); err != nil {
-		return err
-	}
-	log.Info("Decompressed file")
-	return nil
-}
diff --git a/go.mod b/go.mod
index fdb4dcf..155bc1d 100644
--- a/go.mod
+++ b/go.mod
@@ -2,13 +2,13 @@ module github.com/sansmoraxz/godm
 
 go 1.21.3
 
-require golang.org/x/sync v0.7.0
-
 require golang.org/x/sys v0.0.0-20220715151400-c0bba94af5f8 // indirect
 
 require (
+	github.com/dsnet/compress v0.0.1
 	github.com/inconshreveable/mousetrap v1.1.0 // indirect
+	github.com/klauspost/compress v1.17.9
 	github.com/sirupsen/logrus v1.9.3
-	github.com/spf13/cobra v1.8.0 // indirect
+	github.com/spf13/cobra v1.8.0
 	github.com/spf13/pflag v1.0.5 // indirect
 )
diff --git a/go.sum b/go.sum
index 4273234..e1dbe10 100644
--- a/go.sum
+++ b/go.sum
@@ -1,8 +1,17 @@
 github.com/cpuguy83/go-md2man/v2 v2.0.3/go.mod h1:tgQtvFlXSQOSOSIRvRPT7W67SCa46tRHOmNcaadrF8o=
 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/dsnet/compress v0.0.1 h1:PlZu0n3Tuv04TzpfPbrnI0HW/YwodEXDS+oPKahKF0Q=
+github.com/dsnet/compress v0.0.1/go.mod h1:Aw8dCMJ7RioblQeTqt88akK31OvO8Dhf5JflhBbQEHo=
+github.com/dsnet/golib v0.0.0-20171103203638-1ea166775780/go.mod h1:Lj+Z9rebOhdfkVLjJ8T6VcRQv3SXugXy999NBtR9aFY=
 github.com/inconshreveable/mousetrap v1.1.0 h1:wN+x4NVGpMsO7ErUn/mUI3vEoE6Jt13X2s0bqwp9tc8=
 github.com/inconshreveable/mousetrap v1.1.0/go.mod h1:vpF70FUmC8bwa3OWnCshd2FqLfsEA9PFc4w1p2J65bw=
+github.com/klauspost/compress v1.4.1/go.mod h1:RyIbtBH6LamlWaDj8nUwkbUhJ87Yi3uG0guNDohfE1A=
+github.com/klauspost/compress v1.17.9 h1:6KIumPrER1LHsvBVuDa0r5xaG0Es51mhhB9BQB2qeMA=
+github.com/klauspost/compress v1.17.9/go.mod h1:Di0epgTjJY877eYKx5yC51cX2A2Vl2ibi7bDH9ttBbw=
+github.com/klauspost/cpuid v1.2.0/go.mod h1:Pj4uuM528wm8OyEC2QMXAi2YiTZ96dNQPGgoMS4s3ek=
+github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
 github.com/russross/blackfriday/v2 v2.1.0/go.mod h1:+Rmxgy9KzJVeS9/2gXHxylqXiyQDYRxCVz55jmeOWTM=
 github.com/sirupsen/logrus v1.9.3 h1:dueUQJ1C2q9oE3F7wvmSGAaVtTmUizReu6fjN8uqzbQ=
@@ -12,11 +21,12 @@ github.com/spf13/cobra v1.8.0/go.mod h1:WXLWApfZ71AjXPya3WOlMsY9yMs7YeiHhFVlvLyh
 github.com/spf13/pflag v1.0.5 h1:iy+VFUOCP1a+8yFto/drg2CJ5u0yRoB7fZw3DKv/JXA=
 github.com/spf13/pflag v1.0.5/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg=
 github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
+github.com/stretchr/testify v1.7.0 h1:nwc3DEeHmmLAfoZucVR881uASk0Mfjw8xYJ99tb5CcY=
 github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
-golang.org/x/sync v0.7.0 h1:YsImfSBoP9QPYL0xyKJPq0gcaJdG3rInoqxTWbfQu9M=
-golang.org/x/sync v0.7.0/go.mod h1:Czt+wKu1gCyEFDUtn0jG5QVvpJ6rzVqr5aXyt9drQfk=
+github.com/ulikunitz/xz v0.5.6/go.mod h1:2bypXElzHzzJZwzH67Y6wb67pO62Rzfn7BSiF4ABRW8=
 golang.org/x/sys v0.0.0-20220715151400-c0bba94af5f8 h1:0A+M6Uqn+Eje4kHMK80dtF3JCXC4ykBgQG4Fe06QRhQ=
 golang.org/x/sys v0.0.0-20220715151400-c0bba94af5f8/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
+gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
 gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
diff --git a/utils.go b/utils.go
index 85df659..b317390 100644
--- a/utils.go
+++ b/utils.go
@@ -1,11 +1,13 @@
 package godm
 
-import "net/http"
+import (
+	"net/http"
+	"strings"
+)
 
 func getHeaders(client *http.Client, url string, compress bool) (http.Header, error) {
 	// head request to get metadata
 	// Note: the uncompressed payload is considered for Content-Length
-	// Use Accept-Encoding: gzip, deflate, br to get compressed payload size
 
 	req, err := http.NewRequest("HEAD", url, nil)
 	if err != nil {
@@ -13,7 +15,7 @@ func getHeaders(client *http.Client, url string, compress bool) (http.Header, er
 	}
 
 	if compress {
-		req.Header.Set("Accept-Encoding", "gzip, deflate, br")
+		req.Header.Set("Accept-Encoding", strings.Join(supportedEncodings, ", "))
 	}
 	resp, err := client.Do(req)
 	if err != nil {