From d1a4237cb21341b8f2d41ad972ca1860df2d6f22 Mon Sep 17 00:00:00 2001
From: David Goitia <david@claviux.com>
Date: Wed, 21 Oct 2020 21:57:51 +0200
Subject: [PATCH] allow protoset file

---
 core/resource.go      | 46 ++++++++++++++++++++++++++++++++-----------
 core/resource_test.go |  2 +-
 index/js/proto.js     |  2 +-
 3 files changed, 36 insertions(+), 14 deletions(-)

diff --git a/core/resource.go b/core/resource.go
index 04819d1..dfc4b8b 100644
--- a/core/resource.go
+++ b/core/resource.go
@@ -32,6 +32,7 @@ type Resource struct {
 	descSource grpcurl.DescriptorSource
 	refClient  *grpcreflect.Client
 	protos     []Proto
+	protosets  []Proto
 
 	headers []string
 	md      metadata.MD
@@ -44,21 +45,31 @@ func (r *Resource) openDescriptor() error {
 	r.refClient = grpcreflect.NewClient(refCtx, reflectpb.NewServerReflectionClient(r.clientConn))
 
 	// if no protos available use server reflection
-	if r.protos == nil {
+	if r.protos == nil && r.protosets == nil {
 		r.descSource = grpcurl.DescriptorSourceFromServer(ctx, r.refClient)
 		return nil
 	}
 
 	protoPath := filepath.Join(BasePath, r.clientConn.Target())
 
-	// make list of protos name to be used as descriptor
-	protos := make([]string, 0, len(r.protos))
-	for _, proto := range r.protos {
-		protos = append(protos, proto.Name)
-	}
-
 	var err error
-	r.descSource, err = grpcurl.DescriptorSourceFromProtoFiles([]string{protoPath}, protos...)
+	if len(r.protosets) > 0 {
+		// make list of protos name to be used as descriptor
+		protos := make([]string, 0, len(r.protosets))
+		for _, proto := range r.protosets {
+			protos = append(protos, filepath.Join(protoPath, proto.Name))
+		}
+
+		r.descSource, err = grpcurl.DescriptorSourceFromProtoSets(protos...)
+	} else {
+		// make list of protos name to be used as descriptor
+		protos := make([]string, 0, len(r.protos))
+		for _, proto := range r.protos {
+			protos = append(protos, proto.Name)
+		}
+
+		r.descSource, err = grpcurl.DescriptorSourceFromProtoFiles([]string{protoPath}, protos...)
+	}
 	return err
 }
 
@@ -297,16 +308,27 @@ func (r *Resource) AddProtos(protos []Proto) error {
 		return err
 	}
 
+	var protoSlice, protosetSlice []Proto
 	for _, proto := range protos {
-		err := ioutil.WriteFile(filepath.Join(protoPath, "/", proto.Name),
-			prepareImport(proto.Content),
-			0777)
+		var err error
+		if strings.HasSuffix(proto.Name, ".protoset") {
+			protosetSlice = append(protosetSlice, proto)
+			err = ioutil.WriteFile(filepath.Join(protoPath, "/", proto.Name),
+				proto.Content,
+				0777)
+		} else {
+			protoSlice = append(protoSlice, proto)
+			err = ioutil.WriteFile(filepath.Join(protoPath, "/", proto.Name),
+				prepareImport(proto.Content),
+				0777)
+		}
 		if err != nil {
 			return err
 		}
 	}
 
-	r.protos = protos
+	r.protos = protoSlice
+	r.protosets = protosetSlice
 	return nil
 }
 
diff --git a/core/resource_test.go b/core/resource_test.go
index b3764b3..7e053c2 100644
--- a/core/resource_test.go
+++ b/core/resource_test.go
@@ -60,7 +60,7 @@ func Test_prepareImport(t *testing.T) {
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			if got := prepareImport(tt.args.proto); !reflect.DeepEqual(got, tt.want) {
+			if got := prepareImport(tt.args.proto, ""); !reflect.DeepEqual(got, tt.want) {
 				t.Errorf("prepareImport() = %v, want %v",
 					string(got),
 					string(tt.want))
diff --git a/index/js/proto.js b/index/js/proto.js
index 2e9d94a..a5494a2 100644
--- a/index/js/proto.js
+++ b/index/js/proto.js
@@ -74,7 +74,7 @@ function handleProtoUpload() {
     const files = this.files;
 
     for (const file of files) {
-        if (!file.name.endsWith(".proto")) {
+        if (!file.name.endsWith(".proto") && !file.name.endsWith(".protoset")) {
             continue;
         }
         if (!protoMap.has(file.name)) {