Skip to content

Commit

Permalink
feat(cli/ssh): simplify log file flags (coder#7863)
Browse files Browse the repository at this point in the history
And, fix a race condition.
  • Loading branch information
ammario authored Jun 12, 2023
1 parent ec7b117 commit 5de1084
Show file tree
Hide file tree
Showing 5 changed files with 92 additions and 77 deletions.
13 changes: 13 additions & 0 deletions cli/clibase/cmd.go
Original file line number Diff line number Diff line change
Expand Up @@ -448,6 +448,19 @@ func (inv *Invocation) Run() (err error) {
panic(err)
}
}()
// We close Stdin to prevent deadlocks, e.g. when the command
// has ended but an io.Copy is still reading from Stdin.
defer func() {
if inv.Stdin == nil {
return
}
rc, ok := inv.Stdin.(io.ReadCloser)
if !ok {
return
}
e := rc.Close()
err = errors.Join(err, e)
}()
err = inv.run(&runState{
allArgs: inv.Args,
})
Expand Down
109 changes: 68 additions & 41 deletions cli/ssh.go
Original file line number Diff line number Diff line change
Expand Up @@ -10,9 +10,9 @@ import (
"net/url"
"os"
"os/exec"
"path"
"path/filepath"
"strings"
"sync"
"time"

"github.com/gen2brain/beeep"
Expand Down Expand Up @@ -52,8 +52,7 @@ func (r *RootCmd) ssh() *clibase.Cmd {
wsPollInterval time.Duration
waitEnum string
noWait bool
logDir string
logToFile bool
logDirPath string
)
client := new(codersdk.Client)
cmd := &clibase.Cmd{
Expand All @@ -76,24 +75,45 @@ func (r *RootCmd) ssh() *clibase.Cmd {
logger.Error(ctx, "command exit", slog.Error(retErr))
}
}()
if logToFile {
// we need a way to ensure different ssh invocations don't clobber
// each other's logs. Date-time strings will likely have collisions
// in unit tests and/or scripts unless we extend precision out to
// sub-millisecond, which seems unwieldy. A simple 5-character random
// string will do it, since the operating system already tracks
// dates and times for file IO.
qual, err := cryptorand.String(5)

// This WaitGroup solves for a race condition where we were logging
// while closing the log file in a defer. It probably solves
// others too.
var wg sync.WaitGroup
wg.Add(1)
defer wg.Done()

if logDirPath != "" {
nonce, err := cryptorand.StringCharset(cryptorand.Lower, 5)
if err != nil {
return xerrors.Errorf("generate random qualifier: %w", err)
return xerrors.Errorf("generate nonce: %w", err)
}
logPth := path.Join(logDir, fmt.Sprintf("coder-ssh-%s.log", qual))
logFile, err := os.Create(logPth)
logFilePath := filepath.Join(
logDirPath,
fmt.Sprintf(
"coder-ssh-%s-%s.log",
// The time portion makes it easier to find the right
// log file.
time.Now().Format("20060102-150405"),
// The nonce prevents collisions, as SSH invocations
// frequently happen in parallel.
nonce,
),
)
logFile, err := os.OpenFile(
logFilePath,
os.O_CREATE|os.O_APPEND|os.O_WRONLY|os.O_EXCL,
0o600,
)
if err != nil {
return xerrors.Errorf("error opening %s for logging: %w", logPth, err)
return xerrors.Errorf("error opening %s for logging: %w", logDirPath, err)
}
go func() {
wg.Wait()
logFile.Close()
}()

logger = slog.Make(sloghuman.Sink(logFile))
defer logFile.Close()
if r.verbose {
logger = logger.Leveled(slog.LevelDebug)
}
Expand Down Expand Up @@ -192,9 +212,18 @@ func (r *RootCmd) ssh() *clibase.Cmd {
return xerrors.Errorf("connect SSH: %w", err)
}
defer rawSSH.Close()
go watchAndClose(ctx, rawSSH.Close, logger, client, workspace)

wg.Add(1)
go func() {
defer wg.Done()
watchAndClose(ctx, func() error {
return rawSSH.Close()
}, logger, client, workspace)
}()

wg.Add(1)
go func() {
defer wg.Done()
// Ensure stdout copy closes incase stdin is closed
// unexpectedly. Typically we wouldn't worry about
// this since OpenSSH should kill the proxy command.
Expand Down Expand Up @@ -227,19 +256,24 @@ func (r *RootCmd) ssh() *clibase.Cmd {
return xerrors.Errorf("ssh session: %w", err)
}
defer sshSession.Close()
go watchAndClose(
ctx,
func() error {
err := sshSession.Close()
logger.Debug(ctx, "session close", slog.Error(err))
err = sshClient.Close()
logger.Debug(ctx, "client close", slog.Error(err))
return nil
},
logger,
client,
workspace,
)

wg.Add(1)
go func() {
defer wg.Done()
watchAndClose(
ctx,
func() error {
err := sshSession.Close()
logger.Debug(ctx, "session close", slog.Error(err))
err = sshClient.Close()
logger.Debug(ctx, "client close", slog.Error(err))
return nil
},
logger,
client,
workspace,
)
}()

if identityAgent == "" {
identityAgent = os.Getenv("SSH_AUTH_SOCK")
Expand Down Expand Up @@ -389,18 +423,11 @@ func (r *RootCmd) ssh() *clibase.Cmd {
UseInstead: []clibase.Option{waitOption},
},
{
Flag: "log-dir",
Default: os.TempDir(),
Description: "Specify the location for the log files.",
Env: "CODER_SSH_LOG_DIR",
Value: clibase.StringOf(&logDir),
},
{
Flag: "log-to-file",
Flag: "log-dir",
Description: "Specify the directory containing SSH diagnostic log files.",
Env: "CODER_SSH_LOG_DIR",
FlagShorthand: "l",
Env: "CODER_SSH_LOG_TO_FILE",
Description: "Enable diagnostic logging to file.",
Value: clibase.BoolOf(&logToFile),
Value: clibase.StringOf(&logDirPath),
},
}
return cmd
Expand Down
26 changes: 7 additions & 19 deletions cli/ssh_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -261,7 +261,7 @@ func TestSSH(t *testing.T) {
client, workspace, agentToken := setupWorkspaceForAgent(t, nil)
_, _ = tGoContext(t, func(ctx context.Context) {
// Run this async so the SSH command has to wait for
// the build and agent to connect!
// the build and agent to connect.
agentClient := agentsdk.New(client.URL)
agentClient.SetSessionToken(agentToken)
agentCloser := agent.New(agent.Options{
Expand Down Expand Up @@ -411,20 +411,14 @@ func TestSSH(t *testing.T) {
t.Run("FileLogging", func(t *testing.T) {
t.Parallel()

dir := t.TempDir()
logDir := t.TempDir()

client, workspace, agentToken := setupWorkspaceForAgent(t, nil)
inv, root := clitest.New(t, "ssh", workspace.Name, "-l", "--log-dir", dir)
inv, root := clitest.New(t, "ssh", "-l", logDir, workspace.Name)
clitest.SetupConfig(t, client, root)
pty := ptytest.New(t).Attach(inv)
w := clitest.StartWithWaiter(t, inv)

ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
defer cancel()

cmdDone := tGo(t, func() {
err := inv.WithContext(ctx).Run()
assert.NoError(t, err)
})
pty.ExpectMatch("Waiting")

agentClient := agentsdk.New(client.URL)
Expand All @@ -439,17 +433,11 @@ func TestSSH(t *testing.T) {

// Shells on Mac, Windows, and Linux all exit shells with the "exit" command.
pty.WriteLine("exit")
<-cmdDone
w.RequireSuccess()

entries, err := os.ReadDir(dir)
ents, err := os.ReadDir(logDir)
require.NoError(t, err)
for _, e := range entries {
t.Logf("logdir entry: %s", e.Name())
if strings.HasPrefix(e.Name(), "coder-ssh") {
return
}
}
t.Fatal("failed to find ssh logfile")
require.Len(t, ents, 1, "expected one file in logdir %s", logDir)
})
}

Expand Down
7 changes: 2 additions & 5 deletions cli/testdata/coder_ssh_--help.golden
Original file line number Diff line number Diff line change
Expand Up @@ -18,11 +18,8 @@ Start a shell into a workspace
Specifies which identity agent to use (overrides $SSH_AUTH_SOCK),
forward agent must also be enabled.

--log-dir string, $CODER_SSH_LOG_DIR (default: /tmp)
Specify the location for the log files.

-l, --log-to-file bool, $CODER_SSH_LOG_TO_FILE
Enable diagnostic logging to file.
-l, --log-dir string, $CODER_SSH_LOG_DIR
Specify the directory containing SSH diagnostic log files.

--no-wait bool, $CODER_SSH_NO_WAIT
Enter workspace immediately after the agent has connected. This is the
Expand Down
14 changes: 2 additions & 12 deletions docs/cli/ssh.md
Original file line number Diff line number Diff line change
Expand Up @@ -39,24 +39,14 @@ Specifies whether to forward the GPG agent. Unsupported on Windows workspaces, b

Specifies which identity agent to use (overrides $SSH_AUTH_SOCK), forward agent must also be enabled.

### --log-dir
### -l, --log-dir

| | |
| ----------- | ------------------------------- |
| Type | <code>string</code> |
| Environment | <code>$CODER_SSH_LOG_DIR</code> |
| Default | <code>/tmp</code> |

Specify the location for the log files.

### -l, --log-to-file

| | |
| ----------- | ----------------------------------- |
| Type | <code>bool</code> |
| Environment | <code>$CODER_SSH_LOG_TO_FILE</code> |

Enable diagnostic logging to file.
Specify the directory containing SSH diagnostic log files.

### --no-wait

Expand Down

0 comments on commit 5de1084

Please sign in to comment.