Improve 2FA security and user experience

cache/cache.go

@@ -44,27 +44,6 @@ func init() {
 	fileCache = make(map[string]*FileWithExpired)
 	ticker := time.NewTicker(time.Minute)
 
-	go func() {
-		for {
-			<-ticker.C
-			currentTime := time.Now()
-			cacheClean := 0
-			cleanID := utils.GenerateRandomString(10)
-			log.Info(fmt.Sprintf("Cache cleanup [user] [%s] initiated at %02d:%02d:%02d", cleanID, currentTime.Hour(), currentTime.Minute(), currentTime.Second()))
-
-			for _, user := range userCache {
-				user.mu.Lock()
-				if currentTime.Sub(user.AccessAt) > time.Hour*8 {
-					delete(userCache, user.Email)
-					cacheClean++
-				}
-				user.mu.Unlock()
-			}
-
-			log.Info(fmt.Sprintf("Cache cleanup [user] [%s] completed: %d entries removed. Finished at %s", cleanID, cacheClean, time.Since(currentTime)))
-		}
-	}()
-
 	go func() {
 		for {
 			<-ticker.C
@@ -89,35 +68,6 @@ func init() {
 	}()
 }
 
-func GetUser(email string) (*UserWithExpired, error) {
-	if user, ok := userCache[email]; ok {
-		return user, nil
-	}
-
-	userData, err := db.DB.GetUser(email)
-	if err != nil {
-		return nil, err
-	}
-
-	userCache[email] = &UserWithExpired{
-		UserID:   userData.UserID,
-		Username: userData.Username,
-		Email:    userData.Email,
-		Password: userData.Password,
-		Totp:     userData.Totp,
-		AccessAt: time.Now(),
-	}
-
-	return userCache[email], nil
-}
-
-func DeleteUser(email string) {
-	userCache[email].mu.Lock()
-	defer userCache[email].mu.Unlock()
-
-	delete(userCache, email)
-}
-
 func GetFile(id string) (*FileWithExpired, error) {
 	if file, ok := fileCache[id]; ok {
 		file.AccessAt = time.Now()

handler/auth/google/callback/callback.go

@@ -5,7 +5,6 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
-	"github.com/fossyy/filekeeper/cache"
 	"github.com/fossyy/filekeeper/db"
 	googleOauthSetupHandler "github.com/fossyy/filekeeper/handler/auth/google/setup"
 	signinHandler "github.com/fossyy/filekeeper/handler/signin"
@@ -157,7 +156,7 @@ func GET(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	user, err := cache.GetUser(oauthUser.Email)
+	user, err := db.DB.GetUser(oauthUser.Email)
 
 	if err != nil {
 		w.WriteHeader(http.StatusInternalServerError)

handler/auth/totp/totp.go

@@ -1,16 +1,66 @@
 package totpHandler
 
 import (
+	"errors"
 	"fmt"
+	"github.com/fossyy/filekeeper/logger"
 	"github.com/fossyy/filekeeper/session"
 	"github.com/fossyy/filekeeper/types"
+	"github.com/fossyy/filekeeper/utils"
 	totpView "github.com/fossyy/filekeeper/view/totp"
+	"github.com/google/uuid"
 	"github.com/xlzd/gotp"
 	"net/http"
+	"strings"
+	"sync"
 	"time"
 )
 
+type TotpInfo struct {
+	ID         string
+	UserID     uuid.UUID
+	Secret     string
+	Email      string
+	Username   string
+	CreateTime time.Time
+	mu         sync.Mutex
+}
+
+var TotpInfoList = make(map[string]*TotpInfo)
+var log *logger.AggregatedLogger
+
+func init() {
+	log = logger.Logger()
+
+	ticker := time.NewTicker(time.Minute)
+	go func() {
+		for {
+			<-ticker.C
+			currentTime := time.Now()
+			cacheClean := 0
+			cleanID := utils.GenerateRandomString(10)
+			log.Info(fmt.Sprintf("Cache cleanup [TOTP] [%s] initiated at %02d:%02d:%02d", cleanID, currentTime.Hour(), currentTime.Minute(), currentTime.Second()))
+
+			for _, data := range TotpInfoList {
+				data.mu.Lock()
+				if currentTime.Sub(data.CreateTime) > time.Minute*10 {
+					delete(TotpInfoList, data.ID)
+					cacheClean++
+				}
+				data.mu.Unlock()
+			}
+
+			log.Info(fmt.Sprintf("Cache cleanup [TOTP] [%s] completed: %d entries removed. Finished at %s", cleanID, cacheClean, time.Since(currentTime)))
+		}
+	}()
+}
+
 func GET(w http.ResponseWriter, r *http.Request) {
+	_, ok := TotpInfoList[r.PathValue("id")]
+	if !ok {
+		w.WriteHeader(http.StatusNotFound)
+		return
+	}
 	component := totpView.Main("Filekeeper - 2FA Page")
 	err := component.Render(r.Context(), w)
 	if err != nil {
@@ -23,25 +73,111 @@ func GET(w http.ResponseWriter, r *http.Request) {
 func POST(w http.ResponseWriter, r *http.Request) {
 	r.ParseForm()
 	code := r.Form.Get("code")
-	_, user, key := session.GetSession(r)
+	data, ok := TotpInfoList[r.PathValue("id")]
-
+	if !ok {
-	totp := gotp.NewDefaultTOTP(user.Totp)
+		w.WriteHeader(http.StatusNotFound)
-	if totp.Verify(code, time.Now().Unix()) {
-		storeSession, err := session.Get(key)
-		if err != nil {
 		return
 	}
-		fmt.Println(storeSession)
+	fmt.Println(data)
+	totp := gotp.NewDefaultTOTP(data.Secret)
+	if totp.Verify(code, time.Now().Unix()) {
+		storeSession := session.Create()
 		storeSession.Values["user"] = types.User{
-			UserID:        user.UserID,
+			UserID:        data.UserID,
-			Email:         user.Email,
+			Email:         data.Email,
-			Username:      user.Username,
+			Username:      data.Username,
 			Totp:          "",
 			Authenticated: true,
 		}
-		http.Redirect(w, r, "/user", http.StatusFound)
+		userAgent := r.Header.Get("User-Agent")
+		browserInfo, osInfo := ParseUserAgent(userAgent)
+
+		sessionInfo := session.SessionInfo{
+			SessionID: storeSession.ID,
+			Browser:   browserInfo["browser"],
+			Version:   browserInfo["version"],
+			OS:        osInfo["os"],
+			OSVersion: osInfo["version"],
+			IP:        utils.ClientIP(r),
+			Location:  "Indonesia",
+		}
+
+		storeSession.Save(w)
+		session.AddSessionInfo(data.Email, &sessionInfo)
+
+		cookie, err := r.Cookie("redirect")
+		if errors.Is(err, http.ErrNoCookie) {
+			http.Redirect(w, r, "/", http.StatusSeeOther)
+			return
+		}
+		http.SetCookie(w, &http.Cookie{
+			Name:   "redirect",
+			MaxAge: -1,
+		})
+		http.Redirect(w, r, cookie.Value, http.StatusSeeOther)
 		return
 	} else {
 		fmt.Fprint(w, "wrong")
 	}
 }
+
+func ParseUserAgent(userAgent string) (map[string]string, map[string]string) {
+	browserInfo := make(map[string]string)
+	osInfo := make(map[string]string)
+	if strings.Contains(userAgent, "Firefox") {
+		browserInfo["browser"] = "Firefox"
+		parts := strings.Split(userAgent, "Firefox/")
+		if len(parts) > 1 {
+			version := strings.Split(parts[1], " ")[0]
+			browserInfo["version"] = version
+		}
+	} else if strings.Contains(userAgent, "Chrome") {
+		browserInfo["browser"] = "Chrome"
+		parts := strings.Split(userAgent, "Chrome/")
+		if len(parts) > 1 {
+			version := strings.Split(parts[1], " ")[0]
+			browserInfo["version"] = version
+		}
+	} else {
+		browserInfo["browser"] = "Unknown"
+		browserInfo["version"] = "Unknown"
+	}
+
+	if strings.Contains(userAgent, "Windows") {
+		osInfo["os"] = "Windows"
+		parts := strings.Split(userAgent, "Windows ")
+		if len(parts) > 1 {
+			version := strings.Split(parts[1], ";")[0]
+			osInfo["version"] = version
+		}
+	} else if strings.Contains(userAgent, "Macintosh") {
+		osInfo["os"] = "Mac OS"
+		parts := strings.Split(userAgent, "Mac OS X ")
+		if len(parts) > 1 {
+			version := strings.Split(parts[1], ";")[0]
+			osInfo["version"] = version
+		}
+	} else if strings.Contains(userAgent, "Linux") {
+		osInfo["os"] = "Linux"
+		osInfo["version"] = "Unknown"
+	} else if strings.Contains(userAgent, "Android") {
+		osInfo["os"] = "Android"
+		parts := strings.Split(userAgent, "Android ")
+		if len(parts) > 1 {
+			version := strings.Split(parts[1], ";")[0]
+			osInfo["version"] = version
+		}
+	} else if strings.Contains(userAgent, "iPhone") || strings.Contains(userAgent, "iPad") || strings.Contains(userAgent, "iPod") {
+		osInfo["os"] = "iOS"
+		parts := strings.Split(userAgent, "OS ")
+		if len(parts) > 1 {
+			version := strings.Split(parts[1], " ")[0]
+			osInfo["version"] = version
+		}
+	} else {
+		osInfo["os"] = "Unknown"
+		osInfo["version"] = "Unknown"
+	}
+
+	return browserInfo, osInfo
+}

handler/forgotPassword/forgotPassword.go

@@ -5,7 +5,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
-	"github.com/fossyy/filekeeper/cache"
+	"github.com/fossyy/filekeeper/db"
 	"github.com/google/uuid"
 	"net/http"
 	"strconv"
@@ -87,7 +87,7 @@ func POST(w http.ResponseWriter, r *http.Request) {
 
 	emailForm := r.Form.Get("email")
 
-	user, err := cache.GetUser(emailForm)
+	user, err := db.DB.GetUser(emailForm)
 	if errors.Is(err, gorm.ErrRecordNotFound) {
 		component := forgotPasswordView.Main("Filekeeper - Forgot Password Page", types.Message{
 			Code:    0,

handler/forgotPassword/verify/verify.go

@@ -1,7 +1,6 @@
 package forgotPasswordVerifyHandler
 
 import (
-	"github.com/fossyy/filekeeper/cache"
 	"github.com/fossyy/filekeeper/db"
 	forgotPasswordHandler "github.com/fossyy/filekeeper/handler/forgotPassword"
 	"github.com/fossyy/filekeeper/logger"
@@ -98,8 +97,6 @@ func POST(w http.ResponseWriter, r *http.Request) {
 
 	session.RemoveAllSessions(data.User.Email)
 
-	cache.DeleteUser(data.User.Email)
-
 	component := forgotPasswordView.ChangeSuccess("Filekeeper - Forgot Password Page")
 	err = component.Render(r.Context(), w)
 	if err != nil {

handler/signin/signin.go

@@ -3,9 +3,11 @@ package signinHandler
 import (
 	"errors"
 	"github.com/a-h/templ"
-	"github.com/fossyy/filekeeper/cache"
+	"github.com/fossyy/filekeeper/db"
+	totpHandler "github.com/fossyy/filekeeper/handler/auth/totp"
 	"net/http"
 	"strings"
+	"time"
 
 	"github.com/fossyy/filekeeper/logger"
 	"github.com/fossyy/filekeeper/session"
@@ -79,7 +81,7 @@ func POST(w http.ResponseWriter, r *http.Request) {
 	}
 	email := r.Form.Get("email")
 	password := r.Form.Get("password")
-	userData, err := cache.GetUser(email)
+	userData, err := db.DB.GetUser(email)
 	if err != nil {
 		component := signinView.Main("Filekeeper - Sign in Page", types.Message{
 			Code:    0,
@@ -97,16 +99,16 @@ func POST(w http.ResponseWriter, r *http.Request) {
 
 	if email == userData.Email && utils.CheckPasswordHash(password, userData.Password) {
 		if userData.Totp != "" {
-			storeSession := session.Create()
+			id := utils.GenerateRandomString(32)
-			storeSession.Values["user"] = types.User{
+			totpHandler.TotpInfoList[id] = &totpHandler.TotpInfo{
+				ID:         id,
 				UserID:     userData.UserID,
-				Email:         email,
+				Secret:     userData.Totp,
+				Email:      userData.Email,
 				Username:   userData.Username,
-				Totp:          userData.Totp,
+				CreateTime: time.Now(),
-				Authenticated: false,
 			}
-			storeSession.Save(w)
+			http.Redirect(w, r, "/auth/totp/"+id, http.StatusSeeOther)
-			http.Redirect(w, r, "/auth/totp", http.StatusSeeOther)
 			return
 		}
 

routes/routes.go

@@ -50,7 +50,7 @@ func SetupRoutes() *http.ServeMux {
 		}
 	})
 
-	authRouter.HandleFunc("/totp", func(w http.ResponseWriter, r *http.Request) {
+	authRouter.HandleFunc("/totp/{id}", func(w http.ResponseWriter, r *http.Request) {
 		switch r.Method {
 		case http.MethodGet:
 			middleware.Guest(totpHandler.GET, w, r)

view/totp/totp.templ

@@ -15,7 +15,7 @@ templ content(title string) {
                     Enter the 6-digit code sent to your registered device to complete the login process.
                   </p>
                 </div>
-                <form class="space-y-6" action="/auth/totp" method="POST">
+                <form class="space-y-6" method="POST">
                   <div>
                     <label for="code" class="block text-sm font-medium text-muted-foreground">
                       Verification Code