Fix import cycle

Emyrk · Emyrk · commit 15ed4676e99a · 2023-04-04T16:39:10.000-05:00
diff --git a/coderd/coderd.go b/coderd/coderd.go
@@ -356,7 +356,7 @@ func New(options *Options) *API {
 	apiRateLimiter := httpmw.RateLimit(options.APIRateLimit, time.Minute)
 
 	derpHandler := derphttp.Handler(api.DERPServer)
-	derpHandler = tailnet.WithWebsocketSupport(api.WebsocketWatch, api.DERPServer, derpHandler)
+	derpHandler = tailnet.WithWebsocketSupport(api.WebsocketWatch.Accept, api.DERPServer, derpHandler)
 
 	r.Use(
 		httpmw.Recover(api.Logger),
diff --git a/coderd/healthcheck/derp_test.go b/coderd/healthcheck/derp_test.go
@@ -132,7 +132,7 @@ func TestDERP(t *testing.T) {
 		defer derpSrv.Close()
 
 		sockets := activewebsockets.New(ctx)
-		handler := tailnet.WithWebsocketSupport(sockets, derpSrv, derphttp.Handler(derpSrv))
+		handler := tailnet.WithWebsocketSupport(sockets.Accept, derpSrv, derphttp.Handler(derpSrv))
 		defer sockets.Close()
 
 		srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
diff --git a/tailnet/derp.go b/tailnet/derp.go
@@ -8,15 +8,17 @@ import (
 	"nhooyr.io/websocket"
 	"tailscale.com/derp"
 	"tailscale.com/net/wsconn"
-
-	"github.com/coder/coder/coderd/activewebsockets"
 )
 
+type HandleWebsocket func(rw http.ResponseWriter, r *http.Request, options *websocket.AcceptOptions, f func(conn *websocket.Conn))
+
 // WithWebsocketSupport returns an http.Handler that upgrades
 // connections to the "derp" subprotocol to WebSockets and
 // passes them to the DERP server.
 // Taken from: https://github.com/tailscale/tailscale/blob/e3211ff88ba85435f70984cf67d9b353f3d650d8/cmd/derper/websocket.go#L21
-func WithWebsocketSupport(sockets *activewebsockets.Active, s *derp.Server, base http.Handler) http.Handler {
+// The accept function is used to accept the websocket connection and allows the caller to
+// also affect the lifecycle of the websocket connection. (Eg. to close the connection on shutdown)
+func WithWebsocketSupport(accept HandleWebsocket, s *derp.Server, base http.Handler) http.Handler {
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		up := strings.ToLower(r.Header.Get("Upgrade"))
 
@@ -28,7 +30,7 @@ func WithWebsocketSupport(sockets *activewebsockets.Active, s *derp.Server, base
 			return
 		}
 
-		sockets.Accept(w, r, &websocket.AcceptOptions{
+		accept(w, r, &websocket.AcceptOptions{
 			Subprotocols:   []string{"derp"},
 			OriginPatterns: []string{"*"},
 			// Disable compression because we transmit WireGuard messages that
diff --git a/tailnet/tailnettest/tailnettest.go b/tailnet/tailnettest/tailnettest.go
@@ -78,7 +78,7 @@ func RunDERPOnlyWebSockets(t *testing.T) *tailcfg.DERPMap {
 	defer cancel()
 	sockets := activewebsockets.New(ctx)
 
-	handler = tailnet.WithWebsocketSupport(sockets, d, handler)
+	handler = tailnet.WithWebsocketSupport(sockets.Accept, d, handler)
 	server := httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		if r.URL.Path != "/derp" {
 			handler.ServeHTTP(w, r)
