github.com/Azure/aad-pod-identity@v1.8.17/test/image/identityvalidator/identityvalidator.go (about)

     1  package main
     2  
     3  import (
     4  	"flag"
     5  	"os"
     6  	"strings"
     7  	"sync"
     8  	"time"
     9  
    10  	"github.com/Azure/azure-sdk-for-go/services/keyvault/2016-10-01/keyvault"
    11  	"k8s.io/klog/v2"
    12  )
    13  
    14  type assertFunction func() error
    15  
    16  const (
    17  	contextTimeout = 80 * time.Second
    18  )
    19  
    20  var (
    21  	sleep                 bool
    22  	subscriptionID        string
    23  	identityClientID      string
    24  	identityResourceID    string
    25  	keyvaultName          string
    26  	keyvaultSecretName    string
    27  	keyvaultSecretVersion string
    28  	keyvaultSecretValue   string
    29  )
    30  
    31  func init() {
    32  	flag.BoolVar(&sleep, "sleep", false, "Set to true to enter sleep mode")
    33  	flag.StringVar(&subscriptionID, "subscription-id", "", "subscription id for test")
    34  	flag.StringVar(&identityClientID, "identity-client-id", "", "client id for the msi id")
    35  	flag.StringVar(&identityResourceID, "identity-resource-id", "", "resource id for the msi id")
    36  	flag.StringVar(&keyvaultName, "keyvault-name", "", "the name of the keyvault to extract the secret from")
    37  	flag.StringVar(&keyvaultSecretName, "keyvault-secret-name", "", "the name of the keyvault secret we are extracting with pod identity")
    38  	flag.StringVar(&keyvaultSecretVersion, "keyvault-secret-version", "", "the version of the keyvault secret we are extracting with pod identity")
    39  	flag.StringVar(&keyvaultSecretValue, "keyvault-secret-value", "test-value", "the version of the keyvault secret we are extracting with pod identity")
    40  }
    41  
    42  func main() {
    43  	flag.Parse()
    44  
    45  	if sleep {
    46  		klog.Infof("entering sleep mode")
    47  		for {
    48  			select {}
    49  		}
    50  	}
    51  
    52  	podname := os.Getenv("E2E_TEST_POD_NAME")
    53  	podnamespace := os.Getenv("E2E_TEST_POD_NAMESPACE")
    54  	podip := os.Getenv("E2E_TEST_POD_IP")
    55  
    56  	klog.Infof("starting identity validator pod %s/%s with pod IP %s", podnamespace, podname, podip)
    57  
    58  	kvt := &keyvaultTester{
    59  		client:             keyvault.New(),
    60  		subscriptionID:     subscriptionID,
    61  		identityClientID:   identityClientID,
    62  		identityResourceID: identityResourceID,
    63  		keyvaultName:       keyvaultName,
    64  		secretName:         keyvaultSecretName,
    65  		secretVersion:      keyvaultSecretVersion,
    66  		secretValue:        keyvaultSecretValue,
    67  	}
    68  
    69  	var wg sync.WaitGroup
    70  	errCh := make(chan error, 3)
    71  
    72  	for _, assert := range []assertFunction{
    73  		kvt.assertWithIdentityClientID,
    74  		kvt.assertWithIdentityResourceID,
    75  		assertWithSystemAssignedIdentity,
    76  	} {
    77  		wg.Add(1)
    78  		go func(assert assertFunction) {
    79  			defer wg.Done()
    80  			var err error
    81  			// allow at most 10 retries if we encounter "Identity not found" error
    82  			// The 10 retries is because of the IMDS cache bug that resolves sometimes after a minute
    83  			for i := 0; i < 10; i++ {
    84  				err = assert()
    85  				if !isIdentityNotFoundError(err) {
    86  					break
    87  				}
    88  				klog.Infof("got identity not found error, retrying in 10 seconds... (%d/7)", i+1)
    89  				if i < 10 {
    90  					time.Sleep(10 * time.Second)
    91  				}
    92  			}
    93  			errCh <- err
    94  		}(assert)
    95  	}
    96  	wg.Wait()
    97  
    98  	close(errCh)
    99  
   100  	hasError := false
   101  	for err := range errCh {
   102  		if err != nil {
   103  			hasError = true
   104  			klog.Error(err)
   105  		}
   106  	}
   107  
   108  	if hasError {
   109  		os.Exit(1)
   110  	}
   111  }
   112  
   113  func isIdentityNotFoundError(err error) bool {
   114  	return err != nil && strings.Contains(err.Error(), "Identity not found")
   115  }