Add Verify function

Author: iKapitonau

guest-agent/rpc/proto/agent_rpc.proto

@@ -49,8 +49,11 @@ service DstackGuest {
   // Get app info
   rpc Info(google.protobuf.Empty) returns (AppInfo) {}
 
-  // Get app info
+  // Sign a payload
   rpc Sign(SignRequest) returns (SignResponse) {}
+
+  // Verify a signature
+  rpc Verify(VerifyRequest) returns (VerifyResponse) {}
 }
 
 // The request to derive a key
@@ -234,6 +237,17 @@ message SignResponse {
   bytes signature = 1;
 }
 
+message VerifyRequest {
+  string algorithm = 1;
+  bytes data = 2;
+  bytes signature = 3;
+  bytes public_key = 4;
+}
+
+message VerifyResponse {
+  bool valid = 1;
+}
+
 message GetAttestationForAppKeyRequest {
   string algorithm = 1;
 }

guest-agent/src/rpc_service.rs

@@ -13,11 +13,14 @@ use dstack_guest_agent_rpc::{
     worker_server::{WorkerRpc, WorkerServer},
     AppInfo, DeriveK256KeyResponse, DeriveKeyArgs, EmitEventArgs, GetAttestationForAppKeyRequest,
     GetKeyArgs, GetKeyResponse, GetQuoteResponse, GetTlsKeyArgs, GetTlsKeyResponse, RawQuoteArgs,
-    SignRequest, SignResponse, TdxQuoteArgs, TdxQuoteResponse, WorkerVersion,
+    SignRequest, SignResponse, TdxQuoteArgs, TdxQuoteResponse, VerifyRequest, VerifyResponse,
+    WorkerVersion,
 };
 use dstack_types::{AppKeys, SysConfig};
 use ed25519_dalek::ed25519::signature::hazmat::PrehashSigner;
-use ed25519_dalek::{Signer as Ed25519Signer, SigningKey as Ed25519SigningKey};
+use ed25519_dalek::{
+    Signer as Ed25519Signer, SigningKey as Ed25519SigningKey, Verifier as Ed25519Verifier,
+};
 use fs_err as fs;
 use k256::ecdsa::SigningKey;
 use ra_rpc::{Attestation, CallContext, RpcCall};
@@ -305,6 +308,56 @@ impl DstackGuestRpc for InternalRpcHandler {
         };
         Ok(SignResponse { signature })
     }
+
+    async fn verify(self, request: VerifyRequest) -> Result<VerifyResponse> {
+        let valid = match request.algorithm.as_str() {
+            "ed25519" => {
+                let public_key = if request.public_key.is_empty() {
+                    let key_response = self
+                        .get_key(GetKeyArgs {
+                            path: "vms".to_string(),
+                            purpose: "signing".to_string(),
+                            algorithm: "ed25519".to_string(),
+                        })
+                        .await?;
+                    let key_bytes: [u8; 32] =
+                        key_response.key.try_into().expect("Key is incorrect");
+                    Ed25519SigningKey::from_bytes(&key_bytes)
+                        .verifying_key()
+                        .to_bytes()
+                        .to_vec()
+                } else {
+                    request.public_key
+                };
+                let verifying_key = ed25519_dalek::VerifyingKey::from_bytes(
+                    &public_key.as_slice().try_into().unwrap(),
+                )?;
+                let signature = ed25519_dalek::Signature::from_slice(&request.signature)?;
+                verifying_key.verify(&request.data, &signature).is_ok()
+            }
+            "secp256k1" | "secp256k1_prehashed" => {
+                let public_key = if request.public_key.is_empty() {
+                    let key_response = self
+                        .get_key(GetKeyArgs {
+                            path: "vms".to_string(),
+                            purpose: "signing".to_string(),
+                            algorithm: request.algorithm,
+                        })
+                        .await?;
+                    let signing_key = SigningKey::from_slice(&key_response.key)
+                        .context("Failed to parse secp256k1 key")?;
+                    signing_key.verifying_key().to_sec1_bytes().to_vec()
+                } else {
+                    request.public_key
+                };
+                let verifying_key = k256::ecdsa::VerifyingKey::from_sec1_bytes(&public_key)?;
+                let signature = k256::ecdsa::Signature::from_slice(&request.signature)?;
+                verifying_key.verify(&request.data, &signature).is_ok()
+            }
+            _ => return Err(anyhow::anyhow!("Unsupported algorithm")),
+        };
+        Ok(VerifyResponse { valid })
+    }
 }
 
 fn simulate_quote(
@@ -746,6 +799,60 @@ pNs85uhOZE8z2jr8Pg==
         }
     }
 
+    #[tokio::test]
+    async fn test_verify_ed25519_success() {
+        let state = setup_test_state().await;
+        let handler = InternalRpcHandler {
+            state: state.clone(),
+        };
+        let data_to_sign = b"test message for ed25519";
+        let sign_request = SignRequest {
+            algorithm: "ed25519".to_string(),
+            data: data_to_sign.to_vec(),
+        };
+
+        let sign_response = handler.sign(sign_request).await.unwrap();
+
+        let verify_request = VerifyRequest {
+            algorithm: "ed25519".to_string(),
+            data: data_to_sign.to_vec(),
+            signature: sign_response.signature,
+            public_key: vec![],
+        };
+        let handler = InternalRpcHandler {
+            state: state.clone(),
+        };
+        let verify_response = handler.verify(verify_request).await.unwrap();
+        assert!(verify_response.valid);
+    }
+
+    #[tokio::test]
+    async fn test_verify_secp256k1_success() {
+        let state = setup_test_state().await;
+        let handler = InternalRpcHandler {
+            state: state.clone(),
+        };
+        let data_to_sign = b"test message for secp256k1";
+        let sign_request = SignRequest {
+            algorithm: "secp256k1".to_string(),
+            data: data_to_sign.to_vec(),
+        };
+
+        let sign_response = handler.sign(sign_request).await.unwrap();
+
+        let verify_request = VerifyRequest {
+            algorithm: "secp256k1".to_string(),
+            data: data_to_sign.to_vec(),
+            signature: sign_response.signature,
+            public_key: vec![],
+        };
+        let handler = InternalRpcHandler {
+            state: state.clone(),
+        };
+        let verify_response = handler.verify(verify_request).await.unwrap();
+        assert!(verify_response.valid);
+    }
+
     #[tokio::test]
     async fn test_sign_ed25519_success() {
         let state = setup_test_state().await;