From 9979e949b708228b76e026833779acbc113a04f3 Mon Sep 17 00:00:00 2001 From: zhangyue Date: Mon, 17 Sep 2018 18:35:39 +0800 Subject: [PATCH] feature: add cli attach command Signed-off-by: zhangyue --- cli/attach.go | 119 ++++++++++++++++++++++++++++++++++++++++ cli/main.go | 1 + pkg/ioutils/readers.go | 25 +++++++++ pkg/term/ascii.go | 66 ++++++++++++++++++++++ pkg/term/ascii_test.go | 24 ++++++++ pkg/term/proxy.go | 81 +++++++++++++++++++++++++++ pkg/term/proxy_test.go | 114 ++++++++++++++++++++++++++++++++++++++ test/cli_attach_test.go | 64 +++++++++++++++++++++ 8 files changed, 494 insertions(+) create mode 100644 cli/attach.go create mode 100644 pkg/ioutils/readers.go create mode 100644 pkg/term/ascii.go create mode 100644 pkg/term/ascii_test.go create mode 100644 pkg/term/proxy.go create mode 100644 pkg/term/proxy_test.go create mode 100644 test/cli_attach_test.go diff --git a/cli/attach.go b/cli/attach.go new file mode 100644 index 0000000000..f7a7302e12 --- /dev/null +++ b/cli/attach.go @@ -0,0 +1,119 @@ +package main + +import ( + "context" + "fmt" + "io" + "os" + + "github.com/spf13/cobra" + + "github.com/alibaba/pouch/pkg/ioutils" + "github.com/alibaba/pouch/pkg/term" + "golang.org/x/crypto/ssh/terminal" +) + +// AttachDescription is used to describe attach command in detail and auto generate command doc. +var AttachDescription = "Attach local standard input, output, and error streams to a running container" + +// AttachCommand is used to implement 'attach' command. +type AttachCommand struct { + baseCommand + + // flags for attach command + noStdin bool + detachKeys string +} + +// Init initialize "attach" command. +func (ac *AttachCommand) Init(c *Cli) { + ac.cli = c + ac.cmd = &cobra.Command{ + Use: "attach [OPTIONS] CONTAINER", + Short: "Attach local standard input, output, and error streams to a running container", + Long: AttachDescription, + Args: cobra.ExactArgs(1), + RunE: func(cmd *cobra.Command, args []string) error { + return ac.runAttach(args) + }, + Example: ac.example(), + } + ac.addFlags() +} + +// addFlags adds flags for specific command. +func (ac *AttachCommand) addFlags() { + flagSet := ac.cmd.Flags() + flagSet.BoolVar(&ac.noStdin, "no-stdin", false, "Do not attach STDIN") + flagSet.StringVar(&ac.detachKeys, "detach-keys", "", "Override the key sequence for detaching a container") +} + +// runAttach is used to attach a container. +func (ac *AttachCommand) runAttach(args []string) error { + name := args[0] + + ctx := context.Background() + apiClient := ac.cli.Client() + + if terminal.IsTerminal(int(os.Stdin.Fd())) { + in, out, err := setRawMode(!ac.noStdin, false) + if err != nil { + fmt.Fprintf(os.Stderr, "failed to set raw mode %s", err) + return fmt.Errorf("failed to set raw mode") + } + defer func() { + if err := restoreMode(in, out); err != nil { + fmt.Fprintf(os.Stderr, "failed to restore term mode %s", err) + } + }() + } else { + return fmt.Errorf("can't use attach command in not terminal environment") + } + + conn, br, err := apiClient.ContainerAttach(ctx, name, !ac.noStdin) + if err != nil { + return fmt.Errorf("failed to attach container: %v", err) + } + defer conn.Close() + + var inReader io.Reader = os.Stdin + escapeKeys := term.DefaultEscapeKeys + // Wrap the input to detect detach escape sequence. + // Use default escape keys if an invalid sequence is given. + if ac.detachKeys != "" { + customEscapeKeys, err := term.ToBytes(ac.detachKeys) + if err != nil { + return fmt.Errorf("Invalid detach keys (%s) provided", ac.detachKeys) + } + escapeKeys = customEscapeKeys + } + inReader = ioutils.NewReadCloserWrapper(term.NewEscapeProxy(os.Stdin, escapeKeys), os.Stdin.Close) + + wait := make(chan struct{}) + go func() { + io.Copy(os.Stdout, br) + wait <- struct{}{} + }() + + go func() { + if !ac.noStdin { + _, err := io.Copy(conn, inReader) + if _, ok := err.(term.EscapeError); ok { + wait <- struct{}{} + } + conn.Close() + } + + }() + + <-wait + return nil +} + +// example shows examples in attach command, and is used in auto-generated cli docs. +func (ac *AttachCommand) example() string { + return `$ pouch ps +Name ID Status Image Runtime +foo 71b9c1 Running docker.io/library/busybox:latest runc +$ pouch attach foo` +} diff --git a/cli/main.go b/cli/main.go index ec71110c82..3989e6fa01 100644 --- a/cli/main.go +++ b/cli/main.go @@ -17,6 +17,7 @@ func main() { cli.AddCommand(base, &PullCommand{}) cli.AddCommand(base, &CreateCommand{}) cli.AddCommand(base, &StartCommand{}) + cli.AddCommand(base, &AttachCommand{}) cli.AddCommand(base, &StopCommand{}) cli.AddCommand(base, &PsCommand{}) cli.AddCommand(base, &RmCommand{}) diff --git a/pkg/ioutils/readers.go b/pkg/ioutils/readers.go new file mode 100644 index 0000000000..91ae3039a6 --- /dev/null +++ b/pkg/ioutils/readers.go @@ -0,0 +1,25 @@ +package ioutils + +import "io" + +// ReadCloserWrapper wraps an io.Reader, and implements an io.ReadCloser +// It calls the given callback function when closed. It should be constructed +// with NewReadCloserWrapper +type ReadCloserWrapper struct { + io.Reader + closer func() error +} + +// NewReadCloserWrapper returns a new io.ReadCloser. +func NewReadCloserWrapper(r io.Reader, closer func() error) io.ReadCloser { + return &ReadCloserWrapper{ + Reader: r, + closer: closer, + } + +} + +// Close calls back the passed closer function +func (r *ReadCloserWrapper) Close() error { + return r.closer() +} diff --git a/pkg/term/ascii.go b/pkg/term/ascii.go new file mode 100644 index 0000000000..55873c0556 --- /dev/null +++ b/pkg/term/ascii.go @@ -0,0 +1,66 @@ +package term + +import ( + "fmt" + "strings" +) + +// ASCII list the possible supported ASCII key sequence +var ASCII = []string{ + "ctrl-@", + "ctrl-a", + "ctrl-b", + "ctrl-c", + "ctrl-d", + "ctrl-e", + "ctrl-f", + "ctrl-g", + "ctrl-h", + "ctrl-i", + "ctrl-j", + "ctrl-k", + "ctrl-l", + "ctrl-m", + "ctrl-n", + "ctrl-o", + "ctrl-p", + "ctrl-q", + "ctrl-r", + "ctrl-s", + "ctrl-t", + "ctrl-u", + "ctrl-v", + "ctrl-w", + "ctrl-x", + "ctrl-y", + "ctrl-z", + "ctrl-[", + "ctrl-\\", + "ctrl-]", + "ctrl-^", + "ctrl-_", +} + +// ToBytes converts a string representing a suite of key-sequence to the corresponding ASCII code. +func ToBytes(keys string) ([]byte, error) { + codes := []byte{} +next: + for _, key := range strings.Split(keys, ",") { + if len(key) != 1 { + for code, ctrl := range ASCII { + if ctrl == key { + codes = append(codes, byte(code)) + continue next + } + } + if key == "DEL" { + codes = append(codes, 127) + } else { + return nil, fmt.Errorf("Unknown character: '%s'", key) + } + } else { + codes = append(codes, key[0]) + } + } + return codes, nil +} diff --git a/pkg/term/ascii_test.go b/pkg/term/ascii_test.go new file mode 100644 index 0000000000..890a72560a --- /dev/null +++ b/pkg/term/ascii_test.go @@ -0,0 +1,24 @@ +package term + +import ( + "testing" + + "github.com/stretchr/testify/assert" +) + +func TestToBytes(t *testing.T) { + codes, err := ToBytes("ctrl-a,a") + assert.NoError(t, err) + assert.Equal(t, []byte{1, 97}, codes) + + _, err = ToBytes("shift-z") + assert.Error(t, err) + + codes, err = ToBytes("ctrl-@,ctrl-[,~,ctrl-o") + assert.NoError(t, err) + assert.Equal(t, []byte{0, 27, 126, 15}, codes) + + codes, err = ToBytes("DEL,+") + assert.NoError(t, err) + assert.Equal(t, []byte{127, 43}, codes) +} diff --git a/pkg/term/proxy.go b/pkg/term/proxy.go new file mode 100644 index 0000000000..106d148596 --- /dev/null +++ b/pkg/term/proxy.go @@ -0,0 +1,81 @@ +package term + +import ( + "io" +) + +// DefaultEscapeKeys represents the default escape key sequence: ctrl-p, ctrl-q +var DefaultEscapeKeys = []byte{16, 17} + +// EscapeError is special error which returned by a TTY proxy reader's Read() +// method in case its detach escape sequence is read. +type EscapeError struct{} + +func (EscapeError) Error() string { + return "read escape sequence" +} + +// escapeProxy is used only for attaches with a TTY. It is used to proxy +// stdin keypresses from the underlying reader and look for the passed in +// escape key sequence to signal a detach. +type escapeProxy struct { + escapeKeys []byte + escapeKeyPos int + r io.Reader +} + +// NewEscapeProxy returns a new TTY proxy reader which wraps the given reader +// and detects when the specified escape keys are read, in which case the Read +// method will return an error of type EscapeError. +func NewEscapeProxy(r io.Reader, escapeKeys []byte) io.Reader { + return &escapeProxy{ + escapeKeys: escapeKeys, + r: r, + } +} + +func (r *escapeProxy) Read(buf []byte) (int, error) { + nr, err := r.r.Read(buf) + + if len(r.escapeKeys) == 0 { + return nr, err + } + + preserve := func() { + // this preserves the original key presses in the passed in buffer + nr += r.escapeKeyPos + preserve := make([]byte, 0, r.escapeKeyPos+len(buf)) + preserve = append(preserve, r.escapeKeys[:r.escapeKeyPos]...) + preserve = append(preserve, buf...) + r.escapeKeyPos = 0 + copy(buf[0:nr], preserve) + } + + if nr != 1 || err != nil { + if r.escapeKeyPos > 0 { + preserve() + } + return nr, err + } + + if buf[0] != r.escapeKeys[r.escapeKeyPos] { + if r.escapeKeyPos > 0 { + preserve() + } + return nr, nil + } + + if r.escapeKeyPos == len(r.escapeKeys)-1 { + return 0, EscapeError{} + } + + // Looks like we've got an escape key, but we need to match again on the next + // read. + // Store the current escape key we found so we can look for the next one on + // the next read. + // Since this is an escape key, make sure we don't let the caller read it + // If later on we find that this is not the escape sequence, we'll add the + // keys back + r.escapeKeyPos++ + return nr - r.escapeKeyPos, nil +} diff --git a/pkg/term/proxy_test.go b/pkg/term/proxy_test.go new file mode 100644 index 0000000000..59b08ff41d --- /dev/null +++ b/pkg/term/proxy_test.go @@ -0,0 +1,114 @@ +package term + +import ( + "bytes" + "fmt" + "testing" + + "github.com/stretchr/testify/assert" +) + +func TestEscapeProxyRead(t *testing.T) { + escapeKeys, _ := ToBytes("") + keys, _ := ToBytes("a") + reader := NewEscapeProxy(bytes.NewReader(keys), escapeKeys) + buf := make([]byte, len(keys)) + nr, err := reader.Read(buf) + assert.NoError(t, err) + assert.Equal(t, nr, len(keys), fmt.Sprintf("nr %d should be equal to the number of %d", nr, len(keys))) + assert.Equal(t, keys, buf) + + keys, _ = ToBytes("a,b,c") + reader = NewEscapeProxy(bytes.NewReader(keys), escapeKeys) + buf = make([]byte, len(keys)) + nr, err = reader.Read(buf) + assert.NoError(t, err) + assert.Equal(t, nr, len(keys), fmt.Sprintf("nr %d should be equal to the number of %d", nr, len(keys))) + assert.Equal(t, keys, buf) + + keys, _ = ToBytes("") + reader = NewEscapeProxy(bytes.NewReader(keys), escapeKeys) + buf = make([]byte, len(keys)) + nr, err = reader.Read(buf) + assert.Error(t, err, "Should throw error when no keys are to read") + assert.Equal(t, nr, 0, "nr should be zero") + assert.Equal(t, len(keys), 0) + assert.Equal(t, len(buf), 0) + + escapeKeys, _ = ToBytes("DEL") + keys, _ = ToBytes("a,b,c,+") + reader = NewEscapeProxy(bytes.NewReader(keys), escapeKeys) + buf = make([]byte, len(keys)) + nr, err = reader.Read(buf) + assert.NoError(t, err) + assert.Equal(t, nr, len(keys), fmt.Sprintf("nr %d should be equal to the number of %d", nr, len(keys))) + assert.Equal(t, keys, buf) + + keys, _ = ToBytes("") + reader = NewEscapeProxy(bytes.NewReader(keys), escapeKeys) + buf = make([]byte, len(keys)) + nr, err = reader.Read(buf) + assert.Error(t, err, "Should throw error when no keys are to read") + assert.Equal(t, nr, 0, "nr should be zero") + assert.Equal(t, len(keys), 0) + assert.Equal(t, len(buf), 0) + + escapeKeys, _ = ToBytes("ctrl-x,ctrl-@") + keys, _ = ToBytes("DEL") + reader = NewEscapeProxy(bytes.NewReader(keys), escapeKeys) + buf = make([]byte, len(keys)) + nr, err = reader.Read(buf) + assert.NoError(t, err) + assert.Equal(t, nr, 1, fmt.Sprintf("nr %d should be equal to the number of 1", nr)) + assert.Equal(t, keys, buf) + + escapeKeys, _ = ToBytes("ctrl-c") + keys, _ = ToBytes("ctrl-c") + reader = NewEscapeProxy(bytes.NewReader(keys), escapeKeys) + buf = make([]byte, len(keys)) + nr, err = reader.Read(buf) + assert.Error(t, err, "read escape sequence") + assert.Equal(t, nr, 0, "nr should be equal to 0") + assert.Equal(t, keys, buf) + + escapeKeys, _ = ToBytes("ctrl-c,ctrl-z") + keys, _ = ToBytes("ctrl-c,ctrl-z") + reader = NewEscapeProxy(bytes.NewReader(keys), escapeKeys) + buf = make([]byte, 1) + nr, err = reader.Read(buf) + assert.NoError(t, err) + assert.Equal(t, nr, 0, "nr should be equal to 0") + assert.Equal(t, keys[0:1], buf) + nr, err = reader.Read(buf) + assert.Error(t, err, "read escape sequence") + assert.Equal(t, nr, 0, "nr should be equal to 0") + assert.Equal(t, keys[1:], buf) + + escapeKeys, _ = ToBytes("ctrl-c,ctrl-z") + keys, _ = ToBytes("ctrl-c,DEL,+") + reader = NewEscapeProxy(bytes.NewReader(keys), escapeKeys) + buf = make([]byte, 1) + nr, err = reader.Read(buf) + assert.NoError(t, err) + assert.Equal(t, nr, 0, "nr should be equal to 0") + assert.Equal(t, keys[0:1], buf) + buf = make([]byte, len(keys)) + nr, err = reader.Read(buf) + assert.NoError(t, err) + assert.Equal(t, nr, len(keys), fmt.Sprintf("nr should be equal to %d", len(keys))) + assert.Equal(t, keys, buf) + + escapeKeys, _ = ToBytes("ctrl-c,ctrl-z") + keys, _ = ToBytes("ctrl-c,DEL") + reader = NewEscapeProxy(bytes.NewReader(keys), escapeKeys) + buf = make([]byte, 1) + nr, err = reader.Read(buf) + assert.NoError(t, err) + assert.Equal(t, nr, 0, "nr should be equal to 0") + assert.Equal(t, keys[0:1], buf) + buf = make([]byte, len(keys)) + nr, err = reader.Read(buf) + assert.NoError(t, err) + assert.Equal(t, nr, len(keys), fmt.Sprintf("nr should be equal to %d", len(keys))) + assert.Equal(t, keys, buf) +} diff --git a/test/cli_attach_test.go b/test/cli_attach_test.go new file mode 100644 index 0000000000..5a2d3f4e6b --- /dev/null +++ b/test/cli_attach_test.go @@ -0,0 +1,64 @@ +package main + +import ( + "io" + "os/exec" + "strings" + + "github.com/alibaba/pouch/test/command" + "github.com/alibaba/pouch/test/environment" + "github.com/go-check/check" + "github.com/gotestyourself/gotestyourself/icmd" +) + +// PouchCreateSuite is the test suite for attach CLI. +type PouchAttachSuite struct{} + +func init() { + check.Suite(&PouchAttachSuite{}) +} + +// SetUpSuite does common setup in the beginning of each test suite. +func (suite *PouchAttachSuite) SetUpSuite(c *check.C) { + SkipIfFalse(c, environment.IsLinux) + + environment.PruneAllContainers(apiClient) + + PullImage(c, busyboxImage) +} + +// TearDownTest does cleanup work in the end of each test. +func (suite *PouchAttachSuite) TearDownTest(c *check.C) { +} + +// TestPouchAttachRunningContainer is to verify the correctness of attach a running container. +func (suite *PouchAttachSuite) TestPouchAttachRunningContainer(c *check.C) { + name := "TestPouchAttachRunningContainer" + + res := command.PouchRun("run", "-td", "--name", name, busyboxImage, "/bin/sh", "-c", "while true; do echo hello; done") + + //defer DelContainerForceMultyTime(c, name) + res.Assert(c, icmd.Success) + + cmd := exec.Command(environment.PouchBinary, "attach", name) + + out, err := cmd.StdoutPipe() + if err != nil { + c.Fatal(err) + } + defer out.Close() + + if err := cmd.Start(); err != nil { + c.Fatal(err) + } + + buf := make([]byte, 1024) + + if _, err := out.Read(buf); err != nil && err != io.EOF { + c.Fatal(err) + } + + if !strings.Contains(string(buf), "hello") { + c.Fatalf("unexpected output %s expected hello\n", string(buf)) + } +}