From 903ad415c63f368fb7f91b06f44634bede79319b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Maximilian=20Pa=C3=9F?=
 <22845248+mpass99@users.noreply.github.com>
Date: Tue, 24 May 2022 23:48:39 +0200
Subject: [PATCH] Set and unset aws environment variables

in order to provide the CodeOcean context for some scripts and hide AWS credentials from the users.
---
 internal/runner/aws_runner.go      | 10 ++++++++++
 internal/runner/aws_runner_test.go |  5 +++--
 2 files changed, 13 insertions(+), 2 deletions(-)

diff --git a/internal/runner/aws_runner.go b/internal/runner/aws_runner.go
index 037fdb5..a1895a2 100644
--- a/internal/runner/aws_runner.go
+++ b/internal/runner/aws_runner.go
@@ -84,6 +84,7 @@ func (w *AWSFunctionWorkload) ExecuteInteractively(id string, _ io.ReadWriter, s
 	if !ok {
 		return nil, nil, ErrorUnknownExecution
 	}
+	w.prepareEnvironmentVariables(request)
 	command, ctx, cancel := prepareExecution(request)
 	exitInternal := make(chan ExitInfo)
 	exit := make(chan ExitInfo, 1)
@@ -208,3 +209,12 @@ func (w *AWSFunctionWorkload) handleRunnerTimeout(ctx context.Context,
 		exit <- ExitInfo{255, ErrorRunnerInactivityTimeout}
 	}
 }
+
+// prepareEnvironmentVariables sets the CODEOCEAN variable and unsets some AWS variables.
+func (w *AWSFunctionWorkload) prepareEnvironmentVariables(request *dto.ExecutionRequest) {
+	if request.Environment == nil {
+		request.Environment = make(map[string]string)
+	}
+	request.Environment["CODEOCEAN"] = "true"
+	request.Command = "unset \"${!AWS@}\" && " + request.Command
+}
diff --git a/internal/runner/aws_runner_test.go b/internal/runner/aws_runner_test.go
index 1954421..09955e7 100644
--- a/internal/runner/aws_runner_test.go
+++ b/internal/runner/aws_runner_test.go
@@ -92,7 +92,8 @@ func TestAWSFunctionWorkload_ExecuteInteractively(t *testing.T) {
 		cancel()
 
 		expectedRequestData := "{\"action\":\"" + environment.Image() +
-			"\",\"cmd\":[\"env\",\"sh\",\"-c\",\"" + command + "\"],\"files\":{}}"
+			"\",\"cmd\":[\"env\",\"CODEOCEAN=true\",\"sh\",\"-c\",\"unset \\\"${!AWS@}\\\" \\u0026\\u0026 " + command +
+			"\"],\"files\":{}}"
 		assert.Equal(t, expectedRequestData, awsMock.receivedData)
 	})
 }
@@ -124,7 +125,7 @@ func TestAWSFunctionWorkload_UpdateFileSystem(t *testing.T) {
 	execCancel()
 
 	expectedRequestData := "{\"action\":\"" + environment.Image() +
-		"\",\"cmd\":[\"env\",\"sh\",\"-c\",\"" + command + "\"]," +
+		"\",\"cmd\":[\"env\",\"CODEOCEAN=true\",\"sh\",\"-c\",\"unset \\\"${!AWS@}\\\" \\u0026\\u0026 " + command + "\"]," +
 		"\"files\":{\"" + string(myFile.Path) + "\":\"" + base64.StdEncoding.EncodeToString(myFile.Content) + "\"}}"
 	assert.Equal(t, expectedRequestData, awsMock.receivedData)
 }