Add signature verification in go tests

Author: kvinwang

sdk/go/dstack/client_test.go

@@ -6,13 +6,17 @@ import (
 	"encoding/hex"
 	"encoding/json"
 	"fmt"
+	"math/big"
 	"strings"
 	"testing"
 
 	"crypto/x509"
 	"encoding/pem"
 
+	"crypto/ecdsa"
+
 	"github.com/Dstack-TEE/dstack/sdk/go/dstack"
+	"github.com/ethereum/go-ethereum/crypto"
 )
 
 func TestGetKey(t *testing.T) {
@@ -406,3 +410,185 @@ func TestInfo(t *testing.T) {
 		t.Error("expected event log to not be empty")
 	}
 }
+
+func TestGetKeySignatureVerification(t *testing.T) {
+	expectedAppPubkey, _ := hex.DecodeString("02b85cceca0c02d878f0ebcda72a97469a472416eb6faf3c4807642132f9786810")
+	expectedKmsPubkey, _ := hex.DecodeString("02cad3a8bb11c5c0858fb3e402048b5137457039d577986daade678ed4b4ab1b9b")
+
+	client := dstack.NewDstackClient()
+	path := "/test/path"
+	purpose := "test-purpose"
+	resp, err := client.GetKey(context.Background(), path, purpose)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if resp.Key == "" {
+		t.Error("expected key to not be empty")
+	}
+
+	if len(resp.SignatureChain) != 2 {
+		t.Fatalf("expected signature chain to have 2 elements, got %d", len(resp.SignatureChain))
+	}
+
+	// Extract the app signature and KMS signature from the chain
+	appSignatureHex := resp.SignatureChain[0]
+	kmsSignatureHex := resp.SignatureChain[1]
+
+	// Convert hex strings to bytes
+	appSignature, err := hex.DecodeString(appSignatureHex)
+	if err != nil {
+		t.Fatalf("failed to decode app signature: %v", err)
+	}
+
+	kmsSignature, err := hex.DecodeString(kmsSignatureHex)
+	if err != nil {
+		t.Fatalf("failed to decode KMS signature: %v", err)
+	}
+
+	// Verify signatures have the correct format (signature + recovery ID)
+	if len(appSignature) != 65 {
+		t.Errorf("expected app signature to be 65 bytes (64 bytes signature + 1 byte recovery ID), got %d", len(appSignature))
+	}
+
+	if len(kmsSignature) != 65 {
+		t.Errorf("expected KMS signature to be 65 bytes (64 bytes signature + 1 byte recovery ID), got %d", len(kmsSignature))
+	}
+
+	// Get app info to retrieve app ID for verification
+	infoResp, err := client.Info(context.Background())
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// 1. Derive the public key from the private key
+	derivedPrivKey := resp.Key
+	derivedPubKey, err := derivePublicKeyFromPrivate(derivedPrivKey)
+	if err != nil {
+		t.Fatalf("failed to derive public key: %v", err)
+	}
+
+	// 2. Construct the message that was signed
+	message := fmt.Sprintf("%s:%s", purpose, hex.EncodeToString(derivedPubKey))
+
+	// 3. Recover the app's public key from the signature
+	appPubKey, err := recoverPublicKey(message, appSignature)
+	if err != nil {
+		t.Fatalf("failed to recover app public key: %v", err)
+	}
+
+	// Convert the recovered public key to compressed format for comparison
+	appPubKeyCompressed, err := compressPublicKey(appPubKey)
+	if err != nil {
+		t.Fatalf("failed to compress recovered public key: %v", err)
+	}
+
+	if !bytes.Equal(appPubKeyCompressed, expectedAppPubkey) {
+		t.Errorf("app public key mismatch:\nExpected: %s\nActual:   %s",
+			hex.EncodeToString(expectedAppPubkey),
+			hex.EncodeToString(appPubKeyCompressed))
+	}
+
+	// 4. Verify the app ID matches what we expect
+	// The app ID should be derivable from the app's public key
+	// or should match what's returned from the Info endpoint
+	appIDFromInfo, err := hex.DecodeString(infoResp.AppID)
+	if err != nil {
+		t.Fatalf("failed to decode app ID: %v", err)
+	}
+
+	// 5. Construct the message that KMS would have signed
+	// This would typically be something like "dstack-kms-issued:{app_id}{app_public_key}"
+	kmsMessage := fmt.Sprintf("dstack-kms-issued:%s%s", appIDFromInfo, string(appPubKeyCompressed))
+	kmsPubKey, err := recoverPublicKey(kmsMessage, kmsSignature)
+	if err != nil {
+		t.Fatalf("failed to recover KMS public key: %v", err)
+	}
+
+	kmsPubKeyCompressed, err := compressPublicKey(kmsPubKey)
+	if err != nil {
+		t.Fatalf("failed to compress KMS public key: %v", err)
+	}
+
+	if !bytes.Equal(kmsPubKeyCompressed, expectedKmsPubkey) {
+		t.Errorf("KMS public key mismatch:\nExpected: %s\nActual:   %s",
+			hex.EncodeToString(expectedKmsPubkey),
+			hex.EncodeToString(kmsPubKeyCompressed))
+	}
+
+	// Verify that the recovered app public key can verify the app signature
+	verified, err := verifySignature(message, appSignature, appPubKey)
+	if err != nil {
+		t.Fatalf("signature verification error: %v", err)
+	}
+	if !verified {
+		t.Error("app signature verification failed")
+	}
+}
+
+// Helper function to derive a public key from a private key
+func derivePublicKeyFromPrivate(privateKeyHex string) ([]byte, error) {
+	privateKeyBytes, err := hex.DecodeString(privateKeyHex)
+	if err != nil {
+		return nil, fmt.Errorf("failed to decode private key: %w", err)
+	}
+
+	// Import the private key
+	privateKey, err := crypto.ToECDSA(privateKeyBytes)
+	if err != nil {
+		return nil, fmt.Errorf("failed to convert to ECDSA private key: %w", err)
+	}
+
+	// Derive the public key in compressed format
+	publicKey := crypto.CompressPubkey(&privateKey.PublicKey)
+	return publicKey, nil
+}
+
+// Helper function to recover a public key from a signature
+func recoverPublicKey(message string, signature []byte) ([]byte, error) {
+	if len(signature) != 65 {
+		return nil, fmt.Errorf("invalid signature length: expected 65 bytes, got %d", len(signature))
+	}
+
+	// Hash the message using Keccak256
+	messageHash := crypto.Keccak256([]byte(message))
+
+	// Recover the public key
+	pubKey, err := crypto.Ecrecover(messageHash, signature)
+	if err != nil {
+		return nil, fmt.Errorf("failed to recover public key: %w", err)
+	}
+
+	return pubKey, nil
+}
+
+// Helper function to verify a signature
+func verifySignature(message string, signature []byte, publicKey []byte) (bool, error) {
+	if len(signature) != 65 {
+		return false, fmt.Errorf("invalid signature length: expected 65 bytes, got %d", len(signature))
+	}
+
+	// Hash the message using Keccak256
+	messageHash := crypto.Keccak256([]byte(message))
+
+	// The last byte is the recovery ID, we need to remove it for verification
+	signatureWithoutRecoveryID := signature[:64]
+
+	// Verify the signature
+	return crypto.VerifySignature(publicKey, messageHash, signatureWithoutRecoveryID), nil
+}
+
+// Add this helper function to compress a public key
+func compressPublicKey(uncompressedKey []byte) ([]byte, error) {
+	if len(uncompressedKey) < 65 || uncompressedKey[0] != 4 {
+		return nil, fmt.Errorf("invalid uncompressed public key")
+	}
+	x := new(big.Int).SetBytes(uncompressedKey[1:33])
+	y := new(big.Int).SetBytes(uncompressedKey[33:65])
+	pubKey := &ecdsa.PublicKey{
+		Curve: crypto.S256(),
+		X:     x,
+		Y:     y,
+	}
+	return crypto.CompressPubkey(pubKey), nil
+}

