SGX can not read from /dev/urandom (#60368)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/60368
Problem:
The SGX secure enclave does not support reading from /dev/urandom as it is isolated from the OS for greater security. The SGX api provides a way to generate random numbers as a replacment.
Solution:
Conditionally enable SGX api for random number generation when building for it.
Test Plan: Run the PyTorch tests
Reviewed By: malfet, LiJihang
Differential Revision: D29022616
fbshipit-source-id: 1c7115457a2abde682df4d55fa4a8446fc5f8613
diff --git a/c10/core/GeneratorImpl.cpp b/c10/core/GeneratorImpl.cpp
index 7fb5571..720b1b2 100644
--- a/c10/core/GeneratorImpl.cpp
+++ b/c10/core/GeneratorImpl.cpp
@@ -2,6 +2,10 @@
#include <chrono>
#include <random>
+#if defined(__SGX_ENABLED__)
+#include <sgx_trts.h>
+#endif
+
#ifndef _WIN32
#include <fcntl.h>
#include <unistd.h>
@@ -57,7 +61,9 @@
/**
* Gets a non deterministic random number number from either the
* /dev/urandom or the current time. For CUDA, gets random from
- * std::random_device and adds a transformation on it.
+ * std::random_device and adds a transformation on it. For Intel SGX
+ * platform use sgx_read_rand as reading from /dev/urandom is
+ * prohibited on that platfrom.
*
* FIXME: The behavior in this function is from legacy code
* (THRandom_seed/THCRandom_seed) and is probably not the right thing to do,
@@ -76,6 +82,10 @@
s = (uint64_t)std::chrono::high_resolution_clock::now()
.time_since_epoch()
.count();
+#elif defined(__SGX_ENABLED__)
+ TORCH_CHECK(
+ sgx_read_rand(reinterpret_cast<uint8_t*>(&s), sizeof(s)) == SGX_SUCCESS,
+ "Could not generate random number with sgx_read_rand.");
#else
s = readURandomLong();
#endif