From 5c472279f3352f5873b6403a0f5b608c41ccf7b1 Mon Sep 17 00:00:00 2001
From: Dnomd343 <i@343.re>
Date: Sat, 1 Apr 2023 19:32:05 +0800
Subject: [PATCH] feat: support gzip and brotli compress download

---
 cmd/common/file.go | 66 +++++++++++++++++++++++++++++++++++++---------
 go.mod             |  3 +++
 go.sum             |  6 +++++
 3 files changed, 63 insertions(+), 12 deletions(-)

diff --git a/cmd/common/file.go b/cmd/common/file.go
index 10af7b8..6d458af 100644
--- a/cmd/common/file.go
+++ b/cmd/common/file.go
@@ -1,6 +1,10 @@
 package common
 
 import (
+    "github.com/andybalholm/brotli"
+    "github.com/go-http-utils/headers"
+    "github.com/klauspost/compress/flate"
+    "github.com/klauspost/compress/gzip"
     log "github.com/sirupsen/logrus"
     "io"
     "io/ioutil"
@@ -74,11 +78,12 @@ func CopyFile(source string, target string) {
     }
 }
 
-func DownloadFile(fileUrl string, filePath string, proxyUrl string) bool {
-    log.Debugf("File download `%s` => `%s`", fileUrl, filePath)
+func DownloadBytes(fileUrl string, proxyUrl string) ([]byte, error) {
     client := http.Client{}
-    if proxyUrl != "" { // use proxy for download
-        log.Infof("File download via proxy -> %s", proxyUrl)
+    if proxyUrl == "" {
+        log.Infof("Download `%s` without proxy", fileUrl)
+    } else { // use proxy for download
+        log.Infof("Download `%s` via `%s`", fileUrl, proxyUrl)
         rawUrl, _ := url.Parse(proxyUrl)
         client = http.Client{
             Transport: &http.Transport{
@@ -86,22 +91,59 @@ func DownloadFile(fileUrl string, filePath string, proxyUrl string) bool {
             },
         }
     }
-    resp, err := client.Get(fileUrl)
-    defer func() {
-        if resp != nil {
-            resp.Body.Close()
+
+    req, err := http.NewRequest("GET", fileUrl, nil)
+    if err != nil {
+        log.Errorf("Failed to create http request")
+        return nil, err
+    }
+    req.Header.Set(headers.AcceptEncoding, "gzip, deflate, br")
+    resp, err := client.Do(req)
+    if err != nil {
+        log.Errorf("Failed to execute http GET request")
+        return nil, err
+    }
+    if resp != nil {
+        defer resp.Body.Close()
+        log.Debugf("Remote data downloaded successfully")
+    }
+
+    switch resp.Header.Get(headers.ContentEncoding) {
+    case "br":
+        log.Debugf("Downloaded content using brolti encoding")
+        return io.ReadAll(brotli.NewReader(resp.Body))
+    case "gzip":
+        log.Debugf("Downloaded content using gzip encoding")
+        gr, err := gzip.NewReader(resp.Body)
+        if err != nil {
+            return nil, err
         }
-    }()
+        return io.ReadAll(gr)
+    case "deflate":
+        log.Debugf("Downloaded content using deflate encoding")
+        zr := flate.NewReader(resp.Body)
+        defer zr.Close()
+        return io.ReadAll(zr)
+    default:
+        return io.ReadAll(resp.Body)
+    }
+}
+
+func DownloadFile(fileUrl string, filePath string, proxyUrl string) bool {
+    log.Debugf("File download `%s` => `%s`", fileUrl, filePath)
+    data, err := DownloadBytes(fileUrl, proxyUrl)
     if err != nil {
         log.Errorf("Download `%s` error -> %v", fileUrl, err)
         return false
     }
-    output, err := os.OpenFile(filePath, os.O_WRONLY|os.O_TRUNC|os.O_CREATE, 0644)
-    defer output.Close()
+
+    file, err := os.OpenFile(filePath, os.O_WRONLY|os.O_TRUNC|os.O_CREATE, 0644)
+    defer file.Close()
     if err != nil {
         log.Panicf("Open `%s` error -> %v", filePath, err)
+        return false
     }
-    if _, err = io.Copy(output, resp.Body); err != nil {
+    if _, err = file.Write(data); err != nil {
         log.Errorf("File `%s` save error -> %v", filePath, err)
         return false
     }
diff --git a/go.mod b/go.mod
index b399b92..9f9c266 100644
--- a/go.mod
+++ b/go.mod
@@ -4,6 +4,9 @@ go 1.20
 
 require (
 	github.com/BurntSushi/toml v1.2.1
+	github.com/andybalholm/brotli v1.0.5
+	github.com/go-http-utils/headers v0.0.0-20181008091004-fed159eddc2a
+	github.com/klauspost/compress v1.16.3
 	github.com/robfig/cron v1.2.0
 	github.com/sirupsen/logrus v1.9.0
 	gopkg.in/yaml.v3 v3.0.1
diff --git a/go.sum b/go.sum
index 6dd06ed..5b0ebd2 100644
--- a/go.sum
+++ b/go.sum
@@ -1,8 +1,14 @@
 github.com/BurntSushi/toml v1.2.1 h1:9F2/+DoOYIOksmaJFPw1tGFy1eDnIJXg+UHjuD8lTak=
 github.com/BurntSushi/toml v1.2.1/go.mod h1:CxXYINrC8qIiEnFrOxCa7Jy5BFHlXnUU2pbicEuybxQ=
+github.com/andybalholm/brotli v1.0.5 h1:8uQZIdzKmjc/iuPu7O2ioW48L81FgatrcpfFmiq/cCs=
+github.com/andybalholm/brotli v1.0.5/go.mod h1:fO7iG3H7G2nSZ7m0zPUDn85XEX2GTukHGRSepvi9Eig=
 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/go-http-utils/headers v0.0.0-20181008091004-fed159eddc2a h1:v6zMvHuY9yue4+QkG/HQ/W67wvtQmWJ4SDo9aK/GIno=
+github.com/go-http-utils/headers v0.0.0-20181008091004-fed159eddc2a/go.mod h1:I79BieaU4fxrw4LMXby6q5OS9XnoR9UIKLOzDFjUmuw=
+github.com/klauspost/compress v1.16.3 h1:XuJt9zzcnaz6a16/OU53ZjWp/v7/42WcR5t2a0PcNQY=
+github.com/klauspost/compress v1.16.3/go.mod h1:ntbaceVETuRiXiv4DpjP66DpAtAGkEQskQzEyD//IeE=
 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/robfig/cron v1.2.0 h1:ZjScXvvxeQ63Dbyxy76Fj3AT3Ut0aKsyd2/tl3DTMuQ=