Examples with SocketChannelCommunicationsSession org.apache.nifi.remote.io.socket.SocketChannelCommunicationsSession used on opensource projects

Example 1 with SocketChannelCommunicationsSession

use of org.apache.nifi.remote.io.socket.SocketChannelCommunicationsSession in project nifi by apache.

the class TestStandardRemoteGroupPort method testSendRaw.

@Test
public void testSendRaw() throws Exception {
    setupMock(SiteToSiteTransportProtocol.RAW, TransferDirection.SEND);
    setupMockProcessSession();
    final String peerUrl = "nifi://node1.example.com:9090";
    final PeerDescription peerDescription = new PeerDescription("node1.example.com", 9090, true);
    try (final SocketChannel socketChannel = SocketChannel.open()) {
        final CommunicationsSession commsSession = new SocketChannelCommunicationsSession(socketChannel);
        commsSession.setUserDn("nifi.node1.example.com");
        final Peer peer = new Peer(peerDescription, commsSession, peerUrl, REMOTE_CLUSTER_URL);
        doReturn(peer).when(transaction).getCommunicant();
        final MockFlowFile flowFile = processSession.createFlowFile("0123456789".getBytes());
        sessionState.getFlowFileQueue().offer(flowFile);
        port.onTrigger(processContext, processSession);
        // Assert provenance.
        final List<ProvenanceEventRecord> provenanceEvents = sessionState.getProvenanceEvents();
        assertEquals(1, provenanceEvents.size());
        final ProvenanceEventRecord provenanceEvent = provenanceEvents.get(0);
        assertEquals(ProvenanceEventType.SEND, provenanceEvent.getEventType());
        assertEquals(peerUrl + "/" + flowFile.getAttribute(CoreAttributes.UUID.key()), provenanceEvent.getTransitUri());
        assertEquals("Remote DN=nifi.node1.example.com", provenanceEvent.getDetails());
    }
}

Example 2 with SocketChannelCommunicationsSession

use of org.apache.nifi.remote.io.socket.SocketChannelCommunicationsSession in project nifi by apache.

the class TestStandardRemoteGroupPort method testReceiveRaw.

@Test
public void testReceiveRaw() throws Exception {
    setupMock(SiteToSiteTransportProtocol.RAW, TransferDirection.RECEIVE);
    setupMockProcessSession();
    final String peerUrl = "nifi://node1.example.com:9090";
    final PeerDescription peerDescription = new PeerDescription("node1.example.com", 9090, true);
    try (final SocketChannel socketChannel = SocketChannel.open()) {
        final CommunicationsSession commsSession = new SocketChannelCommunicationsSession(socketChannel);
        commsSession.setUserDn("nifi.node1.example.com");
        final Peer peer = new Peer(peerDescription, commsSession, peerUrl, REMOTE_CLUSTER_URL);
        doReturn(peer).when(transaction).getCommunicant();
        final String sourceFlowFileUuid = "flowfile-uuid";
        final Map<String, String> attributes = new HashMap<>();
        attributes.put(CoreAttributes.UUID.key(), sourceFlowFileUuid);
        final byte[] dataPacketContents = "DataPacket Contents".getBytes();
        final ByteArrayInputStream dataPacketInputStream = new ByteArrayInputStream(dataPacketContents);
        final DataPacket dataPacket = new StandardDataPacket(attributes, dataPacketInputStream, dataPacketContents.length);
        // Return null when it gets called second time.
        doReturn(dataPacket).doReturn(null).when(this.transaction).receive();
        port.onTrigger(processContext, processSession);
        // Assert provenance.
        final List<ProvenanceEventRecord> provenanceEvents = sessionState.getProvenanceEvents();
        assertEquals(1, provenanceEvents.size());
        final ProvenanceEventRecord provenanceEvent = provenanceEvents.get(0);
        assertEquals(ProvenanceEventType.RECEIVE, provenanceEvent.getEventType());
        assertEquals(peerUrl + "/" + sourceFlowFileUuid, provenanceEvent.getTransitUri());
        assertEquals("Remote DN=nifi.node1.example.com", provenanceEvent.getDetails());
        // Assert received flow files.
        processSession.assertAllFlowFilesTransferred(Relationship.ANONYMOUS);
        final List<MockFlowFile> flowFiles = processSession.getFlowFilesForRelationship(Relationship.ANONYMOUS);
        assertEquals(1, flowFiles.size());
        final MockFlowFile flowFile = flowFiles.get(0);
        flowFile.assertAttributeEquals(SiteToSiteAttributes.S2S_HOST.key(), peer.getHost());
        flowFile.assertAttributeEquals(SiteToSiteAttributes.S2S_ADDRESS.key(), peer.getHost() + ":" + peer.getPort());
    }
}

