4
4
"context"
5
5
"fmt"
6
6
"os"
7
+ "time"
7
8
8
9
"github.com/pkg/errors"
9
10
"k8s.io/klog"
26
27
keyvaultSecretVersion = pflag .String ("keyvault-secret-version" , "" , "the version of the keyvault secret we are extracting with pod identity" )
27
28
)
28
29
30
+ const (
31
+ contextTimeout = 150 * time .Second
32
+ )
33
+
29
34
func main () {
30
35
pflag .Parse ()
31
36
@@ -41,14 +46,17 @@ func main() {
41
46
}
42
47
klog .Infof ("Successfully obtain MSIEndpoint: %s\n " , msiEndpoint )
43
48
49
+ ctx , ctxCancel := context .WithTimeout (context .Background (), contextTimeout )
50
+ defer ctxCancel ()
51
+
44
52
if * keyvaultName != "" && * keyvaultSecretName != "" {
45
53
// Test if the pod identity is set up correctly
46
- if err := testUserAssignedIdentityOnPod (msiEndpoint , * identityClientID , * keyvaultName , * keyvaultSecretName , * keyvaultSecretVersion ); err != nil {
54
+ if err := testUserAssignedIdentityOnPod (ctx , msiEndpoint , * identityClientID , * keyvaultName , * keyvaultSecretName , * keyvaultSecretVersion ); err != nil {
47
55
klog .Fatalf ("testUserAssignedIdentityOnPod failed, %+v" , err )
48
56
}
49
57
} else {
50
58
// Test if the cluster-wide user assigned identity is set up correctly
51
- if err := testClusterWideUserAssignedIdentity (msiEndpoint , * subscriptionID , * resourceGroup , * identityClientID ); err != nil {
59
+ if err := testClusterWideUserAssignedIdentity (ctx , msiEndpoint , * subscriptionID , * resourceGroup , * identityClientID ); err != nil {
52
60
klog .Fatalf ("testClusterWideUserAssignedIdentity failed, %+v" , err )
53
61
}
54
62
}
@@ -60,7 +68,7 @@ func main() {
60
68
}
61
69
62
70
// testClusterWideUserAssignedIdentity will verify whether cluster-wide user assigned identity is working properly
63
- func testClusterWideUserAssignedIdentity (msiEndpoint , subscriptionID , resourceGroup , identityClientID string ) error {
71
+ func testClusterWideUserAssignedIdentity (ctx context. Context , msiEndpoint , subscriptionID , resourceGroup , identityClientID string ) error {
64
72
os .Setenv ("AZURE_CLIENT_ID" , identityClientID )
65
73
defer os .Unsetenv ("AZURE_CLIENT_ID" )
66
74
token , err := adal .NewServicePrincipalTokenFromMSIWithUserAssignedID (msiEndpoint , azure .PublicCloud .ResourceManagerEndpoint , identityClientID )
@@ -70,7 +78,7 @@ func testClusterWideUserAssignedIdentity(msiEndpoint, subscriptionID, resourceGr
70
78
71
79
vmClient := compute .NewVirtualMachinesClient (subscriptionID )
72
80
vmClient .Authorizer = autorest .NewBearerAuthorizer (token )
73
- vmlist , err := vmClient .List (context . Background () , resourceGroup )
81
+ vmlist , err := vmClient .List (ctx , resourceGroup )
74
82
if err != nil {
75
83
return errors .Wrapf (err , "Failed to verify cluster-wide user assigned identity" )
76
84
}
@@ -80,7 +88,7 @@ func testClusterWideUserAssignedIdentity(msiEndpoint, subscriptionID, resourceGr
80
88
}
81
89
82
90
// testUserAssignedIdentityOnPod will verify whether a pod identity is working properly
83
- func testUserAssignedIdentityOnPod (msiEndpoint , identityClientID , keyvaultName , keyvaultSecretName , keyvaultSecretVersion string ) error {
91
+ func testUserAssignedIdentityOnPod (ctx context. Context , msiEndpoint , identityClientID , keyvaultName , keyvaultSecretName , keyvaultSecretVersion string ) error {
84
92
// When new authorizer is created, azure-sdk-for-go tries to create dataplane authorizer using MSI. It checks the AZURE_CLIENT_ID to get the client id
85
93
// for the user assigned identity. If client id not found, then NewServicePrincipalTokenFromMSI is invoked instead of using the actual
86
94
// user assigned identity. Setting this env var ensures we validate GetSecret using the desired user assigned identity.
@@ -94,7 +102,7 @@ func testUserAssignedIdentityOnPod(msiEndpoint, identityClientID, keyvaultName,
94
102
}
95
103
96
104
klog .Infof ("%s %s %s\n " , keyvaultName , keyvaultSecretName , keyvaultSecretVersion )
97
- secret , err := keyClient .GetSecret (context . Background () , fmt .Sprintf ("https://%s.vault.azure.net" , keyvaultName ), keyvaultSecretName , keyvaultSecretVersion )
105
+ secret , err := keyClient .GetSecret (ctx , fmt .Sprintf ("https://%s.vault.azure.net" , keyvaultName ), keyvaultSecretName , keyvaultSecretVersion )
98
106
if err != nil || * secret .Value == "" {
99
107
return errors .Wrapf (err , "Failed to verify user assigned identity on pod" )
100
108
}
0 commit comments