fixes #2133 websocket: new header iteration options

4 files changed, 130 insertions, 17 deletions

diff --git a/src/sp/transport/ws/ws_test.c b/src/sp/transport/ws/ws_test.c
index 98856d4f..0ebc91ac 100644
--- a/src/sp/transport/ws/ws_test.c
+++ b/src/sp/transport/ws/ws_test.c
@@ -1,5 +1,5 @@
 //
-// Copyright 2024 Staysail Systems, Inc. <info@staysail.tech>
+// Copyright 2025 Staysail Systems, Inc. <info@staysail.tech>
 // Copyright 2018 Cody Piersall <cody.piersall@gmail.com>
 //
 // This software is supplied under the terms of the MIT License, a
@@ -8,8 +8,8 @@
 // found online at https://opensource.org/licenses/MIT.
 //
 
+#include "../../../testing/nuts.h"
 #include "nng/nng.h"
-#include <nuts.h>
 
 static void
 test_ws_url_path_filters(void)
@@ -231,6 +231,46 @@ check_props_v4(nng_msg *msg)
 
 	NUTS_PASS(nng_pipe_get_bool(p, NNG_OPT_TCP_NODELAY, &b));
 	NUTS_TRUE(b); // default
+
+	const char *uri;
+	NUTS_PASS(nng_pipe_get_string(p, NNG_OPT_WS_REQUEST_URI, &uri));
+
+	NUTS_PASS(nng_pipe_get_bool(p, NNG_OPT_WS_HEADER_RESET, &b));
+	NUTS_ASSERT(b == true);
+	for (;;) {
+		const char *k;
+		const char *v;
+		char       *s;
+		char       *full;
+		char        buf[256];
+		size_t      sz;
+		NUTS_PASS(nng_pipe_get_bool(p, NNG_OPT_WS_HEADER_NEXT, &b));
+		if (!b) {
+			break;
+		}
+		// NB: Normally its unsafe for most callers to use this,
+		// because the pipe may be ripped out from underneath us.  But
+		// we are careful in this test to ensure that the pipe is kept
+		// alive.
+		NUTS_PASS(nng_pipe_get_string(p, NNG_OPT_WS_HEADER_KEY, &k));
+		NUTS_PASS(nng_pipe_get_strdup(p, NNG_OPT_WS_HEADER_KEY, &s));
+		NUTS_MATCH(s, k);
+		sz = sizeof(buf);
+		NUTS_PASS(
+		    nng_pipe_get_strcpy(p, NNG_OPT_WS_HEADER_KEY, buf, sz));
+		NUTS_MATCH(s, buf);
+		NUTS_PASS(nng_pipe_get_strlen(p, NNG_OPT_WS_HEADER_KEY, &sz));
+		NUTS_ASSERT(sz == strlen(s));
+		nng_strfree(s);
+
+		// again, not necessarily safe, but good enough
+		NUTS_PASS(nng_pipe_get_string(p, NNG_OPT_WS_HEADER_VALUE, &v));
+
+		snprintf(buf, sizeof(buf), "%s%s", NNG_OPT_WS_HEADER, k);
+		NUTS_PASS(nng_pipe_get_strdup(p, buf, &full));
+		NUTS_MATCH(v, full);
+		nng_strfree(full);
+	}
 }
 
 void
diff --git a/src/supplemental/http/http_api.h b/src/supplemental/http/http_api.h
index b1a8ec84..6acc1ba5 100644
--- a/src/supplemental/http/http_api.h
+++ b/src/supplemental/http/http_api.h
@@ -126,6 +126,8 @@ extern nng_err     nni_http_add_header(nng_http *, const char *, const char *);
 extern nng_err     nni_http_set_header(nng_http *, const char *, const char *);
 extern void        nni_http_del_header(nng_http *, const char *);
 extern const char *nni_http_get_header(nng_http *, const char *);
+extern bool        nni_http_next_header(
+           nng_http *, const char **, const char **, void **);
 
 extern void    nni_http_get_body(nng_http *, void **, size_t *);
 extern void    nni_http_set_body(nng_http *, void *, size_t);
diff --git a/src/supplemental/http/http_public.c b/src/supplemental/http/http_public.c
index 419632fe..e63c5f96 100644
--- a/src/supplemental/http/http_public.c
+++ b/src/supplemental/http/http_public.c
@@ -27,11 +27,11 @@ nng_http_get_header(nng_http *conn, const char *key)
 #endif
 }
 
