Examples with FlowFileUnpackager org.apache.nifi.util.FlowFileUnpackager used on opensource projects

Example 1 with FlowFileUnpackager

use of org.apache.nifi.util.FlowFileUnpackager in project nifi by apache.

the class ListenHTTPServlet method doPost.

@Override
protected void doPost(final HttpServletRequest request, final HttpServletResponse response) throws ServletException, IOException {
    final ProcessContext context = processContext;
    ProcessSessionFactory sessionFactory;
    do {
        sessionFactory = sessionFactoryHolder.get();
        if (sessionFactory == null) {
            try {
                Thread.sleep(10);
            } catch (final InterruptedException e) {
            }
        }
    } while (sessionFactory == null);
    final ProcessSession session = sessionFactory.createSession();
    FlowFile flowFile = null;
    String holdUuid = null;
    String foundSubject = null;
    try {
        final long n = filesReceived.getAndIncrement() % FILES_BEFORE_CHECKING_DESTINATION_SPACE;
        if (n == 0 || !spaceAvailable.get()) {
            if (context.getAvailableRelationships().isEmpty()) {
                spaceAvailable.set(false);
                if (logger.isDebugEnabled()) {
                    logger.debug("Received request from " + request.getRemoteHost() + " but no space available; Indicating Service Unavailable");
                }
                response.sendError(HttpServletResponse.SC_SERVICE_UNAVAILABLE);
                return;
            } else {
                spaceAvailable.set(true);
            }
        }
        response.setHeader("Content-Type", MediaType.TEXT_PLAIN);
        final boolean contentGzipped = Boolean.parseBoolean(request.getHeader(GZIPPED_HEADER));
        final X509Certificate[] certs = (X509Certificate[]) request.getAttribute("javax.servlet.request.X509Certificate");
        foundSubject = DEFAULT_FOUND_SUBJECT;
        if (certs != null && certs.length > 0) {
            for (final X509Certificate cert : certs) {
                foundSubject = cert.getSubjectDN().getName();
                if (authorizedPattern.matcher(foundSubject).matches()) {
                    break;
                } else {
                    logger.warn("Rejecting transfer attempt from " + foundSubject + " because the DN is not authorized, host=" + request.getRemoteHost());
                    response.sendError(HttpServletResponse.SC_FORBIDDEN, "not allowed based on dn");
                    return;
                }
            }
        }
        final String destinationVersion = request.getHeader(PROTOCOL_VERSION_HEADER);
        Integer protocolVersion = null;
        if (destinationVersion != null) {
            try {
                protocolVersion = Integer.valueOf(destinationVersion);
            } catch (final NumberFormatException e) {
            // Value was invalid. Treat as if the header were missing.
            }
        }
        final boolean destinationIsLegacyNiFi = (protocolVersion == null);
        final boolean createHold = Boolean.parseBoolean(request.getHeader(FLOWFILE_CONFIRMATION_HEADER));
        final String contentType = request.getContentType();
        final InputStream unthrottled = contentGzipped ? new GZIPInputStream(request.getInputStream()) : request.getInputStream();
        final InputStream in = (streamThrottler == null) ? unthrottled : streamThrottler.newThrottledInputStream(unthrottled);
        if (logger.isDebugEnabled()) {
            logger.debug("Received request from " + request.getRemoteHost() + ", createHold=" + createHold + ", content-type=" + contentType + ", gzip=" + contentGzipped);
        }
        final AtomicBoolean hasMoreData = new AtomicBoolean(false);
        final FlowFileUnpackager unpackager;
        if (APPLICATION_FLOW_FILE_V3.equals(contentType)) {
            unpackager = new FlowFileUnpackagerV3();
        } else if (APPLICATION_FLOW_FILE_V2.equals(contentType)) {
            unpackager = new FlowFileUnpackagerV2();
        } else if (APPLICATION_FLOW_FILE_V1.equals(contentType)) {
            unpackager = new FlowFileUnpackagerV1();
        } else {
            unpackager = null;
        }
        final Set<FlowFile> flowFileSet = new HashSet<>();
        do {
            final long startNanos = System.nanoTime();
            final Map<String, String> attributes = new HashMap<>();
            flowFile = session.create();
            flowFile = session.write(flowFile, new OutputStreamCallback() {

                @Override
                public void process(final OutputStream rawOut) throws IOException {
                    try (final BufferedOutputStream bos = new BufferedOutputStream(rawOut, 65536)) {
                        if (unpackager == null) {
                            IOUtils.copy(in, bos);
                            hasMoreData.set(false);
                        } else {
                            attributes.putAll(unpackager.unpackageFlowFile(in, bos));
                            if (destinationIsLegacyNiFi) {
                                if (attributes.containsKey("nf.file.name")) {
                                    // for backward compatibility with old nifi...
                                    attributes.put(CoreAttributes.FILENAME.key(), attributes.remove("nf.file.name"));
                                }
                                if (attributes.containsKey("nf.file.path")) {
                                    attributes.put(CoreAttributes.PATH.key(), attributes.remove("nf.file.path"));
                                }
                            }
                            hasMoreData.set(unpackager.hasMoreData());
                        }
                    }
                }
            });
            final long transferNanos = System.nanoTime() - startNanos;
            final long transferMillis = TimeUnit.MILLISECONDS.convert(transferNanos, TimeUnit.NANOSECONDS);
            // put metadata on flowfile
            final String nameVal = request.getHeader(CoreAttributes.FILENAME.key());
            if (StringUtils.isNotBlank(nameVal)) {
                attributes.put(CoreAttributes.FILENAME.key(), nameVal);
            }
            // put arbitrary headers on flow file
            for (Enumeration<String> headerEnum = request.getHeaderNames(); headerEnum.hasMoreElements(); ) {
                String headerName = headerEnum.nextElement();
                if (headerPattern != null && headerPattern.matcher(headerName).matches()) {
                    String headerValue = request.getHeader(headerName);
                    attributes.put(headerName, headerValue);
                }
            }
            String sourceSystemFlowFileIdentifier = attributes.get(CoreAttributes.UUID.key());
            if (sourceSystemFlowFileIdentifier != null) {
                sourceSystemFlowFileIdentifier = "urn:nifi:" + sourceSystemFlowFileIdentifier;
                // If we receveied a UUID, we want to give the FlowFile a new UUID and register the sending system's
                // identifier as the SourceSystemFlowFileIdentifier field in the Provenance RECEIVE event
                attributes.put(CoreAttributes.UUID.key(), UUID.randomUUID().toString());
            }
            flowFile = session.putAllAttributes(flowFile, attributes);
            session.getProvenanceReporter().receive(flowFile, request.getRequestURL().toString(), sourceSystemFlowFileIdentifier, "Remote DN=" + foundSubject, transferMillis);
            flowFile = session.putAttribute(flowFile, "restlistener.remote.source.host", request.getRemoteHost());
            flowFile = session.putAttribute(flowFile, "restlistener.request.uri", request.getRequestURI());
            flowFile = session.putAttribute(flowFile, "restlistener.remote.user.dn", foundSubject);
            flowFileSet.add(flowFile);
            if (holdUuid == null) {
                holdUuid = flowFile.getAttribute(CoreAttributes.UUID.key());
            }
        } while (hasMoreData.get());
        if (createHold) {
            String uuid = (holdUuid == null) ? UUID.randomUUID().toString() : holdUuid;
            if (flowFileMap.containsKey(uuid)) {
                uuid = UUID.randomUUID().toString();
            }
            final FlowFileEntryTimeWrapper wrapper = new FlowFileEntryTimeWrapper(session, flowFileSet, System.currentTimeMillis(), request.getRemoteHost());
            FlowFileEntryTimeWrapper previousWrapper;
            do {
                previousWrapper = flowFileMap.putIfAbsent(uuid, wrapper);
                if (previousWrapper != null) {
                    uuid = UUID.randomUUID().toString();
                }
            } while (previousWrapper != null);
            response.setStatus(HttpServletResponse.SC_SEE_OTHER);
            final String ackUri = "/" + basePath + "/holds/" + uuid;
            response.addHeader(LOCATION_HEADER_NAME, ackUri);
            response.addHeader(LOCATION_URI_INTENT_NAME, LOCATION_URI_INTENT_VALUE);
            response.getOutputStream().write(ackUri.getBytes("UTF-8"));
            if (logger.isDebugEnabled()) {
                logger.debug("Ingested {} from Remote Host: [{}] Port [{}] SubjectDN [{}]; placed hold on these {} files with ID {}", new Object[] { flowFileSet, request.getRemoteHost(), request.getRemotePort(), foundSubject, flowFileSet.size(), uuid });
            }
        } else {
            response.setStatus(this.returnCode);
            logger.info("Received from Remote Host: [{}] Port [{}] SubjectDN [{}]; transferring to 'success' {}", new Object[] { request.getRemoteHost(), request.getRemotePort(), foundSubject, flowFile });
            session.transfer(flowFileSet, ListenHTTP.RELATIONSHIP_SUCCESS);
            session.commit();
        }
    } catch (final Throwable t) {
        session.rollback();
        if (flowFile == null) {
            logger.error("Unable to receive file from Remote Host: [{}] SubjectDN [{}] due to {}", new Object[] { request.getRemoteHost(), foundSubject, t });
        } else {
            logger.error("Unable to receive file {} from Remote Host: [{}] SubjectDN [{}] due to {}", new Object[] { flowFile, request.getRemoteHost(), foundSubject, t });
        }
        response.sendError(HttpServletResponse.SC_INTERNAL_SERVER_ERROR, t.toString());
    }
}