Java Code Examples for org.apache.nifi.stream.io.StreamUtils#fillBuffer()

@Test
public void testSequentialReads() throws IOException {
    FlowFile ff1 = session.write(session.create(), out -> out.write(new byte[] {'A', 'B'}));
    FlowFile ff2 = session.write(session.create(), out -> out.write('C'));

    final byte[] buff1 = new byte[2];
    try (final InputStream in = session.read(ff1)) {
        StreamUtils.fillBuffer(in, buff1);
    }

    final byte[] buff2 = new byte[1];
    try (final InputStream in = session.read(ff2)) {
        StreamUtils.fillBuffer(in, buff2);
    }

    Assert.assertArrayEquals(new byte[] {'A', 'B'}, buff1);
    Assert.assertArrayEquals(new byte[] {'C'}, buff2);
}
 

@Test
public void testStreamCreatedFromRepository() throws IOException {
    final ContentClaimInputStream in = new ContentClaimInputStream(repo, contentClaim, 0L);

    final byte[] buff = new byte[5];
    StreamUtils.fillBuffer(in, buff);

    Mockito.verify(repo, Mockito.times(1)).read(contentClaim);
    Mockito.verifyNoMoreInteractions(repo);

    final String contentRead = new String(buff);
    assertEquals("hello", contentRead);

    assertEquals(5, in.getBytesConsumed());
    assertFalse(closed.get());

    // Ensure that underlying stream is closed
    in.close();
    assertTrue(closed.get());
}
 

@Test
public void testReadFromInputStream() throws IOException {
    FlowFile flowFile = session.create();
    flowFile = session.write(flowFile, new OutputStreamCallback() {
        @Override
        public void process(final OutputStream out) throws IOException {
            out.write("hello, world".getBytes());
        }
    });

    try (InputStream in = session.read(flowFile)) {
        final byte[] buffer = new byte[12];
        StreamUtils.fillBuffer(in, buffer);
        assertEquals("hello, world", new String(buffer));
    }

    session.remove(flowFile);
    session.commit();
}
 

public static List<Template> parseTemplateStream(final byte[] bytes) {
    final List<Template> templates = new ArrayList<>();

    try (final InputStream rawIn = new ByteArrayInputStream(bytes);
        final DataInputStream in = new DataInputStream(rawIn)) {

        while (isMoreData(in)) {
            final int length = in.readInt();
            final byte[] buffer = new byte[length];
            StreamUtils.fillBuffer(in, buffer, true);
            final TemplateDTO dto = TemplateDeserializer.deserialize(new ByteArrayInputStream(buffer));
            templates.add(new Template(dto));
        }
    } catch (final IOException e) {
        throw new RuntimeException("Could not parse bytes", e);  // won't happen because of the types of streams being used
    }

    return templates;
}
 

@Test
public void testThatContentIsSkipped() throws IOException {
    final ContentClaimInputStream in = new ContentClaimInputStream(repo, contentClaim, 3L);

    final byte[] buff = new byte[2];
    StreamUtils.fillBuffer(in, buff);

    Mockito.verify(repo, Mockito.times(1)).read(contentClaim);
    Mockito.verifyNoMoreInteractions(repo);

    final String contentRead = new String(buff);
    assertEquals("lo", contentRead);

    assertEquals(2, in.getBytesConsumed());
    assertFalse(closed.get());

    // Ensure that underlying stream is closed
    in.close();
    assertTrue(closed.get());
}
 

