#include <linux/bpf.h>
 
 #include <bpf/bpf.h>
+#include <bpf/libbpf.h>
 #include "bpf_util.h"
 
 static int map_flags;
        close(fd);
 }
 
+#include <sys/socket.h>
+#include <sys/ioctl.h>
+#include <arpa/inet.h>
+#include <sys/select.h>
+#include <linux/err.h>
+#define SOCKMAP_PARSE_PROG "./sockmap_parse_prog.o"
+#define SOCKMAP_VERDICT_PROG "./sockmap_verdict_prog.o"
+static void test_sockmap(int task, void *data)
+{
+       int ports[] = {50200, 50201, 50202, 50204};
+       int err, i, fd, sfd[6] = {0xdeadbeef};
+       char buf[] = "hello sockmap user\n";
+       int one = 1, map_fd, s, sc, rc;
+       int parse_prog, verdict_prog;
+       struct bpf_map *bpf_map;
+       struct sockaddr_in addr;
+       struct bpf_object *obj;
+       struct timeval to;
+       __u32 key, value;
+       fd_set w;
+
+       /* Create some sockets to use with sockmap */
+       for (i = 0; i < 2; i++) {
+               sfd[i] = socket(AF_INET, SOCK_STREAM, 0);
+               if (sfd[i] < 0)
+                       goto out;
+               err = setsockopt(sfd[i], SOL_SOCKET, SO_REUSEADDR,
+                                (char *)&one, sizeof(one));
+               if (err) {
+                       printf("failed to setsockopt\n");
+                       goto out;
+               }
+               err = ioctl(sfd[i], FIONBIO, (char *)&one);
+               if (err < 0) {
+                       printf("failed to ioctl\n");
+                       goto out;
+               }
+               memset(&addr, 0, sizeof(struct sockaddr_in));
+               addr.sin_family = AF_INET;
+               addr.sin_addr.s_addr = inet_addr("127.0.0.1");
+               addr.sin_port = htons(ports[i]);
+               err = bind(sfd[i], (struct sockaddr *)&addr, sizeof(addr));
+               if (err < 0) {
+                       printf("failed to bind: err %i: %i:%i\n",
+                              err, i, sfd[i]);
+                       goto out;
+               }
+               err = listen(sfd[i], 32);
+               if (err < 0) {
+                       printf("failed to listeen\n");
+                       goto out;
+               }
+       }
+
+       for (i = 2; i < 4; i++) {
+               sfd[i] = socket(AF_INET, SOCK_STREAM, 0);
+               if (sfd[i] < 0)
+                       goto out;
+               err = setsockopt(sfd[i], SOL_SOCKET, SO_REUSEADDR,
+                                (char *)&one, sizeof(one));
+               if (err) {
+                       printf("set sock opt\n");
+                       goto out;
+               }
+               memset(&addr, 0, sizeof(struct sockaddr_in));
+               addr.sin_family = AF_INET;
+               addr.sin_addr.s_addr = inet_addr("127.0.0.1");
+               addr.sin_port = htons(ports[i - 2]);
+               err = connect(sfd[i], (struct sockaddr *)&addr, sizeof(addr));
+               if (err) {
+                       printf("failed to conenct\n");
+                       goto out;
+               }
+       }
+
+
+       for (i = 4; i < 6; i++) {
+               sfd[i] = accept(sfd[i - 4], NULL, NULL);
+               if (sfd[i] < 0) {
+                       printf("accept failed\n");
+                       goto out;
+               }
+       }
+
+       /* Test sockmap with connected sockets */
+       fd = bpf_create_map(BPF_MAP_TYPE_SOCKMAP,
+                           sizeof(key), sizeof(value),
+                           6, 0);
+       if (fd < 0) {
+               printf("Failed to create sockmap %i\n", fd);
+               goto out_sockmap;
+       }
+
+       /* Nothing attached so these should fail */
+       for (i = 0; i < 6; i++) {
+               err = bpf_map_update_elem(fd, &i, &sfd[i], BPF_ANY);
+               if (!err) {
+                       printf("Failed invalid update sockmap '%i:%i'\n",
+                              i, sfd[i]);
+                       goto out_sockmap;
+               }
+       }
+
+       /* Test attaching bad fds */
+       err = __bpf_prog_attach(-1, -2, fd, BPF_CGROUP_SMAP_INGRESS, 0);
+       if (!err) {
+               printf("Failed invalid prog attach\n");
+               goto out_sockmap;
+       }
+
+       /* Load SK_SKB program and Attach */
+       err = bpf_prog_load(SOCKMAP_PARSE_PROG,
+                           BPF_PROG_TYPE_SK_SKB, &obj, &parse_prog);
+       if (err) {
+               printf("Failed to load SK_SKB parse prog\n");
+               goto out_sockmap;
+       }
+
+       err = bpf_prog_load(SOCKMAP_VERDICT_PROG,
+                           BPF_PROG_TYPE_SK_SKB, &obj, &verdict_prog);
+       if (err) {
+               printf("Failed to load SK_SKB verdict prog\n");
+               goto out_sockmap;
+       }
+
+       bpf_map = bpf_object__find_map_by_name(obj, "sock_map");
+       if (IS_ERR(bpf_map)) {
+               printf("Failed to load map from verdict prog\n");
+               goto out_sockmap;
+       }
+
+       map_fd = bpf_map__fd(bpf_map);
+       if (map_fd < 0) {
+               printf("Failed to get map fd\n");
+               goto out_sockmap;
+       }
+
+       err = __bpf_prog_attach(parse_prog, verdict_prog, map_fd,
+                               BPF_CGROUP_SMAP_INGRESS, 0);
+       if (err) {
+               printf("Failed bpf prog attach\n");
+               goto out_sockmap;
+       }
+
+       /* Test map update elem */
+       for (i = 0; i < 6; i++) {
+               err = bpf_map_update_elem(map_fd, &i, &sfd[i], BPF_ANY);
+               if (err) {
+                       printf("Failed map_fd update sockmap %i '%i:%i'\n",
+                              err, i, sfd[i]);
+                       goto out_sockmap;
+               }
+       }
+
+       /* Test map delete elem and remove send/recv sockets */
+       for (i = 2; i < 4; i++) {
+               err = bpf_map_delete_elem(map_fd, &i);
+               if (err) {
+                       printf("Failed delete  sockmap %i '%i:%i'\n",
+                              err, i, sfd[i]);
+                       goto out_sockmap;
+               }
+       }
+
+       /* Test map send/recv */
+       sc = send(sfd[2], buf, 10, 0);
+       if (sc < 0) {
+               printf("Failed sockmap send\n");
+               goto out_sockmap;
+       }
+
+       FD_ZERO(&w);
+       FD_SET(sfd[3], &w);
+       to.tv_sec = 1;
+       to.tv_usec = 0;
+       s = select(sfd[3] + 1, &w, NULL, NULL, &to);
+       if (s == -1) {
+               perror("Failed sockmap select()");
+               goto out_sockmap;
+       } else if (!s) {
+               printf("Failed sockmap unexpected timeout\n");
+               goto out_sockmap;
+       }
+
+       if (!FD_ISSET(sfd[3], &w)) {
+               printf("Failed sockmap select/recv\n");
+               goto out_sockmap;
+       }
+
+       rc = recv(sfd[3], buf, sizeof(buf), 0);
+       if (rc < 0) {
+               printf("Failed sockmap recv\n");
+               goto out_sockmap;
+       }
+
+       /* Delete the reset of the elems include some NULL elems */
+       for (i = 0; i < 6; i++) {
+               err = bpf_map_delete_elem(map_fd, &i);
+               if (err && (i == 0 || i == 1 || i >= 4)) {
+                       printf("Failed delete  sockmap %i '%i:%i'\n",
+                              err, i, sfd[i]);
+                       goto out_sockmap;
+               } else if (!err && (i == 2 || i == 3)) {
+                       printf("Failed null delete sockmap %i '%i:%i'\n",
+                              err, i, sfd[i]);
+                       goto out_sockmap;
+               }
+       }
+
+       /* Test having multiple SMAPs open and active on same fds */
+       err = __bpf_prog_attach(parse_prog, verdict_prog, fd,
+                               BPF_CGROUP_SMAP_INGRESS, 0);
+       if (err) {
+               printf("Failed fd bpf prog attach\n");
+               goto out_sockmap;
+       }
+
+       for (i = 0; i < 6; i++) {
+               err = bpf_map_update_elem(fd, &i, &sfd[i], BPF_ANY);
+               if (err) {
+                       printf("Failed fd update sockmap %i '%i:%i'\n",
+                              err, i, sfd[i]);
+                       goto out_sockmap;
+               }
+       }
+
+       /* Test duplicate socket add of NOEXIST, ANY and EXIST */
+       i = 0;
+       err = bpf_map_update_elem(fd, &i, &sfd[i], BPF_NOEXIST);
+       if (!err) {
+               printf("Failed BPF_NOEXIST create\n");
+               goto out_sockmap;
+       }
+
+       err = bpf_map_update_elem(fd, &i, &sfd[i], BPF_ANY);
+       if (err) {
+               printf("Failed sockmap update BPF_ANY\n");
+               goto out_sockmap;
+       }
+
+       err = bpf_map_update_elem(fd, &i, &sfd[i], BPF_EXIST);
+       if (err) {
+               printf("Failed sockmap update BPF_EXIST\n");
+               goto out_sockmap;
+       }
+
+       /* The above were pushing fd into same slot try different slot now */
+       i = 2;
+       err = bpf_map_update_elem(fd, &i, &sfd[i], BPF_NOEXIST);
+       if (!err) {
+               printf("Failed BPF_NOEXIST create\n");
+               goto out_sockmap;
+       }
+
+       err = bpf_map_update_elem(fd, &i, &sfd[i], BPF_ANY);
+       if (err) {
+               printf("Failed sockmap update BPF_ANY\n");
+               goto out_sockmap;
+       }
+
+       err = bpf_map_update_elem(fd, &i, &sfd[i], BPF_EXIST);
+       if (err) {
+               printf("Failed sockmap update BPF_EXIST\n");
+               goto out_sockmap;
+       }
+
+       /* Try pushing fd into different map, this is not allowed at the
+        * moment. Which programs would we use?
+        */
+       err = bpf_map_update_elem(map_fd, &i, &sfd[i], BPF_NOEXIST);
+       if (!err) {
+               printf("Failed BPF_NOEXIST create\n");
+               goto out_sockmap;
+       }
+
+       err = bpf_map_update_elem(map_fd, &i, &sfd[i], BPF_ANY);
+       if (!err) {
+               printf("Failed sockmap update BPF_ANY\n");
+               goto out_sockmap;
+       }
+
+       err = bpf_map_update_elem(map_fd, &i, &sfd[i], BPF_EXIST);
+       if (!err) {
+               printf("Failed sockmap update BPF_EXIST\n");
+               goto out_sockmap;
+       }
+
+       /* Test map close sockets */
+       for (i = 0; i < 6; i++)
+               close(sfd[i]);
+       close(fd);
+       close(map_fd);
+       bpf_object__close(obj);
+       return;
+out:
+       for (i = 0; i < 6; i++)
+               close(sfd[i]);
+       printf("Failed to create sockmap '%i:%s'!\n", i, strerror(errno));
+       exit(1);
+out_sockmap:
+       for (i = 0; i < 6; i++)
+               close(sfd[i]);
+       close(fd);
+       exit(1);
+}
+
 #define MAP_SIZE (32 * 1024)
 
 static void test_map_large(void)
        test_arraymap_percpu_many_keys();
 
        test_devmap(0, NULL);
