1/* Licensed to the Apache Software Foundation (ASF) under one or more
2 * contributor license agreements.  See the NOTICE file distributed with
3 * this work for additional information regarding copyright ownership.
4 * The ASF licenses this file to You under the Apache License, Version 2.0
5 * (the "License"); you may not use this file except in compliance with
6 * the License.  You may obtain a copy of the License at
7 *
8 *     http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16
17#include "httpd.h"
18#include "http_config.h"
19#include "http_log.h"
20#include "util_filter.h"
21
22#include "mod_ratelimit.h"
23
24#define RATE_LIMIT_FILTER_NAME "RATE_LIMIT"
25#define RATE_INTERVAL_MS (200)
26
27typedef enum rl_state_e
28{
29    RATE_ERROR,
30    RATE_LIMIT,
31    RATE_FULLSPEED
32} rl_state_e;
33
34typedef struct rl_ctx_t
35{
36    int speed;
37    int chunk_size;
38    rl_state_e state;
39    apr_bucket_brigade *tmpbb;
40    apr_bucket_brigade *holdingbb;
41} rl_ctx_t;
42
43#if 0
44static void brigade_dump(request_rec *r, apr_bucket_brigade *bb)
45{
46    apr_bucket *e;
47    int i = 0;
48
49    for (e = APR_BRIGADE_FIRST(bb);
50         e != APR_BRIGADE_SENTINEL(bb); e = APR_BUCKET_NEXT(e), i++) {
51        ap_log_rerror(APLOG_MARK, APLOG_ERR, 0, r,
52                      "brigade: [%d] %s", i, e->type->name);
53
54    }
55}
56#endif
57
58static apr_status_t
59rate_limit_filter(ap_filter_t *f, apr_bucket_brigade *input_bb)
60{
61    apr_status_t rv = APR_SUCCESS;
62    rl_ctx_t *ctx = f->ctx;
63    apr_bucket *fb;
64    int do_sleep = 0;
65    apr_bucket_alloc_t *ba = f->r->connection->bucket_alloc;
66    apr_bucket_brigade *bb = input_bb;
67
68    if (f->c->aborted) {
69        ap_log_rerror(APLOG_MARK, APLOG_ERR, 0, f->r, APLOGNO(01454) "rl: conn aborted");
70        apr_brigade_cleanup(bb);
71        return APR_ECONNABORTED;
72    }
73
74    if (ctx == NULL) {
75
76        const char *rl = NULL;
77        int ratelimit;
78
79        /* no subrequests. */
80        if (f->r->main != NULL) {
81            ap_remove_output_filter(f);
82            return ap_pass_brigade(f->next, bb);
83        }
84
85        rl = apr_table_get(f->r->subprocess_env, "rate-limit");
86
87        if (rl == NULL) {
88            ap_remove_output_filter(f);
89            return ap_pass_brigade(f->next, bb);
90        }
91
92        /* rl is in kilo bytes / second  */
93        ratelimit = atoi(rl) * 1024;
94        if (ratelimit <= 0) {
95            /* remove ourselves */
96            ap_remove_output_filter(f);
97            return ap_pass_brigade(f->next, bb);
98        }
99
100        /* first run, init stuff */
101        ctx = apr_palloc(f->r->pool, sizeof(rl_ctx_t));
102        f->ctx = ctx;
103        ctx->state = RATE_LIMIT;
104        ctx->speed = ratelimit;
105
106        /* calculate how many bytes / interval we want to send */
107        /* speed is bytes / second, so, how many  (speed / 1000 % interval) */
108        ctx->chunk_size = (ctx->speed / (1000 / RATE_INTERVAL_MS));
109        ctx->tmpbb = apr_brigade_create(f->r->pool, ba);
110        ctx->holdingbb = apr_brigade_create(f->r->pool, ba);
111    }
112
113    while (ctx->state != RATE_ERROR &&
114           (!APR_BRIGADE_EMPTY(bb) || !APR_BRIGADE_EMPTY(ctx->holdingbb))) {
115        apr_bucket *e;
116
117        if (!APR_BRIGADE_EMPTY(ctx->holdingbb)) {
118            APR_BRIGADE_CONCAT(bb, ctx->holdingbb);
119            apr_brigade_cleanup(ctx->holdingbb);
120        }
121
122        while (ctx->state == RATE_FULLSPEED && !APR_BRIGADE_EMPTY(bb)) {
123            /* Find where we 'stop' going full speed. */
124            for (e = APR_BRIGADE_FIRST(bb);
125                 e != APR_BRIGADE_SENTINEL(bb); e = APR_BUCKET_NEXT(e)) {
126                if (AP_RL_BUCKET_IS_END(e)) {
127                    apr_bucket *f;
128                    f = APR_RING_LAST(&bb->list);
129                    APR_RING_UNSPLICE(e, f, link);
130                    APR_RING_SPLICE_TAIL(&ctx->holdingbb->list, e, f,
131                                         apr_bucket, link);
132                    ctx->state = RATE_LIMIT;
133                    break;
134                }
135            }
136
137            if (f->c->aborted) {
138                apr_brigade_cleanup(bb);
139                ctx->state = RATE_ERROR;
140                break;
141            }
142
143            fb = apr_bucket_flush_create(ba);
144            APR_BRIGADE_INSERT_TAIL(bb, fb);
145            rv = ap_pass_brigade(f->next, bb);
146
147            if (rv != APR_SUCCESS) {
148                ctx->state = RATE_ERROR;
149                ap_log_rerror(APLOG_MARK, APLOG_ERR, rv, f->r, APLOGNO(01455)
150                              "rl: full speed brigade pass failed.");
151            }
152        }
153
154        while (ctx->state == RATE_LIMIT && !APR_BRIGADE_EMPTY(bb)) {
155            for (e = APR_BRIGADE_FIRST(bb);
156                 e != APR_BRIGADE_SENTINEL(bb); e = APR_BUCKET_NEXT(e)) {
157                if (AP_RL_BUCKET_IS_START(e)) {
158                    apr_bucket *f;
159                    f = APR_RING_LAST(&bb->list);
160                    APR_RING_UNSPLICE(e, f, link);
161                    APR_RING_SPLICE_TAIL(&ctx->holdingbb->list, e, f,
162                                         apr_bucket, link);
163                    ctx->state = RATE_FULLSPEED;
164                    break;
165                }
166            }
167
168            while (!APR_BRIGADE_EMPTY(bb)) {
169                apr_bucket *stop_point;
170                apr_off_t len = 0;
171
172                if (f->c->aborted) {
173                    apr_brigade_cleanup(bb);
174                    ctx->state = RATE_ERROR;
175                    break;
176                }
177
178                if (do_sleep) {
179                    apr_sleep(RATE_INTERVAL_MS * 1000);
180                }
181                else {
182                    do_sleep = 1;
183                }
184
185                apr_brigade_length(bb, 1, &len);
186
187                rv = apr_brigade_partition(bb, ctx->chunk_size, &stop_point);
188                if (rv != APR_SUCCESS && rv != APR_INCOMPLETE) {
189                    ctx->state = RATE_ERROR;
190                    ap_log_rerror(APLOG_MARK, APLOG_ERR, rv, f->r, APLOGNO(01456)
191                                  "rl: partition failed.");
192                    break;
193                }
194
195                if (stop_point != APR_BRIGADE_SENTINEL(bb)) {
196                    apr_bucket *f;
197                    apr_bucket *e = APR_BUCKET_PREV(stop_point);
198                    f = APR_RING_FIRST(&bb->list);
199                    APR_RING_UNSPLICE(f, e, link);
200                    APR_RING_SPLICE_HEAD(&ctx->tmpbb->list, f, e, apr_bucket,
201                                         link);
202                }
203                else {
204                    APR_BRIGADE_CONCAT(ctx->tmpbb, bb);
205                }
206
207                fb = apr_bucket_flush_create(ba);
208
209                APR_BRIGADE_INSERT_TAIL(ctx->tmpbb, fb);
210
211#if 0
212                brigade_dump(f->r, ctx->tmpbb);
213                brigade_dump(f->r, bb);
214#endif
215
216                rv = ap_pass_brigade(f->next, ctx->tmpbb);
217                apr_brigade_cleanup(ctx->tmpbb);
218
219                if (rv != APR_SUCCESS) {
220                    ctx->state = RATE_ERROR;
221                    ap_log_rerror(APLOG_MARK, APLOG_ERR, 0, f->r, APLOGNO(01457)
222                                  "rl: brigade pass failed.");
223                    break;
224                }
225            }
226        }
227    }
228
229    return rv;
230}
231
232
233static apr_status_t
234rl_bucket_read(apr_bucket *b, const char **str,
235               apr_size_t *len, apr_read_type_e block)
236{
237    *str = NULL;
238    *len = 0;
239    return APR_SUCCESS;
240}
241
242AP_RL_DECLARE(apr_bucket *)
243    ap_rl_end_create(apr_bucket_alloc_t *list)
244{
245    apr_bucket *b = apr_bucket_alloc(sizeof(*b), list);
246
247    APR_BUCKET_INIT(b);
248    b->free = apr_bucket_free;
249    b->list = list;
250    b->length = 0;
251    b->start = 0;
252    b->data = NULL;
253    b->type = &ap_rl_bucket_type_end;
254
255    return b;
256}
257
258AP_RL_DECLARE(apr_bucket *)
259    ap_rl_start_create(apr_bucket_alloc_t *list)
260{
261    apr_bucket *b = apr_bucket_alloc(sizeof(*b), list);
262
263    APR_BUCKET_INIT(b);
264    b->free = apr_bucket_free;
265    b->list = list;
266    b->length = 0;
267    b->start = 0;
268    b->data = NULL;
269    b->type = &ap_rl_bucket_type_start;
270
271    return b;
272}
273
274
275
276AP_RL_DECLARE_DATA const apr_bucket_type_t ap_rl_bucket_type_end = {
277    "RL_END", 5, APR_BUCKET_METADATA,
278    apr_bucket_destroy_noop,
279    rl_bucket_read,
280    apr_bucket_setaside_noop,
281    apr_bucket_split_notimpl,
282    apr_bucket_simple_copy
283};
284
285
286AP_RL_DECLARE_DATA const apr_bucket_type_t ap_rl_bucket_type_start = {
287    "RL_START", 5, APR_BUCKET_METADATA,
288    apr_bucket_destroy_noop,
289    rl_bucket_read,
290    apr_bucket_setaside_noop,
291    apr_bucket_split_notimpl,
292    apr_bucket_simple_copy
293};
294
295
296
297
298static void register_hooks(apr_pool_t *p)
299{
300    /* run after mod_deflate etc etc, but not at connection level, ie, mod_ssl. */
301    ap_register_output_filter(RATE_LIMIT_FILTER_NAME, rate_limit_filter,
302                              NULL, AP_FTYPE_PROTOCOL + 3);
303}
304
305AP_DECLARE_MODULE(ratelimit) = {
306    STANDARD20_MODULE_STUFF,
307    NULL,                       /* create per-directory config structure */
308    NULL,                       /* merge per-directory config structures */
309    NULL,                       /* create per-server config structure */
310    NULL,                       /* merge per-server config structures */
311    NULL,                       /* command apr_table_t */
312    register_hooks
313};
314