利用nginx subrequest实现与多个后端server的交互

/*

 * 利用nginx subrequest实现与多个后端server的交互
 * Nginx version: 1.2.2
 * Copyright (C) eling	
 * Email: luxuejun12@gmai.com
 * 
 */


extern "C" {
#include <ngx_config.h>
#include <ngx_core.h>
#include <ngx_http.h>
}

// nginx header files should go before other, because they define 64-bit off_t

#define LAUTH_EXT_CMD_ARGS_NUM 2

typedef struct {
    ngx_uint_t                n;	// 2
    ngx_array_t              *uris;
} ngx_http_lauth_ext_conf_t;

typedef struct {
	ngx_uint_t                done;
    ngx_uint_t                status;
	ngx_str_t				 *content;
	
    ngx_http_request_t       *sr;
} ngx_http_lauth_ext_subrequest;

typedef struct {
	ngx_uint_t				i;
	ngx_array_t				*subrequests;	
} ngx_http_lauth_ext_ctx_t;

static ngx_int_t ngx_http_lauth_ext_handler(ngx_http_request_t *r);
static ngx_int_t ngx_http_lauth_ext_done(ngx_http_request_t *r,
    void *data, ngx_int_t rc);

static void *ngx_http_lauth_ext_create_conf(ngx_conf_t *cf);

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


static ngx_command_t  ngx_http_lauth_ext_commands[] = {

    { ngx_string("lauth_ext"),
      NGX_HTTP_LOC_CONF|NGX_CONF_TAKE2,
      ngx_http_lauth_ext,
      NGX_HTTP_LOC_CONF_OFFSET,
      0,
      NULL },

      ngx_null_command
};


