1 files changed, 127 insertions, 32 deletions
diff --git a/internal/ssh/server/publickeycallback.go b/internal/ssh/server/publickeycallback.go
index bcc9004..8a23384 100644
--- a/internal/ssh/server/publickeycallback.go
+++ b/internal/ssh/server/publickeycallback.go
@@ -1,38 +1,64 @@
 package server
 
 import (
+	"bytes"
+	"errors"
 	"fmt"
+	iofs "io/fs"
 	"os"
 	goUser "os/user"
+	"path/filepath"
 
 	"github.com/mimecast/dtail/internal/config"
 	"github.com/mimecast/dtail/internal/io/dlog"
+	"github.com/mimecast/dtail/internal/io/fs"
 	user "github.com/mimecast/dtail/internal/user/server"
 
 	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) {
+type authorizedKeyParser func([]byte) (gossh.PublicKey, string, []string, []byte, error)
+
+// NewPublicKeyCallback creates an instance-scoped SSH public key callback.
+// keyStore must be non-nil; callers are responsible for constructing and
+// wiring the store. There is no shared package-level fallback.
+func NewPublicKeyCallback(authKeyEnabled bool, cacheDir string,
+	keyStore *AuthKeyStore) func(gossh.ConnMetadata, gossh.PublicKey) (*gossh.Permissions, error) {
 
-	user, err := user.New(c.User(), c.RemoteAddr().String())
+	if keyStore == nil {
+		panic("NewPublicKeyCallback: keyStore must not be nil")
+	}
+	return func(c gossh.ConnMetadata, offeredPubKey gossh.PublicKey) (*gossh.Permissions, error) {
+		return publicKeyCallback(c, offeredPubKey, authKeyEnabled, cacheDir, keyStore)
+	}
+}
+
+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(), nil)
 	if err != nil {
 		return nil, err
 	}
 	dlog.Server.Info(user, "Incoming authorization")
 
-	authorizedKeysFile, err := authorizedKeysFile(user)
+	if authKeyEnabled {
+		if permissions := authKeyStorePermissions(keyStore, user.Name, offeredPubKey); permissions != nil {
+			dlog.Server.Info(user, "Authorized by in-memory auth key store")
+			return permissions, nil
+		}
+	}
+
+	authorizedKeysPath, err := authorizedKeysPathForUser(user, cacheDir)
 	if err != nil {
 		return nil, err
 	}
 
-	dlog.Server.Info(user, "Reading", authorizedKeysFile)
-	authorizedKeysBytes, err := os.ReadFile(authorizedKeysFile)
+	dlog.Server.Info(user, "Reading", authorizedKeysPath.Path())
+	authorizedKeysBytes, err := authorizedKeysPath.ReadFile()
 	if err != nil {
 		return nil, fmt.Errorf("Unable to read authorized keys file|%s|%s|%s",
-			authorizedKeysFile, user, err.Error())
+			authorizedKeysPath.Path(), user, err.Error())
 	}
 
 	return verifyAuthorizedKeys(user, authorizedKeysBytes, offeredPubKey)
@@ -40,57 +66,126 @@ func PublicKeyCallback(c gossh.ConnMetadata,
 
 func verifyAuthorizedKeys(user *user.User, authorizedKeysBytes []byte,
 	offeredPubKey gossh.PublicKey) (*gossh.Permissions, error) {
+	return verifyAuthorizedKeysWithParser(user, authorizedKeysBytes, offeredPubKey, gossh.ParseAuthorizedKey)
+}
+
+func verifyAuthorizedKeysWithParser(user *user.User, authorizedKeysBytes []byte,
+	offeredPubKey gossh.PublicKey, parseAuthorizedKey authorizedKeyParser) (*gossh.Permissions, error) {
 
 	authorizedKeysMap := map[string]bool{}
 	for len(authorizedKeysBytes) > 0 {
-		authorizedPubKey, _, _, restBytes, err := gossh.ParseAuthorizedKey(authorizedKeysBytes)
+		authorizedPubKey, _, _, restBytes, err := parseAuthorizedKey(authorizedKeysBytes)
 		if err != nil {
-			return nil, fmt.Errorf("unable to parse authorized keys bytes|%s|%s",
-				user, err.Error())
+			if dlog.Server != nil {
+				dlog.Server.Warn(user, "Skipping unparseable authorized_keys line", err)
+			}
+			nextAuthorizedKeysBytes, ok := advanceToNextAuthorizedKeysLine(authorizedKeysBytes)
+			if !ok {
+				break
+			}
+			authorizedKeysBytes = nextAuthorizedKeysBytes
+			continue
 		}
 		authorizedKeysMap[string(authorizedPubKey.Marshal())] = true
 		authorizedKeysBytes = restBytes
-		dlog.Server.Debug(user, "Authorized public key fingerprint",
-			gossh.FingerprintSHA256(authorizedPubKey))
+		if dlog.Server != nil {
+			dlog.Server.Debug(user, "Authorized public key fingerprint",
+				gossh.FingerprintSHA256(authorizedPubKey))
+		}
 	}
 
-	dlog.Server.Debug(user, "Offered public key fingerprint", gossh.FingerprintSHA256(offeredPubKey))
+	if dlog.Server != nil {
+		dlog.Server.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)},
-		}, nil
+		return permissionsFromPublicKey(offeredPubKey), nil
 	}
 
 	return nil, fmt.Errorf("%s|public key of user not authorized", user)
 }
 
