backend/internal/transport/grpcserver/server.go

package grpcserver

import (
	"errors"
	"fmt"
	"net"
	"net/http"
	"strconv"
	"strings"

	"github.com/XingfenD/yoresee_doc/internal/config"
	"github.com/XingfenD/yoresee_doc/internal/repository"
	"github.com/XingfenD/yoresee_doc/internal/status"
	pb "github.com/XingfenD/yoresee_doc/pkg/gen/yoresee_doc/v1"
	"github.com/improbable-eng/grpc-web/go/grpcweb"
	"github.com/sirupsen/logrus"
	"google.golang.org/grpc"
)

func Start(grpcPort, grpcWebPort int, repos *repository.Repositories) (*grpc.Server, error) {
	allowUnauth := map[string]struct{}{
		"/yoresee_doc.v1.AuthService/Login":        {},
		"/yoresee_doc.v1.AuthService/Register":     {},
		"/yoresee_doc.v1.SystemService/Health":     {},
		"/yoresee_doc.v1.SystemService/SystemInfo": {},
	}

	grpcServer := grpc.NewServer(
		grpc.UnaryInterceptor(UnaryAuthInterceptor(allowUnauth)),
	)

	pb.RegisterAuthServiceServer(grpcServer, NewAuthServiceServer())
	pb.RegisterDocumentServiceServer(grpcServer, NewDocumentServiceServer())
	pb.RegisterKnowledgeBaseServiceServer(grpcServer, NewKnowledgeBaseServiceServer())
	pb.RegisterSystemServiceServer(grpcServer, NewSystemServiceServer())
	pb.RegisterMembershipServiceServer(grpcServer, NewMembershipServiceServer())
	pb.RegisterInvitationServiceServer(grpcServer, NewInvitationServiceServer(repos.User))

	grpcListener, err := net.Listen("tcp", fmt.Sprintf(":%d", grpcPort))
	if err != nil {
		return nil, status.GenErrWithCustomMsg(status.StatusServiceInternalError, "listen grpc port failed")
	}

	go func() {
		if serveErr := grpcServer.Serve(grpcListener); serveErr != nil {
			logrus.Errorf("grpc server exited with error: %v", serveErr)
		}
	}()

	grpcWebServer := grpcweb.WrapServer(grpcServer,
		grpcweb.WithOriginFunc(isOriginAllowed),
	)

	grpcWebHandler := grpcWebHTTPHandler(grpcWebServer)
	grpcWebAddr := fmt.Sprintf(":%d", grpcWebPort)
	go func() {
		if serveErr := http.ListenAndServe(grpcWebAddr, grpcWebHandler); serveErr != nil && !errors.Is(serveErr, http.ErrServerClosed) {
			logrus.Errorf("grpc-web server exited with error: %v", serveErr)
		}
	}()

	return grpcServer, nil
}

func grpcWebHTTPHandler(wrapped *grpcweb.WrappedGrpcServer) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		setGrpcWebCorsHeaders(w, r)
		if wrapped.IsGrpcWebRequest(r) || wrapped.IsGrpcWebSocketRequest(r) {
			wrapped.ServeHTTP(w, r)
			return
		}
		if r.Method == http.MethodOptions && wrapped.IsAcceptableGrpcCorsRequest(r) {
			w.WriteHeader(http.StatusOK)
			return
		}
		http.NotFound(w, r)
	})
}

func setGrpcWebCorsHeaders(w http.ResponseWriter, r *http.Request) {
	origin := strings.TrimSpace(r.Header.Get("Origin"))
	if allowedOrigin := resolveAllowedOrigin(origin); allowedOrigin != "" {
		w.Header().Set("Access-Control-Allow-Origin", allowedOrigin)
	}
	w.Header().Set("Access-Control-Allow-Credentials", strconv.FormatBool(corsAllowCredentials()))
	w.Header().Set("Access-Control-Allow-Headers", strings.Join(corsAllowedHeaders(), ","))
	w.Header().Set("Access-Control-Allow-Methods", strings.Join(corsAllowedMethods(), ","))
	w.Header().Set("Access-Control-Expose-Headers", "Grpc-Status,Grpc-Message,Grpc-Status-Details-Bin")
	if maxAge := corsMaxAge(); maxAge > 0 {
		w.Header().Set("Access-Control-Max-Age", strconv.Itoa(maxAge))
	}
}

func isOriginAllowed(origin string) bool {
	origin = strings.TrimSpace(origin)
	if origin == "" {
		return false
	}
	allowedOrigins := corsAllowedOrigins()
	if len(allowedOrigins) == 0 {
		return true
	}
	for _, allowed := range allowedOrigins {
		allowed = strings.TrimSpace(allowed)
		if allowed == "*" || strings.EqualFold(allowed, origin) {
			return true
		}
	}
	return false
}

func resolveAllowedOrigin(origin string) string {
	if origin == "" {
		if len(corsAllowedOrigins()) == 0 {
			return "*"
		}
		return ""
	}
	if isOriginAllowed(origin) {
		return origin
	}
	return ""
}

func corsAllowedOrigins() []string {
	if config.GlobalConfig == nil {
		return nil
	}
	return config.GlobalConfig.Backend.Security.CORS.AllowedOrigins
}

func corsAllowedMethods() []string {
	if config.GlobalConfig != nil && len(config.GlobalConfig.Backend.Security.CORS.AllowedMethods) > 0 {
		return config.GlobalConfig.Backend.Security.CORS.AllowedMethods
	}
	return []string{"POST", "GET", "OPTIONS"}
}

func corsAllowedHeaders() []string {
	if config.GlobalConfig != nil && len(config.GlobalConfig.Backend.Security.CORS.AllowedHeaders) > 0 {
		return config.GlobalConfig.Backend.Security.CORS.AllowedHeaders
	}
	return []string{"Accept", "Accept-Language", "Content-Type", "Content-Length", "X-Grpc-Web", "X-User-Agent", "Grpc-Timeout", "Authorization"}
}

func corsAllowCredentials() bool {
	if config.GlobalConfig != nil {
		return config.GlobalConfig.Backend.Security.CORS.AllowCredentials
	}
	return true
}

func corsMaxAge() int {
	if config.GlobalConfig != nil {
		return config.GlobalConfig.Backend.Security.CORS.MaxAge
	}
	return 0
}