Skip to content
This repository has been archived by the owner on Oct 12, 2023. It is now read-only.

Commit

Permalink
feat: option to set Retry-After in NMI responses (#1114)
Browse files Browse the repository at this point in the history
Adds a new feature flag to enable setting Retry-After header in the
error response from NMI. The error is only when the identity is still
being assigned by NMI or no valid AzureAssignedIdentity is found yet.
This enables SDK's to retry based on the http status code 503 and the
retry after header.

Note: When enabling this feature, the default retries in NMI should be
explicitly disabled to rather rely on the SDK for retries.

--retry-attempts-for-created=1, --retry-attempts-for-assigned=1 and
--find-identity-retry-interval=1. This will force NMI to return
immediately with the Retry-After header 20s.

Signed-off-by: Anish Ramasekar <anish.ramasekar@gmail.com>
  • Loading branch information
aramase authored Jul 23, 2021
1 parent 6e42f63 commit 9e31f9c
Show file tree
Hide file tree
Showing 8 changed files with 153 additions and 26 deletions.
3 changes: 2 additions & 1 deletion cmd/nmi/main.go
Original file line number Diff line number Diff line change
Expand Up @@ -50,6 +50,7 @@ var (
operationMode = pflag.String("operation-mode", "standard", "NMI operation mode")
allowNetworkPluginKubenet = pflag.Bool("allow-network-plugin-kubenet", false, "Allow running aad-pod-identity in cluster with kubenet")
kubeletConfig = pflag.String("kubelet-config", "/etc/default/kubelet", "Path to kubelet default config")
setRetryAfterHeader = pflag.Bool("set-retry-after-header", false, "Set Retry-After header in NMI responses")
)

func main() {
Expand Down Expand Up @@ -112,7 +113,7 @@ func main() {
*forceNamespaced = *forceNamespaced || "true" == os.Getenv("FORCENAMESPACED")
klog.Infof("running NMI in namespaced mode: %v", *forceNamespaced)

s := server.NewServer(*micNamespace, *blockInstanceMetadata, *metadataHeaderRequired)
s := server.NewServer(*micNamespace, *blockInstanceMetadata, *metadataHeaderRequired, *setRetryAfterHeader)
s.KubeClient = client
s.MetadataIP = *metadataIP
s.MetadataPort = *metadataPort
Expand Down
1 change: 1 addition & 0 deletions manifest_staging/charts/aad-pod-identity/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -294,6 +294,7 @@ The following tables list the configurable parameters of the aad-pod-identity ch
| `nmi.findIdentityRetryIntervalInSeconds` | Override retry interval to find assigned identities in seconds | If not provided, default is `5` |
| `nmi.allowNetworkPluginKubenet` | Allow running aad-pod-identity in cluster with kubenet | `false` |
| `nmi.kubeletConfig` | Path to kubelet default config | `/etc/default/kubelet` |
| `nmi.setRetryAfterHeader` | Set Retry-After header in NMI responses | `false` |
| `rbac.enabled` | Create and use RBAC for all aad-pod-identity resources | `true` |
| `rbac.pspEnabled` | If `true`, create and use a restricted pod security policy for AAD Pod Identity pod(s) | `false` |
| `rbac.allowAccessToSecrets` | NMI requires permissions to get secrets when service principal (type: 1) is used in AzureIdentity. If using only MSI (type: 0) in AzureIdentity, secret get permission can be disabled by setting this to false. | `true` |
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -111,6 +111,9 @@ spec:
{{- if .Values.customUserAgent }}
- --custom-user-agent={{ .Values.customUserAgent }}
{{- end }}
{{- if .Values.nmi.setRetryAfterHeader }}
- --set-retry-after-header={{ .Values.nmi.setRetryAfterHeader }}
{{- end }}
env:
{{- if semverCompare "<= 1.6.1-0" .Values.nmi.tag }}
- name: HOST_IP
Expand Down
3 changes: 3 additions & 0 deletions manifest_staging/charts/aad-pod-identity/values.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -241,6 +241,9 @@ nmi:
# default is /etc/default/kubelet
kubeletConfig: "/etc/default/kubelet"

# Set retry-after header in the NMI responses when the identity is still being assigned.
setRetryAfterHeader: false

rbac:
enabled: true
# NMI requires permissions to get secrets when service principal (type: 1) is used in AzureIdentity.
Expand Down
14 changes: 12 additions & 2 deletions pkg/nmi/server/server.go
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,7 @@ const (
hostTokenPathPrefix = "/host/token"
// "/metadata" portion is case-insensitive in IMDS
instancePathPrefix = "/{type:(?i:metadata)}/instance" // #nosec
headerRetryAfter = "Retry-After"
)

// Server encapsulates all of the parameters necessary for starting up
Expand All @@ -51,6 +52,7 @@ type Server struct {
Initialized bool
BlockInstanceMetadata bool
MetadataHeaderRequired bool
SetRetryAfterHeader bool
// TokenClient is client that fetches identities and tokens
TokenClient nmi.TokenClient
Reporter *metrics.Reporter
Expand All @@ -70,7 +72,7 @@ type MetadataResponse struct {
}

// NewServer will create a new Server with default values.
func NewServer(micNamespace string, blockInstanceMetadata bool, metadataHeaderRequired bool) *Server {
func NewServer(micNamespace string, blockInstanceMetadata, metadataHeaderRequired, setRetryAfterHeader bool) *Server {
reporter, err := metrics.NewReporter()
if err != nil {
klog.Errorf("failed to create reporter for metrics, error: %+v", err)
Expand All @@ -84,6 +86,7 @@ func NewServer(micNamespace string, blockInstanceMetadata bool, metadataHeaderRe
BlockInstanceMetadata: blockInstanceMetadata,
MetadataHeaderRequired: metadataHeaderRequired,
Reporter: reporter,
SetRetryAfterHeader: setRetryAfterHeader,
}
}

Expand Down Expand Up @@ -373,7 +376,14 @@ func (s *Server) msiHandler(w http.ResponseWriter, r *http.Request) (ns string)
podID, err := s.TokenClient.GetIdentities(r.Context(), podns, podname, tokenRequest.ClientID, tokenRequest.ResourceID)
if err != nil {
klog.Errorf("failed to get matching identities for pod: %s/%s, error: %+v", podns, podname, err)
http.Error(w, err.Error(), http.StatusNotFound)
httpErrorCode := http.StatusNotFound
if s.SetRetryAfterHeader {
httpErrorCode = http.StatusServiceUnavailable
// setting it to 20s to allow MIC to finish processing current cycle and pick up this
// pod in the next sync cycle
w.Header().Set(headerRetryAfter, "20")
}
http.Error(w, err.Error(), httpErrorCode)
return
}

Expand Down
54 changes: 31 additions & 23 deletions test/e2e/framework/config.go
Original file line number Diff line number Diff line change
Expand Up @@ -11,29 +11,33 @@ import (

// Config holds global test configuration translated from environment variables
type Config struct {
SubscriptionID string `envconfig:"SUBSCRIPTION_ID"`
ResourceGroup string `envconfig:"RESOURCE_GROUP"`
IdentityResourceGroup string `envconfig:"IDENTITY_RESOURCE_GROUP"`
NodeResourceGroup string `envconfig:"NODE_RESOURCE_GROUP"`
AzureClientID string `envconfig:"AZURE_CLIENT_ID"`
AzureClientSecret string `envconfig:"AZURE_CLIENT_SECRET"`
AzureTenantID string `envconfig:"AZURE_TENANT_ID"`
KeyvaultName string `envconfig:"KEYVAULT_NAME"`
KeyvaultSecretName string `envconfig:"KEYVAULT_SECRET_NAME"`
KeyvaultSecretVersion string `envconfig:"KEYVAULT_SECRET_VERSION"`
MICVersion string `envconfig:"MIC_VERSION" default:"v1.8.1"`
NMIVersion string `envconfig:"NMI_VERSION" default:"v1.8.1"`
Registry string `envconfig:"REGISTRY" default:"mcr.microsoft.com/oss/azure/aad-pod-identity"`
IdentityValidatorVersion string `envconfig:"IDENTITY_VALIDATOR_VERSION" default:"v1.8.1"`
EnableScaleFeatures bool `envconfig:"ENABLE_SCALE_FEATURES" default:"true"`
ImmutableUserMSIs string `envconfig:"IMMUTABLE_IDENTITY_CLIENT_ID"`
NMIMode string `envconfig:"NMI_MODE" default:"standard"`
BlockInstanceMetadata bool `envconfig:"BLOCK_INSTANCE_METADATA" default:"true"`
IsSoakTest bool `envconfig:"IS_SOAK_TEST" default:"false"`
IdentityReconcileInterval time.Duration `envconfig:"IDENTITY_RECONCILE_INTERVAL" default:"2m"`
ServicePrincipalClientID string `envconfig:"SERVICE_PRINCIPAL_CLIENT_ID"`
ServicePrincipalClientSecret string `envconfig:"SERVICE_PRINCIPAL_CLIENT_SECRET"`
MICSyncInterval time.Duration `envconfig:"MIC_SYNC_INTERVAL" default:"30s"`
SubscriptionID string `envconfig:"SUBSCRIPTION_ID"`
ResourceGroup string `envconfig:"RESOURCE_GROUP"`
IdentityResourceGroup string `envconfig:"IDENTITY_RESOURCE_GROUP"`
NodeResourceGroup string `envconfig:"NODE_RESOURCE_GROUP"`
AzureClientID string `envconfig:"AZURE_CLIENT_ID"`
AzureClientSecret string `envconfig:"AZURE_CLIENT_SECRET"`
AzureTenantID string `envconfig:"AZURE_TENANT_ID"`
KeyvaultName string `envconfig:"KEYVAULT_NAME"`
KeyvaultSecretName string `envconfig:"KEYVAULT_SECRET_NAME"`
KeyvaultSecretVersion string `envconfig:"KEYVAULT_SECRET_VERSION"`
MICVersion string `envconfig:"MIC_VERSION" default:"v1.8.1"`
NMIVersion string `envconfig:"NMI_VERSION" default:"v1.8.1"`
Registry string `envconfig:"REGISTRY" default:"mcr.microsoft.com/oss/azure/aad-pod-identity"`
IdentityValidatorVersion string `envconfig:"IDENTITY_VALIDATOR_VERSION" default:"v1.8.1"`
EnableScaleFeatures bool `envconfig:"ENABLE_SCALE_FEATURES" default:"true"`
ImmutableUserMSIs string `envconfig:"IMMUTABLE_IDENTITY_CLIENT_ID"`
NMIMode string `envconfig:"NMI_MODE" default:"standard"`
BlockInstanceMetadata bool `envconfig:"BLOCK_INSTANCE_METADATA" default:"true"`
IsSoakTest bool `envconfig:"IS_SOAK_TEST" default:"false"`
IdentityReconcileInterval time.Duration `envconfig:"IDENTITY_RECONCILE_INTERVAL" default:"2m"`
ServicePrincipalClientID string `envconfig:"SERVICE_PRINCIPAL_CLIENT_ID"`
ServicePrincipalClientSecret string `envconfig:"SERVICE_PRINCIPAL_CLIENT_SECRET"`
MICSyncInterval time.Duration `envconfig:"MIC_SYNC_INTERVAL" default:"30s"`
SetRetryAfterHeader bool `envconfig:"SET_RETRY_AFTER_HEADER" default:"false"`
RetryAttemptsForCreated int `envconfig:"RETRY_ATTEMPTS_FOR_CREATED" default:"16"`
RetryAttemptsForAssigned int `envconfig:"RETRY_ATTEMPTS_FOR_ASSIGNED" default:"4"`
FindIdentityRetryIntervalInSeconds int `envconfig:"FIND_IDENTITY_RETRY_INTERVAL_IN_SECONDS" default:"5"`
}

func (c *Config) DeepCopy() *Config {
Expand All @@ -59,6 +63,10 @@ func (c *Config) DeepCopy() *Config {
copy.ServicePrincipalClientID = c.ServicePrincipalClientID
copy.ServicePrincipalClientSecret = c.ServicePrincipalClientSecret
copy.MICSyncInterval = c.MICSyncInterval
copy.SetRetryAfterHeader = c.SetRetryAfterHeader
copy.RetryAttemptsForCreated = c.RetryAttemptsForCreated
copy.RetryAttemptsForAssigned = c.RetryAttemptsForAssigned
copy.FindIdentityRetryIntervalInSeconds = c.FindIdentityRetryIntervalInSeconds

return copy
}
Expand Down
7 changes: 7 additions & 0 deletions test/e2e/framework/helm/helm_helpers.go
Original file line number Diff line number Diff line change
Expand Up @@ -101,6 +101,9 @@ func generateValueArgs(config *framework.Config) []string {
fmt.Sprintf("--set=mic.tag=%s", config.MICVersion),
fmt.Sprintf("--set=nmi.tag=%s", config.NMIVersion),
fmt.Sprintf("--set=mic.syncRetryDuration=%s", config.MICSyncInterval),
fmt.Sprintf("--set=nmi.retryAttemptsForCreated=%d", config.RetryAttemptsForCreated),
fmt.Sprintf("--set=nmi.retryAttemptsForAssigned=%d", config.RetryAttemptsForAssigned),
fmt.Sprintf("--set=nmi.findIdentityRetryIntervalInSeconds=%d", config.FindIdentityRetryIntervalInSeconds),
}

// TODO (aramase) bump this to compare against v1.7.3 after next release
Expand All @@ -124,6 +127,10 @@ func generateValueArgs(config *framework.Config) []string {
args = append(args, fmt.Sprintf("--set=mic.identityAssignmentReconcileInterval=%s", config.IdentityReconcileInterval))
}

if config.SetRetryAfterHeader {
args = append(args, fmt.Sprintf("--set=nmi.setRetryAfterHeader=%t", config.SetRetryAfterHeader))
}

return args
}

Expand Down
94 changes: 94 additions & 0 deletions test/e2e/retry_after_header_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,94 @@
// +build e2e

package e2e

import (
aadpodv1 "github.com/Azure/aad-pod-identity/pkg/apis/aadpodidentity/v1"
"github.com/Azure/aad-pod-identity/test/e2e/framework/azureassignedidentity"
"github.com/Azure/aad-pod-identity/test/e2e/framework/azureidentity"
"github.com/Azure/aad-pod-identity/test/e2e/framework/azureidentitybinding"
"github.com/Azure/aad-pod-identity/test/e2e/framework/helm"
"github.com/Azure/aad-pod-identity/test/e2e/framework/identityvalidator"
"github.com/Azure/aad-pod-identity/test/e2e/framework/namespace"

. "github.com/onsi/ginkgo"
corev1 "k8s.io/api/core/v1"
)

var _ = Describe("When SetRetryAfter header is enabled", func() {
var (
specName = "retry-after"
ns *corev1.Namespace
)

BeforeEach(func() {
ns = namespace.Create(namespace.CreateInput{
Creator: kubeClient,
Name: specName,
})
// upgrade pod identity to use the feature flag set-retry-after-header and
// disable the internal retries.
c := *config
c.RetryAttemptsForCreated = 1
c.RetryAttemptsForAssigned = 1
c.FindIdentityRetryIntervalInSeconds = 1
c.SetRetryAfterHeader = true

helm.Upgrade(helm.UpgradeInput{Config: &c})
})

AfterEach(func() {
Cleanup(CleanupInput{
Namespace: ns,
Getter: kubeClient,
Lister: kubeClient,
Deleter: kubeClient,
})
// reset the feature flag
helm.Upgrade(helm.UpgradeInput{Config: config})
})

It("should pass the identity validation with retries from SDK based on Retry-After header", func() {
azureIdentity := azureidentity.Create(azureidentity.CreateInput{
Creator: kubeClient,
Config: config,
AzureClient: azureClient,
Name: keyvaultIdentity,
Namespace: ns.Name,
IdentityType: aadpodv1.UserAssignedMSI,
IdentityName: keyvaultIdentity,
})
azureIdentityBinding := azureidentitybinding.Create(azureidentitybinding.CreateInput{
Creator: kubeClient,
Name: keyvaultIdentityBinding,
Namespace: ns.Name,
AzureIdentityName: azureIdentity.Name,
Selector: keyvaultIdentitySelector,
})

identityValidator := identityvalidator.Create(identityvalidator.CreateInput{
Creator: kubeClient,
Config: config,
Namespace: ns.Name,
IdentityBinding: azureIdentityBinding.Spec.Selector,
})

azureassignedidentity.Wait(azureassignedidentity.WaitInput{
Getter: kubeClient,
PodName: identityValidator.Name,
Namespace: ns.Name,
AzureIdentityName: azureIdentity.Name,
StateToWaitFor: aadpodv1.AssignedIDAssigned,
})

identityvalidator.Validate(identityvalidator.ValidateInput{
Getter: kubeClient,
Config: config,
KubeconfigPath: kubeconfigPath,
PodName: identityValidator.Name,
Namespace: ns.Name,
IdentityClientID: azureIdentity.Spec.ClientID,
IdentityResourceID: azureIdentity.Spec.ResourceID,
})
})
})

0 comments on commit 9e31f9c

Please sign in to comment.