internal/sdk: Fix SDK's UUID utility to handle partial read (aws#536)

Fixes the SDK's UUID utility to correctly handle partial reads from its
crypto rand source. This error was sometimes causing the SDK's
InvocationID value to fail to be obtained, due to a partial read from
crypto.Rand.

Fix aws#534

CHANGELOG_PENDING.md

@@ -36,3 +36,6 @@ SDK Enhancements
 
 SDK Bugs
 ---
+* `internal/sdk`: Fix SDK's UUID utility to handle partial read ([#536](https://github.com/aws/aws-sdk-go-v2/pull/536))
+  * Fixes the SDK's UUID utility to correctly handle partial reads from its crypto rand source. This error was sometimes causing the SDK's InvocationID value to fail to be obtained, due to a partial read from crypto.Rand.
+  * Fix [#534](https://github.com/aws/aws-sdk-go-v2/issues/534)

internal/sdk/uuid.go

@@ -9,10 +9,14 @@ import (
 // UUIDVersion4 returns a Version 4 random UUID from the byte slice provided
 func UUIDVersion4() (string, error) {
 	b := make([]byte, 16)
-	if n, err := rand.Reader.Read(b); err != nil {
-		return "", fmt.Errorf("unable to get random bytes for UUID, %w", err)
-	} else if n != len(b) {
-		return "", fmt.Errorf("unable to get 16 bytes for UUID, got %d", n)
+
+	var offset int
+	for offset < len(b) {
+		n, err := rand.Reader.Read(b[offset:])
+		if err != nil {
+			return "", fmt.Errorf("unable to get random bytes for UUID, %w", err)
+		}
+		offset += n
 	}
 
 	return uuidVersion4(b), nil

internal/sdk/uuid_test.go

@@ -1,19 +1,75 @@
 package sdk
 
-import "testing"
+import (
+	"io"
+	"strings"
+	"testing"
+	"testing/iotest"
 
-func TestUUIDVersion4(t *testing.T) {
-	uuid := uuidVersion4(make([]byte, 16))
-	if e, a := `00000000-0000-4000-8000-000000000000`, uuid; e != a {
-		t.Errorf("expect %v uuid, got %v", e, a)
+	"github.com/aws/aws-sdk-go-v2/internal/rand"
+)
+
+type byteReader byte
+
+func (b byteReader) Read(p []byte) (n int, err error) {
+	for i := 0; i < len(p); i++ {
+		p[i] = byte(b)
 	}
+	return len(p), nil
+}
+
+type errorReader struct{ err error }
+
+func (e errorReader) Read(p []byte) (n int, err error) {
+	return 0, e.err
+}
+
+func TestUUIDVersion4(t *testing.T) {
+	origReader := rand.Reader
+	defer func() { rand.Reader = origReader }()
 
-	b := make([]byte, 16)
-	for i := 0; i < len(b); i++ {
-		b[i] = 1
+	cases := map[string]struct {
+		Expect string
+		Reader io.Reader
+		Err    string
+	}{
+		"0x00": {
+			Expect: `00000000-0000-4000-8000-000000000000`,
+			Reader: byteReader(0),
+		},
+		"0x01": {
+			Expect: `01010101-0101-4101-8101-010101010101`,
+			Reader: byteReader(1),
+		},
+		"partial": {
+			Expect: `01010101-0101-4101-8101-010101010101`,
+			Reader: iotest.HalfReader(byteReader(1)),
+		},
+		"error": {
+			Reader: errorReader{err: io.ErrUnexpectedEOF},
+			Err:    io.ErrUnexpectedEOF.Error(),
+		},
 	}
-	uuid = uuidVersion4(b)
-	if e, a := `01010101-0101-4101-8101-010101010101`, uuid; e != a {
-		t.Errorf("expect %v uuid, got %v", e, a)
+
+	for name, c := range cases {
+		t.Run(name, func(t *testing.T) {
+			rand.Reader = c.Reader
+
+			uuid, err := UUIDVersion4()
+			if len(c.Err) != 0 {
+				if err == nil {
+					t.Fatalf("expect error, got none")
+				}
+				if e, a := c.Err, err.Error(); !strings.Contains(a, e) {
+					t.Fatalf("expect %q in error, %q", e, a)
+				}
+			} else if err != nil {
+				t.Fatalf("expect no error, got %v", err)
+			}
+
+			if e, a := c.Expect, uuid; e != a {
+				t.Errorf("expect %v uuid, got %v", e, a)
+			}
+		})
 	}
 }