<commit_before>// +build usecase

package nsmd_integration_tests

import (
	"strconv"
	"sync"
	"sync/atomic"
	"testing"

	nsapiv1 "github.com/networkservicemesh/networkservicemesh/k8s/pkg/apis/networkservice/v1alpha1"
	"github.com/networkservicemesh/networkservicemesh/test/kubetest"
	"github.com/networkservicemesh/networkservicemesh/test/kubetest/crds"
	. "github.com/onsi/gomega"
	"github.com/sirupsen/logrus"
	v1 "k8s.io/api/core/v1"
)

type Deployment int

const (
	DeployService Deployment = iota
	DeployEndpoint
	DeployClient
	DeployClientWebhook
)

func TestDeploymentOrder2EndpointClient(t *testing.T) {
	testDeploymentOrder(t, []Deployment{
		DeployEndpoint, DeployEndpoint,
		DeployClient})
}

func TestDeploymentOrder2EndpointClientWebhook(t *testing.T) {
	testDeploymentOrder(t, []Deployment{
		DeployEndpoint, DeployEndpoint,
		DeployClientWebhook})
}

func TestDeploymentOrder2EndpointClientAndWebhook(t *testing.T) {
	testDeploymentOrder(t, []Deployment{
		DeployEndpoint, DeployEndpoint,
		DeployClient,
		DeployClientWebhook})
}

func TestDeploymentOrder4EndpointClientAndWebhook(t *testing.T) {
	testDeploymentOrder(t, []Deployment{
		DeployEndpoint, DeployEndpoint, DeployEndpoint, DeployEndpoint,
		DeployClient,
		DeployClientWebhook})
}

func TestDeploymentOrderClientEndpoint(t *testing.T) {
	testDeploymentOrder(t, []Deployment{
		DeployClient,
		DeployEndpoint})
}

func TestDeploymentOrder2ClientEndpoint(t *testing.T) {
	testDeploymentOrder(t, []Deployment{
		DeployClient, DeployClient,
		DeployEndpoint})
}

func TestDeploymentOrder2Client2Endpoint(t *testing.T) {
	testDeploymentOrder(t, []Deployment{
		DeployClient, DeployClient,
		DeployEndpoint, DeployEndpoint})
}

func TestDeploymentOrder4ClientEndpoint(t *testing.T) {
	testDeploymentOrder(t, []Deployment{
		DeployClient, DeployClient, DeployClient, DeployClient,
		DeployEndpoint})
}

func TestDeploymentOrder4Client2Endpoint(t *testing.T) {
	testDeploymentOrder(t, []Deployment{
		DeployClient, DeployClient, DeployClient, DeployClient,
		DeployEndpoint, DeployEndpoint})
}

func TestDeploymentOrderServiceEndpointClient(t *testing.T) {
	testDeploymentOrder(t, []Deployment{
		DeployService,
		DeployEndpoint,
		DeployClient})
}

func TestDeploymentOrderService2Endpoint4Client(t *testing.T) {
	testDeploymentOrder(t, []Deployment{
		DeployService,
		DeployEndpoint, DeployEndpoint,
		DeployClient, DeployClient, DeployClient, DeployClient})
}

func TestDeploymentOrderServiceClientEndpoint(t *testing.T) {
	testDeploymentOrder(t, []Deployment{
		DeployService,
		DeployClient,
		DeployEndpoint})
}

func TestDeploymentOrderService2ClientEndpoint(t *testing.T) {
	testDeploymentOrder(t, []Deployment{
		DeployService,
		DeployClient, DeployClient,
		DeployEndpoint})
}

func TestDeploymentOrderServiceClientEndpointClient(t *testing.T) {
	testDeploymentOrder(t, []Deployment{
		DeployService,
		DeployClient,
		DeployEndpoint,
		DeployClient})
}

func TestDeploymentOrderServiceClientWebhookEndpoint(t *testing.T) {
	testDeploymentOrder(t, []Deployment{
		DeployService,
		DeployClientWebhook,
		DeployEndpoint})
}

func TestDeploymentOrderService4ClientWebhook2Endpoint(t *testing.T) {
	testDeploymentOrder(t, []Deployment{
		DeployService,
		DeployClientWebhook, DeployClientWebhook, DeployClientWebhook, DeployClientWebhook,
		DeployEndpoint, DeployEndpoint})
}

