vitess.io/vitess@v0.16.2/go/vt/mysqlctl/s3backupstorage/s3_test.go (about)

     1  package s3backupstorage
     2  
     3  import (
     4  	"crypto/md5"
     5  	"crypto/rand"
     6  	"encoding/base64"
     7  	"errors"
     8  	"net/http"
     9  	"os"
    10  	"testing"
    11  
    12  	"github.com/aws/aws-sdk-go/aws"
    13  	"github.com/aws/aws-sdk-go/aws/request"
    14  	"github.com/aws/aws-sdk-go/service/s3"
    15  	"github.com/aws/aws-sdk-go/service/s3/s3iface"
    16  	"github.com/stretchr/testify/assert"
    17  	"github.com/stretchr/testify/require"
    18  )
    19  
    20  type s3ErrorClient struct{ s3iface.S3API }
    21  
    22  func (s3errclient *s3ErrorClient) PutObjectRequest(in *s3.PutObjectInput) (*request.Request, *s3.PutObjectOutput) {
    23  	req := request.Request{
    24  		HTTPRequest: &http.Request{},          // without this we segfault \_(ツ)_/¯ (see https://github.com/aws/aws-sdk-go/blob/v1.28.8/aws/request/request_context.go#L13)
    25  		Error:       errors.New("some error"), // this forces req.Send() (which is called by the uploader) to always return non-nil error
    26  	}
    27  
    28  	return &req, &s3.PutObjectOutput{}
    29  }
    30  
    31  func TestAddFileError(t *testing.T) {
    32  	bh := &S3BackupHandle{client: &s3ErrorClient{}, bs: &S3BackupStorage{}, readOnly: false}
    33  
    34  	wc, err := bh.AddFile(aws.BackgroundContext(), "somefile", 100000)
    35  	require.NoErrorf(t, err, "AddFile() expected no error, got %s", err)
    36  	assert.NotNil(t, wc, "AddFile() expected non-nil WriteCloser")
    37  
    38  	n, err := wc.Write([]byte("here are some bytes"))
    39  	require.NoErrorf(t, err, "TestAddFile() could not write to uploader, got %d bytes written, err %s", n, err)
    40  
    41  	err = wc.Close()
    42  	require.NoErrorf(t, err, "TestAddFile() could not close writer, got %s", err)
    43  
    44  	bh.waitGroup.Wait() // wait for the goroutine to finish, at which point it should have recorded an error
    45  
    46  	require.Equal(t, bh.HasErrors(), true, "AddFile() expected bh to record async error but did not")
    47  }
    48  
    49  func TestNoSSE(t *testing.T) {
    50  	sseData := S3ServerSideEncryption{}
    51  	err := sseData.init()
    52  	require.NoErrorf(t, err, "init() expected to succeed")
    53  
    54  	assert.Nil(t, sseData.awsAlg, "awsAlg expected to be nil")
    55  	assert.Nil(t, sseData.customerAlg, "customerAlg expected to be nil")
    56  	assert.Nil(t, sseData.customerKey, "customerKey expected to be nil")
    57  	assert.Nil(t, sseData.customerMd5, "customerMd5 expected to be nil")
    58  
    59  	sseData.reset()
    60  	require.NoErrorf(t, err, "reset() expected to succeed")
    61  }
    62  
    63  func TestSSEAws(t *testing.T) {
    64  	sse = "aws:kms"
    65  	sseData := S3ServerSideEncryption{}
    66  	err := sseData.init()
    67  	require.NoErrorf(t, err, "init() expected to succeed")
    68  
    69  	assert.Equal(t, aws.String("aws:kms"), sseData.awsAlg, "awsAlg expected to be aws:kms")
    70  	assert.Nil(t, sseData.customerAlg, "customerAlg expected to be nil")
    71  	assert.Nil(t, sseData.customerKey, "customerKey expected to be nil")
    72  	assert.Nil(t, sseData.customerMd5, "customerMd5 expected to be nil")
    73  
    74  	sseData.reset()
    75  	require.NoErrorf(t, err, "reset() expected to succeed")
    76  
    77  	assert.Nil(t, sseData.awsAlg, "awsAlg expected to be nil")
    78  	assert.Nil(t, sseData.customerAlg, "customerAlg expected to be nil")
    79  	assert.Nil(t, sseData.customerKey, "customerKey expected to be nil")
    80  	assert.Nil(t, sseData.customerMd5, "customerMd5 expected to be nil")
    81  }
    82  
    83  func TestSSECustomerFileNotFound(t *testing.T) {
    84  	tempFile, err := os.CreateTemp("", "filename")
    85  	require.NoErrorf(t, err, "TempFile() expected to succeed")
    86  	defer os.Remove(tempFile.Name())
    87  
    88  	err = tempFile.Close()
    89  	require.NoErrorf(t, err, "Close() expected to succeed")
    90  
    91  	err = os.Remove(tempFile.Name())
    92  	require.NoErrorf(t, err, "Remove() expected to succeed")
    93  
    94  	sse = sseCustomerPrefix + tempFile.Name()
    95  	sseData := S3ServerSideEncryption{}
    96  	err = sseData.init()
    97  	require.Errorf(t, err, "init() expected to fail")
    98  }
    99  
   100  func TestSSECustomerFileBinaryKey(t *testing.T) {
   101  	tempFile, err := os.CreateTemp("", "filename")
   102  	require.NoErrorf(t, err, "TempFile() expected to succeed")
   103  	defer os.Remove(tempFile.Name())
   104  
   105  	randomKey := make([]byte, 32)
   106  	_, err = rand.Read(randomKey)
   107  	require.NoErrorf(t, err, "Read() expected to succeed")
   108  	_, err = tempFile.Write(randomKey)
   109  	require.NoErrorf(t, err, "Write() expected to succeed")
   110  	err = tempFile.Close()
   111  	require.NoErrorf(t, err, "Close() expected to succeed")
   112  
   113  	sse = sseCustomerPrefix + tempFile.Name()
   114  	sseData := S3ServerSideEncryption{}
   115  	err = sseData.init()
   116  	require.NoErrorf(t, err, "init() expected to succeed")
   117  
   118  	assert.Nil(t, sseData.awsAlg, "awsAlg expected to be nil")
   119  	assert.Equal(t, aws.String("AES256"), sseData.customerAlg, "customerAlg expected to be AES256")
   120  	assert.Equal(t, aws.String(string(randomKey)), sseData.customerKey, "customerKey expected to be equal to the generated randomKey")
   121  	md5Hash := md5.Sum(randomKey)
   122  	assert.Equal(t, aws.String(base64.StdEncoding.EncodeToString(md5Hash[:])), sseData.customerMd5, "customerMd5 expected to be equal to the customerMd5 hash of the generated randomKey")
   123  
   124  	sseData.reset()
   125  	require.NoErrorf(t, err, "reset() expected to succeed")
   126  
   127  	assert.Nil(t, sseData.awsAlg, "awsAlg expected to be nil")
   128  	assert.Nil(t, sseData.customerAlg, "customerAlg expected to be nil")
   129  	assert.Nil(t, sseData.customerKey, "customerKey expected to be nil")
   130  	assert.Nil(t, sseData.customerMd5, "customerMd5 expected to be nil")
   131  }
   132  
   133  func TestSSECustomerFileBase64Key(t *testing.T) {
   134  	tempFile, err := os.CreateTemp("", "filename")
   135  	require.NoErrorf(t, err, "TempFile() expected to succeed")
   136  	defer os.Remove(tempFile.Name())
   137  
   138  	randomKey := make([]byte, 32)
   139  	_, err = rand.Read(randomKey)
   140  	require.NoErrorf(t, err, "Read() expected to succeed")
   141  
   142  	base64Key := base64.StdEncoding.EncodeToString(randomKey[:])
   143  	_, err = tempFile.WriteString(base64Key)
   144  	require.NoErrorf(t, err, "WriteString() expected to succeed")
   145  	err = tempFile.Close()
   146  	require.NoErrorf(t, err, "Close() expected to succeed")
   147  
   148  	sse = sseCustomerPrefix + tempFile.Name()
   149  	sseData := S3ServerSideEncryption{}
   150  	err = sseData.init()
   151  	require.NoErrorf(t, err, "init() expected to succeed")
   152  
   153  	assert.Nil(t, sseData.awsAlg, "awsAlg expected to be nil")
   154  	assert.Equal(t, aws.String("AES256"), sseData.customerAlg, "customerAlg expected to be AES256")
   155  	assert.Equal(t, aws.String(string(randomKey)), sseData.customerKey, "customerKey expected to be equal to the generated randomKey")
   156  	md5Hash := md5.Sum(randomKey)
   157  	assert.Equal(t, aws.String(base64.StdEncoding.EncodeToString(md5Hash[:])), sseData.customerMd5, "customerMd5 expected to be equal to the customerMd5 hash of the generated randomKey")
   158  
   159  	sseData.reset()
   160  	require.NoErrorf(t, err, "reset() expected to succeed")
   161  
   162  	assert.Nil(t, sseData.awsAlg, "awsAlg expected to be nil")
   163  	assert.Nil(t, sseData.customerAlg, "customerAlg expected to be nil")
   164  	assert.Nil(t, sseData.customerKey, "customerKey expected to be nil")
   165  	assert.Nil(t, sseData.customerMd5, "customerMd5 expected to be nil")
   166  }