Fix shadow

shadow/shadow.go

@@ -12,22 +12,26 @@ import (
 )
 
 type conn struct {
+	blockSize    int
 	aead         cipher.AEAD
 	encryptNonce func() []byte
 	decryptNonce func() []byte
-	readBuf      []byte
 	writeBuf     []byte
 	net.Conn
 }
 
 func (c *conn) Read(b []byte) (n int, err error) {
+	bLength := len(b)
+	if bLength%c.blockSize != 0 {
+		return 0, errors.Errorf("aead: block size %d not match", c.blockSize)
+	}
+
 	n, err = c.Conn.Read(b)
 	if err != nil {
 		return n, err
 	}
-
 	_, err = c.aead.Open(b[:0], c.decryptNonce(), b[:n], nil)
-	return n - c.aead.Overhead(), err
+	return bLength - c.aead.Overhead(), err
 }
 func (c *conn) Write(b []byte) (n int, err error) {
 	c.writeBuf = c.aead.Seal(nil, c.encryptNonce(), b, nil)
@@ -41,20 +45,6 @@ func (c *conn) Write(b []byte) (n int, err error) {
 }
 
 func Shadow(c net.Conn, password string) (net.Conn, error) {
-	aead, err := newAEAD(password)
-	if err != nil {
-		return nil, err
-	}
-
-	return &conn{
-		aead:         aead,
-		encryptNonce: newNonce(password, aead.NonceSize()),
-		decryptNonce: newNonce(password, aead.NonceSize()),
-		Conn:         c,
-	}, nil
-}
-
-func newAEAD(password string) (cipher.AEAD, error) {
 	block, err := aes.NewCipher([]byte(password + password)[:16])
 	if err != nil {
 		return nil, errors.Wrap(err, "password too short")
@@ -65,7 +55,13 @@ func newAEAD(password string) (cipher.AEAD, error) {
 		return nil, errors.Wrap(err, "GCM")
 	}
 
-	return aead, nil
+	return &conn{
+		blockSize:    block.BlockSize() * 8,
+		aead:         aead,
+		encryptNonce: newNonce(password, aead.NonceSize()),
+		decryptNonce: newNonce(password, aead.NonceSize()),
+		Conn:         c,
+	}, nil
 }
 
 func newNonce(password string, size int) func() []byte {