From 59ca63268b231ceda8d35404907fd6bd25e19dce Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Maximilian=20Pa=C3=9F?=
 <22845248+mpass99@users.noreply.github.com>
Date: Thu, 9 Jun 2022 23:26:36 +0200
Subject: [PATCH] Add CODEOCEAN environment variable.

---
 internal/runner/aws_runner.go | 9 ++++-----
 pkg/dto/dto.go                | 6 ++++++
 2 files changed, 10 insertions(+), 5 deletions(-)

diff --git a/internal/runner/aws_runner.go b/internal/runner/aws_runner.go
index a1895a2..0400924 100644
--- a/internal/runner/aws_runner.go
+++ b/internal/runner/aws_runner.go
@@ -84,7 +84,7 @@ func (w *AWSFunctionWorkload) ExecuteInteractively(id string, _ io.ReadWriter, s
 	if !ok {
 		return nil, nil, ErrorUnknownExecution
 	}
-	w.prepareEnvironmentVariables(request)
+	hideEnvironmentVariables(request, "AWS")
 	command, ctx, cancel := prepareExecution(request)
 	exitInternal := make(chan ExitInfo)
 	exit := make(chan ExitInfo, 1)
@@ -210,11 +210,10 @@ func (w *AWSFunctionWorkload) handleRunnerTimeout(ctx context.Context,
 	}
 }
 
-// prepareEnvironmentVariables sets the CODEOCEAN variable and unsets some AWS variables.
-func (w *AWSFunctionWorkload) prepareEnvironmentVariables(request *dto.ExecutionRequest) {
+// hideEnvironmentVariables sets the CODEOCEAN variable and unsets all variables starting with the passed prefix.
+func hideEnvironmentVariables(request *dto.ExecutionRequest, unsetPrefix string) {
 	if request.Environment == nil {
 		request.Environment = make(map[string]string)
 	}
-	request.Environment["CODEOCEAN"] = "true"
-	request.Command = "unset \"${!AWS@}\" && " + request.Command
+	request.Command = "unset \"${!" + unsetPrefix + "@}\" && " + request.Command
 }
diff --git a/pkg/dto/dto.go b/pkg/dto/dto.go
index 35d85ba..d3371c2 100644
--- a/pkg/dto/dto.go
+++ b/pkg/dto/dto.go
@@ -25,6 +25,12 @@ type ExecutionRequest struct {
 func (er *ExecutionRequest) FullCommand() []string {
 	command := make([]string, 0)
 	command = append(command, "env")
+
+	if er.Environment == nil {
+		er.Environment = make(map[string]string)
+	}
+	er.Environment["CODEOCEAN"] = "true"
+
 	for variable, value := range er.Environment {
 		command = append(command, fmt.Sprintf("%s=%s", variable, value))
 	}