tb_bool_t tb_async_transfer_init_ostream_from_data(tb_async_transfer_ref_t transfer, tb_byte_t* data, tb_size_t size)
{
// check
tb_async_transfer_impl_t* impl = (tb_async_transfer_impl_t*)transfer;
tb_assert_and_check_return_val(impl && impl->aicp && data && size, tb_false);
// muse be closed
tb_assert_and_check_return_val(TB_STATE_CLOSED == tb_atomic_get(&impl->state), tb_false);
// exit the previous stream first if be not data stream
if (impl->ostream && tb_async_stream_type(impl->ostream) != TB_STREAM_TYPE_DATA)
{
if (impl->oowner) tb_async_stream_exit(impl->ostream);
impl->ostream = tb_null;
}
// using the previous stream?
if (impl->ostream)
{
// ctrl stream
if (!tb_async_stream_ctrl(impl->ostream, TB_STREAM_CTRL_DATA_SET_DATA, data, size)) return tb_false;
}
else
{
// init stream
impl->ostream = tb_async_stream_init_from_data(impl->aicp, data, size);
tb_assert_and_check_return_val(impl->ostream, tb_false);
// init owner
impl->oowner = 1;
}
// ok
return tb_true;
}
static tb_bool_t tb_demo_spider_task_ctrl(tb_async_stream_ref_t istream, tb_async_stream_ref_t ostream, tb_cpointer_t priv)
{
// check
tb_demo_spider_task_t* task = (tb_demo_spider_task_t*)priv;
tb_assert_and_check_return_val(task && task->spider, tb_false);
tb_assert_and_check_return_val(istream && ostream, tb_false);
tb_assert_and_check_return_val(tb_async_stream_type(istream) == TB_STREAM_TYPE_HTTP, tb_false);
// the url
tb_char_t const* url = tb_null;
if (!tb_async_stream_ctrl(istream, TB_STREAM_CTRL_GET_URL, &url)) return tb_false;
// trace
tb_trace_d("ctrl: %s: ..", url);
// set timeout
if (!tb_async_stream_ctrl(istream, TB_STREAM_CTRL_SET_TIMEOUT, task->spider->timeout)) return tb_false;
#if defined(TB_CONFIG_MODULE_HAVE_ZIP) && defined(TB_CONFIG_THIRD_HAVE_ZLIB)
// need gzip
if (!tb_async_stream_ctrl(istream, TB_STREAM_CTRL_HTTP_SET_HEAD, "Accept-Encoding", "gzip,deflate")) return tb_false;
// auto unzip
if (!tb_async_stream_ctrl(istream, TB_STREAM_CTRL_HTTP_SET_AUTO_UNZIP, 1)) return tb_false;
#endif
// user agent
if (!tb_async_stream_ctrl(istream, TB_STREAM_CTRL_HTTP_SET_HEAD, "User-Agent", task->spider->user_agent)) return tb_false;
// enable cookies
if (!tb_async_stream_ctrl(istream, TB_STREAM_CTRL_HTTP_SET_COOKIES, tb_cookies())) return tb_false;
// ok
return tb_true;
}
tb_bool_t tb_async_transfer_init_istream_from_url(tb_async_transfer_ref_t transfer, tb_char_t const* url)
{
// check
tb_async_transfer_impl_t* impl = (tb_async_transfer_impl_t*)transfer;
tb_assert_and_check_return_val(impl && impl->aicp && url, tb_false);
// muse be closed
tb_assert_and_check_return_val(TB_STATE_CLOSED == tb_atomic_get(&impl->state), tb_false);
// check stream type
if (impl->istream)
{
// probe protocol
tb_size_t protocol = tb_url_protocol_probe(url);
tb_assert_static((tb_size_t)TB_URL_PROTOCOL_FILE == (tb_size_t)TB_STREAM_TYPE_FILE);
tb_assert_static((tb_size_t)TB_URL_PROTOCOL_HTTP == (tb_size_t)TB_STREAM_TYPE_HTTP);
tb_assert_static((tb_size_t)TB_URL_PROTOCOL_SOCK == (tb_size_t)TB_STREAM_TYPE_SOCK);
tb_assert_static((tb_size_t)TB_URL_PROTOCOL_DATA == (tb_size_t)TB_STREAM_TYPE_DATA);
// protocol => type
tb_size_t type = protocol;
if (!type || type > TB_STREAM_TYPE_DATA)
{
tb_trace_e("unknown stream for url: %s", url);
return tb_false;
}
// exit the previous stream first if be different stream type
if (tb_async_stream_type(impl->istream) != type)
{
if (impl->iowner) tb_async_stream_exit(impl->istream);
impl->istream = tb_null;
}
}
// using the previous stream?
if (impl->istream)
{
// ctrl stream
if (!tb_async_stream_ctrl(impl->istream, TB_STREAM_CTRL_SET_URL, url)) return tb_false;
}
else
{
// init stream
impl->istream = tb_async_stream_init_from_url(impl->aicp, url);
tb_assert_and_check_return_val(impl->istream, tb_false);
// init owner
impl->iowner = 1;
}
// ok
return tb_true;
}
/* //////////////////////////////////////////////////////////////////////////////////////
* main
*/
tb_int_t tb_demo_stream_async_stream_main(tb_int_t argc, tb_char_t** argv)
{
// done
tb_demo_context_t context = {0};
do
{
// init option
context.option = tb_option_init("astream", "the astream demo", g_options);
tb_assert_and_check_break(context.option);
// done option
if (tb_option_done(context.option, argc - 1, &argv[1]))
{
// debug and verbose
context.debug = tb_option_find(context.option, "debug")? tb_false : tb_true;
context.verbose = tb_option_find(context.option, "no-verbose")? tb_false : tb_true;
// done url
if (tb_option_find(context.option, "url"))
{
// init event
context.event = tb_event_init();
tb_assert_and_check_break(context.event);
// init istream
context.istream = tb_async_stream_init_from_url(tb_null, tb_option_item_cstr(context.option, "url"));
tb_assert_and_check_break(context.istream);
// ctrl http
if (tb_async_stream_type(context.istream) == TB_STREAM_TYPE_HTTP)
{
// enable gzip?
if (tb_option_find(context.option, "gzip"))
{
// auto unzip
if (!tb_async_stream_ctrl(context.istream, TB_STREAM_CTRL_HTTP_SET_AUTO_UNZIP, 1)) break;
// need gzip
if (!tb_async_stream_ctrl(context.istream, TB_STREAM_CTRL_HTTP_SET_HEAD, "Accept-Encoding", "gzip,deflate")) break;
}
// enable debug?
if (!tb_async_stream_ctrl(context.istream, TB_STREAM_CTRL_HTTP_SET_HEAD_FUNC, context.debug? tb_demo_istream_head_func : tb_null)) break;
// custem header?
if (tb_option_find(context.option, "header"))
{
// init
tb_string_t key;
tb_string_t val;
tb_string_init(&key);
tb_string_init(&val);
// done
tb_bool_t k = tb_true;
tb_char_t const* p = tb_option_item_cstr(context.option, "header");
while (*p)
{
// is key?
if (k)
{
if (*p != ':' && !tb_isspace(*p)) tb_string_chrcat(&key, *p++);
else if (*p == ':')
{
// skip ':'
p++;
// skip space
while (*p && tb_isspace(*p)) p++;
// is val now
k = tb_false;
}
else p++;
}
// is val?
else
{
if (*p != ';') tb_string_chrcat(&val, *p++);
else
{
// skip ';'
p++;
// skip space
while (*p && tb_isspace(*p)) p++;
// set header
if (tb_string_size(&key) && tb_string_size(&val))
{
if (context.debug) tb_printf("header: %s: %s\n", tb_string_cstr(&key), tb_string_cstr(&val));
if (!tb_async_stream_ctrl(context.istream, TB_STREAM_CTRL_HTTP_SET_HEAD, tb_string_cstr(&key), tb_string_cstr(&val))) break;
}
// is key now
k = tb_true;
// clear key & val
tb_string_clear(&key);
tb_string_clear(&val);
}
}
}
// set header
if (tb_string_size(&key) && tb_string_size(&val))
{
if (context.debug) tb_printf("header: %s: %s\n", tb_string_cstr(&key), tb_string_cstr(&val));
if (!tb_async_stream_ctrl(context.istream, TB_STREAM_CTRL_HTTP_SET_HEAD, tb_string_cstr(&key), tb_string_cstr(&val))) break;
}
// exit
tb_string_exit(&key);
tb_string_exit(&val);
}
// keep alive?
if (tb_option_find(context.option, "keep-alive"))
{
if (!tb_async_stream_ctrl(context.istream, TB_STREAM_CTRL_HTTP_SET_HEAD, "Connection", "keep-alive")) break;
}
// post-data?
if (tb_option_find(context.option, "post-data"))
{
tb_char_t const* post_data = tb_option_item_cstr(context.option, "post-data");
tb_hize_t post_size = tb_strlen(post_data);
if (!tb_async_stream_ctrl(context.istream, TB_STREAM_CTRL_HTTP_SET_METHOD, TB_HTTP_METHOD_POST)) break;
if (!tb_async_stream_ctrl(context.istream, TB_STREAM_CTRL_HTTP_SET_POST_DATA, post_data, post_size)) break;
if (!tb_async_stream_ctrl(context.istream, TB_STREAM_CTRL_HTTP_SET_POST_FUNC, tb_demo_http_post_func)) break;
if (context.debug) tb_printf("post: %llu\n", post_size);
}
// post-file?
else if (tb_option_find(context.option, "post-file"))
{
tb_char_t const* url = tb_option_item_cstr(context.option, "post-file");
if (!tb_async_stream_ctrl(context.istream, TB_STREAM_CTRL_HTTP_SET_METHOD, TB_HTTP_METHOD_POST)) break;
if (!tb_async_stream_ctrl(context.istream, TB_STREAM_CTRL_HTTP_SET_POST_URL, url)) break;
if (!tb_async_stream_ctrl(context.istream, TB_STREAM_CTRL_HTTP_SET_POST_FUNC, tb_demo_http_post_func)) break;
if (context.debug) tb_printf("post: %s\n", url);
}
}
// set range
if (tb_option_find(context.option, "range"))
{
tb_char_t const* p = tb_option_item_cstr(context.option, "range");
if (p)
{
// the bof
tb_hize_t eof = 0;
tb_hize_t bof = tb_atoll(p);
while (*p && tb_isdigit(*p)) p++;
if (*p == '-')
{
p++;
eof = tb_atoll(p);
}
if (!tb_async_stream_ctrl(context.istream, TB_STREAM_CTRL_HTTP_SET_RANGE, bof, eof)) break;
}
}
// set timeout
if (tb_option_find(context.option, "timeout"))
{
tb_size_t timeout = tb_option_item_uint32(context.option, "timeout");
tb_async_stream_ctrl(context.istream, TB_STREAM_CTRL_SET_TIMEOUT, &timeout);
}
// print verbose info
if (context.verbose) tb_printf("open: %s: ..\n", tb_option_item_cstr(context.option, "url"));
// open istream
if (!tb_async_stream_open(context.istream, tb_demo_istream_open_func, &context))
{
// print verbose info
if (context.verbose) tb_printf("open: failed\n");
break;
}
// wait it
tb_event_wait(context.event, -1);
}
else tb_option_help(context.option);
}
else tb_option_help(context.option);
} while (0);
// exit transfer
if (context.transfer) tb_async_transfer_exit(context.transfer);
context.transfer = tb_null;
// exit istream
if (context.istream) tb_async_stream_exit(context.istream);
context.istream = tb_null;
// exit ostream
if (context.ostream) tb_async_stream_exit(context.ostream);
context.ostream = tb_null;
// exit option
if (context.option) tb_option_exit(context.option);
context.option = tb_null;
// exit event
if (context.event) tb_event_exit(context.event);
context.event = tb_null;
return 0;
}
