From f7a135af0ed5e05fecadee79d38cd4560e7e31ad Mon Sep 17 00:00:00 2001
From: Zhang Minghan <zmh13054618081@dingtalk.com>
Date: Mon, 11 Mar 2024 20:09:34 +0800
Subject: [PATCH] fix: fix sharing message use message action issue

---
 app/src/components/Message.tsx |  6 ++++--
 app/src/events/sharing.ts      | 12 ------------
 app/src/routes/Sharing.tsx     | 20 +++++++++++++------
 middleware/throttle.go         | 36 +++++++++++++++++++++++-----------
 utils/cache.go                 | 12 ++++++------
 utils/websocket.go             |  3 ++-
 6 files changed, 51 insertions(+), 38 deletions(-)
 delete mode 100644 app/src/events/sharing.ts

diff --git a/app/src/components/Message.tsx b/app/src/components/Message.tsx
index bb892f1..33c208d 100644
--- a/app/src/components/Message.tsx
+++ b/app/src/components/Message.tsx
@@ -41,6 +41,7 @@ type MessageProps = {
   index: number;
   message: Message;
   end?: boolean;
+  username?: string;
   onEvent?: (event: string, index?: number, message?: string) => void;
   ref?: Ref<HTMLElement>;
   sharing?: boolean;
@@ -216,10 +217,11 @@ function MessageContent({
   index,
   onEvent,
   selected,
+  username,
 }: MessageProps) {
   const isUser = message.role === "user";
 
-  const username = useSelector(selectUsername);
+  const user = useSelector(selectUsername);
 
   const [open, setOpen] = useState(false);
   const [editedMessage, setEditedMessage] = useState<string | undefined>("");
@@ -239,7 +241,7 @@ function MessageContent({
           isUser ? (
             <Avatar
               className={`message-avatar animate-fade-in`}
-              username={username}
+              username={username ?? user}
             />
           ) : (
             <img
diff --git a/app/src/events/sharing.ts b/app/src/events/sharing.ts
deleted file mode 100644
index 33fd84b..0000000
--- a/app/src/events/sharing.ts
+++ /dev/null
@@ -1,12 +0,0 @@
-import { EventCommitter } from "./struct.ts";
-import { Message } from "@/api/types.tsx";
-
-export type SharingEvent = {
-  refer: string;
-  data: Message[];
-};
-
-export const sharingEvent = new EventCommitter<SharingEvent>({
-  name: "sharing",
-  destroyedAfterTrigger: true,
-});
diff --git a/app/src/routes/Sharing.tsx b/app/src/routes/Sharing.tsx
index fdcf5bf..ce1545d 100644
--- a/app/src/routes/Sharing.tsx
+++ b/app/src/routes/Sharing.tsx
@@ -19,7 +19,6 @@ import MessageSegment from "@/components/Message.tsx";
 import { Button } from "@/components/ui/button.tsx";
 import router from "@/router.tsx";
 import { useToast } from "@/components/ui/use-toast.ts";
-import { sharingEvent } from "@/events/sharing.ts";
 import { Message } from "@/api/types.tsx";
 import Avatar from "@/components/Avatar.tsx";
 import { toJpeg } from "html-to-image";
@@ -28,18 +27,20 @@ import { extractMessage } from "@/utils/processor.ts";
 import { cn } from "@/components/ui/lib/utils.ts";
 import { useMobile } from "@/utils/device.ts";
 import { ScrollArea } from "@/components/ui/scroll-area.tsx";
+import { useConversationActions } from "@/store/chat.ts";
 
 type SharingFormProps = {
   refer?: string;
   data: ViewData | null;
 };
 
-function SharingForm({ refer, data }: SharingFormProps) {
+function SharingForm({ data }: SharingFormProps) {
   if (data === null) return null;
 
   const { t } = useTranslation();
   const { toast } = useToast();
   const mobile = useMobile();
+  const { mask: setMask } = useConversationActions();
   const [maximized, setMaximized] = useState(false);
   const container = useRef<HTMLDivElement>(null);
   const date = new Date(data.time);
@@ -111,7 +112,12 @@ function SharingForm({ refer, data }: SharingFormProps) {
             </div>
             <div className={`shot-content`}>
               {data.messages.map((message, i) => (
-                <MessageSegment message={message} key={i} index={i} />
+                <MessageSegment
+                  message={message}
+                  key={i}
+                  index={i}
+                  username={data.username}
+                />
               ))}
             </div>
           </div>
@@ -151,9 +157,11 @@ function SharingForm({ refer, data }: SharingFormProps) {
         <Button
           variant={`outline`}
           onClick={async () => {
-            sharingEvent.emit({
-              refer: refer as string,
-              data: data?.messages as Message[],
+            const message: Message[] = data?.messages || [];
+            setMask({
+              avatar: "",
+              name: data.name,
+              context: message,
             });
             await router.navigate("/");
           }}
diff --git a/middleware/throttle.go b/middleware/throttle.go
index b4967c3..37c29c1 100644
--- a/middleware/throttle.go
+++ b/middleware/throttle.go
@@ -14,10 +14,10 @@ type Limiter struct {
 	Count    int64
 }
 
-func (l *Limiter) RateLimit(client *redis.Client, ip string, path string) bool {
+func (l *Limiter) RateLimit(client *redis.Client, ip string, path string) (bool, error) {
 	key := fmt.Sprintf("rate:%s:%s", path, ip)
-	rate := utils.IncrWithLimit(client, key, 1, l.Count, int64(l.Duration))
-	return !rate
+	rate, err := utils.IncrWithLimit(client, key, 1, l.Count, int64(l.Duration))
+	return !rate, err
 }
 
 var limits = map[string]Limiter{
@@ -64,14 +64,28 @@ func ThrottleMiddleware() gin.HandlerFunc {
 		cache := utils.GetCacheFromContext(c)
 
 		limiter := GetPrefixMap[Limiter](path, limits)
-		if limiter != nil && limiter.RateLimit(cache, ip, path) {
-			c.JSON(200, gin.H{
-				"status": false,
-				"reason": "You have sent too many requests. Please try again later.",
-				"error":  "request_throttled",
-			})
-			c.Abort()
-			return
+		if limiter != nil {
+			rate, err := limiter.RateLimit(cache, ip, path)
+
+			if err != nil {
+				c.JSON(200, gin.H{
+					"status": false,
+					"reason": err.Error(),
+					"error":  err.Error(),
+				})
+				c.Abort()
+				return
+			}
+
+			if rate {
+				c.JSON(200, gin.H{
+					"status": false,
+					"reason": "You have sent too many requests. Please try again later.",
+					"error":  "request_throttled",
+				})
+				c.Abort()
+				return
+			}
 		}
 		c.Next()
 	}
diff --git a/utils/cache.go b/utils/cache.go
index fdf4134..561df96 100644
--- a/utils/cache.go
+++ b/utils/cache.go
@@ -53,25 +53,25 @@ func SetCache(cache *redis.Client, key string, value string, expiration int64) e
 	return cache.Set(context.Background(), key, value, time.Duration(expiration)*time.Second).Err()
 }
 
-func IncrWithLimit(cache *redis.Client, key string, delta int64, limit int64, expiration int64) bool {
+func IncrWithLimit(cache *redis.Client, key string, delta int64, limit int64, expiration int64) (bool, error) {
 	// not exist
 	if _, err := cache.Get(context.Background(), key).Result(); err != nil {
 		if errors.Is(err, redis.Nil) {
 			cache.Set(context.Background(), key, delta, time.Duration(expiration)*time.Second)
-			return true
+			return true, nil
 		}
-		return false
+		return false, err
 	}
 	res, err := Incr(cache, key, delta)
 	if err != nil {
-		return false
+		return false, err
 	}
 	if res > limit {
 		// reset
 		cache.Set(context.Background(), key, limit, time.Duration(expiration)*time.Second)
-		return false
+		return false, nil
 	}
-	return true
+	return true, nil
 }
 
 func DecrInt(cache *redis.Client, key string, delta int64) bool {
diff --git a/utils/websocket.go b/utils/websocket.go
index e0cc945..f69d436 100644
--- a/utils/websocket.go
+++ b/utils/websocket.go
@@ -145,7 +145,8 @@ func (w *WebSocket) IncrRate(key string) bool {
 
 func (w *WebSocket) IncrRateWithLimit(key string, limit int64, expiration int64) bool {
 	cache := w.GetCache()
-	return IncrWithLimit(cache, key, 1, limit, expiration)
+	state, err := IncrWithLimit(cache, key, 1, limit, expiration)
+	return state && err == nil
 }
 
 func (w *WebSocket) GetCtx() *gin.Context {