stream_test.go (1940B)
1 // Copyright 2019 Google LLC 2 // 3 // Use of this source code is governed by a BSD-style 4 // license that can be found in the LICENSE file or at 5 // https://developers.google.com/open-source/licenses/bsd 6 7 package stream_test 8 9 import ( 10 "bytes" 11 "crypto/rand" 12 "fmt" 13 "testing" 14 15 "github.com/FiloSottile/age/internal/stream" 16 "golang.org/x/crypto/chacha20poly1305" 17 ) 18 19 const cs = stream.ChunkSize 20 21 func TestRoundTrip(t *testing.T) { 22 for _, stepSize := range []int{512, 600, 1000, cs} { 23 for _, length := range []int{0, 1000, cs, cs + 100} { 24 t.Run(fmt.Sprintf("len=%d,step=%d", length, stepSize), 25 func(t *testing.T) { testRoundTrip(t, stepSize, length) }) 26 } 27 } 28 } 29 30 func testRoundTrip(t *testing.T, stepSize, length int) { 31 src := make([]byte, length) 32 if _, err := rand.Read(src); err != nil { 33 t.Fatal(err) 34 } 35 buf := &bytes.Buffer{} 36 key := make([]byte, chacha20poly1305.KeySize) 37 if _, err := rand.Read(key); err != nil { 38 t.Fatal(err) 39 } 40 41 w, err := stream.NewWriter(key, buf) 42 if err != nil { 43 t.Fatal(err) 44 } 45 46 var n int 47 for n < length { 48 b := length - n 49 if b > stepSize { 50 b = stepSize 51 } 52 nn, err := w.Write(src[n : n+b]) 53 if err != nil { 54 t.Fatal(err) 55 } 56 if nn != b { 57 t.Errorf("Write returned %d, expected %d", nn, b) 58 } 59 n += nn 60 61 nn, err = w.Write(src[n:n]) 62 if err != nil { 63 t.Fatal(err) 64 } 65 if nn != 0 { 66 t.Errorf("Write returned %d, expected 0", nn) 67 } 68 } 69 70 if err := w.Close(); err != nil { 71 t.Error("Close returned an error:", err) 72 } 73 74 t.Logf("buffer size: %d", buf.Len()) 75 76 r, err := stream.NewReader(key, buf) 77 if err != nil { 78 t.Fatal(err) 79 } 80 81 n = 0 82 readBuf := make([]byte, stepSize) 83 for n < length { 84 b := length - n 85 if b > stepSize { 86 b = stepSize 87 } 88 nn, err := r.Read(readBuf) 89 if err != nil { 90 t.Fatalf("Read error at index %d: %v", n, err) 91 } 92 93 if !bytes.Equal(readBuf[:nn], src[n:n+nn]) { 94 t.Errorf("wrong data at indexes %d - %d", n, n+nn) 95 } 96 97 n += nn 98 } 99 }