Add TokenCache test

Add function authz::jwt::TokenRef::toStringRef() for token tracing
Add StringRef::toStringView()

fdbclient/NativeAPI.actor.cpp

@@ -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,

fdbrpc/TokenCache.cpp

@@ -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();
+}

fdbrpc/TokenSign.cpp

@@ -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;

fdbrpc/include/fdbrpc/TokenCache.h

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

fdbrpc/include/fdbrpc/TokenSign.h

@@ -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

fdbserver/workloads/Cycle.actor.cpp

@@ -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);
 		}
 	}
 

flow/include/flow/Arena.h

@@ -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;