item_hook_handler.go

/*
Copyright 2020 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 hook

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

	uuid "github.com/gofrs/uuid"
	"github.com/pkg/errors"
	"github.com/sirupsen/logrus"
	corev1api "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/meta"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
	"k8s.io/apimachinery/pkg/labels"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/runtime/schema"

	velerov1api "github.com/adi-bhardwaj/velero-modified/pkg/apis/velero/v1"
	"github.com/adi-bhardwaj/velero-modified/pkg/kuberesource"
	"github.com/adi-bhardwaj/velero-modified/pkg/podexec"
	"github.com/adi-bhardwaj/velero-modified/pkg/restic"
	"github.com/adi-bhardwaj/velero-modified/pkg/util/collections"
	"github.com/adi-bhardwaj/velero-modified/pkg/util/kube"
)

type hookPhase string

const (
	PhasePre  hookPhase = "pre"
	PhasePost hookPhase = "post"
)

const (
	// Backup hook annotations
	podBackupHookContainerAnnotationKey = "hook.backup.velero.io/container"
	podBackupHookCommandAnnotationKey   = "hook.backup.velero.io/command"
	podBackupHookOnErrorAnnotationKey   = "hook.backup.velero.io/on-error"
	podBackupHookTimeoutAnnotationKey   = "hook.backup.velero.io/timeout"

	// Restore hook annotations
	podRestoreHookContainerAnnotationKey            = "post.hook.restore.velero.io/container"
	podRestoreHookCommandAnnotationKey              = "post.hook.restore.velero.io/command"
	podRestoreHookOnErrorAnnotationKey              = "post.hook.restore.velero.io/on-error"
	podRestoreHookTimeoutAnnotationKey              = "post.hook.restore.velero.io/exec-timeout"
	podRestoreHookWaitTimeoutAnnotationKey          = "post.hook.restore.velero.io/wait-timeout"
	podRestoreHookInitContainerImageAnnotationKey   = "init.hook.restore.velero.io/container-image"
	podRestoreHookInitContainerNameAnnotationKey    = "init.hook.restore.velero.io/container-name"
	podRestoreHookInitContainerCommandAnnotationKey = "init.hook.restore.velero.io/command"
	podRestoreHookInitContainerTimeoutAnnotationKey = "init.hook.restore.velero.io/timeout"
)

// ItemHookHandler invokes hooks for an item.
type ItemHookHandler interface {
	// HandleHooks invokes hooks for an item. If the item is a pod and the appropriate annotations exist
	// to specify a hook, that is executed. Otherwise, this looks at the backup context's Backup to
	// determine if there are any hooks relevant to the item, taking into account the hook spec's
	// namespaces, resources, and label selector.
	HandleHooks(
		log logrus.FieldLogger,
		groupResource schema.GroupResource,
		obj runtime.Unstructured,
		resourceHooks []ResourceHook,
		phase hookPhase,
	) error
}

// ItemRestoreHookHandler invokes restore hooks for an item
type ItemRestoreHookHandler interface {
	HandleRestoreHooks(
		log logrus.FieldLogger,
		groupResource schema.GroupResource,
		obj runtime.Unstructured,
		rh []ResourceRestoreHook,
	) (runtime.Unstructured, error)
}

// InitContainerRestoreHookHandler is the restore hook handler to add init containers to restored pods.
type InitContainerRestoreHookHandler struct{}

// HandleRestoreHooks runs the restore hooks for an item.
// If the item is a pod, then hooks are chosen to be run as follows:
// If the pod has the appropriate annotations specifying the hook action, then hooks from the annotation are run
// Otherwise, the supplied ResourceRestoreHooks are applied.
func (i *InitContainerRestoreHookHandler) HandleRestoreHooks(
	log logrus.FieldLogger,
	groupResource schema.GroupResource,
	obj runtime.Unstructured,
	resourceRestoreHooks []ResourceRestoreHook,
) (runtime.Unstructured, error) {
	// We only support hooks on pods right now
	if groupResource != kuberesource.Pods {
		return nil, nil
	}

	metadata, err := meta.Accessor(obj)
	if err != nil {
		return nil, errors.Wrap(err, "unable to get a metadata accessor")
	}
	pod := new(corev1api.Pod)
	if err := runtime.DefaultUnstructuredConverter.FromUnstructured(obj.UnstructuredContent(), pod); err != nil {
		return nil, errors.WithStack(err)
	}

	initContainers := []corev1api.Container{}
	// If this pod had pod volumes backed up using restic, then we want to the pod volumes restored prior to
	// running the restore hook init containers. This allows the restore hook init containers to prepare the
	// restored data to be consumed by the application container(s).
	// So if there is a "restic-wait" init container already on the pod at index 0, we'll preserve that and run
	// it before running any other init container.
	if len(pod.Spec.InitContainers) > 0 && pod.Spec.InitContainers[0].Name == restic.InitContainer {
		initContainers = append(initContainers, pod.Spec.InitContainers[0])
		pod.Spec.InitContainers = pod.Spec.InitContainers[1:]
	}

	hooksFromAnnotations := getInitRestoreHookFromAnnotation(kube.NamespaceAndName(pod), metadata.GetAnnotations(), log)
	if hooksFromAnnotations != nil {
		log.Infof("Handling InitRestoreHooks from pod annotations")
		initContainers = append(initContainers, hooksFromAnnotations.InitContainers...)
	} else {
		log.Infof("Handling InitRestoreHooks from RestoreSpec")
		// pod did not have the annotations appropriate for restore hooks
		// running applicable ResourceRestoreHooks supplied.
		namespace := metadata.GetNamespace()
		labels := labels.Set(metadata.GetLabels())

		for _, rh := range resourceRestoreHooks {
			if !rh.Selector.applicableTo(groupResource, namespace, labels) {
				continue
			}
			for _, hook := range rh.RestoreHooks {
				if hook.Init != nil {
					initContainers = append(initContainers, hook.Init.InitContainers...)
				}
			}
		}
	}

	pod.Spec.InitContainers = append(initContainers, pod.Spec.InitContainers...)
	log.Infof("Returning pod %s/%s with %d init container(s)", pod.Namespace, pod.Name, len(pod.Spec.InitContainers))

	podMap, err := runtime.DefaultUnstructuredConverter.ToUnstructured(&pod)
	if err != nil {
		return nil, errors.WithStack(err)
	}
	return &unstructured.Unstructured{Object: podMap}, nil
}

// DefaultItemHookHandler is the default itemHookHandler.
type DefaultItemHookHandler struct {
	PodCommandExecutor podexec.PodCommandExecutor
}

func (h *DefaultItemHookHandler) HandleHooks(
	log logrus.FieldLogger,
	groupResource schema.GroupResource,
	obj runtime.Unstructured,
	resourceHooks []ResourceHook,
	phase hookPhase,
) error {
	// We only support hooks on pods right now
	if groupResource != kuberesource.Pods {
		return nil
	}

	metadata, err := meta.Accessor(obj)
	if err != nil {
		return errors.Wrap(err, "unable to get a metadata accessor")
	}

	namespace := metadata.GetNamespace()
	name := metadata.GetName()

	// If the pod has the hook specified via annotations, that takes priority.
	hookFromAnnotations := getPodExecHookFromAnnotations(metadata.GetAnnotations(), phase, log)
	if phase == PhasePre && hookFromAnnotations == nil {
		// See if the pod has the legacy hook annotation keys (i.e. without a phase specified)
		hookFromAnnotations = getPodExecHookFromAnnotations(metadata.GetAnnotations(), "", log)
	}
	if hookFromAnnotations != nil {
		hookLog := log.WithFields(
			logrus.Fields{
				"hookSource": "annotation",
				"hookType":   "exec",
				"hookPhase":  phase,
			},
		)
		if err := h.PodCommandExecutor.ExecutePodCommand(hookLog, obj.UnstructuredContent(), namespace, name, "<from-annotation>", hookFromAnnotations); err != nil {
			hookLog.WithError(err).Error("Error executing hook")
			if hookFromAnnotations.OnError == velerov1api.HookErrorModeFail {
				return err
			}
		}

		return nil
	}

	labels := labels.Set(metadata.GetLabels())
	// Otherwise, check for hooks defined in the backup spec.
	for _, resourceHook := range resourceHooks {
		if !resourceHook.Selector.applicableTo(groupResource, namespace, labels) {
			continue
		}

		var hooks []velerov1api.BackupResourceHook
		if phase == PhasePre {
			hooks = resourceHook.Pre
		} else {
			hooks = resourceHook.Post
		}
		for _, hook := range hooks {
			if groupResource == kuberesource.Pods {
				if hook.Exec != nil {
					hookLog := log.WithFields(
						logrus.Fields{
							"hookSource": "backupSpec",
							"hookType":   "exec",
							"hookPhase":  phase,
						},
					)
					err := h.PodCommandExecutor.ExecutePodCommand(hookLog, obj.UnstructuredContent(), namespace, name, resourceHook.Name, hook.Exec)
					if err != nil {
						hookLog.WithError(err).Error("Error executing hook")
						if hook.Exec.OnError == velerov1api.HookErrorModeFail {
							return err
						}
					}
				}
			}
		}
	}

	return nil
}

func phasedKey(phase hookPhase, key string) string {
	if phase != "" {
		return fmt.Sprintf("%v.%v", phase, key)
	}
	return string(key)
}

func getHookAnnotation(annotations map[string]string, key string, phase hookPhase) string {
	return annotations[phasedKey(phase, key)]
}

// getPodExecHookFromAnnotations returns an ExecHook based on the annotations, as long as the
// 'command' annotation is present. If it is absent, this returns nil.
// If there is an error in parsing a supplied timeout, it is logged.
func getPodExecHookFromAnnotations(annotations map[string]string, phase hookPhase, log logrus.FieldLogger) *velerov1api.ExecHook {
	commandValue := getHookAnnotation(annotations, podBackupHookCommandAnnotationKey, phase)
	if commandValue == "" {
		return nil
	}

	container := getHookAnnotation(annotations, podBackupHookContainerAnnotationKey, phase)

	onError := velerov1api.HookErrorMode(getHookAnnotation(annotations, podBackupHookOnErrorAnnotationKey, phase))
	if onError != velerov1api.HookErrorModeContinue && onError != velerov1api.HookErrorModeFail {
		onError = ""
	}

	var timeout time.Duration
	timeoutString := getHookAnnotation(annotations, podBackupHookTimeoutAnnotationKey, phase)
	if timeoutString != "" {
		if temp, err := time.ParseDuration(timeoutString); err == nil {
			timeout = temp
		} else {
			log.Warn(errors.Wrapf(err, "Unable to parse provided timeout %s, using default", timeoutString))
		}
	}

	return &velerov1api.ExecHook{
		Container: container,
		Command:   parseStringToCommand(commandValue),
		OnError:   onError,
		Timeout:   metav1.Duration{Duration: timeout},
	}
}

func parseStringToCommand(commandValue string) []string {
	var command []string
	// check for json array
	if commandValue[0] == '[' {
		if err := json.Unmarshal([]byte(commandValue), &command); err != nil {
			command = []string{commandValue}
		}
	} else {
		command = append(command, commandValue)
	}
	return command
}

type ResourceHookSelector struct {
	Namespaces    *collections.IncludesExcludes
	Resources     *collections.IncludesExcludes
	LabelSelector labels.Selector
}

// ResourceHook is a hook for a given resource.
type ResourceHook struct {
	Name     string
	Selector ResourceHookSelector
	Pre      []velerov1api.BackupResourceHook
	Post     []velerov1api.BackupResourceHook
}

func (r ResourceHookSelector) applicableTo(groupResource schema.GroupResource, namespace string, labels labels.Set) bool {
	if r.Namespaces != nil && !r.Namespaces.ShouldInclude(namespace) {
		return false
	}
	if r.Resources != nil && !r.Resources.ShouldInclude(groupResource.String()) {
		return false
	}
	if r.LabelSelector != nil && !r.LabelSelector.Matches(labels) {
		return false
	}
	return true
}

// ResourceRestoreHook is a restore hook for a given resource.
type ResourceRestoreHook struct {
	Name         string
	Selector     ResourceHookSelector
	RestoreHooks []velerov1api.RestoreResourceHook
}

func getInitRestoreHookFromAnnotation(podName string, annotations map[string]string, log logrus.FieldLogger) *velerov1api.InitRestoreHook {
	containerImage := annotations[podRestoreHookInitContainerImageAnnotationKey]
	containerName := annotations[podRestoreHookInitContainerNameAnnotationKey]
	command := annotations[podRestoreHookInitContainerCommandAnnotationKey]
	if containerImage == "" {
		log.Infof("Pod %s has no %s annotation, no initRestoreHook in annotation", podName, podRestoreHookInitContainerImageAnnotationKey)
		return nil
	}
	if command == "" {
		log.Infof("RestoreHook init container for pod %s is using container's default entrypoint", podName, containerImage)
	}
	if containerName == "" {
		uid, err := uuid.NewV4()
		uuidStr := "deadfeed"
		if err != nil {
			log.Errorf("Failed to generate UUID for container name")
		} else {
			uuidStr = strings.Split(uid.String(), "-")[0]
		}
		containerName = fmt.Sprintf("velero-restore-init-%s", uuidStr)
		log.Infof("Pod %s has no %s annotation, using generated name %s for initContainer", podName, podRestoreHookInitContainerNameAnnotationKey, containerName)
	}

	return &velerov1api.InitRestoreHook{
		InitContainers: []corev1api.Container{
			{
				Image:   containerImage,
				Name:    containerName,
				Command: parseStringToCommand(command),
			},
		},
	}
}

// GetRestoreHooksFromSpec returns a list of ResourceRestoreHooks from the restore Spec.
func GetRestoreHooksFromSpec(hooksSpec *velerov1api.RestoreHooks) ([]ResourceRestoreHook, error) {
	if hooksSpec == nil {
		return []ResourceRestoreHook{}, nil
	}
	restoreHooks := make([]ResourceRestoreHook, 0, len(hooksSpec.Resources))
	for _, rs := range hooksSpec.Resources {
		rh := ResourceRestoreHook{
			Name: rs.Name,
			Selector: ResourceHookSelector{
				Namespaces: collections.NewIncludesExcludes().Includes(rs.IncludedNamespaces...).Excludes(rs.ExcludedNamespaces...),
				// these hooks ara applicable only to pods.
				// TODO: resolve the pods resource via discovery?
				Resources: collections.NewIncludesExcludes().Includes(kuberesource.Pods.Resource),
			},
			// TODO does this work for ExecRestoreHook as well?
			RestoreHooks: rs.PostHooks,
		}

		if rs.LabelSelector != nil {
			ls, err := metav1.LabelSelectorAsSelector(rs.LabelSelector)
			if err != nil {
				return nil, errors.WithStack(err)
			}
			rh.Selector.LabelSelector = ls
		}
		restoreHooks = append(restoreHooks, rh)
	}

	return restoreHooks, nil
}

// getPodExecRestoreHookFromAnnotations returns an ExecRestoreHook based on restore annotations, as
// long as the 'command' annotation is present. If it is absent, this returns nil.
func getPodExecRestoreHookFromAnnotations(annotations map[string]string, log logrus.FieldLogger) *velerov1api.ExecRestoreHook {
	commandValue := annotations[podRestoreHookCommandAnnotationKey]
	if commandValue == "" {
		return nil
	}

	container := annotations[podRestoreHookContainerAnnotationKey]

	onError := velerov1api.HookErrorMode(annotations[podRestoreHookOnErrorAnnotationKey])
	if onError != velerov1api.HookErrorModeContinue && onError != velerov1api.HookErrorModeFail {
		onError = ""
	}

	var execTimeout time.Duration
	execTimeoutString := annotations[podRestoreHookTimeoutAnnotationKey]
	if execTimeoutString != "" {
		if temp, err := time.ParseDuration(execTimeoutString); err == nil {
			execTimeout = temp
		} else {
			log.Warn(errors.Wrapf(err, "Unable to parse exec timeout %s, ignoring", execTimeoutString))
		}
	}

	var waitTimeout time.Duration
	waitTimeoutString := annotations[podRestoreHookWaitTimeoutAnnotationKey]
	if waitTimeoutString != "" {
		if temp, err := time.ParseDuration(waitTimeoutString); err == nil {
			waitTimeout = temp
		} else {
			log.Warn(errors.Wrapf(err, "Unable to parse wait timeout %s, ignoring", waitTimeoutString))
		}
	}

	return &velerov1api.ExecRestoreHook{
		Container:   container,
		Command:     parseStringToCommand(commandValue),
		OnError:     onError,
		ExecTimeout: metav1.Duration{Duration: execTimeout},
		WaitTimeout: metav1.Duration{Duration: waitTimeout},
	}
}

type PodExecRestoreHook struct {
	HookName   string
	HookSource string
	Hook       velerov1api.ExecRestoreHook
	executed   bool
}

// GroupRestoreExecHooks returns a list of hooks to be executed in a pod grouped by
// container name. If an exec hook is defined in annotation that is used, else applicable exec
// hooks from the restore resource are accumulated.
func GroupRestoreExecHooks(
	resourceRestoreHooks []ResourceRestoreHook,
	pod *corev1api.Pod,
	log logrus.FieldLogger,
) (map[string][]PodExecRestoreHook, error) {
	byContainer := map[string][]PodExecRestoreHook{}

	if pod == nil || len(pod.Spec.Containers) == 0 {
		return byContainer, nil
	}
	metadata, err := meta.Accessor(pod)
	if err != nil {
		return nil, errors.WithStack(err)
	}
	hookFromAnnotation := getPodExecRestoreHookFromAnnotations(metadata.GetAnnotations(), log)
	if hookFromAnnotation != nil {
		// default to first container in pod if unset
		if hookFromAnnotation.Container == "" {
			hookFromAnnotation.Container = pod.Spec.Containers[0].Name
		}
		byContainer[hookFromAnnotation.Container] = []PodExecRestoreHook{
			{
				HookName:   "<from-annotation>",
				HookSource: "annotation",
				Hook:       *hookFromAnnotation,
			},
		}
		return byContainer, nil
	}

	// No hook found on pod's annotations so check for applicable hooks from the restore spec
	labels := metadata.GetLabels()
	namespace := metadata.GetNamespace()
	for _, rrh := range resourceRestoreHooks {
		if !rrh.Selector.applicableTo(kuberesource.Pods, namespace, labels) {
			continue
		}
		for _, rh := range rrh.RestoreHooks {
			if rh.Exec == nil {
				continue
			}
			named := PodExecRestoreHook{
				HookName:   rrh.Name,
				Hook:       *rh.Exec,
				HookSource: "backupSpec",
			}
			// default to first container in pod if unset, without mutating resource restore hook
			if named.Hook.Container == "" {
				named.Hook.Container = pod.Spec.Containers[0].Name
			}
			byContainer[named.Hook.Container] = append(byContainer[named.Hook.Container], named)
		}
	}

	return byContainer, nil
}