From 63f4bd7ecc06b05e55b6020fa4237c314db804b2 Mon Sep 17 00:00:00 2001
From: Arjun Mahishi <arjun.mahishi@gmail.com>
Date: Wed, 29 Jan 2025 00:12:19 +0530
Subject: [PATCH] Add support for SafeByte and SafeBytes

Currently, there is no direct way to write a safe byte or a byte slice
to a SafeWriter. They have to be type-casted to either a rune (for byte)
or a string (for []byte). It would be more efficient to add
the types (and helpers) SafeByte and SafeBytes.

Also, UnsafeByte and UnsafeBytes already exist. So, introducing SafeByte and
SafeBytes would also make the API more consistent.
---
 builder/builder.go               | 10 ++++++++++
 builder/builder_test.go          |  6 +++---
 interfaces/interfaces.go         | 18 ++++++++++++++++++
 internal/rfmt/printer_adapter.go | 12 ++++++++++++
 4 files changed, 43 insertions(+), 3 deletions(-)

diff --git a/builder/builder.go b/builder/builder.go
index 9c0ae29..921deb8 100644
--- a/builder/builder.go
+++ b/builder/builder.go
@@ -130,6 +130,16 @@ func (b *StringBuilder) SafeRune(s i.SafeRune) {
 	_ = b.Buffer.WriteRune(rune(s))
 }
 
+func (b *StringBuilder) SafeByte(s i.SafeByte) {
+	b.SetMode(ib.SafeEscaped)
+	_ = b.Buffer.WriteByte(byte(s))
+}
+
+func (b *StringBuilder) SafeBytes(s i.SafeBytes) {
+	b.SetMode(ib.SafeEscaped)
+	_, _ = b.Buffer.Write([]byte(s))
+}
+
 // UnsafeString is part of the SafeWriter interface.
 func (b *StringBuilder) UnsafeString(s string) {
 	b.SetMode(ib.UnsafeEscaped)
diff --git a/builder/builder_test.go b/builder/builder_test.go
index 4620d07..38bb6de 100644
--- a/builder/builder_test.go
+++ b/builder/builder_test.go
@@ -59,13 +59,13 @@ func TestBuilder(t *testing.T) {
 	b.SafeRune('\n')
 
 	b.UnsafeByte('U')
-	b.SafeRune('\n')
+	b.SafeByte('\n')
 
 	b.UnsafeByte(m.StartS[0])
-	b.SafeRune('\n')
+	b.SafeByte('\n')
 
 	b.UnsafeBytes([]byte("UUU"))
-	b.SafeRune('\n')
+	b.SafeBytes([]byte("\n"))
 
 	actualR := b.RedactableString()
 	const expectedR = `‹unsafe›
diff --git a/interfaces/interfaces.go b/interfaces/interfaces.go
index 20fec2e..bd13a94 100644
--- a/interfaces/interfaces.go
+++ b/interfaces/interfaces.go
@@ -65,6 +65,12 @@ type SafeWriter interface {
 	// SafeRune emits a safe rune.
 	SafeRune(SafeRune)
 
+	// SafeByte emits a safe byte.
+	SafeByte(SafeByte)
+
+	// SafeBytes emits a safe byte slice.
+	SafeBytes(SafeBytes)
+
 	// Print emits its arguments separated by spaces.
 	// For each argument it dynamically checks for the SafeFormatter or
 	// SafeValue interface and either use that, or mark the argument
@@ -119,6 +125,18 @@ type SafeRune rune
 // SafeValue makes SafeRune a SafeValue.
 func (SafeRune) SafeValue() {}
 
+// SafeByte represents a byte that is not a sensitive value.
+type SafeByte byte
+
+// SafeValue makes SafeByte a SafeValue.
+func (SafeByte) SafeValue() {}
+
+// SafeBytes represents a byte slice that is not a sensitive value.
+type SafeBytes []byte
+
+// SafeValue makes SafeBytes a SafeValue.
+func (SafeBytes) SafeValue() {}
+
 // SafeValue is a marker interface to be implemented by types that
 // alias base Go types and whose natural representation via Printf is
 // always safe for reporting.
diff --git a/internal/rfmt/printer_adapter.go b/internal/rfmt/printer_adapter.go
index 923d972..5b57e66 100644
--- a/internal/rfmt/printer_adapter.go
+++ b/internal/rfmt/printer_adapter.go
@@ -46,6 +46,18 @@ func (p *pp) SafeRune(r i.SafeRune) {
 	p.buf.WriteRune(rune(r))
 }
 
+// SafeByte implements SafePrinter.
+func (p *pp) SafeByte(r i.SafeByte) {
+	defer p.startSafeOverride().restore()
+	p.buf.WriteByte(byte(r))
+}
+
+// SafeBytes implements SafePrinter.
+func (p *pp) SafeBytes(r i.SafeBytes) {
+	defer p.startSafeOverride().restore()
+	p.buf.Write(r)
+}
+
 func (p *pp) Print(args ...interface{}) {
 	defer p.buf.SetMode(p.buf.GetMode())
 	np := newPrinter()