stable-diffusion.cpp-rest/include/auth_middleware.h

#ifndef AUTH_MIDDLEWARE_H
#define AUTH_MIDDLEWARE_H

#include <string>
#include <vector>
#include <functional>
#include <memory>
#include "jwt_auth.h"
#include "user_manager.h"
#include "server_config.h"

#ifdef ENABLE_PAM_AUTH
#include "pam_auth.h"
#endif

namespace httplib {
    class Request;
    class Response;
}

/**
 * @brief Authentication context structure
 */
struct AuthContext {
    bool authenticated;                    ///< Authentication status
    std::string userId;                     ///< User ID
    std::string username;                   ///< Username
    std::string role;                       ///< User role
    std::vector<std::string> permissions;  ///< User permissions
    std::string authMethod;                 ///< Authentication method used
    std::string errorMessage;              ///< Error message if authentication failed
    std::string errorCode;                 ///< Error code for API responses
};


/**
 * @brief Authentication middleware class
 *
 * This class provides authentication and authorization middleware for HTTP requests.
 * It supports PAM and shared key authentication methods and
 * role-based access control.
 */
class AuthMiddleware {
public:
    /**
     * @brief Authentication handler function type
     *
     * @param req HTTP request
     * @param res HTTP response
     * @param context Authentication context
     */
    using AuthHandler = std::function<void(const httplib::Request& req,
                                           httplib::Response& res,
                                           const AuthContext& context)>;

    /**
     * @brief Construct a new Auth Middleware object
     *
     * @param config Authentication configuration
     * @param userManager User manager instance
     */
    explicit AuthMiddleware(AuthConfig config,
                           std::shared_ptr<UserManager> userManager);

    /**
     * @brief Destroy the Auth Middleware object
     */
    ~AuthMiddleware();

    /**
     * @brief Initialize the authentication middleware
     *
     * @return true if initialization successful, false otherwise
     */
    bool initialize();

    /**
     * @brief Authenticate HTTP request
     *
     * @param req HTTP request
     * @param res HTTP response
     * @return AuthContext Authentication context
     */
    AuthContext authenticate(const httplib::Request& req, httplib::Response& res);

    /**
     * @brief Check if path requires authentication
     *
     * @param path Request path
     * @return true if authentication required, false otherwise
     */
    bool requiresAuthentication(const std::string& path) const;

    /**
     * @brief Check if path requires admin access
     *
     * @param path Request path
     * @return true if admin access required, false otherwise
     */
    bool requiresAdminAccess(const std::string& path) const;

    /**
     * @brief Check if path requires user access (any authenticated user)
     *
     * @param path Request path
     * @return true if user access required, false otherwise
     */
    bool requiresUserAccess(const std::string& path) const;

    /**
     * @brief Check if user has permission for path
     *
     * @param path Request path
     * @param permissions User permissions
     * @return true if user has access, false otherwise
     */
    bool hasPathAccess(const std::string& path,
                      const std::vector<std::string>& permissions) const;

    /**
     * @brief Create authentication middleware handler
     *
     * @param handler Next handler in chain
     * @return AuthHandler Middleware handler function
     */
    AuthHandler createMiddleware(AuthHandler handler);

    /**
     * @brief Send authentication error response
     *
     * @param res HTTP response
     * @param message Error message
     * @param errorCode Error code
     * @param statusCode HTTP status code
     */
    void sendAuthError(httplib::Response& res,
                      const std::string& message,
                      const std::string& errorCode = "AUTH_ERROR",
                      int statusCode = 401);

    /**
     * @brief Send authorization error response
     *
     * @param res HTTP response
     * @param message Error message
     * @param errorCode Error code
     */
    void sendAuthzError(httplib::Response& res,
                       const std::string& message,
                       const std::string& errorCode = "ACCESS_DENIED");

    /**
     * @brief Add public path (no authentication required)
     *
     * @param path Path to add
     */
    void addPublicPath(const std::string& path);

    /**
     * @brief Add admin-only path
     *
     * @param path Path to add
     */
    void addAdminPath(const std::string& path);

    /**
     * @brief Add user-only path
     *
     * @param path Path to add
     */
    void addUserPath(const std::string& path);

    /**
     * @brief Set JWT secret
     *
     * @param secret JWT secret key
     */
    void setJwtSecret(const std::string& secret);

    /**
     * @brief Get JWT secret
     *
     * @return std::string JWT secret key
     */
    std::string getJwtSecret() const;

    /**
     * @brief Set authentication method
     *
     * @param method Authentication method
     */
    void setAuthMethod(UserManager::AuthMethod method);

    /**
     * @brief Get authentication method
     *
     * @return UserManager::AuthMethod Current authentication method
     */
    UserManager::AuthMethod getAuthMethod() const;

