Skip to content

Commit

Permalink
feat(auth): Add admin & early access flags based on orgs membership (#…
Browse files Browse the repository at this point in the history
  • Loading branch information
nsarrazin authored Jul 29, 2024
1 parent 24f7202 commit a3f5e2a
Show file tree
Hide file tree
Showing 5 changed files with 35 additions and 1 deletion.
3 changes: 3 additions & 0 deletions .env
Original file line number Diff line number Diff line change
Expand Up @@ -164,3 +164,6 @@ METRICS_ENABLED=false
METRICS_PORT=5565
LOG_LEVEL=info
BODY_SIZE_LIMIT=15728640

HF_ORG_ADMIN=
HF_ORG_EARLY_ACCESS=
2 changes: 2 additions & 0 deletions chart/env/prod.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -328,6 +328,8 @@ envVars:
}]
WEBSEARCH_BLOCKLIST: '["youtube.com", "twitter.com"]'
XFF_DEPTH: '2'
HF_ORG_ADMIN: '644171cfbd0c97265298aa99'
HF_ORG_EARLY_ACCESS: '5e67bd5b1009063689407478'

infisical:
enabled: true
Expand Down
1 change: 1 addition & 0 deletions src/lib/types/User.ts
Original file line number Diff line number Diff line change
Expand Up @@ -10,4 +10,5 @@ export interface User extends Timestamps {
avatarUrl: string | undefined;
hfUserId: string;
isAdmin?: boolean;
isEarlyAccess?: boolean;
}
1 change: 1 addition & 0 deletions src/routes/+layout.server.ts
Original file line number Diff line number Diff line change
Expand Up @@ -197,6 +197,7 @@ export const load: LayoutServerLoad = async ({ locals, depends }) => {
email: locals.user.email,
logoutDisabled: locals.user.logoutDisabled,
isAdmin: locals.user.isAdmin ?? false,
isEarlyAccess: locals.user.isEarlyAccess ?? false,
},
assistant,
enableAssistants,
Expand Down
29 changes: 28 additions & 1 deletion src/routes/login/callback/updateUser.ts
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@ import crypto from "crypto";
import { sha256 } from "$lib/utils/sha256";
import { addWeeks } from "date-fns";
import { OIDConfig } from "$lib/server/auth";
import { HF_ORG_ADMIN, HF_ORG_EARLY_ACCESS } from "$env/static/private";

export async function updateUser(params: {
userData: UserinfoResponse;
Expand All @@ -31,13 +32,25 @@ export async function updateUser(params: {
email,
picture: avatarUrl,
sub: hfUserId,
orgs,
} = z
.object({
preferred_username: z.string().optional(),
name: z.string(),
picture: z.string().optional(),
sub: z.string(),
email: z.string().email().optional(),
orgs: z
.array(
z.object({
sub: z.string(),
name: z.string(),
picture: z.string(),
preferred_username: z.string(),
isEnterprise: z.boolean(),
})
)
.optional(),
})
.setKey(OIDConfig.NAME_CLAIM, z.string())
.refine((data) => data.preferred_username || data.email, {
Expand All @@ -53,11 +66,23 @@ export async function updateUser(params: {
picture?: string;
sub: string;
name: string;
orgs?: Array<{
sub: string;
name: string;
picture: string;
preferred_username: string;
isEnterprise: boolean;
}>;
} & Record<string, string>;

// Dynamically access user data based on NAME_CLAIM from environment
// This approach allows us to adapt to different OIDC providers flexibly.

// if using huggingface as auth provider, check orgs for earl access and amin rights
const isAdmin = (HF_ORG_ADMIN && orgs?.some((org) => org.sub === HF_ORG_ADMIN)) || false;
const isEarlyAccess =
(HF_ORG_EARLY_ACCESS && orgs?.some((org) => org.sub === HF_ORG_EARLY_ACCESS)) || false;

// check if user already exists
const existingUser = await collections.users.findOne({ hfUserId });
let userId = existingUser?._id;
Expand All @@ -77,7 +102,7 @@ export async function updateUser(params: {
// update existing user if any
await collections.users.updateOne(
{ _id: existingUser._id },
{ $set: { username, name, avatarUrl } }
{ $set: { username, name, avatarUrl, isAdmin, isEarlyAccess } }
);

// remove previous session if it exists and add new one
Expand All @@ -103,6 +128,8 @@ export async function updateUser(params: {
email,
avatarUrl,
hfUserId,
isAdmin,
isEarlyAccess,
});

userId = insertedId;
Expand Down

0 comments on commit a3f5e2a

Please sign in to comment.