feat: add proxy --probe and enable lifecycle postStart hook (#490)

* feat: add proxy --probe and enable lifecycle postStart hook

Signed-off-by: Anish Ramasekar <anish.ramasekar@gmail.com>

* review feedback

Signed-off-by: Anish Ramasekar <anish.ramasekar@gmail.com>

cmd/proxy/main.go

@@ -12,6 +12,7 @@ import (
 
 var (
 	proxyPort int
+	probe     bool
 )
 
 func init() {
@@ -23,8 +24,21 @@ func main() {
 	logger.AddFlags()
 
 	flag.IntVar(&proxyPort, "proxy-port", 8000, "Port for the proxy to listen on")
+	flag.BoolVar(&probe, "probe", false, "Run a readyz probe on the proxy")
 	flag.Parse()
 
+	// when proxy is run with --probe, it will run a readyz probe on the proxy
+	// this is used in the postStart lifecycle hook to verify the proxy is ready
+	// to serve requests
+	if probe {
+		setupLog := logger.Get().WithName("probe")
+		if err := proxy.Probe(proxyPort); err != nil {
+			setupLog.Error(err, "failed to probe")
+			os.Exit(1)
+		}
+		os.Exit(0)
+	}
+
 	setupLog := logger.Get().WithName("setup")
 	p, err := proxy.NewProxy(proxyPort, logger.Get().WithName("proxy"))
 	if err != nil {

pkg/cmd/podidentity/detect.go

@@ -370,13 +370,25 @@ func (dc *detectCmd) addProxyContainer(containers []corev1.Container) []corev1.C
 		Name:            proxyContainerName,
 		Image:           proxyImage,
 		ImagePullPolicy: corev1.PullIfNotPresent,
-		Args:            []string{"--log-encoder=json"},
+		Args: []string{
+			fmt.Sprintf("--proxy-port=%d", dc.proxyPort),
+		},
 		Ports: []corev1.ContainerPort{
 			{
-				Name:          "http",
 				ContainerPort: int32(dc.proxyPort),
 			},
 		},
+		Lifecycle: &corev1.Lifecycle{
+			PostStart: &corev1.LifecycleHandler{
+				Exec: &corev1.ExecAction{
+					Command: []string{
+						"/proxy",
+						fmt.Sprintf("--proxy-port=%d", dc.proxyPort),
+						"--probe",
+					},
+				},
+			},
+		},
 	}
 
 	containers = append(containers, proxyContainer)

pkg/cmd/podidentity/detect_test.go

@@ -1,6 +1,7 @@
 package podidentity
 
 import (
+	"fmt"
 	"os"
 	"reflect"
 	"strings"
@@ -48,13 +49,25 @@ var (
 		Name:            proxyContainerName,
 		Image:           proxyImage,
 		ImagePullPolicy: corev1.PullIfNotPresent,
-		Args:            []string{"--log-encoder=json"},
+		Args: []string{
+			fmt.Sprintf("--proxy-port=%d", 8000),
+		},
 		Ports: []corev1.ContainerPort{
 			{
-				Name:          "http",
 				ContainerPort: 8000,
 			},
 		},
+		Lifecycle: &corev1.Lifecycle{
+			PostStart: &corev1.LifecycleHandler{
+				Exec: &corev1.ExecAction{
+					Command: []string{
+						"/proxy",
+						fmt.Sprintf("--proxy-port=%d", 8000),
+						"--probe",
+					},
+				},
+			},
+		},
 	}
 )
 

pkg/proxy/probe.go

@@ -0,0 +1,45 @@
+package proxy
+
+import (
+	"fmt"
+	"net/http"
+	"time"
+
+	"github.com/pkg/errors"
+)
+
+const (
+	// retryCount is the number of times to retry probing the proxy.
+	retryCount = 7
+	// waitTime is the time to wait between retries.
+	waitTime = time.Second
+	// clientTimeout is the timeout for the client.
+	clientTimeout = time.Second * 5
+)
+
+// Probe checks if the proxy is ready to serve requests.
+func Probe(port int) error {
+	url := fmt.Sprintf("http://%s:%d%s", localhost, port, readyzPathPrefix)
+	return probe(url)
+}
+
+func probe(url string) error {
+	client := &http.Client{
+		Timeout: clientTimeout,
+	}
+	for i := 0; i < retryCount; i++ {
+		req, err := http.NewRequest(http.MethodGet, url, nil)
+		if err != nil {
+			return err
+		}
+		resp, err := client.Do(req)
+		if err != nil {
+			return err
+		}
+		if resp.StatusCode == http.StatusOK {
+			return nil
+		}
+		time.Sleep(waitTime)
+	}
+	return errors.Errorf("failed to probe proxy")
+}

pkg/proxy/probe_test.go

@@ -0,0 +1,28 @@
+package proxy
+
+import (
+	"testing"
+
+	"k8s.io/klog/v2/klogr"
+)
+
+func TestProbe(t *testing.T) {
+	setup()
+	defer teardown()
+
+	p := &proxy{logger: klogr.New()}
+	rtr.PathPrefix("/readyz").HandlerFunc(p.readyzHandler)
+
+	if err := probe(server.URL + "/readyz"); err != nil {
+		t.Errorf("probe() = %v, want nil", err)
+	}
+}
+
+func TestProbeError(t *testing.T) {
+	setup()
+	defer teardown()
+
+	if err := probe(server.URL + "/readyz"); err == nil {
+		t.Errorf("probe() = nil, want error")
+	}
+}

pkg/proxy/proxy.go

@@ -23,10 +23,15 @@ const (
 	// "/metadata" portion is case-insensitive in IMDS
 	tokenPathPrefix = "/{type:(?i:metadata)}/identity/oauth2/token" // #nosec
 
+	// readyzPathPrefix is the path for readiness probe
+	readyzPathPrefix = "/readyz"
+
 	// metadataIPAddress is the IP address of the metadata service
 	metadataIPAddress = "169.254.169.254"
 	// metadataPort is the port of the metadata service
 	metadataPort = 80
+	// localhost is the hostname of the localhost
+	localhost = "localhost"
 )
 
 var (
@@ -85,10 +90,11 @@ func NewProxy(port int, logger logr.Logger) (Proxy, error) {
 func (p *proxy) Run() error {
 	rtr := mux.NewRouter()
 	rtr.PathPrefix(tokenPathPrefix).HandlerFunc(p.msiHandler)
+	rtr.PathPrefix(readyzPathPrefix).HandlerFunc(p.readyzHandler)
 	rtr.PathPrefix("/").HandlerFunc(p.defaultPathHandler)
 
 	p.logger.Info("starting the proxy server", "port", p.port, "userAgent", userAgent)
-	return http.ListenAndServe(fmt.Sprintf("localhost:%d", p.port), rtr)
+	return http.ListenAndServe(fmt.Sprintf("%s:%d", localhost, p.port), rtr)
 }
 
 func (p *proxy) msiHandler(w http.ResponseWriter, r *http.Request) {
@@ -157,6 +163,11 @@ func (p *proxy) defaultPathHandler(w http.ResponseWriter, r *http.Request) {
 	_, _ = w.Write(body)
 }
 
+func (p *proxy) readyzHandler(w http.ResponseWriter, r *http.Request) {
+	p.logger.Info("received readyz request", "method", r.Method, "uri", r.RequestURI)
+	fmt.Fprintf(w, "ok")
+}
+
 func doTokenRequest(ctx context.Context, clientID, resource, tenantID, authorityHost string) (*token, error) {
 	tokenFilePath := os.Getenv(webhook.AzureFederatedTokenFileEnvVar)
 	signedAssertion, err := readJWTFromFS(tokenFilePath)

pkg/proxy/proxy_test.go

@@ -10,10 +10,10 @@ import (
 	"path/filepath"
 	"testing"
 
-	"github.com/Azure/azure-workload-identity/pkg/webhook"
-
 	"github.com/gorilla/mux"
 	"k8s.io/klog/v2/klogr"
+
+	"github.com/Azure/azure-workload-identity/pkg/webhook"
 )
 
 var (
@@ -268,3 +268,43 @@ func testTokenHandler(w http.ResponseWriter, r *http.Request) {
 func testDefaultHandler(w http.ResponseWriter, r *http.Request) {
 	fmt.Fprintf(w, "default_handler")
 }
+
+func TestProxy_ReadyZHandler(t *testing.T) {
+	tests := []struct {
+		name string
+		path string
+		code int
+	}{
+		{
+			name: "readyz",
+			path: "/readyz",
+			code: http.StatusOK,
+		},
+		{
+			name: "readyz",
+			path: "/readyz/",
+			code: http.StatusOK,
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			setup()
+			defer teardown()
+
+			p := &proxy{logger: klogr.New()}
+			rtr.PathPrefix("/readyz").HandlerFunc(p.readyzHandler)
+
+			req, err := http.NewRequest(http.MethodGet, server.URL+test.path, nil)
+			if err != nil {
+				t.Error(err)
+			}
+
+			recorder := httptest.NewRecorder()
+			rtr.ServeHTTP(recorder, req)
+			if recorder.Code != test.code {
+				t.Errorf("Expected code %d, got %d", test.code, recorder.Code)
+			}
+		})
+	}
+}

pkg/webhook/webhook.go

@@ -254,6 +254,17 @@ func (m *podMutator) injectProxySidecarContainer(containers []corev1.Container,
 		Ports: []corev1.ContainerPort{{
 			ContainerPort: proxyPort,
 		}},
+		Lifecycle: &corev1.Lifecycle{
+			PostStart: &corev1.LifecycleHandler{
+				Exec: &corev1.ExecAction{
+					Command: []string{
+						"/proxy",
+						fmt.Sprintf("--proxy-port=%d", proxyPort),
+						"--probe",
+					},
+				},
+			},
+		},
 	})
 
 	return containers

pkg/webhook/webhook_test.go

@@ -1256,6 +1256,17 @@ func TestInjectProxySidecarContainer(t *testing.T) {
 		Ports: []corev1.ContainerPort{{
 			ContainerPort: proxyPort,
 		}},
+		Lifecycle: &corev1.Lifecycle{
+			PostStart: &corev1.LifecycleHandler{
+				Exec: &corev1.ExecAction{
+					Command: []string{
+						"/proxy",
+						fmt.Sprintf("--proxy-port=%d", proxyPort),
+						"--probe",
+					},
+				},
+			},
+		},
 	}
 
 	tests := []struct {

scripts/ci-e2e.sh

@@ -94,7 +94,9 @@ test_helm_chart() {
     --debug \
     -v=5
   poll_webhook_readiness
-  make test-e2e-run
+  # TODO(aramase: remove GINKGO_SKIP once helm chart is updated to use v0.12.0 to include the 
+  # new proxy probe feature.
+  GINKGO_SKIP=Proxy make test-e2e-run
 
   ${HELM} upgrade --install workload-identity-webhook "${REPO_ROOT}/manifest_staging/charts/workload-identity-webhook" \
     --set image.repository="${REGISTRY:-mcr.microsoft.com/oss/azure/workload-identity/webhook}" \

test/e2e/proxy_test.go

@@ -41,10 +41,7 @@ var _ = ginkgo.Describe("Proxy [LinuxOnly] [AKSSoakOnly] [Exclude:Arc]", func()
 			serviceAccount,
 			"mcr.microsoft.com/azure-cli",
 			nil,
-			// az login -i reuses the connection, so we need to make sure the token request is made after the proxy sidecar is started
-			// otherwise the token request will fail. The sleep 15 is a workaround for the issue.
-			// TODO(aramase): remove the sleep after https://github.com/Azure/azure-workload-identity/issues/486 is fixed and v0.12.0 is released.
-			[]string{"/bin/sh", "-c", fmt.Sprintf("sleep 15; az login -i -u %s --allow-no-subscriptions --debug; sleep 3600", clientID)},
+			[]string{"/bin/sh", "-c", fmt.Sprintf("az login -i -u %s --allow-no-subscriptions --debug; sleep 3600", clientID)},
 			nil,
 			proxyAnnotations,
 			true,