Add some more advanced image anti-spam

src/main/java/org/enginehub/discord/module/ImageSpamTracker.java

@@ -0,0 +1,161 @@
+/*
+ * Copyright (c) EngineHub and Contributors
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in all
+ * copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+ * SOFTWARE.
+ */
+
+package org.enginehub.discord.module;
+
+import com.google.common.base.Ticker;
+import com.google.common.cache.CacheBuilder;
+import com.google.common.math.PairedStatsAccumulator;
+import com.google.common.math.StatsAccumulator;
+import org.jetbrains.annotations.Nullable;
+
+import java.awt.Color;
+import java.awt.Graphics2D;
+import java.awt.RenderingHints;
+import java.awt.image.BufferedImage;
+import java.nio.DoubleBuffer;
+import java.time.Duration;
+import java.time.Instant;
+import java.time.InstantSource;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.Set;
+import java.util.concurrent.TimeUnit;
+
+/// Tracks images sent by users and detects if a user is sending the same image repeatedly within a short time frame.
+final class ImageSpamTracker {
+
+    /// Tracking window for similar images.
+    private static final Duration WINDOW = Duration.ofMinutes(1);
+    /// How many messages with similar images must be sent within the window to trigger a spam detection.
+    private static final int TRIGGER_MESSAGES = 3;
+    /// The similarity threshold for two images to be considered similar.
+    private static final double SIMILARITY_THRESHOLD = 0.999;
+    /// The size of the fingerprint used to compare images. A larger size will be more accurate but slower.
+    private static final int FINGERPRINT_SIZE = 32;
+    /// The maximum number of entries to keep in the tracker.
+    private static final int MAX_ENTRIES = 16 * 1024;
+
+    private record Entry(long userId, long messageId, DoubleBuffer fingerprint) {
+    }
+
+    private final Set<Entry> entries;
+
+    ImageSpamTracker(InstantSource clock) {
+        this.entries = Collections.newSetFromMap(
+            CacheBuilder.newBuilder()
+                .expireAfterWrite(WINDOW.toNanos(), TimeUnit.NANOSECONDS)
+                .maximumSize(MAX_ENTRIES)
+                .ticker(new Ticker() {
+                    @Override
+                    public long read() {
+                        Instant now = clock.instant();
+                        return TimeUnit.SECONDS.toNanos(now.getEpochSecond()) + now.getNano();
+                    }
+                })
+                .<Entry, Boolean>build()
+                .asMap()
+        );
+    }
+
+    synchronized boolean isSpammedImage(long userId, long messageId, BufferedImage image) {
+        DoubleBuffer fingerprint = fingerprint(image);
+        if (fingerprint == null) {
+            return false;
+        }
+
+        int repeatedImageMessageCount = getRepeatedImageMessageCount(userId, messageId, fingerprint);
+        // Add the new entry after counting repeated images to avoid counting the new message as a repeat
+        entries.add(new Entry(userId, messageId, fingerprint));
+
+        return repeatedImageMessageCount >= TRIGGER_MESSAGES;
+    }
+
+    private int getRepeatedImageMessageCount(long userId, long messageId, DoubleBuffer fingerprint) {
+        // Although we're doing set-like things, the number of messages is small enough that a linear search is faster.
+        var messageIds = new ArrayList<Long>();
+        messageIds.add(messageId);
+        for (Entry entry : entries) {
+            if (entry.userId() != userId) {
+                continue;
+            }
+            // Avoid spending effort flagging messages that have already been flagged
+            if (messageIds.contains(entry.messageId())) {
+                continue;
+            }
+            if (similarity(fingerprint, entry.fingerprint()) >= SIMILARITY_THRESHOLD) {
+                messageIds.add(entry.messageId());
+            }
+        }
+        return messageIds.size();
+    }
+
+    private static double similarity(DoubleBuffer a, DoubleBuffer b) {
+        var stats = new PairedStatsAccumulator();
+        for (int i = 0; i < a.limit(); i++) {
+            stats.add(a.get(i), b.get(i));
+        }
+        return stats.pearsonsCorrelationCoefficient();
+    }
+
+    /// Generate a fingerprint for the given image, which can be used with [#similarity(DoubleBuffer, DoubleBuffer)] to
+    /// determine if two images are similar. Returns `null` if the image is a solid color.
+    ///
+    /// @param source the image to fingerprint
+    /// @return a fingerprint of the image, or `null` if not fingerprintable
+    private static @Nullable DoubleBuffer fingerprint(BufferedImage source) {
+        BufferedImage scaled = scaleAndCompositeOverWhite(source);
+
+        double[] cells = new double[FINGERPRINT_SIZE * FINGERPRINT_SIZE];
+        var stats = new StatsAccumulator();
+        int count = 0;
+        for (int i : scaled.getRGB(0, 0, FINGERPRINT_SIZE, FINGERPRINT_SIZE, null, 0, FINGERPRINT_SIZE)) {
+            double luma = getLuma(i);
+            cells[count++] = luma;
+            stats.add(luma);
+        }
+        return stats.populationVariance() == 0 ? null : DoubleBuffer.wrap(cells).asReadOnlyBuffer();
+    }
+
+    private static BufferedImage scaleAndCompositeOverWhite(BufferedImage source) {
+        var scaled = new BufferedImage(FINGERPRINT_SIZE, FINGERPRINT_SIZE, BufferedImage.TYPE_INT_RGB);
+        Graphics2D graphics = scaled.createGraphics();
+        graphics.setRenderingHint(RenderingHints.KEY_INTERPOLATION, RenderingHints.VALUE_INTERPOLATION_BILINEAR);
+        // Set white and fill rather than use bgColor overloads of `drawImage` because it does the bgColor prior to
+        // scaling, which is very slow.
+        graphics.setColor(Color.WHITE);
+        graphics.fillRect(0, 0, FINGERPRINT_SIZE, FINGERPRINT_SIZE);
+        graphics.drawImage(source, 0, 0, FINGERPRINT_SIZE, FINGERPRINT_SIZE, null);
+        graphics.dispose();
+        return scaled;
+    }
+
+    /// {@return the luma of the color}
+    ///
+    /// @param rgb the color in RGB format
+    private static double getLuma(int rgb) {
+        double red = (rgb >> 16) & 0xFF;
+        double green = (rgb >> 8) & 0xFF;
+        double blue = rgb & 0xFF;
+        return 0.299 * red + 0.587 * green + 0.114 * blue;
+    }
+}

