Кирилл Горбунов
ce57938ec2
Co-authored-by: Gorbunov Kirill Andreevich <kgorbunov@mtsbank.ru> Reviewed-on: #5 Co-authored-by: Кирилл Горбунов <kirya_gorbunov_2015@mail.ru> Co-committed-by: Кирилл Горбунов <kirya_gorbunov_2015@mail.ru>
85 lines
2.2 KiB
Go
85 lines
2.2 KiB
Go
package protoset
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"fmt"
|
|
"sync"
|
|
|
|
protocodec "go.unistack.org/micro-codec-proto/v3"
|
|
"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
|
|
}
|
|
|
|
func NewProtoSet() *ProtoSet {
|
|
return &ProtoSet{
|
|
mu: sync.Mutex{},
|
|
files: make(map[string]*protoregistry.Files, 0),
|
|
}
|
|
}
|
|
|
|
func (p *ProtoSet) GetMessage(addr, pkg, svc, mth string) (protoreflect.Message, protoreflect.Message, error) {
|
|
if addr == "" || svc == "" || mth == "" || pkg == "" {
|
|
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 {
|
|
return nil, nil, errNotFound
|
|
}
|
|
|
|
pdesc, err := pfile.FindDescriptorByName(protoreflect.FullName(pkg + "." + svc))
|
|
if err != nil {
|
|
return nil, nil, fmt.Errorf("failed to find service %s.%s, err: %w", pkg, svc, err)
|
|
}
|
|
|
|
sdesc, ok := pdesc.(protoreflect.ServiceDescriptor)
|
|
if !ok {
|
|
return nil, nil, fmt.Errorf("failed to find service " + pkg + "." + svc)
|
|
}
|
|
|
|
mdesc := sdesc.Methods().ByName(protoreflect.Name(mth))
|
|
if mdesc == nil {
|
|
return nil, nil, fmt.Errorf("unknown method " + mth)
|
|
}
|
|
|
|
req := dynamicpb.NewMessageType(mdesc.Input()).New()
|
|
rsp := dynamicpb.NewMessageType(mdesc.Output()).New()
|
|
|
|
return req, rsp, nil
|
|
}
|
|
|
|
func (p *ProtoSet) AddProtoset(addr, svc string, data []byte) error {
|
|
fdset := &descriptorpb.FileDescriptorSet{}
|
|
if err := protocodec.NewCodec().Unmarshal(data, fdset); err != nil {
|
|
return fmt.Errorf("failed to unmarshal protoset file: %w", err)
|
|
}
|
|
|
|
pfileoptions := protodesc.FileOptions{AllowUnresolvable: true}
|
|
pfiles, err := pfileoptions.NewFiles(fdset)
|
|
if err != nil {
|
|
return fmt.Errorf("failed to use protoset file, err: %w", 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 {
|
|
return nil
|
|
}
|