Test and improve AIPI scanning

Refresh.Core/Services/AipiService.cs

@@ -1,13 +1,16 @@
 using System.Diagnostics;
 using System.Net.Http.Json;
 using Bunkum.Core.Services;
+using Bunkum.Core.Storage;
 using JetBrains.Annotations;
 using NotEnoughLogs;
 using Refresh.Common;
 using Refresh.Core.Configuration;
 using Refresh.Core.Importing;
 using Refresh.Core.Types.Data;
+using Refresh.Database;
 using Refresh.Database.Models.Assets;
+using Refresh.Database.Models.Users;
 using SixLabors.ImageSharp;
 using SixLabors.ImageSharp.Formats;
 using SixLabors.ImageSharp.Processing;
@@ -17,7 +20,7 @@ namespace Refresh.Core.Services;
 // Referenced from DO.
 public class AipiService : EndpointService
 {
-    private readonly HttpClient _client;
+    protected HttpClient _client { get; init; }
     private readonly IntegrationConfig _config;
     private readonly DiscordStaffService? _discord;
 
@@ -77,7 +80,7 @@ private async Task<TData> PostAsync<TData>(string endpoint, Stream data)
         return aipiResponse.Data!;
     }
 
-    private async Task<Dictionary<string, float>> PredictEvaAsync(Stream data)
+    private async Task<Dictionary<string, float>> PredictEvaAsync(Stream data, string imageHash, GameUser user)
     {
         Stopwatch stopwatch = new();
         this.Logger.LogTrace(RefreshContext.Aipi, "Pre-processing image data...");
@@ -107,7 +110,7 @@ private async Task<Dictionary<string, float>> PredictEvaAsync(Stream data)
 
         float threshold = this._config.AipiThreshold;
 
-        this.Logger.LogDebug(RefreshContext.Aipi, $"Running prediction for image @ threshold={threshold}...");
+        this.Logger.LogInfo(RefreshContext.Aipi, $"Running prediction for image '{imageHash}' @ threshold={threshold} by {user}...");
 
         stopwatch.Start();
         Dictionary<string, float> prediction = await this.PostAsync<Dictionary<string, float>>($"/eva/predict?threshold={threshold}", processedData);    
@@ -118,44 +121,50 @@ private async Task<Dictionary<string, float>> PredictEvaAsync(Stream data)
         return prediction;
     }
 
-    public bool ScanAndHandleAsset(DataContext context, GameAsset asset)
+    public bool ScanAndHandleAsset(DataContext context, GameAsset asset, GameUser user)
+    {
+        return this.ScanAndHandleAsset(context.Database, context.DataStore, asset, user);
+    }
+
+    // Use the passed user instead of the asset's OriginalUploader because the user trying to use this asset
+    // is not necessarily also its uploader. If we really want to, we should auto-punish the user instead of the uploader,
+    // and since OriginalUploader can be null here, punishing the uploader will not always work anyway.
+    // Also, we should expect asset upload endpoints to pass the uploader as parameter anyway.
+    public bool ScanAndHandleAsset(GameDatabaseContext database, IDataStore dataStore, GameAsset asset, GameUser user)
     {
-        // guard the fact that assets have an owner
-        Debug.Assert(asset.OriginalUploader != null, $"Asset {asset.AssetHash} had no original uploader when trying to scan");
-        if (asset.OriginalUploader == null)
-            return false;
-
         // import the asset as png
         bool isPspAsset = asset.AssetHash.StartsWith("psp/");
 
-        if (!context.DataStore.ExistsInStore("png/" + asset.AssetHash))
+        if (!dataStore.ExistsInStore("png/" + asset.AssetHash))
         {
-            this._importer.ImportAsset(asset.AssetHash, isPspAsset, asset.AssetType, context.DataStore);
+            this._importer.ImportAsset(asset.AssetHash, isPspAsset, asset.AssetType, dataStore);
         }
 
         // do actual prediction
-        using Stream stream = context.DataStore.GetStreamFromStore("png/" + asset.AssetHash);
-        Dictionary<string, float> results = this.PredictEvaAsync(stream).Result;
+        using Stream stream = dataStore.GetStreamFromStore("png/" + asset.AssetHash);
+        Dictionary<string, float> results = this.PredictEvaAsync(stream, asset.AssetHash, user).Result;
 
         if (!results.Any(r => this._config.AipiBannedTags.Contains(r.Key)))
             return false;
 
-        this._discord?.PostPredictionResult(results, asset);
+        this._discord?.PostPredictionResult(results, asset, user);
+        // TODO also log this in our own mod log
 
         if (this._config.AipiRestrictAccountOnDetection)
         {
-            const string reason = "Automatic restriction for posting disallowed content. This will usually be undone within 24 hours if this is a mistake.";
-            context.Database.RestrictUser(asset.OriginalUploader, reason, DateTimeOffset.MaxValue);
+            this.Logger.LogInfo(RefreshContext.Aipi, $"Auto-restricting {user} because their image '{asset.AssetHash}' was determined to contain disallowed content.");
+            const string reason = "Automatic restriction for posting or using disallowed content. This will usually be undone within 24 hours if this is a mistake.";
+            database.RestrictUser(user, reason, DateTimeOffset.MaxValue);
         }
 
         return true;
     }
