90CNServer/server/internal/platform/rpc/grpc_transport.go

package rpc

import (
	"context"
	"encoding/json"
	"errors"
	"net"
	"strings"
	"sync"

	"longheng.io/server/internal/platform/runtimeguard"
	"longheng.io/server/pkg/contracts"

	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials"
	"google.golang.org/grpc/credentials/insecure"
)

// ErrTransportClosed 表示 gRPC transport 已关闭或未正确装配。
var ErrTransportClosed = errors.New("rpc transport is closed")

// ErrTransportCredReq 用于 errors.Is 判断生产强制 TLS/mTLS 时仍缺少 gRPC 传输凭据的配置错误。
var ErrTransportCredReq = errors.New("rpc transport credentials are required")

// GRPCTransportOptions 是 gRPC transport 的监听、节点表和安全凭据配置。
type GRPCTransportOptions struct {
	ListenAddress            string
	Environment              string
	Peers                    map[string]string
	ServerCredentials        credentials.TransportCredentials
	DialCredentials          credentials.TransportCredentials
	RequireTransportSecurity bool
	ServerOptions            []grpc.ServerOption
	DialOptions              []grpc.DialOption
}

// GRPCTransport 使用 gRPC 在不同节点之间传递 RPC 请求和响应。
type GRPCTransport struct {
	mu               sync.RWMutex
	listener         net.Listener
	server           *grpc.Server
	listenAddress    string
	peers            map[string]string
	requestHandlers  map[string]RequestHandler
	responseHandlers map[string]ResponseHandler
	conns            map[string]*grpc.ClientConn
	dialCredentials  credentials.TransportCredentials
	dialOptions      []grpc.DialOption
	closed           bool
}

// NewGRPCTransport 创建并启动一个 gRPC transport。
func NewGRPCTransport(options GRPCTransportOptions) (*GRPCTransport, error) {
	address := strings.TrimSpace(options.ListenAddress)
	if address == "" {
		address = "127.0.0.1:0"
	}
	listener, err := net.Listen("tcp", address)
	if err != nil {
		return nil, err
	}
	needSecureTransport := options.RequireTransportSecurity || runtimeguard.IsProductionEnvironment(options.Environment)
	if needSecureTransport && (options.ServerCredentials == nil || options.DialCredentials == nil) {
		_ = listener.Close()
		return nil, ErrTransportCredReq
	}
	serverOptions := []grpc.ServerOption{grpc.ForceServerCodec(rpcJSONCodec{})}
	if options.ServerCredentials != nil {
		serverOptions = append(serverOptions, grpc.Creds(options.ServerCredentials.Clone()))
	}
	serverOptions = append(serverOptions, options.ServerOptions...)
	transport := &GRPCTransport{
		listener:         listener,
		server:           grpc.NewServer(serverOptions...),
		listenAddress:    listener.Addr().String(),
		peers:            cloneStringMap(options.Peers),
		requestHandlers:  make(map[string]RequestHandler),
		responseHandlers: make(map[string]ResponseHandler),
		conns:            make(map[string]*grpc.ClientConn),
		dialCredentials:  cloneTransportCreds(options.DialCredentials),
		dialOptions:      append([]grpc.DialOption(nil), options.DialOptions...),
	}
	regGRPCWire(transport.server, transport)
	go func() {
		_ = transport.server.Serve(listener)
	}()
	return transport, nil
}

// ListenAddress 返回实际监听地址，端口为 0 时可用于发现系统分配端口。
func (t *GRPCTransport) ListenAddress() string {
	if t == nil {
		return ""
	}
	t.mu.RLock()
	defer t.mu.RUnlock()
	return t.listenAddress
}

// SetPeer 设置或删除目标节点到 gRPC 地址的映射。
func (t *GRPCTransport) SetPeer(nodeID, address string) {
	if t == nil {
		return
	}
	nodeID = strings.TrimSpace(nodeID)
	address = strings.TrimSpace(address)
	if nodeID == "" {
		return
	}
	t.mu.Lock()
	if t.peers == nil {
		t.peers = make(map[string]string)
	}
	if address == "" {
		delete(t.peers, nodeID)
	} else {
		t.peers[nodeID] = address
	}
	t.mu.Unlock()
}

