fix(Coder plugin): Update integration with Backstage's identity API (#118)

* fix: update UI code to forward bearer tokens properly

* refactor: consolidate init setup logic

* fix: update error catching logic

* fix: add new mock to get current tests passing

* fix: add mock bearer token

* chore: add test middleware to verify bearer token behavior

* refactor: update variable names for clarity

Author: Parkreiner

plugins/backstage-plugin-coder/src/api.ts

@@ -9,6 +9,7 @@ import {
   WorkspaceAgentStatus,
 } from './typesConstants';
 import { CoderAuth, assertValidCoderAuth } from './components/CoderProvider';
+import { IdentityApi } from '@backstage/core-plugin-api';
 
 export const CODER_QUERY_KEY_PREFIX = 'coder-backstage-plugin';
 
@@ -19,9 +20,31 @@ export const ASSETS_ROUTE_PREFIX = PROXY_ROUTE_PREFIX;
 export const CODER_AUTH_HEADER_KEY = 'Coder-Session-Token';
 export const REQUEST_TIMEOUT_MS = 20_000;
 
-function getCoderApiRequestInit(authToken: string): RequestInit {
+async function getCoderApiRequestInit(
+  authToken: string,
+  identity: IdentityApi,
+): Promise<RequestInit> {
+  const headers: HeadersInit = {
+    [CODER_AUTH_HEADER_KEY]: authToken,
+  };
+
+  try {
+    const credentials = await identity.getCredentials();
+    if (credentials.token) {
+      headers.Authorization = `Bearer ${credentials.token}`;
+    }
+  } catch (err) {
+    if (err instanceof Error) {
+      throw err;
+    }
+
+    throw new Error(
+      "Unable to parse user information for Coder requests. Please ensure that your Backstage deployment is integrated to use Backstage's Identity API",
+    );
+  }
+
   return {
-    headers: { [CODER_AUTH_HEADER_KEY]: authToken },
+    headers,
     signal: AbortSignal.timeout(REQUEST_TIMEOUT_MS),
   };
 }
@@ -53,6 +76,7 @@ export class BackstageHttpError extends Error {
 type FetchInputs = Readonly<{
   auth: CoderAuth;
   baseUrl: string;
+  identity: IdentityApi;
 }>;
 
 type WorkspacesFetchInputs = Readonly<
@@ -64,17 +88,18 @@ type WorkspacesFetchInputs = Readonly<
 async function getWorkspaces(
   fetchInputs: WorkspacesFetchInputs,
 ): Promise<readonly Workspace[]> {
-  const { baseUrl, coderQuery, auth } = fetchInputs;
+  const { baseUrl, coderQuery, auth, identity } = fetchInputs;
   assertValidCoderAuth(auth);
 
   const urlParams = new URLSearchParams({
     q: coderQuery,
     limit: '0',
   });
 
+  const requestInit = await getCoderApiRequestInit(auth.token, identity);
   const response = await fetch(
     `${baseUrl}${API_ROUTE_PREFIX}/workspaces?${urlParams.toString()}`,
-    getCoderApiRequestInit(auth.token),
+    requestInit,
   );
 
   if (!response.ok) {
@@ -116,12 +141,13 @@ type BuildParamsFetchInputs = Readonly<
 >;
 
 async function getWorkspaceBuildParameters(inputs: BuildParamsFetchInputs) {
-  const { baseUrl, auth, workspaceBuildId } = inputs;
+  const { baseUrl, auth, workspaceBuildId, identity } = inputs;
   assertValidCoderAuth(auth);
 
+  const requestInit = await getCoderApiRequestInit(auth.token, identity);
   const res = await fetch(
     `${baseUrl}${API_ROUTE_PREFIX}/workspacebuilds/${workspaceBuildId}/parameters`,
-    getCoderApiRequestInit(auth.token),
+    requestInit,
   );
 
   if (!res.ok) {
@@ -256,16 +282,18 @@ export function workspacesByRepo(
 type AuthValidationInputs = Readonly<{
   baseUrl: string;
   authToken: string;
+  identity: IdentityApi;
 }>;
 
 async function isAuthValid(inputs: AuthValidationInputs): Promise<boolean> {
-  const { baseUrl, authToken } = inputs;
+  const { baseUrl, authToken, identity } = inputs;
 
   // In this case, the request doesn't actually matter. Just need to make any
   // kind of dummy request to validate the auth
+  const requestInit = await getCoderApiRequestInit(authToken, identity);
   const response = await fetch(
     `${baseUrl}${API_ROUTE_PREFIX}/users/me`,
-    getCoderApiRequestInit(authToken),
+    requestInit,
   );
 
   if (response.status >= 400 && response.status !== 401) {

plugins/backstage-plugin-coder/src/components/CoderProvider/CoderAuthProvider.tsx

@@ -19,6 +19,7 @@ import {
   authValidation,
 } from '../../api';
 import { useBackstageEndpoints } from '../../hooks/useBackstageEndpoints';
+import { identityApiRef, useApi } from '@backstage/core-plugin-api';
 
 const TOKEN_STORAGE_KEY = 'coder-backstage-plugin/token';
 
@@ -98,6 +99,7 @@ export function useCoderAuth(): CoderAuth {
 type CoderAuthProviderProps = Readonly<PropsWithChildren<unknown>>;
 
 export const CoderAuthProvider = ({ children }: CoderAuthProviderProps) => {
+  const identity = useApi(identityApiRef);
   const { baseUrl } = useBackstageEndpoints();
   const [isInsideGracePeriod, setIsInsideGracePeriod] = useState(true);
 
@@ -108,7 +110,7 @@ export const CoderAuthProvider = ({ children }: CoderAuthProviderProps) => {
   const [readonlyInitialAuthToken] = useState(authToken);
 
   const authValidityQuery = useQuery({
-    ...authValidation({ baseUrl, authToken }),
+    ...authValidation({ baseUrl, authToken, identity }),
     refetchOnWindowFocus: query => query.state.data !== false,
   });
 

plugins/backstage-plugin-coder/src/components/CoderProvider/CoderProvider.test.tsx

@@ -3,7 +3,11 @@ import { renderHook } from '@testing-library/react';
 import { act, waitFor } from '@testing-library/react';
 
 import { TestApiProvider, wrapInTestApp } from '@backstage/test-utils';
-import { configApiRef, errorApiRef } from '@backstage/core-plugin-api';
+import {
+  configApiRef,
+  errorApiRef,
+  identityApiRef,
+} from '@backstage/core-plugin-api';
 
 import { CoderProvider } from './CoderProvider';
 import { useCoderAppConfig } from './CoderAppConfigProvider';
@@ -12,6 +16,7 @@ import { type CoderAuth, useCoderAuth } from './CoderAuthProvider';
 import {
   getMockConfigApi,
   getMockErrorApi,
+  getMockIdentityApi,
   mockAppConfig,
   mockCoderAuthToken,
 } from '../../testHelpers/mockBackstageData';
@@ -87,6 +92,7 @@ describe(`${CoderProvider.name}`, () => {
           <TestApiProvider
             apis={[
               [errorApiRef, getMockErrorApi()],
+              [identityApiRef, getMockIdentityApi()],
               [configApiRef, getMockConfigApi()],
             ]}
           >

plugins/backstage-plugin-coder/src/hooks/useCoderWorkspacesQuery.ts

@@ -4,6 +4,7 @@ import { workspaces, workspacesByRepo } from '../api';
 import { useCoderAuth } from '../components/CoderProvider/CoderAuthProvider';
 import { useBackstageEndpoints } from './useBackstageEndpoints';
 import { CoderWorkspacesConfig } from './useCoderWorkspacesConfig';
+import { identityApiRef, useApi } from '@backstage/core-plugin-api';
 
 type QueryInput = Readonly<{
   coderQuery: string;
@@ -15,12 +16,19 @@ export function useCoderWorkspacesQuery({
   workspacesConfig,
 }: QueryInput) {
   const auth = useCoderAuth();
+  const identity = useApi(identityApiRef);
   const { baseUrl } = useBackstageEndpoints();
   const hasRepoData = workspacesConfig && workspacesConfig.repoUrl;
 
   const queryOptions = hasRepoData
-    ? workspacesByRepo({ coderQuery, auth, baseUrl, workspacesConfig })
-    : workspaces({ coderQuery, auth, baseUrl });
+    ? workspacesByRepo({
+        coderQuery,
+        identity,
+        auth,
+        baseUrl,
+        workspacesConfig,
+      })
+    : workspaces({ coderQuery, identity, auth, baseUrl });
 
   return useQuery(queryOptions);
 }

plugins/backstage-plugin-coder/src/testHelpers/mockBackstageData.ts

@@ -17,6 +17,7 @@ import {
 import { ScmIntegrationsApi } from '@backstage/integration-react';
 
 import { API_ROUTE_PREFIX, ASSETS_ROUTE_PREFIX } from '../api';
+import { IdentityApi } from '@backstage/core-plugin-api';
 
 /**
  * This is the key that Backstage checks from the entity data to determine the
@@ -57,6 +58,7 @@ export const mockBackstageProxyEndpoint = `${mockBackstageUrlRoot}${API_ROUTE_PR
 
 export const mockBackstageAssetsEndpoint = `${mockBackstageUrlRoot}${ASSETS_ROUTE_PREFIX}`;
 
+export const mockBearerToken = 'This-is-an-opaque-value-by-design';
 export const mockCoderAuthToken = 'ZG0HRy2gGN-mXljc1s5FqtE8WUJ4sUc5X';
 
 export const mockYamlConfig = {
@@ -207,6 +209,33 @@ export function getMockErrorApi() {
   return errorApi;
 }
 
+export function getMockIdentityApi(): IdentityApi {
+  return {
+    signOut: async () => {
+      return void 'Not going to implement this';
+    },
+    getProfileInfo: async () => {
+      return {
+        displayName: 'Dobah',
+        email: 'i-love-my-dog-dobah@dog.ceo',
+        picture: undefined,
+      };
+    },
+    getBackstageIdentity: async () => {
+      return {
+        type: 'user',
+        userEntityRef: 'User:default/Dobah',
+        ownershipEntityRefs: [],
+      };
+    },
+    getCredentials: async () => {
+      return {
+        token: mockBearerToken,
+      };
+    },
+  };
+}
+
 /**
  * Exposes a mock ScmIntegrationRegistry to be used with scmIntegrationsApiRef
  * for mocking out code that relies on source code data.

plugins/backstage-plugin-coder/src/testHelpers/server.ts

@@ -1,5 +1,12 @@
 /* eslint-disable @backstage/no-undeclared-imports -- For test helpers only */
-import { RestHandler, rest } from 'msw';
+import {
+  type DefaultBodyType,
+  type ResponseResolver,
+  type RestContext,
+  type RestHandler,
+  type RestRequest,
+  rest,
+} from 'msw';
 import { setupServer } from 'msw/node';
 /* eslint-enable @backstage/no-undeclared-imports */
 
@@ -8,14 +15,60 @@ import {
   mockWorkspaceBuildParameters,
 } from './mockCoderAppData';
 import {
+  mockBearerToken,
   mockCoderAuthToken,
   mockBackstageProxyEndpoint as root,
 } from './mockBackstageData';
 import type { Workspace, WorkspacesResponse } from '../typesConstants';
 import { CODER_AUTH_HEADER_KEY } from '../api';
 
-const handlers: readonly RestHandler[] = [
-  rest.get(`${root}/workspaces`, (req, res, ctx) => {
+type RestResolver<TBody extends DefaultBodyType = any> = ResponseResolver<
+  RestRequest<TBody>,
+  RestContext,
+  TBody
+>;
+
+export type RestResolverMiddleware<TBody extends DefaultBodyType = any> = (
+  resolver: RestResolver<TBody>,
+) => RestResolver<TBody>;
+
+const defaultMiddleware = [
+  function validateBearerToken(handler) {
+    return (req, res, ctx) => {
+      const tokenRe = /^Bearer (.+)$/;
+      const authHeader = req.headers.get('Authorization') ?? '';
+      const [, bearerToken] = tokenRe.exec(authHeader) ?? [];
+
+      if (bearerToken === mockBearerToken) {
+        return handler(req, res, ctx);
+      }
+
+      return res(ctx.status(401));
+    };
+  },
+] as const satisfies readonly RestResolverMiddleware[];
+
+export function wrapInDefaultMiddleware<TBody extends DefaultBodyType = any>(
+  resolver: RestResolver<TBody>,
+): RestResolver<TBody> {
+  return defaultMiddleware.reduceRight((currentResolver, middleware) => {
+    const recastMiddleware =
+      middleware as unknown as RestResolverMiddleware<TBody>;
+
+    return recastMiddleware(currentResolver);
+  }, resolver);
+}
+
+function wrappedGet<TBody extends DefaultBodyType = any>(
+  path: string,
+  resolver: RestResolver<TBody>,
+): RestHandler {
+  const wrapped = wrapInDefaultMiddleware(resolver);
+  return rest.get(path, wrapped);
+}
+
+const mainTestHandlers: readonly RestHandler[] = [
+  wrappedGet(`${root}/workspaces`, (req, res, ctx) => {
     const queryText = String(req.url.searchParams.get('q'));
 
     let returnedWorkspaces: Workspace[];
@@ -36,7 +89,7 @@ const handlers: readonly RestHandler[] = [
     );
   }),
 
-  rest.get(
+  wrappedGet(
     `${root}/workspacebuilds/:workspaceBuildId/parameters`,
     (req, res, ctx) => {
       const buildId = String(req.params.workspaceBuildId);
@@ -51,7 +104,7 @@ const handlers: readonly RestHandler[] = [
   ),
 
   // This is the dummy request used to verify a user's auth status
-  rest.get(`${root}/users/me`, (req, res, ctx) => {
+  wrappedGet(`${root}/users/me`, (req, res, ctx) => {
     const token = req.headers.get(CODER_AUTH_HEADER_KEY);
     if (token === mockCoderAuthToken) {
       return res(ctx.status(200));
@@ -61,4 +114,4 @@ const handlers: readonly RestHandler[] = [
   }),
 ];
 
-export const server = setupServer(...handlers);
+export const server = setupServer(...mainTestHandlers);

plugins/backstage-plugin-coder/src/testHelpers/setup.tsx

@@ -12,7 +12,11 @@ import {
 import React from 'react';
 import { QueryClient, QueryClientProvider } from '@tanstack/react-query';
 import { scmIntegrationsApiRef } from '@backstage/integration-react';
-import { configApiRef, errorApiRef } from '@backstage/core-plugin-api';
+import {
+  configApiRef,
+  errorApiRef,
+  identityApiRef,
+} from '@backstage/core-plugin-api';
 import { EntityProvider } from '@backstage/plugin-catalog-react';
 import {
   type CoderAuth,
@@ -30,6 +34,7 @@ import {
   getMockConfigApi,
   mockAuthStates,
   BackstageEntity,
+  getMockIdentityApi,
 } from './mockBackstageData';
 import { CoderErrorBoundary } from '../plugin';
 
@@ -159,6 +164,7 @@ export const renderHookAsCoderEntity = async <
   const mockErrorApi = getMockErrorApi();
   const mockSourceControl = getMockSourceControl();
   const mockConfigApi = getMockConfigApi();
+  const mockIdentityApi = getMockIdentityApi();
   const mockQueryClient = getMockQueryClient();
 
   const renderHookValue = renderHook(hook, {
@@ -168,6 +174,7 @@ export const renderHookAsCoderEntity = async <
         <TestApiProvider
           apis={[
             [errorApiRef, mockErrorApi],
+            [identityApiRef, mockIdentityApi],
             [scmIntegrationsApiRef, mockSourceControl],
             [configApiRef, mockConfigApi],
           ]}
@@ -215,11 +222,13 @@ export async function renderInCoderEnvironment({
   const mockErrorApi = getMockErrorApi();
   const mockSourceControl = getMockSourceControl();
   const mockConfigApi = getMockConfigApi();
+  const mockIdentityApi = getMockIdentityApi();
 
   const mainMarkup = (
     <TestApiProvider
       apis={[
         [errorApiRef, mockErrorApi],
+        [identityApiRef, mockIdentityApi],
         [scmIntegrationsApiRef, mockSourceControl],
         [configApiRef, mockConfigApi],
       ]}