-nng_err
+bool
 nng_http_next_header(
     nng_http *conn, const char **key, const char **val, void **ptr)
 {
-#ifdef NNG_SUP_HTTP
+#ifdef NNG_SUPP_HTTP
 	return (nni_http_next_header(conn, key, val, ptr));
 #else
 	NNI_ARG_UNUSED(conn);
diff --git a/src/supplemental/websocket/websocket.c b/src/supplemental/websocket/websocket.c
index 15078704..d78c032b 100644
--- a/src/supplemental/websocket/websocket.c
+++ b/src/supplemental/websocket/websocket.c
@@ -18,6 +18,7 @@
 #include "nng/http.h"
 
 #include "base64.h"
+#include "nng/nng.h"
 #include "sha1.h"
 #include "websocket.h"
 
@@ -85,6 +86,11 @@ struct nni_ws {
 		nni_http_header wsproto;
 		nni_http_header wsversion;
 	} hdrs;
+
+	// these fields are for header iteration
+	const char *curhdrkey;
+	const char *curhdrval;
+	void       *nexthdr;
 };
 
 struct nni_ws_listener {
@@ -2702,6 +2708,60 @@ ws_get_send_text(void *arg, void *buf, size_t *szp, nni_type t)
 	return (nni_copyout_bool(b, buf, szp, t));
 }
 
+static nng_err
+ws_get_next_header(void *arg, void *buf, size_t *szp, nni_type t)
+{
+	nni_ws *ws = arg;
+	bool    b;
+
+	if (t != NNI_TYPE_BOOL) {
+		return (NNG_EBADTYPE);
+	}
+	nni_mtx_lock(&ws->mtx);
+	b = nng_http_next_header(
+	    ws->http, &ws->curhdrkey, &ws->curhdrval, &ws->nexthdr);
+	nni_mtx_unlock(&ws->mtx);
+	return (nni_copyout_bool(b, buf, szp, t));
+}
+
+static nng_err
+ws_get_reset_header(void *arg, void *buf, size_t *szp, nni_type t)
+{
+	nni_ws *ws = arg;
+
+	if (t != NNI_TYPE_BOOL) {
+		return (NNG_EBADTYPE);
+	}
+	nni_mtx_lock(&ws->mtx);
+	ws->nexthdr = NULL;
+	nni_mtx_unlock(&ws->mtx);
+	return (nni_copyout_bool(true, buf, szp, t));
+}
+
+static nng_err
+ws_get_header_key(void *arg, void *buf, size_t *szp, nni_type t)
+{
+	nni_ws     *ws = arg;
+	const char *s;
+
+	nni_mtx_lock(&ws->mtx);
+	s = ws->curhdrkey;
+	nni_mtx_unlock(&ws->mtx);
+	return (nni_copyout_str(s, buf, szp, t));
+}
+
+static nng_err
+ws_get_header_val(void *arg, void *buf, size_t *szp, nni_type t)
+{
+	nni_ws     *ws = arg;
+	const char *s;
+
+	nni_mtx_lock(&ws->mtx);
+	s = ws->curhdrval;
+	nni_mtx_unlock(&ws->mtx);
+	return (nni_copyout_str(s, buf, szp, t));
+}
+
 static const nni_option ws_options[] = {
 	{
 	    .o_name = NNG_OPT_WS_REQUEST_URI,
@@ -2716,23 +2776,27 @@ static const nni_option ws_options[] = {
 	    .o_get  = ws_get_send_text,
 	},
 	{
+	    .o_name = NNG_OPT_WS_HEADER_NEXT,
+	    .o_get  = ws_get_next_header,
+	},
+	{
+	    .o_name = NNG_OPT_WS_HEADER_RESET,
+	    .o_get  = ws_get_reset_header,
+	},
+	{
+	    .o_name = NNG_OPT_WS_HEADER_KEY,
+	    .o_get  = ws_get_header_key,
+	},
+	{
+	    .o_name = NNG_OPT_WS_HEADER_VALUE,
+	    .o_get  = ws_get_header_val,
+	},
+	{
 	    .o_name = NULL,
 	},
 };
 
 static nng_err
-ws_get_header(nni_ws *ws, const char *nm, void *buf, size_t *szp, nni_type t)
-{
-	const char *s;
-	nm += strlen(NNG_OPT_WS_HEADER);
-	s = nni_http_get_header(ws->http, nm);
-	if (s == NULL) {
-		return (NNG_ENOENT);
-	}
-	return (nni_copyout_str(s, buf, szp, t));
-}
-
-static nng_err
 ws_str_get(void *arg, const char *nm, void *buf, size_t *szp, nni_type t)
 {
 	nni_ws *ws = arg;
@@ -2750,8 +2814,15 @@ ws_str_get(void *arg, const char *nm, void *buf, size_t *szp, nni_type t)
 	}
 	// Check for generic headers...
 	if (rv == NNG_ENOTSUP) {
+		const char *val;
+
 		if (startswith(nm, NNG_OPT_WS_HEADER)) {
-			rv = ws_get_header(ws, nm, buf, szp, t);
+			val = nng_http_get_header(
+			    ws->http, nm + strlen(NNG_OPT_WS_HEADER));
+			if (val == NULL) {
+				return (NNG_ENOENT);
+			}
+			return (nni_copyout_str(val, buf, szp, t));
 		}
 	}
 	return (rv);