101 lines
2.7 KiB
Go
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
|
||
|
|
}
|