2023-01-10 04:23:17 +00:00
package cli
import (
"context"
"encoding/json"
"fmt"
"io"
"net/url"
"os"
"path/filepath"
"strconv"
"strings"
"time"
"github.com/spf13/afero"
"golang.org/x/xerrors"
"tailscale.com/tailcfg"
2023-02-10 03:43:18 +00:00
"tailscale.com/types/netlogtype"
2023-01-10 04:23:17 +00:00
2023-06-21 20:22:43 +00:00
"cdr.dev/slog"
"cdr.dev/slog/sloggers/sloghuman"
2024-02-02 11:18:26 +00:00
"github.com/coder/coder/v2/cli/cliui"
2023-11-14 12:38:34 +00:00
"github.com/coder/coder/v2/cli/cliutil"
2023-08-18 18:55:43 +00:00
"github.com/coder/coder/v2/codersdk"
2024-03-26 17:44:31 +00:00
"github.com/coder/coder/v2/codersdk/workspacesdk"
2024-03-15 16:24:38 +00:00
"github.com/coder/serpent"
2023-01-10 04:23:17 +00:00
)
// vscodeSSH is used by the Coder VS Code extension to establish
// a connection to a workspace.
//
// This command needs to remain stable for compatibility with
// various VS Code versions, so it's kept separate from our
// standard SSH command.
2024-03-17 14:45:26 +00:00
func ( r * RootCmd ) vscodeSSH ( ) * serpent . Command {
2023-01-10 04:23:17 +00:00
var (
sessionTokenFile string
urlFile string
2023-11-03 16:21:31 +00:00
logDir string
2023-01-10 04:23:17 +00:00
networkInfoDir string
networkInfoInterval time . Duration
2024-02-02 11:18:26 +00:00
waitEnum string
2023-01-10 04:23:17 +00:00
)
2024-03-17 14:45:26 +00:00
cmd := & serpent . Command {
2023-01-10 04:23:17 +00:00
// A SSH config entry is added by the VS Code extension that
// passes %h to ProxyCommand. The prefix of `coder-vscode--`
2023-09-08 14:01:57 +00:00
// is a magical string represented in our VS Code extension.
2023-01-10 04:23:17 +00:00
// It's not important here, only the delimiter `--` is.
2023-09-08 14:01:57 +00:00
Use : "vscodessh <coder-vscode--<owner>--<workspace>--<agent?>>" ,
2023-03-23 22:42:20 +00:00
Hidden : true ,
2024-03-15 16:24:38 +00:00
Middleware : serpent . RequireNArgs ( 1 ) ,
Handler : func ( inv * serpent . Invocation ) error {
2023-01-10 04:23:17 +00:00
if networkInfoDir == "" {
return xerrors . New ( "network-info-dir must be specified" )
}
if sessionTokenFile == "" {
return xerrors . New ( "session-token-file must be specified" )
}
if urlFile == "" {
return xerrors . New ( "url-file must be specified" )
}
2023-03-23 22:42:20 +00:00
fs , ok := inv . Context ( ) . Value ( "fs" ) . ( afero . Fs )
2023-01-10 04:23:17 +00:00
if ! ok {
fs = afero . NewOsFs ( )
}
sessionToken , err := afero . ReadFile ( fs , sessionTokenFile )
if err != nil {
return xerrors . Errorf ( "read session token: %w" , err )
}
rawURL , err := afero . ReadFile ( fs , urlFile )
if err != nil {
return xerrors . Errorf ( "read url: %w" , err )
}
serverURL , err := url . Parse ( string ( rawURL ) )
if err != nil {
return xerrors . Errorf ( "parse url: %w" , err )
}
2023-03-23 22:42:20 +00:00
ctx , cancel := context . WithCancel ( inv . Context ( ) )
2023-01-10 04:23:17 +00:00
defer cancel ( )
2023-02-19 00:32:09 +00:00
err = fs . MkdirAll ( networkInfoDir , 0 o700 )
2023-01-10 04:23:17 +00:00
if err != nil {
return xerrors . Errorf ( "mkdir: %w" , err )
}
client := codersdk . New ( serverURL )
client . SetSessionToken ( string ( sessionToken ) )
2023-04-18 13:07:10 +00:00
// This adds custom headers to the request!
2024-03-25 19:01:42 +00:00
err = r . configureClient ( ctx , client , serverURL , inv )
2023-04-18 13:07:10 +00:00
if err != nil {
return xerrors . Errorf ( "set client: %w" , err )
}
2023-03-23 22:42:20 +00:00
parts := strings . Split ( inv . Args [ 0 ] , "--" )
2023-01-10 04:23:17 +00:00
if len ( parts ) < 3 {
2023-09-08 14:01:57 +00:00
return xerrors . Errorf ( "invalid argument format. must be: coder-vscode--<owner>--<name>--<agent?>" )
2023-01-10 04:23:17 +00:00
}
owner := parts [ 1 ]
name := parts [ 2 ]
2024-02-02 11:18:26 +00:00
if len ( parts ) > 3 {
name += "." + parts [ 3 ]
}
// Set autostart to false because it's assumed the VS Code extension
// will call this command after the workspace is started.
autostart := false
2023-01-10 04:23:17 +00:00
2024-04-13 18:39:57 +00:00
_ , workspaceAgent , err := getWorkspaceAndAgent ( ctx , inv , client , autostart , fmt . Sprintf ( "%s/%s" , owner , name ) )
2023-01-10 04:23:17 +00:00
if err != nil {
2024-02-02 11:18:26 +00:00
return xerrors . Errorf ( "find workspace and agent: %w" , err )
2023-01-10 04:23:17 +00:00
}
2023-02-10 03:43:18 +00:00
2024-02-02 11:18:26 +00:00
// Select the startup script behavior based on template configuration or flags.
var wait bool
switch waitEnum {
case "yes" :
wait = true
case "no" :
wait = false
case "auto" :
for _ , script := range workspaceAgent . Scripts {
if script . StartBlocksLogin {
wait = true
2023-01-10 04:23:17 +00:00
break
}
}
2024-02-02 11:18:26 +00:00
default :
return xerrors . Errorf ( "unknown wait value %q" , waitEnum )
}
err = cliui . Agent ( ctx , inv . Stderr , workspaceAgent . ID , cliui . AgentOptions {
Fetch : client . WorkspaceAgent ,
FetchLogs : client . WorkspaceAgentLogsAfter ,
Wait : wait ,
} )
if err != nil {
if xerrors . Is ( err , context . Canceled ) {
return cliui . Canceled
2023-01-10 04:23:17 +00:00
}
}
2023-02-10 03:43:18 +00:00
2023-11-03 16:21:31 +00:00
// The VS Code extension obtains the PID of the SSH process to
// read files to display logs and network info.
//
// We get the parent PID because it's assumed `ssh` is calling this
// command via the ProxyCommand SSH option.
pid := os . Getppid ( )
2023-11-14 18:56:27 +00:00
logger := inv . Logger
2023-11-03 16:21:31 +00:00
if logDir != "" {
logFilePath := filepath . Join ( logDir , fmt . Sprintf ( "%d.log" , pid ) )
logFile , err := fs . OpenFile ( logFilePath , os . O_CREATE | os . O_WRONLY , 0 o600 )
if err != nil {
return xerrors . Errorf ( "open log file %q: %w" , logFilePath , err )
}
2023-11-14 12:38:34 +00:00
dc := cliutil . DiscardAfterClose ( logFile )
defer dc . Close ( )
logger = logger . AppendSinks ( sloghuman . Sink ( dc ) ) . Leveled ( slog . LevelDebug )
2023-06-21 20:22:43 +00:00
}
if r . disableDirect {
2023-11-03 16:21:31 +00:00
logger . Info ( ctx , "direct connections disabled" )
2023-06-21 20:22:43 +00:00
}
2024-03-26 17:44:31 +00:00
agentConn , err := workspacesdk . New ( client ) .
DialAgent ( ctx , workspaceAgent . ID , & workspacesdk . DialAgentOptions {
Logger : logger ,
BlockEndpoints : r . disableDirect ,
} )
2023-01-10 04:23:17 +00:00
if err != nil {
return xerrors . Errorf ( "dial workspace agent: %w" , err )
}
defer agentConn . Close ( )
2023-02-10 03:43:18 +00:00
2023-01-10 04:23:17 +00:00
agentConn . AwaitReachable ( ctx )
rawSSH , err := agentConn . SSH ( ctx )
if err != nil {
return err
}
defer rawSSH . Close ( )
2023-02-10 03:43:18 +00:00
2023-01-10 04:23:17 +00:00
// Copy SSH traffic over stdio.
go func ( ) {
2023-03-23 22:42:20 +00:00
_ , _ = io . Copy ( inv . Stdout , rawSSH )
2023-01-10 04:23:17 +00:00
} ( )
go func ( ) {
2023-03-23 22:42:20 +00:00
_ , _ = io . Copy ( rawSSH , inv . Stdin )
2023-01-10 04:23:17 +00:00
} ( )
2023-02-10 03:43:18 +00:00
2023-01-10 04:23:17 +00:00
// The VS Code extension obtains the PID of the SSH process to
// read the file below which contains network information to display.
//
// We get the parent PID because it's assumed `ssh` is calling this
// command via the ProxyCommand SSH option.
2023-11-03 16:21:31 +00:00
networkInfoFilePath := filepath . Join ( networkInfoDir , fmt . Sprintf ( "%d.json" , pid ) )
2023-02-10 03:43:18 +00:00
statsErrChan := make ( chan error , 1 )
cb := func ( start , end time . Time , virtual , _ map [ netlogtype . Connection ] netlogtype . Counts ) {
sendErr := func ( err error ) {
select {
case statsErrChan <- err :
default :
}
2023-01-10 04:23:17 +00:00
}
2023-02-10 03:43:18 +00:00
stats , err := collectNetworkStats ( ctx , agentConn , start , end , virtual )
2023-01-10 04:23:17 +00:00
if err != nil {
2023-02-10 03:43:18 +00:00
sendErr ( err )
return
2023-01-10 04:23:17 +00:00
}
2023-02-10 03:43:18 +00:00
2023-01-10 04:23:17 +00:00
rawStats , err := json . Marshal ( stats )
if err != nil {
2023-02-10 03:43:18 +00:00
sendErr ( err )
return
2023-01-10 04:23:17 +00:00
}
2023-02-19 00:32:09 +00:00
err = afero . WriteFile ( fs , networkInfoFilePath , rawStats , 0 o600 )
2023-01-10 04:23:17 +00:00
if err != nil {
2023-02-10 03:43:18 +00:00
sendErr ( err )
return
2023-01-10 04:23:17 +00:00
}
2023-02-10 03:43:18 +00:00
}
now := time . Now ( )
cb ( now , now . Add ( time . Nanosecond ) , map [ netlogtype . Connection ] netlogtype . Counts { } , map [ netlogtype . Connection ] netlogtype . Counts { } )
agentConn . SetConnStatsCallback ( networkInfoInterval , 2048 , cb )
select {
case <- ctx . Done ( ) :
return nil
case err := <- statsErrChan :
return err
2023-01-10 04:23:17 +00:00
}
} ,
}
2024-03-15 16:24:38 +00:00
cmd . Options = serpent . OptionSet {
2023-03-23 22:42:20 +00:00
{
Flag : "network-info-dir" ,
Description : "Specifies a directory to write network information periodically." ,
2024-03-15 16:24:38 +00:00
Value : serpent . StringOf ( & networkInfoDir ) ,
2023-03-23 22:42:20 +00:00
} ,
2023-11-03 16:21:31 +00:00
{
Flag : "log-dir" ,
Description : "Specifies a directory to write logs to." ,
2024-03-15 16:24:38 +00:00
Value : serpent . StringOf ( & logDir ) ,
2023-11-03 16:21:31 +00:00
} ,
2023-03-23 22:42:20 +00:00
{
Flag : "session-token-file" ,
Description : "Specifies a file that contains a session token." ,
2024-03-15 16:24:38 +00:00
Value : serpent . StringOf ( & sessionTokenFile ) ,
2023-03-23 22:42:20 +00:00
} ,
{
Flag : "url-file" ,
Description : "Specifies a file that contains the Coder URL." ,
2024-03-15 16:24:38 +00:00
Value : serpent . StringOf ( & urlFile ) ,
2023-03-23 22:42:20 +00:00
} ,
{
Flag : "network-info-interval" ,
Description : "Specifies the interval to update network information." ,
Default : "5s" ,
2024-03-15 16:24:38 +00:00
Value : serpent . DurationOf ( & networkInfoInterval ) ,
2023-03-23 22:42:20 +00:00
} ,
2024-02-02 11:18:26 +00:00
{
Flag : "wait" ,
Description : "Specifies whether or not to wait for the startup script to finish executing. Auto means that the agent startup script behavior configured in the workspace template is used." ,
Default : "auto" ,
2024-03-15 16:24:38 +00:00
Value : serpent . EnumOf ( & waitEnum , "yes" , "no" , "auto" ) ,
2024-02-02 11:18:26 +00:00
} ,
2023-03-23 22:42:20 +00:00
}
2023-01-10 04:23:17 +00:00
return cmd
}
type sshNetworkStats struct {
P2P bool ` json:"p2p" `
Latency float64 ` json:"latency" `
PreferredDERP string ` json:"preferred_derp" `
DERPLatency map [ string ] float64 ` json:"derp_latency" `
UploadBytesSec int64 ` json:"upload_bytes_sec" `
DownloadBytesSec int64 ` json:"download_bytes_sec" `
}
2024-03-26 17:44:31 +00:00
func collectNetworkStats ( ctx context . Context , agentConn * workspacesdk . AgentConn , start , end time . Time , counts map [ netlogtype . Connection ] netlogtype . Counts ) ( * sshNetworkStats , error ) {
2023-08-02 15:30:43 +00:00
latency , p2p , pingResult , err := agentConn . Ping ( ctx )
2023-01-10 04:23:17 +00:00
if err != nil {
return nil , err
}
node := agentConn . Node ( )
derpMap := agentConn . DERPMap ( )
derpLatency := map [ string ] float64 { }
// Convert DERP region IDs to friendly names for display in the UI.
for rawRegion , latency := range node . DERPLatency {
regionParts := strings . SplitN ( rawRegion , "-" , 2 )
regionID , err := strconv . Atoi ( regionParts [ 0 ] )
if err != nil {
continue
}
region , found := derpMap . Regions [ regionID ]
if ! found {
// It's possible that a workspace agent is using an old DERPMap
// and reports regions that do not exist. If that's the case,
// report the region as unknown!
region = & tailcfg . DERPRegion {
RegionID : regionID ,
RegionName : fmt . Sprintf ( "Unnamed %d" , regionID ) ,
}
}
// Convert the microseconds to milliseconds.
derpLatency [ region . RegionName ] = latency * 1000
}
totalRx := uint64 ( 0 )
totalTx := uint64 ( 0 )
2023-02-10 03:43:18 +00:00
for _ , stat := range counts {
2023-01-10 04:23:17 +00:00
totalRx += stat . RxBytes
totalTx += stat . TxBytes
}
// Tracking the time since last request is required because
// ExtractTrafficStats() resets its counters after each call.
2023-02-10 03:43:18 +00:00
dur := end . Sub ( start )
2023-01-10 04:23:17 +00:00
uploadSecs := float64 ( totalTx ) / dur . Seconds ( )
downloadSecs := float64 ( totalRx ) / dur . Seconds ( )
2023-08-02 15:30:43 +00:00
// Sometimes the preferred DERP doesn't match the one we're actually
// connected with. Perhaps because the agent prefers a different DERP and
// we're using that server instead.
preferredDerpID := node . PreferredDERP
if pingResult . DERPRegionID != 0 {
preferredDerpID = pingResult . DERPRegionID
}
preferredDerp , ok := derpMap . Regions [ preferredDerpID ]
preferredDerpName := fmt . Sprintf ( "Unnamed %d" , preferredDerpID )
if ok {
preferredDerpName = preferredDerp . RegionName
}
if _ , ok := derpLatency [ preferredDerpName ] ; ! ok {
derpLatency [ preferredDerpName ] = 0
}
2023-01-10 04:23:17 +00:00
return & sshNetworkStats {
P2P : p2p ,
Latency : float64 ( latency . Microseconds ( ) ) / 1000 ,
2023-08-02 15:30:43 +00:00
PreferredDERP : preferredDerpName ,
2023-01-10 04:23:17 +00:00
DERPLatency : derpLatency ,
UploadBytesSec : int64 ( uploadSecs ) ,
DownloadBytesSec : int64 ( downloadSecs ) ,
} , nil
}