src/pkg/fakes/metadata.go

// Copyright 2018 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

import (
	"context"
	"log"
	"net"
	"net/http"
	"net/http/httptest"
	"os"
	"strings"
	"testing"

	"cloud.google.com/go/compute/metadata"
)

// MDS contains data and functionality for a fake compute Metadata Server.
// It is intended to be constructed with NewMetadataServer.
//
// The MDS struct represents the state of the fake MDS instance.
// Fields on this struct can be modified to influence the
// return values of MDS API calls.
//
// The fake MDS server implements a small part of the API discussed here:
// https://cloud.google.com/compute/docs/metadata/overview#instance-metadata.
// Only the parts that we need for testing are implemented here.
//
// This struct should not be considered concurrency safe.
type MDS struct {
	// Attributes represents attributes for the GCE instance stored by the metadata
	// server.
	Attributes map[string]string
	// Client is the client to use when accessing the fake MDS.
	Client *metadata.Client
	// Server is the fake MDS. It uses state from this struct for serving requests.
	Server *httptest.Server
}

func (m *MDS) instanceAttrHandler(w http.ResponseWriter, r *http.Request) {
	switch r.Method {
	case "GET":
		splitPath := strings.Split(r.URL.Path, "/")
		attributeKey := splitPath[len(splitPath)-1]
		data, ok := m.Attributes[attributeKey]
		if !ok {
			writeError(w, r, http.StatusNotFound)
			return
		}
		if _, err := w.Write([]byte(data)); err != nil {
			log.Printf("write %q failed: %v", r.URL.Path, err)
		}
	default:
		writeError(w, r, http.StatusNotFound)
		return
	}
}

// NewMetadataServer constructs a fake metadata server (MDS) implementation.
func NewMetadataServer(ctx context.Context) (*MDS, error) {
	mds := &MDS{make(map[string]string), nil, nil}
	mux := http.NewServeMux()
	mux.HandleFunc("/computeMetadata/v1/instance/attributes/", mds.instanceAttrHandler)
	mds.Server = httptest.NewServer(mux)
	// The MDS uses a fixed IP to serve metadata, set this env variable to redirect
	// to fake server.
	err := os.Setenv("GCE_METADATA_HOST", mds.Server.Listener.Addr().String())
	if err != nil {
		return nil, err
	}

	transport := &http.Transport{
		DialTLSContext: func(ctx context.Context, network, addr string) (net.Conn, error) {
			return net.Dial(network, addr)
		},
	}
	hc := &http.Client{Transport: transport}
	mds.Client = metadata.NewClient(hc)
	return mds, nil
}

func MDSForTest(t *testing.T) *MDS {
	t.Helper()
	mds, err := NewMetadataServer(context.Background())
	if err != nil {
		t.Fatal(err)
	}
	return mds
}