psample: adjust size if rate_as_probability is set

If PSAMPLE_ATTR_SAMPLE_PROBABILITY flag is to be sent, the available
size for the packet data has to be adjusted accordingly.

Also, check the error code returned by nla_put_flag.

Fixes: 7b1b2b60c63f ("net: psample: allow using rate as probability")
Signed-off-by: Adrian Moreno <amorenoz@redhat.com>
Reviewed-by: Aaron Conole <aconole@redhat.com>
Reviewed-by: Ido Schimmel <idosch@nvidia.com>
Link: https://patch.msgid.link/20241217113739.3929300-1-amorenoz@redhat.com
Signed-off-by: Jakub Kicinski <kuba@kernel.org>

diff --git a/net/psample/psample.c b/net/psample/psample.c
index a0ddae8a65f917c49ea243a6f245ca876349145e..25f92ba0840c67d0d7284ab177a427adcfe579e4 100644 (file)
--- a/net/psample/psample.c
+++ b/net/psample/psample.c
@@ -393,7 +393,9 @@ void psample_sample_packet(struct psample_group *group,
                   nla_total_size_64bit(sizeof(u64)) +  /* timestamp */
                   nla_total_size(sizeof(u16)) +        /* protocol */
                   (md->user_cookie_len ?
-                   nla_total_size(md->user_cookie_len) : 0); /* user cookie */
+                   nla_total_size(md->user_cookie_len) : 0) + /* user cookie */
+                  (md->rate_as_probability ?
+                   nla_total_size(0) : 0);     /* rate as probability */
 
 #ifdef CONFIG_INET
        tun_info = skb_tunnel_info(skb);
@@ -498,8 +500,9 @@ void psample_sample_packet(struct psample_group *group,
                    md->user_cookie))
                goto error;
 
-       if (md->rate_as_probability)
-               nla_put_flag(nl_skb, PSAMPLE_ATTR_SAMPLE_PROBABILITY);
+       if (md->rate_as_probability &&
+           nla_put_flag(nl_skb, PSAMPLE_ATTR_SAMPLE_PROBABILITY))
+               goto error;
 
        genlmsg_end(nl_skb, data);
        genlmsg_multicast_netns(&psample_nl_family, group->net, nl_skb, 0,