Refactor cipherrwc to wrappers

proxy/cipherrwc.go

@@ -1,56 +0,0 @@
-package proxy
-
-import (
-	"bytes"
-	"io"
-)
-
-// Cipher is an interface to anything which can encrypt and decrypt
-type Cipher interface {
-	Encrypt([]byte) []byte
-	Decrypt([]byte) []byte
-}
-
-// CipherReadWriteCloser wraps connection for transparent encryption
-type CipherReadWriteCloser struct {
-	crypt Cipher
-	conn  io.ReadWriteCloser
-	rest  *bytes.Buffer
-}
-
-// Read reads from connection
-func (c *CipherReadWriteCloser) Read(p []byte) (n int, err error) {
-	n, err = c.conn.Read(p)
-	copy(p, c.crypt.Decrypt(p[:n]))
-	return
-}
-
-// Write writes into connection.
-func (c *CipherReadWriteCloser) Write(p []byte) (int, error) {
-	encrypted := c.crypt.Encrypt(p)
-	allWritten := 0
-
-	for len(encrypted) > 0 {
-		n, err := c.conn.Write(encrypted)
-		allWritten += n
-		if err != nil {
-			return allWritten, err
-		}
-		encrypted = encrypted[n:]
-	}
-
-	return allWritten, nil
-}
-
-// Close closes underlying connection.
-func (c *CipherReadWriteCloser) Close() error {
-	return c.conn.Close()
-}
-
-func newCipherReadWriteCloser(conn io.ReadWriteCloser, crypt Cipher) *CipherReadWriteCloser {
-	return &CipherReadWriteCloser{
-		conn:  conn,
-		crypt: crypt,
-		rest:  &bytes.Buffer{},
-	}
-}

proxy/server.go

@@ -9,6 +9,7 @@ import (
 	"time"
 
 	"github.com/9seconds/mtg/obfuscated2"
+	"github.com/9seconds/mtg/wrappers"
 	"github.com/juju/errors"
 	uuid "github.com/satori/go.uuid"
 	"go.uber.org/zap"
@@ -124,7 +125,7 @@ func (s *Server) getClientStream(ctx context.Context, cancel context.CancelFunc,
 	}
 
 	wConn = newLogReadWriteCloser(wConn, s.logger, socketID, "client")
-	wConn = newCipherReadWriteCloser(wConn, obfs2)
+	wConn = wrappers.NewStreamCipherRWC(wConn, obfs2.Encryptor, obfs2.Decryptor)
 	wConn = newCtxReadWriteCloser(ctx, cancel, wConn)
 
 	return wConn, dc, nil
@@ -144,7 +145,7 @@ func (s *Server) getTelegramStream(ctx context.Context, cancel context.CancelFun
 	}
 
 	wConn = newLogReadWriteCloser(wConn, s.logger, socketID, "telegram")
-	wConn = newCipherReadWriteCloser(wConn, obfs2)
+	wConn = wrappers.NewStreamCipherRWC(wConn, obfs2.Encryptor, obfs2.Decryptor)
 	wConn = newCtxReadWriteCloser(ctx, cancel, wConn)
 
 	return wConn, nil