Validating Client Request ID and Testing (#246)

* Validating client request id and testing

* Fixed a bug that was causing failing errors

* Resolved comments from PR & made additional fixes

* Fixed bug throwing nil pointer errors

Co-authored-by: Mohit Sharma <65536214+mohsha-msft@users.noreply.github.com>

azblob/zc_policy_unique_request_id.go

@@ -2,7 +2,7 @@ package azblob
 
 import (
 	"context"
-
+	"errors"
 	"github.com/Azure/azure-pipeline-go/pipeline"
 )
 
@@ -14,9 +14,22 @@ func NewUniqueRequestIDPolicyFactory() pipeline.Factory {
 		return func(ctx context.Context, request pipeline.Request) (pipeline.Response, error) {
 			id := request.Header.Get(xMsClientRequestID)
 			if id == "" { // Add a unique request ID if the caller didn't specify one already
-				request.Header.Set(xMsClientRequestID, newUUID().String())
+				id = newUUID().String()
+				request.Header.Set(xMsClientRequestID, id)
 			}
-			return next.Do(ctx, request)
+
+			resp, err := next.Do(ctx, request)
+
+			if err == nil && resp != nil {
+				val := resp.Response().Header.Values(xMsClientRequestID)
+				if len(val) > 0 {
+					if val[0] != id {
+						err = errors.New("client Request ID from request and response does not match")
+					}
+				}
+			}
+
+			return resp, err
 		}
 	})
 }

azblob/zt_policy_request_id_test.go

@@ -0,0 +1,105 @@
+package azblob
+
+import (
+	"context"
+	"errors"
+	"github.com/Azure/azure-pipeline-go/pipeline"
+	chk "gopkg.in/check.v1"
+	"net/http"
+	"net/url"
+)
+
+type requestIDTestScenario int
+
+const (
+	// Testing scenarios for echoing Client Request ID
+	clientRequestIDMissing             requestIDTestScenario = 1
+	errorFromNextPolicy                requestIDTestScenario = 2
+	clientRequestIDMatch               requestIDTestScenario = 3
+	clientRequestIDNoMatch             requestIDTestScenario = 4
+	errorMessageClientRequestIDNoMatch                       = "client Request ID from request and response does not match"
+	errorMessageFromNextPolicy                               = "error is not nil"
+)
+
+type clientRequestIDPolicy struct {
+	matchID  string
+	scenario requestIDTestScenario
+}
+
+func (p clientRequestIDPolicy) Do(ctx context.Context, request pipeline.Request) (pipeline.Response, error) {
+	var header http.Header = make(map[string][]string)
+	var err error
+
+	// Set headers and errors according to each scenario
+	switch p.scenario {
+	case clientRequestIDMissing:
+	case errorFromNextPolicy:
+		err = errors.New(errorMessageFromNextPolicy)
+	case clientRequestIDMatch:
+		header.Add(xMsClientRequestID, request.Header.Get(xMsClientRequestID))
+	case clientRequestIDNoMatch:
+		header.Add(xMsClientRequestID, "fake-client-request-id")
+	default:
+		header.Add(xMsClientRequestID, newUUID().String())
+	}
+
+	response := http.Response{Header: header}
+
+	return pipeline.NewHTTPResponse(&response), err
+}
+
+func (s *aztestsSuite) TestEchoClientRequestIDMissing(c *chk.C) {
+	factory := NewUniqueRequestIDPolicyFactory()
+
+	// Scenario 1: Client Request ID is missing
+	policy := factory.New(clientRequestIDPolicy{scenario: clientRequestIDMissing}, nil)
+	request, _ := pipeline.NewRequest("GET", url.URL{}, nil)
+	resp, err := policy.Do(context.Background(), request)
+
+	c.Assert(err, chk.IsNil)
+	c.Assert(resp, chk.NotNil)
+	c.Assert(resp.Response().Header.Values(xMsClientRequestID), chk.IsNil)
+}
+
+func (s *aztestsSuite) TestEchoClientRequestIDErrorFromNextPolicy(c *chk.C) {
+	factory := NewUniqueRequestIDPolicyFactory()
+
+	// Scenario 2: Do method returns an error
+	policy := factory.New(clientRequestIDPolicy{scenario: errorFromNextPolicy}, nil)
+	request, _ := pipeline.NewRequest("GET", url.URL{}, nil)
+	resp, err := policy.Do(context.Background(), request)
+
+	c.Assert(err, chk.NotNil)
+	c.Assert(err.Error(), chk.Equals, errorMessageFromNextPolicy)
+	c.Assert(resp, chk.NotNil)
+}
+
+func (s *aztestsSuite) TestEchoClientRequestIDMatch(c *chk.C) {
+	factory := NewUniqueRequestIDPolicyFactory()
+
+	// Scenario 3: Client Request ID matches
+	matchRequestID := newUUID().String()
+	policy := factory.New(clientRequestIDPolicy{matchID: matchRequestID, scenario: clientRequestIDMatch}, nil)
+	request, _ := pipeline.NewRequest("GET", url.URL{}, nil)
+	request.Header.Set(xMsClientRequestID, matchRequestID)
+	resp, err := policy.Do(context.Background(), request)
+
+	c.Assert(err, chk.IsNil)
+	c.Assert(resp, chk.NotNil)
+	c.Assert(resp.Response().Header.Get(xMsClientRequestID), chk.Equals, request.Header.Get(xMsClientRequestID))
+}
+
+func (s *aztestsSuite) TestEchoClientRequestIDNoMatch(c *chk.C) {
+	factory := NewUniqueRequestIDPolicyFactory()
+
+	// Scenario 4: Client Request ID does not match
+	matchRequestID := newUUID().String()
+	policy := factory.New(clientRequestIDPolicy{matchID: matchRequestID, scenario: clientRequestIDNoMatch}, nil)
+	request, _ := pipeline.NewRequest("GET", url.URL{}, nil)
+	request.Header.Set(xMsClientRequestID, matchRequestID)
+	resp, err := policy.Do(context.Background(), request)
+
+	c.Assert(err, chk.NotNil)
+	c.Assert(err.Error(), chk.Equals, errorMessageClientRequestIDNoMatch)
+	c.Assert(resp, chk.NotNil)
+}

go.mod

@@ -7,6 +7,6 @@ require (
 	github.com/Azure/go-autorest/autorest/adal v0.9.2
 	github.com/google/uuid v1.1.1
 	github.com/niemeyer/pretty v0.0.0-20200227124842-a10e7caefd8e // indirect
-	golang.org/x/sys v0.0.0-20200828194041-157a740278f4
+	golang.org/x/sys v0.0.0-20200828194041-157a740278f4 // indirect
 	gopkg.in/check.v1 v1.0.0-20200227125254-8fa46927fb4f
 )