pkg/fakes/metadata.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 fakes contains fake implementations to be used in unit tests.
package fakes

import (
	"log"
	"net/http"
	"net/http/httptest"
	"os"
	"strings"
)

// GCEMetadata is a fake GCE implementation. It is intended to be constructed with NewMetadataServer.
type GCEMetadata struct {
	// MetadataKeys represents the collection of objects that exist in the fake GCS server.
	// Keys and values are strings of the form "/instance/attributes".
	Metadata map[string]string
	// Server is an HTTP server that serves fake metadata requests. Requests are served using the state stored in
	// the other struct fields.
	Server *httptest.Server
}

// metadataKeysHandler writes all metadata keys or values to a particular metadata key upon request
func (mtd *GCEMetadata) metadataKeysHandler(w http.ResponseWriter, r *http.Request) {
	metadataKey := strings.TrimPrefix(r.URL.Path, "/computeMetadata/v1/instance/attributes/")

	if metadataKey == "" {
		var res string
		for k := range mtd.Metadata {
			res += strings.TrimSpace(k) + "\n"
		}
		if _, err := w.Write([]byte(res)); err != nil {
			log.Printf("write %q failed: %v", r.URL.Path, err)
		}

	} else {
		res := string(mtd.Metadata[metadataKey])
		if _, err := w.Write([]byte(res)); err != nil {
			log.Printf("write %q failed: %v", r.URL.Path, err)
		}
	}
}

// NewMetadataServer constructs a fake Metadata implementation. This function also sets
// an environment variable GCE_METADATA_HOST which is required by metadata package. The
// use of this environment variable makes it unsafe for concurrency.
func NewMetadataServer() *GCEMetadata {
	mtd := &GCEMetadata{
		Metadata: make(map[string]string),
		Server:   nil,
	}

	mux := http.NewServeMux()
	mux.HandleFunc("/computeMetadata/v1/instance/attributes/", mtd.metadataKeysHandler)
	mtd.Server = httptest.NewServer(mux)

	// Setting metadata host url by skipping the http:// part
	serverURL := mtd.Server.URL[7:]
	os.Setenv("GCE_METADATA_HOST", serverURL)

	return mtd
}