Address review comments.

* Add publicDoc() function to read from bootstrap firestore
* Ensure updated config is set before authentication
* Move auth from global to local scope
* Rename guard function to fsShutdownGuard for clarity

src/commands/__tests__/__snapshots__/login.test.ts.snap

@@ -1,7 +1,14 @@
 // Jest Snapshot v1, https://goo.gl/fbAQLP
 
-exports[`login organization exists should write the user's identity to the file system 1`] = `
+exports[`login organization exists should write the user's identity & config to the file system 1`] = `
 [
+  [
+    "/Users/fabianjoya/.p0/config.json",
+    "{"fs":{"apiKey":"AIzaSyC9A5VXSwDDS-Vp4WH_UIanEqJvv_7XdlQ","authDomain":"p0-gcp-project.firebaseapp.com","projectId":"p0-gcp-project","storageBucket":"p0-gcp-project.appspot.com","messagingSenderId":"398809717501","appId":"1:398809717501:web:6dd1cab893b2faeb06fc94"},"google":{"clientId":"398809717501-j4j8ldjicsspidl4ecj6nj1oomo9eda1.apps.googleusercontent.com","clientSecret":"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"},"appUrl":"http://localhost:8088","environment":"production"}",
+    {
+      "mode": "600",
+    },
+  ],
   [
     "/path/to/home/.p0",
     "{

src/commands/__tests__/login.test.ts

@@ -73,7 +73,7 @@ describe("login", () => {
       await login({ org: "test-org" });
       expect(pluginLoginMap.google).toHaveBeenCalled();
     });
-    it("should write the user's identity to the file system", async () => {
+    it("should write the user's identity & config to the file system", async () => {
       await login({ org: "test-org" });
       expect(mockWriteFile.mock.calls).toMatchSnapshot();
     });

src/commands/allow.ts

@@ -10,7 +10,7 @@ You should have received a copy of the GNU General Public License along with @p0
 **/
 import { fetchCommand } from "../drivers/api";
 import { authenticate } from "../drivers/auth";
-import { guard } from "../drivers/firestore";
+import { fsShutdownGuard } from "../drivers/firestore";
 import { print2 } from "../drivers/stdio";
 import { AllowResponse } from "../types/allow";
 import { Authn } from "../types/identity";
@@ -37,7 +37,7 @@ export const allowCommand = (yargs: yargs.Argv) =>
     "allow [arguments..]",
     "Create standing access for a resource",
     allowArgs,
-    guard(allow)
+    fsShutdownGuard(allow)
   );
 
 export const allow = async (

src/commands/aws/role.ts

@@ -10,7 +10,7 @@ You should have received a copy of the GNU General Public License along with @p0
 **/
 import { parseXml } from "../../common/xml";
 import { authenticate } from "../../drivers/auth";
-import { guard } from "../../drivers/firestore";
+import { fsShutdownGuard } from "../../drivers/firestore";
 import { print1, print2 } from "../../drivers/stdio";
 import { getAwsConfig } from "../../plugins/aws/config";
 import { AwsFederatedLogin, AwsItem } from "../../plugins/aws/types";
@@ -29,7 +29,7 @@ export const role = (yargs: yargs.Argv<{ account: string | undefined }>) =>
         "List available AWS roles",
         identity,
         // TODO: select based on uidLocation
-        guard(oktaAwsListRoles)
+        fsShutdownGuard(oktaAwsListRoles)
       )
       .command(
         "assume <role>",
@@ -41,7 +41,7 @@ export const role = (yargs: yargs.Argv<{ account: string | undefined }>) =>
             describe: "An AWS role name",
           }),
         // TODO: select based on uidLocation
-        guard(oktaAwsAssumeRole)
+        fsShutdownGuard(oktaAwsAssumeRole)
       )
       .demandCommand(1)
   );

src/commands/grant.ts

@@ -8,7 +8,7 @@ This file is part of @p0security/cli
 
 You should have received a copy of the GNU General Public License along with @p0security/cli. If not, see <https://www.gnu.org/licenses/>.
 **/
