chore: use access url from env resource pool (coder#216)

Author: cmoog

coder-sdk/env.go

@@ -3,6 +3,7 @@ package coder
 import (
 	"context"
 	"net/http"
+	"net/url"
 	"time"
 
 	"cdr.dev/wsep"
@@ -32,7 +33,6 @@ type Environment struct {
 	LastOpenedAt     time.Time        `json:"last_opened_at"     table:"-"`
 	LastConnectionAt time.Time        `json:"last_connection_at" table:"-"`
 	AutoOffThreshold Duration         `json:"auto_off_threshold" table:"-"`
-	SSHAvailable     bool             `json:"ssh_available"      table:"-"`
 	UseContainerVM   bool             `json:"use_container_vm"   table:"CVM"`
 	ResourcePoolID   string           `json:"resource_pool_id"   table:"-"`
 }
@@ -147,22 +147,22 @@ func (c Client) EditEnvironment(ctx context.Context, envID string, req UpdateEnv
 
 // DialWsep dials an environments command execution interface
 // See https://github.com/cdr/wsep for details.
-func (c Client) DialWsep(ctx context.Context, envID string) (*websocket.Conn, error) {
-	return c.dialWebsocket(ctx, "/proxy/environments/"+envID+"/wsep")
+func (c Client) DialWsep(ctx context.Context, baseURL *url.URL, envID string) (*websocket.Conn, error) {
+	return c.dialWebsocket(ctx, "/proxy/environments/"+envID+"/wsep", withBaseURL(baseURL))
 }
 
 // DialExecutor gives a remote execution interface for performing commands inside an environment.
-func (c Client) DialExecutor(ctx context.Context, envID string) (wsep.Execer, error) {
-	ws, err := c.DialWsep(ctx, envID)
+func (c Client) DialExecutor(ctx context.Context, baseURL *url.URL, envID string) (wsep.Execer, error) {
+	ws, err := c.DialWsep(ctx, baseURL, envID)
 	if err != nil {
 		return nil, err
 	}
 	return wsep.RemoteExecer(ws), nil
 }
 
 // DialIDEStatus opens a websocket connection for cpu load metrics on the environment.
-func (c Client) DialIDEStatus(ctx context.Context, envID string) (*websocket.Conn, error) {
-	return c.dialWebsocket(ctx, "/proxy/environments/"+envID+"/ide/api/status")
+func (c Client) DialIDEStatus(ctx context.Context, baseURL *url.URL, envID string) (*websocket.Conn, error) {
+	return c.dialWebsocket(ctx, "/proxy/environments/"+envID+"/ide/api/status", withBaseURL(baseURL))
 }
 
 // DialEnvironmentBuildLog opens a websocket connection for the environment build log messages.

coder-sdk/request.go

@@ -12,12 +12,21 @@ import (
 )
 
 // request is a helper to set the cookie, marshal the payload and execute the request.
-func (c Client) request(ctx context.Context, method, path string, in interface{}) (*http.Response, error) {
+func (c Client) request(ctx context.Context, method, path string, in interface{}, options ...requestOption) (*http.Response, error) {
 	// Create a default http client with the auth in the cookie.
 	client, err := c.newHTTPClient()
 	if err != nil {
 		return nil, xerrors.Errorf("new http client: %w", err)
 	}
+	url := *c.BaseURL
+
+	var config requestOptions
+	for _, o := range options {
+		o(&config)
+	}
+	if config.BaseURLOverride != nil {
+		url = *config.BaseURLOverride
+	}
 
 	// If we have incoming data, encode it as json.
 	var payload io.Reader
@@ -30,7 +39,7 @@ func (c Client) request(ctx context.Context, method, path string, in interface{}
 	}
 
 	// Create the http request.
-	req, err := http.NewRequestWithContext(ctx, method, c.BaseURL.String()+path, payload)
+	req, err := http.NewRequestWithContext(ctx, method, url.String()+path, payload)
 	if err != nil {
 		return nil, xerrors.Errorf("create request: %w", err)
 	}

coder-sdk/resourcepools.go

@@ -17,6 +17,8 @@ type ResourcePool struct {
 	DevurlHost          string   `json:"devurl_host"`
 	NamespaceWhitelist  []string `json:"namespace_whitelist"`
 	OrgWhitelist        []string `json:"org_whitelist"`
+	SSHEnabled          bool     `json:"ssh_enabled"`
+	AccessURL           string   `json:"envproxy_access_url"`
 }
 
 // ResourcePoolByID fetches a resource pool entity by its unique ID.

coder-sdk/ws.go

@@ -3,14 +3,35 @@ package coder
 import (
 	"context"
 	"net/http"
+	"net/url"
 
 	"nhooyr.io/websocket"
 )
 
+type requestOptions struct {
+	BaseURLOverride *url.URL
+}
+
+type requestOption func(*requestOptions)
+
+func withBaseURL(base *url.URL) func(o *requestOptions) {
+	return func(o *requestOptions) {
+		o.BaseURLOverride = base
+	}
+}
+
 // dialWebsocket establish the websocket connection while setting the authentication header.
-func (c Client) dialWebsocket(ctx context.Context, path string) (*websocket.Conn, error) {
+func (c Client) dialWebsocket(ctx context.Context, path string, options ...requestOption) (*websocket.Conn, error) {
 	// Make a copy of the url so we can update the scheme to ws(s) without mutating the state.
 	url := *c.BaseURL
+	var config requestOptions
+	for _, o := range options {
+		o(&config)
+	}
+	if config.BaseURLOverride != nil {
+		url = *config.BaseURLOverride
+	}
+
 	if url.Scheme == "https" {
 		url.Scheme = "wss"
 	} else {

internal/cmd/configssh.go

@@ -4,17 +4,16 @@ import (
 	"context"
 	"fmt"
 	"io/ioutil"
-	"net"
 	"net/url"
 	"os"
 	"os/user"
 	"path/filepath"
 	"strings"
-	"time"
 
 	"cdr.dev/coder-cli/pkg/clog"
 
 	"cdr.dev/coder-cli/coder-sdk"
+	"cdr.dev/coder-cli/internal/coderutil"
 	"cdr.dev/coder-cli/internal/config"
 	"github.com/spf13/cobra"
 	"golang.org/x/xerrors"
@@ -103,20 +102,17 @@ func configSSH(configpath *string, remove *bool) func(cmd *cobra.Command, _ []st
 			return xerrors.New("no environments found")
 		}
 
-		if !sshAvailable(envs) {
-			return xerrors.New("SSH is disabled or not available for any environments in your Coder Enterprise deployment.")
-		}
-
-		err = canConnectSSH(ctx)
+		envsWithPools, err := coderutil.EnvsWithPool(ctx, client, envs)
 		if err != nil {
-			return xerrors.Errorf("check if SSH is available: unable to connect to SSH endpoint: %w", err)
+			return xerrors.Errorf("resolve env pools: %w", err)
 		}
 
-		newConfig, err := makeNewConfigs(user.Username, envs, privateKeyFilepath)
-		if err != nil {
-			return xerrors.Errorf("make new ssh configurations: %w", err)
+		if !sshAvailable(envsWithPools) {
+			return xerrors.New("SSH is disabled or not available for any environments in your Coder Enterprise deployment.")
 		}
 
+		newConfig := makeNewConfigs(user.Username, envsWithPools, privateKeyFilepath)
+
 		err = os.MkdirAll(filepath.Dir(*configpath), os.ModePerm)
 		if err != nil {
 			return xerrors.Errorf("make configuration directory: %w", err)
@@ -159,42 +155,15 @@ func removeOldConfig(config string) (string, bool) {
 }
 
 // sshAvailable returns true if SSH is available for at least one environment.
-func sshAvailable(envs []coder.Environment) bool {
+func sshAvailable(envs []coderutil.EnvWithPool) bool {
 	for _, env := range envs {
-		if env.SSHAvailable {
+		if env.Pool.SSHEnabled {
 			return true
 		}
 	}
-
 	return false
 }
 
-// canConnectSSH returns an error if we cannot dial the SSH port.
-func canConnectSSH(ctx context.Context) error {
-	ctx, cancel := context.WithTimeout(ctx, 3*time.Second)
-	defer cancel()
-
-	host, err := configuredHostname()
-	if err != nil {
-		return xerrors.Errorf("get configured manager hostname: %w", err)
-	}
-
-	var (
-		dialer   net.Dialer
-		hostPort = net.JoinHostPort(host, "22")
-	)
-	conn, err := dialer.DialContext(ctx, "tcp", hostPort)
-	if err != nil {
-		if err == context.DeadlineExceeded {
-			err = xerrors.New("timed out after 3 seconds")
-		}
-		return xerrors.Errorf("dial tcp://%v: %w", hostPort, err)
-	}
-	conn.Close()
-
-	return nil
-}
-
 func writeSSHKey(ctx context.Context, client *coder.Client, privateKeyPath string) error {
 	key, err := client.SSHKey(ctx)
 	if err != nil {
@@ -203,23 +172,26 @@ func writeSSHKey(ctx context.Context, client *coder.Client, privateKeyPath strin
 	return ioutil.WriteFile(privateKeyPath, []byte(key.PrivateKey), 0400)
 }
 
-func makeNewConfigs(userName string, envs []coder.Environment, privateKeyFilepath string) (string, error) {
-	hostname, err := configuredHostname()
-	if err != nil {
-		return "", err
-	}
-
+func makeNewConfigs(userName string, envs []coderutil.EnvWithPool, privateKeyFilepath string) string {
 	newConfig := fmt.Sprintf("\n%s\n%s\n\n", sshStartToken, sshStartMessage)
 	for _, env := range envs {
-		if !env.SSHAvailable {
+		if !env.Pool.SSHEnabled {
+			clog.LogWarn(fmt.Sprintf("SSH is not enabled for pool %q", env.Pool.Name),
+				clog.BlankLine,
+				clog.Tipf("ask an infrastructure administrator to enable SSH for this resource pool"),
+			)
 			continue
 		}
-
-		newConfig += makeSSHConfig(hostname, userName, env.Name, privateKeyFilepath)
+		u, err := url.Parse(env.Pool.AccessURL)
+		if err != nil {
+			clog.LogWarn("invalid access url", clog.Causef("malformed url: %q", env.Pool.AccessURL))
+			continue
+		}
+		newConfig += makeSSHConfig(u.Host, userName, env.Env.Name, privateKeyFilepath)
 	}
 	newConfig += fmt.Sprintf("\n%s\n", sshEndToken)
 
-	return newConfig, nil
+	return newConfig
 }
 
 func makeSSHConfig(host, userName, envName, privateKeyFilepath string) string {
@@ -235,6 +207,7 @@ func makeSSHConfig(host, userName, envName, privateKeyFilepath string) string {
 `, envName, host, userName, envName, privateKeyFilepath)
 }
 
+//nolint:deadcode,unused
 func configuredHostname() (string, error) {
 	u, err := config.URL.Read()
 	if err != nil {

internal/cmd/shell.go

@@ -16,6 +16,7 @@ import (
 
 	"cdr.dev/coder-cli/coder-sdk"
 	"cdr.dev/coder-cli/internal/activity"
+	"cdr.dev/coder-cli/internal/coderutil"
 	"cdr.dev/coder-cli/internal/x/xterminal"
 	"cdr.dev/coder-cli/pkg/clog"
 	"cdr.dev/wsep"
@@ -161,9 +162,9 @@ func runCommand(ctx context.Context, envName, command string, args []string) err
 	ctx, cancel := context.WithCancel(ctx)
 	defer cancel()
 
-	conn, err := client.DialWsep(ctx, env.ID)
+	conn, err := coderutil.DialEnvWsep(ctx, client, env)
 	if err != nil {
-		return xerrors.Errorf("dial websocket: %w", err)
+		return xerrors.Errorf("dial executor: %w", err)
 	}
 	go heartbeat(ctx, conn, 15*time.Second)
 

internal/coderutil/doc.go

@@ -0,0 +1,2 @@
+// Package coderutil providers utilities for high-level operations on coder-sdk entities.
+package coderutil

internal/coderutil/env.go

@@ -0,0 +1,59 @@
+package coderutil
+
+import (
+	"context"
+	"net/url"
+
+	"cdr.dev/coder-cli/coder-sdk"
+	"golang.org/x/xerrors"
+	"nhooyr.io/websocket"
+)
+
+// DialEnvWsep dials the executor endpoint using the https://github.com/cdr/wsep message protocol.
+// The proper resource pool access URL is used.
+func DialEnvWsep(ctx context.Context, client *coder.Client, env *coder.Environment) (*websocket.Conn, error) {
+	resourcePool, err := client.ResourcePoolByID(ctx, env.ResourcePoolID)
+	if err != nil {
+		return nil, xerrors.Errorf("get env resource pool: %w", err)
+	}
+	accessURL, err := url.Parse(resourcePool.AccessURL)
+	if err != nil {
+		return nil, xerrors.Errorf("invalid resource pool access url: %w", err)
+	}
+
+	conn, err := client.DialWsep(ctx, accessURL, env.ID)
+	if err != nil {
+		return nil, xerrors.Errorf("dial websocket: %w", err)
+	}
+	return conn, nil
+}
+
+// EnvWithPool composes an Environment entity with its associated ResourcePool.
+type EnvWithPool struct {
+	Env  coder.Environment
+	Pool coder.ResourcePool
+}
+
+// EnvsWithPool performs the composition of each Environment with its associated ResourcePool.
+func EnvsWithPool(ctx context.Context, client *coder.Client, envs []coder.Environment) ([]EnvWithPool, error) {
+	pooledEnvs := make([]EnvWithPool, len(envs))
+	pools, err := client.ResourcePools(ctx)
+	if err != nil {
+		return nil, err
+	}
+	poolMap := make(map[string]coder.ResourcePool, len(pools))
+	for _, p := range pools {
+		poolMap[p.ID] = p
+	}
+	for _, e := range envs {
+		envPool, ok := poolMap[e.ResourcePoolID]
+		if !ok {
+			return nil, xerrors.Errorf("fetch env resource pool: %w", coder.ErrNotFound)
+		}
+		pooledEnvs = append(pooledEnvs, EnvWithPool{
+			Env:  e,
+			Pool: envPool,
+		})
+	}
+	return pooledEnvs, nil
+}

internal/sync/singlefile.go

@@ -10,14 +10,15 @@ import (
 	"strings"
 
 	"cdr.dev/coder-cli/coder-sdk"
+	"cdr.dev/coder-cli/internal/coderutil"
 	"cdr.dev/wsep"
 	"golang.org/x/xerrors"
 	"nhooyr.io/websocket"
 )
 
 // SingleFile copies the given file into the remote dir or remote path of the given coder.Environment.
 func SingleFile(ctx context.Context, local, remoteDir string, env *coder.Environment, client *coder.Client) error {
-	conn, err := client.DialWsep(ctx, env.ID)
+	conn, err := coderutil.DialEnvWsep(ctx, client, env)
 	if err != nil {
 		return xerrors.Errorf("dial remote execer: %w", err)
 	}

internal/sync/sync.go

@@ -24,6 +24,7 @@ import (
 
 	"cdr.dev/coder-cli/coder-sdk"
 	"cdr.dev/coder-cli/internal/activity"
+	"cdr.dev/coder-cli/internal/coderutil"
 	"cdr.dev/coder-cli/pkg/clog"
 	"cdr.dev/wsep"
 )
@@ -89,9 +90,9 @@ func (s Sync) syncPaths(delete bool, local, remote string) error {
 }
 
 func (s Sync) remoteCmd(ctx context.Context, prog string, args ...string) error {
-	conn, err := s.Client.DialWsep(ctx, s.Env.ID)
+	conn, err := coderutil.DialEnvWsep(ctx, s.Client, &s.Env)
 	if err != nil {
-		return xerrors.Errorf("dial websocket: %w", err)
+		return xerrors.Errorf("dial executor: %w", err)
 	}
 	defer func() { _ = conn.Close(websocket.CloseNormalClosure, "") }() // Best effort.
 
@@ -270,9 +271,9 @@ func (s Sync) Version() (string, error) {
 	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
 	defer cancel()
 
-	conn, err := s.Client.DialWsep(ctx, s.Env.ID)
+	conn, err := coderutil.DialEnvWsep(ctx, s.Client, &s.Env)
 	if err != nil {
-		return "", err
+		return "", xerrors.Errorf("dial env executor: %w", err)
 	}
 	defer func() { _ = conn.Close(websocket.CloseNormalClosure, "") }() // Best effort.