From 1f75a9e64b30618c21a049b1527db58578eb0ed8 Mon Sep 17 00:00:00 2001
From: LittleSheep <littlesheep.code@hotmail.com>
Date: Mon, 27 Jan 2025 15:43:24 +0800
Subject: [PATCH] :recycle: Use nats jetstream instead of database to store otp

---
 pkg/internal/gap/server.go       | 39 +++++++++++++++++++++++--
 pkg/internal/services/factors.go | 49 +++++++++++++++++++++++++-------
 2 files changed, 74 insertions(+), 14 deletions(-)

diff --git a/pkg/internal/gap/server.go b/pkg/internal/gap/server.go
index 91843e4..5c96508 100644
--- a/pkg/internal/gap/server.go
+++ b/pkg/internal/gap/server.go
@@ -1,19 +1,32 @@
 package gap
 
 import (
+	"errors"
 	"fmt"
+	"strings"
+	"time"
+
 	"git.solsynth.dev/hypernet/nexus/pkg/nex"
+	"git.solsynth.dev/hypernet/nexus/pkg/nex/rx"
 	"git.solsynth.dev/hypernet/nexus/pkg/proto"
 	"git.solsynth.dev/hypernet/pusher/pkg/pushkit/pushcon"
+	"github.com/nats-io/nats.go"
 	"github.com/rs/zerolog/log"
 	"github.com/samber/lo"
-	"strings"
 
 	"github.com/spf13/viper"
 )
 
-var Nx *nex.Conn
-var Px *pushcon.Conn
+var (
+	Nx *nex.Conn
+	Px *pushcon.Conn
+	Rx *rx.MqConn
+	Jt nats.JetStreamContext
+)
+
+const (
+	FactorOtpPrefix = "passport.otp."
+)
 
 func InitializeToNexus() error {
 	grpcBind := strings.SplitN(viper.GetString("grpc_bind"), ":", 2)
@@ -46,5 +59,25 @@ func InitializeToNexus() error {
 		return fmt.Errorf("error during initialize pushcon: %v", err)
 	}
 
+	Rx, err = rx.NewMqConn(Nx)
+	if err != nil {
+		return fmt.Errorf("error during initialize nexus rx module: %v", err)
+	}
+	Jt, err = Rx.Nt.JetStream()
+	if err != nil {
+		return fmt.Errorf("error during initialize nats jetstream: %v", err)
+	}
+
+	jetstreamCfg := &nats.StreamConfig{
+		Name:     "Passport OTPs",
+		Subjects: []string{FactorOtpPrefix + ">"},
+		Storage:  nats.MemoryStorage,
+		MaxAge:   5 * time.Minute,
+	}
+	_, err = Jt.AddStream(jetstreamCfg)
+	if err != nil && !errors.Is(err, nats.ErrStreamNameAlreadyInUse) {
+		return fmt.Errorf("error during initialize jetstream stream: %v", err)
+	}
+
 	return err
 }
diff --git a/pkg/internal/services/factors.go b/pkg/internal/services/factors.go
index 3928180..4b8ad1b 100644
--- a/pkg/internal/services/factors.go
+++ b/pkg/internal/services/factors.go
@@ -2,15 +2,18 @@ package services
 
 import (
 	"fmt"
+	"strings"
+	"time"
+
 	"git.solsynth.dev/hypernet/passport/pkg/authkit/models"
 	"git.solsynth.dev/hypernet/passport/pkg/internal/database"
 	"git.solsynth.dev/hypernet/passport/pkg/internal/gap"
 	"git.solsynth.dev/hypernet/pusher/pkg/pushkit"
 	"github.com/google/uuid"
+	"github.com/nats-io/nats.go"
 	"github.com/rs/zerolog/log"
 	"github.com/samber/lo"
 	"github.com/spf13/viper"
-	"strings"
 )
 
 const EmailPasswordTemplate = `Dear %s,
@@ -75,15 +78,20 @@ func GetFactorCode(factor models.AuthFactor) (bool, error) {
 			return true, err
 		}
 
-		factor.Secret = uuid.NewString()[:6]
-		if err := database.C.Save(&factor).Error; err != nil {
-			return true, err
+		secret := uuid.NewString()[:6]
+
+		identifier := fmt.Sprintf("%s%d", gap.FactorOtpPrefix, factor.ID)
+		_, err := gap.Jt.Publish(identifier, []byte(secret))
+		if err != nil {
+			return true, fmt.Errorf("error during publish message: %v", err)
+		} else {
+			log.Info().Uint("factor", factor.ID).Str("secret", secret).Msg("Published one-time-password to JetStream...")
 		}
 
 		subject := fmt.Sprintf("[%s] Login verification code", viper.GetString("name"))
-		content := fmt.Sprintf(EmailPasswordTemplate, user.Name, factor.Secret, viper.GetString("maintainer"))
+		content := fmt.Sprintf(EmailPasswordTemplate, user.Name, secret, viper.GetString("maintainer"))
 
-		err := gap.Px.PushEmail(pushkit.EmailDeliverRequest{
+		err = gap.Px.PushEmail(pushkit.EmailDeliverRequest{
 			To: user.GetPrimaryEmail().Content,
 			Email: pushkit.EmailData{
 				Subject: subject,
@@ -110,11 +118,30 @@ func CheckFactor(factor models.AuthFactor, code string) error {
 			fmt.Errorf("invalid password"),
 		)
 	case models.EmailPasswordFactor:
-		return lo.Ternary(
-			strings.ToUpper(code) == strings.ToUpper(factor.Secret),
-			nil,
-			fmt.Errorf("invalid verification code"),
-		)
+		identifier := fmt.Sprintf("%s%d", gap.FactorOtpPrefix, factor.ID)
+		sub, err := gap.Jt.PullSubscribe(identifier, "otp_consumer", nats.Durable("otp_consumer"))
+		if err != nil {
+			log.Error().Err(err).Msg("Error subscribing to subject when validating factor code...")
+			return fmt.Errorf("error subscribing to subject: %v", err)
+		}
+
+		msgs, err := sub.Fetch(1, nats.MaxWait(2*time.Second))
+		if err != nil {
+			log.Error().Err(err).Msg("Error fetching message when validating factor code...")
+			return fmt.Errorf("error fetching message: %v", err)
+		}
+
+		if len(msgs) > 0 {
+			msg := msgs[0]
+			if !strings.EqualFold(code, string(msg.Data)) {
+				return fmt.Errorf("invalid verification code")
+			}
+			log.Info().Uint("factor", factor.ID).Str("secret", code).Msg("Verified one-time-password...")
+			msg.Ack()
+			return nil
+		}
+
+		return fmt.Errorf("one-time-password not found or expired")
 	}
 
 	return nil