full/pkgs/ecmall/internal/server/server.go

package server

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"net"
	"net/http"
	"time"

	"github.com/gin-gonic/gin"
	gwRuntime "github.com/grpc-ecosystem/grpc-gateway/v2/runtime"
	"google.golang.org/grpc"
	"google.golang.org/grpc/reflection"
)

type Server struct {
	GRPC    *grpc.Server
	Gateway *gwRuntime.ServeMux
	HTTP    *gin.Engine
	http    *http.Server
	dynamic *dynamicGateway
	auth    *authorization
}

func New(key string, expireSeconds int64, anonymous []string) (*Server, error) {
	auth, err := newAuthorization(key, expireSeconds, anonymous)
	if err != nil {
		return nil, err
	}
	grpcServer := grpc.NewServer(grpc.UnaryInterceptor(auth.unaryInterceptor))
	reflection.Register(grpcServer)
	engine := gin.New()
	engine.Use(gin.Logger(), gin.Recovery())
	return &Server{
		GRPC:    grpcServer,
		Gateway: gwRuntime.NewServeMux(gwRuntime.WithRoutingErrorHandler(routeMissHandler)),
		HTTP:    engine,
		auth:    auth,
	}, nil
}

// routeMissHandler 是 grpc-gateway 的路由错误处理器：只有"路由未命中"才会被标记，
// 随后仍按默认行为写出错误体。业务 handler 返回的 NotFound 走的是 errorHandler，
// 不会被标记，因而其响应能原样透出而不会被 Gin 的纯文本 404 覆盖。
func routeMissHandler(ctx context.Context, mux *gwRuntime.ServeMux, marshaler gwRuntime.Marshaler, w http.ResponseWriter, r *http.Request, httpStatus int) {
	if httpStatus == http.StatusNotFound {
		if recorder, ok := w.(routingMissRecorder); ok {
			recorder.markRoutingMiss()
		}
	}
	gwRuntime.DefaultRoutingErrorHandler(ctx, mux, marshaler, w, r, httpStatus)
}

func (s *Server) Start(grpcAddr, httpAddr string) error {
	grpcListener, err := net.Listen("tcp", grpcAddr)
	if err != nil {
		return fmt.Errorf("listen gRPC on %s: %w", grpcAddr, err)
	}
	httpListener, err := net.Listen("tcp", httpAddr)
	if err != nil {
		_ = grpcListener.Close()
		return fmt.Errorf("listen HTTP on %s: %w", httpAddr, err)
	}

	s.dynamic, err = newDynamicGateway(grpcAddr)
	if err != nil {
		_ = grpcListener.Close()
		_ = httpListener.Close()
		return err
	}
	s.HTTP.POST("/rpc/:module/:service/:method", s.dynamic.handle)

	s.http = &http.Server{
		Addr:              httpAddr,
		Handler:           s.auth.httpMiddleware(s.handler()),
		ReadHeaderTimeout: 10 * time.Second,
		IdleTimeout:       120 * time.Second,
		MaxHeaderBytes:    1 << 20,
	}
	fmt.Printf("ecmall gRPC services listening on %s\n", grpcAddr)
	fmt.Printf("ecmall HTTP services listening on %s\n", httpAddr)

	errCh := make(chan error, 2)
	go func() { errCh <- s.GRPC.Serve(grpcListener) }()
	go func() { errCh <- s.http.Serve(httpListener) }()
	serveErr := <-errCh
	if errors.Is(serveErr, grpc.ErrServerStopped) || errors.Is(serveErr, http.ErrServerClosed) {
		return http.ErrServerClosed
	}
	return serveErr
}

// handler 先让 grpc-gateway 处理请求，仅当网关确实未命中路由时才回退到 Gin
// （动态 RPC 与原生的 /rest 路由都注册在 Gin 上）。业务返回的 404/NotFound
// 属于网关的正常响应，必须原样透出。
func (s *Server) handler() http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		recorder := newBufferedResponse()
		s.Gateway.ServeHTTP(recorder, r)
		if !recorder.routingMiss {
			recorder.flush(w)
			return
		}
		s.HTTP.ServeHTTP(w, r)
	})
}

func (s *Server) Stop(ctx context.Context) error {
	stopped := make(chan struct{})
	go func() {
		s.GRPC.GracefulStop()
		close(stopped)
	}()
	select {
	case <-stopped:
	case <-ctx.Done():
		s.GRPC.Stop()
	}
	if s.dynamic != nil {
		_ = s.dynamic.Close()
	}
	if s.http == nil {
		return nil
	}
	return s.http.Shutdown(ctx)
}

// routingMissRecorder 由 bufferedResponse 实现，用于区分"网关未命中路由"与业务 NotFound。
type routingMissRecorder interface {
	markRoutingMiss()
}

type bufferedResponse struct {
	header http.Header
	body   bytes.Buffer
	status int
	// routingMiss 表示该 404 来自 grpc-gateway 的路由错误处理，而非业务响应。
	routingMiss bool
}

func newBufferedResponse() *bufferedResponse {
	return &bufferedResponse{header: make(http.Header), status: http.StatusOK}
}

func (r *bufferedResponse) Header() http.Header            { return r.header }
func (r *bufferedResponse) WriteHeader(status int)         { r.status = status }
func (r *bufferedResponse) Write(data []byte) (int, error) { return r.body.Write(data) }
func (r *bufferedResponse) markRoutingMiss()               { r.routingMiss = true }
func (r *bufferedResponse) flush(w http.ResponseWriter) {
	for key, values := range r.header {
		for _, value := range values {
			w.Header().Add(key, value)
		}
	}
	w.WriteHeader(r.status)
	_, _ = w.Write(r.body.Bytes())
}