267 lines
7.2 KiB
C++
267 lines
7.2 KiB
C++
// Copyright Epic Games, Inc. All Rights Reserved.
|
|
|
|
#pragma once
|
|
|
|
#include "UbaFileAccessor.h"
|
|
#include "UbaNetworkBackend.h"
|
|
#include "UbaStringBuffer.h"
|
|
|
|
#define UBA_USE_CLOUD !PLATFORM_MAC
|
|
|
|
#if UBA_USE_CLOUD
|
|
|
|
namespace uba
|
|
{
|
|
class Cloud
|
|
{
|
|
public:
|
|
struct Provider
|
|
{
|
|
const tchar* name;
|
|
const char* tokenCommand;
|
|
const char* tokenImds;
|
|
const char* tokenHeader;
|
|
const tchar* tokenPrefix;
|
|
const char* instanceIdImds;
|
|
const char* instanceLifeCycleImds;
|
|
const char* autoScalingLifeCycleStateImds;
|
|
const char* availabilityZoneImds;
|
|
const char* maintenanceEventImds;
|
|
};
|
|
|
|
static constexpr Provider g_providers[]
|
|
{
|
|
{
|
|
TC("AWS"),
|
|
"PUT",
|
|
"latest/api/token",
|
|
"X-aws-ec2-metadata-token-ttl-seconds: 21600\r\n",
|
|
TC("X-aws-ec2-metadata-token: "),
|
|
"latest/meta-data/instance-id",
|
|
"latest/meta-data/instance-life-cycle",
|
|
"latest/meta-data/autoscaling/target-lifecycle-state",
|
|
"latest/meta-data/placement/availability-zone",
|
|
nullptr, // maintenance
|
|
},
|
|
{
|
|
TC("GCP"),
|
|
"POST",
|
|
"computeMetadata/v1/instance/service-accounts/default/identity?audience=https://example.com",
|
|
"Accept: */*\r\nMetadata-Flavor: Google\r\nContent-Length: 0\r\n",
|
|
TC("Metadata-Flavor: Google\r\nAuthorization: Bearer "),
|
|
"computeMetadata/v1/instance/id",
|
|
nullptr, // Life cycle
|
|
nullptr, // Autoscaling
|
|
"computeMetadata/v1/instance/zone",
|
|
"computeMetadata/v1/instance/maintenance-event",
|
|
}
|
|
};
|
|
|
|
|
|
static constexpr char g_imdsHost[] = "169.254.169.254";
|
|
static constexpr char g_imdsAutoScalingLifeCycleState[] = "latest/meta-data/autoscaling/target-lifecycle-state";
|
|
static constexpr char g_imdsSpotInstanceAction[] = "latest/meta-data/spot/instance-action";
|
|
|
|
bool QueryToken(Logger& logger, const tchar* rootDir)
|
|
{
|
|
for (auto& provider : g_providers)
|
|
{
|
|
if (IsNotCloud(logger, rootDir, provider.name))
|
|
continue;
|
|
|
|
HttpConnection http;
|
|
http.SetConnectTimeout(200);
|
|
|
|
u32 statusCode = 0;
|
|
|
|
//logger.Warning(TC("%hs TOKEN PUT START"), provider.name);
|
|
|
|
StringBuffer<1024> token;
|
|
token.Append(provider.tokenPrefix);
|
|
|
|
if (!http.Query(logger, provider.tokenCommand, token, statusCode, g_imdsHost, provider.tokenImds, provider.tokenHeader, 250))
|
|
{
|
|
WriteIsNotCloud(logger, rootDir, provider.name);
|
|
continue;
|
|
}
|
|
token.Append("\r\n");
|
|
|
|
#if PLATFORM_WINDOWS
|
|
char tokenString[1024];
|
|
size_t tokenStringLen;
|
|
wcstombs_s(&tokenStringLen, tokenString, sizeof_array(tokenString), token.data, _TRUNCATE);
|
|
m_tokenString = tokenString;
|
|
#else
|
|
m_tokenString = token.data;
|
|
#endif
|
|
|
|
m_provider = &provider;
|
|
|
|
return true;
|
|
}
|
|
|
|
return false;
|
|
}
|
|
|
|
bool QueryInformation(Logger& logger, StringBufferBase& outExtraInfo, const tchar* rootDir)
|
|
{
|
|
if (!QueryToken(logger, rootDir))
|
|
return false;
|
|
|
|
HttpConnection http;
|
|
http.SetConnectTimeout(200);
|
|
|
|
u32 statusCode = 0;
|
|
|
|
//logger.Warning(TC("%hs TOKEN PUT START"), provider.name);
|
|
|
|
StringBuffer<512> instanceId;
|
|
if (!http.Query(logger, "GET", instanceId, statusCode, g_imdsHost, m_provider->instanceIdImds, m_tokenString.c_str()))
|
|
return false;
|
|
|
|
outExtraInfo.Append(TCV(", ")).Append(m_provider->name).Append(TCV(": ")).Append(instanceId);
|
|
|
|
if (m_provider->instanceLifeCycleImds)
|
|
{
|
|
StringBuffer<32> instanceLifeCycle;
|
|
if (http.Query(logger, "GET", instanceLifeCycle, statusCode, g_imdsHost, m_provider->instanceLifeCycleImds, m_tokenString.c_str()))
|
|
{
|
|
outExtraInfo.Append(' ').Append(instanceLifeCycle);
|
|
m_isSpot = instanceLifeCycle.Contains(TC("spot"));
|
|
}
|
|
}
|
|
|
|
if (m_provider->autoScalingLifeCycleStateImds)
|
|
{
|
|
StringBuffer<32> autoscaling;
|
|
if (http.Query(logger, "GET", autoscaling, statusCode, g_imdsHost, m_provider->autoScalingLifeCycleStateImds, m_tokenString.c_str()) && statusCode == 200)
|
|
{
|
|
outExtraInfo.Append(m_isSpot ? '/' : ' ').Append(TCV("autoscale"));
|
|
m_isAutoscaling = true;
|
|
}
|
|
}
|
|
|
|
if (!QueryAvailabilityZone(logger, nullptr))
|
|
return false;
|
|
return true;
|
|
}
|
|
|
|
bool QueryAvailabilityZone(Logger& logger, const tchar* rootDir)
|
|
{
|
|
if (rootDir && !QueryToken(logger, rootDir))
|
|
return false;
|
|
if (!m_provider->availabilityZoneImds)
|
|
return false;
|
|
|
|
HttpConnection http;
|
|
http.SetConnectTimeout(200);
|
|
|
|
StringBuffer<128> availabilityZone;
|
|
u32 statusCode = 0;
|
|
if (!http.Query(logger, "GET", availabilityZone, statusCode, g_imdsHost, m_provider->availabilityZoneImds, m_tokenString.c_str()))
|
|
return false;
|
|
|
|
const tchar* zone = availabilityZone.data;
|
|
if (const tchar* lastSlash = TStrrchr(zone, '/'))
|
|
zone = lastSlash + 1;
|
|
|
|
m_availabilityZone = zone;
|
|
return true;
|
|
}
|
|
|
|
// Returns true if we _know_ we are not in this cloud provider
|
|
bool IsNotCloud(Logger& logger, const tchar* rootDir, const tchar* provider)
|
|
{
|
|
StringBuffer<512> file;
|
|
file.Append(rootDir).EnsureEndsWithSlash().Append(TCV(".isNot")).Append(provider);
|
|
if (FileExists(logger, file.data))
|
|
return true;
|
|
return false;
|
|
}
|
|
|
|
bool WriteIsNotCloud(Logger& logger, const tchar* rootDir, const tchar* provider)
|
|
{
|
|
StringBuffer<512> file;
|
|
file.Append(rootDir).EnsureEndsWithSlash().Append(TCV(".isNot")).Append(provider);
|
|
FileAccessor f(logger, file.data);
|
|
if (!f.CreateWrite())
|
|
return false;
|
|
if (!f.Close())
|
|
return false;
|
|
return false;
|
|
}
|
|
|
|
const tchar* GetProvider() const
|
|
{
|
|
if (!m_provider)
|
|
return TC("");
|
|
return m_provider->name;
|
|
}
|
|
|
|
bool IsTerminating(Logger& logger, StringBufferBase& outReason, u64& outTerminationTimeMs)
|
|
{
|
|
HttpConnection http;
|
|
|
|
outTerminationTimeMs = 0;
|
|
if (m_isSpot)
|
|
{
|
|
StringBuffer<1024> content;
|
|
u32 statusCode = 0;
|
|
if (http.Query(logger, "GET", content, statusCode, g_imdsHost, g_imdsSpotInstanceAction, m_tokenString.c_str()) && statusCode == 200)
|
|
{
|
|
outReason.Append(TCV("AWS spot instance interruption"));
|
|
return true;
|
|
}
|
|
}
|
|
|
|
if (m_isAutoscaling)
|
|
{
|
|
StringBuffer<1024> content;
|
|
u32 statusCode = 0;
|
|
if (http.Query(logger, "GET", content, statusCode, g_imdsHost, g_imdsAutoScalingLifeCycleState, m_tokenString.c_str()) && statusCode == 200)
|
|
{
|
|
//if (!content.Equals(L"InService"))
|
|
//{
|
|
// wprintf(L"AWSACTION: AUTOSCALE: %ls\n", content.data);
|
|
//}
|
|
|
|
if (!content.Contains(TC("InService"))) // AWS can return "InServiceI" as well?
|
|
{
|
|
//wprintf(L"AWSACTION: AUTOSCALE REBALANCING!!!! (%ls)\n", content.data);
|
|
outReason.Append(TCV("AWS autoscale rebalancing"));
|
|
return true;
|
|
}
|
|
}
|
|
}
|
|
|
|
if (m_provider && m_provider->maintenanceEventImds)
|
|
{
|
|
StringBuffer<1024> content;
|
|
u32 statusCode = 0;
|
|
if (http.Query(logger, "GET", content, statusCode, g_imdsHost, m_provider->maintenanceEventImds, m_tokenString.c_str()) && statusCode == 200)
|
|
{
|
|
if (!content.IsEmpty() && !content.Equals(TCV("NONE")))
|
|
{
|
|
outReason.Append(TCV("Google cloud instance interruption (")).Append(content).Append(')');
|
|
return true;
|
|
}
|
|
}
|
|
}
|
|
return false;
|
|
}
|
|
|
|
const tchar* GetAvailabilityZone()
|
|
{
|
|
return m_availabilityZone.c_str();
|
|
}
|
|
|
|
TString m_availabilityZone;
|
|
std::string m_tokenString;
|
|
const Provider* m_provider = nullptr;
|
|
bool m_isSpot = false;
|
|
bool m_isAutoscaling = false;
|
|
};
|
|
}
|
|
|
|
#endif
|