Skip to content

Commit

Permalink
feat(crypto): Allow to trim input lines when writing messages
Browse files Browse the repository at this point in the history
  • Loading branch information
lubux committed Jun 14, 2024
1 parent 3484816 commit b224485
Show file tree
Hide file tree
Showing 8 changed files with 215 additions and 15 deletions.
50 changes: 49 additions & 1 deletion crypto/encrypt_decrypt_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -924,6 +924,32 @@ func TestEncryptDecryptKey(t *testing.T) {
}
}

func TestEncryptDecryptStreamTrimmedLines(t *testing.T) {
for _, material := range testMaterialForProfiles {
t.Run(material.profileName, func(t *testing.T) {
encHandle, _ := material.pgp.Encryption().
Recipients(material.keyRingTestPublic).
SigningKeys(material.keyRingTestPrivate).
TrimLines().
New()
decHandle, _ := material.pgp.Decryption().
DecryptionKeys(material.keyRingTestPrivate).
VerificationKeys(material.keyRingTestPublic).
New()
testEncryptDecryptStreamWithExpected(
t,
[]byte("text with \r \t \n trimmed\n \t"),
[]byte("text with\n trimmed\n"),
nil,
encHandle,
decHandle,
len(material.keyRingTestPrivate.entities),
Bytes,
)
})
}
}

