From d650bae0b6212b5a38aef62a6fab13c04b971f45 Mon Sep 17 00:00:00 2001
From: Jason Song <i@wolfogre.com>
Date: Wed, 2 Nov 2022 12:12:45 +0800
Subject: [PATCH] feat: replace ciphertext in log

---
 runtime/reporter.go | 32 ++++++++++++++++++++++----------
 runtime/task.go     |  2 +-
 2 files changed, 23 insertions(+), 11 deletions(-)

diff --git a/runtime/reporter.go b/runtime/reporter.go
index 1fb014a..2d2ecde 100644
--- a/runtime/reporter.go
+++ b/runtime/reporter.go
@@ -25,19 +25,29 @@ type Reporter struct {
 	client  client.Client
 	clientM sync.Mutex
 
-	logOffset int
-	logRows   []*runnerv1.LogRow
-	state     *runnerv1.TaskState
-	stateM    sync.RWMutex
+	logOffset   int
+	logRows     []*runnerv1.LogRow
+	logReplacer *strings.Replacer
+	state       *runnerv1.TaskState
+	stateM      sync.RWMutex
 }
 
-func NewReporter(ctx context.Context, cancel context.CancelFunc, client client.Client, taskID int64) *Reporter {
+func NewReporter(ctx context.Context, cancel context.CancelFunc, client client.Client, task *runnerv1.Task) *Reporter {
+	var oldnew []string
+	if v := task.Context.Fields["token"].GetStringValue(); v != "" {
+		oldnew = append(oldnew, v, "***")
+	}
+	for _, v := range task.Secrets {
+		oldnew = append(oldnew, v, "***")
+	}
+
 	return &Reporter{
-		ctx:    ctx,
-		cancel: cancel,
-		client: client,
+		ctx:         ctx,
+		cancel:      cancel,
+		client:      client,
+		logReplacer: strings.NewReplacer(oldnew...),
 		state: &runnerv1.TaskState{
-			Id: taskID,
+			Id: task.Id,
 		},
 	}
 }
@@ -266,8 +276,10 @@ func (r *Reporter) parseResult(result interface{}) (runnerv1.Result, bool) {
 }
 
 func (r *Reporter) parseLogRow(entry *log.Entry) *runnerv1.LogRow {
+	content := strings.TrimSuffix(entry.Message, "\r\n")
+	content = r.logReplacer.Replace(content)
 	return &runnerv1.LogRow{
 		Time:    timestamppb.New(entry.Time),
-		Content: strings.TrimSuffix(entry.Message, "\r\n"),
+		Content: content,
 	}
 }
diff --git a/runtime/task.go b/runtime/task.go
index d1588ee..d0b5894 100644
--- a/runtime/task.go
+++ b/runtime/task.go
@@ -131,7 +131,7 @@ func (t *Task) Run(ctx context.Context, task *runnerv1.Task) error {
 	defer globalTaskMap.Delete(task.Id)
 
 	lastWords := ""
-	reporter := NewReporter(ctx, cancel, t.client, task.Id)
+	reporter := NewReporter(ctx, cancel, t.client, task)
 	defer func() {
 		_ = reporter.Close(lastWords)
 	}()