static ngx_http_module_t  ngx_http_lauth_ext_module_ctx = {
    NULL,                                  /* preconfiguration */
    NULL,            						/* postconfiguration */

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

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

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


ngx_module_t  ngx_http_lauth_ext_module = {
    NGX_MODULE_V1,
    &ngx_http_lauth_ext_module_ctx,     /* module context */
    ngx_http_lauth_ext_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 ngx_int_t
ngx_http_lauth_ext_handler(ngx_http_request_t *r)
{
	ngx_log_error(NGX_LOG_DEBUG, r->connection->log, 0,
                   "lauth_ext request handler, uri:%V args:%V", &r->uri, &r->args);

	ngx_int_t 					rc;
    ngx_http_lauth_ext_conf_t  *lacf;
	ngx_str_t					*uri;
	ngx_str_t					*args;
	
	ngx_http_lauth_ext_ctx_t   		*ctx;	
	ngx_http_lauth_ext_subrequest  	*lasr;
	ngx_http_post_subrequest_t  	*ps;	
	
    lacf = (ngx_http_lauth_ext_conf_t *)ngx_http_get_module_loc_conf(r, ngx_http_lauth_ext_module); 	
	
    ctx = (ngx_http_lauth_ext_ctx_t *)ngx_http_get_module_ctx(r, ngx_http_lauth_ext_module);

//	args = &r->args;

    if (ctx != NULL) {
		ngx_http_lauth_ext_subrequest* lasr = (ngx_http_lauth_ext_subrequest*)ctx->subrequests->elts + (ctx->i - 1);
		
        if(ctx->i == LAUTH_EXT_CMD_ARGS_NUM){
			if(lasr->status == NGX_HTTP_INTERNAL_SERVER_ERROR){
				ngx_http_finalize_request(r, NGX_HTTP_INTERNAL_SERVER_ERROR);
				return NGX_OK;
			}
			else if(lasr->status != NGX_HTTP_OK){
				r->headers_out.status = lasr->status;
				r->header_only = 1;//???

				rc = ngx_http_send_header(r);
				ngx_http_finalize_request(r, rc);
				return NGX_OK;
			}

			r->headers_out.status = lasr->status;			
			r->headers_out.content_length_n = lasr->content->len;			
			
			rc = ngx_http_send_header(r);
    		if (rc == NGX_ERROR || rc > NGX_OK || r->header_only) {
        		ngx_http_finalize_request(r, rc);
				return NGX_OK;
    		}			

			ngx_chain_t out;
			ngx_buf_t b;

			ngx_memzero(&b, sizeof(ngx_buf_t));
			b.temporary = 1;
			b.last_buf = 1;	//???
			
			b.pos = lasr->content->data;
			b.last = b.pos + lasr->content->len;
			b.start = b.pos;
			b.end = b.last;			
			
			out.buf = &b;
			out.next = NULL;

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

		if(lasr->status == NGX_HTTP_INTERNAL_SERVER_ERROR){
			ngx_http_finalize_request(r, NGX_HTTP_INTERNAL_SERVER_ERROR);
			return NGX_OK;
		}
		else if(lasr->status != NGX_HTTP_OK){
			r->headers_out.status = lasr->status;
			r->header_only = 1;//???

			rc = ngx_http_send_header(r);
			ngx_http_finalize_request(r, rc);
			return NGX_OK;
		}

//		args = lasr->content;

		args = (ngx_str_t *)ngx_palloc(r->pool, sizeof(ngx_str_t));	
		if(args == NULL){
			ngx_http_finalize_request(r, NGX_HTTP_INTERNAL_SERVER_ERROR);
			return NGX_OK;
		}
		
		args->len = r->args.len + 1 + lasr->content->len;
		args->data = (u_char *)ngx_palloc(r->pool, args->len);
		if(args->data == NULL){
			ngx_http_finalize_request(r, NGX_HTTP_INTERNAL_SERVER_ERROR);
			return NGX_OK;
		}
		
		ngx_memcpy(args->data, r->args.data, r->args.len);
		args->data[r->args.len] = '&';
		ngx_memcpy(args->data + r->args.len + 1, lasr->content->data, lasr->content->len);
#if 0
		/*for test 
		*uri dont contain '\r', ...
		*/
		u_char *p = lasr->content->data + lasr->content->len;
		while(*--p == '\n'){
			args->len --;			
		}

		while(*--p == '\n'){
			ngx_http_finalize_request(r, NGX_HTTP_INTERNAL_SERVER_ERROR);
			return NGX_OK;
		}
#endif
		goto subrequest;
    }
	else{
		args = &r->args;
	}

    ctx = (ngx_http_lauth_ext_ctx_t *)ngx_pcalloc(r->pool, sizeof(ngx_http_lauth_ext_ctx_t));
    if (ctx == NULL) {
        return NGX_ERROR;
    }
	
	ctx->subrequests = ngx_array_create(r->pool, LAUTH_EXT_CMD_ARGS_NUM, sizeof(ngx_http_lauth_ext_subrequest));
	if (ctx->subrequests == NULL) {
        return NGX_ERROR;
    }

	ngx_http_set_ctx(r, ctx, ngx_http_lauth_ext_module);

subrequest:
	uri = (ngx_str_t *)lacf->uris->elts + ctx->i;
	ctx->i ++;
	
	lasr = (ngx_http_lauth_ext_subrequest *)ngx_array_push(ctx->subrequests);
	ngx_memzero(lasr, sizeof(ngx_http_lauth_ext_subrequest));

    ps = (ngx_http_post_subrequest_t *)ngx_palloc(r->pool, sizeof(ngx_http_post_subrequest_t));
    if (ps == NULL) {
        return NGX_ERROR;
    }

    ps->handler = ngx_http_lauth_ext_done;
    ps->data = ctx;	 

    if(ngx_http_subrequest(r, uri, args, &lasr->sr, ps, NGX_HTTP_SUBREQUEST_IN_MEMORY/* | NGX_HTTP_SUBREQUEST_WAITED*/) != NGX_OK){
		return NGX_ERROR;
	}

	ngx_log_error(NGX_LOG_DEBUG, r->connection->log, 0,
                   "lauth_ext request handler, subrequest uri:%V args:%V", uri, args); 
	
	return NGX_AGAIN;
}


static ngx_int_t ngx_http_lauth_ext_done(ngx_http_request_t *r,
    void *data, ngx_int_t rc)
{
	ngx_http_lauth_ext_ctx_t *ctx = (ngx_http_lauth_ext_ctx_t *)data;  

	ngx_http_lauth_ext_subrequest* lasr = (ngx_http_lauth_ext_subrequest*)ctx->subrequests->elts + (ctx->i - 1);

	lasr->done = 1;
	lasr->status = r->headers_out.status;

	if(lasr->status == NGX_HTTP_OK){	
		if(r->upstream == NULL){
			ngx_log_error(NGX_LOG_ERR, r->connection->log, 0,
                   "lauth_ext subrequest done, subrequest does not support upstream, uri:%V args:%V", 
                   &r->uri, &r->args);

			lasr->status = NGX_HTTP_INTERNAL_SERVER_ERROR;
			goto end;
		}

		if(lasr->content == NULL){
			lasr->content = (ngx_str_t *)ngx_palloc(r->pool, sizeof(ngx_str_t));
			if(lasr->content == NULL){
				lasr->status = NGX_HTTP_INTERNAL_SERVER_ERROR;
				goto end;
			}
		}
		
		lasr->content->len = r->headers_out.content_length_n;
		lasr->content->data = (u_char *)ngx_palloc(r->pool, lasr->content->len);
		if(lasr->content->data == NULL){
			lasr->status = NGX_HTTP_INTERNAL_SERVER_ERROR;
			goto end;
		}
		
		ngx_memcpy(lasr->content->data, r->upstream->buffer.pos, lasr->content->len);

		goto end;
	}

end:
	if(lasr->content != NULL){
		ngx_log_error(NGX_LOG_DEBUG, r->connection->log, 0,
                   "lauth_ext subrequest done, uri:%V args:%V rc:%i status:%i content:%V", 
                   &r->uri, &r->args, rc, lasr->status, lasr->content);
	}
	else{
		ngx_log_error(NGX_LOG_DEBUG, r->connection->log, 0,
                   "lauth_ext subrequest done, uri:%V args:%V rc:%i status:%i", 
                   &r->uri, &r->args, rc, lasr->status);
	}
	
	r->parent->write_event_handler = (ngx_http_event_handler_pt)ngx_http_lauth_ext_handler;
	
    return rc;
}


static void *
ngx_http_lauth_ext_create_conf(ngx_conf_t *cf)
{
    ngx_http_lauth_ext_conf_t  *conf;

    conf = (ngx_http_lauth_ext_conf_t *)ngx_pcalloc(cf->pool, sizeof(ngx_http_lauth_ext_conf_t));
    if (conf == NULL) {
        return NULL;
    }

    conf->n = LAUTH_EXT_CMD_ARGS_NUM;
	conf->uris = ngx_array_create(cf->pool, LAUTH_EXT_CMD_ARGS_NUM, sizeof(ngx_str_t));
	if(conf->uris == NULL){
		return NGX_CONF_ERROR;
	}

    return conf;
}


static char *
ngx_http_lauth_ext(ngx_conf_t *cf, ngx_command_t *cmd, void *conf)
{
    ngx_http_lauth_ext_conf_t *lacf = (ngx_http_lauth_ext_conf_t *)conf;
	ngx_http_core_loc_conf_t *clcf;

    ngx_str_t     *uri, *value;
    value = (ngx_str_t *)cf->args->elts;	

	for (ngx_uint_t i = 1; i <= lacf->n; i++) {
		uri = (ngx_str_t *)ngx_array_push(lacf->uris);
		*uri = value[i];
		
		ngx_log_error(NGX_LOG_DEBUG, cf->log, 0,
					"lauth_ext args: %V", uri);
	}

	clcf = (ngx_http_core_loc_conf_t *)ngx_http_conf_get_module_loc_conf(cf, ngx_http_core_module);
	clcf->handler = ngx_http_lauth_ext_handler;

    return NGX_CONF_OK;
}