add protection

Author: ZHallen122

backend/src/chat/chat.model.ts

@@ -16,6 +16,7 @@ import {
 import { forwardRef } from '@nestjs/common';
 import { Message } from 'src/chat/message.model';
 import { SystemBaseModel } from 'src/system-base-model/system-base.model';
+import { User } from 'src/user/user.model';
 
 @Entity()
 @ObjectType()
@@ -31,6 +32,10 @@ export class Chat extends SystemBaseModel {
   @Field(() => [Message], { nullable: true })
   @OneToMany(() => Message, (message) => message.chat, { cascade: true })
   messages: Message[];
+
+  @ManyToOne(() => User, (user) => user.chats)
+  @Field(() => User)
+  user: User;
 }
 
 @ObjectType('ChatCompletionDeltaType')

backend/src/chat/chat.module.ts

@@ -8,10 +8,23 @@ import { TypeOrmModule } from '@nestjs/typeorm';
 import { User } from 'src/user/user.model';
 import { Chat } from './chat.model';
 import { Message } from 'src/chat/message.model';
+import { ChatGuard } from '../guard/chat.guard';
+import { AuthModule } from '../auth/auth.module';
+import { UserService } from 'src/user/user.service';
 
 @Module({
-  imports: [HttpModule, TypeOrmModule.forFeature([Chat, User, Message])],
-  providers: [ChatResolver, ChatProxyService, ChatService],
-  exports: [ChatService],
+  imports: [
+    HttpModule,
+    TypeOrmModule.forFeature([Chat, User, Message]),
+    AuthModule,
+  ],
+  providers: [
+    ChatResolver,
+    ChatProxyService,
+    ChatService,
+    ChatGuard,
+    UserService,
+  ],
+  exports: [ChatService, ChatGuard],
 })
 export class ChatModule {}

backend/src/chat/chat.resolver.ts

@@ -1,19 +1,24 @@
 import { Resolver, Subscription, Args, Query, Mutation } from '@nestjs/graphql';
 import { ChatCompletionChunk } from './chat.model';
 import { ChatProxyService, ChatService } from './chat.service';
+import { UserService } from 'src/user/user.service';
 import { Chat } from './chat.model';
 import { Message } from 'src/chat/message.model';
 import {
   NewChatInput,
   UpateChatTitleInput,
   ChatInput,
 } from 'src/chat/dto/chat.input';
+import { UseGuards } from '@nestjs/common';
+import { ChatGuard } from '../guard/chat.guard';
+import { GetUserIdFromToken } from '../decorator/get-auth-token';
 
 @Resolver('Chat')
 export class ChatResolver {
   constructor(
     private chatProxyService: ChatProxyService,
     private chatService: ChatService,
+    private userService: UserService,
   ) {}
 
   @Subscription(() => ChatCompletionChunk, {
@@ -26,7 +31,7 @@ export class ChatResolver {
       for await (const chunk of iterator) {
         if (chunk) {
           await this.chatService.saveMessage(
-            input.id,
+            input.chatId,
             chunk.id,
             chunk.choices[0].delta.content,
           );
@@ -39,8 +44,15 @@ export class ChatResolver {
     }
   }
 
+  @Query(() => [Chat], { nullable: true })
+  async getUserChats(@GetUserIdFromToken() userId: string): Promise<Chat[]> {
+    const user = await this.userService.getUserChats(userId);
+    return user ? user.chats : []; // Return chats if user exists, otherwise return an empty array
+  }
+
   @Query(() => Message, { nullable: true })
   async getMessageDetail(
+    @GetUserIdFromToken() userId: string,
     @Args('messageId') messageId: string,
   ): Promise<Message> {
     return this.chatService.getMessageById(messageId);
@@ -51,6 +63,7 @@ export class ChatResolver {
     return this.chatService.getChatHistory(chatId);
   }
 
+  @UseGuards(ChatGuard)
   @Query(() => Chat, { nullable: true })
   async getChatDetails(@Args('chatId') chatId: string): Promise<Chat> {
     return this.chatService.getChatDetails(chatId);
@@ -63,21 +76,25 @@ export class ChatResolver {
 
   @Mutation(() => Chat)
   async createChat(
+    @GetUserIdFromToken() userId: string,
     @Args('newChatInput') newChatInput: NewChatInput,
   ): Promise<Chat> {
-    return this.chatService.createChat(newChatInput);
+    return this.chatService.createChat(userId, newChatInput);
   }
 
+  @UseGuards(ChatGuard)
   @Mutation(() => Boolean)
   async deleteChat(@Args('chatId') chatId: string): Promise<boolean> {
     return this.chatService.deleteChat(chatId);
   }
 
+  @UseGuards(ChatGuard)
   @Mutation(() => Boolean)
   async clearChatHistory(@Args('chatId') chatId: string): Promise<boolean> {
     return this.chatService.clearChatHistory(chatId);
   }
 
+  @UseGuards(ChatGuard)
   @Mutation(() => Chat, { nullable: true })
   async updateChatTitle(
     @Args('upateChatTitleInput') upateChatTitleInput: UpateChatTitleInput,

backend/src/chat/chat.service.ts

@@ -171,13 +171,22 @@ export class ChatService {
     });
   }
 
-  async createChat(newChatInput: NewChatInput): Promise<Chat> {
+  async createChat(userId: string, newChatInput: NewChatInput): Promise<Chat> {
+    // Fetch the user entity using the userId
+    const user = await this.userRepository.findOne({ where: { id: userId } });
+    if (!user) {
+      throw new Error('User not found');
+    }
+
+    // Create a new chat and associate it with the user
     const newChat = this.chatRepository.create({
       title: newChatInput.title,
       messages: [],
       createdAt: new Date(),
       updatedAt: new Date(),
+      user: user, // Associate the user with the chat
     });
+
     return await this.chatRepository.save(newChat);
   }
 
@@ -204,7 +213,7 @@ export class ChatService {
     upateChatTitleInput: UpateChatTitleInput,
   ): Promise<Chat> {
     const chat = await this.chatRepository.findOne({
-      where: { id: upateChatTitleInput.id },
+      where: { id: upateChatTitleInput.chatId },
     });
     if (chat) {
       chat.title = upateChatTitleInput.title;

backend/src/chat/dto/chat.input.ts

@@ -11,7 +11,7 @@ export class NewChatInput {
 @InputType()
 export class UpateChatTitleInput {
   @Field()
-  id: string;
+  chatId: string;
 
   @Field({ nullable: true })
   title: string;
@@ -20,7 +20,7 @@ export class UpateChatTitleInput {
 @InputType('ChatInputType')
 export class ChatInput {
   @Field()
-  id: string;
+  chatId: string;
 
   @Field()
   message: string;

backend/src/guard/chat.guard.ts

@@ -0,0 +1,55 @@
+import {
+  Injectable,
+  CanActivate,
+  ExecutionContext,
+  UnauthorizedException,
+} from '@nestjs/common';
+import { GqlExecutionContext } from '@nestjs/graphql';
+import { JwtService } from '@nestjs/jwt';
+import { ChatService } from '../chat/chat.service';
+
+@Injectable()
+export class ChatGuard implements CanActivate {
+  constructor(
+    private readonly chatService: ChatService, // Inject ChatService to fetch chat details
+    private readonly jwtService: JwtService, // JWT Service to verify tokens
+  ) {}
+
+  async canActivate(context: ExecutionContext): Promise<boolean> {
+    const gqlContext = GqlExecutionContext.create(context);
+    const request = gqlContext.getContext().req;
+
+    // Extract the authorization header
+    const authHeader = request.headers.authorization;
+    if (!authHeader || !authHeader.startsWith('Bearer ')) {
+      throw new UnauthorizedException('Authorization token is missing');
+    }
+
+    // Decode the token to get user information
+    const token = authHeader.split(' ')[1];
+    let user: any;
+    try {
+      user = this.jwtService.verify(token);
+    } catch (error) {
+      throw new UnauthorizedException('Invalid token');
+    }
+
+    // Extract chatId from the request arguments
+    const args = gqlContext.getArgs();
+    const { chatId } = args;
+
+    // check if the user is part of the chat
+    const chat = await this.chatService.getChatDetails(chatId);
+    if (!chat) {
+      throw new UnauthorizedException('Chat not found');
+    }
+
+    if (chat.user.id !== user.userId) {
+      throw new UnauthorizedException(
+        'User is not authorized to access this chat',
+      );
+    }
+
+    return true;
+  }
+}

backend/src/schema.gql

@@ -7,9 +7,10 @@ type Chat {
   id: ID!
   isActive: Boolean!
   isDeleted: Boolean!
-  messages: [Message!]!
+  messages: [Message!]
   title: String
   updatedAt: Date!
+  user: User!
 }
 
 type ChatCompletionChoiceType {
@@ -32,6 +33,7 @@ type ChatCompletionDeltaType {
 }
 
 input ChatInputType {
+  chatId: String!
   message: String!
 }
 
@@ -75,23 +77,27 @@ type Message {
 
 type Mutation {
   clearChatHistory(chatId: String!): Boolean!
-  createChat: Chat!
+  createChat(newChatInput: NewChatInput!): Chat!
   deleteChat(chatId: String!): Boolean!
   deleteProject(projectId: String!): Boolean!
   login(input: LoginUserInput!): LoginResponse!
   registerUser(input: RegisterUserInput!): User!
   removePackageFromProject(packageId: String!, projectId: String!): Boolean!
-  updateChatTitle(chatId: String!, title: String!): Chat
+  updateChatTitle(upateChatTitleInput: UpateChatTitleInput!): Chat
   updateProjectPath(newPath: String!, projectId: String!): Boolean!
   upsertProject(upsertProjectInput: UpsertProjectInput!): Project!
 }
 
+input NewChatInput {
+  title: String
+}
+
 type Project {
   createdAt: Date!
   id: ID!
   isActive: Boolean!
   isDeleted: Boolean!
-  path: String!
+  path: String
   projectName: String!
   projectPackages: [ProjectPackages!]
   updatedAt: Date!
@@ -112,7 +118,9 @@ type Query {
   checkToken(input: CheckTokenInput!): Boolean!
   getChatDetails(chatId: String!): Chat
   getChatHistory(chatId: String!): [Message!]!
+  getMessageDetail(messageId: String!): Message
   getProjectDetails(projectId: String!): Project!
+  getUserChats: [Chat!]
   getUserProjects: [Project!]!
   logout: Boolean!
 }
@@ -132,6 +140,11 @@ type Subscription {
   chatStream(input: ChatInputType!): ChatCompletionChunkType
 }
 
+input UpateChatTitleInput {
+  chatId: String!
+  title: String
+}
+
 input UpsertProjectInput {
   projectId: ID
   projectName: String!

backend/src/user/user.service.ts

@@ -13,4 +13,17 @@ import { LoginUserInput } from './dto/login-user.input';
 import { ConfigService } from '@nestjs/config';
 
 @Injectable()
-export class UserService {}
+export class UserService {
+  constructor(
+    @InjectRepository(User)
+    private userRepository: Repository<User>,
+  ) {}
+
+  // Method to get all chats of a user
+  async getUserChats(userId: string): Promise<User> {
+    return this.userRepository.findOne({
+      where: { id: userId },
+      relations: ['chats'], // Ensure 'chats' relation is loaded
+    });
+  }
+}