package csrf import ( "crypto/rand" "crypto/subtle" "encoding/base64" "fmt" "time" "dk/internal/auth" "dk/internal/router" "github.com/valyala/fasthttp" ) const ( TokenLength = 32 TokenFieldName = "_csrf_token" SessionKey = "csrf_token" SessionCtxKey = "session" // Same as middleware.SessionKey CookieName = "_csrf" ) // GetCurrentSession retrieves the session from context (mirrors middleware function) func GetCurrentSession(ctx router.Ctx) *auth.Session { if session, ok := ctx.UserValue(SessionCtxKey).(*auth.Session); ok { return session } return nil } // GenerateToken creates a new CSRF token and stores it in the session or cookie func GenerateToken(ctx router.Ctx, authManager *auth.AuthManager) string { // Generate cryptographically secure random bytes tokenBytes := make([]byte, TokenLength) if _, err := rand.Read(tokenBytes); err != nil { // Fallback - this should never happen in practice return "" } token := base64.URLEncoding.EncodeToString(tokenBytes) // Store token in session if user is authenticated, otherwise use cookie if session := GetCurrentSession(ctx); session != nil { StoreToken(session, token) } else { // Store in cookie for guest users StoreTokenInCookie(ctx, token) } return token } // GetToken retrieves the current CSRF token from session or cookie, generating one if needed func GetToken(ctx router.Ctx, authManager *auth.AuthManager) string { session := GetCurrentSession(ctx) if session != nil { // Authenticated user - check session first if existingToken := GetStoredToken(session); existingToken != "" { return existingToken } } else { // Guest user - check cookie first if existingToken := GetTokenFromCookie(ctx); existingToken != "" { return existingToken } } // Generate new token if none exists return GenerateToken(ctx, authManager) } // ValidateToken verifies a CSRF token against the stored session or cookie token func ValidateToken(ctx router.Ctx, authManager *auth.AuthManager, submittedToken string) bool { if submittedToken == "" { return false } var storedToken string session := GetCurrentSession(ctx) if session != nil { // Authenticated user - get token from session storedToken = GetStoredToken(session) } else { // Guest user - get token from cookie storedToken = GetTokenFromCookie(ctx) } if storedToken == "" { return false // No stored token } // Use constant-time comparison to prevent timing attacks return subtle.ConstantTimeCompare([]byte(submittedToken), []byte(storedToken)) == 1 } // StoreToken saves a CSRF token in the session func StoreToken(session *auth.Session, token string) { if session.Data == nil { session.Data = make(map[string]any) } session.Data[SessionKey] = token } // GetStoredToken retrieves the CSRF token from session func GetStoredToken(session *auth.Session) string { if session.Data == nil { return "" } if token, ok := session.Data[SessionKey].(string); ok { return token } return "" } // RotateToken generates a new token and replaces the old one in the session func RotateToken(ctx router.Ctx, authManager *auth.AuthManager) string { session := GetCurrentSession(ctx) if session == nil { return "" } // Generate new token newToken := GenerateToken(ctx, authManager) return newToken } // HiddenField generates an HTML hidden input field with the CSRF token func HiddenField(ctx router.Ctx, authManager *auth.AuthManager) string { token := GetToken(ctx, authManager) if token == "" { return "" // No token available } return fmt.Sprintf(``, TokenFieldName, token) } // TokenMeta generates HTML meta tag for JavaScript access to CSRF token func TokenMeta(ctx router.Ctx, authManager *auth.AuthManager) string { token := GetToken(ctx, authManager) if token == "" { return "" } return fmt.Sprintf(``, token) } // ValidateFormToken is a convenience function to validate CSRF token from form data func ValidateFormToken(ctx router.Ctx, authManager *auth.AuthManager) bool { // Try to get token from form data tokenBytes := ctx.PostArgs().Peek(TokenFieldName) if len(tokenBytes) == 0 { // Try from query parameters as fallback tokenBytes = ctx.QueryArgs().Peek(TokenFieldName) } if len(tokenBytes) == 0 { return false } return ValidateToken(ctx, authManager, string(tokenBytes)) } // StoreTokenInCookie stores a CSRF token in a cookie for guest users func StoreTokenInCookie(ctx router.Ctx, token string) { cookie := &fasthttp.Cookie{} cookie.SetKey(CookieName) cookie.SetValue(token) cookie.SetHTTPOnly(true) cookie.SetSameSite(fasthttp.CookieSameSiteStrictMode) cookie.SetSecure(false) // Set to true in production with HTTPS cookie.SetExpire(time.Now().Add(24 * time.Hour)) // Expire in 24 hours cookie.SetPath("/") ctx.Response.Header.SetCookie(cookie) } // GetTokenFromCookie retrieves a CSRF token from cookie for guest users func GetTokenFromCookie(ctx router.Ctx) string { return string(ctx.Request.Header.Cookie(CookieName)) }