1212# See the License for the specific language governing permissions and
1313# limitations under the License.
1414
15- from typing import Any , Dict , Optional
15+ from typing import Any , Dict , Optional , cast
1616
1717import attr
1818
2121from ._base import Config
2222
2323
24+ @attr .s (slots = True , frozen = True , auto_attribs = True )
2425class RatelimitSettings :
25- def __init__ (
26- self ,
27- config : Dict [str , float ],
26+ key : str
27+ per_second : float
28+ burst_count : int
29+
30+ @classmethod
31+ def parse (
32+ cls ,
33+ config : Dict [str , Any ],
34+ key : str ,
2835 defaults : Optional [Dict [str , float ]] = None ,
29- ):
36+ ) -> "RatelimitSettings" :
37+ """Parse config[key] as a new-style rate limiter config.
38+
39+ The key may refer to a nested dictionary using a full stop (.) to separate
40+ each nested key. For example, use the key "a.b.c" to parse the following:
41+
42+ a:
43+ b:
44+ c:
45+ per_second: 10
46+ burst_count: 200
47+
48+ If this lookup fails, we'll fallback to the defaults.
49+ """
3050 defaults = defaults or {"per_second" : 0.17 , "burst_count" : 3.0 }
3151
32- self .per_second = config .get ("per_second" , defaults ["per_second" ])
33- self .burst_count = int (config .get ("burst_count" , defaults ["burst_count" ]))
52+ rl_config = config
53+ for part in key .split ("." ):
54+ rl_config = rl_config .get (part , {})
55+
56+ # By this point we should have hit the rate limiter parameters.
57+ # We don't actually check this though!
58+ rl_config = cast (Dict [str , float ], rl_config )
59+
60+ return cls (
61+ key = key ,
62+ per_second = rl_config .get ("per_second" , defaults ["per_second" ]),
63+ burst_count = int (rl_config .get ("burst_count" , defaults ["burst_count" ])),
64+ )
3465
3566
3667@attr .s (auto_attribs = True )
@@ -49,15 +80,14 @@ def read_config(self, config: JsonDict, **kwargs: Any) -> None:
4980 # Load the new-style messages config if it exists. Otherwise fall back
5081 # to the old method.
5182 if "rc_message" in config :
52- self .rc_message = RatelimitSettings (
53- config [ "rc_message" ] , defaults = {"per_second" : 0.2 , "burst_count" : 10.0 }
83+ self .rc_message = RatelimitSettings . parse (
84+ config , "rc_message" , defaults = {"per_second" : 0.2 , "burst_count" : 10.0 }
5485 )
5586 else :
5687 self .rc_message = RatelimitSettings (
57- {
58- "per_second" : config .get ("rc_messages_per_second" , 0.2 ),
59- "burst_count" : config .get ("rc_message_burst_count" , 10.0 ),
60- }
88+ key = "rc_messages" ,
89+ per_second = config .get ("rc_messages_per_second" , 0.2 ),
90+ burst_count = config .get ("rc_message_burst_count" , 10.0 ),
6191 )
6292
6393 # Load the new-style federation config, if it exists. Otherwise, fall
@@ -79,51 +109,59 @@ def read_config(self, config: JsonDict, **kwargs: Any) -> None:
79109 }
80110 )
81111
82- self .rc_registration = RatelimitSettings ( config . get ( "rc_registration" , {}) )
112+ self .rc_registration = RatelimitSettings . parse ( config , "rc_registration" , {})
83113
84- self .rc_registration_token_validity = RatelimitSettings (
85- config .get ("rc_registration_token_validity" , {}),
114+ self .rc_registration_token_validity = RatelimitSettings .parse (
115+ config ,
116+ "rc_registration_token_validity" ,
86117 defaults = {"per_second" : 0.1 , "burst_count" : 5 },
87118 )
88119
89120 # It is reasonable to login with a bunch of devices at once (i.e. when
90121 # setting up an account), but it is *not* valid to continually be
91122 # logging into new devices.
92- rc_login_config = config . get ( "rc_login" , {})
93- self . rc_login_address = RatelimitSettings (
94- rc_login_config . get ( " address", {}) ,
123+ self . rc_login_address = RatelimitSettings . parse (
124+ config ,
125+ "rc_login. address" ,
95126 defaults = {"per_second" : 0.003 , "burst_count" : 5 },
96127 )
97- self .rc_login_account = RatelimitSettings (
98- rc_login_config .get ("account" , {}),
128+ self .rc_login_account = RatelimitSettings .parse (
129+ config ,
130+ "rc_login.account" ,
99131 defaults = {"per_second" : 0.003 , "burst_count" : 5 },
100132 )
101- self .rc_login_failed_attempts = RatelimitSettings (
102- rc_login_config .get ("failed_attempts" , {})
133+ self .rc_login_failed_attempts = RatelimitSettings .parse (
134+ config ,
135+ "rc_login.failed_attempts" ,
136+ {},
103137 )
104138
105139 self .federation_rr_transactions_per_room_per_second = config .get (
106140 "federation_rr_transactions_per_room_per_second" , 50
107141 )
108142
109- rc_admin_redaction = config .get ("rc_admin_redaction" )
110143 self .rc_admin_redaction = None
111- if rc_admin_redaction :
112- self .rc_admin_redaction = RatelimitSettings (rc_admin_redaction )
144+ if "rc_admin_redaction" in config :
145+ self .rc_admin_redaction = RatelimitSettings .parse (
146+ config , "rc_admin_redaction" , {}
147+ )
113148
114- self .rc_joins_local = RatelimitSettings (
115- config .get ("rc_joins" , {}).get ("local" , {}),
149+ self .rc_joins_local = RatelimitSettings .parse (
150+ config ,
151+ "rc_joins.local" ,
116152 defaults = {"per_second" : 0.1 , "burst_count" : 10 },
117153 )
118- self .rc_joins_remote = RatelimitSettings (
119- config .get ("rc_joins" , {}).get ("remote" , {}),
154+ self .rc_joins_remote = RatelimitSettings .parse (
155+ config ,
156+ "rc_joins.remote" ,
120157 defaults = {"per_second" : 0.01 , "burst_count" : 10 },
121158 )
122159
123160 # Track the rate of joins to a given room. If there are too many, temporarily
124161 # prevent local joins and remote joins via this server.
125- self .rc_joins_per_room = RatelimitSettings (
126- config .get ("rc_joins_per_room" , {}),
162+ self .rc_joins_per_room = RatelimitSettings .parse (
163+ config ,
164+ "rc_joins_per_room" ,
127165 defaults = {"per_second" : 1 , "burst_count" : 10 },
128166 )
129167
@@ -132,31 +170,37 @@ def read_config(self, config: JsonDict, **kwargs: Any) -> None:
132170 # * For requests received over federation this is keyed by the origin.
133171 #
134172 # Note that this isn't exposed in the configuration as it is obscure.
135- self .rc_key_requests = RatelimitSettings (
136- config .get ("rc_key_requests" , {}),
173+ self .rc_key_requests = RatelimitSettings .parse (
174+ config ,
175+ "rc_key_requests" ,
137176 defaults = {"per_second" : 20 , "burst_count" : 100 },
138177 )
139178
140- self .rc_3pid_validation = RatelimitSettings (
141- config .get ("rc_3pid_validation" ) or {},
179+ self .rc_3pid_validation = RatelimitSettings .parse (
180+ config ,
181+ "rc_3pid_validation" ,
142182 defaults = {"per_second" : 0.003 , "burst_count" : 5 },
143183 )
144184
145- self .rc_invites_per_room = RatelimitSettings (
146- config .get ("rc_invites" , {}).get ("per_room" , {}),
185+ self .rc_invites_per_room = RatelimitSettings .parse (
186+ config ,
187+ "rc_invites.per_room" ,
147188 defaults = {"per_second" : 0.3 , "burst_count" : 10 },
148189 )
149- self .rc_invites_per_user = RatelimitSettings (
150- config .get ("rc_invites" , {}).get ("per_user" , {}),
190+ self .rc_invites_per_user = RatelimitSettings .parse (
191+ config ,
192+ "rc_invites.per_user" ,
151193 defaults = {"per_second" : 0.003 , "burst_count" : 5 },
152194 )
153195
154- self .rc_invites_per_issuer = RatelimitSettings (
155- config .get ("rc_invites" , {}).get ("per_issuer" , {}),
196+ self .rc_invites_per_issuer = RatelimitSettings .parse (
197+ config ,
198+ "rc_invites.per_issuer" ,
156199 defaults = {"per_second" : 0.3 , "burst_count" : 10 },
157200 )
158201
159- self .rc_third_party_invite = RatelimitSettings (
160- config .get ("rc_third_party_invite" , {}),
202+ self .rc_third_party_invite = RatelimitSettings .parse (
203+ config ,
204+ "rc_third_party_invite" ,
161205 defaults = {"per_second" : 0.0025 , "burst_count" : 5 },
162206 )
0 commit comments