Optimize the code

Signed-off-by: Gao Binlong <[email protected]>
opensearch-project · Feb 7, 2024 · a5cf254 · a5cf254
1 parent ad249ea
commit a5cf254
Show file tree

Hide file tree

Showing 4 changed files with 56 additions and 63 deletions.
diff --git a/...h/ingest/common/CommunityIDProcessor.java → ...h/ingest/common/CommunityIdProcessor.java b/...h/ingest/common/CommunityIDProcessor.java → ...h/ingest/common/CommunityIdProcessor.java
@@ -21,7 +21,6 @@
 import java.util.Base64;
 import java.util.Locale;
 import java.util.Map;
-import java.util.Set;
 import java.util.stream.Collectors;
 
 import static org.opensearch.ingest.ConfigurationUtils.newConfigurationException;
@@ -30,7 +29,7 @@
  * Processor that generating community id flow hash for the network flow tuples, the algorithm is defined in
  * <a href="https://github.com/corelight/community-id-spec">Community ID Flow Hashing</a>.
  */
-public class CommunityIDProcessor extends AbstractProcessor {
+public class CommunityIdProcessor extends AbstractProcessor {
     public static final String TYPE = "community_id";
     // the version of the community id flow hashing algorithm
     private static final String COMMUNITY_ID_HASH_VERSION = "1";
@@ -62,7 +61,7 @@ public class CommunityIDProcessor extends AbstractProcessor {
     private final String targetField;
     private final boolean ignoreMissing;
 
-    CommunityIDProcessor(
+    CommunityIdProcessor(
         String tag,
         String description,
         String sourceIPField,
@@ -162,8 +161,7 @@ public IngestDocument execute(IngestDocument document) {
         // exit quietly if either source port or destination port is null nor empty
         Integer sourcePort = null;
         Integer destinationPort = null;
-        final boolean isTransportProtocol = Protocol.isTransportProtocol(protocol.getProtocolCode());
-        if (isTransportProtocol) {
+        if (protocol.isTransportProtocol()) {
             sourcePort = resolvePort(document, sourcePortField);
             if (sourcePort == null) {
                 return document;
@@ -179,8 +177,7 @@ public IngestDocument execute(IngestDocument document) {
         // set source port to icmp type, and set dest port to icmp code, so that we can have a generic way to handle
         // all protocols
         boolean isOneway = true;
-        final boolean isICMPProtocol = Protocol.ICMP.getProtocolCode() == protocol.getProtocolCode()
-            || Protocol.ICMP_V6.getProtocolCode() == protocol.getProtocolCode();
+        final boolean isICMPProtocol = Protocol.ICMP == protocol || Protocol.ICMP_V6 == protocol;
         if (isICMPProtocol) {
             Integer icmpType = resolveICMP(document, icmpTypeField, ICMP_MESSAGE_TYPE);
             if (icmpType == null) {
@@ -558,31 +555,23 @@ public static Byte getEquivalentCode(int type) {
      * An enumeration of the supported network protocols
      */
     enum Protocol {
-        ICMP((byte) 1),
-        TCP((byte) 6),
-        UDP((byte) 17),
-        ICMP_V6((byte) 58),
-        SCTP((byte) 132);
+        ICMP((byte) 1, false),
+        TCP((byte) 6, true),
+        UDP((byte) 17, true),
+        ICMP_V6((byte) 58, false),
+        SCTP((byte) 132, true);
 
         private final byte protocolCode;
+        private final boolean isTransportProtocol;
 
-        Protocol(int ianaNumber) {
+        Protocol(int ianaNumber, boolean isTransportProtocol) {
             this.protocolCode = Integer.valueOf(ianaNumber).byteValue();
+            this.isTransportProtocol = isTransportProtocol;
         }
 
-        private static final Set<Byte> transportProtocolNumbers = Set.of(
-            TCP.getProtocolCode(),
-            UDP.getProtocolCode(),
-            SCTP.getProtocolCode()
-        );
-
         public static final Map<Byte, Protocol> protocolCodeMap = Arrays.stream(values())
             .collect(Collectors.toMap(Protocol::getProtocolCode, p -> p));
 
-        public static boolean isTransportProtocol(byte ianaProtocolNumber) {
-            return transportProtocolNumbers.contains(ianaProtocolNumber);
-        }
-
         public static Protocol fromProtocolName(String protocolName) {
             String name = protocolName.toUpperCase(Locale.ROOT);
             if (name.equals("IPV6-ICMP")) {
@@ -598,11 +587,15 @@ public static Protocol fromProtocolName(String protocolName) {
         public byte getProtocolCode() {
             return this.protocolCode;
         }
+
+        public boolean isTransportProtocol() {
+            return this.isTransportProtocol;
+        }
     }
 
     public static class Factory implements Processor.Factory {
         @Override
-        public CommunityIDProcessor create(
+        public CommunityIdProcessor create(
             Map<String, Processor.Factory> registry,
             String processorTag,
             String description,
@@ -634,7 +627,7 @@ public CommunityIDProcessor create(
             String targetField = ConfigurationUtils.readStringProperty(TYPE, processorTag, config, "target_field", "community_id");
             boolean ignoreMissing = ConfigurationUtils.readBooleanProperty(TYPE, processorTag, config, "ignore_missing", false);
 
-            return new CommunityIDProcessor(
+            return new CommunityIdProcessor(
                 processorTag,
                 description,
                 sourceIPField,

diff --git a/...es/ingest-common/src/main/java/org/opensearch/ingest/common/IngestCommonModulePlugin.java b/...es/ingest-common/src/main/java/org/opensearch/ingest/common/IngestCommonModulePlugin.java
@@ -108,7 +108,7 @@ public Map<String, Processor.Factory> getProcessors(Processor.Parameters paramet
         processors.put(CsvProcessor.TYPE, new CsvProcessor.Factory());
         processors.put(CopyProcessor.TYPE, new CopyProcessor.Factory(parameters.scriptService));
         processors.put(RemoveByPatternProcessor.TYPE, new RemoveByPatternProcessor.Factory());
-        processors.put(CommunityIDProcessor.TYPE, new CommunityIDProcessor.Factory());
+        processors.put(CommunityIdProcessor.TYPE, new CommunityIdProcessor.Factory());
         return Collections.unmodifiableMap(processors);
     }
 

diff --git a/...mon/CommunityIDProcessorFactoryTests.java → ...mon/CommunityIdProcessorFactoryTests.java b/...mon/CommunityIDProcessorFactoryTests.java → ...mon/CommunityIdProcessorFactoryTests.java
@@ -18,12 +18,12 @@
 
 import static org.hamcrest.CoreMatchers.equalTo;
 
-public class CommunityIDProcessorFactoryTests extends OpenSearchTestCase {
-    private CommunityIDProcessor.Factory factory;
+public class CommunityIdProcessorFactoryTests extends OpenSearchTestCase {
+    private CommunityIdProcessor.Factory factory;
 
     @Before
     public void init() {
-        factory = new CommunityIDProcessor.Factory();
+        factory = new CommunityIdProcessor.Factory();
     }
 
     public void testCreate() throws Exception {
@@ -42,7 +42,7 @@ public void testCreate() throws Exception {
         config.put("target_field", "community_id_hash");
         config.put("ignore_missing", ignoreMissing);
         String processorTag = randomAlphaOfLength(10);
-        CommunityIDProcessor communityIDProcessor = factory.create(null, processorTag, null, config);
+        CommunityIdProcessor communityIDProcessor = factory.create(null, processorTag, null, config);
         assertThat(communityIDProcessor.getTag(), equalTo(processorTag));
         assertThat(communityIDProcessor.getSourceIPField(), equalTo("source_ip"));
         assertThat(communityIDProcessor.getSourcePortField(), equalTo("source_port"));