1
# frozen_string_literal: true
2

3
require 'msf/core/mcp'
4
require 'webmock/rspec'
5

6
RSpec.describe 'Rate Limiting Integration' do
7
  # Disable real HTTP connections for integration tests
8
  before(:all) do
9
    WebMock.disable_net_connect!(allow_localhost: false)
10
  end
11

12
  after(:all) do
13
    WebMock.allow_net_connect!
14
  end
15

16
  let(:host) { 'localhost' }
17
  let(:port) { 55553 }
18
  let(:endpoint) { '/api/' }
19
  let(:api_url) { "https://#{host}:#{port}#{endpoint}" }
20
  let(:user) { 'test_user' }
21
  let(:password) { 'test_password' }
22

23
  describe 'Rate Limiting Across Multiple Tool HTTP Requests' do
24
    it 'enforces rate limit across multiple tool calls with HTTP requests' do
25
      # Stub authentication endpoint
26
      stub_request(:post, api_url)
27
        .with(body: ['auth.login', user, password].to_msgpack)
28
        .to_return(
29
          status: 200,
30
          body: { 'result' => 'success', 'token' => 'test_token' }.to_msgpack,
31
          headers: { 'Content-Type' => 'binary/message-pack' }
32
        )
33

34
      # Stub module search endpoint
35
      search_stub = stub_request(:post, api_url)
36
        .with(body: ['module.search', 'test_token', 'smb'].to_msgpack)
37
        .to_return(
38
          status: 200,
39
          body: [{ 'fullname' => 'auxiliary/scanner/smb/smb_version' }].to_msgpack,
40
          headers: { 'Content-Type' => 'binary/message-pack' }
41
        )
42

43
      # Create rate limiter with low limit
44
      limiter = Msf::MCP::Security::RateLimiter.new(requests_per_minute: 3, burst_size: 3)
45

46
      # Create authenticated client
47
      client = Msf::MCP::Metasploit::MessagePackClient.new(
48
        host: host,
49
        port: port,
50
        endpoint: endpoint
51
      )
52
      client.authenticate(user, password)
53

54
      # Create server context
55
      server_context = {
56
        msf_client: client,
57
        rate_limiter: limiter
58
      }
59

60
      # First 3 tool calls should succeed (within rate limit)
61
      3.times do
62
        expect {
63
          Msf::MCP::Tools::SearchModules.call(
64
            query: 'smb',
65
            limit: 10,
66
            server_context: server_context
67
          )
68
        }.not_to raise_error
69
      end
70

71
      # Verify exactly 3 HTTP requests were made (plus 1 for auth)
72
      expect(search_stub).to have_been_requested.times(3)
73

74
      # 4th call should be rate limited before making HTTP request
75
      result = Msf::MCP::Tools::SearchModules.call(query: 'smb', limit: 10, server_context: server_context)
76
      expect(result.error?).to be true
77
      expect(result.content.first[:text]).to match(/Rate limit exceeded/)
78

79
      # Verify still only 3 search requests (4th was blocked by rate limiter)
80
      expect(search_stub).to have_been_requested.times(3)
81
    end
82

83
    it 'applies global rate limit across different tools with HTTP calls' do
84
      # Stub authentication
85
      stub_request(:post, api_url)
86
        .with(body: ['auth.login', user, password].to_msgpack)
87
        .to_return(
88
          status: 200,
89
          body: { 'result' => 'success', 'token' => 'test_token' }.to_msgpack,
90
          headers: { 'Content-Type' => 'binary/message-pack' }
91
        )
92

93
      # Stub search modules
94
      search_stub = stub_request(:post, api_url)
95
        .with(body: ['module.search', 'test_token', 'smb'].to_msgpack)
96
        .to_return(
97
          status: 200,
98
          body: [].to_msgpack,
99
          headers: { 'Content-Type' => 'binary/message-pack' }
100
        )
101

102
      # Stub db.hosts
103
      hosts_stub = stub_request(:post, api_url)
104
        .with(body: ['db.hosts', 'test_token', { workspace: 'default' }].to_msgpack)
105
        .to_return(
106
          status: 200,
107
          body: { 'hosts' => [] }.to_msgpack,
108
          headers: { 'Content-Type' => 'binary/message-pack' }
109
        )
110

111
      limiter = Msf::MCP::Security::RateLimiter.new(requests_per_minute: 5, burst_size: 5)
112
      client = Msf::MCP::Metasploit::MessagePackClient.new(
113
        host: host,
114
        port: port,
115
        endpoint: endpoint
116
      )
117
      client.authenticate(user, password)
118

119
      server_context = {
120
        msf_client: client,
121
        rate_limiter: limiter
122
      }
123

124
      # Make 3 search_modules calls
125
      3.times do
126
        Msf::MCP::Tools::SearchModules.call(
127
          query: 'smb',
128
          limit: 10,
129
          server_context: server_context
130
        )
131
      end
132

133
      # Make 2 host_info calls
134
      2.times do
135
        Msf::MCP::Tools::HostInfo.call(
136
          workspace: 'default',
137
          server_context: server_context
138
        )
139
      end
140

141
      # 6th call (different tool) should be rate limited
142
      result = Msf::MCP::Tools::SearchModules.call(query: 'http', limit: 10, server_context: server_context)
143
      expect(result.error?).to be true
144
      expect(result.content.first[:text]).to match(/Rate limit exceeded/)
145

146
      # Verify HTTP request counts
147
      expect(search_stub).to have_been_requested.times(3)
148
      expect(hosts_stub).to have_been_requested.times(2)
149
    end
150
  end
151
end
152

153