fix: Refactor agent to consume API client (#4715)

* fix: Refactor agent to consume API client

This simplifies a lot of code by creating an interface for
the codersdk client into the agent. It also moves agent
authentication code so instance identity will work between
restarts.

Fixes #3485 and #4082.

* Fix client reconnections
This commit is contained in:
Kyle Carberry
2022-10-23 22:35:08 -05:00
committed by GitHub
parent c9bf2a9099
commit bf3224e373
19 changed files with 379 additions and 364 deletions

View File

@ -471,7 +471,6 @@ func New(options *Options) *API {
r.Post("/google-instance-identity", api.postWorkspaceAuthGoogleInstanceIdentity)
r.Route("/me", func(r chi.Router) {
r.Use(httpmw.ExtractWorkspaceAgent(options.Database))
r.Get("/apps", api.workspaceAgentApps)
r.Get("/metadata", api.workspaceAgentMetadata)
r.Post("/version", api.postWorkspaceAgentVersion)
r.Post("/app-health", api.postWorkspaceAppHealth)

View File

@ -57,7 +57,6 @@ func AGPLRoutes(a *AuthTester) (map[string]string, map[string]RouteCheck) {
"POST:/api/v2/workspaceagents/aws-instance-identity": {NoAuthorize: true},
"POST:/api/v2/workspaceagents/azure-instance-identity": {NoAuthorize: true},
"POST:/api/v2/workspaceagents/google-instance-identity": {NoAuthorize: true},
"GET:/api/v2/workspaceagents/me/apps": {NoAuthorize: true},
"GET:/api/v2/workspaceagents/me/gitsshkey": {NoAuthorize: true},
"GET:/api/v2/workspaceagents/me/metadata": {NoAuthorize: true},
"GET:/api/v2/workspaceagents/me/coordinate": {NoAuthorize: true},

View File

@ -603,10 +603,8 @@ func TestTemplateMetrics(t *testing.T) {
agentClient := codersdk.New(client.URL)
agentClient.SessionToken = authToken
agentCloser := agent.New(agent.Options{
Logger: slogtest.Make(t, nil),
StatsReporter: agentClient.AgentReportStats,
FetchMetadata: agentClient.WorkspaceAgentMetadata,
CoordinatorDialer: agentClient.ListenWorkspaceAgentTailnet,
Logger: slogtest.Make(t, nil),
Client: agentClient,
})
defer func() {
_ = agentCloser.Close()

View File

@ -61,20 +61,6 @@ func (api *API) workspaceAgent(rw http.ResponseWriter, r *http.Request) {
httpapi.Write(ctx, rw, http.StatusOK, apiAgent)
}
func (api *API) workspaceAgentApps(rw http.ResponseWriter, r *http.Request) {
workspaceAgent := httpmw.WorkspaceAgent(r)
dbApps, err := api.Database.GetWorkspaceAppsByAgentID(r.Context(), workspaceAgent.ID)
if err != nil && !xerrors.Is(err, sql.ErrNoRows) {
httpapi.Write(r.Context(), rw, http.StatusInternalServerError, codersdk.Response{
Message: "Internal error fetching workspace agent applications.",
Detail: err.Error(),
})
return
}
httpapi.Write(r.Context(), rw, http.StatusOK, convertApps(dbApps))
}
func (api *API) workspaceAgentMetadata(rw http.ResponseWriter, r *http.Request) {
ctx := r.Context()
workspaceAgent := httpmw.WorkspaceAgent(r)
@ -86,8 +72,17 @@ func (api *API) workspaceAgentMetadata(rw http.ResponseWriter, r *http.Request)
})
return
}
dbApps, err := api.Database.GetWorkspaceAppsByAgentID(r.Context(), workspaceAgent.ID)
if err != nil && !xerrors.Is(err, sql.ErrNoRows) {
httpapi.Write(r.Context(), rw, http.StatusInternalServerError, codersdk.Response{
Message: "Internal error fetching workspace agent applications.",
Detail: err.Error(),
})
return
}
httpapi.Write(ctx, rw, http.StatusOK, codersdk.WorkspaceAgentMetadata{
Apps: convertApps(dbApps),
DERPMap: api.DERPMap,
EnvironmentVariables: apiAgent.EnvironmentVariables,
StartupScript: apiAgent.StartupScript,

View File

@ -111,9 +111,8 @@ func TestWorkspaceAgentListen(t *testing.T) {
agentClient := codersdk.New(client.URL)
agentClient.SessionToken = authToken
agentCloser := agent.New(agent.Options{
FetchMetadata: agentClient.WorkspaceAgentMetadata,
CoordinatorDialer: agentClient.ListenWorkspaceAgentTailnet,
Logger: slogtest.Make(t, nil).Named("agent").Leveled(slog.LevelDebug),
Client: agentClient,
Logger: slogtest.Make(t, nil).Named("agent").Leveled(slog.LevelDebug),
})
defer func() {
_ = agentCloser.Close()
@ -204,7 +203,7 @@ func TestWorkspaceAgentListen(t *testing.T) {
agentClient := codersdk.New(client.URL)
agentClient.SessionToken = authToken
_, err = agentClient.ListenWorkspaceAgentTailnet(ctx)
_, err = agentClient.ListenWorkspaceAgent(ctx)
require.Error(t, err)
require.ErrorContains(t, err, "build is outdated")
})
@ -244,9 +243,8 @@ func TestWorkspaceAgentTailnet(t *testing.T) {
agentClient := codersdk.New(client.URL)
agentClient.SessionToken = authToken
agentCloser := agent.New(agent.Options{
FetchMetadata: agentClient.WorkspaceAgentMetadata,
CoordinatorDialer: agentClient.ListenWorkspaceAgentTailnet,
Logger: slogtest.Make(t, nil).Named("agent").Leveled(slog.LevelDebug),
Client: agentClient,
Logger: slogtest.Make(t, nil).Named("agent").Leveled(slog.LevelDebug),
})
defer agentCloser.Close()
resources := coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID)
@ -311,9 +309,8 @@ func TestWorkspaceAgentPTY(t *testing.T) {
agentClient := codersdk.New(client.URL)
agentClient.SessionToken = authToken
agentCloser := agent.New(agent.Options{
FetchMetadata: agentClient.WorkspaceAgentMetadata,
CoordinatorDialer: agentClient.ListenWorkspaceAgentTailnet,
Logger: slogtest.Make(t, nil).Named("agent").Leveled(slog.LevelDebug),
Client: agentClient,
Logger: slogtest.Make(t, nil).Named("agent").Leveled(slog.LevelDebug),
})
defer func() {
_ = agentCloser.Close()
@ -409,9 +406,8 @@ func TestWorkspaceAgentListeningPorts(t *testing.T) {
agentClient := codersdk.New(client.URL)
agentClient.SessionToken = authToken
agentCloser := agent.New(agent.Options{
FetchMetadata: agentClient.WorkspaceAgentMetadata,
CoordinatorDialer: agentClient.ListenWorkspaceAgentTailnet,
Logger: slogtest.Make(t, nil).Named("agent").Leveled(slog.LevelDebug),
Client: agentClient,
Logger: slogtest.Make(t, nil).Named("agent").Leveled(slog.LevelDebug),
})
t.Cleanup(func() {
_ = agentCloser.Close()
@ -671,10 +667,10 @@ func TestWorkspaceAgentAppHealth(t *testing.T) {
agentClient := codersdk.New(client.URL)
agentClient.SessionToken = authToken
apiApps, err := agentClient.WorkspaceAgentApps(ctx)
metadata, err := agentClient.WorkspaceAgentMetadata(ctx)
require.NoError(t, err)
require.EqualValues(t, codersdk.WorkspaceAppHealthDisabled, apiApps[0].Health)
require.EqualValues(t, codersdk.WorkspaceAppHealthInitializing, apiApps[1].Health)
require.EqualValues(t, codersdk.WorkspaceAppHealthDisabled, metadata.Apps[0].Health)
require.EqualValues(t, codersdk.WorkspaceAppHealthInitializing, metadata.Apps[1].Health)
err = agentClient.PostWorkspaceAgentAppHealth(ctx, codersdk.PostWorkspaceAppHealthsRequest{})
require.Error(t, err)
// empty
@ -708,9 +704,9 @@ func TestWorkspaceAgentAppHealth(t *testing.T) {
},
})
require.NoError(t, err)
apiApps, err = agentClient.WorkspaceAgentApps(ctx)
metadata, err = agentClient.WorkspaceAgentMetadata(ctx)
require.NoError(t, err)
require.EqualValues(t, codersdk.WorkspaceAppHealthHealthy, apiApps[1].Health)
require.EqualValues(t, codersdk.WorkspaceAppHealthHealthy, metadata.Apps[1].Health)
// update to unhealthy
err = agentClient.PostWorkspaceAgentAppHealth(ctx, codersdk.PostWorkspaceAppHealthsRequest{
Healths: map[string]codersdk.WorkspaceAppHealth{
@ -718,7 +714,7 @@ func TestWorkspaceAgentAppHealth(t *testing.T) {
},
})
require.NoError(t, err)
apiApps, err = agentClient.WorkspaceAgentApps(ctx)
metadata, err = agentClient.WorkspaceAgentMetadata(ctx)
require.NoError(t, err)
require.EqualValues(t, codersdk.WorkspaceAppHealthUnhealthy, apiApps[1].Health)
require.EqualValues(t, codersdk.WorkspaceAppHealthUnhealthy, metadata.Apps[1].Health)
}

View File

@ -195,10 +195,8 @@ func createWorkspaceWithApps(t *testing.T, client *codersdk.Client, orgID uuid.U
agentClient := codersdk.New(client.URL)
agentClient.SessionToken = authToken
agentCloser := agent.New(agent.Options{
FetchMetadata: agentClient.WorkspaceAgentMetadata,
CoordinatorDialer: agentClient.ListenWorkspaceAgentTailnet,
Logger: slogtest.Make(t, nil).Named("agent"),
StatsReporter: agentClient.AgentReportStats,
Client: agentClient,
Logger: slogtest.Make(t, nil).Named("agent"),
})
t.Cleanup(func() {
_ = agentCloser.Close()

View File

@ -3,12 +3,14 @@ package wsconncache_test
import (
"context"
"fmt"
"io"
"net"
"net/http"
"net/http/httptest"
"net/http/httputil"
"net/netip"
"net/url"
"strings"
"sync"
"testing"
"time"
@ -148,17 +150,11 @@ func setupAgent(t *testing.T, metadata codersdk.WorkspaceAgentMetadata, ptyTimeo
coordinator := tailnet.NewCoordinator()
agentID := uuid.New()
closer := agent.New(agent.Options{
FetchMetadata: func(ctx context.Context) (codersdk.WorkspaceAgentMetadata, error) {
return metadata, nil
},
CoordinatorDialer: func(ctx context.Context) (net.Conn, error) {
clientConn, serverConn := net.Pipe()
t.Cleanup(func() {
_ = serverConn.Close()
_ = clientConn.Close()
})
go coordinator.ServeAgent(serverConn, agentID)
return clientConn, nil
Client: &client{
t: t,
agentID: agentID,
metadata: metadata,
coordinator: coordinator,
},
Logger: slogtest.Make(t, nil).Named("agent").Leveled(slog.LevelInfo),
ReconnectingPTYTimeout: ptyTimeout,
@ -187,3 +183,41 @@ func setupAgent(t *testing.T, metadata codersdk.WorkspaceAgentMetadata, ptyTimeo
Conn: conn,
}
}
type client struct {
t *testing.T
agentID uuid.UUID
metadata codersdk.WorkspaceAgentMetadata
coordinator tailnet.Coordinator
}
func (c *client) WorkspaceAgentMetadata(_ context.Context) (codersdk.WorkspaceAgentMetadata, error) {
return c.metadata, nil
}
func (c *client) ListenWorkspaceAgent(_ context.Context) (net.Conn, error) {
clientConn, serverConn := net.Pipe()
closed := make(chan struct{})
c.t.Cleanup(func() {
_ = serverConn.Close()
_ = clientConn.Close()
<-closed
})
go func() {
_ = c.coordinator.ServeAgent(serverConn, c.agentID)
close(closed)
}()
return clientConn, nil
}
func (*client) AgentReportStats(_ context.Context, _ slog.Logger, _ func() *codersdk.AgentStats) (io.Closer, error) {
return io.NopCloser(strings.NewReader("")), nil
}
func (*client) PostWorkspaceAgentAppHealth(_ context.Context, _ codersdk.PostWorkspaceAppHealthsRequest) error {
return nil
}
func (*client) PostWorkspaceAgentVersion(_ context.Context, _ string) error {
return nil
}