diff --git a/server/src/main/java/com/mirth/connect/connectors/file/FileDispatcher.java b/server/src/main/java/com/mirth/connect/connectors/file/FileDispatcher.java index 29139928df..41d46eab51 100644 --- a/server/src/main/java/com/mirth/connect/connectors/file/FileDispatcher.java +++ b/server/src/main/java/com/mirth/connect/connectors/file/FileDispatcher.java @@ -19,6 +19,7 @@ import org.apache.logging.log4j.Logger; import com.mirth.connect.connectors.file.filesystems.FileSystemConnection; +import com.mirth.connect.connectors.file.filesystems.S3Connection; import com.mirth.connect.donkey.model.channel.ConnectorProperties; import com.mirth.connect.donkey.model.event.ConnectionStatusEventType; import com.mirth.connect.donkey.model.event.ErrorEventType; @@ -165,6 +166,10 @@ public Response send(ConnectorProperties connectorProperties, ConnectorMessage c ThreadUtils.checkInterruptedStatus(); fileSystemConnection = fileConnector.getConnection(fileSystemOptions); + if (fileSystemConnection instanceof S3Connection && fileDispatcherProperties.getSchemeProperties() instanceof S3SchemeProperties) { + // A pooled connection keeps the headers it was created with; send this message's + ((S3Connection) fileSystemConnection).setCustomHeaders(((S3SchemeProperties) fileDispatcherProperties.getSchemeProperties()).getCustomHeaders()); + } if (fileDispatcherProperties.isErrorOnExists() && fileSystemConnection.exists(filename, path)) { throw new IOException("Destination file already exists, will not overwrite."); } else if (fileDispatcherProperties.isTemporary()) { diff --git a/server/src/main/java/com/mirth/connect/connectors/file/filesystems/S3Connection.java b/server/src/main/java/com/mirth/connect/connectors/file/filesystems/S3Connection.java index f988649dd6..04563e59d0 100644 --- a/server/src/main/java/com/mirth/connect/connectors/file/filesystems/S3Connection.java +++ b/server/src/main/java/com/mirth/connect/connectors/file/filesystems/S3Connection.java @@ -157,6 +157,9 @@ public void populateSourceMap(Map sourceMap) { S3Client client; StsClient sts; int stsDuration; + // Set by the File Writer for each message. schemeProps belongs to the message that created + // the pool and is shared by every connection in it, so per-message headers can't go there. + Map> customHeaders; public S3Connection(FileSystemConnectionOptions fileSystemOptions, int timeout) throws Exception { this.fileSystemOptions = fileSystemOptions; @@ -271,11 +274,16 @@ ListObjectsV2Request.Builder createListRequest(String bucketName, String prefix) return ListObjectsV2Request.builder().bucket(bucketName).prefix(prefix).delimiter(DELIMITER); } + public void setCustomHeaders(Map> customHeaders) { + this.customHeaders = customHeaders; + } + Map getCustomHeaders() { Map headers = new HashMap(); + Map> source = customHeaders != null ? customHeaders : schemeProps.getCustomHeaders(); - if (MapUtils.isNotEmpty(schemeProps.getCustomHeaders())) { - for (Entry> entry : schemeProps.getCustomHeaders().entrySet()) { + if (MapUtils.isNotEmpty(source)) { + for (Entry> entry : source.entrySet()) { for (String value : entry.getValue()) { headers.put(entry.getKey(), value); } diff --git a/server/src/test/java/com/mirth/connect/connectors/file/FileDispatcherTest.java b/server/src/test/java/com/mirth/connect/connectors/file/FileDispatcherTest.java new file mode 100644 index 0000000000..d4870faf9e --- /dev/null +++ b/server/src/test/java/com/mirth/connect/connectors/file/FileDispatcherTest.java @@ -0,0 +1,124 @@ +// SPDX-License-Identifier: MPL-2.0 +// SPDX-FileCopyrightText: Open Integration Engine + +package com.mirth.connect.connectors.file; + +import static org.junit.Assert.assertEquals; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyBoolean; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import java.lang.reflect.Field; +import java.net.URI; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; + +import org.junit.BeforeClass; +import org.junit.Test; +import org.mockito.ArgumentCaptor; + +import com.google.inject.AbstractModule; +import com.google.inject.Guice; +import com.google.inject.Injector; +import com.mirth.connect.connectors.file.filesystems.S3Connection; +import com.mirth.connect.donkey.model.message.ConnectorMessage; +import com.mirth.connect.donkey.model.message.Status; +import com.mirth.connect.donkey.model.message.attachment.AttachmentHandlerProvider; +import com.mirth.connect.server.controllers.ConfigurationController; +import com.mirth.connect.server.controllers.ControllerFactory; +import com.mirth.connect.server.controllers.EventController; + +import software.amazon.awssdk.core.sync.RequestBody; +import software.amazon.awssdk.regions.Region; +import software.amazon.awssdk.services.s3.S3Client; +import software.amazon.awssdk.services.s3.model.PutObjectRequest; +import software.amazon.awssdk.services.s3.model.PutObjectResponse; + +public class FileDispatcherTest { + + @BeforeClass + public static void setUpBeforeClass() throws Exception { + ControllerFactory controllerFactory = mock(ControllerFactory.class); + EventController eventController = mock(EventController.class); + when(controllerFactory.createEventController()).thenReturn(eventController); + ConfigurationController configurationController = mock(ConfigurationController.class); + when(controllerFactory.createConfigurationController()).thenReturn(configurationController); + + Injector injector = Guice.createInjector(new AbstractModule() { + @Override + protected void configure() { + requestStaticInjection(ControllerFactory.class); + bind(ControllerFactory.class).toInstance(controllerFactory); + } + }); + injector.getInstance(ControllerFactory.class); + } + + /* + * The pool hands the same S3Connection to every message, so each write must carry the custom + * headers of the message being written, not those of the message that created the connection. + */ + @Test + public void pooledS3ConnectionSendsEachMessagesHeaders() throws Exception { + S3Connection connection = new S3Connection(new FileSystemConnectionOptions(true, null, null, schemeProperties("1")), 10000); + S3Client client = mock(S3Client.class); + when(client.putObject(any(PutObjectRequest.class), any(RequestBody.class))).thenReturn(PutObjectResponse.builder().build()); + // The client field is package-private to the filesystems package, and FileConnector's + // getConnection is protected to this one, so the test sets the client by reflection. + Field clientField = S3Connection.class.getDeclaredField("client"); + clientField.setAccessible(true); + clientField.set(connection, client); + + FileConnector fileConnector = mock(FileConnector.class); + URI uri = new URI("s3://bucket/dir"); + when(fileConnector.getEndpointURI(any(), any(), any(), anyBoolean())).thenReturn(uri); + when(fileConnector.getPathPart(any())).thenReturn("bucket/dir"); + when(fileConnector.getConnection(any())).thenReturn(connection); + + TestFileDispatcher dispatcher = new TestFileDispatcher(); + dispatcher.setFileConnector(fileConnector); + + assertEquals(Status.SENT, dispatcher.send(dispatcherProperties("1"), new ConnectorMessage()).getStatus()); + assertEquals(Status.SENT, dispatcher.send(dispatcherProperties("2"), new ConnectorMessage()).getStatus()); + + ArgumentCaptor requests = ArgumentCaptor.forClass(PutObjectRequest.class); + verify(client, times(2)).putObject(requests.capture(), any(RequestBody.class)); + List sent = new ArrayList(); + for (PutObjectRequest request : requests.getAllValues()) { + sent.add(request.metadata().get("x-test-id")); + } + assertEquals(List.of("1", "2"), sent); + } + + private static S3SchemeProperties schemeProperties(String headerValue) { + S3SchemeProperties schemeProperties = new S3SchemeProperties(); + schemeProperties.setRegion(Region.US_WEST_2.id()); + schemeProperties.getCustomHeaders().put("x-test-id", Collections.singletonList(headerValue)); + return schemeProperties; + } + + private static FileDispatcherProperties dispatcherProperties(String headerValue) { + FileDispatcherProperties properties = new FileDispatcherProperties(); + properties.setScheme(FileScheme.S3); + properties.setSchemeProperties(schemeProperties(headerValue)); + properties.setHost("bucket/dir"); + properties.setAnonymous(true); + properties.setOutputPattern("message.txt"); + properties.setTemplate("message"); + properties.setKeepConnectionOpen(true); + return properties; + } + + private static class TestFileDispatcher extends FileDispatcher { + @Override + protected AttachmentHandlerProvider getAttachmentHandlerProvider() { + AttachmentHandlerProvider provider = mock(AttachmentHandlerProvider.class); + when(provider.reAttachMessage(any(), any(), any(), anyBoolean(), anyBoolean())).thenReturn("message".getBytes()); + return provider; + } + } +}