Preserve specified cookies between browser and gRPC UI

cmd/grpcui/grpcui.go

@@ -90,6 +90,7 @@ var (
 	rpcHeaders    multiString
 	reflHeaders   multiString
 	prsvHeaders   multiString
+	prsvCookies   multiString
 	defHeaders    multiString
 	expandHeaders = flags.Bool("expand-headers", false, prettify(`
 		If set, headers may use '${NAME}' syntax to reference environment
@@ -203,6 +204,11 @@ func init() {
 		Having gRPC UI preserve these headers means that the JWTs will also be
 		sent to backend gRPC servers. These headers are only sent when RPCs are
 		invoked and are not included for reflection requests.`))
+	flags.Var(&prsvCookies, "preserve-cookie", prettify(`
+        Cookie names (no values) for cookies that should be forwarded between the
+        web browser and the gRPC server. This will forward the relevant "Set-Cookie"
+        headers from the gRPC server to the web browser and the "Cookie" headers from
+        the browser to the gRPC server.`))
 	flags.Var(&defHeaders, "default-header", prettify(`
 		Additional headers to add to metadata in the gRPCui web form. Each value
 		should be in 'name: value' format. May specify more than one via multiple
@@ -641,6 +647,9 @@ func main() {
 	if len(prsvHeaders) > 0 {
 		handlerOpts = append(handlerOpts, standalone.PreserveHeaders(prsvHeaders))
 	}
+	if len(prsvCookies) > 0 {
+		handlerOpts = append(handlerOpts, standalone.PreserveCookies(prsvCookies))
+	}
 	if verbosity > 0 {
 		handlerOpts = append(handlerOpts, standalone.WithInvokeVerbosity(verbosity))
 	}

handlers.go

@@ -58,6 +58,10 @@ type InvokeOptions struct {
 	// includes conflicting metadata, the values in the HTTP request headers
 	// will override, and the values in the request will not be sent.
 	PreserveHeaders []string
+	// A set of cookie names that will be preserved. This means that the Cookie
+	// header will be sent to the gRPC backend if the cookie name is in this set
+	// and the Set-Cookie header will be send back from it.
+	PreserveCookies []string
 	// Whether or not default values should be emitted in the JSON response
 	EmitDefaults bool
 	// If verbosity is greater than zero, the handler may log events, such as
@@ -110,6 +114,22 @@ func RPCInvokeHandlerWithOptions(ch grpc.ClientConnInterface, descs []*desc.Meth
 					return
 				}
 				w.Header().Set("Content-Type", "application/json")
+				if len(options.PreserveCookies) > 0 {
+					for _, h := range results.Headers {
+						if strings.ToLower(h.Name) != "set-cookie" {
+							continue
+						}
+						cn, _, f := strings.Cut(h.Value, "=")
+						if !f {
+							continue
+						}
+						for _, c := range options.PreserveCookies {
+							if c == cn {
+								w.Header().Add("Set-Cookie", h.Value)
+							}
+						}
+					}
+				}
 				enc := json.NewEncoder(w)
 				enc.SetIndent("", "  ")
 				enc.Encode(results)
@@ -472,7 +492,11 @@ func invokeRPC(ctx context.Context, methodName string, ch grpc.ClientConnInterfa
 
 func (opts *InvokeOptions) overrideHeaders(reqHdrs http.Header) metadata.MD {
 	hdrs := grpcurl.MetadataFromHeaders(opts.ExtraMetadata)
+	preserveAllCookies := false
 	for _, name := range opts.PreserveHeaders {
+		if strings.ToLower(name) == "cookie" {
+			preserveAllCookies = true
+		}
 		vals := reqHdrs.Values(name)
 		if opts.Verbosity > 0 {
 			if existing := hdrs.Get(name); len(existing) > 0 {
@@ -481,6 +505,26 @@ func (opts *InvokeOptions) overrideHeaders(reqHdrs http.Header) metadata.MD {
 		}
 		hdrs.Set(name, vals...)
 	}
+	cookieHeaders := reqHdrs.Values("Cookie")
+	if len(cookieHeaders) > 0 && len(opts.PreserveCookies) > 0 && !preserveAllCookies {
+		var cookiesToForward []string
+		for _, chl := range cookieHeaders {
+			for ch := range strings.SplitSeq(chl, ";") {
+				key, _, found := strings.Cut(ch, "=")
+				if !found {
+					continue
+				}
+				for _, cookie := range opts.PreserveCookies {
+					if strings.Trim(key, " ") == cookie {
+						cookiesToForward = append(cookiesToForward, ch)
+					}
+				}
+			}
+		}
+		if len(cookiesToForward) > 0 {
+			hdrs.Append("Cookie", strings.Join(cookiesToForward, "; "))
+		}
+	}
 	return hdrs
 }
 

standalone/opts.go

@@ -188,6 +188,14 @@ func PreserveHeaders(headerNames []string) HandlerOption {
 	})
 }
 
+// PreserveCookies instructs the Handler to preserve Cookie and Set-Cookie
+// HTTP headers for the given cookies.
+func PreserveCookies(cookieNames []string) HandlerOption {
+	return optFunc(func(opts *handlerOptions) {
+		opts.preserveCookies = cookieNames
+	})
+}
+
 // EmitDefaults tells gRPCurl whether or not default values should be emitted
 func EmitDefaults(emit bool) HandlerOption {
 	return optFunc(func(opts *handlerOptions) {
@@ -242,6 +250,7 @@ type handlerOptions struct {
 	defaultMetadata     []string
 	extraMetadata       []string
 	preserveHeaders     []string
+	preserveCookies     []string
 	emitDefaults        bool
 	invokeVerbosity     int
 	debug               *bool

standalone/standalone.go

@@ -96,6 +96,7 @@ func Handler(ch grpcdynamic.Channel, target string, methods []*desc.MethodDescri
 	invokeOpts := grpcui.InvokeOptions{
 		ExtraMetadata:   uiOpts.extraMetadata,
 		PreserveHeaders: uiOpts.preserveHeaders,
+		PreserveCookies: uiOpts.preserveCookies,
 		EmitDefaults:    uiOpts.emitDefaults,
 		Verbosity:       uiOpts.invokeVerbosity,
 	}