Refactor chat permission check and remove unused organization role check

Co-authored-by: dallin <dallin@buster.so>
This commit is contained in:
Cursor Agent 2025-07-04 14:55:08 +00:00
parent 7275f242ee
commit 6fb546a4ad
2 changed files with 165 additions and 195 deletions

View File

@ -1,7 +1,7 @@
import { canUserAccessChatCached } from '@buster/access-controls';
import { import {
type User, type User,
chats, chats,
checkChatPermission,
createMessage, createMessage,
db, db,
generateAssetMessages, generateAssetMessages,
@ -9,7 +9,6 @@ import {
getMessagesForChat, getMessagesForChat,
messages, messages,
} from '@buster/database'; } from '@buster/database';
import { eq, and, gte, isNull } from 'drizzle-orm';
import type { Chat, Message } from '@buster/database'; import type { Chat, Message } from '@buster/database';
import type { import type {
ChatMessage, ChatMessage,
@ -23,6 +22,7 @@ import {
ReasoningMessageSchema, ReasoningMessageSchema,
ResponseMessageSchema, ResponseMessageSchema,
} from '@buster/server-shared/chats'; } from '@buster/server-shared/chats';
import { and, eq, gte, isNull } from 'drizzle-orm';
// Optimized: Generic function to handle both response and reasoning messages // Optimized: Generic function to handle both response and reasoning messages
function buildMessages<T extends { id: string }>( function buildMessages<T extends { id: string }>(
@ -188,7 +188,10 @@ export async function handleExistingChat(
throw new ChatError(ChatErrorCode.CHAT_NOT_FOUND, 'Chat not found', 404); throw new ChatError(ChatErrorCode.CHAT_NOT_FOUND, 'Chat not found', 404);
} }
const hasPermission = await checkChatPermission(chatId, user.id); const hasPermission = await canUserAccessChatCached({
userId: user.id,
chatId,
});
if (!hasPermission) { if (!hasPermission) {
throw new ChatError( throw new ChatError(
ChatErrorCode.PERMISSION_DENIED, ChatErrorCode.PERMISSION_DENIED,

View File

@ -2,7 +2,7 @@ import { and, eq, isNull } from 'drizzle-orm';
import type { InferSelectModel } from 'drizzle-orm'; import type { InferSelectModel } from 'drizzle-orm';
import { z } from 'zod'; import { z } from 'zod';
import { db } from '../connection'; import { db } from '../connection';
import { chats, messages, userFavorites, users, usersToOrganizations } from '../schema'; import { chats, messages, userFavorites, users } from '../schema';
// Type inference from schema // Type inference from schema
export type Chat = InferSelectModel<typeof chats>; export type Chat = InferSelectModel<typeof chats>;
@ -11,28 +11,28 @@ export type User = InferSelectModel<typeof users>;
// Create a type for updateable chat fields by excluding auto-managed fields // Create a type for updateable chat fields by excluding auto-managed fields
type UpdateableChatFields = Partial< type UpdateableChatFields = Partial<
Omit<typeof chats.$inferInsert, 'id' | 'createdAt' | 'deletedAt'> Omit<typeof chats.$inferInsert, 'id' | 'createdAt' | 'deletedAt'>
>; >;
/** /**
* Input/output schemas for type safety * Input/output schemas for type safety
*/ */
export const CreateChatInputSchema = z.object({ export const CreateChatInputSchema = z.object({
title: z.string().min(1), title: z.string().min(1),
userId: z.string().uuid(), userId: z.string().uuid(),
organizationId: z.string().uuid(), organizationId: z.string().uuid(),
}); });
export const GetChatInputSchema = z.object({ export const GetChatInputSchema = z.object({
chatId: z.string().uuid(), chatId: z.string().uuid(),
userId: z.string().uuid(), userId: z.string().uuid(),
}); });
export const CreateMessageInputSchema = z.object({ export const CreateMessageInputSchema = z.object({
chatId: z.string().uuid(), chatId: z.string().uuid(),
content: z.string(), content: z.string(),
userId: z.string().uuid(), userId: z.string().uuid(),
messageId: z.string().uuid().optional(), messageId: z.string().uuid().optional(),
}); });
export type CreateChatInput = z.infer<typeof CreateChatInputSchema>; export type CreateChatInput = z.infer<typeof CreateChatInputSchema>;
@ -43,188 +43,155 @@ export type CreateMessageInput = z.infer<typeof CreateMessageInputSchema>;
* Create a new chat * Create a new chat
*/ */
export async function createChat(input: CreateChatInput): Promise<Chat> { export async function createChat(input: CreateChatInput): Promise<Chat> {
try { try {
const validated = CreateChatInputSchema.parse(input); const validated = CreateChatInputSchema.parse(input);
const [chat] = await db const [chat] = await db
.insert(chats) .insert(chats)
.values({ .values({
title: validated.title, title: validated.title,
organizationId: validated.organizationId, organizationId: validated.organizationId,
createdBy: validated.userId, createdBy: validated.userId,
updatedBy: validated.userId, updatedBy: validated.userId,
publiclyAccessible: false, publiclyAccessible: false,
}) })
.returning(); .returning();
if (!chat) { if (!chat) {
throw new Error('Failed to create chat'); throw new Error('Failed to create chat');
} }
return chat; return chat;
} catch (error) { } catch (error) {
if (error instanceof z.ZodError) { if (error instanceof z.ZodError) {
throw new Error(`Invalid chat input: ${error.errors.map((e) => e.message).join(', ')}`); throw new Error(`Invalid chat input: ${error.errors.map((e) => e.message).join(', ')}`);
} }
throw error; throw error;
} }
} }
/** /**
* Get a chat with user and favorite information * Get a chat with user and favorite information
*/ */
export async function getChatWithDetails(input: GetChatInput): Promise<{ export async function getChatWithDetails(input: GetChatInput): Promise<{
chat: Chat; chat: Chat;
user: User | null; user: User | null;
isFavorited: boolean; isFavorited: boolean;
} | null> { } | null> {
const validated = GetChatInputSchema.parse(input); const validated = GetChatInputSchema.parse(input);
// Get chat with creator info // Get chat with creator info
const result = await db const result = await db
.select({ .select({
chat: chats, chat: chats,
user: users, user: users,
}) })
.from(chats) .from(chats)
.leftJoin(users, eq(chats.createdBy, users.id)) .leftJoin(users, eq(chats.createdBy, users.id))
.where(and(eq(chats.id, validated.chatId), isNull(chats.deletedAt))) .where(and(eq(chats.id, validated.chatId), isNull(chats.deletedAt)))
.limit(1); .limit(1);
if (!result.length || !result[0]?.chat) { if (!result.length || !result[0]?.chat) {
return null; return null;
} }
// Check if favorited // Check if favorited
const favorite = await db const favorite = await db
.select() .select()
.from(userFavorites) .from(userFavorites)
.where( .where(
and( and(
eq(userFavorites.userId, validated.userId), eq(userFavorites.userId, validated.userId),
eq(userFavorites.assetId, validated.chatId), eq(userFavorites.assetId, validated.chatId),
eq(userFavorites.assetType, 'chat') eq(userFavorites.assetType, 'chat'),
) ),
) )
.limit(1); .limit(1);
const { chat, user } = result[0]; const { chat, user } = result[0];
return { return {
chat, chat,
user, user,
isFavorited: favorite.length > 0, isFavorited: favorite.length > 0,
}; };
} }
/** /**
* Create a new message in a chat * Create a new message in a chat
*/ */
export async function createMessage(input: CreateMessageInput): Promise<Message> { export async function createMessage(input: CreateMessageInput): Promise<Message> {
try { try {
const validated = CreateMessageInputSchema.parse(input); const validated = CreateMessageInputSchema.parse(input);
const messageId = validated.messageId || crypto.randomUUID(); const messageId = validated.messageId || crypto.randomUUID();
// Use transaction to ensure atomicity // Use transaction to ensure atomicity
const result = await db.transaction(async (tx) => { const result = await db.transaction(async (tx) => {
const [message] = await tx const [message] = await tx
.insert(messages) .insert(messages)
.values({ .values({
id: messageId, id: messageId,
chatId: validated.chatId, chatId: validated.chatId,
createdBy: validated.userId, createdBy: validated.userId,
requestMessage: validated.content, requestMessage: validated.content,
responseMessages: {}, responseMessages: {},
reasoning: {}, reasoning: {},
title: validated.content.substring(0, 255), // Ensure title fits in database title: validated.content.substring(0, 255), // Ensure title fits in database
rawLlmMessages: {}, rawLlmMessages: {},
isCompleted: false, isCompleted: false,
}) })
.returning(); .returning();
if (!message) { if (!message) {
throw new Error('Failed to create message'); throw new Error('Failed to create message');
} }
// Update chat's updated_at timestamp // Update chat's updated_at timestamp
await tx await tx
.update(chats) .update(chats)
.set({ updatedAt: new Date().toISOString() }) .set({ updatedAt: new Date().toISOString() })
.where(eq(chats.id, validated.chatId)); .where(eq(chats.id, validated.chatId));
return message; return message;
}); });
return result; return result;
} catch (error) { } catch (error) {
if (error instanceof z.ZodError) { if (error instanceof z.ZodError) {
throw new Error(`Invalid message input: ${error.errors.map((e) => e.message).join(', ')}`); throw new Error(`Invalid message input: ${error.errors.map((e) => e.message).join(', ')}`);
} }
throw error; throw error;
} }
} }
/** /**
* Check if a user has permission to access a chat * Check if a user has permission to access a chat
*/ */
export async function checkChatPermission(chatId: string, userId: string): Promise<boolean> { export async function checkChatPermission(chatId: string, userId: string): Promise<boolean> {
const chat = await db const chat = await db
.select() .select()
.from(chats) .from(chats)
.where(and(eq(chats.id, chatId), isNull(chats.deletedAt))) .where(and(eq(chats.id, chatId), isNull(chats.deletedAt)))
.limit(1); .limit(1);
if (!chat.length) { if (!chat.length) {
return false; return false;
} }
// Check if user is the creator // For now, only check if user is the creator
if (chat[0]?.createdBy === userId) { // TODO: Add more sophisticated permission checking with asset_permissions table
return true; return chat[0]?.createdBy === userId;
}
// Check if user is an admin in the same organization as the chat
const chatOrganizationId = chat[0]?.organizationId;
if (!chatOrganizationId) {
return false;
}
// Check user's role in the organization
const userOrgRole = await db
.select({
role: usersToOrganizations.role,
})
.from(usersToOrganizations)
.where(
and(
eq(usersToOrganizations.userId, userId),
eq(usersToOrganizations.organizationId, chatOrganizationId),
isNull(usersToOrganizations.deletedAt)
)
)
.limit(1);
if (userOrgRole.length > 0 && userOrgRole[0]) {
const role = userOrgRole[0].role;
// Check if user has admin role (workspace_admin or data_admin)
if (role === 'workspace_admin' || role === 'data_admin') {
return true;
}
}
// TODO: Add more sophisticated permission checking with asset_permissions table
return false;
} }
/** /**
* Get all messages for a chat * Get all messages for a chat
*/ */
export async function getMessagesForChat(chatId: string): Promise<Message[]> { export async function getMessagesForChat(chatId: string): Promise<Message[]> {
return db return db
.select() .select()
.from(messages) .from(messages)
.where(and(eq(messages.chatId, chatId), isNull(messages.deletedAt))) .where(and(eq(messages.chatId, chatId), isNull(messages.deletedAt)))
.orderBy(messages.createdAt); .orderBy(messages.createdAt);
} }
/** /**
@ -236,50 +203,50 @@ export async function getMessagesForChat(chatId: string): Promise<Message[]> {
* @returns Success status * @returns Success status
*/ */
export async function updateChat( export async function updateChat(
chatId: string, chatId: string,
fields: UpdateableChatFields fields: UpdateableChatFields,
): Promise<{ success: boolean }> { ): Promise<{ success: boolean }> {
try { try {
// First verify the chat exists and is not deleted // First verify the chat exists and is not deleted
const existingChat = await db const existingChat = await db
.select({ id: chats.id }) .select({ id: chats.id })
.from(chats) .from(chats)
.where(and(eq(chats.id, chatId), isNull(chats.deletedAt))) .where(and(eq(chats.id, chatId), isNull(chats.deletedAt)))
.limit(1); .limit(1);
if (existingChat.length === 0) { if (existingChat.length === 0) {
throw new Error(`Chat not found or has been deleted: ${chatId}`); throw new Error(`Chat not found or has been deleted: ${chatId}`);
} }
// Build update object with only provided fields // Build update object with only provided fields
const updateData: any = { const updateData: any = {
updatedAt: new Date().toISOString(), updatedAt: new Date().toISOString(),
}; };
// Only add fields that are actually provided (not undefined) // Only add fields that are actually provided (not undefined)
for (const [key, value] of Object.entries(fields)) { for (const [key, value] of Object.entries(fields)) {
if (value !== undefined && key !== 'id' && key !== 'createdAt' && key !== 'deletedAt') { if (value !== undefined && key !== 'id' && key !== 'createdAt' && key !== 'deletedAt') {
updateData[key] = value; updateData[key] = value;
} }
} }
// If updatedAt was explicitly provided, use that instead // If updatedAt was explicitly provided, use that instead
if ('updatedAt' in fields && fields.updatedAt !== undefined) { if ('updatedAt' in fields && fields.updatedAt !== undefined) {
updateData.updatedAt = fields.updatedAt; updateData.updatedAt = fields.updatedAt;
} }
await db await db
.update(chats) .update(chats)
.set(updateData) .set(updateData)
.where(and(eq(chats.id, chatId), isNull(chats.deletedAt))); .where(and(eq(chats.id, chatId), isNull(chats.deletedAt)));
return { success: true }; return { success: true };
} catch (error) { } catch (error) {
console.error('Failed to update chat fields:', error); console.error('Failed to update chat fields:', error);
// Re-throw our specific validation errors // Re-throw our specific validation errors
if (error instanceof Error && error.message.includes('Chat not found')) { if (error instanceof Error && error.message.includes('Chat not found')) {
throw error; throw error;
} }
throw new Error(`Failed to update chat fields for chat ${chatId}`); throw new Error(`Failed to update chat fields for chat ${chatId}`);
} }
} }