Nginx中发送udp请求

http://blog.aka-cool.net/blog/2014/02/24/udp-in-nginx/

最近简单了解了一下Nginx发送udp请求的过程, 在这里简单记录一下.

主要参考的代码主要有两块, 分别是ngx_resolver.c以及agentzh的ngx_lua模块中的udp相关代码(ngx_http_lua_socket_udp.c). 有兴趣的同学可以看一下.

简单分析

首先看一下ngx_resolver.c这个文件, 主要解决了dns查询的问题, 也就是resolver这个指令. 这个文件长度还是蛮长的, 大概有2000多行, 但是udp部分的代码还是相当简洁的, 主要分为几个步骤: 初始化, 连接, 发送以及回收.

初始化的代码:

ngx_resolver.c

ngx_resolver_t *
ngx_resolver_create(ngx_conf_t *cf, ngx_str_t *names, ngx_uint_t n)
{
    ngx_str_t              s;
    ngx_url_t              u;
    ngx_uint_t             i, j;
    ngx_resolver_t        *r;
    ngx_pool_cleanup_t    *cln;
    ngx_udp_connection_t  *uc;

    /* 设置回收函数 */

    cln = ngx_pool_cleanup_add(cf->pool, 0);
    if (cln == NULL) {
        return NULL;
    }

    cln->handler = ngx_resolver_cleanup;

    r = ngx_calloc(sizeof(ngx_resolver_t), cf->log);
    if (r == NULL) {
        return NULL;
    }

    cln->data = r;

    /*
        省略很多代码, 初始化红黑树以及队列
    */

    for (i = 0; i < n; i++) {
        if (ngx_strncmp(names[i].data, "valid=", 6) == 0) {
            /*
                省略valid参数的设置
            */
        }

        /* 进入正题 */
        ngx_memzero(&u, sizeof(ngx_url_t));

        u.url = names[i];
        u.default_port = 53;

        /* ngx_parse_url用于解析url, 获取ip, port, 地址等信息 */

        if (ngx_parse_url(cf->pool, &u) != NGX_OK) {
            if (u.err) {
                ngx_conf_log_error(NGX_LOG_EMERG, cf, 0,
                                   "%s in resolver \"%V\"",
                                   u.err, &u.url);
            }

            return NULL;
        }

        uc = ngx_array_push_n(&r->udp_connections, u.naddrs);
        if (uc == NULL) {
            return NULL;
        }

        /* 设置各个ngx_udp_connection_t */

        ngx_memzero(uc, u.naddrs * sizeof(ngx_udp_connection_t));

        for (j = 0; j < u.naddrs; j++) {
            /* 看上去只需要设置这几个值;) */
            uc[j].sockaddr = u.addrs[j].sockaddr;
            uc[j].socklen = u.addrs[j].socklen;
            uc[j].server = u.addrs[j].name;
        }
    }

    return r;
}

这里主要就是根据配置文件来初始化ngx_udp_connection_t类型的变量. 接下来是连接和发送请求:

ngx_resolver_send_query

static ngx_int_t
ngx_resolver_send_query(ngx_resolver_t *r, ngx_resolver_node_t *rn)
{
    ssize_t                n;
    ngx_udp_connection_t  *uc;

    uc = r->udp_connections.elts;

    uc = &uc[r->last_connection++];
    if (r->last_connection == r->udp_connections.nelts) {
        r->last_connection = 0;
    }

    if (uc->connection == NULL) {

        uc->log = *r->log;
        uc->log.handler = ngx_resolver_log_error;
        uc->log.data = uc;
        uc->log.action = "resolving";

        /* 进行连接 */
        if (ngx_udp_connect(uc) != NGX_OK) {
            return NGX_ERROR;
        }

        uc->connection->data = r;
        uc->connection->read->handler = ngx_resolver_read_response;
        uc->connection->read->resolver = 1;
    }

    /* 发送数据 */
    n = ngx_send(uc->connection, rn->query, rn->qlen);

    if (n == -1) {
        return NGX_ERROR;
    }

    if ((size_t) n != (size_t) rn->qlen) {
        ngx_log_error(NGX_LOG_CRIT, &uc->log, 0, "send() incomplete");
        return NGX_ERROR;
    }

    return NGX_OK;
}

主要是使用了ngx_udp_connect和ngx_send这两个函数. 用起来还是相当方便的, 里面的实现有兴趣的同学可以看一下. 最后就是关闭连接:

ngx_resolver_cleanup