    /**
     * @brief Enable or disable guest access
     *
     * @param enable Enable guest access
     */
    void setGuestAccessEnabled(bool enable);

    /**
     * @brief Check if guest access is enabled
     *
     * @return true if guest access enabled, false otherwise
     */
    bool isGuestAccessEnabled() const;

    /**
     * @brief Get authentication configuration
     *
     * @return AuthConfig Current configuration
     */
    AuthConfig getConfig() const;

    /**
     * @brief Update authentication configuration
     *
     * @param config New configuration
     */
    void updateConfig(AuthConfig config);

private:
    AuthConfig m_config;                           ///< Authentication configuration
    std::shared_ptr<UserManager> m_userManager;    ///< User manager instance
    std::unique_ptr<JWTAuth> m_jwtAuth;            ///< JWT authentication instance

    /**
     * @brief Authenticate using JWT token
     *
     * @param req HTTP request
     * @return AuthContext Authentication context
     */
    AuthContext authenticateJwt(const httplib::Request& req);


    /**
     * @brief Authenticate using PAM
     *
     * @param req HTTP request
     * @return AuthContext Authentication context
     */
    AuthContext authenticatePam(const httplib::Request& req);

    /**
     * @brief Authenticate using shared key
     *
     * @param req HTTP request
     * @return AuthContext Authentication context
     */
    AuthContext authenticateSharedKey(const httplib::Request& req);

    /**
     * @brief Extract token from request
     *
     * @param req HTTP request
     * @param headerName Header name to check
     * @return std::string Token string, empty if not found
     */
    std::string extractToken(const httplib::Request& req, const std::string& headerName) const;

    /**
     * @brief Create guest authentication context
     *
     * @return AuthContext Guest context
     */
    AuthContext createGuestContext() const;

    /**
     * @brief Check if path matches pattern
     *
     * @param path Request path
     * @param patterns List of patterns to match
     * @return true if path matches any pattern, false otherwise
     */
    static bool pathMatchesPattern(const std::string& path,
                                  const std::vector<std::string>& patterns);

    /**
     * @brief Get required permissions for path
     *
     * @param path Request path
     * @return std::vector<std::string> Required permissions
     */
    std::vector<std::string> getRequiredPermissions(const std::string& path) const;

    /**
     * @brief Log authentication attempt
     *
     * @param req HTTP request
     * @param context Authentication context
     * @param success Authentication success
     */
    void logAuthAttempt(const httplib::Request& req,
                       const AuthContext& context,
                       bool success) const;

    /**
     * @brief Get client IP address from request
     *
     * @param req HTTP request
     * @return std::string Client IP address
     */
    static std::string getClientIp(const httplib::Request& req);

    /**
     * @brief Get user agent from request
     *
     * @param req HTTP request
     * @return std::string User agent string
     */
    static std::string getUserAgent(const httplib::Request& req);

    /**
     * @brief Validate authentication configuration
     *
     * @param config Configuration to validate
     * @return true if valid, false otherwise
     */
    static bool validateConfig(AuthConfig config);

    /**
     * @brief Initialize default paths
     */
    void initializeDefaultPaths();

    /**
     * @brief Check if authentication is completely disabled
     *
     * @return true if authentication disabled, false otherwise
     */
    bool isAuthenticationDisabled() const;
};

/**
 * @brief Authentication middleware factory functions
 */
namespace AuthMiddlewareFactory {
    /**
     * @brief Create authentication middleware with default configuration
     *
     * @param userManager User manager instance
     * @param dataDir Data directory for user storage
     * @return std::unique_ptr<AuthMiddleware> Auth middleware instance
     */
    std::unique_ptr<AuthMiddleware> createDefault(std::shared_ptr<UserManager> userManager,
                                                  const std::string& dataDir);

    /**
     * @brief Create authentication middleware with shared key only
     *
     * @param userManager User manager instance
     * @param sharedKey Shared key for authentication
     * @return std::unique_ptr<AuthMiddleware> Auth middleware instance
     */
    std::unique_ptr<AuthMiddleware> createSharedKeyOnly(std::shared_ptr<UserManager> userManager,
                                                        const std::string& sharedKey);

    /**
     * @brief Create authentication middleware with multiple methods
     *
     * @param userManager User manager instance
     * @param config Authentication configuration
     * @return std::unique_ptr<AuthMiddleware> Auth middleware instance
     */
    std::unique_ptr<AuthMiddleware> createMultiMethod(std::shared_ptr<UserManager> userManager,
                                                      AuthConfig config);

    /**
     * @brief Create authentication middleware for development (no auth required)
     *
     * @return std::unique_ptr<AuthMiddleware> Auth middleware instance
     */
    std::unique_ptr<AuthMiddleware> createDevelopment();
};

#endif // AUTH_MIDDLEWARE_H