From f8ece97211408d8268b60b33fe89b8744f2318e5 Mon Sep 17 00:00:00 2001
From: peartes <kenny.fale.kf@gmail.com>
Date: Mon, 18 Nov 2024 16:10:43 +0100
Subject: [PATCH] chore: added in unite test for revoke public key

---
 x/dkim/keeper/msg_server.go      |   6 +-
 x/dkim/keeper/msg_server_test.go | 120 +++++++++++++++++++++++++++++++
 2 files changed, 125 insertions(+), 1 deletion(-)

diff --git a/x/dkim/keeper/msg_server.go b/x/dkim/keeper/msg_server.go
index 80de419d..d2aba244 100644
--- a/x/dkim/keeper/msg_server.go
+++ b/x/dkim/keeper/msg_server.go
@@ -2,12 +2,13 @@ package keeper
 
 import (
 	"context"
-	"cosmossdk.io/errors"
 	"crypto/rsa"
 	"crypto/x509"
 	"encoding/base64"
 	"encoding/pem"
 
+	"cosmossdk.io/errors"
+
 	govtypes "github.com/cosmos/cosmos-sdk/x/gov/types"
 
 	dkimv1 "github.com/burnt-labs/xion/api/xion/dkim/v1"
@@ -96,6 +97,9 @@ func (ms msgServer) RevokeDkimPubKey(ctx context.Context, msg *types.MsgRevokeDk
 
 	var privateKey *rsa.PrivateKey
 	d, _ := pem.Decode(msg.PrivKey)
+	if d == nil {
+		return nil, errors.Wrap(types.ErrParsingPrivKey, "failed to decode pem private key")
+	}
 	if key, err := x509.ParsePKCS1PrivateKey(d.Bytes); err != nil {
 		if key, err := x509.ParsePKCS8PrivateKey(d.Bytes); err != nil {
 			return nil, errors.Wrap(types.ErrParsingPrivKey, "failed to decode private key")
diff --git a/x/dkim/keeper/msg_server_test.go b/x/dkim/keeper/msg_server_test.go
index ee1b3977..a64ed52e 100644
--- a/x/dkim/keeper/msg_server_test.go
+++ b/x/dkim/keeper/msg_server_test.go
@@ -1,6 +1,11 @@
 package keeper_test
 
 import (
+	"crypto/rand"
+	"crypto/rsa"
+	"crypto/x509"
+	"encoding/base64"
+	"encoding/pem"
 	"testing"
 
 	"github.com/stretchr/testify/require"
@@ -239,3 +244,118 @@ func TestRemoveDkimPubKey(t *testing.T) {
 		})
 	}
 }
+
+func TestRevokeDkimPubKey(t *testing.T) {
+	f := SetupTest(t)
+	// Generate a test RSA private key
+	privateKey, err := rsa.GenerateKey(rand.Reader, 2048)
+	require.NoError(t, err)
+
+	// Encode private key as PEM
+	privKeyPEM := pem.EncodeToMemory(
+		&pem.Block{
+			Type:  "RSA PRIVATE KEY",
+			Bytes: x509.MarshalPKCS1PrivateKey(privateKey),
+		},
+	)
+
+	// Generate corresponding Base64 public key
+	pubKey_1 := base64.StdEncoding.EncodeToString(privateKey.PublicKey.N.Bytes())
+	domain_1 := "x.com"
+
+	privateKey_2, err := rsa.GenerateKey(rand.Reader, 2048)
+	require.NoError(t, err)
+	privKeyPEM_2 := pem.EncodeToMemory(
+		&pem.Block{
+			Type:  "RSA PRIVATE KEY",
+			Bytes: x509.MarshalPKCS1PrivateKey(privateKey_2),
+		},
+	)
+	pubKey_2 := base64.StdEncoding.EncodeToString(privateKey_2.PublicKey.N.Bytes())
+	domain_2 := "y.com"
+
+	// Add in a DKIM public key
+	_, err = f.msgServer.AddDkimPubKeys(f.ctx, &types.MsgAddDkimPubKeys{
+		Authority: f.govModAddr,
+		DkimPubkeys: []types.DkimPubKey{
+			{
+				Domain:   domain_1,
+				PubKey:   pubKey_1,
+				Selector: "dkim-202308",
+			},
+			{
+				Domain:   domain_1,
+				PubKey:   pubKey_2,
+				Selector: "dkim-202310",
+			},
+			{
+				Domain:   domain_2,
+				PubKey:   pubKey_1,
+				Selector: "dkim-202310",
+			},
+			{
+				Domain:   domain_2,
+				PubKey:   pubKey_2,
+				Selector: "dkim-202311",
+			},
+		},
+	})
+	require.NoError(t, err)
+
+	// Define test cases
+	tests := []struct {
+		name           string
+		msg            *types.MsgRevokeDkimPubKey
+		mockError      error
+		expectedError  error
+		expectedLength int8
+	}{
+		{
+			name: "invalid private key",
+			msg: &types.MsgRevokeDkimPubKey{
+				Domain:  domain_1,
+				PrivKey: []byte("invalid_key"),
+			},
+			mockError:     nil,
+			expectedError: types.ErrParsingPrivKey,
+		},
+		{
+			name: "successfully revoke 1 of domain 1 DKIM public key",
+			msg: &types.MsgRevokeDkimPubKey{
+				Domain:  domain_1,
+				PrivKey: privKeyPEM,
+			},
+			mockError:      nil,
+			expectedError:  nil,
+			expectedLength: 1, // domain_1 has 1 keys, 1 with the matching pub key should be revoked, 1 with a different pubkey should be left
+		},
+		{
+			name: "successfully revoke 1 of domain 2 DKIM public key",
+			msg: &types.MsgRevokeDkimPubKey{
+				Domain:  domain_2,
+				PrivKey: privKeyPEM_2,
+			},
+			mockError:      nil,
+			expectedError:  nil,
+			expectedLength: 1, // domain_2 has 1 keys, 1 with the matching pub key should be revoked, 1 with a different pubkey should be left
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			ctx := f.ctx
+
+			// Call the RevokeDkimPubKey method
+			_, err := f.msgServer.RevokeDkimPubKey(ctx, tt.msg)
+
+			// Validate results
+			if tt.expectedError != nil {
+				require.ErrorIs(t, err, tt.expectedError)
+			} else {
+				res, err := f.queryServer.DkimPubKeys(ctx, &types.QueryDkimPubKeysRequest{Domain: tt.msg.Domain})
+				require.NoError(t, err)
+				require.Len(t, res.DkimPubKeys, int(tt.expectedLength))
+			}
+		})
+	}
+}