Merge pull request #5529 from oasisprotocol/peternose/bugfix/km-race-…

…conditions

go/worker/keymanager: Fix race condition when accessing status fields

.changelog/5529.bugfix.md

@@ -0,0 +1 @@
+go/worker/keymanager: Fix race conditions when accessing status fields

go/worker/keymanager/status.go

@@ -4,7 +4,6 @@ import (
 	"github.com/libp2p/go-libp2p/core/peer"
 
 	"github.com/oasisprotocol/oasis-core/go/common"
-	"github.com/oasisprotocol/oasis-core/go/common/version"
 	"github.com/oasisprotocol/oasis-core/go/worker/keymanager/api"
 )
 
@@ -56,27 +55,17 @@ func (w *Worker) GetStatus() (*api.Status, error) {
 		al = append(al, ral)
 	}
 
-	var pc []byte
-	if w.enclaveStatus != nil {
-		pc = w.enclaveStatus.InitResponse.PolicyChecksum
-	}
-
-	var aw *version.Version
-	if w.rtStatus != nil {
-		aw = &w.rtStatus.version
-	}
-
 	gs := w.globalStatus
 	ws := api.WorkerStatus{
 		Status:           ss,
-		ActiveVersion:    aw,
+		ActiveVersion:    w.activeVersion,
 		MayGenerate:      w.mayGenerate,
 		RuntimeID:        &w.runtimeID,
 		ClientRuntimes:   rts,
 		AccessList:       al,
 		PrivatePeers:     ps,
 		Policy:           w.policy,
-		PolicyChecksum:   pc,
+		PolicyChecksum:   w.policyChecksum,
 		MasterSecrets:    w.masterSecretStats,
 		EphemeralSecrets: w.ephemeralSecretStats,
 	}

go/worker/keymanager/worker.go

