ECS Exec と scp で送受信

recv.go

package main

import (
	"bufio"
	"context"
	"encoding/json"
	"flag"
	"fmt"
	"log"
	"os"
	"os/exec"
	"os/signal"
	"syscall"

	"github.com/aws/aws-sdk-go-v2/aws"
	"github.com/aws/aws-sdk-go-v2/config"
	"github.com/aws/aws-sdk-go-v2/service/ecs"
)

const SessionManagerPluginBinary = "session-manager-plugin"

func main() {
	if err := run(context.Background()); err != nil {
		log.Fatal(err)
	}
}

func run(ctx context.Context) error {
	var (
		cluster       = flag.String("cluster", "", "cluster")
		containerName = flag.String("container-name", "", "container name")
		taskID        = flag.String("task-id", "", "task id")
	)
	flag.Parse()

	if flag.NArg() != 1 {
		fmt.Printf("Usage: %s [options] [TARGET_FILE]\n", os.Args[0])
		flag.PrintDefaults()
		return nil
	}

	conf, err := config.LoadDefaultConfig(ctx)
	if err != nil {
		return err
	}
	svc := ecs.NewFromConfig(conf)
	task, err := getTask(ctx, svc, *cluster, *containerName, *taskID)
	if err != nil {
		return err
	}

	out, err := svc.ExecuteCommand(ctx, &ecs.ExecuteCommandInput{
		Cluster:     aws.String(task.clusterArn),
		Interactive: true,
		Task:        aws.String(task.taskArn),
		Command:     aws.String(fmt.Sprintf(`scpf -q %s`, flag.Arg(0))),
		Container:   aws.String(*containerName),
	})
	if err != nil {
		return err
	}
	sess, _ := json.Marshal(out.Session)

	cmd := exec.Command(SessionManagerPluginBinary, string(sess), conf.Region, "StartSession", "", task.targetJson, fmt.Sprintf("ecs.%s.amazonaws.com", conf.Region))
	r, _ := cmd.StdoutPipe()
	defer r.Close()
	signal.Ignore(os.Interrupt, syscall.SIGPIPE)
	defer signal.Reset(os.Interrupt, syscall.SIGPIPE)
	cmd.Start()

	// なぜか 改行が \r\n になってやってくるので bufio.Scanner で \n にしつつ読む
	scanner := bufio.NewScanner(r)

	// session-manager-plugin が標準出力でログ吐くので読み飛ばす
	//   標準出力を読んでゴニョるツールを作るのに邪魔なので session-manager-plugin のログは標準エラーにしないか？というPRは作られているが無視されていて悲しい
	//   https://github.com/aws/session-manager-plugin/pull/20
	scanner.Scan() // \n
	scanner.Scan() // Starting session with SessionId: ecs-execute-command-XXXXXXXXXXXXXXXXX

	for scanner.Scan() {
		fmt.Println(scanner.Text())
	}
	return cmd.Wait()
}

type task struct {
	clusterArn string
	taskArn    string
	target     string
	targetJson string
}

func getTask(ctx context.Context, svc *ecs.Client, cluster, containerName, taskID string) (*task, error) {
	res, err := svc.DescribeTasks(ctx, &ecs.DescribeTasksInput{
		Cluster: aws.String(cluster),
		Tasks:   []string{taskID},
	})
	if err != nil {
		return nil, fmt.Errorf("error describe tasks: %w", err)
	}

	var clusterArn, taskArn, runtimeID *string
	for _, t := range res.Tasks {
		for _, c := range t.Containers {
			if aws.ToString(c.Name) == containerName {
				runtimeID = c.RuntimeId
				taskArn = c.TaskArn
				clusterArn = t.ClusterArn
			}
		}
	}
	if runtimeID == nil {
		return nil, fmt.Errorf("Could not identify runtimeId")
	}

	target := fmt.Sprintf("ecs:%s_%s_%s", cluster, taskID, aws.ToString(runtimeID))

	targetJson, _ := json.Marshal(struct {
		Target string
	}{
		Target: target,
	})

	return &task{
		clusterArn: aws.ToString(clusterArn),
		taskArn:    aws.ToString(taskArn),
		target:     target,
		targetJson: string(targetJson),
	}, nil
}

scpf.go

package main

import (
	"fmt"
	"log"
	"os"
	"os/exec"
	"strings"
)