@Test
public void testInputStreamFromContentRepo() throws IOException {
    final ContentRepository contentRepo = mock(ContentRepository.class);

    final ResourceClaimManager claimManager = new StandardResourceClaimManager();
    final ResourceClaim resourceClaim = new StandardResourceClaim(claimManager, "container", "section", "id", false);
    final ContentClaim contentClaim = new StandardContentClaim(resourceClaim, 5L);

    final FlowFileRecord flowFile = mock(FlowFileRecord.class);
    when(flowFile.getContentClaim()).thenReturn(contentClaim);
    when(flowFile.getSize()).thenReturn(5L);

    final InputStream inputStream = new ByteArrayInputStream("hello".getBytes());
    when(contentRepo.read(contentClaim)).thenReturn(inputStream);

    final ContentRepositoryFlowFileAccess flowAccess = new ContentRepositoryFlowFileAccess(contentRepo);

    final InputStream repoStream = flowAccess.read(flowFile);
    verify(contentRepo, times(1)).read(contentClaim);

    final byte[] buffer = new byte[5];
    StreamUtils.fillBuffer(repoStream, buffer);
    assertEquals(-1, repoStream.read());
    assertArrayEquals("hello".getBytes(), buffer);
}
 

/**
 * Returns the deserialized and decrypted {@link RepositoryRecord} from the input stream.
 *
 * @param in      stream to read from
 * @param version the version of the SerDe that was used to serialize the
 *                record
 * @return the deserialized record
 * @throws IOException if there is a problem reading from the stream
 */
@Override
public SerializedRepositoryRecord deserializeRecord(final DataInputStream in, final int version) throws IOException {
    // Read the expected length of the encrypted record (including the encryption metadata)
    int encryptedRecordLength = in.readInt();
    if (encryptedRecordLength == -1) {
        return null;
    }

    // Read the encrypted record bytes
    byte[] cipherBytes = new byte[encryptedRecordLength];
    StreamUtils.fillBuffer(in, cipherBytes);
    logger.debug("Read {} bytes (encrypted, including length) from actual input stream", encryptedRecordLength + 4);

    // Decrypt the byte[]
    DataInputStream wrappedInputStream = decryptToStream(cipherBytes);

    // Deserialize the plain bytes using the delegate serde
    final SerializedRepositoryRecord deserializedRecord = wrappedSerDe.deserializeRecord(wrappedInputStream, version);
    logger.debug("Deserialized flowfile record {} from temp stream", getRecordIdentifier(deserializedRecord));
    return deserializedRecord;
}
 

@Override
protected void readHeader(final DataInputStream in, final int serializationVersion) throws IOException {
    verifySerializationVersion(serializationVersion);
    final int schemaLength = in.readInt();
    final byte[] buffer = new byte[schemaLength];
    StreamUtils.fillBuffer(in, buffer);

    try (final ByteArrayInputStream bais = new ByteArrayInputStream(buffer)) {
        schema = RecordSchema.readFrom(bais);
    }

    recordReader = SchemaRecordReader.fromSchema(schema);
}
 

@Test
public void testRereadEntireClaim() throws IOException {
    final ContentClaimInputStream in = new ContentClaimInputStream(repo, contentClaim, 0L);

    final byte[] buff = new byte[5];

    final int invocations = 10;
    for (int i=0; i < invocations; i++) {
        in.mark(5);

        StreamUtils.fillBuffer(in, buff, true);

        final String contentRead = new String(buff);
        assertEquals("hello", contentRead);

        assertEquals(5 * (i+1), in.getBytesConsumed());
        assertEquals(5, in.getCurrentOffset());
        assertEquals(-1, in.read());

        in.reset();
    }

    Mockito.verify(repo, Mockito.times(invocations + 1)).read(contentClaim); // Will call reset() 'invocations' times plus the initial read
    Mockito.verifyNoMoreInteractions(repo);

    // Ensure that underlying stream is closed
    in.close();
    assertTrue(closed.get());
}
 

/**
 * Returns the salt provided as part of the cipher stream, or throws an exception if one cannot be detected.
 *
 * @param in the cipher InputStream
 * @return the salt
 */
