Converting to use both inline and shared policy

refactor: standardize dropdowns accross tables

server/db/pg/schema/schema.ts

@@ -1392,6 +1392,16 @@ export type ResourceHeaderAuthExtendedCompatibility = InferSelectModel<
 export type ResourceOtp = InferSelectModel<typeof resourceOtp>;
 export type ResourceAccessToken = InferSelectModel<typeof resourceAccessToken>;
 export type ResourceWhitelist = InferSelectModel<typeof resourceWhitelist>;
+export type ResourcePolicyPincode = InferSelectModel<
+    typeof resourcePolicyPincode
+>;
+export type ResourcePolicyPassword = InferSelectModel<
+    typeof resourcePolicyPassword
+>;
+export type ResourcePolicyHeaderAuth = InferSelectModel<
+    typeof resourcePolicyHeaderAuth
+>;
+
 export type VersionMigration = InferSelectModel<typeof versionMigrations>;
 export type ResourceRule = InferSelectModel<typeof resourceRules>;
 export type Domain = InferSelectModel<typeof domains>;

server/db/queries/verifySessionQueries.ts

@@ -35,6 +35,7 @@ import {
     resourcePolicyHeaderAuth,
     ResourcePolicyHeaderAuth
 } from "@server/db";
+import { alias } from "drizzle-orm/sqlite-core";
 import { and, eq, inArray, or, sql } from "drizzle-orm";
 
 export type ResourceWithAuth = {
@@ -67,6 +68,33 @@ export async function getResourceByDomain(
         wildcardCandidates.push(`*.${parts.slice(i).join(".")}`);
     }
 
+    const sharedPolicy = alias(resourcePolicies, "sharedPolicy");
+    const defaultPolicy = alias(resourcePolicies, "defaultPolicy");
+    const sharedPolicyPincode = alias(
+        resourcePolicyPincode,
+        "sharedPolicyPincode"
+    );
+    const defaultPolicyPincode = alias(
+        resourcePolicyPincode,
+        "defaultPolicyPincode"
+    );
+    const sharedPolicyPassword = alias(
+        resourcePolicyPassword,
+        "sharedPolicyPassword"
+    );
+    const defaultPolicyPassword = alias(
+        resourcePolicyPassword,
+        "defaultPolicyPassword"
+    );
+    const sharedPolicyHeaderAuth = alias(
+        resourcePolicyHeaderAuth,
+        "sharedPolicyHeaderAuth"
+    );
+    const defaultPolicyHeaderAuth = alias(
+        resourcePolicyHeaderAuth,
+        "defaultPolicyHeaderAuth"
+    );
+
     const potentialResults = await db
         .select()
         .from(resources)
@@ -90,28 +118,56 @@ export async function getResourceByDomain(
             )
         )
         .leftJoin(
-            resourcePolicies,
-            eq(resourcePolicies.resourcePolicyId, resources.resourcePolicyId)
+            sharedPolicy,
+            eq(sharedPolicy.resourcePolicyId, resources.resourcePolicyId)
         )
         .leftJoin(
-            resourcePolicyPincode,
+            sharedPolicyPincode,
             eq(
-                resourcePolicyPincode.resourcePolicyId,
-                resourcePolicies.resourcePolicyId
+                sharedPolicyPincode.resourcePolicyId,
+                sharedPolicy.resourcePolicyId
             )
         )
         .leftJoin(
-            resourcePolicyPassword,
+            sharedPolicyPassword,
             eq(
-                resourcePolicyPassword.resourcePolicyId,
-                resourcePolicies.resourcePolicyId
+                sharedPolicyPassword.resourcePolicyId,
+                sharedPolicy.resourcePolicyId
             )
         )
         .leftJoin(
-            resourcePolicyHeaderAuth,
+            sharedPolicyHeaderAuth,
             eq(
-                resourcePolicyHeaderAuth.resourcePolicyId,
-                resourcePolicies.resourcePolicyId
+                sharedPolicyHeaderAuth.resourcePolicyId,
+                sharedPolicy.resourcePolicyId
+            )
+        )
+        .leftJoin(
+            defaultPolicy,
+            eq(
+                defaultPolicy.resourcePolicyId,
+                resources.defaultResourcePolicyId
+            )
+        )
+        .leftJoin(
+            defaultPolicyPincode,
+            eq(
+                defaultPolicyPincode.resourcePolicyId,
+                defaultPolicy.resourcePolicyId
+            )
+        )
+        .leftJoin(
+            defaultPolicyPassword,
+            eq(
+                defaultPolicyPassword.resourcePolicyId,
+                defaultPolicy.resourcePolicyId
+            )
+        )
+        .leftJoin(
+            defaultPolicyHeaderAuth,
+            eq(
+                defaultPolicyHeaderAuth.resourcePolicyId,
+                defaultPolicy.resourcePolicyId
             )
         )
         .innerJoin(orgs, eq(orgs.orgId, resources.orgId))
@@ -143,18 +199,24 @@ export async function getResourceByDomain(
         return null;
     }
 
+    const effectivePolicyPincode =
+        result.sharedPolicyPincode ?? result.defaultPolicyPincode ?? null;
+    const effectivePolicyPassword =
+        result.sharedPolicyPassword ?? result.defaultPolicyPassword ?? null;
+    const effectivePolicyHeaderAuth =
+        result.sharedPolicyHeaderAuth ?? result.defaultPolicyHeaderAuth ?? null;
+
     return {
         resource: result.resources,
-        pincode: result.resourcePolicyPincode ?? result.resourcePincode,
-        password: result.resourcePolicyPassword ?? result.resourcePassword,
-        headerAuth:
-            result.resourcePolicyHeaderAuth ?? result.resourceHeaderAuth,
-        headerAuthExtendedCompatibility: result.resourcePolicyHeaderAuth
+        pincode: effectivePolicyPincode ?? result.resourcePincode,
+        password: effectivePolicyPassword ?? result.resourcePassword,
+        headerAuth: effectivePolicyHeaderAuth ?? result.resourceHeaderAuth,
+        headerAuthExtendedCompatibility: effectivePolicyHeaderAuth
             ? ({
                   headerAuthExtendedCompatibilityId: 0,
                   resourceId: result.resources.resourceId,
                   extendedCompatibilityIsActivated:
-                      result.resourcePolicyHeaderAuth.extendedCompatibility
+                      effectivePolicyHeaderAuth.extendedCompatibility
               } as ResourceHeaderAuthExtendedCompatibility)
             : result.resourceHeaderAuthExtendedCompatibility,
         org: result.orgs

server/private/routers/hybrid.ts

@@ -61,6 +61,7 @@ import {
     roles
 } from "@server/db";
 import { eq, and, inArray, isNotNull, ne, or, sql } from "drizzle-orm";
+import { alias } from "drizzle-orm/sqlite-core";
 import { response } from "@server/lib/response";
 import HttpCode from "@server/types/HttpCode";
 import { NextFunction, Request, Response } from "express";
@@ -514,6 +515,33 @@ hybridRouter.get(
                 wildcardCandidates.push(`*.${domainParts.slice(i).join(".")}`);
             }
 
+            const sharedPolicy = alias(resourcePolicies, "sharedPolicy");
+            const defaultPolicy = alias(resourcePolicies, "defaultPolicy");
+            const sharedPolicyPincode = alias(
+                resourcePolicyPincode,
+                "sharedPolicyPincode"
+            );
+            const defaultPolicyPincode = alias(
+                resourcePolicyPincode,
+                "defaultPolicyPincode"
+            );
+            const sharedPolicyPassword = alias(
+                resourcePolicyPassword,
+                "sharedPolicyPassword"
+            );
+            const defaultPolicyPassword = alias(
+                resourcePolicyPassword,
+                "defaultPolicyPassword"
+            );
+            const sharedPolicyHeaderAuth = alias(
+                resourcePolicyHeaderAuth,
+                "sharedPolicyHeaderAuth"
+            );
+            const defaultPolicyHeaderAuth = alias(
+                resourcePolicyHeaderAuth,
+                "defaultPolicyHeaderAuth"
+            );
+
             const potentialResults = await db
                 .select()
                 .from(resources)
@@ -537,31 +565,59 @@ hybridRouter.get(
                     )
                 )
                 .leftJoin(
-                    resourcePolicies,
+                    sharedPolicy,
                     eq(
-                        resourcePolicies.resourcePolicyId,
+                        sharedPolicy.resourcePolicyId,
                         resources.resourcePolicyId
                     )
                 )
                 .leftJoin(
-                    resourcePolicyPincode,
+                    sharedPolicyPincode,
                     eq(
-                        resourcePolicyPincode.resourcePolicyId,
-                        resourcePolicies.resourcePolicyId
+                        sharedPolicyPincode.resourcePolicyId,
+                        sharedPolicy.resourcePolicyId
                     )
                 )
                 .leftJoin(
-                    resourcePolicyPassword,
+                    sharedPolicyPassword,
                     eq(
-                        resourcePolicyPassword.resourcePolicyId,
-                        resourcePolicies.resourcePolicyId
+                        sharedPolicyPassword.resourcePolicyId,
+                        sharedPolicy.resourcePolicyId
                     )
                 )
                 .leftJoin(
-                    resourcePolicyHeaderAuth,
+                    sharedPolicyHeaderAuth,
                     eq(
-                        resourcePolicyHeaderAuth.resourcePolicyId,
-                        resourcePolicies.resourcePolicyId
+                        sharedPolicyHeaderAuth.resourcePolicyId,
+                        sharedPolicy.resourcePolicyId
+                    )
+                )
+                .leftJoin(
+                    defaultPolicy,
+                    eq(
+                        defaultPolicy.resourcePolicyId,
+                        resources.defaultResourcePolicyId
+                    )
+                )
+                .leftJoin(
+                    defaultPolicyPincode,
+                    eq(
+                        defaultPolicyPincode.resourcePolicyId,
+                        defaultPolicy.resourcePolicyId
+                    )
+                )
+                .leftJoin(
+                    defaultPolicyPassword,
+                    eq(
+                        defaultPolicyPassword.resourcePolicyId,
+                        defaultPolicy.resourcePolicyId
+                    )
+                )
+                .leftJoin(
+                    defaultPolicyHeaderAuth,
+                    eq(
+                        defaultPolicyHeaderAuth.resourcePolicyId,
+                        defaultPolicy.resourcePolicyId
                     )
                 )
                 .innerJoin(orgs, eq(orgs.orgId, resources.orgId))
