spring-ai/image-gen/src/main/java/com/ankurm/imagegen/ImageVerifier.java

package com.ankurm.imagegen;

import java.awt.image.BufferedImage;
import java.io.ByteArrayInputStream;
import java.io.IOException;
import javax.imageio.ImageIO;

/**
 * Applies the policy to bytes: size cap, then sniff, then allowed type, then pixel cap read from
 * the header (before any decoder allocates a bitmap), then a real decode for the formats the JDK
 * can read. WebP is sniffed only, because the JDK has no WebP reader.
 */
public final class ImageVerifier {

    private final ImagePolicy policy;

    public ImageVerifier(ImagePolicy policy) {
        this.policy = policy;
    }

    public Sniffed verify(byte[] bytes) {
        if (bytes.length > policy.maxBytes()) {
            throw new RejectedException("too-large", bytes.length + " bytes > cap " + policy.maxBytes());
        }
        Sniffed s = ImageSniffer.sniff(bytes)
                .orElseThrow(() -> new RejectedException("not-an-image", "no known image signature in the first bytes"));
        if (!policy.allowedTypes().contains(s.contentType())) {
            throw new RejectedException("type-not-allowed", s.contentType());
        }
        if (s.pixels() > policy.maxPixels()) {
            throw new RejectedException("too-many-pixels", s.width() + "x" + s.height() + " = " + s.pixels() + " > cap " + policy.maxPixels());
        }
        if (!s.contentType().equals("image/webp")) {
            try {
                BufferedImage img = ImageIO.read(new ByteArrayInputStream(bytes));
                if (img == null) {
                    throw new RejectedException("undecodable", "signature says " + s.contentType() + " but no decoder accepts it");
                }
            } catch (IOException | RuntimeException e) {
                if (e instanceof RejectedException r) throw r;
                throw new RejectedException("undecodable", e.getClass().getSimpleName());
            }
        }
        return s;
    }
}