From a50990bef2f757d3e7eb2cdcbe55b3ebe6d5a7a6 Mon Sep 17 00:00:00 2001
From: Ju4tCode <42488585+yanyongyu@users.noreply.github.com>
Date: Tue, 15 Nov 2022 10:50:52 +0800
Subject: [PATCH] :bug: fix run pre/post hook not in context (#1391)

---
 nonebot/internal/matcher.py | 66 +++++++++++++++++++++----------------
 nonebot/message.py          | 36 +++++++++++---------
 2 files changed, 57 insertions(+), 45 deletions(-)

diff --git a/nonebot/internal/matcher.py b/nonebot/internal/matcher.py
index 11eda889..18b30adc 100644
--- a/nonebot/internal/matcher.py
+++ b/nonebot/internal/matcher.py
@@ -1,8 +1,8 @@
 from types import ModuleType
 from contextvars import ContextVar
 from collections import defaultdict
-from contextlib import AsyncExitStack
 from datetime import datetime, timedelta
+from contextlib import AsyncExitStack, contextmanager
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -661,6 +661,18 @@ class Matcher(metaclass=MatcherMeta):
         if REJECT_CACHE_TARGET in self.state:
             self.state[REJECT_TARGET] = self.state[REJECT_CACHE_TARGET]
 
+    @contextmanager
+    def ensure_context(self, bot: Bot, event: Event):
+        b_t = current_bot.set(bot)
+        e_t = current_event.set(event)
+        m_t = current_matcher.set(self)
+        try:
+            yield
+        finally:
+            current_bot.reset(b_t)
+            current_event.reset(e_t)
+            current_matcher.reset(m_t)
+
     async def simple_run(
         self,
         bot: Bot,
@@ -673,35 +685,31 @@ class Matcher(metaclass=MatcherMeta):
             f"{self} run with incoming args: "
             f"bot={bot}, event={event!r}, state={state!r}"
         )
-        b_t = current_bot.set(bot)
-        e_t = current_event.set(event)
-        m_t = current_matcher.set(self)
-        try:
-            # Refresh preprocess state
-            self.state.update(state)
 
-            while self.handlers:
-                handler = self.handlers.pop(0)
-                current_handler.set(handler)
-                logger.debug(f"Running handler {handler}")
-                try:
-                    await handler(
-                        matcher=self,
-                        bot=bot,
-                        event=event,
-                        state=self.state,
-                        stack=stack,
-                        dependency_cache=dependency_cache,
-                    )
-                except SkippedException:
-                    logger.debug(f"Handler {handler} skipped")
-        except StopPropagation:
-            self.block = True
-        finally:
-            logger.info(f"{self} running complete")
-            current_bot.reset(b_t)
-            current_event.reset(e_t)
-            current_matcher.reset(m_t)
+        with self.ensure_context(bot, event):
+            try:
+                # Refresh preprocess state
+                self.state.update(state)
+
+                while self.handlers:
+                    handler = self.handlers.pop(0)
+                    current_handler.set(handler)
+                    logger.debug(f"Running handler {handler}")
+                    try:
+                        await handler(
+                            matcher=self,
+                            bot=bot,
+                            event=event,
+                            state=self.state,
+                            stack=stack,
+                            dependency_cache=dependency_cache,
+                        )
+                    except SkippedException:
+                        logger.debug(f"Handler {handler} skipped")
+            except StopPropagation:
+                self.block = True
+            finally:
+                logger.info(f"{self} running complete")
 
     # 运行handlers
     async def run(
diff --git a/nonebot/message.py b/nonebot/message.py
index 0a77a17c..9203d43b 100644
--- a/nonebot/message.py
+++ b/nonebot/message.py
@@ -167,17 +167,19 @@ async def _run_matcher(
         )
         for proc in _run_preprocessors
     ]:
-        try:
-            await asyncio.gather(*coros)
-        except IgnoredException:
-            logger.opt(colors=True).info(f"{matcher} running is <b>cancelled</b>")
-            return
-        except Exception as e:
-            logger.opt(colors=True, exception=e).error(
-                "<r><bg #f8bbd0>Error when running RunPreProcessors. Running cancelled!</bg #f8bbd0></r>"
-            )
+        # ensure matcher function can be correctly called
+        with matcher.ensure_context(bot, event):
+            try:
+                await asyncio.gather(*coros)
+            except IgnoredException:
+                logger.opt(colors=True).info(f"{matcher} running is <b>cancelled</b>")
+                return
+            except Exception as e:
+                logger.opt(colors=True, exception=e).error(
+                    "<r><bg #f8bbd0>Error when running RunPreProcessors. Running cancelled!</bg #f8bbd0></r>"
+                )
 
-            return
+                return
 
     exception = None
 
@@ -205,12 +207,14 @@ async def _run_matcher(
         )
         for proc in _run_postprocessors
     ]:
-        try:
-            await asyncio.gather(*coros)
-        except Exception as e:
-            logger.opt(colors=True, exception=e).error(
-                "<r><bg #f8bbd0>Error when running RunPostProcessors</bg #f8bbd0></r>"
-            )
+        # ensure matcher function can be correctly called
+        with matcher.ensure_context(bot, event):
+            try:
+                await asyncio.gather(*coros)
+            except Exception as e:
+                logger.opt(colors=True, exception=e).error(
+                    "<r><bg #f8bbd0>Error when running RunPostProcessors</bg #f8bbd0></r>"
+                )
 
     if matcher.block:
         raise StopPropagation