Merge pull request #109078 from tallclair/audit-mutex

Audit mutex

plugin/pkg/admission/security/podsecurity/admission.go

@@ -190,8 +190,14 @@ func (p *Plugin) Validate(ctx context.Context, a admission.Attributes, o admissi
 	for _, w := range result.Warnings {
 		warning.AddWarning(ctx, "", w)
 	}
-	for k, v := range result.AuditAnnotations {
-		audit.AddAuditAnnotation(ctx, podsecurityadmissionapi.AuditAnnotationPrefix+k, v)
+	if len(result.AuditAnnotations) > 0 {
+		annotations := make([]string, len(result.AuditAnnotations)*2)
+		i := 0
+		for k, v := range result.AuditAnnotations {
+			annotations[i], annotations[i+1] = podsecurityadmissionapi.AuditAnnotationPrefix+k, v
+			i += 2
+		}
+		audit.AddAuditAnnotations(ctx, annotations...)
 	}
 	if !result.Allowed {
 		// start with a generic forbidden error

staging/src/k8s.io/apiserver/pkg/admission/audit.go

@@ -19,19 +19,13 @@ package admission
 import (
 	"context"
 	"fmt"
-	"sync"
 
-	auditinternal "k8s.io/apiserver/pkg/apis/audit"
 	"k8s.io/apiserver/pkg/audit"
 )
 
 // auditHandler logs annotations set by other admission handlers
 type auditHandler struct {
 	Interface
-	// TODO: move the lock near the Annotations field of the audit event so it is always protected from concurrent access.
-	// to protect the 'Annotations' map of the audit event from concurrent writes
-	mutex sync.Mutex
-	ae    *auditinternal.Event
 }
 
 var _ Interface = &auditHandler{}
@@ -42,11 +36,11 @@ var _ ValidationInterface = &auditHandler{}
 // of attribute into the audit event. Attributes passed to the Admit and
 // Validate function must be instance of privateAnnotationsGetter or
 // AnnotationsGetter, otherwise an error is returned.
-func WithAudit(i Interface, ae *auditinternal.Event) Interface {
-	if i == nil || ae == nil {
+func WithAudit(i Interface) Interface {
+	if i == nil {
 		return i
 	}
-	return &auditHandler{Interface: i, ae: ae}
+	return &auditHandler{Interface: i}
 }
 
 func (handler *auditHandler) Admit(ctx context.Context, a Attributes, o ObjectInterfaces) error {
@@ -59,7 +53,7 @@ func (handler *auditHandler) Admit(ctx context.Context, a Attributes, o ObjectIn
 	var err error
 	if mutator, ok := handler.Interface.(MutationInterface); ok {
 		err = mutator.Admit(ctx, a, o)
-		handler.logAnnotations(a)
+		handler.logAnnotations(ctx, a)
 	}
 	return err
 }
@@ -74,7 +68,7 @@ func (handler *auditHandler) Validate(ctx context.Context, a Attributes, o Objec
 	var err error
 	if validator, ok := handler.Interface.(ValidationInterface); ok {
 		err = validator.Validate(ctx, a, o)
-		handler.logAnnotations(a)
+		handler.logAnnotations(ctx, a)
 	}
 	return err
 }
@@ -88,23 +82,21 @@ func ensureAnnotationGetter(a Attributes) error {
 	return fmt.Errorf("attributes must be an instance of privateAnnotationsGetter or AnnotationsGetter")
 }
 
-func (handler *auditHandler) logAnnotations(a Attributes) {
-	if handler.ae == nil {
+func (handler *auditHandler) logAnnotations(ctx context.Context, a Attributes) {
+	ae := audit.AuditEventFrom(ctx)
+	if ae == nil {
 		return
 	}
-	handler.mutex.Lock()
-	defer handler.mutex.Unlock()
 
+	var annotations map[string]string
 	switch a := a.(type) {
 	case privateAnnotationsGetter:
-		for key, value := range a.getAnnotations(handler.ae.Level) {
-			audit.LogAnnotation(handler.ae, key, value)
-		}
+		annotations = a.getAnnotations(ae.Level)
 	case AnnotationsGetter:
-		for key, value := range a.GetAnnotations(handler.ae.Level) {
-			audit.LogAnnotation(handler.ae, key, value)
-		}
+		annotations = a.GetAnnotations(ae.Level)
 	default:
 		// this will never happen, because we have already checked it in ensureAnnotationGetter
 	}
+
+	audit.AddAuditAnnotationsMap(ctx, annotations)
 }

staging/src/k8s.io/apiserver/pkg/admission/audit_test.go

@@ -24,6 +24,7 @@ import (
 
 	"k8s.io/apimachinery/pkg/runtime/schema"
 	auditinternal "k8s.io/apiserver/pkg/apis/audit"
+	"k8s.io/apiserver/pkg/audit"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -142,7 +143,8 @@ func TestWithAudit(t *testing.T) {
 	for tcName, tc := range testCases {
 		var handler Interface = fakeHandler{tc.admit, tc.admitAnnotations, tc.validate, tc.validateAnnotations, tc.handles}
 		ae := &auditinternal.Event{Level: auditinternal.LevelMetadata}
-		auditHandler := WithAudit(handler, ae)
+		ctx := audit.WithAuditContext(context.Background(), &audit.AuditContext{Event: ae})
+		auditHandler := WithAudit(handler)
 		a := attributes()
 
 		assert.Equal(t, handler.Handles(Create), auditHandler.Handles(Create), tcName+": WithAudit decorator should not effect the return value")
@@ -151,13 +153,13 @@ func TestWithAudit(t *testing.T) {
 		require.True(t, ok)
 		auditMutator, ok := auditHandler.(MutationInterface)
 		require.True(t, ok)
-		assert.Equal(t, mutator.Admit(context.TODO(), a, nil), auditMutator.Admit(context.TODO(), a, nil), tcName+": WithAudit decorator should not effect the return value")
+		assert.Equal(t, mutator.Admit(ctx, a, nil), auditMutator.Admit(ctx, a, nil), tcName+": WithAudit decorator should not effect the return value")
 
 		validator, ok := handler.(ValidationInterface)
 		require.True(t, ok)
 		auditValidator, ok := auditHandler.(ValidationInterface)
 		require.True(t, ok)
-		assert.Equal(t, validator.Validate(context.TODO(), a, nil), auditValidator.Validate(context.TODO(), a, nil), tcName+": WithAudit decorator should not effect the return value")
+		assert.Equal(t, validator.Validate(ctx, a, nil), auditValidator.Validate(ctx, a, nil), tcName+": WithAudit decorator should not effect the return value")
 
 		annotations := make(map[string]string, len(tc.admitAnnotations)+len(tc.validateAnnotations))
 		for k, v := range tc.admitAnnotations {
@@ -183,7 +185,8 @@ func TestWithAuditConcurrency(t *testing.T) {
 	}
 	var handler Interface = fakeHandler{admitAnnotations: admitAnnotations, handles: true}
 	ae := &auditinternal.Event{Level: auditinternal.LevelMetadata}
-	auditHandler := WithAudit(handler, ae)
+	ctx := audit.WithAuditContext(context.Background(), &audit.AuditContext{Event: ae})
+	auditHandler := WithAudit(handler)
 	a := attributes()
 
 	// Simulate the scenario store.DeleteCollection
@@ -197,7 +200,7 @@ func TestWithAuditConcurrency(t *testing.T) {
 			require.True(t, ok)
 			auditMutator, ok := auditHandler.(MutationInterface)
 			require.True(t, ok)
-			assert.Equal(t, mutator.Admit(context.TODO(), a, nil), auditMutator.Admit(context.TODO(), a, nil), "WithAudit decorator should not effect the return value")
+			assert.Equal(t, mutator.Admit(ctx, a, nil), auditMutator.Admit(ctx, a, nil), "WithAudit decorator should not effect the return value")
 		}()
 	}
 	wg.Wait()

staging/src/k8s.io/apiserver/pkg/audit/context.go

@@ -18,9 +18,12 @@ package audit
 
 import (
 	"context"
+	"fmt"
+	"sync"
 
 	auditinternal "k8s.io/apiserver/pkg/apis/audit"
 	genericapirequest "k8s.io/apiserver/pkg/endpoints/request"
+	"k8s.io/klog/v2"
 )
 
 // The key type is unexported to prevent collisions
@@ -28,15 +31,15 @@ type key int
 
 const (
 	// auditAnnotationsKey is the context key for the audit annotations.
-	// TODO: it's wasteful to store the audit annotations under a separate key, we
-	//  copy the request context twice for audit purposes. We should move the audit
-	//  annotations under AuditContext so we can get rid of the additional request
-	//  context copy.
+	// TODO: consolidate all audit info under the AuditContext, rather than storing 3 separate keys.
 	auditAnnotationsKey key = iota
 
 	// auditKey is the context key for storing the audit event that is being
 	// captured and the evaluated policy that applies to the given request.
 	auditKey
+
+	// auditAnnotationsMutexKey is the context key for the audit annotations mutex.
+	auditAnnotationsMutexKey
 )
 
 // annotations = *[]annotation instead of a map to preserve order of insertions
@@ -54,6 +57,7 @@ func WithAuditAnnotations(parent context.Context) context.Context {
 	if _, ok := parent.Value(auditAnnotationsKey).(*[]annotation); ok {
 		return parent
 	}
+	parent = withAuditAnnotationsMutex(parent)
 
 	var annotations []annotation // avoid allocations until we actually need it
 	return genericapirequest.WithValue(parent, auditAnnotationsKey, &annotations)
@@ -67,35 +71,126 @@ func WithAuditAnnotations(parent context.Context) context.Context {
 // Handlers that are unaware of their position in the overall request flow should
 // prefer AddAuditAnnotation over LogAnnotation to avoid dropping annotations.
 func AddAuditAnnotation(ctx context.Context, key, value string) {
-	// use the audit event directly if we have it
-	if ae := AuditEventFrom(ctx); ae != nil {
-		LogAnnotation(ae, key, value)
+	mutex, ok := auditAnnotationsMutex(ctx)
+	if !ok {
+		klog.ErrorS(nil, "Attempted to add audit annotations from unsupported request chain", "annotation", fmt.Sprintf("%s=%s", key, value))
 		return
 	}
 
-	annotations, ok := ctx.Value(auditAnnotationsKey).(*[]annotation)
-	if !ok {
-		return // adding audit annotation is not supported at this call site
+	mutex.Lock()
+	defer mutex.Unlock()
+
+	ae := AuditEventFrom(ctx)
+	var ctxAnnotations *[]annotation
+	if ae == nil {
+		ctxAnnotations, _ = ctx.Value(auditAnnotationsKey).(*[]annotation)
 	}
 
-	*annotations = append(*annotations, annotation{key: key, value: value})
+	addAuditAnnotationLocked(ae, ctxAnnotations, key, value)
+}
+
+// AddAuditAnnotations is a bulk version of AddAuditAnnotation. Refer to AddAuditAnnotation for
+// restrictions on when this can be called.
+// keysAndValues are the key-value pairs to add, and must have an even number of items.
+func AddAuditAnnotations(ctx context.Context, keysAndValues ...string) {
+	mutex, ok := auditAnnotationsMutex(ctx)
+	if !ok {
+		klog.ErrorS(nil, "Attempted to add audit annotations from unsupported request chain", "annotations", keysAndValues)
+		return
+	}
+
+	mutex.Lock()
+	defer mutex.Unlock()
+
+	ae := AuditEventFrom(ctx)
+	var ctxAnnotations *[]annotation
+	if ae == nil {
+		ctxAnnotations, _ = ctx.Value(auditAnnotationsKey).(*[]annotation)
+	}
+
+	if len(keysAndValues)%2 != 0 {
+		klog.Errorf("Dropping mismatched audit annotation %q", keysAndValues[len(keysAndValues)-1])
+	}
+	for i := 0; i < len(keysAndValues); i += 2 {
+		addAuditAnnotationLocked(ae, ctxAnnotations, keysAndValues[i], keysAndValues[i+1])
+	}
+}
+
+// AddAuditAnnotationsMap is a bulk version of AddAuditAnnotation. Refer to AddAuditAnnotation for
+// restrictions on when this can be called.
+func AddAuditAnnotationsMap(ctx context.Context, annotations map[string]string) {
+	mutex, ok := auditAnnotationsMutex(ctx)
+	if !ok {
+		klog.ErrorS(nil, "Attempted to add audit annotations from unsupported request chain", "annotations", annotations)
+		return
+	}
+
+	mutex.Lock()
+	defer mutex.Unlock()
+
+	ae := AuditEventFrom(ctx)
+	var ctxAnnotations *[]annotation
+	if ae == nil {
+		ctxAnnotations, _ = ctx.Value(auditAnnotationsKey).(*[]annotation)
+	}
+
+	for k, v := range annotations {
+		addAuditAnnotationLocked(ae, ctxAnnotations, k, v)
+	}
+}
+
+// addAuditAnnotationLocked is the shared code for recording an audit annotation. This method should
+// only be called while the auditAnnotationsMutex is locked.
+func addAuditAnnotationLocked(ae *auditinternal.Event, annotations *[]annotation, key, value string) {
+	if ae != nil {
+		logAnnotation(ae, key, value)
+	} else if annotations != nil {
+		*annotations = append(*annotations, annotation{key: key, value: value})
+	}
 }
 
 // This is private to prevent reads/write to the slice from outside of this package.
 // The audit event should be directly read to get access to the annotations.
-func auditAnnotationsFrom(ctx context.Context) []annotation {
-	annotations, ok := ctx.Value(auditAnnotationsKey).(*[]annotation)
+func addAuditAnnotationsFrom(ctx context.Context, ev *auditinternal.Event) {
+	mutex, ok := auditAnnotationsMutex(ctx)
 	if !ok {
-		return nil // adding audit annotation is not supported at this call site
+		klog.Errorf("Attempted to copy audit annotations from unsupported request chain")
+		return
 	}
 
-	return *annotations
+	mutex.Lock()
+	defer mutex.Unlock()
+
+	annotations, ok := ctx.Value(auditAnnotationsKey).(*[]annotation)
+	if !ok {
+		return // no annotations to copy
+	}
+
+	for _, kv := range *annotations {
+		logAnnotation(ev, kv.key, kv.value)
+	}
+}
+
+// LogAnnotation fills in the Annotations according to the key value pair.
+func logAnnotation(ae *auditinternal.Event, key, value string) {
+	if ae == nil || ae.Level.Less(auditinternal.LevelMetadata) {
+		return
+	}
+	if ae.Annotations == nil {
+		ae.Annotations = make(map[string]string)
+	}
+	if v, ok := ae.Annotations[key]; ok && v != value {
+		klog.Warningf("Failed to set annotations[%q] to %q for audit:%q, it has already been set to %q", key, value, ae.AuditID, ae.Annotations[key])
+		return
+	}
+	ae.Annotations[key] = value
 }
 
 // WithAuditContext returns a new context that stores the pair of the audit
 // configuration object that applies to the given request and
 // the audit event that is going to be written to the API audit log.
 func WithAuditContext(parent context.Context, ev *AuditContext) context.Context {
+	parent = withAuditAnnotationsMutex(parent)
 	return genericapirequest.WithValue(parent, auditKey, ev)
 }
 
@@ -114,3 +209,18 @@ func AuditContextFrom(ctx context.Context) *AuditContext {
 	ev, _ := ctx.Value(auditKey).(*AuditContext)
 	return ev
 }
+
+// WithAuditAnnotationMutex adds a mutex for guarding context.AddAuditAnnotation.
+func withAuditAnnotationsMutex(parent context.Context) context.Context {
+	if _, ok := parent.Value(auditAnnotationsMutexKey).(*sync.Mutex); ok {
+		return parent
+	}
+	var mutex sync.Mutex
+	return genericapirequest.WithValue(parent, auditAnnotationsMutexKey, &mutex)
+}
+
+// AuditAnnotationsMutex returns the audit annotations mutex from the context.
+func auditAnnotationsMutex(ctx context.Context) (*sync.Mutex, bool) {
+	mutex, ok := ctx.Value(auditAnnotationsMutexKey).(*sync.Mutex)
+	return mutex, ok
+}

staging/src/k8s.io/apiserver/pkg/audit/context_test.go

@@ -0,0 +1,120 @@
+/*
+Copyright 2021 The Kubernetes Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package audit
+
+import (
+	"context"
+	"fmt"
+	"sync"
+	"testing"
+
+	auditinternal "k8s.io/apiserver/pkg/apis/audit"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestAddAuditAnnotation(t *testing.T) {
+	const (
+		annotationKeyTemplate = "test-annotation-%d"
+		annotationValue       = "test-annotation-value"
+		numAnnotations        = 10
+	)
+
+	expectAnnotations := func(t *testing.T, annotations map[string]string) {
+		assert.Len(t, annotations, numAnnotations)
+	}
+	noopValidator := func(_ *testing.T, _ context.Context) {}
+	preEventValidator := func(t *testing.T, ctx context.Context) {
+		ev := auditinternal.Event{
+			Level: auditinternal.LevelMetadata,
+		}
+		addAuditAnnotationsFrom(ctx, &ev)
+		expectAnnotations(t, ev.Annotations)
+	}
+	postEventValidator := func(t *testing.T, ctx context.Context) {
+		ev := AuditEventFrom(ctx)
+		expectAnnotations(t, ev.Annotations)
+	}
+	postEventEmptyValidator := func(t *testing.T, ctx context.Context) {
+		ev := AuditEventFrom(ctx)
+		assert.Empty(t, ev.Annotations)
+	}
+
+	tests := []struct {
+		description string
+		ctx         context.Context
+		validator   func(t *testing.T, ctx context.Context)
+	}{{
+		description: "no audit",
+		ctx:         context.Background(),
+		validator:   noopValidator,
+	}, {
+		description: "no annotations context",
+		ctx:         WithAuditContext(context.Background(), newAuditContext(auditinternal.LevelMetadata)),
+		validator:   postEventValidator,
+	}, {
+		description: "no audit context",
+		ctx:         WithAuditAnnotations(context.Background()),
+		validator:   preEventValidator,
+	}, {
+		description: "both contexts metadata level",
+		ctx:         WithAuditContext(WithAuditAnnotations(context.Background()), newAuditContext(auditinternal.LevelMetadata)),
+		validator:   postEventValidator,
+	}, {
+		description: "both contexts none level",
+		ctx:         WithAuditContext(WithAuditAnnotations(context.Background()), newAuditContext(auditinternal.LevelNone)),
+		validator:   postEventEmptyValidator,
+	}}
+
+	for _, test := range tests {
+		t.Run(test.description, func(t *testing.T) {
+			var wg sync.WaitGroup
+			wg.Add(numAnnotations)
+			for i := 0; i < numAnnotations; i++ {
+				go func(i int) {
+					AddAuditAnnotation(test.ctx, fmt.Sprintf(annotationKeyTemplate, i), annotationValue)
+					wg.Done()
+				}(i)
+			}
+			wg.Wait()
+
+			test.validator(t, test.ctx)
+		})
+	}
+}
+
+func TestLogAnnotation(t *testing.T) {
+	ev := &auditinternal.Event{
+		Level:   auditinternal.LevelMetadata,
+		AuditID: "fake id",
+	}
+	logAnnotation(ev, "foo", "bar")
+	logAnnotation(ev, "foo", "baz")
+	assert.Equal(t, "bar", ev.Annotations["foo"], "audit annotation should not be overwritten.")
+
+	logAnnotation(ev, "qux", "")
+	logAnnotation(ev, "qux", "baz")
+	assert.Equal(t, "", ev.Annotations["qux"], "audit annotation should not be overwritten.")
+}
+
+func newAuditContext(l auditinternal.Level) *AuditContext {
+	return &AuditContext{
+		Event: &auditinternal.Event{
+			Level: l,
+		},
+	}
+}

staging/src/k8s.io/apiserver/pkg/audit/request.go

@@ -87,9 +87,7 @@ func NewEventFromRequest(req *http.Request, requestReceivedTimestamp time.Time,
 		}
 	}
 
-	for _, kv := range auditAnnotationsFrom(req.Context()) {
-		LogAnnotation(ev, kv.key, kv.value)
-	}
+	addAuditAnnotationsFrom(req.Context(), ev)
 
 	return ev, nil
 }
@@ -245,21 +243,6 @@ func encodeObject(obj runtime.Object, gv schema.GroupVersion, serializer runtime
 	}, nil
 }
 
-// LogAnnotation fills in the Annotations according to the key value pair.
-func LogAnnotation(ae *auditinternal.Event, key, value string) {
-	if ae == nil || ae.Level.Less(auditinternal.LevelMetadata) {
-		return
-	}
-	if ae.Annotations == nil {
-		ae.Annotations = make(map[string]string)
-	}
-	if v, ok := ae.Annotations[key]; ok && v != value {
-		klog.Warningf("Failed to set annotations[%q] to %q for audit:%q, it has already been set to %q", key, value, ae.AuditID, ae.Annotations[key])
-		return
-	}
-	ae.Annotations[key] = value
-}
-
 // truncate User-Agent if too long, otherwise return it directly.
 func maybeTruncateUserAgent(req *http.Request) string {
 	ua := req.UserAgent()

staging/src/k8s.io/apiserver/pkg/audit/request_test.go

@@ -25,26 +25,11 @@ import (
 	corev1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/runtime"
-	auditinternal "k8s.io/apiserver/pkg/apis/audit"
 
 	"github.com/google/go-cmp/cmp"
 	"github.com/stretchr/testify/assert"
 )
 
-func TestLogAnnotation(t *testing.T) {
-	ev := &auditinternal.Event{
-		Level:   auditinternal.LevelMetadata,
-		AuditID: "fake id",
-	}
-	LogAnnotation(ev, "foo", "bar")
-	LogAnnotation(ev, "foo", "baz")
-	assert.Equal(t, "bar", ev.Annotations["foo"], "audit annotation should not be overwritten.")
-
-	LogAnnotation(ev, "qux", "")
-	LogAnnotation(ev, "qux", "baz")
-	assert.Equal(t, "", ev.Annotations["qux"], "audit annotation should not be overwritten.")
-}
-
 func TestMaybeTruncateUserAgent(t *testing.T) {
 	req := &http.Request{}
 	req.Header = http.Header{}

staging/src/k8s.io/apiserver/pkg/endpoints/filters/audit.go

@@ -177,10 +177,8 @@ func writeLatencyToAnnotation(ctx context.Context, ev *auditinternal.Event) {
 	}
 
 	// record the total latency for this request, for convenience.
-	audit.LogAnnotation(ev, "apiserver.latency.k8s.io/total", latency.String())
-	for k, v := range layerLatencies {
-		audit.LogAnnotation(ev, k, v)
-	}
+	layerLatencies["apiserver.latency.k8s.io/total"] = latency.String()
+	audit.AddAuditAnnotationsMap(ctx, layerLatencies)
 }
 
 func processAuditEvent(ctx context.Context, sink audit.Sink, ev *auditinternal.Event, omitStages []auditinternal.Stage) bool {

staging/src/k8s.io/apiserver/pkg/endpoints/filters/authorization.go

@@ -49,7 +49,6 @@ func WithAuthorization(handler http.Handler, a authorizer.Authorizer, s runtime.
 	}
 	return http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
 		ctx := req.Context()
-		ae := audit.AuditEventFrom(ctx)
 
 		attributes, err := GetAuthorizerAttributes(ctx)
 		if err != nil {
@@ -59,20 +58,22 @@ func WithAuthorization(handler http.Handler, a authorizer.Authorizer, s runtime.
 		authorized, reason, err := a.Authorize(ctx, attributes)
 		// an authorizer like RBAC could encounter evaluation errors and still allow the request, so authorizer decision is checked before error here.
 		if authorized == authorizer.DecisionAllow {
-			audit.LogAnnotation(ae, decisionAnnotationKey, decisionAllow)
-			audit.LogAnnotation(ae, reasonAnnotationKey, reason)
+			audit.AddAuditAnnotations(ctx,
+				decisionAnnotationKey, decisionAllow,
+				reasonAnnotationKey, reason)
 			handler.ServeHTTP(w, req)
 			return
 		}
 		if err != nil {
-			audit.LogAnnotation(ae, reasonAnnotationKey, reasonError)
+			audit.AddAuditAnnotation(ctx, reasonAnnotationKey, reasonError)
 			responsewriters.InternalError(w, req, err)
 			return
 		}
 
 		klog.V(4).InfoS("Forbidden", "URI", req.RequestURI, "Reason", reason)
-		audit.LogAnnotation(ae, decisionAnnotationKey, decisionForbid)
-		audit.LogAnnotation(ae, reasonAnnotationKey, reason)
+		audit.AddAuditAnnotations(ctx,
+			decisionAnnotationKey, decisionForbid,
+			reasonAnnotationKey, reason)
 		responsewriters.Forbidden(ctx, attributes, w, req, reason, s)
 	})
 }

staging/src/k8s.io/apiserver/pkg/endpoints/handlers/create.go

@@ -157,8 +157,7 @@ func createHandler(r rest.NamedCreater, scope *RequestScope, admit admission.Int
 		}
 		ctx = request.WithNamespace(ctx, namespace)
 
-		ae := audit.AuditEventFrom(ctx)
-		admit = admission.WithAudit(admit, ae)
+		admit = admission.WithAudit(admit)
 		audit.LogRequestObject(req.Context(), obj, objGV, scope.Resource, scope.Subresource, scope.Serializer)
 
 		userInfo, _ := request.UserFrom(ctx)

staging/src/k8s.io/apiserver/pkg/endpoints/handlers/delete.go

@@ -67,8 +67,7 @@ func DeleteResource(r rest.GracefulDeleter, allowsOptions bool, scope *RequestSc
 		defer cancel()
 
 		ctx = request.WithNamespace(ctx, namespace)
-		ae := audit.AuditEventFrom(ctx)
-		admit = admission.WithAudit(admit, ae)
+		admit = admission.WithAudit(admit)
 
 		outputMediaType, _, err := negotiation.NegotiateOutputMediaType(req, scope.Serializer, scope)
 		if err != nil {
@@ -190,7 +189,6 @@ func DeleteCollection(r rest.CollectionDeleter, checkBody bool, scope *RequestSc
 		defer cancel()
 
 		ctx = request.WithNamespace(ctx, namespace)
-		ae := audit.AuditEventFrom(ctx)
 
 		outputMediaType, _, err := negotiation.NegotiateOutputMediaType(req, scope.Serializer, scope)
 		if err != nil {
@@ -268,7 +266,7 @@ func DeleteCollection(r rest.CollectionDeleter, checkBody bool, scope *RequestSc
 		}
 		options.TypeMeta.SetGroupVersionKind(metav1.SchemeGroupVersion.WithKind("DeleteOptions"))
 
-		admit = admission.WithAudit(admit, ae)
+		admit = admission.WithAudit(admit)
 		userInfo, _ := request.UserFrom(ctx)
 		staticAdmissionAttrs := admission.NewAttributesRecord(nil, nil, scope.Kind, namespace, "", scope.Resource, scope.Subresource, admission.Delete, options, dryrun.IsDryRun(options.DryRun), userInfo)
 		result, err := finisher.FinishRequest(ctx, func() (runtime.Object, error) {

staging/src/k8s.io/apiserver/pkg/endpoints/handlers/patch.go

@@ -124,8 +124,7 @@ func PatchResource(r rest.Patcher, scope *RequestScope, admit admission.Interfac
 		}
 		options.TypeMeta.SetGroupVersionKind(metav1.SchemeGroupVersion.WithKind("PatchOptions"))
 
-		ae := audit.AuditEventFrom(ctx)
-		admit = admission.WithAudit(admit, ae)
+		admit = admission.WithAudit(admit)
 
 		audit.LogRequestPatch(req.Context(), patchBytes)
 		trace.Step("Recorded the audit event")

staging/src/k8s.io/apiserver/pkg/endpoints/handlers/rest.go

@@ -39,7 +39,6 @@ import (
 	"k8s.io/apimachinery/pkg/runtime/schema"
 	"k8s.io/apimachinery/pkg/types"
 	"k8s.io/apiserver/pkg/admission"
-	"k8s.io/apiserver/pkg/audit"
 	"k8s.io/apiserver/pkg/authorization/authorizer"
 	"k8s.io/apiserver/pkg/endpoints/handlers/fieldmanager"
 	"k8s.io/apiserver/pkg/endpoints/handlers/responsewriters"
@@ -191,8 +190,7 @@ func ConnectResource(connecter rest.Connecter, scope *RequestScope, admit admiss
 		}
 		ctx := req.Context()
 		ctx = request.WithNamespace(ctx, namespace)
-		ae := audit.AuditEventFrom(ctx)
-		admit = admission.WithAudit(admit, ae)
+		admit = admission.WithAudit(admit)
 
 		opts, subpath, subpathKey := connecter.NewConnectOptions()
 		if err := getRequestOptions(req, scope, opts, subpath, subpathKey, isSubresource); err != nil {

staging/src/k8s.io/apiserver/pkg/endpoints/handlers/update.go

@@ -135,9 +135,8 @@ func UpdateResource(r rest.Updater, scope *RequestScope, admit admission.Interfa
 		}
 		trace.Step("Conversion done")
 
-		ae := audit.AuditEventFrom(ctx)
 		audit.LogRequestObject(req.Context(), obj, objGV, scope.Resource, scope.Subresource, scope.Serializer)
-		admit = admission.WithAudit(admit, ae)
+		admit = admission.WithAudit(admit)
 
 		// if this object supports namespace info
 		if objectMeta, err := meta.Accessor(obj); err == nil {

staging/src/k8s.io/apiserver/pkg/server/config_test.go

@@ -312,9 +312,6 @@ func TestAuthenticationAuditAnnotationsDefaultChain(t *testing.T) {
 			t.Error("unexpected nil audit event")
 		}
 
-		// confirm that the direct way of setting audit annotations later in the chain works as expected
-		audit.LogAnnotation(ae, "snorlax", "is cool too")
-
 		// confirm that the indirect way of setting audit annotations later in the chain also works
 		audit.AddAuditAnnotation(r.Context(), "dogs", "are okay")
 
@@ -334,7 +331,7 @@ func TestAuthenticationAuditAnnotationsDefaultChain(t *testing.T) {
 		t.Error("expected audit events, got none")
 	}
 	// these should all be the same because the handler chain mutates the event in place
-	want := map[string]string{"pandas": "are awesome", "snorlax": "is cool too", "dogs": "are okay"}
+	want := map[string]string{"pandas": "are awesome", "dogs": "are okay"}
 	for _, event := range backend.events {
 		if event.Stage != auditinternal.StageResponseComplete {
 			t.Errorf("expected event stage to be complete, got: %s", event.Stage)

staging/src/k8s.io/apiserver/plugin/pkg/audit/truncate/truncate.go

@@ -126,7 +126,12 @@ func truncate(e *auditinternal.Event) *auditinternal.Event {
 
 	newEvent.RequestObject = nil
 	newEvent.ResponseObject = nil
-	audit.LogAnnotation(newEvent, annotationKey, annotationValue)
+
+	if newEvent.Annotations == nil {
+		newEvent.Annotations = make(map[string]string)
+	}
+	newEvent.Annotations[annotationKey] = annotationValue
+
 	return newEvent
 }