Add latency check to wsproxy

Emyrk · Emyrk · commit 3cae4b06ff2e · 2023-05-10T13:10:54.000-05:00
diff --git a/coderd/coderd.go b/coderd/coderd.go
@@ -805,6 +805,17 @@ func New(options *Options) *API {
 		return []string{}
 	})
 	r.NotFound(cspMW(compressHandler(http.HandlerFunc(api.siteHandler.ServeHTTP))).ServeHTTP)
+
+	// This must be before all middleware to improve the response time.
+	// So make a new router, and mount the old one as the root.
+	rootRouter := chi.NewRouter()
+	// This is the only route we add before all the middleware.
+	// We want to time the latency of the request, so any middleware will
+	// interfere with that timing.
+	rootRouter.Get("/latency-check", LatencyCheck(api.AccessURL.String()))
+	rootRouter.Mount("/", r)
+	api.RootHandler = rootRouter
+
 	return api
 }
 
diff --git a/coderd/coderd_test.go b/coderd/coderd_test.go
@@ -124,6 +124,15 @@ func TestDERPLatencyCheck(t *testing.T) {
 	require.Equal(t, http.StatusOK, res.StatusCode)
 }
 
+func TestFastLatencyCheck(t *testing.T) {
+	t.Parallel()
+	client := coderdtest.New(t, nil)
+	res, err := client.Request(context.Background(), http.MethodGet, "/latency-check", nil)
+	require.NoError(t, err)
+	defer res.Body.Close()
+	require.Equal(t, http.StatusOK, res.StatusCode)
+}
+
 func TestHealthz(t *testing.T) {
 	t.Parallel()
 	client := coderdtest.New(t, nil)
diff --git a/coderd/latencycheck.go b/coderd/latencycheck.go
@@ -0,0 +1,16 @@
+package coderd
+
+import (
+	"net/http"
+	"strings"
+)
+
+func LatencyCheck(allowedOrigins ...string) http.HandlerFunc {
+	origins := strings.Join(allowedOrigins, ",")
+	return func(rw http.ResponseWriter, r *http.Request) {
+		// Allowing timing information to be shared. This allows the browser
+		// to exclude TLS handshake timing.
+		rw.Header().Set("Timing-Allow-Origin", origins)
+		rw.WriteHeader(http.StatusOK)
+	}
+}
diff --git a/enterprise/wsproxy/wsproxy.go b/enterprise/wsproxy/wsproxy.go
@@ -10,8 +10,6 @@ import (
 	"strings"
 	"time"
 
-	"github.com/go-chi/cors"
-
 	"github.com/go-chi/chi/v5"
 	"github.com/go-chi/cors"
 	"github.com/google/uuid"
@@ -21,6 +19,7 @@ import (
 
 	"cdr.dev/slog"
 	"github.com/coder/coder/buildinfo"
+	"github.com/coder/coder/coderd"
 	"github.com/coder/coder/coderd/httpapi"
 	"github.com/coder/coder/coderd/httpmw"
 	"github.com/coder/coder/coderd/tracing"
@@ -262,6 +261,12 @@ func New(ctx context.Context, opts *Options) (*Server, error) {
 		})
 	})
 
+	// See coderd/coderd.go for why we need this.
+	rootRouter := chi.NewRouter()
+	rootRouter.Get("/latency-check", coderd.LatencyCheck(s.DashboardURL.String(), s.AppServer.AccessURL.String()))
+	rootRouter.Mount("/", r)
+	s.Handler = rootRouter
+
 	return s, nil
 }
 
