From 09bbe072cd60fc96d5900310df82e0eb6107d889 Mon Sep 17 00:00:00 2001
From: Nolan Brubaker <nolan@heptio.com>
Date: Tue, 22 May 2018 12:34:10 -0400
Subject: [PATCH] Check errors when closing GCP objects

Writing to GCP's object store is any async operation, so errors need to
be checked both on write and close calls, since errors like permission
violations aren't reported until a close.

Signed-off-by: Nolan Brubaker <nolan@heptio.com>
---
 pkg/cloudprovider/gcp/object_store.go      | 32 +++++++-
 pkg/cloudprovider/gcp/object_store_test.go | 92 ++++++++++++++++++++++
 2 files changed, 120 insertions(+), 4 deletions(-)
 create mode 100644 pkg/cloudprovider/gcp/object_store_test.go

diff --git a/pkg/cloudprovider/gcp/object_store.go b/pkg/cloudprovider/gcp/object_store.go
index 273bf0771..ac8510a3e 100644
--- a/pkg/cloudprovider/gcp/object_store.go
+++ b/pkg/cloudprovider/gcp/object_store.go
@@ -35,10 +35,25 @@ import (
 
 const credentialsEnvVar = "GOOGLE_APPLICATION_CREDENTIALS"
 
+// bucketWriter wraps the GCP SDK functions for accessing object store so they can be faked for testing.
+type bucketWriter interface {
+	// getWriteCloser returns an io.WriteCloser that can be used to upload data to the specified bucket for the specified key.
+	getWriteCloser(bucket, key string) io.WriteCloser
+}
+
+type writer struct {
+	client *storage.Client
+}
+
+func (w *writer) getWriteCloser(bucket, key string) io.WriteCloser {
+	return w.client.Bucket(bucket).Object(key).NewWriter(context.Background())
+}
+
 type objectStore struct {
 	client         *storage.Client
 	googleAccessID string
 	privateKey     []byte
+	bucketWriter   bucketWriter
 }
 
 func NewObjectStore() cloudprovider.ObjectStore {
@@ -76,16 +91,25 @@ func (o *objectStore) Init(config map[string]string) error {
 	}
 	o.client = client
 
+	o.bucketWriter = &writer{client: o.client}
+
 	return nil
 }
 
 func (o *objectStore) PutObject(bucket string, key string, body io.Reader) error {
-	w := o.client.Bucket(bucket).Object(key).NewWriter(context.Background())
-	defer w.Close()
+	w := o.bucketWriter.getWriteCloser(bucket, key)
 
-	_, err := io.Copy(w, body)
+	// The writer returned by NewWriter is asynchronous, so errors aren't guaranteed
+	// until Close() is called
+	_, copyErr := io.Copy(w, body)
 
-	return errors.WithStack(err)
+	// Ensure we close w and report errors properly
+	closeErr := w.Close()
+	if copyErr != nil {
+		return copyErr
+	}
+
+	return closeErr
 }
 
 func (o *objectStore) GetObject(bucket string, key string) (io.ReadCloser, error) {
diff --git a/pkg/cloudprovider/gcp/object_store_test.go b/pkg/cloudprovider/gcp/object_store_test.go
new file mode 100644
index 000000000..3ac766013
--- /dev/null
+++ b/pkg/cloudprovider/gcp/object_store_test.go
@@ -0,0 +1,92 @@
+/*
+Copyright 2019 the Heptio Ark contributors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+package gcp
+
+import (
+	"errors"
+	"io"
+	"strings"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+type mockWriteCloser struct {
+	closeErr error
+	writeErr error
+}
+
+func (m *mockWriteCloser) Close() error {
+	return m.closeErr
+}
+
+func (m *mockWriteCloser) Write(b []byte) (int, error) {
+	return len(b), m.writeErr
+}
+
+func newMockWriteCloser(writeErr, closeErr error) *mockWriteCloser {
+	return &mockWriteCloser{writeErr: writeErr, closeErr: closeErr}
+}
+
+type fakeWriter struct {
+	wc *mockWriteCloser
+}
+
+func newFakeWriter(wc *mockWriteCloser) *fakeWriter {
+	return &fakeWriter{wc: wc}
+}
+
+func (fw *fakeWriter) getWriteCloser(bucket, name string) io.WriteCloser {
+	return fw.wc
+}
+
+func TestPutObject(t *testing.T) {
+	tests := []struct {
+		name        string
+		writeErr    error
+		closeErr    error
+		expectedErr error
+	}{
+		{
+			name:        "Close() errors are returned",
+			closeErr:    errors.New("error closing"),
+			expectedErr: errors.New("error closing"),
+		},
+		{
+			name:        "Write() errors are returned",
+			writeErr:    errors.New("error writing"),
+			expectedErr: errors.New("error writing"),
+		},
+		{
+			name:        "Write errors supercede close errors",
+			writeErr:    errors.New("error writing"),
+			closeErr:    errors.New("error closing"),
+			expectedErr: errors.New("error writing"),
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			wc := newMockWriteCloser(test.writeErr, test.closeErr)
+			o := NewObjectStore().(*objectStore)
+			o.bucketWriter = newFakeWriter(wc)
+
+			err := o.PutObject("bucket", "key", strings.NewReader("contents"))
+
+			assert.Equal(t, test.expectedErr, err)
+		})
+	}
+}