vimalloc / flask-jwt-extended

An open source Flask extension that provides JWT support (with batteries included)!
http://flask-jwt-extended.readthedocs.io/en/stable/
MIT License
1.55k stars 239 forks source link

@jwt.unauthorized_loader error handler doesn't work #555

Open pak-app opened 2 months ago

pak-app commented 2 months ago

Hi. This Poorya. I wrote an error handler for my Flask application. This application uses Flask JwT Extended and Flask Socket-io libraries for authentication real-time communication. My error handler doesn't work and gives an error.(@vimalloc, could you please take a look at this issue?) My code:

from flask import Flask, jsonify, request
from flask_sqlalchemy import SQLAlchemy
from flask_migrate import Migrate
from flask_jwt_extended import JWTManager, jwt_required, create_access_token, get_jwt_identity
from flask import Flask
from flask_socketio import SocketIO, emit, disconnect
from flask_jwt_extended.exceptions import (
    NoAuthorizationError,
    InvalidHeaderError,
)

# Initialize Flask app
app = Flask(__name__)

port = xxxx
host = "xxx.xxx.xxx.xxxx"

# Configurations
app.config['SECRET_KEY'] = "sdkjfh234"
app.config['SQLALCHEMY_DATABASE_URI'] = 'sqlite:///users.db'
app.config['SQLALCHEMY_TRACK_MODIFICATIONS'] = False
app.config['JWT_SECRET_KEY'] = 'sdkjfh234'  # Change this to a random secret
app.config['PROPAGATE_EXCEPTIONS'] = True

# Initialize extensions
db = SQLAlchemy(app)
migrate = Migrate(app, db)
jwt = JWTManager(app)
socketio = SocketIO(app)

# User model
class User(db.Model):
    id = db.Column(db.Integer, primary_key=True)
    username = db.Column(db.String(80), unique=True, nullable=False)
    password = db.Column(db.String(128), nullable=False)

# Route for user registration
@app.route('/signup', methods=['POST'])
def signup():
    data = request.get_json()
    username = data.get('username')
    password = data.get('password')

    # Check if username already exists
    if User.query.filter_by(username=username).first():
        return jsonify({"msg": "Username already exists"}), 400

    # Create a new user
    new_user = User(username=username, password=password)  # Hash the password in a real app
    db.session.add(new_user)
    db.session.commit()
    return jsonify({"msg": "User created successfully"}), 201

# Route for user login
@app.route('/login', methods=['POST'])
def login():
    data = request.get_json()
    username = data.get('username')
    password = data.get('password')

    user = User.query.filter_by(username=username).first()

    # Check if user exists and password is correct
    if user and user.password == password:  # Hash comparison needed in a real app
        access_token = create_access_token(identity={'username': user.username})
        return jsonify(access_token=access_token), 200
    return jsonify({"msg": "Bad username or password"}), 401

# Protected route
@app.route('/protected', methods=['GET'])
@jwt_required
def protected():
    current_user = get_jwt_identity()
    return jsonify(logged_in_as=current_user), 200

# @jwt.unauthorized_loader
# def handle_no_authorization_error(error):
#     print(error)
#     print("***************************************ERROR********************************")

#     return jsonify({"error": "Permission Denied"}), 403

# @jwt.invalid_token_loader
# def handle_no_authorization_error(error):
#     print(error)
#     print("***************************************ERROR********************************")

#     return jsonify({"error": "Permission Denied"}), 403

# @app.errorhandler(InvalidHeaderError)
# def handle_no_authorization_error(error):
#     print(error)
#     print("***************************************ERROR********************************")
#     return jsonify({"error": "Permission Denied"}), 403

# Handle missing or invalid JWTs
@jwt.unauthorized_loader
def handle_no_authorization_error(error):
    return jsonify({"error": "Permission Denied"}), 403

@jwt.invalid_token_loader
def handle_invalid_token_error(error):
    return jsonify({"error": "Invalid Token"}), 403

# SocketIO event for authenticated connections
@socketio.on('connect')
@jwt_required()
def handle_connect():
    try:    
        print("connected!!!")
        identity = get_jwt_identity()
        print(f"User {identity} connected")
        emit('message', {'message': 'User connected'})

    except Exception as e:
        print(e)
        disconnect()
        return False

    except jwt.invalid_token_loader as e:
        print("Error:", e)
        disconnect(namespace="message")
        return jsonify({"Error": "Invalid Token!!!"}), 403

    except jwt.token_verification_failed_loader as e:
        print("Error:", e)
        disconnect(namespace="message")
        return jsonify({"Error": "Token Verification Failed!!!"}), 403

# Websocket on message 
@socketio.on("message")
# @jwt_required()
def handle_message(message):

    print("Recieved message:", message)
    emit('message', message)

# Main entry point
if __name__ == '__main__':
    # db.create_all()  # Create database tables
    # with app.app_context():
    #     app.run(debug=True, host=host, port=port)
    socketio.run(app, host=host, port=port, debug=True)

And the error I got:

message handler error
Traceback (most recent call last):
  File "/root/Workspace/truck-sensors-dashboard/.env/lib/python3.12/site-packages/engineio/server.py", line 434, in run_handler
    return self.handlers[event](*args)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/root/Workspace/truck-sensors-dashboard/.env/lib/python3.12/site-packages/socketio/server.py", line 643, in _handle_eio_message
    self._handle_connect(eio_sid, pkt.namespace, pkt.data)
  File "/root/Workspace/truck-sensors-dashboard/.env/lib/python3.12/site-packages/socketio/server.py", line 537, in _handle_connect
    success = self._trigger_event(
              ^^^^^^^^^^^^^^^^^^^^
  File "/root/Workspace/truck-sensors-dashboard/.env/lib/python3.12/site-packages/socketio/server.py", line 614, in _trigger_event
    return handler(*args)
           ^^^^^^^^^^^^^^
  File "/root/Workspace/truck-sensors-dashboard/.env/lib/python3.12/site-packages/flask_socketio/__init__.py", line 282, in _handler
    return self._handle_event(handler, message, namespace, sid,
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/root/Workspace/truck-sensors-dashboard/.env/lib/python3.12/site-packages/flask_socketio/__init__.py", line 823, in _handle_event
    ret = handler(auth)
          ^^^^^^^^^^^^^
  File "/root/Workspace/truck-sensors-dashboard/.env/lib/python3.12/site-packages/flask_jwt_extended/view_decorators.py", line 167, in decorator
    verify_jwt_in_request(
  File "/root/Workspace/truck-sensors-dashboard/.env/lib/python3.12/site-packages/flask_jwt_extended/view_decorators.py", line 94, in verify_jwt_in_request
    jwt_data, jwt_header, jwt_location = _decode_jwt_from_request(
                                         ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/root/Workspace/truck-sensors-dashboard/.env/lib/python3.12/site-packages/flask_jwt_extended/view_decorators.py", line 358, in _decode_jwt_from_request
    raise NoAuthorizationError(errors[0])
flask_jwt_extended.exceptions.NoAuthorizationError: Missing Authorization Header