Browse Source

feat: add X-Original-Host header handling in endpoint requests and improve request redirection URL construction

Yeuoly 6 months ago
parent
commit
eed91f2359

+ 6 - 1
internal/cluster/redirect.go

@@ -22,10 +22,15 @@ func (c *Cluster) RedirectRequest(
 
 	ip := ips[0]
 
+	url := "http://" + ip.fullAddress() + request.URL.Path
+	if request.URL.RawQuery != "" {
+		url += "?" + request.URL.RawQuery
+	}
+
 	// create a new request
 	redirectedRequest, err := http.NewRequest(
 		request.Method,
-		"http://"+ip.fullAddress()+request.URL.Path,
+		url,
 		request.Body,
 	)
 

+ 6 - 0
internal/server/endpoint.go

@@ -13,6 +13,7 @@ import (
 	"github.com/langgenius/dify-plugin-daemon/internal/types/exception"
 	"github.com/langgenius/dify-plugin-daemon/internal/types/models"
 	"github.com/langgenius/dify-plugin-daemon/internal/utils/log"
+	"github.com/langgenius/dify-plugin-daemon/pkg/entities/endpoint_entities"
 	"github.com/langgenius/dify-plugin-daemon/pkg/entities/plugin_entities"
 )
 
@@ -28,6 +29,11 @@ func (app *App) Endpoint(config *app.Config) func(c *gin.Context) {
 		hookId := c.Param("hook_id")
 		path := c.Param("path")
 
+		// set X-Original-Host
+		if c.Request.Header.Get(endpoint_entities.HeaderXOriginalHost) == "" {
+			c.Request.Header.Set(endpoint_entities.HeaderXOriginalHost, c.Request.Host)
+		}
+
 		if app.endpointHandler != nil {
 			app.endpointHandler(c, hookId, time.Duration(config.PluginMaxExecutionTimeout)*time.Second, path)
 		} else {

+ 7 - 0
internal/service/endpoint.go

@@ -24,6 +24,7 @@ import (
 	"github.com/langgenius/dify-plugin-daemon/internal/utils/encryption"
 	"github.com/langgenius/dify-plugin-daemon/internal/utils/routine"
 	"github.com/langgenius/dify-plugin-daemon/pkg/entities"
+	"github.com/langgenius/dify-plugin-daemon/pkg/entities/endpoint_entities"
 	"github.com/langgenius/dify-plugin-daemon/pkg/entities/plugin_entities"
 	"github.com/langgenius/dify-plugin-daemon/pkg/entities/requests"
 )
@@ -57,6 +58,12 @@ func copyRequest(req *http.Request, hookId string, path string) (*bytes.Buffer,
 	newReq.Header.Del("X-Original-Url")
 	newReq.Header.Del("X-Original-Host")
 
+	// check if X-Original-Host is set
+	if originalHost := req.Header.Get(endpoint_entities.HeaderXOriginalHost); originalHost != "" {
+		// replace host with original host
+		newReq.Host = originalHost
+	}
+
 	// setup hook id to request
 	newReq.Header.Set("Dify-Hook-Id", hookId)
 	// check if Dify-Hook-Url is set

+ 4 - 0
pkg/entities/endpoint_entities/endpoint.go

@@ -5,3 +5,7 @@ type EndpointResponseChunk struct {
 	Headers map[string]string `json:"headers" validate:"omitempty"`
 	Result  *string           `json:"result" validate:"omitempty"`
 }
+
+const (
+	HeaderXOriginalHost = "X-Original-Host"
+)