From 60ad8b983f11db19a0b3908c2491af61110844a5 Mon Sep 17 00:00:00 2001
From: mloufra <mloufra@amazon.com>
Date: Mon, 17 Oct 2022 18:47:12 +0000
Subject: [PATCH] update changelog and narrow type for TransportAction

Signed-off-by: mloufra <mloufra@amazon.com>
---
 CHANGELOG.md                                    |  1 +
 .../RegisterTransportActionsRequest.java        | 17 +++++++++++------
 2 files changed, 12 insertions(+), 6 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 0a2d48bddcd86..fddc9b5616994 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -176,6 +176,7 @@ Inspired from [Keep a Changelog](https://keepachangelog.com/en/1.0.0/)
  - Return consumed params and content from extensions ([#4705](https://github.com/opensearch-project/OpenSearch/pull/4705))
  - Modified EnvironmentSettingsRequest to pass entire Settings object ([#4731](https://github.com/opensearch-project/OpenSearch/pull/4731))
  - Added contentParser method to ExtensionRestRequest ([#4760](https://github.com/opensearch-project/OpenSearch/pull/4760))
+ - Enforce type safety for RegisterTransportActionsRequest([#4796](https://github.com/opensearch-project/OpenSearch/pull/4796))
 
 ## [2.x]
 
diff --git a/server/src/main/java/org/opensearch/extensions/RegisterTransportActionsRequest.java b/server/src/main/java/org/opensearch/extensions/RegisterTransportActionsRequest.java
index bd2535b688c99..3c5f0b84a763d 100644
--- a/server/src/main/java/org/opensearch/extensions/RegisterTransportActionsRequest.java
+++ b/server/src/main/java/org/opensearch/extensions/RegisterTransportActionsRequest.java
@@ -8,6 +8,9 @@
 
 package org.opensearch.extensions;
 
+import org.opensearch.action.ActionRequest;
+import org.opensearch.action.ActionResponse;
+import org.opensearch.action.support.TransportAction;
 import org.opensearch.common.io.stream.StreamInput;
 import org.opensearch.common.io.stream.StreamOutput;
 import org.opensearch.transport.TransportRequest;
@@ -16,6 +19,7 @@
 import java.util.HashMap;
 import java.util.Map;
 import java.util.Objects;
+import java.util.Map.Entry;
 
 /**
  * Request to register extension Transport actions
@@ -24,9 +28,9 @@
  */
 public class RegisterTransportActionsRequest extends TransportRequest {
     private String uniqueId;
-    private Map<String, Class<?>> transportActions;
+    private Map<String, Class<? extends TransportAction<? extends ActionRequest, ? extends ActionResponse>>> transportActions;
 
-    public RegisterTransportActionsRequest(String uniqueId, Map<String, Class<?>> transportActions) {
+    public RegisterTransportActionsRequest(String uniqueId, Map<String, Class<? extends TransportAction<? extends ActionRequest, ? extends ActionResponse>>> transportActions) {
         this.uniqueId = uniqueId;
         this.transportActions = new HashMap<>(transportActions);
     }
@@ -34,12 +38,13 @@ public RegisterTransportActionsRequest(String uniqueId, Map<String, Class<?>> tr
     public RegisterTransportActionsRequest(StreamInput in) throws IOException {
         super(in);
         this.uniqueId = in.readString();
-        Map<String, Class<?>> actions = new HashMap<>();
+        Map<String, Class<? extends TransportAction<? extends ActionRequest, ? extends ActionResponse>>> actions = new HashMap<>();
         int actionCount = in.readVInt();
         for (int i = 0; i < actionCount; i++) {
             try {
                 String actionName = in.readString();
-                Class<?> transportAction = Class.forName(in.readString());
+                @SuppressWarnings("unchecked")
+                Class<? extends TransportAction<? extends ActionRequest, ? extends ActionResponse>> transportAction = (Class<? extends TransportAction<? extends ActionRequest, ? extends ActionResponse>>) Class.forName(in.readString());
                 actions.put(actionName, transportAction);
             } catch (ClassNotFoundException e) {
                 throw new IllegalArgumentException("Could not read transport action");
@@ -52,7 +57,7 @@ public String getUniqueId() {
         return uniqueId;
     }
 
-    public Map<String, Class<?>> getTransportActions() {
+    public Map<String, Class<? extends TransportAction<? extends ActionRequest, ? extends ActionResponse>>> getTransportActions() {
         return transportActions;
     }
 
@@ -61,7 +66,7 @@ public void writeTo(StreamOutput out) throws IOException {
         super.writeTo(out);
         out.writeString(uniqueId);
         out.writeVInt(this.transportActions.size());
-        for (Map.Entry<String, Class<?>> action : transportActions.entrySet()) {
+        for (Entry<String, Class<? extends TransportAction<? extends ActionRequest, ? extends ActionResponse>>> action : transportActions.entrySet()) {
             out.writeString(action.getKey());
             out.writeString(action.getValue().getName());
         }