@Override
public byte[] readSalt(InputStream in) throws IOException {
    if (in == null) {
        throw new IllegalArgumentException("Cannot read salt from null InputStream");
    }

    // The header and salt format is "Salted__salt x8b" in ASCII
    byte[] salt = new byte[DEFAULT_SALT_LENGTH];

    // Try to read the header and salt from the input
    byte[] header = new byte[OPENSSL_EVP_HEADER_SIZE];

    // Mark the stream in case there is no salt
    in.mark(OPENSSL_EVP_HEADER_SIZE + 1);
    StreamUtils.fillBuffer(in, header);

    final byte[] headerMarkerBytes = OPENSSL_EVP_HEADER_MARKER.getBytes(StandardCharsets.US_ASCII);

    if (!Arrays.equals(headerMarkerBytes, header)) {
        // No salt present
        salt = new byte[0];
        // Reset the stream because we skipped 8 bytes of cipher text
        in.reset();
    }

    StreamUtils.fillBuffer(in, salt);
    return salt;
}
 

void publish(final FlowFile flowFile, final InputStream flowFileContent, final byte[] messageKey, final byte[] demarcatorBytes, final String topic, final Integer partition) throws IOException {
    if (tracker == null) {
        tracker = new InFlightMessageTracker(logger);
    }

    try {
        byte[] messageContent;
        if (demarcatorBytes == null || demarcatorBytes.length == 0) {
            if (flowFile.getSize() > maxMessageSize) {
                tracker.fail(flowFile, new TokenTooLargeException("A message in the stream exceeds the maximum allowed message size of " + maxMessageSize + " bytes."));
                return;
            }
            // Send FlowFile content as it is, to support sending 0 byte message.
            messageContent = new byte[(int) flowFile.getSize()];
            StreamUtils.fillBuffer(flowFileContent, messageContent);
            publish(flowFile, messageKey, messageContent, topic, tracker, partition);
            return;
        }

        try (final StreamDemarcator demarcator = new StreamDemarcator(flowFileContent, demarcatorBytes, maxMessageSize)) {
            while ((messageContent = demarcator.nextToken()) != null) {
                publish(flowFile, messageKey, messageContent, topic, tracker, partition);

                if (tracker.isFailed(flowFile)) {
                    // If we have a failure, don't try to send anything else.
                    return;
                }
            }
        } catch (final TokenTooLargeException ttle) {
            tracker.fail(flowFile, ttle);
        }
    } catch (final Exception e) {
        tracker.fail(flowFile, e);
        poison();
        throw e;
    }
}
 

@Test
public void testReadFromInputStreamWithoutClosingThenRemove() throws IOException {
    FlowFile flowFile = session.create();
    flowFile = session.write(flowFile, new OutputStreamCallback() {
        @Override
        public void process(final OutputStream out) throws IOException {
            out.write("hello, world".getBytes());
        }
    });

    InputStream in = session.read(flowFile);
    final byte[] buffer = new byte[12];
    StreamUtils.fillBuffer(in, buffer);
    assertEquals("hello, world", new String(buffer));

    try {
        session.remove(flowFile);
        Assert.fail("Was able to remove FlowFile while an InputStream is open for it");
    } catch (final IllegalStateException e) {
        // expected
    }

    in.close();

    session.remove(flowFile);

    session.commit(); // This should generate a WARN log message. We can't really test this in a unit test but can verify manually.
}
 

@Test
public void testMultipleResetCallsAfterMark() throws IOException {
    final ContentClaimInputStream in = new ContentClaimInputStream(repo, contentClaim, 0L);

    final byte[] buff = new byte[5];

    final int invocations = 10;
    in.mark(5);

    for (int i=0; i < invocations; i++) {
        StreamUtils.fillBuffer(in, buff, true);

        final String contentRead = new String(buff);
        assertEquals("hello", contentRead);

        assertEquals(5 * (i+1), in.getBytesConsumed());
        assertEquals(5, in.getCurrentOffset());
        assertEquals(-1, in.read());

        in.reset();
    }

    Mockito.verify(repo, Mockito.times(invocations + 1)).read(contentClaim); // Will call reset() 'invocations' times plus the initial read
    Mockito.verifyNoMoreInteractions(repo);

    // Ensure that underlying stream is closed
    in.close();
    assertTrue(closed.get());
}
 

