Dont use relaxed SSH Auth mode anymore for integration tests

9 files changed, 71 insertions, 53 deletions

diff --git a/integrationtests/dcat_test.go b/integrationtests/dcat_test.go
index 777e835..6928afa 100644
--- a/integrationtests/dcat_test.go
+++ b/integrationtests/dcat_test.go
@@ -9,7 +9,7 @@ import (
 )
 
 func TestDCat(t *testing.T) {
-	if !config.Env("DTAIL_RUN_INTEGRATION_TESTS") {
+	if !config.Env("DTAIL_INTEGRATION_TEST_RUN_MODE") {
 		t.Log("Skipping")
 		return
 	}
@@ -33,7 +33,7 @@ func TestDCat(t *testing.T) {
 }
 
 func TestDCat2(t *testing.T) {
-	if !config.Env("DTAIL_RUN_INTEGRATION_TESTS") {
+	if !config.Env("DTAIL_INTEGRATION_TEST_RUN_MODE") {
 		return
 	}
 	testdataFile := "dcat2.txt"
@@ -62,7 +62,7 @@ func TestDCat2(t *testing.T) {
 }
 
 func TestDCatColors(t *testing.T) {
-	if !config.Env("DTAIL_RUN_INTEGRATION_TESTS") {
+	if !config.Env("DTAIL_INTEGRATION_TEST_RUN_MODE") {
 		return
 	}
 
diff --git a/integrationtests/dgrep_test.go b/integrationtests/dgrep_test.go
index 26abc2f..35c3ff5 100644
--- a/integrationtests/dgrep_test.go
+++ b/integrationtests/dgrep_test.go
@@ -9,7 +9,7 @@ import (
 )
 
 func TestDGrep(t *testing.T) {
-	if !config.Env("DTAIL_RUN_INTEGRATION_TESTS") {
+	if !config.Env("DTAIL_INTEGRATION_TEST_RUN_MODE") {
 		t.Log("Skipping")
 		return
 	}
@@ -38,7 +38,7 @@ func TestDGrep(t *testing.T) {
 }
 
 func TestDGrep2(t *testing.T) {
-	if !config.Env("DTAIL_RUN_INTEGRATION_TESTS") {
+	if !config.Env("DTAIL_INTEGRATION_TEST_RUN_MODE") {
 		t.Log("Skipping")
 		return
 	}
@@ -68,7 +68,7 @@ func TestDGrep2(t *testing.T) {
 }
 
 func TestDGrepContext(t *testing.T) {
-	if !config.Env("DTAIL_RUN_INTEGRATION_TESTS") {
+	if !config.Env("DTAIL_INTEGRATION_TEST_RUN_MODE") {
 		t.Log("Skipping")
 		return
 	}
@@ -98,7 +98,7 @@ func TestDGrepContext(t *testing.T) {
 }
 
 func TestDGrepContext2(t *testing.T) {
-	if !config.Env("DTAIL_RUN_INTEGRATION_TESTS") {
+	if !config.Env("DTAIL_INTEGRATION_TEST_RUN_MODE") {
 		t.Log("Skipping")
 		return
 	}
diff --git a/integrationtests/dmap_test.go b/integrationtests/dmap_test.go
index 53b8574..6a93b7b 100644
--- a/integrationtests/dmap_test.go
+++ b/integrationtests/dmap_test.go
@@ -10,7 +10,7 @@ import (
 )
 
 func TestDMap(t *testing.T) {
-	if !config.Env("DTAIL_RUN_INTEGRATION_TESTS") {
+	if !config.Env("DTAIL_INTEGRATION_TEST_RUN_MODE") {
 		t.Log("Skipping")
 		return
 	}
@@ -56,7 +56,7 @@ func TestDMap(t *testing.T) {
 }
 
 func TestDMap2(t *testing.T) {
-	if !config.Env("DTAIL_RUN_INTEGRATION_TESTS") {
+	if !config.Env("DTAIL_INTEGRATION_TEST_RUN_MODE") {
 		t.Log("Skipping")
 		return
 	}
@@ -93,7 +93,7 @@ func TestDMap2(t *testing.T) {
 }
 
 func TestDMap3(t *testing.T) {
-	if !config.Env("DTAIL_RUN_INTEGRATION_TESTS") {
+	if !config.Env("DTAIL_INTEGRATION_TEST_RUN_MODE") {
 		t.Log("Skipping")
 		return
 	}
diff --git a/integrationtests/dtail_test.go b/integrationtests/dtail_test.go
index e9cf257..6fa5308 100644
--- a/integrationtests/dtail_test.go
+++ b/integrationtests/dtail_test.go
@@ -12,7 +12,7 @@ import (
 )
 
 func TestDTailWithServer(t *testing.T) {
-	if !config.Env("DTAIL_RUN_INTEGRATION_TESTS") {
+	if !config.Env("DTAIL_INTEGRATION_TEST_RUN_MODE") {
 		t.Log("Skipping")
 		return
 	}
@@ -41,7 +41,6 @@ func TestDTailWithServer(t *testing.T) {
 		"--logLevel", "info",
 		"--bindAddress", bindAddress,
 		"--port", fmt.Sprintf("%d", port),
-		"--relaxedAuth",
 	)
 	if err != nil {
 		t.Error(err)
@@ -131,7 +130,7 @@ func TestDTailWithServer(t *testing.T) {
 }
 
 func TestDTailColorTable(t *testing.T) {
-	if !config.Env("DTAIL_RUN_INTEGRATION_TESTS") {
+	if !config.Env("DTAIL_INTEGRATION_TEST_RUN_MODE") {
 		t.Log("Skipping")
 		return
 	}
diff --git a/integrationtests/dtailhealth_test.go b/integrationtests/dtailhealth_test.go
index 271f11d..b53c425 100644
--- a/integrationtests/dtailhealth_test.go
+++ b/integrationtests/dtailhealth_test.go
@@ -10,7 +10,7 @@ import (
 )
 
 func TestDTailHealthCheck(t *testing.T) {
-	if !config.Env("DTAIL_RUN_INTEGRATION_TESTS") {
+	if !config.Env("DTAIL_INTEGRATION_TEST_RUN_MODE") {
 		t.Log("Skipping")
 		return
 	}
@@ -32,7 +32,7 @@ func TestDTailHealthCheck(t *testing.T) {
 }
 
 func TestDTailHealthCheck2(t *testing.T) {
-	if !config.Env("DTAIL_RUN_INTEGRATION_TESTS") {
+	if !config.Env("DTAIL_INTEGRATION_TEST_RUN_MODE") {
 		t.Log("Skipping")
 		return
 	}
@@ -57,7 +57,7 @@ func TestDTailHealthCheck2(t *testing.T) {
 }
 
 func TestDTailHealthCheck3(t *testing.T) {
-	if !config.Env("DTAIL_RUN_INTEGRATION_TESTS") {
+	if !config.Env("DTAIL_INTEGRATION_TEST_RUN_MODE") {
 		t.Log("Skipping")
 		return
 	}
diff --git a/internal/config/env.go b/internal/config/env.go
index 804a10a..1ccac9c 100644
--- a/internal/config/env.go
+++ b/internal/config/env.go
@@ -16,12 +16,3 @@ func Hostname() (string, error) {
 	}
 	return os.Hostname()
 }
-
-// SSHKnownHostsFile returns the known hosts file path (useful for integration tests)
-func SSHKnownHostsFile() string {
-	if len(os.Getenv("DTAIL_SSH_KNOWN_HOSTS_FILE")) > 0 {
-		return os.Getenv("DTAIL_SSH_KNOWN_HOSTS_FILE")
-	} else {
-		return os.Getenv("HOME") + "/.ssh/known_hosts"
-	}
-}
diff --git a/internal/config/initializer.go b/internal/config/initializer.go
index 024464e..137b831 100644
--- a/internal/config/initializer.go
+++ b/internal/config/initializer.go
@@ -65,7 +65,7 @@ func (in *initializer) parseSpecificConfig(configFile string) error {
 func (in *initializer) transformConfig(sourceProcess source.Source, args *Args,
 	additionalArgs []string) error {
 
-	in.readEnvironmentVars()
+	in.processEnvVars()
 
 	switch sourceProcess {
 	case source.Server:
@@ -81,10 +81,9 @@ func (in *initializer) transformConfig(sourceProcess source.Source, args *Args,
 }
 
 // There are some special options which can be set by environment variable.
-func (in *initializer) readEnvironmentVars() {
-	if Env("DTAIL_RUN_INTEGRATION_TESTS") {
+func (in *initializer) processEnvVars() {
+	if Env("DTAIL_INTEGRATION_TEST_RUN_MODE") {
 		os.Setenv("DTAIL_HOSTNAME_OVERRIDE", "integrationtest")
-		os.Setenv("DTAIL_SSH_KNOWN_HOSTS_FILE", "./known_hosts")
 	}
 }
 
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")
 }