Refactor cipherrwc to wrappers

obfuscated2/obfuscated2.go

@@ -11,22 +11,8 @@ import (
 // Obfuscated2 contains AES CTR encryption and decryption streams
 // for telegram connection.
 type Obfuscated2 struct {
-	decryptor cipher.Stream
-	encryptor cipher.Stream
-}
-
-// Encrypt encrypts given data.
-func (o *Obfuscated2) Encrypt(data []byte) []byte {
-	buf := make([]byte, len(data))
-	o.encryptor.XORKeyStream(buf, data)
-	return buf
-}
-
-// Decrypt decrypts given data.
-func (o *Obfuscated2) Decrypt(data []byte) []byte {
-	buf := make([]byte, len(data))
-	o.decryptor.XORKeyStream(buf, data)
-	return buf
+	Decryptor cipher.Stream
+	Encryptor cipher.Stream
 }
 
 // ParseObfuscated2ClientFrame parses client frame. Please check this link for
@@ -54,8 +40,8 @@ func ParseObfuscated2ClientFrame(secret, data []byte) (*Obfuscated2, int16, erro
 	}
 
 	obfs := &Obfuscated2{
-		decryptor: decryptor,
-		encryptor: encryptor,
+		Decryptor: decryptor,
+		Encryptor: encryptor,
 	}
 
 	return obfs, decryptedFrame.DC(), nil
@@ -77,8 +63,8 @@ func MakeTelegramObfuscated2Frame() (*Obfuscated2, Frame) {
 	copy(frame, copyFrame)
 
 	obfs := &Obfuscated2{
-		decryptor: decryptor,
-		encryptor: encryptor,
+		Decryptor: decryptor,
+		Encryptor: encryptor,
 	}
 
 	return obfs, frame

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

wrappers/streamcipherrwc.go

@@ -0,0 +1,53 @@
+package wrappers
+
+import (
+	"bytes"
+	"crypto/cipher"
+	"io"
+)
+
+type StreamCipherReadWriteCloser struct {
+	encryptor cipher.Stream
+	decryptor cipher.Stream
+	conn      io.ReadWriteCloser
+	rest      *bytes.Buffer
+}
+
+// Read reads from connection
+func (c *StreamCipherReadWriteCloser) Read(p []byte) (n int, err error) {
+	n, err = c.conn.Read(p)
+	c.decryptor.XORKeyStream(p, p[:n])
+	return
+}
+
+// Write writes into connection.
+func (c *StreamCipherReadWriteCloser) Write(p []byte) (int, error) {
+	encrypted := make([]byte, len(p))
+	c.encryptor.XORKeyStream(encrypted, 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 *StreamCipherReadWriteCloser) Close() error {
+	return c.conn.Close()
+}
+
+func NewStreamCipherRWC(conn io.ReadWriteCloser, encryptor, decryptor cipher.Stream) io.ReadWriteCloser {
+	return &StreamCipherReadWriteCloser{
+		conn:      conn,
+		encryptor: encryptor,
+		decryptor: decryptor,
+		rest:      &bytes.Buffer{},
+	}
+}