diff --git a/server/gateway.go b/server/gateway.go index 39ea1d4..5ce4187 100644 --- a/server/gateway.go +++ b/server/gateway.go @@ -1,62 +1,19 @@ package server import ( - "fmt" - "io" "net/http" "net/http/httputil" "net/url" - "time" "git.magicany.cc/black1552/gf-common/log" - "github.com/gogf/gf/v2/net/gclient" "github.com/gogf/gf/v2/net/ghttp" - "github.com/gogf/gf/v2/os/gcache" "github.com/gogf/gf/v2/os/gctx" ) +// BuildRequest 反向代理请求到指定主机 +// 自动支持所有 HTTP 方法及 WebSocket 连接 func BuildRequest(r *ghttp.Request, host string) { - client := gclient.New() - // 构建目标URL而不是直接复制RequestURI - targetURL := host + r.URL.Path - if r.URL.RawQuery != "" { - targetURL += "?" + r.URL.RawQuery - } - // 复制请求头 - for key, values := range r.Header { - for _, value := range values { - client.SetHeader(key, value) - } - } - response, err := client.DoRequest(gctx.New(), r.Method, targetURL, r.GetBody()) - if err != nil { - log.Error(gctx.New(), "request error:", err) - panic(fmt.Sprintf("request error: %v", err)) - } - defer response.Body.Close() - // 读取响应体 - respBody, err := io.ReadAll(response.Body) - if err != nil { - log.Error(gctx.New(), "read response body error:", err) - panic(fmt.Sprintf("read response body error: %v", err)) - } - - // 复制响应头 - for key, values := range response.Header { - for _, value := range values { - r.Response.Header().Add(key, value) - } - } - gcache.Set(gctx.New(), "host", host, 1*time.Second) - // 设置响应状态码并写入响应体 - r.Response.Status = response.StatusCode - r.Response.Write(respBody) -} - -// ProxyWebSocket 处理 WebSocket 连接的代理 -func ProxyWebSocket(r *ghttp.Request, targetHost string) { - // 解析目标主机 URL - targetURL, err := url.Parse(targetHost) + targetURL, err := url.Parse(host) if err != nil { log.Error(gctx.New(), "parse target host error:", err) r.Response.WriteStatus(http.StatusInternalServerError) @@ -64,13 +21,25 @@ func ProxyWebSocket(r *ghttp.Request, targetHost string) { return } - // 创建反向代理 proxy := httputil.NewSingleHostReverseProxy(targetURL) - // 修改请求 URL,保留原始路径和查询参数 - r.URL.Scheme = targetURL.Scheme - r.URL.Host = targetURL.Host - log.Info(gctx.New(), r.GetBodyString()) - // 处理 WebSocket 连接 + // 自定义 Director 来设置目标 URL,保留原始请求的所有参数 + proxy.Director = func(req *http.Request) { + req.URL.Scheme = targetURL.Scheme + req.URL.Host = targetURL.Host + req.URL.Path = r.URL.Path + req.URL.RawQuery = r.URL.RawQuery + // 保留原始 Host 头 + req.Host = r.Host + } + + // 错误处理 + proxy.ErrorHandler = func(w http.ResponseWriter, req *http.Request, err error) { + log.Error(gctx.New(), "proxy error:", err) + w.WriteHeader(http.StatusBadGateway) + w.Write([]byte("Bad Gateway")) + } + + // ServeHTTP 会自动处理 WebSocket 升级和所有 HTTP 方法 proxy.ServeHTTP(r.Response.Writer, r.Request) }