merge

7 files changed, 154 insertions, 42 deletions

diff --git a/internal/clients/baseclient.go b/internal/clients/baseclient.go
index 4a7bd84..3cd85fe 100644
--- a/internal/clients/baseclient.go
+++ b/internal/clients/baseclient.go
@@ -56,7 +56,7 @@ func (c *baseClient) init() {
 	}
 	c.sshAuthMethods, c.hostKeyCallback = client.InitSSHAuthMethods(
 		c.Args.SSHAuthMethods, c.Args.SSHHostKeyCallback, c.Args.TrustAllHosts,
-		c.throttleCh, c.Args.PrivateKeyPathFile)
+		c.throttleCh, c.Args.SSHPrivateKeyPathFile)
 }
 
 func (c *baseClient) makeConnections(maker maker) {
@@ -112,10 +112,19 @@ func (c *baseClient) startConnection(ctx context.Context, i int,
 		// Retrieve status code from handler (dtail client will exit with that status)
 		status = conn.Handler().Status()
 
+		// Do we want to retry?
 		if !c.retry {
+			// No, we don't.
 			return
 		}
+		select {
+		case <-ctx.Done():
+			// No, context is done, so no retry.
+			return
+		default:
+		}
 
+		// Yes, we want to retry.
 		time.Sleep(time.Second * 2)
 		dlog.Client.Debug(conn.Server(), "Reconnecting")
 		conn = c.makeConnection(conn.Server(), c.sshAuthMethods, c.hostKeyCallback)
diff --git a/internal/config/args.go b/internal/config/args.go
index 3d7ac7d..859166d 100644
--- a/internal/config/args.go
+++ b/internal/config/args.go
@@ -15,31 +15,31 @@ import (
 // Args is a helper struct to summarize common client arguments.
 type Args struct {
 	lcontext.LContext
-	Arguments          []string
-	ConfigFile         string
-	ConnectionsPerCPU  int
-	Discovery          string
-	LogDir             string
-	Logger             string
-	LogLevel           string
-	Mode               omode.Mode
-	NoColor            bool
-	PrivateKeyPathFile string
-	QueryStr           string
-	Quiet              bool
-	RegexInvert        bool
-	RegexStr           string
-	Serverless         bool
-	ServersStr         string
-	Spartan            bool
-	SSHAuthMethods     []gossh.AuthMethod
-	SSHBindAddress     string
-	SSHHostKeyCallback gossh.HostKeyCallback
-	SSHPort            int
-	Timeout            int
-	TrustAllHosts      bool
-	UserName           string
-	What               string
+	Arguments             []string
+	ConfigFile            string
+	ConnectionsPerCPU     int
+	Discovery             string
+	LogDir                string
+	Logger                string
+	LogLevel              string
+	Mode                  omode.Mode
+	NoColor               bool
+	QueryStr              string
+	Quiet                 bool
+	RegexInvert           bool
+	RegexStr              string
+	SSHAuthMethods        []gossh.AuthMethod
+	SSHBindAddress        string
+	SSHHostKeyCallback    gossh.HostKeyCallback
+	SSHPort               int
+	SSHPrivateKeyPathFile string
+	Serverless            bool
+	ServersStr            string
+	Spartan               bool
+	Timeout               int
+	TrustAllHosts         bool
+	UserName              string
+	What                  string
 }
 
 func (a *Args) String() string {
@@ -56,7 +56,6 @@ func (a *Args) String() string {
 	sb.WriteString(fmt.Sprintf("%s:%v,", "Logger", a.Logger))
 	sb.WriteString(fmt.Sprintf("%s:%v,", "Mode", a.Mode))
 	sb.WriteString(fmt.Sprintf("%s:%v,", "NoColor", a.NoColor))
-	sb.WriteString(fmt.Sprintf("%s:%v,", "PrivateKeyPathFile", a.PrivateKeyPathFile))
 	sb.WriteString(fmt.Sprintf("%s:%v,", "QueryStr", a.QueryStr))
 	sb.WriteString(fmt.Sprintf("%s:%v,", "Quiet", a.Quiet))
 	sb.WriteString(fmt.Sprintf("%s:%v,", "RegexInvert", a.RegexInvert))
@@ -64,6 +63,7 @@ func (a *Args) String() string {
 	sb.WriteString(fmt.Sprintf("%s:%v,", "SSHAuthMethods", a.SSHAuthMethods))
 	sb.WriteString(fmt.Sprintf("%s:%v,", "SSHBindAddress", a.SSHBindAddress))
 	sb.WriteString(fmt.Sprintf("%s:%v,", "SSHHostKeyCallback", a.SSHHostKeyCallback))
+	sb.WriteString(fmt.Sprintf("%s:%v,", "SSHPrivateKeyPathFile", a.SSHPrivateKeyPathFile))
 	sb.WriteString(fmt.Sprintf("%s:%v,", "SSHPort", a.SSHPort))
 	sb.WriteString(fmt.Sprintf("%s:%v,", "Serverless", a.Serverless))
 	sb.WriteString(fmt.Sprintf("%s:%v,", "ServersStr", a.ServersStr))
diff --git a/internal/config/initializer.go b/internal/config/initializer.go
index 137b831..f1f9b9b 100644
--- a/internal/config/initializer.go
+++ b/internal/config/initializer.go
@@ -80,10 +80,10 @@ func (in *initializer) transformConfig(sourceProcess source.Source, args *Args,
 	}
 }
 
-// There are some special options which can be set by environment variable.
 func (in *initializer) processEnvVars() {
 	if Env("DTAIL_INTEGRATION_TEST_RUN_MODE") {
 		os.Setenv("DTAIL_HOSTNAME_OVERRIDE", "integrationtest")
+		in.Server.HostKeyFile = "./ssh_host_key"
 	}
 }
 
diff --git a/internal/io/fs/readfile.go b/internal/io/fs/readfile.go
index 5815aa3..a42fc53 100644
--- a/internal/io/fs/readfile.go
+++ b/internal/io/fs/readfile.go
@@ -233,7 +233,6 @@ func (f readFile) filterWithoutLContext(ctx context.Context, rawLines <-chan *by
 				return
 			}
 			if filteredLine, ok := f.transmittable(line, len(lines), cap(lines), re); ok {
-				//dlog.Common.Trace("TODO", "lines", lines, len(lines), cap(lines))
 				select {
 				case lines <- filteredLine:
 				case <-ctx.Done():
diff --git a/internal/ssh/client/authmethods.go b/internal/ssh/client/authmethods.go
index 87d40d8..65fec1f 100644
--- a/internal/ssh/client/authmethods.go
+++ b/internal/ssh/client/authmethods.go
@@ -26,6 +26,22 @@ func InitSSHAuthMethods(sshAuthMethods []gossh.AuthMethod,
 	return initKnownHostsAuthMethods(trustAllHosts, throttleCh, privateKeyPath)
 }
 
+func initIntegrationTestKnownHostsAuthMethods() []gossh.AuthMethod {
+	var sshAuthMethods []gossh.AuthMethod
+	privateKeyPath := "./id_rsa"
+
+	GeneratePrivatePublicKeyPairIfNotExists(privateKeyPath, 4096)
+	authMethod, err := ssh.PrivateKey(privateKeyPath)
+	if err != nil {
+		dlog.Client.FatalPanic("Unable to use private SSH key", privateKeyPath, err)
+	}
+
+	sshAuthMethods = append(sshAuthMethods, authMethod)
+	dlog.Client.Debug("initKnownHostsAuthMethods",
+		"Added path to list of auth methods, not adding further methods", privateKeyPath)
+	return sshAuthMethods
+}
+
 func initKnownHostsAuthMethods(trustAllHosts bool, throttleCh chan struct{},
 	privateKeyPath string) ([]gossh.AuthMethod, HostKeyCallback) {
 
@@ -40,17 +56,19 @@ func initKnownHostsAuthMethods(trustAllHosts bool, throttleCh chan struct{},
 	if err != nil {
 		dlog.Client.FatalPanic(knownHostsFile, err)
 	}
-
 	dlog.Client.Debug("initKnownHostsAuthMethods", "Added known hosts file path", knownHostsFile)
 
-	// First try to read custom private key path.
+	if config.Env("DTAIL_INTEGRATION_TEST_RUN_MODE") {
+		return initIntegrationTestKnownHostsAuthMethods(), knownHostsCallback
+	}
+
+	// Try to read custom private key path.
 	if privateKeyPath != "" {
 		authMethod, err := ssh.PrivateKey(privateKeyPath)
 		if err == nil {
 			sshAuthMethods = append(sshAuthMethods, authMethod)
 			dlog.Client.Debug("initKnownHostsAuthMethods",
-				"Added path to list of auth methods, not adding further methods",
-				privateKeyPath)
+				"Added path to list of auth methods, not adding further methods", privateKeyPath)
 			return sshAuthMethods, knownHostsCallback
 		}
 		dlog.Client.FatalPanic("Unable to use private SSH key", privateKeyPath, err)
@@ -64,8 +82,7 @@ func initKnownHostsAuthMethods(trustAllHosts bool, throttleCh chan struct{},
 			"to list of auth methods, not adding further methods")
 		return sshAuthMethods, knownHostsCallback
 	}
-	dlog.Client.Debug("initKnownHostsAuthMethods",
-		"Unable to init SSH Agent auth method", err)
+	dlog.Client.Debug("initKnownHostsAuthMethods", "Unable to init SSH Agent auth method", err)
 
 	// Third, try Linux/UNIX default key paths
 	privateKeyPath = os.Getenv("HOME") + "/.ssh/id_rsa"
@@ -76,8 +93,7 @@ func initKnownHostsAuthMethods(trustAllHosts bool, throttleCh chan struct{},
 			"Added path to list of auth methods, not adding further methods", privateKeyPath)
 		return sshAuthMethods, knownHostsCallback
 	}
-	dlog.Client.Debug("initKnownHostsAuthMethods", "Unable to use private key",
-		privateKeyPath, err)
+	dlog.Client.Debug("initKnownHostsAuthMethods", "Unable to use private key", privateKeyPath, err)
 
 	privateKeyPath = os.Getenv("HOME") + "/.ssh/id_dsa"
 	authMethod, err = ssh.PrivateKey(privateKeyPath)
@@ -97,10 +113,7 @@ func initKnownHostsAuthMethods(trustAllHosts bool, throttleCh chan struct{},
 		return sshAuthMethods, knownHostsCallback
 	}
 
-	dlog.Client.Debug("initKnownHostsAuthMethods", "Unable to use private key",
-		privateKeyPath, err)
-
-	dlog.Client.FatalPanic("Unable to find private SSH key information")
+	dlog.Client.FatalPanic("Unable to find private SSH key information", privateKeyPath, err)
 	// Never reach this point.
 	return sshAuthMethods, knownHostsCallback
 }
diff --git a/internal/ssh/client/clientkeypair.go b/internal/ssh/client/clientkeypair.go
new file mode 100644
index 0000000..0e21d0c
--- /dev/null
+++ b/internal/ssh/client/clientkeypair.go
@@ -0,0 +1,91 @@
+package client
+
+import (
+	"crypto/rand"
+	"crypto/rsa"
+	"crypto/x509"
+	"encoding/pem"
+	"fmt"
+	"io/ioutil"
+	"os"
+
+	"github.com/mimecast/dtail/internal/io/dlog"
+	"golang.org/x/crypto/ssh"
+)
+
+// GeneratePrivatePublicKeyPairIfNotExists generates a SSH key pair (used by the integration tests)
+func GeneratePrivatePublicKeyPairIfNotExists(keyPath string, bitSize int) {
+	if _, err := os.Stat(keyPath); err == nil {
+		dlog.Common.Debug("Private/public key pair already exists", keyPath)
+		return
+	}
+	GeneratePrivatePublicKeyPair(keyPath, bitSize)
+}
+
+// GeneratePrivatePublicKeyPair generates a SSH key pair (used by the integration tests)
+func GeneratePrivatePublicKeyPair(keyPath string, bitSize int) {
+	privateKeyPath := keyPath
+	publicKeyPath := fmt.Sprintf("%s.pub", keyPath)
+
+	dlog.Common.Debug("Generating private/public key pair", privateKeyPath, publicKeyPath)
+
+	privateKey, err := generatePrivateKey(bitSize)
+	if err != nil {
+		dlog.Common.FatalPanic(err)
+	}
+	publicKeyBytes, err := generatePublicKey(&privateKey.PublicKey)
+	if err != nil {
+		dlog.Common.FatalPanic(err)
+	}
+	privateKeyBytes := encodePrivateKeyToPEM(privateKey)
+	err = writeKey(privateKeyBytes, privateKeyPath)
+	if err != nil {
+		dlog.Common.FatalPanic(err)
+	}
+	err = writeKey([]byte(publicKeyBytes), publicKeyPath)
+	if err != nil {
+		dlog.Common.FatalPanic(err)
+	}
+
+	dlog.Common.Debug("Done generating private/public key pair", privateKeyPath, publicKeyPath)
+}
+
+func generatePrivateKey(bitSize int) (*rsa.PrivateKey, error) {
+	privateKey, err := rsa.GenerateKey(rand.Reader, bitSize)
+	if err != nil {
+		return nil, err
+	}
+	err = privateKey.Validate()
+	if err != nil {
+		return nil, err
+	}
+	return privateKey, nil
+}
+
+func encodePrivateKeyToPEM(privateKey *rsa.PrivateKey) []byte {
+	privDER := x509.MarshalPKCS1PrivateKey(privateKey)
+	privBlock := pem.Block{
+		Type:    "RSA PRIVATE KEY",
+		Headers: nil,
+		Bytes:   privDER,
+	}
+	privatePEM := pem.EncodeToMemory(&privBlock)
+	return privatePEM
+}
+
+func generatePublicKey(privatekey *rsa.PublicKey) ([]byte, error) {
+	publicRsaKey, err := ssh.NewPublicKey(privatekey)
+	if err != nil {
+		return nil, err
+	}
+	pubKeyBytes := ssh.MarshalAuthorizedKey(publicRsaKey)
+	return pubKeyBytes, nil
+}
+
+func writeKey(keyBytes []byte, saveFileTo string) error {
+	err := ioutil.WriteFile(saveFileTo, keyBytes, 0600)
+	if err != nil {
+		return err
+	}
+	return nil
+}
diff --git a/internal/version/version.go b/internal/version/version.go
index 68b9e6e..5872980 100644
--- a/internal/version/version.go
+++ b/internal/version/version.go
@@ -13,7 +13,7 @@ const (
 	// Name of DTail.
 	Name string = "DTail"
 	// Version of DTail.
-	Version string = "4.0.0-RC1"
+	Version string = "4.0.0-RC2"
 	// Additional information for DTail
 	Additional string = "Have a lot of fun!"
 )