sig-pub/internal/gateway/gateway.go

package gateway

import (
	"context"
	"encoding/json"
	"io"
	"net/http"
	"net/http/httputil"
	"net/url"
	"runtime/debug"
	"sig-pub/pkg/grpc/generic"
	"sig-pub/pkg/grpc/session"
	"sig-pub/pkg/resp"
	"sig-pub/pkg/utils/strs"
	"sig-pub/pkg/zlog"
	"strings"
	"time"

	"github.com/gin-gonic/gin"
	"google.golang.org/grpc"
	"google.golang.org/protobuf/encoding/protojson"
)

type GateServer struct {
	engine                   *gin.Engine
	sigServerURL             *url.URL
	grpcGenericClientFactory *generic.GrpcGenericClientFactory
}

func NewGateServer(
	sigServerURL *url.URL,
	grpcGenericClientFactory *generic.GrpcGenericClientFactory,
) *GateServer {
	return &GateServer{
		sigServerURL:             sigServerURL,
		grpcGenericClientFactory: grpcGenericClientFactory,
	}
}

func (s *GateServer) Init() (err error) {
	s.initGinServer()
	return
}

func (s *GateServer) initGinServer() {
	s.engine = gin.Default()
	s.engine.Use(func(c *gin.Context) {
		// global recover
		defer func() {
			if r := recover(); r != nil {
				zlog.Error("http server recover error:", r)
				debug.PrintStack()
				c.JSON(http.StatusInternalServerError, resp.Error("server error"))
				c.Abort()
			}
		}()
		c.Next()
	})

	routerGroup := s.engine.Group("/api")
	// routerGroup.Any("/sig", s.handleSigServerProxy(s.sigServerURL))
	routerGroup.Group("/sig").Any("/", s.handleSigServerProxy(s.sigServerURL))
	routerGroup.POST("/v1/:svr/:method", s.handleGrpcGenericCall)
}

func (s *GateServer) Run(addr string) (err error) {
	// fasthttp.ListenAndServe(addr, func(ctx *fasthttp.RequestCtx) {
	// 	ctx.Path()
	// 	ctx.Method()
	// })
	return s.engine.Run(addr)
}

// handleSigServerProxy 代理sig http server请求
func (s *GateServer) handleSigServerProxy(sigServerURL *url.URL) gin.HandlerFunc {
	// 创建反向代理
	proxy := httputil.NewSingleHostReverseProxy(sigServerURL)
	// 修改请求头、Host 等
	proxy.Director = func(req *http.Request) {
		req.URL.Scheme = sigServerURL.Scheme
		req.URL.Host = sigServerURL.Host
		req.Host = sigServerURL.Host
	}
	return func(c *gin.Context) {
		proxy.ServeHTTP(c.Writer, c.Request)
	}
}

// handleGrpcGenericCall 代理各个grpc服务请求
func (s *GateServer) handleGrpcGenericCall(c *gin.Context) {
	svr := strs.UpperInitialLetter(c.Param("svr"))
	method := strs.UpperInitialLetter(c.Param("method"))

	if svr == "" || method == "" {
		c.JSON(http.StatusBadRequest, resp.Error("service not found"))
		return
	}
	if !strings.HasSuffix(svr, "Service") {
		svr += "Service"
	}
	// todo service white list

	// get request body
	jsonBody, err := io.ReadAll(c.Request.Body)
	if err != nil {
		c.JSON(http.StatusBadRequest, resp.Error("parse request body error: "+err.Error()))
		return
	}

	ctx1, cancel := context.WithTimeout(context.Background(), time.Second*5)
	defer cancel()
	grpcGenericClient, err := s.grpcGenericClientFactory.GetClient(ctx1, svr)
	if err != nil {
		zlog.Error(err)
		c.JSON(http.StatusForbidden, resp.Error(err.Error()))
		return
	}

	// put session
	ctx := context.Background()
	ctx = session.PutSubject(ctx, session.NewRpcSubject("123456"))
	ctx, cancel = context.WithTimeout(ctx, time.Second*20)
	defer cancel()

	// todo config call options
	var opts []grpc.CallOption
	if svr == "ExchangeService" && method == "HistoryKline" {
		opts = append(opts, grpc.UseCompressor("snappy"))
	}
	// generic call with json
	rsp, err := grpcGenericClient.InvokeUnaryJsonBytes(ctx, method, jsonBody, opts...)
	if err != nil {
		if err == generic.ErrorMethodNotExists {
			c.JSON(http.StatusNotFound, resp.Error(err.Error()))
			return
		}
		c.JSON(http.StatusInternalServerError, resp.Error(err.Error()))
		return
	}

	// encode response
	bytes, err := protojson.MarshalOptions{
		UseProtoNames:   false, // false:lowerCamelCase, true:snake_case
		EmitUnpopulated: false, // 是否包含默认值
	}.Marshal(rsp)
	if err != nil {
		c.JSON(http.StatusInternalServerError, resp.Error(err.Error()))
		return
	}

	r := json.RawMessage(bytes)
	c.JSON(http.StatusOK, resp.Success(r))
}