added support for W3C padding (in addition to PKCS#7) in the PaddedReader object, based on DRM-Inside's #111 Pull Request (updated unit tests too)

Author: danielweck

crypto/aes_cbc.go

@@ -40,6 +40,7 @@ const (
 )
 
 func (e cbcEncrypter) Signature() string {
+	// W3C padding scheme, not PKCS#7 (see last parameter "insertPadLengthAll" [false] of PaddedReader constructor)
 	return "http://www.w3.org/2001/04/xmlenc#aes256-cbc"
 }
 
@@ -49,7 +50,8 @@ func (e cbcEncrypter) GenerateKey() (ContentKey, error) {
 }
 
 func (e cbcEncrypter) Encrypt(key ContentKey, r io.Reader, w io.Writer) error {
-	r = PaddedReader(r, aes.BlockSize) // PKCS#7
+
+	r = PaddedReader(r, aes.BlockSize, false)
 
 	block, err := aes.NewCipher(key)
 	if err != nil {
@@ -99,7 +101,7 @@ func (c cbcEncrypter) Decrypt(key ContentKey, r io.Reader, w io.Writer) error {
 	mode := cipher.NewCBCDecrypter(block, iv)
 	mode.CryptBlocks(buf[aes.BlockSize:], buf[aes.BlockSize:])
 
-	padding := buf[len(buf)-1] // PKCS#7
+	padding := buf[len(buf)-1] // padding length valid for both PKCS#7 and W3C schemes
 	w.Write(buf[aes.BlockSize : len(buf)-int(padding)])
 
 	return nil

crypto/pad.go

@@ -27,6 +27,8 @@ package crypto
 
 import (
 	"io"
+	"math/rand"
+	"time"
 )
 
 type paddedReader struct {
@@ -35,6 +37,7 @@ type paddedReader struct {
 	count byte
 	left  byte
 	done  bool
+	insertPadLengthAll bool
 }
 
 func (r *paddedReader) Read(buf []byte) (int, error) {
@@ -73,8 +76,21 @@ func (r *paddedReader) Read(buf []byte) (int, error) {
 
 func (r *paddedReader) pad(buf []byte) (i int, err error) {
 	capacity := cap(buf)
+
+	src := rand.New(rand.NewSource(time.Now().UnixNano()))
+
 	for i = 0; capacity > 0 && r.left > 0; i++ {
-		buf[i] = r.count
+
+		if (r.insertPadLengthAll) {
+			buf[i] = r.count
+		} else {
+			if capacity == 1 && r.left == 1 {
+				buf[i] = r.count
+			} else {
+				buf[i] = byte(src.Intn(254) + 1)
+			}
+		}
+
 		capacity--
 		r.left--
 	}
@@ -86,6 +102,9 @@ func (r *paddedReader) pad(buf []byte) (i int, err error) {
 	return
 }
 
-func PaddedReader(r io.Reader, blockSize byte) io.Reader {
-	return &paddedReader{Reader: r, size: blockSize, count: 0, left: 0, done: false}
+
+// insertPadLengthAll = true means PKCS#7 (padding length inserted in each padding slot),
+// otherwise false means padding length inserted only in the last slot (the rest is random bytes)
+func PaddedReader(r io.Reader, blockSize byte, insertPadLengthAll bool) io.Reader {
+	return &paddedReader{Reader: r, size: blockSize, count: 0, left: 0, done: false, insertPadLengthAll: insertPadLengthAll}
 }

crypto/pad_test.go

@@ -33,7 +33,7 @@ import (
 
 func TestOneBlock(t *testing.T) {
 	buf := bytes.NewBufferString("4321")
-	reader := PaddedReader(buf, 6)
+	reader := PaddedReader(buf, 6, true)
 	var out [12]byte
 	n, err := reader.Read(out[:])
 	if err != nil && err != io.EOF {
@@ -50,7 +50,7 @@ func TestOneBlock(t *testing.T) {
 
 func TestFullPadding(t *testing.T) {
 	buf := bytes.NewBufferString("1234")
-	reader := PaddedReader(buf, 4)
+	reader := PaddedReader(buf, 4, true)
 
 	var out [8]byte
 	n, err := io.ReadFull(reader, out[:])
@@ -68,7 +68,7 @@ func TestFullPadding(t *testing.T) {
 
 func TestManyBlocks(t *testing.T) {
 	buf := bytes.NewBufferString("1234")
-	reader := PaddedReader(buf, 3)
+	reader := PaddedReader(buf, 3, true)
 	var out [3]byte
 	n, err := io.ReadFull(reader, out[:])
 	if err != nil {
@@ -91,7 +91,7 @@ func TestManyBlocks(t *testing.T) {
 
 func TestPaddingInMultipleCalls(t *testing.T) {
 	buf := bytes.NewBufferString("1")
-	reader := PaddedReader(buf, 6)
+	reader := PaddedReader(buf, 6, false)
 	var out [3]byte
 	n, err := io.ReadFull(reader, out[:])
 
@@ -122,7 +122,7 @@ func (r failingReader) Read(buf []byte) (int, error) {
 }
 
 func TestFailingReader(t *testing.T) {
-	reader := PaddedReader(failingReader{}, 8)
+	reader := PaddedReader(failingReader{}, 8, false)
 	var out [8]byte
 	_, err := io.ReadFull(reader, out[:])