servicechecker/pkg/protoset/protoset.go
Gorbunov Kirill Andreevich 2b7fc78872 #1-add protoset
2024-11-21 15:03:07 +03:00

101 lines
2.7 KiB
Go

package protoset
import (
"context"
"errors"
"fmt"
"sync"
protocodec "go.unistack.org/micro-codec-proto/v3"
"go.unistack.org/micro/v3/logger"
"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"
)
var errNotFound = errors.New("file descriptor not found")
type ProtoSet struct {
mu sync.Mutex
files map[string]*protoregistry.Files
log logger.Logger
}
func NewProtoSet(log logger.Logger) *ProtoSet {
return &ProtoSet{
mu: sync.Mutex{},
files: make(map[string]*protoregistry.Files, 0),
log: log,
}
}
func (p *ProtoSet) GetMessage(ctx context.Context, addr, pkg, svc, mth string) (protoreflect.Message, protoreflect.Message, error) {
p.log.Debug(ctx, "start of GetMessage")
if addr == "" || svc == "" || mth == "" {
p.log.Error(ctx, "protoset: empty addr or service param")
return nil, nil, errors.New("addr or service name is empty")
}
p.mu.Lock()
pfile, ok := p.files[addr+"|"+svc]
p.mu.Unlock()
if !ok || pfile == nil {
p.log.Error(ctx, "protoset: file desc not found")
return nil, nil, errNotFound
}
pdesc, err := pfile.FindDescriptorByName(protoreflect.FullName(pkg + "." + svc))
if err != nil {
p.log.Error(ctx, "failed to find service "+pkg+"."+svc)
return nil, nil, err
}
sdesc, ok := pdesc.(protoreflect.ServiceDescriptor)
if !ok {
err = fmt.Errorf("failed to find service " + pkg + "." + svc)
p.log.Error(ctx, "unable to find service in protoset", err)
return nil, nil, err
}
mdesc := sdesc.Methods().ByName(protoreflect.Name(mth))
if mdesc == nil {
err = fmt.Errorf("unknown method " + mth)
p.log.Error(ctx, "failed to find method", err)
return nil, nil, err
}
req := dynamicpb.NewMessageType(mdesc.Input()).New()
rsp := dynamicpb.NewMessageType(mdesc.Output()).New()
return req, rsp, nil
}
func (p *ProtoSet) AddProtoset(ctx context.Context, addr, svc string, data []byte) error {
p.log.Debug(ctx, "start of AddProtoset")
fdset := &descriptorpb.FileDescriptorSet{}
if err := protocodec.NewCodec().Unmarshal(data, fdset); err != nil {
p.log.Error(ctx, "failed to unmarshal protoset file", err)
return err
}
pfileoptions := protodesc.FileOptions{AllowUnresolvable: true}
pfiles, err := pfileoptions.NewFiles(fdset)
if err != nil {
p.log.Error(ctx, "failed to use protoset file", err)
return err
}
p.mu.Lock()
p.files[addr+"|"+svc] = pfiles
p.mu.Unlock()
return nil
}
func (p *ProtoSet) AddReflection(ctx context.Context, service string, addr string) error {
p.log.Debug(ctx, "start of AddReflection")
return nil
}