From cc562c7f5806de265de49c97c8cd772d7cbf7e30 Mon Sep 17 00:00:00 2001
From: xzchaoo <xzchaoo@gmail.com>
Date: Mon, 18 Mar 2024 16:58:31 +0800
Subject: [PATCH] fix: wrap exec with timeout and _FROM env

---
 pkg/cri/impl/engine/containerd_engine.go |  8 +++---
 pkg/cri/impl/engine/docker_engine.go     |  8 +++---
 pkg/cri/impl/engine/utils.go             | 32 ++++++++++++++++++++++++
 3 files changed, 40 insertions(+), 8 deletions(-)
 create mode 100644 pkg/cri/impl/engine/utils.go

diff --git a/pkg/cri/impl/engine/containerd_engine.go b/pkg/cri/impl/engine/containerd_engine.go
index 7a4dc99..3e681e1 100644
--- a/pkg/cri/impl/engine/containerd_engine.go
+++ b/pkg/cri/impl/engine/containerd_engine.go
@@ -283,13 +283,13 @@ func (e *ContainerdContainerEngine) Exec(ctx context.Context, c *cri.Container,
 
 	pspec := spec.Process
 	pspec.Terminal = false
-	pspec.Args = req.Cmd
+	pspec.Args = wrapTimeout(req.Cmd)
 	if req.WorkingDir != "" {
 		pspec.Cwd = req.WorkingDir
 	}
 
 	// Append user specified env
-	pspec.Env = append(pspec.Env, req.Env...)
+	pspec.Env = wrapEnv(append(pspec.Env, req.Env...))
 
 	task, err := container.Task(ctx, nil)
 	if err != nil {
@@ -424,13 +424,13 @@ func (e *ContainerdContainerEngine) ExecAsync(ctx context.Context, c *cri.Contai
 
 	pspec := spec.Process
 	pspec.Terminal = false
-	pspec.Args = req.Cmd
+	pspec.Args = wrapTimeout(req.Cmd)
 	if req.WorkingDir != "" {
 		pspec.Cwd = req.WorkingDir
 	}
 
 	// Append user specified env
-	pspec.Env = append(pspec.Env, req.Env...)
+	pspec.Env = wrapEnv(append(pspec.Env, req.Env...))
 
 	task, err := container.Task(ctx, nil)
 	if err != nil {
diff --git a/pkg/cri/impl/engine/docker_engine.go b/pkg/cri/impl/engine/docker_engine.go
index e7a373c..bab04c6 100644
--- a/pkg/cri/impl/engine/docker_engine.go
+++ b/pkg/cri/impl/engine/docker_engine.go
@@ -131,9 +131,9 @@ func (e *DockerContainerEngine) Exec(ctx context.Context, c *cri.Container, req
 		AttachStdout: true,
 		Detach:       false,
 		DetachKeys:   "",
-		Env:          req.Env,
+		Env:          wrapEnv(req.Env),
 		WorkingDir:   req.WorkingDir,
-		Cmd:          req.Cmd,
+		Cmd:          wrapTimeout(req.Cmd),
 	})
 	if err != nil {
 		return invalidResult, err
@@ -239,9 +239,9 @@ func (e *DockerContainerEngine) ExecAsync(ctx context.Context, c *cri.Container,
 		AttachStdout: true,
 		Detach:       false,
 		DetachKeys:   "",
-		Env:          req.Env,
+		Env:          wrapEnv(req.Env),
 		WorkingDir:   req.WorkingDir,
-		Cmd:          hackedCmd,
+		Cmd:          wrapTimeout(hackedCmd),
 	})
 	if err != nil {
 		return invalidResult, err
diff --git a/pkg/cri/impl/engine/utils.go b/pkg/cri/impl/engine/utils.go
new file mode 100644
index 0000000..f9ab966
--- /dev/null
+++ b/pkg/cri/impl/engine/utils.go
@@ -0,0 +1,32 @@
+/*
+ * Copyright 2022 Holoinsight Project Authors. Licensed under Apache-2.0.
+ */
+
+package engine
+
+import (
+	"github.com/spf13/cast"
+	"os"
+)
+
+var (
+	timeout = "180"
+)
+
+func init() {
+	s := os.Getenv("CRI_EXEC_TIMEOUT")
+	if x := cast.ToInt(s); x > 0 {
+		timeout = s
+	}
+}
+
+// wrapTimeout wraps cmd with timeout -s KILL <seconds> to prevent the process from hanging and not exiting for any reason.
+func wrapTimeout(cmd []string) []string {
+	// timeout -s KILL <seconds> cmd...
+	return append([]string{"timeout", "-s", "KILL", timeout}, cmd...)
+}
+
+// wrapEnv wraps envs with _FROM=holoinsight-agent. This env is used to mark the source of the call.
+func wrapEnv(envs []string) []string {
+	return append(envs, "_FROM=holoinsight-agent")
+}