static void
ngx_resolver_cleanup(void *data)
{
    ngx_resolver_t  *r = data;

    ngx_uint_t             i;
    ngx_udp_connection_t  *uc;

    if (r) {
        ngx_log_debug0(NGX_LOG_DEBUG_CORE, ngx_cycle->log, 0,
                       "cleanup resolver");

        ngx_resolver_cleanup_tree(r, &r->name_rbtree);

        ngx_resolver_cleanup_tree(r, &r->addr_rbtree);

        if (r->event) {
            ngx_free(r->event);
        }


        uc = r->udp_connections.elts;

        /* 关闭每一个建立的连接 */
        for (i = 0; i < r->udp_connections.nelts; i++) {
            if (uc[i].connection) {
                ngx_close_connection(uc[i].connection);
            }
        }

        ngx_free(r);
    }
}

主要是调用了ngx_close_connection这个函数. 通过以上操作就完成了一次udp数据的发送. ngx_lua当中也基本是这个步骤.

总结下就是:

1. 初始化ngx_udp_connection_t类型的变量(通常会用到ngx_parse_url)
2. 使用ngx_udp_connect进行连接
3. 使用ngx_send进行数据发送
4. 关闭连接, 可以在数据回收的回调中进行关闭

自己写一个试试

看明白了基本步骤, 自己验证一下, 写一个的对应的简单handler模块. 在配置文件中设置udp_address来设置要发送的地址, nginx接受到请求之后即向这个地址发送信息.

ngx_resolver_cleanup

#include <ngx_config.h>
#include <ngx_core.h>
#include <ngx_http.h>

#include "ddebug.h"

typedef struct {
    ngx_flag_t  enable;

    ngx_udp_connection_t *play_udp_uc;
} ngx_http_play_loc_conf_t;

ngx_int_t ngx_udp_connect(ngx_udp_connection_t *uc);

static void *ngx_http_play_create_loc_conf(ngx_conf_t *cf);

static char *ngx_conf_play(ngx_conf_t *cf, ngx_command_t *cmd, void *conf);

static void ngx_http_upstream_play_cleanup(void *data);

static char *ngx_conf_set_udp_addr(ngx_conf_t *cf, ngx_command_t *cmd, void *conf);

static ngx_int_t ngx_http_play_handler(ngx_http_request_t *r);


static ngx_command_t ngx_http_play_commands[] = {
    {
        ngx_string("play"),
        NGX_HTTP_LOC_CONF | NGX_CONF_TAKE1,
        ngx_conf_play,
        NGX_HTTP_LOC_CONF_OFFSET,
        0,
        NULL
    },
    {
        ngx_string("udp_address"),
        NGX_HTTP_LOC_CONF | NGX_CONF_TAKE1,
        ngx_conf_set_udp_addr,
        NGX_HTTP_LOC_CONF_OFFSET,
        0,
        NULL
    },
    ngx_null_command
};

static ngx_http_module_t ngx_http_play_module_ctx = {
    NULL,                              /* preconfiguration */
    //ngx_http_play_init,              /* postconfiguration */
    NULL,                              /* postconfiguration */

    NULL,                              /* create main configuration */
    NULL,                              /* init main configuration */

    NULL,                              /* create server configuration */
    NULL,                              /* merge server configuration */

    ngx_http_play_create_loc_conf,     /* create location configuration */
    NULL                               /* merge location configuration */
};

ngx_module_t  ngx_http_play_module = {
    NGX_MODULE_V1,
    &ngx_http_play_module_ctx,                      /* module context */
    ngx_http_play_commands,                         /* module directives */
    NGX_HTTP_MODULE,                                /* module type */
    NULL,                                           /* init master */
    NULL,                                           /* init module */
    NULL,                                           /* init process */
    NULL,                                           /* init thread */
    NULL,                                           /* exit thread */
    NULL,                                           /* exit process */
    NULL,                                           /* exit master */
    NGX_MODULE_V1_PADDING
};

static char *
ngx_conf_play(ngx_conf_t *cf, ngx_command_t *cmd, void *conf)
{
    ngx_http_play_loc_conf_t *plcf;
    ngx_http_core_loc_conf_t *clcf;
    ngx_str_t                *value;

    clcf = ngx_http_conf_get_module_loc_conf(cf, ngx_http_core_module);
    plcf = conf;

    value = cf->args->elts;

    ngx_conf_log_error(NGX_LOG_INFO, cf, 0, "play: %s", value[1].data);

    plcf->enable = 0;
    if (ngx_strncmp(value[1].data, "on", 2) == 0) {
        clcf->handler = ngx_http_play_handler;
        plcf->enable  = 1;
    }

    return NGX_CONF_OK;
}

