diff mbox series

[net-next] selftests: add selftest for UDP SO_PEEK_OFF support

Message ID 20240901051821.94956-1-kerneljasonxing@gmail.com (mailing list archive)
State New
Headers show
Series [net-next] selftests: add selftest for UDP SO_PEEK_OFF support | expand

Commit Message

Jason Xing Sept. 1, 2024, 5:18 a.m. UTC
From: Jason Xing <kernelxing@tencent.com>

Add the SO_PEEK_OFF selftest for UDP. In this patch, I mainly do
three things:
1. rename tcp_so_peek_off.c
2. adjust for UDP protocol
3. add selftests into it

Suggested-by: Jon Maloy <jmaloy@redhat.com>
Signed-off-by: Jason Xing <kernelxing@tencent.com>
---
Link: https://lore.kernel.org/all/9f4dd14d-fbe3-4c61-b04c-f0e6b8096d7b@redhat.com/
---
 tools/testing/selftests/net/Makefile          |  2 +-
 .../{tcp_so_peek_off.c => sk_so_peek_off.c}   | 91 +++++++++++--------
 2 files changed, 56 insertions(+), 37 deletions(-)
 rename tools/testing/selftests/net/{tcp_so_peek_off.c => sk_so_peek_off.c} (58%)

Comments

Stanislav Fomichev Sept. 2, 2024, 1:07 a.m. UTC | #1
On 09/01, Jason Xing wrote:
> From: Jason Xing <kernelxing@tencent.com>
> 
> Add the SO_PEEK_OFF selftest for UDP. In this patch, I mainly do
> three things:
> 1. rename tcp_so_peek_off.c
> 2. adjust for UDP protocol
> 3. add selftests into it
> 
> Suggested-by: Jon Maloy <jmaloy@redhat.com>
> Signed-off-by: Jason Xing <kernelxing@tencent.com>
> ---
> Link: https://lore.kernel.org/all/9f4dd14d-fbe3-4c61-b04c-f0e6b8096d7b@redhat.com/
> ---
>  tools/testing/selftests/net/Makefile          |  2 +-
>  .../{tcp_so_peek_off.c => sk_so_peek_off.c}   | 91 +++++++++++--------
>  2 files changed, 56 insertions(+), 37 deletions(-)
>  rename tools/testing/selftests/net/{tcp_so_peek_off.c => sk_so_peek_off.c} (58%)
> 
> diff --git a/tools/testing/selftests/net/Makefile b/tools/testing/selftests/net/Makefile
> index 1179e3261bef..d5029f978aa9 100644
> --- a/tools/testing/selftests/net/Makefile
> +++ b/tools/testing/selftests/net/Makefile
> @@ -80,7 +80,7 @@ TEST_PROGS += io_uring_zerocopy_tx.sh
>  TEST_GEN_FILES += bind_bhash
>  TEST_GEN_PROGS += sk_bind_sendto_listen
>  TEST_GEN_PROGS += sk_connect_zero_addr
> -TEST_GEN_PROGS += tcp_so_peek_off
> +TEST_GEN_PROGS += sk_so_peek_off

