From 76c2cf7eb66c992cd984d8976f02066f65257eb8 Mon Sep 17 00:00:00 2001
From: Sean Sullivan <seans3@gmail.com>
Date: Tue, 9 Feb 2021 23:09:22 -0800
Subject: [PATCH] Kubectl command headers in requests: KEP 859

Kubernetes-commit: 211fc12b67ea2d04a46b24ef245b3513f2cfeb03
---
 pkg/cmd/cmd.go | 46 ++++++++++++++++++++++++++++++++++++++++++++--
 1 file changed, 44 insertions(+), 2 deletions(-)

diff --git a/pkg/cmd/cmd.go b/pkg/cmd/cmd.go
index 94bf82396..1f9484472 100644
--- a/pkg/cmd/cmd.go
+++ b/pkg/cmd/cmd.go
@@ -20,6 +20,7 @@ import (
 	"flag"
 	"fmt"
 	"io"
+	"net/http"
 	"os"
 	"os/exec"
 	"runtime"
@@ -322,6 +323,8 @@ __kubectl_custom_func() {
 `
 )
 
+const kubectlCmdHeaders = "KUBECTL_COMMAND_HEADERS"
+
 var (
 	bashCompletionFlags = map[string]string{
 		"namespace": "__kubectl_get_resource_namespace",
@@ -469,7 +472,6 @@ func HandlePluginCommand(pluginHandler PluginHandler, cmdArgs []string) error {
 func NewKubectlCommand(in io.Reader, out, err io.Writer) *cobra.Command {
 	warningHandler := rest.NewWarningWriter(err, rest.WarningWriterOptions{Deduplicate: true, Color: term.AllowsColorOutput(err)})
 	warningsAsErrors := false
-
 	// Parent command to which all subcommands are added.
 	cmds := &cobra.Command{
 		Use:   "kubectl",
@@ -521,6 +523,8 @@ func NewKubectlCommand(in io.Reader, out, err io.Writer) *cobra.Command {
 	kubeConfigFlags.AddFlags(flags)
 	matchVersionKubeConfigFlags := cmdutil.NewMatchVersionFlags(kubeConfigFlags)
 	matchVersionKubeConfigFlags.AddFlags(cmds.PersistentFlags())
+	// Updates hooks to add kubectl command headers: SIG CLI KEP 859.
+	addCmdHeaderHooks(cmds, kubeConfigFlags)
 
 	cmds.PersistentFlags().AddGoFlagSet(flag.CommandLine)
 
@@ -538,6 +542,12 @@ func NewKubectlCommand(in io.Reader, out, err io.Writer) *cobra.Command {
 
 	ioStreams := genericclioptions.IOStreams{In: in, Out: out, ErrOut: err}
 
+	// Proxy command is incompatible with CommandHeaderRoundTripper, so
+	// clear the WrapConfigFn before running proxy command.
+	proxyCmd := proxy.NewCmdProxy(f, ioStreams)
+	proxyCmd.PreRun = func(cmd *cobra.Command, args []string) {
+		kubeConfigFlags.WrapConfigFn = nil
+	}
 	groups := templates.CommandGroups{
 		{
 			Message: "Basic Commands (Beginner):",
@@ -585,7 +595,7 @@ func NewKubectlCommand(in io.Reader, out, err io.Writer) *cobra.Command {
 				attach.NewCmdAttach(f, ioStreams),
 				cmdexec.NewCmdExec(f, ioStreams),
 				portforward.NewCmdPortForward(f, ioStreams),
-				proxy.NewCmdProxy(f, ioStreams),
+				proxyCmd,
 				cp.NewCmdCp(f, ioStreams),
 				auth.NewCmdAuth(f, ioStreams),
 				debug.NewCmdDebug(f, ioStreams),
@@ -646,6 +656,38 @@ func NewKubectlCommand(in io.Reader, out, err io.Writer) *cobra.Command {
 	return cmds
 }
 
+// addCmdHeaderHooks performs updates on two hooks:
+//   1) Modifies the passed "cmds" persistent pre-run function to parse command headers.
+//      These headers will be subsequently added as X-headers to every
+//      REST call.
+//   2) Adds CommandHeaderRoundTripper as a wrapper around the standard
+//      RoundTripper. CommandHeaderRoundTripper adds X-Headers then delegates
+//      to standard RoundTripper.
+// For alpha, these hooks are only updated if the KUBECTL_COMMAND_HEADERS
+// environment variable is set.
+// See SIG CLI KEP 859 for more information:
+//   https://github.com/kubernetes/enhancements/tree/master/keps/sig-cli/859-kubectl-headers
+func addCmdHeaderHooks(cmds *cobra.Command, kubeConfigFlags *genericclioptions.ConfigFlags) {
+	if _, exists := os.LookupEnv(kubectlCmdHeaders); !exists {
+		return
+	}
+	crt := &genericclioptions.CommandHeaderRoundTripper{}
+	existingPreRunE := cmds.PersistentPreRunE
+	// Add command parsing to the existing persistent pre-run function.
+	cmds.PersistentPreRunE = func(cmd *cobra.Command, args []string) error {
+		crt.ParseCommandHeaders(cmd, args)
+		return existingPreRunE(cmd, args)
+	}
+	// Wraps CommandHeaderRoundTripper around standard RoundTripper.
+	kubeConfigFlags.WrapConfigFn = func(c *rest.Config) *rest.Config {
+		c.Wrap(func(rt http.RoundTripper) http.RoundTripper {
+			crt.Delegate = rt
+			return crt
+		})
+		return c
+	}
+}
+
 func runHelp(cmd *cobra.Command, args []string) {
 	cmd.Help()
 }