<commit_before>/*
Copyright 2015 The Kubernetes Authors.

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 main

import (
	"errors"
	goflag "flag"
	"fmt"
	"math/rand"
	"os"
	"time"

	"github.com/spf13/cobra"
	"github.com/spf13/pflag"
	"k8s.io/klog/v2"

	v1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/resource"
	"k8s.io/apimachinery/pkg/util/sets"
	clientset "k8s.io/client-go/kubernetes"
	restclient "k8s.io/client-go/rest"
	"k8s.io/client-go/tools/clientcmd"
	"k8s.io/client-go/tools/record"
	cliflag "k8s.io/component-base/cli/flag"
	"k8s.io/component-base/logs"
	_ "k8s.io/component-base/metrics/prometheus/restclient" // for client metric registration
	_ "k8s.io/component-base/metrics/prometheus/version"    // for version metric registration
	"k8s.io/component-base/version"
	"k8s.io/component-base/version/verflag"
	"k8s.io/kubernetes/pkg/api/legacyscheme"
	"k8s.io/kubernetes/pkg/apis/core"
	"k8s.io/kubernetes/pkg/cluster/ports"
	cadvisortest "k8s.io/kubernetes/pkg/kubelet/cadvisor/testing"
	"k8s.io/kubernetes/pkg/kubelet/cm"
	"k8s.io/kubernetes/pkg/kubelet/cri/remote"
	fakeremote "k8s.io/kubernetes/pkg/kubelet/cri/remote/fake"
	"k8s.io/kubernetes/pkg/kubemark"
	fakeiptables "k8s.io/kubernetes/pkg/util/iptables/testing"
	fakesysctl "k8s.io/kubernetes/pkg/util/sysctl/testing"
	utiltaints "k8s.io/kubernetes/pkg/util/taints"
	fakeexec "k8s.io/utils/exec/testing"
)

type hollowNodeConfig struct {
	KubeconfigPath       string
	KubeletPort          int
	KubeletReadOnlyPort  int
	Morph                string
	NodeName             string
	ServerPort           int
	ContentType          string
	UseRealProxier       bool
	ProxierSyncPeriod    time.Duration
	ProxierMinSyncPeriod time.Duration
	NodeLabels           map[string]string
	RegisterWithTaints   []core.Taint
	MaxPods              int
	ExtendedResources    map[string]string
}

const (
	maxPods     = 110
	podsPerCore = 0
)

// TODO(#45650): Refactor hollow-node into hollow-kubelet and hollow-proxy
// and make the config driven.
var knownMorphs = sets.NewString("kubelet", "proxy")

func (c *hollowNodeConfig) addFlags(fs *pflag.FlagSet) {
	fs.StringVar(&c.KubeconfigPath, "kubeconfig", "/kubeconfig/kubeconfig", "Path to kubeconfig file.")
	fs.IntVar(&c.KubeletPort, "kubelet-port", ports.KubeletPort, "Port on which HollowKubelet should be listening.")
	fs.IntVar(&c.KubeletReadOnlyPort, "kubelet-read-only-port", ports.KubeletReadOnlyPort, "Read-only port on which Kubelet is listening.")
	fs.StringVar(&c.NodeName, "name", "fake-node", "Name of this Hollow Node.")
	fs.IntVar(&c.ServerPort, "api-server-port", 443, "Port on which API server is listening.")
	fs.StringVar(&c.Morph, "morph", "", fmt.Sprintf("Specifies into which Hollow component this binary should morph. Allowed values: %v", knownMorphs.List()))
	fs.StringVar(&c.ContentType, "kube-api-content-type", "application/vnd.kubernetes.protobuf", "ContentType of requests sent to apiserver.")
	fs.BoolVar(&c.UseRealProxier, "use-real-proxier", true, "Set to true if you want to use real proxier inside hollow-proxy.")
	fs.DurationVar(&c.ProxierSyncPeriod, "proxier-sync-period", 30*time.Second, "Period that proxy rules are refreshed in hollow-proxy.")
	fs.DurationVar(&c.ProxierMinSyncPeriod, "proxier-min-sync-period", 0, "Minimum period that proxy rules are refreshed in hollow-proxy.")
	bindableNodeLabels := cliflag.ConfigurationMap(c.NodeLabels)
	fs.Var(&bindableNodeLabels, "node-labels", "Additional node labels")
	fs.Var(utiltaints.NewTaintsVar(&c.RegisterWithTaints), "register-with-taints", "Register the node with the given list of taints (comma separated \"<key>=<value>:<effect>\"). No-op if register-node is false.")
	fs.IntVar(&c.MaxPods, "max-pods", maxPods, "Number of pods that can run on this Kubelet.")
	bindableExtendedResources := cliflag.ConfigurationMap(c.ExtendedResources)
	fs.Var(&bindableExtendedResources, "extended-resources", "Register the node with extended resources (comma separated \"<name>=<quantity>\")")
}

func (c *hollowNodeConfig) createClientConfigFromFile() (*restclient.Config, error) {
	clientConfig, err := clientcmd.LoadFromFile(c.KubeconfigPath)
	if err != nil {
		return nil, fmt.Errorf("error while loading kubeconfig from file %v: %v", c.KubeconfigPath, err)
	}
	config, err := clientcmd.NewDefaultClientConfig(*clientConfig, &clientcmd.ConfigOverrides{}).ClientConfig()
	if err != nil {
		return nil, fmt.Errorf("error while creating kubeconfig: %v", err)
	}
	config.ContentType = c.ContentType
	config.QPS = 10
	config.Burst = 20
	return config, nil
}

func (c *hollowNodeConfig) createHollowKubeletOptions() *kubemark.HollowKubletOptions {
	return &kubemark.HollowKubletOptions{
		NodeName:            c.NodeName,
		KubeletPort:         c.KubeletPort,
		KubeletReadOnlyPort: c.KubeletReadOnlyPort,
		MaxPods:             c.MaxPods,
		PodsPerCore:         podsPerCore,
		NodeLabels:          c.NodeLabels,
		RegisterWithTaints:  c.RegisterWithTaints,
	}
}

func main() {
	rand.Seed(time.Now().UnixNano())

	command := newHollowNodeCommand()

	// TODO: once we switch everything over to Cobra commands, we can go back to calling
	// cliflag.InitFlags() (by removing its pflag.Parse() call). For now, we have to set the
	// normalize func and add the go flag set by hand.
	pflag.CommandLine.SetNormalizeFunc(cliflag.WordSepNormalizeFunc)
	pflag.CommandLine.AddGoFlagSet(goflag.CommandLine)
	// cliflag.InitFlags()
	logs.InitLogs()
	defer logs.FlushLogs()

	if err := command.Execute(); err != nil {
		os.Exit(1)
	}
}

// newControllerManagerCommand creates a *cobra.Command object with default parameters
func newHollowNodeCommand() *cobra.Command {
	s := &hollowNodeConfig{
		NodeLabels:        make(map[string]string),
		ExtendedResources: make(map[string]string),
	}

	cmd := &cobra.Command{
		Use:  "kubemark",
		Long: "kubemark",
		Run: func(cmd *cobra.Command, args []string) {
			verflag.PrintAndExitIfRequested()
			run(cmd, s)
		},
		Args: func(cmd *cobra.Command, args []string) error {
			for _, arg := range args {
				if len(arg) > 0 {
					return fmt.Errorf("%q does not take any arguments, got %q", cmd.CommandPath(), args)
				}
			}
			return nil
		},
	}
	s.addFlags(cmd.Flags())

	return cmd
}

func run(cmd *cobra.Command, config *hollowNodeConfig) {
	// To help debugging, immediately log version and print flags.
	klog.Infof("Version: %+v", version.Get())
	cliflag.PrintFlags(cmd.Flags())

	if !knownMorphs.Has(config.Morph) {
		klog.Fatalf("Unknown morph: %v. Allowed values: %v", config.Morph, knownMorphs.List())
	}

	// create a client to communicate with API server.
	clientConfig, err := config.createClientConfigFromFile()
	if err != nil {
		klog.Fatalf("Failed to create a ClientConfig: %v. Exiting.", err)
	}

	client, err := clientset.NewForConfig(clientConfig)
	if err != nil {
		klog.Fatalf("Failed to create a ClientSet: %v. Exiting.", err)
	}

	if config.Morph == "kubelet" {
		f, c := kubemark.GetHollowKubeletConfig(config.createHollowKubeletOptions())

		heartbeatClientConfig := *clientConfig
		heartbeatClientConfig.Timeout = c.NodeStatusUpdateFrequency.Duration
		// The timeout is the minimum of the lease duration and status update frequency
		leaseTimeout := time.Duration(c.NodeLeaseDurationSeconds) * time.Second
		if heartbeatClientConfig.Timeout > leaseTimeout {
			heartbeatClientConfig.Timeout = leaseTimeout
		}

		heartbeatClientConfig.QPS = float32(-1)
		heartbeatClient, err := clientset.NewForConfig(&heartbeatClientConfig)
		if err != nil {
			klog.Fatalf("Failed to create a ClientSet: %v. Exiting.", err)
		}

		cadvisorInterface := &cadvisortest.Fake{
			NodeName: config.NodeName,
		}

		var containerManager cm.ContainerManager
		if config.ExtendedResources != nil {
			extendedResources := v1.ResourceList{}
			for k, v := range config.ExtendedResources {
				extendedResources[v1.ResourceName(k)] = resource.MustParse(v)
			}

			containerManager = cm.NewStubContainerManagerWithDevicePluginResource(extendedResources)
		} else {
			containerManager = cm.NewStubContainerManager()
		}

		endpoint, err := fakeremote.GenerateEndpoint()
		if err != nil {
			klog.Fatalf("Failed to generate fake endpoint %v.", err)
		}
		fakeRemoteRuntime := fakeremote.NewFakeRemoteRuntime()
		if err = fakeRemoteRuntime.Start(endpoint); err != nil {
			klog.Fatalf("Failed to start fake runtime %v.", err)
		}
		defer fakeRemoteRuntime.Stop()
		runtimeService, err := remote.NewRemoteRuntimeService(endpoint, 15*time.Second)
		if err != nil {
			klog.Fatalf("Failed to init runtime service %v.", err)
		}

		remoteImageService, err := remote.NewRemoteImageService(f.RemoteImageEndpoint, 15*time.Second)
		if err != nil {
			klog.Fatalf("Failed to init image service %v.", err)
		}

		hollowKubelet := kubemark.NewHollowKubelet(
			f, c,
			client,
			heartbeatClient,
			cadvisorInterface,
			remoteImageService,
			runtimeService,
			containerManager,
		)
		hollowKubelet.Run()
	}

	if config.Morph == "proxy" {
		client, err := clientset.NewForConfig(clientConfig)
		if err != nil {
			klog.Fatalf("Failed to create API Server client: %v", err)
		}
		iptInterface := fakeiptables.NewFake()
		sysctl := fakesysctl.NewFake()
		execer := &fakeexec.FakeExec{
			LookPathFunc: func(_ string) (string, error) { return "", errors.New("fake execer") },
		}
		eventBroadcaster := record.NewBroadcaster()
		recorder := eventBroadcaster.NewRecorder(legacyscheme.Scheme, v1.EventSource{Component: "kube-proxy", Host: config.NodeName})

		hollowProxy, err := kubemark.NewHollowProxyOrDie(
			config.NodeName,
			client,
			client.CoreV1(),
			iptInterface,
			sysctl,
			execer,
			eventBroadcaster,
			recorder,
			config.UseRealProxier,
			config.ProxierSyncPeriod,
			config.ProxierMinSyncPeriod,
		)
		if err != nil {
			klog.Fatalf("Failed to create hollowProxy instance: %v", err)
		}
		hollowProxy.Run()
	}
}
<commit_msg>Kubemark: allow choosing image-service<commit_after>/*
Copyright 2015 The Kubernetes Authors.

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 main

import (
	"errors"
	goflag "flag"
	"fmt"
	"math/rand"
	"os"
	"time"

	"github.com/spf13/cobra"
	"github.com/spf13/pflag"
	internalapi "k8s.io/cri-api/pkg/apis"
	"k8s.io/klog/v2"

	v1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/resource"
	"k8s.io/apimachinery/pkg/util/sets"
	clientset "k8s.io/client-go/kubernetes"
	restclient "k8s.io/client-go/rest"
	"k8s.io/client-go/tools/clientcmd"
	"k8s.io/client-go/tools/record"
	cliflag "k8s.io/component-base/cli/flag"
	"k8s.io/component-base/logs"
	_ "k8s.io/component-base/metrics/prometheus/restclient" // for client metric registration
	_ "k8s.io/component-base/metrics/prometheus/version"    // for version metric registration
	"k8s.io/component-base/version"
	"k8s.io/component-base/version/verflag"
	"k8s.io/kubernetes/pkg/api/legacyscheme"
	"k8s.io/kubernetes/pkg/apis/core"
	"k8s.io/kubernetes/pkg/cluster/ports"
	cadvisortest "k8s.io/kubernetes/pkg/kubelet/cadvisor/testing"
	"k8s.io/kubernetes/pkg/kubelet/cm"
	"k8s.io/kubernetes/pkg/kubelet/cri/remote"
	fakeremote "k8s.io/kubernetes/pkg/kubelet/cri/remote/fake"
	"k8s.io/kubernetes/pkg/kubemark"
	fakeiptables "k8s.io/kubernetes/pkg/util/iptables/testing"
	fakesysctl "k8s.io/kubernetes/pkg/util/sysctl/testing"
	utiltaints "k8s.io/kubernetes/pkg/util/taints"
	fakeexec "k8s.io/utils/exec/testing"
)

type hollowNodeConfig struct {
	KubeconfigPath       string
	KubeletPort          int
	KubeletReadOnlyPort  int
	Morph                string
	NodeName             string
	ServerPort           int
	ContentType          string
	UseRealProxier       bool
	ProxierSyncPeriod    time.Duration
	ProxierMinSyncPeriod time.Duration
	NodeLabels           map[string]string
	RegisterWithTaints   []core.Taint
	MaxPods              int
	ExtendedResources    map[string]string
	UseHostImageService  bool
}

const (
	maxPods     = 110
	podsPerCore = 0
)

// TODO(#45650): Refactor hollow-node into hollow-kubelet and hollow-proxy
// and make the config driven.
var knownMorphs = sets.NewString("kubelet", "proxy")

func (c *hollowNodeConfig) addFlags(fs *pflag.FlagSet) {
	fs.StringVar(&c.KubeconfigPath, "kubeconfig", "/kubeconfig/kubeconfig", "Path to kubeconfig file.")
	fs.IntVar(&c.KubeletPort, "kubelet-port", ports.KubeletPort, "Port on which HollowKubelet should be listening.")
	fs.IntVar(&c.KubeletReadOnlyPort, "kubelet-read-only-port", ports.KubeletReadOnlyPort, "Read-only port on which Kubelet is listening.")
	fs.StringVar(&c.NodeName, "name", "fake-node", "Name of this Hollow Node.")
	fs.IntVar(&c.ServerPort, "api-server-port", 443, "Port on which API server is listening.")
	fs.StringVar(&c.Morph, "morph", "", fmt.Sprintf("Specifies into which Hollow component this binary should morph. Allowed values: %v", knownMorphs.List()))
	fs.StringVar(&c.ContentType, "kube-api-content-type", "application/vnd.kubernetes.protobuf", "ContentType of requests sent to apiserver.")
	fs.BoolVar(&c.UseRealProxier, "use-real-proxier", true, "Set to true if you want to use real proxier inside hollow-proxy.")
	fs.DurationVar(&c.ProxierSyncPeriod, "proxier-sync-period", 30*time.Second, "Period that proxy rules are refreshed in hollow-proxy.")
	fs.DurationVar(&c.ProxierMinSyncPeriod, "proxier-min-sync-period", 0, "Minimum period that proxy rules are refreshed in hollow-proxy.")
	bindableNodeLabels := cliflag.ConfigurationMap(c.NodeLabels)
	fs.Var(&bindableNodeLabels, "node-labels", "Additional node labels")
	fs.Var(utiltaints.NewTaintsVar(&c.RegisterWithTaints), "register-with-taints", "Register the node with the given list of taints (comma separated \"<key>=<value>:<effect>\"). No-op if register-node is false.")
	fs.IntVar(&c.MaxPods, "max-pods", maxPods, "Number of pods that can run on this Kubelet.")
	bindableExtendedResources := cliflag.ConfigurationMap(c.ExtendedResources)
	fs.Var(&bindableExtendedResources, "extended-resources", "Register the node with extended resources (comma separated \"<name>=<quantity>\")")
	fs.BoolVar(&c.UseHostImageService, "use-host-image-service", true, "Set to true if the hollow-kubelet should use the host image service. If set to false the fake image service will be used")
}

func (c *hollowNodeConfig) createClientConfigFromFile() (*restclient.Config, error) {
	clientConfig, err := clientcmd.LoadFromFile(c.KubeconfigPath)
	if err != nil {
		return nil, fmt.Errorf("error while loading kubeconfig from file %v: %v", c.KubeconfigPath, err)
	}
	config, err := clientcmd.NewDefaultClientConfig(*clientConfig, &clientcmd.ConfigOverrides{}).ClientConfig()
	if err != nil {
		return nil, fmt.Errorf("error while creating kubeconfig: %v", err)
	}
	config.ContentType = c.ContentType
	config.QPS = 10
	config.Burst = 20
	return config, nil
}

func (c *hollowNodeConfig) createHollowKubeletOptions() *kubemark.HollowKubletOptions {
	return &kubemark.HollowKubletOptions{
		NodeName:            c.NodeName,
		KubeletPort:         c.KubeletPort,
		KubeletReadOnlyPort: c.KubeletReadOnlyPort,
		MaxPods:             c.MaxPods,
		PodsPerCore:         podsPerCore,
		NodeLabels:          c.NodeLabels,
		RegisterWithTaints:  c.RegisterWithTaints,
	}
}

func main() {
	rand.Seed(time.Now().UnixNano())

	command := newHollowNodeCommand()

	// TODO: once we switch everything over to Cobra commands, we can go back to calling
	// cliflag.InitFlags() (by removing its pflag.Parse() call). For now, we have to set the
	// normalize func and add the go flag set by hand.
	pflag.CommandLine.SetNormalizeFunc(cliflag.WordSepNormalizeFunc)
	pflag.CommandLine.AddGoFlagSet(goflag.CommandLine)
	// cliflag.InitFlags()
	logs.InitLogs()
	defer logs.FlushLogs()

	if err := command.Execute(); err != nil {
		os.Exit(1)
	}
}

// newControllerManagerCommand creates a *cobra.Command object with default parameters
func newHollowNodeCommand() *cobra.Command {
	s := &hollowNodeConfig{
		NodeLabels:        make(map[string]string),
		ExtendedResources: make(map[string]string),
	}

	cmd := &cobra.Command{
		Use:  "kubemark",
		Long: "kubemark",
		Run: func(cmd *cobra.Command, args []string) {
			verflag.PrintAndExitIfRequested()
			run(cmd, s)
		},
		Args: func(cmd *cobra.Command, args []string) error {
			for _, arg := range args {
				if len(arg) > 0 {
					return fmt.Errorf("%q does not take any arguments, got %q", cmd.CommandPath(), args)
				}
			}
			return nil
		},
	}
	s.addFlags(cmd.Flags())

	return cmd
}

func run(cmd *cobra.Command, config *hollowNodeConfig) {
	// To help debugging, immediately log version and print flags.
	klog.Infof("Version: %+v", version.Get())
	cliflag.PrintFlags(cmd.Flags())

	if !knownMorphs.Has(config.Morph) {
		klog.Fatalf("Unknown morph: %v. Allowed values: %v", config.Morph, knownMorphs.List())
	}

	// create a client to communicate with API server.
	clientConfig, err := config.createClientConfigFromFile()
	if err != nil {
		klog.Fatalf("Failed to create a ClientConfig: %v. Exiting.", err)
	}

	client, err := clientset.NewForConfig(clientConfig)
	if err != nil {
		klog.Fatalf("Failed to create a ClientSet: %v. Exiting.", err)
	}

	if config.Morph == "kubelet" {
		f, c := kubemark.GetHollowKubeletConfig(config.createHollowKubeletOptions())

		heartbeatClientConfig := *clientConfig
		heartbeatClientConfig.Timeout = c.NodeStatusUpdateFrequency.Duration
		// The timeout is the minimum of the lease duration and status update frequency
		leaseTimeout := time.Duration(c.NodeLeaseDurationSeconds) * time.Second
		if heartbeatClientConfig.Timeout > leaseTimeout {
			heartbeatClientConfig.Timeout = leaseTimeout
		}

		heartbeatClientConfig.QPS = float32(-1)
		heartbeatClient, err := clientset.NewForConfig(&heartbeatClientConfig)
		if err != nil {
			klog.Fatalf("Failed to create a ClientSet: %v. Exiting.", err)
		}

		cadvisorInterface := &cadvisortest.Fake{
			NodeName: config.NodeName,
		}

		var containerManager cm.ContainerManager
		if config.ExtendedResources != nil {
			extendedResources := v1.ResourceList{}
			for k, v := range config.ExtendedResources {
				extendedResources[v1.ResourceName(k)] = resource.MustParse(v)
			}

			containerManager = cm.NewStubContainerManagerWithDevicePluginResource(extendedResources)
		} else {
			containerManager = cm.NewStubContainerManager()
		}

		endpoint, err := fakeremote.GenerateEndpoint()
		if err != nil {
			klog.Fatalf("Failed to generate fake endpoint %v.", err)
		}
		fakeRemoteRuntime := fakeremote.NewFakeRemoteRuntime()
		if err = fakeRemoteRuntime.Start(endpoint); err != nil {
			klog.Fatalf("Failed to start fake runtime %v.", err)
		}
		defer fakeRemoteRuntime.Stop()
		runtimeService, err := remote.NewRemoteRuntimeService(endpoint, 15*time.Second)
		if err != nil {
			klog.Fatalf("Failed to init runtime service %v.", err)
		}

		var imageService internalapi.ImageManagerService = fakeRemoteRuntime.ImageService
		if config.UseHostImageService {
			imageService, err = remote.NewRemoteImageService(f.RemoteImageEndpoint, 15*time.Second)
			if err != nil {
				klog.Fatalf("Failed to init image service %v.", err)
			}
		}

		hollowKubelet := kubemark.NewHollowKubelet(
			f, c,
			client,
			heartbeatClient,
			cadvisorInterface,
			imageService,
			runtimeService,
			containerManager,
		)
		hollowKubelet.Run()
	}

	if config.Morph == "proxy" {
		client, err := clientset.NewForConfig(clientConfig)
		if err != nil {
			klog.Fatalf("Failed to create API Server client: %v", err)
		}
		iptInterface := fakeiptables.NewFake()
		sysctl := fakesysctl.NewFake()
		execer := &fakeexec.FakeExec{
			LookPathFunc: func(_ string) (string, error) { return "", errors.New("fake execer") },
		}
		eventBroadcaster := record.NewBroadcaster()
		recorder := eventBroadcaster.NewRecorder(legacyscheme.Scheme, v1.EventSource{Component: "kube-proxy", Host: config.NodeName})

		hollowProxy, err := kubemark.NewHollowProxyOrDie(
			config.NodeName,
			client,
			client.CoreV1(),
			iptInterface,
			sysctl,
			execer,
			eventBroadcaster,
			recorder,
			config.UseRealProxier,
			config.ProxierSyncPeriod,
			config.ProxierMinSyncPeriod,
		)
		if err != nil {
			klog.Fatalf("Failed to create hollowProxy instance: %v", err)
		}
		hollowProxy.Run()
	}
}
<|endoftext|>


<commit_before>/*
 *
 * k6 - a next-generation load testing tool
 * Copyright (C) 2018 Load Impact
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as
 * published by the Free Software Foundation, either version 3 of the
 * License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 */