void publish(final FlowFile flowFile, final InputStream flowFileContent, final byte[] messageKey, final byte[] demarcatorBytes, final String topic) throws IOException {
    if (tracker == null) {
        tracker = new InFlightMessageTracker();
    }

    try {
        byte[] messageContent;
        if (demarcatorBytes == null || demarcatorBytes.length == 0) {
            if (flowFile.getSize() > maxMessageSize) {
                tracker.fail(flowFile, new TokenTooLargeException("A message in the stream exceeds the maximum allowed message size of " + maxMessageSize + " bytes."));
                return;
            }
            // Send FlowFile content as it is, to support sending 0 byte message.
            messageContent = new byte[(int) flowFile.getSize()];
            StreamUtils.fillBuffer(flowFileContent, messageContent);
            publish(flowFile, messageKey, messageContent, topic, tracker);
            return;
        }

        try (final StreamDemarcator demarcator = new StreamDemarcator(flowFileContent, demarcatorBytes, maxMessageSize)) {
            while ((messageContent = demarcator.nextToken()) != null) {
                publish(flowFile, messageKey, messageContent, topic, tracker);

                if (tracker.isFailed(flowFile)) {
                    // If we have a failure, don't try to send anything else.
                    return;
                }
            }
            tracker.trackEmpty(flowFile);
        } catch (final TokenTooLargeException ttle) {
            tracker.fail(flowFile, ttle);
        }
    } catch (final Exception e) {
        tracker.fail(flowFile, e);
        poison();
        throw e;
    }
}
 

@Test
@Ignore("For local testing only; not really a unit test but a manual test")
public void testReceive() throws IOException {
    System.setProperty("org.slf4j.simpleLogger.log.org.apache.nifi.remote", "DEBUG");

    final SiteToSiteClient client = new SiteToSiteClient.Builder()
            .url("http://localhost:8080/nifi")
            .portName("cba")
            .requestBatchCount(10)
            .build();

    try {
        for (int i = 0; i < 1000; i++) {
            final Transaction transaction = client.createTransaction(TransferDirection.RECEIVE);
            Assert.assertNotNull(transaction);

            DataPacket packet;
            while (true) {
                packet = transaction.receive();
                if (packet == null) {
                    break;
                }

                final InputStream in = packet.getData();
                final long size = packet.getSize();
                final byte[] buff = new byte[(int) size];

                StreamUtils.fillBuffer(in, buff);
            }

            transaction.confirm();
            transaction.complete();
        }
    } finally {
        client.close();
    }
}
 

@Override
@SuppressWarnings("unchecked")
protected synchronized void readHeader(final DataInputStream in, final int serializationVersion) throws IOException {
    verifySerializationVersion(serializationVersion);
    final int eventSchemaLength = in.readInt();
    final byte[] buffer = new byte[eventSchemaLength];
    StreamUtils.fillBuffer(in, buffer);

    try (final ByteArrayInputStream bais = new ByteArrayInputStream(buffer)) {
        schema = RecordSchema.readFrom(bais);
    }

    recordReader = SchemaRecordReader.fromSchema(schema);

    final int headerSchemaLength = in.readInt();
    final byte[] headerSchemaBuffer = new byte[headerSchemaLength];
    StreamUtils.fillBuffer(in, headerSchemaBuffer);

    final RecordSchema headerSchema;
    try (final ByteArrayInputStream bais = new ByteArrayInputStream(headerSchemaBuffer)) {
        headerSchema = RecordSchema.readFrom(bais);
    }

    final SchemaRecordReader headerReader = SchemaRecordReader.fromSchema(headerSchema);
    final Record headerRecord = headerReader.readRecord(in);
    componentIds = (List<String>) headerRecord.getFieldValue(EventIdFirstHeaderSchema.FieldNames.COMPONENT_IDS);
    componentTypes = (List<String>) headerRecord.getFieldValue(EventIdFirstHeaderSchema.FieldNames.COMPONENT_TYPES);
    queueIds = (List<String>) headerRecord.getFieldValue(EventIdFirstHeaderSchema.FieldNames.QUEUE_IDS);
    eventTypes = (List<String>) headerRecord.getFieldValue(EventIdFirstHeaderSchema.FieldNames.EVENT_TYPES);
    firstEventId = (Long) headerRecord.getFieldValue(EventIdFirstHeaderSchema.FieldNames.FIRST_EVENT_ID);
    systemTimeOffset = (Long) headerRecord.getFieldValue(EventIdFirstHeaderSchema.FieldNames.TIMESTAMP_OFFSET);
}
 

