Golang中如何设置gRPC自定义处理程序

在 gRPC 中设置自定义处理程序需要实现 grpc.ServiceRegistrar 接口。以下是使用 chi 路由器作为 gRPC 网关的示例：

import (
    "net/http"
    
    "github.com/go-chi/chi/v5"
    "github.com/grpc-ecosystem/grpc-gateway/v2/runtime"
    "google.golang.org/grpc"
)

// 自定义 gRPC 网关处理器
type CustomGRPCHandler struct {
    chiRouter chi.Router
    grpcServer *grpc.Server
}

func NewCustomGRPCHandler() *CustomGRPCHandler {
    r := chi.NewRouter()
    
    // 添加中间件
    r.Use(cors.Handler(cors.Options{
        AllowedOrigins:   []string{"https://*", "http://*"},
        AllowedMethods:   []string{"GET", "POST", "PUT", "DELETE", "OPTIONS"},
        AllowedHeaders:   []string{"Accept", "Authorization", "Content-Type", "X-CSRF-Token"},
        ExposedHeaders:   []string{"Link"},
        AllowCredentials: true,
        MaxAge:           300,
    }))
    r.Use(middleware.Heartbeat("/ping"))
    
    return &CustomGRPCHandler{
        chiRouter: r,
        grpcServer: grpc.NewServer(),
    }
}

// 实现 gRPC 服务注册
func (h *CustomGRPCHandler) RegisterService(desc *grpc.ServiceDesc, impl interface{}) {
    h.grpcServer.RegisterService(desc, impl)
}

// 添加 HTTP 路由
func (h *CustomGRPCHandler) AddHTTPRoute(method, path string, handler http.HandlerFunc) {
    switch method {
    case "GET":
        h.chiRouter.Get(path, handler)
    case "POST":
        h.chiRouter.Post(path, handler)
    case "PUT":
        h.chiRouter.Put(path, handler)
    case "DELETE":
        h.chiRouter.Delete(path, handler)
    }
}

// 启动服务
func (h *CustomGRPCHandler) Serve(addr string) error {
    // 创建 gRPC 网关
    gwmux := runtime.NewServeMux()
    
    // 注册 gRPC 网关端点
    h.chiRouter.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
        // 这里可以添加 gRPC 转 HTTP 的逻辑
        gwmux.ServeHTTP(w, r)
    })
    
    // 添加你的自定义路由
    h.chiRouter.Get("/", app.RequestCome)
    h.chiRouter.Post("/authenticate", app.Authenticate)
    
    // 创建 HTTP 服务器
    httpServer := &http.Server{
        Addr:    addr,
        Handler: h.chiRouter,
    }
    
    return httpServer.ListenAndServe()
}

使用示例：

func main() {
    handler := NewCustomGRPCHandler()
    
    // 注册 gRPC 服务
    pb.RegisterYourServiceServer(handler.grpcServer, &yourServiceImpl{})
    
    // 添加 HTTP 路由
    handler.AddHTTPRoute("GET", "/api/users", func(w http.ResponseWriter, r *http.Request) {
        w.Write([]byte("HTTP Handler"))
    })
    
    // 启动服务
    if err := handler.Serve(":8080"); err != nil {
        log.Fatal(err)
    }
}

对于纯 gRPC 自定义处理器：

import (
    "net"
    
    "google.golang.org/grpc"
    "google.golang.org/grpc/codes"
    "google.golang.org/grpc/status"
)

// 自定义流处理器
type CustomStreamHandler struct{}

func (h *CustomStreamHandler) HandleStream(srv interface{}, stream grpc.ServerStream) error {
    // 自定义流处理逻辑
    return status.Errorf(codes.Unimplemented, "method not implemented")
}

// 创建自定义 gRPC 服务器选项
func WithCustomHandler() grpc.ServerOption {
    return grpc.CustomCodec(&customCodec{})
}

// 自定义编解码器
type customCodec struct{}

func (c *customCodec) Marshal(v interface{}) ([]byte, error) {
    // 自定义序列化逻辑
    return []byte{}, nil
}

func (c *customCodec) Unmarshal(data []byte, v interface{}) error {
    // 自定义反序列化逻辑
    return nil
}

func (c *customCodec) Name() string {
    return "custom"
}

这些示例展示了如何在 gRPC 中集成自定义处理逻辑。