@@ -614,21 +670,31 @@ hybridRouter.get(
                 });
             }
 
+            const effectivePolicyPincode =
+                result.sharedPolicyPincode ??
+                result.defaultPolicyPincode ??
+                null;
+            const effectivePolicyPassword =
+                result.sharedPolicyPassword ??
+                result.defaultPolicyPassword ??
+                null;
+            const effectivePolicyHeaderAuth =
+                result.sharedPolicyHeaderAuth ??
+                result.defaultPolicyHeaderAuth ??
+                null;
+
             const resourceWithAuth: ResourceWithAuth = {
                 resource: result.resources,
-                pincode: result.resourcePolicyPincode ?? result.resourcePincode,
-                password:
-                    result.resourcePolicyPassword ?? result.resourcePassword,
+                pincode: effectivePolicyPincode ?? result.resourcePincode,
+                password: effectivePolicyPassword ?? result.resourcePassword,
                 headerAuth:
-                    result.resourcePolicyHeaderAuth ??
-                    result.resourceHeaderAuth,
-                headerAuthExtendedCompatibility: result.resourcePolicyHeaderAuth
+                    effectivePolicyHeaderAuth ?? result.resourceHeaderAuth,
+                headerAuthExtendedCompatibility: effectivePolicyHeaderAuth
                     ? ({
                           headerAuthExtendedCompatibilityId: 0,
                           resourceId: result.resources.resourceId,
                           extendedCompatibilityIsActivated:
-                              result.resourcePolicyHeaderAuth
-                                  .extendedCompatibility
+                              effectivePolicyHeaderAuth.extendedCompatibility
                       } as ResourceHeaderAuthExtendedCompatibility)
                     : result.resourceHeaderAuthExtendedCompatibility,
                 org: result.orgs

server/private/routers/internal.ts

@@ -19,6 +19,7 @@ import * as license from "#private/routers/license";
 import * as resource from "#private/routers/resource";
 import * as browserTarget from "#private/routers/browserGatewayTarget";
 import * as ssh from "#private/routers/ssh";
+import * as ws from "@server/routers/ws";
 
 import {
     verifySessionUserMiddleware,
@@ -52,4 +53,10 @@ internalRouter.post(
     ssh.signSshKey
 );
 
+internalRouter.get(
+    "/ws/round-trip-message/:messageId",
+    verifyUserFromResourceSessionMiddleware,
+    ws.checkRoundTripMessage
+);
+
 internalRouter.get("/resource/browser-target", browserTarget.getBrowserTarget);

server/private/routers/ssh/signSshKey.ts

@@ -19,12 +19,18 @@ import {
     logsDb,
     newts,
     roles,
+    roleResources,
     roleSiteResources,
+    resources,
     roundTripMessageTracker,
     siteResources,
     siteNetworks,
+    targets,
     userOrgs,
-    sites
+    sites,
+    Resource,
+    SiteResource,
+    browserGatewayTarget
 } from "@server/db";
 import { logAccessAudit } from "#private/lib/logAccessAudit";
 import { isLicensedOrSubscribed } from "#private/lib/isLicencedOrSubscribed";
@@ -35,11 +41,13 @@ import createHttpError from "http-errors";
 import logger from "@server/logger";
 import { fromError } from "zod-validation-error";
 import { and, eq, inArray, or } from "drizzle-orm";
+import { canUserAccessResource } from "@server/auth/canUserAccessResource";
 import { canUserAccessSiteResource } from "@server/auth/canUserAccessSiteResource";
 import { signPublicKey, getOrgCAKeys } from "@server/lib/sshCA";
 import config from "@server/lib/config";
 import { sendToClient } from "#private/routers/ws";
 import { ActionsEnum } from "@server/auth/actions";
+import type { SignSshKeyResponse } from "@server/routers/ssh/types";
 
 const paramsSchema = z.strictObject({
     orgId: z.string().nonempty()
@@ -50,7 +58,8 @@ const bodySchema = z
         publicKey: z.string().nonempty(),
         resourceId: z.number().int().positive().optional(),
         resource: z.string().nonempty().optional(), // this is either the nice id or the alias
-        username: z.string().nonempty().optional()
+        username: z.string().nonempty().optional(),
+        type: z.enum(["public", "private"]).default("private")
     })
     .refine(
         (data) => {
@@ -64,23 +73,6 @@ const bodySchema = z
         }
     );
 
-export type SignSshKeyResponse = {
-    certificate?: string;
-    messageIds: number[];
-    messageId?: number;
-    sshUsername: string;
-    sshHost: string;
-    resourceId: number;
-    siteIds: number[];
-    siteId: number;
-    keyId?: string;
-    validPrincipals?: string[];
-    validAfter?: string;
-    validBefore?: string;
-    expiresIn?: number;
-    authDaemonMode: "site" | "remote" | "native" | null;
-};
-
 export async function signSshKey(
     req: Request,
     res: Response,
@@ -111,6 +103,7 @@ export async function signSshKey(
         const {
             publicKey,
             resourceId,
+            type,
             resource: resourceQueryString,
             username
         } = parsedBody.data;
@@ -175,18 +168,25 @@ export async function signSshKey(
             );
         }
 
-        // Verify the resource exists and belongs to the org
-        // Build the where clause dynamically based on which field is provided
+        let matchingResources: SiteResource[] | Resource[] = [];
+        // Verify the resource exists and belongs to the org.
+        // Build the where clause dynamically based on which field is provided.
         let whereClause;
         if (resourceId !== undefined) {
-            whereClause = eq(siteResources.siteResourceId, resourceId);
+            whereClause =
+                type === "private"
+                    ? eq(siteResources.siteResourceId, resourceId)
+                    : eq(resources.resourceId, resourceId);
         } else if (resourceQueryString !== undefined) {
-            whereClause = or(
-                eq(siteResources.niceId, resourceQueryString),
-                eq(siteResources.alias, resourceQueryString)
-            );
+            whereClause =
+                type === "private"
+                    ? or(
+                          eq(siteResources.niceId, resourceQueryString),
+                          eq(siteResources.alias, resourceQueryString)
+                      )
+                    : eq(resources.niceId, resourceQueryString);
         } else {
-            // This should never happen due to the schema validation, but TypeScript doesn't know that
+            // This should never happen due to the schema validation, but TypeScript doesn't know that.
             return next(
                 createHttpError(
                     HttpCode.BAD_REQUEST,
@@ -195,18 +195,25 @@ export async function signSshKey(
             );
         }
 
-        const resources = await db
-            .select()
-            .from(siteResources)
-            .where(and(whereClause, eq(siteResources.orgId, orgId)));
+        if (type === "private") {
+            matchingResources = await db
+                .select()
+                .from(siteResources)
+                .where(and(whereClause, eq(siteResources.orgId, orgId)));
+        } else {
+            matchingResources = await db
+                .select()
+                .from(resources)
+                .where(and(whereClause, eq(resources.orgId, orgId)));
+        }
 
-        if (!resources || resources.length === 0) {
+        if (!matchingResources || matchingResources.length === 0) {
             return next(
                 createHttpError(HttpCode.NOT_FOUND, `Resource not found`)
             );
         }
 
-        if (resources.length > 1) {
+        if (matchingResources.length > 1) {
             // error but this should not happen because the nice id cant contain a dot and the alias has to have a dot and both have to be unique within the org so there should never be multiple matches
             return next(
                 createHttpError(
@@ -216,7 +223,11 @@ export async function signSshKey(
             );
         }
 
-        const resource = resources[0];
+        const resource = matchingResources[0];
+        const normalizedResourceId =
+            type === "private"
+                ? (resource as SiteResource).siteResourceId
+                : (resource as Resource).resourceId;
 
         if (resource.orgId !== orgId) {
             return next(
@@ -237,11 +248,18 @@ export async function signSshKey(
         }
 
         // Check if the user has access to the resource
-        const hasAccess = await canUserAccessSiteResource({
-            userId: userId,
-            resourceId: resource.siteResourceId,
-            roleIds
-        });
+        const hasAccess =
+            type === "private"
+                ? await canUserAccessSiteResource({
+                      userId: userId,
+                      resourceId: (resource as SiteResource).siteResourceId,
+                      roleIds
+                  })
+                : await canUserAccessResource({
+                      userId: userId,
+                      resourceId: (resource as Resource).resourceId,
+                      roleIds
+                  });
 
         if (!hasAccess) {
             return next(
@@ -252,12 +270,56 @@ export async function signSshKey(
             );
         }
 
-        const sitesFromNetworks = await db
-            .select({ siteId: siteNetworks.siteId })
-            .from(siteNetworks)
-            .where(eq(siteNetworks.networkId, resource.networkId!));
+        const siteAgentHostMap = new Map<number, string>();
+        let siteIds: number[] = [];
 
-        const siteIds = sitesFromNetworks.map((site) => site.siteId);
+        if (type === "private") {
+            const privateResource = resource as SiteResource;
+            const sitesFromNetworks = await db
+                .select({ siteId: siteNetworks.siteId })
+                .from(siteNetworks)
+                .where(eq(siteNetworks.networkId, privateResource.networkId!));
+
+            siteIds = sitesFromNetworks.map((site) => site.siteId);
+            for (const siteId of siteIds) {
+                if (privateResource.destination) {
+                    siteAgentHostMap.set(siteId, privateResource.destination);
+                }
+            }
+        } else {
+            const publicResource = resource as Resource;
+            const targetRows = await db
+                .select({
+                    siteId: browserGatewayTarget.siteId,
+                    ip: browserGatewayTarget.destination
+                })
+                .from(browserGatewayTarget)
+                .where(
+                    and(
+                        eq(
+                            browserGatewayTarget.resourceId,
+                            publicResource.resourceId
+                        )
+                    )
+                );
+
+            if (targetRows.length === 0) {
+                return next(
+                    createHttpError(
+                        HttpCode.NOT_FOUND,
+                        "No enabled targets found for the resource"
+                    )
+                );
+            }
+
+            for (const targetRow of targetRows) {
+                if (!siteAgentHostMap.has(targetRow.siteId)) {
+                    siteAgentHostMap.set(targetRow.siteId, targetRow.ip);
+                }
+            }
+
+            siteIds = Array.from(siteAgentHostMap.keys());
+        }
 
         let expiresIn: number | undefined;
         let messageIds: number[] = [];
@@ -374,27 +436,50 @@ export async function signSshKey(
                 usernameToUse = userOrg.pamUsername;
             }
 
-            const roleRows = await db
-                .select({
-                    sshSudoCommands: roles.sshSudoCommands,
-                    sshUnixGroups: roles.sshUnixGroups,
-                    sshCreateHomeDir: roles.sshCreateHomeDir,
-                    sshSudoMode: roles.sshSudoMode
-                })
-                .from(roles)
-                .innerJoin(
-                    roleSiteResources,
-                    eq(roleSiteResources.roleId, roles.roleId)
-                )
-                .where(
-                    and(
-                        inArray(roles.roleId, roleIds),
-                        eq(
-                            roleSiteResources.siteResourceId,
-                            resource.siteResourceId
-                        )
-                    )
-                );
+            const roleRows =
+                type === "private"
+                    ? await db
+                          .select({
+                              sshSudoCommands: roles.sshSudoCommands,
+                              sshUnixGroups: roles.sshUnixGroups,
+                              sshCreateHomeDir: roles.sshCreateHomeDir,
+                              sshSudoMode: roles.sshSudoMode
+                          })
+                          .from(roles)
+                          .innerJoin(
+                              roleSiteResources,
+                              eq(roleSiteResources.roleId, roles.roleId)
+                          )
+                          .where(
+                              and(
+                                  inArray(roles.roleId, roleIds),
+                                  eq(
+                                      roleSiteResources.siteResourceId,
+                                      (resource as SiteResource).siteResourceId
+                                  )
+                              )
+                          )
+                    : await db
+                          .select({
+                              sshSudoCommands: roles.sshSudoCommands,
+                              sshUnixGroups: roles.sshUnixGroups,
+                              sshCreateHomeDir: roles.sshCreateHomeDir,
+                              sshSudoMode: roles.sshSudoMode
+                          })
+                          .from(roles)
+                          .innerJoin(
+                              roleResources,
+                              eq(roleResources.roleId, roles.roleId)
+                          )
+                          .where(
+                              and(
+                                  inArray(roles.roleId, roleIds),
+                                  eq(
+                                      roleResources.resourceId,
+                                      (resource as Resource).resourceId
+                                  )
+                              )
+                          );
 
             const parsedSudoCommands: string[] = [];
             const parsedGroupsSet = new Set<string>();
@@ -480,6 +565,16 @@ export async function signSshKey(
 
                 messageIds.push(message.messageId);
 
+                const agentHost = siteAgentHostMap.get(siteId);
+                if (!agentHost) {
+                    return next(
+                        createHttpError(
+                            HttpCode.INTERNAL_SERVER_ERROR,
+                            `Unable to determine agent host for site ${siteId}`
+                        )
+                    );
+                }
+
                 await sendToClient(newt.newtId, {
                     type: `newt/pam/connection`,
                     data: {
@@ -489,7 +584,7 @@ export async function signSshKey(
                         authDaemonMode: resource.authDaemonMode, // site, remote, native where native is the pty mode
                         externalAuthDaemon:
                             resource.authDaemonMode === "remote", // keep this for backward compatibility but new newts are using the authDaemonMode field
-                        agentHost: resource.destination,
+                        agentHost,
                         caCert: caKeys.publicKeyOpenSSH,
                         username: usernameToUse,
                         niceId: resource.niceId,
@@ -526,10 +621,19 @@ export async function signSshKey(
             resource.authDaemonMode === "site" ||
             resource.authDaemonMode === "remote"
         ) {
-            if (resource.alias && resource.alias != "") {
-                sshHost = resource.alias;
+            if (type === "private") {
+                const privateResource = resource as SiteResource;
+                if (privateResource.alias && privateResource.alias !== "") {
+                    sshHost = privateResource.alias;
+                } else {
+                    sshHost = privateResource.destination || "";
+                }
             } else {
-                sshHost = resource.destination || "";
+                const publicResource = resource as Resource;
+                sshHost =
+                    publicResource.fullDomain ||
+                    publicResource.subdomain ||
+                    publicResource.niceId;
             }
         } else if (resource.authDaemonMode === "native") {
             if (siteIds.length > 1) {
@@ -587,7 +691,8 @@ export async function signSshKey(
             actorId: req.user?.userId ?? "",
             action: ActionsEnum.signSshKey,
             metadata: JSON.stringify({
-                resourceId: resource.siteResourceId,
+                resourceId: normalizedResourceId,
+                resourceType: type,
                 resource: resource.name,
                 siteIds: siteIds
             })
@@ -597,7 +702,14 @@ export async function signSshKey(
             action: true,
             type: "ssh",
             orgId: orgId,
-            siteResourceId: resource.siteResourceId,
+            resourceId:
+                type === "public"
+                    ? (resource as Resource).resourceId
+                    : undefined,
+            siteResourceId:
+                type === "private"
+                    ? (resource as SiteResource).siteResourceId
+                    : undefined,
             user: req.user
                 ? { username: req.user.username ?? "", userId: req.user.userId }
                 : undefined,
@@ -618,7 +730,7 @@ export async function signSshKey(
                 messageId: messageIds[0], // just pick the first one for backward compatibility with older olms
                 sshUsername: usernameToUse,
                 sshHost: sshHost, // just pick the first one for backward compatibility with older olms
-                resourceId: resource.siteResourceId,
+                resourceId: normalizedResourceId,
                 siteIds: siteIds,
                 siteId: siteIds[0], // just pick the first one for backward compatibility with older olms
                 keyId: cert?.keyId,

server/routers/badger/exchangeSession.ts

@@ -188,6 +188,9 @@ export async function exchangeSession(
                 userSessionId: requestSession.userSessionId,
                 whitelistId: requestSession.whitelistId,
                 accessTokenId: requestSession.accessTokenId,
+                policyPasswordId: requestSession.policyPasswordId,
+                policyPincodeId: requestSession.policyPincodeId,
+                policyWhitelistId: requestSession.policyWhitelistId,
                 doNotExtend: false,
                 expiresAt: expires,
                 sessionLength: RESOURCE_SESSION_COOKIE_EXPIRES

server/routers/badger/verifySession.ts

@@ -876,6 +876,10 @@ function allowed(
         message: "Access allowed",
         status: HttpCode.OK
     };
+    logger.debug(
+        "++++++++++++++++++++++++++++++++++Access allowed, response data:",
+        data
+    );
     return response<VerifyUserResponse>(res, data);
 }
 

server/routers/resource/authWithPassword.ts

@@ -1,10 +1,17 @@
 import { verify } from "@node-rs/argon2";
 import { generateSessionToken } from "@server/auth/sessions/app";
 import { db } from "@server/db";
-import { orgs, resourcePassword, resourcePolicies, resourcePolicyPassword, resources } from "@server/db";
+import {
+    orgs,
+    resourcePassword,
+    resourcePolicies,
+    resourcePolicyPassword,
+    resources
+} from "@server/db";
 import HttpCode from "@server/types/HttpCode";
 import response from "@server/lib/response";
 import { eq } from "drizzle-orm";
+import { alias } from "drizzle-orm/sqlite-core";
 import { NextFunction, Request, Response } from "express";
 import createHttpError from "http-errors";
 import { z } from "zod";
@@ -58,17 +65,45 @@ export async function authWithPassword(
     const { password } = parsedBody.data;
 
     try {
+        const sharedPolicy = alias(resourcePolicies, "sharedPolicy");
+        const defaultPolicy = alias(resourcePolicies, "defaultPolicy");
+        const sharedPolicyPassword = alias(
+            resourcePolicyPassword,
+            "sharedPolicyPassword"
+        );
+        const defaultPolicyPassword = alias(
+            resourcePolicyPassword,
+            "defaultPolicyPassword"
+        );
+
         const [result] = await db
             .select()
             .from(resources)
             .leftJoin(orgs, eq(orgs.orgId, resources.orgId))
             .leftJoin(
-                resourcePolicies,
-                eq(resourcePolicies.resourcePolicyId, resources.resourcePolicyId)
+                sharedPolicy,
+                eq(sharedPolicy.resourcePolicyId, resources.resourcePolicyId)
             )
             .leftJoin(
-                resourcePolicyPassword,
-                eq(resourcePolicyPassword.resourcePolicyId, resourcePolicies.resourcePolicyId)
+                sharedPolicyPassword,
+                eq(
+                    sharedPolicyPassword.resourcePolicyId,
+                    sharedPolicy.resourcePolicyId
+                )
+            )
+            .leftJoin(
+                defaultPolicy,
+                eq(
+                    defaultPolicy.resourcePolicyId,
+                    resources.defaultResourcePolicyId
+                )
+            )
+            .leftJoin(
+                defaultPolicyPassword,
+                eq(
+                    defaultPolicyPassword.resourcePolicyId,
+                    defaultPolicy.resourcePolicyId
+                )
             )
             .leftJoin(
                 resourcePassword,
@@ -80,9 +115,13 @@ export async function authWithPassword(
         const resource = result?.resources;
         const org = result?.orgs;
 
-        // Policy password takes precedence over resource-level password
-        const policyPassword = result?.resourcePolicyPassword ?? null;
-        const definedPassword = policyPassword ?? result?.resourcePassword ?? null;
+        // Shared policy takes precedence, then default (inline) policy, then resource-level
+        const policyPassword =
+            result?.sharedPolicyPassword ??
+            result?.defaultPolicyPassword ??
+            null;
+        const definedPassword =
+            policyPassword ?? result?.resourcePassword ?? null;
         const isPolicyPassword = !!policyPassword;
 
         if (!org) {
@@ -136,7 +175,9 @@ export async function authWithPassword(
             resourceId,
             token,
             passwordId: isPolicyPassword ? null : definedPassword.passwordId,
-            policyPasswordId: isPolicyPassword ? definedPassword.passwordId : null,
+            policyPasswordId: isPolicyPassword
+                ? definedPassword.passwordId
+                : null,
             isRequestToken: true,
             expiresAt: Date.now() + 1000 * 30, // 30 seconds
             sessionLength: 1000 * 30,

server/routers/resource/authWithPincode.ts

@@ -1,9 +1,16 @@
 import { generateSessionToken } from "@server/auth/sessions/app";
 import { db } from "@server/db";
-import { orgs, resourcePincode, resourcePolicies, resourcePolicyPincode, resources } from "@server/db";
+import {
+    orgs,
+    resourcePincode,
+    resourcePolicies,
+    resourcePolicyPincode,
+    resources
+} from "@server/db";
 import HttpCode from "@server/types/HttpCode";
 import response from "@server/lib/response";
 import { eq } from "drizzle-orm";
+import { alias } from "drizzle-orm/sqlite-core";
 import { NextFunction, Request, Response } from "express";
 import createHttpError from "http-errors";
 import { z } from "zod";
@@ -57,17 +64,45 @@ export async function authWithPincode(
     const { pincode } = parsedBody.data;
 
     try {
+        const sharedPolicy = alias(resourcePolicies, "sharedPolicy");
+        const defaultPolicy = alias(resourcePolicies, "defaultPolicy");
+        const sharedPolicyPincode = alias(
+            resourcePolicyPincode,
+            "sharedPolicyPincode"
+        );
+        const defaultPolicyPincode = alias(
+            resourcePolicyPincode,
+            "defaultPolicyPincode"
+        );
+
         const [result] = await db
             .select()
             .from(resources)
             .leftJoin(orgs, eq(orgs.orgId, resources.orgId))
             .leftJoin(
-                resourcePolicies,
-                eq(resourcePolicies.resourcePolicyId, resources.resourcePolicyId)
+                sharedPolicy,
+                eq(sharedPolicy.resourcePolicyId, resources.resourcePolicyId)
             )
             .leftJoin(
-                resourcePolicyPincode,
-                eq(resourcePolicyPincode.resourcePolicyId, resourcePolicies.resourcePolicyId)
+                sharedPolicyPincode,
+                eq(
+                    sharedPolicyPincode.resourcePolicyId,
+                    sharedPolicy.resourcePolicyId
+                )
+            )
+            .leftJoin(
+                defaultPolicy,
+                eq(
+                    defaultPolicy.resourcePolicyId,
+                    resources.defaultResourcePolicyId
+                )
+            )
+            .leftJoin(
+                defaultPolicyPincode,
+                eq(
+                    defaultPolicyPincode.resourcePolicyId,
+                    defaultPolicy.resourcePolicyId
+                )
             )
             .leftJoin(
                 resourcePincode,
@@ -79,8 +114,9 @@ export async function authWithPincode(
         const resource = result?.resources;
         const org = result?.orgs;
 
-        // Policy pincode takes precedence over resource-level pincode
-        const policyPincode = result?.resourcePolicyPincode ?? null;
+        // Shared policy takes precedence, then default (inline) policy, then resource-level
+        const policyPincode =
+            result?.sharedPolicyPincode ?? result?.defaultPolicyPincode ?? null;
         const definedPincode = policyPincode ?? result?.resourcePincode ?? null;
         const isPolicyPincode = !!policyPincode;
 

server/routers/resource/authWithWhitelist.ts

@@ -1,6 +1,12 @@
 import { generateSessionToken } from "@server/auth/sessions/app";
 import { db } from "@server/db";
-import { orgs, resourceOtp, resources, resourceWhitelist, resourcePolicyWhiteList } from "@server/db";
+import {
+    orgs,
+    resourceOtp,
+    resources,
+    resourceWhitelist,
+    resourcePolicyWhiteList
+} from "@server/db";
 import HttpCode from "@server/types/HttpCode";
 import response from "@server/lib/response";
 import { eq, and } from "drizzle-orm";
@@ -84,15 +90,21 @@ export async function authWithWhitelist(
 
         const wildcard = "*@" + email.split("@")[1];
 
-        // Check policy whitelist first (policy takes precedence over resource whitelist)
-        let policyWhitelistEntry: { whitelistId: number; email: string } | null = null;
+        // Check shared policy whitelist first, then default (inline) policy whitelist
+        let policyWhitelistEntry: {
+            whitelistId: number;
+            email: string;
+        } | null = null;
         if (resource.resourcePolicyId) {
             const [exact] = await db
                 .select()
                 .from(resourcePolicyWhiteList)
                 .where(
                     and(
-                        eq(resourcePolicyWhiteList.resourcePolicyId, resource.resourcePolicyId),
+                        eq(
+                            resourcePolicyWhiteList.resourcePolicyId,
+                            resource.resourcePolicyId
+                        ),
                         eq(resourcePolicyWhiteList.email, email)
                     )
                 )
@@ -101,13 +113,57 @@ export async function authWithWhitelist(
             if (exact) {
                 policyWhitelistEntry = exact;
             } else {
-                logger.debug("Checking for wildcard email in policy: " + wildcard);
+                logger.debug(
+                    "Checking for wildcard email in shared policy: " + wildcard
+                );
                 const [wildcardMatch] = await db
                     .select()
                     .from(resourcePolicyWhiteList)
                     .where(
                         and(
-                            eq(resourcePolicyWhiteList.resourcePolicyId, resource.resourcePolicyId),
+                            eq(
+                                resourcePolicyWhiteList.resourcePolicyId,
+                                resource.resourcePolicyId
+                            ),
+                            eq(resourcePolicyWhiteList.email, wildcard)
+                        )
+                    )
+                    .limit(1);
+                if (wildcardMatch) policyWhitelistEntry = wildcardMatch;
+            }
+        }
+
+        // Fall back to default (inline) policy whitelist if shared policy didn't match
+        if (!policyWhitelistEntry && resource.defaultResourcePolicyId) {
+            const [exact] = await db
+                .select()
+                .from(resourcePolicyWhiteList)
+                .where(
+                    and(
+                        eq(
+                            resourcePolicyWhiteList.resourcePolicyId,
+                            resource.defaultResourcePolicyId
+                        ),
+                        eq(resourcePolicyWhiteList.email, email)
+                    )
+                )
+                .limit(1);
+
+            if (exact) {
+                policyWhitelistEntry = exact;
+            } else {
+                logger.debug(
+                    "Checking for wildcard email in default policy: " + wildcard
+                );
+                const [wildcardMatch] = await db
+                    .select()
+                    .from(resourcePolicyWhiteList)
+                    .where(
+                        and(
+                            eq(
+                                resourcePolicyWhiteList.resourcePolicyId,
+                                resource.defaultResourcePolicyId
+                            ),
                             eq(resourcePolicyWhiteList.email, wildcard)
                         )
                     )
@@ -117,7 +173,10 @@ export async function authWithWhitelist(
         }
 
         // Fall back to resource whitelist if not found in policy
-        let resourceWhitelistEntry: { whitelistId: number; email: string } | null = null;
+        let resourceWhitelistEntry: {
+            whitelistId: number;
+            email: string;
+        } | null = null;
         if (!policyWhitelistEntry) {
             const [exact] = await db
                 .select()
@@ -241,8 +300,12 @@ export async function authWithWhitelist(
         await createResourceSession({
             resourceId,
             token,
-            whitelistId: isPolicyWhitelist ? null : whitelistedEmail.whitelistId,
-            policyWhitelistId: isPolicyWhitelist ? whitelistedEmail.whitelistId : null,
+            whitelistId: isPolicyWhitelist
+                ? null
+                : whitelistedEmail.whitelistId,
+            policyWhitelistId: isPolicyWhitelist
+                ? whitelistedEmail.whitelistId
+                : null,
             isRequestToken: true,
             expiresAt: Date.now() + 1000 * 30, // 30 seconds
             sessionLength: 1000 * 30,

server/routers/resource/getResourceAuthInfo.ts

@@ -6,9 +6,13 @@ import {
     resourcePolicyHeaderAuth,
     resourcePolicyPassword,
     resourcePolicyPincode,
+    resourcePincode,
+    resourcePassword,
+    resourceHeaderAuth,
     resources
 } from "@server/db";
-import { eq, or } from "drizzle-orm";
+import { eq } from "drizzle-orm";
+import { alias } from "drizzle-orm/sqlite-core";
 import response from "@server/lib/response";
 import HttpCode from "@server/types/HttpCode";
 import createHttpError from "http-errors";
@@ -60,42 +64,103 @@ export async function getResourceAuthInfo(
 
         const isGuidInteger = /^\d+$/.test(resourceGuid);
 
+        const sharedPolicy = alias(resourcePolicies, "sharedPolicy");
+        const defaultPolicy = alias(resourcePolicies, "defaultPolicy");
+        const sharedPolicyPincode = alias(
+            resourcePolicyPincode,
+            "sharedPolicyPincode"
+        );
+        const defaultPolicyPincode = alias(
+            resourcePolicyPincode,
+            "defaultPolicyPincode"
+        );
+        const sharedPolicyPassword = alias(
+            resourcePolicyPassword,
+            "sharedPolicyPassword"
+        );
+        const defaultPolicyPassword = alias(
+            resourcePolicyPassword,
+            "defaultPolicyPassword"
+        );
+        const sharedPolicyHeaderAuth = alias(
+            resourcePolicyHeaderAuth,
+            "sharedPolicyHeaderAuth"
+        );
+        const defaultPolicyHeaderAuth = alias(
+            resourcePolicyHeaderAuth,
+            "defaultPolicyHeaderAuth"
+        );
+
         const buildQuery = (whereClause: ReturnType<typeof eq>) =>
             db
                 .select()
                 .from(resources)
                 .leftJoin(
-                    resourcePolicies,
-                    or(
-                        eq(
-                            resourcePolicies.resourcePolicyId,
-                            resources.resourcePolicyId
-                        ),
-                        eq(
-                            resourcePolicies.resourcePolicyId,
-                            resources.defaultResourcePolicyId
-                        )
+                    resourcePincode,
+                    eq(resourcePincode.resourceId, resources.resourceId)
+                )
+                .leftJoin(
+                    resourcePassword,
+                    eq(resourcePassword.resourceId, resources.resourceId)
+                )
+                .leftJoin(
+                    resourceHeaderAuth,
+                    eq(resourceHeaderAuth.resourceId, resources.resourceId)
+                )
+                .leftJoin(
+                    sharedPolicy,
+                    eq(
+                        sharedPolicy.resourcePolicyId,
+                        resources.resourcePolicyId
                     )
                 )
                 .leftJoin(
-                    resourcePolicyPincode,
+                    sharedPolicyPincode,
                     eq(
-                        resourcePolicyPincode.resourcePolicyId,
-                        resourcePolicies.resourcePolicyId
+                        sharedPolicyPincode.resourcePolicyId,
+                        sharedPolicy.resourcePolicyId
                     )
                 )
                 .leftJoin(
-                    resourcePolicyPassword,
+                    sharedPolicyPassword,
                     eq(
-                        resourcePolicyPassword.resourcePolicyId,
-                        resourcePolicies.resourcePolicyId
+                        sharedPolicyPassword.resourcePolicyId,
+                        sharedPolicy.resourcePolicyId
                     )
                 )
                 .leftJoin(
-                    resourcePolicyHeaderAuth,
+                    sharedPolicyHeaderAuth,
                     eq(
-                        resourcePolicyHeaderAuth.resourcePolicyId,
-                        resourcePolicies.resourcePolicyId
+                        sharedPolicyHeaderAuth.resourcePolicyId,
+                        sharedPolicy.resourcePolicyId
+                    )
+                )
+                .leftJoin(
+                    defaultPolicy,
+                    eq(
+                        defaultPolicy.resourcePolicyId,
+                        resources.defaultResourcePolicyId
+                    )
+                )
+                .leftJoin(
+                    defaultPolicyPincode,
+                    eq(
+                        defaultPolicyPincode.resourcePolicyId,
+                        defaultPolicy.resourcePolicyId
+                    )
+                )
+                .leftJoin(
+                    defaultPolicyPassword,
+                    eq(
+                        defaultPolicyPassword.resourcePolicyId,
+                        defaultPolicy.resourcePolicyId
+                    )
+                )
+                .leftJoin(
+                    defaultPolicyHeaderAuth,
+                    eq(
+                        defaultPolicyHeaderAuth.resourcePolicyId,
+                        defaultPolicy.resourcePolicyId
                     )
                 )
                 .where(whereClause)
@@ -115,10 +180,24 @@ export async function getResourceAuthInfo(
             );
         }
 
-        const policy = result?.resourcePolicies;
-        const pincode = result?.resourcePolicyPincode;
-        const password = result?.resourcePolicyPassword;
-        const headerAuth = result?.resourcePolicyHeaderAuth;
+        // Shared (custom) policy takes precedence over the default policy.
+        // For boolean fields (sso, whitelist), only fall back to defaultPolicy
+        // when there is no shared policy at all.
+        const effectivePolicyPincode =
+            result.sharedPolicyPincode ?? result.defaultPolicyPincode ?? null;
+        const effectivePolicyPassword =
+            result.sharedPolicyPassword ?? result.defaultPolicyPassword ?? null;
+        const effectivePolicyHeaderAuth =
+            result.sharedPolicyHeaderAuth ??
+            result.defaultPolicyHeaderAuth ??
+            null;
+
+        const effectivePolicy = result.sharedPolicy ?? result.defaultPolicy;
+
+        const pincode = effectivePolicyPincode ?? result.resourcePincode;
+        const password = effectivePolicyPassword ?? result.resourcePassword;
+        const headerAuth =
+            effectivePolicyHeaderAuth ?? result.resourceHeaderAuth;
 
         const url = resource.fullDomain
             ? `${resource.ssl ? "https" : "http"}://${resource.fullDomain}`
@@ -134,13 +213,13 @@ export async function getResourceAuthInfo(
                 pincode: pincode !== null,
                 headerAuth: headerAuth !== null,
                 headerAuthExtendedCompatibility:
-                    headerAuth?.extendedCompatibility ?? false,
-                sso: policy?.sso ?? false,
+                    effectivePolicyHeaderAuth?.extendedCompatibility ?? false,
+                sso: effectivePolicy?.sso ?? false,
                 blockAccess: resource.blockAccess,
                 url: url ?? "",
                 wildcard: resource.wildcard ?? false,
                 fullDomain: resource.fullDomain,
-                whitelist: policy?.emailWhitelistEnabled ?? false,
+                whitelist: effectivePolicy?.emailWhitelistEnabled ?? false,
                 skipToIdpId: resource.skipToIdpId,
                 orgId: resource.orgId,
                 postAuthPath: resource.postAuthPath ?? null

server/routers/ssh/types.ts

@@ -0,0 +1,16 @@
+export type SignSshKeyResponse = {
+    certificate?: string;
+    messageIds: number[];
+    messageId?: number;
+    sshUsername: string;
+    sshHost: string;
+    resourceId: number;
+    siteIds: number[];
+    siteId: number;
+    keyId?: string;
+    validPrincipals?: string[];
+    validAfter?: string;
+    validBefore?: string;
+    expiresIn?: number;
+    authDaemonMode: "site" | "remote" | "native" | null;
+};

src/app/ssh/SshClient.tsx

@@ -6,7 +6,6 @@ import { Button } from "@/components/ui/button";
 import { Input } from "@/components/ui/input";
 import { Label } from "@/components/ui/label";
 import { Textarea } from "@/components/ui/textarea";
-import type { SignSshKeyResponse } from "@server/private/routers/ssh";
 import { GetBrowserTargetResponse } from "@server/routers/browserGatewayTarget";
 import {
     Card,
@@ -18,6 +17,7 @@ import {
 import Link from "next/link";
 import { ExternalLink } from "lucide-react";
 import { cn } from "@app/lib/cn";
+import type { SignSshKeyResponse } from "@server/routers/ssh/types";
 
 type AuthTab = "password" | "privateKey";
 
@@ -337,15 +337,6 @@ export default function SshClient({
                 )}
                 {connected && (
                     <div className="fixed inset-0 z-50 flex flex-col bg-neutral-900">
-                        <div className="flex flex-wrap items-center gap-2 bg-black p-2 text-white">
-                            <Button
-                                size="sm"
-                                variant="destructive"
-                                onClick={disconnect}
-                            >
-                                Terminate
-                            </Button>
-                        </div>
                         <div
                             ref={terminalRef}
                             className="flex-1 overflow-hidden"

src/app/ssh/page.tsx

@@ -3,9 +3,65 @@ import { priv } from "@app/lib/api";
 import { AxiosResponse } from "axios";
 import { GetBrowserTargetResponse } from "@server/routers/browserGatewayTarget";
 import SshClient from "./SshClient";
-import { SignSshKeyResponse } from "@server/private/routers/ssh";
 import crypto from "crypto";
 import AuthFooter from "@app/components/AuthFooter";
+import type { SignSshKeyResponse } from "@server/routers/ssh/types";
+
+const pollInitialDelayMs = 250;
+const pollStartIntervalMs = 250;
+const pollBackoffSteps = 6;
+
+type RoundTripMessageResponse = {
+    messageId: number;
+    complete: boolean;
+    sentAt: number | string;
+    receivedAt: number | string | null;
+    error: string | null;
+};
+
+function sleep(ms: number): Promise<void> {
+    return new Promise((resolve) => setTimeout(resolve, ms));
+}
+
+async function waitForRoundTripCompletion(
+    messageIds: number[],
+    cookieHeader: string
+): Promise<void> {
+    if (messageIds.length === 0) {
+        return;
+    }
+
+    await sleep(pollInitialDelayMs);
+
+    let interval = pollStartIntervalMs;
+    for (let i = 0; i <= pollBackoffSteps; i++) {
+        for (const messageId of messageIds) {
+            const res = await priv.get<AxiosResponse<RoundTripMessageResponse>>(
+                `/ws/round-trip-message/${messageId}`,
+                {
+                    headers: {
+                        Cookie: cookieHeader
+                    }
+                }
+            );
+
+            const message = res.data.data;
+            if (message.complete) {
+                if (message.error) {
+                    throw new Error(message.error);
+                }
+                return;
+            }
+        }
+
+        if (i < pollBackoffSteps) {
+            await sleep(interval);
+            interval *= 2;
+        }
+    }
+
+    throw new Error("Timed out waiting for round-trip message completion");
+}
 
 function generateEphemeralKeyPair(): {
     privateKeyPem: string;
@@ -51,6 +107,7 @@ export default async function SshPage() {
     const headersList = await headers();
     const host = headersList.get("host") || "";
     const hostname = host.split(":")[0];
+    const cookieHeader = headersList.get("cookie") || "";
 
     let target: GetBrowserTargetResponse | null = null;
     let signedKeyData: SignSshKeyResponse | null = null;
@@ -72,14 +129,29 @@ export default async function SshPage() {
                     `/org/${target.orgId}/ssh/sign-key`,
                     {
                         publicKey: publicKeyOpenSSH,
-                        resource: target.niceId
+                        resourceId: target.resourceId,
+                        type: "public"
+                    },
+                    {
+                        headers: {
+                            Cookie: cookieHeader
+                        }
                     }
                 );
                 signedKeyData = res.data.data;
-                console.log("Received signed SSH key:", signedKeyData);
+
+                const messageIds =
+                    signedKeyData.messageIds.length > 0
+                        ? signedKeyData.messageIds
+                        : signedKeyData.messageId
+                          ? [signedKeyData.messageId]
+                          : [];
+
+                await waitForRoundTripCompletion(messageIds, cookieHeader);
             } catch (err) {
                 console.error("Error signing SSH key:", err);
-                error = "Failed to sign SSH key for PAM push authentication.";
+                error =
+                    "Failed to sign SSH key for PAM push authentication. Did you sign in as a user?";
             }
         }
     } catch (err) {

src/components/ColumnMultiFilterButton.tsx

@@ -20,6 +20,7 @@ import { CheckIcon, Funnel } from "lucide-react";
 import { cn } from "@app/lib/cn";
 import { dataTableFilterPopoverContentClassName } from "@app/lib/dataTableFilterPopover";
 import { Badge } from "./ui/badge";
+import { Checkbox } from "./ui/checkbox";
 
 type FilterOption = {
     value: string;
@@ -130,13 +131,11 @@ export function ColumnMultiFilterButton({
                                         toggle(option.value);
                                     }}
                                 >
-                                    <CheckIcon
-                                        className={cn(
-                                            "mr-2 h-4 w-4",
-                                            selectedSet.has(option.value)
-                                                ? "opacity-100"
-                                                : "opacity-0"
-                                        )}
+                                    <Checkbox
+                                        className="pointer-events-none shrink-0"
+                                        checked={selectedSet.has(option.value)}
+                                        aria-hidden
+                                        tabIndex={-1}
                                     />
                                     {option.label}
                                 </CommandItem>

src/components/LabelColumnFilterButton.tsx

@@ -25,6 +25,7 @@ import { useDebounce } from "use-debounce";
 import { LabelBadge } from "./label-badge";
 import { LabelOverflowBadge } from "./label-overflow-badge";
 import { LABEL_COLORS } from "./labels-selector";
+import { Checkbox } from "./ui/checkbox";
 
 function areSelectionsEqual(a: string[], b: string[]) {
     if (a.length !== b.length) {
@@ -179,13 +180,11 @@ export function LabelColumnFilterButton({
                                         }}
                                         className="flex items-center gap-2"
                                     >
-                                        <CheckIcon
-                                            className={cn(
-                                                "mr-2 h-4 w-4",
-                                                draftSet.has(label.name)
-                                                    ? "opacity-100"
-                                                    : "opacity-0"
-                                            )}
+                                        <Checkbox
+                                            className="pointer-events-none shrink-0"
+                                            checked={draftSet.has(label.name)}
+                                            aria-hidden
+                                            tabIndex={-1}
                                         />
                                         <div
                                             className="size-2 rounded-full bg-(--color) flex-none"

src/components/LabelsTableCell.tsx

@@ -21,29 +21,32 @@ const MAX_VISIBLE_BEFORE_OVERFLOW = MAX_VISIBLE_LABELS - 1;
 
 type TableLabelsCellProps = {
     orgId: string;
-    localLabels: SelectedLabel[];
-    toggleLabel: (label: SelectedLabel, action: "attach" | "detach") => void;
+    selectedLabels: SelectedLabel[];
+    onToggleLabel: (label: SelectedLabel, action: "attach" | "detach") => void;
+    onClosePopover: () => void;
 };
 
-export function TableLabelsCell({
+export function LabelsTableCell({
     orgId,
-    localLabels,
-    toggleLabel
+    selectedLabels,
+    onToggleLabel,
+    onClosePopover
 }: TableLabelsCellProps) {
     const t = useTranslations();
     const [isPopoverOpen, setIsPopoverOpen] = useState(false);
+
     const triggerRef = useRef<HTMLButtonElement>(null);
     const frozenAnchorRef = useRef<Measurable>({
         getBoundingClientRect: () => new DOMRect()
     });
 
-    const hasOverflow = localLabels.length > MAX_VISIBLE_LABELS;
-    const visibleLabels = localLabels.slice(
+    const hasOverflow = selectedLabels.length > MAX_VISIBLE_LABELS;
+    const visibleLabels = selectedLabels.slice(
         0,
         hasOverflow ? MAX_VISIBLE_BEFORE_OVERFLOW : MAX_VISIBLE_LABELS
     );
     const overflowLabels = hasOverflow
-        ? localLabels.slice(MAX_VISIBLE_BEFORE_OVERFLOW)
+        ? selectedLabels.slice(MAX_VISIBLE_BEFORE_OVERFLOW)
         : [];
 
     function handleOpenChange(open: boolean) {
@@ -54,10 +57,14 @@ export function TableLabelsCell({
             };
         }
         setIsPopoverOpen(open);
+
+        if (!open) {
+            onClosePopover();
+        }
     }
 
     return (
-        <div className="grid w-full min-w-0 grid-cols-[auto_minmax(0,1fr)] items-center gap-1">
+        <div className="flex items-center gap-1">
             <Popover open={isPopoverOpen} onOpenChange={handleOpenChange}>
                 <PopoverAnchor virtualRef={frozenAnchorRef} />
                 <PopoverTrigger asChild>
@@ -80,9 +87,8 @@ export function TableLabelsCell({
                 >
                     <LabelsSelector
                         orgId={orgId}
-                        selectedLabels={localLabels}
-                        toggleLabel={toggleLabel}
-                        onClose={() => handleOpenChange(false)}
+                        selectedLabels={selectedLabels}
+                        toggleLabel={onToggleLabel}
                     />
                 </PopoverContent>
             </Popover>

src/components/MachineClientsTable.tsx

@@ -34,11 +34,12 @@ import { useDebouncedCallback } from "use-debounce";
 import z from "zod";
 import { ColumnFilterButton } from "./ColumnFilterButton";
 import { type SelectedLabel } from "./labels-selector";
-import { TableLabelsCell } from "./TableLabelsCell";
+import { LabelsTableCell } from "./LabelsTableCell";
 import { Badge } from "./ui/badge";
 import { ControlledDataTable } from "./ui/controlled-data-table";
 import { LabelColumnFilterButton } from "./LabelColumnFilterButton";
 import { useLocalLabels } from "@app/hooks/useLocalLabels";
+import { useOptimisticLabels } from "@app/hooks/useOptimisticLabels";
 
 export type ClientRow = {
     id: number;
@@ -607,54 +608,19 @@ function MachineClientLabelCell({
     client,
     orgId
 }: MachineClientLabelCellProps) {
-    const t = useTranslations();
-    const api = createApiClient(useEnvContext());
-    const [localLabels, setLocalLabels] = useLocalLabels(
-        client.labels,
-        client.id
-    );
-
-    function toggleClientLabel(
-        label: SelectedLabel,
-        action: "attach" | "detach"
-    ) {
-        const previousLabels = localLabels;
-
-        void (async () => {
-            try {
-                if (action === "attach") {
-                    setLocalLabels([...previousLabels, label]);
-                    await api.put(
-                        `/org/${orgId}/label/${label.labelId}/attach`,
-                        { clientId: client.id }
-                    );
-                } else {
-                    setLocalLabels(
-                        previousLabels.filter(
-                            (lb) => lb.labelId !== label.labelId
-                        )
-                    );
-                    await api.put(
-                        `/org/${orgId}/label/${label.labelId}/detach`,
-                        { clientId: client.id }
-                    );
-                }
-            } catch (e) {
-                setLocalLabels(previousLabels);
-                toast({
-                    title: t("error"),
-                    description: formatAxiosError(e, t("errorOccurred")),
-                    variant: "destructive"
-                });
-            }
-        })();
-    }
+    const { localLabels, refresh, toggleLabel } = useOptimisticLabels({
+        serverLabels: client.labels,
+        orgId,
+        entityId: client.id,
+        entityIdField: "clientId"
+    });
 
     return (
-        <TableLabelsCell
+        <LabelsTableCell
             orgId={orgId}
-            localLabels={localLabels}
-            toggleLabel={toggleClientLabel}
+            selectedLabels={localLabels}
+            onToggleLabel={toggleLabel}
+            onClosePopover={() => startTransition(refresh)}
         />
     );
 }

src/components/PrivateResourcesTable.tsx

@@ -61,9 +61,10 @@ import { build } from "@server/build";
 import { usePaidStatus } from "@app/hooks/usePaidStatus";
 import { tierMatrix } from "@server/lib/billing/tierMatrix";
 import { type SelectedLabel } from "./labels-selector";
-import { TableLabelsCell } from "./TableLabelsCell";
+import { LabelsTableCell } from "./LabelsTableCell";
 import { LabelColumnFilterButton } from "./LabelColumnFilterButton";
 import { useLocalLabels } from "@app/hooks/useLocalLabels";
+import { useOptimisticLabels } from "@app/hooks/useOptimisticLabels";
 
 export type InternalResourceSiteRow = ResourceSiteRow;
 
@@ -705,54 +706,19 @@ function ClientResourceLabelCell({
     resource,
     orgId
 }: ClientResourceLabelCellProps) {
-    const t = useTranslations();
-    const api = createApiClient(useEnvContext());
-    const [localLabels, setLocalLabels] = useLocalLabels(
-        resource.labels,
-        resource.id
-    );
-
-    function toggleResourceLabel(
-        label: SelectedLabel,
-        action: "attach" | "detach"
-    ) {
-        const previousLabels = localLabels;
-
-        void (async () => {
-            try {
-                if (action === "attach") {
-                    setLocalLabels([...previousLabels, label]);
-                    await api.put(
-                        `/org/${orgId}/label/${label.labelId}/attach`,
-                        { siteResourceId: resource.id }
-                    );
-                } else {
-                    setLocalLabels(
-                        previousLabels.filter(
-                            (lb) => lb.labelId !== label.labelId
-                        )
-                    );
-                    await api.put(
-                        `/org/${orgId}/label/${label.labelId}/detach`,
-                        { siteResourceId: resource.id }
-                    );
-                }
-            } catch (e) {
-                setLocalLabels(previousLabels);
-                toast({
-                    title: t("error"),
-                    description: formatAxiosError(e, t("errorOccurred")),
-                    variant: "destructive"
-                });
-            }
-        })();
-    }
+    const { localLabels, refresh, toggleLabel } = useOptimisticLabels({
+        serverLabels: resource.labels,
+        orgId,
+        entityId: resource.id,
+        entityIdField: "siteResourceId"
+    });
 
     return (
-        <TableLabelsCell
+        <LabelsTableCell
             orgId={orgId}
-            localLabels={localLabels}
-            toggleLabel={toggleResourceLabel}
+            onClosePopover={() => startTransition(refresh)}
+            onToggleLabel={toggleLabel}
+            selectedLabels={localLabels}
         />
     );
 }

src/components/ProxyResourcesTable.tsx

@@ -73,7 +73,9 @@ import UptimeMiniBar from "./UptimeMiniBar";
 import { type SelectedLabel } from "./labels-selector";
 import { LabelColumnFilterButton } from "./LabelColumnFilterButton";
 import { useLocalLabels } from "@app/hooks/useLocalLabels";
-import { TableLabelsCell } from "./TableLabelsCell";
+import { LabelsTableCell } from "./LabelsTableCell";
+import { useOptimisticLabels } from "@app/hooks/useOptimisticLabels";
+import { refresh } from "next/cache";
 
 export type TargetHealth = {
     targetId: number;
@@ -772,57 +774,19 @@ type ResourceLabelCellProps = {
 };
 
 function ResourceLabelCell({ resource, orgId }: ResourceLabelCellProps) {
-    const t = useTranslations();
-
-    const api = createApiClient(useEnvContext());
-
-    const [localLabels, setLocalLabels] = useLocalLabels(
-        resource.labels,
-        resource.id
-    );
-
-    function toggleSiteLabel(
-        label: SelectedLabel,
-        action: "attach" | "detach"
-    ) {
-        const previousLabels = localLabels;
-
-        void (async () => {
-            try {
-                if (action === "attach") {
-                    setLocalLabels([...previousLabels, label]);
-
-                    await api.put(
-                        `/org/${orgId}/label/${label.labelId}/attach`,
-                        { resourceId: resource.id }
-                    );
-                } else {
-                    setLocalLabels(
-                        previousLabels.filter(
-                            (lb) => lb.labelId !== label.labelId
-                        )
-                    );
-                    await api.put(
-                        `/org/${orgId}/label/${label.labelId}/detach`,
-                        { resourceId: resource.id }
-                    );
-                }
-            } catch (e) {
-                setLocalLabels(previousLabels);
-                toast({
-                    title: t("error"),
-                    description: formatAxiosError(e, t("errorOccurred")),
-                    variant: "destructive"
-                });
-            }
-        })();
-    }
+    const { localLabels, refresh, toggleLabel } = useOptimisticLabels({
+        serverLabels: resource.labels,
+        orgId,
+        entityId: resource.id,
+        entityIdField: "resourceId"
+    });
 
     return (
-        <TableLabelsCell
+        <LabelsTableCell
             orgId={orgId}
-            localLabels={localLabels}
-            toggleLabel={toggleSiteLabel}
+            selectedLabels={localLabels}
+            onToggleLabel={toggleLabel}
+            onClosePopover={() => startTransition(refresh)}
         />
     );
 }

src/components/SitesTable.tsx

@@ -41,13 +41,7 @@ import {
 import { useTranslations } from "next-intl";
 import Link from "next/link";
 import { usePathname, useRouter } from "next/navigation";
-import {
-    startTransition,
-    useEffect,
-    useMemo,
-    useState,
-    useTransition
-} from "react";
+import { startTransition, useMemo, useState, useTransition } from "react";
 import { useDebouncedCallback } from "use-debounce";
 import z from "zod";
 import { ColumnFilterButton } from "./ColumnFilterButton";
@@ -56,13 +50,11 @@ import {
     type ExtendedColumnDef
 } from "./ui/controlled-data-table";
 
+import { useOptimisticLabels } from "@app/hooks/useOptimisticLabels";
 import { usePaidStatus } from "@app/hooks/usePaidStatus";
-import { cn } from "@app/lib/cn";
 import { tierMatrix } from "@server/lib/billing/tierMatrix";
-import { type SelectedLabel } from "./labels-selector";
 import { LabelColumnFilterButton } from "./LabelColumnFilterButton";
-import { useLocalLabels } from "@app/hooks/useLocalLabels";
-import { TableLabelsCell } from "./TableLabelsCell";
+import { LabelsTableCell } from "./LabelsTableCell";
 
 export type SiteRow = {
     id: number;
@@ -686,54 +678,19 @@ type SiteLabelCellProps = {
 };
 
 function SiteLabelCell({ site, orgId }: SiteLabelCellProps) {
-    const t = useTranslations();
-
-    const api = createApiClient(useEnvContext());
-
-    const [localLabels, setLocalLabels] = useLocalLabels(site.labels, site.id);
-
-    function toggleSiteLabel(
-        label: SelectedLabel,
-        action: "attach" | "detach"
-    ) {
-        const previousLabels = localLabels;
-
-        void (async () => {
-            try {
-                if (action === "attach") {
-                    setLocalLabels([...previousLabels, label]);
-
-                    await api.put(
-                        `/org/${orgId}/label/${label.labelId}/attach`,
-                        { siteId: site.id }
-                    );
-                } else {
-                    setLocalLabels(
-                        previousLabels.filter(
-                            (lb) => lb.labelId !== label.labelId
-                        )
-                    );
-                    await api.put(
-                        `/org/${orgId}/label/${label.labelId}/detach`,
-                        { siteId: site.id }
-                    );
-                }
-            } catch (e) {
-                setLocalLabels(previousLabels);
-                toast({
-                    title: t("error"),
-                    description: formatAxiosError(e, t("errorOccurred")),
-                    variant: "destructive"
-                });
-            }
-        })();
-    }
+    const { localLabels, refresh, toggleLabel } = useOptimisticLabels({
+        serverLabels: site.labels,
+        orgId,
+        entityId: site.id,
+        entityIdField: "siteId"
+    });
 
     return (
-        <TableLabelsCell
+        <LabelsTableCell
             orgId={orgId}
-            localLabels={localLabels}
-            toggleLabel={toggleSiteLabel}
+            selectedLabels={localLabels}
+            onToggleLabel={toggleLabel}
+            onClosePopover={() => startTransition(refresh)}
         />
     );
 }

src/components/labels-selector.tsx

@@ -36,7 +36,6 @@ export type LabelsSelectorProps = {
     orgId: string;
     selectedLabels: SelectedLabel[];
     toggleLabel: (newlabel: SelectedLabel, action: "detach" | "attach") => void;
-    onClose?: () => void;
 };
 
 export const LABEL_COLORS = {
@@ -52,8 +51,7 @@ export const LABEL_COLORS = {
 export function LabelsSelector({
     orgId,
     selectedLabels,
-    toggleLabel,
-    onClose
+    toggleLabel
 }: LabelsSelectorProps) {
     const t = useTranslations();
     const [labelSearchQuery, setlabelsSearchQuery] = useState("");
@@ -202,7 +200,6 @@ export function LabelsSelector({
                                         ? "detach"
                                         : "attach"
                                 );
-                                onClose?.();
                             }}
                         >
                             <Checkbox

src/components/multi-select/multi-select-content.tsx

@@ -10,6 +10,7 @@ import {
 import { cn } from "@app/lib/cn";
 import { CheckIcon } from "lucide-react";
 import { useTranslations } from "next-intl";
+import { Checkbox } from "../ui/checkbox";
 
 export type TagValue = { text: string; id: string; isAdmin?: boolean };
 
@@ -70,13 +71,11 @@ export function MultiSelectContent<T extends TagValue>({
                                     onChange(newValues);
                                 }}
                             >
-                                <CheckIcon
-                                    className={cn(
-                                        "mr-2 h-4 w-4",
-                                        selectedValues.has(option.id)
-                                            ? "opacity-100"
-                                            : "opacity-0"
-                                    )}
+                                <Checkbox
+                                    className="pointer-events-none shrink-0"
+                                    checked={selectedValues.has(option.id)}
+                                    aria-hidden
+                                    tabIndex={-1}
                                 />
                                 {`${option.text}`}
                             </CommandItem>

src/hooks/useNavigationContext.ts

@@ -1,5 +1,5 @@
 import { useSearchParams, usePathname, useRouter } from "next/navigation";
-import { useTransition } from "react";
+import { useCallback, useMemo, useTransition } from "react";
 
 export function useNavigationContext() {
     const router = useRouter();
@@ -7,29 +7,38 @@ export function useNavigationContext() {
     const path = usePathname();
     const [isNavigating, startTransition] = useTransition();
 
-    function navigate({
-        searchParams: params,
-        pathname = path,
-        replace = false
-    }: {
-        pathname?: string;
-        searchParams?: URLSearchParams;
-        replace?: boolean;
-    }) {
-        startTransition(() => {
-            const fullPath = pathname + (params ? `?${params.toString()}` : "");
+    const navigate = useCallback(
+        function ({
+            searchParams: params,
+            pathname = path,
+            replace = false
+        }: {
+            pathname?: string;
+            searchParams?: URLSearchParams;
+            replace?: boolean;
+        }) {
+            startTransition(() => {
+                const fullPath =
+                    pathname + (params ? `?${params.toString()}` : "");
 
-            if (replace) {
-                router.replace(fullPath);
-            } else {
-                router.push(fullPath);
-            }
-        });
-    }
+                if (replace) {
+                    router.replace(fullPath);
+                } else {
+                    router.push(fullPath);
+                }
+            });
+        },
+        [router]
+    );
+
+    const writableSearchParams = useMemo(
+        () => new URLSearchParams(searchParams),
+        [searchParams]
+    );
 
     return {
         pathname: path,
-        searchParams: new URLSearchParams(searchParams), // we want the search params to be writeable
+        searchParams: writableSearchParams,
         navigate,
         isNavigating
     };

src/hooks/useOptimisticLabels.ts

@@ -0,0 +1,103 @@
+import type { SelectedLabel } from "@app/components/labels-selector";
+import { createApiClient, formatAxiosError } from "@app/lib/api";
+import { useState, useMemo } from "react";
+import { toast } from "./useToast";
+import { useEnvContext } from "./useEnvContext";
+import { useTranslations } from "next-intl";
+import { useRouter } from "next/navigation";
+
+export type LabelToggleAction = {
+    label: SelectedLabel;
+    action: "attach" | "detach";
+};
+
+function computeLabelToggleActions(
+    values: SelectedLabel[],
+    actions: LabelToggleAction[]
+) {
+    let newValues = [...values];
+    for (const { action, label } of actions) {
+        if (action === "attach") {
+            newValues = [...newValues, label];
+        } else {
+            newValues = newValues.filter((lb) => lb.labelId !== label.labelId);
+        }
+    }
+
+    return newValues;
+}
+
+type UseOptimisticLabelsArgs = {
+    serverLabels: SelectedLabel[] | undefined;
+    orgId: string;
+    entityId: number;
+    entityIdField: string;
+};
+
+export function useOptimisticLabels({
+    serverLabels,
+    orgId,
+    entityId,
+    entityIdField
+}: UseOptimisticLabelsArgs) {
+    const router = useRouter();
+    const labels = serverLabels ?? [];
+    const api = createApiClient(useEnvContext());
+    const t = useTranslations();
+
+    const [pendingActions, setPendingActions] = useState<LabelToggleAction[]>(
+        []
+    );
+
+    const localLabels = useMemo(
+        () => computeLabelToggleActions(labels ?? [], pendingActions),
+        [labels, pendingActions]
+    );
+
+    async function toggleLabel(
+        label: SelectedLabel,
+        action: "attach" | "detach"
+    ) {
+        const oppositeAction = action === "attach" ? "detach" : "attach";
+        const existingActionIndex = pendingActions.findIndex(
+            (pending) =>
+                pending.action === oppositeAction &&
+                pending.label.labelId === label.labelId
+        );
+
+        // if there are two actions that cancel each-other
+        // they should just be removed
+        if (existingActionIndex !== -1) {
+            setPendingActions((prevActions) =>
+                prevActions.toSpliced(existingActionIndex, 1)
+            );
+        } else {
+            setPendingActions((actions) => [...actions, { label, action }]);
+        }
+
+        try {
+            if (action === "attach") {
+                await api.put(`/org/${orgId}/label/${label.labelId}/attach`, {
+                    [entityIdField]: entityId
+                });
+            } else {
+                await api.put(`/org/${orgId}/label/${label.labelId}/detach`, {
+                    [entityIdField]: entityId
+                });
+            }
+        } catch (e) {
+            toast({
+                title: t("error"),
+                description: formatAxiosError(e, t("errorOccurred")),
+                variant: "destructive"
+            });
+        }
+    }
+
+    async function refresh() {
+        router.refresh();
+        setPendingActions([]);
+    }
+
+    return { localLabels, toggleLabel, refresh };
+}