Format requests/responses ahead of time in file audit

This was inspired by #7022 but has the advantage of avoiding double-locking and needing to perform lock upgrades while also simplifying the logic and being faster. Original, #7022, this: goos: linux goarch: amd64 pkg: github.com/hashicorp/vault/builtin/audit/file BenchmarkAuditFile_request-4 30000 60734 ns/op PASS ok github.com/hashicorp/vault/builtin/audit/file 2.428s goos: linux goarch: amd64 pkg: github.com/hashicorp/vault/builtin/audit/file BenchmarkAuditFile_request-4 50000 34772 ns/op PASS ok github.com/hashicorp/vault/builtin/audit/file 2.086s goos: linux goarch: amd64 pkg: github.com/hashicorp/vault/builtin/audit/file BenchmarkAuditFile_request-4 50000 25302 ns/op PASS ok github.com/hashicorp/vault/builtin/audit/file 1.542s Fixes #7014 Closes #7022
hashicorp · Jun 30, 2019 · ae47a01 · ae47a01
1 parent ceb4aea
commit ae47a01
Show file tree

Hide file tree

Showing 2 changed files with 117 additions and 41 deletions.
diff --git a/builtin/audit/file/backend.go b/builtin/audit/file/backend.go
@@ -1,14 +1,16 @@
 package file
 
 import (
+	"bytes"
 	"context"
 	"fmt"
-	"io/ioutil"
+	"io"
 	"os"
 	"path/filepath"
 	"strconv"
 	"strings"
 	"sync"
+	"sync/atomic"
 
 	"github.com/hashicorp/errwrap"
 	"github.com/hashicorp/vault/audit"
@@ -87,12 +89,17 @@ func Factory(ctx context.Context, conf *audit.BackendConfig) (audit.Backend, err
 		mode:       mode,
 		saltConfig: conf.SaltConfig,
 		saltView:   conf.SaltView,
+		salt:       new(atomic.Value),
 		formatConfig: audit.FormatterConfig{
 			Raw:          logRaw,
 			HMACAccessor: hmacAccessor,
 		},
 	}
 
+	// Ensure we are working with the right type by explicitly storing a nil of
+	// the right type
+	b.salt.Store((*salt.Salt)(nil))
+
 	switch format {
 	case "json":
 		b.formatter.AuditFormatWriter = &audit.JSONFormatWriter{
@@ -137,100 +144,118 @@ type Backend struct {
 	mode     os.FileMode
 
 	saltMutex  sync.RWMutex
-	salt       *salt.Salt
+	salt       *atomic.Value
 	saltConfig *salt.Config
 	saltView   logical.Storage
 }
 
 var _ audit.Backend = (*Backend)(nil)
 
 func (b *Backend) Salt(ctx context.Context) (*salt.Salt, error) {
-	b.saltMutex.RLock()
-	if b.salt != nil {
-		defer b.saltMutex.RUnlock()
-		return b.salt, nil
+	s := b.salt.Load().(*salt.Salt)
+	if s != nil {
+		return s, nil
 	}
-	b.saltMutex.RUnlock()
+
 	b.saltMutex.Lock()
 	defer b.saltMutex.Unlock()
-	if b.salt != nil {
-		return b.salt, nil
+
+	s = b.salt.Load().(*salt.Salt)
+	if s != nil {
+		return s, nil
 	}
-	salt, err := salt.NewSalt(ctx, b.saltView, b.saltConfig)
+
+	newSalt, err := salt.NewSalt(ctx, b.saltView, b.saltConfig)
 	if err != nil {
+		b.salt.Store((*salt.Salt)(nil))
 		return nil, err
 	}
-	b.salt = salt
-	return salt, nil
+
+	b.salt.Store(newSalt)
+	return newSalt, nil
 }
 
 func (b *Backend) GetHash(ctx context.Context, data string) (string, error) {
 	salt, err := b.Salt(ctx)
 	if err != nil {
 		return "", err
 	}
+
 	return audit.HashString(salt, data), nil
 }
 
 func (b *Backend) LogRequest(ctx context.Context, in *logical.LogInput) error {
-	b.fileLock.Lock()
-	defer b.fileLock.Unlock()
-
+	var writer io.Writer
 	switch b.path {
 	case "stdout":
-		return b.formatter.FormatRequest(ctx, os.Stdout, b.formatConfig, in)
+		writer = os.Stdout
 	case "discard":
-		return b.formatter.FormatRequest(ctx, ioutil.Discard, b.formatConfig, in)
+		return nil
 	}
 
-	if err := b.open(); err != nil {
+	buf := bytes.NewBuffer(make([]byte, 0, 2000))
+	err := b.formatter.FormatRequest(ctx, buf, b.formatConfig, in)
+	if err != nil {
 		return err
 	}
 
-	if err := b.formatter.FormatRequest(ctx, b.f, b.formatConfig, in); err == nil {
+	return b.log(ctx, buf, writer)
+}
+
+func (b *Backend) log(ctx context.Context, buf *bytes.Buffer, writer io.Writer) error {
+	reader := bytes.NewReader(buf.Bytes())
+
+	b.fileLock.Lock()
+
+	if writer == nil {
+		if err := b.open(); err != nil {
+			b.fileLock.Unlock()
+			return err
+		}
+		writer = b.f
+	}
+
+	if _, err := reader.WriteTo(writer); err == nil {
+		b.fileLock.Unlock()
 		return nil
+	} else if b.path == "stdout" {
+		b.fileLock.Unlock()
+		return err
 	}
 
-	// Opportunistically try to re-open the FD, once per call
+	// If writing to stdout there's no real reason to think anything would have
+	// changed so return above. Otherwise, opportunistically try to re-open the
+	// FD, once per call.
 	b.f.Close()
 	b.f = nil
 
 	if err := b.open(); err != nil {
+		b.fileLock.Unlock()
 		return err
 	}
 
-	return b.formatter.FormatRequest(ctx, b.f, b.formatConfig, in)
+	reader.Seek(0, io.SeekStart)
+	_, err := reader.WriteTo(writer)
+	b.fileLock.Unlock()
+	return err
 }
 
 func (b *Backend) LogResponse(ctx context.Context, in *logical.LogInput) error {
-
-	b.fileLock.Lock()
-	defer b.fileLock.Unlock()
-
+	var writer io.Writer
 	switch b.path {
 	case "stdout":
-		return b.formatter.FormatResponse(ctx, os.Stdout, b.formatConfig, in)
+		writer = os.Stdout
 	case "discard":
-		return b.formatter.FormatResponse(ctx, ioutil.Discard, b.formatConfig, in)
-	}
-
-	if err := b.open(); err != nil {
-		return err
-	}
-
-	if err := b.formatter.FormatResponse(ctx, b.f, b.formatConfig, in); err == nil {
 		return nil
 	}
 
-	// Opportunistically try to re-open the FD, once per call
-	b.f.Close()
-	b.f = nil
-
-	if err := b.open(); err != nil {
+	buf := bytes.NewBuffer(make([]byte, 0, 6000))
+	err := b.formatter.FormatResponse(ctx, buf, b.formatConfig, in)
+	if err != nil {
 		return err
 	}
 
-	return b.formatter.FormatResponse(ctx, b.f, b.formatConfig, in)
+	return b.log(ctx, buf, writer)
 }
 
 // The file lock must be held before calling this
@@ -291,5 +316,5 @@ func (b *Backend) Reload(_ context.Context) error {
 func (b *Backend) Invalidate(_ context.Context) {
 	b.saltMutex.Lock()
 	defer b.saltMutex.Unlock()
-	b.salt = nil
+	b.salt.Store((*salt.Salt)(nil))
 }
diff --git a/builtin/audit/file/backend_test.go b/builtin/audit/file/backend_test.go
@@ -7,8 +7,10 @@ import (
 	"path/filepath"
 	"strconv"
 	"testing"
+	"time"
 
 	"github.com/hashicorp/vault/audit"
+	"github.com/hashicorp/vault/helper/namespace"
 	"github.com/hashicorp/vault/sdk/helper/salt"
 	"github.com/hashicorp/vault/sdk/logical"
 )
@@ -90,3 +92,52 @@ func TestAuditFile_fileModeExisting(t *testing.T) {
 		t.Fatalf("File mode does not match.")
 	}
 }
+
+func BenchmarkAuditFile_request(b *testing.B) {
+	config := map[string]string{
+		"path": "/dev/null",
+	}
+	sink, err := Factory(context.Background(), &audit.BackendConfig{
+		Config:     config,
+		SaltConfig: &salt.Config{},
+		SaltView:   &logical.InmemStorage{},
+	})
+	if err != nil {
+		b.Fatal(err)
+	}
+
+	in := &logical.LogInput{
+		Auth: &logical.Auth{
+			ClientToken:     "foo",
+			Accessor:        "bar",
+			EntityID:        "foobarentity",
+			DisplayName:     "testtoken",
+			NoDefaultPolicy: true,
+			Policies:        []string{"root"},
+			TokenType:       logical.TokenTypeService,
+		},
+		Request: &logical.Request{
+			Operation: logical.UpdateOperation,
+			Path:      "/foo",
+			Connection: &logical.Connection{
+				RemoteAddr: "127.0.0.1",
+			},
+			WrapInfo: &logical.RequestWrapInfo{
+				TTL: 60 * time.Second,
+			},
+			Headers: map[string][]string{
+				"foo": []string{"bar"},
+			},
+		},
+	}
+
+	ctx := namespace.RootContext(nil)
+	b.ResetTimer()
+	b.RunParallel(func(pb *testing.PB) {
+		for pb.Next() {
+			if err := sink.LogRequest(ctx, in); err != nil {
+				panic(err)
+			}
+		}
+	})
+}