github.com/sentienttechnologies/studio-go-runner@v0.0.0-20201118202441-6d21f2ced8ee/internal/runner/rsa_test.go (about)

     1  // Copyright 2018-2020 (c) Cognizant Digital Business, Evolutionary AI. All rights reserved. Issued under the Apache 2.0 License.
     2  
     3  package runner
     4  
     5  import (
     6  	"bytes"
     7  	"crypto/rand"
     8  	"crypto/rsa"
     9  	"crypto/sha256"
    10  	"crypto/x509"
    11  	"encoding/pem"
    12  	"testing"
    13  
    14  	"github.com/go-stack/stack"
    15  	"github.com/jjeffery/kv"
    16  )
    17  
    18  // This file contains a number of tests related to handling key files for use in
    19  // Encryption of the messages being used by the runner.
    20  
    21  // TestRSA will test the encryption and decryption of short
    22  // blocks of data, typically used for encryption of symetrics
    23  // keys embeeded within messages etc
    24  //
    25  func TestRSA(t *testing.T) {
    26  	passphrase := RandomString(10)
    27  	privatePEM, publicPEM, err := GenerateKeyPair(passphrase)
    28  	if err != nil {
    29  		t.Fatal(err.With("stack", stack.Trace().TrimRuntime()))
    30  	}
    31  
    32  	// Extract the PEM-encoded data block
    33  	pubBlock, _ := pem.Decode(publicPEM)
    34  	if pubBlock == nil {
    35  		t.Fatal(kv.NewError("public PEM not decoded").With("stack", stack.Trace().TrimRuntime()))
    36  	}
    37  	if got, want := pubBlock.Type, "RSA PUBLIC KEY"; got != want {
    38  		t.Fatal(kv.NewError("unknown block type").With("got", got, "want", want).With("stack", stack.Trace().TrimRuntime()))
    39  	}
    40  
    41  	pub, errGo := x509.ParsePKCS1PublicKey(pubBlock.Bytes)
    42  	if errGo != nil {
    43  		t.Fatal(kv.Wrap(errGo).With("stack", stack.Trace().TrimRuntime()))
    44  	}
    45  
    46  	msg := []byte(RandomString(256))
    47  	encrypted, errGo := rsa.EncryptOAEP(sha256.New(), rand.Reader, pub, msg, nil)
    48  	if errGo != nil {
    49  		t.Fatal(kv.Wrap(errGo).With("stack", stack.Trace().TrimRuntime()))
    50  	}
    51  
    52  	// Now we have the encrypted data, try decrypting it
    53  	prvBlock, _ := pem.Decode(privatePEM)
    54  	if prvBlock == nil {
    55  		t.Fatal(kv.NewError("private PEM not decoded").With("stack", stack.Trace().TrimRuntime()))
    56  	}
    57  	if got, want := prvBlock.Type, "RSA PRIVATE KEY"; got != want {
    58  		t.Fatal(kv.NewError("unknown block type").With("got", got, "want", want).With("stack", stack.Trace().TrimRuntime()))
    59  	}
    60  
    61  	decryptedBlock, errGo := x509.DecryptPEMBlock(prvBlock, []byte(passphrase))
    62  	if errGo != nil {
    63  		t.Fatal(kv.Wrap(errGo).With("stack", stack.Trace().TrimRuntime()))
    64  	}
    65  
    66  	prv, errGo := x509.ParsePKCS1PrivateKey(decryptedBlock)
    67  	if errGo != nil {
    68  		t.Fatal(kv.Wrap(errGo).With("stack", stack.Trace().TrimRuntime()))
    69  	}
    70  
    71  	out, errGo := rsa.DecryptOAEP(sha256.New(), rand.Reader, prv, encrypted, nil)
    72  	if errGo != nil {
    73  		t.Fatal(kv.Wrap(errGo).With("stack", stack.Trace().TrimRuntime()))
    74  	}
    75  
    76  	if 0 != bytes.Compare(msg, out) {
    77  		t.Fatal(kv.NewError("roundtrip failed").With("stack", stack.Trace().TrimRuntime()))
    78  	}
    79  }