feat: added csrf feature

Author: Wong801

src/controllers/root.go

@@ -1,10 +1,13 @@
 package controller
 
 import (
+	"net/http"
+
 	"github.com/Wong801/gin-api/src/api"
 	"github.com/Wong801/gin-api/src/db"
 	service "github.com/Wong801/gin-api/src/services"
 	"github.com/gin-gonic/gin"
+	csrf "github.com/utrack/gin-csrf"
 )
 
 type RootController struct {
@@ -41,3 +44,12 @@ func (rc RootController) Ping() func(c *gin.Context) {
 		c.Next()
 	}
 }
+
+func (rc RootController) GetToken() func(c *gin.Context) {
+	return func(c *gin.Context) {
+		c.Set("status", http.StatusOK)
+		c.Set("data", csrf.GetToken(c))
+
+		c.Next()
+	}
+}

src/controllers/user.go

@@ -7,7 +7,6 @@ import (
 	model "github.com/Wong801/gin-api/src/models"
 	service "github.com/Wong801/gin-api/src/services"
 	"github.com/gin-gonic/gin"
-	csrf "github.com/utrack/gin-csrf"
 )
 
 type UserController struct {
@@ -94,18 +93,25 @@ func (uc UserController) Login() gin.HandlerFunc {
 			return
 		}
 		c.SetCookie("jwt", token.Jwt, token.MaxAge, "/", token.Domain, token.Secure, token.HttpOnly)
-		c.SetCookie("X-CSRF-TOKEN", csrf.GetToken(c), token.MaxAge, "/", token.Domain, false, false)
 		c.Set("data", map[string]string{
 			"message": "Login Success",
 		})
 		c.Next()
 	}
 }
 
+func (uc UserController) CheckLogin() gin.HandlerFunc {
+	return func(c *gin.Context) {
+		c.Set("data", map[string]string{
+			"message": "success",
+		})
+		c.Next()
+	}
+}
+
 func (uc UserController) Logout() gin.HandlerFunc {
 	return func(c *gin.Context) {
 		c.SetCookie("jwt", "", -1, "/", "", true, true)
-		c.SetCookie("X-CSRF-TOKEN", "", -1, "/", "", false, false)
 		c.Next()
 	}
 }

src/models/user.go

@@ -36,3 +36,9 @@ type UserChangePassword struct {
 func (u User) GetBase() *UserBase {
 	return &u.UserBase
 }
+
+func (ub UserBase) GetUser() *User {
+	return &User{
+		UserBase: ub,
+	}
+}

src/routes/main.go

@@ -1,7 +1,6 @@
 package route
 
 import (
-	"errors"
 	"net/http"
 
 	"github.com/Wong801/gin-api/src/config"
@@ -28,8 +27,10 @@ func InitRoutes() handler {
 	r.router.Use(csrf.Middleware(csrf.Options{
 		Secret: config.GetEnv("CSRF_SECRET", "secret"),
 		ErrorFunc: func(c *gin.Context) {
-			c.Set("status", http.StatusBadRequest)
-			c.Set("error", errors.New("CSRF token mismatch"))
+			c.AbortWithStatusJSON(http.StatusBadRequest, &entity.HttpResponse{
+				Success: false,
+				Data:    "CSRF token mismatch",
+			})
 		},
 	}))
 

src/routes/root.go

@@ -10,4 +10,5 @@ func (r handler) addRoot(rg *gin.RouterGroup) {
 
 	rg.GET("/stats", rc.GetStats())
 	rg.GET("/ping", rc.Ping())
+	rg.GET("/csrf", rc.GetToken())
 }

src/routes/user.go

@@ -17,6 +17,7 @@ func (r handler) addUsers(rg *gin.RouterGroup) {
 
 	userRoute.Use(m.Authenticate())
 
+	userRoute.POST("/check-login", userController.CheckLogin())
 	userRoute.POST("/logout", userController.Logout())
 	userRoute.PUT("/profile", userController.UpdateProfile())
 	userRoute.PATCH("/change-password", userController.ChangePassword())

src/services/user.go

@@ -74,19 +74,16 @@ func (us UserService) GetUser() (int, *model.UserBase, error) {
 		return http.StatusNotFound, nil, err
 	}
 
-	return http.StatusFound, user.GetBase(), nil
+	return http.StatusOK, user.GetBase(), nil
 }
 
 func (us UserService) UpdateUser(id int, u *model.UserBase) (int, *model.UserBase, error) {
-	var user model.User
 	db.Open(us.DB)
 
-	us.DB.Database.Where("id = ?", id)
-
-	if err := us.DB.Database.Save(&user).Error; err != nil {
-		return http.StatusInternalServerError, nil, err
+	if err := us.DB.Database.Save(u.GetUser()).Error; err != nil {
+		return http.StatusBadRequest, nil, err
 	}
-	return http.StatusOK, user.GetBase(), nil
+	return http.StatusOK, u.GetUser().GetBase(), nil
 }
 
 func (us UserService) Register(u *model.User) (int, error) {
@@ -109,7 +106,7 @@ func (us UserService) Login(u *model.UserLogin) (int, *entity.Token, error) {
 	db.Open(us.DB)
 
 	if err := us.DB.Database.First(&user, "email = ? OR username = ?", u.Email, u.Username).Error; err != nil {
-		return http.StatusNotFound, nil, errors.New("user not found")
+		return http.StatusBadRequest, nil, errors.New("incorrect username or password")
 	}
 
 	correctPassword := checkPasswordHash(u.Password, user.Password)