coder · Emyrk · Dec 12, 2023 · Dec 11, 2023 · Dec 11, 2023 · Dec 12, 2023
diff --git a/coderd/coderdtest/authorize.go b/coderd/coderdtest/authorize.go
@@ -7,6 +7,7 @@ import (
 	"runtime"
 	"strings"
 	"sync"
+	"sync/atomic"
 	"testing"
 
 	"github.com/google/uuid"
@@ -16,6 +17,7 @@ import (
 	"golang.org/x/xerrors"
 
 	"github.com/coder/coder/v2/coderd"
+	"github.com/coder/coder/v2/coderd/database"
 	"github.com/coder/coder/v2/coderd/database/dbauthz"
 	"github.com/coder/coder/v2/coderd/rbac"
 	"github.com/coder/coder/v2/coderd/rbac/regosql"
@@ -451,3 +453,22 @@ func must[T any](value T, err error) T {
 	}
 	return value
 }
+
+type FakeAccessControlStore struct{}
+
+func (FakeAccessControlStore) GetTemplateAccessControl(t database.Template) dbauthz.TemplateAccessControl {
+	return dbauthz.TemplateAccessControl{
+		RequireActiveVersion: t.RequireActiveVersion,
+	}
+}
+
+func (FakeAccessControlStore) SetTemplateAccessControl(context.Context, database.Store, uuid.UUID, dbauthz.TemplateAccessControl) error {
+	panic("not implemented")
+}
+
+func AccessControlStorePointer() *atomic.Pointer[dbauthz.AccessControlStore] {
+	acs := &atomic.Pointer[dbauthz.AccessControlStore]{}
+	var tacs dbauthz.AccessControlStore = FakeAccessControlStore{}
+	acs.Store(&tacs)
+	return acs
+}
diff --git a/coderd/database/dbauthz/dbauthz_test.go b/coderd/database/dbauthz/dbauthz_test.go
@@ -13,6 +13,7 @@ import (
 	"golang.org/x/xerrors"
 
 	"cdr.dev/slog"
+
 	"github.com/coder/coder/v2/coderd/coderdtest"
 	"github.com/coder/coder/v2/coderd/database"
 	"github.com/coder/coder/v2/coderd/database/dbauthz"
@@ -62,7 +63,7 @@ func TestAsNoActor(t *testing.T) {
 func TestPing(t *testing.T) {
 	t.Parallel()
 
-	q := dbauthz.New(dbmem.New(), &coderdtest.RecordingAuthorizer{}, slog.Make(), accessControlStorePointer())
+	q := dbauthz.New(dbmem.New(), &coderdtest.RecordingAuthorizer{}, slog.Make(), coderdtest.AccessControlStorePointer())
 	_, err := q.Ping(context.Background())
 	require.NoError(t, err, "must not error")
 }
@@ -74,7 +75,7 @@ func TestInTX(t *testing.T) {
 	db := dbmem.New()
 	q := dbauthz.New(db, &coderdtest.RecordingAuthorizer{
 		Wrapped: &coderdtest.FakeAuthorizer{AlwaysReturn: xerrors.New("custom error")},
-	}, slog.Make(), accessControlStorePointer())
+	}, slog.Make(), coderdtest.AccessControlStorePointer())
 	actor := rbac.Subject{
 		ID:     uuid.NewString(),
 		Roles:  rbac.RoleNames{rbac.RoleOwner()},
@@ -110,8 +111,8 @@ func TestNew(t *testing.T) {
 
 	// Double wrap should not cause an actual double wrap. So only 1 rbac call
 	// should be made.
-	az := dbauthz.New(db, rec, slog.Make(), accessControlStorePointer())
-	az = dbauthz.New(az, rec, slog.Make(), accessControlStorePointer())
+	az := dbauthz.New(db, rec, slog.Make(), coderdtest.AccessControlStorePointer())
+	az = dbauthz.New(az, rec, slog.Make(), coderdtest.AccessControlStorePointer())
 
 	w, err := az.GetWorkspaceByID(ctx, exp.ID)
 	require.NoError(t, err, "must not error")
@@ -128,7 +129,7 @@ func TestDBAuthzRecursive(t *testing.T) {
 	t.Parallel()
 	q := dbauthz.New(dbmem.New(), &coderdtest.RecordingAuthorizer{
 		Wrapped: &coderdtest.FakeAuthorizer{AlwaysReturn: nil},
-	}, slog.Make(), accessControlStorePointer())
+	}, slog.Make(), coderdtest.AccessControlStorePointer())
 	actor := rbac.Subject{
 		ID:     uuid.NewString(),
 		Roles:  rbac.RoleNames{rbac.RoleOwner()},

diff --git a/coderd/database/dbauthz/setup_test.go b/coderd/database/dbauthz/setup_test.go
@@ -6,7 +6,6 @@ import (
 	"reflect"
 	"sort"
 	"strings"
-	"sync/atomic"
 	"testing"
 
 	"github.com/golang/mock/gomock"
@@ -17,6 +16,7 @@ import (
 	"golang.org/x/xerrors"
 
 	"cdr.dev/slog"
+
 	"github.com/coder/coder/v2/coderd/coderdtest"
 	"github.com/coder/coder/v2/coderd/database"
 	"github.com/coder/coder/v2/coderd/database/dbauthz"
@@ -60,7 +60,7 @@ func (s *MethodTestSuite) SetupSuite() {
 	mockStore := dbmock.NewMockStore(ctrl)
 	// We intentionally set no expectations apart from this.
 	mockStore.EXPECT().Wrappers().Return([]string{}).AnyTimes()
-	az := dbauthz.New(mockStore, nil, slog.Make(), accessControlStorePointer())
+	az := dbauthz.New(mockStore, nil, slog.Make(), coderdtest.AccessControlStorePointer())
 	// Take the underlying type of the interface.
 	azt := reflect.TypeOf(az).Elem()
 	s.methodAccounting = make(map[string]int)
@@ -111,7 +111,7 @@ func (s *MethodTestSuite) Subtest(testCaseF func(db database.Store, check *expec
 		rec := &coderdtest.RecordingAuthorizer{
 			Wrapped: fakeAuthorizer,
 		}
-		az := dbauthz.New(db, rec, slog.Make(), accessControlStorePointer())
+		az := dbauthz.New(db, rec, slog.Make(), coderdtest.AccessControlStorePointer())
 		actor := rbac.Subject{
 			ID:     uuid.NewString(),
 			Roles:  rbac.RoleNames{rbac.RoleOwner()},
@@ -399,22 +399,3 @@ func (emptyPreparedAuthorized) Authorize(_ context.Context, _ rbac.Object) error
 func (emptyPreparedAuthorized) CompileToSQL(_ context.Context, _ regosql.ConvertConfig) (string, error) {
 	return "", nil
 }
-
-func accessControlStorePointer() *atomic.Pointer[dbauthz.AccessControlStore] {
-	acs := &atomic.Pointer[dbauthz.AccessControlStore]{}
-	var tacs dbauthz.AccessControlStore = fakeAccessControlStore{}
-	acs.Store(&tacs)
-	return acs
-}
-
-type fakeAccessControlStore struct{}
-
-func (fakeAccessControlStore) GetTemplateAccessControl(t database.Template) dbauthz.TemplateAccessControl {
-	return dbauthz.TemplateAccessControl{
-		RequireActiveVersion: t.RequireActiveVersion,
-	}
-}
-
-func (fakeAccessControlStore) SetTemplateAccessControl(context.Context, database.Store, uuid.UUID, dbauthz.TemplateAccessControl) error {
-	panic("not implemented")
-}
diff --git a/coderd/httpmw/workspaceagentparam.go b/coderd/httpmw/workspaceagentparam.go
@@ -2,8 +2,6 @@ package httpmw
 
 import (
 	"context"
-	"database/sql"
-	"errors"
 	"net/http"
 
 	"github.com/go-chi/chi/v5"
@@ -35,9 +33,9 @@ func ExtractWorkspaceAgentParam(db database.Store) func(http.Handler) http.Handl
 			}
 
 			agent, err := db.GetWorkspaceAgentByID(ctx, agentUUID)
-			if errors.Is(err, sql.ErrNoRows) {
+			if httpapi.Is404Error(err) {
 				httpapi.Write(ctx, rw, http.StatusNotFound, codersdk.Response{
-					Message: "Agent doesn't exist with that id.",
+					Message: "Agent doesn't exist with that id, or you do not have access to it.",
 				})
 				return
 			}

diff --git a/coderd/httpmw/workspaceagentparam_test.go b/coderd/httpmw/workspaceagentparam_test.go
@@ -6,11 +6,15 @@ import (
 	"net/http/httptest"
 	"testing"
 
+	"cdr.dev/slog"
 	"github.com/go-chi/chi/v5"
 	"github.com/google/uuid"
 	"github.com/stretchr/testify/require"
+	"golang.org/x/xerrors"
 
+	"github.com/coder/coder/v2/coderd/coderdtest"
 	"github.com/coder/coder/v2/coderd/database"
+	"github.com/coder/coder/v2/coderd/database/dbauthz"
 	"github.com/coder/coder/v2/coderd/database/dbgen"
 	"github.com/coder/coder/v2/coderd/database/dbmem"
 	"github.com/coder/coder/v2/coderd/httpmw"
@@ -26,8 +30,10 @@ func TestWorkspaceAgentParam(t *testing.T) {
 			_, token = dbgen.APIKey(t, db, database.APIKey{
 				UserID: user.ID,
 			})
+			tpl       = dbgen.Template(t, db, database.Template{})
 			workspace = dbgen.Workspace(t, db, database.Workspace{
-				OwnerID: user.ID,
+				OwnerID:    user.ID,
+				TemplateID: tpl.ID,
 			})
 			build = dbgen.WorkspaceBuild(t, db, database.WorkspaceBuild{
 				WorkspaceID: workspace.ID,
@@ -91,6 +97,36 @@ func TestWorkspaceAgentParam(t *testing.T) {
 		require.Equal(t, http.StatusNotFound, res.StatusCode)
 	})
 
+	t.Run("NotAuthorized", func(t *testing.T) {
+		t.Parallel()
+		db := dbmem.New()
+		fakeAuthz := &coderdtest.FakeAuthorizer{AlwaysReturn: xerrors.Errorf("constant failure")}
+		dbFail := dbauthz.New(db, fakeAuthz, slog.Make(), coderdtest.AccessControlStorePointer())
+
+		rtr := chi.NewRouter()
+		rtr.Use(
+			httpmw.ExtractAPIKeyMW(httpmw.ExtractAPIKeyConfig{
+				DB:              db,
+				RedirectToLogin: false,
+			}),
+			// Only fail authz in this middleware
+			httpmw.ExtractWorkspaceAgentParam(dbFail),
+		)
+		rtr.Get("/", func(rw http.ResponseWriter, r *http.Request) {
+			_ = httpmw.WorkspaceAgentParam(r)
+			rw.WriteHeader(http.StatusOK)
+		})
+
+		r, _ := setupAuthentication(db)
+
+		rw := httptest.NewRecorder()
+		rtr.ServeHTTP(rw, r)
+
+		res := rw.Result()
+		defer res.Body.Close()
+		require.Equal(t, http.StatusNotFound, res.StatusCode)
+	})
+
 	t.Run("WorkspaceAgent", func(t *testing.T) {
 		t.Parallel()
 		db := dbmem.New()

diff --git a/coderd/workspaceagents.go b/coderd/workspaceagents.go
@@ -78,6 +78,10 @@ func (api *API) workspaceAgent(rw http.ResponseWriter, r *http.Request) {
 		return err
 	})
 	err := eg.Wait()
+	if httpapi.Is404Error(err) {
+		httpapi.ResourceNotFound(rw)
+		return
+	}
 	if err != nil {
 		httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
 			Message: "Internal error fetching workspace agent.",