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

package com.ankurm.imagegen;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.net.URI;
import java.net.http.HttpClient;
import java.net.http.HttpRequest;
import java.net.http.HttpResponse;
import java.time.Duration;

/**
 * Downloads an image a provider returned as a URL (DALL-E 3 style {@code response_format=url}).
 * The URL comes from a third party, so it is treated as untrusted input: https only, the host must
 * be on a list, redirects are not followed, the read stops at the byte cap, and there is a timeout.
 */
public final class RemoteFetcher {

    private final ImagePolicy policy;
    private final HttpClient http = HttpClient.newBuilder()
            .followRedirects(HttpClient.Redirect.NEVER)
            .connectTimeout(Duration.ofSeconds(5))
            .build();

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

    public byte[] fetch(String url) {
        URI uri;
        try {
            uri = URI.create(url);
        } catch (IllegalArgumentException e) {
            throw new RejectedException("bad-url", "not a URI");
        }
        boolean httpsOk = "https".equalsIgnoreCase(uri.getScheme());
        boolean httpOk = policy.allowInsecure() && "http".equalsIgnoreCase(uri.getScheme());
        if (!httpsOk && !httpOk) {
            throw new RejectedException("bad-scheme", String.valueOf(uri.getScheme()));
        }
        if (uri.getHost() == null || !policy.fetchHosts().contains(uri.getHost().toLowerCase())) {
            throw new RejectedException("host-not-allowed", String.valueOf(uri.getHost()));
        }
        try {
            HttpResponse<InputStream> r = http.send(
                    HttpRequest.newBuilder(uri).timeout(Duration.ofSeconds(15)).GET().build(),
                    HttpResponse.BodyHandlers.ofInputStream());
            if (r.statusCode() != 200) {
                r.body().close();
                throw new RejectedException("fetch-status", String.valueOf(r.statusCode()));
            }
            return readCapped(r.body());
        } catch (IOException e) {
            throw new RejectedException("fetch-failed", e.getClass().getSimpleName());
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
            throw new RejectedException("fetch-failed", "interrupted");
        }
    }

    /** Reads at most maxBytes + 1 bytes so an endless body costs us one byte over the cap, not the heap. */
    private byte[] readCapped(InputStream in) throws IOException {
        try (in) {
            ByteArrayOutputStream out = new ByteArrayOutputStream();
            byte[] buf = new byte[8192];
            long total = 0;
            int n;
            while ((n = in.read(buf)) > 0) {
                total += n;
                if (total > policy.maxBytes()) {
                    throw new RejectedException("too-large", "download passed the " + policy.maxBytes() + "-byte cap");
                }
                out.write(buf, 0, n);
            }
            return out.toByteArray();
        }
    }
}