From 7416964f4312cadf52a81c61f28c88a924c992ce Mon Sep 17 00:00:00 2001
From: Skyler Grey <skyler3665@gmail.com>
Date: Thu, 9 Mar 2023 11:59:09 +0000
Subject: [PATCH] I believe NSFW PFP scanning now works. Please could reviewers
 test this using either server profiles or isolated alts

---
 src/events/messageCreate.ts |  2 +-
 src/reflex/scanners.ts      | 22 +++++++---------------
 src/reflex/verify.ts        |  8 +++-----
 3 files changed, 11 insertions(+), 21 deletions(-)

diff --git a/src/events/messageCreate.ts b/src/events/messageCreate.ts
index 38c7674..cc69bf0 100644
--- a/src/events/messageCreate.ts
+++ b/src/events/messageCreate.ts
@@ -114,7 +114,7 @@ export async function callback(_client: NucleusClient, message: Message) {
                 if (
                     config.filters.images.NSFW &&
                     !(message.channel instanceof ThreadChannel ? message.channel.parent?.nsfw : message.channel.nsfw) &&
-                    (await NSFWCheck(element))
+                    (await NSFWCheck(element.url))
                 ) {
                     messageException(message.guild.id, message.channel.id, message.id);
                     await message.delete();
diff --git a/src/reflex/scanners.ts b/src/reflex/scanners.ts
index 74887d2..8b8d3c8 100644
--- a/src/reflex/scanners.ts
+++ b/src/reflex/scanners.ts
@@ -30,12 +30,8 @@ const clamscanner = await new ClamScan().init({
     }
 });
 
-export async function testNSFW(attachment: {
-    url: string;
-    height: number | null;
-    width: number | null;
-}): Promise<NSFWSchema> {
-    const [fileStream, hash] = await streamAttachment(attachment.url);
+export async function testNSFW(url: string): Promise<NSFWSchema> {
+    const [fileStream, hash] = await streamAttachment(url);
     const alreadyHaveCheck = await client.database.scanCache.read(hash);
     if (alreadyHaveCheck && "nsfw" in alreadyHaveCheck!) {
         return { nsfw: alreadyHaveCheck.nsfw };
@@ -49,11 +45,11 @@ export async function testNSFW(attachment: {
                 resolve(buf);
             })
     )) as Buffer;
-    const array = new Uint8Array(converted);
 
-    const img = tf.node.decodeImage(array) as tf.Tensor3D;
+    const img = tf.node.decodeImage(converted, 3, undefined, false) as tf.Tensor3D;
 
     const predictions = (await nsfw_model.classify(img, 1))[0]!;
+    img.dispose();
     console.log(2, predictions);
 
     const nsfw = predictions.className === "Hentai" || predictions.className === "Porn";
@@ -155,13 +151,9 @@ export async function LinkCheck(message: Discord.Message): Promise<string[]> {
     return detectionsTypes as string[];
 }
 
-export async function NSFWCheck(element: {
-    url: string;
-    height: number | null;
-    width: number | null;
-}): Promise<boolean> {
+export async function NSFWCheck(url: string): Promise<boolean> {
     try {
-        return (await testNSFW(element)).nsfw;
+        return (await testNSFW(url)).nsfw;
     } catch (e) {
         console.log(e);
         return false;
@@ -237,7 +229,7 @@ export async function doMemberChecks(member: Discord.GuildMember): Promise<void>
     console.log(4, avatarTextCheck);
     // Is the profile picture NSFW
     const avatar = member.displayAvatarURL({ extension: "png", size: 1024, forceStatic: true });
-    const avatarCheck = guildData.filters.images.NSFW && (await NSFWCheck({ url: avatar, height: 1024, width: 1024 }));
+    const avatarCheck = guildData.filters.images.NSFW && (await NSFWCheck(avatar));
     console.log(5, avatarCheck);
     // Does the username contain an invite
     const inviteCheck = guildData.filters.invite.enabled && /discord\.gg\/[a-zA-Z0-9]+/gi.test(member.user.username);
diff --git a/src/reflex/verify.ts b/src/reflex/verify.ts
index bdf4aba..540fb4d 100644
--- a/src/reflex/verify.ts
+++ b/src/reflex/verify.ts
@@ -121,11 +121,9 @@ export default async function (interaction: CommandInteraction | ButtonInteracti
             ]
         });
         if (
-            await NSFWCheck({
-                url: (interaction.member as GuildMember).user.displayAvatarURL({ extension: "png", forceStatic: true }),
-                height: 1024,
-                width: 1024
-            })
+            await NSFWCheck(
+                (interaction.member as GuildMember).user.displayAvatarURL({ extension: "png", forceStatic: true })
+            )
         ) {
             return await interaction.editReply({
                 embeds: [