JCS-pub/client/internal/http/proxy/proxy.go

package proxy

import (
	"io"

	"github.com/gin-gonic/gin"
	"gitlink.org.cn/cloudream/common/utils/io2"
	"gitlink.org.cn/cloudream/jcs-pub/client/internal/cluster"
	"gitlink.org.cn/cloudream/jcs-pub/client/internal/http/types"
	"gitlink.org.cn/cloudream/jcs-pub/common/ecode"
	clirpc "gitlink.org.cn/cloudream/jcs-pub/common/pkgs/rpc/client"
)

type ClusterProxy struct {
	clster *cluster.Cluster
}

func NewClusterProxy(clster *cluster.Cluster) *ClusterProxy {
	return &ClusterProxy{
		clster: clster,
	}
}

func (p *ClusterProxy) Proxy(c *gin.Context) {
	if p.clster.IsMaster() {
		c.Next()
		return
	}

	mstCli := p.clster.MasterClient()
	defer mstCli.Release()

	preq := clirpc.HTTPProxyRequest{
		URI:    c.Request.RequestURI,
		Method: c.Request.Method,
	}

	for k, v := range c.Request.Header {
		for _, vv := range v {
			preq.Header = append(preq.Header, &clirpc.HeaderKV{
				Key:   k,
				Value: vv,
			})
		}
	}

	if c.Request.Body != nil {
		data, err := io.ReadAll(c.Request.Body)
		if err != nil {
			c.AbortWithStatusJSON(500, types.Failed(ecode.OperationFailed, "read body: %v", err))
			return
		}

		preq.Body = data
	}

	resp, err := mstCli.HTTPProxy(c.Request.Context(), &preq)
	if err != nil {
		c.AbortWithStatusJSON(500, types.Failed(ecode.OperationFailed, "proxy: %v", err))
		return
	}

	c.Status(int(resp.StatusCode))

	for _, h := range resp.Header {
		c.Header(h.Key, h.Value)
	}

	if len(resp.Body) > 0 {
		io2.WriteAll(c.Writer, resp.Body)
	}

	c.Abort()
}