3 files changed, 16 insertions, 5 deletions

diff --git a/src/c/include/unit/unit-wasm.h b/src/c/include/unit/unit-wasm.h
index dedf15a..cfff906 100644
--- a/src/c/include/unit/unit-wasm.h
+++ b/src/c/include/unit/unit-wasm.h
@@ -118,6 +118,9 @@ typedef struct {
 
 	/* points to the end of ctx->req_buf */
 	u8 *reqp;
+
+	/* tracks the response header index number */
+	s32 resp_hdr_idx;
 } luw_ctx_t;
 
 typedef enum {
@@ -189,7 +192,7 @@ extern size_t luw_mem_fill_buf_from_req(luw_ctx_t *ctx, size_t from);
 extern void luw_mem_reset(luw_ctx_t *ctx);
 extern void luw_http_send_response(const luw_ctx_t *ctx);
 extern void luw_http_init_headers(luw_ctx_t *ctx, size_t nr, size_t offset);
-extern void luw_http_add_header(luw_ctx_t *ctx, u16 idx, const char *name,
+extern void luw_http_add_header(luw_ctx_t *ctx, const char *name,
 				const char *value);
 extern void luw_http_send_headers(const luw_ctx_t *ctx);
 extern void luw_http_response_end(void);
diff --git a/src/c/libunit-wasm.c b/src/c/libunit-wasm.c
index 8cea78a..11272ed 100644
--- a/src/c/libunit-wasm.c
+++ b/src/c/libunit-wasm.c
@@ -83,6 +83,7 @@ void luw_init_ctx(luw_ctx_t *ctx, u8 *addr, size_t offset)
 	ctx->resp_offset = offset;
 	ctx->resp->size = 0;
 	ctx->resp_hdr->nr_fields = 0;
+	ctx->resp_hdr_idx = -1;
 }
 
 /*
@@ -326,6 +327,7 @@ void luw_mem_reset(luw_ctx_t *ctx)
 	ctx->mem = ctx->resp->data;
 	ctx->resp->size = 0;
 	ctx->resp_hdr->nr_fields = 0;
+	ctx->resp_hdr_idx = -1;
 }
 
 void luw_http_send_response(const luw_ctx_t *ctx)
@@ -342,9 +344,14 @@ void luw_http_init_headers(luw_ctx_t *ctx, size_t nr, size_t offset)
 	ctx->resp_hdr->nr_fields = nr;
 }
 
-void luw_http_add_header(luw_ctx_t *ctx, u16 idx, const char *name,
-			 const char *value)
+void luw_http_add_header(luw_ctx_t *ctx, const char *name, const char *value)
 {
+	s32 idx = ctx->resp_hdr_idx;
+
+	idx++;
+	if ((u32)idx == ctx->resp_hdr->nr_fields)
+		return;
+
 	ctx->resp_hdr->fields[idx].name_off = ctx->hdrp - ctx->addr;
 	ctx->resp_hdr->fields[idx].name_len = strlen(name);
 	ctx->hdrp = mempcpy(ctx->hdrp, name, strlen(name));
@@ -352,6 +359,8 @@ void luw_http_add_header(luw_ctx_t *ctx, u16 idx, const char *name,
 	ctx->resp_hdr->fields[idx].value_off = ctx->hdrp - ctx->addr;
 	ctx->resp_hdr->fields[idx].value_len = strlen(value);
 	ctx->hdrp = mempcpy(ctx->hdrp, value, strlen(value));
+
+	ctx->resp_hdr_idx = idx;
 }
 
 void luw_http_send_headers(const luw_ctx_t *ctx)
diff --git a/src/rust/unit-wasm-sys/rusty.rs b/src/rust/unit-wasm-sys/rusty.rs
index 79c38f3..fbfecd3 100644
--- a/src/rust/unit-wasm-sys/rusty.rs
+++ b/src/rust/unit-wasm-sys/rusty.rs
@@ -43,6 +43,7 @@ pub const fn UWR_CTX_INITIALIZER() -> luw_ctx_t {
         req_buf: null_mut(),
         hdrp: null_mut(),
         reqp: null_mut(),
+        resp_hdr_idx: -1,
     }
 }
 
@@ -173,14 +174,12 @@ pub fn uwr_http_init_headers(ctx: *mut luw_ctx_t, nr: usize, offset: usize) {
 
 pub fn uwr_http_add_header(
     ctx: *mut luw_ctx_t,
-    idx: u16,
     name: &str,
     value: &str,
 ) {
     unsafe {
         luw_http_add_header(
             ctx,
-            idx,
             S2C!(name).as_ptr() as *const i8,
             S2C!(value).as_ptr() as *const i8,
         );