From 336730b91974133c6959b2fc46f54d2e43ecab70 Mon Sep 17 00:00:00 2001
From: 9seconds <nineseconds@yandex.ru>
Date: Tue, 5 Jun 2018 11:20:09 +0300
Subject: [PATCH] Add rwc for block cipher mode

---
 wrappers/blockcipherrwc.go  | 73 +++++++++++++++++++++++++++++++++++++
 wrappers/streamcipherrwc.go |  3 --
 2 files changed, 73 insertions(+), 3 deletions(-)
 create mode 100644 wrappers/blockcipherrwc.go

diff --git a/wrappers/blockcipherrwc.go b/wrappers/blockcipherrwc.go
new file mode 100644
index 0000000..fa72ecb
--- /dev/null
+++ b/wrappers/blockcipherrwc.go
@@ -0,0 +1,73 @@
+package wrappers
+
+import (
+	"bytes"
+	"crypto/cipher"
+	"io"
+
+	"github.com/juju/errors"
+)
+
+type BlockCipherReadWriteCloser struct {
+	encryptor cipher.BlockMode
+	decryptor cipher.BlockMode
+	conn      io.ReadWriteCloser
+	buf       *bytes.Buffer
+}
+
+func (c *BlockCipherReadWriteCloser) Read(p []byte) (n int, err error) {
+	blockSize := c.decryptor.BlockSize()
+	if len(p) < blockSize {
+		return 0, errors.New("Cannot read less than blocksize")
+	}
+
+	n, err = c.conn.Read(p)
+	c.buf.Write(p[:n])
+
+	wantToRead := c.getFullBlocks(len(p), blockSize)
+	haveBlocks := c.getFullBlocks(c.buf.Len(), blockSize)
+	if haveBlocks < wantToRead {
+		wantToRead = haveBlocks
+	}
+	wantToRead *= blockSize
+
+	data := c.buf.Bytes()
+	c.decryptor.CryptBlocks(p, data[:wantToRead])
+	c.buf = bytes.NewBuffer(data[wantToRead:])
+
+	return wantToRead, err
+}
+
+func (c *BlockCipherReadWriteCloser) Write(p []byte) (n int, err error) {
+	blockSize := c.encryptor.BlockSize()
+	if len(p)%blockSize != 0 {
+		return 0, errors.New("Write size should be compatible with block size")
+	}
+
+	buf := make([]byte, len(p))
+	c.encryptor.CryptBlocks(buf, p)
+
+	return c.conn.Write(buf)
+}
+
+func (c *BlockCipherReadWriteCloser) Close() error {
+	return c.conn.Close()
+}
+
+func (c *BlockCipherReadWriteCloser) getFullBlocks(number, blockSize int) int {
+	blocks := number / blockSize
+
+	if blocks > 0 && number%blockSize != 0 {
+		blocks--
+	}
+
+	return blocks
+}
+
+func NewBlockCipherRWC(conn io.ReadWriteCloser, encryptor, decryptor cipher.BlockMode) io.ReadWriteCloser {
+	return &BlockCipherReadWriteCloser{
+		conn:      conn,
+		encryptor: encryptor,
+		decryptor: decryptor,
+	}
+}
diff --git a/wrappers/streamcipherrwc.go b/wrappers/streamcipherrwc.go
index c12b7dd..5d7d018 100644
--- a/wrappers/streamcipherrwc.go
+++ b/wrappers/streamcipherrwc.go
@@ -1,7 +1,6 @@
 package wrappers
 
 import (
-	"bytes"
 	"crypto/cipher"
 	"io"
 )
@@ -10,7 +9,6 @@ type StreamCipherReadWriteCloser struct {
 	encryptor cipher.Stream
 	decryptor cipher.Stream
 	conn      io.ReadWriteCloser
-	rest      *bytes.Buffer
 }
 
 // Read reads from connection
@@ -48,6 +46,5 @@ func NewStreamCipherRWC(conn io.ReadWriteCloser, encryptor, decryptor cipher.Str
 		conn:      conn,
 		encryptor: encryptor,
 		decryptor: decryptor,
-		rest:      &bytes.Buffer{},
 	}
 }