From c0b75e4fbd70af0b3d8fb7f987f32c109d498759 Mon Sep 17 00:00:00 2001
From: Ryan Barrett <git@ryanb.org>
Date: Mon, 11 Mar 2024 11:43:58 -0700
Subject: [PATCH] Protocol.load: return None if the entity is too big

fixes https://console.cloud.google.com/errors/detail/CO-fpZfMq_6ktgE;time=P30D?project=bridgy-federated
---
 protocol.py            |  8 ++++++++
 tests/test_protocol.py | 11 ++++++++++-
 2 files changed, 18 insertions(+), 1 deletion(-)

diff --git a/protocol.py b/protocol.py
index 602ad6b..334299a 100644
--- a/protocol.py
+++ b/protocol.py
@@ -10,9 +10,11 @@ from cachetools import cached, LRUCache
 from flask import g, request
 from google.cloud import ndb
 from google.cloud.ndb import OR
+from google.cloud.ndb.model import _entity_to_protobuf
 from granary import as1
 from oauth_dropins.webutil.flask_util import cloud_tasks_only
 from oauth_dropins.webutil import util
+from oauth_dropins.webutil import models
 from oauth_dropins.webutil.util import json_dumps, json_loads
 import werkzeug.exceptions
 
@@ -1160,6 +1162,12 @@ class Protocol:
         if not fetched:
             return None
 
+        # https://stackoverflow.com/a/3042250/186123
+        size = len(_entity_to_protobuf(obj)._pb.SerializeToString())
+        if size > models.MAX_ENTITY_SIZE:
+            logger.warning(f'Object is too big! {size} bytes is over {models.MAX_ENTITY_SIZE}')
+            return None
+
         obj.resolve_ids()
 
         if obj.new is False:
diff --git a/tests/test_protocol.py b/tests/test_protocol.py
index 3f0696c..04114f7 100644
--- a/tests/test_protocol.py
+++ b/tests/test_protocol.py
@@ -9,7 +9,7 @@ from arroba.tests.testutil import dns_answer
 from flask import g
 from google.cloud import ndb
 from granary import as2
-from oauth_dropins.webutil import appengine_info, util
+from oauth_dropins.webutil import appengine_info, models, util
 from oauth_dropins.webutil.flask_util import CLOUD_TASKS_QUEUE_HEADER, NoContent
 from oauth_dropins.webutil.testutil import NOW, requests_response
 import requests
@@ -319,6 +319,7 @@ class ProtocolTest(TestCase):
         follow = {
             'objectType': 'activity',
             'verb': 'follow',
+            'id': 'fake:follow',
             'actor': 'fake:alice',
             'object': 'fake:bob',
         }
@@ -362,6 +363,14 @@ class ProtocolTest(TestCase):
         loaded = Fake.load('foo')
         self.assertEqual({'fetched': 'x', 'id': 'foo'}, loaded.our_as1)
 
+    @patch('oauth_dropins.webutil.models.MAX_ENTITY_SIZE', new=50)
+    def test_load_too_big(self):
+        Fake.fetchable['fake:foo'] = {
+            'objectType': 'note',
+            'content': 'a bit of text that makes sure we end up over the limit ',
+        }
+        self.assertIsNone(Fake.load('fake:foo'))
+
     def test_actor_key(self):
         user = self.make_user(id='fake:a', cls=Fake)
         a_key = user.key