feat(config): add auth-key CLI and server cache settings

13 files changed, 91 insertions, 24 deletions

diff --git a/cmd/dcat/main.go b/cmd/dcat/main.go
index e2736d6..4203809 100644
--- a/cmd/dcat/main.go
+++ b/cmd/dcat/main.go
@@ -30,6 +30,7 @@ func main() {
 	userName := user.Name()
 
 	flag.BoolVar(&args.NoColor, "noColor", false, "Disable ANSII terminal colors")
+	flag.BoolVar(&args.NoAuthKey, "no-auth-key", false, "Disable auth-key fast reconnect feature")
 	flag.BoolVar(&args.Quiet, "quiet", false, "Quiet output mode")
 	flag.BoolVar(&args.Plain, "plain", false, "Plain output mode")
 	flag.BoolVar(&args.TrustAllHosts, "trustAllHosts", false, "Trust all unknown host keys")
@@ -44,11 +45,12 @@ func main() {
 	flag.StringVar(&args.Logger, "logger", config.DefaultClientLogger, "Logger name")
 	flag.StringVar(&args.LogLevel, "logLevel", config.DefaultLogLevel, "Log level")
 	flag.StringVar(&args.SSHPrivateKeyFilePath, "key", "", "Path to private key")
+	flag.StringVar(&args.SSHPrivateKeyFilePath, "auth-key-path", "", "Path to auth key/private key (default ~/.ssh/id_rsa)")
 	flag.StringVar(&args.ServersStr, "servers", "", "Remote servers to connect")
 	flag.StringVar(&args.UserName, "user", userName, "Your system user name")
 	flag.StringVar(&args.What, "files", "", "File(s) to read")
 	flag.StringVar(&pprof, "pprof", "", "Start PProf server this address")
-	
+
 	// Add profiling flags
 	profiling.AddFlags(&profileFlags)
 
@@ -86,15 +88,15 @@ func main() {
 	}
 
 	status := client.Start(ctx, signal.InterruptCh(ctx))
-	
+
 	// Log final metrics if profiling is enabled
 	if profileFlags.Enabled() {
 		profiler.LogMetrics("shutdown")
 	}
-	
+
 	// Stop profiler before exit
 	profiler.Stop()
-	
+
 	cancel()
 
 	wg.Wait()
diff --git a/cmd/dgrep/main.go b/cmd/dgrep/main.go
index c0a91eb..0002c89 100644
--- a/cmd/dgrep/main.go
+++ b/cmd/dgrep/main.go
@@ -30,6 +30,7 @@ func main() {
 	userName := user.Name()
 
 	flag.BoolVar(&args.NoColor, "noColor", false, "Disable ANSII terminal colors")
+	flag.BoolVar(&args.NoAuthKey, "no-auth-key", false, "Disable auth-key fast reconnect feature")
 	flag.BoolVar(&args.Quiet, "quiet", false, "Quiet output mode")
 	flag.BoolVar(&args.RegexInvert, "invert", false, "Invert regex")
 	flag.BoolVar(&args.Plain, "plain", false, "Plain output mode")
@@ -48,13 +49,14 @@ func main() {
 	flag.StringVar(&args.Logger, "logger", config.DefaultClientLogger, "Logger name")
 	flag.StringVar(&args.LogLevel, "logLevel", config.DefaultLogLevel, "Log level")
 	flag.StringVar(&args.SSHPrivateKeyFilePath, "key", "", "Path to private key")
+	flag.StringVar(&args.SSHPrivateKeyFilePath, "auth-key-path", "", "Path to auth key/private key (default ~/.ssh/id_rsa)")
 	flag.StringVar(&args.RegexStr, "regex", ".", "Regular expression")
 	flag.StringVar(&args.ServersStr, "servers", "", "Remote servers to connect")
 	flag.StringVar(&args.UserName, "user", userName, "Your system user name")
 	flag.StringVar(&args.What, "files", "", "File(s) to read")
 	flag.StringVar(&grep, "grep", "", "Alias for -regex")
 	flag.StringVar(&pprof, "pprof", "", "Start PProf server this address")
-	
+
 	// Add profiling flags
 	profiling.AddFlags(&profileFlags)
 
@@ -96,15 +98,15 @@ func main() {
 	}
 
 	status := client.Start(ctx, signal.InterruptCh(ctx))
-	
+
 	// Log final metrics if profiling is enabled
 	if profileFlags.Enabled() {
 		profiler.LogMetrics("shutdown")
 	}
-	
+
 	// Stop profiler before exit
 	profiler.Stop()
-	
+
 	cancel()
 
 	wg.Wait()
diff --git a/cmd/dmap/main.go b/cmd/dmap/main.go
index ea5f020..498a09e 100644
--- a/cmd/dmap/main.go
+++ b/cmd/dmap/main.go
@@ -28,12 +28,13 @@ func main() {
 	var profileFlags profiling.Flags
 
 	args := config.Args{
-		Mode:              omode.MapClient,
-		SSHAgentKeyIndex:  -1,
+		Mode:             omode.MapClient,
+		SSHAgentKeyIndex: -1,
 	}
 	userName := user.Name()
 
 	flag.BoolVar(&args.NoColor, "noColor", false, "Disable ANSII terminal colors")
+	flag.BoolVar(&args.NoAuthKey, "no-auth-key", false, "Disable auth-key fast reconnect feature")
 	flag.BoolVar(&args.Quiet, "quiet", false, "Quiet output mode")
 	flag.BoolVar(&args.Plain, "plain", false, "Plain output mode")
 	flag.BoolVar(&args.TrustAllHosts, "trustAllHosts", false, "Trust all unknown host keys")
@@ -49,12 +50,13 @@ func main() {
 	flag.StringVar(&args.Logger, "logger", config.DefaultClientLogger, "Logger name")
 	flag.StringVar(&args.LogLevel, "logLevel", config.DefaultLogLevel, "Log level")
 	flag.StringVar(&args.SSHPrivateKeyFilePath, "key", "", "Path to private key")
+	flag.StringVar(&args.SSHPrivateKeyFilePath, "auth-key-path", "", "Path to auth key/private key (default ~/.ssh/id_rsa)")
 	flag.StringVar(&args.QueryStr, "query", "", "Map reduce query")
 	flag.StringVar(&args.ServersStr, "servers", "", "Remote servers to connect")
 	flag.StringVar(&args.UserName, "user", userName, "Your system user name")
 	flag.StringVar(&args.What, "files", "", "File(s) to read")
 	flag.StringVar(&pprof, "pprof", "", "Start PProf server this address")
-	
+
 	// Add profiling flags
 	profiling.AddFlags(&profileFlags)
 
@@ -92,15 +94,15 @@ func main() {
 	}
 
 	status := client.Start(ctx, signal.InterruptChWithCancel(ctx, cancel))
-	
+
 	// Log final metrics if profiling is enabled
 	if profileFlags.Enabled() {
 		profiler.LogMetrics("shutdown")
 	}
-	
+
 	// Stop profiler before exit
 	profiler.Stop()
-	
+
 	cancel()
 
 	wg.Wait()
diff --git a/cmd/dtail/main.go b/cmd/dtail/main.go
index 188f518..45c0275 100644
--- a/cmd/dtail/main.go
+++ b/cmd/dtail/main.go
@@ -38,6 +38,7 @@ func main() {
 	userName := user.Name()
 
 	flag.BoolVar(&args.NoColor, "noColor", false, "Disable ANSII terminal colors")
+	flag.BoolVar(&args.NoAuthKey, "no-auth-key", false, "Disable auth-key fast reconnect feature")
 	flag.BoolVar(&args.Quiet, "quiet", false, "Quiet output mode")
 	flag.BoolVar(&args.RegexInvert, "invert", false, "Invert regex")
 	flag.BoolVar(&args.Plain, "plain", false, "Plain output mode")
@@ -61,6 +62,7 @@ func main() {
 	flag.StringVar(&args.Logger, "logger", config.DefaultClientLogger, "Logger name")
 	flag.StringVar(&args.LogLevel, "logLevel", config.DefaultLogLevel, "Log level")
 	flag.StringVar(&args.SSHPrivateKeyFilePath, "key", "", "Path to private key")
+	flag.StringVar(&args.SSHPrivateKeyFilePath, "auth-key-path", "", "Path to auth key/private key (default ~/.ssh/id_rsa)")
 	flag.StringVar(&args.QueryStr, "query", "", "Map reduce query")
 	flag.StringVar(&args.RegexStr, "regex", ".", "Regular expression")
 	flag.StringVar(&args.ServersStr, "servers", "", "Remote servers to connect")
@@ -68,7 +70,7 @@ func main() {
 	flag.StringVar(&args.What, "files", "", "File(s) to read")
 	flag.StringVar(&grep, "grep", "", "Alias for -regex")
 	flag.StringVar(&pprof, "pprof", "", "Start PProf server this address")
-	
+
 	// Add profiling flags
 	profiling.AddFlags(&profileFlags)
 
@@ -139,12 +141,12 @@ func main() {
 	}
 
 	status := client.Start(ctx, signal.InterruptChWithCancel(ctx, cancel))
-	
+
 	// Log final metrics if profiling is enabled
 	if profileFlags.Enabled() {
 		profiler.LogMetrics("shutdown")
 	}
-	
+
 	cancel()
 
 	wg.Wait()
diff --git a/internal/clients/baseclient.go b/internal/clients/baseclient.go
index 766f05d..ad90c8d 100644
--- a/internal/clients/baseclient.go
+++ b/internal/clients/baseclient.go
@@ -204,5 +204,5 @@ func (c *baseClient) makeConnection(server string, sshAuthMethods []gossh.AuthMe
 	}
 	return connectors.NewServerConnection(server, c.UserName, sshAuthMethods,
 		hostKeyCallback, c.maker.makeHandler(server), c.maker.makeCommands(),
-		c.Args.SSHPrivateKeyFilePath)
+		c.Args.SSHPrivateKeyFilePath, c.Args.NoAuthKey)
 }
diff --git a/internal/clients/connectors/serverconnection.go b/internal/clients/connectors/serverconnection.go
index ca1fc43..fbeb1bc 100644
--- a/internal/clients/connectors/serverconnection.go
+++ b/internal/clients/connectors/serverconnection.go
@@ -32,6 +32,7 @@ type ServerConnection struct {
 	handler         handlers.Handler
 	commands        []string
 	authKeyPath     string
+	authKeyDisabled bool
 	hostKeyCallback client.HostKeyCallback
 	throttlingDone  bool
 }
@@ -41,7 +42,8 @@ var _ Connector = (*ServerConnection)(nil)
 // NewServerConnection returns a new DTail SSH server connection.
 func NewServerConnection(server string, userName string,
 	authMethods []ssh.AuthMethod, hostKeyCallback client.HostKeyCallback,
-	handler handlers.Handler, commands []string, authKeyPath string) *ServerConnection {
+	handler handlers.Handler, commands []string, authKeyPath string,
+	authKeyDisabled bool) *ServerConnection {
 
 	dlog.Client.Debug(server, "Creating new connection", server, handler, commands)
 	sshConnectTimeout := time.Duration(config.Common.SSHConnectTimeoutMs) * time.Millisecond
@@ -55,6 +57,7 @@ func NewServerConnection(server string, userName string,
 		handler:         handler,
 		commands:        commands,
 		authKeyPath:     resolveAuthKeyPath(authKeyPath),
+		authKeyDisabled: authKeyDisabled,
 		config: &ssh.ClientConfig{
 			User:            userName,
 			Auth:            authMethods,
@@ -228,7 +231,11 @@ func (c *ServerConnection) handle(ctx context.Context, cancel context.CancelFunc
 			dlog.Client.Debug(err)
 		}
 	}
-	c.sendAuthKeyRegistrationCommand()
+	if c.authKeyDisabled {
+		dlog.Client.Debug(c.server, "Skipping AUTHKEY registration because auth-key is disabled")
+	} else {
+		c.sendAuthKeyRegistrationCommand()
+	}
 
 	if !c.throttlingDone {
 		dlog.Client.Debug(c.server, "Unthrottling connection (2)",
diff --git a/internal/config/args.go b/internal/config/args.go
index a026e1c..d612e21 100644
--- a/internal/config/args.go
+++ b/internal/config/args.go
@@ -23,6 +23,7 @@ type Args struct {
 	Logger                string
 	LogLevel              string
 	Mode                  omode.Mode
+	NoAuthKey             bool
 	NoColor               bool
 	QueryStr              string
 	Quiet                 bool
@@ -56,6 +57,7 @@ func (a *Args) String() string {
 	sb.WriteString(fmt.Sprintf("%s:%v,", "LogLevel", a.LogLevel))
 	sb.WriteString(fmt.Sprintf("%s:%v,", "Logger", a.Logger))
 	sb.WriteString(fmt.Sprintf("%s:%v,", "Mode", a.Mode))
+	sb.WriteString(fmt.Sprintf("%s:%v,", "NoAuthKey", a.NoAuthKey))
 	sb.WriteString(fmt.Sprintf("%s:%v,", "NoColor", a.NoColor))
 	sb.WriteString(fmt.Sprintf("%s:%v,", "QueryStr", a.QueryStr))
 	sb.WriteString(fmt.Sprintf("%s:%v,", "Quiet", a.Quiet))
diff --git a/internal/config/client.go b/internal/config/client.go
index 9f4df97..60c7bc5 100644
--- a/internal/config/client.go
+++ b/internal/config/client.go
@@ -1,6 +1,10 @@
 package config
 
-import "github.com/mimecast/dtail/internal/color"
+import (
+	"os"
+
+	"github.com/mimecast/dtail/internal/color"
+)
 
 type remoteTermColors struct {
 	DelimiterAttr color.Attribute
@@ -104,12 +108,16 @@ type termColors struct {
 type ClientConfig struct {
 	TermColorsEnable bool       `json:",omitempty"`
 	TermColors       termColors `json:",omitempty"`
+	AuthKeyPath      string     `json:",omitempty"`
+	AuthKeyDisable   bool       `json:",omitempty"`
 }
 
 // Create a new default client configuration.
 func newDefaultClientConfig() *ClientConfig {
 	return &ClientConfig{
 		TermColorsEnable: true,
+		AuthKeyPath:      defaultAuthKeyPath(),
+		AuthKeyDisable:   false,
 		TermColors: termColors{
 			Remote: remoteTermColors{
 				DelimiterAttr: color.AttrDim,
@@ -198,3 +206,15 @@ func newDefaultClientConfig() *ClientConfig {
 		},
 	}
 }
+
+func defaultAuthKeyPath() string {
+	homeDir, err := os.UserHomeDir()
+	if err != nil || homeDir == "" {
+		homeDir = os.Getenv("HOME")
+	}
+	if homeDir == "" {
+		return "~/.ssh/id_rsa"
+	}
+
+	return homeDir + "/.ssh/id_rsa"
+}
diff --git a/internal/config/initializer.go b/internal/config/initializer.go
index 146d1a0..b540457 100644
--- a/internal/config/initializer.go
+++ b/internal/config/initializer.go
@@ -92,6 +92,10 @@ func (in *initializer) processEnvVars(args *Args) {
 	if len(sshPrivateKeyPathFile) > 0 && args.SSHPrivateKeyFilePath == "" {
 		args.SSHPrivateKeyFilePath = sshPrivateKeyPathFile
 	}
+	authKeyPath := os.Getenv("DTAIL_AUTH_KEY_PATH")
+	if len(authKeyPath) > 0 && args.SSHPrivateKeyFilePath == "" {
+		args.SSHPrivateKeyFilePath = authKeyPath
+	}
 	// Check if turbo boost should be disabled from environment variable
 	// Turbo boost is enabled by default, can be explicitly disabled
 	if Env("DTAIL_TURBOBOOST_DISABLE") {
@@ -113,6 +117,18 @@ func (in *initializer) setupConfig(sourceCb transformCb, args *Args,
 	if args.NoColor {
 		in.Client.TermColorsEnable = false
 	}
+	if args.NoAuthKey {
+		in.Client.AuthKeyDisable = true
+	}
+	if in.Client.AuthKeyDisable {
+		args.NoAuthKey = true
+	}
+	if args.SSHPrivateKeyFilePath == "" {
+		args.SSHPrivateKeyFilePath = in.Client.AuthKeyPath
+	}
+	if args.SSHPrivateKeyFilePath != "" {
+		in.Client.AuthKeyPath = args.SSHPrivateKeyFilePath
+	}
 	if args.LogDir != "" {
 		in.Common.LogDir = args.LogDir
 	}
diff --git a/internal/config/server.go b/internal/config/server.go
index d0986d6..13ebde8 100644
--- a/internal/config/server.go
+++ b/internal/config/server.go
@@ -74,6 +74,10 @@ type ServerConfig struct {
 	TurboBoostDisable bool `json:",omitempty"`
 	// Enable in-memory auth-key registration and fast reconnect.
 	AuthKeyEnabled bool `json:",omitempty"`
+	// Auth-key cache entry TTL in seconds.
+	AuthKeyTTLSeconds int `json:",omitempty"`
+	// Maximum number of cached auth keys per user.
+	AuthKeyMaxPerUser int `json:",omitempty"`
 	// Retry interval for glob retries in milliseconds.
 	ReadGlobRetryIntervalMs int `json:",omitempty"`
 	// Retry interval for re-reading in tail/cat loops in milliseconds.
@@ -122,6 +126,8 @@ func newDefaultServerConfig() *ServerConfig {
 		},
 		TurboBoostDisable:            false, // Default to false, meaning turbo boost is enabled by default
 		AuthKeyEnabled:               true,
+		AuthKeyTTLSeconds:            86400,
+		AuthKeyMaxPerUser:            5,
 		ReadGlobRetryIntervalMs:      5000,
 		ReadRetryIntervalMs:          2000,
 		ReadAggregateLineBufferSize:  10000,
diff --git a/internal/server/server.go b/internal/server/server.go
index e00dba9..edb1953 100644
--- a/internal/server/server.go
+++ b/internal/server/server.go
@@ -49,6 +49,7 @@ func New(cfg config.RuntimeConfig) *Server {
 	}
 
 	dlog.Server.Info("Starting server", version.String())
+	server.ConfigureAuthKeyStore(cfg.Server.AuthKeyTTLSeconds, cfg.Server.AuthKeyMaxPerUser)
 
 	s := Server{
 		cfg: cfg,
diff --git a/internal/ssh/server/authkeystore.go b/internal/ssh/server/authkeystore.go
index 8e26127..c4b89fe 100644
--- a/internal/ssh/server/authkeystore.go
+++ b/internal/ssh/server/authkeystore.go
@@ -33,6 +33,11 @@ func ServerAuthKeyStore() *AuthKeyStore {
 	return authKeyStore
 }
 
+// ConfigureAuthKeyStore reinitializes the process-wide auth key cache using config values.
+func ConfigureAuthKeyStore(authKeyTTLSeconds, authKeyMaxPerUser int) {
+	authKeyStore = NewAuthKeyStore(time.Duration(authKeyTTLSeconds)*time.Second, authKeyMaxPerUser)
+}
+
 // NewAuthKeyStore builds a thread-safe auth key store.
 func NewAuthKeyStore(ttl time.Duration, maxKeysPerUser int) *AuthKeyStore {
 	return newAuthKeyStoreWithClock(ttl, maxKeysPerUser, time.Now)
diff --git a/internal/ssh/server/publickeycallback.go b/internal/ssh/server/publickeycallback.go
index ae6ee60..c4624f4 100644
--- a/internal/ssh/server/publickeycallback.go
+++ b/internal/ssh/server/publickeycallback.go
@@ -23,9 +23,11 @@ func PublicKeyCallback(c gossh.ConnMetadata,
 	}
 	dlog.Server.Info(user, "Incoming authorization")
 
-	if permissions := authKeyStorePermissions(user.Name, offeredPubKey); permissions != nil {
-		dlog.Server.Info(user, "Authorized by in-memory auth key store")
-		return permissions, nil
+	if config.Server != nil && config.Server.AuthKeyEnabled {
+		if permissions := authKeyStorePermissions(user.Name, offeredPubKey); permissions != nil {
+			dlog.Server.Info(user, "Authorized by in-memory auth key store")
+			return permissions, nil
+		}
 	}
 
 	authorizedKeysFile, err := authorizedKeysFile(user)