From 564f032e3a7db2b83d4a7bfb3e9013d97c992497 Mon Sep 17 00:00:00 2001
From: IndieKKY <indiekky@gmail.com>
Date: Fri, 4 Oct 2024 15:31:51 +0800
Subject: [PATCH] fix

---
 src/hooks/useMessagingService.ts   |  8 +++++---
 src/messaging/ExtensionMessage.ts  |  1 +
 src/messaging/InjectMessage.ts     |  3 +++
 src/messaging/useMessage.ts        |  1 +
 src/messaging/useMessageService.ts | 20 ++++++++++++++------
 src/typings.d.ts                   |  3 +++
 6 files changed, 27 insertions(+), 9 deletions(-)

diff --git a/src/hooks/useMessagingService.ts b/src/hooks/useMessagingService.ts
index 7024de9..7cd34cf 100644
--- a/src/hooks/useMessagingService.ts
+++ b/src/hooks/useMessagingService.ts
@@ -8,15 +8,17 @@ const useMessagingService = () => {
   const dispatch = useAppDispatch()
   
   //methods
-  const methods = useMemo(() => ({
-    [MESSAGE_TO_APP_SET_INFOS]: (params: any, from: string, context: MethodContext) => {
+  const methods: {
+    [key: string]: (params: any, context: MethodContext) => boolean
+  } = useMemo(() => ({
+    [MESSAGE_TO_APP_SET_INFOS]: (params: any, context: MethodContext) => {
       dispatch(setInfos(params.infos))
       dispatch(setCurInfo(undefined))
       dispatch(setCurFetched(false))
       dispatch(setData(undefined))
       return true
     },
-    [MESSAGE_TO_APP_SET_VIDEO_INFO]: (params: any, from: string, context: MethodContext) => {
+    [MESSAGE_TO_APP_SET_VIDEO_INFO]: (params: any, context: MethodContext) => {
       dispatch(setInfos(params.infos))
       dispatch(setUrl(params.url))
       dispatch(setTitle(params.title))
diff --git a/src/messaging/ExtensionMessage.ts b/src/messaging/ExtensionMessage.ts
index 150eb09..622e28c 100644
--- a/src/messaging/ExtensionMessage.ts
+++ b/src/messaging/ExtensionMessage.ts
@@ -26,6 +26,7 @@ class ExtensionMessage {
       const method = this.methods?.[event.method]
       if (method != null) {
         method(event.params, {
+          from: event.from,
           event,
           sender,
         }).then(data => sendResponse({
diff --git a/src/messaging/InjectMessage.ts b/src/messaging/InjectMessage.ts
index 132484a..0788e48 100644
--- a/src/messaging/InjectMessage.ts
+++ b/src/messaging/InjectMessage.ts
@@ -25,6 +25,7 @@ class InjectMessage {
         const method = this.methods?.[event.method]
         if (method != null) {
             method(event.params, {
+                from: event.from,
                 event,
                 sender,
             }).then(data => {
@@ -71,6 +72,7 @@ class InjectMessage {
             this.postMessageToApp = postMessage
             listenMessage((method, params, sendResponse) => {
                 this.messageHandler({
+                    from: 'app',
                     target: MESSAGE_TARGET_INJECT,
                     method,
                     params,
@@ -87,6 +89,7 @@ class InjectMessage {
 
     sendExtension = async <T = any>(method: string, params?: any): Promise<T> => {
         return await chrome.runtime.sendMessage<MessageData, MessageResult>({
+            from: 'inject',
             target: MESSAGE_TARGET_EXTENSION,
             method,
             params: params ?? {},
diff --git a/src/messaging/useMessage.ts b/src/messaging/useMessage.ts
index b021206..b255bfb 100644
--- a/src/messaging/useMessage.ts
+++ b/src/messaging/useMessage.ts
@@ -5,6 +5,7 @@ import { useCallback } from 'react'
 const useMessage = () => {
     const sendExtension = useCallback(async <T = any>(method: string, params?: any) => {
         return await chrome.runtime.sendMessage<MessageData, MessageResult>({
+            from: 'app',
             target: MESSAGE_TARGET_EXTENSION,
             method,
             params: params ?? {},
diff --git a/src/messaging/useMessageService.ts b/src/messaging/useMessageService.ts
index 82f8ae8..b791d95 100644
--- a/src/messaging/useMessageService.ts
+++ b/src/messaging/useMessageService.ts
@@ -19,12 +19,12 @@ export const injectWaiter = new Waiter<typeof postInjectMessage>(() => ({
 }), 100, 15000)
 
 const useMessageService = (methods?: {
-  [key: string]: (params: any, from: string, context: MethodContext) => boolean
+  [key: string]: (params: any, context: MethodContext) => boolean
 }) => {
-  const messageHandler = useCallback((method: string, params: any, from: string, context: any): boolean => {
+  const messageHandler = useCallback((method: string, params: any, context: MethodContext): boolean => {
     const handler = methods?.[method]
     if (handler != null) {
-      return handler(params, from, context)
+      return handler(params, context)
     }else {
       debug('unknown message method: ', method)
       return false
@@ -48,7 +48,13 @@ const useMessageService = (methods?: {
       listenMessage((method, params, sendResponse) => {
         debug('inject => ', method, params)
 
-        const success = messageHandler(method, params, MESSAGE_TARGET_INJECT, {})
+        const success = messageHandler(method, params, {
+          from: 'inject',
+          event: {
+            method,
+            params,
+          },
+        })
         sendResponse({
           success,
           code: success ? 200 : 500
@@ -69,8 +75,10 @@ const useMessageService = (methods?: {
     // check event target
     if (!event || event.target !== MESSAGE_TARGET_APP) return
 
-    messageHandler(event.method, event.params, MESSAGE_TARGET_EXTENSION, {
-      sender
+    messageHandler(event.method, event.params, {
+      from: 'extension',
+      event,
+      sender,
     })
   }, [messageHandler])
 
diff --git a/src/typings.d.ts b/src/typings.d.ts
index 320162e..25183ab 100644
--- a/src/typings.d.ts
+++ b/src/typings.d.ts
@@ -1,5 +1,7 @@
+type MessageFrom = 'extension' | 'inject' | 'app'
 
 interface MessageData {
+  from: MessageFrom
   target: string
   method: string
   params?: any
@@ -14,6 +16,7 @@ interface MessageResult {
 }
 
 interface MethodContext {
+  from: MessageFrom
   event: any
   sender?: chrome.runtime.MessageSender | null
 }