Add support for workload identity

Signed-off-by: Md. Ishtiaq Islam <[email protected]>

Makefile

@@ -56,7 +56,7 @@ ARCH := $(if $(GOARCH),$(GOARCH),$(shell go env GOARCH))
 GO_VERSION       ?= 1.20
 BUILD_IMAGE      ?= ghcr.io/appscode/golang-dev:$(GO_VERSION)
 
-RESTIC_VER  := 0.15.1
+RESTIC_VER  := 0.16.0
 
 OUTBIN = bin/$(BIN)-$(OS)-$(ARCH)
 ifeq ($(OS),windows)

go.mod

@@ -21,9 +21,9 @@ require (
 	k8s.io/kubectl v0.25.1
 	k8s.io/utils v0.0.0-20220823124924-e9cbc92d1a73
 	kmodules.xyz/client-go v0.25.27
-	kmodules.xyz/offshoot-api v0.25.2
+	kmodules.xyz/offshoot-api v0.25.5-0.20231004063620-dff8fb030d57
 	kmodules.xyz/prober v0.25.0
-	kubestash.dev/apimachinery v0.1.1-0.20231010125857-fc7d815460f0
+	kubestash.dev/apimachinery v0.1.1-0.20231016090735-4ae6a0befe35
 	sigs.k8s.io/controller-runtime v0.13.1
 )
 

go.sum

@@ -1067,12 +1067,12 @@ kmodules.xyz/client-go v0.25.27 h1:Ivl054xbXSvMNMKAJtK7TkS0iZX0AHvVQzxtsrIk9ik=
 kmodules.xyz/client-go v0.25.27/go.mod h1:PYfJtJs+AhgfkJNIeUObU4SqAkY85ARTlXxC+2gAsgo=
 kmodules.xyz/objectstore-api v0.25.1 h1:lYQlxk+edgZYakhq+OoRBXTbHbZTGKhatGZWnKixgEQ=
 kmodules.xyz/objectstore-api v0.25.1/go.mod h1:6wBtktN7/EXyE429OTCB9nwEe+d0ADaoCtm6+IZnJso=
-kmodules.xyz/offshoot-api v0.25.2 h1:71x+MB0AwD9UhzlNGRPAFuB3HVMIAjtXi2yB0IIs+1Y=
-kmodules.xyz/offshoot-api v0.25.2/go.mod h1:PUk4EuJFhhyQykCflHj7EgXcljGIqs9vi0IN0RpxtY4=
+kmodules.xyz/offshoot-api v0.25.5-0.20231004063620-dff8fb030d57 h1:4ld/ujZum5JX6cfYrW1yQErAl1AONEfJa8TDyRpnHuU=
+kmodules.xyz/offshoot-api v0.25.5-0.20231004063620-dff8fb030d57/go.mod h1:PUk4EuJFhhyQykCflHj7EgXcljGIqs9vi0IN0RpxtY4=
 kmodules.xyz/prober v0.25.0 h1:R5uRLHJEvEtEoogj+vaTAob0Btph6+PX5IlS6hPh8PA=
 kmodules.xyz/prober v0.25.0/go.mod h1:z4RTnjaajNQa/vPltsiOnO3xI716I/ziD2ac2Exm+1M=
-kubestash.dev/apimachinery v0.1.1-0.20231010125857-fc7d815460f0 h1:gT03pLKkz02sb3FEuJY4gyZpnab78q2kM5QogNzQvvU=
-kubestash.dev/apimachinery v0.1.1-0.20231010125857-fc7d815460f0/go.mod h1:tSEmpY25qz3lyw4kmX8r4ZvtW/G08V/OL1CcFnFetQ0=
+kubestash.dev/apimachinery v0.1.1-0.20231016090735-4ae6a0befe35 h1:j03vCIago5xPegnBZ5YJXzh87tQj+JHqJylnwicyBQ8=
+kubestash.dev/apimachinery v0.1.1-0.20231016090735-4ae6a0befe35/go.mod h1:TBN95XMI/LYDcmRf3kb/lXbfsOx0E0ZKRZDyQKyzAUE=
 rsc.io/binaryregexp v0.2.0/go.mod h1:qTv7/COck+e2FymRvadv62gMdZztPaShugOCi3I+8D8=
 rsc.io/quote/v3 v3.1.0/go.mod h1:yEA65RcK8LyAZtP9Kv3t0HmxON59tX3rD+tICJqUlj0=
 rsc.io/sampler v1.3.0/go.mod h1:T1hPZKmBbMNahiBKFy5HrXp6adAjACjK9JXDnKaTXpA=

pkg/constants.go

