fi.evolver.ai.spring.util.MultiPartBodyPublisher Maven / Gradle / Ivy

Go to download
package fi.evolver.ai.spring.util;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.OutputStreamWriter;
import java.io.UncheckedIOException;
import java.io.Writer;
import java.net.http.HttpRequest;
import java.nio.charset.StandardCharsets;
import java.util.*;
import java.util.function.Supplier;

import fi.evolver.ai.spring.file.AiFile;

public class MultiPartBodyPublisher {
	private final String boundary;

	private final List parts = new ArrayList<>();


	public MultiPartBodyPublisher(String boundary) {
		this.boundary = boundary;
	}


	public MultiPartBodyPublisher() {
		this(UUID.randomUUID().toString());
	}


	public HttpRequest.BodyPublisher build() {
		if (parts.isEmpty())
			throw new IllegalStateException("At least one part required");
		return HttpRequest.BodyPublishers.ofByteArrays(BodyBytesIterator::new);
	}


	Iterator iterator() {
		return new BodyBytesIterator();
	}


	public String getBoundary() {
		return boundary;
	}

	public MultiPartBodyPublisher addPart(String name, String value) {
		parts.add(new Part(
				name,
				Optional.empty(),
				() -> value.getBytes(StandardCharsets.UTF_8),
				Map.of()));
		return this;
	}

	public MultiPartBodyPublisher addPart(String name, Supplier value, String filename, String contentType) {
		parts.add(new Part(
				name,
				Optional.ofNullable(filename),
				value,
				Map.of("Content-Type", contentType)));
		return this;
	}

	public MultiPartBodyPublisher addPart(String name, AiFile file) {
		parts.add(new Part(
				name,
				Optional.ofNullable(file.fileName()),
				file::data,
				Map.of("Content-Type", file.mimeType())));
		return this;
	}


	private static record Part(
			String name,
			Optional filename,
			Supplier dataSource,
			Map headers) {

	}


	private class BodyBytesIterator implements Iterator {
		private static final String LINE_CHANGE = "\r\n";

		private int index = 0;
		private boolean atBoundary = true;


		@Override
		public boolean hasNext() {
			return atBoundary || hasMoreParts();
		}


		private boolean hasMoreParts() {
			return index < parts.size();
		}


		@Override
		public byte[] next() {
			byte[] result;
			if (atBoundary) {
				ByteArrayOutputStream bout = new ByteArrayOutputStream();
				try (Writer writer = new OutputStreamWriter(bout, StandardCharsets.UTF_8)) {
					if (index > 0)
						writer.append(LINE_CHANGE);
					writer.append("--").append(boundary);
					if (!hasMoreParts())
						writer.append("--");
					writer.append(LINE_CHANGE);

					if (hasMoreParts()) {
						Part part = parts.get(index);
						writer.append("Content-Disposition: form-data; name=\"").append(part.name().replaceAll("\"", "%22")).append('"');
						if (part.filename.isPresent())
							writer.append("; filename=\"").append(part.filename().get().replaceAll("\"", "%22")).append('"');
						writer.append(LINE_CHANGE);
						for (var entry: part.headers().entrySet())
							writer.append(entry.getKey()).append(": ").append(entry.getValue()).append(LINE_CHANGE);
						writer.append(LINE_CHANGE);
					}
				}
				catch (IOException e) {
					throw new UncheckedIOException(e);
				}

				result = bout.toByteArray();
				atBoundary = false;
			}
			else {
				result = parts.get(index).dataSource().get();
				atBoundary = true;
				++index;
			}
			return result;
		}

	}

}