Example 3 with SocketChannelCommunicationsSession

use of org.apache.nifi.remote.io.socket.SocketChannelCommunicationsSession in project nifi by apache.

the class SocketRemoteSiteListener method start.

@Override
public void start() throws IOException {
    final boolean secure = (sslContext != null);
    final List<Thread> threads = new ArrayList<Thread>();
    final ServerSocketChannel serverSocketChannel = ServerSocketChannel.open();
    serverSocketChannel.configureBlocking(true);
    serverSocketChannel.bind(new InetSocketAddress(socketPort));
    stopped.set(false);
    final Thread listenerThread = new Thread(new Runnable() {

        private int threadCount = 0;

        @Override
        public void run() {
            while (!stopped.get()) {
                final ProcessGroup processGroup = rootGroup.get();
                // the processGroup.
                if ((nodeInformant == null) && (processGroup == null || (processGroup.getInputPorts().isEmpty() && processGroup.getOutputPorts().isEmpty()))) {
                    try {
                        Thread.sleep(2000L);
                    } catch (final Exception e) {
                    }
                    continue;
                }
                LOG.trace("Accepting Connection...");
                Socket acceptedSocket = null;
                try {
                    serverSocketChannel.configureBlocking(false);
                    final ServerSocket serverSocket = serverSocketChannel.socket();
                    serverSocket.setSoTimeout(2000);
                    while (!stopped.get() && acceptedSocket == null) {
                        try {
                            acceptedSocket = serverSocket.accept();
                        } catch (final SocketTimeoutException ste) {
                            continue;
                        }
                    }
                } catch (final IOException e) {
                    LOG.error("RemoteSiteListener Unable to accept connection due to {}", e.toString());
                    if (LOG.isDebugEnabled()) {
                        LOG.error("", e);
                    }
                    continue;
                }
                LOG.trace("Got connection");
                if (stopped.get()) {
                    break;
                }
                final Socket socket = acceptedSocket;
                final SocketChannel socketChannel = socket.getChannel();
                final Thread thread = new Thread(new Runnable() {

                    @Override
                    public void run() {
                        LOG.debug("{} Determining URL of connection", this);
                        final InetAddress inetAddress = socket.getInetAddress();
                        String clientHostName = inetAddress.getHostName();
                        final int slashIndex = clientHostName.indexOf("/");
                        if (slashIndex == 0) {
                            clientHostName = clientHostName.substring(1);
                        } else if (slashIndex > 0) {
                            clientHostName = clientHostName.substring(0, slashIndex);
                        }
                        final int clientPort = socket.getPort();
                        final String peerUri = "nifi://" + clientHostName + ":" + clientPort;
                        LOG.debug("{} Connection URL is {}", this, peerUri);
                        final CommunicationsSession commsSession;
                        final String dn;
                        try {
                            if (secure) {
                                final SSLSocketChannel sslSocketChannel = new SSLSocketChannel(sslContext, socketChannel, false);
                                LOG.trace("Channel is secure; connecting...");
                                sslSocketChannel.connect();
                                LOG.trace("Channel connected");
                                commsSession = new SSLSocketChannelCommunicationsSession(sslSocketChannel);
                                dn = sslSocketChannel.getDn();
                                commsSession.setUserDn(dn);
                            } else {
                                LOG.trace("{} Channel is not secure", this);
                                commsSession = new SocketChannelCommunicationsSession(socketChannel);
                                dn = null;
                            }
                        } catch (final Exception e) {
                            LOG.error("RemoteSiteListener Unable to accept connection from {} due to {}", socket, e.toString());
                            if (LOG.isDebugEnabled()) {
                                LOG.error("", e);
                            }
                            try {
                                socketChannel.close();
                            } catch (IOException swallow) {
                            }
                            return;
                        }
                        LOG.info("Received connection from {}, User DN: {}", socket.getInetAddress(), dn);
                        final InputStream socketIn;
                        final OutputStream socketOut;
                        try {
                            socketIn = commsSession.getInput().getInputStream();
                            socketOut = commsSession.getOutput().getOutputStream();
                        } catch (final IOException e) {
                            LOG.error("Connection dropped from {} before any data was transmitted", peerUri);
                            try {
                                commsSession.close();
                            } catch (final IOException ioe) {
                            }
                            return;
                        }
                        final DataInputStream dis = new DataInputStream(socketIn);
                        final DataOutputStream dos = new DataOutputStream(socketOut);
                        ServerProtocol protocol = null;
                        Peer peer = null;
                        try {
                            // ensure that we are communicating with another NiFi
                            LOG.debug("Verifying magic bytes...");
                            verifyMagicBytes(dis, peerUri);
                            LOG.debug("Receiving Server Protocol Negotiation");
                            protocol = RemoteResourceFactory.receiveServerProtocolNegotiation(dis, dos);
                            protocol.setRootProcessGroup(rootGroup.get());
                            protocol.setNodeInformant(nodeInformant);
                            final PeerDescription description = new PeerDescription(clientHostName, clientPort, sslContext != null);
                            peer = new Peer(description, commsSession, peerUri, "nifi://localhost:" + getPort());
                            LOG.debug("Handshaking....");
                            protocol.handshake(peer);
                            if (!protocol.isHandshakeSuccessful()) {
                                LOG.error("Handshake failed with {}; closing connection", peer);
                                try {
                                    peer.close();
                                } catch (final IOException e) {
                                    LOG.warn("Failed to close {} due to {}", peer, e);
                                }
                                // no need to shutdown protocol because we failed to perform handshake
                                return;
                            }
                            commsSession.setTimeout((int) protocol.getRequestExpiration());
                            LOG.info("Successfully negotiated ServerProtocol {} Version {} with {}", new Object[] { protocol.getResourceName(), protocol.getVersionNegotiator().getVersion(), peer });
                            try {
                                while (!protocol.isShutdown()) {
                                    LOG.trace("Getting Protocol Request Type...");
                                    int timeoutCount = 0;
                                    RequestType requestType = null;
                                    while (requestType == null) {
                                        try {
                                            requestType = protocol.getRequestType(peer);
                                        } catch (final SocketTimeoutException e) {
                                            // Give the timeout a bit longer (twice as long) to receive the Request Type,
                                            // in order to attempt to receive more data without shutting down the socket if we don't
                                            // have to.
                                            LOG.debug("{} Timed out waiting to receive RequestType using {} with {}", new Object[] { this, protocol, peer });
                                            timeoutCount++;
                                            requestType = null;
                                            if (timeoutCount >= 2) {
                                                throw e;
                                            }
                                        }
                                    }
                                    handleRequest(protocol, peer, requestType);
                                }
                                LOG.debug("Finished communicating with {} ({})", peer, protocol);
                            } catch (final Exception e) {
                                LOG.error("Unable to communicate with remote instance {} ({}) due to {}; closing connection", peer, protocol, e.toString());
                                if (LOG.isDebugEnabled()) {
                                    LOG.error("", e);
                                }
                            }
                        } catch (final IOException e) {
                            LOG.error("Unable to communicate with remote instance {} due to {}; closing connection", peer, e.toString());
                            if (LOG.isDebugEnabled()) {
                                LOG.error("", e);
                            }
                        } catch (final Throwable t) {
                            LOG.error("Handshake failed when communicating with {}; closing connection. Reason for failure: {}", peerUri, t.toString());
                            if (LOG.isDebugEnabled()) {
                                LOG.error("", t);
                            }
                        } finally {
                            LOG.trace("Cleaning up");
                            try {
                                if (protocol != null && peer != null) {
                                    protocol.shutdown(peer);
                                }
                            } catch (final Exception protocolException) {
                                LOG.warn("Failed to shutdown protocol due to {}", protocolException.toString());
                            }
                            try {
                                if (peer != null) {
                                    peer.close();
                                }
                            } catch (final Exception peerException) {
                                LOG.warn("Failed to close peer due to {}; some resources may not be appropriately cleaned up", peerException.toString());
                            }
                            LOG.trace("Finished cleaning up");
                        }
                    }
                });
                thread.setName("Site-to-Site Worker Thread-" + (threadCount++));
                LOG.debug("Handing connection to {}", thread);
                thread.start();
                threads.add(thread);
                threads.removeIf(t -> !t.isAlive());
            }
            for (Thread thread : threads) {
                if (thread != null) {
                    thread.interrupt();
                }
            }
        }
    });
    listenerThread.setName("Site-to-Site Listener");
    listenerThread.start();
}

