package httpmw_test import ( "context" "crypto/sha256" "encoding/json" "fmt" "net/http" "net/http/httptest" "testing" "time" "github.com/go-chi/chi/v5" "github.com/google/uuid" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "github.com/coder/coder/v2/coderd/database" "github.com/coder/coder/v2/coderd/database/dbgen" "github.com/coder/coder/v2/coderd/database/dbmem" "github.com/coder/coder/v2/coderd/database/dbtime" "github.com/coder/coder/v2/coderd/httpmw" "github.com/coder/coder/v2/codersdk" "github.com/coder/coder/v2/cryptorand" ) func TestWorkspaceParam(t *testing.T) { t.Parallel() setup := func(db database.Store) (*http.Request, database.User) { var ( id, secret = randomAPIKeyParts() hashed = sha256.Sum256([]byte(secret)) ) r := httptest.NewRequest("GET", "/", nil) r.Header.Set(codersdk.SessionTokenHeader, 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: dbtime.Now(), UpdatedAt: dbtime.Now(), LoginType: database.LoginTypePassword, }) require.NoError(t, err) user, err = db.UpdateUserStatus(context.Background(), database.UpdateUserStatusParams{ ID: user.ID, Status: database.UserStatusActive, UpdatedAt: dbtime.Now(), }) require.NoError(t, err) _, err = db.InsertAPIKey(r.Context(), database.InsertAPIKeyParams{ ID: id, UserID: user.ID, HashedSecret: hashed[:], LastUsed: dbtime.Now(), ExpiresAt: dbtime.Now().Add(time.Minute), LoginType: database.LoginTypePassword, Scope: database.APIKeyScopeAll, }) require.NoError(t, err) ctx := chi.NewRouteContext() ctx.URLParams.Add("user", "me") r = r.WithContext(context.WithValue(r.Context(), chi.RouteCtxKey, ctx)) return r, user } t.Run("None", func(t *testing.T) { t.Parallel() db := dbmem.New() rtr := chi.NewRouter() rtr.Use(httpmw.ExtractWorkspaceParam(db)) rtr.Get("/", nil) r, _ := setup(db) rw := httptest.NewRecorder() rtr.ServeHTTP(rw, r) res := rw.Result() defer res.Body.Close() require.Equal(t, http.StatusBadRequest, res.StatusCode) }) t.Run("NotFound", func(t *testing.T) { t.Parallel() db := dbmem.New() rtr := chi.NewRouter() rtr.Use(httpmw.ExtractWorkspaceParam(db)) rtr.Get("/", nil) r, _ := setup(db) chi.RouteContext(r.Context()).URLParams.Add("workspace", uuid.NewString()) rw := httptest.NewRecorder() rtr.ServeHTTP(rw, r) res := rw.Result() defer res.Body.Close() require.Equal(t, http.StatusNotFound, res.StatusCode) }) t.Run("Found", func(t *testing.T) { t.Parallel() db := dbmem.New() rtr := chi.NewRouter() rtr.Use( httpmw.ExtractAPIKeyMW(httpmw.ExtractAPIKeyConfig{ DB: db, RedirectToLogin: false, }), httpmw.ExtractWorkspaceParam(db), ) rtr.Get("/", func(rw http.ResponseWriter, r *http.Request) { _ = httpmw.WorkspaceParam(r) rw.WriteHeader(http.StatusOK) }) r, user := setup(db) workspace, err := db.InsertWorkspace(context.Background(), database.InsertWorkspaceParams{ ID: uuid.New(), OwnerID: user.ID, Name: "hello", AutomaticUpdates: database.AutomaticUpdatesNever, }) require.NoError(t, err) chi.RouteContext(r.Context()).URLParams.Add("workspace", workspace.ID.String()) rw := httptest.NewRecorder() rtr.ServeHTTP(rw, r) res := rw.Result() defer res.Body.Close() 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.ExtractAPIKeyMW(httpmw.ExtractAPIKeyConfig{ DB: db, RedirectToLogin: 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 := dbmem.New() var ( user = dbgen.User(t, db, database.User{}) _, token = dbgen.APIKey(t, db, database.APIKey{ UserID: user.ID, }) workspace = dbgen.Workspace(t, db, database.Workspace{ OwnerID: user.ID, Name: cfg.WorkspaceName, }) build = dbgen.WorkspaceBuild(t, db, database.WorkspaceBuild{ WorkspaceID: workspace.ID, Transition: database.WorkspaceTransitionStart, Reason: database.BuildReasonInitiator, }) job = dbgen.ProvisionerJob(t, db, nil, database.ProvisionerJob{ ID: build.JobID, Type: database.ProvisionerJobTypeWorkspaceBuild, Provisioner: database.ProvisionerTypeEcho, StorageMethod: database.ProvisionerStorageMethodFile, }) ) r := httptest.NewRequest("GET", "/", nil) r.Header.Set(codersdk.SessionTokenHeader, token) for resourceName, agentNames := range cfg.Agents { resource := dbgen.WorkspaceResource(t, db, database.WorkspaceResource{ JobID: job.ID, Name: resourceName, Transition: database.WorkspaceTransitionStart, }) for _, name := range agentNames { _ = dbgen.WorkspaceAgent(t, db, database.WorkspaceAgent{ ResourceID: resource.ID, Name: name, }) } } ctx := chi.NewRouteContext() ctx.URLParams.Add("user", codersdk.Me) r = r.WithContext(context.WithValue(r.Context(), chi.RouteCtxKey, ctx)) return db, r }