Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Implement "runme beta session" #683

Open
wants to merge 8 commits into
base: main
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
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.EnvCollectorSessionEnvName, "1")},
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.EnvCollectorSessionEnvName,
command.EnvCollectorSessionPrePathEnvName,
command.EnvCollectorSessionPostPathEnvName,
); err != nil {
logger.Info("session setup is skipped because the environment variable is not set", zap.Error(err))
return writeNoopShellCommand(out)
}
Comment on lines +145 to +152
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

One thing I don't get is how these won't ever produce an error? In other words how are they set in the first place if they are expected by setup @adambabik?

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The order is that beta session starts a shell process with those envs provided first. This shell process starts and executes, in the case of zsh, the .zshrc file inside which beta session setup is called.

It's perfectly normal that it typically will error and output no-op shell command in .zshrc, but from the user point of view it should be invisible. We can recommend protecting source <(/path/to/runme beta session setup) with a conditional checking command.EnvCollectorSessionEnvName to avoid that.


sessionSetupEnabled := os.Getenv(command.EnvCollectorSessionEnvName)
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)
}
Comment on lines +154 to +158
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Also, is RUNME_SESSION expected to be a bool (serialized)? We are using the env var to hold the session ID all over the place so a semantic clash here is bad if bool <> SessionID (both string).

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Oh yeah, that's a good point. I didn't notice the clash. I will change it to something like RUNME_TERM_SESSION_ENABLED.


envSetter := command.NewScriptEnvSetter(
os.Getenv(command.EnvCollectorSessionPrePathEnvName),
os.Getenv(command.EnvCollectorSessionPostPathEnvName),
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
16 changes: 12 additions & 4 deletions internal/command/env_collector.go
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,10 @@ import (
const maxScannerBufferSizeInBytes = 1024 * 1024 * 1024 // 1GB

const (
EnvCollectorSessionEnvName = "RUNME_SESSION"
EnvCollectorSessionPrePathEnvName = "RUNME_SESSION_PREPATH"
EnvCollectorSessionPostPathEnvName = "RUNME_SESSION_POSTPATH"

envCollectorEncKeyEnvName = "RUNME_ENCRYPTION_KEY"
envCollectorEncNonceEnvName = "RUNME_ENCRYPTION_NONCE"
)
Expand All @@ -29,11 +33,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
Loading