modules/filters/mod_ratelimit.c - httpd - Git at Google

 /* Licensed to the Apache Software Foundation (ASF) under one or more
  * contributor license agreements.  See the NOTICE file distributed with
  * this work for additional information regarding copyright ownership.
  * The ASF licenses this file to You 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.
  */

 #include "httpd.h"
 #include "http_config.h"
 #include "http_log.h"
 #include "util_filter.h"

 #include "mod_ratelimit.h"

 #define RATE_LIMIT_FILTER_NAME "RATE_LIMIT"
 #define RATE_INTERVAL_MS (200)

 typedef enum rl_state_e
 {
     RATE_LIMIT,
     RATE_FULLSPEED
 } rl_state_e;

 typedef struct rl_ctx_t
 {
     int speed;
     int chunk_size;
     int burst;
     int do_sleep;
     rl_state_e state;
     apr_bucket_brigade *tmpbb;
     apr_bucket_brigade *holdingbb;
 } rl_ctx_t;

 #if defined(RLFDEBUG)
 static void brigade_dump(request_rec *r, apr_bucket_brigade *bb)
 {
     apr_bucket *e;
     int i = 0;

     for (e = APR_BRIGADE_FIRST(bb);
          e != APR_BRIGADE_SENTINEL(bb); e = APR_BUCKET_NEXT(e), i++) {
         ap_log_rerror(APLOG_MARK, APLOG_ERR, 0, r, APLOGNO(03193)
                       "brigade: [%d] %s", i, e->type->name);

     }
 }
 #endif /* RLFDEBUG */

 static apr_status_t
 rate_limit_filter(ap_filter_t *f, apr_bucket_brigade *bb)
 {
     apr_status_t rv = APR_SUCCESS;
     rl_ctx_t *ctx = f->ctx;
     apr_bucket_alloc_t *ba = f->r->connection->bucket_alloc;

     /* Set up our rl_ctx_t on first use */
     if (ctx == NULL) {
         const char *rl = NULL;
         int ratelimit;
         int burst = 0;

         /* no subrequests. */
         if (f->r->main != NULL) {
             ap_remove_output_filter(f);
             return ap_pass_brigade(f->next, bb);
         }

         /* Configuration: rate limit */
         rl = apr_table_get(f->r->subprocess_env, "rate-limit");

         if (rl == NULL) {
             ap_remove_output_filter(f);
             return ap_pass_brigade(f->next, bb);
         }

         /* rl is in kilo bytes / second  */
         ratelimit = atoi(rl) * 1024;
         if (ratelimit <= 0) {
             /* remove ourselves */
             ap_log_rerror(APLOG_MARK, APLOG_INFO, 0, f->r,
                           APLOGNO(03488) "rl: disabling: rate-limit = %s (too high?)", rl);
             ap_remove_output_filter(f);
             return ap_pass_brigade(f->next, bb);
         }

         /* Configuration: optional initial burst */
         rl = apr_table_get(f->r->subprocess_env, "rate-initial-burst");
         if (rl != NULL) {
             burst = atoi(rl) * 1024;
             if (burst <= 0) {
                ap_log_rerror(APLOG_MARK, APLOG_INFO, 0, f->r,
                              APLOGNO(03489) "rl: disabling burst: rate-initial-burst = %s (too high?)", rl);
                burst = 0;
             }
         }

         /* Set up our context */
         ctx = apr_palloc(f->r->pool, sizeof(rl_ctx_t));
         f->ctx = ctx;
         ctx->state = RATE_LIMIT;
         ctx->speed = ratelimit;
         ctx->burst = burst;
         ctx->do_sleep = 0;

         /* calculate how many bytes / interval we want to send */
         /* speed is bytes / second, so, how many  (speed / 1000 % interval) */
         ctx->chunk_size = (ctx->speed / (1000 / RATE_INTERVAL_MS));
         ctx->tmpbb = apr_brigade_create(f->r->pool, ba);
         ctx->holdingbb = apr_brigade_create(f->r->pool, ba);
     }
     else {
         APR_BRIGADE_PREPEND(bb, ctx->holdingbb);
     }

     while (!APR_BRIGADE_EMPTY(bb)) {
         apr_bucket *e;

         if (ctx->state == RATE_FULLSPEED) {
             /* Find where we 'stop' going full speed. */
             for (e = APR_BRIGADE_FIRST(bb);
                  e != APR_BRIGADE_SENTINEL(bb); e = APR_BUCKET_NEXT(e)) {
                 if (AP_RL_BUCKET_IS_END(e)) {
                     apr_brigade_split_ex(bb, e, ctx->holdingbb);
                     ctx->state = RATE_LIMIT;
                     break;
                 }
             }

             e = apr_bucket_flush_create(ba);
             APR_BRIGADE_INSERT_TAIL(bb, e);
             rv = ap_pass_brigade(f->next, bb);
             apr_brigade_cleanup(bb);

             if (rv != APR_SUCCESS) {
                 ap_log_rerror(APLOG_MARK, APLOG_TRACE1, rv, f->r, APLOGNO(01455)
                               "rl: full speed brigade pass failed.");
                 return rv;
             }
         }
         else {
             for (e = APR_BRIGADE_FIRST(bb);
                  e != APR_BRIGADE_SENTINEL(bb); e = APR_BUCKET_NEXT(e)) {
                 if (AP_RL_BUCKET_IS_START(e)) {
                     apr_brigade_split_ex(bb, e, ctx->holdingbb);
                     ctx->state = RATE_FULLSPEED;
                     break;
                 }
             }

             while (!APR_BRIGADE_EMPTY(bb)) {
                 apr_off_t len = ctx->chunk_size + ctx->burst;

                 APR_BRIGADE_CONCAT(ctx->tmpbb, bb);

                 /*
                  * Pull next chunk of data; the initial amount is our
                  * burst allotment (if any) plus a chunk.  All subsequent
                  * iterations are just chunks with whatever remaining
                  * burst amounts we have left (in case not done in the
                  * first bucket).
                  */
                 rv = apr_brigade_partition(ctx->tmpbb, len, &e);
                 if (rv != APR_SUCCESS && rv != APR_INCOMPLETE) {
                     ap_log_rerror(APLOG_MARK, APLOG_ERR, rv, f->r, APLOGNO(01456)
                                   "rl: partition failed.");
                     return rv;
                 }
                 /* Send next metadata now if any */
                 while (e != APR_BRIGADE_SENTINEL(ctx->tmpbb)
                        && APR_BUCKET_IS_METADATA(e)) {
                     e = APR_BUCKET_NEXT(e);
                 }
                 if (e != APR_BRIGADE_SENTINEL(ctx->tmpbb)) {
                     apr_brigade_split_ex(ctx->tmpbb, e, bb);
                 }
                 else {
                     apr_brigade_length(ctx->tmpbb, 1, &len);
                 }

                 /*
                  * Adjust the burst amount depending on how much
                  * we've done up to now.
                  */
                 if (ctx->burst) {
                     ap_log_rerror(APLOG_MARK, APLOG_DEBUG, 0, f->r,
                         APLOGNO(03485) "rl: burst %d; len %"APR_OFF_T_FMT, ctx->burst, len);
                     if (len < ctx->burst) {
                         ctx->burst -= len;
                     }
                     else {
                         ctx->burst = 0;
                     }
                 }

                 e = APR_BRIGADE_LAST(ctx->tmpbb);
                 if (APR_BUCKET_IS_EOS(e)) {
                     ap_remove_output_filter(f);
                 }
                 else if (!APR_BUCKET_IS_FLUSH(e)) {
                     if (APR_BRIGADE_EMPTY(bb)) {
                         /* Wait for more (or next call) */
                         break;
                     }
                     e = apr_bucket_flush_create(ba);
                     APR_BRIGADE_INSERT_TAIL(ctx->tmpbb, e);
                 }

 #if defined(RLFDEBUG)
                 brigade_dump(f->r, ctx->tmpbb);
                 brigade_dump(f->r, bb);
 #endif /* RLFDEBUG */

                 if (ctx->do_sleep) {
                     apr_sleep(RATE_INTERVAL_MS * 1000);
                 }
                 else {
                     ctx->do_sleep = 1;
                 }

                 rv = ap_pass_brigade(f->next, ctx->tmpbb);
                 apr_brigade_cleanup(ctx->tmpbb);

                 if (rv != APR_SUCCESS) {
                     /* Most often, user disconnects from stream */
                     ap_log_rerror(APLOG_MARK, APLOG_TRACE1, rv, f->r, APLOGNO(01457)
                                   "rl: brigade pass failed.");
                     return rv;
                 }
             }
         }

         if (!APR_BRIGADE_EMPTY(ctx->holdingbb)) {
             /* Any rate-limited data in tmpbb is sent unlimited along
              * with the rest.
              */
             APR_BRIGADE_CONCAT(bb, ctx->tmpbb);
             APR_BRIGADE_CONCAT(bb, ctx->holdingbb);
         }
     }

 #if defined(RLFDEBUG)
     brigade_dump(f->r, ctx->tmpbb);
 #endif /* RLFDEBUG */

     /* Save remaining tmpbb with the correct lifetime for the next call */
     return ap_save_brigade(f, &ctx->holdingbb, &ctx->tmpbb, f->r->pool);
 }


 static apr_status_t
 rl_bucket_read(apr_bucket *b, const char **str,
                apr_size_t *len, apr_read_type_e block)
 {
     *str = NULL;
     *len = 0;
     return APR_SUCCESS;
 }

 AP_RL_DECLARE(apr_bucket *)
     ap_rl_end_create(apr_bucket_alloc_t *list)
 {
     apr_bucket *b = apr_bucket_alloc(sizeof(*b), list);

     APR_BUCKET_INIT(b);
     b->free = apr_bucket_free;
     b->list = list;
     b->length = 0;
     b->start = 0;
     b->data = NULL;
     b->type = &ap_rl_bucket_type_end;

     return b;
 }

 AP_RL_DECLARE(apr_bucket *)
     ap_rl_start_create(apr_bucket_alloc_t *list)
 {
     apr_bucket *b = apr_bucket_alloc(sizeof(*b), list);

     APR_BUCKET_INIT(b);
     b->free = apr_bucket_free;
     b->list = list;
     b->length = 0;
     b->start = 0;
     b->data = NULL;
     b->type = &ap_rl_bucket_type_start;

     return b;
 }


 AP_RL_DECLARE_DATA const apr_bucket_type_t ap_rl_bucket_type_end = {
     "RL_END", 5, APR_BUCKET_METADATA,
     apr_bucket_destroy_noop,
     rl_bucket_read,
     apr_bucket_setaside_noop,
     apr_bucket_split_notimpl,
     apr_bucket_simple_copy
 };


 AP_RL_DECLARE_DATA const apr_bucket_type_t ap_rl_bucket_type_start = {
     "RL_START", 5, APR_BUCKET_METADATA,
     apr_bucket_destroy_noop,
     rl_bucket_read,
     apr_bucket_setaside_noop,
     apr_bucket_split_notimpl,
     apr_bucket_simple_copy
 };


 static void register_hooks(apr_pool_t *p)
 {
     /* run after mod_deflate etc etc, but not at connection level, ie, mod_ssl. */
     ap_register_output_filter(RATE_LIMIT_FILTER_NAME, rate_limit_filter,
                               NULL, AP_FTYPE_CONNECTION - 1);
 }

 AP_DECLARE_MODULE(ratelimit) = {
     STANDARD20_MODULE_STUFF,
     NULL,                       /* create per-directory config structure */
     NULL,                       /* merge per-directory config structures */
     NULL,                       /* create per-server config structure */
     NULL,                       /* merge per-server config structures */
     NULL,                       /* command apr_table_t */
     register_hooks
 };
	/* Licensed to the Apache Software Foundation (ASF) under one or more
	* contributor license agreements. See the NOTICE file distributed with
	* this work for additional information regarding copyright ownership.
	* The ASF licenses this file to You 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.
	*/

	#include "httpd.h"
	#include "http_config.h"
	#include "http_log.h"
	#include "util_filter.h"

	#include "mod_ratelimit.h"

	#define RATE_LIMIT_FILTER_NAME "RATE_LIMIT"
	#define RATE_INTERVAL_MS (200)

	typedef enum rl_state_e
	{
	RATE_LIMIT,
	RATE_FULLSPEED
	} rl_state_e;

	typedef struct rl_ctx_t
	{
	int speed;
	int chunk_size;
	int burst;
	int do_sleep;
	rl_state_e state;
	apr_bucket_brigade *tmpbb;
	apr_bucket_brigade *holdingbb;
	} rl_ctx_t;

	#if defined(RLFDEBUG)
	static void brigade_dump(request_rec r, apr_bucket_brigade bb)
	{
	apr_bucket *e;
	int i = 0;

	for (e = APR_BRIGADE_FIRST(bb);
	e != APR_BRIGADE_SENTINEL(bb); e = APR_BUCKET_NEXT(e), i++) {
	ap_log_rerror(APLOG_MARK, APLOG_ERR, 0, r, APLOGNO(03193)
	"brigade: [%d] %s", i, e->type->name);

	}
	}
	#endif /* RLFDEBUG */

	static apr_status_t
	rate_limit_filter(ap_filter_t f, apr_bucket_brigade bb)
	{
	apr_status_t rv = APR_SUCCESS;
	rl_ctx_t *ctx = f->ctx;
	apr_bucket_alloc_t *ba = f->r->connection->bucket_alloc;

	/* Set up our rl_ctx_t on first use */
	if (ctx == NULL) {
	const char *rl = NULL;
	int ratelimit;
	int burst = 0;

	/* no subrequests. */
	if (f->r->main != NULL) {
	ap_remove_output_filter(f);
	return ap_pass_brigade(f->next, bb);
	}

	/* Configuration: rate limit */
	rl = apr_table_get(f->r->subprocess_env, "rate-limit");

	if (rl == NULL) {
	ap_remove_output_filter(f);
	return ap_pass_brigade(f->next, bb);
	}

	/* rl is in kilo bytes / second */
	ratelimit = atoi(rl) * 1024;
	if (ratelimit <= 0) {
	/* remove ourselves */
	ap_log_rerror(APLOG_MARK, APLOG_INFO, 0, f->r,
	APLOGNO(03488) "rl: disabling: rate-limit = %s (too high?)", rl);
	ap_remove_output_filter(f);
	return ap_pass_brigade(f->next, bb);
	}

	/* Configuration: optional initial burst */
	rl = apr_table_get(f->r->subprocess_env, "rate-initial-burst");
	if (rl != NULL) {
	burst = atoi(rl) * 1024;
	if (burst <= 0) {
	ap_log_rerror(APLOG_MARK, APLOG_INFO, 0, f->r,
	APLOGNO(03489) "rl: disabling burst: rate-initial-burst = %s (too high?)", rl);
	burst = 0;
	}
	}

	/* Set up our context */
	ctx = apr_palloc(f->r->pool, sizeof(rl_ctx_t));
	f->ctx = ctx;
	ctx->state = RATE_LIMIT;
	ctx->speed = ratelimit;
	ctx->burst = burst;
	ctx->do_sleep = 0;

	/* calculate how many bytes / interval we want to send */
	/* speed is bytes / second, so, how many (speed / 1000 % interval) */
	ctx->chunk_size = (ctx->speed / (1000 / RATE_INTERVAL_MS));
	ctx->tmpbb = apr_brigade_create(f->r->pool, ba);
	ctx->holdingbb = apr_brigade_create(f->r->pool, ba);
	}
	else {
	APR_BRIGADE_PREPEND(bb, ctx->holdingbb);
	}

	while (!APR_BRIGADE_EMPTY(bb)) {
	apr_bucket *e;

	if (ctx->state == RATE_FULLSPEED) {
	/* Find where we 'stop' going full speed. */
	for (e = APR_BRIGADE_FIRST(bb);
	e != APR_BRIGADE_SENTINEL(bb); e = APR_BUCKET_NEXT(e)) {
	if (AP_RL_BUCKET_IS_END(e)) {
	apr_brigade_split_ex(bb, e, ctx->holdingbb);
	ctx->state = RATE_LIMIT;
	break;
	}
	}

	e = apr_bucket_flush_create(ba);
	APR_BRIGADE_INSERT_TAIL(bb, e);
	rv = ap_pass_brigade(f->next, bb);
	apr_brigade_cleanup(bb);

	if (rv != APR_SUCCESS) {
	ap_log_rerror(APLOG_MARK, APLOG_TRACE1, rv, f->r, APLOGNO(01455)
	"rl: full speed brigade pass failed.");
	return rv;
	}
	}
	else {
	for (e = APR_BRIGADE_FIRST(bb);
	e != APR_BRIGADE_SENTINEL(bb); e = APR_BUCKET_NEXT(e)) {
	if (AP_RL_BUCKET_IS_START(e)) {
	apr_brigade_split_ex(bb, e, ctx->holdingbb);
	ctx->state = RATE_FULLSPEED;
	break;
	}
	}

	while (!APR_BRIGADE_EMPTY(bb)) {
	apr_off_t len = ctx->chunk_size + ctx->burst;

	APR_BRIGADE_CONCAT(ctx->tmpbb, bb);

	/*
	* Pull next chunk of data; the initial amount is our
	* burst allotment (if any) plus a chunk. All subsequent
	* iterations are just chunks with whatever remaining
	* burst amounts we have left (in case not done in the
	* first bucket).
	*/
	rv = apr_brigade_partition(ctx->tmpbb, len, &e);
	if (rv != APR_SUCCESS && rv != APR_INCOMPLETE) {
	ap_log_rerror(APLOG_MARK, APLOG_ERR, rv, f->r, APLOGNO(01456)
	"rl: partition failed.");
	return rv;
	}
	/* Send next metadata now if any */
	while (e != APR_BRIGADE_SENTINEL(ctx->tmpbb)
	&& APR_BUCKET_IS_METADATA(e)) {
	e = APR_BUCKET_NEXT(e);
	}
	if (e != APR_BRIGADE_SENTINEL(ctx->tmpbb)) {
	apr_brigade_split_ex(ctx->tmpbb, e, bb);
	}
	else {
	apr_brigade_length(ctx->tmpbb, 1, &len);
	}

	/*
	* Adjust the burst amount depending on how much
	* we've done up to now.
	*/
	if (ctx->burst) {
	ap_log_rerror(APLOG_MARK, APLOG_DEBUG, 0, f->r,
	APLOGNO(03485) "rl: burst %d; len %"APR_OFF_T_FMT, ctx->burst, len);
	if (len < ctx->burst) {
	ctx->burst -= len;
	}
	else {
	ctx->burst = 0;
	}
	}

	e = APR_BRIGADE_LAST(ctx->tmpbb);
	if (APR_BUCKET_IS_EOS(e)) {
	ap_remove_output_filter(f);
	}
	else if (!APR_BUCKET_IS_FLUSH(e)) {
	if (APR_BRIGADE_EMPTY(bb)) {
	/* Wait for more (or next call) */
	break;
	}
	e = apr_bucket_flush_create(ba);
	APR_BRIGADE_INSERT_TAIL(ctx->tmpbb, e);
	}

	#if defined(RLFDEBUG)
	brigade_dump(f->r, ctx->tmpbb);
	brigade_dump(f->r, bb);
	#endif /* RLFDEBUG */

	if (ctx->do_sleep) {
	apr_sleep(RATE_INTERVAL_MS * 1000);
	}
	else {
	ctx->do_sleep = 1;
	}

	rv = ap_pass_brigade(f->next, ctx->tmpbb);
	apr_brigade_cleanup(ctx->tmpbb);

	if (rv != APR_SUCCESS) {
	/* Most often, user disconnects from stream */
	ap_log_rerror(APLOG_MARK, APLOG_TRACE1, rv, f->r, APLOGNO(01457)
	"rl: brigade pass failed.");
	return rv;
	}
	}
	}

	if (!APR_BRIGADE_EMPTY(ctx->holdingbb)) {
	/* Any rate-limited data in tmpbb is sent unlimited along
	* with the rest.
	*/
	APR_BRIGADE_CONCAT(bb, ctx->tmpbb);
	APR_BRIGADE_CONCAT(bb, ctx->holdingbb);
	}
	}

	#if defined(RLFDEBUG)
	brigade_dump(f->r, ctx->tmpbb);
	#endif /* RLFDEBUG */

	/* Save remaining tmpbb with the correct lifetime for the next call */
	return ap_save_brigade(f, &ctx->holdingbb, &ctx->tmpbb, f->r->pool);
	}


	static apr_status_t
	rl_bucket_read(apr_bucket b, const char *str,
	apr_size_t *len, apr_read_type_e block)
	{
	*str = NULL;
	*len = 0;
	return APR_SUCCESS;
	}

	AP_RL_DECLARE(apr_bucket *)
	ap_rl_end_create(apr_bucket_alloc_t *list)
	{
	apr_bucket b = apr_bucket_alloc(sizeof(b), list);

	APR_BUCKET_INIT(b);
	b->free = apr_bucket_free;
	b->list = list;
	b->length = 0;
	b->start = 0;
	b->data = NULL;
	b->type = &ap_rl_bucket_type_end;

	return b;
	}

	AP_RL_DECLARE(apr_bucket *)
	ap_rl_start_create(apr_bucket_alloc_t *list)
	{
	apr_bucket b = apr_bucket_alloc(sizeof(b), list);

	APR_BUCKET_INIT(b);
	b->free = apr_bucket_free;
	b->list = list;
	b->length = 0;
	b->start = 0;
	b->data = NULL;
	b->type = &ap_rl_bucket_type_start;

	return b;
	}



	AP_RL_DECLARE_DATA const apr_bucket_type_t ap_rl_bucket_type_end = {
	"RL_END", 5, APR_BUCKET_METADATA,
	apr_bucket_destroy_noop,
	rl_bucket_read,
	apr_bucket_setaside_noop,
	apr_bucket_split_notimpl,
	apr_bucket_simple_copy
	};


	AP_RL_DECLARE_DATA const apr_bucket_type_t ap_rl_bucket_type_start = {
	"RL_START", 5, APR_BUCKET_METADATA,
	apr_bucket_destroy_noop,
	rl_bucket_read,
	apr_bucket_setaside_noop,
	apr_bucket_split_notimpl,
	apr_bucket_simple_copy
	};




	static void register_hooks(apr_pool_t *p)
	{
	/* run after mod_deflate etc etc, but not at connection level, ie, mod_ssl. */
	ap_register_output_filter(RATE_LIMIT_FILTER_NAME, rate_limit_filter,
	NULL, AP_FTYPE_CONNECTION - 1);
	}

	AP_DECLARE_MODULE(ratelimit) = {
	STANDARD20_MODULE_STUFF,
	NULL, /* create per-directory config structure */
	NULL, /* merge per-directory config structures */
	NULL, /* create per-server config structure */
	NULL, /* merge per-server config structures */
	NULL, /* command apr_table_t */
	register_hooks
	};