diff --git a/src/tsg_action.cpp b/src/tsg_action.cpp index 3f4db6d..6293e0a 100644 --- a/src/tsg_action.cpp +++ b/src/tsg_action.cpp @@ -579,7 +579,7 @@ static unsigned char do_action_drop(const struct streaminfo *a_stream, Maat_rule static unsigned char do_action_ratelimit(const struct streaminfo *a_stream, Maat_rule_t *p_result, struct compile_user_region *user_region, enum ACTION_RETURN_TYPE type) { struct tcpall_context *context=NULL; - struct leaky_bucket *bucket=create_bucket((double)((user_region->deny->bps)/1000000), user_region->deny->bps, a_stream->threadnum); + struct leaky_bucket *bucket=create_bucket(user_region->deny->bps, a_stream->threadnum); tsg_set_bucket_to_tcpall(a_stream, &context, bucket, a_stream->threadnum); set_ratelimit_flag(a_stream); diff --git a/src/tsg_leaky_bucket.cpp b/src/tsg_leaky_bucket.cpp index 6dbcaf1..41b6a93 100644 --- a/src/tsg_leaky_bucket.cpp +++ b/src/tsg_leaky_bucket.cpp @@ -5,26 +5,57 @@ #include +enum BUCKET_UNIT +{ + BUCKET_UNIT_S=0, + BUCKET_UNIT_MS, + BUCKET_UNIT_US, + BUCKET_UNIT_MAX, +}; + struct leaky_bucket { double rate; int used_size; int bucket_size; + enum BUCKET_UNIT unit; struct timespec refresh_time; }; +#define BUCKET_UNIT_S_LIMIT_MAX_BIT 32*8*1024 +#define BUCKET_UNIT_MS_LIMIT_MAX_BIT 64*8*1024 + void refresh_bucket(struct leaky_bucket * bucket, int thread_seq) { long interval_us=0; struct timespec end; clock_gettime(CLOCK_MONOTONIC, &end); - interval_us = (end.tv_sec - bucket->refresh_time.tv_sec)*1000000 + (end.tv_nsec - bucket->refresh_time.tv_nsec)/1000; - bucket->used_size=bucket->used_size - interval_us*(bucket->rate); + switch(bucket->unit) + { + case BUCKET_UNIT_S: + interval_us=(end.tv_sec-bucket->refresh_time.tv_sec)+(end.tv_nsec-bucket->refresh_time.tv_nsec)/1000000000; + break; + case BUCKET_UNIT_MS: + interval_us=(end.tv_sec-bucket->refresh_time.tv_sec)*1000 + (end.tv_nsec-bucket->refresh_time.tv_nsec)/1000000; + break; + case BUCKET_UNIT_US: + interval_us=(end.tv_sec-bucket->refresh_time.tv_sec)*1000000 + (end.tv_nsec-bucket->refresh_time.tv_nsec)/1000; + break; + default: + bucket->used_size=bucket->bucket_size; + return ; + } + + + bucket->used_size=bucket->used_size-interval_us*(bucket->rate); bucket->used_size=(bucket->used_size<0) ? 0 : bucket->used_size; - bucket->refresh_time=end; - + if(interval_us>=1) + { + bucket->refresh_time=end; + } + return ; } @@ -40,11 +71,31 @@ int is_permit_pass(int pkt_size, struct leaky_bucket * bucket, int thread_seq) return 0; } -struct leaky_bucket * create_bucket(double rate, int bucket_size, int thread_seq) +struct leaky_bucket * create_bucket(int bucket_size, int thread_seq) { - struct leaky_bucket * bucket = (struct leaky_bucket *)dictator_malloc(thread_seq, sizeof(struct leaky_bucket)); + if(bucket_size<0) + { + return NULL; + } + + struct leaky_bucket * bucket = (struct leaky_bucket *)dictator_malloc(thread_seq, sizeof(struct leaky_bucket)); + + if(bucket_sizerate=(double)bucket_size; + bucket->unit=BUCKET_UNIT_S; + } + else if(bucket_sizerate=(double)bucket_size/(double)1000; + bucket->unit=BUCKET_UNIT_MS; + } + else + { + bucket->rate=(double)bucket_size/(double)1000000; + bucket->unit=BUCKET_UNIT_US; + } - bucket->rate = rate; bucket->used_size = 0; bucket->bucket_size = bucket_size; diff --git a/src/tsg_leaky_bucket.h b/src/tsg_leaky_bucket.h index 42a62da..923312b 100644 --- a/src/tsg_leaky_bucket.h +++ b/src/tsg_leaky_bucket.h @@ -3,7 +3,7 @@ struct leaky_bucket; -struct leaky_bucket *create_bucket(double rate, int bucket_size, int thread_seq); +struct leaky_bucket *create_bucket(int bucket_size, int thread_seq); void destroy_bucket(struct leaky_bucket **bucket, int thread_seq); int is_permit_pass(int pkt_size, struct leaky_bucket * bucket, int thread_seq);