Skip to content

Commit

Permalink
Implementing workload webhook for auto-annotation (#73)
Browse files Browse the repository at this point in the history
  • Loading branch information
mitali-salvi authored Feb 2, 2024
1 parent b1767cc commit cf477d2
Show file tree
Hide file tree
Showing 7 changed files with 301 additions and 26 deletions.
1 change: 1 addition & 0 deletions config/rbac/role.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,7 @@ rules:
- namespaces
verbs:
- list
- update
- watch
- apiGroups:
- apps
Expand Down
22 changes: 22 additions & 0 deletions config/webhook/manifests.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -64,6 +64,28 @@ webhooks:
resources:
- pods
sideEffects: None
- admissionReviewVersions:
- v1
clientConfig:
service:
name: webhook-service
namespace: system
path: /mutate-v1-workload
failurePolicy: Ignore
name: mworkload.kb.io
rules:
- apiGroups:
- apps
apiVersions:
- v1
operations:
- CREATE
- UPDATE
resources:
- daemonsets
- deployments
- statefulsets
sideEffects: None
---
apiVersion: admissionregistration.k8s.io/v1
kind: ValidatingWebhookConfiguration
Expand Down
92 changes: 92 additions & 0 deletions internal/webhook/workloadmutation/webhookhandler.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,92 @@
// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

// Package workloadmutation contains the webhook that injects annotations into daemon-sets, deployments and stateful-sets.
package workloadmutation

import (
"context"
"encoding/json"
"errors"
"net/http"

"github.com/go-logr/logr"
appsv1 "k8s.io/api/apps/v1"
"k8s.io/apimachinery/pkg/runtime"
"sigs.k8s.io/controller-runtime/pkg/client"
"sigs.k8s.io/controller-runtime/pkg/webhook/admission"

"github.com/aws/amazon-cloudwatch-agent-operator/internal/config"
"github.com/aws/amazon-cloudwatch-agent-operator/pkg/instrumentation/auto"
)

// +kubebuilder:webhook:path=/mutate-v1-workload,mutating=true,failurePolicy=ignore,groups="apps",resources=daemonsets;deployments;statefulsets,verbs=create;update,versions=v1,name=mworkload.kb.io,sideEffects=none,admissionReviewVersions=v1
// +kubebuilder:rbac:groups="apps",resources=daemonsets;deployments;statefulsets,verbs=get;list;watch

var _ WebhookHandler = (*workloadMutationWebhook)(nil)

// WebhookHandler is a webhook handler that analyzes new daemon-sets and injects appropriate annotations into it.
type WebhookHandler interface {
admission.Handler
}

// the implementation.
type workloadMutationWebhook struct {
client client.Client
decoder *admission.Decoder
logger logr.Logger
config config.Config
annotationMutator *auto.AnnotationMutators
}

// NewWebhookHandler creates a new WorkloadWebhookHandler.
func NewWebhookHandler(cfg config.Config, logger logr.Logger, decoder *admission.Decoder, cl client.Client, annotationMutation *auto.AnnotationMutators) WebhookHandler {
return &workloadMutationWebhook{
config: cfg,
decoder: decoder,
logger: logger,
client: cl,
annotationMutator: annotationMutation,
}
}

func (p *workloadMutationWebhook) Handle(ctx context.Context, req admission.Request) admission.Response {
var err error
var marshaledObject []byte
var object runtime.Object
switch objectKind := req.Kind.Kind; objectKind {
case "DaemonSet":
ds := appsv1.DaemonSet{}
err = p.decoder.Decode(req, &ds)
if err != nil {
return admission.Errored(http.StatusBadRequest, err)
}
object = &ds
case "Deployment":
d := appsv1.Deployment{}
err = p.decoder.Decode(req, &d)
if err != nil {
return admission.Errored(http.StatusBadRequest, err)
}
object = &d
case "StatefulSet":
ss := appsv1.StatefulSet{}
err = p.decoder.Decode(req, &ss)
if err != nil {
return admission.Errored(http.StatusBadRequest, err)
}
object = &ss
default:
return admission.Errored(http.StatusBadRequest, errors.New("failed to unmarshal request object"))
}

p.annotationMutator.Mutate(object)
marshaledObject, err = json.Marshal(object)
if err != nil {
res := admission.Errored(http.StatusInternalServerError, err)
res.Allowed = true
return res
}

return admission.PatchResponseFromRaw(req.Object.Raw, marshaledObject)
}
149 changes: 149 additions & 0 deletions internal/webhook/workloadmutation/webhookhandler_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,149 @@
// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

package workloadmutation

import (
"context"
"encoding/json"
"github.com/stretchr/testify/require"
admv1 "k8s.io/api/admission/v1"
appsv1 "k8s.io/api/apps/v1"
corev1 "k8s.io/api/core/v1"
metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
"k8s.io/apimachinery/pkg/runtime"
"net/http"
"testing"

"github.com/go-logr/logr"
"github.com/stretchr/testify/assert"
"k8s.io/kubectl/pkg/scheme"
"sigs.k8s.io/controller-runtime/pkg/client"
logf "sigs.k8s.io/controller-runtime/pkg/log"
"sigs.k8s.io/controller-runtime/pkg/webhook/admission"

"github.com/aws/amazon-cloudwatch-agent-operator/internal/config"
"github.com/aws/amazon-cloudwatch-agent-operator/pkg/instrumentation"
"github.com/aws/amazon-cloudwatch-agent-operator/pkg/instrumentation/auto"
)

var (
k8sClient client.Client
logger = logf.Log.WithName("unit-tests")
)

func TestHandle(t *testing.T) {
for _, tt := range []struct {
req admission.Request
name string
expected int32
allowed bool
}{
{
name: "empty payload",
req: admission.Request{},
expected: http.StatusBadRequest,
allowed: false,
},
{
name: "invalid empty daemon-set payload",
req: func() admission.Request {
ds := appsv1.DaemonSet{}
encoded, err := json.Marshal(ds)
require.NoError(t, err)

return admission.Request{
AdmissionRequest: admv1.AdmissionRequest{
Namespace: "testing",
Object: runtime.RawExtension{
Raw: encoded,
},
},
}
}(),
expected: http.StatusBadRequest,
allowed: false,
},
{
name: "invalid pod payload",
req: func() admission.Request {
pod := corev1.Pod{}
encoded, err := json.Marshal(pod)
require.NoError(t, err)

return admission.Request{
AdmissionRequest: admv1.AdmissionRequest{
Namespace: "testing",
Object: runtime.RawExtension{
Raw: encoded,
},
},
}
}(),
expected: http.StatusBadRequest,
allowed: false,
},
{
name: "valid workload payload",
req: func() admission.Request {
ds := appsv1.DaemonSet{
TypeMeta: metav1.TypeMeta{
Kind: "DaemonSet",
},
Spec: appsv1.DaemonSetSpec{
Template: corev1.PodTemplateSpec{
ObjectMeta: metav1.ObjectMeta{
Annotations: map[string]string{},
},
},
},
}
encoded, err := json.Marshal(ds)
require.NoError(t, err)

return admission.Request{
AdmissionRequest: admv1.AdmissionRequest{
Kind: metav1.GroupVersionKind{
Kind: "DaemonSet",
},
Namespace: "testing",
Object: runtime.RawExtension{
Raw: encoded,
},
},
}
}(),
expected: http.StatusOK,
allowed: true,
},
} {
t.Run(tt.name, func(t *testing.T) {
// prepare
cfg := config.New()
decoder := admission.NewDecoder(scheme.Scheme)
autoAnnotationConfig := auto.AnnotationConfig{
Java: auto.AnnotationResources{
Namespaces: []string{"keep-auto-java"},
},
}
mutators := auto.NewAnnotationMutators(
k8sClient,
k8sClient,
logr.Logger{},
autoAnnotationConfig,
instrumentation.NewTypeSet(instrumentation.TypeJava),
)
injector := NewWebhookHandler(cfg, logger, decoder, k8sClient, mutators)

// test
res := injector.Handle(context.Background(), tt.req)

// verify
assert.Equal(t, tt.allowed, res.Allowed)
if !tt.allowed {
assert.NotNil(t, res.AdmissionResponse.Result)
assert.Equal(t, tt.expected, res.AdmissionResponse.Result.Code)
}
})
}
}
45 changes: 24 additions & 21 deletions main.go
Original file line number Diff line number Diff line change
Expand Up @@ -33,6 +33,7 @@ import (
"github.com/aws/amazon-cloudwatch-agent-operator/internal/config"
"github.com/aws/amazon-cloudwatch-agent-operator/internal/version"
"github.com/aws/amazon-cloudwatch-agent-operator/internal/webhook/podmutation"
"github.com/aws/amazon-cloudwatch-agent-operator/internal/webhook/workloadmutation"
"github.com/aws/amazon-cloudwatch-agent-operator/pkg/featuregate"
"github.com/aws/amazon-cloudwatch-agent-operator/pkg/instrumentation"
"github.com/aws/amazon-cloudwatch-agent-operator/pkg/instrumentation/auto"
Expand Down Expand Up @@ -185,6 +186,29 @@ func main() {
os.Exit(1)
}

decoder := admission.NewDecoder(mgr.GetScheme())

if os.Getenv("DISABLE_AUTO_ANNOTATION") != "true" {
var autoAnnotationConfig auto.AnnotationConfig
if err := json.Unmarshal([]byte(autoAnnotationConfigStr), &autoAnnotationConfig); err != nil {
setupLog.Error(err, "unable to unmarshal auto-annotation config")
} else {
autoAnnotationMutator := auto.NewAnnotationMutators(
mgr.GetClient(),
mgr.GetAPIReader(),
logger,
autoAnnotationConfig,
instrumentation.NewTypeSet(
instrumentation.TypeJava,
instrumentation.TypePython,
),
)
mgr.GetWebhookServer().Register("/mutate-v1-workload", &webhook.Admission{
Handler: workloadmutation.NewWebhookHandler(cfg, ctrl.Log.WithName("workload-webhook"), decoder, mgr.GetClient(), autoAnnotationMutator)})
go autoAnnotationMutator.MutateAll(ctx)
}
}

if os.Getenv("ENABLE_WEBHOOKS") != "false" {
if err = otelv1alpha1.SetupCollectorWebhook(mgr, cfg); err != nil {
setupLog.Error(err, "unable to create webhook", "webhook", "AmazonCloudWatchAgent")
Expand All @@ -194,15 +218,13 @@ func main() {
setupLog.Error(err, "unable to create webhook", "webhook", "Instrumentation")
os.Exit(1)
}
decoder := admission.NewDecoder(mgr.GetScheme())
mgr.GetWebhookServer().Register("/mutate-v1-pod", &webhook.Admission{
Handler: podmutation.NewWebhookHandler(cfg, ctrl.Log.WithName("pod-webhook"), decoder, mgr.GetClient(),
[]podmutation.PodMutator{
sidecar.NewMutator(logger, cfg, mgr.GetClient()),
instrumentation.NewMutator(logger, mgr.GetClient(), mgr.GetEventRecorderFor("amazon-cloudwatch-agent-operator")),
}),
})

} else {
ctrl.Log.Info("Webhooks are disabled, operator is running an unsupported mode", "ENABLE_WEBHOOKS", "false")
}
Expand All @@ -217,25 +239,6 @@ func main() {
os.Exit(1)
}

if os.Getenv("DISABLE_AUTO_ANNOTATION") != "true" {
var autoAnnotationConfig auto.AnnotationConfig
if err := json.Unmarshal([]byte(autoAnnotationConfigStr), &autoAnnotationConfig); err != nil {
setupLog.Error(err, "unable to unmarshal auto-annotation config")
} else {
setupLog.Info("starting auto-annotator")
autoAnnotation := auto.NewAnnotationMutators(
mgr.GetClient(),
logger,
autoAnnotationConfig,
instrumentation.NewTypeSet(
instrumentation.TypeJava,
instrumentation.TypePython,
),
)
go autoAnnotation.MutateAll(ctx)
}
}

setupLog.Info("starting manager")
if err := mgr.Start(ctx); err != nil {
setupLog.Error(err, "problem running manager")
Expand Down
Loading

0 comments on commit cf477d2

Please sign in to comment.