From 46e2e8a812f9c34840b1c62b830faaad02687236 Mon Sep 17 00:00:00 2001
From: mik2k2 <44849223+mik2k2@users.noreply.github.com>
Date: Wed, 22 Dec 2021 15:20:18 +0100
Subject: [PATCH] remove some unimportant TODOs; add a message when reverify
 kicks people; add a reverify test

---
 auth.go      |  3 ---
 host.go      | 14 +++++++++-----
 host_test.go | 29 +++++++++++++++++++++++++----
 3 files changed, 34 insertions(+), 12 deletions(-)

diff --git a/auth.go b/auth.go
index 64c8e10..8cb1ba2 100644
--- a/auth.go
+++ b/auth.go
@@ -174,9 +174,6 @@ func (a *Auth) IsOp(key ssh.PublicKey) bool {
 	return a.ops.In(authkey)
 }
 
-// TODO: the *FromFile could be replaced by a single LoadFromFile taking the function (i.e. auth.Op/auth.Whitelist)
-// TODO: consider reloading on empty path
-
 // LoadOpsFromFile reads a file in authorized_keys format and makes public keys operators
 func (a *Auth) LoadOpsFromFile(path string) error {
 	a.opFile = path
diff --git a/host.go b/host.go
index ae267ee..fa93861 100644
--- a/host.go
+++ b/host.go
@@ -801,16 +801,21 @@ func (h *Host) InitCommands(c *chat.Commands) {
 		return h.auth.LoadWhitelistFromFile(h.auth.whitelistFile)
 	}
 
-	allowlistReverify := func() []string {
+	allowlistReverify := func(room *chat.Room) []string {
 		if !h.auth.WhitelistMode() {
 			return []string{"allowlist is disabled, so nobody will be kicked"}
 		}
+		var kicked []string
 		forConnectedUsers(func(user *chat.Member, pk ssh.PublicKey) error {
-			if h.auth.CheckPublicKey(pk) != nil {
-				user.Close() // TODO: some message anywhere?
+			if h.auth.CheckPublicKey(pk) != nil && !user.IsOp { // we do this check here as well for ops without keys
+				kicked = append(kicked, user.Name())
+				user.Close()
 			}
 			return nil
 		})
+		if kicked != nil {
+			room.Send(message.NewAnnounceMsg(fmt.Sprintf("%v were kicked during pubkey reverification.", kicked)))
+		}
 		return nil
 	}
 
@@ -822,7 +827,6 @@ func (h *Host) InitCommands(c *chat.Commands) {
 		}
 		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 {
@@ -880,7 +884,7 @@ func (h *Host) InitCommands(c *chat.Commands) {
 			case "reload":
 				err = allowlistReload(args[1:])
 			case "reverify":
-				replyLines = allowlistReverify()
+				replyLines = allowlistReverify(room)
 			case "status":
 				replyLines = allowlistStatus()
 			default:
diff --git a/host_test.go b/host_test.go
index bad1e82..6494332 100644
--- a/host_test.go
+++ b/host_test.go
@@ -211,7 +211,18 @@ func TestHostAllowlistCommand(t *testing.T) {
 		users <- u
 	}
 
+	kickSignal := make(chan struct{})
+	go sshd.ConnectShell(s.Addr().String(), "toBeKicked", func(r io.Reader, w io.WriteCloser) error {
+		<-kickSignal
+		n, err := w.Write([]byte("alive and well"))
+		if n != 0 || err == nil {
+			t.Error("could write after being kicked")
+		}
+		return nil
+	})
+
 	sshd.ConnectShell(s.Addr().String(), "foo", func(r io.Reader, w io.WriteCloser) error {
+		<-users
 		<-users
 		m, ok := host.MemberByID("foo")
 		if !ok {
@@ -220,6 +231,7 @@ func TestHostAllowlistCommand(t *testing.T) {
 
 		scanner := bufio.NewScanner(r)
 		scanner.Scan() // Joined
+		scanner.Scan()
 
 		assertLineEq := func(expected string) {
 			if !scanner.Scan() {
@@ -281,9 +293,13 @@ func TestHostAllowlistCommand(t *testing.T) {
 
 		// TODO: to test the AGE arg, we need another connection and possibly a sleep
 		sendCmd("/allowlist import")
-		assertLineEq("users without a public key: [foo]\r")
+		if !scanner.Scan() {
+			t.Error("no line available")
+		}
+		if actual := stripPrompt(scanner.Text()); !strings.HasSuffix(actual, "[foo toBeKicked]\r") && !strings.HasSuffix(actual, "[toBeKicked foo]\r") {
+			t.Errorf("expected \"foo\" and \"toBeKicked\", got %q", actual)
+		}
 
-		// TODO: test reload with files?
 		sendCmd("/allowlist reload keep")
 		if !host.auth.whitelist.In(testKey2FP) {
 			t.Error("cleared allowlist to be kept")
@@ -299,10 +315,15 @@ func TestHostAllowlistCommand(t *testing.T) {
 
 		sendCmd("/allowlist reverify")
 		assertLineEq("allowlist is disabled, so nobody will be kicked\r")
+		sendCmd("/allowlist on")
+		sendCmd("/allowlist reverify")
+		kickSignal <- struct{}{}
+		assertLineEq(" * [toBeKicked] were kicked during pubkey reverification.\r")
+		assertLineEq(" * toBeKicked left. (After 0 seconds)\r")
 
 		sendCmd("/allowlist add " + testKey1)
 		sendCmd("/allowlist status")
-		assertLineEq("The allowlist is currently disabled.\r")
+		assertLineEq("The allowlist is currently enabled.\r")
 		assertLineEq(fmt.Sprintf("The following keys of not connected users are on the allowlist: [%s]\r", testKey1FP))
 
 		sendCmd("/allowlist invalidSubcommand")
@@ -451,5 +472,5 @@ func connectUserWithConfig(t *testing.T, name string, envConfig map[string]strin
 		}
 	}
 	t.Fatalf("user %s not found in the host", name)
-    return nil
+	return nil
 }