github.com/xraypb/xray-core@v1.6.6/common/crypto/auth_test.go (about)

     1  package crypto_test
     2  
     3  import (
     4  	"bytes"
     5  	"crypto/aes"
     6  	"crypto/cipher"
     7  	"crypto/rand"
     8  	"io"
     9  	"testing"
    10  
    11  	"github.com/google/go-cmp/cmp"
    12  	"github.com/xraypb/xray-core/common"
    13  	"github.com/xraypb/xray-core/common/buf"
    14  	. "github.com/xraypb/xray-core/common/crypto"
    15  	"github.com/xraypb/xray-core/common/protocol"
    16  )
    17  
    18  func TestAuthenticationReaderWriter(t *testing.T) {
    19  	key := make([]byte, 16)
    20  	rand.Read(key)
    21  	block, err := aes.NewCipher(key)
    22  	common.Must(err)
    23  
    24  	aead, err := cipher.NewGCM(block)
    25  	common.Must(err)
    26  
    27  	const payloadSize = 1024 * 80
    28  	rawPayload := make([]byte, payloadSize)
    29  	rand.Read(rawPayload)
    30  
    31  	payload := buf.MergeBytes(nil, rawPayload)
    32  
    33  	cache := bytes.NewBuffer(nil)
    34  	iv := make([]byte, 12)
    35  	rand.Read(iv)
    36  
    37  	writer := NewAuthenticationWriter(&AEADAuthenticator{
    38  		AEAD:                    aead,
    39  		NonceGenerator:          GenerateStaticBytes(iv),
    40  		AdditionalDataGenerator: GenerateEmptyBytes(),
    41  	}, PlainChunkSizeParser{}, cache, protocol.TransferTypeStream, nil)
    42  
    43  	common.Must(writer.WriteMultiBuffer(payload))
    44  	if cache.Len() <= 1024*80 {
    45  		t.Error("cache len: ", cache.Len())
    46  	}
    47  	common.Must(writer.WriteMultiBuffer(buf.MultiBuffer{}))
    48  
    49  	reader := NewAuthenticationReader(&AEADAuthenticator{
    50  		AEAD:                    aead,
    51  		NonceGenerator:          GenerateStaticBytes(iv),
    52  		AdditionalDataGenerator: GenerateEmptyBytes(),
    53  	}, PlainChunkSizeParser{}, cache, protocol.TransferTypeStream, nil)
    54  
    55  	var mb buf.MultiBuffer
    56  
    57  	for mb.Len() < payloadSize {
    58  		mb2, err := reader.ReadMultiBuffer()
    59  		common.Must(err)
    60  
    61  		mb, _ = buf.MergeMulti(mb, mb2)
    62  	}
    63  
    64  	if mb.Len() != payloadSize {
    65  		t.Error("mb len: ", mb.Len())
    66  	}
    67  
    68  	mbContent := make([]byte, payloadSize)
    69  	buf.SplitBytes(mb, mbContent)
    70  	if r := cmp.Diff(mbContent, rawPayload); r != "" {
    71  		t.Error(r)
    72  	}
    73  
    74  	_, err = reader.ReadMultiBuffer()
    75  	if err != io.EOF {
    76  		t.Error("error: ", err)
    77  	}
    78  }
    79  
    80  func TestAuthenticationReaderWriterPacket(t *testing.T) {
    81  	key := make([]byte, 16)
    82  	common.Must2(rand.Read(key))
    83  	block, err := aes.NewCipher(key)
    84  	common.Must(err)
    85  
    86  	aead, err := cipher.NewGCM(block)
    87  	common.Must(err)
    88  
    89  	cache := buf.New()
    90  	iv := make([]byte, 12)
    91  	rand.Read(iv)
    92  
    93  	writer := NewAuthenticationWriter(&AEADAuthenticator{
    94  		AEAD:                    aead,
    95  		NonceGenerator:          GenerateStaticBytes(iv),
    96  		AdditionalDataGenerator: GenerateEmptyBytes(),
    97  	}, PlainChunkSizeParser{}, cache, protocol.TransferTypePacket, nil)
    98  
    99  	var payload buf.MultiBuffer
   100  	pb1 := buf.New()
   101  	pb1.Write([]byte("abcd"))
   102  	payload = append(payload, pb1)
   103  
   104  	pb2 := buf.New()
   105  	pb2.Write([]byte("efgh"))
   106  	payload = append(payload, pb2)
   107  
   108  	common.Must(writer.WriteMultiBuffer(payload))
   109  	if cache.Len() == 0 {
   110  		t.Error("cache len: ", cache.Len())
   111  	}
   112  
   113  	common.Must(writer.WriteMultiBuffer(buf.MultiBuffer{}))
   114  
   115  	reader := NewAuthenticationReader(&AEADAuthenticator{
   116  		AEAD:                    aead,
   117  		NonceGenerator:          GenerateStaticBytes(iv),
   118  		AdditionalDataGenerator: GenerateEmptyBytes(),
   119  	}, PlainChunkSizeParser{}, cache, protocol.TransferTypePacket, nil)
   120  
   121  	mb, err := reader.ReadMultiBuffer()
   122  	common.Must(err)
   123  
   124  	mb, b1 := buf.SplitFirst(mb)
   125  	if b1.String() != "abcd" {
   126  		t.Error("b1: ", b1.String())
   127  	}
   128  
   129  	mb, b2 := buf.SplitFirst(mb)
   130  	if b2.String() != "efgh" {
   131  		t.Error("b2: ", b2.String())
   132  	}
   133  
   134  	if !mb.IsEmpty() {
   135  		t.Error("not empty")
   136  	}
   137  
   138  	_, err = reader.ReadMultiBuffer()
   139  	if err != io.EOF {
   140  		t.Error("error: ", err)
   141  	}
   142  }