Allow handlers early in the request chain to set audit annotations

This change adds the generic ability for request handlers that run
before WithAudit to set annotations in the audit.Event.Annotations
map.

Note that this change does not use this capability yet. Determining
which handlers should set audit annotations and what keys and values
should be used requires further discussion (this data will become
part of our public API).

Signed-off-by: Monis Khan <[email protected]>

Author: enj

staging/src/k8s.io/apiserver/pkg/audit/BUILD

@@ -9,6 +9,7 @@ load(
 go_library(
     name = "go_default_library",
     srcs = [
+        "context.go",
         "format.go",
         "metrics.go",
         "request.go",
@@ -35,6 +36,7 @@ go_library(
         "//staging/src/k8s.io/apiserver/pkg/apis/audit/v1beta1:go_default_library",
         "//staging/src/k8s.io/apiserver/pkg/authentication/user:go_default_library",
         "//staging/src/k8s.io/apiserver/pkg/authorization/authorizer:go_default_library",
+        "//staging/src/k8s.io/apiserver/pkg/endpoints/request:go_default_library",
         "//staging/src/k8s.io/component-base/metrics:go_default_library",
         "//staging/src/k8s.io/component-base/metrics/legacyregistry:go_default_library",
         "//vendor/github.com/google/uuid:go_default_library",

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

@@ -0,0 +1,84 @@
+/*
+Copyright 2020 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"
+
+	genericapirequest "k8s.io/apiserver/pkg/endpoints/request"
+)
+
+// The key type is unexported to prevent collisions
+type key int
+
+const (
+	// auditAnnotationsKey is the context key for the audit annotations.
+	auditAnnotationsKey key = iota
+)
+
+// annotations = *[]annotation instead of a map to preserve order of insertions
+type annotation struct {
+	key, value string
+}
+
+// WithAuditAnnotations returns a new context that can store audit annotations
+// via the AddAuditAnnotation function.  This function is meant to be called from
+// an early request handler to allow all later layers to set audit annotations.
+// This is required to support flows where handlers that come before WithAudit
+// (such as WithAuthentication) wish to set audit annotations.
+func WithAuditAnnotations(parent context.Context) context.Context {
+	// this should never really happen, but prevent double registration of this slice
+	if _, ok := parent.Value(auditAnnotationsKey).(*[]annotation); ok {
+		return parent
+	}
+
+	var annotations []annotation // avoid allocations until we actually need it
+	return genericapirequest.WithValue(parent, auditAnnotationsKey, &annotations)
+}
+
+// AddAuditAnnotation sets the audit annotation for the given key, value pair.
+// It is safe to call at most parts of request flow that come after WithAuditAnnotations.
+// The notable exception being that this function must not be called via a
+// defer statement (i.e. after ServeHTTP) in a handler that runs before WithAudit
+// as at that point the audit event has already been sent to the audit sink.
+// 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 := genericapirequest.AuditEventFrom(ctx); ae != nil {
+		LogAnnotation(ae, key, value)
+		return
+	}
+
+	annotations, ok := ctx.Value(auditAnnotationsKey).(*[]annotation)
+	if !ok {
+		return // adding audit annotation is not supported at this call site
+	}
+
+	*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)
+	if !ok {
+		return nil // adding audit annotation is not supported at this call site
+	}
+
+	return *annotations
+}

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

@@ -88,6 +88,10 @@ func NewEventFromRequest(req *http.Request, level auditinternal.Level, attribs a
 		}
 	}
 
+	for _, kv := range auditAnnotationsFrom(req.Context()) {
+		LogAnnotation(ev, kv.key, kv.value)
+	}
+
 	return ev, nil
 }
 

staging/src/k8s.io/apiserver/pkg/endpoints/filters/BUILD

@@ -45,6 +45,7 @@ go_library(
     name = "go_default_library",
     srcs = [
         "audit.go",
+        "audit_annotations.go",
         "authentication.go",
         "authn_audit.go",
         "authorization.go",

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

@@ -0,0 +1,39 @@
+/*
+Copyright 2020 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 filters
+
+import (
+	"net/http"
+
+	"k8s.io/apiserver/pkg/audit"
+	"k8s.io/apiserver/pkg/audit/policy"
+)
+
+// WithAuditAnnotations decorates a http.Handler with a []{key, value} that is merged
+// with the audit.Event.Annotations map.  This allows layers that run before WithAudit
+// (such as authentication) to assert annotations.
+// If sink or audit policy is nil, no decoration takes place.
+func WithAuditAnnotations(handler http.Handler, sink audit.Sink, policy policy.Checker) http.Handler {
+	// no need to wrap if auditing is disabled
+	if sink == nil || policy == nil {
+		return handler
+	}
+	return http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
+		req = req.WithContext(audit.WithAuditAnnotations(req.Context()))
+		handler.ServeHTTP(w, req)
+	})
+}

staging/src/k8s.io/apiserver/pkg/server/BUILD

@@ -24,20 +24,28 @@ go_test(
         "//staging/src/k8s.io/apimachinery/pkg/util/json:go_default_library",
         "//staging/src/k8s.io/apimachinery/pkg/util/runtime:go_default_library",
         "//staging/src/k8s.io/apimachinery/pkg/util/sets:go_default_library",
+        "//staging/src/k8s.io/apimachinery/pkg/util/waitgroup:go_default_library",
         "//staging/src/k8s.io/apimachinery/pkg/version:go_default_library",
+        "//staging/src/k8s.io/apiserver/pkg/apis/audit:go_default_library",
         "//staging/src/k8s.io/apiserver/pkg/apis/example:go_default_library",
         "//staging/src/k8s.io/apiserver/pkg/apis/example/v1:go_default_library",
+        "//staging/src/k8s.io/apiserver/pkg/audit:go_default_library",
+        "//staging/src/k8s.io/apiserver/pkg/audit/policy:go_default_library",
+        "//staging/src/k8s.io/apiserver/pkg/authentication/authenticator:go_default_library",
+        "//staging/src/k8s.io/apiserver/pkg/authentication/user:go_default_library",
         "//staging/src/k8s.io/apiserver/pkg/authorization/authorizer:go_default_library",
         "//staging/src/k8s.io/apiserver/pkg/endpoints/discovery:go_default_library",
         "//staging/src/k8s.io/apiserver/pkg/endpoints/filters:go_default_library",
         "//staging/src/k8s.io/apiserver/pkg/endpoints/openapi:go_default_library",
+        "//staging/src/k8s.io/apiserver/pkg/endpoints/request:go_default_library",
         "//staging/src/k8s.io/apiserver/pkg/registry/rest:go_default_library",
         "//staging/src/k8s.io/apiserver/pkg/server/filters:go_default_library",
         "//staging/src/k8s.io/apiserver/pkg/server/healthz:go_default_library",
         "//staging/src/k8s.io/client-go/informers:go_default_library",
         "//staging/src/k8s.io/client-go/kubernetes/fake:go_default_library",
         "//staging/src/k8s.io/client-go/rest:go_default_library",
         "//vendor/github.com/go-openapi/spec:go_default_library",
+        "//vendor/github.com/google/go-cmp/cmp:go_default_library",
         "//vendor/github.com/stretchr/testify/assert:go_default_library",
         "//vendor/k8s.io/kube-openapi/pkg/common:go_default_library",
         "//vendor/k8s.io/utils/net:go_default_library",

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

@@ -676,6 +676,7 @@ func DefaultBuildHandlerChain(apiHandler http.Handler, c *Config) http.Handler {
 	if c.SecureServing != nil && !c.SecureServing.DisableHTTP2 && c.GoawayChance > 0 {
 		handler = genericfilters.WithProbabilisticGoaway(handler, c.GoawayChance)
 	}
+	handler = genericapifilters.WithAuditAnnotations(handler, c.AuditBackend, c.AuditPolicyChecker)
 	handler = genericfilters.WithPanicRecovery(handler)
 	return handler
 }

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

@@ -25,9 +25,18 @@ import (
 	"net/http/httputil"
 	"reflect"
 	"testing"
+	"time"
 
+	"github.com/google/go-cmp/cmp"
 	"k8s.io/apimachinery/pkg/util/json"
 	"k8s.io/apimachinery/pkg/util/sets"
+	"k8s.io/apimachinery/pkg/util/waitgroup"
+	auditinternal "k8s.io/apiserver/pkg/apis/audit"
+	"k8s.io/apiserver/pkg/audit"
+	"k8s.io/apiserver/pkg/audit/policy"
+	"k8s.io/apiserver/pkg/authentication/authenticator"
+	"k8s.io/apiserver/pkg/authentication/user"
+	"k8s.io/apiserver/pkg/endpoints/request"
 	"k8s.io/apiserver/pkg/server/healthz"
 	"k8s.io/client-go/informers"
 	"k8s.io/client-go/kubernetes/fake"
@@ -241,3 +250,84 @@ func checkExpectedPathsAtRoot(url string, expectedPaths []string, t *testing.T)
 		}
 	})
 }
+
+func TestAuthenticationAuditAnnotationsDefaultChain(t *testing.T) {
+	authn := authenticator.RequestFunc(func(req *http.Request) (*authenticator.Response, bool, error) {
+		// confirm that we can set an audit annotation in a handler before WithAudit
+		audit.AddAuditAnnotation(req.Context(), "pandas", "are awesome")
+
+		// confirm that trying to use the audit event directly would never work
+		if ae := request.AuditEventFrom(req.Context()); ae != nil {
+			t.Errorf("expected nil audit event, got %v", ae)
+		}
+
+		return &authenticator.Response{User: &user.DefaultInfo{}}, true, nil
+	})
+	backend := &testBackend{}
+	c := &Config{
+		Authentication:     AuthenticationInfo{Authenticator: authn},
+		AuditBackend:       backend,
+		AuditPolicyChecker: policy.FakeChecker(auditinternal.LevelMetadata, nil),
+
+		// avoid nil panics
+		HandlerChainWaitGroup: &waitgroup.SafeWaitGroup{},
+		RequestInfoResolver:   &request.RequestInfoFactory{},
+		RequestTimeout:        10 * time.Second,
+		LongRunningFunc:       func(_ *http.Request, _ *request.RequestInfo) bool { return false },
+	}
+
+	h := DefaultBuildHandlerChain(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		// confirm this is a no-op
+		if r.Context() != audit.WithAuditAnnotations(r.Context()) {
+			t.Error("unexpected double wrapping of context")
+		}
+
+		// confirm that we have an audit event
+		ae := request.AuditEventFrom(r.Context())
+		if ae == nil {
+			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")
+
+		if _, err := w.Write([]byte("done")); err != nil {
+			t.Errorf("failed to write response: %v", err)
+		}
+	}), c)
+	w := httptest.NewRecorder()
+
+	h.ServeHTTP(w, httptest.NewRequest("GET", "https://ignored.com", nil))
+
+	r := w.Result()
+	if ok := r.StatusCode == http.StatusOK && w.Body.String() == "done" && len(r.Header.Get(auditinternal.HeaderAuditID)) > 0; !ok {
+		t.Errorf("invalid response: %#v", w)
+	}
+	if len(backend.events) == 0 {
+		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"}
+	for _, event := range backend.events {
+		if event.Stage != auditinternal.StageResponseComplete {
+			t.Errorf("expected event stage to be complete, got: %s", event.Stage)
+		}
+		if diff := cmp.Diff(want, event.Annotations); diff != "" {
+			t.Errorf("event has unexpected annotations (-want +got): %s", diff)
+		}
+	}
+}
+
+type testBackend struct {
+	events []*auditinternal.Event
+
+	audit.Backend // nil panic if anything other than ProcessEvents called
+}
+
+func (b *testBackend) ProcessEvents(events ...*auditinternal.Event) bool {
+	b.events = append(b.events, events...)
+	return true
+}