Reduces allocations when masking sensitive values

Author: gdavison

logging/aws.go

@@ -5,6 +5,7 @@ package logging
 
 import (
 	"regexp"
+	"unsafe"
 )
 
 // IAM Unique ID prefixes from
@@ -24,48 +25,45 @@ var UniqueIDRegex = regexp.MustCompile(`(A3T[A-Z0-9]` +
 	`|ASIA` + // STS temporary access key
 	`)[A-Z0-9]{16,}`)
 
-var SensitiveKeyRegex = regexp.MustCompile(`[A-Za-z0-9/+=]{16,}`)
-
 const (
 	unmaskedFirst = 4
 	unmaskedLast  = 4
 )
 
-func MaskAWSAccessKey(field string) string {
-	field = UniqueIDRegex.ReplaceAllStringFunc(field, func(s string) string {
+func MaskAWSAccessKey(field []byte) []byte {
+	field = UniqueIDRegex.ReplaceAllFunc(field, func(s []byte) []byte {
 		return partialMaskString(s, unmaskedFirst, unmaskedLast)
 	})
 	return field
 }
 
 func MaskAWSSensitiveValues(field string) string {
-	field = MaskAWSAccessKey(field)
-	field = MaskAWSSecretKeys(field)
-	return field
+	b := unsafe.Slice(unsafe.StringData(field), len(field))
+	b = MaskAWSAccessKey(b)
+	MaskAWSSecretKeys(b)
+	return unsafe.String(unsafe.SliceData(b), len(b))
 }
 
 // MaskAWSSecretKeys masks likely AWS secret access keys in the input.
 // See https://aws.amazon.com/blogs/security/a-safer-way-to-distribute-aws-credentials-to-ec2/:
 // "Find me 40-character, base-64 strings that don’t have any base 64 characters immediately before or after".
-func MaskAWSSecretKeys(in string) string {
+func MaskAWSSecretKeys(in []byte) {
 	const (
 		secretKeyLen = 40
 	)
 	len := len(in)
-	out := make([]byte, len)
 	base64Characters := 0
 
 	for i := 0; i < len; i++ {
 		b := in[i]
-		out[i] = b
 
 		if (b >= 'A' && b <= 'Z') || (b >= 'a' && b <= 'z') || (b >= '0' && b <= '9') || b == '/' || b == '+' || b == '=' {
 			// base64 character.
 			base64Characters++
 		} else {
 			if base64Characters == secretKeyLen {
 				for j := (i - secretKeyLen) + unmaskedFirst; j < i-unmaskedLast; j++ {
-					out[j] = '*'
+					in[j] = '*'
 				}
 			}
 
@@ -75,9 +73,7 @@ func MaskAWSSecretKeys(in string) string {
 
 	if base64Characters == secretKeyLen {
 		for j := (len - secretKeyLen) + unmaskedFirst; j < len-unmaskedLast; j++ {
-			out[j] = '*'
+			in[j] = '*'
 		}
 	}
-
-	return string(out)
 }

logging/aws_test.go

@@ -5,6 +5,7 @@ package logging
 
 import (
 	"testing"
+	"unsafe"
 )
 
 func TestMaskAWSSensitiveValues(t *testing.T) {
@@ -101,42 +102,38 @@ func TestMaskAWSSensitiveValues(t *testing.T) {
 }
 
 func BenchmarkMaskAWSAccessKey(b *testing.B) {
-	var s string
 	b.ReportAllocs()
 	for n := 0; n < b.N; n++ {
-		s = MaskAWSAccessKey(`
+		MaskAWSAccessKey([]byte(`
 {
 	"AWSSecretKey": "LEfH8nZmFN4BGIJnku6lkChHydRN5B/YlWCIjOte",
 	"BucketName": "test-bucket",
 	"AWSKeyId": "AIDACKCEVSQ6C2EXAMPLE",
 }
-`)
+`))
 	}
-	dump = s
 }
 
 func BenchmarkPartialMaskString(b *testing.B) {
-	var s string
+	var s []byte
 	b.ReportAllocs()
 	for n := 0; n < b.N; n++ {
-		s = partialMaskString("AIDACKCEVSQ6C2EXAMPLE", 4, 4)
+		s = partialMaskString([]byte("AIDACKCEVSQ6C2EXAMPLE"), 4, 4)
 	}
-	dump = s
+	dump = unsafe.String(unsafe.SliceData(s), len(s))
 }
 
 func BenchmarkMaskAWSSecretKeys(b *testing.B) {
-	var s string
 	b.ReportAllocs()
 	for n := 0; n < b.N; n++ {
-		s = MaskAWSSecretKeys(`
+		MaskAWSSecretKeys([]byte(`
 {
 	"AWSSecretKey": "LEfH8nZmFN4BGIJnku6lkChHydRN5B/YlWCIjOte",
 	"BucketName": "test-bucket",
 	"AWSKeyId": "AIDACKCEVSQ6C2EXAMPLE",
 }
-`)
+`))
 	}
-	dump = s
 }
 
 func BenchmarkMaskAWSSensitiveValues(b *testing.B) {

logging/mask.go

@@ -3,18 +3,13 @@
 
 package logging
 
-import (
-	"strings"
-)
-
-func partialMaskString(s string, first, last int) string {
+func partialMaskString(s []byte, first, last int) []byte {
 	l := len(s)
-	result := strings.Builder{}
-	result.Grow(l)
-	result.WriteString(s[0:first])
+	result := make([]byte, 0, l)
+	result = append(result, s[0:first]...)
 	for i := 0; i < l-first-last; i++ {
-		result.WriteByte('*')
+		result = append(result, '*')
 	}
-	result.WriteString(s[l-last:])
-	return result.String()
+	result = append(result, s[l-last:]...)
+	return result
 }