1  package lnencrypt
  2  
  3  import (
  4  	"bytes"
  5  	"testing"
  6  
  7  	"github.com/btcsuite/btcd/btcec/v2"
  8  	"github.com/stretchr/testify/require"
  9  )
 10  
 11  // TestEncryptDecryptPayload tests that given a static key, we're able to
 12  // properly decrypt and encrypted payload. We also test that we'll reject a
 13  // ciphertext that has been modified.
 14  func TestEncryptDecryptPayload(t *testing.T) {
 15  	t.Parallel()
 16  
 17  	payloadCases := []struct {
 18  		// plaintext is the string that we'll be encrypting.
 19  		plaintext []byte
 20  
 21  		// mutator allows a test case to modify the ciphertext before
 22  		// we attempt to decrypt it.
 23  		mutator func(*[]byte)
 24  
 25  		// valid indicates if this test should pass or fail.
 26  		valid bool
 27  	}{
 28  		// Proper payload, should decrypt.
 29  		{
 30  			plaintext: []byte("payload test plain text"),
 31  			mutator:   nil,
 32  			valid:     true,
 33  		},
 34  
 35  		// Mutator modifies cipher text, shouldn't decrypt.
 36  		{
 37  			plaintext: []byte("payload test plain text"),
 38  			mutator: func(p *[]byte) {
 39  				// Flip a byte in the payload to render it
 40  				// invalid.
 41  				(*p)[0] ^= 1
 42  			},
 43  			valid: false,
 44  		},
 45  
 46  		// Cipher text is too small, shouldn't decrypt.
 47  		{
 48  			plaintext: []byte("payload test plain text"),
 49  			mutator: func(p *[]byte) {
 50  				// Modify the cipher text to be zero length.
 51  				*p = []byte{}
 52  			},
 53  			valid: false,
 54  		},
 55  	}
 56  
 57  	keyRing := &MockKeyRing{}
 58  	keyRingEnc, err := KeyRingEncrypter(keyRing)
 59  	require.NoError(t, err)
 60  
 61  	_, pubKey := btcec.PrivKeyFromBytes([]byte{0x01, 0x02, 0x03, 0x04})
 62  
 63  	privKey, err := btcec.NewPrivateKey()
 64  	require.NoError(t, err)
 65  	privKeyEnc, err := ECDHEncrypter(privKey, pubKey)
 66  	require.NoError(t, err)
 67  
 68  	for _, payloadCase := range payloadCases {
 69  		payloadCase := payloadCase
 70  		for _, enc := range []*Encrypter{keyRingEnc, privKeyEnc} {
 71  			enc := enc
 72  
 73  			// First, we'll encrypt the passed payload with our
 74  			// scheme.
 75  			var cipherBuffer bytes.Buffer
 76  			err = enc.EncryptPayloadToWriter(
 77  				payloadCase.plaintext, &cipherBuffer,
 78  			)
 79  			require.NoError(t, err)
 80  
 81  			// If we have a mutator, then we'll wrong the mutator
 82  			// over the cipher text, then reset the main buffer and
 83  			// re-write the new cipher text.
 84  			if payloadCase.mutator != nil {
 85  				cipherText := cipherBuffer.Bytes()
 86  
 87  				payloadCase.mutator(&cipherText)
 88  
 89  				cipherBuffer.Reset()
 90  				cipherBuffer.Write(cipherText)
 91  			}
 92  
 93  			plaintext, err := enc.DecryptPayloadFromReader(
 94  				&cipherBuffer,
 95  			)
 96  
 97  			if !payloadCase.valid {
 98  				require.Error(t, err)
 99  
100  				continue
101  			}
102  
103  			require.NoError(t, err)
104  			require.Equal(
105  				t, plaintext, payloadCase.plaintext,
106  			)
107  		}
108  	}
109  }
110  
111  // TestInvalidKeyGeneration tests that key generation fails when deriving the
112  // key fails.
113  func TestInvalidKeyGeneration(t *testing.T) {
114  	t.Parallel()
115  
116  	_, err := KeyRingEncrypter(&MockKeyRing{true})
117  	require.Error(t, err)
118  }