From 15469fce29616438e53b19d960c674ddbb251731 Mon Sep 17 00:00:00 2001
From: Menghan Li <menghanl@google.com>
Date: Mon, 16 May 2016 16:30:17 -0700
Subject: [PATCH] Add fileDescForType, tests to be added

---
 reflection/serverreflection.go      | 89 +++++++++++++++++++++++++++++
 reflection/serverreflection_test.go | 15 +++++
 2 files changed, 104 insertions(+)
 create mode 100644 reflection/serverreflection.go
 create mode 100644 reflection/serverreflection_test.go

diff --git a/reflection/serverreflection.go b/reflection/serverreflection.go
new file mode 100644
index 000000000..4f434531b
--- /dev/null
+++ b/reflection/serverreflection.go
@@ -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
+}
diff --git a/reflection/serverreflection_test.go b/reflection/serverreflection_test.go
new file mode 100644
index 000000000..47afb1a50
--- /dev/null
+++ b/reflection/serverreflection_test.go
@@ -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)