fix(download): only return a message from the requested chat

cmd/tg/download.go

@@ -49,12 +49,56 @@ func getMessage(ctx context.Context, api *tg.Client, peer tg.InputPeerClass, id
 	if !ok {
 		return nil, errors.Errorf("unexpected messages type %T", res)
 	}
+	// messages.getMessages looks ids up across all private chats and basic
+	// groups, so check the message belongs to the requested chat.
 	for _, m := range msgs.GetMessages() {
-		if msg, ok := m.(*tg.Message); ok && msg.ID == id {
-			return msg, nil
+		msg, ok := m.(*tg.Message)
+		if !ok || msg.ID != id {
+			continue
 		}
+		if !inChat(peer, msg, msgs.GetUsers()) {
+			break
+		}
+		return msg, nil
+	}
+	return nil, errors.Errorf("message #%d not found in this chat", id)
+}
+
+// inChat reports whether msg belongs to the chat peer addresses. Channel
+// lookups are scoped by the server. Saved Messages ("me") is the private chat
+// whose user the response marks as the current user.
+func inChat(peer tg.InputPeerClass, msg *tg.Message, users []tg.UserClass) bool {
+	switch peer.(type) {
+	case *tg.InputPeerChannel:
+		return true
+	case *tg.InputPeerSelf:
+		pu, ok := msg.PeerID.(*tg.PeerUser)
+		if !ok {
+			return false
+		}
+		for _, u := range users {
+			if user, ok := u.(*tg.User); ok && user.ID == pu.UserID {
+				return user.Self
+			}
+		}
+		return false
+	}
+	want := inputPeerID(peer)
+	return want <= 0 || peerIDOf(msg.PeerID) == want
+}
+
+// peerIDOf returns the bare id of a peer.
+func peerIDOf(p tg.PeerClass) int64 {
+	switch v := p.(type) {
+	case *tg.PeerUser:
+		return v.UserID
+	case *tg.PeerChat:
+		return v.ChatID
+	case *tg.PeerChannel:
+		return v.ChannelID
+	default:
+		return 0
 	}
-	return nil, errors.Errorf("message #%d not found", id)
 }
 
 // documentFilename returns the filename attribute of a document, or a fallback.

cmd/tg/download_test.go

@@ -1,9 +1,11 @@
 package main
 
 import (
+	"context"
 	"path/filepath"
 	"testing"
 
+	"github.com/gotd/td/bin"
 	"github.com/gotd/td/tg"
 )
 
@@ -82,3 +84,41 @@ func TestResolveOutPath(t *testing.T) {
 		t.Errorf("dir join = %q", got)
 	}
 }
+
+func TestGetMessageScopedToChat(t *testing.T) {
+	// messages.getMessages is not scoped to a chat: the same id may belong to
+	// another private chat, which must not be returned for this one.
+	api := newFuncAPI(t, func(bin.Encoder) (bin.Encoder, error) {
+		return &tg.MessagesMessages{Messages: []tg.MessageClass{
+			&tg.Message{ID: 5, PeerID: &tg.PeerUser{UserID: 999}},
+		}}, nil
+	})
+	ctx := context.Background()
+	if _, err := getMessage(ctx, api, &tg.InputPeerUser{UserID: 7, AccessHash: 1}, 5); err == nil {
+		t.Error("expected not found for a message from another chat")
+	}
+	if msg, err := getMessage(ctx, api, &tg.InputPeerUser{UserID: 999, AccessHash: 1}, 5); err != nil || msg.ID != 5 {
+		t.Errorf("same chat: %v, %v", msg, err)
+	}
+}
+
+func TestGetMessageSavedMessagesIsolation(t *testing.T) {
+	// For "me", a message id from another private chat must not be returned
+	// as if it were in Saved Messages.
+	users := []tg.UserClass{&tg.User{ID: 42, Self: true}, &tg.User{ID: 999}}
+	respond := func(peerID int64) *tg.Client {
+		return newFuncAPI(t, func(bin.Encoder) (bin.Encoder, error) {
+			return &tg.MessagesMessages{
+				Messages: []tg.MessageClass{&tg.Message{ID: 5, PeerID: &tg.PeerUser{UserID: peerID}}},
+				Users:    users,
+			}, nil
+		})
+	}
+	ctx := context.Background()
+	if _, err := getMessage(ctx, respond(999), &tg.InputPeerSelf{}, 5); err == nil {
+		t.Error("a message from another private chat must not be returned for me")
+	}
+	if msg, err := getMessage(ctx, respond(42), &tg.InputPeerSelf{}, 5); err != nil || msg.ID != 5 {
+		t.Errorf("saved message: %v, %v", msg, err)
+	}
+}