Fix jwt generation and verification to include userid (#125)

* chore: fix jwt generation and verification to include userid * chore: add issued at in jwt * chore: fix userid from context * chore: fix jwt and auth middleware * chore: remove comment * chore: fix tests * chore: fix failing tests * chore: refactor delete url function * chore: fix query syntax * used ParseWithClaims * query adjusted --------- Co-authored-by: Amit Prakash <34869115+iamitprakash@users.noreply.github.com>
Real-Dev-Squad · Oct 4, 2024 · e7a523a · e7a523a
1 parent cad4be6
commit e7a523a
Show file tree

Hide file tree

Showing 6 changed files with 183 additions and 119 deletions.
diff --git a/controllers/url.go b/controllers/url.go
@@ -31,15 +31,21 @@ func CreateTinyURL(ctx *gin.Context, db *bun.DB) {
 		return
 	}
 
+	userID, exists := ctx.Get("userID")
+	if !exists {
+		ctx.JSON(http.StatusUnauthorized, dtos.URLCreationResponse{
+			Message: "User not authenticated",
+		})
+		return
+	}
+
 	var existingOriginalURL models.Tinyurl
 	if err := db.NewSelect().Model(&existingOriginalURL).
-		Where("original_url = ?", body.OriginalUrl).
-		Where("user_id = ?", body.UserID).
-		Where("is_deleted = ?", false).
+		Where("original_url = ? AND user_id = ? AND is_deleted = ?", body.OriginalUrl, userID, false).
 		Scan(ctx); err == nil {
 		ctx.JSON(http.StatusOK, dtos.URLCreationResponse{
-			Message:  "Shortened URL already exists",
-			ShortURL: existingOriginalURL.ShortUrl,
+			Message:   "Shortened URL already exists",
+			ShortURL:  existingOriginalURL.ShortUrl,
 			CreatedAt: existingOriginalURL.CreatedAt,
 		})
 		return
@@ -76,8 +82,7 @@ func CreateTinyURL(ctx *gin.Context, db *bun.DB) {
 
 	count, err := db.NewSelect().
 		Model(&models.Tinyurl{}).
-		Where("user_id = ?", body.UserID).
-		Where("is_deleted = ?", false).
+		Where("user_id = ? AND is_deleted = ?", userID, false).
 		Count(ctx)
 
 	if err != nil {
@@ -94,16 +99,21 @@ func CreateTinyURL(ctx *gin.Context, db *bun.DB) {
 		return
 	}
 
-	body.CreatedAt = time.Now().UTC()
+	newTinyURL := models.Tinyurl{
+		OriginalUrl: body.OriginalUrl,
+		ShortUrl:    body.ShortUrl,
+		UserID:      userID.(int64),
+		CreatedAt:   time.Now().UTC(),
+	}
 
-	if _, err := db.NewInsert().Model(&body).Exec(ctx); err != nil {
+	if _, err := db.NewInsert().Model(&newTinyURL).Exec(ctx); err != nil {
 		ctx.JSON(http.StatusInternalServerError, dtos.URLCreationResponse{
 			Message: "Failed to create tiny URL",
 		})
 		return
 	}
 
-	if err := utils.IncrementURLCount(body.UserID, db, ctx); err != nil {
+	if err := utils.IncrementURLCount(userID.(int64), db, ctx); err != nil {
 		ctx.JSON(http.StatusInternalServerError, dtos.URLCreationResponse{
 			Message: "Failed to increment URL count: " + err.Error(),
 		})
@@ -112,8 +122,7 @@ func CreateTinyURL(ctx *gin.Context, db *bun.DB) {
 
 	updatedCount, err := db.NewSelect().
 		Model(&models.Tinyurl{}).
-		Where("user_id = ?", body.UserID).
-		Where("is_deleted = ?", false).
+		Where("user_id = ? AND is_deleted = ?", userID, false).
 		Count(ctx)
 
 	if err != nil {
@@ -124,9 +133,9 @@ func CreateTinyURL(ctx *gin.Context, db *bun.DB) {
 	}
 
 	ctx.JSON(http.StatusOK, dtos.URLCreationResponse{
-		Message:   "Tiny URL created successfully",
-		ShortURL:  body.ShortUrl,
-		URLCount:  updatedCount,
+		Message:  "Tiny URL created successfully",
+		ShortURL: newTinyURL.ShortUrl,
+		URLCount: updatedCount,
 	})
 }
 
@@ -215,51 +224,55 @@ func GetAllURLs(ctx *gin.Context, db *bun.DB) {
 }
 
 func DeleteURL(ctx *gin.Context, db *bun.DB) {
-    id, _ := ctx.Params.Get("id")
-
-    var body struct {
-        UserID int64 `json:"user_id"`
-    }
-    if err := ctx.BindJSON(&body); err != nil {
-        ctx.JSON(http.StatusBadRequest, dtos.UserURLsResponse{
-            Message: "Invalid Request.",
-        })
-        return
-    }
-
-    _, err := db.NewUpdate().Model(&models.Tinyurl{}).Set("is_deleted=?", true).Set("deleted_at=?", time.Now().UTC()).Where("id = ?", id).Exec(ctx)
-    if err != nil {
-        ctx.JSON(http.StatusNotFound, dtos.UserURLsResponse{
-            Message: "No URLs found",
-        })
-        return
-    }
-
-
-    if err := utils.DecrementURLCount(body.UserID, db, ctx); err != nil {
-        ctx.JSON(http.StatusInternalServerError, dtos.URLCreationResponse{
-            Message: "Failed to decrement URL count: " + err.Error(),
-        })
-        return
-    }
+	id, _ := ctx.Params.Get("id")
+
+	userID, exists := ctx.Get("userID")
+	if !exists {
+		ctx.JSON(http.StatusUnauthorized, dtos.UserURLsResponse{
+			Message: "User not authenticated",
+		})
+		return
+	}
+
+	_, err := db.NewUpdate().
+		Model(&models.Tinyurl{}).
+		Set("is_deleted=?", true).
+		Set("deleted_at=?", time.Now().UTC()).
+		Where("id = ?", id).
+		Where("user_id = ?", userID).
+		Exec(ctx)
+
+	if err != nil {
+		ctx.JSON(http.StatusNotFound, dtos.UserURLsResponse{
+			Message: "No URLs found",
+		})
+		return
+	}
+
+	if err := utils.DecrementURLCount(userID.(int64), db, ctx); err != nil {
+		ctx.JSON(http.StatusInternalServerError, dtos.URLCreationResponse{
+			Message: "Failed to decrement URL count: " + err.Error(),
+		})
+		return
+	}
 
 	updatedCount, err := db.NewSelect().
-	Model(&models.Tinyurl{}).
-	Where("user_id = ?", body.UserID).
-	Where("is_deleted = ?", false).
-	Count(ctx)
-
-    if err != nil {
-        ctx.JSON(http.StatusInternalServerError, dtos.URLCreationResponse{
-            Message: "Failed to fetch updated URL count",
-        })
-        return
-    }
-
-    ctx.JSON(http.StatusOK, dtos.URLDeleteResponse{
-        Message:   "URL deleted",
-        URLCount:  updatedCount,
-    })
+		Model(&models.Tinyurl{}).
+		Where("user_id = ?", userID).
+		Where("is_deleted = ?", false).
+		Count(ctx)
+
+	if err != nil {
+		ctx.JSON(http.StatusInternalServerError, dtos.URLCreationResponse{
+			Message: "Failed to fetch updated URL count",
+		})
+		return
+	}
+
+	ctx.JSON(http.StatusOK, dtos.URLDeleteResponse{
+		Message:  "URL deleted",
+		URLCount: updatedCount,
+	})
 }
 
 func GetURLDetails(ctx *gin.Context, db *bun.DB) {

diff --git a/middlewares/auth.go b/middlewares/auth.go
@@ -18,15 +18,24 @@ func AuthMiddleware() gin.HandlerFunc {
 
 		token := tokenCookie.Value
 
-		email, err := utils.VerifyToken(token)
+		claims, err := utils.VerifyToken(token)
 
 		if err != nil {
 			ctx.JSON(http.StatusUnauthorized, gin.H{"message": "Unauthorized"})
 			ctx.Abort()
 			return
 		}
 
-		ctx.Set("user", email)
+		userID, ok := claims["userID"].(float64)
+
+        if !ok {
+            ctx.JSON(http.StatusUnauthorized, gin.H{"message": "Invalid UserID format"})
+            ctx.Abort()
+            return
+        }
+
+		ctx.Set("user", claims["email"])
+		ctx.Set("userID", int64(userID))
 		ctx.Next()
 	}
-}
+}
diff --git a/routes/url.go b/routes/url.go
@@ -29,7 +29,7 @@ func TinyURLRoutes(rg *gin.RouterGroup, db *bun.DB) {
 	redirect.GET("/:shortURL", func(ctx *gin.Context) {
 		controller.RedirectShortURL(ctx, db)
 	})
-	urls.DELETE("/:id", func(ctx *gin.Context) {
+	urls.DELETE("/:id", middleware.AuthMiddleware(), func(ctx *gin.Context) {
 		controller.DeleteURL(ctx, db)
-	})
+	})	
 }
diff --git a/tests/integration/url_test.go b/tests/integration/url_test.go
@@ -15,6 +15,12 @@ import (
 func (suite *AppTestSuite) TestCreateTinyURLSuccess() {
 	// Setup the router and route for creating a tiny URL
 	router := gin.Default()
+
+	router.Use(func(ctx *gin.Context) {
+		ctx.Set("userID", int64(1))
+		ctx.Next()
+	})
+
 	router.POST("/v1/tinyurl", func(ctx *gin.Context) {
 		controller.CreateTinyURL(ctx, suite.db)
 	})
@@ -81,14 +87,19 @@ func (suite *AppTestSuite) TestCreateTinyURLEmptyOriginalURL() {
 // TestCreateTinyURLCustomShortURL tests the creation of a tiny URL with a custom short URL and expects a successful response.
 func (suite *AppTestSuite) TestCreateTinyURLCustomShortURL() {
 	router := gin.Default()
+
+	router.Use(func(ctx *gin.Context) {
+		ctx.Set("userID", int64(1))
+		ctx.Next()
+	})
+
 	router.POST("/v1/tinyurl", func(ctx *gin.Context) {
 		controller.CreateTinyURL(ctx, suite.db)
 	})
 
 	requestBody := map[string]interface{}{
 		"OriginalUrl": "https://example.com",
 		"ShortUrl":    "short",
-		"UserId":      1,
 	}
 	requestJSON, _ := json.Marshal(requestBody)
 	req, _ := http.NewRequest("POST", "/v1/tinyurl", bytes.NewBuffer(requestJSON))
@@ -102,14 +113,19 @@ func (suite *AppTestSuite) TestCreateTinyURLCustomShortURL() {
 
 func (suite *AppTestSuite) TestCreateTinyURLCustomShortURLExists() {
 	router := gin.Default()
+
+	router.Use(func(ctx *gin.Context) {
+		ctx.Set("userID", int64(1))
+		ctx.Next()
+	})
+
 	router.POST("/v1/tinyurl", func(ctx *gin.Context) {
 		controller.CreateTinyURL(ctx, suite.db)
 	})
 
 	requestBody := map[string]interface{}{
 		"OriginalUrl": "https://rds.com",
 		"ShortUrl":    "37fff",
-		"UserId":      1,
 	}
 	requestJSON, _ := json.Marshal(requestBody)
 	req, _ := http.NewRequest("POST", "/v1/tinyurl", bytes.NewBuffer(requestJSON))
@@ -123,6 +139,12 @@ func (suite *AppTestSuite) TestCreateTinyURLCustomShortURLExists() {
 
 func (suite *AppTestSuite) TestCreateTinyURLExistingOriginalURL() {
     router := gin.Default()
+
+	router.Use(func(ctx *gin.Context) {
+		ctx.Set("userID", int64(1))
+		ctx.Next()
+	})
+
     router.POST("/v1/tinyurl", func(ctx *gin.Context) {
         controller.CreateTinyURL(ctx, suite.db)
     })
@@ -131,7 +153,6 @@ func (suite *AppTestSuite) TestCreateTinyURLExistingOriginalURL() {
 
     requestBody := map[string]interface{}{
         "OriginalUrl": existingOriginalURL,
-        "UserId":      1,
     }
     requestJSON, _ := json.Marshal(requestBody)
     req, _ := http.NewRequest("POST", "/v1/tinyurl", bytes.NewBuffer(requestJSON))