<commit_before>package k8sraw

import (
	"fmt"
	"github.com/Aptomi/aptomi/pkg/config"
	"github.com/Aptomi/aptomi/pkg/event"
	"github.com/Aptomi/aptomi/pkg/lang"
	"github.com/Aptomi/aptomi/pkg/plugin"
	"github.com/Aptomi/aptomi/pkg/plugin/k8s"
	"github.com/Aptomi/aptomi/pkg/util"
	"github.com/Aptomi/aptomi/pkg/util/sync"
	meta "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/client-go/pkg/api/v1"
	"k8s.io/client-go/pkg/apis/apps/v1beta1"
	"k8s.io/helm/pkg/kube"
	"strings"
)

// Plugin represents Kubernetes Raw code plugin that supports deploying specified k8s objects into the cluster
type Plugin struct {
	once          sync.Init
	cluster       *lang.Cluster
	config        config.K8sRaw
	kube          *k8s.Plugin
	dataNamespace string
}

// New returns new instance of the Kubernetes Raw code (objects) plugin for specified Kubernetes cluster plugin and plugins config
func New(clusterPlugin plugin.ClusterPlugin, cfg config.Plugins) (plugin.CodePlugin, error) {
	kubePlugin, ok := clusterPlugin.(*k8s.Plugin)
	if !ok {
		return nil, fmt.Errorf("k8s cluster plugin expected for k8sraw code plugin creation but received: %T", clusterPlugin)
	}

	return &Plugin{
		cluster: kubePlugin.Cluster,
		config:  cfg.K8sRaw,
		kube:    kubePlugin,
	}, nil
}

func (p *Plugin) init() error {
	return p.once.Do(func() error {
		err := p.kube.Init()
		if err != nil {
			return err
		}

		err = p.parseClusterConfig()
		if err != nil {
			return err
		}

		kubeClient, err := p.kube.NewClient()
		if err != nil {
			return err
		}

		return p.kube.EnsureNamespace(kubeClient, p.dataNamespace)
	})
}

// Cleanup implements cleanup phase for the k8s raw plugin
func (p *Plugin) Cleanup() error {
	return nil
}

// Create implements creation of a new component instance in the cloud by deploying raw k8s objects
func (p *Plugin) Create(deployName string, params util.NestedParameterMap, eventLog *event.Log) error {
	err := p.init()
	if err != nil {
		return err
	}

	kubeClient, err := p.kube.NewClient()
	if err != nil {
		return err
	}

	targetManifest, ok := params["manifest"].(string)
	if !ok {
		return fmt.Errorf("manifest is a mandatory parameter")
	}

	client := p.prepareClient(eventLog, deployName)

	err = client.Create(p.kube.Namespace, strings.NewReader(targetManifest), 42, false)
	if err != nil {
		return err
	}

	return p.storeManifest(kubeClient, deployName, targetManifest)
}

// Update implements update of an existing component instance in the cloud by updating raw k8s objects
func (p *Plugin) Update(deployName string, params util.NestedParameterMap, eventLog *event.Log) error {
	err := p.init()
	if err != nil {
		return err
	}

	kubeClient, err := p.kube.NewClient()
	if err != nil {
		return err
	}

	currentManifest, err := p.loadManifest(kubeClient, deployName)
	if err != nil {
		return err
	}

	targetManifest, ok := params["manifest"].(string)
	if !ok {
		return fmt.Errorf("manifest is a mandatory parameter")
	}

	client := p.prepareClient(eventLog, deployName)

	err = client.Update(p.kube.Namespace, strings.NewReader(currentManifest), strings.NewReader(targetManifest), false, false, 42, false)
	if err != nil {
		return err
	}

	return p.storeManifest(kubeClient, deployName, targetManifest)
}

// Destroy implements destruction of an existing component instance in the cloud by deleting raw k8s objects
func (p *Plugin) Destroy(deployName string, params util.NestedParameterMap, eventLog *event.Log) error {
	err := p.init()
	if err != nil {
		return err
	}

	kubeClient, err := p.kube.NewClient()
	if err != nil {
		return err
	}

	deleteManifest, ok := params["manifest"].(string)
	if !ok {
		return fmt.Errorf("manifest is a mandatory parameter")
	}

	client := p.prepareClient(eventLog, deployName)

	err = client.Delete(p.kube.Namespace, strings.NewReader(deleteManifest))
	if err != nil {
		return err
	}

	return p.deleteManifest(kubeClient, deployName)
}

// Endpoints returns map from port type to url for all services of the deployed raw k8s objects
func (p *Plugin) Endpoints(deployName string, params util.NestedParameterMap, eventLog *event.Log) (map[string]string, error) {
	err := p.init()
	if err != nil {
		return nil, err
	}

	kubeClient, err := p.kube.NewClient()
	if err != nil {
		return nil, err
	}

	targetManifest, ok := params["manifest"].(string)
	if !ok {
		return nil, fmt.Errorf("manifest is a mandatory parameter")
	}

	client := p.prepareClient(eventLog, deployName)

	infos, err := client.BuildUnstructured(p.kube.Namespace, strings.NewReader(targetManifest))
	if err != nil {
		return nil, err
	}

	endpoints := make(map[string]string)

	for _, info := range infos {
		if info.Mapping.GroupVersionKind.Kind == "Service" {
			service, getErr := kubeClient.CoreV1().Services(p.kube.Namespace).Get(info.Name, meta.GetOptions{})
			if getErr != nil {
				return nil, getErr
			}

			p.kube.AddEndpointsFromService(service, endpoints)
		}
	}

	return endpoints, nil
}

// Resources returns list of all resources (like services, config maps, etc.) into the cluster by specified component instance
func (p *Plugin) Resources(deployName string, params util.NestedParameterMap, eventLog *event.Log) (plugin.Resources, error) {
	err := p.init()
	if err != nil {
		return nil, err
	}

	kubeClient, err := p.kube.NewClient()
	if err != nil {
		return nil, err
	}

	targetManifest, ok := params["manifest"].(string)
	if !ok {
		return nil, fmt.Errorf("manifest is a mandatory parameter")
	}

	client := p.prepareClient(eventLog, deployName)

	infos, err := client.BuildUnstructured(p.kube.Namespace, strings.NewReader(targetManifest))
	if err != nil {
		return nil, err
	}

	handlers := make(map[string]ResourceTypeHandler)
	handlers["k8s/v1/Service"] = &serviceResourceTypeHandler{}
	// not sure if it's good to have version.... we could have issues with versions in different k8s clusters
	handlers["k8s/v1/Deployment"] = &deploymentResourceTypeHandler{}

	resources := make(plugin.Resources)
	for _, info := range infos {
		gvk := info.ResourceMapping().GroupVersionKind
		resourceType := "k8s/" + gvk.Version + "/" + gvk.Kind

		handler, exist := handlers[resourceType]
		if !exist {
			continue
		}

		table, exist := resources[resourceType]
		if !exist {
			table = &plugin.ResourceTable{}
			resources[resourceType] = table
			table.Headers = handler.Headers()
		}

		if info.Mapping.GroupVersionKind.Kind == "Service" {
			service, getErr := kubeClient.CoreV1().Services(p.kube.Namespace).Get(info.Name, meta.GetOptions{})
			if getErr != nil {
				return nil, getErr
			}
			table.Items = append(table.Items, handler.Columns(service))
		} else if info.Mapping.GroupVersionKind.Kind == "ConfigMap" {
			configMap, getErr := kubeClient.CoreV1().ConfigMaps(p.kube.Namespace).Get(info.Name, meta.GetOptions{})
			if getErr != nil {
				return nil, getErr
			}
			table.Items = append(table.Items, handler.Columns(configMap))
		} else if info.Mapping.GroupVersionKind.Kind == "Secret" {
			secret, getErr := kubeClient.CoreV1().Secrets(p.kube.Namespace).Get(info.Name, meta.GetOptions{})
			if getErr != nil {
				return nil, getErr
			}
			table.Items = append(table.Items, handler.Columns(secret))
		} else if info.Mapping.GroupVersionKind.Kind == "PersistentVolumeClaim" {
			pvc, getErr := kubeClient.CoreV1().PersistentVolumeClaims(p.kube.Namespace).Get(info.Name, meta.GetOptions{})
			if getErr != nil {
				return nil, getErr
			}
			table.Items = append(table.Items, handler.Columns(pvc))
		} else if info.Mapping.GroupVersionKind.Kind == "Deployment" {
			deployment, getErr := kubeClient.AppsV1beta1().Deployments(p.kube.Namespace).Get(info.Name, meta.GetOptions{})
			if getErr != nil {
				return nil, getErr
			}
			table.Items = append(table.Items, handler.Columns(deployment))
		} else if info.Mapping.GroupVersionKind.Kind == "StatefulSet" {
			statefulSet, getErr := kubeClient.AppsV1beta1().StatefulSets(p.kube.Namespace).Get(info.Name, meta.GetOptions{})
			if getErr != nil {
				return nil, getErr
			}
			table.Items = append(table.Items, handler.Columns(statefulSet))
		}
	}

	return resources, nil
}

// ResourceTypeHandler is an interface for handlers that returns list of headers and columns to represent specified
// object.
type ResourceTypeHandler interface {
	Headers() []string
	Columns(interface{}) []string
}

var serviceResourceHeaders = []string{
	"Namespace",
	"Name",
	"Type",
	"Port(s)",
	"Created",
}

type serviceResourceTypeHandler struct {
}

func (*serviceResourceTypeHandler) Headers() []string {
	return serviceResourceHeaders
}

func (*serviceResourceTypeHandler) Columns(obj interface{}) []string {
	service := obj.(*v1.Service)
	parts := make([]string, len(service.Spec.Ports))
	for idx, port := range service.Spec.Ports {
		if port.NodePort > 0 {
			parts[idx] = fmt.Sprintf("%d:%d/%s", port.Port, port.NodePort, port.Protocol)
		} else {
			parts[idx] = fmt.Sprintf("%d/%s", port.Port, port.Protocol)
		}
		if len(port.Name) > 0 {
			parts[idx] += "(" + port.Name + ")"
		}
	}
	ports := strings.Join(parts, ",")

	return []string{service.Namespace, service.Name, string(service.Spec.Type), ports, service.CreationTimestamp.String()}
}

var deploymentResourceHeaders = []string{
	"Namespace",
	"Name",
	"Desired",
	"Current",
	"Up-to-date",
	"Available",
	"Generation",
	"Created",
}

type deploymentResourceTypeHandler struct {
}

func (*deploymentResourceTypeHandler) Headers() []string {
	return deploymentResourceHeaders
}

func (*deploymentResourceTypeHandler) Columns(obj interface{}) []string {
	deployment := obj.(*v1beta1.Deployment)

	desiredReplicas := fmt.Sprintf("%d", *deployment.Spec.Replicas)
	currentReplicas := fmt.Sprintf("%d", deployment.Status.Replicas)
	updatedReplicas := fmt.Sprintf("%d", deployment.Status.UpdatedReplicas)
	availableReplicas := fmt.Sprintf("%d", deployment.Status.AvailableReplicas)
	gen := fmt.Sprintf("%d", deployment.Generation)
	created := deployment.CreationTimestamp.String()

	return []string{deployment.Namespace, deployment.Name, desiredReplicas, currentReplicas, updatedReplicas, availableReplicas, gen, created}
}

func (p *Plugin) prepareClient(eventLog *event.Log, deployName string) *kube.Client {
	client := kube.New(p.kube.ClientConfig)
	client.Log = func(format string, args ...interface{}) {
		eventLog.WithFields(event.Fields{
			"deployName": deployName,
		}).Debugf(fmt.Sprintf("[instance: %s] ", deployName)+format, args...)
	}

	return client
}
<commit_msg>Add k8s jobs to the list of objects to be included into resources<commit_after>package k8sraw

import (
	"fmt"
	"github.com/Aptomi/aptomi/pkg/config"
	"github.com/Aptomi/aptomi/pkg/event"
	"github.com/Aptomi/aptomi/pkg/lang"
	"github.com/Aptomi/aptomi/pkg/plugin"
	"github.com/Aptomi/aptomi/pkg/plugin/k8s"
	"github.com/Aptomi/aptomi/pkg/util"
	"github.com/Aptomi/aptomi/pkg/util/sync"
	meta "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/client-go/pkg/api/v1"
	"k8s.io/client-go/pkg/apis/apps/v1beta1"
	"k8s.io/helm/pkg/kube"
	"strings"
)

// Plugin represents Kubernetes Raw code plugin that supports deploying specified k8s objects into the cluster
type Plugin struct {
	once          sync.Init
	cluster       *lang.Cluster
	config        config.K8sRaw
	kube          *k8s.Plugin
	dataNamespace string
}

// New returns new instance of the Kubernetes Raw code (objects) plugin for specified Kubernetes cluster plugin and plugins config
func New(clusterPlugin plugin.ClusterPlugin, cfg config.Plugins) (plugin.CodePlugin, error) {
	kubePlugin, ok := clusterPlugin.(*k8s.Plugin)
	if !ok {
		return nil, fmt.Errorf("k8s cluster plugin expected for k8sraw code plugin creation but received: %T", clusterPlugin)
	}

	return &Plugin{
		cluster: kubePlugin.Cluster,
		config:  cfg.K8sRaw,
		kube:    kubePlugin,
	}, nil
}

func (p *Plugin) init() error {
	return p.once.Do(func() error {
		err := p.kube.Init()
		if err != nil {
			return err
		}

		err = p.parseClusterConfig()
		if err != nil {
			return err
		}

		kubeClient, err := p.kube.NewClient()
		if err != nil {
			return err
		}

		return p.kube.EnsureNamespace(kubeClient, p.dataNamespace)
	})
}

