Add a timeout for sed

stdplugins/sed.py

@@ -5,6 +5,8 @@ import regex
 from telethon import events, utils
 from telethon.tl import types, functions
 
+from uniborg import util
+
 HEADER = "「sed」\n"
 KNOWN_RE_BOTS = re.compile(
     r'(regex|moku|BananaButler_|rgx|l4mR)bot',
@@ -17,6 +19,7 @@ KNOWN_RE_BOTS = re.compile(
 last_msgs = defaultdict(lambda: deque(maxlen=10))
 
 
+@util.sync_timeout(1)
 def doit(chat_id, match, original):
     fr = match.group(1)
     to = match.group(2)

uniborg/util.py

@@ -2,7 +2,9 @@
 # License, v. 2.0. If a copy of the MPL was not distributed with this
 # file, You can obtain one at http://mozilla.org/MPL/2.0/.
 
+import functools
 import re
+import signal
 
 from telethon import events
 from telethon.tl.functions.messages import GetPeerDialogsRequest
@@ -28,3 +30,20 @@ async def is_read(borg, entity, message, is_out=None):
     dialog = (await borg(GetPeerDialogsRequest([entity]))).dialogs[0]
     max_id = dialog.read_outbox_max_id if is_out else dialog.read_inbox_max_id
     return message_id <= max_id
+
+def _handle_timeout(signum, frame):
+    raise TimeoutError("Execution took too long")
+
+def sync_timeout(seconds):
+    def decorator(func):
+        @functools.wraps(func)
+        def wrapper(*args, **kwargs):
+            signal.signal(signal.SIGALRM, _handle_timeout)
+            signal.setitimer(signal.ITIMER_REAL, seconds)
+            try:
+                r = func(*args, **kwargs)
+            finally:
+                signal.alarm(0)
+            return r
+        return wrapper
+    return decorator