src/main/java/org/enginehub/discord/module/NoMessageSpam.java

@@ -27,6 +27,7 @@
 import com.google.common.cache.LoadingCache;
 import net.dv8tion.jda.api.entities.Guild;
 import net.dv8tion.jda.api.entities.Member;
+import net.dv8tion.jda.api.entities.Message;
 import net.dv8tion.jda.api.entities.channel.ChannelType;
 import net.dv8tion.jda.api.entities.channel.middleman.GuildChannel;
 import net.dv8tion.jda.api.events.message.MessageReceivedEvent;
@@ -35,10 +36,12 @@
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 import org.enginehub.discord.EngineHubBot;
+import org.enginehub.discord.util.AttachmentImages;
 import org.enginehub.discord.util.PermissionRole;
 import org.enginehub.discord.util.PunishmentUtil;
 import org.jetbrains.annotations.NotNull;
 
+import java.time.InstantSource;
 import java.util.Arrays;
 import java.util.HashSet;
 import java.util.Set;
@@ -61,18 +64,18 @@ private record CacheKey(long userId, int messageHash, int attachmentsHash) { }
     private final LoadingCache<CacheKey, AtomicInteger> messageCounts = CacheBuilder.newBuilder()
         .expireAfterAccess(1, TimeUnit.MINUTES)
         .build(CacheLoader.from(() -> new AtomicInteger(0)));
