identities: refactor is_malicious check to read minimal amount of data

activation/handler_v1.go

@@ -286,14 +286,9 @@ func (h *HandlerV1) validateNonInitialAtx(
 
 // cacheAtx caches the atx in the atxsdata cache.
 // Returns true if the atx was cached, false otherwise.
-func (h *HandlerV1) cacheAtx(ctx context.Context, atx *types.ActivationTx) *atxsdata.ATX {
+func (h *HandlerV1) cacheAtx(ctx context.Context, atx *types.ActivationTx, isMalicious bool) *atxsdata.ATX {
 	if !h.atxsdata.IsEvicted(atx.TargetEpoch()) {
-		malicious, err := identities.IsMalicious(h.cdb, atx.SmesherID)
-		if err != nil {
-			h.logger.Error("failed is malicious read", zap.Error(err), log.ZContext(ctx))
-			return nil
-		}
-		return h.atxsdata.AddFromAtx(atx, malicious)
+		return h.atxsdata.AddFromAtx(atx, isMalicious)
 	}
 	return nil
 }
@@ -450,19 +445,23 @@ func (h *HandlerV1) checkMalicious(
 	ctx context.Context,
 	tx sql.Transaction,
 	watx *wire.ActivationTxV1,
-) (*mwire.MalfeasanceProof, error) {
+) (bool, *mwire.MalfeasanceProof, error) {
 	malicious, err := identities.IsMalicious(tx, watx.SmesherID)
 	if err != nil {
-		return nil, fmt.Errorf("checking if node is malicious: %w", err)
+		return false, nil, fmt.Errorf("checking if node is malicious: %w", err)
 	}
 	if malicious {
-		return nil, nil
+		return true, nil, nil
 	}
 	proof, err := h.checkDoublePublish(ctx, tx, watx)
 	if proof != nil || err != nil {
-		return proof, err
+		return true, proof, err
 	}
-	return h.checkWrongPrevAtx(ctx, tx, watx)
+	proof, err = h.checkWrongPrevAtx(ctx, tx, watx)
+	if err != nil {
+		return false, nil, err
+	}
+	return proof != nil, proof, nil
 }
 
 // storeAtx stores an ATX and notifies subscribers of the ATXID.
@@ -471,10 +470,13 @@ func (h *HandlerV1) storeAtx(
 	atx *types.ActivationTx,
 	watx *wire.ActivationTxV1,
 ) (*mwire.MalfeasanceProof, error) {
-	var proof *mwire.MalfeasanceProof
+	var (
+		isMalicious bool
+		proof       *mwire.MalfeasanceProof
+	)
 	if err := h.cdb.WithTx(ctx, func(tx sql.Transaction) error {
 		var err error
-		proof, err = h.checkMalicious(ctx, tx, watx)
+		isMalicious, proof, err = h.checkMalicious(ctx, tx, watx)
 		if err != nil {
 			return fmt.Errorf("check malicious: %w", err)
 		}
@@ -499,7 +501,7 @@ func (h *HandlerV1) storeAtx(
 		h.tortoise.OnMalfeasance(atx.SmesherID)
 	}
 
-	added := h.cacheAtx(ctx, atx)
+	added := h.cacheAtx(ctx, atx, isMalicious)
 	h.beacon.OnAtx(atx)
 	if added != nil {
 		h.tortoise.OnAtx(atx.TargetEpoch(), atx.ID(), added)

activation/handler_v2.go

@@ -682,6 +682,9 @@ func (h *HandlerV2) checkMalicious(ctx context.Context, tx sql.Transaction, atx
 	if malicious {
 		return true, nil
 	}
+	if atx.MarriageATX != nil {
+		malicious, err := identities.IsMarriageMalicious(tx, *atx.MarriageATX)
+	}
 
 	malicious, err = h.checkDoubleMarry(ctx, tx, atx)
 	if err != nil {

datastore/store.go

@@ -202,7 +202,7 @@ func (db *CachedDB) Previous(id types.ATXID) ([]types.ATXID, error) {
 func (db *CachedDB) IterateMalfeasanceProofs(
 	iter func(types.NodeID, *wire.MalfeasanceProof) error,
 ) error {
-	ids, err := identities.GetMalicious(db)
+	ids, err := identities.GetNodeIDsWithProofs(db)
 	if err != nil {
 		return err
 	}

fetch/handler.go

@@ -46,7 +46,7 @@ func newHandler(
 
 // handleMaliciousIDsReq returns the IDs of all known malicious nodes.
 func (h *handler) handleMaliciousIDsReq(ctx context.Context, _ []byte) ([]byte, error) {
-	nodes, err := identities.GetMalicious(h.cdb)
+	nodes, err := identities.GetNodeIDsWithProofs(h.cdb)
 	if err != nil {
 		return nil, fmt.Errorf("getting malicious IDs: %w", err)
 	}
@@ -63,7 +63,7 @@ func (h *handler) handleMaliciousIDsReq(ctx context.Context, _ []byte) ([]byte,
 
 func (h *handler) handleMaliciousIDsReqStream(ctx context.Context, msg []byte, s io.ReadWriter) error {
 	if err := h.streamIDs(ctx, s, func(cbk retrieveCallback) error {
-		nodeIDs, err := identities.GetMalicious(h.cdb)
+		nodeIDs, err := identities.GetNodeIDsWithProofs(h.cdb)
 		if err != nil {
 			return fmt.Errorf("getting malicious IDs: %w", err)
 		}

sql/identities/identities.go

@@ -11,11 +11,12 @@ import (
 	"github.com/spacemeshos/go-spacemesh/sql"
 )
 
-// SetMalicious records identity as malicious.
+// SetMalicious records identity as malicious and stores proof with that identity.
+// And it also updates every other identity with the same marriage atx as malicious, without duplicating the proof.
 func SetMalicious(db sql.Executor, nodeID types.NodeID, proof []byte, received time.Time) error {
-	_, err := db.Exec(`insert into identities (pubkey, proof, received)
-	values (?1, ?2, ?3)
-	ON CONFLICT(pubkey) DO UPDATE SET proof = excluded.proof
+	_, err := db.Exec(`insert into identities (pubkey, proof, received, is_malicious)
+	values (?1, ?2, ?3, true)
+	ON CONFLICT(pubkey) DO UPDATE SET proof = excluded.proof, is_malicious = true
 	WHERE proof IS NULL;`,
 		func(stmt *sql.Statement) {
 			stmt.BindBytes(1, nodeID.Bytes())
@@ -29,15 +30,26 @@ func SetMalicious(db sql.Executor, nodeID types.NodeID, proof []byte, received t
 	return nil
 }
 
+// SetMaliciousByMarriage records identity as malicious.
+func SetMaliciousByMarriage(db sql.Executor, nodeID types.NodeID) error {
+	_, err := db.Exec(`insert into identities (pubkey, is_malicious)
+	values (?1, true)
+	ON CONFLICT(pubkey) DO UPDATE SET is_malicious = true
+	WHERE is_malicious IS FALSE;`,
+		func(stmt *sql.Statement) {
+			stmt.BindBytes(1, nodeID.Bytes())
+		}, nil,
+	)
+	if err != nil {
+		return fmt.Errorf("set malicious %v: %w", nodeID, err)
+	}
+	return nil
+}
+
 // IsMalicious returns true if identity is known to be malicious.
 func IsMalicious(db sql.Executor, nodeID types.NodeID) (bool, error) {
-	rows, err := db.Exec(`
-	SELECT 1 FROM identities
-	WHERE (
-		marriage_atx = (SELECT marriage_atx FROM identities WHERE pubkey = ?1 AND marriage_atx IS NOT NULL)
-		AND proof IS NOT NULL
-	)
-	OR (pubkey = ?1 AND marriage_atx IS NULL AND proof IS NOT NULL);`,
+	// this query should be served from malicious_identities index
+	rows, err := db.Exec(`SELECT 1 FROM identities WHERE pubkey = ?1 AND is_malicious IS TRUE;`,
 		func(stmt *sql.Statement) {
 			stmt.BindBytes(1, nodeID.Bytes())
 		}, nil)
@@ -47,6 +59,18 @@ func IsMalicious(db sql.Executor, nodeID types.NodeID) (bool, error) {
 	return rows > 0, nil
 }
 
+// IsMarriageMalicious returns true if the marriage ATX is known to be malicious.
+func IsMarriageMalicious(db sql.Executor, marriage types.ATXID) (bool, error) {
+	rows, err := db.Exec(`SELECT 1 FROM identities WHERE marriage_atx = ?1 AND is_malicious IS TRUE;`,
+		func(stmt *sql.Statement) {
+			stmt.BindBytes(1, marriage.Bytes())
+		}, nil)
+	if err != nil {
+		return false, fmt.Errorf("is marriage malicious %v: %w", marriage, err)
+	}
+	return rows > 0, nil
+}
+
 // GetBlobSizes returns the sizes of the blobs corresponding to malfeasance proofs for the
 // specified identities. For non-existent proofs, the corresponding items are set to -1.
 func GetBlobSizes(db sql.Executor, ids [][]byte) (sizes []int, err error) {
@@ -89,8 +113,8 @@ func IterateMalicious(
 	return callbackErr
 }
 
-// GetMalicious retrieves malicious node IDs from the database.
-func GetMalicious(db sql.Executor) (ids []types.NodeID, err error) {
+// GetNodeIDsWithProofs retrieves malicious node IDs from the database.
+func GetNodeIDsWithProofs(db sql.Executor) (ids []types.NodeID, err error) {
 	if err = IterateMalicious(db, func(total int, nid types.NodeID) error {
 		if ids == nil {
 			ids = make([]types.NodeID, 0, total)

sql/identities/identities_test.go

@@ -60,7 +60,7 @@ func TestMalicious(t *testing.T) {
 
 func Test_GetMalicious(t *testing.T) {
 	db := statesql.InMemory()
-	got, err := identities.GetMalicious(db)
+	got, err := identities.GetNodeIDsWithProofs(db)
 	require.NoError(t, err)
 	require.Nil(t, got)
 
@@ -71,7 +71,7 @@ func Test_GetMalicious(t *testing.T) {
 		bad = append(bad, nid)
 		require.NoError(t, identities.SetMalicious(db, nid, types.RandomBytes(11), time.Now().Local()))
 	}
-	got, err = identities.GetMalicious(db)
+	got, err = identities.GetNodeIDsWithProofs(db)
 	require.NoError(t, err)
 	require.Equal(t, bad, got)
 }
@@ -250,7 +250,7 @@ func TestEquivocationSet(t *testing.T) {
 		require.ErrorIs(t, err, sql.ErrNotFound)
 		require.Nil(t, blob.Bytes)
 
-		ids, err := identities.GetMalicious(db)
+		ids, err := identities.GetNodeIDsWithProofs(db)
 		require.NoError(t, err)
 		require.Empty(t, ids)
 	})

sql/statesql/schema/migrations/0024_identities_speedup.sql

@@ -0,0 +1,21 @@
+ALTER TABLE identities RENAME TO identities_old;
+CREATE TABLE identities
+(
+    pubkey VARCHAR PRIMARY KEY,
+    is_malicious BOOLEAN DEFAULT FALSE NOT NULL,
+    proof  BLOB,
+    marriage_atx CHAR(32),
+    received INT DEFAULT 0 NOT NULL, 
+    marriage_atx CHAR(32), 
+    marriage_idx INTEGER, 
+    marriage_target CHAR(32), 
+    marriage_signature CHAR(64)
+);
+
+INSERT INTO identities (pubkey, is_malicious, proof, received, marriage_atx, marriage_idx, marriage_target, marriage_signature)
+  SELECT pubkey, proof IS NOT NULL, proof, received, marriage_atx, marriage_idx, marriage_target, marriage_signature FROM identities_old;
+
+DROP TABLE identities_old;
+
+CREATE INDEX malicious_identities ON identities (pubkey) where is_malicious;
+CREATE INDEX malicious_marriages ON identities (marriage_atx) where is_malicious;

sql/statesql/schema/schema.sql

@@ -85,7 +85,7 @@ CREATE TABLE identities
 (
     pubkey VARCHAR PRIMARY KEY,
     proof  BLOB
-, received INT DEFAULT 0 NOT NULL, marriage_atx CHAR(32), marriage_idx INTEGER, marriage_target CHAR(32), marriage_signature CHAR(64)) WITHOUT ROWID;
+, received INT DEFAULT 0 NOT NULL, marriage_atx CHAR(32), marriage_idx INTEGER, marriage_target CHAR(32), marriage_signature CHAR(64));
 CREATE TABLE layers
 (
     id              INT PRIMARY KEY DESC,