From 0516680e92867321f5300f5d49ce3953eef77485 Mon Sep 17 00:00:00 2001
From: Nicholas Walter Knize <nknize@apache.org>
Date: Fri, 21 Jul 2023 19:08:52 -0500
Subject: [PATCH] pr review fixes

Signed-off-by: Nicholas Walter Knize <nknize@apache.org>
---
 .../core/common/io/stream/StreamOutput.java   | 10 ------
 .../xcontent/MediaTypeParserRegistry.java     | 31 +++++++++++++------
 2 files changed, 22 insertions(+), 19 deletions(-)

diff --git a/libs/core/src/main/java/org/opensearch/core/common/io/stream/StreamOutput.java b/libs/core/src/main/java/org/opensearch/core/common/io/stream/StreamOutput.java
index c78f00a2f3dfa..566abf9f08f53 100644
--- a/libs/core/src/main/java/org/opensearch/core/common/io/stream/StreamOutput.java
+++ b/libs/core/src/main/java/org/opensearch/core/common/io/stream/StreamOutput.java
@@ -53,12 +53,10 @@
 import org.opensearch.core.common.text.Text;
 import org.opensearch.common.unit.TimeValue;
 import org.opensearch.core.concurrency.OpenSearchRejectedExecutionException;
-import org.opensearch.core.xcontent.MediaType;
 
 import java.io.EOFException;
 import java.io.FileNotFoundException;
 import java.io.IOException;
-import java.io.NotSerializableException;
 import java.io.OutputStream;
 import java.math.BigInteger;
 import java.nio.file.AccessDeniedException;
@@ -497,14 +495,6 @@ public final void writeBigInteger(BigInteger v) throws IOException {
         writeString(v.toString());
     }
 
-    public final void writeMediaType(final MediaType v) throws IOException {
-        if (v.getClass().isEnum()) {
-            writeString(v.mediaType());
-        } else {
-            throw new NotSerializableException("unable to serialize MediaType [" + v.getClass().getSimpleName() + "]");
-        }
-    }
-
     private static byte ZERO = 0;
     private static byte ONE = 1;
     private static byte TWO = 2;
diff --git a/libs/core/src/main/java/org/opensearch/core/xcontent/MediaTypeParserRegistry.java b/libs/core/src/main/java/org/opensearch/core/xcontent/MediaTypeParserRegistry.java
index 7aec74131fd61..2df3b49a0e39e 100644
--- a/libs/core/src/main/java/org/opensearch/core/xcontent/MediaTypeParserRegistry.java
+++ b/libs/core/src/main/java/org/opensearch/core/xcontent/MediaTypeParserRegistry.java
@@ -42,26 +42,39 @@
  * @opensearch.internal
  */
 public final class MediaTypeParserRegistry {
-    private static Map<String, MediaType> formatToMediaType;
-    private static Map<String, MediaType> typeWithSubtypeToMediaType;
+    private static Map<String, MediaType> formatToMediaType = new HashMap<>();
+    private static Map<String, MediaType> typeWithSubtypeToMediaType = new HashMap<>();
 
     // Default mediaType singleton
     private static MediaType DEFAULT_MEDIA_TYPE;
 
     public static void register(MediaType[] acceptedMediaTypes, Map<String, MediaType> additionalMediaTypes) {
-        final int size = acceptedMediaTypes.length + additionalMediaTypes.size();
-        Map<String, MediaType> typeMap = new HashMap<>(size);
-        Map<String, MediaType> formatMap = new HashMap<>(size);
+        // ensures the map is not overwritten:
+        Map<String, MediaType> typeMap = new HashMap<>(typeWithSubtypeToMediaType);
+        Map<String, MediaType> formatMap = new HashMap<>(formatToMediaType);
         for (MediaType mediaType : acceptedMediaTypes) {
+            if (formatMap.containsKey(mediaType.format())) {
+                throw new IllegalArgumentException("unable to register mediaType: [" + mediaType.format() + "]. Type already exists.");
+            }
             typeMap.put(mediaType.typeWithSubtype(), mediaType);
             formatMap.put(mediaType.format(), mediaType);
         }
         for (Map.Entry<String, MediaType> entry : additionalMediaTypes.entrySet()) {
-            String typeWithSubtype = entry.getKey();
-            MediaType mediaType = entry.getValue();
+            String typeWithSubtype = entry.getKey().toLowerCase(Locale.ROOT);
+            if (typeMap.containsKey(typeWithSubtype)) {
+                throw new IllegalArgumentException(
+                    "unable to register mediaType: ["
+                        + entry.getKey()
+                        + "]. "
+                        + "Type already exists and is mapped to: [."
+                        + entry.getValue().format()
+                        + "]"
+                );
+            }
 
-            typeMap.put(typeWithSubtype.toLowerCase(Locale.ROOT), mediaType);
-            formatMap.put(mediaType.format(), mediaType);
+            MediaType mediaType = entry.getValue();
+            typeMap.put(typeWithSubtype, mediaType);
+            formatMap.putIfAbsent(mediaType.format(), mediaType); // ignore if the additional type mapping already exists
         }
 
         formatToMediaType = Map.copyOf(formatMap);