Send each message's custom headers on pooled S3 connections

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()) {

server/src/main/java/com/mirth/connect/connectors/file/filesystems/S3Connection.java

@@ -157,6 +157,9 @@ public void populateSourceMap(Map<String, Object> 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<String, List<String>> 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<String, List<String>> customHeaders) {
+        this.customHeaders = customHeaders;
+    }
+
     Map<String, String> getCustomHeaders() {
         Map<String, String> headers = new HashMap<String, String>();
+        Map<String, List<String>> source = customHeaders != null ? customHeaders : schemeProps.getCustomHeaders();
 
-        if (MapUtils.isNotEmpty(schemeProps.getCustomHeaders())) {
-            for (Entry<String, List<String>> entry : schemeProps.getCustomHeaders().entrySet()) {
+        if (MapUtils.isNotEmpty(source)) {
+            for (Entry<String, List<String>> entry : source.entrySet()) {
                 for (String value : entry.getValue()) {
                     headers.put(entry.getKey(), value);
                 }

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<PutObjectRequest> requests = ArgumentCaptor.forClass(PutObjectRequest.class);
+        verify(client, times(2)).putObject(requests.capture(), any(RequestBody.class));
+        List<String> sent = new ArrayList<String>();
+        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;
+        }
+    }
+}