coderd/workspaceagentsrpc_test.go
Rama main · réplica de lectura del repositorio original.
package coderd_test
import (
"context"
"database/sql"
"testing"
"time"
"github.com/google/uuid"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"google.golang.org/protobuf/types/known/timestamppb"
"storj.io/drpc/drpcerr"
agentproto "github.com/coder/coder/v2/agent/proto"
"github.com/coder/coder/v2/coderd/coderdtest"
"github.com/coder/coder/v2/coderd/connectionlog"
"github.com/coder/coder/v2/coderd/database"
"github.com/coder/coder/v2/coderd/database/dbauthz"
"github.com/coder/coder/v2/coderd/database/dbfake"
"github.com/coder/coder/v2/coderd/database/dbgen"
"github.com/coder/coder/v2/coderd/database/dbtestutil"
"github.com/coder/coder/v2/coderd/database/dbtime"
"github.com/coder/coder/v2/coderd/rbac"
"github.com/coder/coder/v2/codersdk"
"github.com/coder/coder/v2/codersdk/agentsdk"
"github.com/coder/coder/v2/provisionersdk/proto"
"github.com/coder/coder/v2/tailnet"
tailnetproto "github.com/coder/coder/v2/tailnet/proto"
"github.com/coder/coder/v2/testutil"
)
// Ported to RPC API from coderd/workspaceagents_test.go
func TestWorkspaceAgentReportStats(t *testing.T) {
t.Parallel()
for _, tc := range []struct {
name string
apiKeyScope rbac.ScopeName
}{
{
name: "empty (backwards compat)",
apiKeyScope: "",
},
{
name: "all",
apiKeyScope: rbac.ScopeAll,
},
{
name: "no_user_data",
apiKeyScope: rbac.ScopeNoUserData,
},
{
name: "application_connect",
apiKeyScope: rbac.ScopeApplicationConnect,
},
} {
t.Run(tc.name, func(t *testing.T) {
t.Parallel()
tickCh := make(chan time.Time)
flushCh := make(chan int, 1)
client, db := coderdtest.NewWithDatabase(t, &coderdtest.Options{
WorkspaceUsageTrackerFlush: flushCh,
WorkspaceUsageTrackerTick: tickCh,
})
user := coderdtest.CreateFirstUser(t, client)
r := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{
OrganizationID: user.OrganizationID,
OwnerID: user.UserID,
LastUsedAt: dbtime.Now().Add(-time.Minute),
}).WithAgent(
func(agent []*proto.Agent) []*proto.Agent {
for _, a := range agent {
a.ApiKeyScope = string(tc.apiKeyScope)
}
return agent
},
).Do()
ac := agentsdk.New(client.URL, agentsdk.WithFixedToken(r.AgentToken))
conn, err := ac.ConnectRPC(context.Background())
require.NoError(t, err)
defer func() {
_ = conn.Close()
}()
agentAPI := agentproto.NewDRPCAgentClient(conn)
_, err = agentAPI.UpdateStats(context.Background(), &agentproto.UpdateStatsRequest{
Stats: &agentproto.Stats{
ConnectionsByProto: map[string]int64{"TCP": 1},
ConnectionCount: 1,
RxPackets: 1,
RxBytes: 1,
TxPackets: 1,
TxBytes: 1,
SessionCountVscode: 1,
SessionCountJetbrains: 0,
SessionCountReconnectingPty: 0,
SessionCountSsh: 0,
ConnectionMedianLatencyMs: 10,
},
})
require.NoError(t, err)
tickCh <- dbtime.Now()
count := <-flushCh
require.Equal(t, 1, count, "expected one flush with one id")
newWorkspace, err := client.Workspace(context.Background(), r.Workspace.ID)
require.NoError(t, err)
assert.True(t,
newWorkspace.LastUsedAt.After(r.Workspace.LastUsedAt),
"%s is not after %s", newWorkspace.LastUsedAt, r.Workspace.LastUsedAt,
)
})
}
}
func TestWorkspaceAgentRPC_TailnetMethods(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitLong)
client, db := coderdtest.NewWithDatabase(t, nil)
user := coderdtest.CreateFirstUser(t, client)
workspace := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{
OrganizationID: user.OrganizationID,
OwnerID: user.UserID,
}).WithAgent().Do()
agentClient := agentsdk.New(client.URL, agentsdk.WithFixedToken(workspace.AgentToken))
conn, err := agentClient.ConnectRPC(ctx)
require.NoError(t, err)
t.Cleanup(func() {
_ = conn.Close()
})
tailnetClient := tailnetproto.NewDRPCTailnetClient(conn)
_, err = tailnetClient.RefreshResumeToken(ctx, &tailnetproto.RefreshResumeTokenRequest{})
require.Error(t, err)
require.EqualValues(t, drpcerr.Unimplemented, drpcerr.Code(err))
updates, err := tailnetClient.WorkspaceUpdates(ctx, &tailnetproto.WorkspaceUpdatesRequest{
WorkspaceOwnerId: tailnet.UUIDToByteSlice(user.UserID),
})
if err == nil {
_, err = updates.Recv()
}
require.Error(t, err)
require.EqualValues(t, drpcerr.Unimplemented, drpcerr.Code(err))
telemetry, err := tailnetClient.PostTelemetry(ctx, &tailnetproto.TelemetryRequest{})
require.NoError(t, err)
require.NotNil(t, telemetry)
agentAPI := agentproto.NewDRPCAgentClient(conn)
_, err = agentAPI.GetManifest(ctx, &agentproto.GetManifestRequest{})
require.NoError(t, err)
}
func TestAgentAPI_LargeManifest(t *testing.T) {
t.Parallel()
for _, tc := range []struct {
name string
apiKeyScope rbac.ScopeName
}{
{
name: "empty (backwards compat)",
apiKeyScope: "",
},
{
name: "all",
apiKeyScope: rbac.ScopeAll,
},
{
name: "no_user_data",
apiKeyScope: rbac.ScopeNoUserData,
},
{
name: "application_connect",
apiKeyScope: rbac.ScopeApplicationConnect,
},
} {
t.Run(tc.name, func(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitLong)
client, store := coderdtest.NewWithDatabase(t, nil)
adminUser := coderdtest.CreateFirstUser(t, client)
n := 512000
longScript := make([]byte, n)
for i := range longScript {
longScript[i] = 'q'
}
r := dbfake.WorkspaceBuild(t, store, database.WorkspaceTable{
OrganizationID: adminUser.OrganizationID,
OwnerID: adminUser.UserID,
}).WithAgent(func(agents []*proto.Agent) []*proto.Agent {
agents[0].Scripts = []*proto.Script{
{
Script: string(longScript),
},
}
agents[0].ApiKeyScope = string(tc.apiKeyScope)
return agents
}).Do()
ac := agentsdk.New(client.URL, agentsdk.WithFixedToken(r.AgentToken))
conn, err := ac.ConnectRPC(ctx)
defer func() {
_ = conn.Close()
}()
require.NoError(t, err)
agentAPI := agentproto.NewDRPCAgentClient(conn)
manifest, err := agentAPI.GetManifest(ctx, &agentproto.GetManifestRequest{})
require.NoError(t, err)
require.Len(t, manifest.Scripts, 1)
require.Len(t, manifest.Scripts[0].Script, n)
})
}
}
// End-to-end wiring: the deployment option must reach the manifest the
// agent actually receives over RPC.
func TestWorkspaceAgentRPCUserSecretFilePathDisabled(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitLong)
dv := coderdtest.DeploymentValues(t, func(dv *codersdk.DeploymentValues) {
dv.DisableUserSecretFilePath = true
})
db, ps := dbtestutil.NewDB(t)
client := coderdtest.New(t, &coderdtest.Options{
Database: db, Pubsub: ps, DeploymentValues: dv,
})
owner := coderdtest.CreateFirstUser(t, client)
// Seed rows the way a deployment that allowed file paths would have
// stored them; the API rejects new paths while the policy is on.
for _, seed := range []database.UserSecret{
{Name: "env-only", Value: "env-value", EnvName: "ENV_ONLY"},
{Name: "file-only", Value: "file-value", FilePath: "~/.file-only"},
{Name: "dual-target", Value: "dual-value", EnvName: "DUAL_TARGET", FilePath: "~/.dual-target"},
} {
dbgen.UserSecret(t, db, database.UserSecret{UserID: owner.UserID, Name: seed.Name},
func(p *database.CreateUserSecretParams) {
p.Value, p.EnvName, p.FilePath, p.Enabled = seed.Value, seed.EnvName, seed.FilePath, true
})
}
workspace := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{
OrganizationID: owner.OrganizationID,
OwnerID: owner.UserID,
}).WithAgent().Do()
agentClient := agentsdk.New(client.URL, agentsdk.WithFixedToken(workspace.AgentToken))
conn, err := agentClient.ConnectRPC(ctx)
require.NoError(t, err)
defer conn.Close()
manifest, err := agentproto.NewDRPCAgentClient(conn).GetManifest(ctx, &agentproto.GetManifestRequest{})
require.NoError(t, err)
require.Len(t, manifest.Secrets, 2)
byEnv := make(map[string][]byte, len(manifest.Secrets))
for _, secret := range manifest.Secrets {
require.Empty(t, secret.FilePath)
byEnv[secret.EnvName] = secret.Value
}
require.Equal(t, map[string][]byte{
"ENV_ONLY": []byte("env-value"),
"DUAL_TARGET": []byte("dual-value"),
}, byEnv)
}
func TestWorkspaceAgentRPCRole(t *testing.T) {
t.Parallel()
t.Run("AgentRoleMonitorsConnection", func(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitLong)
client, db := coderdtest.NewWithDatabase(t, nil)
user := coderdtest.CreateFirstUser(t, client)
r := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{
OrganizationID: user.OrganizationID,
OwnerID: user.UserID,
}).WithAgent().Do()
// Connect with role=agent using ConnectRPCWithRole. This is
// how the real workspace agent connects.
ac := agentsdk.New(client.URL, agentsdk.WithFixedToken(r.AgentToken))
conn, err := ac.ConnectRPCWithRole(ctx, "agent")
require.NoError(t, err)
defer func() {
_ = conn.Close()
}()
// The connection monitor updates the database asynchronously,
// so we need to wait for first_connected_at to be set.
var agent database.WorkspaceAgent
require.Eventually(t, func() bool {
agent, err = db.GetWorkspaceAgentByID(dbauthz.AsSystemRestricted(ctx), r.Agents[0].ID)
if err != nil {
return false
}
return agent.FirstConnectedAt.Valid
}, testutil.WaitShort, testutil.IntervalFast)
assert.True(t, agent.LastConnectedAt.Valid,
"last_connected_at should be set for agent role")
})
t.Run("NonAgentRoleSkipsMonitoring", func(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitLong)
client, db := coderdtest.NewWithDatabase(t, nil)
user := coderdtest.CreateFirstUser(t, client)
r := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{
OrganizationID: user.OrganizationID,
OwnerID: user.UserID,
}).WithAgent().Do()
// Connect with a non-agent role using ConnectRPCWithRole.
// This is how coder-logstream-kube should connect.
ac := agentsdk.New(client.URL, agentsdk.WithFixedToken(r.AgentToken))
conn, err := ac.ConnectRPCWithRole(ctx, "logstream-kube")
require.NoError(t, err)
// Send a log to confirm the RPC connection is functional.
agentAPI := agentproto.NewDRPCAgentClient(conn)
_, err = agentAPI.BatchCreateLogs(ctx, &agentproto.BatchCreateLogsRequest{
LogSourceId: []byte{0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0},
})
// We don't care about the log source error, just that the
// RPC is functional.
_ = err
// Close the connection and give the server time to process.
_ = conn.Close()
// Verify that connectivity timestamps were never set
// (first_connected_at, last_connected_at, disconnected_at).
require.Never(t, func() bool {
agent, err := db.GetWorkspaceAgentByID(dbauthz.AsSystemRestricted(ctx), r.Agents[0].ID)
if err != nil {
return false
}
return agent.FirstConnectedAt.Valid || agent.LastConnectedAt.Valid || agent.DisconnectedAt.Valid
}, testutil.IntervalMedium, testutil.IntervalFast, "connectivity timestamps should NOT be set for non-agent role")
})
// NOTE: Backward compatibility (empty role) is implicitly tested by
// existing tests like TestWorkspaceAgentReportStats which use
// ConnectRPC() (no role). The server defaults to monitoring when
// the role query parameter is omitted.
}
func TestWorkspaceAgentRPCConnectionLog(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitLong)
connLogger := connectionlog.NewFake()
client, db := coderdtest.NewWithDatabase(t, &coderdtest.Options{
ConnectionLogger: connLogger,
})
user := coderdtest.CreateFirstUser(t, client)
r := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{
OrganizationID: user.OrganizationID,
OwnerID: user.UserID,
}).WithAgent().Do()
ac := agentsdk.New(client.URL, agentsdk.WithFixedToken(r.AgentToken))
conn, err := ac.ConnectRPC(ctx)
require.NoError(t, err)
defer conn.Close()
// Report over the wire so coderd must accept and log the method and app.
id := uuid.New()
_, err = agentproto.NewDRPCAgentClient(conn).ReportConnection(ctx, &agentproto.ReportConnectionRequest{
Connection: &agentproto.Connection{
Id: id[:],
Action: agentproto.Connection_CONNECT,
Timestamp: timestamppb.Now(),
ConnectionMethod: agentproto.Connection_METHOD_SSH,
AppName: "Some-IDE",
},
})
require.NoError(t, err)
require.True(t, connLogger.Contains(t, database.UpsertConnectionLogParams{
WorkspaceID: r.Workspace.ID,
ConnectionID: uuid.NullUUID{UUID: id, Valid: true},
ConnectionMethod: database.ConnectionLogMethodSSH,
AppNameOrPort: sql.NullString{String: "some_ide", Valid: true},
}))
}