Add TokenCache test

Add function authz::jwt::TokenRef::toStringRef() for token tracing
Add StringRef::toStringView()
This commit is contained in:
Junhyun Shim 2022-07-11 13:40:51 +02:00
parent f8d66c53a3
commit 61f3d14ec6
7 changed files with 141 additions and 12 deletions

View File

@ -2927,7 +2927,7 @@ Future<KeyRangeLocationInfo> getKeyLocation(Reference<TransactionState> trState,
UseTenant useTenant,
Version version) {
auto f = getKeyLocation(trState->cx,
useTenant ? trState->getTenantInfo(true) : TenantInfo(),
useTenant ? trState->getTenantInfo(AllowInvalidTenantID::True) : TenantInfo(),
key,
member,
trState->spanContext,
@ -3066,7 +3066,7 @@ Future<std::vector<KeyRangeLocationInfo>> getKeyRangeLocations(Reference<Transac
UseTenant useTenant,
Version version) {
auto f = getKeyRangeLocations(trState->cx,
useTenant ? trState->getTenantInfo(true) : TenantInfo(),
useTenant ? trState->getTenantInfo(AllowInvalidTenantID::True) : TenantInfo(),
keys,
limit,
reverse,
@ -7482,7 +7482,7 @@ ACTOR Future<TenantMapEntry> blobGranuleGetTenantEntry(Transaction* self, Key ra
self->trState->cx->getCachedLocation(self->getTenant().get(), rangeStartKey, Reverse::False);
if (!cachedLocationInfo.present()) {
KeyRangeLocationInfo l = wait(getKeyLocation_internal(self->trState->cx,
self->trState->getTenantInfo(true),
self->trState->getTenantInfo(AllowInvalidTenantID::True),
rangeStartKey,
self->trState->spanContext,
self->trState->debugID,

View File

@ -1,11 +1,13 @@
#include "fdbrpc/FlowTransport.h"
#include "fdbrpc/TokenCache.h"
#include "fdbrpc/TokenSign.h"
#include "flow/MkCert.h"
#include "flow/UnitTest.h"
#include "flow/network.h"
#include <boost/unordered_map.hpp>
#include <fmt/format.h>
#include <list>
#include <deque>
@ -239,3 +241,62 @@ bool TokenCacheImpl::validate(TenantNameRef name, StringRef token) {
}
return true;
}
namespace authz::jwt {
extern TokenRef makeRandomTokenSpec(Arena&, IRandom&, authz::Algorithm);
}
TEST_CASE("/fdbrpc/authz/TokenCache") {
std::pair<void (*)(Arena&, IRandom&, authz::jwt::TokenRef&), char const*> badMutations[]{
{
[](Arena&, IRandom&, authz::jwt::TokenRef& token) { token.expiresAtUnixTime.reset(); },
"NoExpirationTime",
},
{
[](Arena&, IRandom& rng, authz::jwt::TokenRef& token) {
token.expiresAtUnixTime = uint64_t(g_network->timer() - 10 - rng.random01() * 50);
},
"ExpiredToken",
},
{
[](Arena&, IRandom&, authz::jwt::TokenRef& token) { token.notBeforeUnixTime.reset(); },
"NoNotBefore",
},
{
[](Arena&, IRandom& rng, authz::jwt::TokenRef& token) {
token.notBeforeUnixTime = uint64_t(g_network->timer() + 10 + rng.random01() * 50);
},
"TokenNotYetValid",
},
{
[](Arena& arena, IRandom&, authz::jwt::TokenRef& token) { token.keyId.reset(); },
"UnknownKey",
},
{
[](Arena& arena, IRandom&, authz::jwt::TokenRef& token) { token.tenants.reset(); },
"NoTenants",
},
};
auto const numBadMutations = sizeof(badMutations) / sizeof(badMutations[0]);
for (auto repeat = 0; repeat < 50; repeat++) {
auto arena = Arena();
auto privateKey = mkcert::makeEcP256();
auto& rng = *deterministicRandom();
auto validTokenSpec = authz::jwt::makeRandomTokenSpec(arena, rng, authz::Algorithm::ES256);
for (auto i = 0; i < numBadMutations; i++) {
auto [mutationFn, mutationDesc] = badMutations[i];
auto tmpArena = Arena();
auto mutatedTokenSpec = validTokenSpec;
mutationFn(tmpArena, rng, mutatedTokenSpec);
auto signedToken = authz::jwt::signToken(tmpArena, mutatedTokenSpec, privateKey);
if (TokenCache::instance().validate(validTokenSpec.tenants.get()[0], signedToken)) {
fmt::print("Unexpected successful validation at mutation {}, token spec: {}\n",
mutationDesc,
mutatedTokenSpec.toStringRef(tmpArena).toStringView());
ASSERT(false);
}
}
}
fmt::print("TEST OK\n");
return Void();
}

View File

@ -30,6 +30,7 @@
#include "flow/Trace.h"
#include "flow/UnitTest.h"
#include <fmt/format.h>
#include <iterator>
#include <string_view>
#include <type_traits>
#include <utility>
@ -222,6 +223,44 @@ TokenRef makeRandomTokenSpec(Arena& arena, IRandom& rng) {
namespace authz::jwt {
template <class FieldType, size_t NameLen>
void appendField(fmt::memory_buffer& b, char const (&name)[NameLen], Optional<FieldType> const& field) {
if (!field.present())
return;
auto const& f = field.get();
auto bi = std::back_inserter(b);
if constexpr (std::is_same_v<FieldType, VectorRef<StringRef>>) {
fmt::format_to(bi, " {}=[", name);
for (auto i = 0; i < f.size(); i++) {
if (i)
fmt::format_to(bi, ",");
fmt::format_to(bi, f[i].toStringView());
}
fmt::format_to(bi, "]");
} else if constexpr (std::is_same_v<FieldType, StringRef>) {
fmt::format_to(bi, " {}={}", name, f.toStringView());
} else {
fmt::format_to(bi, " {}={}", name, f);
}
}
StringRef TokenRef::toStringRef(Arena& arena) {
auto buf = fmt::memory_buffer();
fmt::format_to(std::back_inserter(buf), "alg={}", getAlgorithmName(algorithm));
appendField(buf, "iss", issuer);
appendField(buf, "sub", subject);
appendField(buf, "aud", audience);
appendField(buf, "iat", issuedAtUnixTime);
appendField(buf, "exp", expiresAtUnixTime);
appendField(buf, "nbf", notBeforeUnixTime);
appendField(buf, "kid", keyId);
appendField(buf, "jti", tokenId);
appendField(buf, "tenants", tenants);
auto str = new (arena) uint8_t[buf.size()];
memcpy(str, buf.data(), buf.size());
return StringRef(str, buf.size());
}
template <class FieldType, class Writer>
void putField(Optional<FieldType> const& field, Writer& wr, const char* fieldName) {
if (!field.present())
@ -302,10 +341,6 @@ StringRef signToken(Arena& arena, TokenRef tokenSpec, PrivateKey privateKey) {
return StringRef(out, totalLen);
}
StringRef signToken(Arena& arena, TokenRef tokenSpec, StringRef privateKeyDer) {
return signToken(arena, tokenSpec, PrivateKey(DerEncoded{}, privateKeyDer));
}
bool parseHeaderPart(TokenRef& token, StringRef b64urlHeader) {
auto tmpArena = Arena();
auto [header, valid] = base64url::decode(tmpArena, b64urlHeader);
@ -473,7 +508,7 @@ TokenRef makeRandomTokenSpec(Arena& arena, IRandom& rng, Algorithm alg) {
aud[i] = genRandomAlphanumStringRef(arena, rng, MaxTenantNameLenPlus1);
ret.audience = VectorRef<StringRef>(aud, numAudience);
ret.issuedAtUnixTime = timer_int() / 1'000'000'000ul;
ret.notBeforeUnixTime = timer_int() / 1'000'000'000ul;
ret.notBeforeUnixTime = ret.issuedAtUnixTime.get();
ret.expiresAtUnixTime = ret.issuedAtUnixTime.get() + rng.randomInt(360, 1080 + 1);
ret.keyId = genRandomAlphanumStringRef(arena, rng, MaxKeyNameLenPlus1);
auto numTenants = rng.randomInt(1, 3);
@ -554,6 +589,33 @@ TEST_CASE("/fdbrpc/TokenSign/JWT") {
return Void();
}
TEST_CASE("/fdbrpc/TokenSign/JWT/ToStringRef") {
auto t = authz::jwt::TokenRef();
t.algorithm = authz::Algorithm::ES256;
t.issuer = "issuer"_sr;
t.subject = "subject"_sr;
StringRef aud[3]{ "aud1"_sr, "aud2"_sr, "aud3"_sr };
t.audience = VectorRef<StringRef>(aud, 3);
t.issuedAtUnixTime = 123ul;
t.expiresAtUnixTime = 456ul;
t.notBeforeUnixTime = 789ul;
t.keyId = "keyId"_sr;
t.tokenId = "tokenId"_sr;
StringRef tenants[2]{ "tenant1"_sr, "tenant2"_sr };
t.tenants = VectorRef<StringRef>(tenants, 2);
auto arena = Arena();
auto tokenStr = t.toStringRef(arena);
auto tokenStrExpected =
"alg=ES256 iss=issuer sub=subject aud=[aud1,aud2,aud3] iat=123 exp=456 nbf=789 kid=keyId jti=tokenId tenants=[tenant1,tenant2]"_sr;
if (tokenStr != tokenStrExpected) {
fmt::print("Expected: {}\nGot : {}\n", tokenStrExpected.toStringView(), tokenStr.toStringView());
ASSERT(false);
} else {
fmt::print("TEST OK\n");
}
return Void();
}
TEST_CASE("/fdbrpc/TokenSign/bench") {
constexpr auto repeat = 10;
constexpr auto numSamples = 10000;

View File

@ -23,7 +23,7 @@
#include "flow/Arena.h"
#include "fdbrpc/TenantInfo.h"
class TokenCache {
class TokenCache : NonCopyable {
struct TokenCacheImpl* impl;
TokenCache();

View File

@ -97,6 +97,9 @@ struct TokenRef {
Optional<VectorRef<StringRef>> tenants; // tenants
// signature part
StringRef signature;
// print each non-signature field in non-JSON, human-readable format e.g. for trace
StringRef toStringRef(Arena& arena);
};
// Make plain JSON token string with fields (except signature) from passed spec
@ -106,7 +109,7 @@ StringRef makeTokenPart(Arena& arena, TokenRef tokenSpec);
StringRef makePlainSignature(Arena& arena, Algorithm alg, StringRef tokenPart, StringRef privateKeyDer);
// One-stop function to make JWT from spec
StringRef signToken(Arena& arena, TokenRef tokenSpec, StringRef privateKeyDer);
StringRef signToken(Arena& arena, TokenRef tokenSpec, PrivateKey privateKey);
// Parse passed b64url-encoded header part and materialize its contents into tokenOut,
// using memory allocated from arena

View File

@ -80,7 +80,7 @@ struct CycleWorkload : TestWorkload, CycleMembers<MultiTenancy> {
tenants.push_back_deep(this->arena, this->tenant);
this->token.tenants = tenants;
// we currently don't support this workload to be run outside of simulation
this->signedToken = authz::jwt::signToken(this->arena, this->token, k->second.writeDer(this->arena));
this->signedToken = authz::jwt::signToken(this->arena, this->token, k->second);
}
}

View File

@ -34,6 +34,7 @@
#include <algorithm>
#include <boost/functional/hash.hpp>
#include <stdint.h>
#include <string_view>
#include <string>
#include <cstring>
#include <limits>
@ -530,7 +531,9 @@ public:
return substr(0, size() - s.size());
}
std::string toString() const { return std::string((const char*)data, length); }
std::string toString() const { return std::string(reinterpret_cast<const char*>(data), length); }
std::string_view toStringView() const { return std::string_view(reinterpret_cast<const char*>(data), length); }
static bool isPrintable(char c) { return c > 32 && c < 127; }
inline std::string printable() const;