[PATCH net 1/1] net/handshake: Protect request hash lookups

@@ -236,6 +236,7 @@ static void handshake_req_submit_test4(struct kunit *test)
 	/* Assert */
 	KUNIT_EXPECT_NOT_NULL(test, result);
 	KUNIT_EXPECT_PTR_EQ(test, req, result);
+	handshake_req_put(result);
 
 	handshake_req_cancel(sock->sk);
 	KUNIT_EXPECT_EQ(test, file_count(filp), fcount_before);

@@ -398,6 +399,7 @@ static void handshake_req_cancel_test2(struct kunit *test)
 	 * off by handshake_req_next(); see handshake_nl_accept_doit().
 	 */
 	fput(filp);
+	handshake_req_put(next);
 
 	/* Act */
 	result = handshake_req_cancel(sock->sk);

@@ -452,6 +454,7 @@ static void handshake_req_cancel_test3(struct kunit *test)
 	/* Pretend to complete this request */
 	handshake_complete(next, -ETIMEDOUT, NULL);
 	KUNIT_EXPECT_EQ(test, file_count(filp), fcount_before);
+	handshake_req_put(next);
 
 	/* Act */
 	result = handshake_req_cancel(sock->sk);

@@ -10,6 +10,8 @@
 #ifndef _INTERNAL_HANDSHAKE_H
 #define _INTERNAL_HANDSHAKE_H
 
+#include <linux/kref.h>
+
 /* Per-net namespace context */
 struct handshake_net {
 	spinlock_t		hn_lock;	/* protects next 3 fields */

@@ -36,6 +38,8 @@ struct handshake_req {
 	struct file			*hr_file;
 	struct sock			*hr_sk;
 	void				(*hr_odestruct)(struct sock *sk);
+	struct kref			hr_kref;
+	struct rcu_head			hr_rcu;
 
 	/* Always the last field */
 	char				hr_priv[];

@@ -85,6 +89,7 @@ int handshake_req_hash_init(void);
 void handshake_req_hash_destroy(void);
 void *handshake_req_private(struct handshake_req *req);
 struct handshake_req *handshake_req_hash_lookup(struct sock *sk);
+void handshake_req_put(struct handshake_req *req);
 struct handshake_req *handshake_req_next(struct handshake_net *hn, int class);
 int handshake_req_submit(struct socket *sock, struct handshake_req *req,
 			 gfp_t flags);

@@ -119,6 +119,7 @@ int handshake_nl_accept_doit(struct sk_buff *skb, struct genl_info *info)
 
 		trace_handshake_cmd_accept(net, req, req->hr_sk, fd_prepare_fd(fdf));
 		fd_publish(fdf);
+		handshake_req_put(req);
 		return 0;
 	}
 

@@ -127,6 +128,7 @@ int handshake_nl_accept_doit(struct sk_buff *skb, struct genl_info *info)
 		handshake_complete(req, -EIO, NULL);
 out_status:
 	trace_handshake_cmd_accept_err(net, req, NULL, err);
+	handshake_req_put(req);
 	return err;
 }
 

@@ -160,6 +162,7 @@ int handshake_nl_done_doit(struct sk_buff *skb, struct genl_info *info)
 		status = -(int)nla_get_u32(info->attrs[HANDSHAKE_A_DONE_STATUS]);
 
 	handshake_complete(req, status, info);
+	handshake_req_put(req);
 	sockfd_put(sock);
 	return 0;
 }

@@ -59,8 +59,16 @@ void handshake_req_hash_destroy(void)
 
 struct handshake_req *handshake_req_hash_lookup(struct sock *sk)
 {
-	return rhashtable_lookup_fast(&handshake_rhashtbl, &sk,
-				      handshake_rhash_params);
+	struct handshake_req *req;
+
+	rcu_read_lock();
+	req = rhashtable_lookup_fast(&handshake_rhashtbl, &sk,
+				     handshake_rhash_params);
+	if (req && !kref_get_unless_zero(&req->hr_kref))
+		req = NULL;
+	rcu_read_unlock();
+
+	return req;
 }
 EXPORT_SYMBOL_IF_KUNIT(handshake_req_hash_lookup);
 

@@ -74,13 +82,28 @@ static bool handshake_req_hash_add(struct handshake_req *req)
 	return ret == 0;
 }
 
