File: treehash_test.go

package info (click to toggle)
golang-github-aws-aws-sdk-go-v2 1.17.1-3
  • links: PTS, VCS
  • area: main
  • in suites: bookworm
  • size: 384,244 kB
  • sloc: java: 13,538; makefile: 400; sh: 137
file content (127 lines) | stat: -rw-r--r-- 3,564 bytes parent folder | download
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
package customizations

import (
	"bytes"
	"context"
	"github.com/aws/smithy-go/middleware"
	smithyhttp "github.com/aws/smithy-go/transport/http"
	"io"
	"strings"
	"testing"
)

func TestTreeHashMiddleware(t *testing.T) {
	m := TreeHash{}

	cases := map[string]struct {
		Body       io.Reader
		SetRequest func(r *smithyhttp.Request) *smithyhttp.Request
		Linear     string
		Tree       string
		ExpectErr  string
	}{
		"all zeroes": {
			Body: func() io.ReadSeeker {
				buf := make([]byte, 5767168) // 5.5MB buffer
				for i := range buf {
					buf[i] = '0' // Fill with zero characters
				}
				return bytes.NewReader(buf)
			}(),
			Linear: "68aff0c5a91aa0491752bfb96e3fef33eb74953804f6a2f7b708d5bcefa8ff6b",
			Tree:   "154e26c78fd74d0c2c9b3cc4644191619dc4f2cd539ae2a74d5fd07957a3ee6a",
		},
		"non-seekable stream without precomputed hash": {
			Body: func() io.Reader {
				buf := make([]byte, 5767168)
				for i := range buf {
					buf[i] = '0'
				}
				return bytes.NewBuffer(buf)
			}(),
			ExpectErr: "tree hash",
		},
		"non-seekable stream with precomputed hash": {
			Body: func() io.Reader {
				buf := make([]byte, 5767168)
				for i := range buf {
					buf[i] = '0'
				}
				return bytes.NewBuffer(buf)
			}(),
			SetRequest: func(r *smithyhttp.Request) *smithyhttp.Request {
				r.Header.Set("X-Amz-Content-Sha256", "precomputed")
				r.Header.Set("X-Amz-Sha256-Tree-Hash", "precomputed")
				return r
			},
			// These values aren't actually correct, but they are what we're explicitly setting.
			// If the middleware were to try to compute the values anyway, the checks would fail.
			Linear: "precomputed",
			Tree:   "precomputed",
		},
	}

	for name, c := range cases {
		t.Run(name, func(t *testing.T) {
			_, _, err := m.HandleFinalize(context.Background(),
				middleware.FinalizeInput{
					Request: func() *smithyhttp.Request {
						r := smithyhttp.NewStackRequest()
						req, ok := r.(*smithyhttp.Request)
						if !ok || req == nil {
							t.Fatalf("expect smithy request, got %T", r)
						}
						if c.SetRequest != nil {
							req = c.SetRequest(req)
						}
						req, err := req.SetStream(c.Body)
						if err != nil {
							t.Fatalf("expect no error, got %v", err)
						}
						return req
					}(),
				},
				middleware.FinalizeHandlerFunc(
					func(ctx context.Context, input middleware.FinalizeInput) (
						output middleware.FinalizeOutput, metadata middleware.Metadata, err error,
					) {
						req, ok := input.Request.(*smithyhttp.Request)
						if !ok || req == nil {
							t.Fatalf("expect smithy request, got %T", input.Request)
						}

						actualLinear := req.Header.Get("X-Amz-Content-Sha256")
						if actualLinear != c.Linear {
							t.Fatalf("expected linear hash to be \"%s\" but was \"%s\"", c.Linear, actualLinear)
						}

						actualTree := req.Header.Get("X-Amz-Sha256-Tree-Hash")
						if actualTree != c.Tree {
							t.Fatalf("expected tree hash to be \"%s\" but was \"%s\"", c.Tree, actualTree)
						}

						if body, ok := c.Body.(io.ReadSeeker); ok {
							n, _ := body.Seek(0, io.SeekCurrent)
							if n != 0 {
								t.Fatalf("expected body to be rewound, but was at position %d", n)
							}

						}
						return output, metadata, err
					}),
			)
			if len(c.ExpectErr) != 0 {
				if err == nil {
					t.Fatalf("expect error, got none")
				}
				if e, a := c.ExpectErr, err.Error(); !strings.Contains(a, e) {
					t.Fatalf("expect error to contain %v, got %v", e, a)
				}
			} else {
				if err != nil {
					t.Fatalf("expect no error, got %v", err)
				}
			}
		})
	}
}