fix(multipart): omit nil header string values and support additionalValuesToMerge

components/serialization/multipart/src/main/java/com/microsoft/kiota/serialization/MultipartSerializationWriter.java

@@ -11,12 +11,14 @@
 import java.io.IOException;
 import java.io.InputStream;
 import java.io.OutputStreamWriter;
-import java.io.UnsupportedEncodingException;
 import java.math.BigDecimal;
+import java.nio.charset.StandardCharsets;
 import java.time.LocalDate;
 import java.time.LocalTime;
 import java.time.OffsetDateTime;
+import java.util.ArrayList;
 import java.util.EnumSet;
+import java.util.List;
 import java.util.Map;
 import java.util.Objects;
 import java.util.UUID;
@@ -26,23 +28,78 @@
 /** Serialization writer implementation for Multipart encoded payloads */
 public class MultipartSerializationWriter implements SerializationWriter {
     private final ByteArrayOutputStream stream = new ByteArrayOutputStream();
-    private final OutputStreamWriter writer;
+    private final OutputStreamWriter writer =
+            new OutputStreamWriter(stream, StandardCharsets.UTF_8);
 
-    /** Instantiates a new MultipartSerializationWriter. */
-    public MultipartSerializationWriter() {
-        try {
-            this.writer = new OutputStreamWriter(this.stream, "UTF-8");
-        } catch (UnsupportedEncodingException e) {
-            throw new RuntimeException("could not create writer", e);
-        }
+    private enum FramingState {
+        EXPECTING_FIRST_BOUNDARY,
+        IN_HEADERS,
+        EXPECTING_PAYLOAD,
+        EXPECTING_DELIMITER_AFTER_PAYLOAD,
+        EXPECTING_BOUNDARY_OR_CLOSING,
+        FINISHED
     }
 
+    private String currentTargetBoundary;
+    private String currentReplacementBoundary;
+    private boolean isLastBody;
+    private FramingState framingState;
+
     public void writeStringValue(@Nullable final String key, @Nullable final String value) {
         try {
+            if (key != null && !key.isEmpty() && value == null) {
+                return;
+            }
+            String valueToWrite = value;
+            if (currentTargetBoundary != null
+                    && currentReplacementBoundary != null
+                    && framingState != null) {
+                switch (framingState) {
+                    case EXPECTING_FIRST_BOUNDARY:
+                        if ((key == null || key.isEmpty())
+                                && value != null
+                                && value.equals("--" + currentTargetBoundary)) {
+                            valueToWrite = "--" + currentReplacementBoundary;
+                            framingState = FramingState.IN_HEADERS;
+                        }
+                        break;
+                    case IN_HEADERS:
+                        if ((key == null || key.isEmpty()) && (value == null || value.isEmpty())) {
+                            framingState = FramingState.EXPECTING_PAYLOAD;
+                        }
+                        break;
+                    case EXPECTING_PAYLOAD:
+                        framingState = FramingState.EXPECTING_DELIMITER_AFTER_PAYLOAD;
+                        break;
+                    case EXPECTING_DELIMITER_AFTER_PAYLOAD:
+                        if ((key == null || key.isEmpty()) && (value == null || value.isEmpty())) {
+                            framingState = FramingState.EXPECTING_BOUNDARY_OR_CLOSING;
+                        }
+                        break;
+                    case EXPECTING_BOUNDARY_OR_CLOSING:
+                        if ((key == null || key.isEmpty()) && value != null) {
+                            if (value.equals("--" + currentTargetBoundary)) {
+                                valueToWrite = "--" + currentReplacementBoundary;
+                                framingState = FramingState.IN_HEADERS;
+                            } else if (value.equals("--" + currentTargetBoundary + "--")) {
+                                if (isLastBody) {
+                                    valueToWrite = "--" + currentReplacementBoundary + "--";
+                                    framingState = FramingState.FINISHED;
+                                } else {
+                                    framingState = FramingState.FINISHED;
+                                    return;
+                                }
+                            }
+                        }
+                        break;
+                    case FINISHED:
+                        break;
+                }
+            }
             if (key != null && !key.isEmpty()) writer.write(key);
-            if (value != null && !value.isEmpty()) {
+            if (valueToWrite != null && !valueToWrite.isEmpty()) {
                 if (key != null && !key.isEmpty()) writer.write(": ");
-                writer.write(value);
+                writer.write(valueToWrite);
             }
             writer.write("\r\n");
             writer.flush();
@@ -125,21 +182,56 @@ public <T extends Parsable> void writeObjectValue(
             @Nullable final T value,
             @Nonnull final Parsable... additionalValuesToMerge) {
         Objects.requireNonNull(additionalValuesToMerge);
+        final List<MultipartBody> bodies = new ArrayList<>();
         if (value != null) {
-            if (onBeforeObjectSerialization != null) {
-                onBeforeObjectSerialization.accept(value);
-            }
             if (value instanceof MultipartBody) {
-                if (onStartObjectSerialization != null) {
-                    onStartObjectSerialization.accept(value, this);
-                }
-                value.serialize(this);
+                bodies.add((MultipartBody) value);
             } else {
                 throw new RuntimeException(
                         "expected MultipartBody instance but got " + value.getClass().getName());
             }
+        }
+        for (final Parsable additionalValueToMerge : additionalValuesToMerge) {
+            if (additionalValueToMerge != null) {
+                if (additionalValueToMerge instanceof MultipartBody) {
+                    bodies.add((MultipartBody) additionalValueToMerge);
+                } else {
+                    throw new RuntimeException(
+                            "expected MultipartBody instance but got "
+                                    + additionalValueToMerge.getClass().getName());
+                }
+            }
+        }
+
+        if (bodies.isEmpty()) {
+            return;
+        }
+
+        final String sharedBoundary = bodies.get(0).getBoundary();
+        for (int i = 0; i < bodies.size(); i++) {
+            final MultipartBody body = bodies.get(i);
+            if (onBeforeObjectSerialization != null) {
+                onBeforeObjectSerialization.accept(body);
+            }
+            if (onStartObjectSerialization != null) {
+                onStartObjectSerialization.accept(body, this);
+            }
+
+            currentTargetBoundary = body.getBoundary();
+            currentReplacementBoundary = sharedBoundary;
+            isLastBody = (i == bodies.size() - 1);
+            framingState = FramingState.EXPECTING_FIRST_BOUNDARY;
+            try {
+                body.serialize(this);
+            } finally {
+                currentTargetBoundary = null;
+                currentReplacementBoundary = null;
+                isLastBody = false;
+                framingState = null;
+            }
+
             if (onAfterObjectSerialization != null) {
-                onAfterObjectSerialization.accept(value);
+                onAfterObjectSerialization.accept(body);
             }
         }
     }
@@ -204,6 +296,9 @@ public void setOnStartObjectSerialization(
 
     public void writeByteArrayValue(
             @Nullable final String key, @Nullable @Nonnull final byte[] value) {
+        if (framingState == FramingState.EXPECTING_PAYLOAD) {
+            framingState = FramingState.EXPECTING_DELIMITER_AFTER_PAYLOAD;
+        }
         if (value != null)
             try {
                 this.stream.write(value);

components/serialization/multipart/src/test/java/com/microsoft/kiota/serialization/MultipartSerializationWriterTests.java

@@ -1,10 +1,12 @@
 package com.microsoft.kiota.serialization;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 
+import com.microsoft.kiota.Compatibility;
 import com.microsoft.kiota.MultipartBody;
 import com.microsoft.kiota.PeriodAndDuration;
 import com.microsoft.kiota.RequestAdapter;
@@ -110,4 +112,154 @@ void writesAStructuredObject() throws IOException {
             }
         }
     }
+
+    @Test
+    void writesStringValueWithKeyValue() throws IOException {
+        try (final var serializationWriter = new MultipartSerializationWriter()) {
+            serializationWriter.writeStringValue("Content-Type", "text/plain");
+            try (final var result = serializationWriter.getSerializedContent()) {
+                try (final var reader =
+                        new BufferedReader(new InputStreamReader(result, "UTF-8"))) {
+                    final String strResult = reader.lines().collect(Collectors.joining("\r\n"));
+                    assertEquals("Content-Type: text/plain", strResult);
+                }
+            }
+        }
+    }
+
+    @Test
+    void writeStringValueWithNilValueOmitsHeader() throws IOException {
+        try (final var serializationWriter = new MultipartSerializationWriter()) {
+            serializationWriter.writeStringValue("Content-Type", null);
+            try (final var result = serializationWriter.getSerializedContent()) {
+                assertEquals(0, result.available());
+            }
+        }
+    }
+
+    @Test
+    void writeStringValueWithEmptyKeyAndNilValueWritesDelimiter() throws IOException {
+        try (final var serializationWriter = new MultipartSerializationWriter()) {
+            serializationWriter.writeStringValue("", null);
+            try (final var result = serializationWriter.getSerializedContent()) {
+                final byte[] bytes = Compatibility.readAllBytes(result);
+                assertEquals("\r\n", new String(bytes, "UTF-8"));
+            }
+        }
+    }
+
+    @Test
+    void writeStringValueWithEmptyKeyAndValueWritesRawLine() throws IOException {
+        try (final var serializationWriter = new MultipartSerializationWriter()) {
+            serializationWriter.writeStringValue("", "raw line");
+            try (final var result = serializationWriter.getSerializedContent()) {
+                final byte[] bytes = Compatibility.readAllBytes(result);
+                assertEquals("raw line\r\n", new String(bytes, "UTF-8"));
+            }
+        }
+    }
+
+    @Test
+    void writesAStructuredObjectWithAdditionalValuesToMerge() throws IOException {
+        final RequestAdapter requestAdapter = mock(RequestAdapter.class);
+        when(requestAdapter.getSerializationWriterFactory())
+                .thenReturn(new JsonSerializationWriterFactory());
+
+        final MultipartBody part1 = new MultipartBody();
+        part1.requestAdapter = requestAdapter;
+        part1.addOrReplacePart("part1", "application/octet-stream", byteForTest);
+
+        final MultipartBody part2 = new MultipartBody();
+        part2.requestAdapter = requestAdapter;
+        part2.addOrReplacePart("part2", "application/octet-stream", new byte[] {0x04, 0x05});
+
+        try (final var serializationWriter = new MultipartSerializationWriter()) {
+            serializationWriter.writeObjectValue(null, part1, part2);
+
+            try (final var result = serializationWriter.getSerializedContent()) {
+                try (final var reader =
+                        new BufferedReader(new InputStreamReader(result, "UTF-8"))) {
+                    final String strResult = reader.lines().collect(Collectors.joining("\r\n"));
+                    final String expected =
+                            "--"
+                                    + part1.getBoundary()
+                                    + "\r\n"
+                                    + "Content-Type: application/octet-stream\r\n"
+                                    + "Content-Disposition: form-data; name=\"part1\"\r\n\r\n"
+                                    + new String(byteForTest, "UTF-8")
+                                    + "\r\n--"
+                                    + part1.getBoundary()
+                                    + "\r\n"
+                                    + "Content-Type: application/octet-stream\r\n"
+                                    + "Content-Disposition: form-data; name=\"part2\"\r\n\r\n"
+                                    + new String(new byte[] {0x04, 0x05}, "UTF-8")
+                                    + "\r\n--"
+                                    + part1.getBoundary()
+                                    + "--";
+                    assertEquals(expected, strResult);
+                    assertFalse(strResult.contains(part2.getBoundary()));
+                }
+            }
+        }
+    }
+
+    @Test
+    void throwsOnInvalidAdditionalValuesToMerge() throws IOException {
+        final TestEntity nonMultipart = new TestEntity();
+
+        try (final var serializationWriter = new MultipartSerializationWriter()) {
+            final var exception =
+                    assertThrows(
+                            RuntimeException.class,
+                            () -> serializationWriter.writeObjectValue(null, null, nonMultipart));
+            assertEquals(
+                    "expected MultipartBody instance but got " + nonMultipart.getClass().getName(),
+                    exception.getMessage());
+        }
+    }
+
+    @Test
+    void preservesStringPayloadMatchingBoundaryOrClosingDelimiterWhenMerging() throws IOException {
+        final RequestAdapter requestAdapter = mock(RequestAdapter.class);
+        when(requestAdapter.getSerializationWriterFactory())
+                .thenReturn(new JsonSerializationWriterFactory());
+
+        final MultipartBody part1 = new MultipartBody();
+        part1.requestAdapter = requestAdapter;
+        final String payload1 = "--" + part1.getBoundary() + "--";
+        part1.addOrReplacePart("part1", "text/plain", payload1);
+
+        final MultipartBody part2 = new MultipartBody();
+        part2.requestAdapter = requestAdapter;
+        final String payload2 = "--" + part2.getBoundary() + "--";
+        part2.addOrReplacePart("part2", "text/plain", payload2);
+
+        try (final var serializationWriter = new MultipartSerializationWriter()) {
+            serializationWriter.writeObjectValue(null, part1, part2);
+
+            try (final var result = serializationWriter.getSerializedContent()) {
+                try (final var reader =
+                        new BufferedReader(new InputStreamReader(result, "UTF-8"))) {
+                    final String strResult = reader.lines().collect(Collectors.joining("\r\n"));
+                    final String expected =
+                            "--"
+                                    + part1.getBoundary()
+                                    + "\r\n"
+                                    + "Content-Type: text/plain\r\n"
+                                    + "Content-Disposition: form-data; name=\"part1\"\r\n\r\n"
+                                    + payload1
+                                    + "\r\n\r\n--"
+                                    + part1.getBoundary()
+                                    + "\r\n"
+                                    + "Content-Type: text/plain\r\n"
+                                    + "Content-Disposition: form-data; name=\"part2\"\r\n\r\n"
+                                    + payload2
+                                    + "\r\n\r\n--"
+                                    + part1.getBoundary()
+                                    + "--";
+                    assertEquals(expected, strResult);
+                }
+            }
+        }
+    }
 }