nginx handler模块实现流量统计

长理快乐大学生

于 2024-07-16 14:54:36 发布

阅读量669

点赞数 25

分类专栏： nginx 文章标签： nginx

本文链接：https://blog.csdn.net/weixin_43852750/article/details/140451157

版权

nginx 专栏收录该内容

2 篇文章 0 订阅

订阅专栏

nginx handler模块开发

一开发目的
二开发流程
三完整代码及运行

最近在学习nginx，所以参考别人的代码实现一些自定义模块，这篇博客主要是记载自己对代码的理解。以此也进行学习nginx的模块化开发。

由于我的上一篇博客已经讲了过滤器模块开发，这篇博客不会花很多篇幅去讲解模块开发的流程，而重点专注于handler模块开发的区别和流量统计功能的实现。

一开发目的

首先介绍一下handler模块。
handler模块就是客户端到nginx这条蓝线加上nginx到客户端这条绿色的线。
应用场景包括，ip统计，黑白名单之类。
在这里插入图片描述
这篇文章我们想实现的功能就是做一个流量统计功能，对于不同客户端的访问，我们做一个访问次数的统计，然后把信息加入到返回给客户端的html网页内容中。

大概的思路是开辟一个共享内存，使用红黑树来管理ip地址到count的对应关系。然后每次请求的返回我们对红黑树进行查询，并且输出结果。

二开发流程

2.1 定义模块及上下文ctx和cmd命令组

cmd命令
定义了一个location级别的count命令。
ngx_http_pagecount_set 函数中会进行共享内存的设置，我们之后再看。

static ngx_command_t count_commands[] = {
	{
		ngx_string("count"),
		NGX_HTTP_LOC_CONF | NGX_CONF_NOARGS, //location指令 不接受任何参数
		ngx_http_pagecount_set, //这个函数将在解析到count指令时被调用
		NGX_HTTP_LOC_CONF_OFFSET, //表示这是一个HTTP location级别的配置指令。
		0, NULL
	},
	ngx_null_command
};

ctx上下文
这里有两个回调
ngx_http_pagecount_init
ngx_http_pagecount_create_location_conf

static ngx_http_module_t count_ctx = {
	NULL,
	ngx_http_pagecount_init,
	
	NULL,
	NULL,

	NULL,
	NULL,

	ngx_http_pagecount_create_location_conf,
	NULL,
};

ngx_http_pagecount_init 可以做一些初始化工作
ngx_http_pagecount_create_location_conf在解析location时调用我们可以初始化一下我们的自定义的conf结构体。

模块定义

ngx_module_t ngx_http_pagecount_module = {
	NGX_MODULE_V1,
	&count_ctx,
	count_commands,
	NGX_HTTP_MODULE,
	NULL,
	NULL,
	NULL,
	NULL,
	NULL,
	NULL,
	NULL,
	NGX_MODULE_V1_PADDING
};

2.2 定义相关的结构体

//节点data
typedef struct {
	int count; //count
} ngx_http_pagecount_node_t;

typedef struct {

	ngx_rbtree_t rbtree; //红黑树
	ngx_rbtree_node_t sentinel; //每个红黑树都有一个特殊的哨兵节点（sentinel），通常用于表示树的末端 叶子节点的子节点
	
} ngx_http_pagecount_shm_t;

typedef struct
{
    ssize_t shmsize;
	//slab 池 共享内存
    ngx_slab_pool_t *shpool;
	//红黑树节点
    ngx_http_pagecount_shm_t *sh;
} ngx_http_pagecount_conf_t;

这里主要是三个结构体，第一个node结构体我们用来存count变量，到时候会作为红黑树的节点数据，第二个结构体代表一颗红黑树，第三个结构体我们存放共享内存的slab池和红黑树。

这里再实现一下ctx上下文里的两个回调函数
init函数我们在这里是一个空实现，create_conf函数中我们为当前的配置文件ngx_http_pagecount_conf_t
分配空间

ngx_int_t   ngx_http_pagecount_init(ngx_conf_t *cf) {

	return NGX_OK;
}
void  *ngx_http_pagecount_create_location_conf(ngx_conf_t *cf) {

	ngx_http_pagecount_conf_t *conf;
	
	//为结构体分配空间
	//这里创建的conf会传入到count_commands的ngx_http_pagecount_set函数中
	conf = ngx_palloc(cf->pool, sizeof(ngx_http_pagecount_conf_t));
	if (NULL == conf) {
		return NULL;
	}

	conf->shmsize = 0;

	return conf;

}

