From 5d0457df90f46fc0c9f397d016bef9ece287499f Mon Sep 17 00:00:00 2001
From: Yasin Turan <turyasin@amazon.com>
Date: Tue, 15 Nov 2022 17:02:04 +0000
Subject: [PATCH] Refactored ctr to allow for tty allocation when container has
 {terminal:true} within Spec

Signed-off-by: Yasin Turan <turyasin@amazon.com>
---
 cmd/ctr/commands/containers/restore.go | 54 ++++++++++++++++++++++++--
 1 file changed, 50 insertions(+), 4 deletions(-)

diff --git a/cmd/ctr/commands/containers/restore.go b/cmd/ctr/commands/containers/restore.go
index 2847340c9..639451b9a 100644
--- a/cmd/ctr/commands/containers/restore.go
+++ b/cmd/ctr/commands/containers/restore.go
@@ -19,10 +19,13 @@ package containers
 import (
 	"errors"
 
+	"github.com/containerd/console"
 	"github.com/containerd/containerd"
 	"github.com/containerd/containerd/cio"
 	"github.com/containerd/containerd/cmd/ctr/commands"
+	"github.com/containerd/containerd/cmd/ctr/commands/tasks"
 	"github.com/containerd/containerd/errdefs"
+	"github.com/sirupsen/logrus"
 	"github.com/urfave/cli"
 )
 
@@ -81,17 +84,60 @@ var restoreCommand = cli.Command{
 		if err != nil {
 			return err
 		}
-
 		topts := []containerd.NewTaskOpts{}
 		if context.Bool("live") {
 			topts = append(topts, containerd.WithTaskCheckpoint(checkpoint))
 		}
-
-		task, err := ctr.NewTask(ctx, cio.NewCreator(cio.WithStdio), topts...)
+		spec, err := ctr.Spec(ctx)
 		if err != nil {
 			return err
 		}
 
-		return task.Start(ctx)
+		useTTY := spec.Process.Terminal
+
+		var con console.Console
+		if useTTY {
+			con = console.Current()
+			defer con.Reset()
+			if err := con.SetRaw(); err != nil {
+				return err
+			}
+		}
+
+		task, err := tasks.NewTask(ctx, client, ctr, "", con, false, "", []cio.Opt{}, topts...)
+		if err != nil {
+			return err
+		}
+
+		var statusC <-chan containerd.ExitStatus
+		if useTTY {
+			if statusC, err = task.Wait(ctx); err != nil {
+				return err
+			}
+		}
+
+		if err := task.Start(ctx); err != nil {
+			return err
+		}
+		if !useTTY {
+			return nil
+		}
+
+		if err := tasks.HandleConsoleResize(ctx, task, con); err != nil {
+			logrus.WithError(err).Error("console resize")
+		}
+
+		status := <-statusC
+		code, _, err := status.Result()
+		if err != nil {
+			return err
+		}
+		if _, err := task.Delete(ctx); err != nil {
+			return err
+		}
+		if code != 0 {
+			return cli.NewExitError("", int(code))
+		}
+		return nil
 	},
 }