Support multi-part HTTP data with JDK 11 HttpClient (streaming)

3 files changed, 225 insertions, 1 deletions

diff --git a/vespa-maven-plugin/src/main/java/ai/vespa/hosted/api/Method.java b/vespa-maven-plugin/src/main/java/ai/vespa/hosted/api/Method.java
index 17fe2d2640a..ff7c1e4270b 100644
--- a/vespa-maven-plugin/src/main/java/ai/vespa/hosted/api/Method.java
+++ b/vespa-maven-plugin/src/main/java/ai/vespa/hosted/api/Method.java
@@ -1,5 +1,16 @@
 package ai.vespa.hosted.api;
 
-public class Method {
+/**
+ * HTTP methods.
+ *
+ * @author jonmv
+ */
+public enum Method {
+
+    GET,
+    PUT,
+    POST,
+    PATCH,
+    DELETE;
 
 }
diff --git a/vespa-maven-plugin/src/main/java/ai/vespa/hosted/api/MultiPartStreamer.java b/vespa-maven-plugin/src/main/java/ai/vespa/hosted/api/MultiPartStreamer.java
new file mode 100644
index 00000000000..1b030190289
--- /dev/null
+++ b/vespa-maven-plugin/src/main/java/ai/vespa/hosted/api/MultiPartStreamer.java
@@ -0,0 +1,144 @@
+package ai.vespa.hosted.api;
+
+import java.io.BufferedInputStream;
+import java.io.ByteArrayInputStream;
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.SequenceInputStream;
+import java.io.UncheckedIOException;
+import java.net.http.HttpRequest;
+import java.nio.charset.StandardCharsets;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import java.util.UUID;
+import java.util.function.Supplier;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
+
+/**
+ * Used to create builders for multi part http body entities, which stream their data.
+ *
+ * @author jonmv
+ */
+public class MultiPartStreamer {
+
+    private final String boundary;
+    private final List<Supplier<InputStream>> streams;
+
+    MultiPartStreamer(String boundary) {
+        this.boundary = boundary;
+        this.streams = new ArrayList<>();
+    }
+
+    /** Creates a new MultiPartBodyStreamer in which parts can be aggregated, and then streamed. */
+    public MultiPartStreamer() {
+        this(UUID.randomUUID().toString());
+    }
+
+    /** Adds the given data as a named part in this, using the {@code "text/plain"} content type. */
+    public MultiPartStreamer addText(String name, String json) {
+        return addData(name, "text/plain", json);
+    }
+
+    /** Adds the given data as a named part in this, using the {@code "application/json"} content type. */
+    public MultiPartStreamer addJson(String name, String json) {
+        return addData(name, "application/json", json);
+    }
+
+    /** Adds the given data as a named part in this, using the given content type. */
+    public MultiPartStreamer addData(String name, String type, String data) {
+        streams.add(() -> separator(name, type));
+        streams.add(() -> asStream(data));
+
+        return this;
+    }
+
+    /** Adds the contents of the file at the given path as a named part in this. */
+    public MultiPartStreamer addFile(String name, Path path) {
+        streams.add(() -> separator(name, path));
+        streams.add(() -> asStream(path));
+
+        return this;
+    }
+
+    /**
+     * Returns a builder whose data is an aggregate stream of the current parts of this.
+     * Modifications to this streamer after a request builder has been obtained is not reflected in that builder.
+     * This method can be used multiple times, to create new requests.
+     * The request builder's method and content should not be set after it has been obtained.
+     */
+    public HttpRequest.Builder newBuilderFor(Method method) {
+        InputStream aggregate = aggregate(); // Get the streams now, not when the aggregate is used.
+        return HttpRequest.newBuilder()
+                          .setHeader("Content-Type", "multipart/form-data; boundary=" + boundary)
+                          .method(method.name(), HttpRequest.BodyPublishers.ofInputStream(() -> aggregate));
+    }
+
+    /** Returns an input stream which is an aggregate of all current parts in this, plus an end marker. */
+    InputStream aggregate() {
+        InputStream aggregate = new SequenceInputStream(Collections.enumeration(Stream.concat(streams.stream().map(Supplier::get),
+                                                                                              Stream.of(end()))
+                                                                                      .collect(Collectors.toList())));
+
+        try {
+            if (aggregate.skip(2) != 2)// This should never happen, as the first stream is a ByteArrayInputStream.
+                throw new IllegalStateException("Failed skipping extraneous bytes.");
+        }
+        catch (IOException e) { // This should never happen, as the first stream is a ByteArrayInputStream;
+            throw new IllegalStateException("Failed skipping extraneous bytes.", e);
+        }
+        return new BufferedInputStream(aggregate);
+    }
+
+    /** Returns the separator to put between one part and the next, when this is a string. */
+    private InputStream separator(String name, String contentType) {
+        return asStream(disposition(name) + type(contentType));
+    }
+
+    /** Returns the separator to put between one part and the next, when this is a file. */
+    private InputStream separator(String name, Path path) {
+        try {
+            String contentType = Files.probeContentType(path);
+            return asStream(disposition(name) + "; filename=\"" + path.getFileName() + "\"" +
+                            type(contentType != null ? contentType : "application/octet-stream"));
+        }
+        catch (IOException e) {
+            throw new UncheckedIOException(e);
+        }
+    }
+
+    /** Returns the end delimiter of the request, with line breaks prepended. */
+    private InputStream end() {
+        return asStream("\r\n--" + boundary + "--");
+    }
+
+    /** Returns the boundary and disposition header for a part, with line breaks prepended. */
+    private String disposition(String name) {
+        return "\r\n--" + boundary + "\r\n" +
+               "Content-Disposition: form-data; name=\"" + name + "\"";
+    }
+
+    /** Returns the content type header for a part, with line breaks pre- and appended. */
+    private String type(String contentType) {
+        return "\r\nContent-Type: " + contentType + "\r\n\r\n";
+    }
+
+    /** Returns the a ByteArrayInputStream over the given string, UTF-8 encoded. */
+    private static InputStream asStream(String string) {
+        return new ByteArrayInputStream(string.getBytes(StandardCharsets.UTF_8));
+    }
+
+    /** Returns an InputStream over the file at the given path — rethrows any IOException as UncheckedIOException. */
+    private InputStream asStream(Path path) {
+        try {
+            return Files.newInputStream(path);
+        }
+        catch (IOException e) {
+            throw new UncheckedIOException(e);
+        }
+    }
+
+}
diff --git a/vespa-maven-plugin/src/test/java/ai/vespa/hosted/api/MultiPartStreamerTest.java b/vespa-maven-plugin/src/test/java/ai/vespa/hosted/api/MultiPartStreamerTest.java
new file mode 100644
index 00000000000..b3915c51925
--- /dev/null
+++ b/vespa-maven-plugin/src/test/java/ai/vespa/hosted/api/MultiPartStreamerTest.java
@@ -0,0 +1,69 @@
+package ai.vespa.hosted.api;
+
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.TemporaryFolder;
+
+import java.io.IOException;
+import java.net.URI;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.util.List;
+
+import static org.junit.Assert.assertEquals;
+
+public class MultiPartStreamerTest {
+
+    @Rule
+    public TemporaryFolder tmp = new TemporaryFolder();
+
+    @Test
+    public void test() throws IOException {
+        Path file = tmp.newFile().toPath();
+        Files.write(file, new byte[]{0x48, 0x69});
+        MultiPartStreamer streamer = new MultiPartStreamer("My boundary");
+
+        assertEquals("--My boundary--",
+                     new String(streamer.aggregate().readAllBytes()));
+
+        streamer.addData("data", "uss/enterprise", "lore")
+                .addJson("json", "{\"xml\":false}")
+                .addText("text", "Hello!")
+                .addFile("file", file);
+
+        String expected = "--My boundary\r\n" +
+                          "Content-Disposition: form-data; name=\"data\"\r\n" +
+                          "Content-Type: uss/enterprise\r\n" +
+                          "\r\n" +
+                          "lore\r\n" +
+                          "--My boundary\r\n" +
+                          "Content-Disposition: form-data; name=\"json\"\r\n" +
+                          "Content-Type: application/json\r\n" +
+                          "\r\n" +
+                          "{\"xml\":false}\r\n" +
+                          "--My boundary\r\n" +
+                          "Content-Disposition: form-data; name=\"text\"\r\n" +
+                          "Content-Type: text/plain\r\n" +
+                          "\r\n" +
+                          "Hello!\r\n" +
+                          "--My boundary\r\n" +
+                          "Content-Disposition: form-data; name=\"file\"; filename=\"" + file.getFileName() + "\"\r\n" +
+                          "Content-Type: application/octet-stream\r\n" +
+                          "\r\n" +
+                          "Hi\r\n" +
+                          "--My boundary--";
+
+        assertEquals(expected,
+                     new String(streamer.aggregate().readAllBytes()));
+
+        // Verify that all data is read again for a new builder.
+        assertEquals(expected,
+                     new String(streamer.aggregate().readAllBytes()));
+
+        assertEquals(List.of("multipart/form-data; boundary=My boundary"),
+                     streamer.newBuilderFor(Method.POST)
+                             .uri(URI.create("https://uri/path"))
+                             .build().headers().allValues("Content-Type"));
+    }
+
+}