From a896b5c869f15648c3cc341ea91bb403a14d9bb3 Mon Sep 17 00:00:00 2001
From: Ryan Barrett <git@ryanb.org>
Date: Wed, 28 Aug 2024 21:45:09 -0700
Subject: [PATCH] Protocol.check_supported: allow DMs to/from protocol bot user
 copy ids

for #966
---
 common.py              | 14 ++++++++++++++
 protocol.py            |  5 +++--
 tests/test_protocol.py | 20 +++++++++++++++++---
 tests/testutil.py      |  1 +
 4 files changed, 35 insertions(+), 5 deletions(-)

diff --git a/common.py b/common.py
index dd73f71..cb24249 100644
--- a/common.py
+++ b/common.py
@@ -13,6 +13,7 @@ import cachetools
 from Crypto.Util import number
 from flask import abort, g, has_request_context, make_response, request
 from google.cloud.error_reporting.util import build_flask_context
+from google.cloud import ndb
 from google.cloud.ndb.global_cache import _InProcessGlobalCache, MemcacheCache
 from google.cloud.ndb.key import Key
 from google.protobuf.timestamp_pb2 import Timestamp
@@ -108,6 +109,19 @@ else:
     global_cache = MemcacheCache(memcache)
 
 
+@functools.cache
+def protocol_user_copy_ids():
+    """Returns all copy ids for protocol bot users."""
+    ids = []
+
+    from web import Web
+    for user in ndb.get_multi(Web(id=domain).key for domain in PROTOCOL_DOMAINS):
+        if user:
+            ids.extend(copy.uri for copy in user.copies)
+
+    return tuple(ids)
+
+
 def base64_to_long(x):
     """Converts from URL safe base64 encoding to long integer.
 
diff --git a/protocol.py b/protocol.py
index 7a6140a..324ca76 100644
--- a/protocol.py
+++ b/protocol.py
@@ -1614,9 +1614,10 @@ Hi! You <a href="{inner_obj_as1.get('url') or inner_obj_id}">recently replied</a
 
         # DMs are only allowed to/from protocol bot accounts
         if recip := as1.recipient_if_dm(obj.as1):
+            protocol_user_ids = PROTOCOL_DOMAINS + common.protocol_user_copy_ids()
             if (not cls.SUPPORTS_DMS
-                    or (recip not in PROTOCOL_DOMAINS
-                        and as1.get_owner(obj.as1) not in PROTOCOL_DOMAINS)):
+                    or (recip not in protocol_user_ids
+                        and as1.get_owner(obj.as1) not in protocol_user_ids)):
                 error(f"Bridgy Fed doesn't support DMs", status=204)
 
 
diff --git a/tests/test_protocol.py b/tests/test_protocol.py
index 49999af..e7ea653 100644
--- a/tests/test_protocol.py
+++ b/tests/test_protocol.py
@@ -759,13 +759,13 @@ class ProtocolTest(TestCase):
                 Fake.check_supported(Object(our_as1=obj))
 
         # Fake doesn't support DMs, ExplicitEnableFake does
-        for actor, recip in (
+        for author, recip in (
                 ('ap.brid.gy', 'did:bob'),
                 ('did:bob', 'ap.brid.gy'),
         ):
             bot_dm = Object(our_as1={
                 'objectType': 'note',
-                'actor': actor,
+                'author': author,
                 'to': [recip],
                 'content': 'hello world',
             })
@@ -773,9 +773,10 @@ class ProtocolTest(TestCase):
             with self.assertRaises(NoContent):
                 Fake.check_supported(bot_dm)
 
+        # not from or to a protocol bot user
         dm = Object(our_as1={
             'objectType': 'note',
-            'actor': 'did:alice',
+            'author': 'did:alice',
             'to': ['did:bob'],
             'content': 'hello world',
         })
@@ -783,6 +784,19 @@ class ProtocolTest(TestCase):
             with self.subTest(proto=proto), self.assertRaises(NoContent):
                 proto.check_supported(dm)
 
+        # from and to a copy id of a protocol bot user
+        self.make_user(cls=Web, id='ap.brid.gy',
+                       copies=[Target(protocol='fake', uri='fake:ap-bot')])
+        common.protocol_user_copy_ids.cache_clear()
+        dm.our_as1['author'] = 'fake:ap-bot'
+        proto.check_supported(dm)
+
+        dm.our_as1.update({
+            'author': 'did:alice',
+            'to': ['fake:ap-bot'],
+        })
+        proto.check_supported(dm)
+
     def test_bot_follow(self):
         self.make_user(id='fa.brid.gy', cls=Web)
         user = self.make_user(id='fake:user', cls=Fake, obj_id='fake:user')
diff --git a/tests/testutil.py b/tests/testutil.py
index f39923c..45b07ef 100644
--- a/tests/testutil.py
+++ b/tests/testutil.py
@@ -272,6 +272,7 @@ class TestCase(unittest.TestCase, testutil.Asserts):
         protocol.Protocol.for_id.cache.clear()
         protocol.Protocol.for_handle.cache.clear()
         User.count_followers.cache.clear()
+        common.protocol_user_copy_ids.cache_clear()
 
         for cls in ExplicitEnableFake, Fake, OtherFake:
             cls.fetchable = {}