diff --git a/pkg/plugin/runtime/run.go b/pkg/plugin/runtime/run.go index e451019..2323a53 100644 --- a/pkg/plugin/runtime/run.go +++ b/pkg/plugin/runtime/run.go @@ -4,11 +4,9 @@ import ( "context" "errors" "fmt" - "io" "github.com/docker/docker/api/types" "github.com/docker/docker/api/types/container" - "github.com/docker/docker/pkg/stdcopy" "github.com/dodo-cli/dodo-core/pkg/plugin" "github.com/dodo-cli/dodo-core/pkg/plugin/runtime" log "github.com/hashicorp/go-hclog" @@ -57,63 +55,16 @@ func (c *ContainerRuntime) RunAndWaitContainer(id string, height uint32, width u func (c *ContainerRuntime) StreamContainer(id string, stream *plugin.StreamConfig) (*runtime.Result, error) { ctx := context.Background() - client, err := c.ensureClient() + s, err := c.AttachContainer(ctx, id, stream) if err != nil { return nil, err } - config, err := client.ContainerInspect(ctx, id) - if err != nil { - return nil, fmt.Errorf("could not inspect container: %w", err) - } - - attach, err := client.ContainerAttach( - ctx, - id, - types.ContainerAttachOptions{ - Stream: true, - Stdin: true, - Stdout: true, - Stderr: true, - Logs: true, - }, - ) - if err != nil { - return nil, fmt.Errorf("could not attach to container: %w", err) - } - - defer func() { - if err := attach.Conn.Close(); err != nil { - log.L().Warn("could not close streaming connection", "error", err) - } - }() - eg, _ := errgroup.WithContext(ctx) result := &runtime.Result{} - inCopier := plugin.NewCancelCopier(stream.Stdin, attach.Conn) - - eg.Go(func() error { - defer inCopier.Close() - if config.Config.Tty { - if _, err := io.Copy(stream.Stdout, attach.Reader); err != nil { - log.L().Warn("could not copy container output", "error", err) - } - } else { - if _, err := stdcopy.StdCopy(stream.Stdout, stream.Stderr, attach.Reader); err != nil { - log.L().Warn("could not copy container output", "error", err) - } - } - - return nil - }) - eg.Go(func() error { - if err := inCopier.Copy(); err != nil { - log.L().Error("could not copy container input", "error", err) - } - - return nil - }) + eg.Go(s.CopyOutput) + eg.Go(s.CopyInput) eg.Go(func() error { r, err := c.RunAndWaitContainer(id, stream.TerminalHeight, stream.TerminalWidth) diff --git a/pkg/plugin/runtime/stream.go b/pkg/plugin/runtime/stream.go new file mode 100644 index 0000000..bc17061 --- /dev/null +++ b/pkg/plugin/runtime/stream.go @@ -0,0 +1,96 @@ +package runtime + +import ( + "context" + "fmt" + "io" + + "github.com/docker/docker/api/types" + "github.com/docker/docker/pkg/stdcopy" + "github.com/dodo-cli/dodo-core/pkg/plugin" + log "github.com/hashicorp/go-hclog" +) + +type ContainerStream struct { + hasTTY bool + config *plugin.StreamConfig + hijack types.HijackedResponse + inCopier *plugin.CancelCopier +} + +type closeWriter interface { + CloseWrite() error +} + +func (c *ContainerRuntime) AttachContainer( + ctx context.Context, id string, stream *plugin.StreamConfig, +) (*ContainerStream, error) { + client, err := c.ensureClient() + if err != nil { + return nil, err + } + + config, err := client.ContainerInspect(ctx, id) + if err != nil { + return nil, fmt.Errorf("could not inspect container: %w", err) + } + + attach, err := client.ContainerAttach( + ctx, + id, + types.ContainerAttachOptions{ + Stream: true, + Stdin: true, + Stdout: true, + Stderr: true, + Logs: true, + }, + ) + if err != nil { + return nil, fmt.Errorf("could not attach to container: %w", err) + } + + return &ContainerStream{ + hasTTY: config.Config.Tty, + config: stream, + hijack: attach, + inCopier: plugin.NewCancelCopier(stream.Stdin, attach.Conn), + }, nil +} + +func (s *ContainerStream) CopyOutput() error { + defer s.inCopier.Close() + + if s.hasTTY { + if _, err := io.Copy(s.config.Stdout, s.hijack.Reader); err != nil { + log.L().Warn("could not copy container output", "error", err) + } + } else { + if _, err := stdcopy.StdCopy(s.config.Stdout, s.config.Stderr, s.hijack.Reader); err != nil { + log.L().Warn("could not copy container output", "error", err) + } + } + + return nil +} + +func (s *ContainerStream) CopyInput() error { + defer func() { + cw, ok := s.hijack.Conn.(closeWriter) + if ok { + if err := cw.CloseWrite(); err != nil { + log.L().Warn("could not close streaming connection", "error", err) + } + } else { + if err := s.hijack.Conn.Close(); err != nil { + log.L().Warn("could not close streaming connection", "error", err) + } + } + }() + + if err := s.inCopier.Copy(); err != nil { + log.L().Error("could not copy container input", "error", err) + } + + return nil +}