src/policy_manager/policymanagerutil/metadata_retriever_impl_test.go

// Copyright 2021 Google LLC
//
// 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 policymanagerutil

import (
	"fmt"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
)

// isValidMetadataRequest is a generic helper function to verify that the
// request matches the spec set by GCE metadata service and supports the
// behavior specified in the MetadataRetriever implementation.
func isValidMetadataRequest(r *http.Request) error {
	// Check the request is a GET request.
	if r.Method != "GET" {
		return fmt.Errorf("got a %s request, want GET request",
			r.Method)
	}

	// Check request contains the right headers.
	if _, ok := r.Header[requestHeaderKey]; !ok {
		return fmt.Errorf("missing %s header", requestHeaderKey)
	}

	return nil
}

// TestGetEntry tests that the getEntry function for the following things:
//	- makes a GET request with the required headers
//	- raises an error if the status code is not 200
//	- returns the correct response from the server
//	- returns the correct etag from the server
func TestGetEntry(t *testing.T) {
	tests := []struct {
		name           string
		entry          string
		expectResponse string
		expectErr      bool
		expectEtag     string
	}{
		{
			"InstanceID",
			instanceIDEndpoint,
			"1234567890",
			false,
			"",
		},
		{
			"InstanceCustomMetaData",
			instanceCustomDataDirectory,
			`{
				"key1": "value1",
				"key2": "value2"
			}`,
			false,
			"myEtag",
		},
		{
			"ServerError",
			instanceIDEndpoint,
			"",
			true,
			"",
		},
	}

	for _, test := range tests {
		// Set up a test server.
		handler := func(w http.ResponseWriter, r *http.Request) {
			if test.expectErr {
				w.WriteHeader(http.StatusInternalServerError)
				return
			}

			if err := isValidMetadataRequest(r); err != nil {
				t.Errorf("Test %s: %v", test.name, err)
			}

			w.Header().Set("etag", test.expectEtag)

			fmt.Fprint(w, test.expectResponse)
		}
		ts := httptest.NewServer(http.HandlerFunc(handler))

		resp, etag, err := getEntry(test.entry, ts.URL)
		if err != nil {
			if !test.expectErr {
				t.Errorf("Test %s got unexpected error %v",
					test.name,
					err)
			}
		} else {
			if test.expectErr {
				t.Errorf("Test %s got %s, want error",
					test.name,
					string(resp))
			} else {
				if string(resp) != test.expectResponse {
					t.Errorf("Test %s got %s, want %s",
						test.name,
						string(resp),
						test.expectResponse)
				}
				if etag != test.expectEtag {
					t.Errorf("Test %s got etag %s, want %s",
						test.name,
						etag,
						test.expectEtag)
				}
			}
		}

		ts.Close()
	}
}

// TestGetInstanceID tests that the retriever implementation correctly fetches
// the instance id from the metadata server.
func TestGetInstanceID(t *testing.T) {
	expectID := uint64(15134250476052400915)

	// Set up a test server.
	handler := func(w http.ResponseWriter, r *http.Request) {
		if err := isValidMetadataRequest(r); err != nil {
			t.Error(err)
		}

		// Make sure the request is for the right metadata entry.
		if r.URL.Path != instanceIDEndpoint {
			t.Errorf("got URL path %s, want %s",
				r.URL.Path,
				instanceIDEndpoint)
		}

		fmt.Fprintf(w, "%v", expectID)
	}
	ts := httptest.NewServer(http.HandlerFunc(handler))
	defer ts.Close()

	// Set up retriever to talk to the test server.
	retriever := NewMetadataRetriever(ts.URL)

	if id, err := retriever.GetInstanceID(); err != nil {
		t.Error(err)
	} else if id != expectID {
		t.Errorf("got %d, want %d", id, expectID)
	}
}

// TestFetchMetadata tests that the retriever implementation correctly fetches
// the content of instance custom metadata directory.
func TestFetchMetadata(t *testing.T) {
	// Dummy key-value pairs.
	dummyMap := map[string]string{
		"key1": "val1",
		"key2": "val2",
		"key3": "\t val3 \n  ",
	}

	tests := []struct {
		name           string
		lastEtag       string
		expectResponse string
		expectEtag     string
		expectErr      bool
	}{
		{
			"InstanceCustomMetadataWait",
			"lastEtag",
			dummyMap["key1"],
			"myEtag1",
			false,
		},
		{
			"InstanceCustomMetadataNoWait",
			"",
			dummyMap["key2"],
			"myEtag2",
			false,
		},
		{
			"InstanceCustomMetadataWhitespaces",
			"",
			dummyMap["key3"],
			"myEtag3",
			false,
		},
		{
			"InstanceCustomDataServerError",
			"",
			"",
			"",
			true,
		},
	}

	for _, test := range tests {
		// Set up a test server.
		handler := func(w http.ResponseWriter, r *http.Request) {
			if test.expectErr {
				w.WriteHeader(http.StatusInternalServerError)
				return
			}

			if err := isValidMetadataRequest(r); err != nil {
				t.Errorf("Test %s: %v", test.name, err)
			}

			// Make sure the request is for the directory.
			if r.URL.Path != instanceCustomDataDirectory {
				t.Errorf("Test %s got URL path %s, want %s",
					test.name,
					r.URL.Path,
					instanceCustomDataDirectory)
			}

			expectedQuery := "recursive=true"
			if test.lastEtag != "" {
				expectedQuery += "&wait_for_change=true&last_etag=" + test.lastEtag
			}
			if r.URL.RawQuery != expectedQuery {
				t.Errorf("Test %s got HTTP request query %s, want %s",
					test.name,
					r.URL.RawQuery,
					expectedQuery)
			}

			w.Header().Set("etag", test.expectEtag)

			fmt.Fprint(w, test.expectResponse)
		}
		ts := httptest.NewServer(http.HandlerFunc(handler))

		// Create retriever to talk to the test server.
		retriever := NewMetadataRetriever(ts.URL)

		resp, etag, err := retriever.FetchMetadata(test.lastEtag)
		if err != nil {
			if !test.expectErr {
				t.Errorf("Test %s got unexpected error %v",
					test.name,
					err)
			}
		} else {
			if test.expectErr {
				t.Errorf("Test %s got %s, want error",
					test.name,
					resp)
			} else {
				if resp != strings.TrimSpace(test.expectResponse) {
					t.Errorf("Test %s got %v, want %v",
						test.name,
						resp,
						test.expectResponse)
				}
				if etag != test.expectEtag {
					t.Errorf("Test %s got etag %s, want %s",
						test.name,
						etag,
						test.expectEtag)
				}
			}
		}

		ts.Close()
	}
}