Skip to content

Commit 7be5874

Browse files
nicolstnicolst
authored
refactor!: refactor and simplify statefulset handling (#83)
1 parent 3050379 commit 7be5874

File tree

3 files changed

+103
-135
lines changed

3 files changed

+103
-135
lines changed

‎internal/controller/common.go‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,6 @@ import (
88
)
99

1010
const (
11-
enabledssbucketeerAnnotation = "dapla.ssb.no/enable-ssbucketeer"
1211
impersonateGroupAnnotation = "dapla.ssb.no/impersonate-group"
1312
mountBucketsAnnotation = "dapla.ssb.no/mount-buckets"
1413
mountStandardBucketsAnnotation = "dapla.ssb.no/mount-standard-buckets"
@@ -82,6 +81,10 @@ type AccessGroupConfig struct {
8281
ReasonRequired bool `yaml:"reasonRequired"`
8382
}
8483

84+
func (c AccessGroupConfig) ToTeam(group string) string {
85+
return strings.TrimSuffix(strings.TrimSuffix(group, c.Name), "-")
86+
}
87+
8588
type ProjectTemplateData struct {
8689
TeamName string
8790
Stage string

‎internal/controller/serviceaccount_webhook.go‎

Lines changed: 11 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -82,12 +82,12 @@ func (m *ServiceAccountValidator) validate(ctx context.Context, req runtime.Obje
8282
return nil, nil
8383
}
8484

85-
groupType := m.GroupConfigs.GetConfig(group)
86-
if groupType == nil {
85+
groupConfig := m.GroupConfigs.GetConfig(group)
86+
if groupConfig == nil {
8787
return nil, fmt.Errorf("no group config matches %q", group)
8888
}
8989

90-
team := strings.TrimSuffix(strings.TrimSuffix(group, groupType.Name), "-")
90+
team := groupConfig.ToTeam(group)
9191

9292
user := strings.TrimPrefix(sa.Namespace, userNamespacePrefix)
9393
if user == sa.Namespace {
@@ -111,25 +111,26 @@ func (m *ServiceAccountValidator) validate(ctx context.Context, req runtime.Obje
111111
durationStr, hasDuration := sa.Annotations[requestedServiceDurationAnnotation]
112112
reason, hasReason := sa.Annotations[accessReasonAnnotation]
113113

114-
if !groupType.ReasonRequired {
114+
if !groupConfig.ReasonRequired {
115115
return nil, nil
116116
}
117+
117118
if !hasDuration {
118-
return nil, fmt.Errorf("duration required for group %q of type %q", group, groupType.Name)
119+
return nil, fmt.Errorf("duration required for group %q of type %q", group, groupConfig.Name)
119120
}
120121

121122
parsedDuration, err := time.ParseDuration(durationStr)
122123
if err != nil {
123124
log.Error(err, "failed to parse ServiceAccount duration", "duration", durationStr)
124-
return nil, fmt.Errorf("invalid duration: %s", durationStr)
125+
return nil, fmt.Errorf("invalid duration: %q", durationStr)
125126
}
126127

127-
if parsedDuration > groupType.MaxDuration {
128-
return nil, fmt.Errorf("requested duration %q exceeds max allowed %q for group %q", parsedDuration, groupType.MaxDuration, group)
128+
if parsedDuration > groupConfig.MaxDuration {
129+
return nil, fmt.Errorf("requested duration %q exceeds max allowed %q for group %q", parsedDuration, groupConfig.MaxDuration, group)
129130
}
130131

131132
if !hasReason || reason == "" {
132-
return nil, fmt.Errorf("reason required for group %q of type %q", group, groupType.Name)
133+
return nil, fmt.Errorf("reason required for group %q of type %q", group, groupConfig.Name)
133134
}
134135

135136
chartName, chartVersion := getChartNameAndVersion(*sa)
@@ -139,7 +140,7 @@ func (m *ServiceAccountValidator) validate(ctx context.Context, req runtime.Obje
139140
UserPrincipalName: user,
140141
TeamName: team,
141142
AccessGroup: group,
142-
GroupType: groupType.Name,
143+
GroupType: groupConfig.Name,
143144
Reason: reason,
144145
StartTime: sa.CreationTimestamp.Time,
145146
EndTime: sa.CreationTimestamp.Add(parsedDuration),

‎internal/controller/statefulset_webhook.go‎

Lines changed: 88 additions & 124 deletions
Original file line numberDiff line numberDiff line change
@@ -3,13 +3,11 @@ package controller
33
import (
44
"context"
55
"encoding/json"
6-
"errors"
76
"fmt"
87
"maps"
98
"net/http"
109
"slices"
1110
"strings"
12-
"time"
1311

1412
rm "cloud.google.com/go/resourcemanager/apiv3"
1513
rmpb "cloud.google.com/go/resourcemanager/apiv3/resourcemanagerpb"
@@ -70,164 +68,130 @@ func (m *StatefulsetMutator) Handle(ctx context.Context, req admission.Request)
7068
return admission.Allowed("statefulset is being deleted")
7169
}
7270

73-
if sfs.Annotations[enabledssbucketeerAnnotation] != "true" {
71+
group, ok := sfs.Annotations[impersonateGroupAnnotation]
72+
if !ok {
7473
return admission.Allowed("skipping ssbucketeer mutation")
7574
}
7675

7776
ensurePodAnnotations(&sfs.Spec.Template)
7877

79-
// 1. Check if we want to impersonate a group SA
80-
group, hasGroupAnnotation := sfs.Annotations[impersonateGroupAnnotation]
81-
if hasGroupAnnotation {
82-
saAnnotations := map[string]string{
83-
impersonateGroupAnnotation: group,
84-
}
78+
// Find the service container, so we can add volumeMounts
79+
serviceContainer := getServiceContainer(&sfs.Spec.Template, sfs.Annotations[serviceContainerAnnotation])
80+
if serviceContainer == nil {
81+
err := fmt.Errorf("could not find container with name %q", sfs.Annotations[serviceContainerAnnotation])
82+
log.Error(err, "could not find service container")
83+
return admission.Errored(http.StatusBadRequest, err)
84+
}
8585

86-
groupConfig := m.GroupConfigs.GetConfig(group)
87-
if groupConfig == nil {
88-
return admission.Denied(fmt.Sprintf("no configuration found for group: %s", group))
89-
}
86+
bucketMounts := getExtraBucketMounts(sfs.Annotations)
9087

91-
if groupConfig.ReasonRequired {
92-
reason, hasReason := sfs.Annotations[accessReasonAnnotation]
93-
if !hasReason || reason == "" {
94-
return admission.Denied(fmt.Sprintf("reason is required for access group: %q", group))
95-
}
96-
saAnnotations[accessReasonAnnotation] = reason
97-
log.Info("Reason provided for access group", "group", group, "reason", reason)
98-
}
88+
saAnnotations := map[string]string{
89+
impersonateGroupAnnotation: group,
90+
requestedServiceDurationAnnotation: sfs.Annotations[requestedServiceDurationAnnotation],
91+
accessReasonAnnotation: sfs.Annotations[accessReasonAnnotation],
92+
}
9993

100-
if groupConfig.MaxDuration > 0 {
101-
durationString, ok := sfs.Annotations[requestedServiceDurationAnnotation]
102-
if !ok {
103-
return admission.Denied(fmt.Sprintf("access duration required for group %q", group))
104-
}
105-
106-
duration, err := time.ParseDuration(durationString)
107-
if err != nil {
108-
log.Error(err, "failed to parse requested-duration", "duration", durationString)
109-
return admission.Denied(fmt.Sprintf("invalid requested duration %q", durationString))
110-
}
111-
112-
if duration > groupConfig.MaxDuration {
113-
log.Info("attempt to start service with longer than max duration, defaulting to max",
114-
"duration", duration,
115-
"group", group,
116-
"maxDuration", groupConfig.MaxDuration)
117-
duration = groupConfig.MaxDuration
118-
sfs.Annotations[requestedServiceDurationAnnotation] = duration.String()
119-
}
120-
121-
saAnnotations[requestedServiceDurationAnnotation] = duration.String()
94+
groupConfig := m.GroupConfigs.GetConfig(group)
95+
if groupConfig == nil {
96+
return admission.Denied(fmt.Sprintf("no configuration found for group: %s", group))
97+
}
12298

123-
}
99+
// Handle IAM bindings and k8s SA annotations
100+
if err := m.handleServiceAccount(ctx, req.Namespace, sfs.Spec.Template.Spec.ServiceAccountName, saAnnotations); err != nil {
101+
log.Error(err, "handle serviceaccount")
102+
return admission.Denied("error handling service account")
103+
}
124104

125-
// Handle IAM bindings and k8s SA annotations
126-
if err := m.handleServiceAccount(ctx, req.Namespace, sfs.Spec.Template.Spec.ServiceAccountName, saAnnotations); err != nil {
127-
log.Error(err, "handle serviceaccount")
128-
return admission.Denied("error handling service account")
105+
if m.ADCGroupEnvName != "" {
106+
if !slices.ContainsFunc(serviceContainer.Env, func(e corev1.EnvVar) bool {
107+
return e.Name == m.ADCGroupEnvName
108+
}) {
109+
serviceContainer.Env = append(serviceContainer.Env, corev1.EnvVar{Name: m.ADCGroupEnvName, Value: group})
129110
}
130111
}
131112

132-
// 2. Mount buckets
133-
var bucketNames []string
134-
if annotationBuckets, ok := sfs.Annotations[mountBucketsAnnotation]; ok {
135-
bucketNames = strings.Split(annotationBuckets, ",")
113+
// TODO: Use Dapla Team API for this?
114+
if sfs.Annotations[mountStandardBucketsAnnotation] == "true" {
115+
team := groupConfig.ToTeam(group)
116+
if err := m.addStandardBuckets(ctx, team, *groupConfig, bucketMounts); err != nil {
117+
log.Error(err, "failed to add standard buckets")
118+
}
136119
}
137120

138-
// Find the service container, so we can add volumeMounts
139-
containerIndex := slices.IndexFunc(sfs.Spec.Template.Spec.Containers, func(c corev1.Container) bool {
140-
return c.Name == sfs.Annotations[serviceContainerAnnotation]
141-
})
121+
addBucketsToPodSpec(&sfs.Spec.Template.Spec, serviceContainer, bucketMounts, m.PrecreatorImage)
142122

143-
if containerIndex == -1 {
144-
err := fmt.Errorf("could not find container with name %q", sfs.Annotations[serviceContainerAnnotation])
145-
log.Error(err, "could not find service container")
146-
return admission.Errored(http.StatusBadRequest, err)
123+
if m.IamProbeImage != "" && sfs.Annotations[iamProbeStatus] != iamProbeDone {
124+
if err := m.launchIamProbe(ctx, sfs); err != nil {
125+
log.Error(err, "could not start iam probe")
126+
}
147127
}
148128

149-
if hasGroupAnnotation && m.ADCGroupEnvName != "" {
150-
containerEnv := &sfs.Spec.Template.Spec.Containers[containerIndex].Env
151-
if !slices.ContainsFunc(*containerEnv, func(e corev1.EnvVar) bool {
152-
return e.Name == m.ADCGroupEnvName
153-
}) {
154-
*containerEnv = append(*containerEnv, corev1.EnvVar{Name: m.ADCGroupEnvName, Value: group})
155-
}
129+
marshaledStatefulSet, err := json.Marshal(sfs)
130+
if err != nil {
131+
return admission.Errored(http.StatusInternalServerError, err)
132+
}
133+
return admission.PatchResponseFromRaw(req.Object.Raw, marshaledStatefulSet)
134+
}
135+
136+
func getExtraBucketMounts(annotations map[string]string) map[string]string {
137+
var bucketNames []string
138+
if annotationBuckets, ok := annotations[mountBucketsAnnotation]; ok {
139+
bucketNames = strings.Split(annotationBuckets, ",")
156140
}
157141

158142
bucketMounts := make(map[string]string, len(bucketNames))
159143
for _, bucket := range bucketNames {
160144
bucketMounts[bucket] = bucket
161145
}
146+
return bucketMounts
147+
}
162148

163-
// TODO: Use Dapla Team API for this?
164-
if sfs.Annotations[mountStandardBucketsAnnotation] == "true" {
165-
team := group
166-
var gc *AccessGroupConfig
167-
for _, config := range m.GroupConfigs {
168-
if team = strings.TrimSuffix(group, config.Name); team != group {
169-
team = strings.TrimSuffix(team, "-")
170-
gc = &config
171-
break
172-
}
173-
}
174-
175-
if team != group {
176-
if err := m.addStandardBuckets(ctx, team, *gc, bucketMounts); err != nil {
177-
log.Error(err, "failed to add standard buckets")
178-
}
179-
} else {
180-
log.Error(errors.New("could not deduce team from group"), "team could not be deduced from group", "group", group)
181-
}
149+
func getServiceContainer(pod *corev1.PodTemplateSpec, name string) *corev1.Container {
150+
idx := slices.IndexFunc(pod.Spec.Containers, func(c corev1.Container) bool {
151+
return c.Name == name
152+
})
153+
if idx == -1 {
154+
return nil
182155
}
156+
return &pod.Spec.Containers[idx]
157+
}
183158

184-
addBucketsToPodSpec(&sfs.Spec.Template.Spec, &sfs.Spec.Template.Spec.Containers[containerIndex], bucketMounts, m.PrecreatorImage)
159+
func (m *StatefulsetMutator) launchIamProbe(ctx context.Context, sfs *appsv1.StatefulSet) error {
160+
sfs.Annotations[iamProbeStatus] = fmt.Sprintf("%s%d", iamProbeRunningPrefix, *sfs.Spec.Replicas)
161+
sfs.Spec.Replicas = ptr[int32](0)
185162

186-
if m.IamProbeImage != "" && sfs.Annotations[iamProbeStatus] != iamProbeDone {
187-
sfs.Annotations[iamProbeStatus] = fmt.Sprintf("%s%d", iamProbeRunningPrefix, *sfs.Spec.Replicas)
188-
sfs.Spec.Replicas = ptr[int32](0)
189-
190-
probeJob := &batchv1.Job{
191-
ObjectMeta: metav1.ObjectMeta{
192-
Name: fmt.Sprintf("%s-iam-probe", sfs.Name),
193-
Namespace: sfs.Namespace,
194-
Annotations: map[string]string{
195-
probeJobStatefulsetAnnotation: sfs.Name,
196-
},
163+
probeJob := &batchv1.Job{
164+
ObjectMeta: metav1.ObjectMeta{
165+
Name: fmt.Sprintf("%s-iam-probe", sfs.Name),
166+
Namespace: sfs.Namespace,
167+
Annotations: map[string]string{
168+
probeJobStatefulsetAnnotation: sfs.Name,
197169
},
198-
Spec: batchv1.JobSpec{
199-
ActiveDeadlineSeconds: ptr[int64](300),
200-
TTLSecondsAfterFinished: ptr[int32](0),
201-
Template: corev1.PodTemplateSpec{
202-
ObjectMeta: metav1.ObjectMeta{
203-
Annotations: map[string]string{
204-
istioExcludedIpRangesAnnotation: gcsfuseOutboundIPRange,
205-
},
170+
},
171+
Spec: batchv1.JobSpec{
172+
ActiveDeadlineSeconds: ptr[int64](300),
173+
TTLSecondsAfterFinished: ptr[int32](0),
174+
Template: corev1.PodTemplateSpec{
175+
ObjectMeta: metav1.ObjectMeta{
176+
Annotations: map[string]string{
177+
istioExcludedIpRangesAnnotation: gcsfuseOutboundIPRange,
206178
},
207-
Spec: corev1.PodSpec{
208-
ServiceAccountName: sfs.Spec.Template.Spec.ServiceAccountName,
209-
Containers: []corev1.Container{
210-
{
211-
Image: m.IamProbeImage,
212-
Name: "iam-probe",
213-
},
179+
},
180+
Spec: corev1.PodSpec{
181+
ServiceAccountName: sfs.Spec.Template.Spec.ServiceAccountName,
182+
Containers: []corev1.Container{
183+
{
184+
Image: m.IamProbeImage,
185+
Name: "iam-probe",
214186
},
215-
RestartPolicy: corev1.RestartPolicyNever,
216187
},
188+
RestartPolicy: corev1.RestartPolicyNever,
217189
},
218190
},
219-
}
220-
221-
if err := m.Client.Create(ctx, probeJob); err != nil {
222-
log.Error(err, "could not create probe job")
223-
}
191+
},
224192
}
225193

226-
marshaledStatefulSet, err := json.Marshal(sfs)
227-
if err != nil {
228-
return admission.Errored(http.StatusInternalServerError, err)
229-
}
230-
return admission.PatchResponseFromRaw(req.Object.Raw, marshaledStatefulSet)
194+
return m.Client.Create(ctx, probeJob)
231195
}
232196

233197
func (a *StatefulsetMutator) handleServiceAccount(ctx context.Context, namespace, name string, saAnnotations map[string]string) error {

0 commit comments

Comments
 (0)