2022-02-10 14:33:27 +00:00
package cli
2022-03-22 19:17:50 +00:00
import (
2022-03-29 00:19:28 +00:00
"context"
2022-07-25 15:25:34 +00:00
"errors"
2022-05-13 17:09:04 +00:00
"fmt"
2022-03-30 22:59:54 +00:00
"io"
"os"
2022-05-13 17:09:04 +00:00
"path/filepath"
2022-04-11 23:54:30 +00:00
"strings"
2022-05-13 17:09:04 +00:00
"time"
2022-03-22 19:17:50 +00:00
2022-05-13 17:09:04 +00:00
"github.com/gen2brain/beeep"
"github.com/gofrs/flock"
2022-04-11 21:06:15 +00:00
"github.com/google/uuid"
2022-03-30 22:59:54 +00:00
"github.com/mattn/go-isatty"
2022-03-22 19:17:50 +00:00
"github.com/spf13/cobra"
2022-03-29 00:19:28 +00:00
gossh "golang.org/x/crypto/ssh"
2022-05-25 18:28:10 +00:00
gosshagent "golang.org/x/crypto/ssh/agent"
2022-04-01 21:25:46 +00:00
"golang.org/x/term"
2022-03-25 19:48:08 +00:00
"golang.org/x/xerrors"
2022-03-22 19:17:50 +00:00
2022-06-24 21:21:46 +00:00
"cdr.dev/slog"
2022-09-01 01:09:44 +00:00
"github.com/coder/coder/agent"
2022-03-30 22:59:54 +00:00
"github.com/coder/coder/cli/cliflag"
2022-03-29 00:19:28 +00:00
"github.com/coder/coder/cli/cliui"
2022-05-13 17:09:04 +00:00
"github.com/coder/coder/coderd/autobuild/notify"
2022-06-02 10:23:34 +00:00
"github.com/coder/coder/coderd/util/ptr"
2022-03-25 19:48:08 +00:00
"github.com/coder/coder/codersdk"
2022-05-17 22:55:58 +00:00
"github.com/coder/coder/cryptorand"
2022-03-22 19:17:50 +00:00
)
2022-08-25 16:10:42 +00:00
var (
workspacePollInterval = time . Minute
autostopNotifyCountdown = [ ] time . Duration { 30 * time . Minute }
)
2022-05-13 17:09:04 +00:00
2022-03-29 00:19:28 +00:00
func ssh ( ) * cobra . Command {
2022-03-30 22:59:54 +00:00
var (
2022-05-20 10:57:02 +00:00
stdio bool
shuffle bool
2022-05-25 18:28:10 +00:00
forwardAgent bool
2022-06-02 08:13:38 +00:00
identityAgent string
2022-05-20 10:57:02 +00:00
wsPollInterval time . Duration
2022-06-24 21:21:46 +00:00
wireguard bool
2022-03-30 22:59:54 +00:00
)
2022-03-22 19:17:50 +00:00
cmd := & cobra . Command {
2022-05-09 22:42:02 +00:00
Annotations : workspaceCommand ,
Use : "ssh <workspace>" ,
Short : "SSH into a workspace" ,
2022-05-17 22:55:58 +00:00
Args : cobra . ArbitraryArgs ,
2022-03-22 19:17:50 +00:00
RunE : func ( cmd * cobra . Command , args [ ] string ) error {
2022-08-02 14:44:59 +00:00
ctx , cancel := context . WithCancel ( cmd . Context ( ) )
defer cancel ( )
2022-08-23 20:55:39 +00:00
client , err := CreateClient ( cmd )
2022-03-22 19:17:50 +00:00
if err != nil {
return err
}
2022-04-01 19:42:36 +00:00
2022-05-17 22:55:58 +00:00
if shuffle {
err := cobra . ExactArgs ( 0 ) ( cmd , args )
if err != nil {
return err
}
} else {
err := cobra . MinimumNArgs ( 1 ) ( cmd , args )
if err != nil {
return err
}
2022-03-25 19:48:08 +00:00
}
2022-04-01 19:42:36 +00:00
2022-08-02 14:44:59 +00:00
workspace , workspaceAgent , err := getWorkspaceAndAgent ( ctx , cmd , client , codersdk . Me , args [ 0 ] , shuffle )
2022-03-22 19:17:50 +00:00
if err != nil {
return err
}
2022-04-01 19:42:36 +00:00
2022-03-30 22:59:54 +00:00
// OpenSSH passes stderr directly to the calling TTY.
// This is required in "stdio" mode so a connecting indicator can be displayed.
2022-08-02 14:44:59 +00:00
err = cliui . Agent ( ctx , cmd . ErrOrStderr ( ) , cliui . AgentOptions {
2022-03-29 00:19:28 +00:00
WorkspaceName : workspace . Name ,
2022-04-11 21:06:15 +00:00
Fetch : func ( ctx context . Context ) ( codersdk . WorkspaceAgent , error ) {
2022-06-24 21:21:46 +00:00
return client . WorkspaceAgent ( ctx , workspaceAgent . ID )
2022-03-29 00:19:28 +00:00
} ,
} )
if err != nil {
return xerrors . Errorf ( "await agent: %w" , err )
2022-03-25 19:48:08 +00:00
}
2022-03-22 19:17:50 +00:00
2022-09-01 01:09:44 +00:00
var conn agent . Conn
2022-06-24 21:21:46 +00:00
if ! wireguard {
2022-09-01 01:09:44 +00:00
conn , err = client . DialWorkspaceAgent ( ctx , workspaceAgent . ID , nil )
2022-06-24 21:21:46 +00:00
} else {
2022-09-01 01:09:44 +00:00
conn , err = client . DialWorkspaceAgentTailnet ( ctx , slog . Logger { } , workspaceAgent . ID )
}
if err != nil {
return err
}
defer conn . Close ( )
2022-06-24 21:21:46 +00:00
2022-09-01 01:09:44 +00:00
stopPolling := tryPollWorkspaceAutostop ( ctx , client , workspace )
defer stopPolling ( )
2022-06-24 21:21:46 +00:00
2022-09-01 01:09:44 +00:00
if stdio {
rawSSH , err := conn . SSH ( )
if err != nil {
return err
2022-06-24 21:21:46 +00:00
}
2022-09-01 01:09:44 +00:00
defer rawSSH . Close ( )
2022-06-24 21:21:46 +00:00
2022-09-01 01:09:44 +00:00
go func ( ) {
_ , _ = io . Copy ( cmd . OutOrStdout ( ) , rawSSH )
} ( )
_ , _ = io . Copy ( rawSSH , cmd . InOrStdin ( ) )
return nil
2022-08-02 14:44:59 +00:00
}
2022-06-24 21:21:46 +00:00
2022-09-01 01:09:44 +00:00
sshClient , err := conn . SSHClient ( )
2022-08-02 14:44:59 +00:00
if err != nil {
return err
}
defer sshClient . Close ( )
sshSession , err := sshClient . NewSession ( )
if err != nil {
return err
2022-03-25 19:48:08 +00:00
}
2022-08-02 14:44:59 +00:00
defer sshSession . Close ( )
// Ensure context cancellation is propagated to the
// SSH session, e.g. to cancel `Wait()` at the end.
go func ( ) {
<- ctx . Done ( )
_ = sshSession . Close ( )
} ( )
2022-03-29 00:19:28 +00:00
2022-06-02 08:13:38 +00:00
if identityAgent == "" {
identityAgent = os . Getenv ( "SSH_AUTH_SOCK" )
}
if forwardAgent && identityAgent != "" {
err = gosshagent . ForwardToRemote ( sshClient , identityAgent )
2022-05-25 18:28:10 +00:00
if err != nil {
return xerrors . Errorf ( "forward agent failed: %w" , err )
}
err = gosshagent . RequestAgentForwarding ( sshSession )
if err != nil {
return xerrors . Errorf ( "request agent forwarding failed: %w" , err )
}
}
2022-07-26 14:23:32 +00:00
stdoutFile , validOut := cmd . OutOrStdout ( ) . ( * os . File )
stdinFile , validIn := cmd . InOrStdin ( ) . ( * os . File )
if validOut && validIn && isatty . IsTerminal ( stdoutFile . Fd ( ) ) {
state , err := term . MakeRaw ( int ( stdinFile . Fd ( ) ) )
2022-03-30 22:59:54 +00:00
if err != nil {
return err
}
defer func ( ) {
2022-07-26 14:23:32 +00:00
_ = term . Restore ( int ( stdinFile . Fd ( ) ) , state )
2022-03-30 22:59:54 +00:00
} ( )
2022-04-25 03:23:54 +00:00
2022-08-02 14:44:59 +00:00
windowChange := listenWindowSize ( ctx )
2022-04-25 03:23:54 +00:00
go func ( ) {
for {
select {
2022-08-02 14:44:59 +00:00
case <- ctx . Done ( ) :
2022-04-25 03:23:54 +00:00
return
case <- windowChange :
}
2022-08-02 14:44:59 +00:00
width , height , err := term . GetSize ( int ( stdoutFile . Fd ( ) ) )
if err != nil {
continue
}
2022-04-25 03:23:54 +00:00
_ = sshSession . WindowChange ( height , width )
}
} ( )
2022-03-30 22:59:54 +00:00
}
err = sshSession . RequestPty ( "xterm-256color" , 128 , 128 , gossh . TerminalModes { } )
2022-03-25 19:48:08 +00:00
if err != nil {
return err
}
2022-04-01 19:42:36 +00:00
2022-03-29 00:19:28 +00:00
sshSession . Stdin = cmd . InOrStdin ( )
sshSession . Stdout = cmd . OutOrStdout ( )
2022-08-02 14:44:59 +00:00
sshSession . Stderr = cmd . ErrOrStderr ( )
2022-04-01 19:42:36 +00:00
2022-03-25 19:48:08 +00:00
err = sshSession . Shell ( )
if err != nil {
return err
}
2022-04-01 19:42:36 +00:00
2022-08-02 14:44:59 +00:00
// Put cancel at the top of the defer stack to initiate
// shutdown of services.
defer cancel ( )
2022-09-12 16:27:51 +00:00
if validOut {
// Set initial window size.
width , height , err := term . GetSize ( int ( stdoutFile . Fd ( ) ) )
if err == nil {
_ = sshSession . WindowChange ( height , width )
}
}
2022-03-25 19:48:08 +00:00
err = sshSession . Wait ( )
if err != nil {
2022-07-25 15:25:34 +00:00
// If the connection drops unexpectedly, we get an ExitMissingError but no other
// error details, so try to at least give the user a better message
if errors . Is ( err , & gossh . ExitMissingError { } ) {
return xerrors . New ( "SSH connection ended unexpectedly" )
}
2022-03-25 19:48:08 +00:00
return err
2022-03-22 19:17:50 +00:00
}
return nil
} ,
}
2022-03-30 22:59:54 +00:00
cliflag . BoolVarP ( cmd . Flags ( ) , & stdio , "stdio" , "" , "CODER_SSH_STDIO" , false , "Specifies whether to emit SSH output over stdin/stdout." )
2022-05-17 22:55:58 +00:00
cliflag . BoolVarP ( cmd . Flags ( ) , & shuffle , "shuffle" , "" , "CODER_SSH_SHUFFLE" , false , "Specifies whether to choose a random workspace" )
_ = cmd . Flags ( ) . MarkHidden ( "shuffle" )
2022-05-25 18:28:10 +00:00
cliflag . BoolVarP ( cmd . Flags ( ) , & forwardAgent , "forward-agent" , "A" , "CODER_SSH_FORWARD_AGENT" , false , "Specifies whether to forward the SSH agent specified in $SSH_AUTH_SOCK" )
2022-06-02 08:13:38 +00:00
cliflag . StringVarP ( cmd . Flags ( ) , & identityAgent , "identity-agent" , "" , "CODER_SSH_IDENTITY_AGENT" , "" , "Specifies which identity agent to use (overrides $SSH_AUTH_SOCK), forward agent must also be enabled" )
2022-05-25 18:28:10 +00:00
cliflag . DurationVarP ( cmd . Flags ( ) , & wsPollInterval , "workspace-poll-interval" , "" , "CODER_WORKSPACE_POLL_INTERVAL" , workspacePollInterval , "Specifies how often to poll for workspace automated shutdown." )
2022-06-24 21:21:46 +00:00
cliflag . BoolVarP ( cmd . Flags ( ) , & wireguard , "wireguard" , "" , "CODER_SSH_WIREGUARD" , false , "Whether to use Wireguard for SSH tunneling." )
_ = cmd . Flags ( ) . MarkHidden ( "wireguard" )
2022-03-22 19:17:50 +00:00
return cmd
}
2022-05-13 17:09:04 +00:00
2022-05-18 14:10:40 +00:00
// getWorkspaceAgent returns the workspace and agent selected using either the
// `<workspace>[.<agent>]` syntax via `in` or picks a random workspace and agent
// if `shuffle` is true.
2022-08-02 14:44:59 +00:00
func getWorkspaceAndAgent ( ctx context . Context , cmd * cobra . Command , client * codersdk . Client , userID string , in string , shuffle bool ) ( codersdk . Workspace , codersdk . WorkspaceAgent , error ) { //nolint:revive
2022-05-18 14:10:40 +00:00
var (
workspace codersdk . Workspace
workspaceParts = strings . Split ( in , "." )
err error
)
if shuffle {
2022-08-02 14:44:59 +00:00
workspaces , err := client . Workspaces ( ctx , codersdk . WorkspaceFilter {
2022-06-03 19:36:08 +00:00
Owner : codersdk . Me ,
} )
2022-05-18 14:10:40 +00:00
if err != nil {
return codersdk . Workspace { } , codersdk . WorkspaceAgent { } , err
}
if len ( workspaces ) == 0 {
return codersdk . Workspace { } , codersdk . WorkspaceAgent { } , xerrors . New ( "no workspaces to shuffle" )
}
workspace , err = cryptorand . Element ( workspaces )
if err != nil {
return codersdk . Workspace { } , codersdk . WorkspaceAgent { } , err
}
} else {
2022-06-03 19:36:08 +00:00
workspace , err = namedWorkspace ( cmd , client , workspaceParts [ 0 ] )
2022-05-18 14:10:40 +00:00
if err != nil {
return codersdk . Workspace { } , codersdk . WorkspaceAgent { } , err
}
}
2022-05-19 18:04:44 +00:00
if workspace . LatestBuild . Transition != codersdk . WorkspaceTransitionStart {
2022-05-18 14:10:40 +00:00
return codersdk . Workspace { } , codersdk . WorkspaceAgent { } , xerrors . New ( "workspace must be in start transition to ssh" )
}
if workspace . LatestBuild . Job . CompletedAt == nil {
err := cliui . WorkspaceBuild ( ctx , cmd . ErrOrStderr ( ) , client , workspace . LatestBuild . ID , workspace . CreatedAt )
if err != nil {
return codersdk . Workspace { } , codersdk . WorkspaceAgent { } , err
}
}
2022-05-19 18:04:44 +00:00
if workspace . LatestBuild . Transition == codersdk . WorkspaceTransitionDelete {
2022-05-18 14:10:40 +00:00
return codersdk . Workspace { } , codersdk . WorkspaceAgent { } , xerrors . Errorf ( "workspace %q is being deleted" , workspace . Name )
}
resources , err := client . WorkspaceResourcesByBuild ( ctx , workspace . LatestBuild . ID )
if err != nil {
return codersdk . Workspace { } , codersdk . WorkspaceAgent { } , xerrors . Errorf ( "fetch workspace resources: %w" , err )
}
agents := make ( [ ] codersdk . WorkspaceAgent , 0 )
for _ , resource := range resources {
agents = append ( agents , resource . Agents ... )
}
if len ( agents ) == 0 {
return codersdk . Workspace { } , codersdk . WorkspaceAgent { } , xerrors . Errorf ( "workspace %q has no agents" , workspace . Name )
}
2022-09-01 01:09:44 +00:00
var workspaceAgent codersdk . WorkspaceAgent
2022-05-18 14:10:40 +00:00
if len ( workspaceParts ) >= 2 {
for _ , otherAgent := range agents {
if otherAgent . Name != workspaceParts [ 1 ] {
continue
}
2022-09-01 01:09:44 +00:00
workspaceAgent = otherAgent
2022-05-18 14:10:40 +00:00
break
}
2022-09-01 01:09:44 +00:00
if workspaceAgent . ID == uuid . Nil {
2022-05-18 14:10:40 +00:00
return codersdk . Workspace { } , codersdk . WorkspaceAgent { } , xerrors . Errorf ( "agent not found by name %q" , workspaceParts [ 1 ] )
}
}
2022-09-01 01:09:44 +00:00
if workspaceAgent . ID == uuid . Nil {
2022-05-18 14:10:40 +00:00
if len ( agents ) > 1 {
if ! shuffle {
return codersdk . Workspace { } , codersdk . WorkspaceAgent { } , xerrors . New ( "you must specify the name of an agent" )
}
2022-09-01 01:09:44 +00:00
workspaceAgent , err = cryptorand . Element ( agents )
2022-05-18 14:10:40 +00:00
if err != nil {
return codersdk . Workspace { } , codersdk . WorkspaceAgent { } , err
}
} else {
2022-09-01 01:09:44 +00:00
workspaceAgent = agents [ 0 ]
2022-05-18 14:10:40 +00:00
}
}
2022-09-01 01:09:44 +00:00
return workspace , workspaceAgent , nil
2022-05-18 14:10:40 +00:00
}
2022-05-13 17:09:04 +00:00
// Attempt to poll workspace autostop. We write a per-workspace lockfile to
// avoid spamming the user with notifications in case of multiple instances
// of the CLI running simultaneously.
func tryPollWorkspaceAutostop ( ctx context . Context , client * codersdk . Client , workspace codersdk . Workspace ) ( stop func ( ) ) {
lock := flock . New ( filepath . Join ( os . TempDir ( ) , "coder-autostop-notify-" + workspace . ID . String ( ) ) )
condition := notifyCondition ( ctx , client , workspace . ID , lock )
2022-05-20 10:57:02 +00:00
return notify . Notify ( condition , workspacePollInterval , autostopNotifyCountdown ... )
2022-05-13 17:09:04 +00:00
}
// Notify the user if the workspace is due to shutdown.
func notifyCondition ( ctx context . Context , client * codersdk . Client , workspaceID uuid . UUID , lock * flock . Flock ) notify . Condition {
return func ( now time . Time ) ( deadline time . Time , callback func ( ) ) {
// Keep trying to regain the lock.
2022-05-20 10:57:02 +00:00
locked , err := lock . TryLockContext ( ctx , workspacePollInterval )
2022-05-13 17:09:04 +00:00
if err != nil || ! locked {
return time . Time { } , nil
}
ws , err := client . Workspace ( ctx , workspaceID )
if err != nil {
return time . Time { } , nil
}
2022-06-02 10:23:34 +00:00
if ptr . NilOrZero ( ws . TTLMillis ) {
2022-05-13 17:09:04 +00:00
return time . Time { } , nil
}
2022-08-25 16:10:42 +00:00
deadline = ws . LatestBuild . Deadline . Time
2022-05-13 17:09:04 +00:00
callback = func ( ) {
ttl := deadline . Sub ( now )
var title , body string
if ttl > time . Minute {
2022-05-27 19:04:33 +00:00
title = fmt . Sprintf ( ` Workspace %s stopping soon ` , ws . Name )
2022-05-13 17:09:04 +00:00
body = fmt . Sprintf (
2022-05-27 19:04:33 +00:00
` Your Coder workspace %s is scheduled to stop in %.0f mins ` , ws . Name , ttl . Minutes ( ) )
2022-05-13 17:09:04 +00:00
} else {
title = fmt . Sprintf ( "Workspace %s stopping!" , ws . Name )
body = fmt . Sprintf ( "Your Coder workspace %s is stopping any time now!" , ws . Name )
}
// notify user with a native system notification (best effort)
_ = beeep . Notify ( title , body , "" )
}
return deadline . Truncate ( time . Minute ) , callback
}
}