@@ -7,9 +7,13 @@ import (
77 "net/http"
88 "net/netip"
99 "strconv"
10+ "strings"
1011 "sync"
12+ "sync/atomic"
1113 "testing"
1214
15+ "github.com/davecgh/go-spew/spew"
16+ "github.com/google/uuid"
1317 "github.com/stretchr/testify/assert"
1418 "github.com/stretchr/testify/require"
1519 "go.uber.org/goleak"
@@ -18,8 +22,13 @@ import (
1822 "cdr.dev/slog"
1923 "cdr.dev/slog/sloggers/slogtest"
2024
25+ "github.com/coder/coder/v2/agent"
2126 "github.com/coder/coder/v2/buildinfo"
27+ "github.com/coder/coder/v2/coderd"
2228 "github.com/coder/coder/v2/coderd/coderdtest"
29+ "github.com/coder/coder/v2/codersdk"
30+ "github.com/coder/coder/v2/codersdk/agentsdk"
31+ "github.com/coder/coder/v2/provisioner/echo"
2332 "github.com/coder/coder/v2/tailnet"
2433 "github.com/coder/coder/v2/testutil"
2534)
@@ -119,6 +128,91 @@ func TestDERP(t *testing.T) {
119128 w2 .Close ()
120129}
121130
131+ func TestDERPForceWebSockets (t * testing.T ) {
132+ t .Parallel ()
133+
134+ dv := coderdtest .DeploymentValues (t )
135+ dv .DERP .Config .ForceWebSockets = true
136+ dv .DERP .Config .BlockDirect = true // to ensure the test always uses DERP
137+
138+ // Manually create a server so we can influence the HTTP handler.
139+ options := & coderdtest.Options {
140+ DeploymentValues : dv ,
141+ }
142+ setHandler , cancelFunc , serverURL , newOptions := coderdtest .NewOptions (t , options )
143+ coderAPI := coderd .New (newOptions )
144+ t .Cleanup (func () {
145+ cancelFunc ()
146+ _ = coderAPI .Close ()
147+ })
148+
149+ // Set the HTTP handler to a custom one that ensures all /derp calls are
150+ // WebSockets and not `Upgrade: derp`.
151+ var upgradeCount int64
152+ setHandler (http .HandlerFunc (func (rw http.ResponseWriter , r * http.Request ) {
153+ if strings .HasPrefix (r .URL .Path , "/derp" ) {
154+ up := r .Header .Get ("Upgrade" )
155+ if up != "" && up != "websocket" {
156+ t .Errorf ("expected Upgrade: websocket, got %q" , up )
157+ } else {
158+ atomic .AddInt64 (& upgradeCount , 1 )
159+ }
160+ }
161+
162+ coderAPI .RootHandler .ServeHTTP (rw , r )
163+ }))
164+
165+ // Start a provisioner daemon.
166+ provisionerCloser := coderdtest .NewProvisionerDaemon (t , coderAPI )
167+ t .Cleanup (func () {
168+ _ = provisionerCloser .Close ()
169+ })
170+
171+ client := codersdk .New (serverURL )
172+ t .Cleanup (func () {
173+ client .HTTPClient .CloseIdleConnections ()
174+ })
175+ user := coderdtest .CreateFirstUser (t , client )
176+
177+ gen , err := client .WorkspaceAgentConnectionInfoGeneric (context .Background ())
178+ require .NoError (t , err )
179+ t .Log (spew .Sdump (gen ))
180+
181+ authToken := uuid .NewString ()
182+ version := coderdtest .CreateTemplateVersion (t , client , user .OrganizationID , & echo.Responses {
183+ Parse : echo .ParseComplete ,
184+ ProvisionPlan : echo .ProvisionComplete ,
185+ ProvisionApply : echo .ProvisionApplyWithAgent (authToken ),
186+ })
187+ template := coderdtest .CreateTemplate (t , client , user .OrganizationID , version .ID )
188+ coderdtest .AwaitTemplateVersionJob (t , client , version .ID )
189+ workspace := coderdtest .CreateWorkspace (t , client , user .OrganizationID , template .ID )
190+ coderdtest .AwaitWorkspaceBuildJob (t , client , workspace .LatestBuild .ID )
191+
192+ agentClient := agentsdk .New (client .URL )
193+ agentClient .SetSessionToken (authToken )
194+ agentCloser := agent .New (agent.Options {
195+ Client : agentClient ,
196+ Logger : slogtest .Make (t , nil ).Named ("agent" ).Leveled (slog .LevelDebug ),
197+ })
198+ defer func () {
199+ _ = agentCloser .Close ()
200+ }()
201+
202+ ctx , cancel := context .WithTimeout (context .Background (), testutil .WaitLong )
203+ defer cancel ()
204+
205+ resources := coderdtest .AwaitWorkspaceAgents (t , client , workspace .ID )
206+ conn , err := client .DialWorkspaceAgent (ctx , resources [0 ].Agents [0 ].ID , nil )
207+ require .NoError (t , err )
208+ defer func () {
209+ _ = conn .Close ()
210+ }()
211+ conn .AwaitReachable (ctx )
212+
213+ require .GreaterOrEqual (t , atomic .LoadInt64 (& upgradeCount ), int64 (1 ), "expected at least one /derp call" )
214+ }
215+
122216func TestDERPLatencyCheck (t * testing.T ) {
123217 t .Parallel ()
124218 client := coderdtest .New (t , nil )
0 commit comments