full/all/internal/server/dynamic.go

package server

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net"
	"net/http"
	"strings"
	"sync"

	"github.com/gin-gonic/gin"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/credentials/insecure"
	"google.golang.org/grpc/metadata"
	reflectionv1 "google.golang.org/grpc/reflection/grpc_reflection_v1"
	"google.golang.org/grpc/status"
	"google.golang.org/protobuf/encoding/protojson"
	"google.golang.org/protobuf/proto"
	"google.golang.org/protobuf/reflect/protodesc"
	"google.golang.org/protobuf/reflect/protoreflect"
	"google.golang.org/protobuf/reflect/protoregistry"
	"google.golang.org/protobuf/types/descriptorpb"
	"google.golang.org/protobuf/types/dynamicpb"
)

const maxDynamicRPCBody = 4 << 20

type dynamicGateway struct {
	conn  *grpc.ClientConn
	allow map[string]struct{}
	mu    sync.RWMutex
	cache map[string]protoreflect.MethodDescriptor
}

type dynamicRPCResponse struct {
	Code    int32             `json:"code"`
	Message string            `json:"message"`
	Data    json.RawMessage   `json:"data,omitempty"`
	Details []json.RawMessage `json:"details,omitempty"`
}

func newDynamicGateway(grpcAddr string, allow []string) (*dynamicGateway, error) {
	conn, err := grpc.NewClient(reflectionTarget(grpcAddr), grpc.WithTransportCredentials(insecure.NewCredentials()))
	if err != nil {
		return nil, fmt.Errorf("create dynamic gRPC client: %w", err)
	}
	allowed := make(map[string]struct{}, len(allow))
	for _, item := range allow {
		if item = strings.TrimSpace(item); item != "" {
			allowed[item] = struct{}{}
		}
	}
	return &dynamicGateway{conn: conn, allow: allowed, cache: make(map[string]protoreflect.MethodDescriptor)}, nil
}

func reflectionTarget(addr string) string {
	host, port, err := net.SplitHostPort(addr)
	if err != nil {
		return addr
	}
	if host == "" || host == "0.0.0.0" || host == "::" {
		host = "127.0.0.1"
	}
	return net.JoinHostPort(host, port)
}

func (g *dynamicGateway) Close() error { return g.conn.Close() }

func (g *dynamicGateway) handle(c *gin.Context) {
	moduleName := strings.TrimSpace(c.Param("module"))
	serviceShortName := strings.TrimSpace(c.Param("service"))
	methodName := strings.TrimSpace(c.Param("method"))
	if moduleName == "" || serviceShortName == "" || methodName == "" {
		writeDynamicError(c, status.Error(codes.InvalidArgument, "path must be /rpc/{module}/{service}/{method}"))
		return
	}
	serviceName := moduleName + "." + serviceShortName
	fullMethod := serviceName + "." + methodName
	if !g.isAllowed(serviceName, fullMethod) {
		writeDynamicError(c, status.Error(codes.PermissionDenied, "dynamic RPC method is not allowed"))
		return
	}

	descriptor, err := g.resolveMethod(c.Request.Context(), serviceName, methodName)
	if err != nil {
		writeDynamicError(c, err)
		return
	}
	if descriptor.IsStreamingClient() || descriptor.IsStreamingServer() {
		writeDynamicError(c, status.Error(codes.Unimplemented, "streaming RPC methods are not supported"))
		return
	}

	body, err := io.ReadAll(io.LimitReader(c.Request.Body, maxDynamicRPCBody+1))
	if err != nil {
		writeDynamicError(c, status.Error(codes.InvalidArgument, "read request body: "+err.Error()))
		return
	}
	if len(body) > maxDynamicRPCBody {
		writeDynamicError(c, status.Error(codes.ResourceExhausted, "request body exceeds 4 MiB"))
		return
	}

	request := dynamicpb.NewMessage(descriptor.Input())
	if err := (protojson.UnmarshalOptions{DiscardUnknown: false}).Unmarshal(body, request); err != nil {
		writeDynamicError(c, status.Error(codes.InvalidArgument, "invalid protobuf JSON: "+err.Error()))
		return
	}
	response := dynamicpb.NewMessage(descriptor.Output())
	ctx := outgoingMetadata(c.Request)
	grpcMethod := "/" + serviceName + "/" + methodName
	if err := g.conn.Invoke(ctx, grpcMethod, request, response); err != nil {
		writeDynamicError(c, err)
		return
	}

	data, err := (protojson.MarshalOptions{UseProtoNames: false}).Marshal(response)
	if err != nil {
		writeDynamicError(c, status.Error(codes.Internal, "marshal protobuf response: "+err.Error()))
		return
	}
	c.JSON(http.StatusOK, dynamicRPCResponse{Code: int32(codes.OK), Message: codes.OK.String(), Data: data})
}

