mirror of
https://github.com/coder/coder.git
synced 2026-09-24 15:04:27 +08:00
feat: Separate workspace agent for tests (#567)
This adds tests for Google Cloud authentication, and lays the ground-work for future agent auth types in the future.
This commit is contained in:
Kyle Carberry
+63
-47
@@ -1,17 +1,15 @@
|
||||
package cli
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"os"
|
||||
|
||||
"github.com/pion/webrtc/v3"
|
||||
"github.com/spf13/cobra"
|
||||
"golang.org/x/crypto/ssh"
|
||||
"golang.org/x/term"
|
||||
"golang.org/x/xerrors"
|
||||
|
||||
"github.com/coder/coder/agent"
|
||||
"github.com/coder/coder/peer"
|
||||
"github.com/coder/coder/peerbroker"
|
||||
"github.com/coder/coder/codersdk"
|
||||
"github.com/coder/coder/database"
|
||||
)
|
||||
|
||||
func workspaceSSH() *cobra.Command {
|
||||
@@ -26,58 +24,76 @@ func workspaceSSH() *cobra.Command {
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if workspace.LatestBuild.Transition == database.WorkspaceTransitionDelete {
|
||||
return xerrors.New("workspace is deleting...")
|
||||
}
|
||||
resources, err := client.WorkspaceResourcesByBuild(cmd.Context(), workspace.LatestBuild.ID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
resourceByAddress := make(map[string]codersdk.WorkspaceResource)
|
||||
for _, resource := range resources {
|
||||
_, _ = fmt.Printf("Got resource: %+v\n", resource)
|
||||
if resource.Agent == nil {
|
||||
continue
|
||||
}
|
||||
resourceByAddress[resource.Address] = resource
|
||||
}
|
||||
var resourceAddress string
|
||||
if len(args) >= 2 {
|
||||
resourceAddress = args[1]
|
||||
} else {
|
||||
// No resource name was provided!
|
||||
if len(resourceByAddress) > 1 {
|
||||
// List available resources to connect into?
|
||||
return xerrors.Errorf("multiple agents")
|
||||
}
|
||||
for _, resource := range resourceByAddress {
|
||||
resourceAddress = resource.Address
|
||||
break
|
||||
}
|
||||
}
|
||||
resource, exists := resourceByAddress[resourceAddress]
|
||||
if !exists {
|
||||
resourceKeys := make([]string, 0)
|
||||
for resourceKey := range resourceByAddress {
|
||||
resourceKeys = append(resourceKeys, resourceKey)
|
||||
}
|
||||
return xerrors.Errorf("no sshable agent with address %q: %+v", resourceAddress, resourceKeys)
|
||||
}
|
||||
if resource.Agent.LastConnectedAt == nil {
|
||||
return xerrors.Errorf("agent hasn't connected yet")
|
||||
}
|
||||
|
||||
dialed, err := client.DialWorkspaceAgent(cmd.Context(), resource.ID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
stream, err := dialed.NegotiateConnection(cmd.Context())
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
conn, err := peerbroker.Dial(stream, []webrtc.ICEServer{{
|
||||
URLs: []string{"stun:stun.l.google.com:19302"},
|
||||
}}, &peer.ConnOptions{})
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
client, err := agent.DialSSHClient(conn)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
conn, err := client.DialWorkspaceAgent(cmd.Context(), resource.ID, nil, nil)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
sshClient, err := conn.SSHClient()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
session, err := client.NewSession()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
_, _ = term.MakeRaw(int(os.Stdin.Fd()))
|
||||
err = session.RequestPty("xterm-256color", 128, 128, ssh.TerminalModes{
|
||||
ssh.OCRNL: 1,
|
||||
})
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
session.Stdin = os.Stdin
|
||||
session.Stdout = os.Stdout
|
||||
session.Stderr = os.Stderr
|
||||
err = session.Shell()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
err = session.Wait()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
sshSession, err := sshClient.NewSession()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
_, _ = term.MakeRaw(int(os.Stdin.Fd()))
|
||||
err = sshSession.RequestPty("xterm-256color", 128, 128, ssh.TerminalModes{
|
||||
ssh.OCRNL: 1,
|
||||
})
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
sshSession.Stdin = os.Stdin
|
||||
sshSession.Stdout = os.Stdout
|
||||
sshSession.Stderr = os.Stderr
|
||||
err = sshSession.Shell()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
err = sshSession.Wait()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
return nil
|
||||
|
||||
+29
-11
@@ -6,6 +6,7 @@ import (
|
||||
"os"
|
||||
"time"
|
||||
|
||||
"cloud.google.com/go/compute/metadata"
|
||||
"github.com/spf13/cobra"
|
||||
"golang.org/x/xerrors"
|
||||
|
||||
@@ -19,25 +20,24 @@ import (
|
||||
)
|
||||
|
||||
func workspaceAgent() *cobra.Command {
|
||||
return &cobra.Command{
|
||||
var (
|
||||
rawURL string
|
||||
auth string
|
||||
)
|
||||
cmd := &cobra.Command{
|
||||
Use: "agent",
|
||||
// This command isn't useful to manually execute.
|
||||
Hidden: true,
|
||||
RunE: func(cmd *cobra.Command, args []string) error {
|
||||
coderURLRaw, exists := os.LookupEnv("CODER_URL")
|
||||
if !exists {
|
||||
if rawURL == "" {
|
||||
return xerrors.New("CODER_URL must be set")
|
||||
}
|
||||
coderURL, err := url.Parse(coderURLRaw)
|
||||
coderURL, err := url.Parse(rawURL)
|
||||
if err != nil {
|
||||
return xerrors.Errorf("parse %q: %w", coderURLRaw, err)
|
||||
return xerrors.Errorf("parse %q: %w", rawURL, err)
|
||||
}
|
||||
logger := slog.Make(sloghuman.Sink(cmd.OutOrStdout()))
|
||||
logger := slog.Make(sloghuman.Sink(cmd.OutOrStdout())).Leveled(slog.LevelDebug)
|
||||
client := codersdk.New(coderURL)
|
||||
auth, exists := os.LookupEnv("CODER_AUTH")
|
||||
if !exists {
|
||||
auth = "token"
|
||||
}
|
||||
switch auth {
|
||||
case "token":
|
||||
sessionToken, exists := os.LookupEnv("CODER_TOKEN")
|
||||
@@ -46,16 +46,26 @@ func workspaceAgent() *cobra.Command {
|
||||
}
|
||||
client.SessionToken = sessionToken
|
||||
case "google-instance-identity":
|
||||
// This is *only* done for testing to mock client authentication.
|
||||
// This will never be set in a production scenario.
|
||||
var gcpClient *metadata.Client
|
||||
gcpClientRaw := cmd.Context().Value("gcp-client")
|
||||
if gcpClientRaw != nil {
|
||||
gcpClient, _ = gcpClientRaw.(*metadata.Client)
|
||||
}
|
||||
|
||||
ctx, cancelFunc := context.WithTimeout(cmd.Context(), 30*time.Second)
|
||||
defer cancelFunc()
|
||||
for retry.New(100*time.Millisecond, 5*time.Second).Wait(ctx) {
|
||||
var response codersdk.WorkspaceAgentAuthenticateResponse
|
||||
response, err = client.AuthWorkspaceGoogleInstanceIdentity(cmd.Context(), "", nil)
|
||||
|
||||
response, err = client.AuthWorkspaceGoogleInstanceIdentity(ctx, "", gcpClient)
|
||||
if err != nil {
|
||||
logger.Warn(ctx, "authenticate workspace with Google Instance Identity", slog.Error(err))
|
||||
continue
|
||||
}
|
||||
client.SessionToken = response.SessionToken
|
||||
logger.Info(ctx, "authenticated with Google Instance Identity")
|
||||
break
|
||||
}
|
||||
if err != nil {
|
||||
@@ -73,4 +83,12 @@ func workspaceAgent() *cobra.Command {
|
||||
return closer.Close()
|
||||
},
|
||||
}
|
||||
defaultAuth := os.Getenv("CODER_AUTH")
|
||||
if defaultAuth == "" {
|
||||
defaultAuth = "token"
|
||||
}
|
||||
cmd.Flags().StringVarP(&auth, "auth", "", defaultAuth, "Specify the authentication type to use for the agent.")
|
||||
cmd.Flags().StringVarP(&rawURL, "url", "", os.Getenv("CODER_URL"), "Specify the URL to access Coder.")
|
||||
|
||||
return cmd
|
||||
}
|
||||
|
||||
@@ -0,0 +1,70 @@
|
||||
package cli_test
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/coder/coder/cli/clitest"
|
||||
"github.com/coder/coder/coderd/coderdtest"
|
||||
"github.com/coder/coder/provisioner/echo"
|
||||
"github.com/coder/coder/provisionersdk/proto"
|
||||
)
|
||||
|
||||
func TestWorkspaceAgent(t *testing.T) {
|
||||
t.Parallel()
|
||||
t.Run("GoogleCloud", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
instanceID := "instanceidentifier"
|
||||
validator, metadata := coderdtest.NewGoogleInstanceIdentity(t, instanceID, false)
|
||||
client := coderdtest.New(t, &coderdtest.Options{
|
||||
GoogleTokenValidator: validator,
|
||||
})
|
||||
user := coderdtest.CreateFirstUser(t, client)
|
||||
coderdtest.NewProvisionerDaemon(t, client)
|
||||
version := coderdtest.CreateProjectVersion(t, client, user.OrganizationID, &echo.Responses{
|
||||
Parse: echo.ParseComplete,
|
||||
Provision: []*proto.Provision_Response{{
|
||||
Type: &proto.Provision_Response_Complete{
|
||||
Complete: &proto.Provision_Complete{
|
||||
Resources: []*proto.Resource{{
|
||||
Name: "somename",
|
||||
Type: "someinstance",
|
||||
Agent: &proto.Agent{
|
||||
Auth: &proto.Agent_InstanceId{
|
||||
InstanceId: instanceID,
|
||||
},
|
||||
},
|
||||
}},
|
||||
},
|
||||
},
|
||||
}},
|
||||
})
|
||||
project := coderdtest.CreateProject(t, client, user.OrganizationID, version.ID)
|
||||
coderdtest.AwaitProjectVersionJob(t, client, version.ID)
|
||||
workspace := coderdtest.CreateWorkspace(t, client, "me", project.ID)
|
||||
coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID)
|
||||
|
||||
cmd, _ := clitest.New(t, "workspaces", "agent", "--auth", "google-instance-identity", "--url", client.URL.String())
|
||||
ctx, cancelFunc := context.WithCancel(context.Background())
|
||||
defer cancelFunc()
|
||||
go func() {
|
||||
// A linting error occurs for weakly typing the context value here,
|
||||
// but it seems reasonable for a one-off test.
|
||||
// nolint
|
||||
ctx = context.WithValue(ctx, "gcp-client", metadata)
|
||||
err := cmd.ExecuteContext(ctx)
|
||||
require.NoError(t, err)
|
||||
}()
|
||||
coderdtest.AwaitWorkspaceAgents(t, client, workspace.LatestBuild.ID)
|
||||
resources, err := client.WorkspaceResourcesByBuild(ctx, workspace.LatestBuild.ID)
|
||||
require.NoError(t, err)
|
||||
dialer, err := client.DialWorkspaceAgent(ctx, resources[0].ID, nil, nil)
|
||||
require.NoError(t, err)
|
||||
defer dialer.Close()
|
||||
_, err = dialer.Ping()
|
||||
require.NoError(t, err)
|
||||
cancelFunc()
|
||||
})
|
||||
}
|
||||
Reference in New Issue
Block a user