Implement "runme beta session"

internal/cmd/beta/beta_cmd.go

@@ -96,6 +96,7 @@ All commands use the runme.yaml configuration file.`,
 
 	cmd.AddCommand(listCmd(cFlags))
 	cmd.AddCommand(printCmd(cFlags))
+	cmd.AddCommand(sessionCmd(cFlags))
 	cmd.AddCommand(server.Cmd())
 	cmd.AddCommand(runCmd(cFlags))
 	cmd.AddCommand(envCmd(cFlags))

internal/cmd/beta/session_cmd.go

@@ -0,0 +1,197 @@
+package beta
+
+import (
+	"context"
+	"io"
+	"os"
+	"os/exec"
+	"strconv"
+
+	"github.com/pkg/errors"
+	"github.com/spf13/cobra"
+	"go.uber.org/multierr"
+	"go.uber.org/zap"
+
+	"github.com/stateful/runme/v3/internal/command"
+	"github.com/stateful/runme/v3/internal/config/autoconfig"
+	runnerv2 "github.com/stateful/runme/v3/pkg/api/gen/proto/go/runme/runner/v2"
+)
+
+func sessionCmd(*commonFlags) *cobra.Command {
+	cmd := cobra.Command{
+		Use:   "session",
+		Short: "Start shell within a session.",
+		Long: `Start shell within a session.
+
+All exported variables during the session will be available to the subsequent commands.
+`,
+		RunE: func(cmd *cobra.Command, args []string) error {
+			return autoconfig.InvokeForCommand(
+				func(
+					cmdFactory command.Factory,
+					logger *zap.Logger,
+				) error {
+					defer logger.Sync()
+
+					envs, err := executeDefaultShellProgram(
+						cmd.Context(),
+						cmdFactory,
+						cmd.InOrStdin(),
+						cmd.OutOrStdout(),
+						cmd.ErrOrStderr(),
+						nil,
+					)
+					if err != nil {
+						return err
+					}
+
+					// TODO(adamb): currently, the collected env are printed out,
+					// but they could be put in a session.
+					if _, err := cmd.ErrOrStderr().Write([]byte("Collected env during the session:\n")); err != nil {
+						return errors.WithStack(err)
+					}
+
+					for _, env := range envs {
+						_, err := cmd.OutOrStdout().Write([]byte(env + "\n"))
+						if err != nil {
+							return errors.WithStack(err)
+						}
+					}
+
+					return nil
+				},
+			)
+		},
+	}
+
+	cmd.AddCommand(sessionSetupCmd())
+
+	return &cmd
+}
+
+func executeDefaultShellProgram(
+	ctx context.Context,
+	commandFactory command.Factory,
+	stdin io.Reader,
+	stdout io.Writer,
+	stderr io.Writer,
+	additionalEnv []string,
+) ([]string, error) {
+	envCollector, err := command.NewEnvCollectorFactory().Build()
+	if err != nil {
+		return nil, errors.WithStack(err)
+	}
+
+	cfg := &command.ProgramConfig{
+		ProgramName: defaultShell(),
+		Mode:        runnerv2.CommandMode_COMMAND_MODE_CLI,
+		Env: append(
+			[]string{command.CreateEnv(command.EnvNameTerminalSessionEnabled, "true")},
+			append(envCollector.ExtraEnv(), additionalEnv...)...,
+		),
+	}
+	options := command.CommandOptions{
+		NoShell: true,
+		Stdin:   stdin,
+		Stdout:  stdout,
+		Stderr:  stderr,
+	}
+	program, err := commandFactory.Build(cfg, options)
+	if err != nil {
+		return nil, err
+	}
+
+	err = program.Start(ctx)
+	if err != nil {
+		return nil, err
+	}
+
+	err = program.Wait(ctx)
+	if err != nil {
+		return nil, err
+	}
+
+	changed, _, err := envCollector.Diff()
+	return changed, err
+}
+
+func defaultShell() string {
+	shell := os.Getenv("SHELL")
+	if shell == "" {
+		shell, _ = exec.LookPath("bash")
+	}
+	if shell == "" {
+		shell = "/bin/sh"
+	}
+	return shell
+}
+
+func sessionSetupCmd() *cobra.Command {
+	var debug bool
+
+	cmd := cobra.Command{
+		Use:    "setup",
+		Hidden: true,
+		RunE: func(cmd *cobra.Command, args []string) error {
+			return autoconfig.InvokeForCommand(
+				func(
+					cmdFactory command.Factory,
+					logger *zap.Logger,
+				) error {
+					defer logger.Sync()
+
+					out := cmd.OutOrStdout()
+
+					if err := requireEnvs(
+						command.EnvNameTerminalSessionEnabled,
+						command.EnvNameTerminalSessionPrePath,
+						command.EnvNameTerminalSessionPostPath,
+					); err != nil {
+						logger.Info("session setup is skipped because the environment variable is not set", zap.Error(err))
+						return writeNoopShellCommand(out)
+					}
+
+					sessionSetupEnabled := os.Getenv(command.EnvNameTerminalSessionEnabled)
+					if val, err := strconv.ParseBool(sessionSetupEnabled); err != nil || !val {
+						logger.Debug("session setup is skipped", zap.Error(err), zap.Bool("value", val))
+						return writeNoopShellCommand(out)
+					}
+
+					envSetter := command.NewScriptEnvSetter(
+						os.Getenv(command.EnvNameTerminalSessionPrePath),
+						os.Getenv(command.EnvNameTerminalSessionPostPath),
+						debug,
+					)
+					if err := envSetter.SetOnShell(out); err != nil {
+						return err
+					}
+
+					if _, err := cmd.ErrOrStderr().Write([]byte("Runme session active. When you're done, execute \"exit\".\n")); err != nil {
+						return errors.WithStack(err)
+					}
+
+					return nil
+				},
+			)
+		},
+	}
+
+	cmd.Flags().BoolVar(&debug, "debug", false, "Enable debug mode.")
+
+	return &cmd
+}
+
+func requireEnvs(names ...string) error {
+	var err error
+	for _, name := range names {
+		if os.Getenv(name) == "" {
+			err = multierr.Append(err, errors.Errorf("environment variable %q is required", name))
+		}
+	}
+	return err
+}
+
+func writeNoopShellCommand(w io.Writer) error {
+	_, err := w.Write([]byte(":"))
+	return errors.WithStack(err)
+}

internal/command/command_inline_shell.go

@@ -33,11 +33,13 @@ func (c *inlineShellCommand) Start(ctx context.Context) error {
 	if err != nil {
 		return err
 	}
-
 	c.logger.Debug("inline shell script", zap.String("script", script))
 
 	cfg := c.ProgramConfig()
-	cfg.Arguments = append(cfg.Arguments, "-c", script)
+
+	if script != "" {
+		cfg.Arguments = append(cfg.Arguments, "-c", script)
+	}
 
 	if c.envCollector != nil {
 		cfg.Env = append(cfg.Env, c.envCollector.ExtraEnv()...)
@@ -50,9 +52,19 @@ func (c *inlineShellCommand) Wait(ctx context.Context) error {
 	err := c.internalCommand.Wait(ctx)
 
 	if c.envCollector != nil {
-		c.logger.Info("collecting the environment after the script execution")
+		c.logger.Info(
+			"collecting the environment after the script execution",
+			zap.Int("count", len(c.session.GetAllEnv())), // TODO(adamb): change to session.Size()
+		)
+
 		cErr := c.collectEnv(ctx)
-		c.logger.Info("collected the environment after the script execution", zap.Error(cErr))
+
+		c.logger.Info(
+			"collected the environment after the script execution",
+			zap.Int("count", len(c.session.GetAllEnv())), // TODO(adamb): change to session.Size()
+			zap.Error(cErr),
+		)
+
 		if cErr != nil && err == nil {
 			err = cErr
 		}

internal/command/command_terminal.go

@@ -75,6 +75,7 @@ func (c *terminalCommand) Wait(ctx context.Context) (err error) {
 			err = cErr
 		}
 	}
+
 	return err
 }
 

internal/command/command_test.go

@@ -13,13 +13,7 @@ import (
 )
 
 func init() {
-	// Switch from "runme env" to "env -0" for the tests.
-	// This is because the "runme" program is not available
-	// in the test environment.
-	//
-	// TODO(adamb): this can be changed. runme must be built
-	// in the test environment and put into the PATH.
-	SetEnvDumpCommand("env -0")
+	SetEnvDumpCommandForTesting()
 }
 
 func testExecuteCommand(

internal/command/command_unix_test.go

@@ -12,6 +12,7 @@ import (
 	"time"
 	"unicode"
 
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"go.uber.org/zap/zaptest"
 
@@ -341,6 +342,7 @@ func TestCommand_SetWinsize(t *testing.T) {
 				},
 				Interactive: true,
 				Mode:        runnerv2.CommandMode_COMMAND_MODE_INLINE,
+				Env:         []string{"TERM=xterm"},
 			},
 			CommandOptions{Stdout: stdout},
 		)
@@ -351,8 +353,8 @@ func TestCommand_SetWinsize(t *testing.T) {
 		err = SetWinsize(cmd, &Winsize{Rows: 45, Cols: 56, X: 0, Y: 0})
 		require.NoError(t, err)
 		err = cmd.Wait(context.Background())
-		require.NoError(t, err)
-		require.Equal(t, "56\r\n45\r\n", stdout.String())
+		assert.NoError(t, err)
+		assert.Equal(t, "56\r\n45\r\n", stdout.String())
 	})
 
 	t.Run("Terminal", func(t *testing.T) {

internal/command/env_collector.go

@@ -29,11 +29,15 @@ var envDumpCommand = func() string {
 	return strings.Join([]string{path, "env", "dump", "--insecure"}, " ")
 }()
 
-func SetEnvDumpCommand(cmd string) {
-	envDumpCommand = cmd
+// SetEnvDumpCommandForTesting overrides the default command that dumps the environment variables.
+// It is and should be used only for testing purposes.
+// TODO(adamb): this can be made obsolete. runme must be built
+// in the test environment and put into the PATH.
+func SetEnvDumpCommandForTesting() {
+	envDumpCommand = "env -0"
 	// When overriding [envDumpCommand], we disable the encryption.
-	// There is no way to test the encryption if the dump command
-	// is not controlled.
+	// There is no reliable way at the moment to have encryption and
+	// not control the dump command.
 	envCollectorEnableEncryption = false
 }
 

internal/command/env_collector_factory.go

@@ -7,30 +7,37 @@ import (
 	"github.com/pkg/errors"
 )
 
-type envCollectorFactoryOptions struct {
+type EnvCollectorFactory struct {
 	encryptionEnabled bool
 	useFifo           bool
 }
 
-type envCollectorFactory struct {
-	opts envCollectorFactoryOptions
+func NewEnvCollectorFactory() *EnvCollectorFactory {
+	return &EnvCollectorFactory{
+		encryptionEnabled: envCollectorEnableEncryption,
+		useFifo:           envCollectorUseFifo,
+	}
 }
 
-func newEnvCollectorFactory(opts envCollectorFactoryOptions) *envCollectorFactory {
-	return &envCollectorFactory{
-		opts: opts,
-	}
+func (f *EnvCollectorFactory) WithEnryption(value bool) *EnvCollectorFactory {
+	f.encryptionEnabled = value
+	return f
+}
+
+func (f *EnvCollectorFactory) UseFifo(value bool) *EnvCollectorFactory {
+	f.useFifo = value
+	return f
 }
 
-func (f *envCollectorFactory) Build() (envCollector, error) {
+func (f *EnvCollectorFactory) Build() (envCollector, error) {
 	scanner := scanEnv
 
 	var (
 		encKey   []byte
 		encNonce []byte
 	)
 
-	if f.opts.encryptionEnabled {
+	if f.encryptionEnabled {
 		var err error
 
 		encKey, encNonce, err = f.generateEncryptionKeyAndNonce()
@@ -48,14 +55,14 @@ func (f *envCollectorFactory) Build() (envCollector, error) {
 		}
 	}
 
-	if f.opts.useFifo && runtimestd.GOOS != "windows" {
+	if f.useFifo && runtimestd.GOOS != "windows" {
 		return newEnvCollectorFifo(scanner, encKey, encNonce)
 	}
 
 	return newEnvCollectorFile(scanner, encKey, encNonce)
 }
 
-func (f *envCollectorFactory) generateEncryptionKeyAndNonce() ([]byte, []byte, error) {
+func (f *EnvCollectorFactory) generateEncryptionKeyAndNonce() ([]byte, []byte, error) {
 	key, err := createEnvEncryptionKey()
 	if err != nil {
 		return nil, nil, errors.WithMessage(err, "failed to create the encryption key")