-import { guard } from "../drivers/firestore";
+import { fsShutdownGuard } from "../drivers/firestore";
 import { request, requestArgs } from "./shared/request";
 import yargs from "yargs";
 
@@ -17,5 +17,5 @@ export const grantCommand = (yargs: yargs.Argv) =>
     "grant [arguments..]",
     "Grant access to another identity",
     requestArgs,
-    guard(request("grant"))
+    fsShutdownGuard(request("grant"))
   );

src/commands/kubeconfig.ts

@@ -11,7 +11,7 @@ You should have received a copy of the GNU General Public License along with @p0
 import { retryWithSleep } from "../common/retry";
 import { AnsiSgr } from "../drivers/ansi";
 import { authenticate } from "../drivers/auth";
-import { guard } from "../drivers/firestore";
+import { fsShutdownGuard } from "../drivers/firestore";
 import { print2, spinUntil } from "../drivers/stdio";
 import { parseArn } from "../plugins/aws/utils";
 import {
@@ -65,7 +65,7 @@ export const kubeconfigCommand = (yargs: yargs.Argv) =>
           describe:
             "Requested duration for access (format like '10 minutes', '2 hours', '5 days', or '1 week')",
         }),
-    guard(kubeconfigAction)
+    fsShutdownGuard(kubeconfigAction)
   );
 
 const kubeconfigAction = async (

src/commands/login.ts

@@ -13,12 +13,12 @@ import {
   IDENTITY_FILE_PATH,
   loadCredentials,
 } from "../drivers/auth";
-import { saveTenantConfig } from "../drivers/config";
+import { saveConfig } from "../drivers/config";
+import { bootstrapConfig } from "../drivers/env";
 import {
   authenticateToFirebase,
-  doc,
-  guard,
-  initializeFirebase,
+  fsShutdownGuard,
+  publicDoc,
 } from "../drivers/firestore";
 import { print2 } from "../drivers/stdio";
 import { pluginLoginMap } from "../plugins/login";
@@ -38,12 +38,14 @@ export const login = async (
   args: { org: string },
   options?: { skipAuthenticate?: boolean }
 ) => {
-  initializeFirebase({ useBootstrapConfig: true });
-
-  const orgDoc = await getDoc<RawOrgData, object>(doc(`orgs/${args.org}`));
+  const orgDoc = await getDoc<RawOrgData, object>(
+    publicDoc(`orgs/${args.org}`)
+  );
   const orgData = orgDoc.data();
   if (!orgData) throw "Could not find organization";
 
+  await saveConfig(orgData.config ?? bootstrapConfig);
+
   const orgWithSlug: OrgData = { ...orgData, slug: args.org };
 
   const plugin = orgWithSlug?.ssoProvider;
@@ -56,10 +58,6 @@ export const login = async (
   await clearIdentityCache();
   await writeIdentity(orgWithSlug, tokenResponse);
 
-  if (orgData.config) {
-    await saveTenantConfig(orgData.config);
-  }
-
   // validate auth
   if (!options?.skipAuthenticate) {
     const identity = await loadCredentials({ noRefresh: true });
@@ -110,5 +108,5 @@ export const loginCommand = (yargs: yargs.Argv) =>
         type: "string",
         describe: "Your P0 organization ID",
       }),
-    guard(login)
+    fsShutdownGuard(login)
   );

src/commands/ls.ts

@@ -11,7 +11,7 @@ You should have received a copy of the GNU General Public License along with @p0
 import { AnsiSgr } from "../drivers/ansi";
 import { fetchCommand } from "../drivers/api";
 import { authenticate } from "../drivers/auth";
-import { guard } from "../drivers/firestore";
+import { fsShutdownGuard } from "../drivers/firestore";
 import { print2, print1, spinUntil } from "../drivers/stdio";
 import { max, orderBy } from "lodash";
 import pluralize from "pluralize";
@@ -45,7 +45,7 @@ export const lsCommand = (yargs: yargs.Argv) =>
     "ls [arguments..]",
     "List request-command arguments",
     lsArgs,
-    guard(ls)
+    fsShutdownGuard(ls)
   );
 
 const ls = async (

src/commands/request.ts

@@ -8,7 +8,7 @@ This file is part of @p0security/cli
 
 You should have received a copy of the GNU General Public License along with @p0security/cli. If not, see <https://www.gnu.org/licenses/>.
 **/
-import { guard } from "../drivers/firestore";
+import { fsShutdownGuard } from "../drivers/firestore";
 import { request, requestArgs } from "./shared/request";
 import yargs from "yargs";
 
@@ -17,5 +17,5 @@ export const requestCommand = (yargs: yargs.Argv) =>
     "request [arguments..]",
     "Manually request permissions on a resource",
     requestArgs,
-    guard(request("request"))
+    fsShutdownGuard(request("request"))
   );

