Test unsafe rewind input stream

microsoft · Aug 21, 2024 · a777477 · a777477
1 parent 6291fa2
commit a777477
Show file tree

Hide file tree

Showing 2 changed files with 62 additions and 2 deletions.
diff --git a/components/http/okHttp/src/main/java/com/microsoft/kiota/http/OkHttpRequestAdapter.java b/components/http/okHttp/src/main/java/com/microsoft/kiota/http/OkHttpRequestAdapter.java
@@ -26,7 +26,7 @@
 import jakarta.annotation.Nullable;
 
 import okhttp3.*;
-
+import okio.Buffer;
 import okio.BufferedSink;
 import okio.Okio;
 
@@ -923,6 +923,7 @@ public long contentLength() throws IOException {
                                 @Override
                                 public void writeTo(@Nonnull BufferedSink sink) throws IOException {
                                     sink.writeAll(Okio.source(requestInfo.content));
+                                    requestInfo.content.reset(); // unsafe
                                 }
                             };
 

diff --git a/components/http/okHttp/src/test/java/com/microsoft/kiota/http/OkHttpRequestAdapterTest.java b/components/http/okHttp/src/test/java/com/microsoft/kiota/http/OkHttpRequestAdapterTest.java
@@ -25,7 +25,7 @@
 import okhttp3.Request;
 import okhttp3.Response;
 import okhttp3.ResponseBody;
-
+import okio.Buffer;
 import okio.Okio;
 
 import org.junit.jupiter.api.Test;
@@ -401,6 +401,65 @@ void getRequestFromRequestInformationWithoutContentLengthOverrideWithEmptyPayloa
         assertEquals(0, request.body().contentLength());
     }
 
+	@Test
+	void getRequestFromRequestInformationBuildsRequestSupportingMultipleRequestBodyWritesWithByteArrayStream() throws Exception {
+        final var authenticationProviderMock = mock(AuthenticationProvider.class);
+        final var requestInformation = new RequestInformation();
+        requestInformation.setUri(new URI("https://localhost"));
+        var requestBodyJson = "{\"name\":\"value\",\"array\":[\"1\",\"2\",\"3\"]}";
+        ByteArrayInputStream content =
+                new ByteArrayInputStream(requestBodyJson.getBytes(StandardCharsets.UTF_8));
+        requestInformation.setStreamContent(content, "application/json");
+        requestInformation.httpMethod = HttpMethod.PUT;
+
+        final var adapter = new OkHttpRequestAdapter(authenticationProviderMock);
+        final var request =
+                adapter.getRequestFromRequestInformation(
+                        requestInformation, mock(Span.class), mock(Span.class));
+
+        final var requestBody = request.body();
+        assertNotNull(requestBody);
+        var buffer = new Buffer();
+        requestBody.writeTo(buffer);
+        assertEquals(requestBodyJson, buffer.readUtf8());
+
+        // Second write to the buffer to ensure the body is not consumed
+        buffer = new Buffer();
+        requestBody.writeTo(buffer);
+        assertEquals(requestBodyJson, buffer.readUtf8());
+	}
+
+    @Test
+	void getRequestFromRequestInformationBuildsRequestSupportingMultipleRequestBodyWritesWithFileStream() throws Exception {
+        final var authenticationProviderMock = mock(AuthenticationProvider.class);
+        final var testFile = new File("./src/test/resources/helloWorld.txt");
+        final var requestInformation = new RequestInformation();
+
+        requestInformation.setUri(new URI("https://localhost"));
+        requestInformation.httpMethod = HttpMethod.PUT;
+        final var contentLength = testFile.length();
+        requestInformation.headers.add("Content-Length", String.valueOf(contentLength));
+        try (FileInputStream content = new FileInputStream(testFile)) {
+            requestInformation.setStreamContent(content, "application/octet-stream");
+
+            final var adapter = new OkHttpRequestAdapter(authenticationProviderMock);
+            final var request =
+                    adapter.getRequestFromRequestInformation(
+                            requestInformation, mock(Span.class), mock(Span.class));
+
+            final var requestBody = request.body();
+            assertNotNull(requestBody);
+            var buffer = new Buffer();
+            requestBody.writeTo(buffer);
+            assertEquals(contentLength, buffer.size());
+
+            // Second write to the buffer to ensure the body is not consumed
+            buffer = new Buffer();
+            requestBody.writeTo(buffer);
+            assertEquals(contentLength, buffer.size());
+        }
+	}
+
     public static OkHttpClient getMockClient(final Response response) throws IOException {
         final OkHttpClient mockClient = mock(OkHttpClient.class);
         final Call remoteCall = mock(Call.class);