From 775a3b8868c42b5f97645a14a81d4eed3dc35887 Mon Sep 17 00:00:00 2001
From: LittleSheep <littlesheep.code@hotmail.com>
Date: Wed, 7 Feb 2024 23:15:16 +0800
Subject: [PATCH] :sparkles: Subscribable notification

---
 pkg/database/migrator.go        |  1 +
 pkg/models/accounts.go          | 31 ++++++++++++++++---------------
 pkg/models/notifications.go     | 13 +++++++++++++
 pkg/server/notifications_api.go | 26 ++++++++++++++++++++++++++
 pkg/server/startup.go           |  1 +
 pkg/services/notifications.go   | 13 +++++++++++++
 6 files changed, 70 insertions(+), 15 deletions(-)

diff --git a/pkg/database/migrator.go b/pkg/database/migrator.go
index d7b02b0..4e61f43 100644
--- a/pkg/database/migrator.go
+++ b/pkg/database/migrator.go
@@ -17,6 +17,7 @@ func RunMigration(source *gorm.DB) error {
 		&models.ThirdClient{},
 		&models.ActionEvent{},
 		&models.Notification{},
+		&models.NotificationSubscriber{},
 	); err != nil {
 		return err
 	}
diff --git a/pkg/models/accounts.go b/pkg/models/accounts.go
index a6db7c0..7f40f33 100644
--- a/pkg/models/accounts.go
+++ b/pkg/models/accounts.go
@@ -17,21 +17,22 @@ const (
 type Account struct {
 	BaseModel
 
-	Name          string           `json:"name" gorm:"uniqueIndex"`
-	Nick          string           `json:"nick"`
-	Avatar        string           `json:"avatar"`
-	State         AccountState     `json:"state"`
-	Profile       AccountProfile   `json:"profile"`
-	Sessions      []AuthSession    `json:"sessions"`
-	Challenges    []AuthChallenge  `json:"challenges"`
-	Factors       []AuthFactor     `json:"factors"`
-	Contacts      []AccountContact `json:"contacts"`
-	Events        []ActionEvent    `json:"events"`
-	MagicTokens   []MagicToken     `json:"-" gorm:"foreignKey:AssignTo"`
-	ThirdClients  []ThirdClient    `json:"clients"`
-	Notifications []Notification   `json:"notifications" gorm:"foreignKey:RecipientID"`
-	ConfirmedAt   *time.Time       `json:"confirmed_at"`
-	PowerLevel    int              `json:"power_level"`
+	Name              string                   `json:"name" gorm:"uniqueIndex"`
+	Nick              string                   `json:"nick"`
+	Avatar            string                   `json:"avatar"`
+	State             AccountState             `json:"state"`
+	Profile           AccountProfile           `json:"profile"`
+	Sessions          []AuthSession            `json:"sessions"`
+	Challenges        []AuthChallenge          `json:"challenges"`
+	Factors           []AuthFactor             `json:"factors"`
+	Contacts          []AccountContact         `json:"contacts"`
+	Events            []ActionEvent            `json:"events"`
+	MagicTokens       []MagicToken             `json:"-" gorm:"foreignKey:AssignTo"`
+	ThirdClients      []ThirdClient            `json:"clients"`
+	Notifications     []Notification           `json:"notifications" gorm:"foreignKey:RecipientID"`
+	NotifySubscribers []NotificationSubscriber `json:"notify_subscribers"`
+	ConfirmedAt       *time.Time               `json:"confirmed_at"`
+	PowerLevel        int                      `json:"power_level"`
 }
 
 func (v Account) GetPrimaryEmail() AccountContact {
diff --git a/pkg/models/notifications.go b/pkg/models/notifications.go
index 3929d69..0294e8c 100644
--- a/pkg/models/notifications.go
+++ b/pkg/models/notifications.go
@@ -12,3 +12,16 @@ type Notification struct {
 	SenderID    *uint      `json:"sender_id"`
 	RecipientID uint       `json:"recipient_id"`
 }
+
+const (
+	NotifySubscriberFirebase = "firebase"
+)
+
+type NotificationSubscriber struct {
+	BaseModel
+
+	UserAgent string `json:"user_agent"`
+	Provider  string `json:"provider"`
+	DeviceID  string `json:"device_id"`
+	AccountID uint   `json:"account_id"`
+}
diff --git a/pkg/server/notifications_api.go b/pkg/server/notifications_api.go
index 669081f..750c608 100644
--- a/pkg/server/notifications_api.go
+++ b/pkg/server/notifications_api.go
@@ -3,6 +3,7 @@ package server
 import (
 	"code.smartsheep.studio/hydrogen/passport/pkg/database"
 	"code.smartsheep.studio/hydrogen/passport/pkg/models"
+	"code.smartsheep.studio/hydrogen/passport/pkg/services"
 	"github.com/gofiber/fiber/v2"
 	"github.com/samber/lo"
 	"time"
@@ -57,3 +58,28 @@ func markNotificationRead(c *fiber.Ctx) error {
 		return c.SendStatus(fiber.StatusOK)
 	}
 }
+
+func addNotifySubscriber(c *fiber.Ctx) error {
+	user := c.Locals("principal").(models.Account)
+
+	var data struct {
+		Provider string `json:"provider" validate:"required"`
+		DeviceID string `json:"device_id" validate:"required"`
+	}
+
+	if err := BindAndValidate(c, &data); err != nil {
+		return err
+	}
+
+	subscriber, err := services.AddNotifySubscriber(
+		user,
+		data.Provider,
+		data.DeviceID,
+		c.Get(fiber.HeaderUserAgent),
+	)
+	if err != nil {
+		return fiber.NewError(fiber.StatusBadRequest, err.Error())
+	}
+
+	return c.JSON(subscriber)
+}
diff --git a/pkg/server/startup.go b/pkg/server/startup.go
index b47c681..aaafa6c 100644
--- a/pkg/server/startup.go
+++ b/pkg/server/startup.go
@@ -62,6 +62,7 @@ func NewServer() {
 
 		api.Get("/notifications", auth, getNotifications)
 		api.Put("/notifications/:notificationId/read", auth, markNotificationRead)
+		api.Post("/notification/subscribe", auth, addNotifySubscriber)
 
 		api.Get("/users/me", auth, getUserinfo)
 		api.Put("/users/me", auth, editUserinfo)
diff --git a/pkg/services/notifications.go b/pkg/services/notifications.go
index a666764..60d1c32 100644
--- a/pkg/services/notifications.go
+++ b/pkg/services/notifications.go
@@ -5,6 +5,19 @@ import (
 	"code.smartsheep.studio/hydrogen/passport/pkg/models"
 )
 
+func AddNotifySubscriber(user models.Account, provider, device, ua string) (models.NotificationSubscriber, error) {
+	subscriber := models.NotificationSubscriber{
+		UserAgent: ua,
+		Provider:  provider,
+		DeviceID:  ua,
+		AccountID: user.ID,
+	}
+
+	err := database.C.Save(&subscriber).Error
+
+	return subscriber, err
+}
+
 func NewNotification(user models.ThirdClient, target models.Account, subject, content string, important bool) error {
 	notification := models.Notification{
 		Subject:     subject,