fixes #2133 websocket: new header iteration options

1 files changed, 84 insertions, 13 deletions

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);