2.3 ngx_http_pagecount_set和nginx共享内存的使用

前面的工作做完，我们整个模块的框架已经搭好了。剩下的工作就是在ngx_http_pagecount_set函数中实现我们想要实现的功能

我们先来看一下这个函数中我们做了哪些事情

//这里进行初始化工作
//配置和初始化一个共享内存区域，并将其与一个特定的 NGINX 模块相关联，同时设置该模块的请求处理函数。
//conf 这个参数是指向当前模块的配置结构体的指针
static char *ngx_http_pagecount_set(ngx_conf_t *cf, ngx_command_t *cmd, void *conf) {

	ngx_shm_zone_t *shm_zone;	
	ngx_str_t name = ngx_string("pagecount_slab_shm");

	ngx_http_pagecount_conf_t *mconf = (ngx_http_pagecount_conf_t*)conf;
	ngx_http_core_loc_conf_t *corecf;
	
	//ngx_log_error(NGX_LOG_EMERG, cf->log, ngx_errno, "ngx_http_pagecount_set000");
	//分配1M
	mconf->shmsize = 1024*1024;
	
	// 添加共享内存区域 
	shm_zone = ngx_shared_memory_add(cf, &name, mconf->shmsize, &ngx_http_pagecount_module);
	if (NULL == shm_zone) {
		return NGX_CONF_ERROR;
	}

	//设置共享内存区域的初始化函数
	shm_zone->init = ngx_http_pagecount_shm_init;
	//将配置结构体 mconf 赋值给共享内存区域的 data 字段
	shm_zone->data = mconf;
	//获取 HTTP 核心模块的位置配置信息
	corecf = ngx_http_conf_get_module_loc_conf(cf, ngx_http_core_module);
	//设置 HTTP 核心模块的处理函数
	corecf->handler = ngx_http_pagecount_handler;

	return NGX_CONF_OK;
}

1 .首先我们声明了一个ngx_shm_zone_t 这是nginx中定义的一个共享内存的结构体，我们看一下结构，重点关注data放我们的实际数据，init函数用来进行初始化

struct ngx_shm_zone_s {
    void                     *data;          // 指向共享内存区域的数据指针，可以存储模块需要共享的数据
    ngx_shm_t                 shm;           // ngx_shm_t 结构体，用于描述共享内存的详细信息
    ngx_shm_zone_init_pt      init;          // 初始化回调函数指针，用于在共享内存区域首次分配时进行初始化
    void                     *tag;           // 标记，可以用于标识共享内存区域的唯一性或其他用途
    void                     *sync;          // 同步对象，用于多进程/线程环境下的访问同步
    ngx_uint_t                noreuse;       // 表示该共享内存区域是否可以重新使用的标志位
};

2.ngx_http_pagecount_conf_t mconf = (ngx_http_pagecount_conf_t)conf;

这里我们拿到了我们在ngx_http_pagecount_create_location_conf已经分配好空间的结构体。因为ngx_http_pagecount_create_location_conf在解析location时调用，而ngx_http_pagecount_set函数在解析到count命令时才进行调用。所以此时conf已经分配好空间。

mconf->shmsize = 1024*1024; 我们给共享内存分配1M的空间
4.ngx_shared_memory_add 用于在 NGINX 中创建和管理共享内存区域，原型如下

ngx_shm_zone_t *ngx_shared_memory_add(ngx_conf_t *cf, ngx_str_t *name, size_t size, void *tag);

这个函数的作用是：

创建共享内存区域：它会在 NGINX 的共享内存池中分配一块指定大小的内存区域，并且为这个区域分配一个唯一的名字（通过 name 参数）。

返回共享内存管理结构：返回一个 ngx_shm_zone_t 结构体指针，这个结构体用于后续的共享内存管理，包括初始化和访问。

在这个示例中，ngx_shared_memory_add() 函数创建了一个大小为 1MB 的共享内存区域，并与当前的 HTTP 模块（通过 &ngx_http_module）关联起来

5.设置 ngx_shm_zone_t的init函数和data，data我们放入我们自定义的conf结构体。
data 字段在共享内存区域初始化完成后会被设置为指向实际的共享内存数据结构。这个数据结构可以是任何类型，通常由模块定义。
在调用 ngx_shared_memory_add() 创建共享内存区域时，需要指定一个初始化函数（通过 init 字段）。这个初始化函数负责设置 data 字段。

