diff --git a/clients.go b/clients.go index 4912996a..cfa28c93 100644 --- a/clients.go +++ b/clients.go @@ -48,6 +48,32 @@ func (cm *ClientManager) DeleteWhatsmeowClient(userID string) { delete(cm.whatsmeowClients, userID) } +// DeleteSessionIfCurrent drops every per-user entry of a session, but only if +// the whatsmeow client registered for userID is still `client`. It is the +// clientManager counterpart of deleteKillChannel's staleness guard. +// +// A session goroutine passes the client it created at startup. If a newer +// session has replaced the entry in the meantime (a reconnect for the same +// user), the maps hold a different client and this is a no-op. Without the +// guard, a late cleanup from an abandoned session evicts the LIVE session's +// entries: the user stays paired on WhatsApp while wuzapi forgets the client, +// so every subsequent send fails with "no session" until the process restarts. +// +// The three maps are dropped under a single lock so no reader can observe a +// half-torn-down session. +func (cm *ClientManager) DeleteSessionIfCurrent(userID string, client *whatsmeow.Client) bool { + cm.Lock() + defer cm.Unlock() + if cm.whatsmeowClients[userID] != client { + return false + } + delete(cm.whatsmeowClients, userID) + delete(cm.myClients, userID) + delete(cm.pollOptions, userID) + delete(cm.httpClients, userID) + return true +} + func (cm *ClientManager) SetHTTPClient(userID string, client *resty.Client) { cm.Lock() defer cm.Unlock() diff --git a/helpers.go b/helpers.go index c7f87f46..6b645ece 100644 --- a/helpers.go +++ b/helpers.go @@ -274,6 +274,34 @@ func updateUserInfo(values interface{}, field string, value string) interface{} return Values{m: m} } +// setUserInfoField updates ONE field of a token's cached user info, reading the +// current entry immediately before writing it back. Reports whether the entry +// existed. Long-lived goroutines must use this instead of capturing the Values +// once and reusing it: updateUserInfo copies every field from the snapshot it +// is handed, so writing from a stale one silently reverts fields that other +// code changed in the meantime. +// +// The QR loop is the case that motivated it. It captured the user info before +// pairing (when Jid is still empty) and rewrote the cache on every new QR code. +// A code emitted after PairSuccess therefore restored the empty Jid, and the +// next /session/connect found no JID, built a fresh device and asked for a new +// QR — as if the pairing had never happened, while the real device sat intact +// in the store. +// +// This is a read-modify-write and is not atomic: two goroutines updating +// DIFFERENT fields of the same token concurrently can still lose one update. +// That is the pre-existing behaviour of every userinfocache.Set caller; the +// point here is narrowing the window from "the whole life of a goroutine" to +// "two adjacent statements". +func setUserInfoField(token string, field string, value string) bool { + current, found := userinfocache.Get(token) + if !found { + return false + } + userinfocache.Set(token, updateUserInfo(current, field, value), cache.NoExpiration) + return true +} + // webhook for regular messages func callHook(myurl string, payload map[string]string, userID string) { callHookWithHmac(myurl, payload, userID, nil) diff --git a/main.go b/main.go index 0b7711cc..c16b8b3c 100644 --- a/main.go +++ b/main.go @@ -126,6 +126,17 @@ func signalKill(userID string) { log.Debug().Str("userID", userID).Msg("signalKill: no kill channel registered (already cleaned up?)") return } + signalKillChannel(ch) +} + +// signalKillChannel delivers a non-blocking kill to ONE specific channel, +// bypassing the userID lookup. +// +// A session goroutine that wants to end ITSELF must use this and pass its own +// channel. Going through signalKill(userID) would read whatever channel is +// registered NOW, which after a reconnect belongs to a different, live session +// — so an abandoned goroutine would shut down the session that replaced it. +func signalKillChannel(ch chan bool) { select { case ch <- true: default: diff --git a/stale_session_test.go b/stale_session_test.go new file mode 100644 index 00000000..cc10c040 --- /dev/null +++ b/stale_session_test.go @@ -0,0 +1,215 @@ +package main + +import ( + "sync" + "testing" + + "github.com/patrickmn/go-cache" + "go.mau.fi/whatsmeow" +) + +// The scenario behind every test in this file: two /session/connect calls race +// for the same user. Both pass the "already connected" guard (nothing is +// connected yet), so two startClient goroutines run, each with its own device +// and its own QR channel. The user scans one of them. Minutes later the other +// QR expires unscanned and its goroutine cleans up — and that cleanup used to +// address the shared state by bare userID, so it evicted the session that had +// just paired. + +// TestDeleteSessionIfCurrentStaleSession is the clientManager counterpart of +// TestDeleteKillChannelStaleSession: a cleanup from an abandoned session must +// not evict the live session's entries. +func TestDeleteSessionIfCurrentStaleSession(t *testing.T) { + const u = "stale-session-user" + cm := NewClientManager() + + // Two whatsmeow clients, distinct pointers — the two racing sessions. + stale := &whatsmeow.Client{} + live := &whatsmeow.Client{} + + // The stale session registers first; the winner then replaces the entry. + cm.SetWhatsmeowClient(u, stale) + cm.SetWhatsmeowClient(u, live) + cm.SetMyClient(u, &MyClient{userID: u}) + cm.SetPollOptions(u, "msg-1", []string{"yes", "no"}) + + // The abandoned session's QR expires and it cleans up with ITS client. + if cm.DeleteSessionIfCurrent(u, stale) { + t.Error("DeleteSessionIfCurrent reported a delete for a session that is no longer registered") + } + + if got := cm.GetWhatsmeowClient(u); got != live { + t.Fatalf("stale cleanup evicted the live session: GetWhatsmeowClient = %v, want the live client", got) + } + if cm.GetMyClient(u) == nil { + t.Error("stale cleanup removed the live session's MyClient") + } + if opts := cm.GetPollOptions(u, "msg-1"); len(opts) != 2 { + t.Errorf("stale cleanup dropped the live session's poll options: %v", opts) + } +} + +// TestDeleteSessionIfCurrentOwnSession proves the guard does not block the +// legitimate case: the session that IS registered cleans itself up fully. +func TestDeleteSessionIfCurrentOwnSession(t *testing.T) { + const u = "own-session-user" + cm := NewClientManager() + client := &whatsmeow.Client{} + + cm.SetWhatsmeowClient(u, client) + cm.SetMyClient(u, &MyClient{userID: u}) + cm.SetPollOptions(u, "msg-1", []string{"yes"}) + + if !cm.DeleteSessionIfCurrent(u, client) { + t.Fatal("DeleteSessionIfCurrent refused to clean up the session that owns the entry") + } + if cm.GetWhatsmeowClient(u) != nil { + t.Error("whatsmeow client still registered after its own cleanup") + } + if cm.GetMyClient(u) != nil { + t.Error("MyClient still registered after its own cleanup") + } + if opts := cm.GetPollOptions(u, "msg-1"); opts != nil { + t.Errorf("poll options survived the session cleanup: %v", opts) + } +} + +// TestDeleteSessionIfCurrentUnknownUser guards the no-entry path: cleanup after +// the maps were already cleared must be a silent no-op, not a panic. +func TestDeleteSessionIfCurrentUnknownUser(t *testing.T) { + cm := NewClientManager() + if cm.DeleteSessionIfCurrent("never-registered", &whatsmeow.Client{}) { + t.Error("DeleteSessionIfCurrent reported a delete for a user with no entry") + } +} + +// TestDeleteSessionIfCurrentConcurrent hammers the guard from many goroutines. +// The point is the -race build: the three maps are dropped under one lock, so +// no reader can observe a half-torn-down session. +func TestDeleteSessionIfCurrentConcurrent(t *testing.T) { + const u = "concurrent-session-user" + cm := NewClientManager() + live := &whatsmeow.Client{} + cm.SetWhatsmeowClient(u, live) + + var wg sync.WaitGroup + for i := 0; i < 100; i++ { + wg.Add(2) + go func() { + defer wg.Done() + cm.DeleteSessionIfCurrent(u, &whatsmeow.Client{}) // always stale + }() + go func() { + defer wg.Done() + _ = cm.GetWhatsmeowClient(u) + }() + } + wg.Wait() + + if got := cm.GetWhatsmeowClient(u); got != live { + t.Errorf("live client lost under concurrent stale cleanups: got %v", got) + } +} + +// TestSignalKillChannelTargetsOwnChannel proves a goroutine ending itself hits +// its own channel and not whatever session is registered now. signalKill(userID) +// would resolve to the newer session and shut down the wrong one. +func TestSignalKillChannelTargetsOwnChannel(t *testing.T) { + const u = "kill-target-user" + + stale := make(chan bool, 1) + live := make(chan bool, 1) + setKillChannel(u, stale) + setKillChannel(u, live) // a reconnect replaced the entry + defer deleteKillChannel(u, live) + + // The abandoned goroutine ends ITSELF with the channel it captured. + signalKillChannel(stale) + + select { + case <-stale: + default: + t.Error("signalKillChannel did not deliver to the caller's own channel") + } + select { + case <-live: + t.Error("signalKillChannel killed the live session registered for this user") + default: + } +} + +// TestSetUserInfoFieldKeepsConcurrentUpdates is the cache half of the bug. +// The QR loop held a Values captured before pairing (Jid empty) and rewrote the +// cache from it on every new QR code, so a code emitted after PairSuccess +// restored the empty Jid. The next /session/connect then found no JID, built a +// fresh device and asked for a new QR, while the paired device sat intact in +// the store. +func TestSetUserInfoFieldKeepsConcurrentUpdates(t *testing.T) { + const token = "cache-token" + userinfocache.Set(token, Values{m: map[string]string{"Jid": "", "Id": "u1"}}, cache.NoExpiration) + defer userinfocache.Delete(token) + + // PairSuccess writes the JID. + if !setUserInfoField(token, "Jid", "5511999999999:1@s.whatsapp.net") { + t.Fatal("setUserInfoField reported a missing entry for a token that is cached") + } + + // A QR code emitted after pairing updates an unrelated field. + setUserInfoField(token, "Qrcode", "data:image/png;base64,AAAA") + + v, found := userinfocache.Get(token) + if !found { + t.Fatal("user info vanished from the cache") + } + got := v.(Values) + if jid := got.Get("Jid"); jid != "5511999999999:1@s.whatsapp.net" { + t.Errorf("Jid was reverted by a later field update: got %q", jid) + } + if qr := got.Get("Qrcode"); qr != "data:image/png;base64,AAAA" { + t.Errorf("Qrcode not written: got %q", qr) + } + if id := got.Get("Id"); id != "u1" { + t.Errorf("unrelated field lost: Id=%q", id) + } +} + +// TestSetUserInfoFieldMissingEntry documents the return value: no cached entry +// means nothing was written, and callers use that to skip their log line. +func TestSetUserInfoFieldMissingEntry(t *testing.T) { + if setUserInfoField("token-never-cached", "Qrcode", "x") { + t.Error("setUserInfoField reported a write for a token that is not cached") + } +} + +// TestStaleSnapshotRevertsFields pins down WHY setUserInfoField exists, by +// contrasting it with the pattern it replaced. Writing through a Values that +// was captured earlier reverts every field changed since — reading immediately +// before writing does not. Without this contrast the helper looks like a +// pointless wrapper around updateUserInfo. +func TestStaleSnapshotRevertsFields(t *testing.T) { + const token = "snapshot-token" + const jid = "5511999999999:1@s.whatsapp.net" + fresh := func() { userinfocache.Set(token, Values{m: map[string]string{"Jid": ""}}, cache.NoExpiration) } + defer userinfocache.Delete(token) + + // The pattern that caused the bug: capture once, write from it later. + fresh() + snapshot, _ := userinfocache.Get(token) // Jid is still empty here + setUserInfoField(token, "Jid", jid) // pairing writes the JID + userinfocache.Set(token, updateUserInfo(snapshot, "Qrcode", "x"), cache.NoExpiration) + + v, _ := userinfocache.Get(token) + if got := v.(Values).Get("Jid"); got != "" { + t.Fatalf("test is not reproducing the stale-snapshot write: Jid=%q, want it reverted to empty", got) + } + + // Same sequence through the helper: the JID survives. + fresh() + setUserInfoField(token, "Jid", jid) + setUserInfoField(token, "Qrcode", "x") + + v, _ = userinfocache.Get(token) + if got := v.(Values).Get("Jid"); got != jid { + t.Errorf("setUserInfoField reverted the JID: got %q, want %q", got, jid) + } +} diff --git a/wmiau.go b/wmiau.go index 624c7b34..31886ac4 100644 --- a/wmiau.go +++ b/wmiau.go @@ -529,8 +529,6 @@ func (s *server) startClient(userID string, textjid string, token string, kill c return } - myuserinfo, found := userinfocache.Get(token) - for evt := range qrChan { if evt.Event == "code" { // Display QR code in terminal (useful for testing/developing) @@ -546,12 +544,8 @@ func (s *server) startClient(userID string, textjid string, token string, kill c _, err := s.db.Exec(sqlStmt, base64qrcode, userID) if err != nil { log.Error().Err(err).Msg(sqlStmt) - } else { - if found { - v := updateUserInfo(myuserinfo, "Qrcode", base64qrcode) - userinfocache.Set(token, v, cache.NoExpiration) - log.Info().Str("qrcode", base64qrcode).Msg("update cache userinfo with qr code") - } + } else if setUserInfoField(token, "Qrcode", base64qrcode) { + log.Info().Str("qrcode", base64qrcode).Msg("update cache userinfo with qr code") } //send QR code with webhook @@ -575,16 +569,17 @@ func (s *server) startClient(userID string, textjid string, token string, kill c if err != nil { log.Error().Err(err).Msg(sqlStmt) } else { - if found { - v := updateUserInfo(myuserinfo, "Qrcode", "") - userinfocache.Set(token, v, cache.NoExpiration) - } + setUserInfoField(token, "Qrcode", "") } log.Warn().Msg("QR timeout killing channel") - clientManager.DeleteWhatsmeowClient(userID) - clientManager.DeleteMyClient(userID) - clientManager.DeleteHTTPClient(userID) - signalKill(userID) + // Kill OUR OWN channel and let the single cleanup at the + // end of startClient run: it already tears the session down + // under the staleness guard. signalKill(userID) would look + // up whatever channel is registered NOW, which after a + // racing connect belongs to a different, live session — an + // unscanned QR would then shut down the session that won + // the race and paired. + signalKillChannel(kill) } else if evt.Event == "success" { log.Info().Msg("QR pairing ok!") // Clear QR code after pairing @@ -593,10 +588,7 @@ func (s *server) startClient(userID string, textjid string, token string, kill c if err != nil { log.Error().Err(err).Msg(sqlStmt) } else { - if found { - v := updateUserInfo(myuserinfo, "Qrcode", "") - userinfocache.Set(token, v, cache.NoExpiration) - } + setUserInfoField(token, "Qrcode", "") } } else { log.Info().Str("event", evt.Event).Msg("Login event") @@ -675,11 +667,17 @@ func (s *server) startClient(userID string, textjid string, token string, kill c <-kill log.Info().Str("userid", userID).Msg("Received kill signal") client.Disconnect() - clientManager.DeleteWhatsmeowClient(userID) - clientManager.DeleteMyClient(userID) - clientManager.DeleteHTTPClient(userID) - if _, err := s.db.Exec(`UPDATE users SET qrcode='', connected=0 WHERE id=$1`, userID); err != nil { - log.Error().Err(err).Msg("failed to mark user disconnected on kill") + // Only tear down the shared state if WE are still the registered session. + // A reconnect may have replaced us while this goroutine was parked; that + // newer session owns the maps and the users row now, and evicting its + // client (or flipping connected=0 under it) would take down a session that + // is up and paired. Disconnecting our own client above is always safe. + if clientManager.DeleteSessionIfCurrent(userID, client) { + if _, err := s.db.Exec(`UPDATE users SET qrcode='', connected=0 WHERE id=$1`, userID); err != nil { + log.Error().Err(err).Msg("failed to mark user disconnected on kill") + } + } else { + log.Info().Str("userid", userID).Msg("Stale session goroutine exiting; a newer session owns this user") } deleteKillChannel(userID, kill) }