Call SetShmBase() in init_rs

The ShmEventLoop rust test was using /dev/shm/aos and breaking out of
the sandbox. Fix that.

Change-Id: I14db8b500a03a66fef41399c4cea334ff79b217f
Signed-off-by: James Kuszmaul <james.kuszmaul@bluerivertech.com>
diff --git a/aos/BUILD b/aos/BUILD
index 34af003..f1cb6ba 100644
--- a/aos/BUILD
+++ b/aos/BUILD
@@ -198,6 +198,25 @@
     visibility = ["//visibility:public"],
 )
 
+autocxx_library(
+    name = "test_init_rs",
+    testonly = True,
+    srcs = ["test_init.rs"],
+    crate_name = "aos_test_init",
+    libs = [
+        "//aos/testing:tmpdir",
+    ],
+    override_cc_toolchain = "@llvm_toolchain//:cc-clang-x86_64-linux",
+    target_compatible_with = select({
+        "//conditions:default": ["//tools/platforms/rust:has_support"],
+        "//tools:has_msan": ["@platforms//:incompatible"],
+    }),
+    visibility = ["//visibility:public"],
+    deps = [
+        ":init_rs",
+    ],
+)
+
 cc_library(
     name = "realtime",
     srcs = [
diff --git a/aos/events/BUILD b/aos/events/BUILD
index a09c3d8..c398e7b 100644
--- a/aos/events/BUILD
+++ b/aos/events/BUILD
@@ -460,6 +460,7 @@
         "//aos:init",
         "//aos:realtime",
         "//aos/ipc_lib:lockless_queue",
+        "//aos/ipc_lib:shm_base",
         "//aos/ipc_lib:signalfd",
         "//aos/stl_mutex",
         "//aos/util:phased_loop",
@@ -622,7 +623,7 @@
     rustc_flags = ["-Crelocation-model=static"],
     deps = [
         ":ping_rust_fbs",
-        "//aos:init_rs",
+        "//aos:test_init_rs",
         "@crate_index//:futures",
         "@rules_rust//tools/runfiles",
     ],
@@ -672,7 +673,7 @@
     }),
     deps = [
         ":ping_rust_fbs",
-        "//aos:init_rs",
+        "//aos:test_init_rs",
         "@crate_index//:futures",
         "@rules_rust//tools/runfiles",
     ],
diff --git a/aos/events/shm_event_loop.cc b/aos/events/shm_event_loop.cc
index 0dc23d9..3126265 100644
--- a/aos/events/shm_event_loop.cc
+++ b/aos/events/shm_event_loop.cc
@@ -38,8 +38,6 @@
 
 }  // namespace
 
-DEFINE_string(shm_base, "/dev/shm/aos",
-              "Directory to place queue backing mmaped files in.");
 // This value is affected by the umask of the process which is calling it
 // and is set to the user's value by default (check yours running `umask` on
 // the command line).
@@ -59,10 +57,6 @@
 
 using namespace shm_event_loop_internal;
 