这个init函数的实现我们放在后面讲

6.指定handler回调

//获取 HTTP 核心模块的位置配置信息
corecf = ngx_http_conf_get_module_loc_conf(cf, ngx_http_core_module);
//设置 HTTP 核心模块的处理函数
corecf->handler = ngx_http_pagecount_handler;

ngx_http_core_module 是 NGINX 中的一个核心模块，主要负责处理 HTTP 请求的基本配置和请求处理流程

解析配置文件
当 NGINX 启动或重新加载配置文件时，它会解析配置文件并处理其中的指令。对于每一个 location 指令，如果遇到 count 指令，就会调用 ngx_http_pagecount_set 函数来进行相应的设置。
设置 handler
在 ngx_http_pagecount_set 函数中，corecf->handler 被设置为 ngx_http_pagecount_handler。corecf 是通过 ngx_http_conf_get_module_loc_conf 获取的 HTTP 核心模块的位置配置信息。
请求匹配和处理
当 NGINX 接收到一个 HTTP 请求时，它会根据 URI 路径来匹配相应的 server 和 location 配置。如果请求的 URI 匹配到一个包含 count 指令的 location，NGINX 会使用之前在配置阶段设置的 handler 函数来处理这个请求。

ngx_http_pagecount_handler这个回调我们也在后面介绍

2.4 init 共享内存的初始化

//初始化ngx_http_pagecount_conf_t
ngx_int_t ngx_http_pagecount_shm_init (ngx_shm_zone_t *zone, void *data) {

	ngx_http_pagecount_conf_t *conf;
	ngx_http_pagecount_conf_t *oconf = data;

	//拿到自定义结构体对象
	conf = (ngx_http_pagecount_conf_t*)zone->data;
	// 如果 data 参数不为空，表示这是共享内存的重新初始化（如重新加载配置）
	if (oconf) {
		conf->sh = oconf->sh;
		conf->shpool = oconf->shpool;
		return NGX_OK;
	}

	//printf("ngx_http_pagecount_shm_init 0000\n");
	// 第一次初始化共享内存区域 获取slab_pool
	conf->shpool = (ngx_slab_pool_t*)zone->shm.addr;
	// 使用 slab 分配器在共享内存池中分配红黑树结构体
	conf->sh = ngx_slab_alloc(conf->shpool, sizeof(ngx_http_pagecount_shm_t));
	if (conf->sh == NULL) {
		return NGX_ERROR;
	}
	//把slab_pool中的data设为红黑树
	conf->shpool->data = conf->sh;

	//printf("ngx_http_pagecount_shm_init 1111\n");
	
	// 初始化红黑树 传入自定义的插入函数
	ngx_rbtree_init(&conf->sh->rbtree, &conf->sh->sentinel, 
		ngx_http_pagecount_rbtree_insert_value);


	return NGX_OK;

}

在set函数中我们已经设置了shm_zone对象的data等于我们的自定义conf，这里取出来，然后zone对象本身自己内部有一个slab内存分配器，我们用这个分配器给我们conf中的红黑树结构体分配空间，这里只是给结构体分配空间，然后再调用ngx_rbtree_init 进行红黑树的初始化。这里的红黑树初始化我们传入了自定义插入函数ngx_http_pagecount_rbtree_insert_value，其实就是nginx内部提供的插入函数ngx_rbtree_insert_value ，我们在2.5节说明。

2.5 nginx红黑树使用和相关操作

在nginx中我们一般使用红黑树的步骤。
1.定义红黑树节点结构
定义一个结构体，用于作为红黑树节点的数据结构，通常继承自 ngx_rbtree_node_t 结构体，并包含业务数据。

typedef struct {
    ngx_rbtree_node_t node;  // 红黑树节点
    ngx_str_t key;           // 节点的键值
    // 其他业务数据
} ngx_my_node_t;

2.定义红黑树管理结构

typedef struct {
    ngx_rbtree_t rbtree;     // 红黑树的根节点
    ngx_rbtree_node_t sentinel;  // 哨兵节点
    // 其他辅助信息
} ngx_my_tree_t;

3.初始化红黑树
使用 ngx_rbtree_init 函数来初始化红黑树，需要传入红黑树管理结构、哨兵节点以及比较函数。

ngx_my_tree_t my_tree;
ngx_rbtree_node_t *sentinel = &my_tree.sentinel;

