chore: Make workspace and agent extract a middleware

coder · Emyrk · Aug 29, 2022 · Aug 24, 2022 · Aug 24, 2022 · Aug 24, 2022
commit 7d0bc561c266d4a3d20de800c9033cf158f1011f
diff --git a/coderd/coderd.go b/coderd/coderd.go
@@ -138,14 +138,15 @@ func New(options *Options) *API {
 			httpmw.RateLimitPerMinute(options.APIRateLimit),
 			httpmw.ExtractAPIKey(options.Database, oauthConfigs, true),
 			httpmw.ExtractUserParam(api.Database),
+			httpmw.ExtractWorkspaceAndAgentParam(api.Database),
 		)
 		r.HandleFunc("/*", api.workspaceAppsProxyPath)
 	}
 	// %40 is the encoded character of the @ symbol. VS Code Web does
 	// not handle character encoding properly, so it's safe to assume
 	// other applications might not as well.
-	r.Route("/%40{user}/{workspacename}/apps/{workspaceapp}", apps)
-	r.Route("/@{user}/{workspacename}/apps/{workspaceapp}", apps)
+	r.Route("/%40{user}/{workspacename_and_agent}/apps/{workspaceapp}", apps)
+	r.Route("/@{user}/{workspacename_and_agent}/apps/{workspaceapp}", apps)
 
 	r.Route("/api/v2", func(r chi.Router) {
 		r.NotFound(func(rw http.ResponseWriter, r *http.Request) {

diff --git a/coderd/httpmw/chiparams.go b/coderd/httpmw/chiparams.go
diff --git a/coderd/httpmw/workspaceparam.go b/coderd/httpmw/workspaceparam.go
@@ -4,11 +4,15 @@ import (
 	"context"
 	"database/sql"
 	"errors"
+	"fmt"
 	"net/http"
+	"strings"
 
 	"github.com/coder/coder/coderd/database"
 	"github.com/coder/coder/coderd/httpapi"
 	"github.com/coder/coder/codersdk"
+	"github.com/go-chi/chi/v5"
+	"github.com/google/uuid"
 )
 
 type workspaceParamContextKey struct{}
@@ -48,3 +52,106 @@ func ExtractWorkspaceParam(db database.Store) func(http.Handler) http.Handler {
 		})
 	}
 }
+
+// ExtractWorkspaceAndAgentParam grabs a workspace and an agent from the
+// "workspacename_and_agent" URL parameter. `ExtractUserParam` must be called
+// before this.
+// This can be in the form of:
+//	- "<workspace-name>.[workspace-agent]"	: If multiple agents exist
+//	- "<workspace-name>"					: If one agent exists
+func ExtractWorkspaceAndAgentParam(db database.Store) func(http.Handler) http.Handler {
+	return func(next http.Handler) http.Handler {
+		return http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) {
+			user := UserParam(r)
+			workspaceWithAgent := chi.URLParam(r, "workspacename_and_agent")
+			workspaceParts := strings.Split(workspaceWithAgent, ".")
+
+			workspace, err := db.GetWorkspaceByOwnerIDAndName(r.Context(), database.GetWorkspaceByOwnerIDAndNameParams{
+				OwnerID: user.ID,
+				Name:    workspaceParts[0],
+			})
+			if err != nil {
+				if errors.Is(err, sql.ErrNoRows) {
+					httpapi.ResourceNotFound(rw)
+					return
+				}
+				httpapi.Write(rw, http.StatusInternalServerError, codersdk.Response{
+					Message: "Internal error fetching workspace.",
+					Detail:  err.Error(),
+				})
+				return
+			}
+
+			build, err := db.GetLatestWorkspaceBuildByWorkspaceID(r.Context(), workspace.ID)
+			if err != nil {
+				httpapi.Write(rw, http.StatusInternalServerError, codersdk.Response{
+					Message: "Internal error fetching workspace build.",
+					Detail:  err.Error(),
+				})
+				return
+			}
+
+			resources, err := db.GetWorkspaceResourcesByJobID(r.Context(), build.JobID)
+			if err != nil {
+				httpapi.Write(rw, http.StatusInternalServerError, codersdk.Response{
+					Message: "Internal error fetching workspace resources.",
+					Detail:  err.Error(),
+				})
+				return
+			}
+			resourceIDs := make([]uuid.UUID, 0)
+			for _, resource := range resources {
+				resourceIDs = append(resourceIDs, resource.ID)
+			}
+
+			agents, err := db.GetWorkspaceAgentsByResourceIDs(r.Context(), resourceIDs)
+			if err != nil {
+				httpapi.Write(rw, http.StatusInternalServerError, codersdk.Response{
+					Message: "Internal error fetching workspace agents.",
+					Detail:  err.Error(),
+				})
+				return
+			}
+
+			if len(agents) == 0 {
+				httpapi.Write(rw, http.StatusBadRequest, codersdk.Response{
+					Message: "No agents exist for this workspace",
+				})
+				return
+			}
+
+			// If we have more than 1 workspace agent, we need to specify which one to use.
+			if len(agents) > 1 && len(workspaceParts) <= 1 {
+				httpapi.Write(rw, http.StatusBadRequest, codersdk.Response{
+					Message: "More than one agent exists, but no agent specified.",
+				})
+				return
+			}
+
+			// If we have more than 1 workspace agent, we need to specify which one to use.
+			var agent database.WorkspaceAgent
+			var found bool
+			if len(agents) > 1 {
+				for _, otherAgent := range agents {
+					if otherAgent.Name == workspaceParts[1] {
+						agent = otherAgent
+						found = true
+						break
+					}
+				}
+				if !found {
+					httpapi.Write(rw, http.StatusBadRequest, codersdk.Response{
+						Message: fmt.Sprintf("No agent exists with the name %s", workspaceParts[1]),
+					})
+					return
+				}
+			} else {
+				agent = agents[0]
+			}
+
+			ctx := context.WithValue(r.Context(), workspaceParamContextKey{}, workspace)
+			ctx = context.WithValue(r.Context(), workspaceAgentContextKey{}, agent)
+			next.ServeHTTP(rw, r.WithContext(ctx))
+		})
+	}
+}
diff --git a/coderd/workspaceapps.go b/coderd/workspaceapps.go
@@ -10,7 +10,6 @@ import (
 	"strings"
 
 	"github.com/go-chi/chi/v5"
-	"github.com/google/uuid"
 
 	"github.com/coder/coder/coderd/database"
 	"github.com/coder/coder/coderd/httpapi"
@@ -23,95 +22,14 @@ import (
 // workspaceAppsProxyPath proxies requests to a workspace application
 // through a relative URL path.
 func (api *API) workspaceAppsProxyPath(rw http.ResponseWriter, r *http.Request) {
-	user := httpmw.UserParam(r)
-	// This can be in the form of: "<workspace-name>.[workspace-agent]" or "<workspace-name>"
-	workspaceWithAgent := chi.URLParam(r, "workspacename")
-	workspaceParts := strings.Split(workspaceWithAgent, ".")
-
-	workspace, err := api.Database.GetWorkspaceByOwnerIDAndName(r.Context(), database.GetWorkspaceByOwnerIDAndNameParams{
-		OwnerID: user.ID,
-		Name:    workspaceParts[0],
-	})
-	if errors.Is(err, sql.ErrNoRows) {
-		httpapi.ResourceNotFound(rw)
-		return
-	}
-	if err != nil {
-		httpapi.Write(rw, http.StatusInternalServerError, codersdk.Response{
-			Message: "Internal error fetching workspace.",
-			Detail:  err.Error(),
-		})
-		return
-	}
+	workspace := httpmw.WorkspaceParam(r)
+	agent := httpmw.WorkspaceAgentParam(r)
 
 	if !api.Authorize(r, rbac.ActionCreate, workspace.ExecutionRBAC()) {
 		httpapi.ResourceNotFound(rw)
 		return
 	}
 
-	build, err := api.Database.GetLatestWorkspaceBuildByWorkspaceID(r.Context(), workspace.ID)
-	if err != nil {
-		httpapi.Write(rw, http.StatusInternalServerError, codersdk.Response{
-			Message: "Internal error fetching workspace build.",
-			Detail:  err.Error(),
-		})
-		return
-	}
-
-	resources, err := api.Database.GetWorkspaceResourcesByJobID(r.Context(), build.JobID)
-	if err != nil {
-		httpapi.Write(rw, http.StatusInternalServerError, codersdk.Response{
-			Message: "Internal error fetching workspace resources.",
-			Detail:  err.Error(),
-		})
-		return
-	}
-	resourceIDs := make([]uuid.UUID, 0)
-	for _, resource := range resources {
-		resourceIDs = append(resourceIDs, resource.ID)
-	}
-	agents, err := api.Database.GetWorkspaceAgentsByResourceIDs(r.Context(), resourceIDs)
-	if err != nil {
-		httpapi.Write(rw, http.StatusInternalServerError, codersdk.Response{
-			Message: "Internal error fetching workspace agents.",
-			Detail:  err.Error(),
-		})
-		return
-	}
-	if len(agents) == 0 {
-		httpapi.Write(rw, http.StatusBadRequest, codersdk.Response{
-			Message: "No agents exist.",
-		})
-		return
-	}
-
-	// If we have more than 1 workspace agent, we need to specify which one to use.
-	if len(agents) > 1 && len(workspaceParts) <= 1 {
-		httpapi.Write(rw, http.StatusBadRequest, codersdk.Response{
-			Message: "More than one agent exists, but no agent specified.",
-		})
-		return
-	}
-
-	// If we have more than 1 workspace agent, we need to specify which one to use.
-	var agent *database.WorkspaceAgent
-	if len(agents) > 1 {
-		for _, otherAgent := range agents {
-			if otherAgent.Name == workspaceParts[1] {
-				agent = &otherAgent
-				break
-			}
-		}
-		if agent == nil {
-			httpapi.Write(rw, http.StatusBadRequest, codersdk.Response{
-				Message: fmt.Sprintf("No agent exists with the name %s", workspaceParts[1]),
-			})
-			return
-		}
-	} else {
-		agent = &agents[0]
-	}
-
 	app, err := api.Database.GetWorkspaceAppByAgentIDAndName(r.Context(), database.GetWorkspaceAppByAgentIDAndNameParams{
 		AgentID: agent.ID,
 		Name:    chi.URLParam(r, "workspaceapp"),