+}
 
-    private class AipiResponse<TData>
-    {
-        public bool Success { get; set; }
-
-        public TData? Data { get; set; }
-        public string? Reason { get; set; }
-    }
+public class AipiResponse<TData>
+{
+    public bool Success { get; set; }
+
+    public TData? Data { get; set; }
+    public string? Reason { get; set; }
 }

Refresh.Core/Services/DiscordStaffService.cs

@@ -53,12 +53,12 @@ private void PostMessage(string? message = null, IEnumerable<Embed>? embeds = nu
         this.Logger.LogInfo(RefreshContext.Discord, $"Posted webhook {id}: '{message}'");
     }
 
-    public void PostPredictionResult(Dictionary<string, float> results, GameAsset asset)
+    public void PostPredictionResult(Dictionary<string, float> results, GameAsset asset, GameUser user)
     {
-        GameUser author = asset.OriginalUploader!;
+        GameUser author = user;
 
         EmbedBuilder builder = new EmbedBuilder()
-            .WithAuthor($"Image posted by @{author.Username} (id: {author.UserId})", this.GetAssetUrl(author.IconHash))
+            .WithAuthor($"Image posted or used by @{author.Username} (id: {author.UserId})", this.GetAssetUrl(author.IconHash))
             .WithDescription(DefaultResultsDescription)
             .WithUrl(this.GetAssetInfoUrl(asset.AssetHash))
             .WithTitle($"AI Analysis of `{asset.AssetHash}`");

Refresh.Interfaces.APIv3/Endpoints/ResourceApiEndpoints.cs

@@ -218,7 +218,7 @@ IntegrationConfig integration
 
         gameAsset.OriginalUploader = user;
 
-        if (aipi != null && aipi.ScanAndHandleAsset(dataContext, gameAsset))
+        if (aipi != null && aipi.ScanAndHandleAsset(dataContext, gameAsset, user))
         {
             return ApiModerationError.AssetAutoFlaggedError;
         }

RefreshTests.GameServer/AipiServer/TestAipiEndpoints.cs

@@ -0,0 +1,52 @@
+using Bunkum.Core;
+using Bunkum.Core.Endpoints;
+using Bunkum.Core.Responses;
+using Bunkum.Listener.Protocol;
+using Bunkum.Protocols.Http;
+using Refresh.Common;
+using Refresh.Core.Services;
+
+namespace RefreshTests.GameServer.AipiServer;
+
+public class TestAipiEndpoints : EndpointGroup
+{
+    [HttpEndpoint("/"), Authentication(false)]
+    public string TestAipi(RequestContext context)
+    {
+        return "AIPI scanning service";
+    }
+
+    [HttpEndpoint("/eva/predict", HttpMethods.Post, ContentType.BinaryData), Authentication(false)]
+    public Response ScanImage(RequestContext context, Stream body)
+    {
+        // check if we're secretly requesting a failure
+        string? forcedFailureReason = context.RequestHeaders.Get("X-ForcedFailureReason");
+        if (forcedFailureReason != null)
+        {
+            return new(new AipiResponse<Dictionary<string, float>>
+            {
+                Success = false,
+                Reason = forcedFailureReason,
+                Data = null,
+            }, ContentType.Json, NotAcceptable);
+        }
+
+        string? thresholdStr = context.QueryString.Get("threshold");
+        if (thresholdStr == null || !float.TryParse(thresholdStr, out float threshold))
+        {
+            context.Logger.LogTrace(RefreshContext.Aipi, $"Threshold not provided or unparseable, falling back to 0.0");
+            threshold = 0.0f;
+        }
+
+        Dictionary<string, float> tags = [];
+        if (threshold <= 67) tags.Add("sixSeven", 67.0f);
+        if (threshold <= 123.456) tags.Add("hi", 123.456f);
+
+        return new(new AipiResponse<Dictionary<string, float>>
+        {
+            Success = true,
+            Reason = null,
+            Data = tags,
+        }, ContentType.Json);
+    }
+}

RefreshTests.GameServer/GameServer/Services/TestAipiService.cs

@@ -0,0 +1,13 @@
+using NotEnoughLogs;
+using Refresh.Core.Configuration;
+using Refresh.Core.Services;
+
+namespace RefreshTests.GameServer.GameServer.Services;
+
+public class TestAipiService : AipiService
+{
+    public TestAipiService(Logger logger, IntegrationConfig config, ImportService import, DiscordStaffService discord, HttpClient client) : base(logger, config, import, discord)
+    {
+        this._client = client;
+    }
+}

RefreshTests.GameServer/Tests/Assets/AipiTests.cs

@@ -0,0 +1,162 @@
+using System.Reflection;
+using System.Security.Cryptography;
+using Bunkum.Core.Storage;
+using Newtonsoft.Json;
+using Refresh.Common.Helpers;
+using Refresh.Core.Configuration;
+using Refresh.Core.Services;
+using Refresh.Database.Models.Assets;
+using Refresh.Database.Models.Users;
+using RefreshTests.GameServer.AipiServer;
+using RefreshTests.GameServer.GameServer.Services;
+
+namespace RefreshTests.GameServer.Tests.Assets;
+
+public class AipiTests : GameServerTest
+{
+    private static readonly byte[] TestAsset = ResourceHelper.ReadResource("RefreshTests.GameServer.Resources.1x1.png", Assembly.GetExecutingAssembly());
+
+    [Test]
+    public void TestEndpointsWork()
+    {
+        using TestContext context = this.GetServer();
+        TestRefreshGameServer server = context.Server.Value;
+        server.Server.AddEndpointGroup<TestAipiEndpoints>();
+
+        HttpResponseMessage message = context.Http.GetAsync("/").Result;
+        Assert.That(message.StatusCode, Is.EqualTo(OK));
+        Assert.That(message.Content.ReadAsStringAsync().Result, Is.EqualTo("AIPI scanning service"));
+
+        message = context.Http.PostAsync("/eva/predict", new StreamContent(new MemoryStream(TestAsset))).Result;
+        Assert.That(message.StatusCode, Is.EqualTo(OK));
+
+        AipiResponse<Dictionary<string, float>>? result = JsonConvert.DeserializeObject<AipiResponse<Dictionary<string, float>>>(message.Content.ReadAsStringAsync().Result);
+        Assert.That(result, Is.Not.Null);
+        Assert.That(result!.Success, Is.True);
+        Assert.That(result!.Reason, Is.Null);
+        Assert.That(result!.Data, Is.Not.Null);
+        Assert.That(result!.Data!.Count, Is.EqualTo(2));
+        Assert.That(result!.Data!.GetValueOrDefault("sixSeven"), Is.EqualTo(67.0f));
+        Assert.That(result!.Data!.GetValueOrDefault("hi"), Is.EqualTo(123.456f));
+
+        // Run again but with a threshold that will exclude tag sixSeven
+        message = context.Http.PostAsync("/eva/predict?threshold=70.0", new StreamContent(new MemoryStream(TestAsset))).Result;
+        Assert.That(message.StatusCode, Is.EqualTo(OK));
+
+        result = JsonConvert.DeserializeObject<AipiResponse<Dictionary<string, float>>>(message.Content.ReadAsStringAsync().Result);
+        Assert.That(result, Is.Not.Null);
+        Assert.That(result!.Success, Is.True);
+        Assert.That(result!.Reason, Is.Null);
+        Assert.That(result!.Data, Is.Not.Null);
+        Assert.That(result!.Data!.Count, Is.EqualTo(1));
+        Assert.That(result!.Data!.ContainsKey("sixSeven"), Is.False); // was excluded because probability is too low
+        Assert.That(result!.Data!.GetValueOrDefault("hi"), Is.EqualTo(123.456f));
+
+        // Now force a failure
+        context.Http.DefaultRequestHeaders.Add("X-ForcedFailureReason", "real");
+        message = context.Http.PostAsync("/eva/predict?threshold=70.0", new StreamContent(new MemoryStream(TestAsset))).Result;
+        Assert.That(message.StatusCode, Is.EqualTo(NotAcceptable));
+
+        result = JsonConvert.DeserializeObject<AipiResponse<Dictionary<string, float>>>(message.Content.ReadAsStringAsync().Result);
+        Assert.That(result, Is.Not.Null);
+        Assert.That(result!.Success, Is.False);
+        Assert.That(result!.Reason, Is.Not.Null);
+        Assert.That(result!.Reason, Is.EqualTo("real"));
+        Assert.That(result!.Data, Is.Null);
+    }
+
+    private bool ScanImage(float threshold, string[] bannedTags, TestContext context, GameUser? uploader = null, bool autoRestrict = false)
+    {
+        TestRefreshGameServer server = context.Server.Value;
+        server.Server.AddEndpointGroup<TestAipiEndpoints>();
+        ImportService importer = server.GetService<ImportService>();
+        IDataStore dataStore = context.GetDataStore(); // is an InMemoryDataStore here
+
+        string hash = BitConverter.ToString(SHA1.HashData(TestAsset)).Replace("-", "").ToLower();
+        // Can't expect AipiService's ImageImporter to copy the PNG (or save the converted PNG if we weren't using a PNG)
+        // because Bunkum's InMemoryDataStore stubs OpenWriteStream() by throwing a NotImplementedException
+        dataStore.WriteToStore("png/" + hash, TestAsset);
+
+        IntegrationConfig integration = new()
+        {
+            AipiEnabled = true,
+            AipiThreshold = threshold,
+            AipiBannedTags = bannedTags,
+            AipiRestrictAccountOnDetection = autoRestrict,
+            DiscordStaffWebhookEnabled = false,
+        };
+        TestAipiService aipi = new(server.Logger, integration, importer, null!, context.Http);
+        aipi.Initialize();
+
+        uploader ??= context.CreateUser();
+        GameAsset metadata = new()
+        {
+            AssetHash = hash,
+            AssetType = GameAssetType.Png,
+            OriginalUploader = uploader,
+            SizeInBytes = TestAsset.Length,
+            IsPSP = false,
+        };
+        return aipi.ScanAndHandleAsset(context.Database, dataStore, metadata, uploader);
+    }
+
+    [Test]
+    public void ImageIsIgnoredIfNoKnownTagsReturned()
+    {
+        using TestContext context = this.GetServer();
+
+        Assert.That(this.ScanImage(0.0f, [], context), Is.False);
+    }
+
+    [Test]
+    public void ImageIsIgnoredIfBackendFailure()
+    {
+        using TestContext context = this.GetServer();
+        context.Http.DefaultRequestHeaders.Add("X-ForcedFailureReason", "lel");
+
+        bool hasThrown = false;
+        try
+        {
+            this.ScanImage(0.0f, [], context);
+        }
+        catch (Exception ex)
+        {
+            hasThrown = true;
+            Assert.That(ex.Message, Is.EqualTo($"One or more errors occurred. (NotAcceptable: lel)"));
+        }
+        Assert.That(hasThrown, Is.True);
+    }
+
+    [Test]
+    [TestCase("sixSeven", false)]
+    [TestCase("hi", true)]
+    public void ImageIsFlaggedIfAnyKnownTagsReturned(string tagName, bool flaggedOnHigherProbability)
+    {
+        using TestContext context = this.GetServer();
+
+        Assert.That(this.ScanImage(0.0f, [tagName], context), Is.True);
+
+        // Now raise the probability, so that sixSeven won't appear anymore
+        Assert.That(this.ScanImage(90.0f, [tagName], context), Is.EqualTo(flaggedOnHigherProbability));
+    }
+
+    [Test]
+    [TestCase(false)]
+    [TestCase(true)]
+    public void UserGetsAutoRestrictedIfWanted(bool autoRestrict)
+    {
+        using TestContext context = this.GetServer();
+        GameUser uploader = context.CreateUser();
+
+        Assert.That(this.ScanImage(0.0f, ["hi"], context, uploader, autoRestrict), Is.True);
+
+        // re-get user from DB to check their role
+        GameUser? updatedUser = context.Database.GetUserByObjectId(uploader.UserId);
+        Assert.That(updatedUser, Is.Not.Null);
+
+        if (autoRestrict)
+            Assert.That(updatedUser!.Role, Is.EqualTo(GameUserRole.Restricted));
+        else
+            Assert.That(updatedUser!.Role, Is.GreaterThan(GameUserRole.Restricted));
+    }
+}