diff --git a/src/models/Message.ts b/src/models/Message.ts
index 6c47dc64..3763f02e 100644
--- a/src/models/Message.ts
+++ b/src/models/Message.ts
@@ -342,6 +342,15 @@ MessageSchema.virtual("mention_channels", {
autopopulate: { select: { id: true, guild_id: true, type: true, name: true } },
});
+
+MessageSchema.virtual("referenced_message", {
+ ref: "Message",
+ localField: "message_reference.message_id",
+ foreignField: "id",
+ justOne: true,
+ autopopulate: true,
+});
+
MessageSchema.virtual("created_at").get(function (this: MessageDocument) {
return new Date(Snowflake.deconstruct(this.id).timestamp);
});
@@ -358,3 +367,4 @@ MessageSchema.set("removeResponse", ["mention_channel_ids", "mention_role_ids",
// @ts-ignore
export const MessageModel = db.model<MessageDocument>("Message", MessageSchema, "messages");
+
diff --git a/src/models/RateLimit.ts b/src/models/RateLimit.ts
new file mode 100644
index 00000000..6a0e1ffd
--- /dev/null
+++ b/src/models/RateLimit.ts
@@ -0,0 +1,25 @@
+import { Schema, Document, Types } from "mongoose";
+import db from "../util/Database";
+
+export interface Bucket {
+ id: "global" | "error" | string; // channel_239842397 | guild_238927349823 | webhook_238923423498
+ user_id: string;
+ hits: number;
+ blocked: boolean;
+ expires_at: Date;
+}
+
+export interface BucketDocument extends Bucket, Document {
+ id: string;
+}
+
+export const BucketSchema = new Schema({
+ id: { type: String, required: true },
+ user_id: { type: String, required: true }, // bot, user, oauth_application, webhook
+ hits: { type: Number, required: true }, // Number of times the user hit this bucket
+ blocked: { type: Boolean, required: true },
+ expires_at: { type: Date, required: true },
+});
+
+// @ts-ignore
+export const BucketModel = db.model<BucketDocument>("Bucket", BucketSchema, "ratelimits");
diff --git a/src/models/index.ts b/src/models/index.ts
index 004095db..11a6fe37 100644
--- a/src/models/index.ts
+++ b/src/models/index.ts
@@ -1,7 +1,42 @@
-import mongoose from "mongoose";
-import { Schema } from "mongoose";
+import mongoose, { Schema, Document } from "mongoose";
import mongooseAutoPopulate from "mongoose-autopopulate";
+type UpdateWithAggregationPipeline = UpdateAggregationStage[];
+type UpdateAggregationStage =
+ | { $addFields: any }
+ | { $set: any }
+ | { $project: any }
+ | { $unset: any }
+ | { $replaceRoot: any }
+ | { $replaceWith: any };
+type EnforceDocument<T, TMethods> = T extends Document ? T : T & Document & TMethods;
+
+declare module "mongoose" {
+ interface Model<T, TQueryHelpers = {}, TMethods = {}> {
+ // removed null -> always return document -> throw error if it doesn't exist
+ findOne(
+ filter?: FilterQuery<T>,
+ projection?: any | null,
+ options?: QueryOptions | null,
+ callback?: (err: CallbackError, doc: EnforceDocument<T, TMethods>) => void
+ ): QueryWithHelpers<EnforceDocument<T, TMethods>, EnforceDocument<T, TMethods>, TQueryHelpers>;
+ findOneAndUpdate(
+ filter?: FilterQuery<T>,
+ update?: UpdateQuery<T> | UpdateWithAggregationPipeline,
+ options?: QueryOptions | null,
+ callback?: (err: any, doc: EnforceDocument<T, TMethods> | null, res: any) => void
+ ): QueryWithHelpers<EnforceDocument<T, TMethods>, EnforceDocument<T, TMethods>, TQueryHelpers>;
+ }
+}
+
+var HTTPError: any;
+
+try {
+ HTTPError = require("lambert-server").HTTPError;
+} catch (e) {
+ HTTPError = Error;
+}
+
mongoose.plugin(mongooseAutoPopulate);
mongoose.plugin((schema: Schema, opts: any) => {
@@ -17,6 +52,18 @@ mongoose.plugin((schema: Schema, opts: any) => {
});
},
});
+ schema.post("findOne", function (doc, next) {
+ try {
+ // @ts-ignore
+ const isExistsQuery = JSON.stringify(this._userProvidedFields) === JSON.stringify({ _id: 1 });
+ if (!doc && !isExistsQuery) return next(new HTTPError("Not found", 404));
+ // @ts-ignore
+ return next();
+ } catch (error) {
+ // @ts-ignore
+ next();
+ }
+ });
});
export * from "./Activity";
@@ -35,3 +82,4 @@ export * from "./Status";
export * from "./Role";
export * from "./User";
export * from "./VoiceState";
+export * from "./RateLimit";
|
