Receive images in WhatsApp batches

[onechat] / whatsapp_inbox.go

package main

import (
        "context"
        "crypto/sha256"
        "fmt"
        "log/slog"
        "slices"
        "strings"
        "sync"
        "time"

        "go.mau.fi/whatsmeow/store"
        "go.mau.fi/whatsmeow/types"
        "go.mau.fi/whatsmeow/types/events"
)

const (
        // whatsappSource is the source of the work events of the batches.
        whatsappSource = "whatsapp"

        // whatsappCatchupLimit is the longest wait for the end of the catch-up
        // after a connection.
        whatsappCatchupLimit = time.Minute
)

// newWhatsAppInbox returns the inbox of the batches in state, with clock for
// the times of the batch texts and groups for the names of the groups. images
// holds shared images, and logger gets image failures.
func newWhatsAppInbox(state *WhatsAppState, clock func() time.Time, groups *WhatsAppGroups, images *ImageStore, logger *slog.Logger) *WhatsAppInbox {
        return &WhatsAppInbox{
                state:        state,
                clock:        clock,
                groups:       groups,
                images:       images,
                logger:       logger,
                lookupTime:   whatsappLookupTime,
                imageTime:    whatsappImageTime,
                catchupLimit: whatsappCatchupLimit,
                wake:         make(chan struct{}, 1),
                failures:     make(chan error, 1),
        }
}

// WhatsAppInbox puts messages into batches, one chat at a time, as
// WhatsAppState says. deliver logs each closed batch as a work message that
// starts with [WhatsApp], and a batch that needs the user wakes the turn loop.
//
// After each connection, the servers send the messages that came while the
// daemon was away, and then OfflineSyncCompleted. Until then, or until
// catchupLimit passes, the closed batches wait in the state file. At the end of
// the catch-up, deliver closes all batches and logs them. It wakes the root
// once for its batches, so their catch-up gives one turn, not one for each
// chat. The batches of the state file at the start wait for the first catch-up
// in the same way. The event handler acknowledges a message only after the
// state file holds it, so the servers send again a message that a crash lost. A
// batch leaves the state file only after its work message is durable in the log
// or agent queue. A crash between the two can repeat an unrouted batch once. A
// routed batch enters the durable agent queue, which prevents a second log
// entry and wakes only the child. A missing child uses the root.
//
// The state file is the one of WhatsApp, which keeps the send records in it
// too. groups gives the names of the groups. The mutex protects the state of
// the catch-up, which the event handler of the library changes from its own
// goroutines. wake tells Run that a batch can be due or that a catch-up ended,
// and failures gives Run the error of a write of the state file.
type WhatsAppInbox struct {
        state        *WhatsAppState
        clock        func() time.Time
        groups       *WhatsAppGroups
        images       *ImageStore
        logger       *slog.Logger
        lookupTime   time.Duration
        imageTime    time.Duration
        catchupLimit time.Duration
        wake         chan struct{}
        failures     chan error

        mutex        sync.Mutex
        catchingUp   bool
        caughtUp     bool
        catchups     int
        catchupTimer *time.Timer
}

// due returns a channel that gets a value when the next open batch is due, or
// nil when no batch is open.
func (i *WhatsAppInbox) due() <-chan time.Time {
        wait, open := i.state.Wait(time.Now())
        if !open {
                return nil
        }

        return time.After(wait)
}

