diff --git a/components/serialization/multipart/src/main/java/com/microsoft/kiota/serialization/MultipartSerializationWriter.java b/components/serialization/multipart/src/main/java/com/microsoft/kiota/serialization/MultipartSerializationWriter.java index 634c2ab5e..9478d7502 100644 --- a/components/serialization/multipart/src/main/java/com/microsoft/kiota/serialization/MultipartSerializationWriter.java +++ b/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 void writeObjectValue( @Nullable final T value, @Nonnull final Parsable... additionalValuesToMerge) { Objects.requireNonNull(additionalValuesToMerge); + final List 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); diff --git a/components/serialization/multipart/src/test/java/com/microsoft/kiota/serialization/MultipartSerializationWriterTests.java b/components/serialization/multipart/src/test/java/com/microsoft/kiota/serialization/MultipartSerializationWriterTests.java index 3b198794f..ee4c548b5 100644 --- a/components/serialization/multipart/src/test/java/com/microsoft/kiota/serialization/MultipartSerializationWriterTests.java +++ b/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); + } + } + } + } }