wrapper now encapsulates keys for the server, wait on a decision as t…

…o memory protection for storage
leaf-ai · Apr 16, 2020 · fb60cf9 · fb60cf9
1 parent d7a851f
commit fb60cf9
Show file tree

Hide file tree

Showing 38 changed files with 3,304 additions and 153 deletions.
diff --git a/Gopkg.lock b/Gopkg.lock
diff --git a/cmd/runner/main.go b/cmd/runner/main.go
@@ -151,6 +151,9 @@ func resourceLimits() (cores uint, mem uint64, storage uint64, err error) {
 //
 func main() {
 
+	// Allow the enclave for secrets to wipe things
+	runner.StopSecret()
+
 	quitC := make(chan struct{})
 	defer close(quitC)
 

diff --git a/cmd/runner/main_test.go b/cmd/runner/main_test.go
@@ -55,6 +55,9 @@ func cleanup() {
 	for _, tmpDir := range cleanupDirs {
 		os.RemoveAll(tmpDir)
 	}
+
+	// Allow the enclave for secrets to wipe things
+	runner.StopSecret()
 }
 
 // TestRunMain can be used to run the server in production mode as opposed to

diff --git a/internal/runner/envelope.go b/internal/runner/envelope.go
@@ -1,16 +1,9 @@
-// Copyright 2018-2020 (c) Cognizant Digital Business, Evolutionary AI. All rights reserved. Issued under the Apache 2.0 License.
+// Copyright 2020 (c) Cognizant Digital Business, Evolutionary AI. All rights reserved. Issued under the Apache 2.0 License.
 
 package runner
 
 import (
-	"crypto/rand"
-	"crypto/rsa"
-	"crypto/sha256"
-	"crypto/x509"
-	"encoding/base64"
 	"encoding/json"
-	"encoding/pem"
-	"strings"
 
 	"github.com/go-stack/stack"
 	"github.com/jjeffery/kv"
@@ -77,141 +70,3 @@ func UnmarshalEnvelope(data []byte) (e *Envelope, err kv.Error) {
 func (e *Envelope) Marshal() ([]byte, error) {
 	return json.Marshal(e)
 }
-
-type Wrapper struct {
-	publicPEM  []byte
-	privatePEM []byte
-	passphrase string
-}
-
-func (w *Wrapper) Envelope(r *Request) (e *Envelope, err kv.Error) {
-	e = &Envelope{
-		Message: Message{
-			Experiment: OpenExperiment{
-				Status:    r.Experiment.Status,
-				PythonVer: r.Experiment.PythonVer,
-			},
-			TimeAdded:          r.Experiment.TimeAdded,
-			ExperimentLifetime: r.Config.Lifetime,
-			Resource:           r.Experiment.Resource,
-		},
-	}
-
-	e.Message.Payload, err = w.WrapRequest(r)
-	return e, err
-}
-
-func (w *Wrapper) Request(e *Envelope) (r *Request, err kv.Error) {
-	return w.UnwrapRequest(e.Message.Payload)
-}
-
-func (w *Wrapper) WrapRequest(r *Request) (encrypted string, err kv.Error) {
-
-	if w == nil {
-		return "", kv.NewError("wrapper missing").With("stack", stack.Trace().TrimRuntime())
-	}
-
-	// Check to see if we have a public key
-	if len(w.publicPEM) == 0 {
-		return "", kv.NewError("public key missing").With("stack", stack.Trace().TrimRuntime())
-	}
-
-	// Serialize the request
-	buffer, err := r.Marshal()
-	if err != nil {
-		return "", err
-	}
-	pubBlock, _ := pem.Decode(w.publicPEM)
-	pub, errGo := x509.ParsePKCS1PublicKey(pubBlock.Bytes)
-	if errGo != nil {
-		return "", kv.Wrap(errGo).With("stack", stack.Trace().TrimRuntime())
-	}
-
-	// encrypt the data and retrieve a symmetric key
-	asymKey, asymData, err := EncryptBlock(buffer)
-	if err != nil {
-		return "", err
-	}
-	asymDataB64 := base64.StdEncoding.EncodeToString(asymData)
-
-	// encrypt the symmetric key using the public RSA PEM
-	asymEncKey, errGo := rsa.EncryptOAEP(sha256.New(), rand.Reader, pub, asymKey[:], nil)
-	if errGo != nil {
-		return "", kv.Wrap(errGo).With("stack", stack.Trace().TrimRuntime())
-	}
-	asymKeyB64 := base64.StdEncoding.EncodeToString(asymEncKey)
-
-	// append the encrypted semtric key, and the symetrically encrypted data into a BASE64 result
-	return asymKeyB64 + "," + asymDataB64, nil
-}
-
-func (w *Wrapper) UnwrapRequest(encrypted string) (r *Request, err kv.Error) {
-	// Check we have a private key and a passphrase
-	if w == nil {
-		return nil, kv.NewError("wrapper missing").With("stack", stack.Trace().TrimRuntime())
-	}
-
-	// Check to see if we have a private key, and a passphrase
-	if len(w.privatePEM) == 0 {
-		return nil, kv.NewError("private key missing").With("stack", stack.Trace().TrimRuntime())
-	}
-	if len(w.passphrase) == 0 {
-		return nil, kv.NewError("passphrase missing").With("stack", stack.Trace().TrimRuntime())
-	}
-
-	// break off the fixed length symetric but RSA encrypted key using the comma delimiter
-	items := strings.Split(encrypted, ",")
-	if len(items) > 2 {
-		return nil, kv.NewError("too many values in encrypted data").With("stack", stack.Trace().TrimRuntime())
-	}
-	if len(items) < 2 {
-		return nil, kv.NewError("missing values in encrypted data").With("stack", stack.Trace().TrimRuntime())
-	}
-
-	asymKeyDecoded, errGo := base64.StdEncoding.DecodeString(items[0])
-	if errGo != nil {
-		return nil, kv.Wrap(errGo, "asymmetric key bad").With("stack", stack.Trace().TrimRuntime())
-	}
-	asymBodyDecoded, errGo := base64.StdEncoding.DecodeString(items[1])
-	if errGo != nil {
-		return nil, kv.Wrap(errGo, "asymmetric encrypted data bad").With("stack", stack.Trace().TrimRuntime())
-	}
-
-	// Decrypt the RSA encrypted asymmetric key
-	prvBlock, _ := pem.Decode(w.privatePEM)
-	if prvBlock == nil {
-		return nil, kv.NewError("private PEM not decoded").With("stack", stack.Trace().TrimRuntime())
-	}
-	if got, want := prvBlock.Type, "RSA PRIVATE KEY"; got != want {
-		return nil, kv.NewError("unknown block type").With("got", got, "want", want).With("stack", stack.Trace().TrimRuntime())
-	}
-
-	decryptedBlock, errGo := x509.DecryptPEMBlock(prvBlock, []byte(w.passphrase))
-	if errGo != nil {
-		return nil, kv.Wrap(errGo).With("stack", stack.Trace().TrimRuntime())
-	}
-
-	prv, errGo := x509.ParsePKCS1PrivateKey(decryptedBlock)
-	if errGo != nil {
-		return nil, kv.Wrap(errGo).With("stack", stack.Trace().TrimRuntime())
-	}
-
-	asymSliceKey, errGo := rsa.DecryptOAEP(sha256.New(), rand.Reader, prv, asymKeyDecoded, nil)
-	if errGo != nil {
-		return nil, kv.Wrap(errGo).With("stack", stack.Trace().TrimRuntime())
-	}
-	asymKey := [32]byte{}
-	copy(asymKey[:], asymSliceKey[:32])
-
-	// Decrypt the data using the decrypted asymmetric key
-	decryptedBody, err := DecryptBlock(asymKey, asymBodyDecoded)
-	if err != nil {
-		return nil, err
-	}
-	r, errGo = UnmarshalRequest(decryptedBody)
-	if errGo != nil {
-		return nil, kv.Wrap(errGo).With("stack", stack.Trace().TrimRuntime())
-	}
-
-	return r, nil
-}
diff --git a/internal/runner/envelope_test.go b/internal/runner/envelope_test.go
@@ -1,4 +1,4 @@
-// Copyright 2018-2020 (c) Cognizant Digital Business, Evolutionary AI. All rights reserved. Issued under the Apache 2.0 License.
+// Copyright 2020 (c) Cognizant Digital Business, Evolutionary AI. All rights reserved. Issued under the Apache 2.0 License.
 
 package runner
 
@@ -51,14 +51,12 @@ func TestEnvelopeDetectPos(t *testing.T) {
 }
 
 func setupWrapper() (w *Wrapper, err kv.Error) {
-	w = &Wrapper{
-		passphrase: RandomString(64),
-	}
-	w.privatePEM, w.publicPEM, err = GenerateKeyPair(w.passphrase)
+	passphrase := RandomString(64)
+	privatePEM, publicPEM, err := GenerateKeyPair(passphrase)
 	if err != nil {
 		return nil, err
 	}
-	return w, nil
+	return NewWrapper(publicPEM, privatePEM, []byte(passphrase))
 }
 
 func TestEnvelopeConv(t *testing.T) {

diff --git a/internal/runner/main_test.go b/internal/runner/main_test.go
@@ -32,6 +32,9 @@ var (
 )
 
 func TestMain(m *testing.M) {
+	// Allow the enclave for secrets to wipe things
+	defer StopSecret()
+
 	// Only perform this Parsed check inside the test framework. Do not be tempted
 	// to do this in the main of our production package
 	//