package com.mewna.catnip.rest; import lombok.Getter; import lombok.RequiredArgsConstructor; import lombok.Setter; import lombok.experimental.Accessors; import javax.annotation.ParametersAreNonnullByDefault; import java.net.http.HttpRequest.BodyPublisher; import java.net.http.HttpRequest.BodyPublishers; import java.nio.charset.StandardCharsets; import java.util.*; /** * Adapted from https://stackoverflow.com/a/54675316 * Modified to fit the code standards of this project, to compact it, to remove * unnecessary capabilities (file read straight from disc and InputStream * suppliers), and add capabilities necessary to add compatibility. * * @author kjp12 * @since 3/25/2019 */ public class MultipartBodyPublisher { private final Collection<PartsSpecification> partsSpecificationList = new ArrayList<>(); @Getter private final String boundary = UUID.randomUUID().toString(); public BodyPublisher build() { if(partsSpecificationList.isEmpty()) { throw new IllegalStateException("Must have at least one part to build multipart message."); } addFinalBoundary(); return BodyPublishers.ofByteArrays(PartsIterator::new); } @ParametersAreNonnullByDefault public MultipartBodyPublisher addPart(final String name, final String value) { partsSpecificationList.add(new PartsSpecification(Type.STRING, name).value(value.getBytes(StandardCharsets.UTF_8))); return this; } @ParametersAreNonnullByDefault public MultipartBodyPublisher addPart(final String name, final String filename, final byte[] value) { partsSpecificationList.add(new PartsSpecification(Type.FILE, name).filename(filename).value(value)); return this; } private void addFinalBoundary() { partsSpecificationList.add(new PartsSpecification(Type.FINAL_BOUNDARY, null)); } public enum Type { STRING, FILE, FINAL_BOUNDARY } @RequiredArgsConstructor @Setter @Accessors(fluent = true) protected class PartsSpecification { protected final Type type; protected final String name; protected byte[] value; protected String filename; public String toString() { if(type == Type.FINAL_BOUNDARY) { return "--" + boundary + "--"; } if(type == Type.FILE) { return "--" + boundary + "\r\nContent-Disposition: file; name=" + name + "; filename=" + filename + ";\r\nContent-Type:application/octet-stream\r\n\r\n"; } return "--" + boundary + "\r\nContent-Disposition: form-data; name=" + name + ";\r\nContent-Type: text/plain; charset=UTF-8\r\n\r\n"; } } class PartsIterator implements Iterator<byte[]> { private final Iterator<PartsSpecification> parts; private boolean done; private final List<byte[]> next = new ArrayList<>(); PartsIterator() { parts = partsSpecificationList.iterator(); } @Override public boolean hasNext() { if(done) { return false; } if(!next.isEmpty()) { return true; } computeNext(); if(next.isEmpty()) { done = true; return false; } return true; } @Override public byte[] next() { if(!hasNext()) { throw new NoSuchElementException(); } return next.remove(0); } private void computeNext() { if(!parts.hasNext()) { return; } final var part = parts.next(); next.add(part.toString().getBytes(StandardCharsets.UTF_8)); if(part.type != Type.FINAL_BOUNDARY) { next.add(part.value); next.add(new byte[] {'\r', '\n'}); } } } }