Support more content types for the NSFWProtector

Dockerfile

@@ -1,5 +1,8 @@
 # We can't use alpine anymore because crypto has rust deps.
 FROM node:20-slim
+
+COPY --from=mwader/static-ffmpeg:7.1.1 /ffmpeg /usr/local/bin/
+
 COPY . /tmp/src
 RUN cd /tmp/src \
     && yarn install \

config/default.yaml

@@ -277,6 +277,14 @@ displayReports: true
 # and redacted as NSFW
 nsfwSensitivity: .6
 
+# Which model should be in use when processing NSFW media.
+# Available options are specified in https://github.com/infinitered/nsfwjs/tree/v4.2.0?tab=readme-ov-file#load-the-model
+nsfwModel: MobileNetV2
+
+# FFMPEG is used to generate thumbnails of media for processing with the NSFW model. This can be
+# set to `false` to forbid processing media with ffmpeg. 
+ffmpegPath: /usr/bin/ffmpeg
+
 # Set this to true if the synapse this mjolnir is protecting is using Matrix Authentication Service for auth
 # If so, provide the base url and clientId + clientSecret needed to obtain a token from MAS - see
 # https://element-hq.github.io/matrix-authentication-service/index.html for more information about

package.json

@@ -30,7 +30,7 @@
     "@types/jsdom": "^16.2.11",
     "@types/mocha": "^9.0.0",
     "@types/nedb": "^1.8.12",
-    "@types/node": "^18.0.0",
+    "@types/node": "^20",
     "@types/pg": "^8.6.5",
     "@types/request": "^2.48.8",
     "@types/shell-quote": "1.7.1",
@@ -47,7 +47,7 @@
   "dependencies": {
     "@sentry/node": "^7.17.2",
     "@sentry/tracing": "^7.17.2",
-    "@tensorflow/tfjs-node": "^4.21.0",
+    "@tensorflow/tfjs-node": "^4.22.0",
     "@vector-im/matrix-bot-sdk": "^0.7.1-element.6",
     "await-lock": "^2.2.2",
     "axios": "^1.8.2",
@@ -60,7 +60,7 @@
     "js-yaml": "^4.1.0",
     "jsdom": "^16.6.0",
     "matrix-appservice-bridge": "10.3.1",
-    "nsfwjs": "^4.1.0",
+    "nsfwjs": "4.2.0",
     "parse-duration": "^2.1.3",
     "pg": "^8.8.0",
     "prom-client": "^14.1.0",
@@ -69,6 +69,9 @@
     "ulidx": "^0.3.0",
     "yaml": "^2.2.2"
   },
+  "_dependencies_notes": {
+    "nsfwjs": "4.2.1 has broken imports - https://github.com/infinitered/nsfwjs/issues/921"
+  },
   "engines": {
     "node": ">=20.0.0"
   }

src/config.ts

@@ -187,6 +187,8 @@ export interface IConfig {
         };
     };
     nsfwSensitivity: number;