// deliver logs the closed batches in l in order, each as a work message of
// whatsappSource, and wakes the loop once with the batches that start a turn. A
// batch leaves the state file after its message is durable in the log. A group
// with no name gets its name from WhatsApp first, with the client of link.
// During a catch-up it logs nothing, and at its end it closes all batches
// first. With all, it closes and logs all batches also during a catch-up.
// Routed batches enter the durable child queue instead of the root log and wake
// their child even when the batch does not need a root turn.
func (i *WhatsAppInbox) deliver(ctx context.Context, l *TurnLoop, link WhatsAppLink, all bool) error {
        i.mutex.Lock()

        catchingUp := i.catchingUp
        caughtUp := i.caughtUp
        i.caughtUp = false
        i.mutex.Unlock()

        var err error

        switch {
        case all || caughtUp:
                err = i.state.CloseAll()
        case catchingUp:
                return nil
        default:
                _, err = i.state.CloseDue(time.Now())
        }

        if err != nil {
                return err
        }

        batches := i.state.Closed()
        if l.messages != nil {
                err = l.messages.PruneWhatsApp(batches)
                if err != nil {
                        return fmt.Errorf("prune WhatsApp receipts: %w", err)
                }
        }

        var turns []Message

        for _, batch := range batches {
                if batch.Chat.Group && batch.Chat.Name == "" {
                        batch.Chat.Name = i.groups.name(ctx, link, batch.Chat.ID)
                }

                text := batch.Text(i.clock().Location())
                if l.messages != nil {
                        routed, err := l.messages.RouteWhatsApp(batch, text)
                        if err != nil {
                                return fmt.Errorf("route a WhatsApp batch: %w", err)
                        }

                        if routed {
                                err = i.state.Remove(batch.Number)
                                if err != nil {
                                        return fmt.Errorf("remove routed WhatsApp batch: %w", err)
                                }

                                continue
                        }
                }

                message, err := l.Work(whatsappSource, text, "")
                if err != nil {
                        return fmt.Errorf("log a WhatsApp batch: %w", err)
                }

                // A stop of the daemon during the flush is no failure: the
                // batch stays in the state file and repeats once at the next
                // start.
                err = l.Flush(ctx)
                if ctx.Err() != nil {
                        return nil
                }

                if err != nil {
                        return fmt.Errorf("log a WhatsApp batch: %w", err)
                }

                err = i.state.Remove(batch.Number)
                if err != nil {
                        return err
                }

                if batch.Turn() {
                        turns = append(turns, message)
                }
        }

        l.Wake(turns...)

        if l.messages != nil {
                err = l.messages.PruneWhatsApp(i.state.Closed())
                if err != nil {
                        return fmt.Errorf("prune WhatsApp receipts: %w", err)
                }
        }

        return nil
}

// waitForCatchup makes the batches wait for a catch-up, as for the batches of
// the state file at the start, before the first connection.
func (i *WhatsAppInbox) waitForCatchup() {
        i.mutex.Lock()
        defer i.mutex.Unlock()

        i.catchingUp = true
}

// startCatchup starts the catch-up of a new connection, which ends at
// endCatchup or after catchupLimit.
func (i *WhatsAppInbox) startCatchup() {
        i.mutex.Lock()
        defer i.mutex.Unlock()

        i.catchingUp = true
        i.catchups++
        catchups := i.catchups
        if i.catchupTimer != nil {
                i.catchupTimer.Stop()
        }

        i.catchupTimer = time.AfterFunc(i.catchupLimit, func() {
                i.caughtUpBy(catchups)
        })
}

// receive stores an image before its batch and wakes Run. link gives the device
// and the own IDs. It returns false when the state file cannot hold the
// message, so that the library does not acknowledge it, and Run stops the
// daemon. An image that cannot be downloaded or stored gets a text marker. A
// message that the batches do not take is handled with nothing to do.
func (i *WhatsAppInbox) receive(link WhatsAppLink, event *events.Message) bool {
        ctx, cancel := context.WithTimeout(context.Background(), i.lookupTime)
        defer cancel()

        chat, message, found := whatsappMessage(ctx, link.device, link.own, event)
        if !found {
                return true
        }

        if event.Message.GetImageMessage() != nil {
                ctx, cancel := context.WithTimeout(context.Background(), i.imageTime)
                defer cancel()

                message.Text = i.imageText(ctx, link, event, message.Text)
        }

        _, err := i.state.Add(chat, message, time.Now())
        if err != nil {
                select {
                case i.failures <- err:
                default:
                }

                return false
        }

        select {
        case i.wake <- struct{}{}:
        default:
        }

        return true
}

// imageText stores the image of event before returning its marker and caption.
// Failures keep the caption with a text marker. The daemon log identifies the
// chat by a short hash, not its phone number.
func (i *WhatsAppInbox) imageText(ctx context.Context, link WhatsAppLink, event *events.Message, caption string) string {
        name, err := whatsappImage(ctx, link.client, i.images, event.Message.GetImageMessage())
        if err != nil {
                chat := event.Info.Chat.ToNonAD()
                digest := sha256.Sum256([]byte(chat.String()))
                redact := append(slices.Clone(link.own), chat.User)
                i.logger.Error("image receive failed", "chat", fmt.Sprintf("%x", digest[:6]), "err", whatsappRedact(redact, err.Error()))

                text := "(image not downloaded)"
                if caption != "" {
                        text += "\n" + caption
                }

                return text
        }

        return imageText([]string{name}, caption)
}

// caughtUpBy ends the catch-up of the connection with the number catchups, when
// no later connection started a new one.
func (i *WhatsAppInbox) caughtUpBy(catchups int) {
        i.mutex.Lock()
        defer i.mutex.Unlock()

        if i.catchups == catchups {
                i.finishCatchup()
        }
}

// finishCatchup ends a running catch-up and wakes Run, which logs all batches.
// The caller holds the mutex.
func (i *WhatsAppInbox) finishCatchup() {
        if !i.catchingUp {
                return
        }

        i.catchingUp = false
        i.caughtUp = true
        if i.catchupTimer != nil {
                i.catchupTimer.Stop()
        }

        select {
        case i.wake <- struct{}{}:
        default:
        }
}

