Augment session wrapper SSR

front/lib/iam/provider.ts

@@ -33,7 +33,8 @@ function isAuth0Session(session: unknown): session is Session {
 
 // We only expose generic types to ease phasing out.
 
-export type SessionWithUser = Omit<Session, "user"> & { user: ExternalUser };
+// Overrides the Auth0 type definition entirely, to only expose what we need.
+export type SessionWithUser = { user: ExternalUser };
 
 export function isValidSession(
   session: Session | null

front/lib/iam/session.ts

@@ -7,7 +7,7 @@ import type {
 import type { ParsedUrlQuery } from "querystring";
 import { Op } from "sequelize";
 
-import { getSession } from "@app/lib/auth";
+import { Authenticator, getSession } from "@app/lib/auth";
 import type { SessionWithUser } from "@app/lib/iam/provider";
 import { isValidSession } from "@app/lib/iam/provider";
 import {
@@ -83,39 +83,83 @@ export async function getUserFromSession(
   };
 }
 
+export type UserPrivilege = "none" | "user" | "superuser";
+
 interface MakeGetServerSidePropsRequirementsWrapperOptions<
-  R extends boolean = true
+  R extends UserPrivilege = "user"
 > {
   enableLogging?: boolean;
-  requireAuth: R;
+  requireUserPrivilege: R;
 }
 
 export type CustomGetServerSideProps<
   Props extends { [key: string]: any } = { [key: string]: any },
   Params extends ParsedUrlQuery = ParsedUrlQuery,
   Preview extends PreviewData = PreviewData,
-  RequireAuth extends boolean = true
+  RequireUserPrivilege extends UserPrivilege = "user"
 > = (
   context: GetServerSidePropsContext<Params, Preview>,
-  session: RequireAuth extends true ? SessionWithUser : null
+  auth: RequireUserPrivilege extends "none" ? null : Authenticator,
+  session: RequireUserPrivilege extends "none" ? null : SessionWithUser
 ) => Promise<GetServerSidePropsResult<Props>>;
 
+async function getAuthenticator(
+  context: GetServerSidePropsContext<ParsedUrlQuery, PreviewData>,
+  session: SessionWithUser | null,
+  requireUserPrivilege: UserPrivilege
+) {
+  if (!session) {
+    return null;
+  }
+
+  const { wId } = context.params ?? {};
+  const workspaceId = typeof wId === "string" ? wId : null;
+
+  switch (requireUserPrivilege) {
+    case "user":
+      return workspaceId
+        ? Authenticator.fromSession(session, workspaceId)
+        : null;
+
+    case "superuser":
+      return Authenticator.fromSuperUserSession(session, workspaceId);
+
+    default:
+      return null;
+  }
+}
+
 export function makeGetServerSidePropsRequirementsWrapper<
-  RequireAuth extends boolean = true
+  RequireUserPrivilege extends UserPrivilege = "user"
 >({
   enableLogging = true,
-  requireAuth,
-}: MakeGetServerSidePropsRequirementsWrapperOptions<RequireAuth>) {
+  requireUserPrivilege,
+}: MakeGetServerSidePropsRequirementsWrapperOptions<RequireUserPrivilege>) {
   return <T extends { [key: string]: any } = { [key: string]: any }>(
-    getServerSideProps: CustomGetServerSideProps<T, any, any, RequireAuth>
+    getServerSideProps: CustomGetServerSideProps<
+      T,
+      any,
+      any,
+      RequireUserPrivilege
+    >
   ) => {
     return async (
       context: GetServerSidePropsContext<ParsedUrlQuery, PreviewData>
     ) => {
-      const session = requireAuth
-        ? await getSession(context.req, context.res)
-        : null;
-      if (requireAuth && (!session || !isValidSession(session))) {
+      const session =
+        requireUserPrivilege !== "none"
+          ? await getSession(context.req, context.res)
+          : null;
+      const auth = await getAuthenticator(
+        context,
+        session,
+        requireUserPrivilege
+      );
+
+      if (
+        requireUserPrivilege !== "none" &&
+        (!session || !isValidSession(session))
+      ) {
         return {
           redirect: {
             permanent: false,
@@ -125,21 +169,30 @@ export function makeGetServerSidePropsRequirementsWrapper<
         };
       }
 
-      const userSession = session as RequireAuth extends true
-        ? SessionWithUser
-        : null;
+      const userSession = session as RequireUserPrivilege extends "none"
+        ? null
+        : SessionWithUser;
+      const userAuth = auth as RequireUserPrivilege extends "none"
+        ? null
+        : Authenticator;
 
       if (enableLogging) {
         return withGetServerSidePropsLogging(getServerSideProps)(
           context,
+          userAuth,
           userSession
         );
       }
 
-      return getServerSideProps(context, userSession);
+      return getServerSideProps(context, userAuth, userSession);
     };
   };
 }
 
-export const withDefaultGetServerSidePropsRequirements =
-  makeGetServerSidePropsRequirementsWrapper({ requireAuth: true });
+export const withDefaultUserAuthRequirements =
+  makeGetServerSidePropsRequirementsWrapper({ requireUserPrivilege: "user" });
+
+export const withSuperUserAuthRequirements =
+  makeGetServerSidePropsRequirementsWrapper({
+    requireUserPrivilege: "superuser",
+  });

front/logger/withlogging.ts

@@ -6,7 +6,10 @@ import tracer from "dd-trace";
 import StatsD from "hot-shots";
 import type { NextApiRequest, NextApiResponse } from "next";
 
-import type { CustomGetServerSideProps } from "@app/lib/iam/session";
+import type {
+  CustomGetServerSideProps,
+  UserPrivilege,
+} from "@app/lib/iam/session";
 
 import logger from "./logger";
 
@@ -145,11 +148,16 @@ export function apiError<T>(
 
 export function withGetServerSidePropsLogging<
   T extends { [key: string]: any },
-  RequireAuth extends boolean = true
+  RequireUserPrivilege extends UserPrivilege = "user"
 >(
-  getServerSideProps: CustomGetServerSideProps<T, any, any, RequireAuth>
-): CustomGetServerSideProps<T, any, any, RequireAuth> {
-  return async (context, session) => {
+  getServerSideProps: CustomGetServerSideProps<
+    T,
+    any,
+    any,
+    RequireUserPrivilege
+  >
+): CustomGetServerSideProps<T, any, any, RequireUserPrivilege> {
+  return async (context, auth, session) => {
     const now = new Date();
 
     let route = context.resolvedUrl.split("?")[0];
@@ -161,7 +169,7 @@ export function withGetServerSidePropsLogging<
     }
 
     try {
-      const res = await getServerSideProps(context, session);
+      const res = await getServerSideProps(context, auth, session);
 
       const elapsed = new Date().getTime() - now.getTime();
 

front/pages/index.tsx

@@ -55,7 +55,7 @@ import { classNames } from "@app/lib/utils";
 const { GA_TRACKING_ID = "" } = process.env;
 
 export const getServerSideProps = makeGetServerSidePropsRequirementsWrapper({
-  requireAuth: false,
+  requireUserPrivilege: "none",
 })<{
   gaTrackingId: string;
 }>(async (context) => {

front/pages/login-error.tsx

@@ -7,7 +7,7 @@ import { makeGetServerSidePropsRequirementsWrapper } from "@app/lib/iam/session"
 const { GA_TRACKING_ID = "" } = process.env;
 
 export const getServerSideProps = makeGetServerSidePropsRequirementsWrapper({
-  requireAuth: false,
+  requireUserPrivilege: "none",
 })<{
   domain: string | null;
   gaTrackingId: string;

front/pages/no-workspace.tsx

@@ -11,7 +11,7 @@ import { useRouter } from "next/router";
 
 import {
   getUserFromSession,
-  withDefaultGetServerSidePropsRequirements,
+  withDefaultUserAuthRequirements,
 } from "@app/lib/iam/session";
 import { Membership, Workspace, WorkspaceHasDomain } from "@app/lib/models";
 import logger from "@app/logger/logger";
@@ -56,12 +56,12 @@ async function fetchRevokedWorkspace(
   return Workspace.findByPk(revokedWorkspaceId);
 }
 
-export const getServerSideProps = withDefaultGetServerSidePropsRequirements<{
+export const getServerSideProps = withDefaultUserAuthRequirements<{
   status: "auto-join-disabled" | "revoked";
   userFirstName: string;
   workspaceName: string;
   workspaceVerifiedDomain: string | null;
-}>(async (context, session) => {
+}>(async (context, auth, session) => {
   const user = await getUserFromSession(session);
 
   if (!user) {

front/pages/poke/[wId]/assistants/[aId]/index.tsx

@@ -9,17 +9,11 @@ import type { InferGetServerSidePropsType } from "next";
 
 import PokeNavbar from "@app/components/poke/PokeNavbar";
 import { getAgentConfigurations } from "@app/lib/api/assistant/configuration";
-import { Authenticator } from "@app/lib/auth";
-import { withDefaultGetServerSidePropsRequirements } from "@app/lib/iam/session";
+import { withSuperUserAuthRequirements } from "@app/lib/iam/session";
 
-export const getServerSideProps = withDefaultGetServerSidePropsRequirements<{
+export const getServerSideProps = withSuperUserAuthRequirements<{
   agentConfigurations: AgentConfigurationType[];
-}>(async (context, session) => {
-  const auth = await Authenticator.fromSuperUserSession(
-    session,
-    context.params?.wId as string
-  );
-
+}>(async (context, auth) => {
   if (!auth.isDustSuperUser()) {
     return {
       notFound: true,

front/pages/poke/[wId]/data_sources/[name]/index.tsx

@@ -22,17 +22,16 @@ import { useEffect, useState } from "react";
 import { PokePermissionTree } from "@app/components/poke/PokeConnectorPermissionsTree";
 import PokeNavbar from "@app/components/poke/PokeNavbar";
 import { getDataSource } from "@app/lib/api/data_sources";
-import { Authenticator } from "@app/lib/auth";
 import { useSubmitFunction } from "@app/lib/client/utils";
 import { getDisplayNameForDocument } from "@app/lib/data_sources";
-import { withDefaultGetServerSidePropsRequirements } from "@app/lib/iam/session";
+import { withSuperUserAuthRequirements } from "@app/lib/iam/session";
 import { useDocuments } from "@app/lib/swr";
 import { formatTimestampToFriendlyDate, timeAgoFrom } from "@app/lib/utils";
 import logger from "@app/logger/logger";
 
 const { TEMPORAL_CONNECTORS_NAMESPACE = "" } = process.env;
 
-export const getServerSideProps = withDefaultGetServerSidePropsRequirements<{
+export const getServerSideProps = withSuperUserAuthRequirements<{
   owner: WorkspaceType;
   dataSource: DataSourceType;
   coreDataSource: CoreAPIDataSource;
@@ -43,12 +42,7 @@ export const getServerSideProps = withDefaultGetServerSidePropsRequirements<{
     githubCodeSyncEnabled: boolean;
   };
   temporalWorkspace: string;
-}>(async (context, session) => {
-  const auth = await Authenticator.fromSuperUserSession(
-    session,
-    context.params?.wId as string
-  );
-
+}>(async (context, auth) => {
   const owner = auth.workspace();
 
   if (!owner || !auth.isDustSuperUser()) {

front/pages/poke/[wId]/data_sources/[name]/search.tsx

@@ -8,20 +8,14 @@ import { useEffect, useState } from "react";
 
 import PokeNavbar from "@app/components/poke/PokeNavbar";
 import { getDataSource } from "@app/lib/api/data_sources";
-import { Authenticator } from "@app/lib/auth";
 import { getDisplayNameForDocument } from "@app/lib/data_sources";
-import { withDefaultGetServerSidePropsRequirements } from "@app/lib/iam/session";
+import { withSuperUserAuthRequirements } from "@app/lib/iam/session";
 import { classNames, timeAgoFrom } from "@app/lib/utils";
 
-export const getServerSideProps = withDefaultGetServerSidePropsRequirements<{
+export const getServerSideProps = withSuperUserAuthRequirements<{
   owner: WorkspaceType;
   dataSource: DataSourceType;
-}>(async (context, session) => {
-  const auth = await Authenticator.fromSuperUserSession(
-    session,
-    context.params?.wId as string
-  );
-
+}>(async (context, auth) => {
   const owner = auth.workspace();
 
   if (!owner || !auth.isAdmin()) {

front/pages/poke/[wId]/data_sources/[name]/view.tsx

@@ -5,19 +5,13 @@ import type { InferGetServerSidePropsType } from "next";
 
 import PokeNavbar from "@app/components/poke/PokeNavbar";
 import { getDataSource } from "@app/lib/api/data_sources";
-import { Authenticator } from "@app/lib/auth";
-import { withDefaultGetServerSidePropsRequirements } from "@app/lib/iam/session";
+import { withSuperUserAuthRequirements } from "@app/lib/iam/session";
 import { classNames } from "@app/lib/utils";
 import logger from "@app/logger/logger";
 
-export const getServerSideProps = withDefaultGetServerSidePropsRequirements<{
+export const getServerSideProps = withSuperUserAuthRequirements<{
   document: CoreAPIDocument;
-}>(async (context, session) => {
-  const auth = await Authenticator.fromSuperUserSession(
-    session,
-    context.params?.wId as string
-  );
-
+}>(async (context, auth) => {
   if (!auth.isDustSuperUser()) {
     return {
       notFound: true,

front/pages/poke/[wId]/index.tsx

@@ -35,10 +35,9 @@ import {
   GLOBAL_AGENTS_SID,
   orderDatasourceByImportance,
 } from "@app/lib/assistant";
-import { Authenticator } from "@app/lib/auth";
 import { useSubmitFunction } from "@app/lib/client/utils";
 import { isDevelopment } from "@app/lib/development";
-import { withDefaultGetServerSidePropsRequirements } from "@app/lib/iam/session";
+import { withSuperUserAuthRequirements } from "@app/lib/iam/session";
 import {
   FREE_TEST_PLAN_CODE,
   FREE_UPGRADED_PLAN_CODE,
@@ -47,20 +46,15 @@ import {
 import { getPlanInvitation } from "@app/lib/plans/subscription";
 import { usePokePlans } from "@app/lib/swr";
 
-export const getServerSideProps = withDefaultGetServerSidePropsRequirements<{
+export const getServerSideProps = withSuperUserAuthRequirements<{
   owner: WorkspaceType;
   subscription: SubscriptionType;
   planInvitation: PlanInvitationType | null;
   dataSources: DataSourceType[];
   agentConfigurations: AgentConfigurationType[];
   whitelistableFeatures: WhitelistableFeature[];
   registry: typeof DustProdActionRegistry;
-}>(async (context, session) => {
-  const auth = await Authenticator.fromSuperUserSession(
-    session,
-    context.params?.wId as string
-  );
-
+}>(async (context, auth) => {
   const owner = auth.workspace();
   const subscription = auth.subscription();
 

front/pages/poke/[wId]/memberships.tsx

@@ -10,18 +10,12 @@ import React from "react";
 
 import PokeNavbar from "@app/components/poke/PokeNavbar";
 import { getMembers } from "@app/lib/api/workspace";
-import { Authenticator } from "@app/lib/auth";
-import { withDefaultGetServerSidePropsRequirements } from "@app/lib/iam/session";
+import { withSuperUserAuthRequirements } from "@app/lib/iam/session";
 
-export const getServerSideProps = withDefaultGetServerSidePropsRequirements<{
+export const getServerSideProps = withSuperUserAuthRequirements<{
   owner: WorkspaceType;
   members: UserTypeWithWorkspaces[];
-}>(async (context, session) => {
-  const auth = await Authenticator.fromSuperUserSession(
-    session,
-    context.params?.wId as string
-  );
-
+}>(async (context, auth) => {
   const owner = auth.workspace();
 
   if (!owner || !auth.isDustSuperUser()) {