tiku-backend.net/Tiku.Infrastructure/Assets/ClamAvAssetSecurityScanner.cs

using System.Buffers.Binary;
using System.Net.Sockets;
using System.Text;
using Microsoft.Extensions.Options;
using Tiku.Application.Assets;
using Tiku.Infrastructure.Observability;
using System.Diagnostics;

namespace Tiku.Infrastructure.Assets;

public sealed class ClamAvOptions
{
    public const string SectionName = "Security:ClamAV";

    public string Host { get; set; } = "localhost";
    public int Port { get; set; } = 3310;
    public int TimeoutSeconds { get; set; } = 30;
    public int ChunkBytes { get; set; } = 64 * 1024;
    public long StreamMaxLength { get; set; } = 500L * 1024 * 1024;

    public static bool BeValid(ClamAvOptions options) =>
        !string.IsNullOrWhiteSpace(options.Host) &&
        options.Port is > 0 and <= 65535 &&
        options.TimeoutSeconds is >= 1 and <= 600 &&
        options.ChunkBytes is >= 1024 and <= 1024 * 1024 &&
        options.StreamMaxLength > 0;
}

public sealed class ClamAvAssetSecurityScanner(IOptions<ClamAvOptions> options) : IAssetSecurityScanner
{
    private readonly ClamAvOptions settings = options.Value;

    public async Task<AssetSecurityScanResult> ScanAsync(
        Stream content,
        long? declaredLength,
        CancellationToken cancellationToken = default)
    {
        var startedTimestamp = Stopwatch.GetTimestamp();
        if (declaredLength > settings.StreamMaxLength)
        {
            throw new AssetSecurityScannerException(
                "clamav_stream_too_large",
                "Asset exceeds the configured ClamAV StreamMaxLength.");
        }

        using var timeout = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
        timeout.CancelAfter(TimeSpan.FromSeconds(settings.TimeoutSeconds));
        try
        {
            using var client = new TcpClient();
            await client.ConnectAsync(settings.Host, settings.Port, timeout.Token);
            await using var network = client.GetStream();
            await network.WriteAsync("zINSTREAM\0"u8.ToArray(), timeout.Token);
            var buffer = new byte[settings.ChunkBytes];
            var lengthBuffer = new byte[4];
            long total = 0;
            while (true)
            {
                var read = await content.ReadAsync(buffer, timeout.Token);
                if (read == 0) break;
                total += read;
                if (total > settings.StreamMaxLength)
                {
                    throw new AssetSecurityScannerException(
                        "clamav_stream_too_large",
                        "Asset exceeds the configured ClamAV StreamMaxLength.");
                }
                BinaryPrimitives.WriteUInt32BigEndian(lengthBuffer, (uint)read);
                await network.WriteAsync(lengthBuffer, timeout.Token);
                await network.WriteAsync(buffer.AsMemory(0, read), timeout.Token);
            }
            Array.Clear(lengthBuffer);
            await network.WriteAsync(lengthBuffer, timeout.Token);
            await network.FlushAsync(timeout.Token);
            var response = await ReadResponseAsync(network, timeout.Token);
            if (response.EndsWith(": OK", StringComparison.Ordinal))
            {
                WorkerTelemetry.RecordScan("clean", Stopwatch.GetElapsedTime(startedTimestamp).TotalMilliseconds);
                return new AssetSecurityScanResult(AssetSecurityScanVerdict.Clean, "clamav", null, total, response);
            }
            if (response.EndsWith(" FOUND", StringComparison.Ordinal))
            {
                var separator = response.IndexOf(": ", StringComparison.Ordinal);
                var signature = separator >= 0
                    ? response[(separator + 2)..^" FOUND".Length]
                    : "unknown";
                WorkerTelemetry.RecordScan("infected", Stopwatch.GetElapsedTime(startedTimestamp).TotalMilliseconds);
                return new AssetSecurityScanResult(AssetSecurityScanVerdict.Infected, "clamav", signature, total, response);
            }
            throw new AssetSecurityScannerException("clamav_scan_error", $"ClamAV returned an error response: {response}");
        }
        catch (AssetSecurityScannerException)
        {
            WorkerTelemetry.RecordScan("error", Stopwatch.GetElapsedTime(startedTimestamp).TotalMilliseconds);
            throw;
        }
        catch (OperationCanceledException exception) when (!cancellationToken.IsCancellationRequested)
        {
            WorkerTelemetry.RecordScan("timeout", Stopwatch.GetElapsedTime(startedTimestamp).TotalMilliseconds);
            throw new AssetSecurityScannerException("clamav_timeout", "ClamAV scan timed out.", exception);
        }
        catch (Exception exception) when (exception is SocketException or IOException)
        {
            WorkerTelemetry.RecordScan("unavailable", Stopwatch.GetElapsedTime(startedTimestamp).TotalMilliseconds);
            throw new AssetSecurityScannerException("clamav_unavailable", "ClamAV is unavailable.", exception);
        }
    }

    public async Task<bool> CheckHealthAsync(CancellationToken cancellationToken = default)
    {
        using var timeout = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
        timeout.CancelAfter(TimeSpan.FromSeconds(settings.TimeoutSeconds));
        try
        {
            using var client = new TcpClient();
            await client.ConnectAsync(settings.Host, settings.Port, timeout.Token);
            await using var network = client.GetStream();
            await network.WriteAsync("zPING\0"u8.ToArray(), timeout.Token);
            await network.FlushAsync(timeout.Token);
            return string.Equals(await ReadResponseAsync(network, timeout.Token), "PONG", StringComparison.Ordinal);
        }
        catch when (!cancellationToken.IsCancellationRequested)
        {
            return false;
        }
    }

    private static async Task<string> ReadResponseAsync(Stream stream, CancellationToken cancellationToken)
    {
        using var buffer = new MemoryStream();
        var single = new byte[1];
        while (await stream.ReadAsync(single, cancellationToken) == 1 && single[0] != 0)
        {
            buffer.WriteByte(single[0]);
            if (buffer.Length > 4096)
            {
                throw new AssetSecurityScannerException("clamav_response_too_large", "ClamAV response exceeded the safety limit.");
            }
        }
        return Encoding.UTF8.GetString(buffer.ToArray()).Trim();
    }
}