430 lines
16 KiB
Python
430 lines
16 KiB
Python
|
"""
|
||
|
|
Copyright 2024-2026 ReplyPier
|
||
|
|
|
||
|
|
Licensed under the Apache License, Version 2.0 (the "License");
|
||
|
|
you may not use this file except in compliance with the License.
|
||
|
|
You may obtain a copy of the License at
|
||
|
|
|
||
|
|
http://www.apache.org/licenses/LICENSE-2.0
|
||
|
|
|
||
|
|
Unless required by applicable law or agreed to in writing, software
|
||
|
|
distributed under the License is distributed on an "AS IS" BASIS,
|
||
|
|
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||
|
|
See the License for the specific language governing permissions and
|
||
|
|
limitations under the License.
|
||
|
|
"""
|
||
|
|
|
||
|
|
import pytest
|
||
|
|
import redis
|
||
|
|
import asyncio
|
||
|
|
from unittest.mock import patch, MagicMock, AsyncMock
|
||
|
|
from app.services.socket_rate_limit import socket_rate_limit
|
||
|
|
from app.core.config import settings
|
||
|
|
from app.core.socketio import sio
|
||
|
|
|
||
|
|
# Test constants
|
||
|
|
TEST_SID = "test_session_id"
|
||
|
|
TEST_NAMESPACE = "/widget"
|
||
|
|
TEST_IP = "192.168.1.1"
|
||
|
|
TEST_AGENT_ID = "test_agent_123"
|
||
|
|
|
||
|
|
@pytest.fixture
|
||
|
|
def mock_redis():
|
||
|
|
with patch('app.services.socket_rate_limit.redis_client') as mock_client:
|
||
|
|
# Setup Redis mock to work with asyncio.wait_for
|
||
|
|
mock_client.get = MagicMock()
|
||
|
|
mock_client.setex = MagicMock()
|
||
|
|
mock_client.incr = MagicMock()
|
||
|
|
mock_client.ttl = MagicMock()
|
||
|
|
yield mock_client
|
||
|
|
|
||
|
|
@pytest.fixture
|
||
|
|
def mock_sio():
|
||
|
|
with patch('app.services.socket_rate_limit.sio') as mock_sio:
|
||
|
|
# Setup session data
|
||
|
|
mock_sio.get_session = AsyncMock(return_value={
|
||
|
|
'enable_rate_limiting': True,
|
||
|
|
'overall_limit_per_ip': 100,
|
||
|
|
'requests_per_sec': 1.0,
|
||
|
|
'agent_id': TEST_AGENT_ID
|
||
|
|
})
|
||
|
|
# Setup environment data with load balancer headers
|
||
|
|
mock_sio.get_environ.return_value = {
|
||
|
|
'REMOTE_ADDR': '10.0.0.1', # Load balancer IP
|
||
|
|
'HTTP_X_FORWARDED_FOR': f'{TEST_IP}, 10.0.0.1' # Original client IP, load balancer IP
|
||
|
|
}
|
||
|
|
# Make emit async
|
||
|
|
mock_sio.emit = AsyncMock()
|
||
|
|
yield mock_sio
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
async def test_socket_rate_limit_disabled():
|
||
|
|
"""Test when rate limiting is disabled"""
|
||
|
|
# Create a mock handler function
|
||
|
|
mock_handler = AsyncMock()
|
||
|
|
|
||
|
|
# Create a mock session without rate limiting
|
||
|
|
mock_session = {
|
||
|
|
'enable_rate_limiting': False
|
||
|
|
}
|
||
|
|
|
||
|
|
# Mock sio.get_session
|
||
|
|
with patch('app.services.socket_rate_limit.sio.get_session', AsyncMock(return_value=mock_session)):
|
||
|
|
# Apply the decorator
|
||
|
|
decorated_handler = socket_rate_limit()(mock_handler)
|
||
|
|
|
||
|
|
# Call the decorated handler
|
||
|
|
await decorated_handler('test_sid', 'test_arg')
|
||
|
|
|
||
|
|
# Verify the handler was called
|
||
|
|
mock_handler.assert_called_once_with('test_sid', 'test_arg')
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
async def test_socket_rate_limit_redis_disabled():
|
||
|
|
"""Test when Redis is disabled"""
|
||
|
|
# Create a mock handler function
|
||
|
|
mock_handler = AsyncMock()
|
||
|
|
|
||
|
|
# Create a mock session with rate limiting
|
||
|
|
mock_session = {
|
||
|
|
'enable_rate_limiting': True,
|
||
|
|
'overall_limit_per_ip': 100,
|
||
|
|
'requests_per_sec': 1.0,
|
||
|
|
'agent_id': TEST_AGENT_ID
|
||
|
|
}
|
||
|
|
|
||
|
|
with patch('app.services.socket_rate_limit.settings.REDIS_ENABLED', False), \
|
||
|
|
patch('app.services.socket_rate_limit.sio.get_session', AsyncMock(return_value=mock_session)):
|
||
|
|
# Apply the decorator
|
||
|
|
decorated_handler = socket_rate_limit()(mock_handler)
|
||
|
|
|
||
|
|
# Call the decorated handler
|
||
|
|
await decorated_handler('test_sid', 'test_arg')
|
||
|
|
|
||
|
|
# Verify the handler was called
|
||
|
|
mock_handler.assert_called_once_with('test_sid', 'test_arg')
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
async def test_socket_rate_limit_with_forwarded_ip(mock_redis):
|
||
|
|
"""Test rate limiting with X-Forwarded-For header"""
|
||
|
|
# Create a mock handler function
|
||
|
|
mock_handler = AsyncMock()
|
||
|
|
|
||
|
|
# Create a mock session with rate limiting
|
||
|
|
mock_session = {
|
||
|
|
'enable_rate_limiting': True,
|
||
|
|
'overall_limit_per_ip': 100,
|
||
|
|
'requests_per_sec': 1.0,
|
||
|
|
'agent_id': TEST_AGENT_ID
|
||
|
|
}
|
||
|
|
|
||
|
|
# Mock environment with X-Forwarded-For
|
||
|
|
mock_environ = {
|
||
|
|
'HTTP_X_FORWARDED_FOR': '192.168.1.1, 10.0.0.1'
|
||
|
|
}
|
||
|
|
|
||
|
|
# Configure Redis mock
|
||
|
|
mock_redis.get.return_value = None # No previous requests
|
||
|
|
mock_redis.setex.return_value = True
|
||
|
|
mock_redis.incr.return_value = 1 # First request
|
||
|
|
mock_redis.expire.return_value = True
|
||
|
|
|
||
|
|
# Mock asyncio.wait_for to return Redis results directly
|
||
|
|
async def mock_wait_for(coro, timeout):
|
||
|
|
if isinstance(coro, asyncio.Future):
|
||
|
|
return await coro
|
||
|
|
return coro
|
||
|
|
|
||
|
|
with patch('app.services.socket_rate_limit.settings.REDIS_ENABLED', True), \
|
||
|
|
patch('app.services.socket_rate_limit.sio.get_session', AsyncMock(return_value=mock_session)), \
|
||
|
|
patch('app.services.socket_rate_limit.sio.get_environ', return_value=mock_environ), \
|
||
|
|
patch('asyncio.wait_for', side_effect=mock_wait_for):
|
||
|
|
# Apply the decorator
|
||
|
|
decorated_handler = socket_rate_limit()(mock_handler)
|
||
|
|
|
||
|
|
# Call the decorated handler
|
||
|
|
await decorated_handler('test_sid', 'test_arg')
|
||
|
|
|
||
|
|
# Verify the handler was called
|
||
|
|
mock_handler.assert_called_once_with('test_sid', 'test_arg')
|
||
|
|
|
||
|
|
# Verify Redis interactions
|
||
|
|
assert mock_redis.get.call_count >= 2 # Called for both overall and rate limits
|
||
|
|
assert mock_redis.setex.call_count >= 2 # Called for both overall and rate limits
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
async def test_socket_rate_limit_exceeded(mock_redis):
|
||
|
|
"""Test when rate limit is exceeded"""
|
||
|
|
# Create a mock handler function
|
||
|
|
mock_handler = AsyncMock()
|
||
|
|
|
||
|
|
# Create a mock session with rate limiting
|
||
|
|
mock_session = {
|
||
|
|
'enable_rate_limiting': True,
|
||
|
|
'overall_limit_per_ip': 100,
|
||
|
|
'requests_per_sec': 1.0,
|
||
|
|
'agent_id': TEST_AGENT_ID
|
||
|
|
}
|
||
|
|
|
||
|
|
# Mock environment
|
||
|
|
mock_environ = {
|
||
|
|
'REMOTE_ADDR': '192.168.1.1'
|
||
|
|
}
|
||
|
|
|
||
|
|
# Configure Redis mock to simulate exceeded limit
|
||
|
|
mock_redis.get.return_value = b'101' # Over the limit
|
||
|
|
mock_redis.ttl.return_value = 3600 # 1 hour remaining
|
||
|
|
|
||
|
|
# Mock asyncio.wait_for to return Redis results directly
|
||
|
|
async def mock_wait_for(coro, timeout):
|
||
|
|
if isinstance(coro, asyncio.Future):
|
||
|
|
return await coro
|
||
|
|
return coro
|
||
|
|
|
||
|
|
# Mock sio.emit to capture error messages
|
||
|
|
mock_emit = AsyncMock()
|
||
|
|
|
||
|
|
with patch('app.services.socket_rate_limit.settings.REDIS_ENABLED', True), \
|
||
|
|
patch('app.services.socket_rate_limit.sio.get_session', AsyncMock(return_value=mock_session)), \
|
||
|
|
patch('app.services.socket_rate_limit.sio.get_environ', return_value=mock_environ), \
|
||
|
|
patch('app.services.socket_rate_limit.sio.emit', mock_emit), \
|
||
|
|
patch('asyncio.wait_for', side_effect=mock_wait_for):
|
||
|
|
# Apply the decorator
|
||
|
|
decorated_handler = socket_rate_limit()(mock_handler)
|
||
|
|
|
||
|
|
# Call the decorated handler
|
||
|
|
result = await decorated_handler('test_sid', 'test_arg')
|
||
|
|
|
||
|
|
# Verify the handler was not called
|
||
|
|
mock_handler.assert_not_called()
|
||
|
|
|
||
|
|
# Verify error was emitted
|
||
|
|
mock_emit.assert_called_once()
|
||
|
|
args = mock_emit.call_args[0]
|
||
|
|
assert args[0] == 'error'
|
||
|
|
assert 'Daily request limit reached' in args[1]['error']
|
||
|
|
assert result is None
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
async def test_socket_rate_limit_with_real_ip(mock_redis):
|
||
|
|
"""Test rate limiting with X-Real-IP header"""
|
||
|
|
# Create a mock handler function
|
||
|
|
mock_handler = AsyncMock()
|
||
|
|
|
||
|
|
# Create a mock session with rate limiting
|
||
|
|
mock_session = {
|
||
|
|
'enable_rate_limiting': True,
|
||
|
|
'overall_limit_per_ip': 100,
|
||
|
|
'requests_per_sec': 1.0,
|
||
|
|
'agent_id': TEST_AGENT_ID
|
||
|
|
}
|
||
|
|
|
||
|
|
# Mock environment with X-Real-IP
|
||
|
|
mock_environ = {
|
||
|
|
'HTTP_X_REAL_IP': '192.168.1.1'
|
||
|
|
}
|
||
|
|
|
||
|
|
# Configure Redis mock
|
||
|
|
mock_redis.get.return_value = None # No previous requests
|
||
|
|
mock_redis.setex.return_value = True
|
||
|
|
mock_redis.incr.return_value = 1 # First request
|
||
|
|
mock_redis.expire.return_value = True
|
||
|
|
|
||
|
|
# Mock asyncio.wait_for to return Redis results directly
|
||
|
|
async def mock_wait_for(coro, timeout):
|
||
|
|
if isinstance(coro, asyncio.Future):
|
||
|
|
return await coro
|
||
|
|
return coro
|
||
|
|
|
||
|
|
with patch('app.services.socket_rate_limit.settings.REDIS_ENABLED', True), \
|
||
|
|
patch('app.services.socket_rate_limit.sio.get_session', AsyncMock(return_value=mock_session)), \
|
||
|
|
patch('app.services.socket_rate_limit.sio.get_environ', return_value=mock_environ), \
|
||
|
|
patch('asyncio.wait_for', side_effect=mock_wait_for):
|
||
|
|
# Apply the decorator
|
||
|
|
decorated_handler = socket_rate_limit()(mock_handler)
|
||
|
|
|
||
|
|
# Call the decorated handler
|
||
|
|
await decorated_handler('test_sid', 'test_arg')
|
||
|
|
|
||
|
|
# Verify the handler was called
|
||
|
|
mock_handler.assert_called_once_with('test_sid', 'test_arg')
|
||
|
|
|
||
|
|
# Verify Redis interactions
|
||
|
|
assert mock_redis.get.call_count >= 2 # Called for both overall and rate limits
|
||
|
|
assert mock_redis.setex.call_count >= 2 # Called for both overall and rate limits
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
async def test_socket_rate_limit_redis_error(mock_redis):
|
||
|
|
"""Test handling of Redis errors"""
|
||
|
|
# Create a mock handler function
|
||
|
|
mock_handler = AsyncMock()
|
||
|
|
|
||
|
|
# Create a mock session with rate limiting
|
||
|
|
mock_session = {
|
||
|
|
'enable_rate_limiting': True,
|
||
|
|
'overall_limit_per_ip': 100,
|
||
|
|
'requests_per_sec': 1.0,
|
||
|
|
'agent_id': TEST_AGENT_ID
|
||
|
|
}
|
||
|
|
|
||
|
|
# Mock environment
|
||
|
|
mock_environ = {
|
||
|
|
'REMOTE_ADDR': '192.168.1.1'
|
||
|
|
}
|
||
|
|
|
||
|
|
# Configure Redis mock to raise an error
|
||
|
|
mock_redis.get.side_effect = redis.RedisError("Connection error")
|
||
|
|
|
||
|
|
with patch('app.services.socket_rate_limit.settings.REDIS_ENABLED', True), \
|
||
|
|
patch('app.services.socket_rate_limit.sio.get_session', AsyncMock(return_value=mock_session)), \
|
||
|
|
patch('app.services.socket_rate_limit.sio.get_environ', return_value=mock_environ):
|
||
|
|
# Apply the decorator
|
||
|
|
decorated_handler = socket_rate_limit()(mock_handler)
|
||
|
|
|
||
|
|
# Call the decorated handler
|
||
|
|
await decorated_handler('test_sid', 'test_arg')
|
||
|
|
|
||
|
|
# Verify the handler was called (should proceed on Redis error)
|
||
|
|
mock_handler.assert_called_once_with('test_sid', 'test_arg')
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
async def test_localhost_bypass(mock_redis):
|
||
|
|
"""Test that localhost requests bypass rate limiting"""
|
||
|
|
# Test different localhost variations and headers
|
||
|
|
test_cases = [
|
||
|
|
# (header_name, is_list, localhost_ip)
|
||
|
|
('HTTP_X_FORWARDED_FOR', True, '127.0.0.1'),
|
||
|
|
('HTTP_X_FORWARDED_FOR', True, 'localhost'),
|
||
|
|
('HTTP_X_FORWARDED_FOR', True, '::1'),
|
||
|
|
('HTTP_X_REAL_IP', False, '127.0.0.1'),
|
||
|
|
('HTTP_X_REAL_IP', False, 'localhost'),
|
||
|
|
('HTTP_CF_CONNECTING_IP', False, '127.0.0.1'),
|
||
|
|
('REMOTE_ADDR', False, '127.0.0.1')
|
||
|
|
]
|
||
|
|
|
||
|
|
mock_session = {
|
||
|
|
'enable_rate_limiting': True,
|
||
|
|
'overall_limit_per_ip': 100,
|
||
|
|
'requests_per_sec': 1.0,
|
||
|
|
'agent_id': TEST_AGENT_ID
|
||
|
|
}
|
||
|
|
|
||
|
|
for header, is_list, ip in test_cases:
|
||
|
|
# Setup environment data
|
||
|
|
environ = {'REMOTE_ADDR': '10.0.0.1'} # Default remote addr
|
||
|
|
if is_list:
|
||
|
|
environ[header] = f"{ip}, 10.0.0.1" # Add load balancer IP
|
||
|
|
else:
|
||
|
|
environ[header] = ip
|
||
|
|
|
||
|
|
# Create and decorate handler
|
||
|
|
mock_handler = AsyncMock()
|
||
|
|
decorated_handler = socket_rate_limit()(mock_handler)
|
||
|
|
|
||
|
|
with patch('app.services.socket_rate_limit.settings.REDIS_ENABLED', True), \
|
||
|
|
patch('app.services.socket_rate_limit.sio.get_session', AsyncMock(return_value=mock_session)), \
|
||
|
|
patch('app.services.socket_rate_limit.sio.get_environ', return_value=environ):
|
||
|
|
# Call handler
|
||
|
|
await decorated_handler(TEST_SID)
|
||
|
|
|
||
|
|
# Verify handler was called without checking Redis
|
||
|
|
mock_handler.assert_called_once_with(TEST_SID)
|
||
|
|
mock_redis.get.assert_not_called()
|
||
|
|
mock_redis.setex.assert_not_called()
|
||
|
|
|
||
|
|
# Reset mocks for next iteration
|
||
|
|
mock_handler.reset_mock()
|
||
|
|
mock_redis.get.reset_mock()
|
||
|
|
mock_redis.setex.reset_mock()
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
async def test_redis_timeout(mock_redis):
|
||
|
|
"""Test behavior when Redis operations timeout"""
|
||
|
|
# Create a mock handler function
|
||
|
|
mock_handler = AsyncMock()
|
||
|
|
|
||
|
|
# Create a mock session with rate limiting
|
||
|
|
mock_session = {
|
||
|
|
'enable_rate_limiting': True,
|
||
|
|
'overall_limit_per_ip': 100,
|
||
|
|
'requests_per_sec': 1.0,
|
||
|
|
'agent_id': TEST_AGENT_ID
|
||
|
|
}
|
||
|
|
|
||
|
|
# Mock environment
|
||
|
|
mock_environ = {
|
||
|
|
'REMOTE_ADDR': '192.168.1.1'
|
||
|
|
}
|
||
|
|
|
||
|
|
# Configure Redis mock to timeout
|
||
|
|
mock_redis.get.side_effect = asyncio.TimeoutError()
|
||
|
|
|
||
|
|
# Mock asyncio.wait_for to raise TimeoutError
|
||
|
|
async def mock_wait_for(coro, timeout):
|
||
|
|
raise asyncio.TimeoutError()
|
||
|
|
|
||
|
|
with patch('app.services.socket_rate_limit.settings.REDIS_ENABLED', True), \
|
||
|
|
patch('app.services.socket_rate_limit.sio.get_session', AsyncMock(return_value=mock_session)), \
|
||
|
|
patch('app.services.socket_rate_limit.sio.get_environ', return_value=mock_environ), \
|
||
|
|
patch('asyncio.wait_for', side_effect=mock_wait_for):
|
||
|
|
# Apply the decorator
|
||
|
|
decorated_handler = socket_rate_limit()(mock_handler)
|
||
|
|
|
||
|
|
# Call the decorated handler
|
||
|
|
await decorated_handler('test_sid', 'test_arg')
|
||
|
|
|
||
|
|
# Verify handler was called despite timeout
|
||
|
|
mock_handler.assert_called_once_with('test_sid', 'test_arg')
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
async def test_rate_limit_per_second_exceeded(mock_redis):
|
||
|
|
"""Test when rate limit per second is exceeded"""
|
||
|
|
# Create a mock handler function
|
||
|
|
mock_handler = AsyncMock()
|
||
|
|
|
||
|
|
# Create a mock session with rate limiting
|
||
|
|
mock_session = {
|
||
|
|
'enable_rate_limiting': True,
|
||
|
|
'overall_limit_per_ip': 100,
|
||
|
|
'requests_per_sec': 1.0,
|
||
|
|
'agent_id': TEST_AGENT_ID
|
||
|
|
}
|
||
|
|
|
||
|
|
# Mock environment
|
||
|
|
mock_environ = {
|
||
|
|
'REMOTE_ADDR': '192.168.1.1'
|
||
|
|
}
|
||
|
|
|
||
|
|
# Configure Redis mock for rate limit check
|
||
|
|
mock_redis.get.side_effect = [b'50', b'1'] # Under overall limit but over rate limit
|
||
|
|
mock_redis.ttl.return_value = 1 # 1 second remaining
|
||
|
|
|
||
|
|
# Mock sio.emit to capture error messages
|
||
|
|
mock_emit = AsyncMock()
|
||
|
|
|
||
|
|
# Mock asyncio.wait_for to return Redis results directly
|
||
|
|
async def mock_wait_for(coro, timeout):
|
||
|
|
if isinstance(coro, asyncio.Future):
|
||
|
|
return await coro
|
||
|
|
return coro
|
||
|
|
|
||
|
|
with patch('app.services.socket_rate_limit.settings.REDIS_ENABLED', True), \
|
||
|
|
patch('app.services.socket_rate_limit.sio.get_session', AsyncMock(return_value=mock_session)), \
|
||
|
|
patch('app.services.socket_rate_limit.sio.get_environ', return_value=mock_environ), \
|
||
|
|
patch('app.services.socket_rate_limit.sio.emit', mock_emit), \
|
||
|
|
patch('asyncio.wait_for', side_effect=mock_wait_for):
|
||
|
|
# Apply the decorator
|
||
|
|
decorated_handler = socket_rate_limit()(mock_handler)
|
||
|
|
|
||
|
|
# Call the decorated handler
|
||
|
|
result = await decorated_handler('test_sid', 'test_arg')
|
||
|
|
|
||
|
|
# Verify the handler was not called
|
||
|
|
mock_handler.assert_not_called()
|
||
|
|
|
||
|
|
# Verify error was emitted
|
||
|
|
mock_emit.assert_called_once()
|
||
|
|
args = mock_emit.call_args[0]
|
||
|
|
assert args[0] == 'error'
|
||
|
|
assert 'You\'re sending messages too quickly' in args[1]['error']
|
||
|
|
assert result is None
|