sdk/go/go.mod

@@ -1,3 +1,13 @@
 module github.com/Dstack-TEE/dstack/sdk/go
 
-go 1.21.6
+go 1.23.0
+
+toolchain go1.23.8
+
+require (
+	github.com/decred/dcrd/dcrec/secp256k1/v4 v4.0.1 // indirect
+	github.com/ethereum/go-ethereum v1.15.7 // indirect
+	github.com/holiman/uint256 v1.3.2 // indirect
+	golang.org/x/crypto v0.35.0 // indirect
+	golang.org/x/sys v0.30.0 // indirect
+)

sdk/go/go.sum

@@ -0,0 +1,11 @@
+github.com/decred/dcrd/crypto/blake256 v1.0.0/go.mod h1:sQl2p6Y26YV+ZOcSTP6thNdn47hh8kt6rqSlvmrXFAc=
+github.com/decred/dcrd/dcrec/secp256k1/v4 v4.0.1 h1:YLtO71vCjJRCBcrPMtQ9nqBsqpA1m5sE92cU+pd5Mcc=
+github.com/decred/dcrd/dcrec/secp256k1/v4 v4.0.1/go.mod h1:hyedUtir6IdtD/7lIxGeCxkaw7y45JueMRL4DIyJDKs=
+github.com/ethereum/go-ethereum v1.15.7 h1:vm1XXruZVnqtODBgqFaTclzP0xAvCvQIDKyFNUA1JpY=
+github.com/ethereum/go-ethereum v1.15.7/go.mod h1:+S9k+jFzlyVTNcYGvqFhzN/SFhI6vA+aOY4T5tLSPL0=
+github.com/holiman/uint256 v1.3.2 h1:a9EgMPSC1AAaj1SZL5zIQD3WbwTuHrMGOerLjGmM/TA=
+github.com/holiman/uint256 v1.3.2/go.mod h1:EOMSn4q6Nyt9P6efbI3bueV4e1b3dGlUCXeiRV4ng7E=
+golang.org/x/crypto v0.35.0 h1:b15kiHdrGCHrP6LvwaQ3c03kgNhhiMgvlhxHQhmg2Xs=
+golang.org/x/crypto v0.35.0/go.mod h1:dy7dXNW32cAb/6/PRuTNsix8T+vJAqvuIy5Bli/x0YQ=
+golang.org/x/sys v0.30.0 h1:QjkSwP/36a20jFYWkSue1YwXzLmsV5Gfq7Eiy72C1uc=
+golang.org/x/sys v0.30.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=