Move LRO and payload helpers to the internal module (Azure#20539)

* Move LRO and payload helpers to the internal module

This simplifies the creation of custom PollerHandler[T] by SDK authors.
The payload helpers were also moved as the LRO helpers depend on them.
The contents are directly copied from azcore with the exception of the
Payload func which adds an options parameter.

* fix up changelog

Author: jhendrixMSFT

sdk/internal/CHANGELOG.md

@@ -1,14 +1,10 @@
 # Release History
 
-## 1.2.1 (Unreleased)
+## 1.3.0 (2023-04-04)
 
 ### Features Added
-
-### Breaking Changes
-
-### Bugs Fixed
-
-### Other Changes
+* Added package `poller` which exports various LRO helpers to aid in the creation of custom `PollerHandler[T]`.
+* Added package `exported` which contains payload helpers needed by the `poller` package and exported in `azcore`.
 
 ## 1.2.0 (2023-03-02)
 

sdk/internal/exported/exported.go

@@ -0,0 +1,124 @@
+//go:build go1.18
+// +build go1.18
+
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package exported
+
+import (
+	"errors"
+	"io"
+	"net/http"
+)
+
+// HasStatusCode returns true if the Response's status code is one of the specified values.
+// Exported as runtime.HasStatusCode().
+func HasStatusCode(resp *http.Response, statusCodes ...int) bool {
+	if resp == nil {
+		return false
+	}
+	for _, sc := range statusCodes {
+		if resp.StatusCode == sc {
+			return true
+		}
+	}
+	return false
+}
+
+// PayloadOptions contains the optional values for the Payload func.
+// NOT exported but used by azcore.
+type PayloadOptions struct {
+	// BytesModifier receives the downloaded byte slice and returns an updated byte slice.
+	// Use this to modify the downloaded bytes in a payload (e.g. removing a BOM).
+	BytesModifier func([]byte) []byte
+}
+
+// Payload reads and returns the response body or an error.
+// On a successful read, the response body is cached.
+// Subsequent reads will access the cached value.
+// Exported as runtime.Payload() WITHOUT the opts parameter.
+func Payload(resp *http.Response, opts *PayloadOptions) ([]byte, error) {
+	modifyBytes := func(b []byte) []byte { return b }
+	if opts != nil && opts.BytesModifier != nil {
+		modifyBytes = opts.BytesModifier
+	}
+
+	// r.Body won't be a nopClosingBytesReader if downloading was skipped
+	if buf, ok := resp.Body.(*nopClosingBytesReader); ok {
+		bytesBody := modifyBytes(buf.Bytes())
+		buf.Set(bytesBody)
+		return bytesBody, nil
+	}
+
+	bytesBody, err := io.ReadAll(resp.Body)
+	resp.Body.Close()
+	if err != nil {
+		return nil, err
+	}
+
+	bytesBody = modifyBytes(bytesBody)
+	resp.Body = &nopClosingBytesReader{s: bytesBody}
+	return bytesBody, nil
+}
+
+// PayloadDownloaded returns true if the response body has already been downloaded.
+// This implies that the Payload() func above has been previously called.
+// NOT exported but used by azcore.
+func PayloadDownloaded(resp *http.Response) bool {
+	_, ok := resp.Body.(*nopClosingBytesReader)
+	return ok
+}
+
+// nopClosingBytesReader is an io.ReadSeekCloser around a byte slice.
+// It also provides direct access to the byte slice to avoid rereading.
+type nopClosingBytesReader struct {
+	s []byte
+	i int64
+}
+
+// Bytes returns the underlying byte slice.
+func (r *nopClosingBytesReader) Bytes() []byte {
+	return r.s
+}
+
+// Close implements the io.Closer interface.
+func (*nopClosingBytesReader) Close() error {
+	return nil
+}
+
+// Read implements the io.Reader interface.
+func (r *nopClosingBytesReader) Read(b []byte) (n int, err error) {
+	if r.i >= int64(len(r.s)) {
+		return 0, io.EOF
+	}
+	n = copy(b, r.s[r.i:])
+	r.i += int64(n)
+	return
+}
+
+// Set replaces the existing byte slice with the specified byte slice and resets the reader.
+func (r *nopClosingBytesReader) Set(b []byte) {
+	r.s = b
+	r.i = 0
+}
+
+// Seek implements the io.Seeker interface.
+func (r *nopClosingBytesReader) Seek(offset int64, whence int) (int64, error) {
+	var i int64
+	switch whence {
+	case io.SeekStart:
+		i = offset
+	case io.SeekCurrent:
+		i = r.i + offset
+	case io.SeekEnd:
+		i = int64(len(r.s)) + offset
+	default:
+		return 0, errors.New("nopClosingBytesReader: invalid whence")
+	}
+	if i < 0 {
+		return 0, errors.New("nopClosingBytesReader: negative position")
+	}
+	r.i = i
+	return i, nil
+}

sdk/internal/exported/exported_test.go

@@ -0,0 +1,94 @@
+//go:build go1.18
+// +build go1.18
+
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package exported
+
+import (
+	"io"
+	"net/http"
+	"strings"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestHasStatusCode(t *testing.T) {
+	require.False(t, HasStatusCode(nil, http.StatusAccepted))
+	require.False(t, HasStatusCode(&http.Response{}))
+	require.False(t, HasStatusCode(&http.Response{StatusCode: http.StatusBadGateway}, http.StatusBadRequest))
+	require.True(t, HasStatusCode(&http.Response{StatusCode: http.StatusOK}, http.StatusAccepted, http.StatusOK, http.StatusNoContent))
+}
+
+func TestPayload(t *testing.T) {
+	const val = "payload"
+	resp := &http.Response{
+		Body: io.NopCloser(strings.NewReader(val)),
+	}
+	b, err := Payload(resp, nil)
+	require.NoError(t, err)
+	if string(b) != val {
+		t.Fatalf("got %s, want %s", string(b), val)
+	}
+	b, err = Payload(resp, nil)
+	require.NoError(t, err)
+	if string(b) != val {
+		t.Fatalf("got %s, want %s", string(b), val)
+	}
+}
+
+func TestPayloadDownloaded(t *testing.T) {
+	resp := &http.Response{
+		Body: io.NopCloser(strings.NewReader("payload")),
+	}
+	require.False(t, PayloadDownloaded(resp))
+	_, err := Payload(resp, nil)
+	require.NoError(t, err)
+	require.True(t, PayloadDownloaded((resp)))
+}
+
+func TestPayloadBytesModifier(t *testing.T) {
+	resp := &http.Response{
+		Body: io.NopCloser(strings.NewReader("oldpayload")),
+	}
+	const newPayload = "newpayload"
+	b, err := Payload(resp, &PayloadOptions{
+		BytesModifier: func(b []byte) []byte { return []byte(newPayload) },
+	})
+	require.NoError(t, err)
+	require.EqualValues(t, newPayload, string(b))
+}
+
+func TestNopClosingBytesReader(t *testing.T) {
+	const val1 = "the data"
+	ncbr := &nopClosingBytesReader{s: []byte(val1)}
+	require.NotNil(t, ncbr.Bytes())
+	b, err := io.ReadAll(ncbr)
+	require.NoError(t, err)
+	require.EqualValues(t, val1, b)
+	const val2 = "something else"
+	ncbr.Set([]byte(val2))
+	b, err = io.ReadAll(ncbr)
+	require.NoError(t, err)
+	require.EqualValues(t, val2, b)
+	require.NoError(t, ncbr.Close())
+	// seek to beginning and read again
+	i, err := ncbr.Seek(0, io.SeekStart)
+	require.NoError(t, err)
+	require.Zero(t, i)
+	b, err = io.ReadAll(ncbr)
+	require.NoError(t, err)
+	require.EqualValues(t, val2, b)
+	// seek to middle from the end
+	i, err = ncbr.Seek(-4, io.SeekEnd)
+	require.NoError(t, err)
+	require.EqualValues(t, i, len(val2)-4)
+	b, err = io.ReadAll(ncbr)
+	require.NoError(t, err)
+	require.EqualValues(t, "else", b)
+	// underflow
+	_, err = ncbr.Seek(-int64(len(val2)+1), io.SeekCurrent)
+	require.Error(t, err)
+}

sdk/internal/poller/util.go

@@ -0,0 +1,155 @@
+//go:build go1.18
+// +build go1.18
+
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package poller
+
+import (
+	"encoding/json"
+	"errors"
+	"fmt"
+	"net/http"
+	"net/url"
+	"strings"
+
+	"github.com/Azure/azure-sdk-for-go/sdk/internal/exported"
+)
+
+// the well-known set of LRO status/provisioning state values.
+const (
+	StatusSucceeded  = "Succeeded"
+	StatusCanceled   = "Canceled"
+	StatusFailed     = "Failed"
+	StatusInProgress = "InProgress"
+)
+
+// these are non-conformant states that we've seen in the wild.
+// we support them for back-compat.
+const (
+	StatusCancelled = "Cancelled"
+	StatusCompleted = "Completed"
+)
+
+// IsTerminalState returns true if the LRO's state is terminal.
+func IsTerminalState(s string) bool {
+	return Failed(s) || Succeeded(s)
+}
+
+// Failed returns true if the LRO's state is terminal failure.
+func Failed(s string) bool {
+	return strings.EqualFold(s, StatusFailed) || strings.EqualFold(s, StatusCanceled) || strings.EqualFold(s, StatusCancelled)
+}
+
+// Succeeded returns true if the LRO's state is terminal success.
+func Succeeded(s string) bool {
+	return strings.EqualFold(s, StatusSucceeded) || strings.EqualFold(s, StatusCompleted)
+}
+
+// returns true if the LRO response contains a valid HTTP status code
+func StatusCodeValid(resp *http.Response) bool {
+	return exported.HasStatusCode(resp, http.StatusOK, http.StatusAccepted, http.StatusCreated, http.StatusNoContent)
+}
+
+// IsValidURL verifies that the URL is valid and absolute.
+func IsValidURL(s string) bool {
+	u, err := url.Parse(s)
+	return err == nil && u.IsAbs()
+}
+
+// ErrNoBody is returned if the response didn't contain a body.
+var ErrNoBody = errors.New("the response did not contain a body")
+
+// GetJSON reads the response body into a raw JSON object.
+// It returns ErrNoBody if there was no content.
+func GetJSON(resp *http.Response) (map[string]any, error) {
+	body, err := exported.Payload(resp, nil)
+	if err != nil {
+		return nil, err
+	}
+	if len(body) == 0 {
+		return nil, ErrNoBody
+	}
+	// unmarshall the body to get the value
+	var jsonBody map[string]any
+	if err = json.Unmarshal(body, &jsonBody); err != nil {
+		return nil, err
+	}
+	return jsonBody, nil
+}
+
+// provisioningState returns the provisioning state from the response or the empty string.
+func provisioningState(jsonBody map[string]any) string {
+	jsonProps, ok := jsonBody["properties"]
+	if !ok {
+		return ""
+	}
+	props, ok := jsonProps.(map[string]any)
+	if !ok {
+		return ""
+	}
+	rawPs, ok := props["provisioningState"]
+	if !ok {
+		return ""
+	}
+	ps, ok := rawPs.(string)
+	if !ok {
+		return ""
+	}
+	return ps
+}
+
+// status returns the status from the response or the empty string.
+func status(jsonBody map[string]any) string {
+	rawStatus, ok := jsonBody["status"]
+	if !ok {
+		return ""
+	}
+	status, ok := rawStatus.(string)
+	if !ok {
+		return ""
+	}
+	return status
+}
+
+// GetStatus returns the LRO's status from the response body.
+// Typically used for Azure-AsyncOperation flows.
+// If there is no status in the response body the empty string is returned.
+func GetStatus(resp *http.Response) (string, error) {
+	jsonBody, err := GetJSON(resp)
+	if err != nil {
+		return "", err
+	}
+	return status(jsonBody), nil
+}
+
+// GetProvisioningState returns the LRO's state from the response body.
+// If there is no state in the response body the empty string is returned.
+func GetProvisioningState(resp *http.Response) (string, error) {
+	jsonBody, err := GetJSON(resp)
+	if err != nil {
+		return "", err
+	}
+	return provisioningState(jsonBody), nil
+}
+
+// GetResourceLocation returns the LRO's resourceLocation value from the response body.
+// Typically used for Operation-Location flows.
+// If there is no resourceLocation in the response body the empty string is returned.
+func GetResourceLocation(resp *http.Response) (string, error) {
+	jsonBody, err := GetJSON(resp)
+	if err != nil {
+		return "", err
+	}
+	v, ok := jsonBody["resourceLocation"]
+	if !ok {
+		// it might be ok if the field doesn't exist, the caller must make that determination
+		return "", nil
+	}
+	vv, ok := v.(string)
+	if !ok {
+		return "", fmt.Errorf("the resourceLocation value %v was not in string format", v)
+	}
+	return vv, nil
+}