go.fuchsia.dev/infra@v0.0.0-20240507153436-9b593402251b/cmd/gcs-util/lib/lib_test.go (about)

     1  // Copyright 2022 The Fuchsia Authors. All rights reserved.
     2  // Use of this source code is governed by a BSD-style license that can be
     3  // found in the LICENSE file.
     4  
     5  package lib
     6  
     7  import (
     8  	"context"
     9  	"crypto/ed25519"
    10  	"crypto/x509"
    11  	"encoding/base64"
    12  	"encoding/pem"
    13  	"os"
    14  	"path/filepath"
    15  	"reflect"
    16  	"testing"
    17  	"time"
    18  
    19  	"cloud.google.com/go/storage"
    20  
    21  	"go.chromium.org/luci/common/errors"
    22  	"go.chromium.org/luci/common/retry"
    23  	"go.chromium.org/luci/common/retry/transient"
    24  
    25  	"go.fuchsia.dev/infra/cmd/gcs-util/types"
    26  )
    27  
    28  func TestSign(t *testing.T) {
    29  	dir := t.TempDir()
    30  	var pkey ed25519.PrivateKey
    31  	dataFile := filepath.Join(dir, "data")
    32  
    33  	uploads := []types.Upload{{
    34  		Source:      dataFile,
    35  		Destination: "data",
    36  		Signed:      true,
    37  	}}
    38  	actual, err := Sign(uploads, pkey)
    39  	if err != nil {
    40  		t.Errorf("failed to sign uploads: %v", err)
    41  	}
    42  	if !reflect.DeepEqual(actual, uploads) {
    43  		t.Errorf("missing pkey should return unmodified uploads; got %v", actual)
    44  	}
    45  
    46  	_, pkey, err = ed25519.GenerateKey(nil)
    47  	if err != nil {
    48  		t.Errorf("failed to generate key: %v", err)
    49  	}
    50  	actual, err = Sign(uploads, pkey)
    51  	if err != nil {
    52  		t.Errorf("failed to sign uploads: %v", err)
    53  	}
    54  	if !reflect.DeepEqual(actual, uploads) {
    55  		t.Errorf("missing data file should return unmodified uploads; got %v", actual)
    56  	}
    57  
    58  	err = os.WriteFile(dataFile, []byte("data"), 0o400)
    59  	if err != nil {
    60  		t.Errorf("failed to write data file: %v", err)
    61  	}
    62  	expectedSignature := base64.StdEncoding.EncodeToString(ed25519.Sign(pkey, []byte("data")))
    63  	expected := []types.Upload{{
    64  		Source:      dataFile,
    65  		Destination: "data",
    66  		Signed:      true,
    67  		Metadata: map[string]string{
    68  			signatureKey: expectedSignature,
    69  		},
    70  	}}
    71  	actual, err = Sign(uploads, pkey)
    72  	if err != nil {
    73  		t.Errorf("failed to sign uploads: %v", err)
    74  	}
    75  	if !reflect.DeepEqual(actual, expected) {
    76  		t.Errorf("expected: %v, actual: %v", expected, actual)
    77  	}
    78  }
    79  
    80  func TestPublicKeyUpload(t *testing.T) {
    81  	upload, err := PublicKeyUpload([]byte{})
    82  	if err == nil {
    83  		t.Errorf("nil public key should err")
    84  	}
    85  	if upload != nil {
    86  		t.Errorf("nil public key should return nil pubkey upload; got: %v", upload)
    87  	}
    88  
    89  	expectedPubkey, pkey, err := ed25519.GenerateKey(nil)
    90  	if err != nil {
    91  		t.Errorf("failed to generate key: %v", err)
    92  	}
    93  
    94  	upload, err = PublicKeyUpload(pkey.Public().(ed25519.PublicKey))
    95  	if err != nil {
    96  		t.Errorf("failed to derive public key: %v", err)
    97  	}
    98  	if upload == nil || len(upload.Contents) == 0 {
    99  		t.Errorf("got empty pubkey data")
   100  	}
   101  	if upload.Destination != releasePubkeyFilename {
   102  		t.Errorf("incorrect destination; got: %s, expected: %s", upload.Destination, releasePubkeyFilename)
   103  	}
   104  	block, _ := pem.Decode(upload.Contents)
   105  	if block.Bytes == nil {
   106  		t.Errorf("failed to decode public key from pem")
   107  	}
   108  	pubkey, err := x509.ParsePKIXPublicKey(block.Bytes)
   109  	if err != nil {
   110  		t.Errorf("failed to parse public key from DER bytes")
   111  	}
   112  	if string(pubkey.(ed25519.PublicKey)) != string(expectedPubkey) {
   113  		t.Errorf("got: %s, expected: %s", string(pubkey.(ed25519.PublicKey)), string(expectedPubkey))
   114  	}
   115  }
   116  
   117  func TestRetry(t *testing.T) {
   118  	ctx := context.Background()
   119  
   120  	maxRetries := 10 // arbitrary
   121  	policy := transient.Only(func() retry.Iterator {
   122  		return &retry.ExponentialBackoff{
   123  			Limited: retry.Limited{
   124  				Delay:   0 * time.Second,
   125  				Retries: maxRetries,
   126  			},
   127  			Multiplier: 2,
   128  		}
   129  	})
   130  
   131  	tests := []struct {
   132  		name             string
   133  		err              error
   134  		expectErr        bool
   135  		expectedAttempts int
   136  	}{
   137  		{
   138  			name:             "pass",
   139  			err:              nil,
   140  			expectErr:        false,
   141  			expectedAttempts: 1,
   142  		},
   143  		{
   144  			name:             "transient error",
   145  			err:              errors.New("something transient"),
   146  			expectErr:        true,
   147  			expectedAttempts: maxRetries + 1,
   148  		},
   149  		{
   150  			name:             "nonexistent bucket",
   151  			err:              storage.ErrBucketNotExist,
   152  			expectErr:        true,
   153  			expectedAttempts: 1,
   154  		},
   155  		{
   156  			name:             "nonexistent object",
   157  			err:              storage.ErrObjectNotExist,
   158  			expectErr:        true,
   159  			expectedAttempts: 1,
   160  		},
   161  	}
   162  	for _, test := range tests {
   163  		t.Run(test.name, func(t *testing.T) {
   164  			var attempts int
   165  			err := retryWithPolicy(ctx, policy, func() error {
   166  				attempts++
   167  				return test.err
   168  			})
   169  			if err != nil && !test.expectErr {
   170  				t.Errorf("Unexpected error from Retry(): %s", err)
   171  			} else if err == nil && test.expectErr {
   172  				t.Errorf("Expected Retry to return an error, but got nil")
   173  			}
   174  			if attempts != test.expectedAttempts {
   175  				t.Errorf("Got %d attempts, expected %d", attempts, test.expectedAttempts)
   176  			}
   177  		})
   178  	}
   179  }