Fix login redirect on Hugging Face

app.py

@@ -7,6 +7,7 @@ from uuid import uuid4
 import numpy as np
 from flask import Flask, flash, redirect, render_template, request, session, url_for
 from PIL import Image
+from werkzeug.middleware.proxy_fix import ProxyFix
 from werkzeug.security import check_password_hash, generate_password_hash
 from werkzeug.utils import secure_filename
 
@@ -72,6 +73,14 @@ TRAFFIC_SIGN_CLASSES = [
 app = Flask(__name__)
 app.config["SECRET_KEY"] = os.environ.get("SECRET_KEY", "dev-change-me")
 app.config["MAX_CONTENT_LENGTH"] = 8 * 1024 * 1024
+app.wsgi_app = ProxyFix(app.wsgi_app, x_for=1, x_proto=1, x_host=1, x_prefix=1)
+
+if os.environ.get("SPACE_ID"):
+    app.config["SESSION_COOKIE_SECURE"] = True
+    app.config["SESSION_COOKIE_SAMESITE"] = "None"
+else:
+    app.config["SESSION_COOKIE_SAMESITE"] = "Lax"
+
 UPLOAD_DIR.mkdir(parents=True, exist_ok=True)
 DATABASE_PATH.parent.mkdir(parents=True, exist_ok=True)
 
@@ -138,7 +147,7 @@ def login_required(view):
     def wrapped(*args, **kwargs):
         if not current_user():
             flash("Please log in to access the classifier.", "warning")
-            return redirect(url_for("login"))
+            return redirect(url_for("login", next=request.path))
         return view(*args, **kwargs)
 
     wrapped.__name__ = view.__name__
@@ -178,6 +187,9 @@ def welcome():
 
 @app.route("/register", methods=["GET", "POST"])
 def register():
+    if current_user():
+        return redirect(url_for("predict"))
+
     if request.method == "POST":
         name = request.form.get("name", "").strip()
         email = request.form.get("email", "").strip().lower()
@@ -201,8 +213,10 @@ def register():
                         (name, email, generate_password_hash(password), datetime.utcnow().isoformat()),
                     )
                     session["user_id"] = cursor.lastrowid
+                    session.permanent = True
+                    session.modified = True
                 flash("Account created. Welcome to the classifier.", "success")
-                return redirect(url_for("predict"))
+                return redirect(request.args.get("next") or url_for("predict"))
             except sqlite3.IntegrityError:
                 flash("This email is already registered.", "danger")
 
@@ -211,6 +225,9 @@ def register():
 
 @app.route("/login", methods=["GET", "POST"])
 def login():
+    if current_user():
+        return redirect(url_for("predict"))
+
     if request.method == "POST":
         email = request.form.get("email", "").strip().lower()
         password = request.form.get("password", "")
@@ -218,8 +235,10 @@ def login():
             user = conn.execute("SELECT * FROM users WHERE email = ?", (email,)).fetchone()
         if user and check_password_hash(user["password_hash"], password):
             session["user_id"] = user["id"]
+            session.permanent = True
+            session.modified = True
             flash("Connection established.", "success")
-            return redirect(url_for("predict"))
+            return redirect(request.args.get("next") or url_for("predict"))
         flash("Invalid email or password.", "danger")
     return render_template("login.html")