Skip to content

Commit a902abc

Browse files
skattojuskattoju
committed
Adding a token getter to get service account tokens
1 parent 872b7f7 commit a902abc

File tree

2 files changed

+217
-0
lines changed

2 files changed

+217
-0
lines changed

‎internal/authentication/tokengetter.go

+105
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,105 @@
1+
package authentication
2+
3+
import (
4+
"context"
5+
"sync"
6+
"time"
7+
8+
authenticationv1 "k8s.io/api/authentication/v1"
9+
metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
10+
"k8s.io/apimachinery/pkg/types"
11+
corev1 "k8s.io/client-go/kubernetes/typed/core/v1"
12+
"k8s.io/utils/ptr"
13+
)
14+
15+
type TokenGetter struct {
16+
client corev1.ServiceAccountsGetter
17+
expirationSeconds int64
18+
rotationThreshold time.Duration
19+
tokens map[types.NamespacedName]*authenticationv1.TokenRequestStatus
20+
mu sync.RWMutex
21+
}
22+
23+
type TokenGetterOption func(*TokenGetter)
24+
25+
// Returns a token getter that can fetch tokens given a service account.
26+
// The token getter also caches tokens which helps reduce the number of requests to the API Server.
27+
// In case a cached token is expiring a fresh token is created.
28+
func NewTokenGetter(client corev1.ServiceAccountsGetter, options ...TokenGetterOption) *TokenGetter {
29+
tokenGetter := &TokenGetter{
30+
client: client,
31+
expirationSeconds: int64(5 * time.Minute), // default token ttl
32+
rotationThreshold: 1 * time.Minute, // default rotation threshold
33+
tokens: map[types.NamespacedName]*authenticationv1.TokenRequestStatus{},
34+
}
35+
36+
for _, opt := range options {
37+
opt(tokenGetter)
38+
}
39+
40+
return tokenGetter
41+
}
42+
43+
func WithExpirationSeconds(expirationSeconds int64) TokenGetterOption {
44+
return func(tg *TokenGetter) {
45+
tg.expirationSeconds = expirationSeconds
46+
}
47+
}
48+
49+
func WithRotationThreshold(rotationThreshold time.Duration) TokenGetterOption {
50+
return func(tg *TokenGetter) {
51+
tg.rotationThreshold = rotationThreshold
52+
}
53+
}
54+
55+
// Get returns a token from the cache if available and not expiring, otherwise creates a new token
56+
func (t *TokenGetter) Get(ctx context.Context, key types.NamespacedName) (string, error) {
57+
t.mu.RLock()
58+
token, ok := t.tokens[key]
59+
t.mu.RUnlock()
60+
61+
expireTime := time.Time{}
62+
if ok {
63+
expireTime = token.ExpirationTimestamp.Time
64+
}
65+
66+
// Create a new token if the cached token expires within rotationThreshold seconds from now
67+
rotationThresholdAfterNow := metav1.Now().Add(t.rotationThreshold)
68+
if expireTime.Before(rotationThresholdAfterNow) {
69+
var err error
70+
token, err = t.getToken(ctx, key)
71+
if err != nil {
72+
return "", err
73+
}
74+
t.mu.Lock()
75+
t.tokens[key] = token
76+
t.mu.Unlock()
77+
}
78+
79+
return token.Token, nil
80+
}
81+
82+
func (t *TokenGetter) getToken(ctx context.Context, key types.NamespacedName) (*authenticationv1.TokenRequestStatus, error) {
83+
req, err := t.client.ServiceAccounts(key.Namespace).CreateToken(ctx,
84+
key.Name,
85+
&authenticationv1.TokenRequest{
86+
Spec: authenticationv1.TokenRequestSpec{ExpirationSeconds: ptr.To[int64](t.expirationSeconds)},
87+
}, metav1.CreateOptions{})
88+
if err != nil {
89+
return nil, err
90+
}
91+
return &req.Status, nil
92+
}
93+
94+
func (t *TokenGetter) Clean(ctx context.Context, key types.NamespacedName) {
95+
t.mu.RLock()
96+
defer t.mu.RUnlock()
97+
delete(t.tokens, key)
98+
}
99+
100+
func (t *TokenGetter) TokenExists(ctx context.Context, key types.NamespacedName) bool {
101+
t.mu.RLock()
102+
defer t.mu.RUnlock()
103+
_, ok := t.tokens[key]
104+
return ok
105+
}

‎internal/authentication/tokengetter_test.go

