diff --git a/cli/portforward.go b/cli/portforward.go index cd7160e31f..c321041827 100644 --- a/cli/portforward.go +++ b/cli/portforward.go @@ -11,8 +11,11 @@ import ( "strconv" "strings" "sync" + "sync/atomic" "syscall" + "time" + "github.com/google/uuid" "golang.org/x/xerrors" "cdr.dev/slog/v3" @@ -132,6 +135,7 @@ func (r *RootCmd) portForward() *serpent.Command { var ( wg = new(sync.WaitGroup) listeners = make([]net.Listener, 0, len(specs)*2) + usage = new(portForwardUsage) closeAllListeners = func() { logger.Debug(ctx, "closing all listeners") for _, l := range listeners { @@ -149,7 +153,7 @@ func (r *RootCmd) portForward() *serpent.Command { // first, opportunistically try to listen on IPv6 spec6 := spec spec6.listenHost = ipv6Loopback - l6, err6 := listenAndPortForward(ctx, inv, conn, wg, spec6, logger) + l6, err6 := listenAndPortForward(ctx, inv, conn, wg, spec6, logger, usage) if err6 != nil { logger.Info(ctx, "failed to opportunistically listen on IPv6", slog.F("spec", spec), slog.Error(err6)) } else { @@ -157,7 +161,7 @@ func (r *RootCmd) portForward() *serpent.Command { } spec.listenHost = ipv4Loopback } - l, err := listenAndPortForward(ctx, inv, conn, wg, spec, logger) + l, err := listenAndPortForward(ctx, inv, conn, wg, spec, logger, usage) if err != nil { logger.Error(ctx, "failed to listen", slog.F("spec", spec), slog.Error(err)) return err @@ -165,7 +169,7 @@ func (r *RootCmd) portForward() *serpent.Command { listeners = append(listeners, l) } - stopUpdating := client.UpdateWorkspaceUsageContext(ctx, workspace.ID) + stopUpdating := reportPortForwardUsage(ctx, client, workspace.ID, workspaceAgent.ID, usage) // Wait for the context to be canceled or for a signal and close // all listeners. @@ -219,6 +223,67 @@ func (r *RootCmd) portForward() *serpent.Command { return cmd } +// portForwardUsage counts forwarded TCP connections for usage reporting. +type portForwardUsage struct { + active atomic.Int64 + seen atomic.Bool // Set when a connection opens, cleared by inUse. +} + +// track records an open connection. Call the returned func when it closes. +func (u *portForwardUsage) track() func() { + u.active.Add(1) + u.seen.Store(true) + return func() { u.active.Add(-1) } +} + +// inUse reports whether a connection is open or opened since the last call. +func (u *portForwardUsage) inUse() bool { + // Swap first so a connection that opened and closed between calls counts. + return u.seen.Swap(false) || u.active.Load() > 0 +} + +// reportPortForwardUsage reports usage now, every minute, and once more on +// stop. Call the returned func to stop and wait for the final report. +func reportPortForwardUsage(ctx context.Context, client *codersdk.Client, workspaceID, agentID uuid.UUID, usage *portForwardUsage) func() { + report := func(ctx context.Context) { + var req codersdk.PostWorkspaceUsageRequest + if usage.inUse() { + req = codersdk.PostWorkspaceUsageRequest{ + AgentID: agentID, + AppName: string(codersdk.UsageAppNamePortForwarding), + } + } + if err := client.PostWorkspaceUsageWithBody(ctx, workspaceID, req); err != nil { + client.Logger().Warn(ctx, "failed to post workspace usage", slog.Error(err)) + } + } + + ctx, cancel := context.WithCancel(ctx) + report(ctx) + done := make(chan struct{}) + go func() { + defer close(done) + ticker := time.NewTicker(time.Minute) + defer ticker.Stop() + for { + select { + case <-ticker.C: + report(ctx) + case <-ctx.Done(): + // ctx is canceled, so send the final report on a detached one. + finalCtx, finalCancel := context.WithTimeout(context.WithoutCancel(ctx), 5*time.Second) + report(finalCtx) + finalCancel() + return + } + } + }() + return func() { + cancel() + <-done + } +} + func listenAndPortForward( ctx context.Context, inv *serpent.Invocation, @@ -226,6 +291,7 @@ func listenAndPortForward( wg *sync.WaitGroup, spec portForwardSpec, logger slog.Logger, + usage *portForwardUsage, ) (net.Listener, error) { logger = logger.With( slog.F("network", spec.network), @@ -276,6 +342,12 @@ func listenAndPortForward( logger.Debug(ctx, "dialed remote", slog.F("remote_addr", netConn.RemoteAddr())) + // UDP forwards never see EOF, so they would count as active + // until exit. + if spec.network == "tcp" { + defer usage.track()() + } + agentssh.Bicopy(ctx, netConn, remoteConn) logger.Debug(ctx, "connection closing", slog.F("remote_addr", netConn.RemoteAddr())) diff --git a/cli/portforward_internal_test.go b/cli/portforward_internal_test.go index 5698363f95..9f74d7c70d 100644 --- a/cli/portforward_internal_test.go +++ b/cli/portforward_internal_test.go @@ -115,3 +115,60 @@ func Test_parsePortForwards(t *testing.T) { }) } } + +func TestPortForwardUsageRequest(t *testing.T) { + t.Parallel() + + for _, tc := range []struct { + name string + // steps runs before each request. Each entry is one reporting tick. + steps []func(u *portForwardUsage, done *func()) + want []bool + }{ + { + name: "Idle", + steps: []func(*portForwardUsage, *func()){nil, nil}, + want: []bool{false, false}, + }, + { + name: "OpenConnection", + steps: []func(*portForwardUsage, *func()){ + func(u *portForwardUsage, done *func()) { *done = u.track() }, + nil, + }, + want: []bool{true, true}, + }, + { + name: "GoesIdleAfterClose", + steps: []func(*portForwardUsage, *func()){ + func(u *portForwardUsage, done *func()) { *done = u.track() }, + nil, + func(_ *portForwardUsage, done *func()) { (*done)() }, + }, + want: []bool{true, true, false}, + }, + { + name: "OpenOnlyBetweenTicks", + steps: []func(*portForwardUsage, *func()){ + func(u *portForwardUsage, _ *func()) { u.track()() }, + nil, + }, + want: []bool{true, false}, + }, + } { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + var ( + u portForwardUsage + done func() + ) + for i, step := range tc.steps { + if step != nil { + step(&u, &done) + } + require.Equal(t, tc.want[i], u.inUse(), "tick %d", i) + } + }) + } +} diff --git a/cli/portforward_test.go b/cli/portforward_test.go index fd693120c3..e3df92afe4 100644 --- a/cli/portforward_test.go +++ b/cli/portforward_test.go @@ -19,11 +19,13 @@ import ( "github.com/coder/coder/v2/agent" "github.com/coder/coder/v2/agent/agenttest" + agentproto "github.com/coder/coder/v2/agent/proto" "github.com/coder/coder/v2/cli/clitest" "github.com/coder/coder/v2/coderd/coderdtest" "github.com/coder/coder/v2/coderd/database" "github.com/coder/coder/v2/coderd/database/dbfake" "github.com/coder/coder/v2/coderd/database/dbtime" + "github.com/coder/coder/v2/coderd/workspacestats/workspacestatstest" "github.com/coder/coder/v2/codersdk" "github.com/coder/coder/v2/testutil" "github.com/coder/coder/v2/testutil/expecter" @@ -500,3 +502,69 @@ func assertWritePayload(t *testing.T, w io.Writer, prefix []byte) { assert.NoError(t, err, "write payload") assert.Equal(t, len(payload), n, "payload length does not match") } + +// TestPortForward_UpdateUsage checks that a forwarded connection is reported +// under the port_forwarding app for the workspace's agent. +func TestPortForward_UpdateUsage(t *testing.T) { + t.Parallel() + + dv := coderdtest.DeploymentValues(t) + dv.Experiments = []string{string(codersdk.ExperimentWorkspaceUsage)} + batcher := &workspacestatstest.StatsBatcher{ + LastStats: &agentproto.Stats{}, + } + client, db := coderdtest.NewWithDatabase(t, &coderdtest.Options{ + DeploymentValues: dv, + StatsBatcher: batcher, + }) + owner := coderdtest.CreateFirstUser(t, client) + member, memberUser := coderdtest.CreateAnotherUser(t, client, owner.OrganizationID) + workspace := runAgent(t, client, memberUser.ID, db) + + ws, err := client.Workspace(context.Background(), workspace.ID) + require.NoError(t, err) + var agentID uuid.UUID + for _, res := range ws.LatestBuild.Resources { + for _, a := range res.Agents { + agentID = a.ID + } + } + require.NotEqual(t, uuid.Nil, agentID, "workspace has no agent") + + prefix := generateRandomPrefix(t) + remoteLis, err := net.Listen("tcp", "127.0.0.1:0") + require.NoError(t, err, "create TCP listener") + p1 := setupTestListener(t, remoteLis, prefix) + + flag := fmt.Sprintf("--tcp=5555:%v", p1) + inv, root := clitest.New(t, "port-forward", workspace.Name, flag) + clitest.SetupConfig(t, member, root) + stdout := expecter.NewAttachedToInvocation(t, inv) + + iNet := testutil.NewInProcNet() + inv.Net = iNet + ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong) + defer cancel() + errC := make(chan error) + go func() { + errC <- inv.WithContext(ctx).Run() + }() + stdout.ExpectMatch(ctx, "Ready!") + + // Open and use a forwarded connection so the session is reported. + dialCtx, dialCtxCancel := context.WithTimeout(ctx, testutil.WaitShort) + defer dialCtxCancel() + c1, err := iNet.Dial(dialCtx, testutil.NewAddr("tcp", "127.0.0.1:5555")) + require.NoError(t, err, "open connection to 'local' listener") + testDial(t, c1, prefix) + require.NoError(t, c1.Close()) + + cancel() + err = <-errC + require.ErrorIs(t, err, context.Canceled) + + batcher.Mu.Lock() + defer batcher.Mu.Unlock() + require.Equal(t, map[string]int64{"port_forwarding": 1}, batcher.LastStats.GetSessionCounts()) + require.Equal(t, agentID, batcher.LastAgentID) +} diff --git a/coderd/apidoc/docs.go b/coderd/apidoc/docs.go index 105b173043..7f650e6a8d 100644 --- a/coderd/apidoc/docs.go +++ b/coderd/apidoc/docs.go @@ -19664,6 +19664,7 @@ const docTemplate = `{ "jetbrains", "ssh", "reconnecting_pty", + "port_forwarding", "sftp", "unknown" ], @@ -19672,6 +19673,7 @@ const docTemplate = `{ "AppFamilyJetBrains", "AppFamilySSH", "AppFamilyReconnectingPTY", + "AppFamilyPortForwarding", "AppFamilySFTP", "AppFamilyUnknown" ] diff --git a/coderd/apidoc/swagger.json b/coderd/apidoc/swagger.json index b265722006..778f3cc1fa 100644 --- a/coderd/apidoc/swagger.json +++ b/coderd/apidoc/swagger.json @@ -17665,6 +17665,7 @@ "jetbrains", "ssh", "reconnecting_pty", + "port_forwarding", "sftp", "unknown" ], @@ -17673,6 +17674,7 @@ "AppFamilyJetBrains", "AppFamilySSH", "AppFamilyReconnectingPTY", + "AppFamilyPortForwarding", "AppFamilySFTP", "AppFamilyUnknown" ] diff --git a/codersdk/appname.go b/codersdk/appname.go index 13cc0f13c7..9a06784a44 100644 --- a/codersdk/appname.go +++ b/codersdk/appname.go @@ -26,6 +26,7 @@ const ( AppFamilyJetBrains AppFamilyName = "jetbrains" AppFamilySSH AppFamilyName = "ssh" AppFamilyReconnectingPTY AppFamilyName = "reconnecting_pty" + AppFamilyPortForwarding AppFamilyName = "port_forwarding" // AppFamilySFTP only comes from history the sftp_mins column recorded. AppFamilySFTP AppFamilyName = "sftp" AppFamilyUnknown AppFamilyName = "unknown" @@ -95,6 +96,7 @@ var sessionApps = map[string]sessionApp{ "zed": {AppFamilySSH, "Zed", "/icon/zed.svg"}, "ssh": {AppFamilySSH, "SSH", "/icon/terminal.svg"}, "reconnecting_pty": {AppFamilyReconnectingPTY, "Web Terminal", ""}, + "port_forwarding": {AppFamilyPortForwarding, "Port Forwarding", ""}, } // SessionCountApps pairs each count with its presentation, keyed by app name. diff --git a/codersdk/appname_test.go b/codersdk/appname_test.go index b40bed8f97..ffeaf62d9b 100644 --- a/codersdk/appname_test.go +++ b/codersdk/appname_test.go @@ -85,6 +85,7 @@ func TestAppNameFamily(t *testing.T) { {"CaseInsensitive", "JetBrains", codersdk.AppFamilyJetBrains}, {"SSHClientJoinsSSHFamily", "zed", codersdk.AppFamilySSH}, {"Alias", "reconnecting-pty", codersdk.AppFamilyReconnectingPTY}, + {"PortForwarding", "port_forwarding", codersdk.AppFamilyPortForwarding}, {"Unknown", "SomeFutureIDE", codersdk.AppFamilyUnknown}, } { t.Run(tc.name, func(t *testing.T) { diff --git a/codersdk/workspaces.go b/codersdk/workspaces.go index 0b06494e93..d928579b39 100644 --- a/codersdk/workspaces.go +++ b/codersdk/workspaces.go @@ -405,6 +405,7 @@ const ( UsageAppNameJetbrains UsageAppName = "jetbrains" UsageAppNameReconnectingPty UsageAppName = "reconnecting-pty" UsageAppNameSSH UsageAppName = "ssh" + UsageAppNamePortForwarding UsageAppName = "port_forwarding" ) // PostWorkspaceUsage marks the workspace as having been used recently and records an app stat. diff --git a/docs/reference/api/schemas.md b/docs/reference/api/schemas.md index 643ffde327..e422aba9c3 100644 --- a/docs/reference/api/schemas.md +++ b/docs/reference/api/schemas.md @@ -1615,9 +1615,9 @@ None #### Enumerated Values -| Value(s) | -|---------------------------------------------------------------------| -| `jetbrains`, `reconnecting_pty`, `sftp`, `ssh`, `unknown`, `vscode` | +| Value(s) | +|----------------------------------------------------------------------------------------| +| `jetbrains`, `port_forwarding`, `reconnecting_pty`, `sftp`, `ssh`, `unknown`, `vscode` | ## codersdk.AppHostResponse diff --git a/site/src/api/typesGenerated.ts b/site/src/api/typesGenerated.ts index e4f4d7c9fd..51abdc08de 100644 --- a/site/src/api/typesGenerated.ts +++ b/site/src/api/typesGenerated.ts @@ -1613,6 +1613,7 @@ export const AnthropicInlineImageCapBytes = 5242880; // From codersdk/appname.go export type AppFamilyName = | "jetbrains" + | "port_forwarding" | "reconnecting_pty" | "sftp" | "ssh" @@ -1621,6 +1622,7 @@ export type AppFamilyName = export const AppFamilyNames: AppFamilyName[] = [ "jetbrains", + "port_forwarding", "reconnecting_pty", "sftp", "ssh", @@ -11139,10 +11141,16 @@ export interface UpsertWorkspaceAgentPortShareRequest { } // From codersdk/workspaces.go -export type UsageAppName = "jetbrains" | "reconnecting-pty" | "ssh" | "vscode"; +export type UsageAppName = + | "jetbrains" + | "port_forwarding" + | "reconnecting-pty" + | "ssh" + | "vscode"; export const UsageAppNames: UsageAppName[] = [ "jetbrains", + "port_forwarding", "reconnecting-pty", "ssh", "vscode", diff --git a/site/src/modules/dashboard/DeploymentBanner/DeploymentBannerView.tsx b/site/src/modules/dashboard/DeploymentBanner/DeploymentBannerView.tsx index fbcfca557c..97be029b02 100644 --- a/site/src/modules/dashboard/DeploymentBanner/DeploymentBannerView.tsx +++ b/site/src/modules/dashboard/DeploymentBanner/DeploymentBannerView.tsx @@ -308,6 +308,7 @@ const SESSION_FAMILIES = { }, unknown: { name: "Other", icon: }, sftp: null, + port_forwarding: null, } satisfies Record< AppFamilyName, { name: string; icon: React.ReactNode } | null