Extend Avro Converter to support Cross-Account schema registry access (#376)

* Ability to assume role for cross account schema access in avro-kafkaconnect-converter

* Update readme with cross account details

* Simplify logic by dynamically assuming role in existing converter class if arn is provided

* Review changes

Author: mohamedq-mo

README.md

@@ -635,6 +635,30 @@ repository for the latest support: [Avro SerializationSchema and Deserialization
             properties);
     DataStream<GenericRecord> stream = env.addSource(consumer);
 ```
+
+## Cross-Account Avro Converter Support
+
+The `AWSKafkaAvroConverter` Avro converter is able to assume an IAM role in a different AWS account before accessing Glue Schema Registry. You can configure the role ARN and an optional session name.
+
+If `assumeRoleArn` is not provided, the converter will fallback to using the default credentials associated to the host.
+
+### Connector configuration
+
+Include these properties in your Kafka Connect worker or connector config:
+
+```properties
+# Define converter
+key.converter=com.amazonaws.services.schemaregistry.kafkaconnect.AWSKafkaAvroConverter
+value.converter=com.amazonaws.services.schemaregistry.kafkaconnect.AWSKafkaAvroConverter
+
+# Specify cross-account role arn
+key.converter.assumeRoleArn="arn:aws:iam::123456789012:role/my-role"
+value.converter.assumeRoleArn="arn:aws:iam::123456789012:role/my-role"
+
+# Override default session name (optional; default is "kafka-connect-session")
+key.converter.assumeRoleSessionName=my-custom-session
+value.converter.assumeRoleSessionName=my-custom-session
+```
  
  ## Security issue notifications
 If you discover a potential security issue in this project we ask that you notify AWS/Amazon Security via our [vulnerability reporting page](http://aws.amazon.com/security/vulnerability-reporting/). Please do **not** create a public github issue.

avro-kafkaconnect-converter/src/main/java/com/amazonaws/services/schemaregistry/kafkaconnect/AWSKafkaAvroConverter.java

@@ -21,14 +21,20 @@
 import com.amazonaws.services.schemaregistry.kafkaconnect.avrodata.AvroData;
 import com.amazonaws.services.schemaregistry.kafkaconnect.avrodata.AvroDataConfig;
 import com.amazonaws.services.schemaregistry.serializers.avro.AWSKafkaAvroSerializer;
-
+import com.amazonaws.services.schemaregistry.utils.AWSSchemaRegistryConstants;
+import com.google.common.annotations.VisibleForTesting;
 import lombok.Data;
 import lombok.extern.slf4j.Slf4j;
 import org.apache.kafka.common.errors.SerializationException;
 import org.apache.kafka.connect.data.Schema;
 import org.apache.kafka.connect.data.SchemaAndValue;
 import org.apache.kafka.connect.errors.DataException;
 import org.apache.kafka.connect.storage.Converter;
+import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider;
+import software.amazon.awssdk.http.urlconnection.UrlConnectionHttpClient;
+import software.amazon.awssdk.regions.Region;
+import software.amazon.awssdk.services.sts.StsClient;
+import software.amazon.awssdk.services.sts.auth.StsAssumeRoleCredentialsProvider;
 
 import java.util.Map;
 
@@ -75,6 +81,21 @@ public void configure(Map<String, ?> configs, boolean isKey) {
         this.isKey = isKey;
         new AWSKafkaAvroConverterConfig(configs);
 
+        //TODO: add this feature to all other converters
+        String roleToAssume = (String) configs.get(AWSSchemaRegistryConstants.ASSUME_ROLE_ARN);
+        if (roleToAssume != null && !roleToAssume.isEmpty()) {
+            String sessionName = configs.get(AWSSchemaRegistryConstants.ASSUME_ROLE_SESSION_NAME) != null
+                    ? configs.get(AWSSchemaRegistryConstants.ASSUME_ROLE_SESSION_NAME).toString()
+                    : "kafka-connect-session";
+
+            String region = configs.get(AWSSchemaRegistryConstants.AWS_REGION).toString();
+
+            AwsCredentialsProvider credentialsProvider = getCredentialsProvider(roleToAssume, sessionName, region);
+
+            deserializer = new AWSKafkaAvroDeserializer(credentialsProvider, configs);
+            serializer = new AWSKafkaAvroSerializer(credentialsProvider, configs);
+        }
+
         serializer.configure(configs, this.isKey);
         deserializer.configure(configs, this.isKey);
 
@@ -123,4 +144,19 @@ public SchemaAndValue toConnectData(String topic, byte[] value) {
 
         return avroData.toConnectData(avroSchema, deserialized);
     }
+
+    @VisibleForTesting
+    protected AwsCredentialsProvider getCredentialsProvider(String roleArn, String sessionName, String region) {
+        UrlConnectionHttpClient.Builder urlConnectionHttpClientBuilder = UrlConnectionHttpClient.builder();
+        StsClient stsClient = StsClient.builder()
+                .httpClient(urlConnectionHttpClientBuilder.build())
+                .region(Region.of(region))
+                .build();
+        return StsAssumeRoleCredentialsProvider.builder()
+                .refreshRequest(assumeRoleRequest -> assumeRoleRequest
+                        .roleArn(roleArn)
+                        .roleSessionName(sessionName))
+                .stsClient(stsClient)
+                .build();
+    }
 }

avro-kafkaconnect-converter/src/test/java/com/amazonaws/services/schemaregistry/kafkaconnect/AWSKafkaAvroConverterTest.java

@@ -46,13 +46,13 @@
 import java.util.Map;
 import java.util.UUID;
 
-import static org.junit.jupiter.api.Assertions.assertEquals;
-import static org.junit.jupiter.api.Assertions.assertNotNull;
-import static org.junit.jupiter.api.Assertions.assertThrows;
-import static org.mockito.ArgumentMatchers.anyMap;
-import static org.mockito.ArgumentMatchers.eq;
-import static org.mockito.Mockito.mock;
-import static org.mockito.Mockito.when;
+import static com.amazonaws.services.schemaregistry.utils.AWSSchemaRegistryConstants.ASSUME_ROLE_ARN;
+import static com.amazonaws.services.schemaregistry.utils.AWSSchemaRegistryConstants.ASSUME_ROLE_SESSION_NAME;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+import static org.junit.jupiter.api.Assertions.*;
+import static org.mockito.ArgumentMatchers.*;
+import static org.mockito.Mockito.*;
 
 /**
  * Unit tests for testing AWSKafkaAvroConverter class.
@@ -78,6 +78,8 @@ public class AWSKafkaAvroConverterTest {
     private static final UUID schemaVersionIdForTesting = UUID.fromString("b7b4a7f0-9c96-4e4a-a687-fb5de9ef0c63");
     private static final byte[] genericBytes = new byte[] {3, 0, -73, -76, -89, -16, -100, -106, 78, 74, -90, -121, -5,
             93, -23, -17, 12, 99, 10, 115, 97, 110, 115, 97, -58, 1, 6, 114, 101, 100};
+    private static final String ROLE_ARN = "arn:aws:iam::123456789012:role/my-role";
+    private static final String REGION = "us-west-2";
 
     @BeforeEach
     public void setUp() {
@@ -149,6 +151,81 @@ public void testConverter_toConnectData_throwsException() {
         assertThrows(DataException.class, () -> converter.toConnectData(testTopic, serializedData));
     }
 
+    /**
+     * Test AWSKafkaAvroConverter when value is null.
+     */
+    @Test
+    void testConverter_toConnectData_NullValue() {
+        converter = spy(new AWSKafkaAvroConverter());
+        assertEquals(SchemaAndValue.NULL, converter.toConnectData(testTopic, null));
+    }
+
+    /**
+     * Test AWSKafkaAvroConverter with assume role.
+     */
+    @Test
+    void testConverter_configure_invokeAssumeRoleWithCustomSession() {
+        configs.put(ASSUME_ROLE_ARN, ROLE_ARN);
+        configs.put(ASSUME_ROLE_SESSION_NAME, "my-session");
+
+        converter = spy(new AWSKafkaAvroConverter());
+        doReturn(mockCredProvider)
+                .when(converter)
+                .getCredentialsProvider(anyString(), anyString(), anyString());
+
+        converter.configure(configs, true);
+
+        verify(converter).getCredentialsProvider(ROLE_ARN, "my-session", REGION);
+        assertTrue(converter.isKey());
+        assertNotNull(converter.getSerializer());
+        assertNotNull(converter.getDeserializer());
+        assertNotNull(converter.getAvroData());
+    }
+
+    /**
+     * Test AWSKafkaAvroConverter assume role, default session name.
+     */
+    @Test
+    void testConverter_configure_defaultSessionNameForAssumeRole() {
+        configs.put(ASSUME_ROLE_ARN, ROLE_ARN);
+
+        converter = spy(new AWSKafkaAvroConverter());
+        doReturn(mockCredProvider)
+                .when(converter)
+                .getCredentialsProvider(anyString(), anyString(), anyString());
+
+        converter.configure(configs, false);
+
+        verify(converter).getCredentialsProvider(ROLE_ARN, "kafka-connect-session", REGION);
+        assertFalse(converter.isKey());
+    }
+
+    /**
+     * Test AWSKafkaAvroConverter assume role empty.
+     */
+    @Test
+    void testConverter_configure_noAssumeRoleIfArnIsEmpty() {
+        configs.put(ASSUME_ROLE_ARN, "");
+
+        converter = spy(new AWSKafkaAvroConverter());
+        converter.configure(configs, false);
+
+        verify(converter, never())
+                .getCredentialsProvider(anyString(), anyString(), anyString());
+    }
+
+    /**
+     * Test AWSKafkaAvroConverter assume role null.
+     */
+    @Test
+    void testConverter_configure_noAssumeRoleIfArnIsNotProvided() {
+        converter = spy(new AWSKafkaAvroConverter());
+        converter.configure(getProperties(), false);
+
+        verify(converter, never())
+                .getCredentialsProvider(anyString(), anyString(), anyString());
+    }
+
     /**
      * To create a AWSKafkaAvroSerializer instance with mocked parameters.
      *
@@ -226,7 +303,7 @@ private Struct createStructRecord() {
     private Map<String, Object> getProperties() {
         Map<String, Object> props = new HashMap<>();
 
-        props.put(AWSSchemaRegistryConstants.AWS_REGION, "us-west-2");
+        props.put(AWSSchemaRegistryConstants.AWS_REGION, REGION);
         props.put(AWSSchemaRegistryConstants.AWS_ENDPOINT, "https://test");
         props.put(AWSSchemaRegistryConstants.SCHEMA_AUTO_REGISTRATION_SETTING, true);
         props.put(AWSSchemaRegistryConstants.AVRO_RECORD_TYPE, AvroRecordType.GENERIC_RECORD.getName());

common/src/main/java/com/amazonaws/services/schemaregistry/utils/AWSSchemaRegistryConstants.java

@@ -172,6 +172,16 @@ public final class AWSSchemaRegistryConstants {
      */
     public static final String USER_AGENT_APP = "userAgentApp";
 
+    /**
+     * IAM Role ARN to assume for accessing the registry
+     */
+    public static final String ASSUME_ROLE_ARN = "assumeRoleArn";
+
+    /**
+     * IAM Role session name for accessing the registry
+     */
+    public static final String ASSUME_ROLE_SESSION_NAME = "assumeRoleSessionName";
+
     /**
      * Private constructor to avoid initialization of the class.
      */