Golang AES CBC PKCS7 io.Reader

aes_reader.go

package hls

import (
	"crypto/aes"
	"crypto/cipher"
	"io"

	"github.com/Sonelli/gojuice/crypto/pkcs7"
)

type aesReader struct {
	src   io.Reader
	cbc   cipher.BlockMode
	bsize int
	peek  []byte
	buf   []byte
}

func NewAesReader(src io.Reader, key []byte, iv []byte) (*aesReader, error) {
	cipherBlock, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}
	cbc := cipher.NewCBCDecrypter(cipherBlock, iv)
	return &aesReader{
		src:   src,
		cbc:   cbc,
		bsize: cbc.BlockSize(),
		peek:  make([]byte, cbc.BlockSize()),
	}, nil
}

func (a *aesReader) Read(out []byte) (int, error) {
	if a.buf == nil {
		a.buf = make([]byte, len(out))
	}
	decoded := 0
	copy(a.buf, a.peek)
	n, err := a.src.Read(a.buf[len(a.peek):])
	if len(a.peek) > 0 {
		n += len(a.peek)
		a.peek = a.peek[:0]
	}
	if n > 0 {
		o := 0
		for o = 0; o < n; o += a.bsize {
			if o+a.bsize > n {
				break
			}
			a.cbc.CryptBlocks(out[decoded:], a.buf[o:o+a.bsize])
			decoded += a.bsize
		}
		a.peek = a.peek[:n-o]
		copy(a.peek, a.buf[o:n])
		if len(a.peek) == 0 {
			a.peek = a.peek[0:1]
			_, err = a.src.Read(a.peek)
			if err == io.EOF {
				a.peek = a.peek[:0]
				depad, err := pkcs7.UnpadPKCS7(out[decoded-a.bsize:decoded], a.bsize)
				if err != nil {
					return decoded, err
				}
				padding := a.bsize - len(depad)
				decoded -= padding
			}
		}
	}
	return decoded, err
}