Add client for agent

coder · kylecarbs · Mar 3, 2022 · Mar 3, 2022 · Mar 3, 2022 · Mar 4, 2022
commit 051eddab13d6056cb9ceed77e9d884fb36df5298
diff --git a/agent/agent.go b/agent/agent.go
@@ -59,9 +59,9 @@ type Options struct {
 	Logger slog.Logger
 }
 
-type Dialer func(ctx context.Context) (*peerbroker.Listener, error)
+type Dialer func(ctx context.Context, options *peer.ConnOptions) (*peerbroker.Listener, error)
 
-func New(dialer Dialer, options *Options) io.Closer {
+func New(dialer Dialer, options *peer.ConnOptions) io.Closer {
 	ctx, cancelFunc := context.WithCancel(context.Background())
 	server := &server{
 		clientDialer: dialer,
@@ -75,7 +75,7 @@ func New(dialer Dialer, options *Options) io.Closer {
 
 type server struct {
 	clientDialer Dialer
-	options      *Options
+	options      *peer.ConnOptions
 
 	closeCancel context.CancelFunc
 	closeMutex  sync.Mutex
@@ -249,7 +249,7 @@ func (s *server) run(ctx context.Context) {
 	// An exponential back-off occurs when the connection is failing to dial.
 	// This is to prevent server spam in case of a coderd outage.
 	for retrier := retry.New(50*time.Millisecond, 10*time.Second); retrier.Wait(ctx); {
-		peerListener, err = s.clientDialer(ctx)
+		peerListener, err = s.clientDialer(ctx, s.options)
 		if err != nil {
 			if errors.Is(err, context.Canceled) {
 				return

diff --git a/agent/agent_test.go b/agent/agent_test.go
@@ -94,11 +94,9 @@ func TestAgent(t *testing.T) {
 
 func setup(t *testing.T) proto.DRPCPeerBrokerClient {
 	client, server := provisionersdk.TransportPipe()
-	closer := agent.New(func(ctx context.Context) (*peerbroker.Listener, error) {
-		return peerbroker.Listen(server, &peer.ConnOptions{
-			Logger: slogtest.Make(t, nil),
-		})
-	}, &agent.Options{
+	closer := agent.New(func(ctx context.Context, opts *peer.ConnOptions) (*peerbroker.Listener, error) {
+		return peerbroker.Listen(server, opts)
+	}, &peer.ConnOptions{
 		Logger: slogtest.Make(t, nil).Leveled(slog.LevelDebug),
 	})
 	t.Cleanup(func() {

diff --git a/cli/workspaceagent.go b/cli/workspaceagent.go
@@ -0,0 +1,56 @@
+package cli
+
+import (
+	"net/url"
+	"os"
+
+	"github.com/coder/coder/agent"
+	"github.com/coder/coder/codersdk"
+	"github.com/powersj/whatsthis/pkg/cloud"
+	"github.com/spf13/cobra"
+	"golang.org/x/xerrors"
+)
+
+func workspaceAgent() *cobra.Command {
+	return &cobra.Command{
+		Use: "agent",
+		// This command isn't useful for users, and seems
+		// more likely to confuse.
+		Hidden: true,
+		RunE: func(cmd *cobra.Command, args []string) error {
+			coderURLRaw, exists := os.LookupEnv("CODER_URL")
+			if !exists {
+				return xerrors.New("CODER_URL must be set")
+			}
+			coderURL, err := url.Parse(coderURLRaw)
+			if err != nil {
+				return xerrors.Errorf("parse %q: %w", coderURLRaw, err)
+			}
+			client := codersdk.New(coderURL)
+			sessionToken, exists := os.LookupEnv("CODER_TOKEN")
+			if !exists {
+				probe, err := cloud.New()
+				if err != nil {
+					return xerrors.Errorf("probe cloud: %w", err)
+				}
+				if !probe.Detected {
+					return xerrors.Errorf("no valid authentication method found; set \"CODER_TOKEN\"")
+				}
+				switch {
+				case probe.GCP():
+					response, err := client.AuthenticateWorkspaceAgentUsingGoogleCloudIdentity(cmd.Context(), "", nil)
+					if err != nil {
+						return xerrors.Errorf("authenticate workspace with gcp: %w", err)
+					}
+					sessionToken = response.SessionToken
+				default:
+					return xerrors.Errorf("%q authentication not supported; set \"CODER_TOKEN\" instead", probe.Name)
+				}
+			}
+			client.SessionToken = sessionToken
+			closer := agent.New(client.WorkspaceAgentServe, nil)
+			<-cmd.Context().Done()
+			return closer.Close()
+		},
+	}
+}
diff --git a/coderd/cmd/root.go b/coderd/cmd/root.go
@@ -98,7 +98,7 @@ func newProvisionerDaemon(ctx context.Context, client *codersdk.Client, logger s
 	if err != nil {
 		return nil, err
 	}
-	return provisionerd.New(client.ProvisionerDaemonClient, &provisionerd.Options{
+	return provisionerd.New(client.ProvisionerDaemonServe, &provisionerd.Options{
 		Logger:         logger,
 		PollInterval:   50 * time.Millisecond,
 		UpdateInterval: 50 * time.Millisecond,

diff --git a/coderd/coderd.go b/coderd/coderd.go
@@ -116,6 +116,10 @@ func New(options *Options) (http.Handler, func()) {
 			r.Route("/authenticate", func(r chi.Router) {
 				r.Post("/google-instance-identity", api.postAuthenticateWorkspaceAgentUsingGoogleInstanceIdentity)
 			})
+			r.Group(func(r chi.Router) {
+				r.Use(httpmw.ExtractWorkspaceAgent(options.Database))
+				r.Get("/serve", api.workspaceAgentServe)
+			})
 		})
 
 		r.Route("/upload", func(r chi.Router) {
@@ -134,7 +138,7 @@ func New(options *Options) (http.Handler, func()) {
 				r.Get("/", api.provisionerJobByID)
 				r.Get("/schemas", api.projectImportJobSchemasByID)
 				r.Get("/parameters", api.projectImportJobParametersByID)
-				r.Get("/resources", api.projectImportJobResourcesByID)
+				r.Get("/resources", api.provisionerJobResourcesByID)
 				r.Get("/logs", api.provisionerJobLogsByID)
 			})
 		})
@@ -148,6 +152,13 @@ func New(options *Options) (http.Handler, func()) {
 				r.Use(httpmw.ExtractProvisionerJobParam(options.Database))
 				r.Get("/", api.provisionerJobByID)
 				r.Get("/logs", api.provisionerJobLogsByID)
+				r.Route("/resources", func(r chi.Router) {
+					r.Get("/", api.provisionerJobResourcesByID)
+					r.Route("/{workspaceresource}", func(r chi.Router) {
+						r.Use(httpmw.ExtractWorkspaceResourceParam(options.Database))
+						r.Get("/agent", api.workspaceAgentConnectByResource)
+					})
+				})
 			})
 		})
 

diff --git a/coderd/coderdtest/coderdtest.go b/coderd/coderdtest/coderdtest.go
@@ -125,7 +125,7 @@ func NewProvisionerDaemon(t *testing.T, client *codersdk.Client) io.Closer {
 		require.NoError(t, err)
 	}()
 
-	closer := provisionerd.New(client.ProvisionerDaemonClient, &provisionerd.Options{
+	closer := provisionerd.New(client.ProvisionerDaemonServe, &provisionerd.Options{
 		Logger:         slogtest.Make(t, nil).Named("provisionerd").Leveled(slog.LevelDebug),
 		PollInterval:   50 * time.Millisecond,
 		UpdateInterval: 50 * time.Millisecond,

diff --git a/coderd/projectimport.go b/coderd/projectimport.go
@@ -154,29 +154,3 @@ func (api *api) projectImportJobParametersByID(rw http.ResponseWriter, r *http.R
 	render.Status(r, http.StatusOK)
 	render.JSON(rw, r, values)
 }
-
-// Returns resources for an import job by ID.
-func (api *api) projectImportJobResourcesByID(rw http.ResponseWriter, r *http.Request) {
-	job := httpmw.ProvisionerJobParam(r)
-	if !convertProvisionerJob(job).Status.Completed() {
-		httpapi.Write(rw, http.StatusPreconditionFailed, httpapi.Response{
-			Message: "Job hasn't completed!",
-		})
-		return
-	}
-	resources, err := api.Database.GetProvisionerJobResourcesByJobID(r.Context(), job.ID)
-	if errors.Is(err, sql.ErrNoRows) {
-		err = nil
-	}
-	if err != nil {
-		httpapi.Write(rw, http.StatusInternalServerError, httpapi.Response{
-			Message: fmt.Sprintf("get project import job resources: %s", err),
-		})
-		return
-	}
-	if resources == nil {
-		resources = []database.ProvisionerJobResource{}
-	}
-	render.Status(r, http.StatusOK)
-	render.JSON(rw, r, resources)
-}
diff --git a/coderd/provisionerdaemons.go b/coderd/provisionerdaemons.go
@@ -590,12 +590,19 @@ func insertProvisionerJobResource(ctx context.Context, db database.Store, jobID
 				Valid:      true,
 			}
 		}
+		authToken := uuid.New()
+		if protoResource.Agent.GetToken() != "" {
+			authToken, err = uuid.Parse(protoResource.Agent.GetToken())
+			if err != nil {
+				return xerrors.Errorf("invalid auth token format; must be uuid: %w", err)
+			}
+		}
 
 		_, err := db.InsertProvisionerJobAgent(ctx, database.InsertProvisionerJobAgentParams{
 			ID:                   resource.AgentID.UUID,
 			CreatedAt:            database.Now(),
 			ResourceID:           resource.ID,
-			AuthToken:            uuid.New(),
+			AuthToken:            authToken,
 			AuthInstanceID:       instanceID,
 			EnvironmentVariables: env,
 			StartupScript: sql.NullString{

diff --git a/coderd/provisionerjobs.go b/coderd/provisionerjobs.go
@@ -12,6 +12,7 @@ import (
 
 	"github.com/go-chi/render"
 	"github.com/google/uuid"
+	"golang.org/x/xerrors"
 
 	"cdr.dev/slog"
 
@@ -64,6 +65,7 @@ type ProvisionerJobResource struct {
 	Transition database.WorkspaceTransition `json:"workspace_transition"`
 	Type       string                       `json:"type"`
 	Name       string                       `json:"name"`
+	Agent      *ProvisionerJobAgent         `json:"agent,omitempty"`
 }
 
 type ProvisionerJobAgent struct {
@@ -238,6 +240,49 @@ func (api *api) provisionerJobLogsByID(rw http.ResponseWriter, r *http.Request)
 	}
 }
 
+func (api *api) provisionerJobResourcesByID(rw http.ResponseWriter, r *http.Request) {
+	job := httpmw.ProvisionerJobParam(r)
+	if !convertProvisionerJob(job).Status.Completed() {
+		httpapi.Write(rw, http.StatusPreconditionFailed, httpapi.Response{
+			Message: "Job hasn't completed!",
+		})
+		return
+	}
+	resources, err := api.Database.GetProvisionerJobResourcesByJobID(r.Context(), job.ID)
+	if err != nil {
+		httpapi.Write(rw, http.StatusInternalServerError, httpapi.Response{
+			Message: fmt.Sprintf("get provisioner job resources: %s", err),
+		})
+		return
+	}
+	apiResources := make([]ProvisionerJobResource, 0)
+	for _, resource := range resources {
+		if !resource.AgentID.Valid {
+			apiResources = append(apiResources, convertProvisionerJobResource(resource, nil))
+			continue
+		}
+		// TODO: This should be combined.
+		agents, err := api.Database.GetProvisionerJobAgentsByResourceIDs(r.Context(), []uuid.UUID{resource.ID})
+		if err != nil {
+			httpapi.Write(rw, http.StatusInternalServerError, httpapi.Response{
+				Message: fmt.Sprintf("get provisioner job agent: %s", err),
+			})
+			return
+		}
+		agent := agents[0]
+		apiAgent, err := convertProvisionerJobAgent(agent)
+		if err != nil {
+			httpapi.Write(rw, http.StatusInternalServerError, httpapi.Response{
+				Message: fmt.Sprintf("convert provisioner job agent: %s", err),
+			})
+			return
+		}
+		apiResources = append(apiResources, convertProvisionerJobResource(resource, &apiAgent))
+	}
+	render.Status(r, http.StatusOK)
+	render.JSON(rw, r, apiResources)
+}
+
 func convertProvisionerJobLog(provisionerJobLog database.ProvisionerJobLog) ProvisionerJobLog {
 	return ProvisionerJobLog{
 		ID:        provisionerJobLog.ID,
@@ -291,6 +336,37 @@ func convertProvisionerJob(provisionerJob database.ProvisionerJob) ProvisionerJo
 	return job
 }
 
+func convertProvisionerJobResource(resource database.ProvisionerJobResource, agent *ProvisionerJobAgent) ProvisionerJobResource {
+	return ProvisionerJobResource{
+		ID:         resource.ID,
+		CreatedAt:  resource.CreatedAt,
+		JobID:      resource.JobID,
+		Transition: resource.Transition,
+		Type:       resource.Type,
+		Name:       resource.Name,
+		Agent:      agent,
+	}
+}
+
+func convertProvisionerJobAgent(agent database.ProvisionerJobAgent) (ProvisionerJobAgent, error) {
+	var envs map[string]string
+	if agent.EnvironmentVariables.Valid {
+		err := json.Unmarshal(agent.EnvironmentVariables.RawMessage, &envs)
+		if err != nil {
+			return ProvisionerJobAgent{}, xerrors.Errorf("unmarshal: %w", err)
+		}
+	}
+	return ProvisionerJobAgent{
+		ID:                   agent.ID,
+		CreatedAt:            agent.CreatedAt,
+		UpdatedAt:            agent.UpdatedAt.Time,
+		ResourceID:           agent.ResourceID,
+		InstanceID:           agent.AuthInstanceID.String,
+		StartupScript:        agent.StartupScript.String,
+		EnvironmentVariables: envs,
+	}, nil
+}
+
 func provisionerJobLogsChannel(jobID uuid.UUID) string {
 	return fmt.Sprintf("provisioner-log-logs:%s", jobID)
 }