// Cleanup implements cleanup phase for the k8s raw plugin
func (p *Plugin) Cleanup() error {
	return nil
}

// Create implements creation of a new component instance in the cloud by deploying raw k8s objects
func (p *Plugin) Create(deployName string, params util.NestedParameterMap, eventLog *event.Log) error {
	err := p.init()
	if err != nil {
		return err
	}

	kubeClient, err := p.kube.NewClient()
	if err != nil {
		return err
	}

	targetManifest, ok := params["manifest"].(string)
	if !ok {
		return fmt.Errorf("manifest is a mandatory parameter")
	}

	client := p.prepareClient(eventLog, deployName)

	err = client.Create(p.kube.Namespace, strings.NewReader(targetManifest), 42, false)
	if err != nil {
		return err
	}

	return p.storeManifest(kubeClient, deployName, targetManifest)
}

// Update implements update of an existing component instance in the cloud by updating raw k8s objects
func (p *Plugin) Update(deployName string, params util.NestedParameterMap, eventLog *event.Log) error {
	err := p.init()
	if err != nil {
		return err
	}

	kubeClient, err := p.kube.NewClient()
	if err != nil {
		return err
	}

	currentManifest, err := p.loadManifest(kubeClient, deployName)
	if err != nil {
		return err
	}

	targetManifest, ok := params["manifest"].(string)
	if !ok {
		return fmt.Errorf("manifest is a mandatory parameter")
	}

	client := p.prepareClient(eventLog, deployName)

	err = client.Update(p.kube.Namespace, strings.NewReader(currentManifest), strings.NewReader(targetManifest), false, false, 42, false)
	if err != nil {
		return err
	}

	return p.storeManifest(kubeClient, deployName, targetManifest)
}

// Destroy implements destruction of an existing component instance in the cloud by deleting raw k8s objects
func (p *Plugin) Destroy(deployName string, params util.NestedParameterMap, eventLog *event.Log) error {
	err := p.init()
	if err != nil {
		return err
	}

	kubeClient, err := p.kube.NewClient()
	if err != nil {
		return err
	}

	deleteManifest, ok := params["manifest"].(string)
	if !ok {
		return fmt.Errorf("manifest is a mandatory parameter")
	}

	client := p.prepareClient(eventLog, deployName)

	err = client.Delete(p.kube.Namespace, strings.NewReader(deleteManifest))
	if err != nil {
		return err
	}

	return p.deleteManifest(kubeClient, deployName)
}

// Endpoints returns map from port type to url for all services of the deployed raw k8s objects
func (p *Plugin) Endpoints(deployName string, params util.NestedParameterMap, eventLog *event.Log) (map[string]string, error) {
	err := p.init()
	if err != nil {
		return nil, err
	}

	kubeClient, err := p.kube.NewClient()
	if err != nil {
		return nil, err
	}

	targetManifest, ok := params["manifest"].(string)
	if !ok {
		return nil, fmt.Errorf("manifest is a mandatory parameter")
	}

	client := p.prepareClient(eventLog, deployName)

	infos, err := client.BuildUnstructured(p.kube.Namespace, strings.NewReader(targetManifest))
	if err != nil {
		return nil, err
	}

	endpoints := make(map[string]string)

	for _, info := range infos {
		if info.Mapping.GroupVersionKind.Kind == "Service" {
			service, getErr := kubeClient.CoreV1().Services(p.kube.Namespace).Get(info.Name, meta.GetOptions{})
			if getErr != nil {
				return nil, getErr
			}

			p.kube.AddEndpointsFromService(service, endpoints)
		}
	}

	return endpoints, nil
}

// Resources returns list of all resources (like services, config maps, etc.) into the cluster by specified component instance
func (p *Plugin) Resources(deployName string, params util.NestedParameterMap, eventLog *event.Log) (plugin.Resources, error) {
	err := p.init()
	if err != nil {
		return nil, err
	}

	kubeClient, err := p.kube.NewClient()
	if err != nil {
		return nil, err
	}

	targetManifest, ok := params["manifest"].(string)
	if !ok {
		return nil, fmt.Errorf("manifest is a mandatory parameter")
	}

	client := p.prepareClient(eventLog, deployName)

	infos, err := client.BuildUnstructured(p.kube.Namespace, strings.NewReader(targetManifest))
	if err != nil {
		return nil, err
	}

	handlers := make(map[string]ResourceTypeHandler)
	handlers["k8s/v1/Service"] = &serviceResourceTypeHandler{}
	// not sure if it's good to have version.... we could have issues with versions in different k8s clusters
	handlers["k8s/v1/Deployment"] = &deploymentResourceTypeHandler{}

	resources := make(plugin.Resources)
	for _, info := range infos {
		gvk := info.ResourceMapping().GroupVersionKind
		resourceType := "k8s/" + gvk.Version + "/" + gvk.Kind

		handler, exist := handlers[resourceType]
		if !exist {
			continue
		}

		table, exist := resources[resourceType]
		if !exist {
			table = &plugin.ResourceTable{}
			resources[resourceType] = table
			table.Headers = handler.Headers()
		}

		if info.Mapping.GroupVersionKind.Kind == "Service" {
			service, getErr := kubeClient.CoreV1().Services(p.kube.Namespace).Get(info.Name, meta.GetOptions{})
			if getErr != nil {
				return nil, getErr
			}
			table.Items = append(table.Items, handler.Columns(service))
		} else if info.Mapping.GroupVersionKind.Kind == "ConfigMap" {
			configMap, getErr := kubeClient.CoreV1().ConfigMaps(p.kube.Namespace).Get(info.Name, meta.GetOptions{})
			if getErr != nil {
				return nil, getErr
			}
			table.Items = append(table.Items, handler.Columns(configMap))
		} else if info.Mapping.GroupVersionKind.Kind == "Secret" {
			secret, getErr := kubeClient.CoreV1().Secrets(p.kube.Namespace).Get(info.Name, meta.GetOptions{})
			if getErr != nil {
				return nil, getErr
			}
			table.Items = append(table.Items, handler.Columns(secret))
		} else if info.Mapping.GroupVersionKind.Kind == "PersistentVolumeClaim" {
			pvc, getErr := kubeClient.CoreV1().PersistentVolumeClaims(p.kube.Namespace).Get(info.Name, meta.GetOptions{})
			if getErr != nil {
				return nil, getErr
			}
			table.Items = append(table.Items, handler.Columns(pvc))
		} else if info.Mapping.GroupVersionKind.Kind == "Deployment" {
			deployment, getErr := kubeClient.AppsV1beta1().Deployments(p.kube.Namespace).Get(info.Name, meta.GetOptions{})
			if getErr != nil {
				return nil, getErr
			}
			table.Items = append(table.Items, handler.Columns(deployment))
		} else if info.Mapping.GroupVersionKind.Kind == "StatefulSet" {
			statefulSet, getErr := kubeClient.AppsV1beta1().StatefulSets(p.kube.Namespace).Get(info.Name, meta.GetOptions{})
			if getErr != nil {
				return nil, getErr
			}
			table.Items = append(table.Items, handler.Columns(statefulSet))
		} else if info.Mapping.GroupVersionKind.Kind == "Job" {
			job, getErr := kubeClient.BatchV1().Jobs(p.kube.Namespace).Get(info.Name, meta.GetOptions{})
			if getErr != nil {
				return nil, getErr
			}
			table.Items = append(table.Items, handler.Columns(job))
		}
	}

	return resources, nil
}

// ResourceTypeHandler is an interface for handlers that returns list of headers and columns to represent specified
// object.
type ResourceTypeHandler interface {
	Headers() []string
	Columns(interface{}) []string
}

var serviceResourceHeaders = []string{
	"Namespace",
	"Name",
	"Type",
	"Port(s)",
	"Created",
}

type serviceResourceTypeHandler struct {
}

func (*serviceResourceTypeHandler) Headers() []string {
	return serviceResourceHeaders
}

func (*serviceResourceTypeHandler) Columns(obj interface{}) []string {
	service := obj.(*v1.Service)
	parts := make([]string, len(service.Spec.Ports))
	for idx, port := range service.Spec.Ports {
		if port.NodePort > 0 {
			parts[idx] = fmt.Sprintf("%d:%d/%s", port.Port, port.NodePort, port.Protocol)
		} else {
			parts[idx] = fmt.Sprintf("%d/%s", port.Port, port.Protocol)
		}
		if len(port.Name) > 0 {
			parts[idx] += "(" + port.Name + ")"
		}
	}
	ports := strings.Join(parts, ",")

	return []string{service.Namespace, service.Name, string(service.Spec.Type), ports, service.CreationTimestamp.String()}
}

var deploymentResourceHeaders = []string{
	"Namespace",
	"Name",
	"Desired",
	"Current",
	"Up-to-date",
	"Available",
	"Generation",
	"Created",
}

type deploymentResourceTypeHandler struct {
}

func (*deploymentResourceTypeHandler) Headers() []string {
	return deploymentResourceHeaders
}

func (*deploymentResourceTypeHandler) Columns(obj interface{}) []string {
	deployment := obj.(*v1beta1.Deployment)

	desiredReplicas := fmt.Sprintf("%d", *deployment.Spec.Replicas)
	currentReplicas := fmt.Sprintf("%d", deployment.Status.Replicas)
	updatedReplicas := fmt.Sprintf("%d", deployment.Status.UpdatedReplicas)
	availableReplicas := fmt.Sprintf("%d", deployment.Status.AvailableReplicas)
	gen := fmt.Sprintf("%d", deployment.Generation)
	created := deployment.CreationTimestamp.String()

	return []string{deployment.Namespace, deployment.Name, desiredReplicas, currentReplicas, updatedReplicas, availableReplicas, gen, created}
}

func (p *Plugin) prepareClient(eventLog *event.Log, deployName string) *kube.Client {
	client := kube.New(p.kube.ClientConfig)
	client.Log = func(format string, args ...interface{}) {
		eventLog.WithFields(event.Fields{
			"deployName": deployName,
		}).Debugf(fmt.Sprintf("[instance: %s] ", deployName)+format, args...)
	}

	return client
}
<|endoftext|>