Example 4 with SocketChannelCommunicationsSession

use of org.apache.nifi.remote.io.socket.SocketChannelCommunicationsSession in project nifi by apache.

the class TestSocketFlowFileServerProtocol method getDefaultPeer.

private Peer getDefaultPeer(final HandshakeProperties handshakeProperties, final OutputStream outputStream) throws IOException {
    final PeerDescription description = new PeerDescription("peer-host", 8080, false);
    final byte[] inputBytes;
    try (final ByteArrayOutputStream bos = new ByteArrayOutputStream();
        final DataOutputStream dos = new DataOutputStream(bos)) {
        dos.writeUTF(handshakeProperties.getCommsIdentifier());
        dos.writeUTF(handshakeProperties.getTransitUriPrefix());
        // num of properties
        dos.writeInt(1);
        dos.writeUTF(HandshakeProperty.GZIP.name());
        dos.writeUTF(String.valueOf(handshakeProperties.isUseGzip()));
        dos.flush();
        inputBytes = bos.toByteArray();
    }
    final InputStream inputStream = new ByteArrayInputStream(inputBytes);
    final SocketChannelCommunicationsSession commsSession = mock(SocketChannelCommunicationsSession.class);
    final SocketChannelInput channelInput = mock(SocketChannelInput.class);
    final SocketChannelOutput channelOutput = mock(SocketChannelOutput.class);
    when(commsSession.getInput()).thenReturn(channelInput);
    when(commsSession.getOutput()).thenReturn(channelOutput);
    when(channelInput.getInputStream()).thenReturn(inputStream);
    when(channelOutput.getOutputStream()).thenReturn(outputStream);
    final String peerUrl = "http://peer-host:8080/";
    final String clusterUrl = "cluster-url";
    return new Peer(description, commsSession, peerUrl, clusterUrl);
}

