check for ops in Auth.CheckPublicKey and move /allowlist handling to helper functions

auth.go

@@ -53,16 +53,16 @@ func newAuthAddr(addr net.Addr) string {
 // Auth stores lookups for bans, whitelists, and ops. It implements the sshd.Auth interface.
 // If the contained passphrase is not empty, it complements a whitelist.
 type Auth struct {
-	passphraseHash []byte
+	passphraseHash  []byte
 	whitelistModeMu sync.RWMutex
-	whitelistMode  bool
+	whitelistMode   bool
-	bannedAddr     *set.Set
+	bannedAddr      *set.Set
-	bannedClient   *set.Set
+	bannedClient    *set.Set
-	banned         *set.Set
+	banned          *set.Set
-	whitelist      *set.Set
+	whitelist       *set.Set
-	ops            *set.Set
+	ops             *set.Set
-	opFile         string
+	opFile          string
-	whitelistFile  string
+	whitelistFile   string
 }
 
 // NewAuth creates a new empty Auth.
@@ -76,13 +76,13 @@ func NewAuth() *Auth {
 	}
 }
 
-func (a *Auth)WhitelistMode() bool{
+func (a *Auth) WhitelistMode() bool {
 	a.whitelistModeMu.RLock()
 	defer a.whitelistModeMu.RUnlock()
 	return a.whitelistMode
 }
 
-func (a *Auth) SetWhitelistMode(value bool){
+func (a *Auth) SetWhitelistMode(value bool) {
 	a.whitelistModeMu.Lock()
 	defer a.whitelistModeMu.Unlock()
 	a.whitelistMode = value
@@ -135,7 +135,7 @@ func (a *Auth) CheckBans(addr net.Addr, key ssh.PublicKey, clientVersion string)
 func (a *Auth) CheckPublicKey(key ssh.PublicKey) error {
 	authkey := newAuthKey(key)
 	whitelisted := a.whitelist.In(authkey)
-	if a.AllowAnonymous() || whitelisted {
+	if a.AllowAnonymous() || whitelisted || a.IsOp(key) {
 		return nil
 	} else {
 		return ErrNotWhitelisted
@@ -170,9 +170,6 @@ func (a *Auth) Op(key ssh.PublicKey, d time.Duration) {
 
 // IsOp checks if a public key is an op.
 func (a *Auth) IsOp(key ssh.PublicKey) bool {
-	if key == nil {
-		return false
-	}
 	authkey := newAuthKey(key)
 	return a.ops.In(authkey)
 }

host.go

@@ -754,13 +754,104 @@ func (h *Host) InitCommands(c *chat.Commands) {
 		return
 	}
 
+	allowlistHelptext := []string{
+		"Usage: /allowlist help | on | off | add {PUBKEY|USER}... | remove {PUBKEY|USER}... | import [AGE] | reload {keep|flush} | reverify | status",
+		"help: this help message",
+		"on, off: set allowlist mode (applies to new connections)",
+		"add, remove: add or remove keys from the allowlist",
+		"import: add all keys of users connected since AGE (default 0) ago to the allowlist",
+		"reload: re-read the allowlist file and keep or discard entries in the current allowlist but not in the file",
+		"reverify: kick all users not in the allowlist if allowlisting is enabled",
+		"status: show status information",
+	}
+
+	allowlistImport := func(args []string) (msgs []string, err error) {
+		var since time.Duration
+		if len(args) > 0 {
+			since, err = time.ParseDuration(args[0])
+			if err != nil {
+				return
+			}
+		}
+		cutoff := time.Now().Add(-since)
+		noKeyUsers := []string{}
+		forConnectedUsers(func(user *chat.Member, pk ssh.PublicKey) error {
+			if user.Joined().Before(cutoff) {
+				if pk == nil {
+					noKeyUsers = append(noKeyUsers, user.Identifier.Name())
+				} else {
+					h.auth.Whitelist(pk, 0)
+				}
+			}
+			return nil
+		})
+		if len(noKeyUsers) != 0 {
+			msgs = []string{fmt.Sprintf("users without a public key: %v", noKeyUsers)}
+		}
+		return
+	}
+
+	allowlistReload := func(args []string) error {
+		if !(len(args) > 0 && (args[0] == "keep" || args[0] == "flush")) {
+			return errors.New("must specify whether to keep or flush current entries")
+		}
+		if args[0] == "flush" {
+			h.auth.whitelist.Clear()
+		}
+		return h.auth.LoadWhitelistFromFile(h.auth.whitelistFile)
+	}
+
+	allowlistReverify := func() []string {
+		if !h.auth.WhitelistMode() {
+			return []string{"allowlist is disabled, so nobody will be kicked"}
+		}
+		forConnectedUsers(func(user *chat.Member, pk ssh.PublicKey) error {
+			if h.auth.CheckPublicKey(pk) != nil {
+				user.Close() // TODO: some message anywhere?
+			}
+			return nil
+		})
+		return nil
+	}
+
+	allowlistStatus := func() (msgs []string) {
+		if h.auth.WhitelistMode() {
+			msgs = []string{"The allowlist is currently enabled."}
+		} else {
+			msgs = []string{"The allowlist is currently disabled."}
+		}
+		whitelistedUsers := []string{}
+		whitelistedKeys := []string{}
+		// TODO: this can probably be optimized
+		h.auth.whitelist.Each(func(key string, item set.Item) error {
+			keyFP := item.Key()
+			if forConnectedUsers(func(user *chat.Member, pk ssh.PublicKey) error {
+				if pk != nil && sshd.Fingerprint(pk) == keyFP {
+					whitelistedUsers = append(whitelistedUsers, user.Name())
+					return errors.New("not an actual error, but exit early because we found the key")
+				}
+				return nil
+			}) == nil {
+				// if we land here, the key matches no users
+				whitelistedKeys = append(whitelistedKeys, keyFP)
+			}
+			return nil
+		})
+		if len(whitelistedUsers) != 0 {
+			msgs = append(msgs, fmt.Sprintf("The following connected users are on the allowlist: %v", whitelistedUsers))
+		}
+		if len(whitelistedKeys) != 0 {
+			msgs = append(msgs, fmt.Sprintf("The following keys of not connected users are on the allowlist: %v", whitelistedKeys))
+		}
+		return
+	}
+
 	c.Add(chat.Command{
-		// TODO: default for reload
 		Op:         true,
 		Prefix:     "/allowlist",
 		PrefixHelp: "COMMAND [ARGS...]",
 		Help:       "Modify the allowlist or allowlist state. See /allowlist help for subcommands",
-		Handler: func(room *chat.Room, msg message.CommandMsg) error {
+		Handler: func(room *chat.Room, msg message.CommandMsg) (err error) {
 			if !room.IsOp(msg.From()) {
 				return errors.New("must be op")
 			}
@@ -771,115 +862,34 @@ func (h *Host) InitCommands(c *chat.Commands) {
 			}
 
 			// send exactly one message to preserve order
-			replyLines := []string{}
+			var replyLines []string
-			sendMsg := func(content string, formatting ...interface{}) {
-				replyLines = append(replyLines, fmt.Sprintf(content, formatting...))
-			}
 
 			switch args[0] {
 			case "help":
-				sendMsg("Usage: /allowlist help | on | off | add {PUBKEY|USER}... | remove {PUBKEY|USER}... | import [AGE] | reload {keep|flush} | reverify | status")
+				replyLines = allowlistHelptext
-				sendMsg("help: this help message")
-				sendMsg("on, off: set allowlist mode (applies to new connections)")
-				sendMsg("add, remove: add or remove keys from the allowlist")
-				sendMsg("import: add all keys of users connected since AGE (default 0) ago to the allowlist")
-				sendMsg("reload: re-read the allowlist file and keep or discard entries in the current allowlist but not in the file")
-				sendMsg("reverify: kick all users not in the allowlist if allowlisting is enabled")
-				sendMsg("status: show status information")
 			case "on":
 				h.auth.SetWhitelistMode(true)
 			case "off":
 				h.auth.SetWhitelistMode(false)
 			case "add":
-				for _, errLine := range forPubkeyUser(args[1:], func(pk ssh.PublicKey) { h.auth.Whitelist(pk, 0) }) {
+				replyLines = forPubkeyUser(args[1:], func(pk ssh.PublicKey) { h.auth.Whitelist(pk, 0) })
-					sendMsg(errLine)
-				}
 			case "remove":
-				for _, errLine := range forPubkeyUser(args[1:], func(pk ssh.PublicKey) { h.auth.Whitelist(pk, 1) }) {
+				replyLines = forPubkeyUser(args[1:], func(pk ssh.PublicKey) { h.auth.Whitelist(pk, 1) })
-					sendMsg(errLine)
-				}
 			case "import":
-				var since time.Duration
+				replyLines, err = allowlistImport(args[1:])
-				if len(args) > 1 {
-					var err error
-					since, err = time.ParseDuration(args[1])
-					if err != nil {
-						return err
-					}
-				}
-				cutoff := time.Now().Add(-since)
-				noKeyUsers := []string{}
-				forConnectedUsers(func(user *chat.Member, pk ssh.PublicKey) error {
-					if user.Joined().Before(cutoff) {
-						if pk == nil {
-							noKeyUsers = append(noKeyUsers, user.Identifier.Name())
-						} else {
-							h.auth.Whitelist(pk, 0)
-						}
-					}
-					return nil
-				})
-				if len(noKeyUsers) != 0 {
-					sendMsg("users without a public key: %v", noKeyUsers)
-				}
 			case "reload":
-				if !(len(args) > 1 && (args[1] == "keep" || args[1] == "flush")) {
+				err = allowlistReload(args[1:])
-					return errors.New("must specify whether to keep or flush current entries")
-				}
-				if args[1] == "flush" {
-					h.auth.whitelist.Clear()
-				}
-				err := h.auth.LoadWhitelistFromFile(h.auth.whitelistFile)
-				if err != nil {
-					return err
-				}
 			case "reverify":
-				if !h.auth.WhitelistMode() {
+				replyLines = allowlistReverify()
-					sendMsg("allowlist is disabled, so nobody will be kicked")
-					break
-				}
-				forConnectedUsers(func(user *chat.Member, pk ssh.PublicKey) error {
-					if !h.auth.IsOp(pk) && h.auth.CheckPublicKey(pk) != nil { // TODO: why doesn't CheckPublicKey do this?
-						user.Close() // TODO: some message anywhere?
-					}
-					return nil
-				})
 			case "status":
-				if h.auth.WhitelistMode() {
+				replyLines = allowlistStatus()
-					sendMsg("The allowlist is currently enabled.")
-				} else {
-					sendMsg("The allowlist is currently disabled.")
-				}
-				whitelistedUsers := []string{}
-				whitelistedKeys := []string{}
-				// TODO: this can probably be optimized
-				h.auth.whitelist.Each(func(key string, item set.Item) error {
-					keyFP := item.Key()
-					if forConnectedUsers(func(user *chat.Member, pk ssh.PublicKey) error {
-						if pk != nil && sshd.Fingerprint(pk) == keyFP {
-							whitelistedUsers = append(whitelistedUsers, user.Name())
-							return errors.New("not an actual error, but exit early because we found the key")
-						}
-						return nil
-					}) == nil {
-						// if we land here, the key matches no users
-						whitelistedKeys = append(whitelistedKeys, keyFP)
-					}
-					return nil
-				})
-				if len(whitelistedUsers) != 0 {
-					sendMsg("The following connected users are on the allowlist: %v", whitelistedUsers)
-				}
-				if len(whitelistedKeys) != 0 {
-					sendMsg("The following keys of not connected users are on the allowlist: %v", whitelistedKeys)
-				}
 			default:
-				return errors.New("invalid subcommand: " + args[0])
+				err = errors.New("invalid subcommand: " + args[0])
 			}
-			if len(replyLines) != 0 {
+			if err == nil && replyLines != nil {
 				room.Send(message.NewSystemMsg(strings.Join(replyLines, "\r\n"), msg.From()))
 			}
-			return nil
+			return
 		},
 	})
 }