Upgrade connection to websocket and provision new connection

fw/connection.h

@@ -51,6 +51,14 @@ enum {
 	/* HTTPS */
 	Conn_HttpsClnt	= Conn_Clnt | TFW_FSM_HTTPS,
 	Conn_HttpsSrv	= Conn_Srv | TFW_FSM_HTTPS,
+
+	/* Websocket plain */
+	Conn_WsClnt	= Conn_HttpClnt | TFW_FSM_WS,
+	Conn_WsSrv	= Conn_HttpSrv | TFW_FSM_WS,
+
+	/* Websocket secure */
+	Conn_WssClnt	= Conn_HttpsClnt | TFW_FSM_WS,
+	Conn_WssSrv	= Conn_HttpsSrv | TFW_FSM_WS,
 };
 
 #define TFW_CONN_TYPE2IDX(t)	TFW_FSM_TYPE(t)
@@ -100,7 +108,7 @@ enum {
 	struct sock		*sk;			\
 	void			(*destructor)(void *);
 
-typedef struct TfwConn {
+typedef struct tfw_conn_t {
 	TFW_CONN_COMMON;
 } TfwConn;
 
@@ -199,7 +207,9 @@ enum {
 	/* Connection is in use or at least scheduled to be established. */
 	TFW_CONN_B_ACTIVE,
 	/* Connection is disconnected and stopped. */
-	TFW_CONN_B_STOPPED
+	TFW_CONN_B_STOPPED,
+	/* Mark connection as unavailable to schedulers */
+	TFW_CONN_B_UNSCHED
 };
 
 /**
@@ -297,6 +307,15 @@ tfw_srv_conn_restricted(TfwSrvConn *srv_conn)
 	return test_bit(TFW_CONN_B_RESEND, &srv_conn->flags);
 }
 
+/*
+ * Connection is unavailable to scheduler and may be removed from it
+ */
+static inline bool
+tfw_srv_conn_unscheduled(TfwSrvConn *srv_conn)
+{
+	return test_bit(TFW_CONN_B_UNSCHED, &srv_conn->flags);
+}
+
 /*
  * Tell if a connection has non-idempotent requests.
  */

fw/gfsm.h

@@ -84,6 +84,8 @@ enum {
 	/* Protocols */
 	TFW_FSM_HTTP,
 	TFW_FSM_HTTPS,
+	/* Not really a FSM */
+	TFW_FSM_WS,
 
 	/* Security rules enforcement. */
 	TFW_FSM_FRANG_REQ,
@@ -181,7 +183,7 @@ typedef struct {
 			  & ((TFW_GFSM_FSM_MASK << TFW_GFSM_FSM_SHIFT)	\
 			    | TFW_GFSM_STATE_MASK))
 
-typedef struct TfwConn TfwConn;
+typedef struct tfw_conn_t TfwConn;
 typedef int (*tfw_gfsm_handler_t)(TfwConn *conn, TfwFsmData *data);
 
 void tfw_gfsm_state_init(TfwGState *st, void *obj, int st0);

fw/http.c

@@ -2339,7 +2339,6 @@ static int
 tfw_http_conn_init(TfwConn *conn)
 {
 	T_DBG2("%s: conn=[%p]\n", __func__, conn);
-
 	if (TFW_CONN_TYPE(conn) & Conn_Srv) {
 		TfwSrvConn *srv_conn = (TfwSrvConn *)conn;
 		if (!list_empty(&srv_conn->fwd_queue)) {
@@ -5920,6 +5919,42 @@ tfw_http_resp_terminate(TfwHttpMsg *hm)
 	tfw_http_resp_cache(hm);
 }
 
+/**
+ * Does websocket upgrade procedure.
+ *
+ * Marks current server and client connection as websocket connection. Starts
+ * reconnection with backend to restore full backend connection count.
+ *
+ * @return zero on success and negative otherwise
+ */
+static int
+tfw_http_websocket_upgrade(TfwHttpResp *resp)
+{
+	TfwServer *srv = (TfwServer *)resp->conn->peer;
+	TfwSrvConn *srv_conn;
+
+	/* Cannot proceed with upgrade websocket due to error
+	* in creation of new http connection. While it will not be
+	* inherently erroneous to upgrade existing connection, but
+	* we would pay for it with essentially dropping connection with
+	* server. Better just drop upgrade request and reestablish connection.
+	*/
+	if (!(srv_conn = tfw_sock_srv_new_conn(srv))) {
+		tfw_http_conn_error_log(resp->conn,
+					"Can't create new connection for "
+					"websocket upgrade response");
+		return TFW_BLOCK;
+	}
+
+	set_bit(TFW_CONN_B_UNSCHED, &((TfwSrvConn *)resp->conn)->flags);
+
+	tfw_sock_srv_connect_one(srv, srv_conn);
+
+	srv->sg->sched->upd_srv(srv);
+
+	return TFW_PASS;
+}
+
 /**
  * @return zero on success and negative value otherwise.
  * TODO enter the function depending on current GFSM state.
@@ -5931,6 +5966,7 @@ tfw_http_resp_process(TfwConn *conn, TfwStream *stream, struct sk_buff *skb)
 	unsigned int chunks_unused, parsed;
 	TfwHttpReq *bad_req;
 	TfwHttpMsg *hmresp, *hmsib;
+	TfwHttpResp *resp;
 	TfwFsmData data_up;
 	bool conn_stop, filtout = false;
 
@@ -5953,6 +5989,7 @@ tfw_http_resp_process(TfwConn *conn, TfwStream *stream, struct sk_buff *skb)
 	parsed = 0;
 	hmsib = NULL;
 	hmresp = (TfwHttpMsg *)stream->msg;
+	resp = (TfwHttpResp *)hmresp;
 
 	r = ss_skb_process(skb, tfw_http_parse_resp, hmresp, &chunks_unused,
 			   &parsed);
@@ -6098,6 +6135,23 @@ tfw_http_resp_process(TfwConn *conn, TfwStream *stream, struct sk_buff *skb)
 		r = TFW_PASS;
 		goto next_resp;
 	}
+
+	/*
+	 * Upgrade client and server connection to websocket, remove it
+	 * from scheduler and provision new connection.
+	 *
+	 * TODO #755: set existent client and server connection to Conn_Ws*
+	 * when websocket proxing protocol will be implemented
+	 */
+	if (unlikely(test_bit(TFW_HTTP_B_CONN_UPGRADE, hmresp->flags)
+		     && test_bit(TFW_HTTP_B_UPGRADE_WEBSOCKET, hmresp->flags)
+		     && resp->status == 101))
+	{
+		r = tfw_http_websocket_upgrade(resp);
+		if (unlikely(r < TFW_PASS))
+			return TFW_BLOCK;
+	}
+
 	/*
 	 * Pass the response to cache for further processing.
 	 * In the end, the response is sent on to the client.

fw/http_frame.h

@@ -1,7 +1,7 @@
 /**
  *		Tempesta FW
  *
- * Copyright (C) 2019 Tempesta Technologies, Inc.
+ * Copyright (C) 2022 Tempesta Technologies, Inc.
  *
  * This program is free software; you can redistribute it and/or modify it
  * under the terms of the GNU General Public License as published by
@@ -200,7 +200,7 @@ typedef struct {
 	unsigned char	data_off;
 } TfwH2Ctx;
 
-typedef struct TfwConn TfwConn;
+typedef struct tfw_conn_t TfwConn;
 
 int tfw_h2_init(void);
 void tfw_h2_cleanup(void);

fw/http_sched_hash.c

@@ -19,7 +19,7 @@
  * server unless it is offline.
  *
  * Copyright (C) 2014 NatSys Lab. (info@natsys-lab.com).
- * Copyright (C) 2015-2021 Tempesta Technologies, Inc.
+ * Copyright (C) 2015-2022 Tempesta Technologies, Inc.
  *
  * This program is free software; you can redistribute it and/or modify it
  * under the terms of the GNU General Public License as published by
@@ -129,6 +129,7 @@ __is_conn_suitable(TfwSrvConn *conn, bool hmonitor)
 {
 	return (hmonitor || !tfw_srv_suspended((TfwServer *)conn->peer))
 		&& !tfw_srv_conn_restricted(conn)
+		&& !tfw_srv_conn_unscheduled(conn)
 		&& !tfw_srv_conn_busy(conn)
 		&& !tfw_srv_conn_queue_full(conn)
 		&& tfw_srv_conn_get_if_live(conn);
@@ -333,6 +334,8 @@ tfw_sched_hash_add_conns(TfwServer *srv, TfwHashConnList *cl, size_t *seed,
 
 	list_for_each_entry(conn, &srv->conn_list, list) {
 		unsigned long hash;
+		if (tfw_srv_conn_unscheduled(conn))
+			continue;
 		do {
 			hash = __hash_64(srv_hash ^ *seed);
 			*seed += seed_inc;
@@ -392,34 +395,46 @@ tfw_sched_hash_add_grp(TfwSrvGroup *sg, void *data)
 	return 0;
 }
 
+static void
+tfw_sched_hash_put_srv_data(struct rcu_head *rcu)
+{
+	TfwHashConnList *cl = container_of(rcu, TfwHashConnList, rcu);
+	kfree(cl);
+}
+
 static int
-tfw_sched_hash_add_srv(TfwServer *srv)
+tfw_sched_hash_srv_setup(TfwServer *srv)
 {
 	size_t size, seed, seed_inc = 0;
-	TfwHashConnList *cl = rcu_dereference_check(srv->sched_data, 1);
-
-	if (unlikely(cl))
-		return -EEXIST;
+	TfwHashConnList *cl = rcu_dereference_bh_check(srv->sched_data, 1);
+	TfwHashConnList *cl_copy;
 
 	seed = get_random_long();
 	seed_inc = get_random_int();
 
 	size = sizeof(TfwHashConnList) + srv->conn_n * sizeof(TfwHashConn);
-	if (!(cl = kzalloc(size, GFP_KERNEL)))
+	if (!(cl_copy = kzalloc(size, in_task() ? GFP_KERNEL : GFP_ATOMIC)))
 		return -ENOMEM;
 
-	tfw_sched_hash_add_conns(srv, cl, &seed, seed_inc);
+	tfw_sched_hash_add_conns(srv, cl_copy, &seed, seed_inc);
 
-	rcu_assign_pointer(srv->sched_data, cl);
+	rcu_assign_pointer(srv->sched_data, cl_copy);
+
+	if (cl)
+		call_rcu(&cl->rcu, tfw_sched_hash_put_srv_data);
 
 	return 0;
 }
 
-static void
-tfw_sched_hash_put_srv_data(struct rcu_head *rcu)
+static int
+tfw_sched_hash_add_srv(TfwServer *srv)
 {
-	TfwHashConnList *cl = container_of(rcu, TfwHashConnList, rcu);
-	kfree(cl);
+	TfwHashConnList *cl = rcu_dereference_check(srv->sched_data, 1);
+
+	if (unlikely(cl))
+		return -EEXIST;
+
+	return tfw_sched_hash_srv_setup(srv);
 }
 
 static void
@@ -432,13 +447,20 @@ tfw_sched_hash_del_srv(TfwServer *srv)
 		call_rcu(&cl->rcu, tfw_sched_hash_put_srv_data);
 }
 
+static int
+tfw_sched_hash_upd_srv(TfwServer *srv)
+{
+	return tfw_sched_hash_srv_setup(srv);
+}
+
 static TfwScheduler tfw_sched_hash = {
 	.name		= "hash",
 	.list		= LIST_HEAD_INIT(tfw_sched_hash.list),
 	.add_grp	= tfw_sched_hash_add_grp,
 	.del_grp	= tfw_sched_hash_del_grp,
 	.add_srv	= tfw_sched_hash_add_srv,
 	.del_srv	= tfw_sched_hash_del_srv,
+	.upd_srv	= tfw_sched_hash_upd_srv,
 	.sched_sg_conn	= tfw_sched_hash_get_sg_conn,
 	.sched_srv_conn	= tfw_sched_hash_get_srv_conn,
 };