1 files changed, 25 insertions, 4 deletions

diff --git a/net/ipv4/af_inet.c b/net/ipv4/af_inet.c
index 193dcd6ed64f..c5376c725503 100644
--- a/net/ipv4/af_inet.c
+++ b/net/ipv4/af_inet.c
@@ -419,6 +419,8 @@ int inet_release(struct socket *sock)
         if (sk) {
                 long timeout;
 
+                inet_rps_reset_flow(sk);
+
                 /* Applications forget to leave groups before exiting */
                 ip_mc_drop_socket(sk);
 
@@ -720,6 +722,8 @@ int inet_sendmsg(struct kiocb *iocb, struct socket *sock, struct msghdr *msg,
 {
         struct sock *sk = sock->sk;
 
+        inet_rps_record_flow(sk);
+
         /* We may need to bind the socket. */
         if (!inet_sk(sk)->inet_num && inet_autobind(sk))
                 return -EAGAIN;
@@ -728,12 +732,13 @@ int inet_sendmsg(struct kiocb *iocb, struct socket *sock, struct msghdr *msg,
 }
 EXPORT_SYMBOL(inet_sendmsg);
 
-
 static ssize_t inet_sendpage(struct socket *sock, struct page *page, int offset,
                              size_t size, int flags)
 {
         struct sock *sk = sock->sk;
 
+        inet_rps_record_flow(sk);
+
         /* We may need to bind the socket. */
         if (!inet_sk(sk)->inet_num && inet_autobind(sk))
                 return -EAGAIN;
@@ -743,6 +748,22 @@ static ssize_t inet_sendpage(struct socket *sock, struct page *page, int offset,
         return sock_no_sendpage(sock, page, offset, size, flags);
 }
 
+int inet_recvmsg(struct kiocb *iocb, struct socket *sock, struct msghdr *msg,
+                 size_t size, int flags)
+{
+        struct sock *sk = sock->sk;
+        int addr_len = 0;
+        int err;
+
+        inet_rps_record_flow(sk);
+
+        err = sk->sk_prot->recvmsg(iocb, sk, msg, size, flags & MSG_DONTWAIT,
+                                   flags & ~MSG_DONTWAIT, &addr_len);
+        if (err >= 0)
+                msg->msg_namelen = addr_len;
+        return err;
+}
+EXPORT_SYMBOL(inet_recvmsg);
 
 int inet_shutdown(struct socket *sock, int how)
 {
@@ -872,7 +893,7 @@ const struct proto_ops inet_stream_ops = {
         .setsockopt        = sock_common_setsockopt,
         .getsockopt        = sock_common_getsockopt,
         .sendmsg           = tcp_sendmsg,
-        .recvmsg           = sock_common_recvmsg,
+        .recvmsg           = inet_recvmsg,
         .mmap              = sock_no_mmap,
         .sendpage          = tcp_sendpage,
         .splice_read       = tcp_splice_read,
@@ -899,7 +920,7 @@ const struct proto_ops inet_dgram_ops = {
         .setsockopt        = sock_common_setsockopt,
         .getsockopt        = sock_common_getsockopt,
         .sendmsg           = inet_sendmsg,
-        .recvmsg           = sock_common_recvmsg,
+        .recvmsg           = inet_recvmsg,
         .mmap              = sock_no_mmap,
         .sendpage          = inet_sendpage,
 #ifdef CONFIG_COMPAT
@@ -929,7 +950,7 @@ static const struct proto_ops inet_sockraw_ops = {
         .setsockopt        = sock_common_setsockopt,
         .getsockopt        = sock_common_getsockopt,
         .sendmsg           = inet_sendmsg,
-        .recvmsg           = sock_common_recvmsg,
+        .recvmsg           = inet_recvmsg,
         .mmap              = sock_no_mmap,
         .sendpage          = inet_sendpage,
 #ifdef CONFIG_COMPAT