Skip to content

Commit 051edda

Browse files
kylecarbskylecarbs
committed
Add client for agent
1 parent 73cb452 commit 051edda

30 files changed

+1124
-270
lines changed

‎agent/agent.go

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -59,9 +59,9 @@ type Options struct {
5959
Logger slog.Logger
6060
}
6161

62-
type Dialer func(ctx context.Context) (*peerbroker.Listener, error)
62+
type Dialer func(ctx context.Context, options *peer.ConnOptions) (*peerbroker.Listener, error)
6363

64-
func New(dialer Dialer, options *Options) io.Closer {
64+
func New(dialer Dialer, options *peer.ConnOptions) io.Closer {
6565
ctx, cancelFunc := context.WithCancel(context.Background())
6666
server := &server{
6767
clientDialer: dialer,
@@ -75,7 +75,7 @@ func New(dialer Dialer, options *Options) io.Closer {
7575

7676
type server struct {
7777
clientDialer Dialer
78-
options *Options
78+
options *peer.ConnOptions
7979

8080
closeCancel context.CancelFunc
8181
closeMutex sync.Mutex
@@ -249,7 +249,7 @@ func (s *server) run(ctx context.Context) {
249249
// An exponential back-off occurs when the connection is failing to dial.
250250
// This is to prevent server spam in case of a coderd outage.
251251
for retrier := retry.New(50*time.Millisecond, 10*time.Second); retrier.Wait(ctx); {
252-
peerListener, err = s.clientDialer(ctx)
252+
peerListener, err = s.clientDialer(ctx, s.options)
253253
if err != nil {
254254
if errors.Is(err, context.Canceled) {
255255
return

‎agent/agent_test.go

Lines changed: 3 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -94,11 +94,9 @@ func TestAgent(t *testing.T) {
9494

9595
func setup(t *testing.T) proto.DRPCPeerBrokerClient {
9696
client, server := provisionersdk.TransportPipe()
97-
closer := agent.New(func(ctx context.Context) (*peerbroker.Listener, error) {
98-
return peerbroker.Listen(server, &peer.ConnOptions{
99-
Logger: slogtest.Make(t, nil),
100-
})
101-
}, &agent.Options{
97+
closer := agent.New(func(ctx context.Context, opts *peer.ConnOptions) (*peerbroker.Listener, error) {
98+
return peerbroker.Listen(server, opts)
99+
}, &peer.ConnOptions{
102100
Logger: slogtest.Make(t, nil).Leveled(slog.LevelDebug),
103101
})
104102
t.Cleanup(func() {

‎cli/workspaceagent.go

Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,56 @@
1+
package cli
2+
3+
import (
4+
"net/url"
5+
"os"
6+
7+
"github.com/coder/coder/agent"
8+
"github.com/coder/coder/codersdk"
9+
"github.com/powersj/whatsthis/pkg/cloud"
10+
"github.com/spf13/cobra"
11+
"golang.org/x/xerrors"
12+
)
13+
14+
func workspaceAgent() *cobra.Command {
15+
return &cobra.Command{
16+
Use: "agent",
17+
// This command isn't useful for users, and seems
18+
// more likely to confuse.
19+
Hidden: true,
20+
RunE: func(cmd *cobra.Command, args []string) error {
21+
coderURLRaw, exists := os.LookupEnv("CODER_URL")
22+
if !exists {
23+
return xerrors.New("CODER_URL must be set")
24+
}
25+
coderURL, err := url.Parse(coderURLRaw)
26+
if err != nil {
27+
return xerrors.Errorf("parse %q: %w", coderURLRaw, err)
28+
}
29+
client := codersdk.New(coderURL)
30+
sessionToken, exists := os.LookupEnv("CODER_TOKEN")
31+
if !exists {
32+
probe, err := cloud.New()
33+
if err != nil {
34+
return xerrors.Errorf("probe cloud: %w", err)
35+
}
36+
if !probe.Detected {
37+
return xerrors.Errorf("no valid authentication method found; set \"CODER_TOKEN\"")
38+
}
39+
switch {
40+
case probe.GCP():
41+
response, err := client.AuthenticateWorkspaceAgentUsingGoogleCloudIdentity(cmd.Context(), "", nil)
42+
if err != nil {
43+
return xerrors.Errorf("authenticate workspace with gcp: %w", err)
44+
}
45+
sessionToken = response.SessionToken
46+
default:
47+
return xerrors.Errorf("%q authentication not supported; set \"CODER_TOKEN\" instead", probe.Name)
48+
}
49+
}
50+
client.SessionToken = sessionToken
51+
closer := agent.New(client.WorkspaceAgentServe, nil)
52+
<-cmd.Context().Done()
53+
return closer.Close()
54+
},
55+
}
56+
}

‎coderd/cmd/root.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -98,7 +98,7 @@ func newProvisionerDaemon(ctx context.Context, client *codersdk.Client, logger s
9898
if err != nil {
9999
return nil, err
100100
}
101-
return provisionerd.New(client.ProvisionerDaemonClient, &provisionerd.Options{
101+
return provisionerd.New(client.ProvisionerDaemonServe, &provisionerd.Options{
102102
Logger: logger,
103103
PollInterval: 50 * time.Millisecond,
104104
UpdateInterval: 50 * time.Millisecond,

‎coderd/coderd.go

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -116,6 +116,10 @@ func New(options *Options) (http.Handler, func()) {
116116
r.Route("/authenticate", func(r chi.Router) {
117117
r.Post("/google-instance-identity", api.postAuthenticateWorkspaceAgentUsingGoogleInstanceIdentity)
118118
})
119+
r.Group(func(r chi.Router) {
120+
r.Use(httpmw.ExtractWorkspaceAgent(options.Database))
121+
r.Get("/serve", api.workspaceAgentServe)
122+
})
119123
})
120124

121125
r.Route("/upload", func(r chi.Router) {
@@ -134,7 +138,7 @@ func New(options *Options) (http.Handler, func()) {
134138
r.Get("/", api.provisionerJobByID)
135139
r.Get("/schemas", api.projectImportJobSchemasByID)
136140
r.Get("/parameters", api.projectImportJobParametersByID)
137-
r.Get("/resources", api.projectImportJobResourcesByID)
141+
r.Get("/resources", api.provisionerJobResourcesByID)
138142
r.Get("/logs", api.provisionerJobLogsByID)
139143
})
140144
})
@@ -148,6 +152,13 @@ func New(options *Options) (http.Handler, func()) {
148152
r.Use(httpmw.ExtractProvisionerJobParam(options.Database))
149153
r.Get("/", api.provisionerJobByID)
150154
r.Get("/logs", api.provisionerJobLogsByID)
155+
r.Route("/resources", func(r chi.Router) {
156+
r.Get("/", api.provisionerJobResourcesByID)
157+
r.Route("/{workspaceresource}", func(r chi.Router) {
158+
r.Use(httpmw.ExtractWorkspaceResourceParam(options.Database))
159+
r.Get("/agent", api.workspaceAgentConnectByResource)
160+
})
161+
})
151162
})
152163
})
153164

‎coderd/coderdtest/coderdtest.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -125,7 +125,7 @@ func NewProvisionerDaemon(t *testing.T, client *codersdk.Client) io.Closer {
125125
require.NoError(t, err)
126126
}()
127127

128-
closer := provisionerd.New(client.ProvisionerDaemonClient, &provisionerd.Options{
128+
closer := provisionerd.New(client.ProvisionerDaemonServe, &provisionerd.Options{
129129
Logger: slogtest.Make(t, nil).Named("provisionerd").Leveled(slog.LevelDebug),
130130
PollInterval: 50 * time.Millisecond,
131131
UpdateInterval: 50 * time.Millisecond,

‎coderd/projectimport.go

Lines changed: 0 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -154,29 +154,3 @@ func (api *api) projectImportJobParametersByID(rw http.ResponseWriter, r *http.R
154154
render.Status(r, http.StatusOK)
155155
render.JSON(rw, r, values)
156156
}
157-
158-
// Returns resources for an import job by ID.
159-
func (api *api) projectImportJobResourcesByID(rw http.ResponseWriter, r *http.Request) {
160-
job := httpmw.ProvisionerJobParam(r)
161-
if !convertProvisionerJob(job).Status.Completed() {
162-
httpapi.Write(rw, http.StatusPreconditionFailed, httpapi.Response{
163-
Message: "Job hasn't completed!",
164-
})
165-
return
166-
}
167-
resources, err := api.Database.GetProvisionerJobResourcesByJobID(r.Context(), job.ID)
168-
if errors.Is(err, sql.ErrNoRows) {
169-
err = nil
170-
}
171-
if err != nil {
172-
httpapi.Write(rw, http.StatusInternalServerError, httpapi.Response{
173-
Message: fmt.Sprintf("get project import job resources: %s", err),
174-
})
175-
return
176-
}
177-
if resources == nil {
178-
resources = []database.ProvisionerJobResource{}
179-
}
180-
render.Status(r, http.StatusOK)
181-
render.JSON(rw, r, resources)
182-
}

‎coderd/provisionerdaemons.go

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -590,12 +590,19 @@ func insertProvisionerJobResource(ctx context.Context, db database.Store, jobID
590590
Valid: true,
591591
}
592592
}
593+
authToken := uuid.New()
594+
if protoResource.Agent.GetToken() != "" {
595+
authToken, err = uuid.Parse(protoResource.Agent.GetToken())
596+
if err != nil {
597+
return xerrors.Errorf("invalid auth token format; must be uuid: %w", err)
598+
}
599+
}
593600

594601
_, err := db.InsertProvisionerJobAgent(ctx, database.InsertProvisionerJobAgentParams{
595602
ID: resource.AgentID.UUID,
596603
CreatedAt: database.Now(),
597604
ResourceID: resource.ID,
598-
AuthToken: uuid.New(),
605+
AuthToken: authToken,
599606
AuthInstanceID: instanceID,
600607
EnvironmentVariables: env,
601608
StartupScript: sql.NullString{

‎coderd/provisionerjobs.go

Lines changed: 76 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@ import (
1212

1313
"github.com/go-chi/render"
1414
"github.com/google/uuid"
15+
"golang.org/x/xerrors"
1516

1617
"cdr.dev/slog"
1718

@@ -64,6 +65,7 @@ type ProvisionerJobResource struct {
6465
Transition database.WorkspaceTransition `json:"workspace_transition"`
6566
Type string `json:"type"`
6667
Name string `json:"name"`
68+
Agent *ProvisionerJobAgent `json:"agent,omitempty"`
6769
}
6870

6971
type ProvisionerJobAgent struct {
@@ -238,6 +240,49 @@ func (api *api) provisionerJobLogsByID(rw http.ResponseWriter, r *http.Request)
238240
}
239241
}
240242

243+
func (api *api) provisionerJobResourcesByID(rw http.ResponseWriter, r *http.Request) {
244+
job := httpmw.ProvisionerJobParam(r)
245+
if !convertProvisionerJob(job).Status.Completed() {
246+
httpapi.Write(rw, http.StatusPreconditionFailed, httpapi.Response{
247+
Message: "Job hasn't completed!",
248+
})
249+
return
250+
}
251+
resources, err := api.Database.GetProvisionerJobResourcesByJobID(r.Context(), job.ID)
252+
if err != nil {
253+
httpapi.Write(rw, http.StatusInternalServerError, httpapi.Response{
254+
Message: fmt.Sprintf("get provisioner job resources: %s", err),
255+
})
256+
return
257+
}
258+
apiResources := make([]ProvisionerJobResource, 0)
259+
for _, resource := range resources {
260+
if !resource.AgentID.Valid {
261+
apiResources = append(apiResources, convertProvisionerJobResource(resource, nil))
262+
continue
263+
}
264+
// TODO: This should be combined.
265+
agents, err := api.Database.GetProvisionerJobAgentsByResourceIDs(r.Context(), []uuid.UUID{resource.ID})
266+
if err != nil {
267+
httpapi.Write(rw, http.StatusInternalServerError, httpapi.Response{
268+
Message: fmt.Sprintf("get provisioner job agent: %s", err),
269+
})
270+
return
271+
}
272+
agent := agents[0]
273+
apiAgent, err := convertProvisionerJobAgent(agent)
274+
if err != nil {
275+
httpapi.Write(rw, http.StatusInternalServerError, httpapi.Response{
276+
Message: fmt.Sprintf("convert provisioner job agent: %s", err),
277+
})
278+
return
279+
}
280+
apiResources = append(apiResources, convertProvisionerJobResource(resource, &apiAgent))
281+
}
282+
render.Status(r, http.StatusOK)
283+
render.JSON(rw, r, apiResources)
284+
}
285+
241286
func convertProvisionerJobLog(provisionerJobLog database.ProvisionerJobLog) ProvisionerJobLog {
242287
return ProvisionerJobLog{
243288
ID: provisionerJobLog.ID,
@@ -291,6 +336,37 @@ func convertProvisionerJob(provisionerJob database.ProvisionerJob) ProvisionerJo
291336
return job
292337
}
293338

339+
func convertProvisionerJobResource(resource database.ProvisionerJobResource, agent *ProvisionerJobAgent) ProvisionerJobResource {
340+
return ProvisionerJobResource{
341+
ID: resource.ID,
342+
CreatedAt: resource.CreatedAt,
343+
JobID: resource.JobID,
344+
Transition: resource.Transition,
345+
Type: resource.Type,
346+
Name: resource.Name,
347+
Agent: agent,
348+
}
349+
}
350+
351+
func convertProvisionerJobAgent(agent database.ProvisionerJobAgent) (ProvisionerJobAgent, error) {
352+
var envs map[string]string
353+
if agent.EnvironmentVariables.Valid {
354+
err := json.Unmarshal(agent.EnvironmentVariables.RawMessage, &envs)
355+
if err != nil {
356+
return ProvisionerJobAgent{}, xerrors.Errorf("unmarshal: %w", err)
357+
}
358+
}
359+
return ProvisionerJobAgent{
360+
ID: agent.ID,
361+
CreatedAt: agent.CreatedAt,
362+
UpdatedAt: agent.UpdatedAt.Time,
363+
ResourceID: agent.ResourceID,
364+
InstanceID: agent.AuthInstanceID.String,
365+
StartupScript: agent.StartupScript.String,
366+
EnvironmentVariables: envs,
367+
}, nil
368+
}
369+
294370
func provisionerJobLogsChannel(jobID uuid.UUID) string {
295371
return fmt.Sprintf("provisioner-log-logs:%s", jobID)
296372
}

0 commit comments

Comments
 (0)