Fixes issue #79: do not close user SSH connections

The issue was fixed in the way proposed in #79, an additional field
called `CloseHandler` was added to the `Client` struct. This handler can
either be `EmptyHandler` which is equivalent to a no-op, or a
`CloseSSHClient` which closes the passed SSH client when executed.

The `EmptyHandler` is used by default unless `Client.Connect()` is
called to establish the SSH connection. The reasoning is that whenever
someone passes their own `ssh.Client` to the library, `Connect()` will
not be called, thus signaling that we should not manage the
`ssh.Client`.

To ensure the correctness of this fix, two additional test cases were added
(1) `TestUserSuppliedSSHClientDoesNotClose` which creates an `scp`
client using a existing `ssh.Client` and ensures that the client is not
closed by attempting to create a new session from it.
(2) `TestSSHClientNoLeak` which uses `Connect()` to establish the SSH
connection and ensures that its underlying `ssh.Client` is no longer
functioning using the same mechanism as (1).

Author: bramvdbogaerde

client.go

@@ -1,4 +1,4 @@
-/* Copyright (c) 2021 Bram Vandenbogaerde And Contributors
+/* Copyright (c) 2024 Bram Vandenbogaerde And Contributors
  * You may use, distribute or modify this code under the
  * terms of the Mozilla Public License 2.0, which is distributed
  * along with the source code.
@@ -21,6 +21,27 @@ import (
 	"golang.org/x/crypto/ssh"
 )
 
+// Callback for freeing managed resources
+type ICloseHandler interface {
+	Close()
+}
+
+// Close handler equivalent to a no-op. Used by default
+// when no resources have to be cleaned.
+type EmptyHandler struct{}
+
+func (EmptyHandler) Close() {}
+
+// Close handler to close an SSH client
+type CloseSSHCLient struct {
+	// Reference to the used SSH client
+	sshClient *ssh.Client
+}
+
+func (scp CloseSSHCLient) Close() {
+	scp.sshClient.Close()
+}
+
 type PassThru func(r io.Reader, total int64) io.Reader
 
 type Client struct {
@@ -39,6 +60,10 @@ type Client struct {
 
 	// RemoteBinary the absolute path to the remote SCP binary.
 	RemoteBinary string
+
+	// Handler called when calling `Close` to clean up any remaining
+	// resources managed by `Client`.
+	closeHandler ICloseHandler
 }
 
 // Connect connects to the remote SSH server, returns error if it couldn't establish a session to the SSH server.
@@ -49,9 +74,16 @@ func (a *Client) Connect() error {
 	}
 
 	a.sshClient = client
+	a.closeHandler = CloseSSHCLient{sshClient: client}
 	return nil
 }
 
+// Returns the underlying SSH client, this should be used carefully as
+// it will be closed by `client.Close`.
+func (a *Client) SSHClient() *ssh.Client {
+	return a.sshClient
+}
+
 // CopyFromFile copies the contents of an os.File to a remote location, it will get the length of the file by looking it up from the filesystem.
 func (a *Client) CopyFromFile(ctx context.Context, file os.File, remotePath string, permissions string) error {
 	return a.CopyFromFilePassThru(ctx, file, remotePath, permissions, nil)
@@ -347,7 +379,5 @@ func (a *Client) CopyFromRemotePassThru(ctx context.Context, w io.Writer, remote
 }
 
 func (a *Client) Close() {
-	if a.sshClient != nil {
-		a.sshClient.Close()
-	}
+	a.closeHandler.Close()
 }

configurer.go

@@ -78,5 +78,6 @@ func (c *ClientConfigurer) Create() Client {
 		Timeout:      c.timeout,
 		RemoteBinary: c.remoteBinary,
 		sshClient:    c.sshClient,
+                closeHandler: EmptyHandler{},
 	}
 }

tests/basic_test.go

@@ -1,4 +1,4 @@
-package tests
+package scp
 
 import (
 	"context"
@@ -317,3 +317,75 @@ func TestFileNotFound(t *testing.T) {
 		t.Errorf("Expected %v, got %v", expected, err.Error())
 	}
 }
+
+func TestUserSuppliedSSHClientDoesNotClose(t *testing.T) {
+	// create the SSH connection
+	clientConfig, err := buildClientConfig()
+	if err != nil {
+		t.Error("Could not build client config", clientConfig)
+	}
+
+	sshClient, err := ssh.Dial("tcp", "127.0.0.1:2244", &clientConfig)
+	if err != nil {
+		t.Error("Could not establish SSH connection", err)
+	}
+	defer sshClient.Close()
+
+	// create the SCP client
+	client, err := scp.NewClientBySSH(sshClient)
+	if err != nil {
+		t.Error("Could not create SCP client", err)
+	}
+
+	// copy a file for good measure
+
+	f, _ := os.Open("./data/upload_file.txt")
+	defer f.Close()
+
+	err = client.CopyFile(context.Background(), f, "/data/test.txt", "0777")
+
+	if err != nil {
+		t.Error("Could not copy file to remote", err)
+	}
+
+	// then close the SCP client
+	client.Close()
+
+	var session *ssh.Session
+
+	// ensure that the SSH client is still opened
+	// we do so by creating a new session, if this fails
+	// the SSH connection was already closed
+	if session, err = sshClient.NewSession(); err != nil {
+		t.Fatal("SSH session was already closed.")
+	}
+
+	session.Close()
+}
+
+// Ensure that the underlying SSH client managed by the library is correctly closed
+// after closing the SCP connection
+func TestSSHClientNoLeak(t *testing.T) {
+	client := establishConnection(t)
+
+	// copy a file for good measure
+	f, _ := os.Open("./data/upload_file.txt")
+	defer f.Close()
+
+	err := client.CopyFile(context.Background(), f, "/data/test.txt", "0777")
+
+	if err != nil {
+		t.Error("Could not copy file to remote", err)
+	}
+
+	// then close the SCP client
+	client.Close()
+
+	// ensure that the SSH client is still opened
+	// we do so by creating a new session, if this fails
+	// the SSH connection was already closed
+	if session, err := client.SSHClient().NewSession(); err == nil {
+		session.Close()
+		t.Fatal("SSH session was not closed.")
+	}
+}