NpgsqlCommandRuntimePatcher.cs

using System;
using System.Linq;
using System.Reflection;
using HarmonyLib;

namespace PostgresRuntimePatcher;

internal sealed class NpgsqlCommandRuntimePatcher : IDisposable
{
    private const string HarmonyId = "com.cenizalunar.tshock.postgres-runtime-patcher";

    private static readonly string[] TargetMethodNames =
    [
        "ExecuteNonQuery",
        "ExecuteScalar",
        "ExecuteNonQueryAsync",
        "ExecuteScalarAsync",
        "ExecuteDbDataReader",
        "ExecuteDbDataReaderAsync"
    ];

    private readonly Harmony harmony = new(HarmonyId);
    private bool applied;

    public int Apply()
    {
        if (this.applied)
        {
            return 0;
        }

        Type? commandType = GetNpgsqlCommandType();
        if (commandType is null)
        {
            return 0;
        }

        var prefixMethod = typeof(NpgsqlCommandRuntimePatcher).GetMethod(nameof(RewriteCommandTextPrefix), BindingFlags.NonPublic | BindingFlags.Static);
        if (prefixMethod is null)
        {
            return 0;
        }

        var prefix = new HarmonyMethod(prefixMethod);
        int patched = 0;

        foreach (var method in commandType.GetMethods(BindingFlags.Public | BindingFlags.NonPublic | BindingFlags.Instance | BindingFlags.DeclaredOnly).Where(IsPatchableCommandMethod).Distinct())
        {
            try
            {
                this.harmony.Patch(method, prefix: prefix);
                patched++;
            }
            catch
            {
            }
        }

        this.applied = patched > 0;
        return patched;
    }

    public void Dispose()
    {
        if (!this.applied)
        {
            return;
        }

        this.harmony.UnpatchAll(HarmonyId);
        this.applied = false;
    }

    private static Type? GetNpgsqlCommandType()
    {
        var assembly = AppDomain.CurrentDomain.GetAssemblies().FirstOrDefault(x => string.Equals(x.GetName().Name, "Npgsql", StringComparison.OrdinalIgnoreCase));

        if (assembly is null)
        {
            try
            {
                assembly = Assembly.Load("Npgsql");
            }
            catch
            {
                return null;
            }
        }

        return assembly.GetType("Npgsql.NpgsqlCommand", false);
    }

    private static bool IsPatchableCommandMethod(MethodInfo method)
    {
        if (method.IsAbstract || method.ContainsGenericParameters)
        {
            return false;
        }

        return TargetMethodNames.Contains(method.Name, StringComparer.Ordinal);
    }

    private static void RewriteCommandTextPrefix(object __instance)
    {
        var property = __instance.GetType().GetProperty("CommandText", BindingFlags.Public | BindingFlags.Instance);
        if (property is null || !property.CanRead || !property.CanWrite)
        {
            return;
        }

        if (property.GetValue(__instance) is not string sql || string.IsNullOrWhiteSpace(sql))
        {
            return;
        }

        string rewritten = PostgresSqlRewriter.Rewrite(sql);
        if (!string.Equals(sql, rewritten, StringComparison.Ordinal))
        {
            property.SetValue(__instance, rewritten);
        }
    }
}