Skip to content

Commit

Permalink
add protection
Browse files Browse the repository at this point in the history
  • Loading branch information
ZHallen122 committed Oct 26, 2024
1 parent aaea83e commit a6fec39
Show file tree
Hide file tree
Showing 8 changed files with 139 additions and 14 deletions.
5 changes: 5 additions & 0 deletions backend/src/chat/chat.model.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand All @@ -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')
Expand Down
19 changes: 16 additions & 3 deletions backend/src/chat/chat.module.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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 {}
21 changes: 19 additions & 2 deletions backend/src/chat/chat.resolver.ts
Original file line number Diff line number Diff line change
@@ -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, {
Expand All @@ -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,
);
Expand All @@ -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);
Expand All @@ -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);
Expand All @@ -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,
Expand Down
13 changes: 11 additions & 2 deletions backend/src/chat/chat.service.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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);
}

Expand All @@ -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;
Expand Down
4 changes: 2 additions & 2 deletions backend/src/chat/dto/chat.input.ts
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@ export class NewChatInput {
@InputType()
export class UpateChatTitleInput {
@Field()
id: string;
chatId: string;

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

@Field()
message: string;
Expand Down
55 changes: 55 additions & 0 deletions backend/src/guard/chat.guard.ts
Original file line number Diff line number Diff line change
@@ -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;
}
}
21 changes: 17 additions & 4 deletions backend/src/schema.gql
Original file line number Diff line number Diff line change
Expand Up @@ -7,9 +7,10 @@ type Chat {
id: ID!
isActive: Boolean!
isDeleted: Boolean!
messages: [Message!]!
messages: [Message!]
title: String
updatedAt: Date!
user: User!
}

type ChatCompletionChoiceType {
Expand All @@ -32,6 +33,7 @@ type ChatCompletionDeltaType {
}

input ChatInputType {
chatId: String!
message: String!
}

Expand Down Expand Up @@ -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!
Expand All @@ -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!
}
Expand All @@ -132,6 +140,11 @@ type Subscription {
chatStream(input: ChatInputType!): ChatCompletionChunkType
}

input UpateChatTitleInput {
chatId: String!
title: String
}

input UpsertProjectInput {
projectId: ID
projectName: String!
Expand Down
15 changes: 14 additions & 1 deletion backend/src/user/user.service.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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
});
}
}

0 comments on commit a6fec39

Please sign in to comment.