From c5185d51077fa64d16ba6b93c8909a8daa0f90ed Mon Sep 17 00:00:00 2001 From: Bagas Aulia Rezki Date: Thu, 20 Jun 2024 08:49:49 +0700 Subject: [PATCH] Finalize 2FA implementation --- cache/cache.go | 50 +++++++++++++++++ handler/auth/google/callback/callback.go | 3 +- handler/auth/totp/totp.go | 70 +++++------------------- handler/forgotPassword/forgotPassword.go | 4 +- handler/forgotPassword/verify/verify.go | 3 + handler/signin/signin.go | 29 +++++----- handler/user/totp/setup.go | 3 + routes/routes.go | 2 +- 8 files changed, 87 insertions(+), 77 deletions(-) diff --git a/cache/cache.go b/cache/cache.go index 8f2eb75..d859b71 100644 --- a/cache/cache.go +++ b/cache/cache.go @@ -44,6 +44,27 @@ 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 @@ -68,6 +89,35 @@ 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() diff --git a/handler/auth/google/callback/callback.go b/handler/auth/google/callback/callback.go index d5b42db..e765b8f 100644 --- a/handler/auth/google/callback/callback.go +++ b/handler/auth/google/callback/callback.go @@ -5,6 +5,7 @@ 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" @@ -156,7 +157,7 @@ func GET(w http.ResponseWriter, r *http.Request) { return } - user, err := db.DB.GetUser(oauthUser.Email) + user, err := cache.GetUser(oauthUser.Email) if err != nil { w.WriteHeader(http.StatusInternalServerError) diff --git a/handler/auth/totp/totp.go b/handler/auth/totp/totp.go index 846006c..e095fb5 100644 --- a/handler/auth/totp/totp.go +++ b/handler/auth/totp/totp.go @@ -3,61 +3,19 @@ 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 { + _, user, _ := session.GetSession(r) + if user.Authenticated || user.Totp == "" { w.WriteHeader(http.StatusNotFound) return } @@ -73,20 +31,18 @@ func GET(w http.ResponseWriter, r *http.Request) { func POST(w http.ResponseWriter, r *http.Request) { r.ParseForm() code := r.Form.Get("code") - data, ok := TotpInfoList[r.PathValue("id")] - if !ok { - w.WriteHeader(http.StatusNotFound) - return - } - fmt.Println(data) - totp := gotp.NewDefaultTOTP(data.Secret) + _, user, key := session.GetSession(r) + totp := gotp.NewDefaultTOTP(user.Totp) + if totp.Verify(code, time.Now().Unix()) { - storeSession := session.Create() + storeSession, err := session.Get(key) + if err != nil { + return + } storeSession.Values["user"] = types.User{ - UserID: data.UserID, - Email: data.Email, - Username: data.Username, - Totp: "", + UserID: user.UserID, + Email: user.Email, + Username: user.Username, Authenticated: true, } userAgent := r.Header.Get("User-Agent") @@ -103,7 +59,7 @@ func POST(w http.ResponseWriter, r *http.Request) { } storeSession.Save(w) - session.AddSessionInfo(data.Email, &sessionInfo) + session.AddSessionInfo(user.Email, &sessionInfo) cookie, err := r.Cookie("redirect") if errors.Is(err, http.ErrNoCookie) { diff --git a/handler/forgotPassword/forgotPassword.go b/handler/forgotPassword/forgotPassword.go index 02c3fbe..d5c9df9 100644 --- a/handler/forgotPassword/forgotPassword.go +++ b/handler/forgotPassword/forgotPassword.go @@ -5,7 +5,7 @@ import ( "context" "errors" "fmt" - "github.com/fossyy/filekeeper/db" + "github.com/fossyy/filekeeper/cache" "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 := db.DB.GetUser(emailForm) + user, err := cache.GetUser(emailForm) if errors.Is(err, gorm.ErrRecordNotFound) { component := forgotPasswordView.Main("Filekeeper - Forgot Password Page", types.Message{ Code: 0, diff --git a/handler/forgotPassword/verify/verify.go b/handler/forgotPassword/verify/verify.go index 0680b59..b1c7dbc 100644 --- a/handler/forgotPassword/verify/verify.go +++ b/handler/forgotPassword/verify/verify.go @@ -1,6 +1,7 @@ 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" @@ -97,6 +98,8 @@ 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 { diff --git a/handler/signin/signin.go b/handler/signin/signin.go index ff17a26..edcde52 100644 --- a/handler/signin/signin.go +++ b/handler/signin/signin.go @@ -3,17 +3,14 @@ package signinHandler import ( "errors" "github.com/a-h/templ" - "github.com/fossyy/filekeeper/db" - totpHandler "github.com/fossyy/filekeeper/handler/auth/totp" - "net/http" - "strings" - "time" - + "github.com/fossyy/filekeeper/cache" "github.com/fossyy/filekeeper/logger" "github.com/fossyy/filekeeper/session" "github.com/fossyy/filekeeper/types" "github.com/fossyy/filekeeper/utils" signinView "github.com/fossyy/filekeeper/view/signin" + "net/http" + "strings" ) var log *logger.AggregatedLogger @@ -81,7 +78,7 @@ func POST(w http.ResponseWriter, r *http.Request) { } email := r.Form.Get("email") password := r.Form.Get("password") - userData, err := db.DB.GetUser(email) + userData, err := cache.GetUser(email) if err != nil { component := signinView.Main("Filekeeper - Sign in Page", types.Message{ Code: 0, @@ -99,16 +96,16 @@ func POST(w http.ResponseWriter, r *http.Request) { if email == userData.Email && utils.CheckPasswordHash(password, userData.Password) { if userData.Totp != "" { - id := utils.GenerateRandomString(32) - totpHandler.TotpInfoList[id] = &totpHandler.TotpInfo{ - ID: id, - UserID: userData.UserID, - Secret: userData.Totp, - Email: userData.Email, - Username: userData.Username, - CreateTime: time.Now(), + storeSession := session.Create() + storeSession.Values["user"] = types.User{ + UserID: userData.UserID, + Email: email, + Username: userData.Username, + Totp: userData.Totp, + Authenticated: false, } - http.Redirect(w, r, "/auth/totp/"+id, http.StatusSeeOther) + storeSession.Save(w) + http.Redirect(w, r, "/auth/totp", http.StatusSeeOther) return } diff --git a/handler/user/totp/setup.go b/handler/user/totp/setup.go index 2af54ec..065d8c6 100644 --- a/handler/user/totp/setup.go +++ b/handler/user/totp/setup.go @@ -4,6 +4,7 @@ import ( "bytes" "encoding/base64" "fmt" + "github.com/fossyy/filekeeper/cache" userTotpSetupView "github.com/fossyy/filekeeper/view/user/totp" "image/png" "net/http" @@ -63,7 +64,9 @@ func POST(w http.ResponseWriter, r *http.Request) { w.WriteHeader(http.StatusInternalServerError) return } + cache.DeleteUser(userSession.Email) fmt.Fprint(w, "Authentication successful! Access granted.") + return } else { uri := totp.ProvisioningUri(userSession.Email, "filekeeper") diff --git a/routes/routes.go b/routes/routes.go index 879809d..8948521 100644 --- a/routes/routes.go +++ b/routes/routes.go @@ -50,7 +50,7 @@ func SetupRoutes() *http.ServeMux { } }) - authRouter.HandleFunc("/totp/{id}", func(w http.ResponseWriter, r *http.Request) { + authRouter.HandleFunc("/totp", func(w http.ResponseWriter, r *http.Request) { switch r.Method { case http.MethodGet: middleware.Guest(totpHandler.GET, w, r)