Skip to content

Commit af0bb7a

Browse files
azkrishpyazkrishpy
committed
add option to elasticurl
1 parent b08bcaf commit af0bb7a

1 file changed

Lines changed: 159 additions & 2 deletions

File tree

‎bin/elasticurl/main.c‎

Lines changed: 159 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -64,6 +64,11 @@ struct elasticurl_ctx {
6464
enum aws_log_level log_level;
6565
enum aws_http_version required_http_version;
6666
bool exchange_completed;
67+
bool manual_write;
68+
bool manual_write_chunked;
69+
int64_t manual_write_content_length;
70+
struct aws_http_stream *stream;
71+
bool stream_ready;
6772
};
6873

6974
static void s_usage(int exit_code) {
@@ -96,6 +101,7 @@ static void s_usage(int exit_code) {
96101
fprintf(stderr, " --version: print the version of elasticurl.\n");
97102
fprintf(stderr, " --http2: HTTP/2 connection required\n");
98103
fprintf(stderr, " --http1_1: HTTP/1.1 connection required\n");
104+
fprintf(stderr, " --manual-write: interactively write request body via stdin\n");
99105
fprintf(stderr, " -h, --help\n");
100106
fprintf(stderr, " Display this message and quit.\n");
101107
exit(exit_code);
@@ -125,6 +131,7 @@ static struct aws_cli_option s_long_options[] = {
125131
{"version", AWS_CLI_OPTIONS_NO_ARGUMENT, NULL, 'V'},
126132
{"http2", AWS_CLI_OPTIONS_NO_ARGUMENT, NULL, 'w'},
127133
{"http1_1", AWS_CLI_OPTIONS_NO_ARGUMENT, NULL, 'W'},
134+
{"manual-write", AWS_CLI_OPTIONS_NO_ARGUMENT, NULL, 'n'},
128135
{"help", AWS_CLI_OPTIONS_NO_ARGUMENT, NULL, 'h'},
129136
/* Per getopt(3) the last element of the array has to be filled with all zeros */
130137
{NULL, AWS_CLI_OPTIONS_NO_ARGUMENT, NULL, 0},
@@ -162,7 +169,7 @@ static void s_parse_options(int argc, char **argv, struct elasticurl_ctx *ctx) {
162169
while (true) {
163170
int option_index = 0;
164171
int c =
165-
aws_cli_getopt_long(argc, argv, "a:b:c:e:f:H:d:g:j:l:m:M:GPHiko:t:v:VwWh", s_long_options, &option_index);
172+
aws_cli_getopt_long(argc, argv, "a:b:c:e:f:H:d:g:j:l:m:M:GPHiko:t:v:VwWnh", s_long_options, &option_index);
166173
if (c == -1) {
167174
break;
168175
}
@@ -276,6 +283,9 @@ static void s_parse_options(int argc, char **argv, struct elasticurl_ctx *ctx) {
276283
ctx->alpn = "http/1.1";
277284
ctx->required_http_version = AWS_HTTP_VERSION_1_1;
278285
break;
286+
case 'n':
287+
ctx->manual_write = true;
288+
break;
279289
case 'h':
280290
s_usage(0);
281291
break;
@@ -432,7 +442,26 @@ static struct aws_http_message *s_build_http_request(
432442
};
433443
aws_http_message_add_header(request, user_agent_header);
434444

435-
if (app_ctx->input_body) {
445+
if (app_ctx->manual_write) {
446+
/* Manual write mode: set headers but no body stream.
447+
* H2 doesn't use Transfer-Encoding — just send DATA frames. */
448+
if (app_ctx->manual_write_chunked && protocol_version != AWS_HTTP_VERSION_2) {
449+
struct aws_http_header te_header = {
450+
.name = aws_byte_cursor_from_c_str("transfer-encoding"),
451+
.value = aws_byte_cursor_from_c_str("chunked"),
452+
};
453+
aws_http_message_add_header(request, te_header);
454+
} else if (!app_ctx->manual_write_chunked) {
455+
char content_length[64];
456+
AWS_ZERO_ARRAY(content_length);
457+
snprintf(content_length, sizeof(content_length), "%" PRIi64, app_ctx->manual_write_content_length);
458+
struct aws_http_header cl_header = {
459+
.name = aws_byte_cursor_from_c_str("content-length"),
460+
.value = aws_byte_cursor_from_c_str(content_length),
461+
};
462+
aws_http_message_add_header(request, cl_header);
463+
}
464+
} else if (app_ctx->input_body) {
436465
int64_t data_len = 0;
437466
if (aws_input_stream_get_length(app_ctx->input_body, &data_len)) {
438467
fprintf(stderr, "failed to get length of input stream.\n");
@@ -522,6 +551,7 @@ static void s_on_signing_complete(struct aws_http_message *request, int error_co
522551
.on_response_header_block_done = s_on_incoming_header_block_done_fn,
523552
.on_response_body = s_on_incoming_body_fn,
524553
.on_complete = s_on_stream_complete_fn,
554+
.use_manual_data_writes = app_ctx->manual_write,
525555
};
526556

527557
app_ctx->response_code_written = false;
@@ -533,6 +563,15 @@ static void s_on_signing_complete(struct aws_http_message *request, int error_co
533563
}
534564
aws_http_stream_activate(stream);
535565

566+
if (app_ctx->manual_write) {
567+
/* Store stream and signal main thread to begin interactive writes */
568+
app_ctx->stream = stream;
569+
aws_mutex_lock(&app_ctx->mutex);
570+
app_ctx->stream_ready = true;
571+
aws_mutex_unlock(&app_ctx->mutex);
572+
aws_condition_variable_notify_all(&app_ctx->c_var);
573+
}
574+
536575
/* Connection will stay alive until stream completes */
537576
aws_http_connection_release(app_ctx->connection);
538577
app_ctx->connection = NULL;
@@ -554,6 +593,103 @@ static bool s_completion_predicate(void *arg) {
554593
return app_ctx->exchange_completed;
555594
}
556595

596+
static bool s_stream_ready_predicate(void *arg) {
597+
struct elasticurl_ctx *app_ctx = arg;
598+
return app_ctx->stream_ready || app_ctx->exchange_completed;
599+
}
600+
601+
struct manual_write_ctx {
602+
struct aws_allocator *allocator;
603+
uint8_t *data;
604+
struct aws_input_stream *stream;
605+
};
606+
607+
static void s_on_manual_write_complete(struct aws_http_stream *stream, int error_code, void *user_data) {
608+
(void)stream;
609+
(void)error_code;
610+
struct manual_write_ctx *ctx = user_data;
611+
aws_input_stream_release(ctx->stream);
612+
aws_mem_release(ctx->allocator, ctx->data);
613+
aws_mem_release(ctx->allocator, ctx);
614+
}
615+
616+
static void s_manual_write_loop(struct elasticurl_ctx *app_ctx) {
617+
/* Wait for stream to be activated */
618+
aws_mutex_lock(&app_ctx->mutex);
619+
aws_condition_variable_wait_pred(&app_ctx->c_var, &app_ctx->mutex, s_stream_ready_predicate, app_ctx);
620+
aws_mutex_unlock(&app_ctx->mutex);
621+
622+
if (app_ctx->exchange_completed) {
623+
return;
624+
}
625+
626+
int64_t bytes_sent = 0;
627+
char line_buf[4096];
628+
629+
fprintf(stderr, "Enter data (empty line to finish):\n");
630+
while (fgets(line_buf, sizeof(line_buf), stdin)) {
631+
/* Strip trailing newline */
632+
size_t len = strlen(line_buf);
633+
if (len > 0 && line_buf[len - 1] == '\n') {
634+
line_buf[--len] = '\0';
635+
}
636+
637+
/* Empty line = done */
638+
if (len == 0) {
639+
break;
640+
}
641+
642+
/* Heap-allocate data so it outlives this stack frame */
643+
uint8_t *heap_data = aws_mem_calloc(app_ctx->allocator, 1, len);
644+
memcpy(heap_data, line_buf, len);
645+
646+
struct aws_byte_cursor data_cursor = aws_byte_cursor_from_array(heap_data, len);
647+
struct aws_input_stream *data_stream =
648+
aws_input_stream_new_from_cursor(app_ctx->allocator, &data_cursor);
649+
650+
struct manual_write_ctx *write_ctx = aws_mem_calloc(app_ctx->allocator, 1, sizeof(struct manual_write_ctx));
651+
write_ctx->allocator = app_ctx->allocator;
652+
write_ctx->data = heap_data;
653+
write_ctx->stream = data_stream;
654+
655+
struct aws_http_stream_write_data_options write_opts = {
656+
.data = data_stream,
657+
.end_stream = false,
658+
.on_complete = s_on_manual_write_complete,
659+
.user_data = write_ctx,
660+
};
661+
662+
if (aws_http_stream_write_data(app_ctx->stream, &write_opts)) {
663+
fprintf(stderr, "write_data failed: %s\n", aws_error_debug_str(aws_last_error()));
664+
aws_input_stream_release(data_stream);
665+
aws_mem_release(app_ctx->allocator, heap_data);
666+
aws_mem_release(app_ctx->allocator, write_ctx);
667+
break;
668+
}
669+
670+
bytes_sent += (int64_t)len;
671+
fprintf(stderr, "Sent %zu bytes (total: %" PRIi64 ")\n", len, bytes_sent);
672+
}
673+
674+
/* Send final write */
675+
struct aws_byte_cursor empty_cursor = aws_byte_cursor_from_c_str("");
676+
struct aws_input_stream *empty_stream =
677+
aws_input_stream_new_from_cursor(app_ctx->allocator, &empty_cursor);
678+
679+
struct aws_http_stream_write_data_options final_opts = {
680+
.data = empty_stream,
681+
.end_stream = true,
682+
};
683+
684+
if (aws_http_stream_write_data(app_ctx->stream, &final_opts)) {
685+
fprintf(stderr, "final write_data failed: %s\n", aws_error_debug_str(aws_last_error()));
686+
} else {
687+
fprintf(stderr, "Stream complete. Sent %" PRIi64 " bytes.\n", bytes_sent);
688+
}
689+
690+
aws_input_stream_release(empty_stream);
691+
}
692+
557693
int main(int argc, char **argv) {
558694
struct aws_allocator *allocator = aws_default_allocator();
559695

@@ -579,6 +715,24 @@ int main(int argc, char **argv) {
579715

580716
s_parse_options(argc, argv, &app_ctx);
581717

718+
/* Interactive prompt for manual-write mode */
719+
if (app_ctx.manual_write) {
720+
if (!strcmp(app_ctx.verb, "GET")) {
721+
app_ctx.verb = "POST";
722+
}
723+
fprintf(stderr, "Manual write mode enabled.\n");
724+
fprintf(stderr, "Content-Length (leave empty for chunked transfer encoding): ");
725+
char cl_buf[64];
726+
if (fgets(cl_buf, sizeof(cl_buf), stdin) && cl_buf[0] != '\n') {
727+
app_ctx.manual_write_content_length = (int64_t)atoll(cl_buf);
728+
app_ctx.manual_write_chunked = false;
729+
fprintf(stderr, "Using Content-Length: %" PRIi64 "\n", app_ctx.manual_write_content_length);
730+
} else {
731+
app_ctx.manual_write_chunked = true;
732+
fprintf(stderr, "Using chunked transfer encoding.\n");
733+
}
734+
}
735+
582736
struct aws_logger logger;
583737
AWS_ZERO_STRUCT(logger);
584738

@@ -728,6 +882,9 @@ int main(int argc, char **argv) {
728882
http_client_options.prior_knowledge_http2 = true;
729883
}
730884
aws_http_client_connect(&http_client_options);
885+
if (app_ctx.manual_write) {
886+
s_manual_write_loop(&app_ctx);
887+
}
731888
aws_mutex_lock(&app_ctx.mutex);
732889
aws_condition_variable_wait_pred(&app_ctx.c_var, &app_ctx.mutex, s_completion_predicate, &app_ctx);
733890
aws_mutex_unlock(&app_ctx.mutex);

0 commit comments

Comments
 (0)