1 files changed, 45 insertions, 8 deletions

diff --git a/samples/sockmap/sockmap_kern.c b/samples/sockmap/sockmap_kern.c
index 205ec36449d1..735226794ce1 100644
--- a/samples/sockmap/sockmap_kern.c
+++ b/samples/sockmap/sockmap_kern.c
@@ -64,6 +64,13 @@ struct bpf_map_def SEC("maps") sock_apply_bytes = {
 	.max_entries = 1
 };
 
+struct bpf_map_def SEC("maps") sock_cork_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)
 {
@@ -135,6 +142,9 @@ int bpf_prog4(struct sk_msg_md *msg)
 	bytes = bpf_map_lookup_elem(&sock_apply_bytes, &zero);
 	if (bytes)
 		bpf_msg_apply_bytes(msg, *bytes);
+	bytes = bpf_map_lookup_elem(&sock_cork_bytes, &zero);
+	if (bytes)
+		bpf_msg_cork_bytes(msg, *bytes);
 	return SK_PASS;
 }
 
@@ -143,13 +153,16 @@ 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;
+	int *bytes, err1 = -1, err2 = -1, zero = 0;
 
 	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);
+		err1 = bpf_msg_apply_bytes(msg, *bytes);
+	bytes = bpf_map_lookup_elem(&sock_cork_bytes, &zero);
+	if (bytes)
+		err2 = bpf_msg_cork_bytes(msg, *bytes);
+	bpf_printk("sk_msg2: data length %i err1 %i err2 %i\n",
+		   (__u64)data_end - (__u64)data, err1, err2);
 	return SK_PASS;
 }
 
@@ -163,6 +176,9 @@ int bpf_prog6(struct sk_msg_md *msg)
 	bytes = bpf_map_lookup_elem(&sock_apply_bytes, &zero);
 	if (bytes)
 		bpf_msg_apply_bytes(msg, *bytes);
+	bytes = bpf_map_lookup_elem(&sock_cork_bytes, &zero);
+	if (bytes)
+		bpf_msg_cork_bytes(msg, *bytes);
 	return bpf_msg_redirect_map(msg, &sock_map_redir, zero, 0);
 }
 
@@ -171,13 +187,17 @@ int bpf_prog7(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;
+	int *bytes, err1 = 0, err2 = 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);
+		err1 = bpf_msg_apply_bytes(msg, *bytes);
+	bytes = bpf_map_lookup_elem(&sock_cork_bytes, &zero);
+	if (bytes)
+		err2 = bpf_msg_cork_bytes(msg, *bytes);
+
+	bpf_printk("sk_msg3: redirect(%iB) err1=%i err2=%i\n",
+		   (__u64)data_end - (__u64)data, err1, err2);
 	return bpf_msg_redirect_map(msg, &sock_map_redir, zero, 0);
 }
 
@@ -198,5 +218,22 @@ int bpf_prog8(struct sk_msg_md *msg)
 	}
 	return SK_PASS;
 }
+SEC("sk_msg6")
+int bpf_prog9(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_cork_bytes, &zero);
+	if (bytes) {
+		if (((__u64)data_end - (__u64)data) >= *bytes)
+			return SK_PASS;
+		ret = bpf_msg_cork_bytes(msg, *bytes);
+		if (ret)
+			return SK_DROP;
+	}
+	return SK_PASS;
+}
 
 char _license[] SEC("license") = "GPL";