func TestEncryptCompressionApplied(t *testing.T) {
const numReplicas = 10
builder := strings.Builder{}
Expand Down Expand Up @@ -1169,6 +1195,28 @@ func testEncryptDecryptStream(
decHandle PGPDecryption,
numberOfSigsToVerify int,
encoding int8,
) {
testEncryptDecryptStreamWithExpected(
t,
messageBytes,
messageBytes,
metadata,
encHandle,
decHandle,
numberOfSigsToVerify,
encoding,
)
}

func testEncryptDecryptStreamWithExpected(
t *testing.T,
messageBytes []byte,
expected []byte,
metadata *LiteralMetadata,
encHandle PGPEncryption,
decHandle PGPDecryption,
numberOfSigsToVerify int,
encoding int8,
) {
messageReader := bytes.NewReader(messageBytes)
var ciphertextBuf bytes.Buffer
Expand Down Expand Up @@ -1200,7 +1248,7 @@ func testEncryptDecryptStream(
if err != nil {
t.Fatal("Expected no error while reading the decrypted data, got:", err)
}
if !bytes.Equal(decryptedBytes, messageBytes) {
if !bytes.Equal(decryptedBytes, expected) {
t.Fatalf("Expected the decrypted data to be %s got %s", string(decryptedBytes), string(messageBytes))
}
if numberOfSigsToVerify > 0 {
Expand Down
6 changes: 6 additions & 0 deletions crypto/encryption_handle.go
Original file line number Diff line number Diff line change
Expand Up @@ -55,6 +55,9 @@ type encryptionHandle struct {
// Is only considered if DetachedSignature is not set.
PlainDetachedSignature bool
IsUTF8 bool
// TrimLines trims each end of the line in the input message before encryption.
// Remove trailing spaces, carriage returns and tabs from each line (separated by \n characters).
TrimLines bool
// ExternalSignature allows to include an external signature into
// the encrypted message.
ExternalSignature []byte
Expand Down Expand Up @@ -332,6 +335,9 @@ func (eh *encryptionHandle) encryptingWriters(keys, data, detachedSignature Writ
openpgp.NewCanonicalTextWriteCloser(messageWriter),
)
}
if eh.TrimLines {
messageWriter = internal.NewTrimWriteCloser(messageWriter)
}
return messageWriter, nil
}

Expand Down
7 changes: 7 additions & 0 deletions crypto/encryption_handle_builder.go
Original file line number Diff line number Diff line change
Expand Up @@ -148,6 +148,13 @@ func (ehb *EncryptionHandleBuilder) Utf8() *EncryptionHandleBuilder {
return ehb
}

// TrimLines enables that each line in the input message is trimmed before encryption.
// Trim removes trailing spaces, carriage returns and tabs from each line (separated by \n characters).
func (ehb *EncryptionHandleBuilder) TrimLines() *EncryptionHandleBuilder {
ehb.handle.TrimLines = true
return ehb
}

// DetachedSignature indicates that the message should be signed,
// but the signature should not be included in the same pgp message as the input data.
// Instead the detached signature is encrypted in a separate pgp message.
Expand Down
14 changes: 10 additions & 4 deletions crypto/sign_handle.go
Original file line number Diff line number Diff line change
Expand Up @@ -16,10 +16,13 @@ import (
)

type signatureHandle struct {
SignKeyRing *KeyRing
SignContext *SigningContext
IsUTF8 bool
Detached bool
SignKeyRing *KeyRing
SignContext *SigningContext
IsUTF8 bool
Detached bool
// TrimLines trims each end of the line in the input message before encryption.
// Remove trailing spaces, carriage returns and tabs from each line (separated by \n characters).
TrimLines bool
ArmorHeaders map[string]string
profile SignProfile
clock Clock
Expand Down Expand Up @@ -87,6 +90,9 @@ func (sh *signatureHandle) SigningWriter(outputWriter Writer, encoding int8) (me
openpgp.NewCanonicalTextWriteCloser(messageWriter),
)
}
if sh.TrimLines {
messageWriter = internal.NewTrimWriteCloser(messageWriter)
}
return messageWriter, nil
}

Expand Down
7 changes: 7 additions & 0 deletions crypto/sign_handle_builder.go
Original file line number Diff line number Diff line change
Expand Up @@ -68,6 +68,13 @@ func (shb *SignHandleBuilder) Utf8() *SignHandleBuilder {
return shb
}

// TrimLines enables that each line in the input message is trimmed before encryption.
// Trim removes trailing spaces, carriage returns and tabs from each line (separated by \n characters).
func (shb *SignHandleBuilder) TrimLines() *SignHandleBuilder {
shb.handle.TrimLines = true
return shb
}

// SignTime sets the internal clock to always return
// the supplied unix time for signing instead of the device time.
func (shb *SignHandleBuilder) SignTime(unixTime int64) *SignHandleBuilder {
Expand Down
89 changes: 89 additions & 0 deletions internal/trim_lines_writer.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,89 @@
package internal

import (
"bytes"
"io"
)

func trim(p []byte) []byte {
return bytes.TrimRight(p, " \t\r")
}

func NewTrimWriteCloser(internal io.WriteCloser) *TrimWriteCloser {
return NewTrimWriteCloserWithBufferSize(internal, 256)
}

func NewTrimWriteCloserWithBufferSize(internal io.WriteCloser, size int) *TrimWriteCloser {
return &TrimWriteCloser{
internal: internal,
whitespace: bytes.NewBuffer(make([]byte, 0, size)),
err: nil,
}
}

type TrimWriteCloser struct {
internal io.WriteCloser
whitespace *bytes.Buffer
err error
}

func (w *TrimWriteCloser) Write(p []byte) (n int, err error) {
n = len(p)
if w.err != nil {
return 0, err
}
for index := bytes.IndexByte(p, '\n'); index != -1; index = bytes.IndexByte(p, '\n') {
trimmedSuffixLine := trim(p[:index])
bufferWhitespace := w.whitespace.Bytes()
if len(bufferWhitespace) > 0 {
if len(trimmedSuffixLine) != 0 {
if _, err = w.internal.Write(bufferWhitespace); err != nil {
w.err = err
return 0, err
}
}
w.whitespace.Reset()
}
if len(trimmedSuffixLine) < len(p[:index]) {
if _, err = w.internal.Write(trimmedSuffixLine); err != nil {
w.err = err
return index, err
}
if _, err = w.internal.Write([]byte("\n")); err != nil {
w.err = err
return index + 1, err
}
} else {
if _, err = w.internal.Write(p[:index+1]); err != nil {
w.err = err
return index + 1, err
}
}
p = p[index+1:]
}

if len(p) > 0 {
nonWhitespace := trim(p)
if len(nonWhitespace) > 0 && w.whitespace.Len() > 0 {
if _, err = w.internal.Write(w.whitespace.Bytes()); err != nil {
w.err = err
return n - len(p), err
}
w.whitespace.Reset()
}
if _, err = w.internal.Write(nonWhitespace); err != nil {
w.err = err
return n - len(p), err
}

if _, err = w.whitespace.Write(p[len(nonWhitespace):]); err != nil {
w.err = err
return n - len(p), err
}
}
return n, nil
}

func (w *TrimWriteCloser) Close() error {
return w.internal.Close()
}
41 changes: 41 additions & 0 deletions internal/trim_lines_writer_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,41 @@
package internal

import (
"bytes"
"strings"
"testing"

"github.com/stretchr/testify/assert"
)

func testTrimWriteCloser(t *testing.T, test string) {
testBytes := []byte(test)
for _, batchSize := range []int{1, 2, 3, 7, 11} {
var buff bytes.Buffer
w := &noOpWriteCloser{
writer: &buff,
}
trimWriter := NewTrimWriteCloser(w)
for ind := 0; ind < len(testBytes); ind += batchSize {
end := ind + batchSize
if end > len(testBytes) {
end = len(testBytes)
}
if _, err := trimWriter.Write(testBytes[ind:end]); err != nil {
t.Fatal(err)
}
}
if err := trimWriter.Close(); err != nil {
t.Fatal(err)
}
assert.Equal(t, TrimEachLine(test), buff.String())
}
}

func TestTrimWriteCloser(t *testing.T) {
testTrimWriteCloser(t, "\n \t \r")
testTrimWriteCloser(t, "this is a test \n \t \n\n")
testTrimWriteCloser(t, "saf\n \t sdf\n \r\rsd \t fsdf\n")
testTrimWriteCloser(t, "BEGIN:VCARD\r\nVERSION:4.0\r\nFN;PREF=1: \r\nEND:VCARD")
testTrimWriteCloser(t, strings.Repeat("\r \nthis is a test \n \t \n\n)", 10000))
}
16 changes: 6 additions & 10 deletions internal/utf8.go
Original file line number Diff line number Diff line change
Expand Up @@ -164,19 +164,15 @@ func NewUtf8CheckWriteCloser(wrap io.WriteCloser) *Utf8CheckWriteCloser {
}

func (cw *Utf8CheckWriteCloser) Write(p []byte) (n int, err error) {
err = cw.check(p)
if err != nil {
return
if err = cw.check(p); err != nil {
return 0, err
}
n, err = cw.internal.Write(p)
return
return cw.internal.Write(p)
}

func (cw *Utf8CheckWriteCloser) Close() (err error) {
err = cw.close()
if err != nil {
return
if err = cw.close(); err != nil {
return err
}
err = cw.internal.Close()
return
return cw.internal.Close()
}

0 comments on commit b224485

Please sign in to comment.