src/commands/scp.ts

@@ -9,7 +9,7 @@ This file is part of @p0security/cli
 You should have received a copy of the GNU General Public License along with @p0security/cli. If not, see <https://www.gnu.org/licenses/>.
 **/
 import { authenticate } from "../drivers/auth";
-import { guard } from "../drivers/firestore";
+import { fsShutdownGuard } from "../drivers/firestore";
 import { sshOrScp } from "../plugins/ssh";
 import { SshRequest, SupportedSshProviders } from "../types/ssh";
 import { prepareRequest, ScpCommandArgs } from "./shared/ssh";
@@ -69,7 +69,7 @@ export const scpCommand = (yargs: yargs.Argv) =>
   The '--' argument must be specified between P0-specific args on the left and SCP_ARGS on the right.`
         ),
 
-    guard(scpAction)
+    fsShutdownGuard(scpAction)
   );
 
 /** Transfers files between a local and remote hosts using SSH.

src/commands/ssh.ts

@@ -9,7 +9,7 @@ This file is part of @p0security/cli
 You should have received a copy of the GNU General Public License along with @p0security/cli. If not, see <https://www.gnu.org/licenses/>.
 **/
 import { authenticate } from "../drivers/auth";
-import { guard } from "../drivers/firestore";
+import { fsShutdownGuard } from "../drivers/firestore";
 import { sshOrScp } from "../plugins/ssh";
 import { SshCommandArgs, prepareRequest } from "./shared/ssh";
 import yargs from "yargs";
@@ -69,7 +69,7 @@ export const sshCommand = (yargs: yargs.Argv) =>
   $ p0 ssh example-instance --provider gcloud -- -NR '*:8080:localhost:8088' -o 'GatewayPorts yes'`
         ),
 
-    guard(sshAction)
+    fsShutdownGuard(sshAction)
   );
 
 /** Connect to an SSH backend

src/drivers/auth.ts

@@ -11,7 +11,7 @@ You should have received a copy of the GNU General Public License along with @p0
 import { login } from "../commands/login";
 import { Authn, Identity } from "../types/identity";
 import { P0_PATH } from "../util";
-import { loadTenantConfig } from "./config";
+import { loadConfig } from "./config";
 import { authenticateToFirebase, initializeFirebase } from "./firestore";
 import { print2 } from "./stdio";
 import * as fs from "fs/promises";
@@ -92,13 +92,10 @@ export const loadCredentials = async (options?: {
 export const authenticate = async (options?: {
   noRefresh?: boolean;
 }): Promise<Authn> => {
-  const identity = await loadCredentials(options);
-
-  await loadTenantConfig();
-
-  // Re-initialize Firebase with the loaded config
+  await loadConfig();
   initializeFirebase();
 
+  const identity = await loadCredentials(options);
   const userCredential = await authenticateToFirebase(identity);
 
   return { userCredential, identity };

src/drivers/config.ts

@@ -10,49 +10,22 @@ You should have received a copy of the GNU General Public License along with @p0
 **/
 import { Config } from "../types/org";
 import { P0_PATH } from "../util";
-import { bootstrapConfig } from "./env";
 import { print2 } from "./stdio";
 import fs from "fs/promises";
 import path from "path";
 
 export const CONFIG_FILE_PATH = path.join(P0_PATH, "config.json");
 
