test_map_in_map_kern.c 3.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176
  1. /*
  2. * Copyright (c) 2017 Facebook
  3. *
  4. * This program is free software; you can redistribute it and/or
  5. * modify it under the terms of version 2 of the GNU General Public
  6. * License as published by the Free Software Foundation.
  7. */
  8. #define KBUILD_MODNAME "foo"
  9. #include <linux/ptrace.h>
  10. #include <linux/version.h>
  11. #include <uapi/linux/bpf.h>
  12. #include <uapi/linux/in6.h>
  13. #include <bpf/bpf_helpers.h>
  14. #include <bpf/bpf_tracing.h>
  15. #include <bpf/bpf_core_read.h>
  16. #include "trace_common.h"
  17. #define MAX_NR_PORTS 65536
  18. /* map #0 */
  19. struct inner_a {
  20. __uint(type, BPF_MAP_TYPE_ARRAY);
  21. __type(key, u32);
  22. __type(value, int);
  23. __uint(max_entries, MAX_NR_PORTS);
  24. } port_a SEC(".maps");
  25. /* map #1 */
  26. struct inner_h {
  27. __uint(type, BPF_MAP_TYPE_HASH);
  28. __type(key, u32);
  29. __type(value, int);
  30. __uint(max_entries, 1);
  31. } port_h SEC(".maps");
  32. /* map #2 */
  33. struct {
  34. __uint(type, BPF_MAP_TYPE_HASH);
  35. __type(key, u32);
  36. __type(value, int);
  37. __uint(max_entries, 1);
  38. } reg_result_h SEC(".maps");
  39. /* map #3 */
  40. struct {
  41. __uint(type, BPF_MAP_TYPE_HASH);
  42. __type(key, u32);
  43. __type(value, int);
  44. __uint(max_entries, 1);
  45. } inline_result_h SEC(".maps");
  46. /* map #4 */ /* Test case #0 */
  47. struct {
  48. __uint(type, BPF_MAP_TYPE_ARRAY_OF_MAPS);
  49. __uint(max_entries, MAX_NR_PORTS);
  50. __uint(key_size, sizeof(u32));
  51. __array(values, struct inner_a); /* use inner_a as inner map */
  52. } a_of_port_a SEC(".maps");
  53. /* map #5 */ /* Test case #1 */
  54. struct {
  55. __uint(type, BPF_MAP_TYPE_HASH_OF_MAPS);
  56. __uint(max_entries, 1);
  57. __uint(key_size, sizeof(u32));
  58. __array(values, struct inner_a); /* use inner_a as inner map */
  59. } h_of_port_a SEC(".maps");
  60. /* map #6 */ /* Test case #2 */
  61. struct {
  62. __uint(type, BPF_MAP_TYPE_HASH_OF_MAPS);
  63. __uint(max_entries, 1);
  64. __uint(key_size, sizeof(u32));
  65. __array(values, struct inner_h); /* use inner_h as inner map */
  66. } h_of_port_h SEC(".maps");
  67. static __always_inline int do_reg_lookup(void *inner_map, u32 port)
  68. {
  69. int *result;
  70. result = bpf_map_lookup_elem(inner_map, &port);
  71. return result ? *result : -ENOENT;
  72. }
  73. static __always_inline int do_inline_array_lookup(void *inner_map, u32 port)
  74. {
  75. int *result;
  76. if (inner_map != &port_a)
  77. return -EINVAL;
  78. result = bpf_map_lookup_elem(&port_a, &port);
  79. return result ? *result : -ENOENT;
  80. }
  81. static __always_inline int do_inline_hash_lookup(void *inner_map, u32 port)
  82. {
  83. int *result;
  84. if (inner_map != &port_h)
  85. return -EINVAL;
  86. result = bpf_map_lookup_elem(&port_h, &port);
  87. return result ? *result : -ENOENT;
  88. }
  89. SEC("kprobe/__sys_connect")
  90. int trace_sys_connect(struct pt_regs *ctx)
  91. {
  92. struct sockaddr_in6 *in6;
  93. u16 test_case, port, dst6[8];
  94. int addrlen, ret, inline_ret, ret_key = 0;
  95. u32 port_key;
  96. void *outer_map, *inner_map;
  97. bool inline_hash = false;
  98. in6 = (struct sockaddr_in6 *)PT_REGS_PARM2_CORE(ctx);
  99. addrlen = (int)PT_REGS_PARM3_CORE(ctx);
  100. if (addrlen != sizeof(*in6))
  101. return 0;
  102. ret = bpf_probe_read_user(dst6, sizeof(dst6), &in6->sin6_addr);
  103. if (ret) {
  104. inline_ret = ret;
  105. goto done;
  106. }
  107. if (dst6[0] != 0xdead || dst6[1] != 0xbeef)
  108. return 0;
  109. test_case = dst6[7];
  110. ret = bpf_probe_read_user(&port, sizeof(port), &in6->sin6_port);
  111. if (ret) {
  112. inline_ret = ret;
  113. goto done;
  114. }
  115. port_key = port;
  116. ret = -ENOENT;
  117. if (test_case == 0) {
  118. outer_map = &a_of_port_a;
  119. } else if (test_case == 1) {
  120. outer_map = &h_of_port_a;
  121. } else if (test_case == 2) {
  122. outer_map = &h_of_port_h;
  123. } else {
  124. ret = __LINE__;
  125. inline_ret = ret;
  126. goto done;
  127. }
  128. inner_map = bpf_map_lookup_elem(outer_map, &port_key);
  129. if (!inner_map) {
  130. ret = __LINE__;
  131. inline_ret = ret;
  132. goto done;
  133. }
  134. ret = do_reg_lookup(inner_map, port_key);
  135. if (test_case == 0 || test_case == 1)
  136. inline_ret = do_inline_array_lookup(inner_map, port_key);
  137. else
  138. inline_ret = do_inline_hash_lookup(inner_map, port_key);
  139. done:
  140. bpf_map_update_elem(&reg_result_h, &ret_key, &ret, BPF_ANY);
  141. bpf_map_update_elem(&inline_result_h, &ret_key, &inline_ret, BPF_ANY);
  142. return 0;
  143. }
  144. char _license[] SEC("license") = "GPL";
  145. u32 _version SEC("version") = LINUX_VERSION_CODE;