+       test_sockmap(0, NULL);
 
        test_map_large();
        test_map_parallel();
 
        __ret;                                                          \
 })
 
-static int bpf_prog_load(const char *file, enum bpf_prog_type type,
-                        struct bpf_object **pobj, int *prog_fd)
-{
-       struct bpf_program *prog;
-       struct bpf_object *obj;
-       int err;
-
-       obj = bpf_object__open(file);
-       if (IS_ERR(obj)) {
-               error_cnt++;
-               return -ENOENT;
-       }
-
-       prog = bpf_program__next(NULL, obj);
-       if (!prog) {
-               bpf_object__close(obj);
-               error_cnt++;
-               return -ENOENT;
-       }
-
-       bpf_program__set_type(prog, type);
-       err = bpf_object__load(obj);
-       if (err) {
-               bpf_object__close(obj);
-               error_cnt++;
-               return -EINVAL;
-       }
-
-       *pobj = obj;
-       *prog_fd = bpf_program__fd(prog);
-       return 0;
-}
-
 static int bpf_find_map(const char *test, struct bpf_object *obj,
                        const char *name)
 {
        int err, prog_fd;
 
        err = bpf_prog_load(file, BPF_PROG_TYPE_SCHED_CLS, &obj, &prog_fd);
-       if (err)
+       if (err) {
+               error_cnt++;
                return;
+       }
 
        err = bpf_prog_test_run(prog_fd, 100000, &pkt_v4, sizeof(pkt_v4),
                                NULL, NULL, &retval, &duration);
        int err, prog_fd, map_fd;
 
        err = bpf_prog_load(file, BPF_PROG_TYPE_XDP, &obj, &prog_fd);
-       if (err)
+       if (err) {
+               error_cnt++;
                return;
+       }
 
        map_fd = bpf_find_map(__func__, obj, "vip2tnl");
        if (map_fd < 0)
        u32 *magic = (u32 *)buf;
 
        err = bpf_prog_load(file, BPF_PROG_TYPE_SCHED_CLS, &obj, &prog_fd);
-       if (err)
+       if (err) {
+               error_cnt++;
                return;
+       }
 
        map_fd = bpf_find_map(__func__, obj, "vip_map");
        if (map_fd < 0)
 
        err = bpf_prog_load(file, BPF_PROG_TYPE_TRACEPOINT, &obj, &prog_fd);
        CHECK(err, "", "err %d errno %d\n", err, errno);
-       if (err)
+       if (err) {
+               error_cnt++;
                return;
+       }
 
        bpf_object__close(obj);
 }
                /* test_obj_id.o is a dumb prog. It should never fail
                 * to load.
                 */
+               if (err)
+                       error_cnt++;
                assert(!err);
 
                /* Check getting prog info */
        int err, prog_fd;
 
        err = bpf_prog_load(file, BPF_PROG_TYPE_SCHED_CLS, &obj, &prog_fd);
-       if (err)
+       if (err) {
+               error_cnt++;
                return;
+       }
 
        err = bpf_prog_test_run(prog_fd, 10, &pkt_v4, sizeof(pkt_v4),
                                NULL, NULL, &retval, &duration);