-void SetShmBase(const std::string_view base) {
-  FLAGS_shm_base = std::string(base) + "/aos";
-}
-
 namespace {
 
 const Node *MaybeMyNode(const Configuration *configuration) {
diff --git a/aos/events/shm_event_loop.h b/aos/events/shm_event_loop.h
index 7b7e68f..0e71f96 100644
--- a/aos/events/shm_event_loop.h
+++ b/aos/events/shm_event_loop.h
@@ -8,11 +8,11 @@
 #include "aos/events/epoll.h"
 #include "aos/events/event_loop.h"
 #include "aos/events/event_loop_generated.h"
+#include "aos/ipc_lib/shm_base.h"
 #include "aos/ipc_lib/signalfd.h"
 #include "aos/stl_mutex/stl_mutex.h"
 
 DECLARE_string(application_name);
-DECLARE_string(shm_base);
 
 namespace aos {
 namespace shm_event_loop_internal {
diff --git a/aos/events/shm_event_loop.rs b/aos/events/shm_event_loop.rs
index 3e387ef..880f72b 100644
--- a/aos/events/shm_event_loop.rs
+++ b/aos/events/shm_event_loop.rs
@@ -233,7 +233,7 @@
 
     use aos_configuration::read_config_from;
     use aos_events_event_loop_runtime::{Sender, Watcher};
-    use aos_init::test_init;
+    use aos_test_init::test_init;
     use ping_rust_fbs::aos::examples as ping;
     use std::sync::atomic::{AtomicUsize, Ordering};
     use std::sync::Barrier;
diff --git a/aos/events/simulated_event_loop.rs b/aos/events/simulated_event_loop.rs
index ea3a154..90c4c86 100644
--- a/aos/events/simulated_event_loop.rs
+++ b/aos/events/simulated_event_loop.rs
@@ -143,7 +143,7 @@
     use runfiles::Runfiles;
 
     use aos_configuration::read_config_from;
-    use aos_init::test_init;
+    use aos_test_init::test_init;
     use ping_rust_fbs::aos::examples::PingBuilder;
 
     // A really basic test of the functionality here.
diff --git a/aos/init.rs b/aos/init.rs
index 9ac62f1..13066f1 100644
--- a/aos/init.rs
+++ b/aos/init.rs
@@ -8,19 +8,7 @@
 generate!("aos::InitFromRust")
 );
 
-/// Initializes things for a test.
-///
-/// TODO(Brian): Should we provide a proc macro attribute that handles calling this?
-///
-/// # Panics
-///
-/// Panics if non-test initialization has already been performed.
-pub fn test_init() {
-    init();
-    // TODO(Brian): Do we want any of the other stuff that `:gtest_main` has?
-    // TODO(Brian): Call `aos::SetShmBase` like `:gtest_main` does.
-}
-
+/// Initializes AOS.
 pub fn init() {
     static ONCE: Once = Once::new();
     ONCE.call_once(|| {
diff --git a/aos/ipc_lib/BUILD b/aos/ipc_lib/BUILD
index 53d30cd..471b75d 100644
--- a/aos/ipc_lib/BUILD
+++ b/aos/ipc_lib/BUILD
@@ -404,3 +404,14 @@
         "shm_observers.h",
     ],
 )
+
+cc_library(
+    name = "shm_base",
+    srcs = ["shm_base.cc"],
+    hdrs = ["shm_base.h"],
+    target_compatible_with = ["@platforms//os:linux"],
+    visibility = ["//visibility:public"],
+    deps = [
+        "@com_github_gflags_gflags//:gflags",
+    ],
+)
diff --git a/aos/ipc_lib/shm_base.cc b/aos/ipc_lib/shm_base.cc
new file mode 100644
index 0000000..22bf915
--- /dev/null
+++ b/aos/ipc_lib/shm_base.cc
@@ -0,0 +1,9 @@
+#include "aos/ipc_lib/shm_base.h"
+
+DEFINE_string(shm_base, "/dev/shm/aos",
+              "Directory to place queue backing mmaped files in.");
+namespace aos::testing {
+void SetShmBase(const std::string_view base) {
+  FLAGS_shm_base = std::string(base) + "/aos";
+}
+}  // namespace aos::testing
diff --git a/aos/ipc_lib/shm_base.h b/aos/ipc_lib/shm_base.h
new file mode 100644
index 0000000..b70df06
--- /dev/null
+++ b/aos/ipc_lib/shm_base.h
@@ -0,0 +1,11 @@
+#ifndef AOS_IPC_LIB_SHM_BASE_H_
+#define AOS_IPC_LIB_SHM_BASE_H_
+
+#include "gflags/gflags.h"
+
+DECLARE_string(shm_base);
+
+namespace aos::testing {
+void SetShmBase(const std::string_view base);
+}
+#endif  // AOS_IPC_LIB_SHM_BASE_H_
diff --git a/aos/network/message_bridge_test_lib.cc b/aos/network/message_bridge_test_lib.cc
index d5a5f8f..3c08644 100644
--- a/aos/network/message_bridge_test_lib.cc
+++ b/aos/network/message_bridge_test_lib.cc
@@ -3,8 +3,6 @@
 DECLARE_string(boot_uuid);
 
 namespace aos {
-void SetShmBase(const std::string_view base);
-
 namespace message_bridge::testing {
 
 namespace chrono = std::chrono;
@@ -20,7 +18,7 @@
 }
 
 void DoSetShmBase(const std::string_view node) {
-  aos::SetShmBase(ShmBase(node));
+  aos::testing::SetShmBase(ShmBase(node));
 }
 
 ThreadedEventLoopRunner::ThreadedEventLoopRunner(aos::ShmEventLoop *event_loop)
diff --git a/aos/test_init.rs b/aos/test_init.rs
new file mode 100644
index 0000000..a8380e2
--- /dev/null
+++ b/aos/test_init.rs
@@ -0,0 +1,21 @@
+use aos_init::init;
+
+autocxx::include_cpp! (
+#include "aos/testing/tmpdir.h"
+
+safety!(unsafe)
+
+generate!("aos::testing::SetTestShmBase")
+);
+
+// TODO(Brian): Should we provide a proc macro attribute that handles calling this?
+/// Initializes things for a test.
+///
+/// # Panics
+///
+/// Panics if non-test initialization has already been performed.
+pub fn test_init() {
+    init();
+    ffi::aos::testing::SetTestShmBase();
+    // TODO(Brian): Do we want any of the other stuff that `:gtest_main` has?
+}
diff --git a/aos/testing/BUILD b/aos/testing/BUILD
index 787fdda..02459e4 100644
--- a/aos/testing/BUILD
+++ b/aos/testing/BUILD
@@ -8,6 +8,7 @@
     visibility = ["//visibility:public"],
     deps = [
         "//aos:init",
+        "//aos/testing:tmpdir",
         "@com_github_gflags_gflags//:gflags",
         "@com_github_google_glog//:glog",
         "@com_google_googletest//:gtest",
@@ -100,6 +101,9 @@
     hdrs = ["tmpdir.h"],
     target_compatible_with = ["@platforms//os:linux"],
     visibility = ["//visibility:public"],
+    deps = [
+        "//aos/ipc_lib:shm_base",
+    ],
 )
 
 cc_library(
diff --git a/aos/testing/gtest_main.cc b/aos/testing/gtest_main.cc
index 21f141f..2fb77fe 100644
--- a/aos/testing/gtest_main.cc
+++ b/aos/testing/gtest_main.cc
@@ -7,6 +7,7 @@
 #include "gtest/gtest.h"
 
 #include "aos/init.h"
+#include "aos/testing/tmpdir.h"
 
 DEFINE_bool(print_logs, false,
             "Print the log messages as they are being generated.");
@@ -14,7 +15,6 @@
               "Print all log messages to FILE instead of standard output.");
 
 namespace aos {
-void SetShmBase(const std::string_view base) __attribute__((weak));
 
 namespace testing {
 
@@ -48,12 +48,7 @@
 
   // Point shared memory away from /dev/shm if we are testing.  We don't care
   // about RT in this case, so if it is backed by disk, we are fine.
-  if (::aos::SetShmBase) {
-    const char *tmpdir_c_str = getenv("TEST_TMPDIR");
-    if (tmpdir_c_str != nullptr) {
-      aos::SetShmBase(tmpdir_c_str);
-    }
-  }
+  aos::testing::SetTestShmBase();
 
   return RUN_ALL_TESTS();
 }
diff --git a/aos/testing/test_shm.cc b/aos/testing/test_shm.cc
index 0787e5a..8dc189a 100644
--- a/aos/testing/test_shm.cc
+++ b/aos/testing/test_shm.cc
@@ -2,6 +2,8 @@
 
 #include <sys/mman.h>
 
+#include "gflags/gflags.h"
+
 #include "aos/logging/logging.h"
 #include "aos/testing/test_logging.h"
 
diff --git a/aos/testing/tmpdir.cc b/aos/testing/tmpdir.cc
index 15e3c13..7287603 100644
--- a/aos/testing/tmpdir.cc
+++ b/aos/testing/tmpdir.cc
@@ -3,16 +3,24 @@
 #include <cstdlib>
 #include <string>
 
+#include "aos/ipc_lib/shm_base.h"
+
 namespace aos {
 namespace testing {
 
-std::string TestTmpDir() {
+namespace {
+std::string TestTmpDirOr(std::string fallback) {
   const char *tmp_dir = std::getenv("TEST_TMPDIR");
   if (tmp_dir != nullptr) {
     return tmp_dir;
   }
-  return "/tmp";
+  return fallback;
 }
+}  // namespace
+
+std::string TestTmpDir() { return TestTmpDirOr("/tmp"); }
+
+void SetTestShmBase() { SetShmBase(TestTmpDirOr(FLAGS_shm_base)); }
 
 }  // namespace testing
 }  // namespace aos
diff --git a/aos/testing/tmpdir.h b/aos/testing/tmpdir.h
index 7e64342..8eabf86 100644
--- a/aos/testing/tmpdir.h
+++ b/aos/testing/tmpdir.h
@@ -9,6 +9,10 @@
 // Returns a usable temporary directory.
 std::string TestTmpDir();
 
+// Sets shm_base to a folder inside of TEST_TMPDIR if set, or --shm_base
+// otherwise.
+void SetTestShmBase();
+
 }  // namespace testing
 }  // namespace aos