Skip to content

Commit

Permalink
Implement "runme beta session" (#683)
Browse files Browse the repository at this point in the history
It implements a new command that allows starting a session and
collecting all exported environment variables throughout its lifetime.

## Testing

As a prerequisite, add the following to your shell's startup script
(`.bashrc` or `.zshrc`):

```
# runme
source <(runme beta session setup)
```

Then, you can start a new session in the terminal:

```
runme beta session
```

The session is a regular shell of your choice, inferred from `$SHELL`.
You can use it like a normal shell, but any exported environment
variables will be returned upon exiting the session via `exit`.

Currently, the collected variables are printed out, but in the future,
they will be routed to a proper store.

---------

Co-authored-by: Sebastian (Tiedtke) Huckleberry <[email protected]>
  • Loading branch information
adambabik and sourishkrout authored Dec 2, 2024
1 parent a38548e commit 69ba348
Show file tree
Hide file tree
Showing 23 changed files with 484 additions and 96 deletions.
1 change: 1 addition & 0 deletions internal/cmd/beta/beta_cmd.go
Original file line number Diff line number Diff line change
Expand Up @@ -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))
Expand Down
197 changes: 197 additions & 0 deletions internal/cmd/beta/session_cmd.go
Original file line number Diff line number Diff line change
@@ -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)
}
20 changes: 16 additions & 4 deletions internal/command/command_inline_shell.go
Original file line number Diff line number Diff line change
Expand Up @@ -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()...)
Expand All @@ -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
}
Expand Down
1 change: 1 addition & 0 deletions internal/command/command_terminal.go
Original file line number Diff line number Diff line change
Expand Up @@ -75,6 +75,7 @@ func (c *terminalCommand) Wait(ctx context.Context) (err error) {
err = cErr
}
}

return err
}

Expand Down
8 changes: 1 addition & 7 deletions internal/command/command_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down
6 changes: 4 additions & 2 deletions internal/command/command_unix_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@ import (
"time"
"unicode"

"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"go.uber.org/zap/zaptest"

Expand Down Expand Up @@ -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},
)
Expand All @@ -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) {
Expand Down
12 changes: 8 additions & 4 deletions internal/command/env_collector.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
}

Expand Down
29 changes: 18 additions & 11 deletions internal/command/env_collector_factory.go
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand All @@ -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")
Expand Down
Loading

0 comments on commit 69ba348

Please sign in to comment.