+    private final ImageSpamTracker imageSpamTracker = new ImageSpamTracker(InstantSource.system());
     private final HashSet<Long> hasPingedBefore = new HashSet<>();
     private volatile Set<Long> guildsForPunish;
 
     @Override
     public void onMessageReceived(@Nonnull MessageReceivedEvent event) {
-        // Don't check for people who are trusted not to spam
-        if (EngineHubBot.isAuthorised(event.getMember(), PermissionRole.TRUSTED)) {
+        if (isSpamExempt(event)) {
             return;
         }
 
         if (event.getChannel() instanceof GuildChannel) {
-            if (checkForGeneralSpam(event)) {
+            if (checkForGeneralSpam(event) || checkForImageSpam(event)) {
                 // Skip further checks if they were already banned
                 return;
             }
@@ -81,6 +84,11 @@ public void onMessageReceived(@Nonnull MessageReceivedEvent event) {
         checkForAtEveryone(event);
     }
 
+    private boolean isSpamExempt(@NotNull MessageReceivedEvent event) {
+        return event.getAuthor().isBot() || event.getAuthor().isSystem() || event.isWebhookMessage()
+            || EngineHubBot.isAuthorised(event.getMember(), PermissionRole.TRUSTED);
+    }
+
     private void checkForAtEveryone(@NotNull MessageReceivedEvent event) {
         if (!(event.getMessage().getMentions().mentionsEveryone() || event.getMessage().getContentRaw().contains("@everyone"))) {
             return;
@@ -125,6 +133,27 @@ private boolean checkForGeneralSpam(@NotNull MessageReceivedEvent event) {
         return false;
     }
 
+    private boolean checkForImageSpam(@NotNull MessageReceivedEvent event) {
+        long userId = event.getAuthor().getIdLong();
+        long messageId = event.getMessageIdLong();
+        for (Message.Attachment attachment : event.getMessage().getAttachments()) {
+            if (!attachment.isImage()) {
+                continue;
+            }
+
+            try {
+                if (imageSpamTracker.isSpammedImage(userId, messageId, AttachmentImages.fetch(attachment))) {
+                    PunishmentUtil.banUser(event.getGuild(), event.getAuthor(), "Image spam", true);
+                    return true;
+                }
+            } catch (Exception e) {
+                LOGGER.warn("Failed to check image", e);
+            }
+        }
+
+        return false;
+    }
+
     // TODO: Consider extracting out?
     private interface Punishment {
         void enact(Guild guild, Member member, String reason);

src/main/java/org/enginehub/discord/module/errorhelper/ErrorHelper.java

@@ -46,13 +46,13 @@
 import org.enginehub.discord.module.errorhelper.resolver.MCLogsResolver;
 import org.enginehub.discord.module.errorhelper.resolver.RawSubdirectoryUrlResolver;
 import org.enginehub.discord.module.errorhelper.resolver.RawSubdomainUrlResolver;
+import org.enginehub.discord.util.AttachmentImages;
 import org.enginehub.discord.util.HttpUtil;
 import org.enginehub.discord.util.PasteUtil;
 
 import java.awt.image.BufferedImage;
 import java.io.BufferedReader;
 import java.io.IOException;
-import java.io.InputStream;
 import java.io.InputStreamReader;
 import java.nio.charset.StandardCharsets;
 import java.nio.file.Files;
@@ -69,7 +69,6 @@
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
 import javax.annotation.Nonnull;
-import javax.imageio.ImageIO;
 
 public class ErrorHelper extends ListenerAdapter implements Module {
 
@@ -109,8 +108,8 @@ public void scanMessage(Message message, User author, MessageChannel channel) {
                     continue;
                 }
 
-                try (InputStream is = attachment.getProxy().download().get()) {
-                    BufferedImage image = ImageIO.read(is);
+                try {
+                    BufferedImage image = AttachmentImages.fetch(attachment);
                     messageText.append(tesseract.doOCR(image));
                     if (EngineHubBot.isBotOwner(author) && message.getChannel() instanceof PrivateChannel) {
                         // If it's a bot developer, send OCR debug text.

src/main/java/org/enginehub/discord/util/AttachmentImages.java

@@ -0,0 +1,98 @@
+/*
+ * Copyright (c) EngineHub and Contributors
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in all
+ * copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+ * SOFTWARE.
+ */
+
+package org.enginehub.discord.util;
+
+import com.google.common.cache.Cache;
+import com.google.common.cache.CacheBuilder;
+import com.google.common.util.concurrent.UncheckedExecutionException;
+import net.dv8tion.jda.api.entities.Message;
+import okhttp3.HttpUrl;
+import okhttp3.Request;
+import okhttp3.Response;
+
+import java.awt.image.BufferedImage;
+import java.io.IOException;
+import java.util.concurrent.ExecutionException;
+import javax.imageio.ImageIO;
+
+/// Fetches images attached to Discord messages and caches them in memory. This is used to avoid repeatedly downloading
+/// the same image when various modules process the same message.
+///
+/// We might want to consider changing message events to instead carry an object that can hold the reference, so we're
+/// not reliant on GC to clean up the cache.
+public final class AttachmentImages {
+
+    // We probably don't need more than this many pixels, and it bounds how large an image can be.
+    private static final int MAX_SIDE = 4096;
+
+    private static final Cache<Message.Attachment, BufferedImage> IMAGES = CacheBuilder.newBuilder()
+        .weakKeys()
+        .build();
+
+    public static BufferedImage fetch(Message.Attachment attachment) throws IOException {
+        try {
+            return IMAGES.get(attachment, () -> load(attachment));
+        } catch (ExecutionException | UncheckedExecutionException e) {
+            Throwable cause = e.getCause();
+            if (cause instanceof IOException ioException) {
+                throw ioException;
+            }
+            throw new RuntimeException("Unexpected error while loading image: " + attachment.getFileName(), cause);
+        }
+    }
+
+    private static BufferedImage load(Message.Attachment attachment) throws IOException {
+        Request request = new Request.Builder().url(getPngUrl(attachment)).build();
+        try (Response response = HttpUtil.getClient().newCall(request).execute()) {
+            if (!response.isSuccessful()) {
+                throw new IOException(
+                    "HTTP " + response.code() + " while downloading image: " + attachment.getFileName()
+                );
+            }
+            BufferedImage image = ImageIO.read(response.body().byteStream());
+            if (image == null) {
+                throw new IOException("Unreadable image: " + attachment.getFileName());
+            }
+            return image;
+        }
+    }
+
+    private static HttpUrl getPngUrl(Message.Attachment attachment) {
+        HttpUrl.Builder url = HttpUrl.get(attachment.getProxyUrl()).newBuilder()
+            .setQueryParameter("format", "png");
+        int largestSide = Math.max(attachment.getWidth(), attachment.getHeight());
+        if (largestSide > MAX_SIDE) {
+            double scale = (double) MAX_SIDE / largestSide;
+            url.setQueryParameter("width", String.valueOf(scaleSide(attachment.getWidth(), scale)));
+            url.setQueryParameter("height", String.valueOf(scaleSide(attachment.getHeight(), scale)));
+        }
+        return url.build();
+    }
+
+    private static int scaleSide(int side, double scale) {
+        return Math.max(1, (int) Math.round(side * scale));
+    }
+
+    private AttachmentImages() {
+    }
+}