-func authorizedKeysFile(user *user.User) (string, error) {
+func advanceToNextAuthorizedKeysLine(authorizedKeysBytes []byte) ([]byte, bool) {
+	lineEnd := bytes.IndexByte(authorizedKeysBytes, '\n')
+	if lineEnd == -1 {
+		return nil, false
+	}
+
+	nextBytes := authorizedKeysBytes[lineEnd+1:]
+	return nextBytes, true
+}
+
+func authKeyStorePermissions(keyStore *AuthKeyStore, userName string,
+	offeredPubKey gossh.PublicKey) *gossh.Permissions {
+
+	if keyStore == nil || !keyStore.Has(userName, offeredPubKey) {
+		return nil
+	}
+
+	return permissionsFromPublicKey(offeredPubKey)
+}
+
+func permissionsFromPublicKey(offeredPubKey gossh.PublicKey) *gossh.Permissions {
+	return &gossh.Permissions{
+		Extensions: map[string]string{"pubkey-fp": gossh.FingerprintSHA256(offeredPubKey)},
+	}
+}
+
+type userLookupFunc func(string) (*goUser.User, error)
+
+func authorizedKeysPathForUser(user *user.User, cacheDir string) (fs.RootedPath, 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
+		return fs.NewRootedPath("./id_rsa.pub")
 	}
 
 	cwd, err := os.Getwd()
 	if err != nil {
-		return "", err
+		return fs.RootedPath{}, 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
+	return findAuthorizedKeysPath(user, cacheDir, cwd, goUser.Lookup)
+}
+
+func findAuthorizedKeysPath(user *user.User, cacheDir, cwd string,
+	lookupUser userLookupFunc) (fs.RootedPath, error) {
+
+	// Check for cached version in the dserver directory. An absolute
+	// CacheDir (as used by the BSD packages, e.g. /var/run/dserver/cache)
+	// must be used as-is: joining it with the CWD would break auth
+	// whenever dserver is started from a directory other than / (e.g. a
+	// manual rc.d restart from a home directory). Relative CacheDirs stay
+	// relative to the CWD as before.
+	if cacheDir != "" {
+		cacheBase := cacheDir
+		if !filepath.IsAbs(cacheDir) {
+			cacheBase = filepath.Join(cwd, cacheDir)
+		}
+		cachePath := filepath.Join(cacheBase, fmt.Sprintf("%s.authorized_keys", user.Name))
+		rootedCachePath, err := fs.NewRootedPath(cachePath)
+		if err != nil {
+			return fs.RootedPath{}, err
+		}
+		if _, err := rootedCachePath.Stat(); err == nil {
+			return rootedCachePath, nil
+		}
 	}
 
 	// As the last option, check the regular SSH path.
-	osUser, err := goUser.Lookup(user.Name)
+	osUser, err := lookupUser(user.Name)
+	if err != nil {
+		return fs.RootedPath{}, err
+	}
+	authorizedKeysPath := filepath.Join(osUser.HomeDir, ".ssh", "authorized_keys")
+	rootedAuthorizedKeysPath, err := fs.NewRootedPath(authorizedKeysPath)
 	if err != nil {
-		return "", err
+		return fs.RootedPath{}, err
+	}
+	if _, err = rootedAuthorizedKeysPath.Stat(); err == nil {
+		return rootedAuthorizedKeysPath, nil
 	}
-	authorizedKeysFile = fmt.Sprintf("%s/.ssh/authorized_keys", osUser.HomeDir)
-	if _, err = os.Stat(authorizedKeysFile); err == nil {
-		return authorizedKeysFile, nil
+	if !errors.Is(err, iofs.ErrNotExist) {
+		return fs.RootedPath{}, err
 	}
 
-	return "", fmt.Errorf("unable to find a any authorized keys file")
+	return fs.RootedPath{}, fmt.Errorf("unable to find any authorized keys file")
 }