ClamAvChatAttachmentMalwareScanner.java

package com.ecommerce.chat.infrastructure.storage;

import com.ecommerce.chat.application.port.ChatAttachmentMalwareScanner;
import org.springframework.stereotype.Component;

import java.io.ByteArrayOutputStream;
import java.io.DataOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.net.InetSocketAddress;
import java.net.Socket;
import java.nio.charset.StandardCharsets;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.util.HexFormat;

@Component
public class ClamAvChatAttachmentMalwareScanner implements ChatAttachmentMalwareScanner {

    private static final byte[] INSTREAM_COMMAND =
            "zINSTREAM\0".getBytes(StandardCharsets.US_ASCII);
    private static final int BUFFER_SIZE = 8192;

    private final ChatAttachmentScanProperties properties;

    public ClamAvChatAttachmentMalwareScanner(ChatAttachmentScanProperties properties) {
        this.properties = properties;
    }

    @Override
    public ScanResult scan(InputStream content, long maximumBytes) {
        try (Socket socket = new Socket()) {
            socket.connect(
                    new InetSocketAddress(properties.host(), properties.port()),
                    Math.toIntExact(properties.connectTimeout().toMillis()));
            socket.setSoTimeout(Math.toIntExact(properties.readTimeout().toMillis()));
            MessageDigest digest = MessageDigest.getInstance("SHA-256");
            long sizeBytes = sendContent(socket, content, maximumBytes, digest);
            String response = readResponse(socket.getInputStream());
            return parseResponse(
                    response,
                    sizeBytes,
                    HexFormat.of().formatHex(digest.digest()));
        } catch (IOException exception) {
            throw new IllegalStateException("ClamAV scan transport failed", exception);
        } catch (NoSuchAlgorithmException exception) {
            throw new IllegalStateException("SHA-256 is unavailable", exception);
        }
    }

    private long sendContent(
            Socket socket,
            InputStream content,
            long maximumBytes,
            MessageDigest digest) throws IOException {
        DataOutputStream output = new DataOutputStream(socket.getOutputStream());
        output.write(INSTREAM_COMMAND);
        byte[] buffer = new byte[BUFFER_SIZE];
        long totalBytes = 0L;
        int read;
        while ((read = content.read(buffer)) != -1) {
            totalBytes += read;
            if (totalBytes > maximumBytes) {
                throw new IllegalStateException("Attachment exceeded the scan size limit");
            }
            digest.update(buffer, 0, read);
            output.writeInt(read);
            output.write(buffer, 0, read);
        }
        output.writeInt(0);
        output.flush();
        return totalBytes;
    }

    private String readResponse(InputStream input) throws IOException {
        ByteArrayOutputStream response = new ByteArrayOutputStream();
        while (response.size() < properties.maximumResponseBytes()) {
            int value = input.read();
            if (value == -1 || value == 0 || value == '\n') {
                break;
            }
            response.write(value);
        }
        if (response.size() == 0) {
            throw new IllegalStateException("ClamAV returned an empty response");
        }
        if (response.size() >= properties.maximumResponseBytes()) {
            throw new IllegalStateException("ClamAV response exceeded the configured limit");
        }
        return response.toString(StandardCharsets.UTF_8).trim();
    }

    private ScanResult parseResponse(String response, long sizeBytes, String sha256) {
        int separator = response.indexOf(": ");
        if (separator < 0) {
            throw new IllegalStateException("Unsupported ClamAV response: " + response);
        }
        String result = response.substring(separator + 2);
        if ("OK".equals(result)) {
            return new ScanResult(Verdict.CLEAN, "ClamAV", null, sizeBytes, sha256);
        }
        if (result.endsWith(" FOUND")) {
            String signature = result.substring(0, result.length() - " FOUND".length()).trim();
            if (signature.isEmpty()) {
                throw new IllegalStateException("ClamAV omitted the malware signature");
            }
            return new ScanResult(Verdict.INFECTED, "ClamAV", signature, sizeBytes, sha256);
        }
        throw new IllegalStateException("ClamAV scan failed: " + result);
    }
}