From 80ef2d48e6b8d01227ff5c4cb6036e7282643253 Mon Sep 17 00:00:00 2001
From: Petra <petragabriela2701@gmail.com>
Date: Sun, 16 Jul 2023 15:02:43 +0700
Subject: [PATCH] =?UTF-8?q?fix:=20=F0=9F=90=9B=20prevent=20reflection=20he?=
 =?UTF-8?q?aders=20to=20be=20stored=20in=20the=20server?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 core/resource.go   | 21 +++++++++++++--------
 handler/handler.go |  6 ++----
 2 files changed, 15 insertions(+), 12 deletions(-)

diff --git a/core/resource.go b/core/resource.go
index 33d5c38..cc27a8c 100644
--- a/core/resource.go
+++ b/core/resource.go
@@ -38,10 +38,15 @@ type Resource struct {
 	md      metadata.MD
 }
 
-//openDescriptor - use it to reflect server descriptor
-func (r *Resource) openDescriptor() error {
+// openDescriptor - use it to reflect server descriptor
+func (r *Resource) openDescriptor(reflectHeaders []string) error {
 	ctx := context.Background()
-	refCtx := metadata.NewOutgoingContext(ctx, r.md)
+	outgoingMD := r.md.Copy()
+	md := grpcurl.MetadataFromHeaders(reflectHeaders)
+	for k, v := range md {
+		outgoingMD.Append(k, v...)
+	}
+	refCtx := metadata.NewOutgoingContext(ctx, outgoingMD)
 	r.refClient = grpcreflect.NewClient(refCtx, reflectpb.NewServerReflectionClient(r.clientConn))
 
 	// if no protos available use server reflection
@@ -73,7 +78,7 @@ func (r *Resource) openDescriptor() error {
 	return err
 }
 
-//closeDescriptor - please ensure to always close after open in the same flow
+// closeDescriptor - please ensure to always close after open in the same flow
 func (r *Resource) closeDescriptor() {
 	done := make(chan int)
 	go func() {
@@ -95,8 +100,8 @@ func (r *Resource) closeDescriptor() {
 // List - To list all services exposed by a server
 // symbol can be "" to list all available services
 // symbol also can be service name to list all available method
-func (r *Resource) List(symbol string) ([]string, error) {
-	err := r.openDescriptor()
+func (r *Resource) List(symbol string, reflectHeaders []string) ([]string, error) {
+	err := r.openDescriptor(reflectHeaders)
 	if err != nil {
 		return nil, err
 	}
@@ -137,7 +142,7 @@ func (r *Resource) List(symbol string) ([]string, error) {
 // It also prints a description of that symbol, in the form of snippets of proto source.
 // It won't necessarily be the original source that defined the element, but it will be equivalent.
 func (r *Resource) Describe(symbol string) (string, string, error) {
-	err := r.openDescriptor()
+	err := r.openDescriptor(nil)
 	if err != nil {
 		return "", "", err
 	}
@@ -196,7 +201,7 @@ func (r *Resource) Describe(symbol string) (string, string, error) {
 
 // Invoke - invoking gRPC function
 func (r *Resource) Invoke(ctx context.Context, metadata []string, symbol string, in io.Reader) (string, time.Duration, error) {
-	err := r.openDescriptor()
+	err := r.openDescriptor(nil)
 	if err != nil {
 		return "", 0, err
 	}
diff --git a/handler/handler.go b/handler/handler.go
index dbd11e9..a25d261 100644
--- a/handler/handler.go
+++ b/handler/handler.go
@@ -88,15 +88,13 @@ func (h *Handler) getLists(w http.ResponseWriter, r *http.Request) {
 		metadataStr = fmt.Sprintf("%s:", m)
 	}
 
-	h.g.SetReflectHeaders(metadata...)
-
 	res, err := h.g.GetResource(context.Background(), host, !useTLS, restart)
 	if err != nil {
 		writeError(w, err)
 		return
 	}
 
-	result, err := res.List(service)
+	result, err := res.List(service, metadata)
 	if err != nil {
 		writeError(w, err)
 		return
@@ -155,7 +153,7 @@ func (h *Handler) getListsWithProto(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	result, err := res.List(service)
+	result, err := res.List(service, nil)
 	if err != nil {
 		writeError(w, err)
 		return