func main() {
	args := append([]string{"-rf"}, os.Args[1:]...)
	cmd := exec.Command("scp", args...)
	w, err := cmd.StdinPipe()
	if err != nil {
		log.Fatal(err)
	}
	cmd.Stderr = os.Stderr
	cmd.Stdout = os.Stdout
	fmt.Fprint(w, strings.Repeat("\x00", 7))
	if err := cmd.Run(); err != nil {
		log.Fatal(err)
	}
}

send.go

package main

import (
	"context"
	"encoding/json"
	"flag"
	"fmt"
	"log"
	"os"
	"os/exec"
	"os/signal"

	"github.com/aws/aws-sdk-go-v2/aws"
	"github.com/aws/aws-sdk-go-v2/config"
	"github.com/aws/aws-sdk-go-v2/service/ecs"
)

const SessionManagerPluginBinary = "session-manager-plugin"

func main() {
	if err := run(context.Background()); err != nil {
		log.Fatal(err)
	}
}

func run(ctx context.Context) error {
	var (
		cluster       = flag.String("cluster", "", "cluster")
		containerName = flag.String("container-name", "", "container name")
		taskID        = flag.String("task-id", "", "task id")
	)
	flag.Parse()

	if flag.NArg() != 2 {
		fmt.Printf("Usage: %s [options] [SRC] [DST]\n", os.Args[0])
		flag.PrintDefaults()
		return nil
	}

	stat, err := os.Stat(flag.Arg(0))
	if err != nil {
		return err
	}
	f, err := os.ReadFile(flag.Arg(0))
	if err != nil {
		return err
	}

	conf, err := config.LoadDefaultConfig(ctx)
	if err != nil {
		return err
	}
	svc := ecs.NewFromConfig(conf)
	task, err := getTask(ctx, svc, *cluster, *containerName, *taskID)
	if err != nil {
		return err
	}

	out, err := svc.ExecuteCommand(ctx, &ecs.ExecuteCommandInput{
		Cluster:     aws.String(task.clusterArn),
		Interactive: true,
		Task:        aws.String(task.taskArn),
		Command:     aws.String(fmt.Sprintf(`scp -qrt %s`, flag.Arg(1))),
		Container:   aws.String(*containerName),
	})
	if err != nil {
		return err
	}
	sess, _ := json.Marshal(out.Session)

	cmd := exec.Command(SessionManagerPluginBinary, string(sess), conf.Region, "StartSession", "", task.targetJson, fmt.Sprintf("ecs.%s.amazonaws.com", conf.Region))

	w, err := cmd.StdinPipe()
	if err != nil {
		return err
	}
	signal.Ignore(os.Interrupt)
	defer signal.Reset(os.Interrupt)
	cmd.Stderr = os.Stderr
	cmd.Stdout = os.Stdout
	cmd.Start()

	fmt.Fprintln(w, fmt.Sprintf("C%04o", stat.Mode()), len(f), stat.Name())
	w.Write(f)
	fmt.Fprint(w, "\x00")
	fmt.Fprint(w, "\n") // -t で待ち受けてる scp を直接 kill することが出来ないので改行ぶち込んで自死してもらう

	return cmd.Wait()
}

type task struct {
	clusterArn string
	taskArn    string
	target     string
	targetJson string
}

func getTask(ctx context.Context, svc *ecs.Client, cluster, containerName, taskID string) (*task, error) {
	res, err := svc.DescribeTasks(ctx, &ecs.DescribeTasksInput{
		Cluster: aws.String(cluster),
		Tasks:   []string{taskID},
	})
	if err != nil {
		return nil, fmt.Errorf("error describe tasks: %w", err)
	}

	var clusterArn, taskArn, runtimeID *string
	for _, t := range res.Tasks {
		for _, c := range t.Containers {
			if aws.ToString(c.Name) == containerName {
				runtimeID = c.RuntimeId
				taskArn = c.TaskArn
				clusterArn = t.ClusterArn
			}
		}
	}
	if runtimeID == nil {
		return nil, fmt.Errorf("Could not identify runtimeId")
	}

	target := fmt.Sprintf("ecs:%s_%s_%s", cluster, taskID, aws.ToString(runtimeID))

	targetJson, _ := json.Marshal(struct {
		Target string
	}{
		Target: target,
	})

	return &task{
		clusterArn: aws.ToString(clusterArn),
		taskArn:    aws.ToString(taskArn),
		target:     target,
		targetJson: string(targetJson),
	}, nil
}