fix: remove sensitive informations from original endpoint request

langgenius · Dec 15, 2024 · 53bd801 · 53bd801
1 parent 1974c68
commit 53bd801
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 2 deletions.
diff --git a/internal/core/plugin_daemon/endpoint_service.go b/internal/core/plugin_daemon/endpoint_service.go
@@ -69,6 +69,7 @@ func InvokeEndpoint(
 				for resp.Next() {
 					chunk, err := resp.Read()
 					if err != nil {
+						response.WriteError(err)
 						return
 					}
 

diff --git a/internal/service/endpoint.go b/internal/service/endpoint.go
@@ -48,10 +48,20 @@ func Endpoint(
 		return
 	}
 
-	var buffer bytes.Buffer
+	// replace with a new reader
 	req.Body = io.NopCloser(bytes.NewReader(body))
 	req.ContentLength = int64(len(body))
 	req.TransferEncoding = nil
+
+	// remove ip traces for security
+	req.Header.Del("X-Forwarded-For")
+	req.Header.Del("X-Real-IP")
+	req.Header.Del("X-Forwarded")
+	req.Header.Del("X-Original-Forwarded-For")
+	req.Header.Del("X-Original-Url")
+	req.Header.Del("X-Original-Host")
+
+	var buffer bytes.Buffer
 	err = req.Write(&buffer)
 	if err != nil {
 		ctx.JSON(500, exception.InternalServerError(err).ToResponse())
@@ -157,7 +167,8 @@ func Endpoint(
 		for response.Next() {
 			chunk, err := response.Read()
 			if err != nil {
-				ctx.JSON(500, exception.InternalServerError(err).ToResponse())
+				ctx.Writer.Write([]byte(err.Error()))
+				ctx.Writer.Flush()
 				return
 			}
 			ctx.Writer.Write(chunk)