From ae203388e17cd02c04dc429a49f62131243cec9c Mon Sep 17 00:00:00 2001
From: Gregory Anders <greg@gpanders.com>
Date: Sun, 30 Jun 2024 16:31:16 -0500
Subject: [PATCH] Fix channel membership prefixes in cached WHO replies

Channel membership prefixes in WHO replies (RPL_WHOREPLY and
RPL_WHOSPCRPL) were cached in the user's flags, which meant those same
prefixes were returned on future cache hits, even though the flags are
channel specific.

Strip the channel membership prefixes from the user's flags before
adding a user to the cache and add the prefixes back when reading from
the cache (using the membership info from the NAMES reply).
---
 downstream.go | 17 +++++++++++++++++
 irc.go        | 13 +++++++++++++
 upstream.go   |  5 +++--
 3 files changed, 33 insertions(+), 2 deletions(-)

diff --git a/downstream.go b/downstream.go
index f48653e..0c02201 100644
--- a/downstream.go
+++ b/downstream.go
@@ -2230,6 +2230,23 @@ func (dc *downstreamConn) handleMessageRegistered(ctx context.Context, msg *irc.
 				}
 				if uc.isChannel(mask) {
 					info.Channel = mask
+
+					// Set channel membership prefixes from cached NAMES reply
+					ch := uc.channels.Get(info.Channel)
+					memberships := ch.Members.Get(info.Nickname)
+					prefixes := formatMemberPrefix(*memberships, dc)
+
+					// Channel membership prefixes are listed after away status ('G'/'H')
+					// and optional server operator indicator ('*')
+					i := strings.IndexFunc(info.Flags, func(f rune) bool {
+						return f != 'G' && f != 'H' && f != '*'
+					})
+
+					if i == -1 {
+						info.Flags = prefixes + info.Flags
+					} else {
+						info.Flags = info.Flags[:i] + prefixes + info.Flags[i:]
+					}
 				}
 				dc.SendMessage(ctx, xirc.GenerateWHOXReply(fields, &info))
 			}
diff --git a/irc.go b/irc.go
index 12732c4..9a7735b 100644
--- a/irc.go
+++ b/irc.go
@@ -194,6 +194,19 @@ func formatMemberPrefix(ms xirc.MembershipSet, dc *downstreamConn) string {
 	return string(prefixes)
 }
 
+// Remove channel membership prefixes from flags
+func stripMemberPrefixes(flags string, uc *upstreamConn) string {
+	return strings.Map(func(r rune) rune {
+		for _, v := range uc.availableMemberships {
+			if byte(r) == v.Prefix {
+				return -1
+			}
+		}
+
+		return r
+	}, flags)
+}
+
 func parseMessageParams(msg *irc.Message, out ...*string) error {
 	if len(msg.Params) < len(out) {
 		return newNeedMoreParamsError(msg.Command)
diff --git a/upstream.go b/upstream.go
index 9173c3a..025646e 100644
--- a/upstream.go
+++ b/upstream.go
@@ -1541,7 +1541,7 @@ func (uc *upstreamConn) handleMessage(ctx context.Context, msg *irc.Message) err
 				Hostname: host,
 				Server:   server,
 				Nickname: nick,
-				Flags:    flags,
+				Flags:    stripMemberPrefixes(flags, uc),
 				Realname: realname,
 			})
 		}
@@ -1564,13 +1564,14 @@ func (uc *upstreamConn) handleMessage(ctx context.Context, msg *irc.Message) err
 			if err != nil {
 				return err
 			}
+
 			if uc.shouldCacheUserInfo(info.Nickname) {
 				uc.cacheUserInfo(info.Nickname, &upstreamUser{
 					Nickname: info.Nickname,
 					Username: info.Username,
 					Hostname: info.Hostname,
 					Server:   info.Server,
-					Flags:    info.Flags,
+					Flags:    stripMemberPrefixes(info.Flags, uc),
 					Account:  info.Account,
 					Realname: info.Realname,
 				})