func testDeploymentOrder(t *testing.T, order []Deployment) {
	RegisterTestingT(t)

	if testing.Short() {
		t.Skip("Skip, please run without -short")
		return
	}

	k8s, err := kubetest.NewK8s(true)
	defer k8s.Cleanup()

	Expect(err).To(BeNil())

	for _, deploy := range order {
		if deploy == DeployClientWebhook {
			awc, awDeployment, awService := kubetest.DeployAdmissionWebhook(k8s, "nsm-admission-webhook", "networkservicemesh/admission-webhook", k8s.GetK8sNamespace())
			defer kubetest.DeleteAdmissionWebhook(k8s, "nsm-admission-webhook-certs", awc, awDeployment, awService, k8s.GetK8sNamespace())
			break
		}
	}

	_, err = kubetest.SetupNodes(k8s, 1, defaultTimeout)
	Expect(err).To(BeNil())

	var nseCount uint64
	var nscPods []*v1.Pod
	var nscCount uint64
	var waitgroup sync.WaitGroup
	var scheduled = make(chan interface{})

	for _, deploy := range order {
		waitgroup.Add(1)

		go func() {
			scheduled <- true
			defer func() { waitgroup.Done() }()
			switch deploy {
			case DeployService:
				nscrd, err := crds.NewNSCRD(k8s.GetK8sNamespace())
				Expect(err).To(BeNil())
				nsIcmpResponder := crds.IcmpResponder(map[string]string{}, map[string]string{"app": "icmp"})
				logrus.Printf("About to insert: %v", nsIcmpResponder)
				var result *nsapiv1.NetworkService
				result, err = nscrd.Create(nsIcmpResponder)
				Expect(err).To(BeNil())
				logrus.Printf("CRD applied with result: %v", result)
				result, err = nscrd.Get(nsIcmpResponder.ObjectMeta.Name)
				Expect(err).To(BeNil())
				logrus.Printf("Registered CRD is: %v", result)
			case DeployEndpoint:
				kubetest.DeployICMP(k8s, nil, "nse-"+strconv.FormatUint(atomic.AddUint64(&nseCount, 1), 10), defaultTimeout)
			case DeployClient:
				nscPods = append(nscPods, kubetest.DeployNSC(k8s, nil, "nsc-"+strconv.FormatUint(atomic.AddUint64(&nscCount, 1), 10), defaultTimeout))
			case DeployClientWebhook:
				nscPods = append(nscPods, kubetest.DeployNSCWebhook(k8s, nil, "nsc-webhook-"+strconv.FormatUint(atomic.AddUint64(&nscCount, 1), 10), defaultTimeout))
			}
		}()
		// ensure the deplyment routine is scheduled
		<-scheduled
	}

	// wait for all deployment routines to end
	waitgroup.Wait()

	for _, p := range nscPods {
		kubetest.CheckNSC(k8s, p)
	}
}
<commit_msg>Disable TestDeploymentOrder* (#1373)<commit_after>// +build usecase

package nsmd_integration_tests

import (
	"strconv"
	"sync"
	"sync/atomic"
	"testing"

	nsapiv1 "github.com/networkservicemesh/networkservicemesh/k8s/pkg/apis/networkservice/v1alpha1"
	"github.com/networkservicemesh/networkservicemesh/test/kubetest"
	"github.com/networkservicemesh/networkservicemesh/test/kubetest/crds"
	. "github.com/onsi/gomega"
	"github.com/sirupsen/logrus"
	"k8s.io/api/core/v1"
)

type Deployment int

const (
	DeployService Deployment = iota
	DeployEndpoint
	DeployClient
	DeployClientWebhook
)

func TestDeploymentOrder2EndpointClient(t *testing.T) {
	testDeploymentOrder(t, []Deployment{
		DeployEndpoint, DeployEndpoint,
		DeployClient})
}

func TestDeploymentOrder2EndpointClientWebhook(t *testing.T) {
	testDeploymentOrder(t, []Deployment{
		DeployEndpoint, DeployEndpoint,
		DeployClientWebhook})
}

func TestDeploymentOrder2EndpointClientAndWebhook(t *testing.T) {
	testDeploymentOrder(t, []Deployment{
		DeployEndpoint, DeployEndpoint,
		DeployClient,
		DeployClientWebhook})
}

func TestDeploymentOrder4EndpointClientAndWebhook(t *testing.T) {
	testDeploymentOrder(t, []Deployment{
		DeployEndpoint, DeployEndpoint, DeployEndpoint, DeployEndpoint,
		DeployClient,
		DeployClientWebhook})
}

func TestDeploymentOrderClientEndpoint(t *testing.T) {
	testDeploymentOrder(t, []Deployment{
		DeployClient,
		DeployEndpoint})
}

func TestDeploymentOrder2ClientEndpoint(t *testing.T) {
	testDeploymentOrder(t, []Deployment{
		DeployClient, DeployClient,
		DeployEndpoint})
}

func TestDeploymentOrder2Client2Endpoint(t *testing.T) {
	testDeploymentOrder(t, []Deployment{
		DeployClient, DeployClient,
		DeployEndpoint, DeployEndpoint})
}

func TestDeploymentOrder4ClientEndpoint(t *testing.T) {
	testDeploymentOrder(t, []Deployment{
		DeployClient, DeployClient, DeployClient, DeployClient,
		DeployEndpoint})
}

func TestDeploymentOrder4Client2Endpoint(t *testing.T) {
	testDeploymentOrder(t, []Deployment{
		DeployClient, DeployClient, DeployClient, DeployClient,
		DeployEndpoint, DeployEndpoint})
}

func TestDeploymentOrderServiceEndpointClient(t *testing.T) {
	testDeploymentOrder(t, []Deployment{
		DeployService,
		DeployEndpoint,
		DeployClient})
}

func TestDeploymentOrderService2Endpoint4Client(t *testing.T) {
	testDeploymentOrder(t, []Deployment{
		DeployService,
		DeployEndpoint, DeployEndpoint,
		DeployClient, DeployClient, DeployClient, DeployClient})
}

func TestDeploymentOrderServiceClientEndpoint(t *testing.T) {
	testDeploymentOrder(t, []Deployment{
		DeployService,
		DeployClient,
		DeployEndpoint})
}

func TestDeploymentOrderService2ClientEndpoint(t *testing.T) {
	testDeploymentOrder(t, []Deployment{
		DeployService,
		DeployClient, DeployClient,
		DeployEndpoint})
}

func TestDeploymentOrderServiceClientEndpointClient(t *testing.T) {
	testDeploymentOrder(t, []Deployment{
		DeployService,
		DeployClient,
		DeployEndpoint,
		DeployClient})
}

func TestDeploymentOrderServiceClientWebhookEndpoint(t *testing.T) {
	testDeploymentOrder(t, []Deployment{
		DeployService,
		DeployClientWebhook,
		DeployEndpoint})
}

func TestDeploymentOrderService4ClientWebhook2Endpoint(t *testing.T) {
	testDeploymentOrder(t, []Deployment{
		DeployService,
		DeployClientWebhook, DeployClientWebhook, DeployClientWebhook, DeployClientWebhook,
		DeployEndpoint, DeployEndpoint})
}

func testDeploymentOrder(t *testing.T, order []Deployment) {
	RegisterTestingT(t)

	if !kubetest.IsBrokeTestsEnabled() {
		t.Skip("Skip, issue - https://github.com/networkservicemesh/networkservicemesh/issues/1372")
		return
	}

	if testing.Short() {
		t.Skip("Skip, please run without -short")
		return
	}

	k8s, err := kubetest.NewK8s(true)
	defer k8s.Cleanup()

	Expect(err).To(BeNil())

	for _, deploy := range order {
		if deploy == DeployClientWebhook {
			awc, awDeployment, awService := kubetest.DeployAdmissionWebhook(k8s, "nsm-admission-webhook", "networkservicemesh/admission-webhook", k8s.GetK8sNamespace())
			defer kubetest.DeleteAdmissionWebhook(k8s, "nsm-admission-webhook-certs", awc, awDeployment, awService, k8s.GetK8sNamespace())
			break
		}
	}

	_, err = kubetest.SetupNodes(k8s, 1, defaultTimeout)
	Expect(err).To(BeNil())

	var nseCount uint64
	var nscPods []*v1.Pod
	var nscCount uint64
	var waitgroup sync.WaitGroup
	var scheduled = make(chan interface{})

	for _, deploy := range order {
		waitgroup.Add(1)

		go func() {
			scheduled <- true
			defer func() { waitgroup.Done() }()
			switch deploy {
			case DeployService:
				nscrd, err := crds.NewNSCRD(k8s.GetK8sNamespace())
				Expect(err).To(BeNil())
				nsIcmpResponder := crds.IcmpResponder(map[string]string{}, map[string]string{"app": "icmp"})
				logrus.Printf("About to insert: %v", nsIcmpResponder)
				var result *nsapiv1.NetworkService
				result, err = nscrd.Create(nsIcmpResponder)
				Expect(err).To(BeNil())
				logrus.Printf("CRD applied with result: %v", result)
				result, err = nscrd.Get(nsIcmpResponder.ObjectMeta.Name)
				Expect(err).To(BeNil())
				logrus.Printf("Registered CRD is: %v", result)
			case DeployEndpoint:
				kubetest.DeployICMP(k8s, nil, "nse-"+strconv.FormatUint(atomic.AddUint64(&nseCount, 1), 10), defaultTimeout)
			case DeployClient:
				nscPods = append(nscPods, kubetest.DeployNSC(k8s, nil, "nsc-"+strconv.FormatUint(atomic.AddUint64(&nscCount, 1), 10), defaultTimeout))
			case DeployClientWebhook:
				nscPods = append(nscPods, kubetest.DeployNSCWebhook(k8s, nil, "nsc-webhook-"+strconv.FormatUint(atomic.AddUint64(&nscCount, 1), 10), defaultTimeout))
			}
		}()
		// ensure the deplyment routine is scheduled
		<-scheduled
	}

	// wait for all deployment routines to end
	waitgroup.Wait()

	for _, p := range nscPods {
		kubetest.CheckNSC(k8s, p)
	}
}
<|endoftext|>


<commit_before>// +build windows

// Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License"). You may
// not use this file except in compliance with the License. A copy of the
// License is located at
//
//	http://aws.amazon.com/apache2.0/
//
// or in the "license" file accompanying this file. This file 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 ecscni

import (
	"github.com/containernetworking/cni/pkg/types"

	"github.com/aws/amazon-ecs-agent/agent/api/eni"
	"github.com/containernetworking/cni/libcni"
	"github.com/pkg/errors"
)

// NewVPCENIPluginConfigForTaskNSSetup is used to create the configuration of vpc-eni plugin for task namespace setup.
func NewVPCENIPluginConfigForTaskNSSetup(eni *eni.ENI, cfg *Config) (*libcni.NetworkConfig, error) {
	dns := types.DNS{
		Nameservers: eni.DomainNameServers,
	}

	if len(eni.DomainNameSearchList) == 0 && cfg.InstanceENIDNSServerList != nil {
		dns.Nameservers = cfg.InstanceENIDNSServerList
	}

	eniConf := VPCENIPluginConfig{
		Type:               ECSVPCENIPluginName,
		DNS:                dns,
		ENIName:            eni.GetLinkName(),
		ENIMACAddress:      eni.MacAddress,
		ENIIPAddresses:     []string{eni.GetPrimaryIPv4AddressWithPrefixLength()},
		GatewayIPAddresses: []string{eni.GetSubnetGatewayIPv4Address()},
		UseExistingNetwork: false,
	}

	networkConfig, err := newNetworkConfig(eniConf, ECSVPCENIPluginExecutable, cfg.MinSupportedCNIVersion)
	if err != nil {
		return nil, errors.Wrap(err, "failed to create vpc-eni plugin configuration for setting up task network namespace")
	}

	networkConfig.Network.Name = TaskHNSNetworkNamePrefix
	return networkConfig, nil
}

// NewVPCENIPluginConfigForECSBridgeSetup creates the configuration required by vpc-eni plugin to setup ecs-bridge endpoint for the task.
func NewVPCENIPluginConfigForECSBridgeSetup(cfg *Config) (*libcni.NetworkConfig, error) {
	bridgeConf := VPCENIPluginConfig{
		Type:               ECSVPCENIPluginName,
		UseExistingNetwork: true,
	}

	networkConfig, err := newNetworkConfig(bridgeConf, ECSVPCENIPluginExecutable, cfg.MinSupportedCNIVersion)
	if err != nil {
		return nil, errors.Wrap(err, "failed to create vpc-eni plugin configuration for setting up ecs-bridge endpoint of the task")
	}

	networkConfig.Network.Name = ECSBridgeNetworkName
	return networkConfig, nil
}
<commit_msg>Removed redundant checks for DNS server information<commit_after>// +build windows

// Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License"). You may
// not use this file except in compliance with the License. A copy of the
// License is located at
//
//	http://aws.amazon.com/apache2.0/
//
// or in the "license" file accompanying this file. This file 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 ecscni

import (
	"github.com/containernetworking/cni/pkg/types"

	"github.com/aws/amazon-ecs-agent/agent/api/eni"
	"github.com/containernetworking/cni/libcni"
	"github.com/pkg/errors"
)

// NewVPCENIPluginConfigForTaskNSSetup is used to create the configuration of vpc-eni plugin for task namespace setup.
func NewVPCENIPluginConfigForTaskNSSetup(eni *eni.ENI, cfg *Config) (*libcni.NetworkConfig, error) {
	// Use the DNS server addresses of the instance ENI it would belong in the same VPC as
	// the task ENI and therefore, have same DNS configuration.
	dns := types.DNS{
		Nameservers: cfg.InstanceENIDNSServerList,
	}

	eniConf := VPCENIPluginConfig{
		Type:               ECSVPCENIPluginName,
		DNS:                dns,
		ENIName:            eni.GetLinkName(),
		ENIMACAddress:      eni.MacAddress,
		ENIIPAddresses:     []string{eni.GetPrimaryIPv4AddressWithPrefixLength()},
		GatewayIPAddresses: []string{eni.GetSubnetGatewayIPv4Address()},
		UseExistingNetwork: false,
	}

	networkConfig, err := newNetworkConfig(eniConf, ECSVPCENIPluginExecutable, cfg.MinSupportedCNIVersion)
	if err != nil {
		return nil, errors.Wrap(err, "failed to create vpc-eni plugin configuration for setting up task network namespace")
	}

	networkConfig.Network.Name = TaskHNSNetworkNamePrefix
	return networkConfig, nil
}

// NewVPCENIPluginConfigForECSBridgeSetup creates the configuration required by vpc-eni plugin to setup ecs-bridge endpoint for the task.
func NewVPCENIPluginConfigForECSBridgeSetup(cfg *Config) (*libcni.NetworkConfig, error) {
	bridgeConf := VPCENIPluginConfig{
		Type:               ECSVPCENIPluginName,
		UseExistingNetwork: true,
	}

	networkConfig, err := newNetworkConfig(bridgeConf, ECSVPCENIPluginExecutable, cfg.MinSupportedCNIVersion)
	if err != nil {
		return nil, errors.Wrap(err, "failed to create vpc-eni plugin configuration for setting up ecs-bridge endpoint of the task")
	}

	networkConfig.Network.Name = ECSBridgeNetworkName
	return networkConfig, nil
}
<|endoftext|>


<commit_before>// +build integration
// Copyright 2014-2017 Amazon.com, Inc. or its affiliates. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License"). You may
// not use this file except in compliance with the License. A copy of the
// License is located at
//
//	http://aws.amazon.com/apache2.0/
//
// or in the "license" file accompanying this file. This file 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 engine

import (
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"
	"testing"
	"time"

	"github.com/aws/amazon-ecs-agent/agent/api"
	"github.com/aws/amazon-ecs-agent/agent/config"
	"github.com/aws/amazon-ecs-agent/agent/credentials"
	"github.com/aws/amazon-ecs-agent/agent/ec2"
	"github.com/aws/amazon-ecs-agent/agent/engine/dockerclient"
	"github.com/aws/amazon-ecs-agent/agent/engine/dockerstate"
	"github.com/aws/amazon-ecs-agent/agent/eventstream"
	"github.com/aws/amazon-ecs-agent/agent/statemanager"
	docker "github.com/fsouza/go-dockerclient"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"golang.org/x/net/context"
)

const (
	testDockerStopTimeout  = 2 * time.Second
	credentialsIDIntegTest = "credsid"
)

func init() {
	// Set this very low for integ tests only
	_stoppedSentWaitInterval = 1 * time.Second
}

func createTestTask(arn string) *api.Task {
	return &api.Task{
		Arn:                 arn,
		Family:              arn,
		Version:             "1",
		DesiredStatusUnsafe: api.TaskRunning,
		Containers:          []*api.Container{createTestContainer()},
	}
}

func defaultTestConfigIntegTest() *config.Config {
	cfg, _ := config.NewConfig(ec2.NewBlackholeEC2MetadataClient())
	return cfg
}

func setupWithDefaultConfig(t *testing.T) (TaskEngine, func(), credentials.Manager) {
	return setup(defaultTestConfigIntegTest(), t)
}

func setup(cfg *config.Config, t *testing.T) (TaskEngine, func(), credentials.Manager) {
	if os.Getenv("ECS_SKIP_ENGINE_INTEG_TEST") != "" {
		t.Skip("ECS_SKIP_ENGINE_INTEG_TEST")
	}
	if !isDockerRunning() {
		t.Skip("Docker not running")
	}
	clientFactory := dockerclient.NewFactory(dockerEndpoint)
	dockerClient, err := NewDockerGoClient(clientFactory, cfg)
	if err != nil {
		t.Fatalf("Error creating Docker client: %v", err)
	}
	credentialsManager := credentials.NewManager()
	state := dockerstate.NewTaskEngineState()
	imageManager := NewImageManager(cfg, dockerClient, state)
	imageManager.SetSaver(statemanager.NewNoopStateManager())
	taskEngine := NewDockerTaskEngine(cfg, dockerClient, credentialsManager,
		eventstream.NewEventStream("ENGINEINTEGTEST", context.Background()), imageManager, state)
	taskEngine.Init(context.TODO())
	return taskEngine, func() {
		taskEngine.Shutdown()
	}, credentialsManager
}

func discardEvents(from interface{}) func() {
	done := make(chan bool)

	go func() {
		for {
			ndx, _, _ := reflect.Select([]reflect.SelectCase{
				{
					Dir:  reflect.SelectRecv,
					Chan: reflect.ValueOf(from),
				},
				{
					Dir:  reflect.SelectRecv,
					Chan: reflect.ValueOf(done),
				},
			})
			if ndx == 1 {
				break
			}
		}
	}()
	return func() {
		done <- true
	}
}

// TestDockerStateToContainerState tests convert the container status from
// docker inspect to the status defined in agent
func TestDockerStateToContainerState(t *testing.T) {
	taskEngine, done, _ := setupWithDefaultConfig(t)
	defer done()

	testTask := createTestTask("test_task")
	container := testTask.Containers[0]

	client, err := docker.NewClientFromEnv()
	require.NoError(t, err, "Creating go docker client failed")

	containerMetadata := taskEngine.(*DockerTaskEngine).pullContainer(testTask, container)
	assert.NoError(t, containerMetadata.Error)

	containerMetadata = taskEngine.(*DockerTaskEngine).createContainer(testTask, container)
	assert.NoError(t, containerMetadata.Error)
	state, _ := client.InspectContainer(containerMetadata.DockerID)
	assert.Equal(t, api.ContainerCreated, dockerStateToState(state.State))

	containerMetadata = taskEngine.(*DockerTaskEngine).startContainer(testTask, container)
	assert.NoError(t, containerMetadata.Error)
	state, _ = client.InspectContainer(containerMetadata.DockerID)
	assert.Equal(t, api.ContainerRunning, dockerStateToState(state.State))

	containerMetadata = taskEngine.(*DockerTaskEngine).stopContainer(testTask, container)
	assert.NoError(t, containerMetadata.Error)
	state, _ = client.InspectContainer(containerMetadata.DockerID)
	assert.Equal(t, api.ContainerStopped, dockerStateToState(state.State))

	// clean up the container
	err = taskEngine.(*DockerTaskEngine).removeContainer(testTask, container)
	assert.NoError(t, err, "remove the created container failed")

	// Start the container failed
	testTask = createTestTask("test_task2")
	testTask.Containers[0].EntryPoint = &[]string{"non-existed"}
	container = testTask.Containers[0]
	containerMetadata = taskEngine.(*DockerTaskEngine).createContainer(testTask, container)
	assert.NoError(t, containerMetadata.Error)
	containerMetadata = taskEngine.(*DockerTaskEngine).startContainer(testTask, container)
	assert.Error(t, containerMetadata.Error)
	state, _ = client.InspectContainer(containerMetadata.DockerID)
	assert.Equal(t, api.ContainerStopped, dockerStateToState(state.State))

	// clean up the container
	err = taskEngine.(*DockerTaskEngine).removeContainer(testTask, container)
	assert.NoError(t, err, "remove the created container failed")
}

func TestHostVolumeMount(t *testing.T) {
	taskEngine, done, _ := setupWithDefaultConfig(t)
	defer done()

	stateChangeEvents := taskEngine.StateChangeEvents()

	tmpPath, _ := ioutil.TempDir("", "ecs_volume_test")
	defer os.RemoveAll(tmpPath)
	ioutil.WriteFile(filepath.Join(tmpPath, "test-file"), []byte("test-data"), 0644)

	testTask := createTestHostVolumeMountTask(tmpPath)

	go taskEngine.AddTask(testTask)

	event := <-stateChangeEvents
	assert.Equal(t, event.(api.ContainerStateChange).Status, api.ContainerRunning, "Expected container to be RUNNING")

	event = <-stateChangeEvents
	assert.Equal(t, event.(api.TaskStateChange).Status, api.TaskRunning, "Expected task to be RUNNING")

	event = <-stateChangeEvents
	assert.Equal(t, event.(api.ContainerStateChange).Status, api.ContainerStopped, "Expected container to be STOPPED")

	event = <-stateChangeEvents
	assert.Equal(t, event.(api.TaskStateChange).Status, api.TaskStopped, "Expected task to be STOPPED")

	assert.NotNil(t, testTask.Containers[0].GetKnownExitCode(), "No exit code found")
	assert.Equal(t, 42, *testTask.Containers[0].GetKnownExitCode(), "Wrong exit code")

	data, err := ioutil.ReadFile(filepath.Join(tmpPath, "hello-from-container"))
	assert.Nil(t, err, "Unexpected error")
	assert.Equal(t, "hi", strings.TrimSpace(string(data)), "Incorrect file contents")
}

func TestEmptyHostVolumeMount(t *testing.T) {
	taskEngine, done, _ := setupWithDefaultConfig(t)
	defer done()

	stateChangeEvents := taskEngine.StateChangeEvents()

	// creates a task with two containers
	testTask := createTestEmptyHostVolumeMountTask()

	go taskEngine.AddTask(testTask)

	event := <-stateChangeEvents
	assert.Equal(t, event.(api.ContainerStateChange).Status, api.ContainerRunning, "Expected container to be RUNNING")

	event = <-stateChangeEvents
	assert.Equal(t, event.(api.ContainerStateChange).Status, api.ContainerRunning, "Expected container to be RUNNING")

	event = <-stateChangeEvents
	assert.Equal(t, event.(api.TaskStateChange).Status, api.TaskRunning, "Expected task to be RUNNING")

	event = <-stateChangeEvents
	assert.Equal(t, event.(api.ContainerStateChange).Status, api.ContainerStopped, "Expected container to be STOPPED")

	event = <-stateChangeEvents
	assert.Equal(t, event.(api.ContainerStateChange).Status, api.ContainerStopped, "Expected container to be STOPPED")

	event = <-stateChangeEvents
	assert.Equal(t, event.(api.TaskStateChange).Status, api.TaskStopped, "Expected task to be STOPPED")

	assert.NotNil(t, testTask.Containers[0].GetKnownExitCode(), "No exit code found")
	assert.Equal(t, 42, *testTask.Containers[0].GetKnownExitCode(), "Wrong exit code, file probably wasn't present")
}

func TestSweepContainer(t *testing.T) {
	cfg := defaultTestConfigIntegTest()
	cfg.TaskCleanupWaitDuration = 1 * time.Minute
	taskEngine, done, _ := setup(cfg, t)
	defer done()

	stateChangeEvents := taskEngine.StateChangeEvents()

	testTask := createTestTask("testSweepContainer")

	go taskEngine.AddTask(testTask)

	event := <-stateChangeEvents
	assert.Equal(t, event.(api.ContainerStateChange).Status, api.ContainerRunning, "Expected container to be RUNNING")

	event = <-stateChangeEvents
	assert.Equal(t, event.(api.TaskStateChange).Status, api.TaskRunning, "Expected task to be RUNNING")

	event = <-stateChangeEvents
	assert.Equal(t, event.(api.ContainerStateChange).Status, api.ContainerStopped, "Expected container to be STOPPED")

	event = <-stateChangeEvents
	assert.Equal(t, event.(api.TaskStateChange).Status, api.TaskStopped, "Expected task to be STOPPED")

	// Should be stopped, let's verify it's still listed...
	task, ok := taskEngine.(*DockerTaskEngine).State().TaskByArn("testSweepContainer")
	assert.True(t, ok, "Expected task to be present still, but wasn't")
	task.SetSentStatus(api.TaskStopped) // cleanupTask waits for TaskStopped to be sent before cleaning
	time.Sleep(1 * time.Minute)
	for i := 0; i < 60; i++ {
		_, ok = taskEngine.(*DockerTaskEngine).State().TaskByArn("testSweepContainer")
		if !ok {
			break
		}
		time.Sleep(1 * time.Second)
	}
	assert.False(t, ok, "Expected container to have been swept but was not")
}

// TestStartStopWithCredentials starts and stops a task for which credentials id
// has been set
func TestStartStopWithCredentials(t *testing.T) {
	taskEngine, done, credentialsManager := setupWithDefaultConfig(t)
	defer done()

	testTask := createTestTask("testStartWithCredentials")
	taskCredentials := credentials.TaskIAMRoleCredentials{
		IAMRoleCredentials: credentials.IAMRoleCredentials{CredentialsID: credentialsIDIntegTest},
	}
	credentialsManager.SetTaskCredentials(taskCredentials)
	testTask.SetCredentialsID(credentialsIDIntegTest)

	stateChangeEvents := taskEngine.StateChangeEvents()

	go taskEngine.AddTask(testTask)

	event := <-stateChangeEvents
	assert.Equal(t, event.(api.ContainerStateChange).Status, api.ContainerRunning, "Expected container to be RUNNING")

	event = <-stateChangeEvents
	assert.Equal(t, event.(api.TaskStateChange).Status, api.TaskRunning, "Expected task to be RUNNING")

	event = <-stateChangeEvents
	assert.Equal(t, event.(api.ContainerStateChange).Status, api.ContainerStopped, "Expected container to be STOPPED")

	event = <-stateChangeEvents
	assert.Equal(t, event.(api.TaskStateChange).Status, api.TaskStopped, "Expected task to be STOPPED")

	// When task is stopped, credentials should have been removed for the
	// credentials id set in the task
	_, ok := credentialsManager.GetTaskCredentials(credentialsIDIntegTest)
	assert.False(t, ok, "Credentials not removed from credentials manager for stopped task")
}
<commit_msg>test: fix integration test failure<commit_after>// +build integration
// Copyright 2014-2017 Amazon.com, Inc. or its affiliates. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License"). You may
// not use this file except in compliance with the License. A copy of the
// License is located at
//
//	http://aws.amazon.com/apache2.0/
//
// or in the "license" file accompanying this file. This file 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 engine

import (
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"
	"testing"
	"time"

	"github.com/aws/amazon-ecs-agent/agent/api"
	"github.com/aws/amazon-ecs-agent/agent/config"
	"github.com/aws/amazon-ecs-agent/agent/credentials"
	"github.com/aws/amazon-ecs-agent/agent/ec2"
	"github.com/aws/amazon-ecs-agent/agent/engine/dockerclient"
	"github.com/aws/amazon-ecs-agent/agent/engine/dockerstate"
	"github.com/aws/amazon-ecs-agent/agent/eventstream"
	"github.com/aws/amazon-ecs-agent/agent/statemanager"
	docker "github.com/fsouza/go-dockerclient"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"golang.org/x/net/context"
)

const (
	testDockerStopTimeout  = 2 * time.Second
	credentialsIDIntegTest = "credsid"
)

func init() {
	// Set this very low for integ tests only
	_stoppedSentWaitInterval = 1 * time.Second
}

func createTestTask(arn string) *api.Task {
	return &api.Task{
		Arn:                 arn,
		Family:              arn,
		Version:             "1",
		DesiredStatusUnsafe: api.TaskRunning,
		Containers:          []*api.Container{createTestContainer()},
	}
}

func defaultTestConfigIntegTest() *config.Config {
	cfg, _ := config.NewConfig(ec2.NewBlackholeEC2MetadataClient())
	return cfg
}

func setupWithDefaultConfig(t *testing.T) (TaskEngine, func(), credentials.Manager) {
	return setup(defaultTestConfigIntegTest(), t)
}

func setup(cfg *config.Config, t *testing.T) (TaskEngine, func(), credentials.Manager) {
	if os.Getenv("ECS_SKIP_ENGINE_INTEG_TEST") != "" {
		t.Skip("ECS_SKIP_ENGINE_INTEG_TEST")
	}
	if !isDockerRunning() {
		t.Skip("Docker not running")
	}
	clientFactory := dockerclient.NewFactory(dockerEndpoint)
	dockerClient, err := NewDockerGoClient(clientFactory, cfg)
	if err != nil {
		t.Fatalf("Error creating Docker client: %v", err)
	}
	credentialsManager := credentials.NewManager()
	state := dockerstate.NewTaskEngineState()
	imageManager := NewImageManager(cfg, dockerClient, state)
	imageManager.SetSaver(statemanager.NewNoopStateManager())
	taskEngine := NewDockerTaskEngine(cfg, dockerClient, credentialsManager,
		eventstream.NewEventStream("ENGINEINTEGTEST", context.Background()), imageManager, state)
	taskEngine.Init(context.TODO())
	return taskEngine, func() {
		taskEngine.Shutdown()
	}, credentialsManager
}

/*
func discardEvents(from interface{}) func() {
	done := make(chan bool)

	go func() {
		for {
			ndx, _, _ := reflect.Select([]reflect.SelectCase{
				{
					Dir:  reflect.SelectRecv,
					Chan: reflect.ValueOf(from),
				},
				{
					Dir:  reflect.SelectRecv,
					Chan: reflect.ValueOf(done),
				},
			})
			if ndx == 1 {
				break
			}
		}
	}()
	return func() {
		done <- true
	}
}
*/
// TestDockerStateToContainerState tests convert the container status from
// docker inspect to the status defined in agent
func TestDockerStateToContainerState(t *testing.T) {
	taskEngine, done, _ := setupWithDefaultConfig(t)
	defer done()

	testTask := createTestTask("test_task")
	container := testTask.Containers[0]

	client, err := docker.NewClientFromEnv()
	require.NoError(t, err, "Creating go docker client failed")

	containerMetadata := taskEngine.(*DockerTaskEngine).pullContainer(testTask, container)
	assert.NoError(t, containerMetadata.Error)

	containerMetadata = taskEngine.(*DockerTaskEngine).createContainer(testTask, container)
	assert.NoError(t, containerMetadata.Error)
	state, _ := client.InspectContainer(containerMetadata.DockerID)
	assert.Equal(t, api.ContainerCreated, dockerStateToState(state.State))

	containerMetadata = taskEngine.(*DockerTaskEngine).startContainer(testTask, container)
	assert.NoError(t, containerMetadata.Error)
	state, _ = client.InspectContainer(containerMetadata.DockerID)
	assert.Equal(t, api.ContainerRunning, dockerStateToState(state.State))

	containerMetadata = taskEngine.(*DockerTaskEngine).stopContainer(testTask, container)
	assert.NoError(t, containerMetadata.Error)
	state, _ = client.InspectContainer(containerMetadata.DockerID)
	assert.Equal(t, api.ContainerStopped, dockerStateToState(state.State))

	// clean up the container
	err = taskEngine.(*DockerTaskEngine).removeContainer(testTask, container)
	assert.NoError(t, err, "remove the created container failed")

	// Start the container failed
	testTask = createTestTask("test_task2")
	testTask.Containers[0].EntryPoint = &[]string{"non-existed"}
	container = testTask.Containers[0]
	containerMetadata = taskEngine.(*DockerTaskEngine).createContainer(testTask, container)
	assert.NoError(t, containerMetadata.Error)
	containerMetadata = taskEngine.(*DockerTaskEngine).startContainer(testTask, container)
	assert.Error(t, containerMetadata.Error)
	state, _ = client.InspectContainer(containerMetadata.DockerID)
	assert.Equal(t, api.ContainerStopped, dockerStateToState(state.State))

	// clean up the container
	err = taskEngine.(*DockerTaskEngine).removeContainer(testTask, container)
	assert.NoError(t, err, "remove the created container failed")
}

func TestHostVolumeMount(t *testing.T) {
	taskEngine, done, _ := setupWithDefaultConfig(t)
	defer done()

	stateChangeEvents := taskEngine.StateChangeEvents()

	tmpPath, _ := ioutil.TempDir("", "ecs_volume_test")
	defer os.RemoveAll(tmpPath)
	ioutil.WriteFile(filepath.Join(tmpPath, "test-file"), []byte("test-data"), 0644)

	testTask := createTestHostVolumeMountTask(tmpPath)

	go taskEngine.AddTask(testTask)

	event := <-stateChangeEvents
	assert.Equal(t, event.(api.ContainerStateChange).Status, api.ContainerRunning, "Expected container to be RUNNING")

	event = <-stateChangeEvents
	assert.Equal(t, event.(api.TaskStateChange).Status, api.TaskRunning, "Expected task to be RUNNING")

	event = <-stateChangeEvents
	assert.Equal(t, event.(api.ContainerStateChange).Status, api.ContainerStopped, "Expected container to be STOPPED")

	event = <-stateChangeEvents
	assert.Equal(t, event.(api.TaskStateChange).Status, api.TaskStopped, "Expected task to be STOPPED")

	assert.NotNil(t, testTask.Containers[0].GetKnownExitCode(), "No exit code found")
	assert.Equal(t, 42, *testTask.Containers[0].GetKnownExitCode(), "Wrong exit code")

	data, err := ioutil.ReadFile(filepath.Join(tmpPath, "hello-from-container"))
	assert.Nil(t, err, "Unexpected error")
	assert.Equal(t, "hi", strings.TrimSpace(string(data)), "Incorrect file contents")
}

func TestEmptyHostVolumeMount(t *testing.T) {
	taskEngine, done, _ := setupWithDefaultConfig(t)
	defer done()

	stateChangeEvents := taskEngine.StateChangeEvents()

	// creates a task with two containers
	testTask := createTestEmptyHostVolumeMountTask()

	go taskEngine.AddTask(testTask)

	event := <-stateChangeEvents
	assert.Equal(t, event.(api.ContainerStateChange).Status, api.ContainerRunning, "Expected container to be RUNNING")

	event = <-stateChangeEvents
	assert.Equal(t, event.(api.ContainerStateChange).Status, api.ContainerRunning, "Expected container to be RUNNING")

	event = <-stateChangeEvents
	assert.Equal(t, event.(api.TaskStateChange).Status, api.TaskRunning, "Expected task to be RUNNING")

	event = <-stateChangeEvents
	assert.Equal(t, event.(api.ContainerStateChange).Status, api.ContainerStopped, "Expected container to be STOPPED")

	event = <-stateChangeEvents
	assert.Equal(t, event.(api.ContainerStateChange).Status, api.ContainerStopped, "Expected container to be STOPPED")

	event = <-stateChangeEvents
	assert.Equal(t, event.(api.TaskStateChange).Status, api.TaskStopped, "Expected task to be STOPPED")

	assert.NotNil(t, testTask.Containers[0].GetKnownExitCode(), "No exit code found")
	assert.Equal(t, 42, *testTask.Containers[0].GetKnownExitCode(), "Wrong exit code, file probably wasn't present")
}

func TestSweepContainer(t *testing.T) {
	cfg := defaultTestConfigIntegTest()
	cfg.TaskCleanupWaitDuration = 1 * time.Minute
	taskEngine, done, _ := setup(cfg, t)
	defer done()

	stateChangeEvents := taskEngine.StateChangeEvents()

	testTask := createTestTask("testSweepContainer")

	go taskEngine.AddTask(testTask)

	event := <-stateChangeEvents
	assert.Equal(t, event.(api.ContainerStateChange).Status, api.ContainerRunning, "Expected container to be RUNNING")

	event = <-stateChangeEvents
	assert.Equal(t, event.(api.TaskStateChange).Status, api.TaskRunning, "Expected task to be RUNNING")

	event = <-stateChangeEvents
	assert.Equal(t, event.(api.ContainerStateChange).Status, api.ContainerStopped, "Expected container to be STOPPED")

	event = <-stateChangeEvents
	assert.Equal(t, event.(api.TaskStateChange).Status, api.TaskStopped, "Expected task to be STOPPED")

	// Should be stopped, let's verify it's still listed...
	task, ok := taskEngine.(*DockerTaskEngine).State().TaskByArn("testSweepContainer")
	assert.True(t, ok, "Expected task to be present still, but wasn't")
	task.SetSentStatus(api.TaskStopped) // cleanupTask waits for TaskStopped to be sent before cleaning
	time.Sleep(1 * time.Minute)
	for i := 0; i < 60; i++ {
		_, ok = taskEngine.(*DockerTaskEngine).State().TaskByArn("testSweepContainer")
		if !ok {
			break
		}
		time.Sleep(1 * time.Second)
	}
	assert.False(t, ok, "Expected container to have been swept but was not")
}

// TestStartStopWithCredentials starts and stops a task for which credentials id
// has been set
func TestStartStopWithCredentials(t *testing.T) {
	taskEngine, done, credentialsManager := setupWithDefaultConfig(t)
	defer done()

	testTask := createTestTask("testStartWithCredentials")
	taskCredentials := credentials.TaskIAMRoleCredentials{
		IAMRoleCredentials: credentials.IAMRoleCredentials{CredentialsID: credentialsIDIntegTest},
	}
	credentialsManager.SetTaskCredentials(taskCredentials)
	testTask.SetCredentialsID(credentialsIDIntegTest)

	stateChangeEvents := taskEngine.StateChangeEvents()

	go taskEngine.AddTask(testTask)

	event := <-stateChangeEvents
	assert.Equal(t, event.(api.ContainerStateChange).Status, api.ContainerRunning, "Expected container to be RUNNING")

	event = <-stateChangeEvents
	assert.Equal(t, event.(api.TaskStateChange).Status, api.TaskRunning, "Expected task to be RUNNING")

	event = <-stateChangeEvents
	assert.Equal(t, event.(api.ContainerStateChange).Status, api.ContainerStopped, "Expected container to be STOPPED")

	event = <-stateChangeEvents
	assert.Equal(t, event.(api.TaskStateChange).Status, api.TaskStopped, "Expected task to be STOPPED")

	// When task is stopped, credentials should have been removed for the
	// credentials id set in the task
	_, ok := credentialsManager.GetTaskCredentials(credentialsIDIntegTest)
	assert.False(t, ok, "Credentials not removed from credentials manager for stopped task")
}
<|endoftext|>


<commit_before>package bongo

import (
	"encoding/json"
	"errors"
	"fmt"
	"strconv"
	"strings"

	"github.com/jinzhu/gorm"
)

// Fetch fetches the data from db by given parameters(fields of the struct)
func (b *Bongo) Fetch(i Modellable) error {
	if i.GetId() == 0 {
		return errors.New(fmt.Sprintf("Id is not set for %s", i.TableName()))
	}

	if err := b.DB.Table(i.TableName()).
		Find(i).
		Error; err != nil {
		return err
	}

	return nil
}

// ById Fetches data from db by it's id
func (b *Bongo) ById(i Modellable, id int64) error {
	if err := b.DB.
		Table(i.TableName()).
		Where("id = ?", id).
		Find(i).
		Error; err != nil {
		return err
	}

	return nil
}

// Creates a new record with the given struct and its fields
func (b *Bongo) Create(i Modellable) error {
	if err := b.DB.Save(i).Error; err != nil {
		return err
	}

	return nil
}

// Update updates all fields of a struct with assigned data
func (b *Bongo) Update(i Modellable) error {
	if i.GetId() == 0 {
		return errors.New(fmt.Sprintf("Id is not set for %s", i.TableName()))
	}

	// Update and Create is using the Save method, so they are
	// same functions but GORM handles, AfterCreate and AfterUpdate
	// in correct manner
	return b.Create(i)
}

// Delete deletes the data by it's id, it doesnt take any other fields
// into consideration
func (b *Bongo) Delete(i Modellable) error {
	if i.GetId() == 0 {
		return errors.New(fmt.Sprintf("Id is not set for %s", i.TableName()))
	}

	if err := b.DB.Delete(i).Error; err != nil {
		return err
	}

	return nil
}

// FetchByIds fetches data from db by their ids in ordered fashion,
// if non-found this function doesnt return any error.
func (b *Bongo) FetchByIds(i Modellable, data interface{}, ids []int64) error {
	if len(ids) == 0 {
		return nil
	}

	orderByQuery := ""
	comma := ""
	for _, id := range ids {
		orderByQuery = orderByQuery + comma + " id = " + strconv.FormatInt(id, 10) + " desc"
		comma = ","
	}

	// init query
	query := b.DB.Model(i)

	// add table name
	query = query.Table(i.TableName())

	query = query.Order(orderByQuery)

	query = query.Where(ids)

	query = query.Find(data)

	// supress not found errors
	return CheckErr(query)

}

func (b *Bongo) UpdatePartial(i Modellable, set map[string]interface{}) error {
	if i.GetId() == 0 {
		return errors.New(fmt.Sprintf("Id is not set for %s", i.TableName()))
	}

	// init query
	query := b.DB.Model(i)

	query = query.Table(i.TableName())

	query = query.Where(i.GetId())

	if err := query.Update(set).Error; err != nil {
		return err
	}

	if err := b.Fetch(i); err != nil {
		return err
	}

	b.AfterUpdate(i)
	return nil
}

// selector, set
func (b *Bongo) UpdateMulti(i Modellable, rest ...map[string]interface{}) error {
	var set, selector map[string]interface{}

	switch len(rest) {
	case 1:
		set = rest[0]
		selector = nil
	case 2:
		selector = rest[0]
		set = rest[1]
	default:
		return errors.New("Update partial parameter list is wrong")
	}

	query := b.DB.Table(i.TableName())

	//add selector
	query = addWhere(query, selector)

	if err := query.Updates(set).Error; err != nil {
		return err
	}

	return nil
}

func (b *Bongo) Count(i Modellable, where ...interface{}) (int, error) {
	var count int

	// init query
	query := b.DB.Model(i)

	// add table name
	query = query.Table(i.TableName())

	// add query
	query = query.Where(where[0], where[1:len(where)]...)

	return count, query.Count(&count).Error
}

func (b *Bongo) CountWithQuery(i Modellable, q *Query) (int, error) {
	query := b.BuildQuery(i, q)
	var count int
	return count, query.Count(&count).Error
}

type Scope func(d *gorm.DB) *gorm.DB

type Query struct {
	Selector   map[string]interface{}
	Sort       map[string]string
	Pluck      string
	Pagination Pagination
	Scopes     []Scope
}

func (q *Query) AddScope(scope Scope) {
	if q.Scopes == nil {
		q.Scopes = make([]Scope, 0)
	}

	q.Scopes = append(q.Scopes, scope)
}

type Pagination struct {
	Limit int
	Skip  int
}

func NewPagination(limit int, skip int) *Pagination {
	return &Pagination{
		Limit: limit,
		Skip:  skip,
	}
}

func NewQS(selector map[string]interface{}) *Query {
	return &Query{
		Selector: selector,
	}
}

// selector, sort, limit, pluck,
func (b *Bongo) Some(i Modellable, data interface{}, q *Query) error {
	err := b.executeQuery(i, data, q)
	if err == gorm.RecordNotFound {
		return nil
	}
	return err
}

func (b *Bongo) One(i Modellable, data interface{}, q *Query) error {
	q.Pagination.Limit = 1
	return b.executeQuery(i, data, q)
}

func (b *Bongo) BuildQuery(i Modellable, q *Query) *gorm.DB {
	// init query
	query := b.DB.Model(i)

	// add table name
	query = query.Table(i.TableName())

	// add sort options
	query = addSort(query, q.Sort)

	query = addSkip(query, q.Pagination.Skip)

	query = addLimit(query, q.Pagination.Limit)

	// add selector
	query = addWhere(query, q.Selector)

	// put scopes
	if q.Scopes != nil && len(q.Scopes) > 0 {
		for _, scope := range q.Scopes {
			query = query.Scopes(scope)
		}
	}

	return query
}

func (b *Bongo) executeQuery(i Modellable, data interface{}, q *Query) error {
	// init query
	query := b.BuildQuery(i, q)

	var err error
	// TODO refactor this part
	if q.Pluck != "" {
		if strings.Contains(q.Pluck, ",") {
			// add pluck data
			query = addPluck(query, q.Pluck)

			err = query.Find(data).Error
		} else {
			err = query.Pluck(q.Pluck, data).Error
		}
	} else {
		err = query.Find(data).Error
	}

	return err
}

func (b *Bongo) PublishEvent(eventName string, i Modellable) error {
	data, err := json.Marshal(i)
	if err != nil {
		b.log.Error("Error while marshalling for publish %s", err)
		return err
	}

	err = b.Broker.Publish(i.TableName()+"_"+eventName, data)
	if err != nil {
		b.log.Error("Error while publishing %s", err)
		return err
	}

	return nil
}

func (b *Bongo) AfterCreate(i Modellable) {
	b.PublishEvent("created", i)
}

func (b *Bongo) AfterUpdate(i Modellable) {
	b.PublishEvent("updated", i)
}

func (b *Bongo) AfterDelete(i Modellable) {
	b.PublishEvent("deleted", i)
}

// addSort injects sort parameters into query
func addSort(query *gorm.DB, options map[string]string) *gorm.DB {

	if options == nil {
		return query
	}

	if len(options) == 0 {
		return query
	}

	var opts []string
	for key, val := range options {
		opts = append(opts, fmt.Sprintf("%s %v", key, val))
	}
	return query.Order(strings.Join(opts, ","))
}

// addPluck basically adds select statement for
// only required fields
func addPluck(query *gorm.DB, plucked string) *gorm.DB {
	if plucked == "" {
		return query
	}

	return query.Select(plucked)
}

// addWhere adds where query
func addWhere(query *gorm.DB, selector map[string]interface{}) *gorm.DB {
	if selector == nil {
		return query
	}

	// instead sending one selector, do chaining here
	return query.Where(selector)
}

// addSkip adds skip parameter into sql query
func addSkip(query *gorm.DB, skip int) *gorm.DB {
	if skip > 0 {
		return query.Offset(skip)
	}

	return query
}

// addLimit adds limit into query if set
func addLimit(query *gorm.DB, limit int) *gorm.DB {
	// if limit is minus or 0 ignore
	if limit > 0 {
		return query.Limit(limit)
	}

	return query
}

// CheckErr checks error exitence and returns
// if found, but this function suppress RecordNotFound errors
func CheckErr(res *gorm.DB) error {
	if res == nil {
		return nil
	}

	if res.Error == nil {
		return nil
	}

	if res.Error == gorm.RecordNotFound {
		return nil
	}

	return res.Error
}
<commit_msg>Social: model is problematic, use it with care, for this case, it wasn't updating the record itself<commit_after>package bongo

import (
	"encoding/json"
	"errors"
	"fmt"
	"strconv"
	"strings"

	"github.com/jinzhu/gorm"
)

// Fetch fetches the data from db by given parameters(fields of the struct)
func (b *Bongo) Fetch(i Modellable) error {
	if i.GetId() == 0 {
		return errors.New(fmt.Sprintf("Id is not set for %s", i.TableName()))
	}

	if err := b.DB.Table(i.TableName()).
		Find(i).
		Error; err != nil {
		return err
	}

	return nil
}

// ById Fetches data from db by it's id
func (b *Bongo) ById(i Modellable, id int64) error {
	if err := b.DB.
		Table(i.TableName()).
		Where("id = ?", id).
		Find(i).
		Error; err != nil {
		return err
	}

	return nil
}

// Creates a new record with the given struct and its fields
func (b *Bongo) Create(i Modellable) error {
	if err := b.DB.Save(i).Error; err != nil {
		return err
	}

	return nil
}

// Update updates all fields of a struct with assigned data
func (b *Bongo) Update(i Modellable) error {
	if i.GetId() == 0 {
		return errors.New(fmt.Sprintf("Id is not set for %s", i.TableName()))
	}

	// Update and Create is using the Save method, so they are
	// same functions but GORM handles, AfterCreate and AfterUpdate
	// in correct manner
	return b.Create(i)
}

// Delete deletes the data by it's id, it doesnt take any other fields
// into consideration
func (b *Bongo) Delete(i Modellable) error {
	if i.GetId() == 0 {
		return errors.New(fmt.Sprintf("Id is not set for %s", i.TableName()))
	}

	if err := b.DB.Delete(i).Error; err != nil {
		return err
	}

	return nil
}

// FetchByIds fetches data from db by their ids in ordered fashion,
// if non-found this function doesnt return any error.
func (b *Bongo) FetchByIds(i Modellable, data interface{}, ids []int64) error {
	if len(ids) == 0 {
		return nil
	}

	orderByQuery := ""
	comma := ""
	for _, id := range ids {
		orderByQuery = orderByQuery + comma + " id = " + strconv.FormatInt(id, 10) + " desc"
		comma = ","
	}

	// init query
	query := b.DB.Model(i)

	// add table name
	query = query.Table(i.TableName())

	query = query.Order(orderByQuery)

	query = query.Where(ids)

	query = query.Find(data)

	// supress not found errors
	return CheckErr(query)

}

func (b *Bongo) UpdatePartial(i Modellable, set map[string]interface{}) error {
	if i.GetId() == 0 {
		return errors.New(fmt.Sprintf("Id is not set for %s", i.TableName()))
	}

	// init query
	query := b.DB

	query = query.Table(i.TableName())

	query = query.Where(i.GetId())

	if err := query.Update(set).Error; err != nil {
		return err
	}

	if err := b.Fetch(i); err != nil {
		return err
	}

	b.AfterUpdate(i)
	return nil
}

// selector, set
func (b *Bongo) UpdateMulti(i Modellable, rest ...map[string]interface{}) error {
	var set, selector map[string]interface{}

	switch len(rest) {
	case 1:
		set = rest[0]
		selector = nil
	case 2:
		selector = rest[0]
		set = rest[1]
	default:
		return errors.New("Update partial parameter list is wrong")
	}

	query := b.DB.Table(i.TableName())

	//add selector
	query = addWhere(query, selector)

	if err := query.Updates(set).Error; err != nil {
		return err
	}

	return nil
}

func (b *Bongo) Count(i Modellable, where ...interface{}) (int, error) {
	var count int

	// init query
	query := b.DB.Model(i)

	// add table name
	query = query.Table(i.TableName())

	// add query
	query = query.Where(where[0], where[1:len(where)]...)

	return count, query.Count(&count).Error
}

func (b *Bongo) CountWithQuery(i Modellable, q *Query) (int, error) {
	query := b.BuildQuery(i, q)
	var count int
	return count, query.Count(&count).Error
}

type Scope func(d *gorm.DB) *gorm.DB

type Query struct {
	Selector   map[string]interface{}
	Sort       map[string]string
	Pluck      string
	Pagination Pagination
	Scopes     []Scope
}

func (q *Query) AddScope(scope Scope) {
	if q.Scopes == nil {
		q.Scopes = make([]Scope, 0)
	}

	q.Scopes = append(q.Scopes, scope)
}

type Pagination struct {
	Limit int
	Skip  int
}

func NewPagination(limit int, skip int) *Pagination {
	return &Pagination{
		Limit: limit,
		Skip:  skip,
	}
}

func NewQS(selector map[string]interface{}) *Query {
	return &Query{
		Selector: selector,
	}
}

// selector, sort, limit, pluck,
func (b *Bongo) Some(i Modellable, data interface{}, q *Query) error {
	err := b.executeQuery(i, data, q)
	if err == gorm.RecordNotFound {
		return nil
	}
	return err
}

func (b *Bongo) One(i Modellable, data interface{}, q *Query) error {
	q.Pagination.Limit = 1
	return b.executeQuery(i, data, q)
}

func (b *Bongo) BuildQuery(i Modellable, q *Query) *gorm.DB {
	// init query
	query := b.DB.Model(i)

	// add table name
	query = query.Table(i.TableName())

	// add sort options
	query = addSort(query, q.Sort)

	query = addSkip(query, q.Pagination.Skip)

	query = addLimit(query, q.Pagination.Limit)

	// add selector
	query = addWhere(query, q.Selector)

	// put scopes
	if q.Scopes != nil && len(q.Scopes) > 0 {
		for _, scope := range q.Scopes {
			query = query.Scopes(scope)
		}
	}

	return query
}

func (b *Bongo) executeQuery(i Modellable, data interface{}, q *Query) error {
	// init query
	query := b.BuildQuery(i, q)

	var err error
	// TODO refactor this part
	if q.Pluck != "" {
		if strings.Contains(q.Pluck, ",") {
			// add pluck data
			query = addPluck(query, q.Pluck)

			err = query.Find(data).Error
		} else {
			err = query.Pluck(q.Pluck, data).Error
		}
	} else {
		err = query.Find(data).Error
	}

	return err
}

func (b *Bongo) PublishEvent(eventName string, i Modellable) error {
	data, err := json.Marshal(i)
	if err != nil {
		b.log.Error("Error while marshalling for publish %s", err)
		return err
	}

	err = b.Broker.Publish(i.TableName()+"_"+eventName, data)
	if err != nil {
		b.log.Error("Error while publishing %s", err)
		return err
	}

	return nil
}

func (b *Bongo) AfterCreate(i Modellable) {
	b.PublishEvent("created", i)
}

func (b *Bongo) AfterUpdate(i Modellable) {
	b.PublishEvent("updated", i)
}

func (b *Bongo) AfterDelete(i Modellable) {
	b.PublishEvent("deleted", i)
}

// addSort injects sort parameters into query
func addSort(query *gorm.DB, options map[string]string) *gorm.DB {

	if options == nil {
		return query
	}

	if len(options) == 0 {
		return query
	}

	var opts []string
	for key, val := range options {
		opts = append(opts, fmt.Sprintf("%s %v", key, val))
	}
	return query.Order(strings.Join(opts, ","))
}

// addPluck basically adds select statement for
// only required fields
func addPluck(query *gorm.DB, plucked string) *gorm.DB {
	if plucked == "" {
		return query
	}

	return query.Select(plucked)
}

// addWhere adds where query
func addWhere(query *gorm.DB, selector map[string]interface{}) *gorm.DB {
	if selector == nil {
		return query
	}

	// instead sending one selector, do chaining here
	return query.Where(selector)
}

// addSkip adds skip parameter into sql query
func addSkip(query *gorm.DB, skip int) *gorm.DB {
	if skip > 0 {
		return query.Offset(skip)
	}

	return query
}

// addLimit adds limit into query if set
func addLimit(query *gorm.DB, limit int) *gorm.DB {
	// if limit is minus or 0 ignore
	if limit > 0 {
		return query.Limit(limit)
	}

	return query
}

// CheckErr checks error exitence and returns
// if found, but this function suppress RecordNotFound errors
func CheckErr(res *gorm.DB) error {
	if res == nil {
		return nil
	}

	if res.Error == nil {
		return nil
	}

	if res.Error == gorm.RecordNotFound {
		return nil
	}

	return res.Error
}
<|endoftext|>


<commit_before>// Copyright 2021 Google Inc. All Rights Reserved.
//
// 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.

// Provides integration tests for write flows with implicit_dir flag set.
package implicitdir_test

import (
	"os"
	"testing"
)

func TestWriteAtEndOfFile(t *testing.T) {
	fileName := createTempFile()
	f, err := os.OpenFile(fileName, os.O_APPEND|os.O_WRONLY, 0600)
	if err != nil {
		t.Errorf("Open file for append: %v", err)
	}

	if _, err = f.WriteString("line 3\n"); err != nil {
		t.Errorf("AppendString: %v", err)
	}

	f.Close()

	compareFileContents(t, fileName, "line 1\nline 2\nline 3\n")
}

func TestWriteAtStartOfFile(t *testing.T) {
	fileName := createTempFile()
	f, err := os.OpenFile(fileName, os.O_WRONLY, 0600)
	if err != nil {
		t.Errorf("Open file for write at start: %v", err)
	}

	if _, err = f.WriteAt([]byte("line 4\n"), 0); err != nil {
		t.Errorf("WriteString-Start: %v", err)
	}

	f.Close()

	compareFileContents(t, fileName, "line 4\nline 2\n")
}

func TestWriteAtRandom(t *testing.T){
	fileName := createTempFile()

	f, err := os.OpenFile(fileName, os.O_WRONLY, 0600)
	if err != nil {
	    t.Errorf("Open file for write at random: %v", err)
	}

	if _, err = f.WriteAt([]byte("line 5\n"),7); err != nil {
		t.Errorf("WriteString-Random: %v", err)
	}

	f.Close()

	compareFileContents(t, fileName, "line 1\nline 5\n")
}
<commit_msg>goimports error fixed<commit_after>// Copyright 2021 Google Inc. All Rights Reserved.
//
// 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.

// Provides integration tests for write flows with implicit_dir flag set.
package implicitdir_test

import (
	"os"
	"testing"
)

func TestWriteAtEndOfFile(t *testing.T) {
	fileName := createTempFile()
	f, err := os.OpenFile(fileName, os.O_APPEND|os.O_WRONLY, 0600)
	if err != nil {
		t.Errorf("Open file for append: %v", err)
	}

	if _, err = f.WriteString("line 3\n"); err != nil {
		t.Errorf("AppendString: %v", err)
	}

	f.Close()

	compareFileContents(t, fileName, "line 1\nline 2\nline 3\n")
}

func TestWriteAtStartOfFile(t *testing.T) {
	fileName := createTempFile()
	f, err := os.OpenFile(fileName, os.O_WRONLY, 0600)
	if err != nil {
		t.Errorf("Open file for write at start: %v", err)
	}

	if _, err = f.WriteAt([]byte("line 4\n"), 0); err != nil {
		t.Errorf("WriteString-Start: %v", err)
	}

	f.Close()

	compareFileContents(t, fileName, "line 4\nline 2\n")
}

func TestWriteAtRandom(t *testing.T) {
	fileName := createTempFile()

	f, err := os.OpenFile(fileName, os.O_WRONLY, 0600)
	if err != nil {
		t.Errorf("Open file for write at random: %v", err)
	}

	if _, err = f.WriteAt([]byte("line 5\n"), 7); err != nil {
		t.Errorf("WriteString-Random: %v", err)
	}

	f.Close()

	compareFileContents(t, fileName, "line 1\nline 5\n")
}
<|endoftext|>


<commit_before>/*
Copyright 2019 The Rook Authors. All rights reserved.

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 crash

import (
	"context"

	"github.com/pkg/errors"
	"github.com/rook/rook/pkg/operator/ceph/cluster/mgr"
	"github.com/rook/rook/pkg/operator/ceph/cluster/mon"
	"github.com/rook/rook/pkg/operator/ceph/cluster/rbd"
	appsv1 "k8s.io/api/apps/v1"
	kerrors "k8s.io/apimachinery/pkg/api/errors"

	"github.com/rook/rook/pkg/operator/ceph/file/mds"
	"github.com/rook/rook/pkg/operator/ceph/file/mirror"
	"github.com/rook/rook/pkg/operator/ceph/object"

	"github.com/coreos/pkg/capnslog"

	"github.com/rook/rook/pkg/operator/ceph/cluster/osd"
	"github.com/rook/rook/pkg/operator/k8sutil"

	"k8s.io/apimachinery/pkg/runtime"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/reconcile"

	cephv1 "github.com/rook/rook/pkg/apis/ceph.rook.io/v1"
	opcontroller "github.com/rook/rook/pkg/operator/ceph/controller"
	"github.com/rook/rook/pkg/operator/ceph/disruption/controllerconfig"
	cephver "github.com/rook/rook/pkg/operator/ceph/version"
	"k8s.io/api/batch/v1beta1"
	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
)

var (
	logger = capnslog.NewPackageLogger("github.com/rook/rook", controllerName)
	// Implement reconcile.Reconciler so the controller can reconcile objects
	_ reconcile.Reconciler = &ReconcileNode{}
)

// ReconcileNode reconciles ReplicaSets
type ReconcileNode struct {
	// client can be used to retrieve objects from the APIServer.
	scheme *runtime.Scheme
	client client.Client
}

// Reconcile reconciles a node and ensures that it has a crashcollector deployment
// attached to it.
// The Controller will requeue the Request to be processed again if an error is non-nil or
// Result.Requeue is true, otherwise upon completion it will remove the work from the queue.
func (r *ReconcileNode) Reconcile(context context.Context, request reconcile.Request) (reconcile.Result, error) {
	// workaround because the rook logging mechanism is not compatible with the controller-runtime logging interface
	result, err := r.reconcile(request)
	if err != nil {
		logger.Error(err)
	}
	return result, err
}

func (r *ReconcileNode) reconcile(request reconcile.Request) (reconcile.Result, error) {
	logger.Debugf("reconciling node: %q", request.Name)

	// get the node object
	node := &corev1.Node{ObjectMeta: metav1.ObjectMeta{Name: request.Name}}
	err := r.client.Get(context.TODO(), request.NamespacedName, node)
	if err != nil {
		if kerrors.IsNotFound(err) {
			// if a node is not present, check if there are any crashcollector deployment for that node and delete it.
			err := r.listCrashCollectorAndDelete(request.Name, request.Namespace)
			if err != nil {
				logger.Errorf("failed to list and delete crash collector deployment on node %q; user should delete them manually. %v", request.Name, err)
			}
		} else {
			return reconcile.Result{}, errors.Wrapf(err, "could not get node %q", request.Name)
		}
	}

	// Get the list of all the Ceph pods
	cephPods, err := r.cephPodList()
	if err != nil {
		if len(cephPods) == 0 {
			return reconcile.Result{}, nil
		}
		return reconcile.Result{}, errors.Wrap(err, "failed to list all ceph pods")
	}

	namespaceToPodList := make(map[string][]corev1.Pod)
	for _, cephPod := range cephPods {
		podNamespace := cephPod.GetNamespace()
		podList, ok := namespaceToPodList[podNamespace]
		if !ok {
			// initialize list
			namespaceToPodList[podNamespace] = []corev1.Pod{cephPod}
		} else {
			// append cephPod to namespace's pod list
			namespaceToPodList[podNamespace] = append(podList, cephPod)
		}
	}

	for namespace, cephPods := range namespaceToPodList {
		// get dataDirHostPath from the CephCluster
		cephClusters := &cephv1.CephClusterList{}
		err := r.client.List(context.TODO(), cephClusters, client.InNamespace(namespace))
		if err != nil {
			return reconcile.Result{}, errors.Wrapf(err, "could not get cephcluster in namespaces %q", namespace)
		}
		if len(cephClusters.Items) < 1 {
			logger.Debugf("no CephCluster found in the namespace %q", namespace)
			return reconcile.Result{}, nil
		}

		cephCluster := cephClusters.Items[0]
		if len(cephClusters.Items) > 1 {
			logger.Errorf("more than one CephCluster found in the namespace %q, choosing the first one %q", namespace, cephCluster.GetName())
		}

		// If the crash controller is disabled in the spec let's do a noop
		if cephCluster.Spec.CrashCollector.Disable {
			deploymentList := &appsv1.DeploymentList{}
			namespaceListOpts := client.InNamespace(request.Namespace)

			// Try to fetch the list of existing deployment and remove them
			err := r.client.List(context.TODO(), deploymentList, client.MatchingLabels{k8sutil.AppAttr: AppName}, namespaceListOpts)
			if err != nil {
				logger.Errorf("failed to list crash collector deployments, delete it/them manually. %v", err)
				return reconcile.Result{}, nil
			}

			//  Try to delete all the crash deployments
			for _, d := range deploymentList.Items {
				err := r.deleteCrashCollector(d)
				if err != nil {
					logger.Errorf("failed to delete crash collector deployment %q, delete it manually. %v", d.Name, err)
					continue
				}
				logger.Infof("crash collector deployment %q successfully removed", d.Name)
			}

			return reconcile.Result{}, nil
		}

		clusterImage := cephCluster.Spec.CephVersion.Image
		cephVersion, err := opcontroller.GetImageVersion(cephCluster)
		if err != nil {
			logger.Errorf("ceph version not found for image %q used by cluster %q. %v", clusterImage, cephCluster.Name, err)
			return reconcile.Result{}, nil
		}

		uniqueTolerations := controllerconfig.TolerationSet{}
		hasCephPods := false
		for _, cephPod := range cephPods {
			if cephPod.Spec.NodeName == request.Name {
				hasCephPods = true
				for _, podToleration := range cephPod.Spec.Tolerations {
					// Add toleration to the map
					uniqueTolerations.Add(podToleration)
				}
			}
		}

		// If the node has Ceph pods we create a crash collector
		if hasCephPods {
			tolerations := uniqueTolerations.ToList()
			op, err := r.createOrUpdateCephCrash(*node, tolerations, cephCluster, cephVersion)
			if err != nil {
				return reconcile.Result{}, errors.Wrapf(err, "node reconcile failed on op %q", op)
			}
			logger.Debugf("deployment successfully reconciled for node %q. operation: %q", request.Name, op)
			// If there are no Ceph pods, check that there are no crash collector pods in case Ceph pods moved to another node
			// Thus the crash collector must be removed from that node
		} else {
			err := r.listCrashCollectorAndDelete(request.Name, request.Namespace)
			if err != nil {
				return reconcile.Result{}, errors.Wrapf(err, "failed to list and delete crash collector deployments on node %q", request.Name)
			}
		}

		if err := r.reconcileCrashRetention(namespace, cephCluster, cephVersion); err != nil {
			return reconcile.Result{}, err
		}
	}

	return reconcile.Result{}, nil
}

func (r *ReconcileNode) cephPodList() ([]corev1.Pod, error) {
	cephPods := make([]corev1.Pod, 0)
	cephAppNames := []string{mon.AppName, mgr.AppName, osd.AppName, object.AppName, mds.AppName, rbd.AppName, mirror.AppName}

	for _, app := range cephAppNames {
		podList := &corev1.PodList{}
		err := r.client.List(context.TODO(), podList, client.MatchingLabels{k8sutil.AppAttr: app})
		if err != nil {
			return cephPods, errors.Wrapf(err, "could not list the %q pods", app)
		}

		cephPods = append(cephPods, podList.Items...)
	}

	return cephPods, nil
}

func (r *ReconcileNode) listCrashCollectorAndDelete(nodeName, ns string) error {
	deploymentList := &appsv1.DeploymentList{}
	namespaceListOpts := client.InNamespace(ns)
	err := r.client.List(context.TODO(), deploymentList, client.MatchingLabels{k8sutil.AppAttr: AppName, NodeNameLabel: nodeName}, namespaceListOpts)
	if err != nil {
		return errors.Wrap(err, "failed to list crash collector deployments")
	}
	for _, d := range deploymentList.Items {
		logger.Infof("deleting deployment %q for node %q", d.ObjectMeta.Name, nodeName)
		err := r.deleteCrashCollector(d)
		if err != nil {
			return errors.Wrapf(err, "failed to delete crash collector deployment %q", d.Name)
		}
		logger.Infof("successfully removed crash collector deployment %q from node %q", d.Name, nodeName)
	}

	return nil
}

func (r *ReconcileNode) deleteCrashCollector(deployment appsv1.Deployment) error {
	deploymentName := deployment.ObjectMeta.Name
	namespace := deployment.ObjectMeta.Namespace
	dep := &appsv1.Deployment{
		ObjectMeta: metav1.ObjectMeta{
			Name:      deploymentName,
			Namespace: namespace,
		},
	}

	err := r.client.Delete(context.TODO(), dep)
	if err != nil && !kerrors.IsNotFound(err) {
		return errors.Wrapf(err, "could not delete crash collector deployment %q", deploymentName)
	}

	return nil
}

func (r *ReconcileNode) reconcileCrashRetention(namespace string, cephCluster cephv1.CephCluster, cephVersion *cephver.CephVersion) error {
	if cephCluster.Spec.CrashCollector.DaysToRetain == 0 {
		logger.Debug("deleting cronjob if it exists...")
		cronJob := &v1beta1.CronJob{
			ObjectMeta: metav1.ObjectMeta{
				Name:      prunerName,
				Namespace: namespace,
			},
		}

		err := r.client.Delete(context.TODO(), cronJob)
		if err != nil {
			if kerrors.IsNotFound(err) {
				logger.Debug("cronJob resource not found. Ignoring since object must be deleted.")
			} else {
				return err
			}
		} else {
			logger.Debug("successfully deleted crash pruner cronjob.")
		}
	} else {
		logger.Debugf("daysToRetain set to: %d", cephCluster.Spec.CrashCollector.DaysToRetain)
		op, err := r.createOrUpdateCephCron(cephCluster, cephVersion)
		if err != nil {
			return errors.Wrapf(err, "node reconcile failed on op %q", op)
		}
		logger.Debugf("cronjob successfully reconciled. operation: %q", op)
	}
	return nil
}
<commit_msg>ceph: create crash collector pods after secret created<commit_after>/*
Copyright 2019 The Rook Authors. All rights reserved.

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 crash

import (
	"context"
	"time"

	"github.com/pkg/errors"
	"github.com/rook/rook/pkg/operator/ceph/cluster/mgr"
	"github.com/rook/rook/pkg/operator/ceph/cluster/mon"
	"github.com/rook/rook/pkg/operator/ceph/cluster/rbd"
	appsv1 "k8s.io/api/apps/v1"
	kerrors "k8s.io/apimachinery/pkg/api/errors"
	"k8s.io/apimachinery/pkg/types"

	"github.com/rook/rook/pkg/operator/ceph/file/mds"
	"github.com/rook/rook/pkg/operator/ceph/file/mirror"
	"github.com/rook/rook/pkg/operator/ceph/object"

	"github.com/coreos/pkg/capnslog"

	"github.com/rook/rook/pkg/operator/ceph/cluster/osd"
	"github.com/rook/rook/pkg/operator/k8sutil"

	"k8s.io/apimachinery/pkg/runtime"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/reconcile"

	cephv1 "github.com/rook/rook/pkg/apis/ceph.rook.io/v1"
	opcontroller "github.com/rook/rook/pkg/operator/ceph/controller"
	"github.com/rook/rook/pkg/operator/ceph/disruption/controllerconfig"
	cephver "github.com/rook/rook/pkg/operator/ceph/version"
	"k8s.io/api/batch/v1beta1"
	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
)

var (
	logger = capnslog.NewPackageLogger("github.com/rook/rook", controllerName)
	// Implement reconcile.Reconciler so the controller can reconcile objects
	_ reconcile.Reconciler = &ReconcileNode{}

	// wait for secret "rook-ceph-crash-collector-keyring" to be created
	waitForRequeueIfSecretNotCreated = reconcile.Result{Requeue: true, RequeueAfter: 30 * time.Second}
)

// ReconcileNode reconciles ReplicaSets
type ReconcileNode struct {
	// client can be used to retrieve objects from the APIServer.
	scheme *runtime.Scheme
	client client.Client
}

// Reconcile reconciles a node and ensures that it has a crashcollector deployment
// attached to it.
// The Controller will requeue the Request to be processed again if an error is non-nil or
// Result.Requeue is true, otherwise upon completion it will remove the work from the queue.
func (r *ReconcileNode) Reconcile(context context.Context, request reconcile.Request) (reconcile.Result, error) {
	// workaround because the rook logging mechanism is not compatible with the controller-runtime logging interface
	result, err := r.reconcile(request)
	if err != nil {
		logger.Error(err)
	}
	return result, err
}

func (r *ReconcileNode) reconcile(request reconcile.Request) (reconcile.Result, error) {
	logger.Debugf("reconciling node: %q", request.Name)

	// get the node object
	node := &corev1.Node{ObjectMeta: metav1.ObjectMeta{Name: request.Name}}
	err := r.client.Get(context.TODO(), request.NamespacedName, node)
	if err != nil {
		if kerrors.IsNotFound(err) {
			// if a node is not present, check if there are any crashcollector deployment for that node and delete it.
			err := r.listCrashCollectorAndDelete(request.Name, request.Namespace)
			if err != nil {
				logger.Errorf("failed to list and delete crash collector deployment on node %q; user should delete them manually. %v", request.Name, err)
			}
		} else {
			return reconcile.Result{}, errors.Wrapf(err, "could not get node %q", request.Name)
		}
	}

	// Get the list of all the Ceph pods
	cephPods, err := r.cephPodList()
	if err != nil {
		if len(cephPods) == 0 {
			return reconcile.Result{}, nil
		}
		return reconcile.Result{}, errors.Wrap(err, "failed to list all ceph pods")
	}

	namespaceToPodList := make(map[string][]corev1.Pod)
	for _, cephPod := range cephPods {
		podNamespace := cephPod.GetNamespace()
		podList, ok := namespaceToPodList[podNamespace]
		if !ok {
			// initialize list
			namespaceToPodList[podNamespace] = []corev1.Pod{cephPod}
		} else {
			// append cephPod to namespace's pod list
			namespaceToPodList[podNamespace] = append(podList, cephPod)
		}
	}

	for namespace, cephPods := range namespaceToPodList {
		// get dataDirHostPath from the CephCluster
		cephClusters := &cephv1.CephClusterList{}
		err := r.client.List(context.TODO(), cephClusters, client.InNamespace(namespace))
		if err != nil {
			return reconcile.Result{}, errors.Wrapf(err, "could not get cephcluster in namespaces %q", namespace)
		}
		if len(cephClusters.Items) < 1 {
			logger.Debugf("no CephCluster found in the namespace %q", namespace)
			return reconcile.Result{}, nil
		}

		cephCluster := cephClusters.Items[0]
		if len(cephClusters.Items) > 1 {
			logger.Errorf("more than one CephCluster found in the namespace %q, choosing the first one %q", namespace, cephCluster.GetName())
		}

		// If the crash controller is disabled in the spec let's do a noop
		if cephCluster.Spec.CrashCollector.Disable {
			deploymentList := &appsv1.DeploymentList{}
			namespaceListOpts := client.InNamespace(request.Namespace)

			// Try to fetch the list of existing deployment and remove them
			err := r.client.List(context.TODO(), deploymentList, client.MatchingLabels{k8sutil.AppAttr: AppName}, namespaceListOpts)
			if err != nil {
				logger.Errorf("failed to list crash collector deployments, delete it/them manually. %v", err)
				return reconcile.Result{}, nil
			}

			//  Try to delete all the crash deployments
			for _, d := range deploymentList.Items {
				err := r.deleteCrashCollector(d)
				if err != nil {
					logger.Errorf("failed to delete crash collector deployment %q, delete it manually. %v", d.Name, err)
					continue
				}
				logger.Infof("crash collector deployment %q successfully removed", d.Name)
			}

			return reconcile.Result{}, nil
		}

		// checking if secret "rook-ceph-crash-collector-keyring" is present which is required to create crashcollector pods
		secret := &corev1.Secret{}
		err = r.client.Get(context.TODO(), types.NamespacedName{Name: crashCollectorKeyName, Namespace: namespace}, secret)
		if err != nil {
			if kerrors.IsNotFound(err) {
				logger.Debugf("secret %q not found. retrying in %q. %v", crashCollectorKeyName, waitForRequeueIfSecretNotCreated.RequeueAfter.String(), err)
				return waitForRequeueIfSecretNotCreated, nil
			}

			return reconcile.Result{}, errors.Wrapf(err, "failed to list the %q secret.", crashCollectorKeyName)
		}

		clusterImage := cephCluster.Spec.CephVersion.Image
		cephVersion, err := opcontroller.GetImageVersion(cephCluster)
		if err != nil {
			logger.Errorf("ceph version not found for image %q used by cluster %q. %v", clusterImage, cephCluster.Name, err)
			return reconcile.Result{}, nil
		}

		uniqueTolerations := controllerconfig.TolerationSet{}
		hasCephPods := false
		for _, cephPod := range cephPods {
			if cephPod.Spec.NodeName == request.Name {
				hasCephPods = true
				for _, podToleration := range cephPod.Spec.Tolerations {
					// Add toleration to the map
					uniqueTolerations.Add(podToleration)
				}
			}
		}

		// If the node has Ceph pods we create a crash collector
		if hasCephPods {
			tolerations := uniqueTolerations.ToList()
			op, err := r.createOrUpdateCephCrash(*node, tolerations, cephCluster, cephVersion)
			if err != nil {
				return reconcile.Result{}, errors.Wrapf(err, "node reconcile failed on op %q", op)
			}
			logger.Debugf("deployment successfully reconciled for node %q. operation: %q", request.Name, op)
			// If there are no Ceph pods, check that there are no crash collector pods in case Ceph pods moved to another node
			// Thus the crash collector must be removed from that node
		} else {
			err := r.listCrashCollectorAndDelete(request.Name, request.Namespace)
			if err != nil {
				return reconcile.Result{}, errors.Wrapf(err, "failed to list and delete crash collector deployments on node %q", request.Name)
			}
		}

		if err := r.reconcileCrashRetention(namespace, cephCluster, cephVersion); err != nil {
			return reconcile.Result{}, err
		}
	}

	return reconcile.Result{}, nil
}

func (r *ReconcileNode) cephPodList() ([]corev1.Pod, error) {
	cephPods := make([]corev1.Pod, 0)
	cephAppNames := []string{mon.AppName, mgr.AppName, osd.AppName, object.AppName, mds.AppName, rbd.AppName, mirror.AppName}

	for _, app := range cephAppNames {
		podList := &corev1.PodList{}
		err := r.client.List(context.TODO(), podList, client.MatchingLabels{k8sutil.AppAttr: app})
		if err != nil {
			return cephPods, errors.Wrapf(err, "could not list the %q pods", app)
		}

		cephPods = append(cephPods, podList.Items...)
	}

	return cephPods, nil
}

func (r *ReconcileNode) listCrashCollectorAndDelete(nodeName, ns string) error {
	deploymentList := &appsv1.DeploymentList{}
	namespaceListOpts := client.InNamespace(ns)
	err := r.client.List(context.TODO(), deploymentList, client.MatchingLabels{k8sutil.AppAttr: AppName, NodeNameLabel: nodeName}, namespaceListOpts)
	if err != nil {
		return errors.Wrap(err, "failed to list crash collector deployments")
	}
	for _, d := range deploymentList.Items {
		logger.Infof("deleting deployment %q for node %q", d.ObjectMeta.Name, nodeName)
		err := r.deleteCrashCollector(d)
		if err != nil {
			return errors.Wrapf(err, "failed to delete crash collector deployment %q", d.Name)
		}
		logger.Infof("successfully removed crash collector deployment %q from node %q", d.Name, nodeName)
	}

	return nil
}

func (r *ReconcileNode) deleteCrashCollector(deployment appsv1.Deployment) error {
	deploymentName := deployment.ObjectMeta.Name
	namespace := deployment.ObjectMeta.Namespace
	dep := &appsv1.Deployment{
		ObjectMeta: metav1.ObjectMeta{
			Name:      deploymentName,
			Namespace: namespace,
		},
	}

	err := r.client.Delete(context.TODO(), dep)
	if err != nil && !kerrors.IsNotFound(err) {
		return errors.Wrapf(err, "could not delete crash collector deployment %q", deploymentName)
	}

	return nil
}

func (r *ReconcileNode) reconcileCrashRetention(namespace string, cephCluster cephv1.CephCluster, cephVersion *cephver.CephVersion) error {
	if cephCluster.Spec.CrashCollector.DaysToRetain == 0 {
		logger.Debug("deleting cronjob if it exists...")
		cronJob := &v1beta1.CronJob{
			ObjectMeta: metav1.ObjectMeta{
				Name:      prunerName,
				Namespace: namespace,
			},
		}

		err := r.client.Delete(context.TODO(), cronJob)
		if err != nil {
			if kerrors.IsNotFound(err) {
				logger.Debug("cronJob resource not found. Ignoring since object must be deleted.")
			} else {
				return err
			}
		} else {
			logger.Debug("successfully deleted crash pruner cronjob.")
		}
	} else {
		logger.Debugf("daysToRetain set to: %d", cephCluster.Spec.CrashCollector.DaysToRetain)
		op, err := r.createOrUpdateCephCron(cephCluster, cephVersion)
		if err != nil {
			return errors.Wrapf(err, "node reconcile failed on op %q", op)
		}
		logger.Debugf("cronjob successfully reconciled. operation: %q", op)
	}
	return nil
}
<|endoftext|>


<commit_before>/*
 * This file is part of the KubeVirt project
 *
 * 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.
 *
 * Copyright 2018 Red Hat, Inc.
 *
 */

//go:generate mockgen -source $GOFILE -package=$GOPACKAGE -destination=generated_mock_$GOFILE

package network

import (
	"crypto/rand"
	"encoding/json"
	"fmt"

	"io/ioutil"
	"net"
	"os"

	"github.com/coreos/go-iptables/iptables"

	lmf "github.com/subgraph/libmacouflage"
	"github.com/vishvananda/netlink"

	v1 "kubevirt.io/kubevirt/pkg/api/v1"
	"kubevirt.io/kubevirt/pkg/log"
	"kubevirt.io/kubevirt/pkg/virt-launcher/virtwrap/api"
	"kubevirt.io/kubevirt/pkg/virt-launcher/virtwrap/network/dhcp"
)

const randomMacGenerationAttempts = 10

type VIF struct {
	Name    string
	IP      netlink.Addr
	MAC     net.HardwareAddr
	Gateway net.IP
	Routes  *[]netlink.Route
	Mtu     uint16
}

func (vif VIF) String() string {
	return fmt.Sprintf(
		"VIF: { Name: %s, IP: %s, Mask: %s, MAC: %s, Gateway: %s, MTU: %d}",
		vif.Name,
		vif.IP.IP,
		vif.IP.Mask,
		vif.MAC,
		vif.Gateway,
		vif.Mtu,
	)
}

type NetworkHandler interface {
	LinkByName(name string) (netlink.Link, error)
	AddrList(link netlink.Link, family int) ([]netlink.Addr, error)
	RouteList(link netlink.Link, family int) ([]netlink.Route, error)
	AddrDel(link netlink.Link, addr *netlink.Addr) error
	AddrAdd(link netlink.Link, addr *netlink.Addr) error
	LinkSetDown(link netlink.Link) error
	LinkSetUp(link netlink.Link) error
	LinkAdd(link netlink.Link) error
	LinkSetLearningOff(link netlink.Link) error
	ParseAddr(s string) (*netlink.Addr, error)
	GetHostAndGwAddressesFromCIDR(s string) (string, string, error)
	SetRandomMac(iface string) (net.HardwareAddr, error)
	GenerateRandomMac() (net.HardwareAddr, error)
	GetMacDetails(iface string) (net.HardwareAddr, error)
	LinkSetMaster(link netlink.Link, master *netlink.Bridge) error
	StartDHCP(nic *VIF, serverAddr *netlink.Addr, bridgeInterfaceName string, dhcpOptions *v1.DHCPOptions)
	IptablesNewChain(table, chain string) error
	IptablesAppendRule(table, chain string, rulespec ...string) error
}

type NetworkUtilsHandler struct{}

var Handler NetworkHandler

func (h *NetworkUtilsHandler) LinkByName(name string) (netlink.Link, error) {
	return netlink.LinkByName(name)
}
func (h *NetworkUtilsHandler) AddrList(link netlink.Link, family int) ([]netlink.Addr, error) {
	return netlink.AddrList(link, family)
}
func (h *NetworkUtilsHandler) RouteList(link netlink.Link, family int) ([]netlink.Route, error) {
	return netlink.RouteList(link, family)
}
func (h *NetworkUtilsHandler) AddrDel(link netlink.Link, addr *netlink.Addr) error {
	return netlink.AddrDel(link, addr)
}
func (h *NetworkUtilsHandler) LinkSetDown(link netlink.Link) error {
	return netlink.LinkSetDown(link)
}
func (h *NetworkUtilsHandler) LinkSetUp(link netlink.Link) error {
	return netlink.LinkSetUp(link)
}
func (h *NetworkUtilsHandler) LinkAdd(link netlink.Link) error {
	return netlink.LinkAdd(link)
}
func (h *NetworkUtilsHandler) LinkSetLearningOff(link netlink.Link) error {
	return netlink.LinkSetLearning(link, false)
}
func (h *NetworkUtilsHandler) ParseAddr(s string) (*netlink.Addr, error) {
	return netlink.ParseAddr(s)
}
func (h *NetworkUtilsHandler) AddrAdd(link netlink.Link, addr *netlink.Addr) error {
	return netlink.AddrAdd(link, addr)
}
func (h *NetworkUtilsHandler) LinkSetMaster(link netlink.Link, master *netlink.Bridge) error {
	return netlink.LinkSetMaster(link, master)
}
func (h *NetworkUtilsHandler) IptablesNewChain(table, chain string) error {
	iptablesObject, err := iptables.New()
	if err != nil {
		return err
	}

	return iptablesObject.NewChain(table, chain)
}
func (h *NetworkUtilsHandler) IptablesAppendRule(table, chain string, rulespec ...string) error {
	iptablesObject, err := iptables.New()
	if err != nil {
		return err
	}

	return iptablesObject.Append(table, chain, rulespec...)
}
func (h *NetworkUtilsHandler) GetHostAndGwAddressesFromCIDR(s string) (string, string, error) {
	ip, ipnet, err := net.ParseCIDR(s)
	if err != nil {
		return "", "", err
	}

	subnet, _ := ipnet.Mask.Size()
	var ips []string
	for ip := ip.Mask(ipnet.Mask); ipnet.Contains(ip); inc(ip) {
		ips = append(ips, fmt.Sprintf("%s/%d", ip.String(), subnet))

		if len(ips) == 4 {
			// remove network address and broadcast address
			return ips[1], ips[2], nil
		}
	}

	return "", "", fmt.Errorf("less than 4 addresses on network")
}

func inc(ip net.IP) {
	for j := len(ip) - 1; j >= 0; j-- {
		ip[j]++
		if ip[j] > 0 {
			break
		}
	}
}

// GetMacDetails from an interface
func (h *NetworkUtilsHandler) GetMacDetails(iface string) (net.HardwareAddr, error) {
	currentMac, err := lmf.GetCurrentMac(iface)
	if err != nil {
		log.Log.Reason(err).Errorf("failed to get mac information for interface: %s", iface)
		return nil, err
	}
	return currentMac, nil
}

// SetRandomMac changes the MAC address for a given interface to a randomly generated, preserving the vendor prefix
func (h *NetworkUtilsHandler) SetRandomMac(iface string) (net.HardwareAddr, error) {
	var mac net.HardwareAddr

	currentMac, err := Handler.GetMacDetails(iface)
	if err != nil {
		return nil, err
	}

	changed := false

	for i := 0; i < randomMacGenerationAttempts; i++ {
		changed, err = lmf.SpoofMacSameVendor(iface, false)
		if err != nil {
			log.Log.Reason(err).Errorf("failed to spoof MAC for an interface: %s", iface)
			return nil, err
		}

		if changed {
			mac, err = Handler.GetMacDetails(iface)
			if err != nil {
				return nil, err
			}
			log.Log.Infof("updated MAC for interface: %s - %s -> %s", iface, currentMac, mac)
			break
		}
	}
	if !changed {
		err := fmt.Errorf("failed to spoof MAC for an interface %s after %d attempts", iface, randomMacGenerationAttempts)
		log.Log.Reason(err)
		return nil, err
	}
	return currentMac, nil
}

func (h *NetworkUtilsHandler) StartDHCP(nic *VIF, serverAddr *netlink.Addr, bridgeInterfaceName string, dhcpOptions *v1.DHCPOptions) {
	log.Log.V(4).Infof("StartDHCP network Nic: %+v", nic)
	nameservers, searchDomains, err := api.GetResolvConfDetailsFromPod()
	if err != nil {
		log.Log.Errorf("Failed to get DNS servers from resolv.conf: %v", err)
		panic(err)
	}

	// panic in case the DHCP server failed during the vm creation
	// but ignore dhcp errors when the vm is destroyed or shutting down
	go func() {
		if err = DHCPServer(
			nic.MAC,
			nic.IP.IP,
			nic.IP.Mask,
			bridgeInterfaceName,
			serverAddr.IP,
			nic.Gateway,
			nameservers,
			nic.Routes,
			searchDomains,
			nic.Mtu,
			dhcpOptions,
		); err != nil {
			log.Log.Errorf("failed to run DHCP: %v", err)
			panic(err)
		}
	}()
}

// Generate a random mac for interface
// Avoid MAC address starting with reserved value 0xFE (https://github.com/kubevirt/kubevirt/issues/1494)
func (h *NetworkUtilsHandler) GenerateRandomMac() (net.HardwareAddr, error) {
	prefix := []byte{0x02, 0x00, 0x00} // local unicast prefix
	suffix := make([]byte, 3)
	_, err := rand.Read(suffix)
	if err != nil {
		return nil, err
	}
	return net.HardwareAddr(append(prefix, suffix...)), nil
}

// Allow mocking for tests
var SetupPodNetwork = SetupNetworkInterfaces
var DHCPServer = dhcp.SingleClientDHCPServer

func initHandler() {
	if Handler == nil {
		Handler = &NetworkUtilsHandler{}
	}
}

func writeToCachedFile(inter interface{}, fileName, name string) error {
	buf, err := json.MarshalIndent(&inter, "", "  ")
	if err != nil {
		return fmt.Errorf("error marshaling cached object: %v", err)
	}
	err = ioutil.WriteFile(getInterfaceCacheFile(fileName, name), buf, 0644)
	if err != nil {
		return fmt.Errorf("error writing cached object: %v", err)
	}
	return nil
}

func readFromCachedFile(name, fileName string, inter interface{}) (bool, error) {
	buf, err := ioutil.ReadFile(getInterfaceCacheFile(fileName, name))
	if err != nil {
		if os.IsNotExist(err) {
			return false, nil
		}
		return false, err
	}

	err = json.Unmarshal(buf, &inter)
	if err != nil {
		return false, fmt.Errorf("error unmarshaling cached object: %v", err)
	}
	return true, nil
}

func getInterfaceCacheFile(filePath, name string) string {
	return fmt.Sprintf(filePath, name)
}

// filter out irrelevant routes
func filterPodNetworkRoutes(routes []netlink.Route, nic *VIF) (filteredRoutes []netlink.Route) {
	for _, route := range routes {
		// don't create empty static routes
		if route.Dst == nil && route.Src.Equal(nil) && route.Gw.Equal(nil) {
			continue
		}

		// don't create static route for src == nic
		if route.Src != nil && route.Src.Equal(nic.IP.IP) {
			continue
		}

		filteredRoutes = append(filteredRoutes, route)
	}
	return
}

// only used by unit test suite
func setInterfaceCacheFile(path string) {
	interfaceCacheFile = path
}
<commit_msg>Clarify log message about MAC address changed<commit_after>/*
 * This file is part of the KubeVirt project
 *
 * 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.
 *
 * Copyright 2018 Red Hat, Inc.
 *
 */

//go:generate mockgen -source $GOFILE -package=$GOPACKAGE -destination=generated_mock_$GOFILE

package network

import (
	"crypto/rand"
	"encoding/json"
	"fmt"

	"io/ioutil"
	"net"
	"os"

	"github.com/coreos/go-iptables/iptables"

	lmf "github.com/subgraph/libmacouflage"
	"github.com/vishvananda/netlink"

	v1 "kubevirt.io/kubevirt/pkg/api/v1"
	"kubevirt.io/kubevirt/pkg/log"
	"kubevirt.io/kubevirt/pkg/virt-launcher/virtwrap/api"
	"kubevirt.io/kubevirt/pkg/virt-launcher/virtwrap/network/dhcp"
)

const randomMacGenerationAttempts = 10

type VIF struct {
	Name    string
	IP      netlink.Addr
	MAC     net.HardwareAddr
	Gateway net.IP
	Routes  *[]netlink.Route
	Mtu     uint16
}

func (vif VIF) String() string {
	return fmt.Sprintf(
		"VIF: { Name: %s, IP: %s, Mask: %s, MAC: %s, Gateway: %s, MTU: %d}",
		vif.Name,
		vif.IP.IP,
		vif.IP.Mask,
		vif.MAC,
		vif.Gateway,
		vif.Mtu,
	)
}

type NetworkHandler interface {
	LinkByName(name string) (netlink.Link, error)
	AddrList(link netlink.Link, family int) ([]netlink.Addr, error)
	RouteList(link netlink.Link, family int) ([]netlink.Route, error)
	AddrDel(link netlink.Link, addr *netlink.Addr) error
	AddrAdd(link netlink.Link, addr *netlink.Addr) error
	LinkSetDown(link netlink.Link) error
	LinkSetUp(link netlink.Link) error
	LinkAdd(link netlink.Link) error
	LinkSetLearningOff(link netlink.Link) error
	ParseAddr(s string) (*netlink.Addr, error)
	GetHostAndGwAddressesFromCIDR(s string) (string, string, error)
	SetRandomMac(iface string) (net.HardwareAddr, error)
	GenerateRandomMac() (net.HardwareAddr, error)
	GetMacDetails(iface string) (net.HardwareAddr, error)
	LinkSetMaster(link netlink.Link, master *netlink.Bridge) error
	StartDHCP(nic *VIF, serverAddr *netlink.Addr, bridgeInterfaceName string, dhcpOptions *v1.DHCPOptions)
	IptablesNewChain(table, chain string) error
	IptablesAppendRule(table, chain string, rulespec ...string) error
}

type NetworkUtilsHandler struct{}

var Handler NetworkHandler

func (h *NetworkUtilsHandler) LinkByName(name string) (netlink.Link, error) {
	return netlink.LinkByName(name)
}
func (h *NetworkUtilsHandler) AddrList(link netlink.Link, family int) ([]netlink.Addr, error) {
	return netlink.AddrList(link, family)
}
func (h *NetworkUtilsHandler) RouteList(link netlink.Link, family int) ([]netlink.Route, error) {
	return netlink.RouteList(link, family)
}
func (h *NetworkUtilsHandler) AddrDel(link netlink.Link, addr *netlink.Addr) error {
	return netlink.AddrDel(link, addr)
}
func (h *NetworkUtilsHandler) LinkSetDown(link netlink.Link) error {
	return netlink.LinkSetDown(link)
}
func (h *NetworkUtilsHandler) LinkSetUp(link netlink.Link) error {
	return netlink.LinkSetUp(link)
}
func (h *NetworkUtilsHandler) LinkAdd(link netlink.Link) error {
	return netlink.LinkAdd(link)
}
func (h *NetworkUtilsHandler) LinkSetLearningOff(link netlink.Link) error {
	return netlink.LinkSetLearning(link, false)
}
func (h *NetworkUtilsHandler) ParseAddr(s string) (*netlink.Addr, error) {
	return netlink.ParseAddr(s)
}
func (h *NetworkUtilsHandler) AddrAdd(link netlink.Link, addr *netlink.Addr) error {
	return netlink.AddrAdd(link, addr)
}
func (h *NetworkUtilsHandler) LinkSetMaster(link netlink.Link, master *netlink.Bridge) error {
	return netlink.LinkSetMaster(link, master)
}
func (h *NetworkUtilsHandler) IptablesNewChain(table, chain string) error {
	iptablesObject, err := iptables.New()
	if err != nil {
		return err
	}

	return iptablesObject.NewChain(table, chain)
}
func (h *NetworkUtilsHandler) IptablesAppendRule(table, chain string, rulespec ...string) error {
	iptablesObject, err := iptables.New()
	if err != nil {
		return err
	}

	return iptablesObject.Append(table, chain, rulespec...)
}
func (h *NetworkUtilsHandler) GetHostAndGwAddressesFromCIDR(s string) (string, string, error) {
	ip, ipnet, err := net.ParseCIDR(s)
	if err != nil {
		return "", "", err
	}

	subnet, _ := ipnet.Mask.Size()
	var ips []string
	for ip := ip.Mask(ipnet.Mask); ipnet.Contains(ip); inc(ip) {
		ips = append(ips, fmt.Sprintf("%s/%d", ip.String(), subnet))

		if len(ips) == 4 {
			// remove network address and broadcast address
			return ips[1], ips[2], nil
		}
	}

	return "", "", fmt.Errorf("less than 4 addresses on network")
}

func inc(ip net.IP) {
	for j := len(ip) - 1; j >= 0; j-- {
		ip[j]++
		if ip[j] > 0 {
			break
		}
	}
}

// GetMacDetails from an interface
func (h *NetworkUtilsHandler) GetMacDetails(iface string) (net.HardwareAddr, error) {
	currentMac, err := lmf.GetCurrentMac(iface)
	if err != nil {
		log.Log.Reason(err).Errorf("failed to get mac information for interface: %s", iface)
		return nil, err
	}
	return currentMac, nil
}

// SetRandomMac changes the MAC address for a given interface to a randomly generated, preserving the vendor prefix
func (h *NetworkUtilsHandler) SetRandomMac(iface string) (net.HardwareAddr, error) {
	var mac net.HardwareAddr

	currentMac, err := Handler.GetMacDetails(iface)
	if err != nil {
		return nil, err
	}

	changed := false

	for i := 0; i < randomMacGenerationAttempts; i++ {
		changed, err = lmf.SpoofMacSameVendor(iface, false)
		if err != nil {
			log.Log.Reason(err).Errorf("failed to spoof MAC for an interface: %s", iface)
			return nil, err
		}

		if changed {
			mac, err = Handler.GetMacDetails(iface)
			if err != nil {
				return nil, err
			}
			log.Log.Infof("updated MAC for %s interface: old: %s -> new: %s", iface, currentMac, mac)
			break
		}
	}
	if !changed {
		err := fmt.Errorf("failed to spoof MAC for an interface %s after %d attempts", iface, randomMacGenerationAttempts)
		log.Log.Reason(err)
		return nil, err
	}
	return currentMac, nil
}

func (h *NetworkUtilsHandler) StartDHCP(nic *VIF, serverAddr *netlink.Addr, bridgeInterfaceName string, dhcpOptions *v1.DHCPOptions) {
	log.Log.V(4).Infof("StartDHCP network Nic: %+v", nic)
	nameservers, searchDomains, err := api.GetResolvConfDetailsFromPod()
	if err != nil {
		log.Log.Errorf("Failed to get DNS servers from resolv.conf: %v", err)
		panic(err)
	}

	// panic in case the DHCP server failed during the vm creation
	// but ignore dhcp errors when the vm is destroyed or shutting down
	go func() {
		if err = DHCPServer(
			nic.MAC,
			nic.IP.IP,
			nic.IP.Mask,
			bridgeInterfaceName,
			serverAddr.IP,
			nic.Gateway,
			nameservers,
			nic.Routes,
			searchDomains,
			nic.Mtu,
			dhcpOptions,
		); err != nil {
			log.Log.Errorf("failed to run DHCP: %v", err)
			panic(err)
		}
	}()
}

// Generate a random mac for interface
// Avoid MAC address starting with reserved value 0xFE (https://github.com/kubevirt/kubevirt/issues/1494)
func (h *NetworkUtilsHandler) GenerateRandomMac() (net.HardwareAddr, error) {
	prefix := []byte{0x02, 0x00, 0x00} // local unicast prefix
	suffix := make([]byte, 3)
	_, err := rand.Read(suffix)
	if err != nil {
		return nil, err
	}
	return net.HardwareAddr(append(prefix, suffix...)), nil
}

// Allow mocking for tests
var SetupPodNetwork = SetupNetworkInterfaces
var DHCPServer = dhcp.SingleClientDHCPServer

func initHandler() {
	if Handler == nil {
		Handler = &NetworkUtilsHandler{}
	}
}

func writeToCachedFile(inter interface{}, fileName, name string) error {
	buf, err := json.MarshalIndent(&inter, "", "  ")
	if err != nil {
		return fmt.Errorf("error marshaling cached object: %v", err)
	}
	err = ioutil.WriteFile(getInterfaceCacheFile(fileName, name), buf, 0644)
	if err != nil {
		return fmt.Errorf("error writing cached object: %v", err)
	}
	return nil
}

func readFromCachedFile(name, fileName string, inter interface{}) (bool, error) {
	buf, err := ioutil.ReadFile(getInterfaceCacheFile(fileName, name))
	if err != nil {
		if os.IsNotExist(err) {
			return false, nil
		}
		return false, err
	}

	err = json.Unmarshal(buf, &inter)
	if err != nil {
		return false, fmt.Errorf("error unmarshaling cached object: %v", err)
	}
	return true, nil
}

func getInterfaceCacheFile(filePath, name string) string {
	return fmt.Sprintf(filePath, name)
}

// filter out irrelevant routes
func filterPodNetworkRoutes(routes []netlink.Route, nic *VIF) (filteredRoutes []netlink.Route) {
	for _, route := range routes {
		// don't create empty static routes
		if route.Dst == nil && route.Src.Equal(nil) && route.Gw.Equal(nil) {
			continue
		}

		// don't create static route for src == nic
		if route.Src != nil && route.Src.Equal(nic.IP.IP) {
			continue
		}

		filteredRoutes = append(filteredRoutes, route)
	}
	return
}

// only used by unit test suite
func setInterfaceCacheFile(path string) {
	interfaceCacheFile = path
}
<|endoftext|>