@@ -91,9 +91,10 @@ type Worker struct { // nolint: maligned
 	roleProvider registration.RoleProvider
 	backend      api.Backend
 
-	globalStatus  *api.Status
-	enclaveStatus *api.SignedInitResponse
-	policy        *api.SignedPolicySGX
+	globalStatus   *api.Status
+	policy         *api.SignedPolicySGX
+	policyChecksum []byte
+	activeVersion  *version.Version
 
 	masterSecretStats    workerKeymanager.MasterSecretStats
 	ephemeralSecretStats workerKeymanager.EphemeralSecretStats
@@ -202,8 +203,6 @@ func (w *Worker) CallEnclave(ctx context.Context, data []byte, kind enclaverpc.K
 		switch frame.UntrustedPlaintext {
 		case "":
 			// Anyone can connect.
-		case api.RPCMethodGetPublicKey, api.RPCMethodGetPublicEphemeralKey:
-			// Anyone can get public keys.
 		default:
 			if _, privatePeered := w.privatePeers[peerID]; !privatePeered {
 				// Defer to access control to check the policy.
@@ -359,13 +358,13 @@ func (w *Worker) initEnclave(kmStatus *api.Status, rtStatus *runtimeStatus) (*ap
 
 	// Update metrics.
 	enclaveMasterSecretGenerationNumber.WithLabelValues(w.runtimeLabel).Set(float64(kmStatus.Generation))
-	if w.enclaveStatus == nil || !bytes.Equal(w.enclaveStatus.InitResponse.PolicyChecksum, signedInitResp.InitResponse.PolicyChecksum) {
+	if !bytes.Equal(w.policyChecksum, signedInitResp.InitResponse.PolicyChecksum) {
 		policyUpdateCount.WithLabelValues(w.runtimeLabel).Inc()
 	}
 
-	// Cache the key manager enclave status and the currently active policy.
-	w.enclaveStatus = &signedInitResp
+	// Cache the currently active policy and its checksum.
 	w.policy = kmStatus.Policy
+	w.policyChecksum = signedInitResp.InitResponse.PolicyChecksum
 
 	return &signedInitResp, nil
 }
@@ -410,6 +409,13 @@ func (w *Worker) setStatus(status *api.Status) {
 	w.globalStatus = status
 }
 
+func (w *Worker) setVersion(v *version.Version) {
+	w.Lock()
+	defer w.Unlock()
+
+	w.activeVersion = v
+}
+
 func (w *Worker) setLastGeneratedMasterSecretGeneration(generation uint64) {
 	w.Lock()
 	defer w.Unlock()
@@ -645,21 +651,12 @@ func (w *Worker) generateMasterSecret(runtimeID common.Namespace, height int64,
 		return fmt.Errorf("failed to generate master secret: %w", err)
 	}
 
-	// Fetch key manager runtime details.
-	kmRt, err := w.commonWorker.Consensus.Registry().GetRuntime(w.ctx, &registry.GetRuntimeQuery{
-		Height: consensus.HeightLatest,
-		ID:     kmStatus.ID,
-	})
+	rak, err := w.runtimeAttestationKey(rtStatus)
 	if err != nil {
 		return err
 	}
 
-	rak, err := w.runtimeAttestationKey(rtStatus, kmRt)
-	if err != nil {
-		return err
-	}
-
-	reks, err := w.runtimeEncryptionKeys(kmStatus, kmRt)
+	reks, err := w.runtimeEncryptionKeys(kmStatus)
 	if err != nil {
 		return err
 	}
@@ -722,21 +719,12 @@ func (w *Worker) generateEphemeralSecret(runtimeID common.Namespace, height int6
 		return fmt.Errorf("failed to generate ephemeral secret: %w", err)
 	}
 
-	// Fetch key manager runtime details.
-	kmRt, err := w.commonWorker.Consensus.Registry().GetRuntime(w.ctx, &registry.GetRuntimeQuery{
-		Height: consensus.HeightLatest,
-		ID:     kmStatus.ID,
-	})
-	if err != nil {
-		return err
-	}
-
-	rak, err := w.runtimeAttestationKey(rtStatus, kmRt)
+	rak, err := w.runtimeAttestationKey(rtStatus)
 	if err != nil {
 		return err
 	}
 
-	reks, err := w.runtimeEncryptionKeys(kmStatus, kmRt)
+	reks, err := w.runtimeEncryptionKeys(kmStatus)
 	if err != nil {
 		return err
 	}
@@ -759,7 +747,12 @@ func (w *Worker) generateEphemeralSecret(runtimeID common.Namespace, height int6
 	return err
 }
 
-func (w *Worker) runtimeAttestationKey(rtStatus *runtimeStatus, kmRt *registry.Runtime) (*signature.PublicKey, error) {
+func (w *Worker) runtimeAttestationKey(rtStatus *runtimeStatus) (*signature.PublicKey, error) {
+	kmRt, err := w.runtime.RegistryDescriptor(w.ctx)
+	if err != nil {
+		return nil, err
+	}
+
 	var rak *signature.PublicKey
 	switch kmRt.TEEHardware {
 	case node.TEEHardwareInvalid:
@@ -776,7 +769,12 @@ func (w *Worker) runtimeAttestationKey(rtStatus *runtimeStatus, kmRt *registry.R
 	return rak, nil
 }
 
-func (w *Worker) runtimeEncryptionKeys(kmStatus *api.Status, kmRt *registry.Runtime) (map[x25519.PublicKey]struct{}, error) {
+func (w *Worker) runtimeEncryptionKeys(kmStatus *api.Status) (map[x25519.PublicKey]struct{}, error) {
+	kmRt, err := w.runtime.RegistryDescriptor(w.ctx)
+	if err != nil {
+		return nil, err
+	}
+
 	reks := make(map[x25519.PublicKey]struct{})
 	for _, id := range kmStatus.Nodes {
 		var n *node.Node
@@ -1039,15 +1037,22 @@ func (w *Worker) handleRuntimeHostEvent(ev *host.Event) {
 		default:
 			return
 		}
+		w.setVersion(&w.rtStatus.version)
 
 		if w.kmStatus == nil {
 			return
 		}
 
+		// Check whether the enclave has been initialized at least once.
+		// If true, preregistration is not required.
+		w.RLock()
+		initialized := w.policyChecksum != nil
+		w.RUnlock()
+
 		// Send a node preregistration, so that other nodes know to update their access
 		// control. Without it, the enclave won't be able to replicate the master secrets
 		// needed for initialization.
-		if w.enclaveStatus == nil {
+		if !initialized {
 			rtStatus := w.rtStatus
 			w.roleProvider.SetAvailableWithCallback(func(n *node.Node) error {
 				rt := n.AddOrUpdateRuntime(w.runtime.ID(), rtStatus.version)
@@ -1065,6 +1070,7 @@ func (w *Worker) handleRuntimeHostEvent(ev *host.Event) {
 	case ev.FailedToStart != nil, ev.Stopped != nil:
 		// Worker failed to start or was stopped -- we can no longer service requests.
 		w.rtStatus = nil
+		w.setVersion(nil)
 		w.roleProvider.SetUnavailable()
 	default:
 		// Unknown event.

keymanager/src/runtime/methods.rs

@@ -66,7 +66,7 @@ const MAX_EPHEMERAL_KEY_AGE: EpochTime = 10;
 const MAX_FRESH_HEIGHT_AGE: u64 = 50;
 
 /// Initialize the Kdf.
-pub fn init_kdf(ctx: &mut RpcContext, req: &InitRequest) -> Result<SignedInitResponse> {
+pub fn init_kdf(ctx: &RpcContext, req: &InitRequest) -> Result<SignedInitResponse> {
     let rctx = runtime_context!(ctx, KmContext);
     let runtime_id = rctx.runtime_id;
     let storage = ctx.untrusted_local_storage;
@@ -96,7 +96,7 @@ pub fn init_kdf(ctx: &mut RpcContext, req: &InitRequest) -> Result<SignedInitRes
 }
 
 /// See `Kdf::get_or_create_keys`.
-pub fn get_or_create_keys(ctx: &mut RpcContext, req: &LongTermKeyRequest) -> Result<KeyPair> {
+pub fn get_or_create_keys(ctx: &RpcContext, req: &LongTermKeyRequest) -> Result<KeyPair> {
     authorize_private_key_generation(ctx, &req.runtime_id)?;
     validate_height_freshness(ctx, req.height)?;
 
@@ -109,7 +109,7 @@ pub fn get_or_create_keys(ctx: &mut RpcContext, req: &LongTermKeyRequest) -> Res
 }
 
 /// See `Kdf::get_public_key`.
-pub fn get_public_key(ctx: &mut RpcContext, req: &LongTermKeyRequest) -> Result<SignedPublicKey> {
+pub fn get_public_key(ctx: &RpcContext, req: &LongTermKeyRequest) -> Result<SignedPublicKey> {
     // No authentication or authorization.
     // Absolutely anyone is allowed to query public long-term keys.
 
@@ -126,7 +126,7 @@ pub fn get_public_key(ctx: &mut RpcContext, req: &LongTermKeyRequest) -> Result<
 
 /// See `Kdf::get_or_create_ephemeral_keys`.
 pub fn get_or_create_ephemeral_keys(
-    ctx: &mut RpcContext,
+    ctx: &RpcContext,
     req: &EphemeralKeyRequest,
 ) -> Result<KeyPair> {
     authorize_private_key_generation(ctx, &req.runtime_id)?;
@@ -138,7 +138,7 @@ pub fn get_or_create_ephemeral_keys(
 
 /// See `Kdf::get_public_ephemeral_key`.
 pub fn get_public_ephemeral_key(
-    ctx: &mut RpcContext,
+    ctx: &RpcContext,
     req: &EphemeralKeyRequest,
 ) -> Result<SignedPublicKey> {
     // No authentication or authorization.
@@ -163,7 +163,7 @@ pub fn get_public_ephemeral_key(
 
 /// See `Kdf::replicate_master_secret`.
 pub fn replicate_master_secret(
-    ctx: &mut RpcContext,
+    ctx: &RpcContext,
     req: &ReplicateMasterSecretRequest,
 ) -> Result<ReplicateMasterSecretResponse> {
     authorize_secret_replication(ctx)?;
@@ -181,7 +181,7 @@ pub fn replicate_master_secret(
 
 /// See `Kdf::replicate_ephemeral_secret`.
 pub fn replicate_ephemeral_secret(
-    ctx: &mut RpcContext,
+    ctx: &RpcContext,
     req: &ReplicateEphemeralSecretRequest,
 ) -> Result<ReplicateEphemeralSecretResponse> {
     authorize_secret_replication(ctx)?;
@@ -193,7 +193,7 @@ pub fn replicate_ephemeral_secret(
 
 /// Generate a master secret and encrypt it using the key manager's REK keys.
 pub fn generate_master_secret(
-    ctx: &mut RpcContext,
+    ctx: &RpcContext,
     req: &GenerateMasterSecretRequest,
 ) -> Result<GenerateMasterSecretResponse> {
     let kdf = Kdf::global();
@@ -228,7 +228,7 @@ pub fn generate_master_secret(
 
 /// Generate an ephemeral secret and encrypt it using the key manager's REK keys.
 pub fn generate_ephemeral_secret(
-    ctx: &mut RpcContext,
+    ctx: &RpcContext,
     req: &GenerateEphemeralSecretRequest,
 ) -> Result<GenerateEphemeralSecretResponse> {
     let kdf = Kdf::global();
@@ -260,7 +260,7 @@ pub fn generate_ephemeral_secret(
 
 /// Encrypt a secret using the Deoxys-II MRAE algorithm and the key manager's REK keys.
 fn encrypt_secret(
-    ctx: &mut RpcContext,
+    ctx: &RpcContext,
     secret: Secret,
     checksum: Vec<u8>,
     additional_data: Vec<u8>,
@@ -302,7 +302,7 @@ fn encrypt_secret(
 }
 
 /// Decrypt and store a proposal for the next master secret.
-pub fn load_master_secret(ctx: &mut RpcContext, req: &LoadMasterSecretRequest) -> Result<()> {
+pub fn load_master_secret(ctx: &RpcContext, req: &LoadMasterSecretRequest) -> Result<()> {
     let signed_secret = validate_signed_master_secret(ctx, &req.signed_secret)?;
 
     let secret = match decrypt_master_secret(ctx, &signed_secret)? {
@@ -321,7 +321,7 @@ pub fn load_master_secret(ctx: &mut RpcContext, req: &LoadMasterSecretRequest) -
 
 /// Decrypt and store an ephemeral secret. If decryption fails, try to replicate the secret
 /// from another key manager enclave.
-pub fn load_ephemeral_secret(ctx: &mut RpcContext, req: &LoadEphemeralSecretRequest) -> Result<()> {
+pub fn load_ephemeral_secret(ctx: &RpcContext, req: &LoadEphemeralSecretRequest) -> Result<()> {
     let signed_secret = validate_signed_ephemeral_secret(ctx, &req.signed_secret)?;
 
     let secret = match decrypt_ephemeral_secret(ctx, &signed_secret)? {
@@ -356,7 +356,7 @@ pub fn load_ephemeral_secret(ctx: &mut RpcContext, req: &LoadEphemeralSecretRequ
 
 /// Decrypt master secret with local REK key.
 fn decrypt_master_secret(
-    ctx: &mut RpcContext,
+    ctx: &RpcContext,
     secret: &EncryptedMasterSecret,
 ) -> Result<Option<Secret>> {
     let generation = secret.generation;
@@ -390,7 +390,7 @@ fn decrypt_master_secret(
 
 /// Decrypt ephemeral secret with local REK key.
 fn decrypt_ephemeral_secret(
-    ctx: &mut RpcContext,
+    ctx: &RpcContext,
     secret: &EncryptedEphemeralSecret,
 ) -> Result<Option<Secret>> {
     let epoch = secret.epoch;
@@ -422,7 +422,7 @@ fn decrypt_ephemeral_secret(
 }
 
 /// Key manager client for master and ephemeral secret replication.
-fn key_manager_client_for_replication(ctx: &mut RpcContext) -> RemoteClient {
+fn key_manager_client_for_replication(ctx: &RpcContext) -> RemoteClient {
     let rctx = runtime_context!(ctx, KmContext);
     let protocol = rctx.protocol.clone();
     let runtime_id = rctx.runtime_id;
@@ -442,7 +442,7 @@ fn key_manager_client_for_replication(ctx: &mut RpcContext) -> RemoteClient {
 
 /// Create init response and sign it with RAK.
 fn sign_init_response(
-    ctx: &mut RpcContext,
+    ctx: &RpcContext,
     state: State,
     policy_checksum: Vec<u8>,
 ) -> Result<SignedInitResponse> {

runtime/src/enclave_rpc/dispatcher.rs

@@ -47,16 +47,16 @@ pub struct MethodDescriptor {
 /// Handler for a RPC method.
 pub trait MethodHandler<Rq, Rsp> {
     /// Invoke the method implementation and return a response.
-    fn handle(&self, ctx: &mut Context, request: &Rq) -> Result<Rsp>;
+    fn handle(&self, ctx: &Context, request: &Rq) -> Result<Rsp>;
 }
 
 impl<Rq, Rsp, F> MethodHandler<Rq, Rsp> for F
 where
     Rq: 'static,
     Rsp: 'static,
-    F: Fn(&mut Context, &Rq) -> Result<Rsp> + 'static,
+    F: Fn(&Context, &Rq) -> Result<Rsp> + 'static,
 {
-    fn handle(&self, ctx: &mut Context, request: &Rq) -> Result<Rsp> {
+    fn handle(&self, ctx: &Context, request: &Rq) -> Result<Rsp> {
         (*self)(ctx, request)
     }
 }
@@ -67,7 +67,7 @@ pub trait MethodHandlerDispatch {
     fn get_descriptor(&self) -> &MethodDescriptor;
 
     /// Dispatch request.
-    fn dispatch(&self, ctx: &mut Context, request: Request) -> Result<Response>;
+    fn dispatch(&self, ctx: &Context, request: Request) -> Result<Response>;
 }
 
 struct MethodHandlerDispatchImpl<Rq, Rsp> {
@@ -86,7 +86,7 @@ where
         &self.descriptor
     }
 
-    fn dispatch(&self, ctx: &mut Context, request: Request) -> Result<Response> {
+    fn dispatch(&self, ctx: &Context, request: Request) -> Result<Response> {
         let request = cbor::from_value(request.args)?;
         let response = self.handler.handle(ctx, &request)?;
 

runtime/src/macros.rs

@@ -13,7 +13,7 @@
 macro_rules! runtime_context {
     ($ctx:ident, $type:ty) => {
         $ctx.runtime
-            .downcast_mut::<$type>()
+            .downcast_ref::<$type>()
             .expect("invalid runtime context")
     };
 }