+    nsfwModelName: string;
+    ffmpegPath: string | false;
 
     /**
      * Config options only set at runtime. Try to avoid using the objects
@@ -274,6 +276,8 @@ const defaultConfig: IConfig = {
         },
     },
     nsfwSensitivity: 0.6,
+    nsfwModelName: "MobileNetV2",
+    ffmpegPath: "/usr/bin/ffmpeg",
     // Needed to make the interface happy.
     RUNTIME: {},
 };

src/protections/NsfwProtection.ts

@@ -17,20 +17,113 @@ limitations under the License.
 import { Protection } from "./IProtection";
 import { Mjolnir } from "../Mjolnir";
 import * as nsfw from "nsfwjs";
-import { LogLevel, LogService } from "@vector-im/matrix-bot-sdk";
-import { node } from "@tensorflow/tfjs-node";
+import { LogLevel, LogService, MXCUrl } from "@vector-im/matrix-bot-sdk";
+import { node, Tensor3D } from "@tensorflow/tfjs-node";
+import { spawn } from "node:child_process";
+import { ReadableStream } from "node:stream/web";
+
+const TENSOR_SUPPORTED_TYPES = [
+    "image/png",
+    "image/apng",
+    "image/bmp",
+    "image/x-bmp",
+    "image/gif",
+    "image/jpeg",
+    "image/jp2",
+    "image/jpx",
+    "image/jpm",
+];
+
+const FFMPEG_SUPPORTED_TYPES = ["video/", "image/"];
+
+const FFMPEG_EXTRA_ARGS: Record<string, string[]> = {
+    // Extra params needed to extract svg thumbnails.
+    "image/svg+xml": ["-f", "svg_pipe", "-frame_size", "10000", "-video_size", "512x512"],
+};
 
 export class NsfwProtection extends Protection {
     settings = {};
     // @ts-ignore
-    private model: any;
+    private model: nsfw.NSFWJS;
+
+    private classificationCache = new Map<string, boolean>();
+
+    /**
+     * Extract the first frame from a video or image source.
+     * @param ffmpegPath The path to the `ffmpeg` binary.
+     * @param stream The stream containing the source.
+     * @param mimetype The mimetype provided by the source.
+     *
+     * @returns A byte array containing the thumbnail in JPEG format.
+     */
+    static async extractFrame(
+        ffmpegPath: string,
+        stream: ReadableStream<Uint8Array>,
+        mimetype: string,
+    ): Promise<Uint8Array> {
+        const errData: Buffer[] = [];
+        const imageData: Buffer[] = [];
+        const extraArgs = FFMPEG_EXTRA_ARGS[mimetype] ?? [];
+        const cmd = spawn(ffmpegPath, [
+            ...extraArgs,
+            "-i",
+            "-",
+            "-update",
+            "true",
+            "-frames:v",
+            "1",
+            "-f",
+            "image2",
+            "-",
+        ]);
+        let stdinErrorFinished!: Promise<void>;
+        const p = new Promise<Uint8Array>((resolve, reject) => {
+            cmd.once("exit", (code) => {
+                if (code !== 0) {
+                    reject(new Error(`FFMPEG failed to run: ${code}, ${Buffer.concat(errData).toString()}`));
+                } else {
+                    resolve(Buffer.concat(imageData));
+                }
+            });
+            cmd.stderr.on("data", (b) => {
+                errData.push(b);
+            });
+            cmd.stdout.on("data", (b) => {
+                imageData.push(b);
+            });
+            // EPIPE is "normal" for ffmpeg to emit after it's finished processing an input.
+            stdinErrorFinished = new Promise((res, rej) => {
+                cmd.stdin.once("error", (e: { code: string }) => {
+                    if (e.code !== "EPIPE") {
+                        LogService.debug("NsfwProtection", "Unexpected error from ffmpeg", e);
+                        rej(e);
+                    }
+                    res();
+                });
+            });
+        });
+        for await (const element of stream) {
+            if (cmd.stdin.write(element) === false) {
+                // Wait for either a drain, or the whole stream to complete
+                await Promise.race([stdinErrorFinished, new Promise((r) => cmd.stdin.once("drain", r))]);
+            }
+        }
+        if (!cmd.stdin.writableEnded) {
+            cmd.stdin.end();
+        }
+        try {
+            return await p;
+        } finally {
+            LogService.debug("NsfwProtection", `Generated thumbnail`);
+        }
+    }
 
     constructor() {
         super();
     }
 
