nfeeder/web/handlers_auth.go

package web

import (
	"context"
	"crypto/rand"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"log"
	"net/http"
	"nfeeder/db"
	"strconv"
	"time"

	"github.com/golang-jwt/jwt/v5"
	"github.com/jackc/pgx/v5/pgtype"
	"golang.org/x/crypto/bcrypt"
)

func (s *Server) handleRegister() http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		email := r.FormValue("email")
		password := r.FormValue("password")

		if email == "" || password == "" {
			WriteJSON(w, http.StatusBadRequest, ErrorResponse{Error: "Email and password required"})
			return
		}

		// hash password
		hashpw, err := bcrypt.GenerateFromPassword([]byte(password), bcrypt.DefaultCost)
		if err != nil {
			WriteJSON(w, http.StatusInternalServerError, ErrorResponse{Error: "Internal Error"})
			return
		}

		// create the user
		user, err := s.store.CreateUser(r.Context(), db.CreateUserParams{
			Email:    email,
			Password: string(hashpw),
		})
		if err != nil {
			// Log the actual error for yourself, send a generic one to the user
			log.Printf("Failed to create user: %v", err)
			WriteJSON(w, http.StatusBadRequest, ErrorResponse{Error: "Failed to create user"})
			return
		}

		s.issueToken(w, r, user.ID, false)
	}
}

func (s *Server) handleLogin() http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		email := r.FormValue("email")
		password := r.FormValue("password")

		user, err := s.store.GetUserByEmail(r.Context(), email)
		if err != nil {
			log.Printf("failed to get user: %v", err)
			WriteJSON(w, http.StatusUnauthorized, ErrorResponse{Error: "Invalid credentials"})
			return
		}

		// compare passwords
		err = bcrypt.CompareHashAndPassword([]byte(user.Password), []byte(password))
		if err != nil {
			log.Printf("incorrect password entered: %v", err)
			WriteJSON(w, http.StatusUnauthorized, ErrorResponse{Error: "Invalid credentials"})
			return
		}

		s.issueToken(w, r, user.ID, true)
	}
}

func (s *Server) handleLogout() http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		var req RefreshTokenReq

		err := json.NewDecoder(r.Body).Decode(&req)
		if err != nil {
			log.Printf("handleRefresh failed to parse request body: %v", err)
			WriteJSON(w, http.StatusBadRequest, ErrorResponse{Error: "Invalid JSON"})
			return
		}

		hash := sha256.Sum256([]byte(req.RefreshToken))
		hex_token := hex.EncodeToString(hash[:])

		tokenRecord, err := s.store.GetValidRefreshToken(r.Context(), hex_token)
		if err != nil {
			log.Printf("Logout succes, warning: token not found or expired: %v", err)
			WriteJSON(w, http.StatusOK, SuccessResponse{Message: "logout success"})
			return
		}

		err = s.store.DeleteRefreshToken(r.Context(), db.DeleteRefreshTokenParams{
			TokenHash: hex_token,
			UserID:    tokenRecord.UserID,
		})
		if err != nil {
			log.Printf("Failed to delete old refresh token: %v", err)
		}

		WriteJSON(w, http.StatusOK, SuccessResponse{
			Message: "logout success",
		})
	}
}

func (s *Server) handleRefresh() http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		var req RefreshTokenReq

		err := json.NewDecoder(r.Body).Decode(&req)
		if err != nil {
			log.Printf("handleRefresh failed to parse request body: %v", err)
			WriteJSON(w, http.StatusBadRequest, ErrorResponse{Error: "Invalid JSON"})
			return
		}

		hash := sha256.Sum256([]byte(req.RefreshToken))
		hex_token := hex.EncodeToString(hash[:])

		tokenRecord, err := s.store.GetValidRefreshToken(r.Context(), hex_token)
		if err != nil {
			log.Printf("Refresh failed - token not found or expired: %v", err)
			WriteJSON(w, http.StatusUnauthorized, ErrorResponse{Error: "Invalid or expired session"})
			return
		}

		err = s.store.DeleteRefreshToken(r.Context(), db.DeleteRefreshTokenParams{
			TokenHash: hex_token,
			UserID:    tokenRecord.UserID,
		})
		if err != nil {
			log.Printf("Failed to delete old refresh token: %v", err)
		}

		s.issueToken(w, r, tokenRecord.UserID, false)
	}
}

func (s *Server) issueToken(w http.ResponseWriter, r *http.Request, userID int64, checkLimit bool) {
	nowTime := time.Now()
	jwtExpireTime := nowTime.Add(24 * time.Hour)

	claims := jwt.RegisteredClaims{
		Issuer:    ISSUER,
		Subject:   strconv.FormatInt(userID, 10),
		ExpiresAt: jwt.NewNumericDate(jwtExpireTime),
		IssuedAt:  jwt.NewNumericDate(nowTime),
	}

	token := jwt.NewWithClaims(jwt.SigningMethodHS256, claims)
	tokenString, err := token.SignedString(s.jwtSecret)
	if err != nil {
		WriteJSON(w, http.StatusInternalServerError, ErrorResponse{
			Error: "Internal Server Error",
		})
		return
	}

	refreshToken := s.issueRefreshToken(r.Context(), userID, nowTime, checkLimit)
	if refreshToken == "" {
		WriteJSON(w, http.StatusInternalServerError, ErrorResponse{Error: "Internal Server Error"})
		return
	}

	WriteJSON(w, http.StatusOK, AuthResponse{
		AccessToken:  tokenString,
		RefreshToken: refreshToken,
	})
}

func (s *Server) issueRefreshToken(ctx context.Context, userID int64, issuedTime time.Time, checkLimit bool) string {
	refreshExpireTime := issuedTime.Add((24 * time.Hour) * 3)

	// Generate refresh token
	rawToken := rand.Text()
	hash := sha256.Sum256([]byte(rawToken))
	hex_token := hex.EncodeToString(hash[:])

	// Check the user has not got the session limit.
	if checkLimit {
		sessionCount, err := s.store.CountUserRefreshTokens(ctx, userID)
		if err != nil {
			log.Printf("failed to count refresh token: %v", err)
			return ""
		}

		// Each user can only have 3 sessions, if they have 3 refresh tokens
		// remove the oldest one and then create a new refresh token
		if sessionCount >= int64(MAX_USER_SESSIONS) {
			err = s.store.DeleteOldestRefreshToken(ctx, userID)
			if err != nil {
				log.Printf("failed to delete oldest refresh token: %v", err)
				return ""
			}
		}
	}

	_, err := s.store.CreateRefreshToken(ctx, db.CreateRefreshTokenParams{
		UserID:    userID,
		TokenHash: hex_token,
		ExpiresAt: pgtype.Timestamptz{Time: refreshExpireTime, Valid: true},
	})
	if err != nil {
		// Log the actual error for yourself, send a generic one to the user
		log.Printf("failed to create refresh token: %v", err)
		return ""
	}

	return rawToken
}