@@ -20,9 +20,12 @@ import "time"
 
 // These variables will be set during build time
 const (
-	ScratchDir     = "/tmp/scratch"
-	DestinationDir = "/tmp/destination"
-	configDirName  = "config"
+	CmdKubectl = "kubectl"
+
+	ScratchDir          = "/tmp/scratch"
+	DestinationDir      = "/tmp/destination"
+	SnapshotDownloadDir = "/tmp/snapshot"
+	configDirName       = "config"
 
 	ResticEnvs     = "restic-envs"
 	ResticRegistry = "restic"
@@ -55,7 +58,3 @@ const (
 	PVCSchedule    = "*/59 * * * *"
 	LatestSnapshot = "latest"
 )
-
-const (
-	CmdKubectl = "kubectl"
-)

pkg/download.go

@@ -17,7 +17,6 @@ limitations under the License.
 package pkg
 
 import (
-	"context"
 	"fmt"
 	"os"
 	"os/exec"
@@ -28,7 +27,6 @@ import (
 	"github.com/pkg/errors"
 	"github.com/spf13/cobra"
 	core "k8s.io/api/core/v1"
-	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/cli-runtime/pkg/genericclioptions"
 	"k8s.io/client-go/rest"
 	"k8s.io/klog/v2"
@@ -39,7 +37,6 @@ import (
 	"kubestash.dev/apimachinery/apis"
 	storageapi "kubestash.dev/apimachinery/apis/storage/v1alpha1"
 	"kubestash.dev/apimachinery/pkg/restic"
-	"sigs.k8s.io/controller-runtime/pkg/client"
 )
 
 type downloadOptions struct {
@@ -81,17 +78,26 @@ func NewCmdDownload(clientGetter genericclioptions.RESTClientGetter) *cobra.Comm
 				return err
 			}
 
-			snapshot, err := downloadOpt.getSnapshot(snapshotName)
+			snapshot, err := getSnapshot(kmapi.ObjectReference{
+				Name:      snapshotName,
+				Namespace: srcNamespace,
+			})
 			if err != nil {
 				return err
 			}
 
-			repository, err := downloadOpt.getRepository(snapshot.Spec.Repository)
+			repository, err := getRepository(kmapi.ObjectReference{
+				Name:      snapshot.Spec.Repository,
+				Namespace: srcNamespace,
+			})
 			if err != nil {
 				return err
 			}
 
-			backupStorage, err := downloadOpt.getBackupStorage(repository.Spec.StorageRef)
+			backupStorage, err := getBackupStorage(kmapi.ObjectReference{
+				Name:      repository.Spec.StorageRef.Name,
+				Namespace: repository.Spec.StorageRef.Namespace,
+			})
 			if err != nil {
 				return err
 			}
@@ -110,20 +116,21 @@ func NewCmdDownload(clientGetter genericclioptions.RESTClientGetter) *cobra.Comm
 					return err
 				}
 
-				if err := downloadOpt.runRestoreViaPod(accessorPod, snapshotName); err != nil {
-					return err
-				}
+				return downloadOpt.runRestoreViaPod(accessorPod, snapshotName)
+			}
 
-				if err := downloadOpt.copyDownloadedDataToDestination(accessorPod); err != nil {
-					return err
-				}
+			operatorPod, err := getOperatorPod()
+			if err != nil {
+				return err
+			}
 
-				if err := downloadOpt.clearDataFromPod(accessorPod); err != nil {
-					return err
-				}
+			yes, err := isWorkloadIdentity(operatorPod)
+			if err != nil {
+				return err
+			}
 
-				klog.Infof("Snapshot %s/%s restored in path %s", srcNamespace, snapshotName, downloadOpt.destinationDir)
-				return nil
+			if yes {
+				return downloadOpt.runRestoreViaPod(&operatorPod, snapshotName)
 			}
 
 			if err = os.MkdirAll(ScratchDir, 0o755); err != nil {
@@ -208,51 +215,29 @@ func NewCmdDownload(clientGetter genericclioptions.RESTClientGetter) *cobra.Comm
 	return cmd
 }
 
-func (opt *downloadOptions) getSnapshot(snapshotName string) (*storageapi.Snapshot, error) {
-	snapshot := &storageapi.Snapshot{
-		ObjectMeta: metav1.ObjectMeta{
-			Name:      snapshotName,
-			Namespace: srcNamespace,
-		},
-	}
-	if err := klient.Get(context.Background(), client.ObjectKeyFromObject(snapshot), snapshot); err != nil {
-		return nil, err
+func (opt *downloadOptions) runRestoreViaPod(pod *core.Pod, snapshotName string) error {
+	if err := opt.runProbe(pod, snapshotName); err != nil {
+		return err
 	}
-	return snapshot, nil
-}
 
-func (opt *downloadOptions) getRepository(repoName string) (*storageapi.Repository, error) {
-	repository := &storageapi.Repository{
-		ObjectMeta: metav1.ObjectMeta{
-			Name:      repoName,
-			Namespace: srcNamespace,
-		},
-	}
-	if err := klient.Get(context.Background(), client.ObjectKeyFromObject(repository), repository); err != nil {
-		return nil, err
+	if err := opt.copyDownloadedDataToDestination(pod); err != nil {
+		return err
 	}
-	return repository, nil
-}
 
-func (opt *downloadOptions) getBackupStorage(storage kmapi.TypedObjectReference) (*storageapi.BackupStorage, error) {
-	backupStorage := &storageapi.BackupStorage{
-		ObjectMeta: metav1.ObjectMeta{
-			Name:      storage.Name,
-			Namespace: storage.Namespace,
-		},
-	}
-	if err := klient.Get(context.Background(), client.ObjectKeyFromObject(backupStorage), backupStorage); err != nil {
-		return nil, err
+	if err := opt.clearDataFromPod(pod); err != nil {
+		return err
 	}
-	return backupStorage, nil
+
+	klog.Infof("Snapshot %s/%s restored in path %s", srcNamespace, snapshotName, opt.destinationDir)
+	return nil
 }
 
-func (opt *downloadOptions) runRestoreViaPod(pod *core.Pod, snapshotName string) error {
+func (opt *downloadOptions) runProbe(pod *core.Pod, snapshotName string) error {
 	command := []string{
 		"/kubestash",
 		"download", snapshotName,
 		"--namespace", srcNamespace,
-		"--destination", getPodDirForSnapshot(),
+		"--destination", SnapshotDownloadDir,
 	}
 
 	if len(opt.components) != 0 {
@@ -269,33 +254,29 @@ func (opt *downloadOptions) runRestoreViaPod(pod *core.Pod, snapshotName string)
 
 	action := &prober.Handler{
 		Exec:          &core.ExecAction{Command: command},
-		ContainerName: apis.AccessorContainerName,
+		ContainerName: apis.OperatorContainer,
 	}
 
 	return probe.RunProbe(opt.restConfig, action, pod.Name, pod.Namespace)
 }
 
 func (opt *downloadOptions) copyDownloadedDataToDestination(pod *core.Pod) error {
-	_, err := exec.Command(CmdKubectl, "cp", "--namespace", pod.Namespace, fmt.Sprintf("%s/%s:%s", pod.Namespace, pod.Name, getPodDirForSnapshot()), opt.destinationDir).CombinedOutput()
+	_, err := exec.Command(CmdKubectl, "cp", "--namespace", pod.Namespace, fmt.Sprintf("%s/%s:%s", pod.Namespace, pod.Name, SnapshotDownloadDir), opt.destinationDir).CombinedOutput()
 	if err != nil {
 		return err
 	}
 	return nil
 }
 
 func (opt *downloadOptions) clearDataFromPod(pod *core.Pod) error {
-	cmd := []string{"rm", "-rf", getPodDirForSnapshot()}
+	cmd := []string{"rm", "-rf", SnapshotDownloadDir}
 	action := &prober.Handler{
 		Exec:          &core.ExecAction{Command: cmd},
-		ContainerName: apis.AccessorContainerName, // TODO: need to change for different pod
+		ContainerName: apis.OperatorContainer,
 	}
 	return probe.RunProbe(opt.restConfig, action, pod.Name, pod.Namespace)
 }
 
-func getPodDirForSnapshot() string {
-	return filepath.Join(apis.ScratchDirMountPath, apis.SnapshotDownloadDir)
-}
-
 func (opt *downloadOptions) prepareDestinationDir() (err error) {
 	// if destination flag is not specified, restore in current directory
 	if opt.destinationDir == "" {

pkg/unlock.go

@@ -103,10 +103,24 @@ func NewCmdUnlockRepository(clientGetter genericclioptions.RESTClientGetter) *co
 					return err
 				}
 
-				return unlockOpt.unlockLocalRepository(accessorPod)
+				return unlockOpt.unlockRepositoryViaPod(accessorPod)
 			}
 
-			return unlockOpt.unlockRepository()
+			operatorPod, err := getOperatorPod()
+			if err != nil {
+				return err
+			}
+
+			yes, err := isWorkloadIdentity(operatorPod)
+			if err != nil {
+				return err
+			}
+
+			if yes {
+				return unlockOpt.unlockRepositoryViaPod(&operatorPod)
+			}
+
+			return unlockOpt.unlockRepositoryViaDocker()
 		},
 	}
 
@@ -115,7 +129,7 @@ func NewCmdUnlockRepository(clientGetter genericclioptions.RESTClientGetter) *co
 	return cmd
 }
 
-func (opt *unlockOptions) unlockLocalRepository(pod *core.Pod) error {
+func (opt *unlockOptions) unlockRepositoryViaPod(pod *core.Pod) error {
 	command := []string{
 		"/kubestash",
 		"unlock", opt.repo.Name,
@@ -125,13 +139,13 @@ func (opt *unlockOptions) unlockLocalRepository(pod *core.Pod) error {
 
 	action := &prober.Handler{
 		Exec:          &core.ExecAction{Command: command},
-		ContainerName: apis.AccessorContainerName,
+		ContainerName: apis.OperatorContainer,
 	}
 
 	return probe.RunProbe(opt.restConfig, action, pod.Name, pod.Namespace)
 }
 
-func (opt *unlockOptions) unlockRepository() error {
+func (opt *unlockOptions) unlockRepositoryViaDocker() error {
 	for _, path := range opt.paths {
 		setupOptions := restic.SetupOptions{
 			Client:           klient,
@@ -207,9 +221,6 @@ func (opt *unlockOptions) runCmdViaDocker() error {
 	args = append(args, opt.extraArgs...)
 	klog.Infoln("Running docker with args:", args)
 	out, err := exec.Command("docker", args...).CombinedOutput()
-	if len(out) == 0 {
-		return fmt.Errorf("lock not stale")
-	}
 	klog.Infoln("Output:", string(out))
 	return err
 }

pkg/util.go

@@ -20,6 +20,7 @@ import (
 	"context"
 	"fmt"
 	"os"
+	"strconv"
 	"strings"
 
 	vsapi "github.com/kubernetes-csi/external-snapshotter/client/v4/apis/volumesnapshot/v1"
@@ -93,6 +94,21 @@ func getSecret(ref kmapi.ObjectReference) (*core.Secret, error) {
 	return secret, nil
 }
 
+func getServiceAccount(ref kmapi.ObjectReference) (*core.ServiceAccount, error) {
+	sa := &core.ServiceAccount{
+		ObjectMeta: metav1.ObjectMeta{
+			Name:      ref.Name,
+			Namespace: ref.Namespace,
+		},
+	}
+
+	if err := klient.Get(context.Background(), client.ObjectKeyFromObject(sa), sa); err != nil {
+		return nil, err
+	}
+
+	return sa, nil
+}
+
 func getPVC(ref kmapi.ObjectReference) (*core.PersistentVolumeClaim, error) {
 	pvc := &core.PersistentVolumeClaim{
 		ObjectMeta: metav1.ObjectMeta{
@@ -153,6 +169,21 @@ func getRepository(ref kmapi.ObjectReference) (*storageapi.Repository, error) {
 	return repo, nil
 }
 
+func getSnapshot(ref kmapi.ObjectReference) (*storageapi.Snapshot, error) {
+	snap := &storageapi.Snapshot{
+		ObjectMeta: metav1.ObjectMeta{
+			Name:      ref.Name,
+			Namespace: ref.Namespace,
+		},
+	}
+
+	if err := klient.Get(context.Background(), client.ObjectKeyFromObject(snap), snap); err != nil {
+		return nil, err
+	}
+
+	return snap, nil
+}
+
 func getBackupConfiguration(ref kmapi.ObjectReference) (*coreapi.BackupConfiguration, error) {
 	bc := &coreapi.BackupConfiguration{
 		ObjectMeta: metav1.ObjectMeta{
@@ -308,3 +339,33 @@ func hasVolumeMount(mounts []core.VolumeMount, name string) bool {
 	}
 	return false
 }
+
+func isWorkloadIdentity(pod core.Pod) (bool, error) {
+	azureLabel := "azure.workload.identity/use"
+	googleAnnotation := "iam.gke.io/gcp-service-account"
+	awsAnnotation := "eks.amazonaws.com/role-arn"
+
+	if value, exists := pod.Labels[azureLabel]; exists {
+		boolValue, err := strconv.ParseBool(value)
+		if err != nil {
+			return false, err
+		}
+		return boolValue, nil
+	}
+
+	sa, err := getServiceAccount(kmapi.ObjectReference{
+		Name:      pod.Spec.ServiceAccountName,
+		Namespace: pod.Namespace,
+	})
+	if err != nil {
+		return false, err
+	}
+
+	if _, exists := sa.Annotations[googleAnnotation]; exists {
+		return true, nil
+	} else if _, exists := sa.Annotations[awsAnnotation]; exists {
+		return true, nil
+	}
+
+	return false, nil
+}