-    async initialize() {
-        this.model = await nsfw.load();
+    async initialize(modelName: string) {
+        this.model = await nsfw.load(modelName);
     }
 
     public get name(): string {
@@ -44,6 +137,57 @@ export class NsfwProtection extends Protection {
         );
     }
 
+    /**
+     * Given a MXC URL, attempt to extract a image supported by our NSFW model.
+     * @param mjolnir The mjolnir instance.
+     * @param mxc The MXC url.
+     * @returns Bytes of an image processable by the model.
+     */
+    private async determineImageFromMedia(mjolnir: Mjolnir, mxc: MXCUrl): Promise<Uint8Array | null> {
+        const res = await fetch(
+            `${mjolnir.client.homeserverUrl}/_matrix/client/v1/media/download/${encodeURIComponent(mxc.domain)}/${encodeURIComponent(mxc.mediaId)}`,
+            {
+                headers: {
+                    Authorization: `Bearer ${mjolnir.client.accessToken}`,
+                },
+            },
+        );
+        if (!res.body || res.status !== 200) {
+            LogService.error("NsfwProtection", `Could not fetch mxc ${mxc}: ${res.status}`);
+            return null;
+        }
+        const contentType = res.headers.get("content-type")?.split(";")[0];
+        if (!contentType) {
+            LogService.warn("NsfwProtection", `No content type header specified`);
+            return null;
+        }
+        console.log(contentType);
+        if (TENSOR_SUPPORTED_TYPES.includes(contentType)) {
+            return new Uint8Array(await res.arrayBuffer());
+        }
+        // Why do we do this?
+        // - We don't want to trust client thumbnails, which might not match the content. Or they might
+        //   not exist at all (which forces clients to generate their own)
+        // - We also don't want to make our homeserver generate thumbnails of potentially
+        //   harmful images, so this locally generates a thumbnail of a range of types in memory.
+        if (mjolnir.config.ffmpegPath && FFMPEG_SUPPORTED_TYPES.some((mt) => contentType.startsWith(mt))) {
+            const stream = res.body as ReadableStream<Uint8Array>;
+            try {
+                LogService.debug(
+                    "NsfwProtection",
+                    `Image type ${contentType} is unsupported by model, attempting to generate thumbnail`,
+                );
+                return await NsfwProtection.extractFrame(mjolnir.config.ffmpegPath, stream, contentType);
+            } catch (ex) {
+                LogService.warn("NsfwProtection", "Could not extract thumbnail from image", ex);
+                return null;
+            }
+        }
+
+        LogService.debug("NsfwProtection", `Unsupported file type ${contentType}`);
+        return null;
+    }
+
     public async handleEvent(mjolnir: Mjolnir, roomId: string, event: any): Promise<any> {
         if (event.type !== "m.room.message" && event.type !== "m.sticker") {
             return;
@@ -66,27 +210,46 @@ export class NsfwProtection extends Protection {
         }
 
         for (const mxc of mxcs) {
-            const image = await mjolnir.client.downloadContent(mxc);
+            // If we've already scanned this media, return early.
+            if (this.classificationCache.has(mxc)) {
+                if (this.classificationCache.get(mxc)) {
+                    await this.redactEvent(mjolnir, roomId, event, room);
+                    break;
+                }
+                continue;
+            }
 
-            let decodedImage;
+            const data = await this.determineImageFromMedia(mjolnir, MXCUrl.parse(mxc));
+            if (!data) {
+                // Couldn't extract an image, skip.
+                continue;
+            } else {
+                LogService.debug("NsfwProtection", `Thumbnail generated for ${mxc}`);
+            }
+            let decodedImage: Tensor3D | undefined;
             try {
-                decodedImage = await node.decodeImage(image.data, 3);
+                decodedImage = (await node.decodeImage(data, 3)) as Tensor3D;
+                const predictions = await this.model.classify(decodedImage);
+                LogService.debug("NsfwProtection", `Classified ${mxc} as`, predictions);
+                const isNsfw = predictions.some(
+                    (prediction) =>
+                        ["Hentai", "Porn"].includes(prediction.className) &&
+                        prediction.probability > mjolnir.config.nsfwSensitivity,
+                );
+                this.classificationCache.set(mxc, isNsfw);
+                if (isNsfw) {
+                    await this.redactEvent(mjolnir, roomId, event, room);
+                    // Stop scanning media once we've redacted.
+                    break;
+                }
             } catch (e) {
                 LogService.error("NsfwProtection", `There was an error processing an image: ${e}`);
                 continue;
-            }
-
-            const predictions = await this.model.classify(decodedImage);
-
-            for (const prediction of predictions) {
-                if (["Hentai", "Porn"].includes(prediction["className"])) {
-                    if (prediction["probability"] > mjolnir.config.nsfwSensitivity) {
-                        await this.redactEvent(mjolnir, roomId, event, room);
-                        break;
-                    }
+            } finally {
+                if (decodedImage) {
+                    decodedImage.dispose();
                 }
             }
-            decodedImage.dispose();
         }
     }
 

src/protections/ProtectionManager.ts

@@ -108,8 +108,8 @@ export class ProtectionManager {
             protection.settings[key].setValue(value);
         }
         if (protection.enabled) {
-            if (protection.name === "NsfwProtection") {
-                (protection as NsfwProtection).initialize();
+            if (protection instanceof NsfwProtection) {
+                await protection.initialize(this.mjolnir.config.nsfwModelName);
             }
             for (let roomId of this.mjolnir.protectedRoomsTracker.getProtectedRooms()) {
                 await protection.startProtectingRoom(this.mjolnir, roomId);