From d40e0c981919457a7644fc2185e3d9c76bf15cf9 Mon Sep 17 00:00:00 2001
From: Ivaylo Novakov <inovakov@gmail.com>
Date: Thu, 24 Mar 2022 12:00:13 +0100
Subject: [PATCH] More reliable pubkey deletion.

---
 database/user.go           | 12 +++++++-----
 test/database/user_test.go |  3 +--
 2 files changed, 8 insertions(+), 7 deletions(-)

diff --git a/database/user.go b/database/user.go
index d8e04a56..9018eff8 100644
--- a/database/user.go
+++ b/database/user.go
@@ -5,7 +5,6 @@ import (
 	"crypto/subtle"
 	"fmt"
 	"net/mail"
-	"strings"
 	"sync"
 	"time"
 
@@ -504,16 +503,19 @@ func (db *DB) UserPubKeyAdd(ctx context.Context, u User, pk PubKey) (err error)
 
 // UserPubKeyRemove removes a PubKey from the given user's set.
 func (db *DB) UserPubKeyRemove(ctx context.Context, u User, pk PubKey) error {
-	filter := bson.M{"_id": u.ID}
+	filter := bson.M{
+		"_id":      u.ID,
+		"pub_keys": bson.M{"$ne": nil},
+	}
 	update := bson.M{
 		"$pull": bson.M{"pub_keys": pk},
 	}
 	opts := options.UpdateOptions{
 		Upsert: &False,
 	}
-	_, err := db.staticUsers.UpdateOne(ctx, filter, update, &opts)
-	if err != nil && strings.Contains(err.Error(), "Cannot apply $pull to a non-array value") {
-		return mongo.ErrNoDocuments
+	ur, err := db.staticUsers.UpdateOne(ctx, filter, update, &opts)
+	if err == nil && ur.MatchedCount == 0 {
+		err = mongo.ErrNoDocuments
 	}
 	return err
 }
diff --git a/test/database/user_test.go b/test/database/user_test.go
index dba767f9..2614ca04 100644
--- a/test/database/user_test.go
+++ b/test/database/user_test.go
@@ -14,7 +14,6 @@ import (
 	"gitlab.com/NebulousLabs/errors"
 	"gitlab.com/NebulousLabs/fastrand"
 	"go.mongodb.org/mongo-driver/bson/primitive"
-	"go.mongodb.org/mongo-driver/mongo"
 	"go.sia.tech/siad/crypto"
 )
 
@@ -462,7 +461,7 @@ func TestUserPubKey(t *testing.T) {
 	copy(pk1[:], fastrand.Bytes(database.PubKeySize))
 	// Try to remove a pubkey. Expect this to fail.
 	err = db.UserPubKeyRemove(ctx, *u, pk)
-	if err == nil || !errors.Contains(err, mongo.ErrNoDocuments) {
+	if err == nil {
 		t.Fatal(err)
 	}
 	// Add a pubkey.