1 files changed, 57 insertions, 12 deletions

diff --git a/src/pkg/crypto/rc4/rc4_test.go b/src/pkg/crypto/rc4/rc4_test.go
index 9e12789f7..7b4df6791 100644
--- a/src/pkg/crypto/rc4/rc4_test.go
+++ b/src/pkg/crypto/rc4/rc4_test.go
@@ -5,6 +5,8 @@
 package rc4
 
 import (
+	"bytes"
+	"fmt"
 	"testing"
 )
 
@@ -72,25 +74,68 @@ var golden = []rc4Test{
 	},
 }
 
+func testEncrypt(t *testing.T, desc string, c *Cipher, src, expect []byte) {
+	dst := make([]byte, len(src))
+	c.XORKeyStream(dst, src)
+	for i, v := range dst {
+		if v != expect[i] {
+			t.Fatalf("%s: mismatch at byte %d:\nhave %x\nwant %x", desc, i, dst, expect)
+		}
+	}
+}
+
 func TestGolden(t *testing.T) {
-	for i := 0; i < len(golden); i++ {
-		g := golden[i]
-		c, err := NewCipher(g.key)
-		if err != nil {
-			t.Errorf("Failed to create cipher at golden index %d", i)
-			return
+	for gi, g := range golden {
+		data := make([]byte, len(g.keystream))
+		for i := range data {
+			data[i] = byte(i)
 		}
-		keystream := make([]byte, len(g.keystream))
-		c.XORKeyStream(keystream, keystream)
-		for j, v := range keystream {
-			if g.keystream[j] != v {
-				t.Errorf("Failed at golden index %d:\n%x\nvs\n%x", i, keystream, g.keystream)
-				break
+
+		expect := make([]byte, len(g.keystream))
+		for i := range expect {
+			expect[i] = byte(i) ^ g.keystream[i]
+		}
+
+		for size := 1; size <= len(g.keystream); size++ {
+			c, err := NewCipher(g.key)
+			if err != nil {
+				t.Fatalf("#%d: NewCipher: %v", gi, err)
+			}
+
+			off := 0
+			for off < len(g.keystream) {
+				n := len(g.keystream) - off
+				if n > size {
+					n = size
+				}
+				desc := fmt.Sprintf("#%d@[%d:%d]", gi, off, off+n)
+				testEncrypt(t, desc, c, data[off:off+n], expect[off:off+n])
+				off += n
 			}
 		}
 	}
 }
 
+func TestBlock(t *testing.T) {
+	c1a, _ := NewCipher(golden[0].key)
+	c1b, _ := NewCipher(golden[1].key)
+	data1 := make([]byte, 1<<20)
+	for i := range data1 {
+		c1a.XORKeyStream(data1[i:i+1], data1[i:i+1])
+		c1b.XORKeyStream(data1[i:i+1], data1[i:i+1])
+	}
+
+	c2a, _ := NewCipher(golden[0].key)
+	c2b, _ := NewCipher(golden[1].key)
+	data2 := make([]byte, 1<<20)
+	c2a.XORKeyStream(data2, data2)
+	c2b.XORKeyStream(data2, data2)
+
+	if !bytes.Equal(data1, data2) {
+		t.Fatalf("bad block")
+	}
+}
+
 func benchmark(b *testing.B, size int64) {
 	buf := make([]byte, size)
 	c, err := NewCipher(golden[0].key)