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 }