gateway/core/util/http.go

package util

import (
	log "app.yhyue.com/moapp/jylog"
	"bufio"
	"bytes"
	"compress/flate"
	"compress/gzip"
	"encoding/json"
	"fmt"
	"github.com/gogf/gf/v2/frame/g"
	"github.com/gogf/gf/v2/net/ghttp"
	"io"
	"net/http"
	"net/url"
	"strings"
)

func CopyReqAndFormData(req *http.Request, copy ...bool) {
	//获取请求体内容
	bodyBytes, _ := io.ReadAll(req.Body)
	_ = req.Body.Close()
	//将请求体内容重新写入请求体
	req.Body = io.NopCloser(bytes.NewBuffer(bodyBytes))
	if strings.Index(req.Header.Get("Content-Type"), "json") > -1 {
		_ = ParseJson(req)
	} else {
		_ = req.ParseForm() //格式化请求内容
	}
	//将请求体内容重新写入请求体
	if len(copy) == 0 || copy[0] == false {
		req.Body = io.NopCloser(bytes.NewBuffer(bodyBytes))
	}
}

func ParseJson(req *http.Request) error {
	bytes, err := io.ReadAll(req.Body)
	if err != nil {
		return err
	}
	jsonMap := map[string]interface{}{}
	err = json.Unmarshal(bytes, &jsonMap)
	if err != nil {
		return err
	}
	req.Form = make(url.Values)
	for k, v := range jsonMap {
		req.Form.Set(k, fmt.Sprint(v))
	}
	return nil
}

// ChangeGResponse 修改相应内容
// resp 相应体
// fn 修改相应内容方法
func ChangeGResponse(r *ghttp.Request, fn func([]byte) error) (err error) {
	encoding, content := r.Response.Header().Get("Content-Encoding"), []byte{}
	if encoding == "gzip" { //处理解压
		body, err := gzip.NewReader(r.Body)
		if err != nil {
			log.WithContext(r.Context()).Error(r.Context(), "ChangeResponse gzip err", err)
			return err
		}
		content, err = io.ReadAll(body)
	} else if encoding == "deflate" {
		content, err = io.ReadAll(flate.NewReader(r.Body))
	} else { //无设置压缩
		content, err = io.ReadAll(r.Body)
	}
	if err != nil {
		return err
	}

	// 修改相应内容
	err = fn(content)
	if err != nil {
		return err
	}

	var zBuf bytes.Buffer
	var outputWriter io.Writer
	if encoding == "gzip" || encoding == "deflate" { //压缩
		if encoding == "gzip" {
			outputWriter, _ = gzip.NewWriterLevel(&zBuf, gzip.BestSpeed)
		} else {
			outputWriter, _ = flate.NewWriter(&zBuf, flate.BestSpeed)
		}
	} else {
		outputWriter = bufio.NewWriter(&zBuf)
		r.Response.Header().Set("content-length", fmt.Sprint(len(content)))
	}
	if _, err = outputWriter.Write(content); err != nil {
		return err
	}
	switch outputWriter.(type) {
	case *gzip.Writer:
		err = outputWriter.(*gzip.Writer).Close()
	case *flate.Writer:
		err = outputWriter.(*flate.Writer).Close()
	case *bufio.Writer:
		err = outputWriter.(*bufio.Writer).Flush()
	}
	if err != nil {
		return err
	}
	_, err = fmt.Fprint(r.Response.ResponseWriter, zBuf.Bytes())
	return nil
}

func SingleJoiningSlash(a, b string) string {
	aslash := strings.HasSuffix(a, "/")
	bslash := strings.HasPrefix(b, "/")
	switch {
	case aslash && bslash:
		return a + b[1:]
	case !aslash && !bslash:
		return a + "/" + b
	}
	return a + b
}

func ChangeResponse(resp *http.Response, fn func([]byte) ([]byte, error)) (err error) {
	var (
		encoding = resp.Header.Get("Content-Encoding")
		content  []byte
		ctx      = resp.Request.Context()
	)
	if encoding == "gzip" { //处理解压
		body, err := gzip.NewReader(resp.Body)
		if err != nil {
			log.WithContext(ctx).Error(ctx, "ChangeResponse gzip err", err)
			return err
		}
		content, err = io.ReadAll(body)
	} else if encoding == "deflate" {
		content, err = io.ReadAll(flate.NewReader(resp.Body))
	} else { //无设置压缩
		content, err = io.ReadAll(resp.Body)
	}
	if err != nil {
		return err
	}

	// 修改相应内容
	var newContent []byte
	newContent, err = fn(content)
	if err != nil {
		g.Log().Errorf(ctx, "changeRes func %v err %v", fn, err)
		newContent = content
	}

	var zBuf bytes.Buffer
	var outputWriter io.Writer
	if encoding == "gzip" || encoding == "deflate" { //压缩
		if encoding == "gzip" {
			outputWriter, _ = gzip.NewWriterLevel(&zBuf, gzip.BestSpeed)
		} else {
			outputWriter, _ = flate.NewWriter(&zBuf, flate.BestSpeed)
		}
	} else {
		outputWriter = bufio.NewWriter(&zBuf)
		resp.Header.Set("content-length", fmt.Sprint(len(newContent)))
	}
	if _, err = outputWriter.Write(newContent); err != nil {
		return err
	}
	switch outputWriter.(type) {
	case *gzip.Writer:
		err = outputWriter.(*gzip.Writer).Close()
	case *flate.Writer:
		err = outputWriter.(*flate.Writer).Close()
	case *bufio.Writer:
		err = outputWriter.(*bufio.Writer).Flush()
	}
	if err != nil {
		return err
	}
	resp.Body = io.NopCloser(bytes.NewReader(zBuf.Bytes()))
	return nil
}