ngx_rbtree_init(&my_tree.rbtree, sentinel, ngx_rbtree_insert_value);

4.插入节点
使用 ngx_rbtree_insert 函数来插入一个新节点到红黑树中。

ngx_my_node_t *node = ngx_pcalloc(pool, sizeof(ngx_my_node_t));
node->key = ngx_string("some_key");

ngx_rbtree_insert(&my_tree.rbtree, &node->node);

5.查找节点
在 NGINX 中，通常需要自行实现查找节点的功能。这通常涉及从根节点开始，根据节点的键值递归地遍历红黑树，直到找到匹配的节点或遍历到哨兵节点为止。

ngx_my_node_t *ngx_my_tree_lookup(ngx_my_tree_t *tree, ngx_str_t *key) {
    ngx_rbtree_node_t *node = tree->rbtree.root;
    ngx_rbtree_node_t *sentinel = tree->rbtree.sentinel;

    while (node != sentinel) {
        ngx_my_node_t *my_node = (ngx_my_node_t *) node;
        int cmp = ngx_strncmp(key->data, my_node->key.data, key->len);

        if (cmp == 0) {
            return my_node;  // 找到了匹配的节点
        } else if (cmp < 0) {
            node = node->left;
        } else {
            node = node->right;
        }
    }

    return NULL;  // 没有找到匹配的节点
}

6.删除节点

void ngx_my_tree_delete(ngx_my_tree_t *tree, ngx_my_node_t *node) {
    ngx_rbtree_delete(&tree->rbtree, &node->node);
}

这里我们只需要用到插入以及查询
插入的代码和nginx 提供的 ngx_rbtree_insert_value 一致

static void
ngx_http_pagecount_rbtree_insert_value(ngx_rbtree_node_t *temp,
        ngx_rbtree_node_t *node, ngx_rbtree_node_t *sentinel)
{
   ngx_rbtree_node_t **p;
   //ngx_http_testslab_node_t *lrn, *lrnt;
 
    for (;;)
    {
        if (node->key < temp->key)
        {
            p = &temp->left;
        }
        else if (node->key > temp->key) {
           	p = &temp->right;
        }
        else
        {
          	return ;
        }
 
        if (*p == sentinel)
        {
            break;
        }
 
        temp = *p;
    }
 
    *p = node;
 
    node->parent = temp;
    node->left = sentinel;
    node->right = sentinel;
    ngx_rbt_red(node);
}

查找函数我们要自己手动实现，这里如果找到了对应的键，说明是一个修改操作，我们就把data++，如果没有找到对应的键说明是一个插入操作，我们就执行插入。

//查找函数需要自己手动实现
static ngx_int_t ngx_http_pagecount_lookup(ngx_http_request_t *r, ngx_http_pagecount_conf_t *conf, ngx_uint_t key) {

    //根节点和哨兵节点
	ngx_rbtree_node_t *node, *sentinel;

	node = conf->sh->rbtree.root;
	sentinel = conf->sh->rbtree.sentinel;

    //记录日志，显示正在查找的键值 key。
	ngx_log_error(NGX_LOG_EMERG, r->connection->log, ngx_errno, " ngx_http_pagecount_lookup 111 --> %x\n", key);
	
	while (node != sentinel) {

		if (key < node->key) {
			node = node->left;
			continue;
		} else if (key > node->key) {
			node = node->right;
			continue;
		} else { // key == node //找到对应的节点 让count++
			node->data ++;
			return NGX_OK;
		}
	}
	
	ngx_log_error(NGX_LOG_EMERG, r->connection->log, ngx_errno, " ngx_http_pagecount_lookup 222 --> %x\n", key);
	
	// insert rbtree
	//没找到的话就执行插入
	
	node = ngx_slab_alloc_locked(conf->shpool, sizeof(ngx_rbtree_node_t));
	if (NULL == node) {
		return NGX_ERROR;
	}

	node->key = key;
	node->data = 1;

    //插入节点
	ngx_rbtree_insert(&conf->sh->rbtree, node);

	ngx_log_error(NGX_LOG_EMERG, r->connection->log, ngx_errno, " insert success\n");
	

	return NGX_OK;
}

2.6 ngx_http_pagecount_handler 及相关操作

首先我们提供一个辅助函数，帮助将红黑树中所有数据转为html网页

