From 4d0d18931d27ac5c7379170eb6f098c123b78c56 Mon Sep 17 00:00:00 2001
From: CalciumIon <1808837298@qq.com>
Date: Wed, 17 Jul 2024 16:38:56 +0800
Subject: [PATCH] fix: try to fix panic #369

---
 relay/channel/openai/relay-openai.go |  5 ++++-
 service/relay.go                     | 20 ++++++++++++--------
 2 files changed, 16 insertions(+), 9 deletions(-)

diff --git a/relay/channel/openai/relay-openai.go b/relay/channel/openai/relay-openai.go
index b71fcce..b6418cf 100644
--- a/relay/channel/openai/relay-openai.go
+++ b/relay/channel/openai/relay-openai.go
@@ -53,7 +53,10 @@ func OpenaiStreamHandler(c *gin.Context, resp *http.Response, info *relaycommon.
 			}
 			data = data[6:]
 			if !strings.HasPrefix(data, "[DONE]") {
-				service.StringData(c, data)
+				err := service.StringData(c, data)
+				if err != nil {
+					common.LogError(c, "streaming error: "+err.Error())
+				}
 				streamItems = append(streamItems, data)
 			}
 		}
diff --git a/service/relay.go b/service/relay.go
index 22f9ce3..4f5ff8d 100644
--- a/service/relay.go
+++ b/service/relay.go
@@ -2,10 +2,10 @@ package service
 
 import (
 	"encoding/json"
+	"errors"
 	"fmt"
 	"github.com/gin-gonic/gin"
 	"one-api/common"
-	"strings"
 )
 
 func SetEventStreamHeaders(c *gin.Context) {
@@ -16,11 +16,16 @@ func SetEventStreamHeaders(c *gin.Context) {
 	c.Writer.Header().Set("X-Accel-Buffering", "no")
 }
 
-func StringData(c *gin.Context, str string) {
-	str = strings.TrimPrefix(str, "data: ")
-	str = strings.TrimSuffix(str, "\r")
+func StringData(c *gin.Context, str string) error {
+	//str = strings.TrimPrefix(str, "data: ")
+	//str = strings.TrimSuffix(str, "\r")
 	c.Render(-1, common.CustomEvent{Data: "data: " + str})
-	c.Writer.Flush()
+	if c.Writer != nil {
+		c.Writer.Flush()
+	} else {
+		return errors.New("writer is nil")
+	}
+	return nil
 }
 
 func ObjectData(c *gin.Context, object interface{}) error {
@@ -28,12 +33,11 @@ func ObjectData(c *gin.Context, object interface{}) error {
 	if err != nil {
 		return fmt.Errorf("error marshalling object: %w", err)
 	}
-	StringData(c, string(jsonData))
-	return nil
+	return StringData(c, string(jsonData))
 }
 
 func Done(c *gin.Context) {
-	StringData(c, "[DONE]")
+	_ = StringData(c, "[DONE]")
 }
 
 func GetResponseID(c *gin.Context) string {