// endCatchup ends a running catch-up at OfflineSyncCompleted.
func (i *WhatsAppInbox) endCatchup() {
        i.mutex.Lock()
        defer i.mutex.Unlock()

        i.finishCatchup()
}

// stopCatchup ends a running catch-up with no wake, after the session is
// removed.
func (i *WhatsAppInbox) stopCatchup() {
        i.mutex.Lock()
        defer i.mutex.Unlock()

        i.catchingUp = false
}

// close stops the timer of the catch-up.
func (i *WhatsAppInbox) close() {
        i.mutex.Lock()
        defer i.mutex.Unlock()

        if i.catchupTimer != nil {
                i.catchupTimer.Stop()
        }
}

// whatsappMessage returns the chat and the batch message of event, or false
// when the batches do not take it. device is the linked device, and own holds
// the own IDs. The batches take text, images, and media markers from a private
// chat or a group, from another person or from the user. They do not take
// reactions, edits, deletions, view-once and disappearing messages, old
// history, status broadcasts, channels, and the chat of the user with itself.
//
// A mention of the user becomes @you, another mention @NAME when the name is
// known, and the own number and LID become OWN, so that they never go to the
// log.
func whatsappMessage(ctx context.Context, device *store.Device, own []string, event *events.Message) (WhatsAppChat, WhatsAppMessage, bool) {
        info := event.Info
        chatJID := info.Chat.ToNonAD()
        switch {
        case device == nil:
        case event.IsEdit, info.Edit != "", event.IsEphemeral, event.IsViewOnce, event.Message.GetImageMessage().GetViewOnce(), event.Message.GetVideoMessage().GetViewOnce(), event.SourceWebMsg != nil:
        case chatJID.Server != types.DefaultUserServer && chatJID.Server != types.HiddenUserServer && chatJID.Server != types.GroupServer:
        case slices.Contains(own, chatJID.User):
        default:
                return whatsappText(ctx, device, own, event)
        }

        return WhatsAppChat{}, WhatsAppMessage{}, false
}

// whatsappText returns the chat and the batch message of event, a message that
// the batches take. An image returns its caption until receive stores it.
func whatsappText(ctx context.Context, device *store.Device, own []string, event *events.Message) (WhatsAppChat, WhatsAppMessage, bool) {
        info := event.Info
        text := event.Message.GetConversation()
        extended := event.Message.GetExtendedTextMessage()
        if text == "" {
                text = extended.GetText()
        }

        contextInfo := extended.GetContextInfo()
        switch {
        case event.Message.GetImageMessage() != nil:
                image := event.Message.GetImageMessage()
                text = image.GetCaption()
                contextInfo = image.GetContextInfo()
        case event.Message.GetStickerMessage() != nil:
                text = "(sticker)"
                contextInfo = event.Message.GetStickerMessage().GetContextInfo()
        case event.Message.GetVideoMessage() != nil:
                video := event.Message.GetVideoMessage()
                text = "(video)"
                if video.GetGifPlayback() {
                        text = "(GIF)"
                }

                if video.GetCaption() != "" {
                        text += "\n" + video.GetCaption()
                }

                contextInfo = video.GetContextInfo()
        default:
                if strings.TrimSpace(text) == "" {
                        return WhatsAppChat{}, WhatsAppMessage{}, false
                }
        }

        message := WhatsAppMessage{
                ID:   info.ID,
                Time: info.Timestamp,
                Mine: info.IsFromMe,
        }

        if !message.Mine {
                message.Sender = whatsappContactName(ctx, device, info.Sender, info.PushName)
        }

        for _, mentioned := range contextInfo.GetMentionedJID() {
                jid, err := types.ParseJID(mentioned)
                switch {
                case err != nil:
                case slices.Contains(own, jid.User):
                        message.Mention = true
                        text = strings.ReplaceAll(text, "@"+jid.User, "@you")
                default:
                        text = strings.ReplaceAll(text, "@"+jid.User, "@"+whatsappContactName(ctx, device, jid, ""))
                }
        }

        participant, err := types.ParseJID(contextInfo.GetParticipant())
        message.Reply = err == nil && slices.Contains(own, participant.User)

        message.Text = whatsappRedact(own, text)
        chatJID := info.Chat.ToNonAD()
        chat := WhatsAppChat{
                ID:    chatJID.String(),
                Group: chatJID.Server == types.GroupServer,
        }

        if !chat.Group {
                push := info.PushName
                if message.Mine {
                        push = ""
                }

                chat.Name = whatsappContactName(ctx, device, chatJID, push)
        }

        return chat, message, true
}