Should we also add sk_so_peek_off to gitignore?
Jason Xing Sept. 2, 2024, 2:54 a.m. UTC | #2
On Mon, Sep 2, 2024 at 9:07 AM Stanislav Fomichev <sdf@fomichev.me> wrote:
>
> On 09/01, Jason Xing wrote:
> > From: Jason Xing <kernelxing@tencent.com>
> >
> > Add the SO_PEEK_OFF selftest for UDP. In this patch, I mainly do
> > three things:
> > 1. rename tcp_so_peek_off.c
> > 2. adjust for UDP protocol
> > 3. add selftests into it
> >
> > Suggested-by: Jon Maloy <jmaloy@redhat.com>
> > Signed-off-by: Jason Xing <kernelxing@tencent.com>
> > ---
> > Link: https://lore.kernel.org/all/9f4dd14d-fbe3-4c61-b04c-f0e6b8096d7b@redhat.com/
> > ---
> >  tools/testing/selftests/net/Makefile          |  2 +-
> >  .../{tcp_so_peek_off.c => sk_so_peek_off.c}   | 91 +++++++++++--------
> >  2 files changed, 56 insertions(+), 37 deletions(-)
> >  rename tools/testing/selftests/net/{tcp_so_peek_off.c => sk_so_peek_off.c} (58%)
> >
> > diff --git a/tools/testing/selftests/net/Makefile b/tools/testing/selftests/net/Makefile
> > index 1179e3261bef..d5029f978aa9 100644
> > --- a/tools/testing/selftests/net/Makefile
> > +++ b/tools/testing/selftests/net/Makefile
> > @@ -80,7 +80,7 @@ TEST_PROGS += io_uring_zerocopy_tx.sh
> >  TEST_GEN_FILES += bind_bhash
> >  TEST_GEN_PROGS += sk_bind_sendto_listen
> >  TEST_GEN_PROGS += sk_connect_zero_addr
> > -TEST_GEN_PROGS += tcp_so_peek_off
> > +TEST_GEN_PROGS += sk_so_peek_off
>
> Should we also add sk_so_peek_off to gitignore?

Good catch. Will update it in the next version :)

Thanks,
Jason
Willem de Bruijn Sept. 2, 2024, 3:02 p.m. UTC | #3
Jason Xing wrote:
> From: Jason Xing <kernelxing@tencent.com>
> 
> Add the SO_PEEK_OFF selftest for UDP. In this patch, I mainly do
> three things:
> 1. rename tcp_so_peek_off.c
> 2. adjust for UDP protocol
> 3. add selftests into it
> 
> Suggested-by: Jon Maloy <jmaloy@redhat.com>
> Signed-off-by: Jason Xing <kernelxing@tencent.com>

A few minor comments. Nothing important.

Subject to Stan's point about .gitignore:

Reviewed-by: Willem de Bruijn <willemb@google.com>

