From d1f54099c4d777f0aa43ca179ce78990be5f41ac Mon Sep 17 00:00:00 2001
From: LittleSheep <littlesheep.code@hotmail.com>
Date: Sat, 6 Apr 2024 23:34:02 +0800
Subject: [PATCH] :bug: Fixed friendship check

---
 pkg/models/accounts.go        | 18 ------------------
 pkg/models/friendships.go     | 20 ++++++++++++++++++++
 pkg/server/friendships_api.go |  7 ++-----
 pkg/services/friendships.go   | 16 ++++++++++++++++
 4 files changed, 38 insertions(+), 23 deletions(-)
 create mode 100644 pkg/models/friendships.go

diff --git a/pkg/models/accounts.go b/pkg/models/accounts.go
index e9e2e77..6d43937 100644
--- a/pkg/models/accounts.go
+++ b/pkg/models/accounts.go
@@ -72,21 +72,3 @@ type AccountContact struct {
 	VerifiedAt *time.Time `json:"verified_at"`
 	AccountID  uint       `json:"account_id"`
 }
-
-type FriendshipStatus = int8
-
-const (
-	FriendshipPending = FriendshipStatus(iota)
-	FriendshipActive
-	FriendshipBlocked
-)
-
-type AccountFriendship struct {
-	BaseModel
-
-	AccountID uint             `json:"account_id"`
-	RelatedID uint             `json:"related_id"`
-	Account   Account          `json:"account"`
-	Related   Account          `json:"related"`
-	Status    FriendshipStatus `json:"status"`
-}
diff --git a/pkg/models/friendships.go b/pkg/models/friendships.go
new file mode 100644
index 0000000..da5d1ca
--- /dev/null
+++ b/pkg/models/friendships.go
@@ -0,0 +1,20 @@
+package models
+
+type FriendshipStatus = int8
+
+const (
+	FriendshipPending = FriendshipStatus(iota)
+	FriendshipActive
+	FriendshipBlocked
+)
+
+type AccountFriendship struct {
+	BaseModel
+
+	AccountID uint             `json:"account_id"`
+	RelatedID uint             `json:"related_id"`
+	BlockedBy *uint            `json:"blocked_by"`
+	Account   Account          `json:"account"`
+	Related   Account          `json:"related"`
+	Status    FriendshipStatus `json:"status"`
+}
diff --git a/pkg/server/friendships_api.go b/pkg/server/friendships_api.go
index aac0b6a..0831a30 100644
--- a/pkg/server/friendships_api.go
+++ b/pkg/server/friendships_api.go
@@ -81,15 +81,12 @@ func editFriendship(c *fiber.Ctx) error {
 	friendship, err := services.GetFriendWithTwoSides(user.ID, related.ID)
 	if err != nil {
 		return fiber.NewError(fiber.StatusNotFound, err.Error())
-	} else if friendship.Status == models.FriendshipPending && data.Status == uint8(models.FriendshipActive) {
-		if friendship.RelatedID != user.ID {
-			return fiber.NewError(fiber.StatusNotFound, "only related person can accept friendship")
-		}
 	}
 
+	originalStatus := friendship.Status
 	friendship.Status = models.FriendshipStatus(data.Status)
 
-	if friendship, err := services.EditFriend(friendship); err != nil {
+	if friendship, err := services.EditFriendWithCheck(friendship, user, originalStatus); err != nil {
 		return fiber.NewError(fiber.StatusBadRequest, err.Error())
 	} else {
 		return c.JSON(friendship)
diff --git a/pkg/services/friendships.go b/pkg/services/friendships.go
index 64cae36..301b946 100644
--- a/pkg/services/friendships.go
+++ b/pkg/services/friendships.go
@@ -80,6 +80,22 @@ func NewFriend(user models.Account, related models.Account, status models.Friend
 	return relationship, nil
 }
 
+func EditFriendWithCheck(relationship models.AccountFriendship, user models.Account, originalStatus models.FriendshipStatus) (models.AccountFriendship, error) {
+	if relationship.Status != originalStatus {
+		if originalStatus == models.FriendshipBlocked && relationship.BlockedBy != nil && user.ID != *relationship.BlockedBy {
+			return relationship, fmt.Errorf("the friendship has been blocked by the otherside, you cannot modify it status")
+		}
+		if relationship.Status == models.FriendshipPending && relationship.RelatedID != user.ID {
+			return relationship, fmt.Errorf("only related person can accept friendship")
+		}
+	}
+	if originalStatus != models.FriendshipBlocked && relationship.Status == models.FriendshipBlocked {
+		relationship.BlockedBy = &user.ID
+	}
+
+	return EditFriend(relationship)
+}
+
 func EditFriend(relationship models.AccountFriendship) (models.AccountFriendship, error) {
 	if err := database.C.Save(&relationship).Error; err != nil {
 		return relationship, err