Add a timeout for sed

This reverts commit dc2eefd581.

stdplugins/sed.py

@@ -1,9 +1,12 @@
 from collections import defaultdict, deque
 import re
 
+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',
@@ -16,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)
@@ -33,7 +37,7 @@ def doit(chat_id, match, original):
     flags = 0
     for f in fl:
         if f == 'i':
-            flags |= re.IGNORECASE
+            flags |= regex.IGNORECASE
         elif f == 'g':
             count = 0
         else:
@@ -44,7 +48,7 @@ def doit(chat_id, match, original):
             s = original.message
             if s.startswith(HEADER):
                 s = s[len(HEADER):]
-            s, i = re.subn(fr, to, s, count=count, flags=flags)
+            s, i = regex.subn(fr, to, s, count=count, flags=flags)
             if i > 0:
                 return original, s
         except Exception as e:

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