//将红黑树中的数据转换为 HTML 格式，以便在网页中显示
static int ngx_encode_http_page_rb(ngx_http_pagecount_conf_t *conf, char *html) {

	sprintf(html, "<h1>Source Insight </h1>");
	strcat(html, "<h2>");

	//使用 ngx_rbtree_min 查找红黑树中的最小节点并将其赋值给 node。
	//ngx_rbtree_traversal(&ngx_pv_tree, ngx_pv_tree.root, ngx_http_count_rbtree_iterator, html);
	ngx_rbtree_node_t *node = ngx_rbtree_min(conf->sh->rbtree.root, conf->sh->rbtree.sentinel);

	do {
		// 声明并初始化两个缓冲区 str 和 buffer。
		char str[INET_ADDRSTRLEN] = {0};
		char buffer[128] = {0};
		//将节点的键（ IP 地址）转换为字符串形式并存储在 str 中
		sprintf(buffer, "req from : %s, count: %d <br/>",
			inet_ntop(AF_INET, &node->key, str, sizeof(str)), node->data);

		strcat(html, buffer);
		// 使用 ngx_rbtree_next 获取红黑树中的下一个节点。
		node = ngx_rbtree_next(&conf->sh->rbtree, node);

	} while (node);
	
	
	strcat(html, "</h2>");

	return NGX_OK;
}

然后看handler函数

static ngx_int_t ngx_http_pagecount_handler(ngx_http_request_t *r) {

	//用于存放生成的HTML内容。
	u_char html[1024] = {0};
	int len = sizeof(html);
	
	ngx_rbtree_key_t key = 0;


	struct sockaddr_in *client_addr =  (struct sockaddr_in*)r->connection->sockaddr;
	
	//使用 ngx_http_get_module_loc_conf 函数获取当前 HTTP 请求的模块配置信息
	ngx_http_pagecount_conf_t *conf = ngx_http_get_module_loc_conf(r, ngx_http_pagecount_module);
	//将客户端IP地址的32位表示设置为关键字 key
	key = (ngx_rbtree_key_t)client_addr->sin_addr.s_addr;
	
	//记录日志信息：
	ngx_log_error(NGX_LOG_EMERG, r->connection->log, ngx_errno, " ngx_http_pagecount_handler --> %x\n", key);
	//加锁共享内存：
	ngx_shmtx_lock(&conf->shpool->mutex);
	//调用页面计数模块的查找函数，传递当前请求、配置信息和关键字参数。这个函数可能会在共享内存中更新计数器或其他相关数据。
	ngx_http_pagecount_lookup(r, conf, key);	
	//解锁
	ngx_shmtx_unlock(&conf->shpool->mutex);
	//生成HTTP响应内容
	ngx_encode_http_page_rb(conf, (char*)html);

	//header 设置HTTP响应头部：
	r->headers_out.status = 200;
	ngx_str_set(&r->headers_out.content_type, "text/html");
	ngx_http_send_header(r);

	//body 准备输出缓冲区：
	ngx_buf_t *b = ngx_pcalloc(r->pool,  sizeof(ngx_buf_t));

	ngx_chain_t out;
	out.buf = b;
	out.next = NULL;

	b->pos = html;
	b->last = html+len;
	b->memory = 1;
	b->last_buf = 1;

	return ngx_http_output_filter(r, &out);
	
	
}

这里首先我们拿到用户请求的ip地址作为key，然后加锁访问共享内存。
如果用户IP已经存在我们就做++，如果不存在我们做插入。
最后分别输出header和body。
源代码中ngx_encode_http_page_rb(conf, (char*)html); 是放在锁外面的，但是个人觉得应该将两步操作都进行加锁，否则可能会引起进程安全问题。读者可以自行尝试

//加锁共享内存：
	ngx_shmtx_lock(&conf->shpool->mutex);
	//调用页面计数模块的查找函数，传递当前请求、配置信息和关键字参数。这个函数可能会在共享内存中更新计数器或其他相关数据。
	ngx_http_pagecount_lookup(r, conf, key);	
	//生成HTTP响应内容
	ngx_encode_http_page_rb(conf, (char*)html);
	//解锁
	ngx_shmtx_unlock(&conf->shpool->mutex);

NGINX 提供了跨进程互斥锁 (ngx_shmtx_t) 来保护共享内存。使用 ngx_shmtx_lock 和 ngx_shmtx_unlock 函数，可以确保只有一个进程在任何时候访问共享内存的临界区。