// SubscribeRequests 注册当前节点的请求处理函数。
func (t *GRPCTransport) SubscribeRequests(_ context.Context, nodeID string, handler RequestHandler) (TransportSubscription, error) {
	nodeID = strings.TrimSpace(nodeID)
	if nodeID == "" {
		return nil, ErrInvalidNodeID
	}
	if handler == nil {
		return nil, ErrInvalidRoute
	}
	if t == nil {
		return nil, ErrTransportClosed
	}
	t.mu.Lock()
	defer t.mu.Unlock()
	if t.closed {
		return nil, ErrTransportClosed
	}
	t.requestHandlers[nodeID] = handler
	return transportSub(func() {
		t.mu.Lock()
		delete(t.requestHandlers, nodeID)
		t.mu.Unlock()
	}), nil
}

// SubscribeResponses 注册当前节点的响应处理函数。
func (t *GRPCTransport) SubscribeResponses(_ context.Context, nodeID string, handler ResponseHandler) (TransportSubscription, error) {
	nodeID = strings.TrimSpace(nodeID)
	if nodeID == "" {
		return nil, ErrInvalidNodeID
	}
	if handler == nil {
		return nil, ErrInvalidRoute
	}
	if t == nil {
		return nil, ErrTransportClosed
	}
	t.mu.Lock()
	defer t.mu.Unlock()
	if t.closed {
		return nil, ErrTransportClosed
	}
	t.responseHandlers[nodeID] = handler
	return transportSub(func() {
		t.mu.Lock()
		delete(t.responseHandlers, nodeID)
		t.mu.Unlock()
	}), nil
}

// PublishRequest 通过 gRPC 把请求发送到目标节点。
func (t *GRPCTransport) PublishRequest(ctx context.Context, targetNodeID string, request contracts.RPCRequest) error {
	request.TargetNodeID = strings.TrimSpace(targetNodeID)
	return t.invoke(ctx, request.TargetNodeID, "/longheng.rpc.Transport/Request", &request, &grpcAck{})
}

// PublishResponse 通过 gRPC 把响应发送到目标节点。
func (t *GRPCTransport) PublishResponse(ctx context.Context, targetNodeID string, response contracts.RPCResponse) error {
	response.TargetNodeID = strings.TrimSpace(targetNodeID)
	return t.invoke(ctx, response.TargetNodeID, "/longheng.rpc.Transport/Response", &response, &grpcAck{})
}

// Close 停止 gRPC server 并关闭所有已缓存连接。
func (t *GRPCTransport) Close() error {
	if t == nil {
		return nil
	}
	t.mu.Lock()
	if t.closed {
		t.mu.Unlock()
		return nil
	}
	t.closed = true
	server := t.server
	conns := t.conns
	t.conns = nil
	t.requestHandlers = nil
	t.responseHandlers = nil
	t.mu.Unlock()
	if server != nil {
		server.Stop()
	}
	var firstErr error
	for _, conn := range conns {
		if conn == nil {
			continue
		}
		if err := conn.Close(); err != nil && firstErr == nil {
			firstErr = err
		}
	}
	return firstErr
}

func (t *GRPCTransport) request(ctx context.Context, request *contracts.RPCRequest) (*grpcAck, error) {
	if request == nil {
		return nil, ErrInvalidRequestID
	}
	t.mu.RLock()
	handler := t.requestHandlers[strings.TrimSpace(request.TargetNodeID)]
	t.mu.RUnlock()
	if handler == nil {
		return nil, ErrNoTargetNodes
	}
	return &grpcAck{}, handler(ctx, *request)
}

func (t *GRPCTransport) response(ctx context.Context, response *contracts.RPCResponse) (*grpcAck, error) {
	if response == nil {
		return nil, ErrInvalidRequestID
	}
	t.mu.RLock()
	handler := t.responseHandlers[strings.TrimSpace(response.TargetNodeID)]
	t.mu.RUnlock()
	if handler == nil {
		return nil, ErrNoTargetNodes
	}
	return &grpcAck{}, handler(ctx, *response)
}

func (t *GRPCTransport) invoke(ctx context.Context, nodeID, method string, in any, out any) error {
	if ctx == nil {
		ctx = context.Background()
	}
	if strings.TrimSpace(nodeID) == "" {
		return ErrInvalidTarget
	}
	conn, err := t.clientConn(nodeID)
	if err != nil {
		return err
	}
	return conn.Invoke(ctx, method, in, out, grpc.ForceCodec(rpcJSONCodec{}))
}