package cmd

import (
	"bytes"
	"fmt"
	"net/url"
	"strings"
	"testing"

	"gopkg.in/guregu/null.v3"

	"github.com/loadimpact/k6/lib"
	"github.com/loadimpact/k6/loader"
	"github.com/spf13/afero"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

type runtimeOptionsTestCase struct {
	useSysEnv     bool // Whether to include the system env vars by default (run) or not (cloud/archive/inspect)
	expErr        bool
	cliFlags      []string
	systemEnv     map[string]string
	expEnv        map[string]string
	expCompatMode null.String
}

//nolint:gochecknoglobals
var (
	defaultCompatMode  = null.NewString("extended", false)
	baseCompatMode     = null.NewString("base", true)
	extendedCompatMode = null.NewString("extended", true)
)

var runtimeOptionsTestCases = map[string]runtimeOptionsTestCase{ //nolint:gochecknoglobals
	"empty env": {
		useSysEnv: true,
		// everything else is empty
		expCompatMode: defaultCompatMode,
	},
	"disabled sys env by default": {
		useSysEnv:     false,
		systemEnv:     map[string]string{"test1": "val1"},
		expEnv:        map[string]string{},
		expCompatMode: defaultCompatMode,
	},
	"disabled sys env by default with ext compat mode": {
		useSysEnv:     false,
		systemEnv:     map[string]string{"test1": "val1", "K6_COMPATIBILITY_MODE": "extended"},
		expEnv:        map[string]string{},
		expCompatMode: extendedCompatMode,
	},
	"disabled sys env by cli 1": {
		useSysEnv:     true,
		systemEnv:     map[string]string{"test1": "val1", "K6_COMPATIBILITY_MODE": "base"},
		cliFlags:      []string{"--include-system-env-vars=false"},
		expEnv:        map[string]string{},
		expCompatMode: baseCompatMode,
	},
	"disabled sys env by cli 2": {
		useSysEnv:     true,
		systemEnv:     map[string]string{"K6_INCLUDE_SYSTEM_ENV_VARS": "true", "K6_COMPATIBILITY_MODE": "extended"},
		cliFlags:      []string{"--include-system-env-vars=0", "--compatibility-mode=base"},
		expEnv:        map[string]string{},
		expCompatMode: baseCompatMode,
	},
	"disabled sys env by env": {
		useSysEnv:     true,
		systemEnv:     map[string]string{"K6_INCLUDE_SYSTEM_ENV_VARS": "false", "K6_COMPATIBILITY_MODE": "extended"},
		expEnv:        map[string]string{},
		expCompatMode: extendedCompatMode,
	},
	"enabled sys env by env": {
		useSysEnv:     false,
		systemEnv:     map[string]string{"K6_INCLUDE_SYSTEM_ENV_VARS": "true", "K6_COMPATIBILITY_MODE": "extended"},
		expEnv:        map[string]string{"K6_INCLUDE_SYSTEM_ENV_VARS": "true", "K6_COMPATIBILITY_MODE": "extended"},
		expCompatMode: extendedCompatMode,
	},
	"enabled sys env by default": {
		useSysEnv:     true,
		systemEnv:     map[string]string{"test1": "val1"},
		cliFlags:      []string{},
		expEnv:        map[string]string{"test1": "val1"},
		expCompatMode: defaultCompatMode,
	},
	"enabled sys env by cli 1": {
		useSysEnv:     false,
		systemEnv:     map[string]string{"test1": "val1"},
		cliFlags:      []string{"--include-system-env-vars"},
		expEnv:        map[string]string{"test1": "val1"},
		expCompatMode: defaultCompatMode,
	},
	"enabled sys env by cli 2": {
		useSysEnv:     false,
		systemEnv:     map[string]string{"test1": "val1"},
		cliFlags:      []string{"--include-system-env-vars=true"},
		expEnv:        map[string]string{"test1": "val1"},
		expCompatMode: defaultCompatMode,
	},
	"run only system env": {
		useSysEnv:     true,
		systemEnv:     map[string]string{"test1": "val1"},
		cliFlags:      []string{},
		expEnv:        map[string]string{"test1": "val1"},
		expCompatMode: defaultCompatMode,
	},
	"mixed system and cli env": {
		useSysEnv:     true,
		systemEnv:     map[string]string{"test1": "val1", "test2": ""},
		cliFlags:      []string{"--env", "test3=val3", "-e", "test4", "-e", "test5="},
		expEnv:        map[string]string{"test1": "val1", "test2": "", "test3": "val3", "test4": "", "test5": ""},
		expCompatMode: defaultCompatMode,
	},
	"mixed system and cli env 2": {
		useSysEnv:     false,
		systemEnv:     map[string]string{"test1": "val1", "test2": ""},
		cliFlags:      []string{"--env", "test3=val3", "-e", "test4", "-e", "test5=", "--include-system-env-vars=1"},
		expEnv:        map[string]string{"test1": "val1", "test2": "", "test3": "val3", "test4": "", "test5": ""},
		expCompatMode: defaultCompatMode,
	},
	"disabled system env with cli params": {
		useSysEnv:     false,
		systemEnv:     map[string]string{"test1": "val1"},
		cliFlags:      []string{"-e", "test2=overwriten", "-e", "test2=val2"},
		expEnv:        map[string]string{"test2": "val2"},
		expCompatMode: defaultCompatMode,
	},
	"overwriting system env with cli param": {
		useSysEnv:     true,
		systemEnv:     map[string]string{"test1": "val1sys"},
		cliFlags:      []string{"--env", "test1=val1cli"},
		expEnv:        map[string]string{"test1": "val1cli"},
		expCompatMode: defaultCompatMode,
	},
	"error wrong compat mode env var value": {
		systemEnv: map[string]string{"K6_COMPATIBILITY_MODE": "asdf"},
		expErr:    true,
	},
	"error wrong compat mode cli flag value": {
		cliFlags: []string{"--compatibility-mode", "whatever"},
		expErr:   true,
	},
	"error invalid cli var name 1": {
		useSysEnv:     true,
		systemEnv:     map[string]string{},
		cliFlags:      []string{"--env", "test a=error"},
		expErr:        true,
		expEnv:        map[string]string{},
		expCompatMode: defaultCompatMode,
	},
	"error invalid cli var name 2": {
		useSysEnv:     true,
		systemEnv:     map[string]string{},
		cliFlags:      []string{"--env", "1var=error"},
		expErr:        true,
		expEnv:        map[string]string{},
		expCompatMode: defaultCompatMode,
	},
	"error invalid cli var name 3": {
		useSysEnv:     true,
		systemEnv:     map[string]string{},
		cliFlags:      []string{"--env", "уникод=unicode-disabled"},
		expErr:        true,
		expEnv:        map[string]string{},
		expCompatMode: defaultCompatMode,
	},
	"valid env vars with spaces": {
		useSysEnv:     true,
		systemEnv:     map[string]string{"test1": "value 1"},
		cliFlags:      []string{"--env", "test2=value 2"},
		expEnv:        map[string]string{"test1": "value 1", "test2": "value 2"},
		expCompatMode: defaultCompatMode,
	},
	"valid env vars with special chars": {
		useSysEnv:     true,
		systemEnv:     map[string]string{"test1": "value 1"},
		cliFlags:      []string{"--env", "test2=value,2", "-e", `test3= ,  ,,, value, ,, 2!'@#,"`},
		expEnv:        map[string]string{"test1": "value 1", "test2": "value,2", "test3": ` ,  ,,, value, ,, 2!'@#,"`},
		expCompatMode: defaultCompatMode,
	},
}

func testRuntimeOptionsCase(t *testing.T, tc runtimeOptionsTestCase) {
	flags := runtimeOptionFlagSet(tc.useSysEnv)
	require.NoError(t, flags.Parse(tc.cliFlags))

	rtOpts, err := getRuntimeOptions(flags, tc.systemEnv)
	if tc.expErr {
		require.Error(t, err)
		return
	}
	require.NoError(t, err)
	require.EqualValues(t, tc.expEnv, rtOpts.Env)
	assert.Equal(t, tc.expCompatMode, rtOpts.CompatibilityMode)

	compatMode, err := lib.ValidateCompatibilityMode(rtOpts.CompatibilityMode.String)
	require.NoError(t, err)

	jsCodeLines := make([]string, 0, len(tc.expEnv)+2)
	if compatMode == lib.CompatibilityModeExtended {
		jsCodeLines = append(jsCodeLines, "export default function() {")
	} else {
		jsCodeLines = append(jsCodeLines, "module.exports.default = function() {")
	}

	for key, val := range tc.expEnv {
		jsCodeLines = append(jsCodeLines, fmt.Sprintf(
			"if (__ENV.%s !== `%s`) { throw new Error('Invalid %s: ' + __ENV.%s); }",
			key, val, key, key,
		))
	}
	jsCodeLines = append(jsCodeLines, "}")
	jsCode := []byte(strings.Join(jsCodeLines, "\n"))

	fs := afero.NewMemMapFs()
	require.NoError(t, afero.WriteFile(fs, "/script.js", jsCode, 0644))
	runner, err := newRunner(
		&loader.SourceData{Data: jsCode, URL: &url.URL{Path: "/script.js", Scheme: "file"}},
		typeJS,
		map[string]afero.Fs{"file": fs},
		rtOpts,
	)
	require.NoError(t, err)

	archive := runner.MakeArchive()
	archiveBuf := &bytes.Buffer{}
	assert.NoError(t, archive.Write(archiveBuf))

	getRunnerErr := func(rtOpts lib.RuntimeOptions) (lib.Runner, error) {
		return newRunner(
			&loader.SourceData{
				Data: archiveBuf.Bytes(),
				URL:  &url.URL{Path: "/script.js"},
			},
			typeArchive,
			nil,
			rtOpts,
		)
	}

	_, err = getRunnerErr(lib.RuntimeOptions{})
	require.NoError(t, err)
	for key, val := range tc.expEnv {
		r, err := getRunnerErr(lib.RuntimeOptions{Env: map[string]string{key: "almost " + val}})
		assert.NoError(t, err)
		assert.Equal(t, r.MakeArchive().Env[key], "almost "+val)
	}
}

func TestRuntimeOptions(t *testing.T) {
	for name, tc := range runtimeOptionsTestCases {
		tc := tc
		t.Run(fmt.Sprintf("RuntimeOptions test '%s'", name), func(t *testing.T) {
			t.Parallel()
			testRuntimeOptionsCase(t, tc)
		})
	}
}
<commit_msg>Fix minor issues in the runtime options test<commit_after>/*
 *
 * k6 - a next-generation load testing tool
 * Copyright (C) 2018 Load Impact
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as
 * published by the Free Software Foundation, either version 3 of the
 * License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 */

package cmd

import (
	"bytes"
	"fmt"
	"net/url"
	"testing"

	"gopkg.in/guregu/null.v3"

	"github.com/loadimpact/k6/lib"
	"github.com/loadimpact/k6/loader"
	"github.com/spf13/afero"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

type runtimeOptionsTestCase struct {
	useSysEnv     bool // Whether to include the system env vars by default (run) or not (cloud/archive/inspect)
	expErr        bool
	cliFlags      []string
	systemEnv     map[string]string
	expEnv        map[string]string
	expCompatMode null.String
}

//nolint:gochecknoglobals
var (
	defaultCompatMode  = null.NewString("extended", false)
	baseCompatMode     = null.NewString("base", true)
	extendedCompatMode = null.NewString("extended", true)
)

var runtimeOptionsTestCases = map[string]runtimeOptionsTestCase{ //nolint:gochecknoglobals
	"empty env": {
		useSysEnv: true,
		// everything else is empty
		expCompatMode: defaultCompatMode,
	},
	"disabled sys env by default": {
		useSysEnv:     false,
		systemEnv:     map[string]string{"test1": "val1"},
		expEnv:        map[string]string{},
		expCompatMode: defaultCompatMode,
	},
	"disabled sys env by default with ext compat mode": {
		useSysEnv:     false,
		systemEnv:     map[string]string{"test1": "val1", "K6_COMPATIBILITY_MODE": "extended"},
		expEnv:        map[string]string{},
		expCompatMode: extendedCompatMode,
	},
	"disabled sys env by cli 1": {
		useSysEnv:     true,
		systemEnv:     map[string]string{"test1": "val1", "K6_COMPATIBILITY_MODE": "base"},
		cliFlags:      []string{"--include-system-env-vars=false"},
		expEnv:        map[string]string{},
		expCompatMode: baseCompatMode,
	},
	"disabled sys env by cli 2": {
		useSysEnv:     true,
		systemEnv:     map[string]string{"K6_INCLUDE_SYSTEM_ENV_VARS": "true", "K6_COMPATIBILITY_MODE": "extended"},
		cliFlags:      []string{"--include-system-env-vars=0", "--compatibility-mode=base"},
		expEnv:        map[string]string{},
		expCompatMode: baseCompatMode,
	},
	"disabled sys env by env": {
		useSysEnv:     true,
		systemEnv:     map[string]string{"K6_INCLUDE_SYSTEM_ENV_VARS": "false", "K6_COMPATIBILITY_MODE": "extended"},
		expEnv:        map[string]string{},
		expCompatMode: extendedCompatMode,
	},
	"enabled sys env by env": {
		useSysEnv:     false,
		systemEnv:     map[string]string{"K6_INCLUDE_SYSTEM_ENV_VARS": "true", "K6_COMPATIBILITY_MODE": "extended"},
		expEnv:        map[string]string{"K6_INCLUDE_SYSTEM_ENV_VARS": "true", "K6_COMPATIBILITY_MODE": "extended"},
		expCompatMode: extendedCompatMode,
	},
	"enabled sys env by default": {
		useSysEnv:     true,
		systemEnv:     map[string]string{"test1": "val1"},
		cliFlags:      []string{},
		expEnv:        map[string]string{"test1": "val1"},
		expCompatMode: defaultCompatMode,
	},
	"enabled sys env by cli 1": {
		useSysEnv:     false,
		systemEnv:     map[string]string{"test1": "val1"},
		cliFlags:      []string{"--include-system-env-vars"},
		expEnv:        map[string]string{"test1": "val1"},
		expCompatMode: defaultCompatMode,
	},
	"enabled sys env by cli 2": {
		useSysEnv:     false,
		systemEnv:     map[string]string{"test1": "val1"},
		cliFlags:      []string{"--include-system-env-vars=true"},
		expEnv:        map[string]string{"test1": "val1"},
		expCompatMode: defaultCompatMode,
	},
	"run only system env": {
		useSysEnv:     true,
		systemEnv:     map[string]string{"test1": "val1"},
		cliFlags:      []string{},
		expEnv:        map[string]string{"test1": "val1"},
		expCompatMode: defaultCompatMode,
	},
	"mixed system and cli env": {
		useSysEnv:     true,
		systemEnv:     map[string]string{"test1": "val1", "test2": ""},
		cliFlags:      []string{"--env", "test3=val3", "-e", "test4", "-e", "test5="},
		expEnv:        map[string]string{"test1": "val1", "test2": "", "test3": "val3", "test4": "", "test5": ""},
		expCompatMode: defaultCompatMode,
	},
	"mixed system and cli env 2": {
		useSysEnv:     false,
		systemEnv:     map[string]string{"test1": "val1", "test2": ""},
		cliFlags:      []string{"--env", "test3=val3", "-e", "test4", "-e", "test5=", "--include-system-env-vars=1"},
		expEnv:        map[string]string{"test1": "val1", "test2": "", "test3": "val3", "test4": "", "test5": ""},
		expCompatMode: defaultCompatMode,
	},
	"disabled system env with cli params": {
		useSysEnv:     false,
		systemEnv:     map[string]string{"test1": "val1"},
		cliFlags:      []string{"-e", "test2=overwriten", "-e", "test2=val2"},
		expEnv:        map[string]string{"test2": "val2"},
		expCompatMode: defaultCompatMode,
	},
	"overwriting system env with cli param": {
		useSysEnv:     true,
		systemEnv:     map[string]string{"test1": "val1sys"},
		cliFlags:      []string{"--env", "test1=val1cli"},
		expEnv:        map[string]string{"test1": "val1cli"},
		expCompatMode: defaultCompatMode,
	},
	"error wrong compat mode env var value": {
		systemEnv: map[string]string{"K6_COMPATIBILITY_MODE": "asdf"},
		expErr:    true,
	},
	"error wrong compat mode cli flag value": {
		cliFlags: []string{"--compatibility-mode", "whatever"},
		expErr:   true,
	},
	"error invalid cli var name 1": {
		useSysEnv:     true,
		systemEnv:     map[string]string{},
		cliFlags:      []string{"--env", "test a=error"},
		expErr:        true,
		expEnv:        map[string]string{},
		expCompatMode: defaultCompatMode,
	},
	"error invalid cli var name 2": {
		useSysEnv:     true,
		systemEnv:     map[string]string{},
		cliFlags:      []string{"--env", "1var=error"},
		expErr:        true,
		expEnv:        map[string]string{},
		expCompatMode: defaultCompatMode,
	},
	"error invalid cli var name 3": {
		useSysEnv:     true,
		systemEnv:     map[string]string{},
		cliFlags:      []string{"--env", "уникод=unicode-disabled"},
		expErr:        true,
		expEnv:        map[string]string{},
		expCompatMode: defaultCompatMode,
	},
	"valid env vars with spaces": {
		useSysEnv:     true,
		systemEnv:     map[string]string{"test1": "value 1"},
		cliFlags:      []string{"--env", "test2=value 2"},
		expEnv:        map[string]string{"test1": "value 1", "test2": "value 2"},
		expCompatMode: defaultCompatMode,
	},
	"valid env vars with special chars": {
		useSysEnv:     true,
		systemEnv:     map[string]string{"test1": "value 1"},
		cliFlags:      []string{"--env", "test2=value,2", "-e", `test3= ,  ,,, value, ,, 2!'@#,"`},
		expEnv:        map[string]string{"test1": "value 1", "test2": "value,2", "test3": ` ,  ,,, value, ,, 2!'@#,"`},
		expCompatMode: defaultCompatMode,
	},
}

func testRuntimeOptionsCase(t *testing.T, tc runtimeOptionsTestCase) {
	flags := runtimeOptionFlagSet(tc.useSysEnv)
	require.NoError(t, flags.Parse(tc.cliFlags))

	rtOpts, err := getRuntimeOptions(flags, tc.systemEnv)
	if tc.expErr {
		require.Error(t, err)
		return
	}
	require.NoError(t, err)
	require.EqualValues(t, tc.expEnv, rtOpts.Env)
	assert.Equal(t, tc.expCompatMode, rtOpts.CompatibilityMode)

	compatMode, err := lib.ValidateCompatibilityMode(rtOpts.CompatibilityMode.String)
	require.NoError(t, err)

	jsCode := new(bytes.Buffer)
	if compatMode == lib.CompatibilityModeExtended {
		fmt.Fprint(jsCode, "export default function() {")
	} else {
		fmt.Fprint(jsCode, "module.exports.default = function() {")
	}

	for key, val := range tc.expEnv {
		fmt.Fprintf(jsCode,
			"if (__ENV.%s !== `%s`) { throw new Error('Invalid %s: ' + __ENV.%s); }",
			key, val, key, key,
		)
	}
	fmt.Fprint(jsCode, "}")

	fs := afero.NewMemMapFs()
	require.NoError(t, afero.WriteFile(fs, "/script.js", jsCode.Bytes(), 0644))
	runner, err := newRunner(
		&loader.SourceData{Data: jsCode.Bytes(), URL: &url.URL{Path: "/script.js", Scheme: "file"}},
		typeJS,
		map[string]afero.Fs{"file": fs},
		rtOpts,
	)
	require.NoError(t, err)

	archive := runner.MakeArchive()
	archiveBuf := &bytes.Buffer{}
	require.NoError(t, archive.Write(archiveBuf))

	getRunnerErr := func(rtOpts lib.RuntimeOptions) (lib.Runner, error) {
		return newRunner(
			&loader.SourceData{
				Data: archiveBuf.Bytes(),
				URL:  &url.URL{Path: "/script.js"},
			},
			typeArchive,
			nil,
			rtOpts,
		)
	}

	_, err = getRunnerErr(lib.RuntimeOptions{})
	require.NoError(t, err)
	for key, val := range tc.expEnv {
		r, err := getRunnerErr(lib.RuntimeOptions{Env: map[string]string{key: "almost " + val}})
		assert.NoError(t, err)
		assert.Equal(t, r.MakeArchive().Env[key], "almost "+val)
	}
}

func TestRuntimeOptions(t *testing.T) {
	for name, tc := range runtimeOptionsTestCases {
		tc := tc
		t.Run(fmt.Sprintf("RuntimeOptions test '%s'", name), func(t *testing.T) {
			t.Parallel()
			testRuntimeOptionsCase(t, tc)
		})
	}
}
<|endoftext|>


<commit_before>package api

import (
	"io"
	"net"
	"os"
	"syscall"

	"sync"
	"time"

	"github.com/denkhaus/logging"
	"github.com/juju/errors"
	"github.com/mitchellh/mapstructure"
	"github.com/pquerna/ffjson/ffjson"
	"github.com/tevino/abool"
	"golang.org/x/net/websocket"
)

var (
	DialerTimeout    = time.Duration(5 * time.Second)
	ReadWriteTimeout = time.Duration(10 * time.Second)
	ErrShutdown      = errors.New("connection is shut down")
)

type wsClient struct {
	*ffjson.Decoder
	*ffjson.Encoder
	conn        *websocket.Conn
	url         string
	resp        rpcResponse // unmarshal target
	notify      rpcNotify   // unmarshal target
	onError     ErrorFunc
	errors      chan error
	closing     *abool.AtomicBool
	shutdown    *abool.AtomicBool
	currentID   uint64
	wg          sync.WaitGroup
	mutex       sync.Mutex // protects the following
	pending     map[uint64]*RPCCall
	mutexNotify sync.Mutex // protects the following
	notifyFns   map[int]NotifyFunc
}

func NewWebsocketClient(endpointURL string) WebsocketClient {
	cli := wsClient{
		closing:   abool.NewBool(false),
		shutdown:  abool.NewBool(false),
		pending:   make(map[uint64]*RPCCall),
		notifyFns: make(map[int]NotifyFunc),
		currentID: 1,
		url:       endpointURL,
	}

	return &cli
}

func (p *wsClient) Connect() error {
	config, err := websocket.NewConfig(p.url, "http://localhost/")
	if err != nil {
		return errors.Annotate(err, "NewConfig")
	}

	config.Dialer = &net.Dialer{
		Timeout: DialerTimeout,
	}

	conn, err := websocket.DialConfig(config)
	if err != nil {
		return errors.Annotate(err, "Dial")
	}

	p.shutdown.UnSet()
	p.closing.UnSet()

	p.errors = make(chan error, 10)
	p.Decoder = ffjson.NewDecoder()
	p.Encoder = ffjson.NewEncoder(conn)
	p.conn = conn

	p.wg.Add(1)
	go p.monitor()

	p.wg.Add(1)
	go p.receive()

	return nil
}

func (p *wsClient) Close() error {
	if p.conn != nil {
		p.closing.Set()
		if !p.shutdown.IsSet() {
			if err := p.conn.SetDeadline(time.Now().Add(ReadWriteTimeout)); err != nil {
				return errors.Annotate(err, "SetDeadline")
			}
			if err := p.conn.Close(); err != nil {
				return errors.Annotate(err, "Close [conn]")
			}
		}

		p.wg.Wait()
		close(p.errors)
		p.conn = nil
	}

	return nil
}

func (p *wsClient) IsConnected() bool {
	if p.shutdown.IsSet() || p.closing.IsSet() {
		return false
	}

	return p.conn != nil
}

func (p *wsClient) monitor() {
	defer p.wg.Done()

	for !p.shutdown.IsSet() {
		select {
		case err := <-p.errors:
			if err != nil {
				if p.onError != nil {
					p.onError(err)
				} else {
					logging.Errorf("WebsocketClient error: %s", err)
					logging.Warn("please set the API OnError hook to avoid this message")
				}
			}
		default:
		}
	}
}

func (p *wsClient) handleCustomData(data map[string]interface{}) error {
	logging.DDumpJSON("ws notify <", data)

	switch {
	case p.notify.Is(data):
		p.notify.reset()
		err := mapstructure.Decode(data, &p.notify)
		if err != nil {
			return errors.Annotate(err, "Decode [notify]")
		}

		params := p.notify.Params.([]interface{})
		subscriberID := int(params[0].(float64))

		var fn NotifyFunc
		p.mutexNotify.Lock()
		fn = p.notifyFns[subscriberID]
		p.mutexNotify.Unlock()

		if fn != nil {
			if err := fn(params[1]); err != nil {
				return errors.Annotate(err, "handle notify")
			}
		}
	default:
		return errors.Errorf("unhandled custom data: %v", data)
	}

	return nil
}

func (p *wsClient) mustEndReceive(err error) bool {
	if e, ok := err.(*net.OpError); ok {
		if e.Err.Error() == "use of closed network connection" {
			return true
		}

		if syscallErr, ok := e.Err.(*os.SyscallError); ok {
			if syscallErr.Err == syscall.ECONNRESET {
				return true
			}
		}
	}

	if e, ok := err.(net.Error); ok {
		if e.Timeout() {
			return true
		}
	}

	return false
}

func (p *wsClient) receive() {
	defer p.wg.Done()

	for !p.closing.IsSet() {
		//TODO: is there a faster way to distinguish between RPCResponse and RPCNotify data
		var data map[string]interface{}
		if err := p.DecodeReader(p.conn, &data); err != nil {
			if p.mustEndReceive(err) {
				p.errors <- errors.Annotate(err, "mustEndReceive")
				break
			}

			//report all errors but EOF
			if err != io.EOF {
				p.errors <- errors.Annotate(err, "DecodeReader")
			}

			continue
		}

		if p.resp.Is(data) {
			p.resp.reset()
			if err := mapstructure.Decode(data, &p.resp); err != nil {
				p.errors <- errors.Annotate(err, "Decode [resp]")
				continue
			}

			logging.DDumpJSON("ws resp <", data)

			if call, ok := p.pending[p.resp.ID]; ok {
				p.mutex.Lock()
				delete(p.pending, p.resp.ID)
				p.mutex.Unlock()

				call.Reply = p.resp.Result
				if p.resp.HasError() {
					call.Error = p.resp.Error
				}

				call.done()
			} else {
				p.errors <- errors.Errorf("no corresponding call found for incoming rpc data %v", p.resp)
				continue
			}
		} else if err := p.handleCustomData(data); err != nil {
			p.errors <- errors.Annotate(err, "handleCustomData")
			continue
		}
	}

	// Terminate pending calls
	p.mutex.Lock()
	defer p.mutex.Unlock()

	p.shutdown.Set()
	for _, call := range p.pending {
		call.Error = ErrShutdown
		call.done()
	}
}

func (p *wsClient) OnNotify(subscriberID int, fn NotifyFunc) error {
	if _, ok := p.notifyFns[subscriberID]; ok {
		return errors.Errorf("a notify hook for subscriberID %d is already defined", subscriberID)
	}

	p.mutexNotify.Lock()
	p.notifyFns[subscriberID] = fn
	p.mutexNotify.Unlock()

	return nil
}

func (p *wsClient) OnError(fn ErrorFunc) {
	p.onError = fn
}

func (p *wsClient) CallAPI(apiID int, method string, args ...interface{}) (interface{}, error) {
	param := []interface{}{
		apiID,
		method,
		args,
	}

	call, err := p.Call("call", param)
	if err != nil {
		return nil, errors.Annotate(err, "Call")
	}

	<-call.Done
	return call.Reply, call.Error
}

func (p *wsClient) Call(method string, args []interface{}) (*RPCCall, error) {
	if !p.IsConnected() {
		return nil, ErrShutdown
	}

	call := &RPCCall{
		Request: rpcRequest{
			Method: method,
			Params: args,
			ID:     p.currentID,
		},
		Done: make(chan *RPCCall, 10),
	}

	p.mutex.Lock()
	p.currentID++
	p.pending[call.Request.ID] = call
	p.mutex.Unlock()

	logging.DDumpJSON("ws req >", call.Request)

	if err := p.conn.SetDeadline(time.Now().Add(ReadWriteTimeout)); err != nil {
		return nil, errors.Annotate(err, "SetDeadline")
	}

	if err := p.Encode(call.Request); err != nil {
		p.mutex.Lock()
		delete(p.pending, call.Request.ID)
		p.mutex.Unlock()

		return nil, errors.Annotate(err, "Encode [req]")
	}

	return call, nil
}
<commit_msg>don't report error<commit_after>package api

import (
	"io"
	"net"
	"os"
	"syscall"

	"sync"
	"time"

	"github.com/denkhaus/logging"
	"github.com/juju/errors"
	"github.com/mitchellh/mapstructure"
	"github.com/pquerna/ffjson/ffjson"
	"github.com/tevino/abool"
	"golang.org/x/net/websocket"
)

var (
	DialerTimeout    = time.Duration(5 * time.Second)
	ReadWriteTimeout = time.Duration(10 * time.Second)
	ErrShutdown      = errors.New("connection is shut down")
)

type wsClient struct {
	*ffjson.Decoder
	*ffjson.Encoder
	conn        *websocket.Conn
	url         string
	resp        rpcResponse // unmarshal target
	notify      rpcNotify   // unmarshal target
	onError     ErrorFunc
	errors      chan error
	closing     *abool.AtomicBool
	shutdown    *abool.AtomicBool
	currentID   uint64
	wg          sync.WaitGroup
	mutex       sync.Mutex // protects the following
	pending     map[uint64]*RPCCall
	mutexNotify sync.Mutex // protects the following
	notifyFns   map[int]NotifyFunc
}

func NewWebsocketClient(endpointURL string) WebsocketClient {
	cli := wsClient{
		closing:   abool.NewBool(false),
		shutdown:  abool.NewBool(false),
		pending:   make(map[uint64]*RPCCall),
		notifyFns: make(map[int]NotifyFunc),
		currentID: 1,
		url:       endpointURL,
	}

	return &cli
}

func (p *wsClient) Connect() error {
	config, err := websocket.NewConfig(p.url, "http://localhost/")
	if err != nil {
		return errors.Annotate(err, "NewConfig")
	}

	config.Dialer = &net.Dialer{
		Timeout: DialerTimeout,
	}

	conn, err := websocket.DialConfig(config)
	if err != nil {
		return errors.Annotate(err, "Dial")
	}

	p.shutdown.UnSet()
	p.closing.UnSet()

	p.errors = make(chan error, 10)
	p.Decoder = ffjson.NewDecoder()
	p.Encoder = ffjson.NewEncoder(conn)
	p.conn = conn

	p.wg.Add(1)
	go p.monitor()

	p.wg.Add(1)
	go p.receive()

	return nil
}

func (p *wsClient) Close() error {
	if p.conn != nil {
		p.closing.Set()
		if !p.shutdown.IsSet() {
			if err := p.conn.SetDeadline(time.Now().Add(ReadWriteTimeout)); err != nil {
				return errors.Annotate(err, "SetDeadline")
			}
			if err := p.conn.Close(); err != nil {
				return errors.Annotate(err, "Close [conn]")
			}
		}

		p.wg.Wait()
		close(p.errors)
		p.conn = nil
	}

	return nil
}

func (p *wsClient) IsConnected() bool {
	if p.shutdown.IsSet() || p.closing.IsSet() {
		return false
	}

	return p.conn != nil
}

func (p *wsClient) monitor() {
	defer p.wg.Done()

	for !p.shutdown.IsSet() {
		select {
		case err := <-p.errors:
			if err != nil {
				if p.onError != nil {
					p.onError(err)
				} else {
					logging.Errorf("WebsocketClient error: %s", err)
					logging.Warn("please set the API OnError hook to avoid this message")
				}
			}
		default:
		}
	}
}

func (p *wsClient) handleCustomData(data map[string]interface{}) error {
	logging.DDumpJSON("ws notify <", data)

	switch {
	case p.notify.Is(data):
		p.notify.reset()
		err := mapstructure.Decode(data, &p.notify)
		if err != nil {
			return errors.Annotate(err, "Decode [notify]")
		}

		params := p.notify.Params.([]interface{})
		subscriberID := int(params[0].(float64))

		var fn NotifyFunc
		p.mutexNotify.Lock()
		fn = p.notifyFns[subscriberID]
		p.mutexNotify.Unlock()

		if fn != nil {
			if err := fn(params[1]); err != nil {
				return errors.Annotate(err, "handle notify")
			}
		}
	default:
		return errors.Errorf("unhandled custom data: %v", data)
	}

	return nil
}

func (p *wsClient) mustEndReceive(err error) bool {
	if e, ok := err.(*net.OpError); ok {
		if e.Err.Error() == "use of closed network connection" {
			return true
		}

		if syscallErr, ok := e.Err.(*os.SyscallError); ok {
			if syscallErr.Err == syscall.ECONNRESET {
				return true
			}
		}
	}

	if e, ok := err.(net.Error); ok {
		if e.Timeout() {
			return true
		}
	}

	return false
}

func (p *wsClient) receive() {
	defer p.wg.Done()

	for !p.closing.IsSet() {
		//TODO: is there a faster way to distinguish between RPCResponse and RPCNotify data
		var data map[string]interface{}
		if err := p.DecodeReader(p.conn, &data); err != nil {
			if p.mustEndReceive(err) {
				break
			}

			//report all errors but EOF
			if err != io.EOF {
				p.errors <- errors.Annotate(err, "DecodeReader")
			}

			continue
		}

		if p.resp.Is(data) {
			p.resp.reset()
			if err := mapstructure.Decode(data, &p.resp); err != nil {
				p.errors <- errors.Annotate(err, "Decode [resp]")
				continue
			}

			logging.DDumpJSON("ws resp <", data)

			if call, ok := p.pending[p.resp.ID]; ok {
				p.mutex.Lock()
				delete(p.pending, p.resp.ID)
				p.mutex.Unlock()

				call.Reply = p.resp.Result
				if p.resp.HasError() {
					call.Error = p.resp.Error
				}

				call.done()
			} else {
				p.errors <- errors.Errorf("no corresponding call found for incoming rpc data %v", p.resp)
				continue
			}
		} else if err := p.handleCustomData(data); err != nil {
			p.errors <- errors.Annotate(err, "handleCustomData")
			continue
		}
	}

	// Terminate pending calls
	p.mutex.Lock()
	defer p.mutex.Unlock()

	p.shutdown.Set()
	for _, call := range p.pending {
		call.Error = ErrShutdown
		call.done()
	}
}

func (p *wsClient) OnNotify(subscriberID int, fn NotifyFunc) error {
	if _, ok := p.notifyFns[subscriberID]; ok {
		return errors.Errorf("a notify hook for subscriberID %d is already defined", subscriberID)
	}

	p.mutexNotify.Lock()
	p.notifyFns[subscriberID] = fn
	p.mutexNotify.Unlock()

	return nil
}

func (p *wsClient) OnError(fn ErrorFunc) {
	p.onError = fn
}

func (p *wsClient) CallAPI(apiID int, method string, args ...interface{}) (interface{}, error) {
	param := []interface{}{
		apiID,
		method,
		args,
	}

	call, err := p.Call("call", param)
	if err != nil {
		return nil, errors.Annotate(err, "Call")
	}

	<-call.Done
	return call.Reply, call.Error
}

func (p *wsClient) Call(method string, args []interface{}) (*RPCCall, error) {
	if !p.IsConnected() {
		return nil, ErrShutdown
	}

	call := &RPCCall{
		Request: rpcRequest{
			Method: method,
			Params: args,
			ID:     p.currentID,
		},
		Done: make(chan *RPCCall, 10),
	}

	p.mutex.Lock()
	p.currentID++
	p.pending[call.Request.ID] = call
	p.mutex.Unlock()

	logging.DDumpJSON("ws req >", call.Request)

	if err := p.conn.SetDeadline(time.Now().Add(ReadWriteTimeout)); err != nil {
		return nil, errors.Annotate(err, "SetDeadline")
	}

	if err := p.Encode(call.Request); err != nil {
		p.mutex.Lock()
		delete(p.pending, call.Request.ID)
		p.mutex.Unlock()

		return nil, errors.Annotate(err, "Encode [req]")
	}

	return call, nil
}
<|endoftext|>


<commit_before>// Package kitchen implements a driver based on Test Kitchen.
package kitchen

import (
	"errors"
	"fmt"
	"io/ioutil"
	"path"
	"strconv"
	"strings"

	"github.com/mlafeldt/chef-runner/log"
	"github.com/mlafeldt/chef-runner/openssh"
	"github.com/mlafeldt/chef-runner/rsync"
	"gopkg.in/yaml.v2"
)

// Driver is a driver based on Test Kitchen.
type Driver struct {
	Instance    string
	SSHClient   *openssh.Client
	RsyncClient *rsync.Client
}

// This is what `vagrant ssh` uses
var defaultSSHOptions = [...]string{
	"UserKnownHostsFile /dev/null",
	"StrictHostKeyChecking no",
	"PasswordAuthentication no",
	"IdentitiesOnly yes",
	"LogLevel FATAL",
}

type kitchenConfigEntry struct {
	Name string `yaml:"name"`
}

type kitchenConfig struct {
	Platforms []kitchenConfigEntry `yaml:"platforms"`
	Suites    []kitchenConfigEntry `yaml:"suites"`
}

type instanceConfig struct {
	Hostname string `yaml:"hostname"`
	Username string `yaml:"username"`
	SSHKey   string `yaml:"ssh_key"`
	Port     string `yaml:"port"`
}

func readInstanceNames() ([]string, error) {
	data, err := ioutil.ReadFile(".kitchen.yml")
	if err != nil {
		return nil, err
	}

	var config kitchenConfig
	if err := yaml.Unmarshal(data, &config); err != nil {
		return nil, err
	}

	var names []string
	for _, suite := range config.Suites {
		for _, platform := range config.Platforms {
			s := suite.Name + "-" + platform.Name
			s = strings.Replace(s, "_", "-", -1)
			s = strings.Replace(s, ".", "", -1)
			names = append(names, s)
		}
	}
	return names, nil
}

func findInstanceName(instance string) (string, error) {
	instanceNames, err := readInstanceNames()
	if err != nil {
		return "", err
	}
	log.Debugf("Kitchen instances = %s\n", instanceNames)

	for _, name := range instanceNames {
		// Return first match
		if strings.Contains(name, instance) {
			return name, nil
		}
	}
	return "", errors.New("Kitchen instance not found")
}

func readInstanceConfig(instance string) (*instanceConfig, error) {
	configFile := path.Join(".kitchen", instance+".yml")
	log.Debugf("Kitchen config file = %s\n", configFile)

	data, err := ioutil.ReadFile(configFile)
	if err != nil {
		return nil, err
	}

	var config instanceConfig
	if err := yaml.Unmarshal(data, &config); err != nil {
		return nil, err
	}
	log.Debugf("Kitchen config = %+v\n", config)

	if config.Hostname == "" {
		return nil, errors.New(configFile + ": invalid `hostname`")
	}
	if config.Username == "" {
		return nil, errors.New(configFile + ": invalid `username`")
	}
	if config.SSHKey == "" {
		return nil, errors.New(configFile + ": invalid `ssh_key`")
	}
	if _, err := strconv.Atoi(config.Port); err != nil {
		return nil, errors.New(configFile + ": invalid `port`")
	}

	return &config, nil
}

// NewDriver creates a new Test Kitchen driver that communicates with the given
// Test Kitchen instance. Under the hood the instance's YAML configuration is
// parsed to get a working SSH configuration.
func NewDriver(instance string, sshOptions, rsyncOptions []string) (*Driver, error) {
	instance, err := findInstanceName(instance)
	if err != nil {
		return nil, err
	}

	config, err := readInstanceConfig(instance)
	if err != nil {
		return nil, err
	}

	// Test Kitchen stores the port as an string
	port, _ := strconv.Atoi(config.Port)

	sshOpts := make([]string, len(defaultSSHOptions))
	copy(sshOpts, defaultSSHOptions[:])
	for _, o := range sshOptions {
		sshOpts = append(sshOpts, o)
	}

	sshClient := &openssh.Client{
		Host:        config.Hostname,
		User:        config.Username,
		Port:        port,
		PrivateKeys: []string{config.SSHKey},
		Options:     sshOpts,
	}

	rsyncClient := *rsync.MirrorClient
	rsyncClient.RemoteHost = config.Hostname
	rsyncClient.RemoteShell = sshClient.Shell()
	rsyncClient.Options = rsyncOptions

	return &Driver{instance, sshClient, &rsyncClient}, nil
}

// RunCommand runs the specified command on the Test Kitchen instance.
func (drv Driver) RunCommand(args []string) error {
	return drv.SSHClient.RunCommand(args)
}

// Upload copies files to the Test Kitchen instance.
func (drv Driver) Upload(dst string, src ...string) error {
	return drv.RsyncClient.Copy(dst, src...)
}

// String returns the driver's name.
func (drv Driver) String() string {
	return fmt.Sprintf("Test Kitchen driver (instance: %s)", drv.Instance)
}
<commit_msg>Use auxiliary struct type for Kitchen instance config<commit_after>// Package kitchen implements a driver based on Test Kitchen.
package kitchen

import (
	"errors"
	"fmt"
	"io/ioutil"
	"path"
	"strconv"
	"strings"

	"github.com/mlafeldt/chef-runner/log"
	"github.com/mlafeldt/chef-runner/openssh"
	"github.com/mlafeldt/chef-runner/rsync"
	"gopkg.in/yaml.v2"
)

// Driver is a driver based on Test Kitchen.
type Driver struct {
	Instance    string
	SSHClient   *openssh.Client
	RsyncClient *rsync.Client
}

// This is what `vagrant ssh` uses
var defaultSSHOptions = [...]string{
	"UserKnownHostsFile /dev/null",
	"StrictHostKeyChecking no",
	"PasswordAuthentication no",
	"IdentitiesOnly yes",
	"LogLevel FATAL",
}

type kitchenConfigEntry struct {
	Name string `yaml:"name"`
}

type kitchenConfig struct {
	Platforms []kitchenConfigEntry `yaml:"platforms"`
	Suites    []kitchenConfigEntry `yaml:"suites"`
}

func readInstanceNames() ([]string, error) {
	data, err := ioutil.ReadFile(".kitchen.yml")
	if err != nil {
		return nil, err
	}

	var config kitchenConfig
	if err := yaml.Unmarshal(data, &config); err != nil {
		return nil, err
	}

	var names []string
	for _, suite := range config.Suites {
		for _, platform := range config.Platforms {
			s := suite.Name + "-" + platform.Name
			s = strings.Replace(s, "_", "-", -1)
			s = strings.Replace(s, ".", "", -1)
			names = append(names, s)
		}
	}
	return names, nil
}

func findInstanceName(instance string) (string, error) {
	instanceNames, err := readInstanceNames()
	if err != nil {
		return "", err
	}
	log.Debugf("Kitchen instances = %s\n", instanceNames)

	for _, name := range instanceNames {
		// Return first match
		if strings.Contains(name, instance) {
			return name, nil
		}
	}
	return "", errors.New("Kitchen instance not found")
}

type instanceConfig struct {
	Hostname string
	Username string
	SSHKey   string
	Port     int
}

func (c *instanceConfig) Parse(data []byte) error {
	var aux struct {
		Hostname string `yaml:"hostname"`
		Username string `yaml:"username"`
		SSHKey   string `yaml:"ssh_key"`
		Port     string `yaml:"port"`
	}

	if err := yaml.Unmarshal(data, &aux); err != nil {
		return err
	}
	if aux.Hostname == "" {
		return errors.New("Kitchen config: invalid `hostname`")
	}
	if aux.Username == "" {
		return errors.New("Kitchen config: invalid `username`")
	}
	if aux.SSHKey == "" {
		return errors.New("Kitchen config: invalid `ssh_key`")
	}
	// Test Kitchen stores the port as an string
	port, err := strconv.Atoi(aux.Port)
	if err != nil {
		return errors.New("Kitchen config: invalid `port`")
	}

	c.Hostname = aux.Hostname
	c.Username = aux.Username
	c.SSHKey = aux.SSHKey
	c.Port = port
	return nil
}

func readInstanceConfig(instance string) (*instanceConfig, error) {
	configFile := path.Join(".kitchen", instance+".yml")
	log.Debugf("Kitchen config file = %s\n", configFile)

	data, err := ioutil.ReadFile(configFile)
	if err != nil {
		return nil, err
	}

	var config instanceConfig
	if err := config.Parse(data); err != nil {
		return nil, err
	}
	log.Debugf("Kitchen config = %+v\n", config)

	return &config, nil
}

// NewDriver creates a new Test Kitchen driver that communicates with the given
// Test Kitchen instance. Under the hood the instance's YAML configuration is
// parsed to get a working SSH configuration.
func NewDriver(instance string, sshOptions, rsyncOptions []string) (*Driver, error) {
	instance, err := findInstanceName(instance)
	if err != nil {
		return nil, err
	}

	config, err := readInstanceConfig(instance)
	if err != nil {
		return nil, err
	}

	sshOpts := make([]string, len(defaultSSHOptions))
	copy(sshOpts, defaultSSHOptions[:])
	for _, o := range sshOptions {
		sshOpts = append(sshOpts, o)
	}

	sshClient := &openssh.Client{
		Host:        config.Hostname,
		User:        config.Username,
		Port:        config.Port,
		PrivateKeys: []string{config.SSHKey},
		Options:     sshOpts,
	}

	rsyncClient := *rsync.MirrorClient
	rsyncClient.RemoteHost = config.Hostname
	rsyncClient.RemoteShell = sshClient.Shell()
	rsyncClient.Options = rsyncOptions

	return &Driver{instance, sshClient, &rsyncClient}, nil
}

// RunCommand runs the specified command on the Test Kitchen instance.
func (drv Driver) RunCommand(args []string) error {
	return drv.SSHClient.RunCommand(args)
}

// Upload copies files to the Test Kitchen instance.
func (drv Driver) Upload(dst string, src ...string) error {
	return drv.RsyncClient.Copy(dst, src...)
}

// String returns the driver's name.
func (drv Driver) String() string {
	return fmt.Sprintf("Test Kitchen driver (instance: %s)", drv.Instance)
}
<|endoftext|>


<commit_before>package terraform

import (
	"fmt"
	"os"
	"strings"
	"sync"

	"github.com/hashicorp/terraform/config"
	"github.com/hashicorp/terraform/config/lang/ast"
	"github.com/hashicorp/terraform/config/module"
)

const (
	// VarEnvPrefix is the prefix of variables that are read from
	// the environment to set variables here.
	VarEnvPrefix = "TF_VAR_"
)

// Interpolater is the structure responsible for determining the values
// for interpolations such as `aws_instance.foo.bar`.
type Interpolater struct {
	Operation walkOperation
	Module    *module.Tree
	State     *State
	StateLock *sync.RWMutex
	Variables map[string]string
}

// InterpolationScope is the current scope of execution. This is required
// since some variables which are interpolated are dependent on what we're
// operating on and where we are.
type InterpolationScope struct {
	Path     []string
	Resource *Resource
}

// Values returns the values for all the variables in the given map.
func (i *Interpolater) Values(
	scope *InterpolationScope,
	vars map[string]config.InterpolatedVariable) (map[string]ast.Variable, error) {
	result := make(map[string]ast.Variable, len(vars))

	// Copy the default variables
	if i.Module != nil && scope != nil {
		mod := i.Module
		if len(scope.Path) > 1 {
			mod = i.Module.Child(scope.Path[1:])
		}
		for _, v := range mod.Config().Variables {
			for k, val := range v.DefaultsMap() {
				envKey := VarEnvPrefix + strings.TrimPrefix(k, "var.")
				if v := os.Getenv(envKey); v != "" {
					val = v
				}

				result[k] = ast.Variable{
					Value: val,
					Type:  ast.TypeString,
				}
			}
		}
	}

	for n, rawV := range vars {
		var err error
		switch v := rawV.(type) {
		case *config.CountVariable:
			err = i.valueCountVar(scope, n, v, result)
		case *config.ModuleVariable:
			err = i.valueModuleVar(scope, n, v, result)
		case *config.PathVariable:
			err = i.valuePathVar(scope, n, v, result)
		case *config.ResourceVariable:
			err = i.valueResourceVar(scope, n, v, result)
		case *config.SelfVariable:
			err = i.valueSelfVar(scope, n, v, result)
		case *config.UserVariable:
			err = i.valueUserVar(scope, n, v, result)
		default:
			err = fmt.Errorf("%s: unknown variable type: %T", n, rawV)
		}

		if err != nil {
			return nil, err
		}
	}

	return result, nil
}

func (i *Interpolater) valueCountVar(
	scope *InterpolationScope,
	n string,
	v *config.CountVariable,
	result map[string]ast.Variable) error {
	switch v.Type {
	case config.CountValueIndex:
		if scope.Resource == nil {
			return fmt.Errorf("%s: count.index is only valid within resources", n)
		}
		result[n] = ast.Variable{
			Value: scope.Resource.CountIndex,
			Type:  ast.TypeInt,
		}
		return nil
	default:
		return fmt.Errorf("%s: unknown count type: %#v", n, v.Type)
	}
}

func (i *Interpolater) valueModuleVar(
	scope *InterpolationScope,
	n string,
	v *config.ModuleVariable,
	result map[string]ast.Variable) error {
	// If we're computing all dynamic fields, then module vars count
	// and we mark it as computed.
	if i.Operation == walkValidate {
		result[n] = ast.Variable{
			Value: config.UnknownVariableValue,
			Type:  ast.TypeString,
		}
		return nil
	}

	// Build the path to the child module we want
	path := make([]string, len(scope.Path), len(scope.Path)+1)
	copy(path, scope.Path)
	path = append(path, v.Name)

	// Grab the lock so that if other interpolations are running or
	// state is being modified, we'll be safe.
	i.StateLock.RLock()
	defer i.StateLock.RUnlock()

	// Get the module where we're looking for the value
	var value string
	mod := i.State.ModuleByPath(path)
	if mod == nil {
		// If the module doesn't exist, then we can return an empty string.
		// This happens usually only in Refresh() when we haven't populated
		// a state. During validation, we semantically verify that all
		// modules reference other modules, and graph ordering should
		// ensure that the module is in the state, so if we reach this
		// point otherwise it really is a panic.
		value = config.UnknownVariableValue
	} else {
		// Get the value from the outputs
		var ok bool
		value, ok = mod.Outputs[v.Field]
		if !ok {
			// Same reasons as the comment above.
			value = config.UnknownVariableValue
		}
	}

	result[n] = ast.Variable{
		Value: value,
		Type:  ast.TypeString,
	}
	return nil
}

func (i *Interpolater) valuePathVar(
	scope *InterpolationScope,
	n string,
	v *config.PathVariable,
	result map[string]ast.Variable) error {
	switch v.Type {
	case config.PathValueCwd:
		wd, err := os.Getwd()
		if err != nil {
			return fmt.Errorf(
				"Couldn't get cwd for var %s: %s",
				v.FullKey(), err)
		}

		result[n] = ast.Variable{
			Value: wd,
			Type:  ast.TypeString,
		}
	case config.PathValueModule:
		if t := i.Module.Child(scope.Path[1:]); t != nil {
			result[n] = ast.Variable{
				Value: t.Config().Dir,
				Type:  ast.TypeString,
			}
		}
	case config.PathValueRoot:
		result[n] = ast.Variable{
			Value: i.Module.Config().Dir,
			Type:  ast.TypeString,
		}
	default:
		return fmt.Errorf("%s: unknown path type: %#v", n, v.Type)
	}

	return nil

}

func (i *Interpolater) valueResourceVar(
	scope *InterpolationScope,
	n string,
	v *config.ResourceVariable,
	result map[string]ast.Variable) error {
	// If we're computing all dynamic fields, then module vars count
	// and we mark it as computed.
	if i.Operation == walkValidate {
		result[n] = ast.Variable{
			Value: config.UnknownVariableValue,
			Type:  ast.TypeString,
		}
		return nil
	}

	var attr string
	var err error
	if v.Multi && v.Index == -1 {
		attr, err = i.computeResourceMultiVariable(scope, v)
	} else {
		attr, err = i.computeResourceVariable(scope, v)
	}
	if err != nil {
		return err
	}

	result[n] = ast.Variable{
		Value: attr,
		Type:  ast.TypeString,
	}
	return nil
}

func (i *Interpolater) valueSelfVar(
	scope *InterpolationScope,
	n string,
	v *config.SelfVariable,
	result map[string]ast.Variable) error {
	rv, err := config.NewResourceVariable(fmt.Sprintf(
		"%s.%s.%d.%s",
		scope.Resource.Type,
		scope.Resource.Name,
		scope.Resource.CountIndex,
		v.Field))
	if err != nil {
		return err
	}

	return i.valueResourceVar(scope, n, rv, result)
}

func (i *Interpolater) valueUserVar(
	scope *InterpolationScope,
	n string,
	v *config.UserVariable,
	result map[string]ast.Variable) error {
	val, ok := i.Variables[v.Name]
	if ok {
		result[n] = ast.Variable{
			Value: val,
			Type:  ast.TypeString,
		}
		return nil
	}

	if _, ok := result[n]; !ok && i.Operation == walkValidate {
		result[n] = ast.Variable{
			Value: config.UnknownVariableValue,
			Type:  ast.TypeString,
		}
		return nil
	}

	// Look up if we have any variables with this prefix because
	// those are map overrides. Include those.
	for k, val := range i.Variables {
		if strings.HasPrefix(k, v.Name+".") {
			result["var."+k] = ast.Variable{
				Value: val,
				Type:  ast.TypeString,
			}
		}
	}

	return nil
}

func (i *Interpolater) computeResourceVariable(
	scope *InterpolationScope,
	v *config.ResourceVariable) (string, error) {
	id := v.ResourceId()
	if v.Multi {
		id = fmt.Sprintf("%s.%d", id, v.Index)
	}

	i.StateLock.RLock()
	defer i.StateLock.RUnlock()

	// Get the information about this resource variable, and verify
	// that it exists and such.
	module, _, err := i.resourceVariableInfo(scope, v)
	if err != nil {
		return "", err
	}

	// If we have no module in the state yet or count, return empty
	if module == nil || len(module.Resources) == 0 {
		return "", nil
	}

	// Get the resource out from the state. We know the state exists
	// at this point and if there is a state, we expect there to be a
	// resource with the given name.
	r, ok := module.Resources[id]
	if !ok && v.Multi && v.Index == 0 {
		r, ok = module.Resources[v.ResourceId()]
	}
	if !ok {
		r = nil
	}
	if r == nil {
		return "", fmt.Errorf(
			"Resource '%s' not found for variable '%s'",
			id,
			v.FullKey())
	}

	if r.Primary == nil {
		goto MISSING
	}

	if attr, ok := r.Primary.Attributes[v.Field]; ok {
		return attr, nil
	}

	// At apply time, we can't do the "maybe has it" check below
	// that we need for plans since parent elements might be computed.
	// Therefore, it is an error and we're missing the key.
	//
	// TODO: test by creating a state and configuration that is referencing
	// a non-existent variable "foo.bar" where the state only has "foo"
	// and verify plan works, but apply doesn't.
	if i.Operation == walkApply {
		goto MISSING
	}

	// We didn't find the exact field, so lets separate the dots
	// and see if anything along the way is a computed set. i.e. if
	// we have "foo.0.bar" as the field, check to see if "foo" is
	// a computed list. If so, then the whole thing is computed.
	if parts := strings.Split(v.Field, "."); len(parts) > 1 {
		for i := 1; i < len(parts); i++ {
			// Lists and sets make this
			key := fmt.Sprintf("%s.#", strings.Join(parts[:i], "."))
			if attr, ok := r.Primary.Attributes[key]; ok {
				return attr, nil
			}

			// Maps make this
			key = fmt.Sprintf("%s", strings.Join(parts[:i], "."))
			if attr, ok := r.Primary.Attributes[key]; ok {
				return attr, nil
			}
		}
	}

MISSING:
	// If the operation is refresh, it isn't an error for a value to
	// be unknown. Instead, we return that the value is computed so
	// that the graph can continue to refresh other nodes. It doesn't
	// matter because the config isn't interpolated anyways.
	if i.Operation == walkRefresh {
		return config.UnknownVariableValue, nil
	}

	return "", fmt.Errorf(
		"Resource '%s' does not have attribute '%s' "+
			"for variable '%s'",
		id,
		v.Field,
		v.FullKey())
}

func (i *Interpolater) computeResourceMultiVariable(
	scope *InterpolationScope,
	v *config.ResourceVariable) (string, error) {
	i.StateLock.RLock()
	defer i.StateLock.RUnlock()

	// Get the information about this resource variable, and verify
	// that it exists and such.
	module, cr, err := i.resourceVariableInfo(scope, v)
	if err != nil {
		return "", err
	}

	// Get the count so we know how many to iterate over
	count, err := cr.Count()
	if err != nil {
		return "", fmt.Errorf(
			"Error reading %s count: %s",
			v.ResourceId(),
			err)
	}

	// If we have no module in the state yet or count, return empty
	if module == nil || len(module.Resources) == 0 || count == 0 {
		return "", nil
	}

	var values []string
	for i := 0; i < count; i++ {
		id := fmt.Sprintf("%s.%d", v.ResourceId(), i)

		// If we're dealing with only a single resource, then the
		// ID doesn't have a trailing index.
		if count == 1 {
			id = v.ResourceId()
		}

		r, ok := module.Resources[id]
		if !ok {
			continue
		}

		if r.Primary == nil {
			continue
		}

		attr, ok := r.Primary.Attributes[v.Field]
		if !ok {
			continue
		}

		values = append(values, attr)
	}

	if len(values) == 0 {
		return "", fmt.Errorf(
			"Resource '%s' does not have attribute '%s' "+
				"for variable '%s'",
			v.ResourceId(),
			v.Field,
			v.FullKey())
	}

	return strings.Join(values, config.InterpSplitDelim), nil
}

func (i *Interpolater) resourceVariableInfo(
	scope *InterpolationScope,
	v *config.ResourceVariable) (*ModuleState, *config.Resource, error) {
	// Get the module tree that contains our current path. This is
	// either the current module (path is empty) or a child.
	modTree := i.Module
	if len(scope.Path) > 1 {
		modTree = i.Module.Child(scope.Path[1:])
	}

	// Get the resource from the configuration so we can verify
	// that the resource is in the configuration and so we can access
	// the configuration if we need to.
	var cr *config.Resource
	for _, r := range modTree.Config().Resources {
		if r.Id() == v.ResourceId() {
			cr = r
			break
		}
	}
	if cr == nil {
		return nil, nil, fmt.Errorf(
			"Resource '%s' not found for variable '%s'",
			v.ResourceId(),
			v.FullKey())
	}

	// Get the relevant module
	module := i.State.ModuleByPath(scope.Path)
	return module, cr, nil
}
<commit_msg>terraform: return unknown if resource not found<commit_after>package terraform

import (
	"fmt"
	"os"
	"strings"
	"sync"

	"github.com/hashicorp/terraform/config"
	"github.com/hashicorp/terraform/config/lang/ast"
	"github.com/hashicorp/terraform/config/module"
)

const (
	// VarEnvPrefix is the prefix of variables that are read from
	// the environment to set variables here.
	VarEnvPrefix = "TF_VAR_"
)

// Interpolater is the structure responsible for determining the values
// for interpolations such as `aws_instance.foo.bar`.
type Interpolater struct {
	Operation walkOperation
	Module    *module.Tree
	State     *State
	StateLock *sync.RWMutex
	Variables map[string]string
}

// InterpolationScope is the current scope of execution. This is required
// since some variables which are interpolated are dependent on what we're
// operating on and where we are.
type InterpolationScope struct {
	Path     []string
	Resource *Resource
}

// Values returns the values for all the variables in the given map.
func (i *Interpolater) Values(
	scope *InterpolationScope,
	vars map[string]config.InterpolatedVariable) (map[string]ast.Variable, error) {
	result := make(map[string]ast.Variable, len(vars))

	// Copy the default variables
	if i.Module != nil && scope != nil {
		mod := i.Module
		if len(scope.Path) > 1 {
			mod = i.Module.Child(scope.Path[1:])
		}
		for _, v := range mod.Config().Variables {
			for k, val := range v.DefaultsMap() {
				envKey := VarEnvPrefix + strings.TrimPrefix(k, "var.")
				if v := os.Getenv(envKey); v != "" {
					val = v
				}

				result[k] = ast.Variable{
					Value: val,
					Type:  ast.TypeString,
				}
			}
		}
	}

	for n, rawV := range vars {
		var err error
		switch v := rawV.(type) {
		case *config.CountVariable:
			err = i.valueCountVar(scope, n, v, result)
		case *config.ModuleVariable:
			err = i.valueModuleVar(scope, n, v, result)
		case *config.PathVariable:
			err = i.valuePathVar(scope, n, v, result)
		case *config.ResourceVariable:
			err = i.valueResourceVar(scope, n, v, result)
		case *config.SelfVariable:
			err = i.valueSelfVar(scope, n, v, result)
		case *config.UserVariable:
			err = i.valueUserVar(scope, n, v, result)
		default:
			err = fmt.Errorf("%s: unknown variable type: %T", n, rawV)
		}

		if err != nil {
			return nil, err
		}
	}

	return result, nil
}

func (i *Interpolater) valueCountVar(
	scope *InterpolationScope,
	n string,
	v *config.CountVariable,
	result map[string]ast.Variable) error {
	switch v.Type {
	case config.CountValueIndex:
		if scope.Resource == nil {
			return fmt.Errorf("%s: count.index is only valid within resources", n)
		}
		result[n] = ast.Variable{
			Value: scope.Resource.CountIndex,
			Type:  ast.TypeInt,
		}
		return nil
	default:
		return fmt.Errorf("%s: unknown count type: %#v", n, v.Type)
	}
}

func (i *Interpolater) valueModuleVar(
	scope *InterpolationScope,
	n string,
	v *config.ModuleVariable,
	result map[string]ast.Variable) error {
	// If we're computing all dynamic fields, then module vars count
	// and we mark it as computed.
	if i.Operation == walkValidate {
		result[n] = ast.Variable{
			Value: config.UnknownVariableValue,
			Type:  ast.TypeString,
		}
		return nil
	}

	// Build the path to the child module we want
	path := make([]string, len(scope.Path), len(scope.Path)+1)
	copy(path, scope.Path)
	path = append(path, v.Name)

	// Grab the lock so that if other interpolations are running or
	// state is being modified, we'll be safe.
	i.StateLock.RLock()
	defer i.StateLock.RUnlock()

	// Get the module where we're looking for the value
	var value string
	mod := i.State.ModuleByPath(path)
	if mod == nil {
		// If the module doesn't exist, then we can return an empty string.
		// This happens usually only in Refresh() when we haven't populated
		// a state. During validation, we semantically verify that all
		// modules reference other modules, and graph ordering should
		// ensure that the module is in the state, so if we reach this
		// point otherwise it really is a panic.
		value = config.UnknownVariableValue
	} else {
		// Get the value from the outputs
		var ok bool
		value, ok = mod.Outputs[v.Field]
		if !ok {
			// Same reasons as the comment above.
			value = config.UnknownVariableValue
		}
	}

	result[n] = ast.Variable{
		Value: value,
		Type:  ast.TypeString,
	}
	return nil
}

func (i *Interpolater) valuePathVar(
	scope *InterpolationScope,
	n string,
	v *config.PathVariable,
	result map[string]ast.Variable) error {
	switch v.Type {
	case config.PathValueCwd:
		wd, err := os.Getwd()
		if err != nil {
			return fmt.Errorf(
				"Couldn't get cwd for var %s: %s",
				v.FullKey(), err)
		}

		result[n] = ast.Variable{
			Value: wd,
			Type:  ast.TypeString,
		}
	case config.PathValueModule:
		if t := i.Module.Child(scope.Path[1:]); t != nil {
			result[n] = ast.Variable{
				Value: t.Config().Dir,
				Type:  ast.TypeString,
			}
		}
	case config.PathValueRoot:
		result[n] = ast.Variable{
			Value: i.Module.Config().Dir,
			Type:  ast.TypeString,
		}
	default:
		return fmt.Errorf("%s: unknown path type: %#v", n, v.Type)
	}

	return nil

}

func (i *Interpolater) valueResourceVar(
	scope *InterpolationScope,
	n string,
	v *config.ResourceVariable,
	result map[string]ast.Variable) error {
	// If we're computing all dynamic fields, then module vars count
	// and we mark it as computed.
	if i.Operation == walkValidate {
		result[n] = ast.Variable{
			Value: config.UnknownVariableValue,
			Type:  ast.TypeString,
		}
		return nil
	}

	var attr string
	var err error
	if v.Multi && v.Index == -1 {
		attr, err = i.computeResourceMultiVariable(scope, v)
	} else {
		attr, err = i.computeResourceVariable(scope, v)
	}
	if err != nil {
		return err
	}

	result[n] = ast.Variable{
		Value: attr,
		Type:  ast.TypeString,
	}
	return nil
}

func (i *Interpolater) valueSelfVar(
	scope *InterpolationScope,
	n string,
	v *config.SelfVariable,
	result map[string]ast.Variable) error {
	rv, err := config.NewResourceVariable(fmt.Sprintf(
		"%s.%s.%d.%s",
		scope.Resource.Type,
		scope.Resource.Name,
		scope.Resource.CountIndex,
		v.Field))
	if err != nil {
		return err
	}

	return i.valueResourceVar(scope, n, rv, result)
}

func (i *Interpolater) valueUserVar(
	scope *InterpolationScope,
	n string,
	v *config.UserVariable,
	result map[string]ast.Variable) error {
	val, ok := i.Variables[v.Name]
	if ok {
		result[n] = ast.Variable{
			Value: val,
			Type:  ast.TypeString,
		}
		return nil
	}

	if _, ok := result[n]; !ok && i.Operation == walkValidate {
		result[n] = ast.Variable{
			Value: config.UnknownVariableValue,
			Type:  ast.TypeString,
		}
		return nil
	}

	// Look up if we have any variables with this prefix because
	// those are map overrides. Include those.
	for k, val := range i.Variables {
		if strings.HasPrefix(k, v.Name+".") {
			result["var."+k] = ast.Variable{
				Value: val,
				Type:  ast.TypeString,
			}
		}
	}

	return nil
}

func (i *Interpolater) computeResourceVariable(
	scope *InterpolationScope,
	v *config.ResourceVariable) (string, error) {
	id := v.ResourceId()
	if v.Multi {
		id = fmt.Sprintf("%s.%d", id, v.Index)
	}

	i.StateLock.RLock()
	defer i.StateLock.RUnlock()

	// Get the information about this resource variable, and verify
	// that it exists and such.
	module, _, err := i.resourceVariableInfo(scope, v)
	if err != nil {
		return "", err
	}

	// If we have no module in the state yet or count, return empty
	if module == nil || len(module.Resources) == 0 {
		return "", nil
	}

	// Get the resource out from the state. We know the state exists
	// at this point and if there is a state, we expect there to be a
	// resource with the given name.
	r, ok := module.Resources[id]
	if !ok && v.Multi && v.Index == 0 {
		r, ok = module.Resources[v.ResourceId()]
	}
	if !ok {
		r = nil
	}
	if r == nil {
		goto MISSING
	}

	if r.Primary == nil {
		goto MISSING
	}

	if attr, ok := r.Primary.Attributes[v.Field]; ok {
		return attr, nil
	}

	// At apply time, we can't do the "maybe has it" check below
	// that we need for plans since parent elements might be computed.
	// Therefore, it is an error and we're missing the key.
	//
	// TODO: test by creating a state and configuration that is referencing
	// a non-existent variable "foo.bar" where the state only has "foo"
	// and verify plan works, but apply doesn't.
	if i.Operation == walkApply {
		goto MISSING
	}

	// We didn't find the exact field, so lets separate the dots
	// and see if anything along the way is a computed set. i.e. if
	// we have "foo.0.bar" as the field, check to see if "foo" is
	// a computed list. If so, then the whole thing is computed.
	if parts := strings.Split(v.Field, "."); len(parts) > 1 {
		for i := 1; i < len(parts); i++ {
			// Lists and sets make this
			key := fmt.Sprintf("%s.#", strings.Join(parts[:i], "."))
			if attr, ok := r.Primary.Attributes[key]; ok {
				return attr, nil
			}

			// Maps make this
			key = fmt.Sprintf("%s", strings.Join(parts[:i], "."))
			if attr, ok := r.Primary.Attributes[key]; ok {
				return attr, nil
			}
		}
	}

MISSING:
	// Validation for missing interpolations should happen at a higher
	// semantic level. If we reached this point and don't have variables,
	// just return the computed value.
	if scope == nil || scope.Resource == nil {
		return config.UnknownVariableValue, nil
	}

	// If the operation is refresh, it isn't an error for a value to
	// be unknown. Instead, we return that the value is computed so
	// that the graph can continue to refresh other nodes. It doesn't
	// matter because the config isn't interpolated anyways.
	if i.Operation == walkRefresh {
		return config.UnknownVariableValue, nil
	}

	return "", fmt.Errorf(
		"Resource '%s' does not have attribute '%s' "+
			"for variable '%s'",
		id,
		v.Field,
		v.FullKey())
}

func (i *Interpolater) computeResourceMultiVariable(
	scope *InterpolationScope,
	v *config.ResourceVariable) (string, error) {
	i.StateLock.RLock()
	defer i.StateLock.RUnlock()

	// Get the information about this resource variable, and verify
	// that it exists and such.
	module, cr, err := i.resourceVariableInfo(scope, v)
	if err != nil {
		return "", err
	}

	// Get the count so we know how many to iterate over
	count, err := cr.Count()
	if err != nil {
		return "", fmt.Errorf(
			"Error reading %s count: %s",
			v.ResourceId(),
			err)
	}

	// If we have no module in the state yet or count, return empty
	if module == nil || len(module.Resources) == 0 || count == 0 {
		return "", nil
	}

	var values []string
	for i := 0; i < count; i++ {
		id := fmt.Sprintf("%s.%d", v.ResourceId(), i)

		// If we're dealing with only a single resource, then the
		// ID doesn't have a trailing index.
		if count == 1 {
			id = v.ResourceId()
		}

		r, ok := module.Resources[id]
		if !ok {
			continue
		}

		if r.Primary == nil {
			continue
		}

		attr, ok := r.Primary.Attributes[v.Field]
		if !ok {
			continue
		}

		values = append(values, attr)
	}

	if len(values) == 0 {
		return "", fmt.Errorf(
			"Resource '%s' does not have attribute '%s' "+
				"for variable '%s'",
			v.ResourceId(),
			v.Field,
			v.FullKey())
	}

	return strings.Join(values, config.InterpSplitDelim), nil
}

func (i *Interpolater) resourceVariableInfo(
	scope *InterpolationScope,
	v *config.ResourceVariable) (*ModuleState, *config.Resource, error) {
	// Get the module tree that contains our current path. This is
	// either the current module (path is empty) or a child.
	modTree := i.Module
	if len(scope.Path) > 1 {
		modTree = i.Module.Child(scope.Path[1:])
	}

	// Get the resource from the configuration so we can verify
	// that the resource is in the configuration and so we can access
	// the configuration if we need to.
	var cr *config.Resource
	for _, r := range modTree.Config().Resources {
		if r.Id() == v.ResourceId() {
			cr = r
			break
		}
	}
	if cr == nil {
		return nil, nil, fmt.Errorf(
			"Resource '%s' not found for variable '%s'",
			v.ResourceId(),
			v.FullKey())
	}

	// Get the relevant module
	module := i.State.ModuleByPath(scope.Path)
	return module, cr, nil
}
<|endoftext|>