scoped_set_env.cc 2.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081
  1. // Copyright 2019 The Abseil Authors.
  2. //
  3. // Licensed under the Apache License, Version 2.0 (the "License");
  4. // you may not use this file except in compliance with the License.
  5. // You may obtain a copy of the License at
  6. //
  7. // https://www.apache.org/licenses/LICENSE-2.0
  8. //
  9. // Unless required by applicable law or agreed to in writing, software
  10. // distributed under the License is distributed on an "AS IS" BASIS,
  11. // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. // See the License for the specific language governing permissions and
  13. // limitations under the License.
  14. #include "absl/base/internal/scoped_set_env.h"
  15. #ifdef _WIN32
  16. #include <windows.h>
  17. #endif
  18. #include <cstdlib>
  19. #include "absl/base/internal/raw_logging.h"
  20. namespace absl {
  21. ABSL_NAMESPACE_BEGIN
  22. namespace base_internal {
  23. namespace {
  24. #ifdef _WIN32
  25. const int kMaxEnvVarValueSize = 1024;
  26. #endif
  27. void SetEnvVar(const char* name, const char* value) {
  28. #ifdef _WIN32
  29. SetEnvironmentVariableA(name, value);
  30. #else
  31. if (value == nullptr) {
  32. ::unsetenv(name);
  33. } else {
  34. ::setenv(name, value, 1);
  35. }
  36. #endif
  37. }
  38. } // namespace
  39. ScopedSetEnv::ScopedSetEnv(const char* var_name, const char* new_value)
  40. : var_name_(var_name), was_unset_(false) {
  41. #ifdef _WIN32
  42. char buf[kMaxEnvVarValueSize];
  43. auto get_res = GetEnvironmentVariableA(var_name_.c_str(), buf, sizeof(buf));
  44. ABSL_INTERNAL_CHECK(get_res < sizeof(buf), "value exceeds buffer size");
  45. if (get_res == 0) {
  46. was_unset_ = (GetLastError() == ERROR_ENVVAR_NOT_FOUND);
  47. } else {
  48. old_value_.assign(buf, get_res);
  49. }
  50. SetEnvironmentVariableA(var_name_.c_str(), new_value);
  51. #else
  52. const char* val = ::getenv(var_name_.c_str());
  53. if (val == nullptr) {
  54. was_unset_ = true;
  55. } else {
  56. old_value_ = val;
  57. }
  58. #endif
  59. SetEnvVar(var_name_.c_str(), new_value);
  60. }
  61. ScopedSetEnv::~ScopedSetEnv() {
  62. SetEnvVar(var_name_.c_str(), was_unset_ ? nullptr : old_value_.c_str());
  63. }
  64. } // namespace base_internal
  65. ABSL_NAMESPACE_END
  66. } // namespace absl