Add support for Sign/Verify to SDK

Author: iKapitonau

sdk/curl/api.md

@@ -71,6 +71,7 @@ Generates an ECDSA key using the k256 elliptic curve, derived from the applicati
 |-------|------|-------------|----------|
 | `path` | string | Path for the key | `"my/key/path"` |
 | `purpose` | string | Purpose for the key. Can be any string. This is used in the signature chain. | `"signing"` | `"encryption"` |
+| `algorithm` | string | Either `secp256k1` or `ed25519`. Defaults to `secp256k1` | `ed25519` |
 
 **Example:**
 ```bash
@@ -79,14 +80,15 @@ curl --unix-socket /var/run/dstack.sock -X POST \
   -H 'Content-Type: application/json' \
   -d '{
     "path": "my/key/path",
-    "purpose": "signing"
+    "purpose": "signing",
+    "algorithm": "ed25519",
   }'
 ```
 
 Or
 
 ```bash
-curl --unix-socket /var/run/dstack.sock http://dstack/GetKey?path=my/key/path&purpose=signing
+curl --unix-socket /var/run/dstack.sock http://dstack/GetKey?path=my/key/path&purpose=signing&algorithm=ed25519
 ```
 
 **Response:**
@@ -189,6 +191,80 @@ curl --unix-socket /var/run/dstack.sock -X POST \
 **Response:**
 Empty response with HTTP 200 status code on success.
 
