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

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))
+	}
 }