go/streamingaead/decrypt_reader_test.go - platform/external/tink - Git at Google

 // Copyright 2022 Google LLC
 //
 // 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 streamingaead

 import (
 	"bytes"
 	"crypto/rand"
 	"fmt"
 	"io"
 	"strings"
 	"testing"

 	"github.com/google/tink/go/subtle/random"
 	"github.com/google/tink/go/testkeyset"
 	"github.com/google/tink/go/testutil"
 	tinkpb "github.com/google/tink/go/proto/tink_go_proto"
 )

 func BenchmarkDecryptReader(b *testing.B) {
 	b.ReportAllocs()

 	// Create a Streaming AEAD primitive using a full keyset.
 	decKeyset := testutil.NewTestAESGCMHKDFKeyset()
 	decKeysetHandle, err := testkeyset.NewHandle(decKeyset)
 	if err != nil {
 		b.Fatalf("Failed creating keyset handle: %v", err)
 	}
 	decCipher, err := New(decKeysetHandle)
 	if err != nil {
 		b.Errorf("streamingaead.New failed: %v", err)
 	}

 	// Extract the raw key from the keyset and create a Streaming AEAD primitive
 	// using only that key.
 	//
 	// testutil.NewTestAESGCMHKDFKeyset() places a raw key at position 1.
 	rawKey := decKeyset.Key[1]
 	if rawKey.OutputPrefixType != tinkpb.OutputPrefixType_RAW {
 		b.Fatalf("Expected a raw key.")
 	}
 	encKeyset := testutil.NewKeyset(rawKey.KeyId, []*tinkpb.Keyset_Key{rawKey})
 	encKeysetHandle, err := testkeyset.NewHandle(encKeyset)
 	if err != nil {
 		b.Fatalf("Failed creating keyset handle: %v", err)
 	}
 	encCipher, err := New(encKeysetHandle)
 	if err != nil {
 		b.Fatalf("streamingaead.New failed: %v", err)
 	}

 	plaintext := random.GetRandomBytes(8)
 	associatedData := random.GetRandomBytes(32)

 	b.ResetTimer()
 	for i := 0; i < b.N; i++ {
 		// Create a pipe for communication between the encrypting writer and
 		// decrypting reader.
 		r, w := io.Pipe()
 		defer r.Close()

 		// Repeatedly encrypt the plaintext and write the ciphertext to a pipe.
 		go func() {
 			const writeAtLeast = 1 << 30 // 1 GiB

 			enc, err := encCipher.NewEncryptingWriter(w, associatedData)
 			if err != nil {
 				b.Errorf("Cannot create encrypt writer: %v", err)
 				return
 			}

 			for i := 0; i < writeAtLeast; i += len(plaintext) {
 				if _, err := enc.Write(plaintext); err != nil {
 					b.Errorf("Error encrypting data: %v", err)
 					return
 				}
 			}
 			if err := enc.Close(); err != nil {
 				b.Errorf("Error closing encrypting writer: %v", err)
 				return
 			}
 			if err := w.Close(); err != nil {
 				b.Errorf("Error closing pipe: %v", err)
 				return
 			}
 		}()

 		// Decrypt the ciphertext in small chunks.
 		dec, err := decCipher.NewDecryptingReader(r, associatedData)
 		if err != nil {
 			b.Fatalf("Cannot create decrypt reader: %v", err)
 		}
 		buf := make([]byte, 16384) // 16 KiB
 		for {
 			_, err := dec.Read(buf)
 			if err == io.EOF {
 				break
 			}
 			if err != nil {
 				b.Fatalf("Error decrypting data: %v", err)
 			}
 		}
 	}
 }

 func TestUnreaderUnread(t *testing.T) {
 	original := make([]byte, 4096)
 	if _, err := io.ReadFull(rand.Reader, original); err != nil {
 		t.Fatalf("Failed to fill buffer with random bytes: %v", err)
 	}

 	u := &unreader{r: bytes.NewReader(original)}
 	got, err := io.ReadAll(u)
 	if err != nil {
 		t.Errorf("First io.ReadAll(%T) failed unexpectedly: %v", u, err)
 	}
 	if !bytes.Equal(got, original) {
 		t.Errorf("First io.ReadAll(%T) got %d bytes, want %d bytes that match the original random data.\nGot: %X\nWant: %X", u, len(got), len(original), got, original)
 	}

 	u.unread()
 	got, err = io.ReadAll(u)
 	if err != nil {
 		t.Errorf("After %T.unread(), io.ReadAll(%T) failed unexpectedly: %v", u, u, err)
 	}
 	if !bytes.Equal(got, original) {
 		t.Errorf("After %T.unread(), io.ReadAll(%T) got %d bytes, want %d bytes that match the original random data.\nGot: %X\nWant: %X", u, u, len(got), len(original), got, original)
 	}
 }

 func TestUnreader(t *testing.T) {
 	// Repeating sequence of characters '0' through '9' makes it easy to see
 	// holes or repeated data.
 	original := make([]byte, 100)
 	for i := range original {
 		original[i] = '0' + byte(i%10)
 	}

 	type step struct {
 		read    int  // If set, read the given number of bytes exactly.
 		unread  bool // If true, call unread().
 		disable bool // If true, call disable().
 	}
 	tcs := []struct {
 		name  string
 		steps []step
 	}{
 		{"Read2UnreadRead4Unread", []step{{read: 2}, {unread: true}, {read: 4}, {unread: true}}},
 		{"Read4UnreadRead2Unread", []step{{read: 4}, {unread: true}, {read: 2}, {unread: true}}},
 		{"Read3UnreadRead3Unread", []step{{read: 3}, {unread: true}, {read: 3}, {unread: true}}},
 		{"Read3Disable", []step{{read: 3}, {disable: true}}},
 		{"Read2UnreadRead4Disable", []step{{read: 2}, {unread: true}, {read: 4}, {disable: true}}},
 		{"Read4UnreadRead2Disable", []step{{read: 4}, {unread: true}, {read: 2}, {disable: true}}},
 		{"Read3UnreadRead3Disable", []step{{read: 3}, {unread: true}, {read: 3}, {disable: true}}},
 		{"Read2UnreadDisable", []step{{read: 2}, {unread: true}, {disable: true}}},
 		{"Read4UnreadDisable", []step{{read: 4}, {unread: true}, {disable: true}}},
 		{"ReadAllUnread", []step{{read: len(original)}, {unread: true}}},
 		{"ReadAllDisable", []step{{read: len(original)}, {disable: true}}},
 		{"Unread", []step{{unread: true}}},
 		{"Disable", []step{{disable: true}}},
 		{"UnreadDisable", []step{{unread: true}, {disable: true}}},
 	}

 	for _, tc := range tcs {
 		t.Run(tc.name, func(t *testing.T) {
 			u := &unreader{r: bytes.NewReader(original)}
 			var (
 				after []string
 				pos   int
 			)
 			// Explains what happened before the failure.
 			prefix := func() string {
 				if after == nil {
 					return ""
 				}
 				return fmt.Sprintf("After %s, ", strings.Join(after, "+"))
 			}
 			for _, s := range tc.steps {
 				if s.read != 0 {
 					buf := make([]byte, s.read)
 					if _, err := io.ReadFull(u, buf); err != nil {
 						t.Fatalf("%sio.ReadFull(%T, %d byte buffer) failed unexpectedly: %v", prefix(), u, s.read, err)
 					}
 					if want := original[pos : pos+s.read]; !bytes.Equal(buf, want) {
 						t.Fatalf("%sio.ReadFull(%T, %d byte buffer) got %q, want %q", prefix(), u, s.read, buf, want)
 					}
 					after = append(after, fmt.Sprintf("Read(%d bytes)", s.read))
 					pos += s.read
 				}
 				if s.disable {
 					u.disable()
 					after = append(after, "disable()")
 				}
 				if s.unread {
 					u.unread()
 					after = append(after, "unread()")
 					pos = 0
 				}
 			}
 			got, err := io.ReadAll(u)
 			if err != nil {
 				t.Fatalf("%sio.ReadAll(%T) failed unexpectedly: %v", prefix(), u, err)
 			}
 			if want := original[pos:]; !bytes.Equal(want, got) {
 				t.Errorf("%sio.ReadAll(%T) got %q, want %q", prefix(), u, got, want)
 			}
 		})
 	}
 }
	// Copyright 2022 Google LLC
	//
	// 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 streamingaead

	import (
	"bytes"
	"crypto/rand"
	"fmt"
	"io"
	"strings"
	"testing"

	"github.com/google/tink/go/subtle/random"
	"github.com/google/tink/go/testkeyset"
	"github.com/google/tink/go/testutil"
	tinkpb "github.com/google/tink/go/proto/tink_go_proto"
	)

	func BenchmarkDecryptReader(b *testing.B) {
	b.ReportAllocs()

	// Create a Streaming AEAD primitive using a full keyset.
	decKeyset := testutil.NewTestAESGCMHKDFKeyset()
	decKeysetHandle, err := testkeyset.NewHandle(decKeyset)
	if err != nil {
	b.Fatalf("Failed creating keyset handle: %v", err)
	}
	decCipher, err := New(decKeysetHandle)
	if err != nil {
	b.Errorf("streamingaead.New failed: %v", err)
	}

	// Extract the raw key from the keyset and create a Streaming AEAD primitive
	// using only that key.
	//
	// testutil.NewTestAESGCMHKDFKeyset() places a raw key at position 1.
	rawKey := decKeyset.Key[1]
	if rawKey.OutputPrefixType != tinkpb.OutputPrefixType_RAW {
	b.Fatalf("Expected a raw key.")
	}
	encKeyset := testutil.NewKeyset(rawKey.KeyId, []*tinkpb.Keyset_Key{rawKey})
	encKeysetHandle, err := testkeyset.NewHandle(encKeyset)
	if err != nil {
	b.Fatalf("Failed creating keyset handle: %v", err)
	}
	encCipher, err := New(encKeysetHandle)
	if err != nil {
	b.Fatalf("streamingaead.New failed: %v", err)
	}

	plaintext := random.GetRandomBytes(8)
	associatedData := random.GetRandomBytes(32)

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
	// Create a pipe for communication between the encrypting writer and
	// decrypting reader.
	r, w := io.Pipe()
	defer r.Close()

	// Repeatedly encrypt the plaintext and write the ciphertext to a pipe.
	go func() {
	const writeAtLeast = 1 << 30 // 1 GiB

	enc, err := encCipher.NewEncryptingWriter(w, associatedData)
	if err != nil {
	b.Errorf("Cannot create encrypt writer: %v", err)
	return
	}

	for i := 0; i < writeAtLeast; i += len(plaintext) {
	if _, err := enc.Write(plaintext); err != nil {
	b.Errorf("Error encrypting data: %v", err)
	return
	}
	}
	if err := enc.Close(); err != nil {
	b.Errorf("Error closing encrypting writer: %v", err)
	return
	}
	if err := w.Close(); err != nil {
	b.Errorf("Error closing pipe: %v", err)
	return
	}
	}()

	// Decrypt the ciphertext in small chunks.
	dec, err := decCipher.NewDecryptingReader(r, associatedData)
	if err != nil {
	b.Fatalf("Cannot create decrypt reader: %v", err)
	}
	buf := make([]byte, 16384) // 16 KiB
	for {
	_, err := dec.Read(buf)
	if err == io.EOF {
	break
	}
	if err != nil {
	b.Fatalf("Error decrypting data: %v", err)
	}
	}
	}
	}

	func TestUnreaderUnread(t *testing.T) {
	original := make([]byte, 4096)
	if _, err := io.ReadFull(rand.Reader, original); err != nil {
	t.Fatalf("Failed to fill buffer with random bytes: %v", err)
	}

	u := &unreader{r: bytes.NewReader(original)}
	got, err := io.ReadAll(u)
	if err != nil {
	t.Errorf("First io.ReadAll(%T) failed unexpectedly: %v", u, err)
	}
	if !bytes.Equal(got, original) {
	t.Errorf("First io.ReadAll(%T) got %d bytes, want %d bytes that match the original random data.\nGot: %X\nWant: %X", u, len(got), len(original), got, original)
	}

	u.unread()
	got, err = io.ReadAll(u)
	if err != nil {
	t.Errorf("After %T.unread(), io.ReadAll(%T) failed unexpectedly: %v", u, u, err)
	}
	if !bytes.Equal(got, original) {
	t.Errorf("After %T.unread(), io.ReadAll(%T) got %d bytes, want %d bytes that match the original random data.\nGot: %X\nWant: %X", u, u, len(got), len(original), got, original)
	}
	}

	func TestUnreader(t *testing.T) {
	// Repeating sequence of characters '0' through '9' makes it easy to see
	// holes or repeated data.
	original := make([]byte, 100)
	for i := range original {
	original[i] = '0' + byte(i%10)
	}

	type step struct {
	read int // If set, read the given number of bytes exactly.
	unread bool // If true, call unread().
	disable bool // If true, call disable().
	}
	tcs := []struct {
	name string
	steps []step
	}{
	{"Read2UnreadRead4Unread", []step{{read: 2}, {unread: true}, {read: 4}, {unread: true}}},
	{"Read4UnreadRead2Unread", []step{{read: 4}, {unread: true}, {read: 2}, {unread: true}}},
	{"Read3UnreadRead3Unread", []step{{read: 3}, {unread: true}, {read: 3}, {unread: true}}},
	{"Read3Disable", []step{{read: 3}, {disable: true}}},
	{"Read2UnreadRead4Disable", []step{{read: 2}, {unread: true}, {read: 4}, {disable: true}}},
	{"Read4UnreadRead2Disable", []step{{read: 4}, {unread: true}, {read: 2}, {disable: true}}},
	{"Read3UnreadRead3Disable", []step{{read: 3}, {unread: true}, {read: 3}, {disable: true}}},
	{"Read2UnreadDisable", []step{{read: 2}, {unread: true}, {disable: true}}},
	{"Read4UnreadDisable", []step{{read: 4}, {unread: true}, {disable: true}}},
	{"ReadAllUnread", []step{{read: len(original)}, {unread: true}}},
	{"ReadAllDisable", []step{{read: len(original)}, {disable: true}}},
	{"Unread", []step{{unread: true}}},
	{"Disable", []step{{disable: true}}},
	{"UnreadDisable", []step{{unread: true}, {disable: true}}},
	}

	for _, tc := range tcs {
	t.Run(tc.name, func(t *testing.T) {
	u := &unreader{r: bytes.NewReader(original)}
	var (
	after []string
	pos int
	)
	// Explains what happened before the failure.
	prefix := func() string {
	if after == nil {
	return ""
	}
	return fmt.Sprintf("After %s, ", strings.Join(after, "+"))
	}
	for _, s := range tc.steps {
	if s.read != 0 {
	buf := make([]byte, s.read)
	if _, err := io.ReadFull(u, buf); err != nil {
	t.Fatalf("%sio.ReadFull(%T, %d byte buffer) failed unexpectedly: %v", prefix(), u, s.read, err)
	}
	if want := original[pos : pos+s.read]; !bytes.Equal(buf, want) {
	t.Fatalf("%sio.ReadFull(%T, %d byte buffer) got %q, want %q", prefix(), u, s.read, buf, want)
	}
	after = append(after, fmt.Sprintf("Read(%d bytes)", s.read))
	pos += s.read
	}
	if s.disable {
	u.disable()
	after = append(after, "disable()")
	}
	if s.unread {
	u.unread()
	after = append(after, "unread()")
	pos = 0
	}
	}
	got, err := io.ReadAll(u)
	if err != nil {
	t.Fatalf("%sio.ReadAll(%T) failed unexpectedly: %v", prefix(), u, err)
	}
	if want := original[pos:]; !bytes.Equal(want, got) {
	t.Errorf("%sio.ReadAll(%T) got %q, want %q", prefix(), u, got, want)
	}
	})
	}
	}