all: add use_gzip configure and some improvements, see #28

src/frp/utils/conn/conn.go

@@ -16,17 +16,12 @@ package conn
 
 import (
 	"bufio"
-	"bytes"
-	"encoding/binary"
 	"fmt"
 	"io"
 	"net"
 	"strings"
 	"sync"
 	"time"
-
-	"frp/utils/log"
-	"frp/utils/pcrypto"
 )
 
 type Listener struct {
@@ -169,146 +164,3 @@ func (c *Conn) IsClosed() (closeFlag bool) {
 	c.mutex.RUnlock()
 	return
 }
-
-// will block until connection close
-func Join(c1 *Conn, c2 *Conn) {
-	var wait sync.WaitGroup
-	pipe := func(to *Conn, from *Conn) {
-		defer to.Close()
-		defer from.Close()
-		defer wait.Done()
-
-		var err error
-		_, err = io.Copy(to.TcpConn, from.TcpConn)
-		if err != nil {
-			log.Warn("join connections error, %v", err)
-		}
-	}
-
-	wait.Add(2)
-	go pipe(c1, c2)
-	go pipe(c2, c1)
-	wait.Wait()
-	return
-}
-
-// messages from c1 to c2 will be encrypted
-// and from c2 to c1 will be decrypted
-func JoinMore(c1 *Conn, c2 *Conn, cryptKey string, ptype int) {
-	var wait sync.WaitGroup
-	encryptPipe := func(from *Conn, to *Conn, key string, ttype int) {
-		defer from.Close()
-		defer to.Close()
-		defer wait.Done()
-
-		// we don't care about errors here
-		PipeEncrypt(from.TcpConn, to.TcpConn, key, ttype)
-	}
-
-	decryptPipe := func(to *Conn, from *Conn, key string, ttype int) {
-		defer from.Close()
-		defer to.Close()
-		defer wait.Done()
-
-		// we don't care about errors here
-		PipeDecrypt(to.TcpConn, from.TcpConn, key, ttype)
-	}
-
-	wait.Add(2)
-	go encryptPipe(c1, c2, cryptKey, ptype)
-
-	go decryptPipe(c2, c1, cryptKey, ptype)
-	wait.Wait()
-	log.Debug("One tunnel stopped")
-	return
-}
-
-func unpkgMsg(data []byte) (int, []byte, []byte) {
-	if len(data) < 4 {
-		return -1, nil, nil
-	}
-	llen := int(binary.BigEndian.Uint32(data[0:4]))
-	// no complete
-	if len(data) < llen+4 {
-		return -1, nil, nil
-	}
-
-	return 0, data[4 : llen+4], data[llen+4:]
-}
-
-// decrypt msg from reader, then write into writer
-func PipeDecrypt(r net.Conn, w net.Conn, key string, ptype int) error {
-	laes := new(pcrypto.Pcrypto)
-	if err := laes.Init([]byte(key), ptype); err != nil {
-		log.Error("Pcrypto Init error: %v", err)
-		return fmt.Errorf("Pcrypto Init error: %v", err)
-	}
-
-	buf := make([]byte, 10*1024)
-	var left []byte
-	nreader := bufio.NewReader(r)
-	for {
-		n, err := nreader.Read(buf)
-		if err != nil {
-			return err
-		}
-		left := append(left, buf[:n]...)
-		cnt, buf, left := unpkgMsg(left)
-
-		if cnt < 0 {
-			continue
-		}
-
-		res, err := laes.Decrypt(buf)
-		if err != nil {
-			log.Error("Decrypt [%s] error, %v", string(buf), err)
-			return fmt.Errorf("Decrypt [%s] error: %v", string(buf), err)
-		}
-
-		_, err = w.Write(res)
-		if err != nil {
-			return err
-		}
-	}
-	return nil
-}
-
-func pkgMsg(data []byte) []byte {
-	llen := uint32(len(data))
-	buf := new(bytes.Buffer)
-	binary.Write(buf, binary.BigEndian, llen)
-	buf.Write(data)
-	return buf.Bytes()
-}
-
-// recvive msg from reader, then encrypt msg into write
-func PipeEncrypt(r net.Conn, w net.Conn, key string, ptype int) error {
-	laes := new(pcrypto.Pcrypto)
-	if err := laes.Init([]byte(key), ptype); err != nil {
-		log.Error("Pcrypto Init error: %v", err)
-		return fmt.Errorf("Pcrypto Init error: %v", err)
-	}
-
-	nreader := bufio.NewReader(r)
-	buf := make([]byte, 10*1024)
-
-	for {
-		n, err := nreader.Read(buf)
-		if err != nil {
-			return err
-		}
-		res, err := laes.Encrypt(buf[:n])
-		if err != nil {
-			log.Error("Encrypt error: %v", err)
-			return fmt.Errorf("Encrypt error: %v", err)
-		}
-
-		res = pkgMsg(res)
-		_, err = w.Write(res)
-		if err != nil {
-			return err
-		}
-	}
-
-	return nil
-}

src/frp/utils/pcrypto/pcrypto.go

@@ -29,79 +29,61 @@ import (
 type Pcrypto struct {
 	pkey []byte
 	paes cipher.Block
-	// 0: nono; 1:compress; 2: encrypt; 3: compress and encrypt
-	ptyp int
 }
 
-func (pc *Pcrypto) Init(key []byte, ptyp int) error {
+func (pc *Pcrypto) Init(key []byte) error {
 	var err error
 	pc.pkey = pKCS7Padding(key, aes.BlockSize)
 	pc.paes, err = aes.NewCipher(pc.pkey)
-	if ptyp == 1 || ptyp == 2 || ptyp == 3 {
-		pc.ptyp = ptyp
-	} else {
-		pc.ptyp = 0
-	}
-
 	return err
 }
 
 func (pc *Pcrypto) Encrypt(src []byte) ([]byte, error) {
-	var zbuf bytes.Buffer
-
-	// gzip
-	if pc.ptyp == 1 || pc.ptyp == 3 {
-		zwr, err := gzip.NewWriterLevel(&zbuf, gzip.DefaultCompression)
-		if err != nil {
-			return nil, err
-		}
-		defer zwr.Close()
-		zwr.Write(src)
-		zwr.Flush()
-		src = zbuf.Bytes()
-	}
-
 	// aes
-	if pc.ptyp == 2 || pc.ptyp == 3 {
-		src = pKCS7Padding(src, aes.BlockSize)
-		blockMode := cipher.NewCBCEncrypter(pc.paes, pc.pkey)
-		crypted := make([]byte, len(src))
-		blockMode.CryptBlocks(crypted, src)
-		src = crypted
-	}
-
-	return src, nil
+	src = pKCS7Padding(src, aes.BlockSize)
+	blockMode := cipher.NewCBCEncrypter(pc.paes, pc.pkey)
+	crypted := make([]byte, len(src))
+	blockMode.CryptBlocks(crypted, src)
+	return crypted, nil
 }
 
 func (pc *Pcrypto) Decrypt(str []byte) ([]byte, error) {
 	// aes
-	if pc.ptyp == 2 || pc.ptyp == 3 {
-		decryptText, err := hex.DecodeString(fmt.Sprintf("%x", str))
-		if err != nil {
-			return nil, err
-		}
-
-		if len(decryptText)%aes.BlockSize != 0 {
-			return nil, errors.New("crypto/cipher: ciphertext is not a multiple of the block size")
-		}
-
-		blockMode := cipher.NewCBCDecrypter(pc.paes, pc.pkey)
-
-		blockMode.CryptBlocks(decryptText, decryptText)
-		str = pKCS7UnPadding(decryptText)
+	decryptText, err := hex.DecodeString(fmt.Sprintf("%x", str))
+	if err != nil {
+		return nil, err
 	}
 
-	// gunzip
-	if pc.ptyp == 1 || pc.ptyp == 3 {
-		zbuf := bytes.NewBuffer(str)
-		zrd, err := gzip.NewReader(zbuf)
-		if err != nil {
-			return nil, err
-		}
-		defer zrd.Close()
-		str, _ = ioutil.ReadAll(zrd)
+	if len(decryptText)%aes.BlockSize != 0 {
+		return nil, errors.New("crypto/cipher: ciphertext is not a multiple of the block size")
 	}
 
+	blockMode := cipher.NewCBCDecrypter(pc.paes, pc.pkey)
+
+	blockMode.CryptBlocks(decryptText, decryptText)
+	return pKCS7UnPadding(decryptText), nil
+}
+
+func (pc *Pcrypto) Compression(src []byte) ([]byte, error) {
+	var zbuf bytes.Buffer
+	zwr, err := gzip.NewWriterLevel(&zbuf, gzip.DefaultCompression)
+	if err != nil {
+		return nil, err
+	}
+	defer zwr.Close()
+	zwr.Write(src)
+	zwr.Flush()
+	return zbuf.Bytes(), nil
+}
+
+func (pc *Pcrypto) Decompression(src []byte) ([]byte, error) {
+	zbuf := bytes.NewBuffer(src)
+	zrd, err := gzip.NewReader(zbuf)
+	if err != nil {
+		return nil, err
+	}
+	defer zrd.Close()
+	str, _ := ioutil.ReadAll(zrd)
 	return str, nil
 }
 

src/frp/utils/pcrypto/pcrypto_test.go

@@ -15,83 +15,48 @@
 package pcrypto
 
 import (
-	"fmt"
 	"testing"
 )
 
+var (
+	pp *Pcrypto
+)
+
+func init() {
+	pp = &Pcrypto{}
+	pp.Init([]byte("Hana"))
+}
+
 func TestEncrypt(t *testing.T) {
-	return
-	pp := new(Pcrypto)
-	pp.Init([]byte("Hana"), 1)
-	res, err := pp.Encrypt([]byte("Test Encrypt!"))
+	testStr := "Test Encrypt!"
+	res, err := pp.Encrypt([]byte(testStr))
 	if err != nil {
-		t.Fatal(err)
+		t.Fatalf("encrypt error: %v", err)
 	}
 
-	fmt.Printf("Encrypt: len %d, [%x]\n", len(res), res)
+	res, err = pp.Decrypt([]byte(res))
+	if err != nil {
+		t.Fatalf("decrypt error: %v", err)
+	}
+
+	if string(res) != testStr {
+		t.Fatalf("test encrypt error, from [%s] to [%s]", testStr, string(res))
+	}
 }
 
-func TestDecrypt(t *testing.T) {
-	fmt.Println("*****************************************************")
-	{
-		pp := new(Pcrypto)
-		pp.Init([]byte("Hana"), 0)
-		res, err := pp.Encrypt([]byte("Test Decrypt! 0"))
-		if err != nil {
-			t.Fatal(err)
-		}
-
-		res, err = pp.Decrypt(res)
-		if err != nil {
-			t.Fatal(err)
-		}
-
-		fmt.Printf("[%s]\n", string(res))
-	}
-	{
-		pp := new(Pcrypto)
-		pp.Init([]byte("Hana"), 1)
-		res, err := pp.Encrypt([]byte("Test Decrypt! 1"))
-		if err != nil {
-			t.Fatal(err)
-		}
-
-		res, err = pp.Decrypt(res)
-		if err != nil {
-			t.Fatal(err)
-		}
-
-		fmt.Printf("[%s]\n", string(res))
-	}
-	{
-		pp := new(Pcrypto)
-		pp.Init([]byte("Hana"), 2)
-		res, err := pp.Encrypt([]byte("Test Decrypt! 2"))
-		if err != nil {
-			t.Fatal(err)
-		}
-
-		res, err = pp.Decrypt(res)
-		if err != nil {
-			t.Fatal(err)
-		}
-
-		fmt.Printf("[%s]\n", string(res))
-	}
-	{
-		pp := new(Pcrypto)
-		pp.Init([]byte("Hana"), 3)
-		res, err := pp.Encrypt([]byte("Test Decrypt! 3"))
-		if err != nil {
-			t.Fatal(err)
-		}
-
-		res, err = pp.Decrypt(res)
-		if err != nil {
-			t.Fatal(err)
-		}
-
-		fmt.Printf("[%s]\n", string(res))
+func TestCompression(t *testing.T) {
+	testStr := "Test Compression!"
+	res, err := pp.Compression([]byte(testStr))
+	if err != nil {
+		t.Fatalf("compression error: %v", err)
 	}
 
+	res, err = pp.Decompression(res)
+	if err != nil {
+		t.Fatalf("decompression error: %v", err)
+	}
+
+	if string(res) != testStr {
+		t.Fatalf("test compression error, from [%s] to [%s]", testStr, string(res))
+	}
 }