2 files changed, 14 insertions, 17 deletions
diff --git a/internal/ssh/server/hostkey.go b/internal/ssh/server/hostkey.go
index b2d4569..809a870 100644
--- a/internal/ssh/server/hostkey.go
+++ b/internal/ssh/server/hostkey.go
@@ -8,9 +8,19 @@ import (
 	"github.com/mimecast/dtail/internal/ssh"
 )
 
+const (
+	defaultHostKeyBits = 4096
+	defaultHostKeyFile = "./cache/ssh_host_key"
+)
+
 // PrivateHostKey retrieves the private server RSA host key.
-func PrivateHostKey() []byte {
-	hostKeyFile := config.Server.HostKeyFile
+func PrivateHostKey(hostKeyFile string, hostKeyBits int) []byte {
+	if hostKeyFile == "" {
+		hostKeyFile = defaultHostKeyFile
+	}
+	if hostKeyBits <= 0 {
+		hostKeyBits = defaultHostKeyBits
+	}
 	if config.Env("DTAIL_INTEGRATION_TEST_RUN_MODE") {
 		hostKeyFile = "./ssh_host_key"
 	}
@@ -18,7 +28,7 @@ func PrivateHostKey() []byte {
 
 	if os.IsNotExist(err) {
 		dlog.Server.Info("Generating private server RSA host key")
-		privateKey, err := ssh.GeneratePrivateRSAKey(config.Server.HostKeyBits)
+		privateKey, err := ssh.GeneratePrivateRSAKey(hostKeyBits)
 
 		if err != nil {
 			dlog.Server.FatalPanic("Failed to generate private server RSA host key", err)
diff --git a/internal/ssh/server/publickeycallback.go b/internal/ssh/server/publickeycallback.go
index ccf9111..d4e328b 100644
--- a/internal/ssh/server/publickeycallback.go
+++ b/internal/ssh/server/publickeycallback.go
@@ -12,19 +12,6 @@ import (
 	gossh "golang.org/x/crypto/ssh"
 )
 
-// PublicKeyCallback is for the server to check whether a public SSH key is
-// authorized ot not.
-func PublicKeyCallback(c gossh.ConnMetadata,
-	offeredPubKey gossh.PublicKey) (*gossh.Permissions, error) {
-
-	authKeyEnabled := config.Server != nil && config.Server.AuthKeyEnabled
-	cacheDir := ""
-	if config.Common != nil {
-		cacheDir = config.Common.CacheDir
-	}
-	return publicKeyCallback(c, offeredPubKey, authKeyEnabled, cacheDir, authKeyStore)
-}
-
 // NewPublicKeyCallback creates an instance-scoped SSH public key callback.
 // It avoids relying on package-level mutable configuration/state.
 func NewPublicKeyCallback(authKeyEnabled bool, cacheDir string,
@@ -41,7 +28,7 @@ func NewPublicKeyCallback(authKeyEnabled bool, cacheDir string,
 func publicKeyCallback(c gossh.ConnMetadata, offeredPubKey gossh.PublicKey,
 	authKeyEnabled bool, cacheDir string, keyStore *AuthKeyStore) (*gossh.Permissions, error) {
 
-	user, err := user.New(c.User(), c.RemoteAddr().String())
+	user, err := user.New(c.User(), c.RemoteAddr().String(), nil)
 	if err != nil {
 		return nil, err
 	}