Dont use relaxed SSH Auth mode anymore for integration tests

2 files changed, 53 insertions, 24 deletions

diff --git a/internal/ssh/client/authmethods.go b/internal/ssh/client/authmethods.go
index 2ee32ad..87d40d8 100644
--- a/internal/ssh/client/authmethods.go
+++ b/internal/ssh/client/authmethods.go
@@ -1,6 +1,7 @@
 package client
 
 import (
+	"fmt"
 	"os"
 
 	"github.com/mimecast/dtail/internal/config"
@@ -29,9 +30,13 @@ func initKnownHostsAuthMethods(trustAllHosts bool, throttleCh chan struct{},
 	privateKeyPath string) ([]gossh.AuthMethod, HostKeyCallback) {
 
 	var sshAuthMethods []gossh.AuthMethod
-	knownHostsFile := config.SSHKnownHostsFile()
-	knownHostsCallback, err := NewKnownHostsCallback(knownHostsFile, trustAllHosts,
-		throttleCh)
+	knownHostsFile := fmt.Sprintf("%s/.ssh/known_hosts", os.Getenv("HOME"))
+	if config.Env("DTAIL_INTEGRATION_TEST_RUN_MODE") {
+		// In case of integration test, override known hosts file path.
+		knownHostsFile = "./known_hosts"
+	}
+
+	knownHostsCallback, err := NewKnownHostsCallback(knownHostsFile, trustAllHosts, throttleCh)
 	if err != nil {
 		dlog.Client.FatalPanic(knownHostsFile, err)
 	}
diff --git a/internal/ssh/server/publickeycallback.go b/internal/ssh/server/publickeycallback.go
index ebc428a..585469f 100644
--- a/internal/ssh/server/publickeycallback.go
+++ b/internal/ssh/server/publickeycallback.go
@@ -4,7 +4,7 @@ import (
 	"fmt"
 	"io/ioutil"
 	"os"
-	osUser "os/user"
+	goUser "os/user"
 
 	"github.com/mimecast/dtail/internal/config"
 	"github.com/mimecast/dtail/internal/io/dlog"
@@ -24,24 +24,14 @@ func PublicKeyCallback(c gossh.ConnMetadata,
 	}
 
 	dlog.Common.Info(user, "Incoming authorization")
-	cwd, err := os.Getwd()
-	if err != nil {
-		return nil, fmt.Errorf("Unable to get current working directory|%s|", err.Error())
-	}
 	if config.ServerRelaxedAuthEnable {
 		dlog.Common.Fatal(user, "Granting permissions via relaxed-auth")
 		return nil, nil
 	}
 
-	authorizedKeysFile := fmt.Sprintf("%s/%s/%s.authorized_keys", cwd,
-		config.Common.CacheDir, user.Name)
-	if _, err := os.Stat(authorizedKeysFile); os.IsNotExist(err) {
-		user, err := osUser.Lookup(user.Name)
-		if err != nil {
-			return nil, fmt.Errorf("Unable to authorize|%s|%s|", user, err.Error())
-		}
-		// Fallback to ~
-		authorizedKeysFile = user.HomeDir + "/.ssh/authorized_keys"
+	authorizedKeysFile, err := authorizedKeysFile(user)
+	if err != nil {
+		return nil, err
 	}
 
 	dlog.Common.Info(user, "Reading", authorizedKeysFile)
@@ -51,11 +41,17 @@ func PublicKeyCallback(c gossh.ConnMetadata,
 			authorizedKeysFile, user, err.Error())
 	}
 
+	return verifyAuthorizedKeys(user, authorizedKeysBytes, offeredPubKey)
+}
+
+func verifyAuthorizedKeys(user *user.User, authorizedKeysBytes []byte,
+	offeredPubKey gossh.PublicKey) (*gossh.Permissions, error) {
+
 	authorizedKeysMap := map[string]bool{}
 	for len(authorizedKeysBytes) > 0 {
 		authorizedPubKey, _, _, restBytes, err := gossh.ParseAuthorizedKey(authorizedKeysBytes)
 		if err != nil {
-			return nil, fmt.Errorf("Unable to parse authorized keys bytes|%s|%s",
+			return nil, fmt.Errorf("unable to parse authorized keys bytes|%s|%s",
 				user, err.Error())
 		}
 		authorizedKeysMap[string(authorizedPubKey.Marshal())] = true
@@ -64,15 +60,43 @@ func PublicKeyCallback(c gossh.ConnMetadata,
 			gossh.FingerprintSHA256(authorizedPubKey))
 	}
 
-	dlog.Common.Debug(user, "Offered public key fingerprint",
-		gossh.FingerprintSHA256(offeredPubKey))
+	dlog.Common.Debug(user, "Offered public key fingerprint", gossh.FingerprintSHA256(offeredPubKey))
 	if authorizedKeysMap[string(offeredPubKey.Marshal())] {
 		return &gossh.Permissions{
-			Extensions: map[string]string{
-				"pubkey-fp": gossh.FingerprintSHA256(offeredPubKey),
-			},
+			Extensions: map[string]string{"pubkey-fp": gossh.FingerprintSHA256(offeredPubKey)},
 		}, nil
 	}
 
-	return nil, fmt.Errorf("%s|Public key of user not authorized", user)
+	return nil, fmt.Errorf("%s|public key of user not authorized", user)
+}
+
+func authorizedKeysFile(user *user.User) (string, error) {
+	if config.Env("DTAIL_INTEGRATION_TEST_RUN_MODE") {
+		// In this case, we expect a pub key in the current directory.
+		return "./id_rsa.pub", nil
+	}
+
+	cwd, err := os.Getwd()
+	if err != nil {
+		return "", err
+	}
+
+	// Check for cached version in the dserver directory.
+	authorizedKeysFile := fmt.Sprintf("%s/%s/%s.authorized_keys", cwd,
+		config.Common.CacheDir, user.Name)
+	if _, err = os.Stat(authorizedKeysFile); err == nil {
+		return authorizedKeysFile, nil
+	}
+
+	// As the last option, check the regular SSH path.
+	osUser, err := goUser.Lookup(user.Name)
+	if err != nil {
+		return "", err
+	}
+	authorizedKeysFile = fmt.Sprintf("%s/.ssh/authorized_keys", osUser.HomeDir)
+	if _, err = os.Stat(authorizedKeysFile); err == nil {
+		return authorizedKeysFile, nil
+	}
+
+	return "", fmt.Errorf("unable to find a any authorized keys file")
 }