-static void handshake_req_destroy(struct handshake_req *req)
+static void handshake_req_free(struct kref *kref)
 {
+	struct handshake_req *req =
+		container_of(kref, struct handshake_req, hr_kref);
+
 	if (req->hr_proto->hp_destroy)
 		req->hr_proto->hp_destroy(req);
+	kfree_rcu(req, hr_rcu);
+}
+
+void handshake_req_put(struct handshake_req *req)
+{
+	if (req)
+		kref_put(&req->hr_kref, handshake_req_free);
+}
+EXPORT_SYMBOL_IF_KUNIT(handshake_req_put);
+
+static void handshake_req_destroy(struct handshake_req *req)
+{
 	rhashtable_remove_fast(&handshake_rhashtbl, &req->hr_rhash,
 			       handshake_rhash_params);
-	kfree(req);
+	handshake_req_put(req);
 }
 
 static void handshake_sk_destruct(struct sock *sk)

@@ -95,6 +118,7 @@ static void handshake_sk_destruct(struct sock *sk)
 	trace_handshake_destruct(sock_net(sk), req, sk);
 	sk_destruct = req->hr_odestruct;
 	handshake_req_destroy(req);
+	handshake_req_put(req);
 	if (sk_destruct)
 		sk_destruct(sk);
 }

@@ -124,6 +148,7 @@ struct handshake_req *handshake_req_alloc(const struct handshake_proto *proto,
 	if (!req)
 		return NULL;
 
+	kref_init(&req->hr_kref);
 	INIT_LIST_HEAD(&req->hr_list);
 	req->hr_proto = proto;
 	return req;

@@ -187,8 +212,10 @@ static bool remove_pending(struct handshake_net *hn, struct handshake_req *req)
  * @class: handler class to match
  *
  * On a non-NULL return, the caller owns an extra reference
- * on @req->hr_file.  FD_PREPARE() consumes it on success; on
- * the FD_PREPARE() failure path the caller must fput() it.
+ * on @req->hr_file and a reference on @req.  FD_PREPARE() consumes
+ * the file reference on success; on the FD_PREPARE() failure path
+ * the caller must fput() it.  The caller must release the request
+ * reference using handshake_req_put().
  *
  * Return: pointer to a removed handshake_req, or NULL.
  */

@@ -209,6 +236,7 @@ struct handshake_req *handshake_req_next(struct handshake_net *hn, int class)
 		 * ownership.
 		 */
 		get_file(pos->hr_file);
+		kref_get(&pos->hr_kref);
 		req = pos;
 		break;
 	}

@@ -367,6 +395,7 @@ bool handshake_req_cancel(struct sock *sk)
 		/* Request hadn't been accepted - mark cancelled */
 		if (test_and_set_bit(HANDSHAKE_F_REQ_COMPLETED, &req->hr_flags)) {
 			trace_handshake_cancel_busy(net, req, sk);
+			handshake_req_put(req);
 			return false;
 		}
 		goto out_true;

@@ -374,6 +403,7 @@ bool handshake_req_cancel(struct sock *sk)
 	if (test_and_set_bit(HANDSHAKE_F_REQ_COMPLETED, &req->hr_flags)) {
 		/* Request already completed */
 		trace_handshake_cancel_busy(net, req, sk);
+		handshake_req_put(req);
 		return false;
 	}
 

@@ -381,6 +411,7 @@ bool handshake_req_cancel(struct sock *sk)
 	trace_handshake_cancel(net, req, sk);
 
 	fput(req->hr_file);
+	handshake_req_put(req);
 	return true;
 }
 EXPORT_SYMBOL(handshake_req_cancel);

@@ -449,9 +449,9 @@ void tls_handshake_close(struct socket *sock)
 	req = handshake_req_hash_lookup(sock->sk);
 	if (!req)
 		return;
-	if (!test_and_clear_bit(HANDSHAKE_F_REQ_SESSION, &req->hr_flags))
-		return;
-	tls_alert_send(sock, TLS_ALERT_LEVEL_WARNING,
-		       TLS_ALERT_DESC_CLOSE_NOTIFY);
+	if (test_and_clear_bit(HANDSHAKE_F_REQ_SESSION, &req->hr_flags))
+		tls_alert_send(sock, TLS_ALERT_LEVEL_WARNING,
+			       TLS_ALERT_DESC_CLOSE_NOTIFY);
+	handshake_req_put(req);
 }
 EXPORT_SYMBOL(tls_handshake_close);