<commit_before>/*
Copyright 2018 the Velero contributors.

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 restic

import (
	"bytes"
	"encoding/json"
	"fmt"
	"time"

	"github.com/pkg/errors"
	"github.com/sirupsen/logrus"

	velerov1api "github.com/vmware-tanzu/velero/pkg/apis/velero/v1"
	"github.com/vmware-tanzu/velero/pkg/util/exec"
	"github.com/vmware-tanzu/velero/pkg/util/filesystem"
)

const restoreProgressCheckInterval = 10 * time.Second
const backupProgressCheckInterval = 10 * time.Second

var fileSystem = filesystem.NewFileSystem()

type backupStatusLine struct {
	MessageType string `json:"message_type"`
	// seen in status lines
	TotalBytes int64 `json:"total_bytes"`
	BytesDone  int64 `json:"bytes_done"`
	// seen in summary line at the end
	TotalBytesProcessed int64 `json:"total_bytes_processed"`
}

// GetSnapshotID runs a 'restic snapshots' command to get the ID of the snapshot
// in the specified repo matching the set of provided tags, or an error if a
// unique snapshot cannot be identified.
func GetSnapshotID(repoIdentifier, passwordFile string, tags map[string]string, env []string, caCertFile string) (string, error) {
	cmd := GetSnapshotCommand(repoIdentifier, passwordFile, tags)
	if len(env) > 0 {
		cmd.Env = env
	}
	cmd.CACertFile = caCertFile

	stdout, stderr, err := exec.RunCommand(cmd.Cmd())
	if err != nil {
		return "", errors.Wrapf(err, "error running command, stderr=%s", stderr)
	}

	type snapshotID struct {
		ShortID string `json:"short_id"`
	}

	var snapshots []snapshotID
	if err := json.Unmarshal([]byte(stdout), &snapshots); err != nil {
		return "", errors.Wrap(err, "error unmarshalling restic snapshots result")
	}

	if len(snapshots) != 1 {
		return "", errors.Errorf("expected one matching snapshot, got %d", len(snapshots))
	}

	return snapshots[0].ShortID, nil
}

// RunBackup runs a `restic backup` command and watches the output to provide
// progress updates to the caller.
func RunBackup(backupCmd *Command, log logrus.FieldLogger, updateFunc func(velerov1api.PodVolumeOperationProgress)) (string, string, error) {
	// buffers for copying command stdout/err output into
	stdoutBuf := new(bytes.Buffer)
	stderrBuf := new(bytes.Buffer)

	// create a channel to signal when to end the goroutine scanning for progress
	// updates
	quit := make(chan struct{})

	cmd := backupCmd.Cmd()
	cmd.Stdout = stdoutBuf
	cmd.Stderr = stderrBuf

	cmd.Start()

	go func() {
		ticker := time.NewTicker(backupProgressCheckInterval)
		for {
			select {
			case <-ticker.C:
				lastLine := getLastLine(stdoutBuf.Bytes())
				if len(lastLine) > 0 {
					stat, err := decodeBackupStatusLine(lastLine)
					if err != nil {
						log.WithError(err).Errorf("error getting restic backup progress")
					}

					// if the line contains a non-empty bytes_done field, we can update the
					// caller with the progress
					if stat.BytesDone != 0 {
						updateFunc(velerov1api.PodVolumeOperationProgress{
							TotalBytes: stat.TotalBytes,
							BytesDone:  stat.BytesDone,
						})
					}
				}
			case <-quit:
				ticker.Stop()
				return
			}
		}
	}()

	cmd.Wait()
	quit <- struct{}{}

	summary, err := getSummaryLine(stdoutBuf.Bytes())
	if err != nil {
		return stdoutBuf.String(), stderrBuf.String(), err
	}
	stat, err := decodeBackupStatusLine(summary)
	if err != nil {
		return stdoutBuf.String(), stderrBuf.String(), err
	}
	if stat.MessageType != "summary" {
		return stdoutBuf.String(), stderrBuf.String(), errors.WithStack(fmt.Errorf("error getting restic backup summary: %s", string(summary)))
	}

	// update progress to 100%
	updateFunc(velerov1api.PodVolumeOperationProgress{
		TotalBytes: stat.TotalBytesProcessed,
		BytesDone:  stat.TotalBytesProcessed,
	})

	return string(summary), stderrBuf.String(), nil
}

func decodeBackupStatusLine(lastLine []byte) (backupStatusLine, error) {
	var stat backupStatusLine
	if err := json.Unmarshal(lastLine, &stat); err != nil {
		return stat, errors.Wrapf(err, "unable to decode backup JSON line: %s", string(lastLine))
	}
	return stat, nil
}

// getLastLine returns the last line of a byte array. The string is assumed to
// have a newline at the end of it, so this returns the substring between the
// last two newlines.
func getLastLine(b []byte) []byte {
	if b == nil || len(b) == 0 {
		return []byte("")
	}
	// subslice the byte array to ignore the newline at the end of the string
	lastNewLineIdx := bytes.LastIndex(b[:len(b)-1], []byte("\n"))
	return b[lastNewLineIdx+1 : len(b)-1]
}

// getSummaryLine looks for the summary JSON line
// (`{"message_type:"summary",...`) in the restic backup command output. Due to
// an issue in Restic, this might not always be the last line
// (https://github.com/restic/restic/issues/2389). It returns an error if it
// can't be found.
func getSummaryLine(b []byte) ([]byte, error) {
	summaryLineIdx := bytes.LastIndex(b, []byte(`{"message_type":"summary"`))
	if summaryLineIdx < 0 {
		return nil, errors.New("unable to find summary in restic backup command output")
	}
	// find the end of the summary line
	newLineIdx := bytes.Index(b[summaryLineIdx:], []byte("\n"))
	if newLineIdx < 0 {
		return nil, errors.New("unable to get summary line from restic backup command output")
	}
	return b[summaryLineIdx : summaryLineIdx+newLineIdx], nil
}

// RunRestore runs a `restic restore` command and monitors the volume size to
// provide progress updates to the caller.
func RunRestore(restoreCmd *Command, log logrus.FieldLogger, updateFunc func(velerov1api.PodVolumeOperationProgress)) (string, string, error) {
	snapshotSize, err := getSnapshotSize(restoreCmd.RepoIdentifier, restoreCmd.PasswordFile, restoreCmd.Args[0], restoreCmd.Env)
	if err != nil {
		return "", "", errors.Wrap(err, "error getting snapshot size")
	}

	updateFunc(velerov1api.PodVolumeOperationProgress{
		TotalBytes: snapshotSize,
	})

	// create a channel to signal when to end the goroutine scanning for progress
	// updates
	quit := make(chan struct{})

	go func() {
		ticker := time.NewTicker(restoreProgressCheckInterval)
		for {
			select {
			case <-ticker.C:
				volumeSize, err := getVolumeSize(restoreCmd.Dir)
				if err != nil {
					log.WithError(err).Errorf("error getting restic restore progress")
				}

				updateFunc(velerov1api.PodVolumeOperationProgress{
					TotalBytes: snapshotSize,
					BytesDone:  volumeSize,
				})
			case <-quit:
				ticker.Stop()
				return
			}
		}
	}()

	stdout, stderr, err := exec.RunCommand(restoreCmd.Cmd())
	quit <- struct{}{}

	// update progress to 100%
	updateFunc(velerov1api.PodVolumeOperationProgress{
		TotalBytes: snapshotSize,
		BytesDone:  snapshotSize,
	})

	return stdout, stderr, err
}

func getSnapshotSize(repoIdentifier, passwordFile, snapshotID string, env []string) (int64, error) {
	cmd := StatsCommand(repoIdentifier, passwordFile, snapshotID)
	cmd.Env = env

	stdout, stderr, err := exec.RunCommand(cmd.Cmd())
	if err != nil {
		return 0, errors.Wrapf(err, "error running command, stderr=%s", stderr)
	}

	var snapshotStats struct {
		TotalSize int64 `json:"total_size"`
	}

	if err := json.Unmarshal([]byte(stdout), &snapshotStats); err != nil {
		return 0, errors.Wrap(err, "error unmarshalling restic stats result")
	}

	return snapshotStats.TotalSize, nil
}

func getVolumeSize(path string) (int64, error) {
	var size int64

	files, err := fileSystem.ReadDir(path)
	if err != nil {
		return 0, errors.Wrapf(err, "error reading directory %s", path)
	}

	for _, file := range files {
		if file.IsDir() {
			s, err := getVolumeSize(fmt.Sprintf("%s/%s", path, file.Name()))
			if err != nil {
				return 0, err
			}
			size += s
		} else {
			size += file.Size()
		}
	}

	return size, nil
}
<commit_msg>bug fix: pass ca cert file to restic stats command on restore<commit_after>/*
Copyright 2018 the Velero contributors.

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 restic

import (
	"bytes"
	"encoding/json"
	"fmt"
	"time"

	"github.com/pkg/errors"
	"github.com/sirupsen/logrus"

	velerov1api "github.com/vmware-tanzu/velero/pkg/apis/velero/v1"
	"github.com/vmware-tanzu/velero/pkg/util/exec"
	"github.com/vmware-tanzu/velero/pkg/util/filesystem"
)

const restoreProgressCheckInterval = 10 * time.Second
const backupProgressCheckInterval = 10 * time.Second

var fileSystem = filesystem.NewFileSystem()

type backupStatusLine struct {
	MessageType string `json:"message_type"`
	// seen in status lines
	TotalBytes int64 `json:"total_bytes"`
	BytesDone  int64 `json:"bytes_done"`
	// seen in summary line at the end
	TotalBytesProcessed int64 `json:"total_bytes_processed"`
}

// GetSnapshotID runs a 'restic snapshots' command to get the ID of the snapshot
// in the specified repo matching the set of provided tags, or an error if a
// unique snapshot cannot be identified.
func GetSnapshotID(repoIdentifier, passwordFile string, tags map[string]string, env []string, caCertFile string) (string, error) {
	cmd := GetSnapshotCommand(repoIdentifier, passwordFile, tags)
	if len(env) > 0 {
		cmd.Env = env
	}
	cmd.CACertFile = caCertFile

	stdout, stderr, err := exec.RunCommand(cmd.Cmd())
	if err != nil {
		return "", errors.Wrapf(err, "error running command, stderr=%s", stderr)
	}

	type snapshotID struct {
		ShortID string `json:"short_id"`
	}

	var snapshots []snapshotID
	if err := json.Unmarshal([]byte(stdout), &snapshots); err != nil {
		return "", errors.Wrap(err, "error unmarshalling restic snapshots result")
	}

	if len(snapshots) != 1 {
		return "", errors.Errorf("expected one matching snapshot, got %d", len(snapshots))
	}

	return snapshots[0].ShortID, nil
}

// RunBackup runs a `restic backup` command and watches the output to provide
// progress updates to the caller.
func RunBackup(backupCmd *Command, log logrus.FieldLogger, updateFunc func(velerov1api.PodVolumeOperationProgress)) (string, string, error) {
	// buffers for copying command stdout/err output into
	stdoutBuf := new(bytes.Buffer)
	stderrBuf := new(bytes.Buffer)

	// create a channel to signal when to end the goroutine scanning for progress
	// updates
	quit := make(chan struct{})

	cmd := backupCmd.Cmd()
	cmd.Stdout = stdoutBuf
	cmd.Stderr = stderrBuf

	cmd.Start()

	go func() {
		ticker := time.NewTicker(backupProgressCheckInterval)
		for {
			select {
			case <-ticker.C:
				lastLine := getLastLine(stdoutBuf.Bytes())
				if len(lastLine) > 0 {
					stat, err := decodeBackupStatusLine(lastLine)
					if err != nil {
						log.WithError(err).Errorf("error getting restic backup progress")
					}

					// if the line contains a non-empty bytes_done field, we can update the
					// caller with the progress
					if stat.BytesDone != 0 {
						updateFunc(velerov1api.PodVolumeOperationProgress{
							TotalBytes: stat.TotalBytes,
							BytesDone:  stat.BytesDone,
						})
					}
				}
			case <-quit:
				ticker.Stop()
				return
			}
		}
	}()

	cmd.Wait()
	quit <- struct{}{}

	summary, err := getSummaryLine(stdoutBuf.Bytes())
	if err != nil {
		return stdoutBuf.String(), stderrBuf.String(), err
	}
	stat, err := decodeBackupStatusLine(summary)
	if err != nil {
		return stdoutBuf.String(), stderrBuf.String(), err
	}
	if stat.MessageType != "summary" {
		return stdoutBuf.String(), stderrBuf.String(), errors.WithStack(fmt.Errorf("error getting restic backup summary: %s", string(summary)))
	}

	// update progress to 100%
	updateFunc(velerov1api.PodVolumeOperationProgress{
		TotalBytes: stat.TotalBytesProcessed,
		BytesDone:  stat.TotalBytesProcessed,
	})

	return string(summary), stderrBuf.String(), nil
}

func decodeBackupStatusLine(lastLine []byte) (backupStatusLine, error) {
	var stat backupStatusLine
	if err := json.Unmarshal(lastLine, &stat); err != nil {
		return stat, errors.Wrapf(err, "unable to decode backup JSON line: %s", string(lastLine))
	}
	return stat, nil
}

// getLastLine returns the last line of a byte array. The string is assumed to
// have a newline at the end of it, so this returns the substring between the
// last two newlines.
func getLastLine(b []byte) []byte {
	if b == nil || len(b) == 0 {
		return []byte("")
	}
	// subslice the byte array to ignore the newline at the end of the string
	lastNewLineIdx := bytes.LastIndex(b[:len(b)-1], []byte("\n"))
	return b[lastNewLineIdx+1 : len(b)-1]
}

// getSummaryLine looks for the summary JSON line
// (`{"message_type:"summary",...`) in the restic backup command output. Due to
// an issue in Restic, this might not always be the last line
// (https://github.com/restic/restic/issues/2389). It returns an error if it
// can't be found.
func getSummaryLine(b []byte) ([]byte, error) {
	summaryLineIdx := bytes.LastIndex(b, []byte(`{"message_type":"summary"`))
	if summaryLineIdx < 0 {
		return nil, errors.New("unable to find summary in restic backup command output")
	}
	// find the end of the summary line
	newLineIdx := bytes.Index(b[summaryLineIdx:], []byte("\n"))
	if newLineIdx < 0 {
		return nil, errors.New("unable to get summary line from restic backup command output")
	}
	return b[summaryLineIdx : summaryLineIdx+newLineIdx], nil
}

// RunRestore runs a `restic restore` command and monitors the volume size to
// provide progress updates to the caller.
func RunRestore(restoreCmd *Command, log logrus.FieldLogger, updateFunc func(velerov1api.PodVolumeOperationProgress)) (string, string, error) {
	snapshotSize, err := getSnapshotSize(restoreCmd.RepoIdentifier, restoreCmd.PasswordFile, restoreCmd.CACertFile, restoreCmd.Args[0], restoreCmd.Env)
	if err != nil {
		return "", "", errors.Wrap(err, "error getting snapshot size")
	}

	updateFunc(velerov1api.PodVolumeOperationProgress{
		TotalBytes: snapshotSize,
	})

	// create a channel to signal when to end the goroutine scanning for progress
	// updates
	quit := make(chan struct{})

	go func() {
		ticker := time.NewTicker(restoreProgressCheckInterval)
		for {
			select {
			case <-ticker.C:
				volumeSize, err := getVolumeSize(restoreCmd.Dir)
				if err != nil {
					log.WithError(err).Errorf("error getting restic restore progress")
				}

				updateFunc(velerov1api.PodVolumeOperationProgress{
					TotalBytes: snapshotSize,
					BytesDone:  volumeSize,
				})
			case <-quit:
				ticker.Stop()
				return
			}
		}
	}()

	stdout, stderr, err := exec.RunCommand(restoreCmd.Cmd())
	quit <- struct{}{}

	// update progress to 100%
	updateFunc(velerov1api.PodVolumeOperationProgress{
		TotalBytes: snapshotSize,
		BytesDone:  snapshotSize,
	})

	return stdout, stderr, err
}

func getSnapshotSize(repoIdentifier, passwordFile, caCertFile, snapshotID string, env []string) (int64, error) {
	cmd := StatsCommand(repoIdentifier, passwordFile, snapshotID)
	cmd.Env = env
	cmd.CACertFile = caCertFile

	stdout, stderr, err := exec.RunCommand(cmd.Cmd())
	if err != nil {
		return 0, errors.Wrapf(err, "error running command, stderr=%s", stderr)
	}

	var snapshotStats struct {
		TotalSize int64 `json:"total_size"`
	}

	if err := json.Unmarshal([]byte(stdout), &snapshotStats); err != nil {
		return 0, errors.Wrap(err, "error unmarshalling restic stats result")
	}

	return snapshotStats.TotalSize, nil
}

func getVolumeSize(path string) (int64, error) {
	var size int64

	files, err := fileSystem.ReadDir(path)
	if err != nil {
		return 0, errors.Wrapf(err, "error reading directory %s", path)
	}

	for _, file := range files {
		if file.IsDir() {
			s, err := getVolumeSize(fmt.Sprintf("%s/%s", path, file.Name()))
			if err != nil {
				return 0, err
			}
			size += s
		} else {
			size += file.Size()
		}
	}

	return size, nil
}
<|endoftext|>


<commit_before>package server

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"strings"
	"time"

	"github.com/facette/facette/pkg/config"
	"github.com/facette/facette/pkg/connector"
	"github.com/facette/facette/pkg/library"
	"github.com/facette/facette/pkg/logger"
	"github.com/facette/facette/pkg/types"
	"github.com/facette/facette/pkg/utils"
	"github.com/facette/facette/thirdparty/github.com/fatih/set"
)

func (server *Server) serveGraph(writer http.ResponseWriter, request *http.Request) {
	graphID := strings.TrimPrefix(request.URL.Path, urlLibraryPath+"graphs/")

	switch request.Method {
	case "DELETE":
		if graphID == "" {
			server.serveResponse(writer, serverResponse{mesgMethodNotAllowed}, http.StatusMethodNotAllowed)
			return
		}

		err := server.Library.DeleteItem(graphID, library.LibraryItemGraph)
		if os.IsNotExist(err) {
			server.serveResponse(writer, serverResponse{mesgResourceNotFound}, http.StatusNotFound)
			return
		} else if err != nil {
			logger.Log(logger.LevelError, "server", "%s", err)
			server.serveResponse(writer, serverResponse{mesgUnhandledError}, http.StatusInternalServerError)
			return
		}

		server.serveResponse(writer, nil, http.StatusOK)

	case "GET", "HEAD":
		if graphID == "" {
			server.serveGraphList(writer, request)
			return
		}

		item, err := server.Library.GetItem(graphID, library.LibraryItemGraph)
		if os.IsNotExist(err) {
			server.serveResponse(writer, serverResponse{mesgResourceNotFound}, http.StatusNotFound)
			return
		} else if err != nil {
			logger.Log(logger.LevelError, "server", "%s", err)
			server.serveResponse(writer, serverResponse{mesgUnhandledError}, http.StatusInternalServerError)
			return
		}

		server.serveResponse(writer, item, http.StatusOK)

	case "POST", "PUT":
		var graph *library.Graph

		if response, status := server.parseStoreRequest(writer, request, graphID); status != http.StatusOK {
			server.serveResponse(writer, response, status)
			return
		}

		if request.Method == "POST" && request.FormValue("inherit") != "" {
			// Get graph from library
			item, err := server.Library.GetItem(request.FormValue("inherit"), library.LibraryItemGraph)
			if os.IsNotExist(err) {
				server.serveResponse(writer, serverResponse{mesgResourceNotFound}, http.StatusNotFound)
				return
			} else if err != nil {
				logger.Log(logger.LevelError, "server", "%s", err)
				server.serveResponse(writer, serverResponse{mesgUnhandledError}, http.StatusInternalServerError)
				return
			}

			graph = &library.Graph{}
			utils.Clone(item.(*library.Graph), graph)

			graph.ID = ""
		} else {
			// Create a new graph instance
			graph = &library.Graph{Item: library.Item{ID: graphID}}
		}

		graph.Modified = time.Now()

		// Parse input JSON for graph data
		body, _ := ioutil.ReadAll(request.Body)

		if err := json.Unmarshal(body, graph); err != nil {
			logger.Log(logger.LevelError, "server", "%s", err)
			server.serveResponse(writer, serverResponse{mesgResourceInvalid}, http.StatusBadRequest)
			return
		}

		err := server.Library.StoreItem(graph, library.LibraryItemGraph)
		if response, status := server.parseError(writer, request, err); status != http.StatusOK {
			logger.Log(logger.LevelError, "server", "%s", err)
			server.serveResponse(writer, response, status)
			return
		}

		if request.Method == "POST" {
			writer.Header().Add("Location", strings.TrimRight(request.URL.Path, "/")+"/"+graph.ID)
			server.serveResponse(writer, nil, http.StatusCreated)
		} else {
			server.serveResponse(writer, nil, http.StatusOK)
		}

	default:
		server.serveResponse(writer, serverResponse{mesgMethodNotAllowed}, http.StatusMethodNotAllowed)
	}
}

func (server *Server) serveGraphList(writer http.ResponseWriter, request *http.Request) {
	var offset, limit int

	if response, status := server.parseListRequest(writer, request, &offset, &limit); status != http.StatusOK {
		server.serveResponse(writer, response, status)
		return
	}

	graphSet := set.New(set.ThreadSafe)

	// Filter on collection if any
	if request.FormValue("collection") != "" {
		item, err := server.Library.GetItem(request.FormValue("collection"), library.LibraryItemCollection)
		if os.IsNotExist(err) {
			server.serveResponse(writer, serverResponse{mesgResourceNotFound}, http.StatusNotFound)
			return
		} else if err != nil {
			logger.Log(logger.LevelError, "server", "%s", err)
			server.serveResponse(writer, serverResponse{mesgUnhandledError}, http.StatusInternalServerError)
			return
		}

		collection := item.(*library.Collection)

		for _, graph := range collection.Entries {
			graphSet.Add(graph.ID)
		}
	}

	// Fill graphs list
	items := make(ItemListResponse, 0)

	for _, graph := range server.Library.Graphs {
		if !graphSet.IsEmpty() && !graphSet.Has(graph.ID) {
			continue
		}

		if request.FormValue("filter") != "" && !utils.FilterMatch(request.FormValue("filter"), graph.Name) {
			continue
		}

		items = append(items, &ItemResponse{
			ID:          graph.ID,
			Name:        graph.Name,
			Description: graph.Description,
			Modified:    graph.Modified.Format(time.RFC3339),
		})
	}

	response := &listResponse{
		list:   items,
		offset: offset,
		limit:  limit,
	}

	server.applyResponseLimit(writer, request, response)

	server.serveResponse(writer, response.list, http.StatusOK)
}

func (server *Server) serveGraphPlots(writer http.ResponseWriter, request *http.Request) {
	var (
		err                error
		graph              *library.Graph
		item               interface{}
		startTime, endTime time.Time
	)

	if request.Method != "POST" && request.Method != "HEAD" {
		server.serveResponse(writer, serverResponse{mesgMethodNotAllowed}, http.StatusMethodNotAllowed)
		return
	} else if utils.HTTPGetContentType(request) != "application/json" {
		server.serveResponse(writer, serverResponse{mesgUnsupportedMediaType}, http.StatusUnsupportedMediaType)
		return
	}

	// Parse input JSON for graph data
	body, _ := ioutil.ReadAll(request.Body)

	plotReq := PlotRequest{}

	if err := json.Unmarshal(body, &plotReq); err != nil {
		logger.Log(logger.LevelError, "server", "%s", err)
		server.serveResponse(writer, serverResponse{mesgResourceInvalid}, http.StatusBadRequest)
		return
	}

	if plotReq.Time == "" {
		endTime = time.Now()
	} else if strings.HasPrefix(strings.Trim(plotReq.Range, " "), "-") {
		if endTime, err = time.Parse(time.RFC3339, plotReq.Time); err != nil {
			logger.Log(logger.LevelError, "server", "%s", err)
			server.serveResponse(writer, serverResponse{mesgResourceInvalid}, http.StatusBadRequest)
			return
		}
	} else {
		if startTime, err = time.Parse(time.RFC3339, plotReq.Time); err != nil {
			logger.Log(logger.LevelError, "server", "%s", err)
			server.serveResponse(writer, serverResponse{mesgResourceInvalid}, http.StatusBadRequest)
			return
		}
	}

	if startTime.IsZero() {
		if startTime, err = utils.TimeApplyRange(endTime, plotReq.Range); err != nil {
			logger.Log(logger.LevelError, "server", "%s", err)
			server.serveResponse(writer, serverResponse{mesgResourceInvalid}, http.StatusBadRequest)
			return
		}
	} else if endTime, err = utils.TimeApplyRange(startTime, plotReq.Range); err != nil {
		logger.Log(logger.LevelError, "server", "%s", err)
		server.serveResponse(writer, serverResponse{mesgResourceInvalid}, http.StatusBadRequest)
		return
	}

	if plotReq.Sample == 0 {
		plotReq.Sample = config.DefaultPlotSample
	}

	// Get graph from library
	graph = plotReq.Graph

	if plotReq.ID != "" {
		if item, err = server.Library.GetItem(plotReq.ID, library.LibraryItemGraph); err == nil {
			graph = item.(*library.Graph)
		}
	}

	if graph == nil {
		err = os.ErrNotExist
	}

	if err != nil {
		if os.IsNotExist(err) {
			server.serveResponse(writer, serverResponse{mesgResourceNotFound}, http.StatusNotFound)
		} else {
			logger.Log(logger.LevelError, "server", "%s", err)
			server.serveResponse(writer, serverResponse{mesgUnhandledError}, http.StatusInternalServerError)
		}

		return
	}

	step := endTime.Sub(startTime) / time.Duration(plotReq.Sample)

	// Get plots data
	groupOptions := make(map[string]map[string]interface{})

	data := make([]map[string]*types.PlotResult, 0)

	for _, groupItem := range graph.Groups {
		groupOptions[groupItem.Name] = groupItem.Options

		query, providerConnector, err := server.preparePlotQuery(&plotReq, groupItem)
		if err != nil {
			if err != os.ErrInvalid {
				logger.Log(logger.LevelError, "server", "%s", err)
			}

			data = append(data, nil)
			continue
		}

		plotResult, err := providerConnector.GetPlots(&types.PlotQuery{query, startTime, endTime, step,
			plotReq.Percentiles})
		if err != nil {
			logger.Log(logger.LevelError, "server", "%s", err)
		}

		data = append(data, plotResult)
	}

	response := &PlotResponse{
		ID:          graph.ID,
		Start:       startTime.Format(time.RFC3339),
		End:         endTime.Format(time.RFC3339),
		Step:        step.Seconds(),
		Name:        graph.Name,
		Description: graph.Description,
		Type:        graph.Type,
		StackMode:   graph.StackMode,
		UnitLabel:   graph.UnitLabel,
		UnitType:    graph.UnitType,
		Modified:    graph.Modified,
	}

	if len(data) == 0 {
		server.serveResponse(writer, serverResponse{mesgEmptyData}, http.StatusOK)
		return
	}

	plotMax := 0

	for _, groupItem := range graph.Groups {
		var plotResult map[string]*types.PlotResult

		plotResult, data = data[0], data[1:]

		for serieName, serieResult := range plotResult {
			if len(serieResult.Plots) > plotMax {
				plotMax = len(serieResult.Plots)
			}

			response.Series = append(response.Series, &SerieResponse{
				Name:    serieName,
				Plots:   serieResult.Plots,
				Info:    serieResult.Info,
				Options: groupOptions[groupItem.Name],
			})
		}
	}

	if plotMax > 0 {
		response.Step = (endTime.Sub(startTime) / time.Duration(plotMax)).Seconds()
	}

	server.serveResponse(writer, response, http.StatusOK)
}

func (server *Server) preparePlotQuery(plotReq *PlotRequest, groupItem *library.OperGroup) (*types.GroupQuery,
	connector.Connector, error) {
	var providerConnector connector.Connector

	query := &types.GroupQuery{
		Name:  groupItem.Name,
		Type:  groupItem.Type,
		Scale: groupItem.Scale,
	}

	for _, serieItem := range groupItem.Series {
		// Check for connectors errors or conflicts
		if _, ok := server.Catalog.Origins[serieItem.Origin]; !ok {
			return nil, nil, fmt.Errorf("unknown serie origin `%s'", serieItem.Origin)
		}

		serieSources := make([]string, 0)

		if strings.HasPrefix(serieItem.Source, library.LibraryGroupPrefix) {
			serieSources = server.Library.ExpandGroup(
				strings.TrimPrefix(serieItem.Source, library.LibraryGroupPrefix),
				library.LibraryItemSourceGroup,
			)
		} else {
			serieSources = []string{serieItem.Source}
		}

		index := 0

		for _, serieSource := range serieSources {
			if strings.HasPrefix(serieItem.Metric, library.LibraryGroupPrefix) {
				for _, serieChunk := range server.Library.ExpandGroup(
					strings.TrimPrefix(serieItem.Metric, library.LibraryGroupPrefix),
					library.LibraryItemMetricGroup,
				) {
					metric := server.Catalog.GetMetric(serieItem.Origin, serieSource, serieChunk)

					if metric == nil {
						logger.Log(
							logger.LevelError,
							"server",
							"unknown metric `%s' for source `%s' (origin: %s)",
							serieChunk,
							serieSource,
							serieItem.Origin,
						)

						continue
					}

					if providerConnector == nil {
						providerConnector = metric.Connector.(connector.Connector)
					} else if providerConnector != metric.Connector.(connector.Connector) {
						return nil, nil, fmt.Errorf("connectors differ between series")
					}

					query.Series = append(query.Series, &types.SerieQuery{
						Name: fmt.Sprintf("%s-%d", serieItem.Name, index),
						Metric: &types.MetricQuery{
							Name:   metric.OriginalName,
							Origin: metric.Source.Origin.OriginalName,
							Source: metric.Source.OriginalName,
						},
						Scale: serieItem.Scale,
					})

					index += 1
				}
			} else {
				metric := server.Catalog.GetMetric(serieItem.Origin, serieSource, serieItem.Metric)

				if metric == nil {
					logger.Log(
						logger.LevelError,
						"server",
						"unknown metric `%s' for source `%s' (origin: %s)",
						serieItem.Metric,
						serieSource,
						serieItem.Origin,
					)

					continue
				}

				if providerConnector == nil {
					providerConnector = metric.Connector.(connector.Connector)
				} else if providerConnector != metric.Connector.(connector.Connector) {
					return nil, nil, fmt.Errorf("connectors differ between series")
				}

				serie := &types.SerieQuery{
					Metric: &types.MetricQuery{
						Name:   metric.OriginalName,
						Origin: metric.Source.Origin.OriginalName,
						Source: metric.Source.OriginalName,
					},
					Scale: serieItem.Scale,
				}

				if len(serieSources) > 1 {
					serie.Name = fmt.Sprintf("%s-%d", serieItem.Name, index)
				} else {
					serie.Name = serieItem.Name
				}

				query.Series = append(query.Series, serie)

				index += 1
			}
		}
	}

	if len(query.Series) == 0 {
		return nil, nil, os.ErrInvalid
	}

	return query, providerConnector, nil
}
<commit_msg>Lower unknown metric messages criticity<commit_after>package server

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"strings"
	"time"

	"github.com/facette/facette/pkg/config"
	"github.com/facette/facette/pkg/connector"
	"github.com/facette/facette/pkg/library"
	"github.com/facette/facette/pkg/logger"
	"github.com/facette/facette/pkg/types"
	"github.com/facette/facette/pkg/utils"
	"github.com/facette/facette/thirdparty/github.com/fatih/set"
)

func (server *Server) serveGraph(writer http.ResponseWriter, request *http.Request) {
	graphID := strings.TrimPrefix(request.URL.Path, urlLibraryPath+"graphs/")

	switch request.Method {
	case "DELETE":
		if graphID == "" {
			server.serveResponse(writer, serverResponse{mesgMethodNotAllowed}, http.StatusMethodNotAllowed)
			return
		}

		err := server.Library.DeleteItem(graphID, library.LibraryItemGraph)
		if os.IsNotExist(err) {
			server.serveResponse(writer, serverResponse{mesgResourceNotFound}, http.StatusNotFound)
			return
		} else if err != nil {
			logger.Log(logger.LevelError, "server", "%s", err)
			server.serveResponse(writer, serverResponse{mesgUnhandledError}, http.StatusInternalServerError)
			return
		}

		server.serveResponse(writer, nil, http.StatusOK)

	case "GET", "HEAD":
		if graphID == "" {
			server.serveGraphList(writer, request)
			return
		}

		item, err := server.Library.GetItem(graphID, library.LibraryItemGraph)
		if os.IsNotExist(err) {
			server.serveResponse(writer, serverResponse{mesgResourceNotFound}, http.StatusNotFound)
			return
		} else if err != nil {
			logger.Log(logger.LevelError, "server", "%s", err)
			server.serveResponse(writer, serverResponse{mesgUnhandledError}, http.StatusInternalServerError)
			return
		}

		server.serveResponse(writer, item, http.StatusOK)

	case "POST", "PUT":
		var graph *library.Graph

		if response, status := server.parseStoreRequest(writer, request, graphID); status != http.StatusOK {
			server.serveResponse(writer, response, status)
			return
		}

		if request.Method == "POST" && request.FormValue("inherit") != "" {
			// Get graph from library
			item, err := server.Library.GetItem(request.FormValue("inherit"), library.LibraryItemGraph)
			if os.IsNotExist(err) {
				server.serveResponse(writer, serverResponse{mesgResourceNotFound}, http.StatusNotFound)
				return
			} else if err != nil {
				logger.Log(logger.LevelError, "server", "%s", err)
				server.serveResponse(writer, serverResponse{mesgUnhandledError}, http.StatusInternalServerError)
				return
			}

			graph = &library.Graph{}
			utils.Clone(item.(*library.Graph), graph)

			graph.ID = ""
		} else {
			// Create a new graph instance
			graph = &library.Graph{Item: library.Item{ID: graphID}}
		}

		graph.Modified = time.Now()

		// Parse input JSON for graph data
		body, _ := ioutil.ReadAll(request.Body)

		if err := json.Unmarshal(body, graph); err != nil {
			logger.Log(logger.LevelError, "server", "%s", err)
			server.serveResponse(writer, serverResponse{mesgResourceInvalid}, http.StatusBadRequest)
			return
		}

		err := server.Library.StoreItem(graph, library.LibraryItemGraph)
		if response, status := server.parseError(writer, request, err); status != http.StatusOK {
			logger.Log(logger.LevelError, "server", "%s", err)
			server.serveResponse(writer, response, status)
			return
		}

		if request.Method == "POST" {
			writer.Header().Add("Location", strings.TrimRight(request.URL.Path, "/")+"/"+graph.ID)
			server.serveResponse(writer, nil, http.StatusCreated)
		} else {
			server.serveResponse(writer, nil, http.StatusOK)
		}

	default:
		server.serveResponse(writer, serverResponse{mesgMethodNotAllowed}, http.StatusMethodNotAllowed)
	}
}

func (server *Server) serveGraphList(writer http.ResponseWriter, request *http.Request) {
	var offset, limit int

	if response, status := server.parseListRequest(writer, request, &offset, &limit); status != http.StatusOK {
		server.serveResponse(writer, response, status)
		return
	}

	graphSet := set.New(set.ThreadSafe)

	// Filter on collection if any
	if request.FormValue("collection") != "" {
		item, err := server.Library.GetItem(request.FormValue("collection"), library.LibraryItemCollection)
		if os.IsNotExist(err) {
			server.serveResponse(writer, serverResponse{mesgResourceNotFound}, http.StatusNotFound)
			return
		} else if err != nil {
			logger.Log(logger.LevelError, "server", "%s", err)
			server.serveResponse(writer, serverResponse{mesgUnhandledError}, http.StatusInternalServerError)
			return
		}

		collection := item.(*library.Collection)

		for _, graph := range collection.Entries {
			graphSet.Add(graph.ID)
		}
	}

	// Fill graphs list
	items := make(ItemListResponse, 0)

	for _, graph := range server.Library.Graphs {
		if !graphSet.IsEmpty() && !graphSet.Has(graph.ID) {
			continue
		}

		if request.FormValue("filter") != "" && !utils.FilterMatch(request.FormValue("filter"), graph.Name) {
			continue
		}

		items = append(items, &ItemResponse{
			ID:          graph.ID,
			Name:        graph.Name,
			Description: graph.Description,
			Modified:    graph.Modified.Format(time.RFC3339),
		})
	}

	response := &listResponse{
		list:   items,
		offset: offset,
		limit:  limit,
	}

	server.applyResponseLimit(writer, request, response)

	server.serveResponse(writer, response.list, http.StatusOK)
}

func (server *Server) serveGraphPlots(writer http.ResponseWriter, request *http.Request) {
	var (
		err                error
		graph              *library.Graph
		item               interface{}
		startTime, endTime time.Time
	)

	if request.Method != "POST" && request.Method != "HEAD" {
		server.serveResponse(writer, serverResponse{mesgMethodNotAllowed}, http.StatusMethodNotAllowed)
		return
	} else if utils.HTTPGetContentType(request) != "application/json" {
		server.serveResponse(writer, serverResponse{mesgUnsupportedMediaType}, http.StatusUnsupportedMediaType)
		return
	}

	// Parse input JSON for graph data
	body, _ := ioutil.ReadAll(request.Body)

	plotReq := PlotRequest{}

	if err := json.Unmarshal(body, &plotReq); err != nil {
		logger.Log(logger.LevelError, "server", "%s", err)
		server.serveResponse(writer, serverResponse{mesgResourceInvalid}, http.StatusBadRequest)
		return
	}

	if plotReq.Time == "" {
		endTime = time.Now()
	} else if strings.HasPrefix(strings.Trim(plotReq.Range, " "), "-") {
		if endTime, err = time.Parse(time.RFC3339, plotReq.Time); err != nil {
			logger.Log(logger.LevelError, "server", "%s", err)
			server.serveResponse(writer, serverResponse{mesgResourceInvalid}, http.StatusBadRequest)
			return
		}
	} else {
		if startTime, err = time.Parse(time.RFC3339, plotReq.Time); err != nil {
			logger.Log(logger.LevelError, "server", "%s", err)
			server.serveResponse(writer, serverResponse{mesgResourceInvalid}, http.StatusBadRequest)
			return
		}
	}

	if startTime.IsZero() {
		if startTime, err = utils.TimeApplyRange(endTime, plotReq.Range); err != nil {
			logger.Log(logger.LevelError, "server", "%s", err)
			server.serveResponse(writer, serverResponse{mesgResourceInvalid}, http.StatusBadRequest)
			return
		}
	} else if endTime, err = utils.TimeApplyRange(startTime, plotReq.Range); err != nil {
		logger.Log(logger.LevelError, "server", "%s", err)
		server.serveResponse(writer, serverResponse{mesgResourceInvalid}, http.StatusBadRequest)
		return
	}

	if plotReq.Sample == 0 {
		plotReq.Sample = config.DefaultPlotSample
	}

	// Get graph from library
	graph = plotReq.Graph

	if plotReq.ID != "" {
		if item, err = server.Library.GetItem(plotReq.ID, library.LibraryItemGraph); err == nil {
			graph = item.(*library.Graph)
		}
	}

	if graph == nil {
		err = os.ErrNotExist
	}

	if err != nil {
		if os.IsNotExist(err) {
			server.serveResponse(writer, serverResponse{mesgResourceNotFound}, http.StatusNotFound)
		} else {
			logger.Log(logger.LevelError, "server", "%s", err)
			server.serveResponse(writer, serverResponse{mesgUnhandledError}, http.StatusInternalServerError)
		}

		return
	}

	step := endTime.Sub(startTime) / time.Duration(plotReq.Sample)

	// Get plots data
	groupOptions := make(map[string]map[string]interface{})

	data := make([]map[string]*types.PlotResult, 0)

	for _, groupItem := range graph.Groups {
		groupOptions[groupItem.Name] = groupItem.Options

		query, providerConnector, err := server.preparePlotQuery(&plotReq, groupItem)
		if err != nil {
			if err != os.ErrInvalid {
				logger.Log(logger.LevelError, "server", "%s", err)
			}

			data = append(data, nil)
			continue
		}

		plotResult, err := providerConnector.GetPlots(&types.PlotQuery{query, startTime, endTime, step,
			plotReq.Percentiles})
		if err != nil {
			logger.Log(logger.LevelError, "server", "%s", err)
		}

		data = append(data, plotResult)
	}

	response := &PlotResponse{
		ID:          graph.ID,
		Start:       startTime.Format(time.RFC3339),
		End:         endTime.Format(time.RFC3339),
		Step:        step.Seconds(),
		Name:        graph.Name,
		Description: graph.Description,
		Type:        graph.Type,
		StackMode:   graph.StackMode,
		UnitLabel:   graph.UnitLabel,
		UnitType:    graph.UnitType,
		Modified:    graph.Modified,
	}

	if len(data) == 0 {
		server.serveResponse(writer, serverResponse{mesgEmptyData}, http.StatusOK)
		return
	}

	plotMax := 0

	for _, groupItem := range graph.Groups {
		var plotResult map[string]*types.PlotResult

		plotResult, data = data[0], data[1:]

		for serieName, serieResult := range plotResult {
			if len(serieResult.Plots) > plotMax {
				plotMax = len(serieResult.Plots)
			}

			response.Series = append(response.Series, &SerieResponse{
				Name:    serieName,
				Plots:   serieResult.Plots,
				Info:    serieResult.Info,
				Options: groupOptions[groupItem.Name],
			})
		}
	}

	if plotMax > 0 {
		response.Step = (endTime.Sub(startTime) / time.Duration(plotMax)).Seconds()
	}

	server.serveResponse(writer, response, http.StatusOK)
}

func (server *Server) preparePlotQuery(plotReq *PlotRequest, groupItem *library.OperGroup) (*types.GroupQuery,
	connector.Connector, error) {
	var providerConnector connector.Connector

	query := &types.GroupQuery{
		Name:  groupItem.Name,
		Type:  groupItem.Type,
		Scale: groupItem.Scale,
	}

	for _, serieItem := range groupItem.Series {
		// Check for connectors errors or conflicts
		if _, ok := server.Catalog.Origins[serieItem.Origin]; !ok {
			return nil, nil, fmt.Errorf("unknown serie origin `%s'", serieItem.Origin)
		}

		serieSources := make([]string, 0)

		if strings.HasPrefix(serieItem.Source, library.LibraryGroupPrefix) {
			serieSources = server.Library.ExpandGroup(
				strings.TrimPrefix(serieItem.Source, library.LibraryGroupPrefix),
				library.LibraryItemSourceGroup,
			)
		} else {
			serieSources = []string{serieItem.Source}
		}

		index := 0

		for _, serieSource := range serieSources {
			if strings.HasPrefix(serieItem.Metric, library.LibraryGroupPrefix) {
				for _, serieChunk := range server.Library.ExpandGroup(
					strings.TrimPrefix(serieItem.Metric, library.LibraryGroupPrefix),
					library.LibraryItemMetricGroup,
				) {
					metric := server.Catalog.GetMetric(serieItem.Origin, serieSource, serieChunk)

					if metric == nil {
						logger.Log(
							logger.LevelWarning,
							"server",
							"unknown metric `%s' for source `%s' (origin: %s)",
							serieChunk,
							serieSource,
							serieItem.Origin,
						)

						continue
					}

					if providerConnector == nil {
						providerConnector = metric.Connector.(connector.Connector)
					} else if providerConnector != metric.Connector.(connector.Connector) {
						return nil, nil, fmt.Errorf("connectors differ between series")
					}

					query.Series = append(query.Series, &types.SerieQuery{
						Name: fmt.Sprintf("%s-%d", serieItem.Name, index),
						Metric: &types.MetricQuery{
							Name:   metric.OriginalName,
							Origin: metric.Source.Origin.OriginalName,
							Source: metric.Source.OriginalName,
						},
						Scale: serieItem.Scale,
					})

					index += 1
				}
			} else {
				metric := server.Catalog.GetMetric(serieItem.Origin, serieSource, serieItem.Metric)

				if metric == nil {
					logger.Log(
						logger.LevelWarning,
						"server",
						"unknown metric `%s' for source `%s' (origin: %s)",
						serieItem.Metric,
						serieSource,
						serieItem.Origin,
					)

					continue
				}

				if providerConnector == nil {
					providerConnector = metric.Connector.(connector.Connector)
				} else if providerConnector != metric.Connector.(connector.Connector) {
					return nil, nil, fmt.Errorf("connectors differ between series")
				}

				serie := &types.SerieQuery{
					Metric: &types.MetricQuery{
						Name:   metric.OriginalName,
						Origin: metric.Source.Origin.OriginalName,
						Source: metric.Source.OriginalName,
					},
					Scale: serieItem.Scale,
				}

				if len(serieSources) > 1 {
					serie.Name = fmt.Sprintf("%s-%d", serieItem.Name, index)
				} else {
					serie.Name = serieItem.Name
				}

				query.Series = append(query.Series, serie)

				index += 1
			}
		}
	}

	if len(query.Series) == 0 {
		return nil, nil, os.ErrInvalid
	}

	return query, providerConnector, nil
}
<|endoftext|>


<commit_before>package muta

import (
	"errors"
	"fmt"
	"io"
	"os"
	"path/filepath"
)

type FileInfo struct {
	Name string
	Path string
	Ctx  *interface{}
}

type SrcOpts struct {
	ReadSize uint
}

type Streamer func(*FileInfo, []byte) (*FileInfo, []byte, error)

// A convenience function to let functions that return Streamers
// "return an error". Ie, the following syntax:
//
// ```golang
// err := doSomething()
// if err != nil {
//   return ErrorStreamer(err)
// }
// ```
//
// ErrorStreamer will simply return a Streamer that will return an
// error when called.
func ErrorStreamer(err error) Streamer {
	return func(fi *FileInfo, chunk []byte) (*FileInfo, []byte, error) {
		return fi, chunk, err
	}
}

func SrcStreamer(ps []string, opts SrcOpts) Streamer {
	if opts.ReadSize == 0 {
		opts.ReadSize = 50
	}

	// Setup our channels
	fi := make(chan *FileInfo)
	chunk := make(chan []byte)
	err := make(chan error)
	read := make(chan bool)

	// This method of reading files needs to be abstracted further
	// to ensure that the file closing is deferred. In this
	// implementation i can't think of a way to test that.
	// Also, moving it out would let us ensure closing of the files
	// in tests
	go func() {
		loadFile := func(p string) {
			pchunks := make([]byte, opts.ReadSize)
			pfi := &FileInfo{
				Name: filepath.Base(p),
				Path: filepath.Dir(p),
			}

			f, ferr := os.Open(p)
			defer f.Close()
			if ferr != nil {
				fi <- pfi
				chunk <- nil
				err <- ferr
				return
			}

			// Wait for a read request
			for <-read {

				// Read
				count, ferr := f.Read(pchunks)
				if ferr != nil && ferr == io.EOF {
					break
				}

				// Send
				fi <- pfi
				chunk <- pchunks[0:count]
				err <- ferr
				if ferr != nil {
					return
				}
			}

			fi <- pfi
			chunk <- nil
			err <- nil
		}

		for _, p := range ps {
			loadFile(p)
		}
	}()

	return func(inFi *FileInfo, inC []byte) (*FileInfo, []byte, error) {
		// If there is an incoming file pass the data along unmodified. This
		// func doesn't care to modify the data in any way
		if inFi != nil {
			return inFi, inC, nil
		}

		read <- true
		return <-fi, <-chunk, <-err
	}
}

type DestOpts struct {
	// Not implemented
	Clean bool
	// Not implemented
	Overwrite bool
}

func Dest(d string, args ...interface{}) Streamer {
	var opts DestOpts
	if len(args) == 0 {
		opts = DestOpts{
			Clean:     false,
			Overwrite: true,
		}
	} else if len(args) == 1 {
		_opts, ok := args[0].(DestOpts)
		opts = _opts
		if !ok {
			return ErrorStreamer(errors.New(
				"Unrecognized type in Dest(string, ...interface{}). " +
					"Use DestOpts()",
			))
		}
	}

	if opts.Clean {
		err := os.RemoveAll(d)
		if err != nil {
			return ErrorStreamer(err)
		}
	}

	// Make the destination if needed
	if err := os.MkdirAll(d, 0755); err != nil {
		return ErrorStreamer(err)
	}

	// A staging variable for the currently working file.
	var f *os.File
	return func(fi *FileInfo, chunk []byte) (*FileInfo, []byte, error) {
		// If fi is nil, then this func is now the generator. Dest() has no
		// need to generate, so signal EOS
		if fi == nil {
			return nil, chunk, nil
		}

		destPath := filepath.Join(d, fi.Path)
		destFilepath := filepath.Join(destPath, fi.Name)
		// MkdirAll checks if the given path is a dir, and exists. So
		// i believe there is no reason for us to bother checking.
		err := os.MkdirAll(destPath, 0755)
		if err != nil {
			return fi, chunk, err
		}

		if chunk == nil && f != nil {
			// f is open for writing, but chunk is nil, we're at EOF.
			// Close f, and set it to nil
			err = f.Close()
			f = nil
			return nil, nil, err
		}

		if f == nil {
			osFi, err := os.Stat(destFilepath)
			if err == nil && osFi.IsDir() {
				return fi, chunk, errors.New(fmt.Sprintf(
					"Cannot write to '%s', path is directory.",
					destFilepath,
				))
			}

			// This area is a bit of a cluster f*ck. In short:
			//
			// 1. If there is an error, and the error is that the file
			// does not exist, create it.
			// 2. If it's not a file does not exist error, return it.
			// 3. If there is no error, and the filepath is a directory,
			// return an error.
			// 4. If it's not a directory, and we're not allowed to overwrite
			// it, return an error.
			// 5. If we are allowed to overwrite it, open it up.
			//
			// Did i drink too much while writing this? It feels so messy.
			if err != nil {
				if os.IsNotExist(err) {
					f, err = os.Create(destFilepath)
					if err != nil {
						// Failed to create file, return
						return fi, chunk, err
					}
				} else {
					// Stat() error is unknown, return
					return fi, chunk, err
				}
			} else {
				// There was no error Stating path, it exist
				if osFi.IsDir() {
					// The file path is a dir, return error
					return fi, chunk, errors.New(fmt.Sprintf(
						"Cannot write to '%s', path is directory.",
						destFilepath,
					))
				} else if !opts.Overwrite {
					// We're not allowed to overwrite. Return error.
					return fi, chunk, errors.New(fmt.Sprintf(
						"Cannot write to '%s', path exists and Overwrite is set "+
							"to false.",
						destFilepath,
					))
				} else {
					f, err = os.Open(destFilepath)
					if err != nil {
						// Failed to open file for writing.
						return fi, chunk, err
					}
				}
			}
		}

		// lenth written can be ignored, because Write() returns an error
		// if len(chunk) != n
		_, err = f.Write(chunk)

		// Return EOS always. Dest() writes everything, like a boss..?
		return nil, nil, err
	}
}
<commit_msg>Fixed SrcStreamer() not waiting for `read<-true` before signalling EOS<commit_after>package muta

import (
	"errors"
	"fmt"
	"io"
	"os"
	"path/filepath"
)

type FileInfo struct {
	Name string
	Path string
	Ctx  *interface{}
}

type SrcOpts struct {
	ReadSize uint
}

type Streamer func(*FileInfo, []byte) (*FileInfo, []byte, error)

// A convenience function to let functions that return Streamers
// "return an error". Ie, the following syntax:
//
// ```golang
// err := doSomething()
// if err != nil {
//   return ErrorStreamer(err)
// }
// ```
//
// ErrorStreamer will simply return a Streamer that will return an
// error when called.
func ErrorStreamer(err error) Streamer {
	return func(fi *FileInfo, chunk []byte) (*FileInfo, []byte, error) {
		return fi, chunk, err
	}
}

func SrcStreamer(ps []string, opts SrcOpts) Streamer {
	if opts.ReadSize == 0 {
		opts.ReadSize = 50
	}

	// Setup our channels
	fi := make(chan *FileInfo)
	chunk := make(chan []byte)
	err := make(chan error)
	read := make(chan bool)

	// This method of reading files needs to be abstracted further
	// to ensure that the file closing is deferred. In this
	// implementation i can't think of a way to test that.
	// Also, moving it out would let us ensure closing of the files
	// in tests
	go func() {
		loadFile := func(p string) error {
			pchunks := make([]byte, opts.ReadSize)
			pfi := &FileInfo{
				Name: filepath.Base(p),
				Path: filepath.Dir(p),
			}

			f, ferr := os.Open(p)
			defer f.Close()
			if ferr != nil {
				fi <- pfi
				chunk <- nil
				err <- ferr
				return ferr
			}

			// Wait for a read request
			for <-read {
				// Read
				count, ferr := f.Read(pchunks)
				if ferr != nil && ferr == io.EOF {
					break
				}

				// Send
				fi <- pfi
				chunk <- pchunks[0:count]
				err <- ferr
				if ferr != nil {
					return ferr
				}
			}

			// The for loop stopped, send EOF
			fi <- pfi
			chunk <- nil
			err <- nil
			return nil
		}

		for _, p := range ps {
			err := loadFile(p)
			if err != nil {
				return
			}
		}

		<-read
		// send EOS
		fi <- nil
		chunk <- nil
		err <- nil
	}()

	return func(inFi *FileInfo, inC []byte) (*FileInfo, []byte, error) {
		// If there is an incoming file pass the data along unmodified. This
		// func doesn't care to modify the data in any way
		if inFi != nil {
			return inFi, inC, nil
		}

		read <- true
		return <-fi, <-chunk, <-err
	}
}

type DestOpts struct {
	// Not implemented
	Clean bool
	// Not implemented
	Overwrite bool
}

func Dest(d string, args ...interface{}) Streamer {
	var opts DestOpts
	if len(args) == 0 {
		opts = DestOpts{
			Clean:     false,
			Overwrite: true,
		}
	} else if len(args) == 1 {
		_opts, ok := args[0].(DestOpts)
		opts = _opts
		if !ok {
			return ErrorStreamer(errors.New(
				"Unrecognized type in Dest(string, ...interface{}). " +
					"Use DestOpts()",
			))
		}
	}

	if opts.Clean {
		err := os.RemoveAll(d)
		if err != nil {
			return ErrorStreamer(err)
		}
	}

	// Make the destination if needed
	if err := os.MkdirAll(d, 0755); err != nil {
		return ErrorStreamer(err)
	}

	// A staging variable for the currently working file.
	var f *os.File
	return func(fi *FileInfo, chunk []byte) (*FileInfo, []byte, error) {
		// If fi is nil, then this func is now the generator. Dest() has no
		// need to generate, so signal EOS
		if fi == nil {
			return nil, chunk, nil
		}

		destPath := filepath.Join(d, fi.Path)
		destFilepath := filepath.Join(destPath, fi.Name)
		// MkdirAll checks if the given path is a dir, and exists. So
		// i believe there is no reason for us to bother checking.
		err := os.MkdirAll(destPath, 0755)
		if err != nil {
			return fi, chunk, err
		}

		if chunk == nil && f != nil {
			// f is open for writing, but chunk is nil, we're at EOF.
			// Close f, and set it to nil
			err = f.Close()
			f = nil
			return nil, nil, err
		}

		if f == nil {
			osFi, err := os.Stat(destFilepath)
			if err == nil && osFi.IsDir() {
				return fi, chunk, errors.New(fmt.Sprintf(
					"Cannot write to '%s', path is directory.",
					destFilepath,
				))
			}

			// This area is a bit of a cluster f*ck. In short:
			//
			// 1. If there is an error, and the error is that the file
			// does not exist, create it.
			// 2. If it's not a file does not exist error, return it.
			// 3. If there is no error, and the filepath is a directory,
			// return an error.
			// 4. If it's not a directory, and we're not allowed to overwrite
			// it, return an error.
			// 5. If we are allowed to overwrite it, open it up.
			//
			// Did i drink too much while writing this? It feels so messy.
			if err != nil {
				if os.IsNotExist(err) {
					f, err = os.Create(destFilepath)
					if err != nil {
						// Failed to create file, return
						return fi, chunk, err
					}
				} else {
					// Stat() error is unknown, return
					return fi, chunk, err
				}
			} else {
				// There was no error Stating path, it exist
				if osFi.IsDir() {
					// The file path is a dir, return error
					return fi, chunk, errors.New(fmt.Sprintf(
						"Cannot write to '%s', path is directory.",
						destFilepath,
					))
				} else if !opts.Overwrite {
					// We're not allowed to overwrite. Return error.
					return fi, chunk, errors.New(fmt.Sprintf(
						"Cannot write to '%s', path exists and Overwrite is set "+
							"to false.",
						destFilepath,
					))
				} else {
					f, err = os.Open(destFilepath)
					if err != nil {
						// Failed to open file for writing.
						return fi, chunk, err
					}
				}
			}
		}

		// lenth written can be ignored, because Write() returns an error
		// if len(chunk) != n
		_, err = f.Write(chunk)

		// Return EOS always. Dest() writes everything, like a boss..?
		return nil, nil, err
	}
}
<|endoftext|>


<commit_before>package servicelb

import (
	"context"
	"fmt"
	"sort"
	"strconv"

	appclient "github.com/rancher/k3s/types/apis/apps/v1"
	coreclient "github.com/rancher/k3s/types/apis/core/v1"
	"github.com/rancher/norman/condition"
	"github.com/rancher/norman/pkg/changeset"
	"github.com/rancher/norman/pkg/objectset"
	"github.com/rancher/norman/types/slice"
	"github.com/sirupsen/logrus"
	apps "k8s.io/api/apps/v1"
	core "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/errors"
	meta "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/labels"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/util/intstr"
	"k8s.io/client-go/kubernetes"
	coregetter "k8s.io/client-go/kubernetes/typed/core/v1"
)

const (
	image        = "rancher/klipper-lb:v0.1.1"
	svcNameLabel = "svccontroller.k3s.cattle.io/svcname"
	Ready        = condition.Cond("Ready")
)

var (
	trueVal = true
)

func Register(ctx context.Context, kubernetes kubernetes.Interface, enabled bool) error {
	clients := coreclient.ClientsFrom(ctx)
	appClients := appclient.ClientsFrom(ctx)

	h := &handler{
		enabled:   enabled,
		nodeCache: clients.Node.Cache(),
		podCache:  clients.Pod.Cache(),
		processor: objectset.NewProcessor("svccontroller").
			Client(appClients.Deployment),
		serviceCache: clients.Service.Cache(),
		services:     kubernetes.CoreV1(),
	}

	clients.Service.OnChange(ctx, "svccontroller", h.onChange)
	changeset.Watch(ctx, "svccontroller-watcher",
		h.onResourceChange,
		clients.Service,
		clients.Pod,
		clients.Endpoints)

	return nil
}

type handler struct {
	enabled      bool
	nodeCache    coreclient.NodeClientCache
	podCache     coreclient.PodClientCache
	processor    *objectset.Processor
	serviceCache coreclient.ServiceClientCache
	services     coregetter.ServicesGetter
}

func (h *handler) onResourceChange(name, namespace string, obj runtime.Object) ([]changeset.Key, error) {
	if ep, ok := obj.(*core.Endpoints); ok {
		return []changeset.Key{
			{
				Name:      ep.Name,
				Namespace: ep.Namespace,
			},
		}, nil
	}

	pod, ok := obj.(*core.Pod)
	if !ok {
		return nil, nil
	}

	serviceName := pod.Labels[svcNameLabel]
	if serviceName == "" {
		return nil, nil
	}

	if pod.Status.PodIP == "" {
		return nil, nil
	}

	return []changeset.Key{
		{
			Name:      serviceName,
			Namespace: pod.Namespace,
		},
	}, nil
}

func (h *handler) onChange(svc *core.Service) (runtime.Object, error) {
	if svc.Spec.Type != core.ServiceTypeLoadBalancer || svc.Spec.ClusterIP == "" ||
		svc.Spec.ClusterIP == "None" {
		return svc, nil
	}

	if err := h.deployPod(svc); err != nil {
		return svc, err
	}

	return h.updateService(svc)
}

func (h *handler) updateService(svc *core.Service) (runtime.Object, error) {
	pods, err := h.podCache.List(svc.Namespace, labels.SelectorFromSet(map[string]string{
		svcNameLabel: svc.Name,
	}))

	if err != nil {
		return svc, err
	}

	existingIPs := serviceIPs(svc)
	expectedIPs, err := h.podIPs(pods)
	if err != nil {
		return svc, err
	}

	sort.Strings(expectedIPs)
	sort.Strings(existingIPs)

	if slice.StringsEqual(expectedIPs, existingIPs) {
		return svc, nil
	}

	svc = svc.DeepCopy()
	svc.Status.LoadBalancer.Ingress = nil
	for _, ip := range expectedIPs {
		svc.Status.LoadBalancer.Ingress = append(svc.Status.LoadBalancer.Ingress, core.LoadBalancerIngress{
			IP: ip,
		})
	}

	logrus.Debugf("Setting service loadbalancer %s/%s to IPs %v", svc.Namespace, svc.Name, expectedIPs)
	return h.services.Services(svc.Namespace).UpdateStatus(svc)
}

func serviceIPs(svc *core.Service) []string {
	var ips []string

	for _, ingress := range svc.Status.LoadBalancer.Ingress {
		if ingress.IP != "" {
			ips = append(ips, ingress.IP)
		}
	}

	return ips
}

func (h *handler) podIPs(pods []*core.Pod) ([]string, error) {
	ips := map[string]bool{}

	for _, pod := range pods {
		if pod.Spec.NodeName == "" || pod.Status.PodIP == "" {
			continue
		}
		if !Ready.IsTrue(pod) {
			continue
		}

		node, err := h.nodeCache.Get("", pod.Spec.NodeName)
		if errors.IsNotFound(err) {
			continue
		} else if err != nil {
			return nil, err
		}

		for _, addr := range node.Status.Addresses {
			if addr.Type == core.NodeInternalIP {
				ips[addr.Address] = true
			}
		}
	}

	var ipList []string
	for k := range ips {
		ipList = append(ipList, k)
	}
	return ipList, nil
}

func (h *handler) deployPod(svc *core.Service) error {
	objs := objectset.NewObjectSet()
	if !h.enabled {
		return h.processor.NewDesiredSet(svc, objs).Apply()
	}

	dep, err := h.newDeployment(svc)
	if err != nil {
		return err
	}
	if dep != nil {
		objs.Add(dep)
	}

	return h.processor.NewDesiredSet(svc, objs).Apply()
}

func (h *handler) resolvePort(svc *core.Service, targetPort core.ServicePort) (int32, error) {
	if len(svc.Spec.Selector) == 0 {
		return 0, nil
	}

	if targetPort.TargetPort.IntVal != 0 {
		return targetPort.TargetPort.IntVal, nil
	}

	pods, err := h.podCache.List(svc.Namespace, labels.SelectorFromSet(svc.Spec.Selector))
	if err != nil {
		return 0, err
	}

	for _, pod := range pods {
		if !Ready.IsTrue(pod) {
			continue
		}
		for _, container := range pod.Spec.Containers {
			for _, port := range container.Ports {
				if port.Name == targetPort.TargetPort.StrVal {
					return port.ContainerPort, nil
				}
			}
		}
	}

	return 0, nil
}

func (h *handler) newDeployment(svc *core.Service) (*apps.Deployment, error) {
	name := fmt.Sprintf("svclb-%s", svc.Name)
	zeroInt := intstr.FromInt(0)
	oneInt := intstr.FromInt(1)
	replicas := int32(0)

	nodes, err := h.nodeCache.List("", labels.Everything())
	if err != nil {
		return nil, err
	}

	for _, node := range nodes {
		if Ready.IsTrue(node) {
			replicas += 1
		}
		if replicas >= 2 {
			break
		}
	}

	dep := &apps.Deployment{
		ObjectMeta: meta.ObjectMeta{
			Name:      name,
			Namespace: svc.Namespace,
			OwnerReferences: []meta.OwnerReference{
				{
					Name:       svc.Name,
					APIVersion: "v1",
					Kind:       "Service",
					UID:        svc.UID,
					Controller: &trueVal,
				},
			},
		},
		TypeMeta: meta.TypeMeta{
			Kind:       "Deployment",
			APIVersion: "apps/v1",
		},
		Spec: apps.DeploymentSpec{
			Replicas: &replicas,
			Selector: &meta.LabelSelector{
				MatchLabels: map[string]string{
					"app": name,
				},
			},
			Template: core.PodTemplateSpec{
				ObjectMeta: meta.ObjectMeta{
					Labels: map[string]string{
						"app":        name,
						svcNameLabel: svc.Name,
					},
				},
			},
			Strategy: apps.DeploymentStrategy{
				Type: apps.RollingUpdateDeploymentStrategyType,
				RollingUpdate: &apps.RollingUpdateDeployment{
					MaxSurge:       &zeroInt,
					MaxUnavailable: &oneInt,
				},
			},
		},
	}

	for _, port := range svc.Spec.Ports {
		targetPort, err := h.resolvePort(svc, port)
		if err != nil {
			return nil, err
		}

		container := core.Container{
			Name:            fmt.Sprintf("port-%s", port.Name),
			Image:           image,
			ImagePullPolicy: core.PullIfNotPresent,
			Ports: []core.ContainerPort{
				{
					Name:          port.Name,
					ContainerPort: port.Port,
					HostPort:      port.Port,
				},
			},
			Env: []core.EnvVar{
				{
					Name:  "SRC_PORT",
					Value: strconv.Itoa(int(port.Port)),
				},
				{
					Name:  "DEST_PROTO",
					Value: string(port.Protocol),
				},
				{
					Name:  "DEST_PORT",
					Value: strconv.Itoa(int(targetPort)),
				},
				{
					Name:  "DEST_IP",
					Value: svc.Spec.ClusterIP,
				},
			},
			SecurityContext: &core.SecurityContext{
				Capabilities: &core.Capabilities{
					Add: []core.Capability{
						"NET_ADMIN",
					},
				},
			},
		}

		dep.Spec.Template.Spec.Containers = append(dep.Spec.Template.Spec.Containers, container)
	}

	return dep, nil
}
<commit_msg>Fix dest port so it's the same as src port<commit_after>package servicelb

import (
	"context"
	"fmt"
	"sort"
	"strconv"

	appclient "github.com/rancher/k3s/types/apis/apps/v1"
	coreclient "github.com/rancher/k3s/types/apis/core/v1"
	"github.com/rancher/norman/condition"
	"github.com/rancher/norman/pkg/changeset"
	"github.com/rancher/norman/pkg/objectset"
	"github.com/rancher/norman/types/slice"
	"github.com/sirupsen/logrus"
	apps "k8s.io/api/apps/v1"
	core "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/errors"
	meta "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/labels"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/util/intstr"
	"k8s.io/client-go/kubernetes"
	coregetter "k8s.io/client-go/kubernetes/typed/core/v1"
)

const (
	image        = "rancher/klipper-lb:v0.1.1"
	svcNameLabel = "svccontroller.k3s.cattle.io/svcname"
	Ready        = condition.Cond("Ready")
)

var (
	trueVal = true
)

func Register(ctx context.Context, kubernetes kubernetes.Interface, enabled bool) error {
	clients := coreclient.ClientsFrom(ctx)
	appClients := appclient.ClientsFrom(ctx)

	h := &handler{
		enabled:   enabled,
		nodeCache: clients.Node.Cache(),
		podCache:  clients.Pod.Cache(),
		processor: objectset.NewProcessor("svccontroller").
			Client(appClients.Deployment),
		serviceCache: clients.Service.Cache(),
		services:     kubernetes.CoreV1(),
	}

	clients.Service.OnChange(ctx, "svccontroller", h.onChange)
	changeset.Watch(ctx, "svccontroller-watcher",
		h.onResourceChange,
		clients.Service,
		clients.Pod,
		clients.Endpoints)

	return nil
}

type handler struct {
	enabled      bool
	nodeCache    coreclient.NodeClientCache
	podCache     coreclient.PodClientCache
	processor    *objectset.Processor
	serviceCache coreclient.ServiceClientCache
	services     coregetter.ServicesGetter
}

func (h *handler) onResourceChange(name, namespace string, obj runtime.Object) ([]changeset.Key, error) {
	if ep, ok := obj.(*core.Endpoints); ok {
		return []changeset.Key{
			{
				Name:      ep.Name,
				Namespace: ep.Namespace,
			},
		}, nil
	}

	pod, ok := obj.(*core.Pod)
	if !ok {
		return nil, nil
	}

	serviceName := pod.Labels[svcNameLabel]
	if serviceName == "" {
		return nil, nil
	}

	if pod.Status.PodIP == "" {
		return nil, nil
	}

	return []changeset.Key{
		{
			Name:      serviceName,
			Namespace: pod.Namespace,
		},
	}, nil
}

func (h *handler) onChange(svc *core.Service) (runtime.Object, error) {
	if svc.Spec.Type != core.ServiceTypeLoadBalancer || svc.Spec.ClusterIP == "" ||
		svc.Spec.ClusterIP == "None" {
		return svc, nil
	}

	if err := h.deployPod(svc); err != nil {
		return svc, err
	}

	return h.updateService(svc)
}

func (h *handler) updateService(svc *core.Service) (runtime.Object, error) {
	pods, err := h.podCache.List(svc.Namespace, labels.SelectorFromSet(map[string]string{
		svcNameLabel: svc.Name,
	}))

	if err != nil {
		return svc, err
	}

	existingIPs := serviceIPs(svc)
	expectedIPs, err := h.podIPs(pods)
	if err != nil {
		return svc, err
	}

	sort.Strings(expectedIPs)
	sort.Strings(existingIPs)

	if slice.StringsEqual(expectedIPs, existingIPs) {
		return svc, nil
	}

	svc = svc.DeepCopy()
	svc.Status.LoadBalancer.Ingress = nil
	for _, ip := range expectedIPs {
		svc.Status.LoadBalancer.Ingress = append(svc.Status.LoadBalancer.Ingress, core.LoadBalancerIngress{
			IP: ip,
		})
	}

	logrus.Debugf("Setting service loadbalancer %s/%s to IPs %v", svc.Namespace, svc.Name, expectedIPs)
	return h.services.Services(svc.Namespace).UpdateStatus(svc)
}

func serviceIPs(svc *core.Service) []string {
	var ips []string

	for _, ingress := range svc.Status.LoadBalancer.Ingress {
		if ingress.IP != "" {
			ips = append(ips, ingress.IP)
		}
	}

	return ips
}

func (h *handler) podIPs(pods []*core.Pod) ([]string, error) {
	ips := map[string]bool{}

	for _, pod := range pods {
		if pod.Spec.NodeName == "" || pod.Status.PodIP == "" {
			continue
		}
		if !Ready.IsTrue(pod) {
			continue
		}

		node, err := h.nodeCache.Get("", pod.Spec.NodeName)
		if errors.IsNotFound(err) {
			continue
		} else if err != nil {
			return nil, err
		}

		for _, addr := range node.Status.Addresses {
			if addr.Type == core.NodeInternalIP {
				ips[addr.Address] = true
			}
		}
	}

	var ipList []string
	for k := range ips {
		ipList = append(ipList, k)
	}
	return ipList, nil
}

func (h *handler) deployPod(svc *core.Service) error {
	objs := objectset.NewObjectSet()
	if !h.enabled {
		return h.processor.NewDesiredSet(svc, objs).Apply()
	}

	dep, err := h.newDeployment(svc)
	if err != nil {
		return err
	}
	if dep != nil {
		objs.Add(dep)
	}

	return h.processor.NewDesiredSet(svc, objs).Apply()
}

func (h *handler) newDeployment(svc *core.Service) (*apps.Deployment, error) {
	name := fmt.Sprintf("svclb-%s", svc.Name)
	zeroInt := intstr.FromInt(0)
	oneInt := intstr.FromInt(1)
	replicas := int32(0)

	nodes, err := h.nodeCache.List("", labels.Everything())
	if err != nil {
		return nil, err
	}

	for _, node := range nodes {
		if Ready.IsTrue(node) {
			replicas += 1
		}
		if replicas >= 2 {
			break
		}
	}

	dep := &apps.Deployment{
		ObjectMeta: meta.ObjectMeta{
			Name:      name,
			Namespace: svc.Namespace,
			OwnerReferences: []meta.OwnerReference{
				{
					Name:       svc.Name,
					APIVersion: "v1",
					Kind:       "Service",
					UID:        svc.UID,
					Controller: &trueVal,
				},
			},
		},
		TypeMeta: meta.TypeMeta{
			Kind:       "Deployment",
			APIVersion: "apps/v1",
		},
		Spec: apps.DeploymentSpec{
			Replicas: &replicas,
			Selector: &meta.LabelSelector{
				MatchLabels: map[string]string{
					"app": name,
				},
			},
			Template: core.PodTemplateSpec{
				ObjectMeta: meta.ObjectMeta{
					Labels: map[string]string{
						"app":        name,
						svcNameLabel: svc.Name,
					},
				},
			},
			Strategy: apps.DeploymentStrategy{
				Type: apps.RollingUpdateDeploymentStrategyType,
				RollingUpdate: &apps.RollingUpdateDeployment{
					MaxSurge:       &zeroInt,
					MaxUnavailable: &oneInt,
				},
			},
		},
	}

	for i, port := range svc.Spec.Ports {
		portName := port.Name
		if portName == "" {
			portName = fmt.Sprintf("port-%d", i)
		}
		container := core.Container{
			Name:            portName,
			Image:           image,
			ImagePullPolicy: core.PullIfNotPresent,
			Ports: []core.ContainerPort{
				{
					Name:          portName,
					ContainerPort: port.Port,
					HostPort:      port.Port,
				},
			},
			Env: []core.EnvVar{
				{
					Name:  "SRC_PORT",
					Value: strconv.Itoa(int(port.Port)),
				},
				{
					Name:  "DEST_PROTO",
					Value: string(port.Protocol),
				},
				{
					Name:  "DEST_PORT",
					Value: strconv.Itoa(int(port.Port)),
				},
				{
					Name:  "DEST_IP",
					Value: svc.Spec.ClusterIP,
				},
			},
			SecurityContext: &core.SecurityContext{
				Capabilities: &core.Capabilities{
					Add: []core.Capability{
						"NET_ADMIN",
					},
				},
			},
		}

		dep.Spec.Template.Spec.Containers = append(dep.Spec.Template.Spec.Containers, container)
	}

	return dep, nil
}
<|endoftext|>


<commit_before>/*
Copyright 2018 Google LLC

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package deploy

import (
	"context"
	"fmt"
	"io"
	"os/exec"

	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/build"
	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/schema/v1alpha2"
	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/util"
	"github.com/pkg/errors"
	"github.com/sirupsen/logrus"
)

type HelmDeployer struct {
	*v1alpha2.DeployConfig
	kubeContext string
}

// NewHelmDeployer returns a new HelmDeployer for a DeployConfig filled
// with the needed configuration for `helm`
func NewHelmDeployer(cfg *v1alpha2.DeployConfig, kubeContext string) *HelmDeployer {
	return &HelmDeployer{
		DeployConfig: cfg,
		kubeContext:  kubeContext,
	}
}

func (h *HelmDeployer) Deploy(ctx context.Context, out io.Writer, b *build.BuildResult) (*Result, error) {
	for _, r := range h.HelmDeploy.Releases {
		if err := h.deployRelease(out, r, b); err != nil {
			return nil, errors.Wrapf(err, "deploying %s", r.Name)
		}
	}
	return nil, nil
}

// Not implemented
func (k *HelmDeployer) Dependencies() ([]string, error) {
	return nil, nil
}

// Cleanup deletes what was deployed by calling Deploy.
func (h *HelmDeployer) Cleanup(ctx context.Context, out io.Writer) error {
	for _, r := range h.HelmDeploy.Releases {
		if err := h.deleteRelease(out, r); err != nil {
			return errors.Wrapf(err, "deploying %s", r.Name)
		}
	}
	return nil
}

func (h *HelmDeployer) args(moreArgs ...string) []string {
	return append([]string{"--kube-context", h.kubeContext}, moreArgs...)
}

func (h *HelmDeployer) deployRelease(out io.Writer, r v1alpha2.HelmRelease, b *build.BuildResult) error {
	isInstalled := true
	getCmd := exec.Command("helm", h.args("get", r.Name)...)
	if stdout, stderr, err := util.RunCommand(getCmd, nil); err != nil {
		logrus.Debugf("Error getting release %s: %s stdout: %s stderr: %s", r.Name, err, string(stdout), string(stderr))
		fmt.Fprintf(out, "Helm release %s not installed. Installing...\n", r.Name)
		isInstalled = false
	}
	params, err := JoinTagsToBuildResult(b.Builds, r.Values)
	if err != nil {
		return errors.Wrap(err, "matching build results to chart values")
	}

	var setOpts []string
	for k, v := range params {
		setOpts = append(setOpts, "--set")
		setOpts = append(setOpts, fmt.Sprintf("%s=%s", k, v.Tag))
	}

	// First build dependencies.
	logrus.Infof("Building helm dependencies...")
	depCmd := exec.Command("helm", h.args("dep", "build", r.ChartPath)...)
	stdout, stderr, err := util.RunCommand(depCmd, nil)
	if err != nil {
		return errors.Wrapf(err, "helm dep build stdout: %s, stderr: %s", string(stdout), string(stderr))
	}
	out.Write(stdout)

	args := h.args()
	if !isInstalled {
		args = append(args, "install", "--name", r.Name, r.ChartPath)
	} else {
		args = append(args, "upgrade", r.Name, r.ChartPath)
	}

	if r.Namespace != "" {
		args = append(args, "--namespace", r.Namespace)
	}
	if r.ValuesFilePath != "" {
		args = append(args, "-f", r.ValuesFilePath)
	}
	if r.Version != "" {
		args = append(args, "--version", r.Version)
	}

	if len(r.SetValues) != 0 {
		for k, v := range r.SetValues {
			setOpts = append(setOpts, "--set")
			setOpts = append(setOpts, fmt.Sprintf("%s=%s", k, v))
		}
	}
	args = append(args, setOpts...)

	execCmd := exec.Command("helm", args...)
	stdout, stderr, err = util.RunCommand(execCmd, nil)
	if err != nil {
		return errors.Wrapf(err, "helm updater stdout: %s, stderr: %s", string(stdout), string(stderr))
	}

	out.Write(stdout)
	return nil
}

func (h *HelmDeployer) deleteRelease(out io.Writer, r v1alpha2.HelmRelease) error {
	getCmd := exec.Command("helm", h.args("delete", r.Name, "--purge")...)
	stdout, stderr, err := util.RunCommand(getCmd, nil)
	if err != nil {
		logrus.Debugf("running helm delete %s: %v stdout: %s stderr: %s", r.Name, err, string(stdout), string(stderr))
	}

	out.Write(stdout)
	return nil
}
<commit_msg>fix: allow an environment variable to default the deploy namespace<commit_after>/*
Copyright 2018 Google LLC

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package deploy

import (
	"context"
	"fmt"
	"io"
	"os"
	"os/exec"

	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/build"
	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/schema/v1alpha2"
	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/util"
	"github.com/pkg/errors"
	"github.com/sirupsen/logrus"
)

type HelmDeployer struct {
	*v1alpha2.DeployConfig
	kubeContext string
}

// NewHelmDeployer returns a new HelmDeployer for a DeployConfig filled
// with the needed configuration for `helm`
func NewHelmDeployer(cfg *v1alpha2.DeployConfig, kubeContext string) *HelmDeployer {
	return &HelmDeployer{
		DeployConfig: cfg,
		kubeContext:  kubeContext,
	}
}

func (h *HelmDeployer) Deploy(ctx context.Context, out io.Writer, b *build.BuildResult) (*Result, error) {
	for _, r := range h.HelmDeploy.Releases {
		if err := h.deployRelease(out, r, b); err != nil {
			return nil, errors.Wrapf(err, "deploying %s", r.Name)
		}
	}
	return nil, nil
}

// Not implemented
func (k *HelmDeployer) Dependencies() ([]string, error) {
	return nil, nil
}

// Cleanup deletes what was deployed by calling Deploy.
func (h *HelmDeployer) Cleanup(ctx context.Context, out io.Writer) error {
	for _, r := range h.HelmDeploy.Releases {
		if err := h.deleteRelease(out, r); err != nil {
			return errors.Wrapf(err, "deploying %s", r.Name)
		}
	}
	return nil
}

func (h *HelmDeployer) args(moreArgs ...string) []string {
	return append([]string{"--kube-context", h.kubeContext}, moreArgs...)
}

func (h *HelmDeployer) deployRelease(out io.Writer, r v1alpha2.HelmRelease, b *build.BuildResult) error {
	isInstalled := true
	getCmd := exec.Command("helm", h.args("get", r.Name)...)
	if stdout, stderr, err := util.RunCommand(getCmd, nil); err != nil {
		logrus.Debugf("Error getting release %s: %s stdout: %s stderr: %s", r.Name, err, string(stdout), string(stderr))
		fmt.Fprintf(out, "Helm release %s not installed. Installing...\n", r.Name)
		isInstalled = false
	}
	params, err := JoinTagsToBuildResult(b.Builds, r.Values)
	if err != nil {
		return errors.Wrap(err, "matching build results to chart values")
	}

	var setOpts []string
	for k, v := range params {
		setOpts = append(setOpts, "--set")
		setOpts = append(setOpts, fmt.Sprintf("%s=%s", k, v.Tag))
	}

	// First build dependencies.
	logrus.Infof("Building helm dependencies...")
	depCmd := exec.Command("helm", h.args("dep", "build", r.ChartPath)...)
	stdout, stderr, err := util.RunCommand(depCmd, nil)
	if err != nil {
		return errors.Wrapf(err, "helm dep build stdout: %s, stderr: %s", string(stdout), string(stderr))
	}
	out.Write(stdout)

	args := h.args()
	if !isInstalled {
		args = append(args, "install", "--name", r.Name, r.ChartPath)
	} else {
		args = append(args, "upgrade", r.Name, r.ChartPath)
	}

	ns := r.Namespace
	if ns == "" {
		ns = os.Getenv("SKAFFOLD_DEPLOY_NAMESPACE")
	}
	if ns != "" {
		args = append(args, "--namespace", ns)
	}
	if r.ValuesFilePath != "" {
		args = append(args, "-f", r.ValuesFilePath)
	}
	if r.Version != "" {
		args = append(args, "--version", r.Version)
	}

	if len(r.SetValues) != 0 {
		for k, v := range r.SetValues {
			setOpts = append(setOpts, "--set")
			setOpts = append(setOpts, fmt.Sprintf("%s=%s", k, v))
		}
	}
	args = append(args, setOpts...)

	execCmd := exec.Command("helm", args...)
	stdout, stderr, err = util.RunCommand(execCmd, nil)
	if err != nil {
		return errors.Wrapf(err, "helm updater stdout: %s, stderr: %s", string(stdout), string(stderr))
	}

	out.Write(stdout)
	return nil
}

func (h *HelmDeployer) deleteRelease(out io.Writer, r v1alpha2.HelmRelease) error {
	getCmd := exec.Command("helm", h.args("delete", r.Name, "--purge")...)
	stdout, stderr, err := util.RunCommand(getCmd, nil)
	if err != nil {
		logrus.Debugf("running helm delete %s: %v stdout: %s stderr: %s", r.Name, err, string(stdout), string(stderr))
	}

	out.Write(stdout)
	return nil
}
<|endoftext|>