From eaaff4ede9c4352a5f2f1817237c34679a13d72d Mon Sep 17 00:00:00 2001 From: Peter Turcsanyi Date: Wed, 17 Aug 2022 13:08:07 +0200 Subject: [PATCH] NIFI-10364: Simplified connection/session handling in SmbjClientService This closes #6307. Signed-off-by: Tamas Palfy --- .../apache/nifi/processors/smb/ListSmb.java | 2 +- .../smb/SmbjClientProviderService.java | 62 ++++++++++++++++-- .../nifi/services/smb/SmbjClientService.java | 65 ++++--------------- .../nifi/services/smb/NiFiSmbjClientTest.java | 19 +----- 4 files changed, 73 insertions(+), 75 deletions(-) diff --git a/nifi-nar-bundles/nifi-smb-bundle/nifi-smb-processors/src/main/java/org/apache/nifi/processors/smb/ListSmb.java b/nifi-nar-bundles/nifi-smb-bundle/nifi-smb-processors/src/main/java/org/apache/nifi/processors/smb/ListSmb.java index f9eab575d4..3388d02ce7 100644 --- a/nifi-nar-bundles/nifi-smb-bundle/nifi-smb-processors/src/main/java/org/apache/nifi/processors/smb/ListSmb.java +++ b/nifi-nar-bundles/nifi-smb-bundle/nifi-smb-processors/src/main/java/org/apache/nifi/processors/smb/ListSmb.java @@ -96,7 +96,7 @@ import org.apache.nifi.services.smb.SmbListableEntity; + "Share root directory. For example, for a given remote location" + "smb://HOSTNAME:PORT/SHARE/DIRECTORY, and a file is being listed from " + "smb://HOSTNAME:PORT/SHARE/DIRECTORY/sub/folder/file then the path attribute will be set to " - + "\"DIRECTORY/sub/folder/file\"."), + + "\"DIRECTORY/sub/folder\"."), @WritesAttribute(attribute = SERVICE_LOCATION, description = "The SMB URL of the share."), @WritesAttribute(attribute = LAST_MODIFIED_TIME, description = diff --git a/nifi-nar-bundles/nifi-smb-bundle/nifi-smb-smbj-client/src/main/java/org/apache/nifi/services/smb/SmbjClientProviderService.java b/nifi-nar-bundles/nifi-smb-bundle/nifi-smb-smbj-client/src/main/java/org/apache/nifi/services/smb/SmbjClientProviderService.java index 3ee35d6289..e1de5e5bcf 100644 --- a/nifi-nar-bundles/nifi-smb-bundle/nifi-smb-smbj-client/src/main/java/org/apache/nifi/services/smb/SmbjClientProviderService.java +++ b/nifi-nar-bundles/nifi-smb-bundle/nifi-smb-smbj-client/src/main/java/org/apache/nifi/services/smb/SmbjClientProviderService.java @@ -30,6 +30,11 @@ import java.io.IOException; import java.net.URI; import java.util.Collections; import java.util.List; + +import com.hierynomus.smbj.connection.Connection; +import com.hierynomus.smbj.session.Session; +import com.hierynomus.smbj.share.DiskShare; +import com.hierynomus.smbj.share.Share; import org.apache.nifi.annotation.documentation.CapabilityDescription; import org.apache.nifi.annotation.documentation.Tags; import org.apache.nifi.annotation.lifecycle.OnDisabled; @@ -116,16 +121,63 @@ public class SmbjClientProviderService extends AbstractControllerService impleme @Override public SmbClientService getClient() throws IOException { - final SmbjClientService client = new SmbjClientService(smbClient, authenticationContext, getServiceLocation()); + Connection connection = smbClient.connect(hostname, port); try { - client.connectToShare(hostname, port, shareName); + return connectToShare(connection); } catch (IOException e) { - client.forceFullyCloseConnection(); - client.connectToShare(hostname, port, shareName); + getLogger().debug("Closing stale connection and trying to create a new one for share " + getServiceLocation()); + + closeConnection(connection); + + connection = smbClient.connect(hostname, port); + return connectToShare(connection); + } + } + + private SmbjClientService connectToShare(Connection connection) throws IOException { + final Session session; + final Share share; + + try { + session = connection.authenticate(authenticationContext); + } catch (Exception e) { + throw new IOException("Could not create session for share " + getServiceLocation(), e); } - return client; + try { + share = session.connectShare(shareName); + } catch (Exception e) { + closeSession(session); + throw new IOException("Could not connect to share " + getServiceLocation(), e); + } + + if (!(share instanceof DiskShare)) { + closeSession(session); + throw new IllegalArgumentException("DiskShare not found. Share " + share.getClass().getSimpleName() + " found on " + getServiceLocation()); + } + + return new SmbjClientService(session, (DiskShare) share, getServiceLocation()); + } + + private void closeConnection(Connection connection) { + try { + if (connection != null) { + connection.close(true); + } + } catch (Exception e) { + getLogger().error("Could not close connection to {}", getServiceLocation(), e); + } + } + + private void closeSession(Session session) { + try { + if (session != null) { + session.close(); + } + } catch (Exception e) { + getLogger().error("Could not close session to {}", getServiceLocation(), e); + } } @Override diff --git a/nifi-nar-bundles/nifi-smb-bundle/nifi-smb-smbj-client/src/main/java/org/apache/nifi/services/smb/SmbjClientService.java b/nifi-nar-bundles/nifi-smb-bundle/nifi-smb-smbj-client/src/main/java/org/apache/nifi/services/smb/SmbjClientService.java index a103fec359..ae9307ea64 100644 --- a/nifi-nar-bundles/nifi-smb-bundle/nifi-smb-smbj-client/src/main/java/org/apache/nifi/services/smb/SmbjClientService.java +++ b/nifi-nar-bundles/nifi-smb-bundle/nifi-smb-smbj-client/src/main/java/org/apache/nifi/services/smb/SmbjClientService.java @@ -16,7 +16,6 @@ */ package org.apache.nifi.services.smb; -import static java.lang.String.format; import static java.util.Arrays.asList; import static java.util.stream.StreamSupport.stream; @@ -27,14 +26,13 @@ import com.hierynomus.mssmb2.SMB2CreateDisposition; import com.hierynomus.mssmb2.SMB2CreateOptions; import com.hierynomus.mssmb2.SMB2ShareAccess; import com.hierynomus.mssmb2.SMBApiException; -import com.hierynomus.smbj.SMBClient; -import com.hierynomus.smbj.auth.AuthenticationContext; -import com.hierynomus.smbj.connection.Connection; import com.hierynomus.smbj.session.Session; import com.hierynomus.smbj.share.Directory; import com.hierynomus.smbj.share.DiskShare; import com.hierynomus.smbj.share.File; -import com.hierynomus.smbj.share.Share; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + import java.io.IOException; import java.io.OutputStream; import java.net.URI; @@ -42,66 +40,31 @@ import java.util.EnumSet; import java.util.List; import java.util.stream.Stream; -public class SmbjClientService implements SmbClientService { +class SmbjClientService implements SmbClientService { + + private final static Logger LOGGER = LoggerFactory.getLogger(SmbjClientService.class); private static final List SPECIAL_DIRECTORIES = asList(".", ".."); private static final long UNCATEGORISED_ERROR = -1L; - final private AuthenticationContext authenticationContext; - final private SMBClient smbClient; - final private URI serviceLocation; + private final Session session; + private final DiskShare share; + private final URI serviceLocation; - private Connection connection; - private Session session; - private DiskShare share; - - public SmbjClientService(SMBClient smbClient, AuthenticationContext authenticationContext, URI serviceLocation) { - this.smbClient = smbClient; - this.authenticationContext = authenticationContext; + SmbjClientService(Session session, DiskShare share, URI serviceLocation) { + this.session = session; + this.share = share; this.serviceLocation = serviceLocation; } - public void connectToShare(String hostname, int port, String shareName) throws IOException { - Share share; - try { - connection = smbClient.connect(hostname, port); - session = connection.authenticate(authenticationContext); - share = session.connectShare(shareName); - } catch (Exception e) { - close(); - throw new IOException("Could not connect to share " + format("%s:%d/%s", hostname, port, shareName), e); - } - if (share instanceof DiskShare) { - this.share = (DiskShare) share; - } else { - close(); - throw new IllegalArgumentException("DiskShare not found. Share " + - share.getClass().getSimpleName() + " found on " + format("%s:%d/%s", hostname, port, - shareName)); - } - } - - public void forceFullyCloseConnection() { - try { - if (connection != null) { - connection.close(true); - } - } catch (IOException ignore) { - } finally { - connection = null; - } - } - @Override public void close() { try { if (session != null) { session.close(); } - } catch (IOException ignore) { - - } finally { - session = null; + } catch (Exception e) { + LOGGER.error("Could not close session to {}", serviceLocation, e); } } diff --git a/nifi-nar-bundles/nifi-smb-bundle/nifi-smb-smbj-client/src/test/java/org/apache/nifi/services/smb/NiFiSmbjClientTest.java b/nifi-nar-bundles/nifi-smb-bundle/nifi-smb-smbj-client/src/test/java/org/apache/nifi/services/smb/NiFiSmbjClientTest.java index abf032072d..f00b505bea 100644 --- a/nifi-nar-bundles/nifi-smb-bundle/nifi-smb-smbj-client/src/test/java/org/apache/nifi/services/smb/NiFiSmbjClientTest.java +++ b/nifi-nar-bundles/nifi-smb-bundle/nifi-smb-smbj-client/src/test/java/org/apache/nifi/services/smb/NiFiSmbjClientTest.java @@ -20,9 +20,6 @@ import static org.mockito.ArgumentMatchers.anyString; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; -import com.hierynomus.smbj.SMBClient; -import com.hierynomus.smbj.auth.AuthenticationContext; -import com.hierynomus.smbj.connection.Connection; import com.hierynomus.smbj.session.Session; import com.hierynomus.smbj.share.DiskShare; import org.junit.jupiter.api.BeforeEach; @@ -33,20 +30,11 @@ import org.mockito.MockitoAnnotations; class NiFiSmbjClientTest { - @Mock - DiskShare share; - - @Mock - SMBClient smbClient; - - @Mock - AuthenticationContext authenticationContext; - @Mock Session session; @Mock - Connection connection; + DiskShare share; @InjectMocks SmbjClientService underTest; @@ -58,17 +46,12 @@ class NiFiSmbjClientTest { @Test public void shouldCreateDirectoriesRecursively() throws Exception { - - when(smbClient.connect("hostname", 445)) - .thenReturn(connection); - when(connection.authenticate(authenticationContext)).thenReturn(session); when(session.connectShare(anyString())).thenReturn(share); when(share.fileExists("directory")).thenReturn(true); when(share.fileExists("path")).thenReturn(false); when(share.fileExists("to")).thenReturn(false); when(share.fileExists("create")).thenReturn(false); - underTest.connectToShare("hostname", 445, "share"); underTest.createDirectory("directory/path/to/create"); verify(share).mkdir("directory/path");