From b20e89cd72125bae991b2a277b60f44165d266d4 Mon Sep 17 00:00:00 2001
From: dnomd343 <i@343.re>
Date: Sat, 27 Aug 2022 17:44:38 +0800
Subject: [PATCH] feat: download asset via proxy

---
 cmd/asset/update.go  | 11 ++++++-----
 cmd/common/file.go   | 33 ++++++++++++++++++++++++---------
 cmd/config/decode.go |  8 ++++++++
 cmd/xproxy.go        |  1 -
 4 files changed, 38 insertions(+), 15 deletions(-)

diff --git a/cmd/asset/update.go b/cmd/asset/update.go
index 437efc1..55aca79 100644
--- a/cmd/asset/update.go
+++ b/cmd/asset/update.go
@@ -8,11 +8,12 @@ import (
 )
 
 type Config struct {
-    Cron string            `yaml:"cron" json:"cron"`
-    Url  map[string]string `yaml:"url" json:"url"`
+    Proxy string            `yaml:"proxy" json:"proxy"`
+    Cron  string            `yaml:"cron" json:"cron"`
+    Url   map[string]string `yaml:"url" json:"url"`
 }
 
-func updateAsset(urls map[string]string, assetDir string) { // download new assets
+func updateAsset(urls map[string]string, assetDir string, updateProxy string) { // download new assets
     defer func() {
         if err := recover(); err != nil {
             log.Errorf("Update failed -> %v", err)
@@ -21,7 +22,7 @@ func updateAsset(urls map[string]string, assetDir string) { // download new asse
     if len(urls) != 0 {
         log.Info("Start update assets")
         for file, url := range urls {
-            common.DownloadFile(url, path.Join(assetDir, file)) // maybe override old asset
+            common.DownloadFile(url, path.Join(assetDir, file), updateProxy) // maybe override old asset
         }
     }
 }
@@ -30,7 +31,7 @@ func AutoUpdate(update *Config, assetDir string) { // set cron task for auto upd
     if update.Cron != "" {
         autoUpdate := cron.New()
         _ = autoUpdate.AddFunc(update.Cron, func() { // cron function
-            updateAsset(update.Url, assetDir)
+            updateAsset(update.Url, assetDir, update.Proxy)
         })
         autoUpdate.Start()
     }
diff --git a/cmd/common/file.go b/cmd/common/file.go
index 92744a6..06590cd 100644
--- a/cmd/common/file.go
+++ b/cmd/common/file.go
@@ -5,6 +5,7 @@ import (
     "io"
     "io/ioutil"
     "net/http"
+    "net/url"
     "os"
     "strings"
 )
@@ -69,21 +70,35 @@ func CopyFile(source string, target string) {
     }
 }
 
-func DownloadFile(url string, file string) {
-    log.Debugf("File download `%s` => `%s`", url, file)
-    resp, err := http.Get(url)
-    defer resp.Body.Close()
+func DownloadFile(fileUrl string, filePath string, proxyUrl string) {
+    log.Debugf("File download `%s` => `%s`", fileUrl, filePath)
+    client := http.Client{}
+    if proxyUrl != "" { // use proxy for download
+        log.Infof("File download via proxy -> %s", proxyUrl)
+        rawUrl, _ := url.Parse(proxyUrl)
+        client = http.Client{
+            Transport: &http.Transport{
+                Proxy: http.ProxyURL(rawUrl),
+            },
+        }
+    }
+    resp, err := client.Get(fileUrl)
+    defer func() {
+        if resp != nil {
+            resp.Body.Close()
+        }
+    }()
     if err != nil {
-        log.Errorf("Download `%s` error -> %v", url, err)
+        log.Errorf("Download `%s` error -> %v", fileUrl, err)
         return
     }
-    output, err := os.OpenFile(file, os.O_WRONLY|os.O_TRUNC|os.O_CREATE, 0644)
+    output, err := os.OpenFile(filePath, os.O_WRONLY|os.O_TRUNC|os.O_CREATE, 0644)
     defer output.Close()
     if err != nil {
-        log.Panicf("Open `%s` error -> %v", file, err)
+        log.Panicf("Open `%s` error -> %v", filePath, err)
     }
     if _, err = io.Copy(output, resp.Body); err != nil {
-        log.Panicf("File `%s` save error -> %v", file, err)
+        log.Panicf("File `%s` save error -> %v", filePath, err)
     }
-    log.Infof("Download success `%s` => `%s`", url, file)
+    log.Infof("Download success `%s` => `%s`", fileUrl, filePath)
 }
diff --git a/cmd/config/decode.go b/cmd/config/decode.go
index 1692575..9b75f5c 100644
--- a/cmd/config/decode.go
+++ b/cmd/config/decode.go
@@ -8,6 +8,7 @@ import (
     "encoding/json"
     log "github.com/sirupsen/logrus"
     "gopkg.in/yaml.v3"
+    "net/url"
 )
 
 type NetConfig struct {
@@ -146,6 +147,13 @@ func decodeRadvd(rawConfig *RawConfig, config *Config) {
 
 func decodeUpdate(rawConfig *RawConfig, config *Config) {
     config.Update = rawConfig.Update
+    if config.Update.Proxy != "" {
+        _, err := url.Parse(config.Update.Proxy) // check proxy info
+        if err != nil {
+            log.Panicf("Invalid update proxy -> %s", config.Update.Proxy)
+        }
+    }
+    log.Debugf("Update proxy -> %s", config.Update.Proxy)
     log.Debugf("Update cron -> %s", config.Update.Cron)
     log.Debugf("Update urls -> %v", config.Update.Url)
 }
diff --git a/cmd/xproxy.go b/cmd/xproxy.go
index 9844c67..c8fa252 100644
--- a/cmd/xproxy.go
+++ b/cmd/xproxy.go
@@ -90,7 +90,6 @@ func main() {
     config.Load(configFile, &settings)
     loadNetwork(&settings)
     loadProxy(&settings)
-    // TODO: update assets via proxy
     loadAsset(&settings)
     loadRadvd(&settings)