From 8aaa3d2f10205425be94fc9f97c499eab9224b30 Mon Sep 17 00:00:00 2001
From: Tamas Palfy <tamas.bertalan.palfy@gmail.com>
Date: Thu, 5 May 2022 17:45:26 +0200
Subject: [PATCH] NIFI-9991 Fix: The asn1 StandardRecordModelIteratorProvider's
 interator returns only the first object read from the input stream.

This closes #6019.

Signed-off-by: Peter Turcsanyi <turcsanyi@apache.org>
---
 .../apache/nifi/jasn1/JASN1RecordReader.java  |  1 -
 .../StandardRecordModelIteratorProvider.java  | 65 +++++++++++++------
 2 files changed, 46 insertions(+), 20 deletions(-)

diff --git a/nifi-nar-bundles/nifi-asn1-bundle/nifi-asn1-services/src/main/java/org/apache/nifi/jasn1/JASN1RecordReader.java b/nifi-nar-bundles/nifi-asn1-bundle/nifi-asn1-services/src/main/java/org/apache/nifi/jasn1/JASN1RecordReader.java
index 0111b70abc..d189efcb64 100644
--- a/nifi-nar-bundles/nifi-asn1-bundle/nifi-asn1-services/src/main/java/org/apache/nifi/jasn1/JASN1RecordReader.java
+++ b/nifi-nar-bundles/nifi-asn1-bundle/nifi-asn1-services/src/main/java/org/apache/nifi/jasn1/JASN1RecordReader.java
@@ -144,7 +144,6 @@ public class JASN1RecordReader implements RecordReader {
             } else {
                 return null;
             }
-
         });
     }
 
diff --git a/nifi-nar-bundles/nifi-asn1-bundle/nifi-asn1-services/src/main/java/org/apache/nifi/jasn1/StandardRecordModelIteratorProvider.java b/nifi-nar-bundles/nifi-asn1-bundle/nifi-asn1-services/src/main/java/org/apache/nifi/jasn1/StandardRecordModelIteratorProvider.java
index ea73c89389..8e4c106f62 100644
--- a/nifi-nar-bundles/nifi-asn1-bundle/nifi-asn1-services/src/main/java/org/apache/nifi/jasn1/StandardRecordModelIteratorProvider.java
+++ b/nifi-nar-bundles/nifi-asn1-bundle/nifi-asn1-services/src/main/java/org/apache/nifi/jasn1/StandardRecordModelIteratorProvider.java
@@ -35,7 +35,52 @@ public class StandardRecordModelIteratorProvider implements RecordModelIteratorP
     @Override
     @SuppressWarnings("unchecked")
     public Iterator<BerType> iterator(InputStream inputStream, ComponentLog logger, Class<? extends BerType> rootClass, String recordField, Field seqOfField) {
+        if (StringUtils.isEmpty(recordField)) {
+            return new Iterator<BerType>() {
+                @Override
+                public boolean hasNext() {
+                    boolean hasNext;
+
+                    try {
+                        hasNext = inputStream.available() > 0;
+                    } catch (IOException e) {
+                        hasNext = false;
+                    }
+
+                    return hasNext;
+                }
+
+                @Override
+                public BerType next() {
+                    final BerType model = decode(inputStream, logger, rootClass);
+
+                    return model;
+                }
+            };
+        } else {
+            final BerType model = decode(inputStream, logger, rootClass);
+
+            final List<BerType> recordModels;
+            try {
+                final Method recordModelGetter = rootClass.getMethod(JASN1Utils.toGetterMethod(recordField));
+                final BerType readPointModel = (BerType) recordModelGetter.invoke(model);
+                if (seqOfField != null) {
+                    final Class seqOf = JASN1Utils.getSeqOfElementType(seqOfField);
+                    recordModels = (List<BerType>) invokeGetter(readPointModel, JASN1Utils.toGetterMethod(seqOf.getSimpleName()));
+                } else {
+                    recordModels = Collections.singletonList(readPointModel);
+                }
+            } catch (ReflectiveOperationException e) {
+                throw new RuntimeException("Failed to get record models due to " + e, e);
+            }
+
+            return recordModels.iterator();
+        }
+    }
+
+    private BerType decode(InputStream inputStream, ComponentLog logger, Class<? extends BerType> rootClass) {
         final BerType model;
+
         try {
             model = rootClass.getDeclaredConstructor().newInstance();
         } catch (ReflectiveOperationException e) {
@@ -49,24 +94,6 @@ public class StandardRecordModelIteratorProvider implements RecordModelIteratorP
             throw new RuntimeException("Failed to decode " + rootClass.getCanonicalName(), e);
         }
 
-        final List<BerType> recordModels;
-        if (StringUtils.isEmpty(recordField)) {
-            recordModels = Collections.singletonList(model);
-        } else {
-            try {
-                final Method recordModelGetter = rootClass.getMethod(JASN1Utils.toGetterMethod(recordField));
-                final BerType readPointModel = (BerType) recordModelGetter.invoke(model);
-                if (seqOfField != null) {
-                    final Class seqOf = JASN1Utils.getSeqOfElementType(seqOfField);
-                    recordModels = (List<BerType>) invokeGetter(readPointModel, JASN1Utils.toGetterMethod(seqOf.getSimpleName()));
-                } else {
-                    recordModels = Collections.singletonList(readPointModel);
-                }
-            } catch (ReflectiveOperationException e) {
-                throw new RuntimeException("Failed to get record models due to " + e, e);
-            }
-        }
-
-        return recordModels.iterator();
+        return model;
     }
 }