+112
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,112 @@
1+
package authentication
2+
3+
import (
4+
"context"
5+
"fmt"
6+
"testing"
7+
"time"
8+
9+
"github.com/stretchr/testify/assert"
10+
authenticationv1 "k8s.io/api/authentication/v1"
11+
metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
12+
"k8s.io/apimachinery/pkg/runtime"
13+
"k8s.io/apimachinery/pkg/types"
14+
"k8s.io/client-go/kubernetes/fake"
15+
ctest "k8s.io/client-go/testing"
16+
)
17+
18+
func TestNewTokenGetterGet(t *testing.T) {
19+
fakeClient := fake.NewSimpleClientset()
20+
fakeClient.PrependReactor("create", "serviceaccounts/token",
21+
func(action ctest.Action) (bool, runtime.Object, error) {
22+
act, ok := action.(ctest.CreateActionImpl)
23+
if !ok {
24+
return false, nil, nil
25+
}
26+
tokenRequest := act.GetObject().(*authenticationv1.TokenRequest)
27+
var err error
28+
if act.Name == "test-service-account-1" {
29+
tokenRequest.Status = authenticationv1.TokenRequestStatus{
30+
Token: "test-token-1",
31+
ExpirationTimestamp: metav1.NewTime(metav1.Now().Add(5 * time.Minute)),
32+
}
33+
}
34+
if act.Name == "test-service-account-2" {
35+
tokenRequest.Status = authenticationv1.TokenRequestStatus{
36+
Token: "test-token-2",
37+
ExpirationTimestamp: metav1.NewTime(metav1.Now().Add(1 * time.Second)),
38+
}
39+
}
40+
if act.Name == "test-service-account-3" {
41+
tokenRequest = nil
42+
err = fmt.Errorf("error when fetching token")
43+
}
44+
return true, tokenRequest, err
45+
})
46+
47+
tg := NewTokenGetter(fakeClient.CoreV1(),
48+
WithExpirationSeconds(int64(5*time.Minute)),
49+
WithRotationThreshold(1*time.Minute))
50+
51+
tests := []struct {
52+
testName string
53+
serviceAccountName string
54+
namespace string
55+
want string
56+
errorMsg string
57+
}{
58+
{"Testing getting token with fake client", "test-service-account-1",
59+
"test-namespace-1", "test-token-1", "failed to get token"},
60+
{"Testing getting token from cache", "test-service-account-1",
61+
"test-namespace-1", "test-token-1", "failed to get token from cache"},
62+
{"Testing getting short lived token from fake client", "test-service-account-2",
63+
"test-namespace-2", "test-token-2", "failed to get token"},
64+
{"Testing getting expired token from cache", "test-service-account-2",
65+
"test-namespace-2", "test-token-2", "failed to refresh token"},
66+
{"Testing error when getting token from fake client", "test-service-account-3",
67+
"test-namespace-3", "error when fetching token", "error when fetching token"},
68+
}
69+
70+
for _, tc := range tests {
71+
got, err := tg.Get(context.Background(), types.NamespacedName{Namespace: tc.namespace, Name: tc.serviceAccountName})
72+
if err != nil {
73+
t.Logf("%s: expected: %v, got: %v", tc.testName, tc.want, err)
74+
assert.EqualError(t, err, tc.errorMsg)
75+
} else {
76+
t.Logf("%s: expected: %v, got: %v", tc.testName, tc.want, got)
77+
assert.Equal(t, tc.want, got, tc.errorMsg)
78+
}
79+
}
80+
}
81+
82+
func TestNewTokenGetterClean(t *testing.T) {
83+
fakeClient := fake.NewSimpleClientset()
84+
fakeClient.PrependReactor("create", "serviceaccounts/token",
85+
func(action ctest.Action) (bool, runtime.Object, error) {
86+
act, ok := action.(ctest.CreateActionImpl)
87+
if !ok {
88+
return false, nil, nil
89+
}
90+
tokenRequest := act.GetObject().(*authenticationv1.TokenRequest)
91+
if act.Name == "test-service-account-1" {
92+
tokenRequest.Status = authenticationv1.TokenRequestStatus{
93+
Token: "test-token-1",
94+
ExpirationTimestamp: metav1.NewTime(metav1.Now().Add(5 * time.Minute)),
95+
}
96+
}
97+
return true, tokenRequest, nil
98+
})
99+
100+
tg := NewTokenGetter(fakeClient.CoreV1(),
101+
WithExpirationSeconds(int64(5*time.Minute)),
102+
WithRotationThreshold(1*time.Minute))
103+
104+
_, err := tg.Get(context.Background(), types.NamespacedName{Namespace: "test-namespace-1", Name: "test-service-account-1"})
105+
if err != nil {
106+
t.Fatalf("failed to get token: %v", err)
107+
}
108+
assert.True(t, tg.TokenExists(context.Background(), types.NamespacedName{Namespace: "test-namespace-1", Name: "test-service-account-1"}))
109+
t.Logf("Testing removing token from cache")
110+
tg.Clean(context.Background(), types.NamespacedName{Namespace: "test-namespace-1", Name: "test-service-account-1"})
111+
assert.False(t, tg.TokenExists(context.Background(), types.NamespacedName{Namespace: "test-namespace-1", Name: "test-service-account-1"}))
112+
}

0 commit comments

Comments
 (0)