Skip to content

Commit 8477393

Browse files
tvernumtvernum
authored
Add ThreadContextTransient utility class (#134278)
This class is a little bit like `ThreadLocal` and a little bit like `Setting`. It allows for type-safe getting and setting of _transient_ values insode of a `ThreadContext`, eliminating the need to work with `String` keys and type-inferred values.
1 parent d915090 commit 8477393

File tree

32 files changed

+310
-129
lines changed

32 files changed

+310
-129
lines changed

‎server/src/main/java/org/elasticsearch/common/util/concurrent/ThreadContextTransient.java‎

Lines changed: 104 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,104 @@
1+
/*
2+
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
3+
* or more contributor license agreements. Licensed under the "Elastic License
4+
* 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
5+
* Public License v 1"; you may not use this file except in compliance with, at
6+
* your election, the "Elastic License 2.0", the "GNU Affero General Public
7+
* License v3.0 only", or the "Server Side Public License, v 1".
8+
*/
9+
10+
package org.elasticsearch.common.util.concurrent;
11+
12+
import org.elasticsearch.common.Strings;
13+
import org.elasticsearch.core.Nullable;
14+
15+
/**
16+
* A utility class for resolving/setting values in the {@link ThreadContext} in a typesafe way
17+
* @see ThreadContext#getTransient(String)
18+
* @see ThreadContext#putTransient(String, Object)
19+
*/
20+
public final class ThreadContextTransient<T> {
21+
22+
private final String key;
23+
private final Class<T> type;
24+
25+
private ThreadContextTransient(String key, Class<T> type) {
26+
this.key = key;
27+
this.type = type;
28+
}
29+
30+
/**
31+
* @return The key/name of the transient header
32+
*/
33+
public String getKey() {
34+
return key;
35+
}
36+
37+
/**
38+
* @return {@code true} if the thread context contains a non-null value for this {@link #getKey() key}
39+
*/
40+
public boolean exists(ThreadContext threadContext) {
41+
return threadContext.getTransient(key) != null;
42+
}
43+
44+
/**
45+
* @return The current value for this {@link #getKey() key}. May be {@code null}.
46+
*/
47+
@Nullable
48+
public T get(ThreadContext threadContext) {
49+
final Object val = threadContext.getTransient(key);
50+
if (val == null) {
51+
return null;
52+
}
53+
if (this.type.isInstance(val)) {
54+
return this.type.cast(val);
55+
} else {
56+
final String message = Strings.format(
57+
"Found object of type [%s] as transient [%s] in thread-context, but expected it to be [%s]",
58+
val.getClass(),
59+
key,
60+
type
61+
);
62+
assert false : message;
63+
throw new IllegalStateException(message);
64+
}
65+
}
66+
67+
/**
68+
* @return The current value for this {@link #getKey() key}. May not be {@code null}
69+
* @throws IllegalStateException if the thread context does not contain a value (or contains {@code null}).
70+
*/
71+
public T require(ThreadContext threadContext) {
72+
final T value = get(threadContext);
73+
if (value == null) {
74+
throw new IllegalStateException("Cannot find value for [" + key + "] in thread-context");
75+
}
76+
return value;
77+
}
78+
79+
/**
80+
* Set the value for the this {@link #getKey() key}.
81+
* Because transient headers cannot be overwritten, this method will throw an exception if a value already exists
82+
* @see ThreadContext#putTransient(String, Object)
83+
*/
84+
public void set(ThreadContext threadContext, T value) {
85+
threadContext.putTransient(this.key, value);
86+
}
87+
88+
/**
89+
* Set the value for the this {@link #getKey() key}, if and only if there is no current value
90+
* @return {@code true} if the value was set, {@code false} otherwise
91+
*/
92+
public boolean setIfEmpty(ThreadContext threadContext, T value) {
93+
if (exists(threadContext) == false) {
94+
set(threadContext, value);
95+
return true;
96+
} else {
97+
return false;
98+
}
99+
}
100+
101+
public static <T> ThreadContextTransient<T> transientValue(String key, Class<T> type) {
102+
return new ThreadContextTransient<>(key, type);
103+
}
104+
}

‎server/src/test/java/org/elasticsearch/common/util/concurrent/ThreadContextTransientTests.java‎

Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,56 @@
1+
/*
2+
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
3+
* or more contributor license agreements. Licensed under the "Elastic License
4+
* 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
5+
* Public License v 1"; you may not use this file except in compliance with, at
6+
* your election, the "Elastic License 2.0", the "GNU Affero General Public
7+
* License v3.0 only", or the "Server Side Public License, v 1".
8+
*/
9+
10+
package org.elasticsearch.common.util.concurrent;
11+
12+
import org.elasticsearch.common.settings.SecureString;
13+
import org.elasticsearch.common.settings.Settings;
14+
import org.elasticsearch.test.ESTestCase;
15+
16+
import static org.hamcrest.Matchers.is;
17+
import static org.hamcrest.Matchers.nullValue;
18+
import static org.hamcrest.Matchers.sameInstance;
19+
20+
public class ThreadContextTransientTests extends ESTestCase {
21+
22+
public void testSetAndGetSecureStringValue() {
23+
final String key = randomAlphanumericOfLength(12);
24+
final ThreadContextTransient<SecureString> tcv = ThreadContextTransient.transientValue(key, SecureString.class);
25+
26+
final ThreadContext threadContext = new ThreadContext(Settings.EMPTY);
27+
28+
assertThat(tcv.exists(threadContext), is(false));
29+
assertThat(tcv.get(threadContext), nullValue());
30+
expectThrows(IllegalStateException.class, () -> tcv.require(threadContext));
31+
32+
final SecureString value = new SecureString(randomAlphanumericOfLength(8).toCharArray());
33+
tcv.set(threadContext, value);
34+
assertThat(tcv.exists(threadContext), is(true));
35+
assertThat(tcv.get(threadContext), sameInstance(value));
36+
assertThat(tcv.require(threadContext), sameInstance(value));
37+
38+
final SecureString value2 = new SecureString(randomAlphanumericOfLength(10).toCharArray());
39+
assertThat(tcv.setIfEmpty(threadContext, value2), is(false));
40+
assertThat(tcv.get(threadContext), sameInstance(value));
41+
assertThat(tcv.require(threadContext), sameInstance(value));
42+
43+
try (var restore = threadContext.stashContext()) {
44+
assertThat(tcv.exists(threadContext), is(false));
45+
assertThat(tcv.get(threadContext), nullValue());
46+
47+
assertThat(tcv.setIfEmpty(threadContext, value2), is(true));
48+
assertThat(tcv.get(threadContext), sameInstance(value2));
49+
assertThat(tcv.require(threadContext), sameInstance(value2));
50+
}
51+
52+
assertThat(tcv.get(threadContext), sameInstance(value));
53+
assertThat(tcv.require(threadContext), sameInstance(value));
54+
}
55+
56+
}

‎x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/security/SecurityContext.java‎

Lines changed: 9 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -37,7 +37,7 @@
3737

3838
import static org.elasticsearch.xpack.core.security.authc.Authentication.getAuthenticationFromCrossClusterAccessMetadata;
3939
import static org.elasticsearch.xpack.core.security.authc.AuthenticationField.AUTHENTICATION_KEY;
40-
import static org.elasticsearch.xpack.core.security.authz.AuthorizationServiceField.AUTHORIZATION_INFO_KEY;
40+
import static org.elasticsearch.xpack.core.security.authz.AuthorizationServiceField.AUTHORIZATION_INFO_VALUE;
4141

4242
/**
4343
* A lightweight utility that can find the current user and authentication information for the local thread.
@@ -92,7 +92,7 @@ public Authentication getAuthentication() {
9292
}
9393

9494
public AuthorizationEngine.AuthorizationInfo getAuthorizationInfoFromContext() {
95-
return Objects.requireNonNull(threadContext.getTransient(AUTHORIZATION_INFO_KEY), "authorization info is missing from context");
95+
return Objects.requireNonNull(AUTHORIZATION_INFO_VALUE.get(threadContext), "authorization info is missing from context");
9696
}
9797

9898
@Nullable
@@ -135,20 +135,22 @@ public void putIndicesAccessControl(@Nullable IndicesAccessControl indicesAccess
135135
if (indicesAccessControl.isGranted() == false) {
136136
throw new IllegalStateException("Unexpected unauthorized access control :" + indicesAccessControl);
137137
}
138-
threadContext.putTransient(AuthorizationServiceField.INDICES_PERMISSIONS_KEY, indicesAccessControl);
138+
AuthorizationServiceField.INDICES_PERMISSIONS_VALUE.set(threadContext, indicesAccessControl);
139139
}
140140
}
141141

142142
public void copyIndicesAccessControlToReaderContext(ReaderContext readerContext) {
143-
IndicesAccessControl indicesAccessControl = getThreadContext().getTransient(AuthorizationServiceField.INDICES_PERMISSIONS_KEY);
143+
IndicesAccessControl indicesAccessControl = AuthorizationServiceField.INDICES_PERMISSIONS_VALUE.get(getThreadContext());
144144
assert indicesAccessControl != null : "thread context does not contain index access control";
145-
readerContext.putInContext(AuthorizationServiceField.INDICES_PERMISSIONS_KEY, indicesAccessControl);
145+
readerContext.putInContext(AuthorizationServiceField.INDICES_PERMISSIONS_VALUE.getKey(), indicesAccessControl);
146146
}
147147

148148
public void copyIndicesAccessControlFromReaderContext(ReaderContext readerContext) {
149-
IndicesAccessControl scrollIndicesAccessControl = readerContext.getFromContext(AuthorizationServiceField.INDICES_PERMISSIONS_KEY);
149+
IndicesAccessControl scrollIndicesAccessControl = readerContext.getFromContext(
150+
AuthorizationServiceField.INDICES_PERMISSIONS_VALUE.getKey()
151+
);
150152
assert scrollIndicesAccessControl != null : "scroll does not contain index access control";
151-
getThreadContext().putTransient(AuthorizationServiceField.INDICES_PERMISSIONS_KEY, scrollIndicesAccessControl);
153+
AuthorizationServiceField.INDICES_PERMISSIONS_VALUE.set(getThreadContext(), scrollIndicesAccessControl);
152154
}
153155

154156
/**

‎x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/security/authc/AuthenticationResult.java‎

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,8 @@
66
*/
77
package org.elasticsearch.xpack.core.security.authc;
88

9+
import org.elasticsearch.common.util.concurrent.ThreadContext;
10+
import org.elasticsearch.common.util.concurrent.ThreadContextTransient;
911
import org.elasticsearch.core.Nullable;
1012
import org.elasticsearch.xpack.core.security.user.User;
1113

@@ -26,7 +28,16 @@
2628
public final class AuthenticationResult<T> {
2729
private static final AuthenticationResult<?> NOT_HANDLED = new AuthenticationResult<>(Status.CONTINUE, null, null, null, null);
2830

29-
public static final String THREAD_CONTEXT_KEY = "_xpack_security_auth_result";
31+
@SuppressWarnings("rawtypes")
32+
public static final ThreadContextTransient<AuthenticationResult> THREAD_CONTEXT_VALUE = ThreadContextTransient.transientValue(
33+
"_xpack_security_auth_result",
34+
AuthenticationResult.class
35+
);
36+
37+
@SuppressWarnings("unchecked")
38+
public static <T> AuthenticationResult<T> get(ThreadContext threadContext) {
39+
return (AuthenticationResult<T>) AuthenticationResult.THREAD_CONTEXT_VALUE.get(threadContext);
40+
}
3041

3142
public enum Status {
3243
/**

‎x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/security/authz/AuthorizationServiceField.java‎

Lines changed: 17 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -6,18 +6,31 @@
66
*/
77
package org.elasticsearch.xpack.core.security.authz;
88

9+
import org.elasticsearch.common.util.concurrent.ThreadContextTransient;
10+
import org.elasticsearch.xpack.core.security.authz.accesscontrol.IndicesAccessControl;
11+
912
import java.util.Collection;
1013
import java.util.List;
1114

1215
public final class AuthorizationServiceField {
1316

14-
public static final String INDICES_PERMISSIONS_KEY = "_indices_permissions";
15-
public static final String ORIGINATING_ACTION_KEY = "_originating_action_name";
16-
public static final String AUTHORIZATION_INFO_KEY = "_authz_info";
17+
public static final ThreadContextTransient<IndicesAccessControl> INDICES_PERMISSIONS_VALUE = ThreadContextTransient.transientValue(
18+
"_indices_permissions",
19+
IndicesAccessControl.class
20+
);
21+
public static final ThreadContextTransient<String> ORIGINATING_ACTION_VALUE = ThreadContextTransient.transientValue(
22+
"_originating_action_name",
23+
String.class
24+
);
25+
public static final ThreadContextTransient<AuthorizationEngine.AuthorizationInfo> AUTHORIZATION_INFO_VALUE = ThreadContextTransient
26+
.transientValue("_authz_info", AuthorizationEngine.AuthorizationInfo.class);
1727

1828
// Most often, transient authorisation headers are scoped (i.e. set, read and cleared) for the authorisation and execution
1929
// of individual actions (i.e. there is a different scope between the parent and the child actions)
20-
public static final Collection<String> ACTION_SCOPE_AUTHORIZATION_KEYS = List.of(INDICES_PERMISSIONS_KEY, AUTHORIZATION_INFO_KEY);
30+
public static final Collection<String> ACTION_SCOPE_AUTHORIZATION_KEYS = List.of(
31+
INDICES_PERMISSIONS_VALUE.getKey(),
32+
AUTHORIZATION_INFO_VALUE.getKey()
33+
);
2134

2235
private AuthorizationServiceField() {}
2336
}

‎x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/security/authz/accesscontrol/SecurityIndexReaderWrapper.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -107,7 +107,7 @@ public DirectoryReader apply(final DirectoryReader reader) {
107107

108108
protected IndicesAccessControl getIndicesAccessControl() {
109109
final ThreadContext threadContext = securityContext.getThreadContext();
110-
IndicesAccessControl indicesAccessControl = threadContext.getTransient(AuthorizationServiceField.INDICES_PERMISSIONS_KEY);
110+
IndicesAccessControl indicesAccessControl = AuthorizationServiceField.INDICES_PERMISSIONS_VALUE.get(threadContext);
111111
if (indicesAccessControl == null) {
112112
throw Exceptions.authorizationError("no indices permissions found");
113113
}

‎x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/termsenum/action/TransportTermsEnumAction.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -428,7 +428,7 @@ private boolean canAccess(
428428
ThreadContext threadContext
429429
) {
430430
if (XPackSettings.SECURITY_ENABLED.get(settings)) {
431-
IndicesAccessControl indicesAccessControl = threadContext.getTransient(AuthorizationServiceField.INDICES_PERMISSIONS_KEY);
431+
IndicesAccessControl indicesAccessControl = AuthorizationServiceField.INDICES_PERMISSIONS_VALUE.get(threadContext);
432432
IndicesAccessControl.IndexAccessControl indexAccessControl = indicesAccessControl.getIndexPermissions(shardId.getIndexName());
433433

434434
if (indexAccessControl != null

‎x-pack/plugin/downsample/src/main/java/org/elasticsearch/xpack/downsample/TransportDownsampleAction.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -212,7 +212,7 @@ protected void masterOperation(
212212
String sourceIndexName = request.getSourceIndex();
213213
IndexNameExpressionResolver.assertExpressionHasNullOrDataSelector(sourceIndexName);
214214
IndexNameExpressionResolver.assertExpressionHasNullOrDataSelector(request.getTargetIndex());
215-
final IndicesAccessControl indicesAccessControl = threadContext.getTransient(AuthorizationServiceField.INDICES_PERMISSIONS_KEY);
215+
final IndicesAccessControl indicesAccessControl = AuthorizationServiceField.INDICES_PERMISSIONS_VALUE.get(threadContext);
216216
if (indicesAccessControl != null) {
217217
final IndicesAccessControl.IndexAccessControl indexPermissions = indicesAccessControl.getIndexPermissions(sourceIndexName);
218218
if (indexPermissions != null) {

0 commit comments

Comments
 (0)