fix(storage): fix endpoint selection logic (#3172)

storage/storage.go

@@ -43,6 +43,7 @@ import (
 	"cloud.google.com/go/internal/version"
 	"google.golang.org/api/googleapi"
 	"google.golang.org/api/option"
+	"google.golang.org/api/option/internaloption"
 	raw "google.golang.org/api/storage/v1"
 	htransport "google.golang.org/api/transport/http"
 )
@@ -105,41 +106,48 @@ type Client struct {
 func NewClient(ctx context.Context, opts ...option.ClientOption) (*Client, error) {
 	var host, readHost, scheme string
 
+	// In general, it is recommended to use raw.NewService instead of htransport.NewClient
+	// since raw.NewService configures the correct default endpoints when initializing the
+	// internal http client. However, in our case, "NewRangeReader" in reader.go needs to
+	// access the http client directly to make requests, so we create the client manually
+	// here so it can be re-used by both reader.go and raw.NewService. This means we need to
+	// manually configure the default endpoint options on the http client. Furthermore, we
+	// need to account for STORAGE_EMULATOR_HOST override when setting the default endpoints.
 	if host = os.Getenv("STORAGE_EMULATOR_HOST"); host == "" {
 		scheme = "https"
 		readHost = "storage.googleapis.com"
 
 		// Prepend default options to avoid overriding options passed by the user.
 		opts = append([]option.ClientOption{option.WithScopes(ScopeFullControl), option.WithUserAgent(userAgent)}, opts...)
+
+		opts = append(opts, internaloption.WithDefaultEndpoint("https://storage.googleapis.com/storage/v1/"))
+		opts = append(opts, internaloption.WithDefaultMTLSEndpoint("https://storage.mtls.googleapis.com/storage/v1/"))
 	} else {
 		scheme = "http"
 		readHost = host
 
 		opts = append([]option.ClientOption{option.WithoutAuthentication()}, opts...)
+
+		opts = append(opts, internaloption.WithDefaultEndpoint(host))
+		opts = append(opts, internaloption.WithDefaultMTLSEndpoint(host))
 	}
 
+	// htransport selects the correct endpoint among WithEndpoint (user override), WithDefaultEndpoint, and WithDefaultMTLSEndpoint.
 	hc, ep, err := htransport.NewClient(ctx, opts...)
 	if err != nil {
 		return nil, fmt.Errorf("dialing: %v", err)
 	}
-	rawService, err := raw.NewService(ctx, option.WithHTTPClient(hc))
+	// RawService should be created with the chosen endpoint to take account of user override.
+	rawService, err := raw.NewService(ctx, option.WithEndpoint(ep), option.WithHTTPClient(hc))
 	if err != nil {
 		return nil, fmt.Errorf("storage client: %v", err)
 	}
-	if ep == "" {
-		// Override the default value for BasePath from the raw client.
-		// TODO: remove when the raw client uses this endpoint as its default (~end of 2020)
-		rawService.BasePath = "https://storage.googleapis.com/storage/v1/"
-	} else {
-		// If the endpoint has been set explicitly, use this for the BasePath
-		// as well as readHost
-		rawService.BasePath = ep
-		u, err := url.Parse(ep)
-		if err != nil {
-			return nil, fmt.Errorf("supplied endpoint %v is not valid: %v", ep, err)
-		}
-		readHost = u.Host
+	// Update readHost with the chosen endpoint.
+	u, err := url.Parse(ep)
+	if err != nil {
+		return nil, fmt.Errorf("supplied endpoint %q is not valid: %v", ep, err)
 	}
+	readHost = u.Host
 
 	return &Client{
 		hc:       hc,

storage/storage_test.go

@@ -26,6 +26,7 @@ import (
 	"net/http"
 	"net/http/httptest"
 	"net/url"
+	"os"
 	"reflect"
 	"regexp"
 	"sort"
@@ -1220,22 +1221,67 @@ func TestAttrToFieldMapCoverage(t *testing.T) {
 	}
 }
 
-// Create a client using a custom endpoint, and verify that raw.BasePath (used
-// for writes) and readHost (used for reads) are both set correctly.
+// Create a client using a combination of custom endpoint and
+// STORAGE_EMULATOR_HOST env variable and verify that raw.BasePath (used
+// for writes) and readHost and scheme (used for reads) are all set correctly.
 func TestWithEndpoint(t *testing.T) {
-	ctx := context.Background()
-	endpoint := "https://fake.gcs.com:8080/storage/v1"
-	c, err := NewClient(ctx, option.WithEndpoint(endpoint))
-	if err != nil {
-		t.Fatalf("error creating client: %v", err)
-	}
-
-	if c.raw.BasePath != endpoint {
-		t.Errorf("raw.BasePath not set correctly: got %v, want %v", c.raw.BasePath, endpoint)
+	originalStorageEmulatorHost := os.Getenv("STORAGE_EMULATOR_HOST")
+	testCases := []struct {
+		CustomEndpoint      string
+		StorageEmulatorHost string
+		WantRawBasePath     string
+		WantReadHost        string
+		WantScheme          string
+	}{
+		{
+			CustomEndpoint:      "",
+			StorageEmulatorHost: "",
+			WantRawBasePath:     "https://storage.googleapis.com/storage/v1/",
+			WantReadHost:        "storage.googleapis.com",
+			WantScheme:          "https",
+		},
+		{
+			CustomEndpoint:      "https://fake.gcs.com:8080/storage/v1",
+			StorageEmulatorHost: "",
+			WantRawBasePath:     "https://fake.gcs.com:8080/storage/v1",
+			WantReadHost:        "fake.gcs.com:8080",
+			WantScheme:          "https",
+		},
+		{
+			CustomEndpoint:      "",
+			StorageEmulatorHost: "http://emu.com",
+			WantRawBasePath:     "http://emu.com",
+			WantReadHost:        "emu.com",
+			WantScheme:          "http",
+		},
+		{
+			CustomEndpoint:      "https://fake.gcs.com:8080/storage/v1",
+			StorageEmulatorHost: "http://emu.com",
+			WantRawBasePath:     "https://fake.gcs.com:8080/storage/v1",
+			WantReadHost:        "fake.gcs.com:8080",
+			WantScheme:          "http",
+		},
 	}
+	ctx := context.Background()
+	for _, tc := range testCases {
+		os.Setenv("STORAGE_EMULATOR_HOST", tc.StorageEmulatorHost)
+		c, err := NewClient(ctx, option.WithEndpoint(tc.CustomEndpoint))
+		if err != nil {
+			t.Fatalf("error creating client: %v", err)
+		}
+		if err != nil {
+			t.Fatalf("error creating client: %v", err)
+		}
 
-	want := "fake.gcs.com:8080"
-	if c.readHost != want {
-		t.Errorf("readHost not set correctly: got %v, want %v", c.readHost, want)
+		if c.raw.BasePath != tc.WantRawBasePath {
+			t.Errorf("raw.BasePath not set correctly: got %v, want %v", c.raw.BasePath, tc.WantRawBasePath)
+		}
+		if c.readHost != tc.WantReadHost {
+			t.Errorf("readHost not set correctly: got %v, want %v", c.readHost, tc.WantReadHost)
+		}
+		if c.scheme != tc.WantScheme {
+			t.Errorf("scheme not set correctly: got %v, want %v", c.scheme, tc.WantScheme)
+		}
 	}
+	os.Setenv("STORAGE_EMULATOR_HOST", originalStorageEmulatorHost)
 }