> -int tcp_peek_offset_probe(sa_family_t af)
> +int sk_peek_offset_probe(sa_family_t af, int proto)
>  {
> +	int type = (proto == IPPROTO_TCP ? SOCK_STREAM : SOCK_DGRAM);
>  	int optv = 0;
>  	int ret = 0;
>  	int s;
>  
> -	s = socket(af, SOCK_STREAM | SOCK_CLOEXEC, IPPROTO_TCP);
> +	s = socket(af, type, proto);

Removing the SOCK_CLOEXEC because not relevant to this single thread
process, I suppose?

Not important, but no need for proto, can just be 0.

>  	if (s < 0) {
>  		ksft_perror("Temporary TCP socket creation failed");
>  	} else {
>  		if (!setsockopt(s, SOL_SOCKET, SO_PEEK_OFF, &optv, sizeof(int)))
>  			ret = 1;
>  		else
> -			printf("%s does not support SO_PEEK_OFF\n", afstr(af));
> +			printf("%s does not support SO_PEEK_OFF\n", afstr(af, proto));
>  		close(s);
>  	}
>  	return ret;
>  }
>  
> -static void tcp_peek_offset_set(int s, int offset)
> +static void sk_peek_offset_set(int s, int offset)
>  {
>  	if (setsockopt(s, SOL_SOCKET, SO_PEEK_OFF, &offset, sizeof(offset)))
>  		ksft_perror("Failed to set SO_PEEK_OFF value\n");
>  }
>  
> -static int tcp_peek_offset_get(int s)
> +static int sk_peek_offset_get(int s)
>  {
>  	int offset;
>  	socklen_t len = sizeof(offset);
> @@ -50,8 +54,9 @@ static int tcp_peek_offset_get(int s)
>  	return offset;
>  }
>  
> -static int tcp_peek_offset_test(sa_family_t af)
> +static int sk_peek_offset_test(sa_family_t af, int proto)
>  {
> +	int type = (proto == IPPROTO_TCP ? SOCK_STREAM : SOCK_DGRAM);
>  	union {
>  		struct sockaddr sa;
>  		struct sockaddr_in a4;
> @@ -62,13 +67,13 @@ static int tcp_peek_offset_test(sa_family_t af)
>  	int recv_sock = 0;
>  	int offset = 0;
>  	ssize_t len;
> -	char buf;
> +	char buf[2];
>  
>  	memset(&a, 0, sizeof(a));
>  	a.sa.sa_family = af;
>  
> -	s[0] = socket(af, SOCK_STREAM, IPPROTO_TCP);
> -	s[1] = socket(af, SOCK_STREAM | SOCK_NONBLOCK, IPPROTO_TCP);
> +	s[0] = recv_sock = socket(af, type, proto);
> +	s[1] = socket(af, type, proto);

Same

>  
>  	if (s[0] < 0 || s[1] < 0) {
>  		ksft_perror("Temporary socket creation failed\n");
> @@ -82,76 +87,78 @@ static int tcp_peek_offset_test(sa_family_t af)
>  		ksft_perror("Temporary socket getsockname() failed\n");
>  		goto out;
>  	}
> -	if (listen(s[0], 0) < 0) {
> +	if (proto == IPPROTO_TCP && listen(s[0], 0) < 0) {
>  		ksft_perror("Temporary socket listen() failed\n");
>  		goto out;
>  	}
> -	if (connect(s[1], &a.sa, sizeof(a)) >= 0 || errno != EINPROGRESS) {
> +	if (connect(s[1], &a.sa, sizeof(a))) {
>  		ksft_perror("Temporary socket connect() failed\n");
>  		goto out;
>  	}

Changed due to the removal of SOCK_NONBLOCK above. Definitely
simplifies the test.

Just note that error test is == -1 or < 0, also for consistency with
the rest of the file.

> -	recv_sock = accept(s[0], NULL, NULL);
> -	if (recv_sock <= 0) {
> -		ksft_perror("Temporary socket accept() failed\n");
> -		goto out;
> +	if (proto == IPPROTO_TCP) {
> +		recv_sock = accept(s[0], NULL, NULL);
> +		if (recv_sock <= 0) {
> +			ksft_perror("Temporary socket accept() failed\n");
> +			goto out;
> +		}
Jason Xing Sept. 2, 2024, 3:33 p.m. UTC | #4
On Mon, Sep 2, 2024 at 11:02 PM Willem de Bruijn
<willemdebruijn.kernel@gmail.com> wrote:
>
> Jason Xing wrote:
> > From: Jason Xing <kernelxing@tencent.com>
> >
> > Add the SO_PEEK_OFF selftest for UDP. In this patch, I mainly do
> > three things:
> > 1. rename tcp_so_peek_off.c
> > 2. adjust for UDP protocol
> > 3. add selftests into it
> >
> > Suggested-by: Jon Maloy <jmaloy@redhat.com>
> > Signed-off-by: Jason Xing <kernelxing@tencent.com>
>
> A few minor comments. Nothing important.
>
> Subject to Stan's point about .gitignore:
>
> Reviewed-by: Willem de Bruijn <willemb@google.com>

Thanks for your review!

>
> > -int tcp_peek_offset_probe(sa_family_t af)
> > +int sk_peek_offset_probe(sa_family_t af, int proto)
> >  {
> > +     int type = (proto == IPPROTO_TCP ? SOCK_STREAM : SOCK_DGRAM);
> >       int optv = 0;
> >       int ret = 0;
> >       int s;
> >
> > -     s = socket(af, SOCK_STREAM | SOCK_CLOEXEC, IPPROTO_TCP);
> > +     s = socket(af, type, proto);
>
> Removing the SOCK_CLOEXEC because not relevant to this single thread
> process, I suppose?

Yep. We don't need this one.

>
> Not important, but no need for proto, can just be 0.

You're right. I wonder if it is better if we explicitly pass the proto
here? I would like not to touch it here.

>
> >       if (s < 0) {
> >               ksft_perror("Temporary TCP socket creation failed");
> >       } else {
> >               if (!setsockopt(s, SOL_SOCKET, SO_PEEK_OFF, &optv, sizeof(int)))
> >                       ret = 1;
> >               else
> > -                     printf("%s does not support SO_PEEK_OFF\n", afstr(af));
> > +                     printf("%s does not support SO_PEEK_OFF\n", afstr(af, proto));
> >               close(s);
> >       }
> >       return ret;
> >  }
> >
> > -static void tcp_peek_offset_set(int s, int offset)
> > +static void sk_peek_offset_set(int s, int offset)
> >  {
> >       if (setsockopt(s, SOL_SOCKET, SO_PEEK_OFF, &offset, sizeof(offset)))
> >               ksft_perror("Failed to set SO_PEEK_OFF value\n");
> >  }
> >
> > -static int tcp_peek_offset_get(int s)
> > +static int sk_peek_offset_get(int s)
> >  {
> >       int offset;
> >       socklen_t len = sizeof(offset);
> > @@ -50,8 +54,9 @@ static int tcp_peek_offset_get(int s)
> >       return offset;
> >  }
> >
> > -static int tcp_peek_offset_test(sa_family_t af)
> > +static int sk_peek_offset_test(sa_family_t af, int proto)
> >  {
> > +     int type = (proto == IPPROTO_TCP ? SOCK_STREAM : SOCK_DGRAM);
> >       union {
> >               struct sockaddr sa;
> >               struct sockaddr_in a4;
> > @@ -62,13 +67,13 @@ static int tcp_peek_offset_test(sa_family_t af)
> >       int recv_sock = 0;
> >       int offset = 0;
> >       ssize_t len;
> > -     char buf;
> > +     char buf[2];
> >
> >       memset(&a, 0, sizeof(a));
> >       a.sa.sa_family = af;
> >
> > -     s[0] = socket(af, SOCK_STREAM, IPPROTO_TCP);
> > -     s[1] = socket(af, SOCK_STREAM | SOCK_NONBLOCK, IPPROTO_TCP);
> > +     s[0] = recv_sock = socket(af, type, proto);
> > +     s[1] = socket(af, type, proto);
>
> Same

I think we don't need this one, either.
As we can see, there are already some existing test files without the
SOCK_NONBLOCK flag.

>
> >
> >       if (s[0] < 0 || s[1] < 0) {
> >               ksft_perror("Temporary socket creation failed\n");
> > @@ -82,76 +87,78 @@ static int tcp_peek_offset_test(sa_family_t af)
> >               ksft_perror("Temporary socket getsockname() failed\n");
> >               goto out;
> >       }
> > -     if (listen(s[0], 0) < 0) {
> > +     if (proto == IPPROTO_TCP && listen(s[0], 0) < 0) {
> >               ksft_perror("Temporary socket listen() failed\n");
> >               goto out;
> >       }
> > -     if (connect(s[1], &a.sa, sizeof(a)) >= 0 || errno != EINPROGRESS) {
> > +     if (connect(s[1], &a.sa, sizeof(a))) {
> >               ksft_perror("Temporary socket connect() failed\n");
> >               goto out;
> >       }
>
> Changed due to the removal of SOCK_NONBLOCK above. Definitely
> simplifies the test.

Yep.

>
> Just note that error test is == -1 or < 0, also for consistency with
> the rest of the file.

I will add "< 0" here as you said.

Thanks,
Jason
Willem de Bruijn Sept. 2, 2024, 4:45 p.m. UTC | #5
Jason Xing wrote:
> On Mon, Sep 2, 2024 at 11:02 PM Willem de Bruijn
> <willemdebruijn.kernel@gmail.com> wrote:
> >
> > Jason Xing wrote:
> > > From: Jason Xing <kernelxing@tencent.com>
> > >
> > > Add the SO_PEEK_OFF selftest for UDP. In this patch, I mainly do
> > > three things:
> > > 1. rename tcp_so_peek_off.c
> > > 2. adjust for UDP protocol
> > > 3. add selftests into it
> > >
> > > Suggested-by: Jon Maloy <jmaloy@redhat.com>
> > > Signed-off-by: Jason Xing <kernelxing@tencent.com>
> >
> > A few minor comments. Nothing important.
> >
> > Subject to Stan's point about .gitignore:
> >
> > Reviewed-by: Willem de Bruijn <willemb@google.com>
> 
> Thanks for your review!
> 
> >
> > > -int tcp_peek_offset_probe(sa_family_t af)
> > > +int sk_peek_offset_probe(sa_family_t af, int proto)
> > >  {
> > > +     int type = (proto == IPPROTO_TCP ? SOCK_STREAM : SOCK_DGRAM);
> > >       int optv = 0;
> > >       int ret = 0;
> > >       int s;
> > >
> > > -     s = socket(af, SOCK_STREAM | SOCK_CLOEXEC, IPPROTO_TCP);
> > > +     s = socket(af, type, proto);
> >
> > Removing the SOCK_CLOEXEC because not relevant to this single thread
> > process, I suppose?
> 
> Yep. We don't need this one.
> 
> >
> > Not important, but no need for proto, can just be 0.
> 
> You're right. I wonder if it is better if we explicitly pass the proto
> here? I would like not to touch it here.

It's not better or worse. Just not needed. So either way.
diff mbox series

Patch

diff --git a/tools/testing/selftests/net/Makefile b/tools/testing/selftests/net/Makefile
index 1179e3261bef..d5029f978aa9 100644
--- a/tools/testing/selftests/net/Makefile
+++ b/tools/testing/selftests/net/Makefile
@@ -80,7 +80,7 @@  TEST_PROGS += io_uring_zerocopy_tx.sh
 TEST_GEN_FILES += bind_bhash
 TEST_GEN_PROGS += sk_bind_sendto_listen
 TEST_GEN_PROGS += sk_connect_zero_addr
-TEST_GEN_PROGS += tcp_so_peek_off
+TEST_GEN_PROGS += sk_so_peek_off
 TEST_PROGS += test_ingress_egress_chaining.sh
 TEST_GEN_PROGS += so_incoming_cpu
 TEST_PROGS += sctp_vrf.sh
diff --git a/tools/testing/selftests/net/tcp_so_peek_off.c b/tools/testing/selftests/net/sk_so_peek_off.c
similarity index 58%
rename from tools/testing/selftests/net/tcp_so_peek_off.c
rename to tools/testing/selftests/net/sk_so_peek_off.c
index df8a39d9d3c3..870a890138c4 100644
--- a/tools/testing/selftests/net/tcp_so_peek_off.c
+++ b/tools/testing/selftests/net/sk_so_peek_off.c
@@ -10,37 +10,41 @@ 
 #include <arpa/inet.h>
 #include "../kselftest.h"
 
-static char *afstr(int af)
+static char *afstr(int af, int proto)
 {
-	return af == AF_INET ? "TCP/IPv4" : "TCP/IPv6";
+	if (proto == IPPROTO_TCP)
+		return af == AF_INET ? "TCP/IPv4" : "TCP/IPv6";
+	else
+		return af == AF_INET ? "UDP/IPv4" : "UDP/IPv6";
 }
 
-int tcp_peek_offset_probe(sa_family_t af)
+int sk_peek_offset_probe(sa_family_t af, int proto)
 {
+	int type = (proto == IPPROTO_TCP ? SOCK_STREAM : SOCK_DGRAM);
 	int optv = 0;
 	int ret = 0;
 	int s;
 
-	s = socket(af, SOCK_STREAM | SOCK_CLOEXEC, IPPROTO_TCP);
+	s = socket(af, type, proto);
 	if (s < 0) {
 		ksft_perror("Temporary TCP socket creation failed");
 	} else {
 		if (!setsockopt(s, SOL_SOCKET, SO_PEEK_OFF, &optv, sizeof(int)))
 			ret = 1;
 		else
-			printf("%s does not support SO_PEEK_OFF\n", afstr(af));
+			printf("%s does not support SO_PEEK_OFF\n", afstr(af, proto));
 		close(s);
 	}
 	return ret;
 }
 
-static void tcp_peek_offset_set(int s, int offset)
+static void sk_peek_offset_set(int s, int offset)
 {
 	if (setsockopt(s, SOL_SOCKET, SO_PEEK_OFF, &offset, sizeof(offset)))
 		ksft_perror("Failed to set SO_PEEK_OFF value\n");
 }
 
-static int tcp_peek_offset_get(int s)
+static int sk_peek_offset_get(int s)
 {
 	int offset;
 	socklen_t len = sizeof(offset);
@@ -50,8 +54,9 @@  static int tcp_peek_offset_get(int s)
 	return offset;
 }
 
-static int tcp_peek_offset_test(sa_family_t af)
+static int sk_peek_offset_test(sa_family_t af, int proto)
 {
+	int type = (proto == IPPROTO_TCP ? SOCK_STREAM : SOCK_DGRAM);
 	union {
 		struct sockaddr sa;
 		struct sockaddr_in a4;
@@ -62,13 +67,13 @@  static int tcp_peek_offset_test(sa_family_t af)
 	int recv_sock = 0;
 	int offset = 0;
 	ssize_t len;
-	char buf;
+	char buf[2];
 
 	memset(&a, 0, sizeof(a));
 	a.sa.sa_family = af;
 
-	s[0] = socket(af, SOCK_STREAM, IPPROTO_TCP);
-	s[1] = socket(af, SOCK_STREAM | SOCK_NONBLOCK, IPPROTO_TCP);
+	s[0] = recv_sock = socket(af, type, proto);
+	s[1] = socket(af, type, proto);
 
 	if (s[0] < 0 || s[1] < 0) {
 		ksft_perror("Temporary socket creation failed\n");
@@ -82,76 +87,78 @@  static int tcp_peek_offset_test(sa_family_t af)
 		ksft_perror("Temporary socket getsockname() failed\n");
 		goto out;
 	}
-	if (listen(s[0], 0) < 0) {
+	if (proto == IPPROTO_TCP && listen(s[0], 0) < 0) {
 		ksft_perror("Temporary socket listen() failed\n");
 		goto out;
 	}
-	if (connect(s[1], &a.sa, sizeof(a)) >= 0 || errno != EINPROGRESS) {
+	if (connect(s[1], &a.sa, sizeof(a))) {
 		ksft_perror("Temporary socket connect() failed\n");
 		goto out;
 	}
-	recv_sock = accept(s[0], NULL, NULL);
-	if (recv_sock <= 0) {
-		ksft_perror("Temporary socket accept() failed\n");
-		goto out;
+	if (proto == IPPROTO_TCP) {
+		recv_sock = accept(s[0], NULL, NULL);
+		if (recv_sock <= 0) {
+			ksft_perror("Temporary socket accept() failed\n");
+			goto out;
+		}
 	}
 
 	/* Some basic tests of getting/setting offset */
-	offset = tcp_peek_offset_get(recv_sock);
+	offset = sk_peek_offset_get(recv_sock);
 	if (offset != -1) {
 		ksft_perror("Initial value of socket offset not -1\n");
 		goto out;
 	}
-	tcp_peek_offset_set(recv_sock, 0);
-	offset = tcp_peek_offset_get(recv_sock);
+	sk_peek_offset_set(recv_sock, 0);
+	offset = sk_peek_offset_get(recv_sock);
 	if (offset != 0) {
 		ksft_perror("Failed to set socket offset to 0\n");
 		goto out;
 	}
 
 	/* Transfer a message */
-	if (send(s[1], (char *)("ab"), 2, 0) <= 0 || errno != EINPROGRESS) {
+	if (send(s[1], (char *)("ab"), 2, 0) != 2) {
 		ksft_perror("Temporary probe socket send() failed\n");
 		goto out;
 	}
 	/* Read first byte */
-	len = recv(recv_sock, &buf, 1, MSG_PEEK);
-	if (len != 1 || buf != 'a') {
+	len = recv(recv_sock, buf, 1, MSG_PEEK);
+	if (len != 1 || buf[0] != 'a') {
 		ksft_perror("Failed to read first byte of message\n");
 		goto out;
 	}
-	offset = tcp_peek_offset_get(recv_sock);
+	offset = sk_peek_offset_get(recv_sock);
 	if (offset != 1) {
 		ksft_perror("Offset not forwarded correctly at first byte\n");
 		goto out;
 	}
 	/* Try to read beyond last byte */
-	len = recv(recv_sock, &buf, 2, MSG_PEEK);
-	if (len != 1 || buf != 'b') {
+	len = recv(recv_sock, buf, 2, MSG_PEEK);
+	if (len != 1 || buf[0] != 'b') {
 		ksft_perror("Failed to read last byte of message\n");
 		goto out;
 	}
-	offset = tcp_peek_offset_get(recv_sock);
+	offset = sk_peek_offset_get(recv_sock);
 	if (offset != 2) {
 		ksft_perror("Offset not forwarded correctly at last byte\n");
 		goto out;
 	}
 	/* Flush message */
-	len = recv(recv_sock, NULL, 2, MSG_TRUNC);
+	len = recv(recv_sock, buf, 2, MSG_TRUNC);
 	if (len != 2) {
 		ksft_perror("Failed to flush message\n");
 		goto out;
 	}
-	offset = tcp_peek_offset_get(recv_sock);
+	offset = sk_peek_offset_get(recv_sock);
 	if (offset != 0) {
 		ksft_perror("Offset not reverted correctly after flush\n");
 		goto out;
 	}
 
-	printf("%s with MSG_PEEK_OFF works correctly\n", afstr(af));
+	printf("%s with MSG_PEEK_OFF works correctly\n", afstr(af, proto));
 	res = 1;
 out:
-	if (recv_sock >= 0)
+	if (proto == IPPROTO_TCP && recv_sock >= 0)
 		close(recv_sock);
 	if (s[1] >= 0)
 		close(s[1]);
@@ -160,24 +167,36 @@  static int tcp_peek_offset_test(sa_family_t af)
 	return res;
 }
 
-int main(void)
+static int do_test(int proto)
 {
 	int res4, res6;
 
-	res4 = tcp_peek_offset_probe(AF_INET);
-	res6 = tcp_peek_offset_probe(AF_INET6);
+	res4 = sk_peek_offset_probe(AF_INET, proto);
+	res6 = sk_peek_offset_probe(AF_INET6, proto);
 
 	if (!res4 && !res6)
 		return KSFT_SKIP;
 
 	if (res4)
-		res4 = tcp_peek_offset_test(AF_INET);
+		res4 = sk_peek_offset_test(AF_INET, proto);
 
 	if (res6)
-		res6 = tcp_peek_offset_test(AF_INET6);
+		res6 = sk_peek_offset_test(AF_INET6, proto);
 
 	if (!res4 || !res6)
 		return KSFT_FAIL;
 
 	return KSFT_PASS;
 }
+
+int main(void)
+{
+	int restcp, resudp;
+
+	restcp = do_test(IPPROTO_TCP);
+	resudp = do_test(IPPROTO_UDP);
+	if (restcp == KSFT_FAIL || resudp == KSFT_FAIL)
+		return KSFT_FAIL;
+
+	return KSFT_PASS;
+}