2 files changed, 66 insertions, 7 deletions

diff --git a/samples/sockmap/sockmap_kern.c b/samples/sockmap/sockmap_kern.c
index 75edb2f151d4..205ec36449d1 100644
--- a/samples/sockmap/sockmap_kern.c
+++ b/samples/sockmap/sockmap_kern.c
@@ -57,6 +57,13 @@ struct bpf_map_def SEC("maps") sock_map_redir = {
 	.max_entries = 1,
 };
 
+struct bpf_map_def SEC("maps") sock_apply_bytes = {
+	.type = BPF_MAP_TYPE_ARRAY,
+	.key_size = sizeof(int),
+	.value_size = sizeof(int),
+	.max_entries = 1
+};
+
 SEC("sk_skb1")
 int bpf_prog1(struct __sk_buff *skb)
 {
@@ -123,6 +130,11 @@ int bpf_sockmap(struct bpf_sock_ops *skops)
 SEC("sk_msg1")
 int bpf_prog4(struct sk_msg_md *msg)
 {
+	int *bytes, zero = 0;
+
+	bytes = bpf_map_lookup_elem(&sock_apply_bytes, &zero);
+	if (bytes)
+		bpf_msg_apply_bytes(msg, *bytes);
 	return SK_PASS;
 }
 
@@ -131,8 +143,13 @@ int bpf_prog5(struct sk_msg_md *msg)
 {
 	void *data_end = (void *)(long) msg->data_end;
 	void *data = (void *)(long) msg->data;
+	int *bytes, err = 0, zero = 0;
 
-	bpf_printk("sk_msg2: data length %i\n", (__u32)data_end - (__u32)data);
+	bytes = bpf_map_lookup_elem(&sock_apply_bytes, &zero);
+	if (bytes)
+		err = bpf_msg_apply_bytes(msg, *bytes);
+	bpf_printk("sk_msg2: data length %i err %i\n",
+		   (__u64)data_end - (__u64)data, err);
 	return SK_PASS;
 }
 
@@ -141,9 +158,12 @@ int bpf_prog6(struct sk_msg_md *msg)
 {
 	void *data_end = (void *)(long) msg->data_end;
 	void *data = (void *)(long) msg->data;
-	int ret = 0;
+	int *bytes, zero = 0;
 
-	return bpf_msg_redirect_map(msg, &sock_map_redir, ret, 0);
+	bytes = bpf_map_lookup_elem(&sock_apply_bytes, &zero);
+	if (bytes)
+		bpf_msg_apply_bytes(msg, *bytes);
+	return bpf_msg_redirect_map(msg, &sock_map_redir, zero, 0);
 }
 
 SEC("sk_msg4")
@@ -151,10 +171,32 @@ int bpf_prog7(struct sk_msg_md *msg)
 {
 	void *data_end = (void *)(long) msg->data_end;
 	void *data = (void *)(long) msg->data;
-	int ret = 0;
+	int *bytes, err = 0, zero = 0;
+
+	bytes = bpf_map_lookup_elem(&sock_apply_bytes, &zero);
+	if (bytes)
+		err = bpf_msg_apply_bytes(msg, *bytes);
+	bpf_printk("sk_msg3: redirect(%iB) err=%i\n",
+		   (__u64)data_end - (__u64)data, err);
+	return bpf_msg_redirect_map(msg, &sock_map_redir, zero, 0);
+}
 
-	bpf_printk("sk_msg3: redirect(%iB)\n", (__u32)data_end - (__u32)data);
-	return bpf_msg_redirect_map(msg, &sock_map_redir, ret, 0);
+SEC("sk_msg5")
+int bpf_prog8(struct sk_msg_md *msg)
+{
+	void *data_end = (void *)(long) msg->data_end;
+	void *data = (void *)(long) msg->data;
+	int ret = 0, *bytes, zero = 0;
+
+	bytes = bpf_map_lookup_elem(&sock_apply_bytes, &zero);
+	if (bytes) {
+		ret = bpf_msg_apply_bytes(msg, *bytes);
+		if (ret)
+			return SK_DROP;
+	} else {
+		return SK_DROP;
+	}
+	return SK_PASS;
 }
 
 char _license[] SEC("license") = "GPL";
diff --git a/samples/sockmap/sockmap_user.c b/samples/sockmap/sockmap_user.c
index 8017ad7afea9..41774ec0f0b3 100644
--- a/samples/sockmap/sockmap_user.c
+++ b/samples/sockmap/sockmap_user.c
@@ -59,6 +59,7 @@ int txmsg_pass;
 int txmsg_noisy;
 int txmsg_redir;
 int txmsg_redir_noisy;
+int txmsg_apply;
 
 static const struct option long_options[] = {
 	{"help",	no_argument,		NULL, 'h' },
@@ -73,6 +74,7 @@ static const struct option long_options[] = {
 	{"txmsg_noisy",		no_argument,	&txmsg_noisy, 1  },
 	{"txmsg_redir",		no_argument,	&txmsg_redir, 1  },
 	{"txmsg_redir_noisy",	no_argument,	&txmsg_redir_noisy, 1},
+	{"txmsg_apply",	required_argument,	NULL, 'a'},
 	{0, 0, NULL, 0 }
 };
 
@@ -546,7 +548,9 @@ int main(int argc, char **argv)
 	while ((opt = getopt_long(argc, argv, ":dhvc:r:i:l:t:",
 				  long_options, &longindex)) != -1) {
 		switch (opt) {
-		/* Cgroup configuration */
+		case 'a':
+			txmsg_apply = atoi(optarg);
+			break;
 		case 'c':
 			cg_fd = open(optarg, O_DIRECTORY, O_RDONLY);
 			if (cg_fd < 0) {
@@ -665,6 +669,8 @@ run:
 		tx_prog_fd = prog_fd[5];
 	else if (txmsg_redir_noisy)
 		tx_prog_fd = prog_fd[6];
+	else if (txmsg_apply)
+		tx_prog_fd = prog_fd[7];
 	else
 		tx_prog_fd = 0;
 
@@ -699,6 +705,17 @@ run:
 				err, strerror(errno));
 			return err;
 		}
+
+		if (txmsg_apply) {
+			err = bpf_map_update_elem(map_fd[3],
+						  &i, &txmsg_apply, BPF_ANY);
+			if (err) {
+				fprintf(stderr,
+					"ERROR: bpf_map_update_elem (apply_bytes):  %d (%s\n",
+					err, strerror(errno));
+				return err;
+			}
+		}
 	}
 	if (test == PING_PONG)
 		err = forever_ping_pong(rate, &options);