func (g *dynamicGateway) isAllowed(serviceName, fullMethod string) bool {
	for _, key := range []string{"*", serviceName, fullMethod} {
		if _, ok := g.allow[key]; ok {
			return true
		}
	}
	return false
}

func (g *dynamicGateway) resolveMethod(ctx context.Context, serviceName, methodName string) (protoreflect.MethodDescriptor, error) {
	cacheKey := serviceName + "." + methodName
	g.mu.RLock()
	method := g.cache[cacheKey]
	g.mu.RUnlock()
	if method != nil {
		return method, nil
	}

	stream, err := reflectionv1.NewServerReflectionClient(g.conn).ServerReflectionInfo(ctx)
	if err != nil {
		return nil, status.Error(codes.Unavailable, "open gRPC reflection stream: "+err.Error())
	}
	if err := stream.Send(&reflectionv1.ServerReflectionRequest{
		MessageRequest: &reflectionv1.ServerReflectionRequest_FileContainingSymbol{FileContainingSymbol: serviceName},
	}); err != nil {
		return nil, status.Error(codes.Unavailable, "query gRPC reflection: "+err.Error())
	}
	reflectionResponse, err := stream.Recv()
	if err != nil {
		return nil, status.Error(codes.Unavailable, "read gRPC reflection response: "+err.Error())
	}
	fileResponse := reflectionResponse.GetFileDescriptorResponse()
	if fileResponse == nil {
		if reflectionErr := reflectionResponse.GetErrorResponse(); reflectionErr != nil {
			return nil, status.Error(codes.Code(reflectionErr.ErrorCode), reflectionErr.ErrorMessage)
		}
		return nil, status.Error(codes.NotFound, "service descriptor not found")
	}

	set := &descriptorpb.FileDescriptorSet{}
	for _, encoded := range fileResponse.FileDescriptorProto {
		file := &descriptorpb.FileDescriptorProto{}
		if err := proto.Unmarshal(encoded, file); err != nil {
			return nil, status.Error(codes.Internal, "decode reflected descriptor: "+err.Error())
		}
		set.File = append(set.File, file)
	}
	files, err := protodesc.NewFiles(set)
	if err != nil {
		return nil, status.Error(codes.Internal, "build reflected descriptors: "+err.Error())
	}
	descriptor, err := files.FindDescriptorByName(protoreflect.FullName(serviceName))
	if err != nil {
		if err == protoregistry.NotFound {
			return nil, status.Error(codes.NotFound, "service not found")
		}
		return nil, status.Error(codes.Internal, "resolve service descriptor: "+err.Error())
	}
	service, ok := descriptor.(protoreflect.ServiceDescriptor)
	if !ok {
		return nil, status.Error(codes.NotFound, "symbol is not a gRPC service")
	}
	method = service.Methods().ByName(protoreflect.Name(methodName))
	if method == nil {
		return nil, status.Error(codes.NotFound, "method not found")
	}
	g.mu.Lock()
	g.cache[cacheKey] = method
	g.mu.Unlock()
	return method, nil
}

func outgoingMetadata(request *http.Request) context.Context {
	pairs := make([]string, 0)
	for name, values := range request.Header {
		lower := strings.ToLower(name)
		if lower != "authorization" && lower != "x-request-id" && !strings.HasPrefix(lower, "x-") {
			continue
		}
		for _, value := range values {
			pairs = append(pairs, lower, value)
		}
	}
	return metadata.NewOutgoingContext(request.Context(), metadata.Pairs(pairs...))
}

func writeDynamicError(c *gin.Context, err error) {
	grpcStatus := status.Convert(err)
	details := make([]json.RawMessage, 0, len(grpcStatus.Details()))
	for _, detail := range grpcStatus.Details() {
		message, ok := detail.(proto.Message)
		if !ok {
			continue
		}
		if encoded, marshalErr := protojson.Marshal(message); marshalErr == nil {
			details = append(details, encoded)
		}
	}
	c.JSON(http.StatusOK, dynamicRPCResponse{
		Code:    int32(grpcStatus.Code()),
		Message: grpcStatus.Message(),
		Details: details,
	})
}