diff --git a/README.md b/README.md index 22707a5..0dd6877 100644 --- a/README.md +++ b/README.md @@ -115,4 +115,5 @@ Or you can start the service with Docker Compose. This can be combined with the - `/start `: Start the bot with your Memos access token. - Send text messages: Save the message content as a memo. - Send files (photos, documents): Save the files as resources in a memo. +- Edit a sent message: Update the memo that message originally created. Edits to messages Memogram never captured (for example, messages sent before this feature was available) are ignored. - `/search `: Search for the memos. diff --git a/memogram.go b/memogram.go index c753d3e..7ee530f 100644 --- a/memogram.go +++ b/memogram.go @@ -167,6 +167,10 @@ func (s *Service) handler(ctx context.Context, b *bot.Bot, m *models.Update) { fmt.Println("Service or config is nil") return } + if m != nil && m.EditedMessage != nil { + s.editedMessageHandler(ctx, b, m) + return + } if m == nil || m.Message == nil || m.Message.From == nil { s.sendError(b, 0, errors.New("invalid message structure: missing required fields")) return @@ -272,6 +276,7 @@ func (s *Service) handler(ctx context.Context, b *bot.Bot, m *models.Update) { }) return } + s.store.SetMemoForMessage(message.Chat.ID, int64(message.ID), memo.Name) if message.Document != nil { s.processFileMessage(ctx, authClient, b, m, message.Document.FileID, memo) @@ -341,6 +346,54 @@ func (s *Service) startHandler(ctx context.Context, b *bot.Bot, m *models.Update }) } +// Edits of messages with no known memo (sent before this feature existed) are ignored. +func (s *Service) editedMessageHandler(ctx context.Context, b *bot.Bot, m *models.Update) { + message := m.EditedMessage + if message == nil || message.From == nil || message.Chat.ID == 0 { + return + } + if !s.isUserAllowed(message.From.Username) { + return + } + + userID := message.From.ID + accessToken, ok := s.store.GetUserAccessToken(userID) + if !ok { + return + } + + memoName, ok := s.store.GetMemoForMessage(message.Chat.ID, int64(message.ID)) + if !ok { + return + } + + content := message.Text + contentEntities := message.Entities + if message.Caption != "" { + content = message.Caption + contentEntities = message.CaptionEntities + } + if len(contentEntities) > 0 { + content = formatContent(content, contentEntities) + } + if content == "" { + return + } + + authClient := s.client.NewAuthenticatedClient(accessToken) + _, err := authClient.MemoService.UpdateMemo(ctx, connect.NewRequest(&v1pb.UpdateMemoRequest{ + Memo: &v1pb.Memo{ + Name: memoName, + Content: content, + }, + UpdateMask: &fieldmaskpb.FieldMask{Paths: []string{"content"}}, + })) + if err != nil { + slog.Error("failed to update memo on edit", slog.Any("err", err)) + s.sendError(b, message.Chat.ID, fmt.Errorf("failed to update memo: %w", err)) + } +} + func (s *Service) keyboard(memo *v1pb.Memo) *models.InlineKeyboardMarkup { // add inline keyboard to edit memo's visibility or pinned status. return &models.InlineKeyboardMarkup{ diff --git a/store/message.go b/store/message.go new file mode 100644 index 0000000..535aecb --- /dev/null +++ b/store/message.go @@ -0,0 +1,153 @@ +package store + +import ( + "bufio" + "fmt" + "log/slog" + "os" + "path/filepath" + "sort" + "strconv" + "strings" +) + +// messageKey scopes a Telegram message ID to its chat: message IDs are only +// unique within a chat, not globally. +type messageKey struct { + chatID int64 + messageID int64 +} + +func (s *Store) GetMemoForMessage(chatID, messageID int64) (string, bool) { + memoName, ok := s.messageMemoCache.Load(messageKey{chatID: chatID, messageID: messageID}) + if !ok { + return "", false + } + return memoName.(string), true +} + +func (s *Store) SetMemoForMessage(chatID, messageID int64, memoName string) { + s.saveMu.Lock() + defer s.saveMu.Unlock() + + s.messageMemoCache.Store(messageKey{chatID: chatID, messageID: messageID}, memoName) + if err := s.saveMessageMemoMapToFile(); err != nil { + slog.Error("failed to save message memo map to file", "error", err) + } +} + +func (s *Store) messageMemoMapDataPath() string { + return s.Data + ".messages" +} + +func (s *Store) saveMessageMemoMapToFile() error { + entries := s.snapshotMessageMemoMap() + dataPath := s.messageMemoMapDataPath() + dataDir := filepath.Dir(dataPath) + tmpFile, err := os.CreateTemp(dataDir, "memogram-messages-*.tmp") + if err != nil { + return fmt.Errorf("create temp file: %w", err) + } + defer os.Remove(tmpFile.Name()) + + writer := bufio.NewWriter(tmpFile) + for _, entry := range entries { + if _, err := fmt.Fprintf(writer, "%d:%d:%s\n", entry.key.chatID, entry.key.messageID, entry.memoName); err != nil { + tmpFile.Close() + return fmt.Errorf("write data file: %w", err) + } + } + if err := writer.Flush(); err != nil { + tmpFile.Close() + return fmt.Errorf("flush data file: %w", err) + } + if err := tmpFile.Sync(); err != nil { + tmpFile.Close() + return fmt.Errorf("sync data file: %w", err) + } + if err := tmpFile.Close(); err != nil { + return fmt.Errorf("close data file: %w", err) + } + + if err := os.Rename(tmpFile.Name(), dataPath); err != nil { + return fmt.Errorf("replace data file: %w", err) + } + return nil +} + +func (s *Store) loadMessageMemoMapFromFile() error { + dataPath := s.messageMemoMapDataPath() + if _, err := os.Stat(dataPath); os.IsNotExist(err) { + file, err := os.Create(dataPath) + if err != nil { + return err + } + defer file.Close() + } + + file, err := os.Open(dataPath) + if err != nil { + return err + } + defer file.Close() + + scanner := bufio.NewScanner(file) + for scanner.Scan() { + line := strings.TrimSpace(scanner.Text()) + if line == "" || strings.HasPrefix(line, "#") { + continue + } + key, memoName := parseMessageMemoLine(line) + if key.messageID == 0 || memoName == "" { + continue + } + s.messageMemoCache.Store(key, memoName) + } + return scanner.Err() +} + +func parseMessageMemoLine(line string) (messageKey, string) { + parts := strings.SplitN(line, ":", 3) + if len(parts) != 3 { + return messageKey{}, "" + } + chatID, err := strconv.ParseInt(parts[0], 10, 64) + if err != nil { + return messageKey{}, "" + } + messageID, err := strconv.ParseInt(parts[1], 10, 64) + if err != nil { + return messageKey{}, "" + } + return messageKey{chatID: chatID, messageID: messageID}, parts[2] +} + +type messageMemoEntry struct { + key messageKey + memoName string +} + +func (s *Store) snapshotMessageMemoMap() []messageMemoEntry { + entries := make([]messageMemoEntry, 0) + s.messageMemoCache.Range(func(key, value interface{}) bool { + messageKey, ok := key.(messageKey) + if !ok { + return true + } + memoName, ok := value.(string) + if !ok { + return true + } + entries = append(entries, messageMemoEntry{key: messageKey, memoName: memoName}) + return true + }) + + sort.Slice(entries, func(i, j int) bool { + if entries[i].key.chatID != entries[j].key.chatID { + return entries[i].key.chatID < entries[j].key.chatID + } + return entries[i].key.messageID < entries[j].key.messageID + }) + + return entries +} diff --git a/store/message_test.go b/store/message_test.go new file mode 100644 index 0000000..6daa9ca --- /dev/null +++ b/store/message_test.go @@ -0,0 +1,77 @@ +package store + +import ( + "path/filepath" + "testing" +) + +func TestSaveAndLoadMessageMemoMap(t *testing.T) { + dataPath := filepath.Join(t.TempDir(), "data.txt") + + store := NewStore(dataPath) + if err := store.Init(); err != nil { + t.Fatalf("init store: %v", err) + } + + store.SetMemoForMessage(1, 101, "memos/abc123") + store.SetMemoForMessage(1, 202, "memos/def456") + + reloaded := NewStore(dataPath) + if err := reloaded.Init(); err != nil { + t.Fatalf("init reloaded store: %v", err) + } + + memoName, ok := reloaded.GetMemoForMessage(1, 101) + if !ok || memoName != "memos/abc123" { + t.Fatalf("expected memos/abc123 for message 101, got %q", memoName) + } + + memoName, ok = reloaded.GetMemoForMessage(1, 202) + if !ok || memoName != "memos/def456" { + t.Fatalf("expected memos/def456 for message 202, got %q", memoName) + } +} + +func TestGetMemoForMessageMissing(t *testing.T) { + dataPath := filepath.Join(t.TempDir(), "data.txt") + + store := NewStore(dataPath) + if err := store.Init(); err != nil { + t.Fatalf("init store: %v", err) + } + + if _, ok := store.GetMemoForMessage(1, 999); ok { + t.Fatalf("expected no mapping for message 999") + } +} + +func TestMessageMemoMapScopedByChat(t *testing.T) { + dataPath := filepath.Join(t.TempDir(), "data.txt") + + store := NewStore(dataPath) + if err := store.Init(); err != nil { + t.Fatalf("init store: %v", err) + } + + // Same message ID in two different chats must not collide: Telegram + // message IDs are only unique within a chat. + store.SetMemoForMessage(1, 5, "memos/chat-one") + store.SetMemoForMessage(2, 5, "memos/chat-two") + + // Reload from disk to also cover the chatID:messageID:memoName persisted format, + // not just the in-memory cache. + reloaded := NewStore(dataPath) + if err := reloaded.Init(); err != nil { + t.Fatalf("init reloaded store: %v", err) + } + + memoName, ok := reloaded.GetMemoForMessage(1, 5) + if !ok || memoName != "memos/chat-one" { + t.Fatalf("expected memos/chat-one for chat 1 message 5, got %q", memoName) + } + + memoName, ok = reloaded.GetMemoForMessage(2, 5) + if !ok || memoName != "memos/chat-two" { + t.Fatalf("expected memos/chat-two for chat 2 message 5, got %q", memoName) + } +} diff --git a/store/store.go b/store/store.go index 61469d3..59909d2 100644 --- a/store/store.go +++ b/store/store.go @@ -8,7 +8,12 @@ import ( type Store struct { Data string + // saveMu serializes cache-mutation-then-file-save sequences so an older + // snapshot can't win the rename race and discard a newer entry. + saveMu sync.Mutex + userAccessTokenCache sync.Map // map[int64]string + messageMemoCache sync.Map // map[messageKey]string } func NewStore(data string) *Store { @@ -16,6 +21,7 @@ func NewStore(data string) *Store { Data: data, userAccessTokenCache: sync.Map{}, + messageMemoCache: sync.Map{}, } } @@ -23,6 +29,9 @@ func (s *Store) Init() error { if err := s.loadUserAccessTokenMapFromFile(); err != nil { return fmt.Errorf("failed to load user access token map from file: %w", err) } + if err := s.loadMessageMemoMapFromFile(); err != nil { + return fmt.Errorf("failed to load message memo map from file: %w", err) + } return nil } diff --git a/store/user.go b/store/user.go index 18227a8..70d12c1 100644 --- a/store/user.go +++ b/store/user.go @@ -22,6 +22,9 @@ func (s *Store) GetUserAccessToken(userID int64) (string, bool) { // SetUserAccessToken sets the access token for the user. func (s *Store) SetUserAccessToken(userID int64, accessToken string) { + s.saveMu.Lock() + defer s.saveMu.Unlock() + s.userAccessTokenCache.Store(userID, accessToken) if err := s.SaveUserAccessTokenMapToFile(); err != nil { slog.Error("failed to save user access token map to file", "error", err)