granicus.if.org Git - apache/blob - modules/filters/mod_ratelimit.c

   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
  27 typedef enum rl_state_e
  28 {
  29     RATE_ERROR,
  30     RATE_LIMIT,
  31     RATE_FULLSPEED
  32 } rl_state_e;
  33
  34 typedef 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
  44 static 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, APLOGNO(03193)
  52                       "brigade: [%d] %s", i, e->type->name);
  53
  54     }
  55 }
  56 #endif
  57
  58 static apr_status_t
  59 rate_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         }
 120
 121         while (ctx->state == RATE_FULLSPEED && !APR_BRIGADE_EMPTY(bb)) {
 122             /* Find where we 'stop' going full speed. */
 123             for (e = APR_BRIGADE_FIRST(bb);
 124                  e != APR_BRIGADE_SENTINEL(bb); e = APR_BUCKET_NEXT(e)) {
 125                 if (AP_RL_BUCKET_IS_END(e)) {
 126                     apr_bucket *f;
 127                     f = APR_RING_LAST(&bb->list);
 128                     APR_RING_UNSPLICE(e, f, link);
 129                     APR_RING_SPLICE_TAIL(&ctx->holdingbb->list, e, f,
 130                                          apr_bucket, link);
 131                     ctx->state = RATE_LIMIT;
 132                     break;
 133                 }
 134             }
 135
 136             if (f->c->aborted) {
 137                 apr_brigade_cleanup(bb);
 138                 ctx->state = RATE_ERROR;
 139                 break;
 140             }
 141
 142             fb = apr_bucket_flush_create(ba);
 143             APR_BRIGADE_INSERT_TAIL(bb, fb);
 144             rv = ap_pass_brigade(f->next, bb);
 145
 146             if (rv != APR_SUCCESS) {
 147                 ctx->state = RATE_ERROR;
 148                 ap_log_rerror(APLOG_MARK, APLOG_TRACE1, rv, f->r, APLOGNO(01455)
 149                               "rl: full speed brigade pass failed.");
 150             }
 151         }
 152
 153         while (ctx->state == RATE_LIMIT && !APR_BRIGADE_EMPTY(bb)) {
 154             for (e = APR_BRIGADE_FIRST(bb);
 155                  e != APR_BRIGADE_SENTINEL(bb); e = APR_BUCKET_NEXT(e)) {
 156                 if (AP_RL_BUCKET_IS_START(e)) {
 157                     apr_bucket *f;
 158                     f = APR_RING_LAST(&bb->list);
 159                     APR_RING_UNSPLICE(e, f, link);
 160                     APR_RING_SPLICE_TAIL(&ctx->holdingbb->list, e, f,
 161                                          apr_bucket, link);
 162                     ctx->state = RATE_FULLSPEED;
 163                     break;
 164                 }
 165             }
 166
 167             while (!APR_BRIGADE_EMPTY(bb)) {
 168                 apr_bucket *stop_point;
 169                 apr_off_t len = 0;
 170
 171                 if (f->c->aborted) {
 172                     apr_brigade_cleanup(bb);
 173                     ctx->state = RATE_ERROR;
 174                     break;
 175                 }
 176
 177                 if (do_sleep) {
 178                     apr_sleep(RATE_INTERVAL_MS * 1000);
 179                 }
 180                 else {
 181                     do_sleep = 1;
 182                 }
 183
 184                 apr_brigade_length(bb, 1, &len);
 185
 186                 rv = apr_brigade_partition(bb, ctx->chunk_size, &stop_point);
 187                 if (rv != APR_SUCCESS && rv != APR_INCOMPLETE) {
 188                     ctx->state = RATE_ERROR;
 189                     ap_log_rerror(APLOG_MARK, APLOG_ERR, rv, f->r, APLOGNO(01456)
 190                                   "rl: partition failed.");
 191                     break;
 192                 }
 193
 194                 if (stop_point != APR_BRIGADE_SENTINEL(bb)) {
 195                     apr_bucket *f;
 196                     apr_bucket *e = APR_BUCKET_PREV(stop_point);
 197                     f = APR_RING_FIRST(&bb->list);
 198                     APR_RING_UNSPLICE(f, e, link);
 199                     APR_RING_SPLICE_HEAD(&ctx->tmpbb->list, f, e, apr_bucket,
 200                                          link);
 201                 }
 202                 else {
 203                     APR_BRIGADE_CONCAT(ctx->tmpbb, bb);
 204                 }
 205
 206                 fb = apr_bucket_flush_create(ba);
 207
 208                 APR_BRIGADE_INSERT_TAIL(ctx->tmpbb, fb);
 209
 210 #if 0
 211                 brigade_dump(f->r, ctx->tmpbb);
 212                 brigade_dump(f->r, bb);
 213 #endif
 214
 215                 rv = ap_pass_brigade(f->next, ctx->tmpbb);
 216                 apr_brigade_cleanup(ctx->tmpbb);
 217
 218                 if (rv != APR_SUCCESS) {
 219                     ctx->state = RATE_ERROR;
 220                     ap_log_rerror(APLOG_MARK, APLOG_TRACE1, rv, f->r, APLOGNO(01457)
 221                                   "rl: brigade pass failed.");
 222                     break;
 223                 }
 224             }
 225         }
 226     }
 227
 228     return rv;
 229 }
 230
 231
 232 static apr_status_t
 233 rl_bucket_read(apr_bucket *b, const char **str,
 234                apr_size_t *len, apr_read_type_e block)
 235 {
 236     *str = NULL;
 237     *len = 0;
 238     return APR_SUCCESS;
 239 }
 240
 241 AP_RL_DECLARE(apr_bucket *)
 242     ap_rl_end_create(apr_bucket_alloc_t *list)
 243 {
 244     apr_bucket *b = apr_bucket_alloc(sizeof(*b), list);
 245
 246     APR_BUCKET_INIT(b);
 247     b->free = apr_bucket_free;
 248     b->list = list;
 249     b->length = 0;
 250     b->start = 0;
 251     b->data = NULL;
 252     b->type = &ap_rl_bucket_type_end;
 253
 254     return b;
 255 }
 256
 257 AP_RL_DECLARE(apr_bucket *)
 258     ap_rl_start_create(apr_bucket_alloc_t *list)
 259 {
 260     apr_bucket *b = apr_bucket_alloc(sizeof(*b), list);
 261
 262     APR_BUCKET_INIT(b);
 263     b->free = apr_bucket_free;
 264     b->list = list;
 265     b->length = 0;
 266     b->start = 0;
 267     b->data = NULL;
 268     b->type = &ap_rl_bucket_type_start;
 269
 270     return b;
 271 }
 272
 273
 274
 275 AP_RL_DECLARE_DATA const apr_bucket_type_t ap_rl_bucket_type_end = {
 276     "RL_END", 5, APR_BUCKET_METADATA,
 277     apr_bucket_destroy_noop,
 278     rl_bucket_read,
 279     apr_bucket_setaside_noop,
 280     apr_bucket_split_notimpl,
 281     apr_bucket_simple_copy
 282 };
 283
 284
 285 AP_RL_DECLARE_DATA const apr_bucket_type_t ap_rl_bucket_type_start = {
 286     "RL_START", 5, APR_BUCKET_METADATA,
 287     apr_bucket_destroy_noop,
 288     rl_bucket_read,
 289     apr_bucket_setaside_noop,
 290     apr_bucket_split_notimpl,
 291     apr_bucket_simple_copy
 292 };
 293
 294
 295
 296
 297 static void register_hooks(apr_pool_t *p)
 298 {
 299     /* run after mod_deflate etc etc, but not at connection level, ie, mod_ssl. */
 300     ap_register_output_filter(RATE_LIMIT_FILTER_NAME, rate_limit_filter,
 301                               NULL, AP_FTYPE_PROTOCOL + 3);
 302 }
 303
 304 AP_DECLARE_MODULE(ratelimit) = {
 305     STANDARD20_MODULE_STUFF,
 306     NULL,                       /* create per-directory config structure */
 307     NULL,                       /* merge per-directory config structures */
 308     NULL,                       /* create per-server config structure */
 309     NULL,                       /* merge per-server config structures */
 310     NULL,                       /* command apr_table_t */
 311     register_hooks
 312 };