static char *
ngx_conf_set_udp_addr(ngx_conf_t *cf, ngx_command_t *cmd, void *conf)
{
    ngx_http_play_loc_conf_t   *plcf;
    ngx_str_t                  *value;
    ngx_url_t                   u;
    ngx_pool_cleanup_t         *cln;
    ngx_udp_connection_t       *uc;

    plcf = conf;

    ngx_memzero(&u, sizeof(ngx_url_t));

    value = cf->args->elts;

    /* resolve url */
    u.url          = value[1];
    u.default_port = 12345;
    u.no_resolve   = 0;

    if (ngx_parse_url(cf->pool, &u) != NGX_OK) {
        ngx_conf_log_error(NGX_LOG_EMERG, cf, 0, "parse udp addr failed, %s", value[1].data);
        return NGX_CONF_ERROR;
    }

    /* init udp connection */
    plcf->play_udp_uc = ngx_pcalloc(cf->pool, sizeof(ngx_udp_connection_t));
    uc                = plcf->play_udp_uc;

    if (uc == NULL) {
        return NGX_CONF_ERROR;
    }

    /* implement udp connection */
    uc->sockaddr = u.addrs[0].sockaddr;
    uc->socklen  = u.addrs[0].socklen;
    uc->server   = u.addrs[0].name;
    uc->log      = cf->cycle->new_log;

    /* implement clean handler */
    cln = ngx_pool_cleanup_add(cf->pool, 0);
    if (cln == NULL) {
        return NGX_CONF_ERROR;
    }

    cln->data    = uc;
    cln->handler = ngx_http_upstream_play_cleanup;   // set the cleanup handler

    return NGX_OK;
}

static void
ngx_http_upstream_play_cleanup(void *data)
{
    ngx_udp_connection_t  *uc = data;

    ngx_log_debug0(NGX_LOG_DEBUG_CORE, ngx_cycle->log, 0,
                   "cleanup http_upstream_keepalive_fluentd");

    if (uc->connection) {
        ngx_close_connection(uc->connection);
    }
}

static void ngx_http_play_dummy_handler(ngx_event_t *ev)
{
    /* just empty */
}


static void
ngx_http_play_send_udp(ngx_http_request_t *r, ngx_udp_connection_t *uc)
{
    if (uc->connection == NULL) {
        if (ngx_udp_connect(uc) != NGX_OK) {
            return;
        }

        uc->connection->data = NULL;
        uc->connection->read->handler = ngx_http_play_dummy_handler;
        uc->connection->read->resolver = 1;
    }

    ngx_send(uc->connection, (u_char*)"hello", 5);

    return;
}

/* the main function */
static ngx_int_t
ngx_http_play_handler(ngx_http_request_t *r)
{
    ngx_int_t    rc;
    ngx_buf_t   *b;
    ngx_chain_t  out;
    ngx_variable_value_t *var;

    ngx_http_play_loc_conf_t *plcf;

    plcf = ngx_http_get_module_loc_conf(r, ngx_http_play_module);

    /* here we go ;) */
    rc = ngx_http_discard_request_body(r);
    if (rc != NGX_OK) {
        return rc;
    }

    ngx_str_set(&r->headers_out.content_type, "text/html");

    b = ngx_pcalloc(r->pool, sizeof(ngx_buf_t));
    if (b == NULL) {
        return NGX_HTTP_INTERNAL_SERVER_ERROR;
    }

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

    /*
     *
     * I want to change here ;)
     *
     * */

    ngx_http_play_send_udp(r, plcf->play_udp_uc);

    b->pos  = (u_char *)"send";
    b->last = b->pos + 4;


    ///

    b->memory   = 1;
    b->last_buf = 1;

    r->headers_out.status = NGX_HTTP_OK;
    //r->headers_out.content_length_n = var.len;

    rc = ngx_http_send_header(r);
    return ngx_http_output_filter(r, &out);
}


static void *
ngx_http_play_create_loc_conf(ngx_conf_t *cf)
{
    ngx_http_play_loc_conf_t *plcf =  NULL;

    plcf = ngx_pcalloc(cf->pool, sizeof(ngx_http_play_loc_conf_t));

    if (plcf == NULL) {
        return NULL;
    }

    return plcf;
}

使用nc监听一下发现可以成功~