1 files changed, 17 insertions, 12 deletions
diff --git a/internal/net/server.go b/internal/net/server.go
index 6367419..71e34c8 100644
--- a/internal/net/server.go
+++ b/internal/net/server.go
@@ -331,17 +331,11 @@ func sessionIdleTimeout(sess *Session) time.Duration {
 	}
 }
 
-// armReadDeadline sets an idle read deadline on the session's transport when
-// it supports one. tcpConn and wsConn (both in this package) expose their
-// underlying connection, so telnet, telnet-TLS, and websocket are all covered.
+// armReadDeadline sets an idle read deadline on the session's transport. The
+// Conn interface requires SetReadDeadline, so every transport (telnet,
+// telnet-TLS, websocket) is covered and a new one cannot silently miss it.
 func armReadDeadline(sess *Session, d time.Duration) {
-	deadline := time.Now().Add(d)
-	switch c := sess.Conn.(type) {
-	case *tcpConn:
-		_ = c.conn.SetReadDeadline(deadline)
-	case *wsConn:
-		_ = c.conn.SetReadDeadline(deadline)
-	}
+	_ = sess.Conn.SetReadDeadline(time.Now().Add(d))
 }
 
 // StartWriter puts the session into queued mode: writes are enqueued and a
@@ -559,8 +553,11 @@ func (h *Hub) LeaveRoom(s *Session) {
 }
 
 func (h *Hub) leaveRoomLocked(s *Session) {
-	for _, room := range h.rooms {
+	for id, room := range h.rooms {
 		delete(room, s)
+		if len(room) == 0 {
+			delete(h.rooms, id)
+		}
 	}
 }
 
@@ -582,7 +579,15 @@ func (h *Hub) RekeyRooms(idMap map[int]int) {
 		}
 	}
 	for newID, set := range moved {
-		h.rooms[newID] = set
+		if existing := h.rooms[newID]; existing != nil {
+			// Merge instead of clobbering a target bucket that exists but is
+			// not itself in idMap; assigning would drop its sessions.
+			for s := range set {
+				existing[s] = true
+			}
+		} else {
+			h.rooms[newID] = set
+		}
 	}
 }