Hard & simplify CORORS handling logic

dannykopping · dannykopping · commit 3e4ba618e012 · 2024-11-27T09:10:28.000Z
Signed-off-by: Danny Kopping &lt;danny@coder.com&gt;
diff --git a/coderd/httpmw/cors.go b/coderd/httpmw/cors.go
@@ -8,7 +8,6 @@ import (
 	"github.com/go-chi/cors"
 
 	"github.com/coder/coder/v2/coderd/workspaceapps/appurl"
-	ws_cors "github.com/coder/coder/v2/coderd/workspaceapps/cors"
 )
 
 const (
@@ -48,11 +47,6 @@ func Cors(allowAll bool, origins ...string) func(next http.Handler) http.Handler
 func WorkspaceAppCors(regex *regexp.Regexp, app appurl.ApplicationURL) func(next http.Handler) http.Handler {
 	return cors.Handler(cors.Options{
 		AllowOriginFunc: func(r *http.Request, rawOrigin string) bool {
-			// If passthru behavior is set, disable our simplified CORS handling.
-			if ws_cors.HasBehavior(r.Context(), ws_cors.AppCORSBehaviorPassthru) {
-				return true
-			}
-
 			origin, err := url.Parse(rawOrigin)
 			if rawOrigin == "" || origin.Host == "" || err != nil {
 				return false
diff --git a/coderd/workspaceapps/proxy.go b/coderd/workspaceapps/proxy.go
@@ -424,24 +424,42 @@ func (s *Server) HandleSubdomain(middlewares ...func(http.Handler) http.Handler)
 				return
 			}
 
-			// Use the passed in app middlewares and CORS middleware with the token
-			mws := chi.Middlewares(append(middlewares, s.injectCORSBehavior(token), httpmw.WorkspaceAppCors(s.HostnameRegex, app)))
+			// Proxy the request (possibly with the CORS middleware).
+			mws := chi.Middlewares(append(middlewares, s.determineCORSBehavior(token, app)))
 			mws.Handler(http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) {
 				s.proxyWorkspaceApp(rw, r, *token, r.URL.Path, app)
 			})).ServeHTTP(rw, r.WithContext(ctx))
 		})
 	}
 }
 
-func (s *Server) injectCORSBehavior(token *SignedToken) func(http.Handler) http.Handler {
+// determineCORSBehavior examines the given token and conditionally applies
+// CORS middleware if the token specifies that behavior.
+func (s *Server) determineCORSBehavior(token *SignedToken, app appurl.ApplicationURL) func(http.Handler) http.Handler {
 	return func(next http.Handler) http.Handler {
+		// Create the CORS middleware handler upfront.
+		corsHandler := httpmw.WorkspaceAppCors(s.HostnameRegex, app)(next)
+
 		return http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) {
 			var behavior cors.AppCORSBehavior
 			if token != nil {
 				behavior = token.CORSBehavior
 			}
 
-			next.ServeHTTP(rw, r.WithContext(cors.WithBehavior(r.Context(), behavior)))
+			// Add behavior to context regardless of which handler we use,
+			// since we will use this later on to determine if we should strip
+			// CORS headers in the response.
+			r = r.WithContext(cors.WithBehavior(r.Context(), behavior))
+
+			switch behavior {
+			case cors.AppCORSBehaviorPassthru:
+				// Bypass the CORS middleware.
+				next.ServeHTTP(rw, r)
+				return
+			default:
+				// Apply the CORS middleware.
+				corsHandler.ServeHTTP(rw, r)
+			}
 		})
 	}
 }
