mirror of
https://github.com/coder/coder.git
synced 2025-07-06 15:41:45 +00:00
fix: Fix properly selecting workspace apps by agent (#3684)
This commit is contained in:
Steven Masley
GitHub
@ -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,109 @@ func ExtractWorkspaceParam(db database.Store) func(http.Handler) http.Handler {
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// ExtractWorkspaceAndAgentParam grabs a workspace and an agent from the
|
||||
// "workspace_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, "workspace_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
|
||||
}
|
||||
|
||||
var agent database.WorkspaceAgent
|
||||
var found bool
|
||||
// If we have more than 1 workspace agent, we need to specify which one to use.
|
||||
// If the user specified an agent, we need to make sure that agent
|
||||
// actually exists.
|
||||
if len(workspaceParts) > 1 || 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 %q", workspaceParts[1]),
|
||||
})
|
||||
return
|
||||
}
|
||||
} else {
|
||||
agent = agents[0]
|
||||
}
|
||||
|
||||
ctx := r.Context()
|
||||
ctx = context.WithValue(ctx, workspaceParamContextKey{}, workspace)
|
||||
ctx = context.WithValue(ctx, workspaceAgentParamContextKey{}, agent)
|
||||
next.ServeHTTP(rw, r.WithContext(ctx))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@ -3,12 +3,15 @@ package httpmw_test
|
||||
import (
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
|
||||
"github.com/go-chi/chi/v5"
|
||||
"github.com/google/uuid"
|
||||
"github.com/stretchr/testify/require"
|
||||
@ -122,3 +125,291 @@ func TestWorkspaceParam(t *testing.T) {
|
||||
require.Equal(t, http.StatusOK, res.StatusCode)
|
||||
})
|
||||
}
|
||||
|
||||
func TestWorkspaceAgentByNameParam(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
testCases := []struct {
|
||||
Name string
|
||||
// Agents are mapped to a resource
|
||||
Agents map[string][]string
|
||||
URLParam string
|
||||
WorkspaceName string
|
||||
ExpectedAgent string
|
||||
ExpectedStatusCode int
|
||||
ExpectedError string
|
||||
}{
|
||||
{
|
||||
Name: "NoAgents",
|
||||
WorkspaceName: "dev",
|
||||
Agents: map[string][]string{},
|
||||
URLParam: "dev",
|
||||
ExpectedError: "No agents exist",
|
||||
ExpectedStatusCode: http.StatusBadRequest,
|
||||
},
|
||||
{
|
||||
Name: "NoAgentsSpecify",
|
||||
WorkspaceName: "dev",
|
||||
Agents: map[string][]string{},
|
||||
URLParam: "dev.agent",
|
||||
ExpectedError: "No agents exist",
|
||||
ExpectedStatusCode: http.StatusBadRequest,
|
||||
},
|
||||
{
|
||||
Name: "MultipleAgents",
|
||||
WorkspaceName: "dev",
|
||||
Agents: map[string][]string{
|
||||
"resource-a": {
|
||||
"agent-one",
|
||||
"agent-two",
|
||||
},
|
||||
},
|
||||
URLParam: "dev",
|
||||
ExpectedStatusCode: http.StatusBadRequest,
|
||||
ExpectedError: "More than one agent exists, but no agent specified",
|
||||
},
|
||||
{
|
||||
Name: "MultipleResources",
|
||||
WorkspaceName: "dev",
|
||||
Agents: map[string][]string{
|
||||
"resource-a": {
|
||||
"agent-one",
|
||||
},
|
||||
"resource-b": {
|
||||
"agent-two",
|
||||
},
|
||||
},
|
||||
URLParam: "dev",
|
||||
ExpectedStatusCode: http.StatusBadRequest,
|
||||
ExpectedError: "More than one agent exists, but no agent specified",
|
||||
},
|
||||
{
|
||||
Name: "NotExistsOneAgent",
|
||||
WorkspaceName: "dev",
|
||||
Agents: map[string][]string{
|
||||
"resource-a": {
|
||||
"agent-one",
|
||||
},
|
||||
},
|
||||
URLParam: "dev.not-exists",
|
||||
ExpectedStatusCode: http.StatusBadRequest,
|
||||
ExpectedError: "No agent exists with the name",
|
||||
},
|
||||
{
|
||||
Name: "NotExistsMultipleAgents",
|
||||
WorkspaceName: "dev",
|
||||
Agents: map[string][]string{
|
||||
"resource-a": {
|
||||
"agent-one",
|
||||
},
|
||||
"resource-b": {
|
||||
"agent-two",
|
||||
},
|
||||
"resource-c": {
|
||||
"agent-three",
|
||||
},
|
||||
},
|
||||
URLParam: "dev.not-exists",
|
||||
ExpectedStatusCode: http.StatusBadRequest,
|
||||
ExpectedError: "No agent exists with the name",
|
||||
},
|
||||
|
||||
// OKs
|
||||
{
|
||||
Name: "MultipleResourcesOneAgent",
|
||||
WorkspaceName: "dev",
|
||||
Agents: map[string][]string{
|
||||
"resource-a": {},
|
||||
"resource-b": {
|
||||
"agent-one",
|
||||
},
|
||||
},
|
||||
URLParam: "dev",
|
||||
ExpectedAgent: "agent-one",
|
||||
ExpectedStatusCode: http.StatusOK,
|
||||
},
|
||||
{
|
||||
Name: "OneAgent",
|
||||
WorkspaceName: "dev",
|
||||
Agents: map[string][]string{
|
||||
"resource-a": {
|
||||
"agent-one",
|
||||
},
|
||||
},
|
||||
URLParam: "dev",
|
||||
ExpectedAgent: "agent-one",
|
||||
ExpectedStatusCode: http.StatusOK,
|
||||
},
|
||||
{
|
||||
Name: "OneAgentSelected",
|
||||
WorkspaceName: "dev",
|
||||
Agents: map[string][]string{
|
||||
"resource-a": {
|
||||
"agent-one",
|
||||
},
|
||||
},
|
||||
URLParam: "dev.agent-one",
|
||||
ExpectedAgent: "agent-one",
|
||||
ExpectedStatusCode: http.StatusOK,
|
||||
},
|
||||
{
|
||||
Name: "MultipleAgentSelectOne",
|
||||
WorkspaceName: "dev",
|
||||
Agents: map[string][]string{
|
||||
"resource-a": {
|
||||
"agent-one",
|
||||
"agent-two",
|
||||
"agent-selected",
|
||||
},
|
||||
},
|
||||
URLParam: "dev.agent-selected",
|
||||
ExpectedAgent: "agent-selected",
|
||||
ExpectedStatusCode: http.StatusOK,
|
||||
},
|
||||
{
|
||||
Name: "MultipleResourcesSelectOne",
|
||||
WorkspaceName: "dev",
|
||||
Agents: map[string][]string{
|
||||
"resource-a": {
|
||||
"agent-one",
|
||||
},
|
||||
"resource-b": {
|
||||
"agent-two",
|
||||
},
|
||||
"resource-c": {
|
||||
"agent-selected",
|
||||
"agent-three",
|
||||
},
|
||||
},
|
||||
URLParam: "dev.agent-selected",
|
||||
ExpectedAgent: "agent-selected",
|
||||
ExpectedStatusCode: http.StatusOK,
|
||||
},
|
||||
}
|
||||
|
||||
for _, c := range testCases {
|
||||
c := c
|
||||
t.Run(c.Name, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
db, r := setupWorkspaceWithAgents(t, setupConfig{
|
||||
WorkspaceName: c.WorkspaceName,
|
||||
Agents: c.Agents,
|
||||
})
|
||||
|
||||
chi.RouteContext(r.Context()).URLParams.Add("workspace_and_agent", c.URLParam)
|
||||
|
||||
rtr := chi.NewRouter()
|
||||
rtr.Use(
|
||||
httpmw.ExtractAPIKey(db, nil, true),
|
||||
httpmw.ExtractUserParam(db),
|
||||
httpmw.ExtractWorkspaceAndAgentParam(db),
|
||||
)
|
||||
rtr.Get("/", func(w http.ResponseWriter, r *http.Request) {
|
||||
workspace := httpmw.WorkspaceParam(r)
|
||||
agent := httpmw.WorkspaceAgentParam(r)
|
||||
|
||||
assert.Equal(t, c.ExpectedAgent, agent.Name, "expected agent name")
|
||||
assert.Equal(t, c.WorkspaceName, workspace.Name, "expected workspace name")
|
||||
})
|
||||
|
||||
rw := httptest.NewRecorder()
|
||||
rtr.ServeHTTP(rw, r)
|
||||
res := rw.Result()
|
||||
var coderResp codersdk.Response
|
||||
_ = json.NewDecoder(res.Body).Decode(&coderResp)
|
||||
res.Body.Close()
|
||||
require.Equal(t, c.ExpectedStatusCode, res.StatusCode)
|
||||
if c.ExpectedError != "" {
|
||||
require.Contains(t, coderResp.Message, c.ExpectedError)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
type setupConfig struct {
|
||||
WorkspaceName string
|
||||
// Agents are mapped to a resource
|
||||
Agents map[string][]string
|
||||
}
|
||||
|
||||
func setupWorkspaceWithAgents(t testing.TB, cfg setupConfig) (database.Store, *http.Request) {
|
||||
t.Helper()
|
||||
db := databasefake.New()
|
||||
var (
|
||||
id, secret = randomAPIKeyParts()
|
||||
hashed = sha256.Sum256([]byte(secret))
|
||||
)
|
||||
r := httptest.NewRequest("GET", "/", nil)
|
||||
r.AddCookie(&http.Cookie{
|
||||
Name: codersdk.SessionTokenKey,
|
||||
Value: fmt.Sprintf("%s-%s", id, secret),
|
||||
})
|
||||
|
||||
userID := uuid.New()
|
||||
username, err := cryptorand.String(8)
|
||||
require.NoError(t, err)
|
||||
user, err := db.InsertUser(r.Context(), database.InsertUserParams{
|
||||
ID: userID,
|
||||
Email: "testaccount@coder.com",
|
||||
HashedPassword: hashed[:],
|
||||
Username: username,
|
||||
CreatedAt: database.Now(),
|
||||
UpdatedAt: database.Now(),
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
_, err = db.InsertAPIKey(r.Context(), database.InsertAPIKeyParams{
|
||||
ID: id,
|
||||
UserID: user.ID,
|
||||
HashedSecret: hashed[:],
|
||||
LastUsed: database.Now(),
|
||||
ExpiresAt: database.Now().Add(time.Minute),
|
||||
LoginType: database.LoginTypePassword,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
workspace, err := db.InsertWorkspace(context.Background(), database.InsertWorkspaceParams{
|
||||
ID: uuid.New(),
|
||||
TemplateID: uuid.New(),
|
||||
OwnerID: user.ID,
|
||||
Name: cfg.WorkspaceName,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
build, err := db.InsertWorkspaceBuild(context.Background(), database.InsertWorkspaceBuildParams{
|
||||
ID: uuid.New(),
|
||||
WorkspaceID: workspace.ID,
|
||||
JobID: uuid.New(),
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
job, err := db.InsertProvisionerJob(context.Background(), database.InsertProvisionerJobParams{
|
||||
ID: build.JobID,
|
||||
Type: database.ProvisionerJobTypeWorkspaceBuild,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
for resourceName, agentNames := range cfg.Agents {
|
||||
resource, err := db.InsertWorkspaceResource(context.Background(), database.InsertWorkspaceResourceParams{
|
||||
ID: uuid.New(),
|
||||
JobID: job.ID,
|
||||
Name: resourceName,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
for _, name := range agentNames {
|
||||
_, err = db.InsertWorkspaceAgent(context.Background(), database.InsertWorkspaceAgentParams{
|
||||
ID: uuid.New(),
|
||||
ResourceID: resource.ID,
|
||||
Name: name,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
}
|
||||
}
|
||||
|
||||
ctx := chi.NewRouteContext()
|
||||
ctx.URLParams.Add("user", codersdk.Me)
|
||||
r = r.WithContext(context.WithValue(r.Context(), chi.RouteCtxKey, ctx))
|
||||
|
||||
return db, r
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user