2 files changed, 78 insertions, 11 deletions

diff --git a/rest-api-templates/ari_resource.h.mustache b/rest-api-templates/ari_resource.h.mustache
index 3a20776a7..d3f40b6bd 100644
--- a/rest-api-templates/ari_resource.h.mustache
+++ b/rest-api-templates/ari_resource.h.mustache
@@ -89,6 +89,23 @@ int ast_ari_{{c_name}}_{{c_nickname}}_parse_body(
 void ast_ari_{{c_name}}_{{c_nickname}}(struct ast_variable *headers, struct ast_ari_{{c_name}}_{{c_nickname}}_args *args, struct ast_ari_response *response);
 {{/is_req}}
 {{#is_websocket}}
+
+/*!
+ * \brief {{summary}}
+{{#notes}}
+ *
+ * {{{notes}}}
+{{/notes}}
+ *
+ * \param ser HTTP TCP/TLS Server Session
+ * \param headers HTTP headers
+ * \param args Swagger parameters
+ *
+ * \retval 0 success
+ * \retval non-zero error
+ */
+int ast_ari_websocket_{{c_name}}_{{c_nickname}}_attempted(struct ast_tcptls_session_instance *ser, struct ast_variable *headers, struct ast_ari_{{c_name}}_{{c_nickname}}_args *args);
+
 /*!
  * \brief {{summary}}
 {{#notes}}
@@ -100,7 +117,7 @@ void ast_ari_{{c_name}}_{{c_nickname}}(struct ast_variable *headers, struct ast_
  * \param headers HTTP headers.
  * \param args Swagger parameters.
  */
-void ast_ari_websocket_{{c_name}}_{{c_nickname}}(struct ast_ari_websocket_session *session, struct ast_variable *headers, struct ast_ari_{{c_name}}_{{c_nickname}}_args *args);
+void ast_ari_websocket_{{c_name}}_{{c_nickname}}_established(struct ast_ari_websocket_session *session, struct ast_variable *headers, struct ast_ari_{{c_name}}_{{c_nickname}}_args *args);
 {{/is_websocket}}
 {{/operations}}
 {{/apis}}
diff --git a/rest-api-templates/res_ari_resource.c.mustache b/rest-api-templates/res_ari_resource.c.mustache
index 4cc9e0db0..08f6204be 100644
--- a/rest-api-templates/res_ari_resource.c.mustache
+++ b/rest-api-templates/res_ari_resource.c.mustache
@@ -137,7 +137,52 @@ fin: __attribute__((unused))
 }
 {{/is_req}}
 {{#is_websocket}}
-static void ast_ari_{{c_name}}_{{c_nickname}}_ws_cb(struct ast_websocket *ws_session,
+static int ast_ari_{{c_name}}_{{c_nickname}}_ws_attempted_cb(struct ast_tcptls_session_instance *ser, struct ast_variable *get_params, struct ast_variable *headers)
+{
+	struct ast_ari_{{c_name}}_{{c_nickname}}_args args = {};
+{{#has_parameters}}
+	int res = 0;
+	RAII_VAR(struct ast_ari_response *, response, NULL, ast_free);
+	struct ast_variable *i;
+{{/has_parameters}}
+
+{{#has_parameters}}
+	response = ast_calloc(1, sizeof(*response));
+	if (!response) {
+		ast_log(LOG_ERROR, "Failed to create response.\n");
+		goto fin;
+	}
+{{/has_parameters}}
+
+{{> param_parsing}}
+
+	res = ast_ari_websocket_{{c_name}}_{{c_nickname}}_attempted(ser, headers, &args);
+
+fin: __attribute__((unused))
+	if (!response) {
+		ast_http_error(ser, 500, "Server Error", "Memory allocation error");
+		res = -1;
+	} else if (response->response_code != 0) {
+		/* Param parsing failure */
+		RAII_VAR(char *, msg, NULL, ast_json_free);
+		if (response->message) {
+			msg = ast_json_dump_string(response->message);
+		} else {
+			ast_log(LOG_ERROR, "Missing response message\n");
+		}
+
+		if (msg) {
+			ast_http_error(ser, response->response_code, response->response_text, msg);
+		}
+		res = -1;
+	}
+{{> param_cleanup}}
+{{#has_parameters}}
+	return res;
+{{/has_parameters}}
+}
+
+static void ast_ari_{{c_name}}_{{c_nickname}}_ws_established_cb(struct ast_websocket *ws_session,
 	struct ast_variable *get_params, struct ast_variable *headers)
 {
 	struct ast_ari_{{c_name}}_{{c_nickname}}_args args = {};
@@ -175,16 +220,11 @@ static void ast_ari_{{c_name}}_{{c_nickname}}_ws_cb(struct ast_websocket *ws_ses
 
 {{> param_parsing}}
 
-	ast_ari_websocket_{{c_name}}_{{c_nickname}}(session, headers, &args);
+	ast_ari_websocket_{{c_name}}_{{c_nickname}}_established(session, headers, &args);
 
 fin: __attribute__((unused))
 	if (response && response->response_code != 0) {
 		/* Param parsing failure */
-		/* TODO - ideally, this would return the error code to the
-		 * HTTP client; but we've already done the WebSocket
-		 * negotiation. Param parsing should happen earlier, but we
-		 * need a way to pass it through the WebSocket code to the
-		 * callback */
 		RAII_VAR(char *, msg, NULL, ast_json_free);
 		if (response->message) {
 			msg = ast_json_dump_string(response->message);
@@ -211,16 +251,26 @@ static int load_module(void)
 {
 	int res = 0;
 {{#apis}}
+{{#operations}}
 {{#has_websocket}}
+	struct ast_websocket_protocol *protocol;
+
 	{{full_name}}.ws_server = ast_websocket_server_create();
 	if (!{{full_name}}.ws_server) {
 		return AST_MODULE_LOAD_FAILURE;
 	}
+
+	protocol = ast_websocket_sub_protocol_alloc("{{websocket_protocol}}");
+	if (!protocol) {
+		ao2_ref({{full_name}}.ws_server, -1);
+		{{full_name}}.ws_server = NULL;
+		return AST_MODULE_LOAD_FAILURE;
+	}
+	protocol->session_attempted = ast_ari_{{c_name}}_{{c_nickname}}_ws_attempted_cb;
+	protocol->session_established = ast_ari_{{c_name}}_{{c_nickname}}_ws_established_cb;
 {{/has_websocket}}
-{{#operations}}
 {{#is_websocket}}
-	res |= ast_websocket_server_add_protocol({{full_name}}.ws_server,
-		"{{websocket_protocol}}", ast_ari_{{c_name}}_{{c_nickname}}_ws_cb);
+	res |= ast_websocket_server_add_protocol2({{full_name}}.ws_server, protocol);
 {{/is_websocket}}
 {{/operations}}
 {{/apis}}