From a5e761f57dc42e91a2a460b8c1af39bda58956e8 Mon Sep 17 00:00:00 2001
From: Tom Bamford <tom@bamford.io>
Date: Tue, 19 Mar 2024 17:45:30 +0000
Subject: [PATCH] improve handling of nil responses and responses with
 malformed headers

---
 sdk/client/client.go      | 33 ++++++++++++++++++++++-----------
 sdk/client/client_test.go | 32 ++++++++++++++++++++++++++++++++
 2 files changed, 54 insertions(+), 11 deletions(-)

diff --git a/sdk/client/client.go b/sdk/client/client.go
index 48b577d899b..07ef77247bf 100644
--- a/sdk/client/client.go
+++ b/sdk/client/client.go
@@ -167,22 +167,33 @@ func (r *Response) Unmarshal(model interface{}) error {
 	if model == nil {
 		return fmt.Errorf("model was nil")
 	}
+	if r.Response == nil {
+		return fmt.Errorf("could not unmarshal as the HTTP response was nil")
+	}
+
+	var contentType string
+	if r.Response.Header != nil {
+		contentType = strings.ToLower(r.Response.Header.Get("Content-Type"))
 
-	contentType := strings.ToLower(r.Header.Get("Content-Type"))
-	if contentType == "" {
-		// some APIs (e.g. Storage Data Plane) don't return a content type... so we'll assume from the Accept header
-		acc, err := accept.FromString(r.Request.Header.Get("Accept"))
-		if err != nil {
-			if preferred := acc.FirstChoice(); preferred != nil {
-				contentType = preferred.ContentType
-			}
-		}
 		if contentType == "" {
-			// fall back on request media type
-			contentType = strings.ToLower(r.Request.Header.Get("Content-Type"))
+			// some APIs (e.g. Storage Data Plane) don't return a content type... so we'll assume from the Accept header
+			acc, err := accept.FromString(r.Request.Header.Get("Accept"))
+			if err != nil {
+				if preferred := acc.FirstChoice(); preferred != nil {
+					contentType = preferred.ContentType
+				}
+			}
+			if contentType == "" {
+				// fall back on request media type
+				contentType = strings.ToLower(r.Request.Header.Get("Content-Type"))
+			}
 		}
 	}
 
+	if contentType == "" {
+		return fmt.Errorf("could not determine Content-Type for response")
+	}
+
 	// Some APIs (e.g. Maintenance) return 200 without a body, don't unmarshal these
 	if r.ContentLength == 0 && (r.Body == nil || r.Body == http.NoBody) {
 		return nil
diff --git a/sdk/client/client_test.go b/sdk/client/client_test.go
index 668edf4c841..d344531b8c1 100644
--- a/sdk/client/client_test.go
+++ b/sdk/client/client_test.go
@@ -471,6 +471,38 @@ func TestUnmarshalXml(t *testing.T) {
 	}
 }
 
+func TestUnmarshalNilHeaders(t *testing.T) {
+	expected := []byte("any payload")
+	r := &Response{
+		Response: &http.Response{
+			Header: nil,
+			Body:   io.NopCloser(bytes.NewReader(expected)),
+		},
+	}
+	var unmarshaled []byte
+	if err := r.Unmarshal(&unmarshaled); err != nil {
+		if err.Error() != "could not determine Content-Type for response" {
+			t.Fatalf("unexpected error when unmarshaling: %+v", err)
+		}
+	} else {
+		t.Fatalf("expected an error but got no error")
+	}
+}
+
+func TestUnmarshalNilResponse(t *testing.T) {
+	r := &Response{
+		Response: nil,
+	}
+	var unmarshaled = make([]byte, 0)
+	if err := r.Unmarshal(&unmarshaled); err != nil {
+		if err.Error() != "could not unmarshal as the HTTP response was nil" {
+			t.Fatalf("unexpected error when unmarshaling: %+v", err)
+		}
+	} else {
+		t.Fatalf("expected an error but got no error")
+	}
+}
+
 func unmarshalResponse(body io.ReadCloser, unmarshal func(in []byte) error) error {
 	respBody, err := io.ReadAll(body)
 	if err != nil {