@Test
public void testMigrateWithAppendableStream() throws IOException {
    FlowFile flowFile = session.create();
    flowFile = session.append(flowFile, out -> out.write("1".getBytes()));
    flowFile = session.append(flowFile, out -> out.write("2".getBytes()));

    final StandardProcessSession newSession = new StandardProcessSession(context);

    assertTrue(session.isFlowFileKnown(flowFile));
    assertFalse(newSession.isFlowFileKnown(flowFile));

    session.migrate(newSession, Collections.singleton(flowFile));

    assertFalse(session.isFlowFileKnown(flowFile));
    assertTrue(newSession.isFlowFileKnown(flowFile));

    flowFile = newSession.append(flowFile, out -> out.write("3".getBytes()));

    final byte[] buff = new byte[3];
    try (final InputStream in = newSession.read(flowFile)) {
        StreamUtils.fillBuffer(in, buff, true);
        assertEquals(-1, in.read());
    }

    assertTrue(Arrays.equals(new byte[] {'1', '2', '3'}, buff));

    newSession.remove(flowFile);
    newSession.commit();
    session.commit();
}
 

/**
 * Returns the salt provided as part of the cipher stream, or throws an exception if one cannot be detected.
 * This method is only implemented by {@link NiFiLegacyCipherProvider} because the legacy salt generation was dependent on the cipher block size.
 *
 * @param encryptionMethod the encryption method
 * @param in the cipher InputStream
 * @return the salt
 */
public byte[] readSalt(EncryptionMethod encryptionMethod, InputStream in) throws IOException {
    if (in == null) {
        throw new IllegalArgumentException("Cannot read salt from null InputStream");
    }

    // The first 8-16 bytes (depending on the cipher blocksize) of the input stream are the salt
    final int saltLength = calculateSaltLength(encryptionMethod);
    if (in.available() < saltLength) {
        throw new ProcessException("The cipher stream is too small to contain the salt");
    }
    byte[] salt = new byte[saltLength];
    StreamUtils.fillBuffer(in, salt);
    return salt;
}
 

@Test
public void testWriteToOutputStream() throws IOException {
    final FlowFileRecord flowFileRecord = new StandardFlowFileRecord.Builder()
        .addAttribute("uuid", "12345678-1234-1234-1234-123456789012")
        .entryDate(System.currentTimeMillis())
        .size(12L)
        .build();

    flowFileQueue.put(flowFileRecord);

    FlowFile flowFile = session.get();
    try (final OutputStream out = session.write(flowFile)) {
        out.write("hello, world".getBytes());
    }

    // Call putAllAttributes, because this will return to us the most recent version
    // of the FlowFile. In a Processor, we wouldn't need this, but for testing purposes
    // we need it in order to get the Content Claim.
    flowFile = session.putAllAttributes(flowFile, Collections.emptyMap());
    assertEquals(12L, flowFile.getSize());

    final byte[] buffer = new byte[(int) flowFile.getSize()];
    try (final InputStream in = session.read(flowFile)) {
        StreamUtils.fillBuffer(in, buffer);
    }

    assertEquals(new String(buffer), "hello, world");
}
 

private String readString(final DataInputStream in) throws IOException {
    final int numBytes = in.readInt();
    final byte[] bytes = new byte[numBytes];
    StreamUtils.fillBuffer(in, bytes, true);
    return new String(bytes, "UTF-8");
}