*: use policy.Policy construct for logging

Previously, we were substituting the entire http transport just to add some client-side middleware. The Azure SDK for Go ships with default transport parameters that we lost by doing this; furthermore, we were not using the per-call policy options to model our middleware, as the SDK expected. Signed-off-by: Steve Kuznetsov <[email protected]>
Azure · Feb 6, 2025 · e47aef0 · e47aef0
1 parent 5831670
commit e47aef0
Show file tree

Hide file tree

Showing 8 changed files with 73 additions and 115 deletions.
diff --git a/pkg/env/prod.go b/pkg/env/prod.go
@@ -410,54 +410,52 @@ func (p *prod) MsiDataplaneClientOptions() (*policy.ClientOptions, error) {
 	return &clientOptions, nil
 }
 
-func ClientDebugLoggerMiddleware(log *logrus.Entry) azureclient.Middleware {
-	return func(delegate http.RoundTripper) http.RoundTripper {
-		return azureclient.RoundTripperFunc(func(req *http.Request) (*http.Response, error) {
-			log := log.WithFields(logrus.Fields{
-				"method": req.Method,
-				"url":    req.URL,
-			})
-			if req.Body != nil {
-				body, err := io.ReadAll(req.Body)
-				if err != nil {
-					log.WithError(err).Error("error reading request body")
-				}
-				if err := req.Body.Close(); err != nil {
-					log.WithError(err).Error("error closing request body")
-				}
-				log = log.WithField("body", string(body))
-				req.Body = io.NopCloser(bytes.NewBuffer(body)) // reset body so the delegate can use it
+func ClientDebugLoggerMiddleware(log *logrus.Entry) policy.Policy {
+	return azureclient.PolicyFunc(func(req *policy.Request) (*http.Response, error) {
+		log := log.WithFields(logrus.Fields{
+			"method": req.Raw().Method,
+			"url":    req.Raw().URL,
+		})
+		if req.Raw().Body != nil {
+			body, err := io.ReadAll(req.Raw().Body)
+			if err != nil {
+				log.WithError(err).Error("error reading request body")
+			}
+			if err := req.Raw().Body.Close(); err != nil {
+				log.WithError(err).Error("error closing request body")
 			}
-			log.Info("Sending request.")
-			resp, err := delegate.RoundTrip(req)
+			log = log.WithField("body", string(body))
+			req.Raw().Body = io.NopCloser(bytes.NewBuffer(body)) // reset body so the delegate can use it
+		}
+		log.Info("Sending request.")
+		resp, err := req.Next()
+		if err != nil {
+			log.WithError(err).Error("Request errored.")
+		} else if resp != nil {
+			log = log.WithFields(logrus.Fields{
+				"status": resp.StatusCode,
+			})
+			body, err := io.ReadAll(resp.Body)
 			if err != nil {
-				log.WithError(err).Error("Request errored.")
-			} else if resp != nil {
-				log = log.WithFields(logrus.Fields{
-					"status": resp.StatusCode,
-				})
-				body, err := io.ReadAll(resp.Body)
-				if err != nil {
-					log.WithError(err).Error("error reading response body")
-				}
-				if err := resp.Body.Close(); err != nil {
-					log.WithError(err).Error("error closing response body")
-				}
-				// n.b.: we only send one request now, this is best-effort but would need to be updated if we use other methods
-				response := dataplane.ManagedIdentityCredentials{}
-				if err := json.Unmarshal(body, &response); err != nil {
-					log.WithError(err).Error("error unmarshalling response body")
-				} else {
-					censorCredentials(&response)
-					log = log.WithField("body", string(body))
-				}
-				resp.Body = io.NopCloser(bytes.NewBuffer(body)) // reset body so the upstream round-trippers can use it
+				log.WithError(err).Error("error reading response body")
 			}
-			log.Info("Received response.")
+			if err := resp.Body.Close(); err != nil {
+				log.WithError(err).Error("error closing response body")
+			}
+			// n.b.: we only send one request now, this is best-effort but would need to be updated if we use other methods
+			response := dataplane.ManagedIdentityCredentials{}
+			if err := json.Unmarshal(body, &response); err != nil {
+				log.WithError(err).Error("error unmarshalling response body")
+			} else {
+				censorCredentials(&response)
+				log = log.WithField("body", string(body))
+			}
+			resp.Body = io.NopCloser(bytes.NewBuffer(body)) // reset body so the upstream round-trippers can use it
+		}
+		log.Info("Received response.")
 
-			return resp, err
-		})
-	}
+		return resp, err
+	})
 }
 
 func censorCredentials(input *dataplane.ManagedIdentityCredentials) {

diff --git a/pkg/util/azureclient/applens/applens_client_options.go b/pkg/util/azureclient/applens/applens_client_options.go
@@ -32,10 +32,8 @@ func NewClientOptions(certPool *x509.CertPool) *ClientOptions {
 		},
 	}
 
-	customRoundTripper := azureclient.NewCustomRoundTripper(httpTransport)
-
 	httpClient := &http.Client{
-		Transport: customRoundTripper,
+		Transport: httpTransport,
 	}
 
 	return &ClientOptions{
@@ -54,6 +52,7 @@ func NewClientOptions(certPool *x509.CertPool) *ClientOptions {
 				ApplicationID: userAgent,
 				Disabled:      false,
 			},
+			PerCallPolicies: []policy.Policy{azureclient.NewLoggingPolicy()},
 			Transport: httpClient,
 		},
 	}

diff --git a/pkg/util/azureclient/decorators.go b/pkg/util/azureclient/decorators.go
@@ -21,19 +21,10 @@ func DecorateSenderWithLogging(sender autorest.Sender) autorest.Sender {
 // in order to intercept http calls using our custom RoundTripper (through the adapter).
 func loggingDecorator() autorest.SendDecorator {
 	return func(s autorest.Sender) autorest.Sender {
-		rt := NewCustomRoundTripper(
-			&roundTripperAdapter{Sender: s},
-		)
-		return autorest.SenderFunc(rt.RoundTrip)
+		return autorest.SenderFunc(func(req *http.Request) (*http.Response, error) {
+			return loggingRoundTripper(req, func() (*http.Response, error) {
+				return s.Do(req)
+			})
+		})
 	}
 }
-
-// roundTripperAdapter converts from autorest.Sender (internal field)
-// to http.RoundTripper (external method).
-type roundTripperAdapter struct {
-	Sender autorest.Sender
-}
-
-func (rta *roundTripperAdapter) RoundTrip(req *http.Request) (*http.Response, error) {
-	return rta.Sender.Do(req)
-}
diff --git a/pkg/util/azureclient/environments.go b/pkg/util/azureclient/environments.go
@@ -5,12 +5,12 @@ package azureclient
 
 import (
 	"fmt"
-	"net/http"
 	"strings"
 
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore"
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore/arm"
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore/cloud"
+	"github.com/Azure/azure-sdk-for-go/sdk/azcore/policy"
 	"github.com/Azure/azure-sdk-for-go/sdk/azidentity"
 	"github.com/Azure/go-autorest/autorest/azure"
 	"github.com/Azure/msi-dataplane/pkg/dataplane"
@@ -105,40 +105,14 @@ func EnvironmentFromName(name string) (AROEnvironment, error) {
 	return AROEnvironment{}, fmt.Errorf("cloud environment %q is unsupported by ARO", name)
 }
 
-// RoundTripperFunc allows a function to implement http.RoundTripper
-type RoundTripperFunc func(*http.Request) (*http.Response, error)
-
-func (rt RoundTripperFunc) RoundTrip(req *http.Request) (*http.Response, error) {
-	return rt(req)
-}
-
-// Middleware closes over any client-side middleware
-type Middleware func(http.RoundTripper) http.RoundTripper
-
-// Chain is a handy function to wrap a base RoundTripper (optional) with the middlewares.
-func Chain(rt http.RoundTripper, middlewares ...Middleware) http.RoundTripper {
-	if rt == nil {
-		rt = http.DefaultTransport
-	}
-
-	for _, m := range middlewares {
-		rt = m(rt)
-	}
-
-	return rt
-}
-
 // ArmClientOptions returns an arm.ClientOptions to be passed in when instantiating
 // Azure SDK for Go clients.
-func (e *AROEnvironment) ArmClientOptions(middlewares ...Middleware) *arm.ClientOptions {
-	customRoundTripper := Chain(http.DefaultTransport, append([]Middleware{NewCustomRoundTripper}, middlewares...)...)
+func (e *AROEnvironment) ArmClientOptions(middlewares ...policy.Policy) *arm.ClientOptions {
 	return &arm.ClientOptions{
 		ClientOptions: azcore.ClientOptions{
 			Cloud: e.Cloud,
 			Retry: common.RetryOptions,
-			Transport: &http.Client{
-				Transport: customRoundTripper,
-			},
+			PerCallPolicies: append([]policy.Policy{NewLoggingPolicy()}, middlewares...),
 		},
 	}
 }

diff --git a/pkg/util/azureclient/roundtripper.go b/pkg/util/azureclient/roundtripper.go
@@ -7,6 +7,7 @@ import (
 	"net/http"
 	"time"
 
+	"github.com/Azure/azure-sdk-for-go/sdk/azcore/policy"
 	"github.com/sirupsen/logrus"
 
 	"github.com/Azure/ARO-RP/pkg/api"
@@ -25,17 +26,21 @@ const (
 	correlationIdHeader  = "X-Ms-Correlation-Request-Id"
 )
 
-func NewCustomRoundTripper(next http.RoundTripper) http.RoundTripper {
-	return &customRoundTripper{
-		next: next,
-	}
+type PolicyFunc func(req *policy.Request) (*http.Response, error)
+
+func (p PolicyFunc) Do(req *policy.Request) (*http.Response, error) {
+	return p(req)
 }
 
-type customRoundTripper struct {
-	next http.RoundTripper
+var _ policy.Policy = PolicyFunc(nil)
+
+func NewLoggingPolicy() policy.Policy {
+	return PolicyFunc(func(req *policy.Request) (*http.Response, error) {
+		return loggingRoundTripper(req.Raw(), req.Next)
+	})
 }
 
-func (crt *customRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
+func loggingRoundTripper(req *http.Request, next func() (*http.Response, error)) (*http.Response, error) {
 	correlationData := api.GetCorrelationDataFromCtx(req.Context())
 	if correlationData == nil {
 		correlationData = api.CreateCorrelationDataFromReq(req)
@@ -48,7 +53,7 @@ func (crt *customRoundTripper) RoundTrip(req *http.Request) (*http.Response, err
 
 	l.Info("HttpRequestStart")
 
-	res, err := crt.next.RoundTrip(req)
+	res, err := next()
 
 	l = updateCorrelationDataAndEnrichLogWithResponse(correlationData, l, res, requestTime)
 	l.Info("HttpRequestEnd")

diff --git a/pkg/util/purge/purge.go b/pkg/util/purge/purge.go
@@ -6,10 +6,9 @@ package purge
 // all the purge functions are located here
 
 import (
-	"net/http"
-
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore"
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore/arm"
+	"github.com/Azure/azure-sdk-for-go/sdk/azcore/policy"
 	"github.com/Azure/azure-sdk-for-go/sdk/azidentity"
 	mgmtfeatures "github.com/Azure/azure-sdk-for-go/services/resources/mgmt/2019-07-01/features"
 	"github.com/jongio/azidext/go/azidext"
@@ -49,14 +48,11 @@ func NewResourceCleaner(log *logrus.Entry, env env.Core, shouldDelete checkFn, d
 	scopes := []string{env.Environment().ResourceManagerScope}
 	authorizer := azidext.NewTokenCredentialAdapter(spTokenCredential, scopes)
 
-	customRoundTripper := azureclient.NewCustomRoundTripper(http.DefaultTransport)
 	clientOptions := &arm.ClientOptions{
 		ClientOptions: azcore.ClientOptions{
 			Cloud: env.Environment().Cloud,
 			Retry: common.RetryOptions,
-			Transport: &http.Client{
-				Transport: customRoundTripper,
-			},
+			PerCallPolicies: []policy.Policy{azureclient.NewLoggingPolicy()},
 		},
 	}
 

diff --git a/pkg/validate/openshiftcluster_validatedynamic.go b/pkg/validate/openshiftcluster_validatedynamic.go
@@ -14,7 +14,7 @@ import (
 	"github.com/Azure/checkaccess-v2-go-sdk/client"
 	"github.com/Azure/go-autorest/autorest"
 	"github.com/Azure/go-autorest/autorest/to"
-	jwt "github.com/golang-jwt/jwt/v4"
+	"github.com/golang-jwt/jwt/v4"
 	"github.com/sirupsen/logrus"
 
 	"github.com/Azure/ARO-RP/pkg/api"
@@ -139,9 +139,7 @@ func (dv *openShiftClusterDynamicValidator) Dynamic(ctx context.Context) error {
 
 	aroEnv := dv.env.Environment()
 	clientOptions := &azcore.ClientOptions{
-		Transport: &http.Client{
-			Transport: azureclient.NewCustomRoundTripper(http.DefaultTransport),
-		},
+		PerCallPolicies: []policy.Policy{azureclient.NewLoggingPolicy()},
 	}
 	pdpClient, err := client.NewRemotePDPClient(
 		fmt.Sprintf(aroEnv.Endpoint, dv.env.Location()),

diff --git a/test/e2e/setup.go b/test/e2e/setup.go
@@ -8,14 +8,14 @@ import (
 	"embed"
 	"fmt"
 	"math"
-	"net/http"
 	"net/url"
 	"os"
 	"os/exec"
 	"path/filepath"
 	"regexp"
 	"time"
 
+	"github.com/Azure/azure-sdk-for-go/sdk/azcore/policy"
 	. "github.com/onsi/ginkgo/v2"
 	. "github.com/onsi/gomega"
 
@@ -399,14 +399,11 @@ func newClientSet(ctx context.Context) (*clientSet, error) {
 		}
 	}
 
-	customRoundTripper := azureclient.NewCustomRoundTripper(http.DefaultTransport)
 	clientOptions := &arm.ClientOptions{
 		ClientOptions: azcore.ClientOptions{
-			Cloud: _env.Environment().Cloud,
-			Retry: common.RetryOptions,
-			Transport: &http.Client{
-				Transport: customRoundTripper,
-			},
+			Cloud:           _env.Environment().Cloud,
+			Retry:           common.RetryOptions,
+			PerCallPolicies: []policy.Policy{azureclient.NewLoggingPolicy()},
 		},
 	}