+### 6. Sign
+
+Signs a payload.
+
+**Endpoint:** `/Sign`
+
+**Request Parameters:**
+
+| Field | Type | Description | Example |
+|-------|------|-------------|----------|
+| `algorithm` | string | `ed25519`, `secp256k1_prehashed` or `secp256k1`| `ed25519` |
+| `data` | string | Hex-encoded payload data | `deadbeef` |
+
+**Example:**
+```bash
+curl --unix-socket /var/run/dstack.sock -X POST \
+  http://dstack/Sign \
+  -H 'Content-Type: application/json' \
+  -d '{
+    "algorithm": "ed25519",
+    "data": "deadbeef"
+  }'
+```
+
+**Response:**
+```json
+{
+  "signature": "<hex-encoded-signature>",
+  "signature_chain": [
+    "<hex-encoded-signature-1>",
+    "<hex-encoded-signature-2>",
+    "<hex-encoded-signature-3>"
+  ]
+  "public_key": "<hex-encoded-public-key>"
+}
+```
+
+### 7. Verify
+
+Verifies a signature.
+
+**Endpoint:** `/Verify`
+
+**Request Parameters:**
+
+| Field | Type | Description | Example |
+|-------|------|-------------|----------|
+| `algorithm` | string | `ed25519`, `secp256k1_prehashed` or `secp256k1`| `ed25519` |
+| `data` | string | Hex-encoded payload data | `deadbeef` |
+| `signature` | string | Hex-encoded signature | `deadbeef` |
+| `public_key` | string | Hex-encoded public key | `deadbeef` |
+
+**Example:**
+```bash
+curl --unix-socket /var/run/dstack.sock -X POST \
+  http://dstack/Verify \
+  -H 'Content-Type: application/json' \
+  -d '{
+    "algorithm": "ed25519",
+    "data": "deadbeef",
+    "signature": "deadbeef",
+    "public_key": "deadbeef"
+  }'
+```
+
+**Response:**
+```json
+{
+  "valid": "<true|false>"
+}
+```
+
+```
+
 ## Error Responses
 
 All endpoints may return the following HTTP status codes:

sdk/go/README.md

@@ -91,9 +91,11 @@ NOTE: Leave endpoint empty in production. You only need to add `volumes` in your
 #### Methods
 
 - `Info(ctx context.Context) (*InfoResponse, error)`: Retrieves information about the CVM instance.
-- `GetKey(ctx context.Context, path string, purpose string) (*GetKeyResponse, error)`: Derives a key for the given path and purpose.
+- `GetKey(ctx context.Context, path string, purpose string, algorithm string) (*GetKeyResponse, error)`: Derives a key for the given path, purpose and algorithm.
 - `GetQuote(ctx context.Context, reportData []byte) (*GetQuoteResponse, error)`: Generates a TDX quote using SHA512 as the hash algorithm.
 - `GetTlsKey(ctx context.Context, path string, subject string, altNames []string, usageRaTls bool, usageServerAuth bool, usageClientAuth bool, randomSeed bool) (*GetTlsKeyResponse, error)`: Derives a key for the given path and purpose.
+- `Sign(ctx context.Context, algorithm string, data []byte) (*SignResponse, error)`: Signs a payload
+- `Verify(ctx context.Context, algorithm string, data []byte, signature []byte, public_key []byte) (*VerifyResponse, error)`: Verifies a payload
 
 ## Development
 

sdk/go/dstack/client.go

@@ -348,10 +348,11 @@ func (c *DstackClient) GetTlsKey(
 }
 
 // Gets a key from the dstack service.
-func (c *DstackClient) GetKey(ctx context.Context, path string, purpose string) (*GetKeyResponse, error) {
+func (c *DstackClient) GetKey(ctx context.Context, path string, purpose string, algorithm string) (*GetKeyResponse, error) {
 	payload := map[string]interface{}{
-		"path":    path,
-		"purpose": purpose,
+		"path":      path,
+		"purpose":   purpose,
+		"algorithm": algorithm,
 	}
 
 	data, err := c.sendRPCRequest(ctx, "/GetKey", payload)
@@ -422,6 +423,83 @@ func (c *DstackClient) Info(ctx context.Context) (*InfoResponse, error) {
 	return &response, nil
 }
 
+type SignResponse struct {
+	Signature      []byte
+	SignatureChain [][]byte
+	PublicKey      []byte
+}
+
+// Signs a payload.
+func (c *DstackClient) Sign(ctx context.Context, algorithm string, data []byte) (*SignResponse, error) {
+	payload := map[string]interface{}{
+		"algorithm": algorithm,
+		"data":      hex.EncodeToString(data),
+	}
+
+	respData, err := c.sendRPCRequest(ctx, "/Sign", payload)
+	if err != nil {
+		return nil, err
+	}
+
+	var response struct {
+		Signature      string   `json:"signature"`
+		SignatureChain []string `json:"signature_chain"`
+		PublicKey      string   `json:"public_key"`
+	}
+	if err := json.Unmarshal(respData, &response); err != nil {
+		return nil, fmt.Errorf("failed to unmarshal sign response: %w", err)
+	}
+
+	sig, err := hex.DecodeString(response.Signature)
+	if err != nil {
+		return nil, fmt.Errorf("failed to decode signature: %w", err)
+	}
+	pubKey, err := hex.DecodeString(response.PublicKey)
+	if err != nil {
+		return nil, fmt.Errorf("failed to decode public key: %w", err)
+	}
+
+	sigChain := make([][]byte, len(response.SignatureChain))
+	for i, s := range response.SignatureChain {
+		sigChain[i], err = hex.DecodeString(s)
+		if err != nil {
+			return nil, fmt.Errorf("failed to decode signature chain element %d: %w", i, err)
+		}
+	}
+
+	return &SignResponse{
+		Signature:      sig,
+		SignatureChain: sigChain,
+		PublicKey:      pubKey,
+	}, nil
+}
+
+type VerifyResponse struct {
+	Valid bool `json:"valid"`
+}
+
+// Verifies a payload.
+func (c *DstackClient) Verify(ctx context.Context, algorithm string, data []byte, signature []byte, publicKey []byte) (*VerifyResponse, error) {
+	payload := map[string]interface{}{
+		"algorithm":  algorithm,
+		"data":       hex.EncodeToString(data),
+		"signature":  hex.EncodeToString(signature),
+		"public_key": hex.EncodeToString(publicKey),
+	}
+
+	respData, err := c.sendRPCRequest(ctx, "/Verify", payload)
+	if err != nil {
+		return nil, err
+	}
+
+	var response VerifyResponse
+	if err := json.Unmarshal(respData, &response); err != nil {
+		return nil, fmt.Errorf("failed to unmarshal verify response: %w", err)
+	}
+
+	return &response, nil
+}
+
 // EmitEvent sends an event to be extended to RTMR3 on TDX platform.
 // The event will be extended to RTMR3 with the provided name and payload.
 //

sdk/go/dstack/client_test.go

@@ -7,6 +7,7 @@ package dstack_test
 import (
 	"bytes"
 	"context"
+	"crypto/sha256"
 	"encoding/hex"
 	"encoding/json"
 	"fmt"
@@ -25,7 +26,7 @@ import (
 
 func TestGetKey(t *testing.T) {
 	client := dstack.NewDstackClient()
-	resp, err := client.GetKey(context.Background(), "/", "test")
+	resp, err := client.GetKey(context.Background(), "/", "test", "ed25519")
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -434,7 +435,7 @@ func TestGetKeySignatureVerification(t *testing.T) {
 	client := dstack.NewDstackClient()
 	path := "/test/path"
 	purpose := "test-purpose"
-	resp, err := client.GetKey(context.Background(), path, purpose)
+	resp, err := client.GetKey(context.Background(), path, purpose, "secp256k1")
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -608,3 +609,111 @@ func compressPublicKey(uncompressedKey []byte) ([]byte, error) {
 	}
 	return crypto.CompressPubkey(pubKey), nil
 }
+
+func TestSignAndVerifyEd25519(t *testing.T) {
+	client := dstack.NewDstackClient()
+	dataToSign := []byte("test message for ed25519")
+	algorithm := "ed25519"
+
+	signResp, err := client.Sign(context.Background(), algorithm, dataToSign)
+	if err != nil {
+		t.Fatalf("Sign() error = %v", err)
+	}
+
+	if len(signResp.Signature) == 0 {
+		t.Error("expected signature to not be empty")
+	}
+	if len(signResp.PublicKey) == 0 {
+		t.Error("expected public key to not be empty")
+	}
+	if len(signResp.SignatureChain) != 3 {
+		t.Errorf("expected signature chain to have 3 elements, got %d", len(signResp.SignatureChain))
+	}
+	if !bytes.Equal(signResp.Signature, signResp.SignatureChain[0]) {
+		t.Error("expected Signature to be the same as SignatureChain[0]")
+	}
+
+	verifyResp, err := client.Verify(context.Background(), algorithm, dataToSign, signResp.Signature, signResp.PublicKey)
+	if err != nil {
+		t.Fatalf("Verify() error = %v", err)
+	}
+
+	if !verifyResp.Valid {
+		t.Error("expected verification to be valid")
+	}
+
+	badData := []byte("wrong message")
+	verifyResp, err = client.Verify(context.Background(), algorithm, badData, signResp.Signature, signResp.PublicKey)
+	if err != nil {
+		t.Fatalf("Verify() with bad data error = %v", err)
+	}
+
+	if verifyResp.Valid {
+		t.Error("expected verification with bad data to be invalid")
+	}
+}
+
+func TestSignAndVerifySecp256k1(t *testing.T) {
+	client := dstack.NewDstackClient()
+	dataToSign := []byte("test message for secp256k1")
+	algorithm := "secp256k1"
+
+	signResp, err := client.Sign(context.Background(), algorithm, dataToSign)
+	if err != nil {
+		t.Fatalf("Sign() error = %v", err)
+	}
+
+	if len(signResp.Signature) == 0 {
+		t.Error("expected signature to not be empty")
+	}
+	if len(signResp.PublicKey) == 0 {
+		t.Error("expected public key to not be empty")
+	}
+	if len(signResp.SignatureChain) != 3 {
+		t.Errorf("expected signature chain to have 3 elements, got %d", len(signResp.SignatureChain))
+	}
+
+	verifyResp, err := client.Verify(context.Background(), algorithm, dataToSign, signResp.Signature, signResp.PublicKey)
+	if err != nil {
+		t.Fatalf("Verify() error = %v", err)
+	}
+
+	if !verifyResp.Valid {
+		t.Error("expected verification to be valid")
+	}
+}
+
+func TestSignAndVerifySecp256k1Prehashed(t *testing.T) {
+	client := dstack.NewDstackClient()
+	dataToSign := []byte("test message for secp256k1 prehashed")
+	digest := sha256.Sum256(dataToSign)
+	algorithm := "secp256k1_prehashed"
+
+	signResp, err := client.Sign(context.Background(), algorithm, digest[:])
+	if err != nil {
+		t.Fatalf("Sign() error = %v", err)
+	}
+
+	if len(signResp.Signature) == 0 {
+		t.Error("expected signature to not be empty")
+	}
+
+	verifyResp, err := client.Verify(context.Background(), algorithm, digest[:], signResp.Signature, signResp.PublicKey)
+	if err != nil {
+		t.Fatalf("Verify() error = %v", err)
+	}
+
+	if !verifyResp.Valid {
+		t.Error("expected verification to be valid")
+	}
+
+	// Test invalid digest length for signing
+	invalidDigest := []byte{1, 2, 3}
+	_, err = client.Sign(context.Background(), algorithm, invalidDigest)
+	if err == nil {
+		t.Fatal("expected error for invalid digest length, got nil")
+	}
+	if !strings.Contains(err.Error(), "32-byte digest") {
+		t.Errorf("expected error to mention '32-byte digest', got: %v", err)
+	}
+}