From 033aa60540de29d1275eaad5b21e9b3353670ea9 Mon Sep 17 00:00:00 2001
From: Shengqi Chen <harry-chen@outlook.com>
Date: Fri, 28 Feb 2025 14:43:51 +0800
Subject: [PATCH] Implement mirror.success_exit_codes and
 global.dangerous_global_success_exit_codes

Signed-off-by: Harry Chen <i@harrychen.xyz>
---
 worker/base_provider.go | 22 +++++++++++++++++++---
 worker/config.go        |  6 ++++++
 worker/provider.go      | 16 ++++++++++++++++
 worker/runner.go        | 14 ++++++++++++--
 4 files changed, 53 insertions(+), 5 deletions(-)

diff --git a/worker/base_provider.go b/worker/base_provider.go
index 5c690cf..cb18eb7 100644
--- a/worker/base_provider.go
+++ b/worker/base_provider.go
@@ -19,9 +19,10 @@ type baseProvider struct {
 	timeout  time.Duration
 	isMaster bool
 
-	cmd       *cmdJob
-	logFileFd *os.File
-	isRunning atomic.Value
+	cmd              *cmdJob
+	logFileFd        *os.File
+	isRunning        atomic.Value
+	successExitCodes []int
 
 	cgroup *cgroupHook
 	zfs    *zfsHook
@@ -186,3 +187,18 @@ func (p *baseProvider) Terminate() error {
 func (p *baseProvider) DataSize() string {
 	return ""
 }
+
+func (p *baseProvider) SetSuccessExitCodes(codes []int) {
+	if codes == nil {
+		p.successExitCodes = []int{}
+	} else {
+		p.successExitCodes = codes
+	}
+}
+
+func (p *baseProvider) GetSuccessExitCodes() []int {
+	if p.successExitCodes == nil {
+		return []int{}
+	}
+	return p.successExitCodes
+}
diff --git a/worker/config.go b/worker/config.go
index 152f056..dc91c79 100644
--- a/worker/config.go
+++ b/worker/config.go
@@ -63,6 +63,9 @@ type globalConfig struct {
 
 	ExecOnSuccess []string `toml:"exec_on_success"`
 	ExecOnFailure []string `toml:"exec_on_failure"`
+
+	// merged with mirror-specific options. make sure you know what you are doing!
+	SuccessExitCodes []int `toml:"dangerous_global_success_exit_codes"`
 }
 
 type managerConfig struct {
@@ -169,6 +172,9 @@ type mirrorConfig struct {
 	ExecOnSuccessExtra []string `toml:"exec_on_success_extra"`
 	ExecOnFailureExtra []string `toml:"exec_on_failure_extra"`
 
+	// will be merged with global option
+	SuccessExitCodes []int `toml:"success_exit_codes"`
+
 	Command           string   `toml:"command"`
 	FailOnMatch       string   `toml:"fail_on_match"`
 	SizePattern       string   `toml:"size_pattern"`
diff --git a/worker/provider.go b/worker/provider.go
index d45b0ed..dada7e3 100644
--- a/worker/provider.go
+++ b/worker/provider.go
@@ -60,6 +60,10 @@ type mirrorProvider interface {
 	ExitContext() *Context
 	// return context
 	Context() *Context
+
+	// set in newMirrorProvider, used by cmdJob.Wait
+	SetSuccessExitCodes(codes []int)
+	GetSuccessExitCodes() []int
 }
 
 // newProvider creates a mirrorProvider instance
@@ -249,5 +253,17 @@ func newMirrorProvider(mirror mirrorConfig, cfg *Config) mirrorProvider {
 	}
 	addHookFromCmdList(mirror.ExecOnFailureExtra, execOnFailure)
 
+	successExitCodes := []int{}
+	if cfg.Global.SuccessExitCodes != nil {
+		successExitCodes = append(successExitCodes, cfg.Global.SuccessExitCodes...)
+	}
+	if mirror.SuccessExitCodes != nil {
+		successExitCodes = append(successExitCodes, mirror.SuccessExitCodes...)
+	}
+	if len(successExitCodes) > 0 {
+		logger.Infof("Non-zero success exit codes set for mirror %s: %v", mirror.Name, successExitCodes)
+		provider.SetSuccessExitCodes(successExitCodes)
+	}
+
 	return provider
 }
diff --git a/worker/runner.go b/worker/runner.go
index 50f0378..ebf2d74 100644
--- a/worker/runner.go
+++ b/worker/runner.go
@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"os"
 	"os/exec"
+	"slices"
 	"strings"
 	"sync"
 	"syscall"
@@ -171,9 +172,18 @@ func (c *cmdJob) Wait() error {
 		return c.retErr
 	default:
 		err := c.cmd.Wait()
-		c.retErr = err
 		close(c.finished)
-		return err
+		if err != nil {
+			code := err.(*exec.ExitError).ExitCode()
+			allowedCodes := c.provider.GetSuccessExitCodes()
+			if slices.Contains(allowedCodes, code) {
+				// process exited with non-success status
+				logger.Infof("Command %s exited with code %d: treated as success (allowed: %v)", c.cmd.Args, code, allowedCodes)
+			} else {
+				c.retErr = err
+			}
+		}
+		return c.retErr
 	}
 }