Merge pull request #2 from wim-web/dev

Dev

internal/handler/exec.go

@@ -27,12 +27,19 @@ func ExecHandler(cmd string) error {
 		return err
 	}
 
-	return command.ExecCommand(
+	exeCmd, err := command.ExecCommand(
+		context.Background(),
 		ecsService,
 		cluster,
 		*task.TaskArn,
 		cmd,
 		container.Name,
 		cfg.Region,
 	)
+
+	if err != nil {
+		return err
+	}
+
+	return exeCmd.Run()
 }

internal/handler/portforward.go

@@ -32,7 +32,8 @@ func PortforwardHandler(doc command.DocumentName, params map[string][]string) er
 
 	taskId := strings.Split(*task.TaskArn, "/")[2]
 
-	return command.PortForwardCommand(
+	cmd, err := command.PortForwardCommand(
+		context.Background(),
 		ssmService,
 		cluster,
 		taskId,
@@ -41,4 +42,10 @@ func PortforwardHandler(doc command.DocumentName, params map[string][]string) er
 		doc,
 		params,
 	)
+
+	if err != nil {
+		return err
+	}
+
+	return cmd.Run()
 }

internal/session_manager/session.go

@@ -1,14 +1,18 @@
 package session_manager
 
-import "os/exec"
+import (
+	"context"
+	"os/exec"
+)
 
 const SESSION_MANAGER_COMMAND = "session-manager-plugin"
 
-func MakeStartSessionCmd(response string, region string) *exec.Cmd {
+func MakeStartSessionCmd(ctx context.Context, response string, region string) *exec.Cmd {
 	const OperationName = "StartSession"
 
 	// https://github.com/aws/session-manager-plugin/blob/1.2.463.0/src/sessionmanagerplugin/session/session.go#L163-L178
-	return exec.Command(
+	return exec.CommandContext(
+		ctx,
 		SESSION_MANAGER_COMMAND,
 		response,
 		region,

pkg/command/exec.go

@@ -4,13 +4,14 @@ import (
 	"context"
 	"encoding/json"
 	"os"
+	"os/exec"
 
 	"github.com/aws/aws-sdk-go-v2/service/ecs"
 	"github.com/aws/aws-sdk-go/aws"
 	"github.com/wim-web/tonneeeeel/internal/session_manager"
 )
 
-func ExecCommand(c *ecs.Client, cluster string, task string, command string, container *string, region string) error {
+func ExecCommand(ctx context.Context, c *ecs.Client, cluster string, task string, command string, container *string, region string) (*exec.Cmd, error) {
 	input := &ecs.ExecuteCommandInput{
 		Cluster:     aws.String(cluster),
 		Task:        aws.String(task),
@@ -22,20 +23,20 @@ func ExecCommand(c *ecs.Client, cluster string, task string, command string, con
 	res, err := c.ExecuteCommand(context.Background(), input)
 
 	if err != nil {
-		return err
+		return nil, err
 	}
 
 	r, err := json.Marshal(res.Session)
 
 	if err != nil {
-		return err
+		return nil, err
 	}
 
-	cmd := session_manager.MakeStartSessionCmd(string(r), region)
+	cmd := session_manager.MakeStartSessionCmd(ctx, string(r), region)
 
 	cmd.Stdout = os.Stdout
 	cmd.Stdin = os.Stdin
 	cmd.Stderr = os.Stderr
 
-	return cmd.Run()
+	return cmd, nil
 }

pkg/command/portforward.go

@@ -5,6 +5,7 @@ import (
 	"encoding/json"
 	"fmt"
 	"os"
+	"os/exec"
 
 	"github.com/aws/aws-sdk-go-v2/service/ssm"
 	"github.com/aws/aws-sdk-go/aws"
@@ -18,7 +19,7 @@ const (
 	REMOTE_PORT_FORWARD_DOCUMENT_NAME DocumentName = "AWS-StartPortForwardingSessionToRemoteHost"
 )
 
-func PortForwardCommand(c *ssm.Client, cluster string, taskId string, containerId string, region string, doc DocumentName, params map[string][]string) error {
+func PortForwardCommand(ctx context.Context, c *ssm.Client, cluster string, taskId string, containerId string, region string, doc DocumentName, params map[string][]string) (*exec.Cmd, error) {
 	input := &ssm.StartSessionInput{
 		Target:       aws.String(fmt.Sprintf("ecs:%s_%s_%s", cluster, taskId, containerId)),
 		DocumentName: aws.String(string(doc)),
@@ -28,20 +29,20 @@ func PortForwardCommand(c *ssm.Client, cluster string, taskId string, containerI
 	res, err := c.StartSession(context.Background(), input)
 
 	if err != nil {
-		return err
+		return nil, err
 	}
 
 	r, err := json.Marshal(res)
 
 	if err != nil {
-		return err
+		return nil, err
 	}
 
-	cmd := session_manager.MakeStartSessionCmd(string(r), region)
+	cmd := session_manager.MakeStartSessionCmd(ctx, string(r), region)
 
 	cmd.Stdout = os.Stdout
 	cmd.Stdin = os.Stdin
 	cmd.Stderr = os.Stderr
 
-	return cmd.Run()
+	return cmd, nil
 }