Example 5 with SocketChannelCommunicationsSession

use of org.apache.nifi.remote.io.socket.SocketChannelCommunicationsSession in project nifi by apache.

the class EndpointConnectionPool method establishSiteToSiteConnection.

private CommunicationsSession establishSiteToSiteConnection(final String hostname, final int port) throws IOException {
    final boolean siteToSiteSecure = siteInfoProvider.isSecure();
    CommunicationsSession commsSession = null;
    try {
        if (siteToSiteSecure) {
            if (sslContext == null) {
                throw new IOException("Unable to communicate with " + hostname + ":" + port + " because it requires Secure Site-to-Site communications, but this instance is not configured for secure communications");
            }
            final SSLSocketChannel socketChannel = new SSLSocketChannel(sslContext, hostname, port, localAddress, true);
            socketChannel.connect();
            commsSession = new SSLSocketChannelCommunicationsSession(socketChannel);
            try {
                commsSession.setUserDn(socketChannel.getDn());
            } catch (final CertificateException ex) {
                throw new IOException(ex);
            }
        } else {
            final SocketChannel socketChannel = SocketChannel.open();
            if (localAddress != null) {
                final SocketAddress localSocketAddress = new InetSocketAddress(localAddress, 0);
                socketChannel.socket().bind(localSocketAddress);
            }
            socketChannel.socket().connect(new InetSocketAddress(hostname, port), commsTimeout);
            socketChannel.socket().setSoTimeout(commsTimeout);
            commsSession = new SocketChannelCommunicationsSession(socketChannel);
        }
        commsSession.getOutput().getOutputStream().write(CommunicationsSession.MAGIC_BYTES);
    } catch (final IOException ioe) {
        if (commsSession != null) {
            commsSession.close();
        }
        throw ioe;
    }
    return commsSession;
}