-export let tenantConfig = bootstrapConfig;
-
-/**
- * Configures the CLI to use a tenant-specific configuration instead of the bootstrap configuration.
- * The tenant-specific config is also written to the local filesystem for future use.
- *
- * @param config the tenant-specific config to use
- */
-export const saveTenantConfig = async (config: Config) => {
-  tenantConfig = config;
-  await writeConfigToFile(config);
-};
-
-/**
- * Configures the CLI to use a tenant-specific configuration, if present.
- *
- * Loads the tenant-specific config from the local filesystem, if present.
- * If not present, it will keep using the bootstrap config.
- */
-export const loadTenantConfig = async () => {
-  const config = await loadConfigFromFile();
-  if (config) {
-    tenantConfig = config;
-  }
-};
-
-const loadConfigFromFile = async (): Promise<Config> => {
-  print2(`Loading config from ${CONFIG_FILE_PATH}.`);
-  const buffer = await fs.readFile(CONFIG_FILE_PATH);
-  const config: Config = JSON.parse(buffer.toString());
-  return config;
-};
+export let tenantConfig: Config;
 
-const writeConfigToFile = async (config: Config) => {
+export async function saveConfig(config: Config) {
   print2(`Saving config to ${CONFIG_FILE_PATH}.`);
   const dir = path.dirname(CONFIG_FILE_PATH);
   await fs.mkdir(dir, { recursive: true });
   await fs.writeFile(CONFIG_FILE_PATH, JSON.stringify(config), { mode: "600" });
-};
+}
+
+export async function loadConfig() {
+  const buffer = await fs.readFile(CONFIG_FILE_PATH);
+  tenantConfig = JSON.parse(buffer.toString());
+}

src/drivers/firestore.ts

@@ -11,9 +11,8 @@ You should have received a copy of the GNU General Public License along with @p0
 import { Identity } from "../types/identity";
 import { tenantConfig } from "./config";
 import { bootstrapConfig } from "./env";
-import { FirebaseOptions, initializeApp } from "firebase/app";
+import { FirebaseApp, initializeApp } from "firebase/app";
 import {
-  Auth,
   getAuth,
   OAuthProvider,
   SignInMethod,
@@ -29,17 +28,15 @@ import {
   Firestore,
 } from "firebase/firestore";
 
-let firestore: Firestore;
-let auth: Auth;
+const bootstrapApp = initializeApp(bootstrapConfig.fs, "bootstrapApp");
+const bootstrapFirestore = getFirestore(bootstrapApp);
 
-export function initializeFirebase(options?: { useBootstrapConfig: boolean }) {
-  const config: FirebaseOptions = options?.useBootstrapConfig
-    ? bootstrapConfig.fs
-    : tenantConfig.fs;
-  const app = initializeApp(config);
+let app: FirebaseApp;
+let firestore: Firestore;
 
+export function initializeFirebase() {
+  app = initializeApp(tenantConfig.fs, "authFirebase");
   firestore = getFirestore(app);
-  auth = getAuth();
 }
 
 export async function authenticateToFirebase(identity: Identity) {
@@ -58,6 +55,7 @@ export async function authenticateToFirebase(identity: Identity) {
     idToken: credential.id_token,
   });
 
+  const auth = getAuth(app);
   auth.tenantId = tenantId;
 
   const userCredential = await signInWithCredential(auth, firebaseCredential);
@@ -81,16 +79,21 @@ export const doc = <T>(path: string) => {
   return fsDoc(firestore, path) as DocumentReference<T>;
 };
 
+export const publicDoc = <T>(path: string) => {
+  return fsDoc(bootstrapFirestore, path) as DocumentReference<T>;
+};
+
 /** Ensures that Firestore is shutdown at command termination
  *
  * This prevents Firestore from holding the command on execution completion or failure.
  */
-export const guard =
+export const fsShutdownGuard =
   <P, T>(cb: (args: P) => Promise<T>) =>
   async (args: P) => {
     try {
       await cb(args);
     } finally {
+      void terminate(bootstrapFirestore);
       void terminate(firestore);
     }
   };