From 365f49e6d3abbb7f6411f9f5c2a3613b5cbd5754 Mon Sep 17 00:00:00 2001
From: z4yx <z4yx@users.noreply.github.com>
Date: Mon, 16 Mar 2020 20:59:08 +0800
Subject: [PATCH] add support of env config for rsync provider

---
 worker/provider.go                 |  2 ++
 worker/rsync_provider.go           | 16 ++++++++--------
 worker/two_stage_rsync_provider.go | 17 ++++++++---------
 3 files changed, 18 insertions(+), 17 deletions(-)

diff --git a/worker/provider.go b/worker/provider.go
index 9024d82..00c333d 100644
--- a/worker/provider.go
+++ b/worker/provider.go
@@ -135,6 +135,7 @@ func newMirrorProvider(mirror mirrorConfig, cfg *Config) mirrorProvider {
 			excludeFile:       mirror.ExcludeFile,
 			extraOptions:      mirror.RsyncOptions,
 			overriddenOptions: mirror.RsyncOverride,
+			rsyncEnv:          mirror.Env,
 			workingDir:        mirrorDir,
 			logDir:            logDir,
 			logFile:           filepath.Join(logDir, "latest.log"),
@@ -159,6 +160,7 @@ func newMirrorProvider(mirror mirrorConfig, cfg *Config) mirrorProvider {
 			password:      mirror.Password,
 			excludeFile:   mirror.ExcludeFile,
 			extraOptions:  mirror.RsyncOptions,
+			rsyncEnv:      mirror.Env,
 			workingDir:    mirrorDir,
 			logDir:        logDir,
 			logFile:       filepath.Join(logDir, "latest.log"),
diff --git a/worker/rsync_provider.go b/worker/rsync_provider.go
index cc06ee8..96c5316 100644
--- a/worker/rsync_provider.go
+++ b/worker/rsync_provider.go
@@ -15,6 +15,7 @@ type rsyncConfig struct {
 	upstreamURL, username, password, excludeFile string
 	extraOptions                                 []string
 	overriddenOptions                            []string
+	rsyncEnv                                     map[string]string
 	workingDir, logDir, logFile                  string
 	useIPv6, useIPv4                             bool
 	interval                                     time.Duration
@@ -50,6 +51,12 @@ func newRsyncProvider(c rsyncConfig) (*rsyncProvider, error) {
 	if c.rsyncCmd == "" {
 		provider.rsyncCmd = "rsync"
 	}
+	if c.username != "" {
+		c.rsyncEnv["USER"] = c.username
+	}
+	if c.password != "" {
+		c.rsyncEnv["RSYNC_PASSWORD"] = c.password
+	}
 
 	options := []string{
 		"-aHvh", "--no-o", "--no-g", "--stats",
@@ -116,18 +123,11 @@ func (p *rsyncProvider) Start() error {
 		return errors.New("provider is currently running")
 	}
 
-	env := map[string]string{}
-	if p.username != "" {
-		env["USER"] = p.username
-	}
-	if p.password != "" {
-		env["RSYNC_PASSWORD"] = p.password
-	}
 	command := []string{p.rsyncCmd}
 	command = append(command, p.options...)
 	command = append(command, p.upstreamURL, p.WorkingDir())
 
-	p.cmd = newCmdJob(p, command, p.WorkingDir(), env)
+	p.cmd = newCmdJob(p, command, p.WorkingDir(), p.rsyncEnv)
 	if err := p.prepareLogFile(false); err != nil {
 		return err
 	}
diff --git a/worker/two_stage_rsync_provider.go b/worker/two_stage_rsync_provider.go
index e865490..bbade4b 100644
--- a/worker/two_stage_rsync_provider.go
+++ b/worker/two_stage_rsync_provider.go
@@ -16,6 +16,7 @@ type twoStageRsyncConfig struct {
 	stage1Profile                                string
 	upstreamURL, username, password, excludeFile string
 	extraOptions                                 []string
+	rsyncEnv                                     map[string]string
 	workingDir, logDir, logFile                  string
 	useIPv6                                      bool
 	interval                                     time.Duration
@@ -69,6 +70,12 @@ func newTwoStageRsyncProvider(c twoStageRsyncConfig) (*twoStageRsyncProvider, er
 		},
 	}
 
+	if c.username != "" {
+		c.rsyncEnv["USER"] = c.username
+	}
+	if c.password != "" {
+		c.rsyncEnv["RSYNC_PASSWORD"] = c.password
+	}
 	if c.rsyncCmd == "" {
 		provider.rsyncCmd = "rsync"
 	}
@@ -132,14 +139,6 @@ func (p *twoStageRsyncProvider) Run() error {
 		return errors.New("provider is currently running")
 	}
 
-	env := map[string]string{}
-	if p.username != "" {
-		env["USER"] = p.username
-	}
-	if p.password != "" {
-		env["RSYNC_PASSWORD"] = p.password
-	}
-
 	p.dataSize = ""
 	stages := []int{1, 2}
 	for _, stage := range stages {
@@ -151,7 +150,7 @@ func (p *twoStageRsyncProvider) Run() error {
 		command = append(command, options...)
 		command = append(command, p.upstreamURL, p.WorkingDir())
 
-		p.cmd = newCmdJob(p, command, p.WorkingDir(), env)
+		p.cmd = newCmdJob(p, command, p.WorkingDir(), p.rsyncEnv)
 		if err := p.prepareLogFile(stage > 1); err != nil {
 			return err
 		}