From bb2d989bc51680a29e8c7ebb1ea014ae647b32ac Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Maximilian=20Pa=C3=9F?=
 <22845248+mpass99@users.noreply.github.com>
Date: Thu, 21 Apr 2022 19:43:17 +0200
Subject: [PATCH] Add the missing newline of AWS.

---
 .../src/main/java/poseidon/App.java           | 18 ++++---
 .../src/test/java/poseidon/AppTest.java       | 50 +++++++++++++++++--
 2 files changed, 58 insertions(+), 10 deletions(-)

diff --git a/deploy/aws/java11Exec/src/main/java/poseidon/App.java b/deploy/aws/java11Exec/src/main/java/poseidon/App.java
index fa3237f..acb4fa7 100644
--- a/deploy/aws/java11Exec/src/main/java/poseidon/App.java
+++ b/deploy/aws/java11Exec/src/main/java/poseidon/App.java
@@ -19,6 +19,7 @@ import java.nio.charset.StandardCharsets;
 import java.nio.file.Files;
 import java.util.Base64;
 import java.util.Map;
+import java.util.NoSuchElementException;
 import java.util.Scanner;
 
 // AwsFunctionRequest contains the java files that needs to be executed.
@@ -153,15 +154,14 @@ public class App implements RequestHandler<APIGatewayV2WebSocketEvent, APIGatewa
     private void scanForOutput(Process p, InputStream stream, WebSocketMessageType type) {
         Scanner outputScanner = new Scanner(stream);
         while (p.isAlive() || outputScanner.hasNextLine()) {
-            this.sendMessage(type, outputScanner.nextLine(), null);
+            try {
+                this.sendMessage(type, outputScanner.nextLine() + "\n", null);
+            } catch (NoSuchElementException ignored) {}
         }
     }
 
     // sendMessage sends WebSocketMessage objects back to the requester of this Lambda function.
     private void sendMessage(WebSocketMessageType type, String data, Integer exitCode) {
-        if (this.disableOutput) {
-            return;
-        }
         JsonObject msg = new JsonObject();
         msg.addProperty("type", type.toString());
         if (type == WebSocketMessageType.WebSocketExit) {
@@ -170,8 +170,12 @@ public class App implements RequestHandler<APIGatewayV2WebSocketEvent, APIGatewa
             msg.addProperty("data", data);
         }
 
-        this.gwClient.postToConnection(new PostToConnectionRequest()
-                .withConnectionId(this.connectionID)
-                .withData(ByteBuffer.wrap(gson.toJson(msg).getBytes(StandardCharsets.UTF_8))));
+        if (this.disableOutput) {
+            System.out.println(gson.toJson(msg));
+        } else {
+            this.gwClient.postToConnection(new PostToConnectionRequest()
+                    .withConnectionId(this.connectionID)
+                    .withData(ByteBuffer.wrap(gson.toJson(msg).getBytes(StandardCharsets.UTF_8))));
+        }
     }
 }
diff --git a/deploy/aws/java11Exec/src/test/java/poseidon/AppTest.java b/deploy/aws/java11Exec/src/test/java/poseidon/AppTest.java
index 42a0de1..5c4f5bb 100644
--- a/deploy/aws/java11Exec/src/test/java/poseidon/AppTest.java
+++ b/deploy/aws/java11Exec/src/test/java/poseidon/AppTest.java
@@ -4,6 +4,8 @@ import com.amazonaws.services.lambda.runtime.events.APIGatewayProxyResponseEvent
 import com.amazonaws.services.lambda.runtime.events.APIGatewayV2WebSocketEvent;
 import org.junit.Test;
 
+import java.io.ByteArrayOutputStream;
+import java.io.PrintStream;
 import java.nio.charset.StandardCharsets;
 import java.util.Base64;
 import java.util.HashMap;
@@ -28,8 +30,51 @@ public class AppTest {
           "    }\n" +
           "}").getBytes(StandardCharsets.UTF_8));
 
+  static final String MultilineMathContent = Base64.getEncoder().encodeToString(
+          ("package org.example;\n" +
+          "\n" +
+          "public class RecursiveMath {\n" +
+          "\n" +
+          "    public static void main(String[] args) {\n" +
+          "        System.out.println(\"Mein Text\");\n" +
+          "        System.out.println(\"Mein Text\");\n" +
+          "    }\n" +
+          "}").getBytes(StandardCharsets.UTF_8));
+
   @Test
   public void successfulResponse() {
+    APIGatewayProxyResponseEvent result = getApiGatewayProxyResponse(RecursiveMathContent);
+    assertEquals(200, result.getStatusCode().intValue());
+  }
+
+  // Scanner.nextLine() consumes the new line. Therefore, we need to add it again.
+  @Test
+  public void successfulMultilineResponse() {
+    ByteArrayOutputStream out = setupStdOutLogs();
+    APIGatewayProxyResponseEvent result = getApiGatewayProxyResponse(MultilineMathContent);
+    restoreStdOutLogs();
+
+    assertEquals(200, result.getStatusCode().intValue());
+    String expectedOutput =
+            "{\"type\":\"stdout\",\"data\":\"Mein Text\\n\"}\n" +
+            "{\"type\":\"stdout\",\"data\":\"Mein Text\\n\"}\n" +
+            "{\"type\":\"exit\",\"data\":0}\n";
+    assertEquals(expectedOutput, out.toString());
+  }
+
+  private PrintStream originalOut;
+  private ByteArrayOutputStream setupStdOutLogs() {
+    originalOut = System.out;
+    ByteArrayOutputStream out = new ByteArrayOutputStream();
+    System.setOut(new PrintStream(out));
+    return out;
+  }
+
+  private void restoreStdOutLogs() {
+    System.setOut(originalOut);
+  }
+
+  private APIGatewayProxyResponseEvent getApiGatewayProxyResponse(String content) {
     App app = new App();
     APIGatewayV2WebSocketEvent input = new APIGatewayV2WebSocketEvent();
     APIGatewayV2WebSocketEvent.RequestContext ctx = new APIGatewayV2WebSocketEvent.RequestContext();
@@ -40,8 +85,7 @@ public class AppTest {
     headers.put(App.disableOutputHeaderKey, "True");
     input.setHeaders(headers);
     input.setBody("{\"action\":\"java11Exec\",\"cmd\":[\"sh\",\"-c\",\"javac org/example/RecursiveMath.java && java org/example/RecursiveMath\"]," +
-            "\"files\":{\"org/example/RecursiveMath.java\":\"" + RecursiveMathContent + "\"}}");
-    APIGatewayProxyResponseEvent result = app.handleRequest(input, null);
-    assertEquals(200, result.getStatusCode().intValue());
+            "\"files\":{\"org/example/RecursiveMath.java\":\"" + content + "\"}}");
+    return app.handleRequest(input, null);
   }
 }