From 3306a1f277cd7d6affbf4b439c7a9cbf541c75b6 Mon Sep 17 00:00:00 2001
From: Saw-jan <saw.jan.grg3e@gmail.com>
Date: Fri, 5 Jul 2024 14:55:03 +0545
Subject: [PATCH] timeout command execution

---
 tests/ociswrapper/ocis/ocis.go                | 17 ++++++--
 tests/ociswrapper/wrapper/handlers/handler.go | 43 +++++++++----------
 2 files changed, 34 insertions(+), 26 deletions(-)

diff --git a/tests/ociswrapper/ocis/ocis.go b/tests/ociswrapper/ocis/ocis.go
index 84d8d309f6..32d51602ed 100644
--- a/tests/ociswrapper/ocis/ocis.go
+++ b/tests/ociswrapper/ocis/ocis.go
@@ -2,6 +2,7 @@ package ocis
 
 import (
 	"bufio"
+	"context"
 	"crypto/tls"
 	"fmt"
 	"net/http"
@@ -219,7 +220,17 @@ func waitUntilCompleteShutdown() (bool, string) {
 }
 
 func RunCommand(command string) (int, string) {
-	c := exec.Command(config.Get("bin"), command)
-	out, _ := c.CombinedOutput()
-	return c.ProcessState.ExitCode(), string(out)
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+
+	c := exec.CommandContext(ctx, config.Get("bin"), command)
+	output, err := c.CombinedOutput()
+	if err != nil {
+		if ctx.Err() == context.DeadlineExceeded {
+			message := "Command timed out:\n" + string(output)
+			return c.ProcessState.ExitCode(), message
+		}
+	}
+
+	return c.ProcessState.ExitCode(), string(output)
 }
diff --git a/tests/ociswrapper/wrapper/handlers/handler.go b/tests/ociswrapper/wrapper/handlers/handler.go
index 2531612afa..308bc23437 100644
--- a/tests/ociswrapper/wrapper/handlers/handler.go
+++ b/tests/ociswrapper/wrapper/handlers/handler.go
@@ -52,18 +52,20 @@ func sendResponse(res http.ResponseWriter, success bool, message string) {
 }
 
 func sendCmdResponse(res http.ResponseWriter, exitCode int, message string) {
-	var resBody CommandResponse
+	resBody := CommandResponse{
+		BasicResponse: &BasicResponse{
+			Message: message,
+		},
+		ExitCode: exitCode,
+	}
+
+	if exitCode == 0 {
+		resBody.BasicResponse.Status = "OK"
+	} else {
+		resBody.BasicResponse.Status = "ERROR"
+	}
 
 	res.WriteHeader(http.StatusOK)
-	if exitCode == 0 {
-		resBody.Status = "OK"
-		resBody.ExitCode = exitCode
-		resBody.Message = message
-	} else {
-		resBody.Status = "ERROR"
-		resBody.ExitCode = exitCode
-		resBody.Message = message
-	}
 	res.Header().Set("Content-Type", "application/json")
 
 	jsonResponse, _ := json.Marshal(resBody)
@@ -128,23 +130,18 @@ func CommandHandler(res http.ResponseWriter, req *http.Request) {
 		return
 	}
 
-	body, err := parseJsonBody(req.Body)
+	if req.Body == nil {
+		http.Error(res, "Bad request", http.StatusBadRequest)
+		return
+	}
+
+	body, err := io.ReadAll(req.Body)
 	if err != nil {
 		http.Error(res, "Bad request", http.StatusBadRequest)
 		return
 	}
-	if body["command"] == nil {
-		http.Error(res, "Bad request", http.StatusBadRequest)
-		return
-	}
 
-	command, ok := body["command"].(string)
-	if !ok || command == "" {
-		http.Error(res, "Bad request", http.StatusBadRequest)
-		return
-	}
+	exitCode, output := ocis.RunCommand(string(body))
 
-	exitCode, out := ocis.RunCommand(command)
-
-	sendCmdResponse(res, exitCode, out)
+	sendCmdResponse(res, exitCode, output)
 }