diff --git a/nifi-commons/nifi-flowfile-packager/src/main/java/org/apache/nifi/util/FlowFileUnpackagerV1.java b/nifi-commons/nifi-flowfile-packager/src/main/java/org/apache/nifi/util/FlowFileUnpackagerV1.java index a067fc2f98a3..7629d6b32922 100644 --- a/nifi-commons/nifi-flowfile-packager/src/main/java/org/apache/nifi/util/FlowFileUnpackagerV1.java +++ b/nifi-commons/nifi-flowfile-packager/src/main/java/org/apache/nifi/util/FlowFileUnpackagerV1.java @@ -22,10 +22,9 @@ import java.io.IOException; import java.io.InputStream; import java.io.OutputStream; -import java.util.HashMap; import java.util.Map; -import java.util.Map.Entry; import java.util.Properties; +import java.util.stream.Collectors; public class FlowFileUnpackagerV1 implements FlowFileUnpackager { @@ -70,29 +69,12 @@ public Map unpackageFlowFile(final InputStream in, final OutputS } protected Map getAttributes(final TarArchiveInputStream stream) throws IOException { - final Properties props = new Properties(); props.loadFromXML(new NonCloseableInputStream(stream)); + final Map attributes = props.stringPropertyNames().stream() + .collect(Collectors.toMap(key -> key, props::getProperty)); - final Map result = new HashMap<>(); - for (final Entry entry : props.entrySet()) { - final Object keyObject = entry.getKey(); - final Object valueObject = entry.getValue(); - if (!(keyObject instanceof final String key)) { - throw new IOException("Flow file attributes object contains key of type " - + keyObject.getClass().getCanonicalName() - + " but expected java.lang.String"); - } else if (!(keyObject instanceof String)) { - throw new IOException("Flow file attributes object contains value of type " - + keyObject.getClass().getCanonicalName() - + " but expected java.lang.String"); - } - - final String value = (String) valueObject; - result.put(key, value); - } - - return result; + return attributes; } @Override diff --git a/nifi-commons/nifi-flowfile-packager/src/test/java/org/apache/nifi/util/TestPackageUnpackageV1.java b/nifi-commons/nifi-flowfile-packager/src/test/java/org/apache/nifi/util/TestPackageUnpackageV1.java new file mode 100644 index 000000000000..2f6e1b9c1989 --- /dev/null +++ b/nifi-commons/nifi-flowfile-packager/src/test/java/org/apache/nifi/util/TestPackageUnpackageV1.java @@ -0,0 +1,77 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.nifi.util; + +import org.junit.jupiter.api.Test; + +import java.io.ByteArrayInputStream; +import java.io.ByteArrayOutputStream; +import java.nio.charset.StandardCharsets; +import java.util.HashMap; +import java.util.Map; + +import static org.junit.jupiter.api.Assertions.assertArrayEquals; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; + +public class TestPackageUnpackageV1 { + @Test + @SuppressWarnings("unchecked") + void testPackageWithNonStringAttributes() { + final FlowFilePackager packager = new FlowFilePackagerV1(); + final byte[] data = "Hello, World!".getBytes(StandardCharsets.UTF_8); + final Map map = new HashMap<>(); + map.put(12, 34); + map.put(56, 78); + map.put(9, 10); + + final Map cast = (Map) map; + + final ByteArrayOutputStream baos = new ByteArrayOutputStream(); + final ByteArrayInputStream in = new ByteArrayInputStream(data); + + assertThrows(ClassCastException.class, () -> packager.packageFlowFile(in, baos, cast, data.length)); + } + + @Test + void testPackageAndUnpackage() throws Exception { + final FlowFilePackager packager = new FlowFilePackagerV1(); + final FlowFileUnpackager unpackager = new FlowFileUnpackagerV1(); + + final byte[] data = "Hello, World!".getBytes(StandardCharsets.UTF_8); + final Map map = new HashMap<>(); + map.put("abc", "cba"); + map.put("123", null); + + final ByteArrayOutputStream baos = new ByteArrayOutputStream(); + final ByteArrayInputStream in = new ByteArrayInputStream(data); + packager.packageFlowFile(in, baos, map, data.length); + + final byte[] encoded = baos.toByteArray(); + final ByteArrayInputStream encodedIn = new ByteArrayInputStream(encoded); + final ByteArrayOutputStream decodedOut = new ByteArrayOutputStream(); + final Map unpackagedAttributes = unpackager.unpackageFlowFile(encodedIn, decodedOut); + final byte[] decoded = decodedOut.toByteArray(); + + /*Since the properties serialized as XML has the value null specified between the entry tags, + when loaded into a java.util.Properties it is read as string whose value is "null" hence for + verification must change to the expected value.*/ + map.put("123", "null"); + assertEquals(map, unpackagedAttributes); + assertArrayEquals(data, decoded); + } +}