@@ -4,9 +4,11 @@ import (
44 "bytes"
55 "context"
66 "errors"
7- "github.com/aws/aws-sdk-go/aws"
8- "github.com/aws/aws-sdk-go/aws/awserr"
9- "github.com/aws/aws-sdk-go/service/s3"
7+ "github.com/aws/aws-sdk-go-v2/aws"
8+ "github.com/aws/aws-sdk-go-v2/service/s3"
9+ "github.com/aws/aws-sdk-go-v2/service/s3/types"
10+ "github.com/aws/smithy-go"
11+ smithyhttp "github.com/aws/smithy-go/transport/http"
1012 "github.com/cirruslabs/cirrus-cli/pkg/api"
1113 "github.com/samber/lo"
1214 "github.com/stretchr/testify/require"
@@ -24,20 +26,20 @@ import (
2426)
2527
2628type cirrusCIMock struct {
27- s3Client * s3.S3
29+ s3Client * s3.Client
2830 s3Bucket * string
2931
3032 api.UnimplementedCirrusCIServiceServer
3133 bytestream.UnimplementedByteStreamServer
3234}
3335
34- func newCirrusCIMock (t * testing.T , s3Client * s3.S3 ) * cirrusCIMock {
36+ func newCirrusCIMock (t * testing.T , s3Client * s3.Client ) * cirrusCIMock {
3537 mock := & cirrusCIMock {
3638 s3Client : s3Client ,
3739 s3Bucket : aws .String ("test" ),
3840 }
3941
40- _ , err := mock .s3Client .CreateBucket (& s3.CreateBucketInput {
42+ _ , err := mock .s3Client .CreateBucket (context . Background (), & s3.CreateBucketInput {
4143 Bucket : mock .s3Bucket ,
4244 })
4345 require .NoError (t , err )
@@ -113,7 +115,7 @@ func (mock *cirrusCIMock) Write(stream bytestream.ByteStream_WriteServer) error
113115 return status .Error (codes .FailedPrecondition , "attempted to upload cache without specifying resource name" )
114116 }
115117
116- _ , err := mock .s3Client .PutObjectWithContext (stream .Context (), & s3.PutObjectInput {
118+ _ , err := mock .s3Client .PutObject (stream .Context (), & s3.PutObjectInput {
117119 Bucket : mock .s3Bucket ,
118120 Key : aws .String (resourceName ),
119121 Body : bytes .NewReader (buf .Bytes ()),
@@ -128,13 +130,12 @@ func (mock *cirrusCIMock) Write(stream bytestream.ByteStream_WriteServer) error
128130}
129131
130132func (mock * cirrusCIMock ) Read (request * bytestream.ReadRequest , stream bytestream.ByteStream_ReadServer ) error {
131- result , err := mock .s3Client .GetObjectWithContext (stream .Context (), & s3.GetObjectInput {
133+ result , err := mock .s3Client .GetObject (stream .Context (), & s3.GetObjectInput {
132134 Bucket : mock .s3Bucket ,
133135 Key : aws .String (request .ResourceName ),
134136 })
135137 if err != nil {
136- var aerr awserr.Error
137- if errors .As (err , & aerr ) && aerr .Code () == s3 .ErrCodeNoSuchKey {
138+ if isS3NotFound (err ) {
138139 return status .Errorf (codes .NotFound , "cache entry for key %s is not found" ,
139140 request .ResourceName )
140141 }
@@ -165,54 +166,51 @@ func (mock *cirrusCIMock) Read(request *bytestream.ReadRequest, stream bytestrea
165166}
166167
167168func (mock * cirrusCIMock ) GenerateCacheUploadURL (ctx context.Context , request * api.CacheKey ) (* api.GenerateURLResponse , error ) {
168- putObjectRequest , _ := mock .s3Client .PutObjectRequest (& s3.PutObjectInput {
169+ presignClient := s3 .NewPresignClient (mock .s3Client )
170+ presignResult , err := presignClient .PresignPutObject (ctx , & s3.PutObjectInput {
169171 Bucket : mock .s3Bucket ,
170172 Key : aws .String (request .CacheKey ),
171- })
172-
173- url , _ , err := putObjectRequest .PresignRequest (10 * time .Minute )
173+ }, s3 .WithPresignExpires (10 * time .Minute ))
174174 if err != nil {
175175 return nil , err
176176 }
177177
178178 return & api.GenerateURLResponse {
179- Url : url ,
179+ Url : presignResult . URL ,
180180 }, nil
181181}
182182
183183func (mock * cirrusCIMock ) GenerateCacheDownloadURLs (
184- _ context.Context ,
184+ ctx context.Context ,
185185 request * api.CacheKey ,
186186) (* api.GenerateURLsResponse , error ) {
187- getObjectRequest , _ := mock .s3Client .GetObjectRequest (& s3.GetObjectInput {
187+ presignClient := s3 .NewPresignClient (mock .s3Client )
188+ presignResult , err := presignClient .PresignGetObject (ctx , & s3.GetObjectInput {
188189 Bucket : mock .s3Bucket ,
189190 Key : aws .String (request .CacheKey ),
190- })
191-
192- url , _ , err := getObjectRequest .PresignRequest (10 * time .Minute )
191+ }, s3 .WithPresignExpires (10 * time .Minute ))
193192 if err != nil {
194193 return nil , err
195194 }
196195
197196 return & api.GenerateURLsResponse {
198- Urls : []string {url },
197+ Urls : []string {presignResult . URL },
199198 }, nil
200199}
201200
202201func (mock * cirrusCIMock ) CacheInfo (ctx context.Context , request * api.CacheInfoRequest ) (* api.CacheInfoResponse , error ) {
203- result , err := mock .s3Client .HeadObjectWithContext (ctx , & s3.HeadObjectInput {
202+ result , err := mock .s3Client .HeadObject (ctx , & s3.HeadObjectInput {
204203 Bucket : mock .s3Bucket ,
205204 Key : aws .String (request .CacheKey ),
206205 })
207206 if err != nil {
208- var requestFailure awserr.RequestFailure
209- if errors .As (err , & requestFailure ) && requestFailure .StatusCode () == http .StatusNotFound {
207+ if isS3NotFound (err ) {
210208 // Try to match cache entry by key prefixes as a fallback
211209 for _ , cacheKeyPrefix := range request .CacheKeyPrefixes {
212- result , err := mock .s3Client .ListObjectsWithContext (ctx , & s3.ListObjectsInput {
210+ result , err := mock .s3Client .ListObjects (ctx , & s3.ListObjectsInput {
213211 Bucket : mock .s3Bucket ,
214212 Prefix : aws .String (cacheKeyPrefix ),
215- MaxKeys : aws .Int64 (1 ),
213+ MaxKeys : aws .Int32 (1 ),
216214 })
217215 if err != nil {
218216 return nil , err
@@ -246,7 +244,7 @@ func (mock *cirrusCIMock) CacheInfo(ctx context.Context, request *api.CacheInfoR
246244}
247245
248246func (mock * cirrusCIMock ) MultipartCacheUploadCreate (ctx context.Context , request * api.CacheKey ) (* api.MultipartCacheUploadCreateResponse , error ) {
249- result , err := mock .s3Client .CreateMultipartUploadWithContext (ctx , & s3.CreateMultipartUploadInput {
247+ result , err := mock .s3Client .CreateMultipartUpload (ctx , & s3.CreateMultipartUploadInput {
250248 Bucket : mock .s3Bucket ,
251249 Key : aws .String (request .CacheKey ),
252250 })
@@ -260,41 +258,40 @@ func (mock *cirrusCIMock) MultipartCacheUploadCreate(ctx context.Context, reques
260258}
261259
262260func (mock * cirrusCIMock ) MultipartCacheUploadPart (ctx context.Context , request * api.MultipartCacheUploadPartRequest ) (* api.GenerateURLResponse , error ) {
263- uploadPartRequest , _ := mock .s3Client .UploadPartRequest (& s3.UploadPartInput {
261+ presignClient := s3 .NewPresignClient (mock .s3Client )
262+ presignResult , err := presignClient .PresignUploadPart (ctx , & s3.UploadPartInput {
264263 Bucket : mock .s3Bucket ,
265264 Key : aws .String (request .CacheKey .CacheKey ),
266265 UploadId : aws .String (request .UploadId ),
267- PartNumber : aws .Int64 (int64 (request .PartNumber )),
268- })
269-
270- url , headers , err := uploadPartRequest .PresignRequest (10 * time .Minute )
266+ PartNumber : aws .Int32 (int32 (request .PartNumber )),
267+ }, s3 .WithPresignExpires (10 * time .Minute ))
271268 if err != nil {
272269 return nil , err
273270 }
274271
275272 return & api.GenerateURLResponse {
276- Url : url ,
277- ExtraHeaders : lo .MapEntries (headers , func (key string , value []string ) (string , string ) {
273+ Url : presignResult . URL ,
274+ ExtraHeaders : lo .MapEntries (presignResult . SignedHeader , func (key string , value []string ) (string , string ) {
278275 return key , value [0 ]
279276 }),
280277 }, nil
281278}
282279
283280func (mock * cirrusCIMock ) MultipartCacheUploadCommit (ctx context.Context , request * api.MultipartCacheUploadCommitRequest ) (* emptypb.Empty , error ) {
284- var parts []* s3 .CompletedPart
281+ var parts []types .CompletedPart
285282
286283 for _ , part := range request .Parts {
287- parts = append (parts , & s3 .CompletedPart {
288- PartNumber : aws .Int64 ( int64 (part .PartNumber )),
284+ parts = append (parts , types .CompletedPart {
285+ PartNumber : aws .Int32 ( int32 (part .PartNumber )),
289286 ETag : aws .String (part .Etag ),
290287 })
291288 }
292289
293- _ , err := mock .s3Client .CompleteMultipartUploadWithContext (ctx , & s3.CompleteMultipartUploadInput {
290+ _ , err := mock .s3Client .CompleteMultipartUpload (ctx , & s3.CompleteMultipartUploadInput {
294291 Bucket : mock .s3Bucket ,
295292 Key : aws .String (request .CacheKey .CacheKey ),
296293 UploadId : aws .String (request .UploadId ),
297- MultipartUpload : & s3 .CompletedMultipartUpload {
294+ MultipartUpload : & types .CompletedMultipartUpload {
298295 Parts : parts ,
299296 },
300297 })
@@ -304,3 +301,30 @@ func (mock *cirrusCIMock) MultipartCacheUploadCommit(ctx context.Context, reques
304301
305302 return & emptypb.Empty {}, nil
306303}
304+
305+ func isS3NotFound (err error ) bool {
306+ var noSuchKey * types.NoSuchKey
307+ if errors .As (err , & noSuchKey ) {
308+ return true
309+ }
310+
311+ var notFound * types.NotFound
312+ if errors .As (err , & notFound ) {
313+ return true
314+ }
315+
316+ var apiErr smithy.APIError
317+ if errors .As (err , & apiErr ) {
318+ switch apiErr .ErrorCode () {
319+ case "NoSuchKey" , "NotFound" :
320+ return true
321+ }
322+ }
323+
324+ var responseErr * smithyhttp.ResponseError
325+ if errors .As (err , & responseErr ) && responseErr .HTTPStatusCode () == http .StatusNotFound {
326+ return true
327+ }
328+
329+ return false
330+ }
0 commit comments