func (t *GRPCTransport) clientConn(nodeID string) (*grpc.ClientConn, error) {
	if t == nil {
		return nil, ErrTransportClosed
	}
	t.mu.RLock()
	if t.closed {
		t.mu.RUnlock()
		return nil, ErrTransportClosed
	}
	address := strings.TrimSpace(t.peers[strings.TrimSpace(nodeID)])
	if address == "" {
		t.mu.RUnlock()
		return nil, ErrNoTargetNodes
	}
	if conn := t.conns[address]; conn != nil {
		t.mu.RUnlock()
		return conn, nil
	}
	t.mu.RUnlock()

	options := []grpc.DialOption{
		grpc.WithDefaultCallOptions(grpc.ForceCodec(rpcJSONCodec{})),
	}
	if t.dialCredentials != nil {
		options = append(options, grpc.WithTransportCredentials(t.dialCredentials.Clone()))
	} else {
		options = append(options, grpc.WithTransportCredentials(insecure.NewCredentials()))
	}
	t.mu.RLock()
	options = append(options, t.dialOptions...)
	t.mu.RUnlock()
	conn, err := grpc.NewClient(grpcTarget(address), options...)
	if err != nil {
		return nil, err
	}
	t.mu.Lock()
	if t.closed {
		t.mu.Unlock()
		_ = conn.Close()
		return nil, ErrTransportClosed
	}
	if existing := t.conns[address]; existing != nil {
		t.mu.Unlock()
		_ = conn.Close()
		return existing, nil
	}
	t.conns[address] = conn
	t.mu.Unlock()
	return conn, nil
}

func cloneTransportCreds(creds credentials.TransportCredentials) credentials.TransportCredentials {
	if creds == nil {
		return nil
	}
	return creds.Clone()
}

type transportSub func()

// Close 注销 transport 订阅。
func (s transportSub) Close() error {
	if s != nil {
		s()
	}
	return nil
}

type rpcJSONCodec struct{}

// Name 返回 gRPC codec 名称。
func (rpcJSONCodec) Name() string {
	return "json"
}

// Marshal 把 RPC 结构编码为 JSON 字节。
func (rpcJSONCodec) Marshal(v any) ([]byte, error) {
	return json.Marshal(v)
}

// Unmarshal 把 JSON 字节解码为 RPC 结构。
func (rpcJSONCodec) Unmarshal(data []byte, v any) error {
	return json.Unmarshal(data, v)
}

type grpcAck struct{}

type grpcTransportService interface {
	request(context.Context, *contracts.RPCRequest) (*grpcAck, error)
	response(context.Context, *contracts.RPCResponse) (*grpcAck, error)
}

var grpcServiceDesc = grpc.ServiceDesc{
	ServiceName: "longheng.rpc.Transport",
	HandlerType: (*grpcTransportService)(nil),
	Methods: []grpc.MethodDesc{
		{MethodName: "Request", Handler: grpcRequestHandler},
		{MethodName: "Response", Handler: grpcResponseHandler},
	},
}

func grpcRequestHandler(srv any, ctx context.Context, dec func(any) error, interceptor grpc.UnaryServerInterceptor) (any, error) {
	in := new(contracts.RPCRequest)
	if err := dec(in); err != nil {
		return nil, err
	}
	if interceptor == nil {
		return srv.(grpcTransportService).request(ctx, in)
	}
	info := &grpc.UnaryServerInfo{Server: srv, FullMethod: "/longheng.rpc.Transport/Request"}
	handler := func(ctx context.Context, req any) (any, error) {
		return srv.(grpcTransportService).request(ctx, req.(*contracts.RPCRequest))
	}
	return interceptor(ctx, in, info, handler)
}

func grpcResponseHandler(srv any, ctx context.Context, dec func(any) error, interceptor grpc.UnaryServerInterceptor) (any, error) {
	in := new(contracts.RPCResponse)
	if err := dec(in); err != nil {
		return nil, err
	}
	if interceptor == nil {
		return srv.(grpcTransportService).response(ctx, in)
	}
	info := &grpc.UnaryServerInfo{Server: srv, FullMethod: "/longheng.rpc.Transport/Response"}
	handler := func(ctx context.Context, req any) (any, error) {
		return srv.(grpcTransportService).response(ctx, req.(*contracts.RPCResponse))
	}
	return interceptor(ctx, in, info, handler)
}

func regGRPCWire(server *grpc.Server, transport *GRPCTransport) {
	server.RegisterService(&grpcServiceDesc, transport)
}

func grpcTarget(address string) string {
	address = strings.TrimSpace(address)
	if strings.Contains(address, "://") {
		return address
	}
	return "passthrough:///" + address
}

func cloneStringMap(values map[string]string) map[string]string {
	if len(values) == 0 {
		return nil
	}
	out := make(map[string]string, len(values))
	for key, value := range values {
		key = strings.TrimSpace(key)
		value = strings.TrimSpace(value)
		if key != "" && value != "" {
			out[key] = value
		}
	}
	if len(out) == 0 {
		return nil
	}
	return out
}