Add fileDescForType, tests to be added

This commit is contained in:
Menghan Li 2016-05-16 16:30:17 -07:00
Родитель 5144c679d0
Коммит 15469fce29
2 изменённых файлов: 104 добавлений и 0 удалений

Просмотреть файл

@ -0,0 +1,89 @@
package main
import (
"bytes"
"compress/gzip"
"fmt"
"io/ioutil"
"reflect"
"sync"
"github.com/golang/protobuf/proto"
dpb "github.com/golang/protobuf/protoc-gen-go/descriptor"
)
type serverReflectionServer struct {
mu sync.Mutex
typeToNameMap map[reflect.Type]string
nameToTypeMap map[string]reflect.Type
typeToFileDescMap map[reflect.Type]*dpb.FileDescriptorProto
filenameToDescMap map[string]*dpb.FileDescriptorProto
}
func newServerReflectionServer() *serverReflectionServer {
return &serverReflectionServer{
typeToNameMap: make(map[reflect.Type]string),
nameToTypeMap: make(map[string]reflect.Type),
typeToFileDescMap: make(map[reflect.Type]*dpb.FileDescriptorProto),
filenameToDescMap: make(map[string]*dpb.FileDescriptorProto),
}
}
type protoMessage interface {
Descriptor() ([]byte, []int)
}
// TODO return an error rather than a bool
func (s *serverReflectionServer) fileDescForType(st reflect.Type) (*dpb.FileDescriptorProto, []int, error) {
m, ok := reflect.Zero(reflect.PtrTo(st)).Interface().(protoMessage)
if !ok {
// TODO better print content.
return nil, nil, fmt.Errorf("failed to create message from type: %v", st)
}
enc, idxs := m.Descriptor()
// TODO Check cache first.
// if fn, ok := s.typeToFilenameMap[st]; ok {
// if fd, ok := s.filenameToDescMap[fn]; ok {
// return fd, idxs, ok
// }
// }
fd, err := s.decodeFileDesc(enc)
if err != nil {
return nil, nil, err
}
// TODO Cache missed, add to cache.
// s.typeToFilenameMap[st] = fd.GetName()
return fd, idxs, nil
}
func (s *serverReflectionServer) decodeFileDesc(enc []byte) (*dpb.FileDescriptorProto, error) {
raw := decompress(enc)
if raw == nil {
return nil, fmt.Errorf("failed to decompress enc")
}
fd := new(dpb.FileDescriptorProto)
if err := proto.Unmarshal(raw, fd); err != nil {
return nil, fmt.Errorf("bad descriptor: %v", err)
}
// TODO If decodeFileDesc is called, it's the first time this file is seen.
// Add it to cache.
// s.filenameToDescMap[fd.GetName()] = fd
return fd, nil
}
func decompress(b []byte) []byte {
r, err := gzip.NewReader(bytes.NewReader(b))
if err != nil {
fmt.Printf("bad gzipped descriptor: %v\n", err)
return nil
}
out, err := ioutil.ReadAll(r)
if err != nil {
fmt.Printf("bad gzipped descriptor: %v\n", err)
return nil
}
return out
}

Просмотреть файл

@ -0,0 +1,15 @@
package main
import (
"reflect"
"sort"
"testing"
pb "google.golang.org/grpc/reflection/grpc_testing"
)
var (
s = newServerReflectionServer()
)
// TODO TestFileDescForType(t *testing.T)