From b7c731a37ebb5f8c1cd99c6a420b5a5b6de19e89 Mon Sep 17 00:00:00 2001
From: Preetinder Singh <preetindersingh13per@gmail.com>
Date: Thu, 20 Jun 2024 01:03:41 +0530
Subject: [PATCH] Auth Updated

---
 backend/main.go          |  5 ++---
 backend/pkg/auth/auth.go | 47 +++++++++++++++++++++++-----------------
 2 files changed, 29 insertions(+), 23 deletions(-)

diff --git a/backend/main.go b/backend/main.go
index 065786d..8e025ce 100644
--- a/backend/main.go
+++ b/backend/main.go
@@ -31,10 +31,9 @@ func main() {
 	fmt.Println("Redis [PING]: ", res)
 
 	r := mux.NewRouter()
-
 	r.HandleFunc("/links/all", router.GetAllLinks).Methods(http.MethodOptions, http.MethodGet)
-	r.HandleFunc("/generate/jwt", auth.GenerateJWT).Methods(http.MethodOptions, http.MethodGet)
-	r.HandleFunc("/validate/jwt", auth.ValidateJWT).Methods(http.MethodOptions, http.MethodGet)
+	r.HandleFunc("/generate-token", auth.GenerateJWT).Methods(http.MethodOptions, http.MethodGet)
+	r.Handle("/login", auth.TokenRequired(http.HandlerFunc(auth.ProtectedRoute))).Methods(http.MethodOptions, http.MethodGet)
 	r.HandleFunc("/register", auth.Register).Methods(http.MethodOptions, http.MethodPost)
 	r.HandleFunc("/show/users", auth.ShowUsers).Methods(http.MethodOptions, http.MethodGet)
 	r.HandleFunc("/health", func(w http.ResponseWriter, r *http.Request) {
diff --git a/backend/pkg/auth/auth.go b/backend/pkg/auth/auth.go
index 9ef9c48..c1ae036 100644
--- a/backend/pkg/auth/auth.go
+++ b/backend/pkg/auth/auth.go
@@ -42,40 +42,47 @@ func GenerateJWT(w http.ResponseWriter, r *http.Request) {
 	}
 
 	w.Header().Set("Content-Type", "application/json")
-	w.Write([]byte(`{"token": "` + tokenString + `"}`))
-	// return tokenString
+	json.NewEncoder(w).Encode(map[string]string{"token": tokenString})
 }
 
-func ValidateJWT(w http.ResponseWriter, r *http.Request) {
-	tokenString := r.URL.Query().Get("token")
+func ValidateJWT(tokenString string) (string, error) {
 	claims := &Claims{}
 	token, err := jwt.ParseWithClaims(tokenString, claims, func(token *jwt.Token) (interface{}, error) {
 		return jwtKey, nil
 	})
-
 	if err != nil {
 		if err == jwt.ErrSignatureInvalid {
-			http.Error(w, "Invalid token signature", http.StatusUnauthorized)
-			return
+			return "", http.ErrBodyNotAllowed
 		}
-		http.Error(w, "Invalid token", http.StatusBadRequest)
-		return
+		return "", err
 	}
-
 	if !token.Valid {
-		http.Error(w, "Invalid token", http.StatusUnauthorized)
-		return
+		return "", http.ErrBodyNotAllowed
 	}
-	w.Header().Set("Content-Type", "application/json")
-	json.NewEncoder(w).Encode(map[string]interface{}{
-		"username":  claims.Username,
-		"id":        claims.Id,
-		"issuer":    claims.Issuer,
-		"expiresAt": claims.ExpiresAt,
-		"issuedAt":  claims.IssuedAt,
+	return claims.Username, nil
+}
+
+func TokenRequired(next http.Handler) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		tokenString := r.Header.Get("Authorization")
+		if tokenString == "" {
+			http.Error(w, "Token is missing", http.StatusForbidden)
+			return
+		}
+
+		username, err := ValidateJWT(tokenString)
+		if err != nil {
+			http.Error(w, err.Error(), http.StatusForbidden)
+			return
+		}
+
+		r.Header.Set("username", username)
+		next.ServeHTTP(w, r)
 	})
+}
 
-	// return claims, nil
+func ProtectedRoute(w http.ResponseWriter, r *http.Request) {
+	http.Redirect(w, r, "http://localhost:4000/admin", http.StatusSeeOther)
 }
 
 func Register(w http.ResponseWriter, r *http.Request) {