Lifecycle management for Account objects reused accross websocket requests

websocket-resources/src/main/java/org/whispersystems/websocket/ReusableAuth.java

@@ -0,0 +1,182 @@
+/*
+ * Copyright 2024 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+package org.whispersystems.websocket;
+
+import java.security.Principal;
+import java.util.Optional;
+import java.util.concurrent.atomic.AtomicBoolean;
+import org.whispersystems.websocket.auth.PrincipalSupplier;
+
+/**
+ * This class holds a principal that can be reused across requests on a websocket. Since two requests may operate
+ * concurrently on the same principal, and some principals contain non thread-safe mutable state,  appropriate use of
+ * this class ensures that no data races occur. It also ensures that after a principal is modified, a subsequent request
+ * gets the up-to-date principal
+ *
+ * @param <T> The underlying principal type
+ * @see PrincipalSupplier
+ */
+public abstract sealed class ReusableAuth<T extends Principal> {
+
+  /**
+   * Get a reference to the underlying principal that callers pledge not to modify.
+   * <p>
+   * The reference returned will potentially be provided to many threads concurrently accessing the principal. Callers
+   * should use this method only if they can ensure that they will not modify the in-memory principal object AND they do
+   * not intend to modify the underlying canonical representation of the principal.
+   * <p>
+   * For example, if a caller retrieves a reference to a principal, does not modify the in memory state, but updates a
+   * field on a database that should be reflected in subsequent retrievals of the principal, they will have met the
+   * first criteria, but not the second. In that case they should instead use {@link #mutableRef()}.
+   * <p>
+   * If other callers have modified the underlying principal by using {@link #mutableRef()}, this method may need to
+   * refresh the principal via {@link PrincipalSupplier#refresh} which could be a blocking operation.
+   *
+   * @return If authenticated, a reference to the underlying principal that should not be modified
+   */
+  public abstract Optional<T> ref();
+
+
+  public interface MutableRef<T> {
+
+    T ref();
+
+    void close();
+  }
+
+  /**
+   * Get a reference to the underlying principal that may be modified.
+   * <p>
+   * The underlying principal can be safely modified. Multiple threads may operate on the same {@link ReusableAuth} so
+   * long as they each have their own {@link MutableRef}. After any modifications, the caller must call
+   * {@link MutableRef#close} to notify the principal has become dirty. Close should be called after modifications but
+   * before sending a response on the websocket. This ensures that a request that comes in after a modification response
+   * is received is guaranteed to see the modification.
+   *
+   * @return If authenticated, a reference to the underlying principal that may be modified
+   */
+  public abstract Optional<MutableRef<T>> mutableRef();
+
+  public boolean invalidCredentialsProvided() {
+    return switch (this) {
+      case Invalid<T> ignored -> true;
+      case ReusableAuth.Anonymous<T> ignored -> false;
+      case ReusableAuth.Authenticated<T> ignored-> false;
+    };
+  }
+
+  /**
+   * @return A {@link ReusableAuth} indicating no credential were provided
+   */
+  public static <T extends Principal> ReusableAuth<T> anonymous() {
+    //noinspection unchecked
+    return (ReusableAuth<T>) Anonymous.ANON_RESULT;
+  }
+
+  /**
+   * @return A {@link ReusableAuth} indicating that invalid credentials were provided
+   */
+  public static <T extends Principal> ReusableAuth<T> invalid() {
+    //noinspection unchecked
+    return (ReusableAuth<T>) Invalid.INVALID_RESULT;
+  }
+
+  /**
+   * Create a successfully authenticated {@link ReusableAuth}
+   *
+   * @param principal         The authenticated principal
+   * @param principalSupplier Instructions for how to refresh or copy this principal
+   * @param <T>               The principal type
+   * @return A {@link ReusableAuth} for a successfully authenticated principal
+   */
+  public static <T extends Principal> ReusableAuth<T> authenticated(T principal,
+      PrincipalSupplier<T> principalSupplier) {
+    return new Authenticated<>(principal, principalSupplier);
+  }
+
+
+  private static final class Invalid<T extends Principal> extends ReusableAuth<T> {
+
+    @SuppressWarnings({"rawtypes"})
+    private static final ReusableAuth INVALID_RESULT = new Invalid();
+
+    @Override
+    public Optional<T> ref() {
+      return Optional.empty();
+    }
+
+    @Override
+    public Optional<MutableRef<T>> mutableRef() {
+      return Optional.empty();
+    }
+  }
+
+  private static final class Anonymous<T extends Principal> extends ReusableAuth<T> {
+
+    @SuppressWarnings({"rawtypes"})
+    private static final ReusableAuth ANON_RESULT = new Anonymous();
+
+    @Override
+    public Optional<T> ref() {
+      return Optional.empty();
+    }
+
+    @Override
+    public Optional<MutableRef<T>> mutableRef() {
+      return Optional.empty();
+    }
+  }
+
+  private static final class Authenticated<T extends Principal> extends ReusableAuth<T> {
+
+    private T basePrincipal;
+    private final AtomicBoolean needRefresh = new AtomicBoolean(false);
+    private final PrincipalSupplier<T> principalSupplier;
+
+    Authenticated(final T basePrincipal, PrincipalSupplier<T> principalSupplier) {
+      this.basePrincipal = basePrincipal;
+      this.principalSupplier = principalSupplier;
+
+    }
+
+    @Override
+    public Optional<T> ref() {
+      maybeRefresh();
+      return Optional.of(basePrincipal);
+    }
+
+    @Override
+    public Optional<MutableRef<T>> mutableRef() {
+      maybeRefresh();
+      return Optional.of(new AuthenticatedMutableRef(principalSupplier.deepCopy(basePrincipal)));
+    }
+
+    private void maybeRefresh() {
+      if (needRefresh.compareAndSet(true, false)) {
+        basePrincipal = principalSupplier.refresh(basePrincipal);
+      }
+    }
+
+    private class AuthenticatedMutableRef implements MutableRef<T> {
+
+      final T ref;
+
+      private AuthenticatedMutableRef(T ref) {
+        this.ref = ref;
+      }
+
+      public T ref() {
+        return ref;
+      }
+
+      public void close() {
+        needRefresh.set(true);
+      }
+    }
+  }
+
+  private ReusableAuth() {
+  }
+}

websocket-resources/src/main/java/org/whispersystems/websocket/WebSocketResourceProvider.java

@@ -58,7 +58,7 @@ public class WebSocketResourceProvider<T extends Principal> implements WebSocket
 
   private final Map<Long, CompletableFuture<WebSocketResponseMessage>> requestMap = new ConcurrentHashMap<>();
 
-  private final T authenticated;
+  private final ReusableAuth<T> reusableAuth;
   private final WebSocketMessageFactory messageFactory;
   private final Optional<WebSocketConnectListener> connectListener;
   private final ApplicationHandler jerseyHandler;
@@ -77,7 +77,7 @@ public class WebSocketResourceProvider<T extends Principal> implements WebSocket
       String remoteAddressPropertyName,
       ApplicationHandler jerseyHandler,
       WebsocketRequestLog requestLog,
-      T authenticated,
+      ReusableAuth<T> authenticated,
       WebSocketMessageFactory messageFactory,
       Optional<WebSocketConnectListener> connectListener,
       Duration idleTimeout) {
@@ -85,7 +85,7 @@ public class WebSocketResourceProvider<T extends Principal> implements WebSocket
     this.remoteAddressPropertyName = remoteAddressPropertyName;
     this.jerseyHandler = jerseyHandler;
     this.requestLog = requestLog;
-    this.authenticated = authenticated;
+    this.reusableAuth = authenticated;
     this.messageFactory = messageFactory;
     this.connectListener = connectListener;
     this.idleTimeout = idleTimeout;
@@ -97,7 +97,7 @@ public class WebSocketResourceProvider<T extends Principal> implements WebSocket
     this.remoteEndpoint = session.getRemote();
     this.context = new WebSocketSessionContext(
         new WebSocketClient(session, remoteEndpoint, messageFactory, requestMap));
-    this.context.setAuthenticated(authenticated);
+    this.context.setAuthenticated(reusableAuth.ref().orElse(null));
     this.session.setIdleTimeout(idleTimeout);
 
     connectListener.ifPresent(listener -> listener.onWebSocketConnect(this.context));
@@ -162,6 +162,17 @@ public class WebSocketResourceProvider<T extends Principal> implements WebSocket
     logger.debug("onWebSocketText!");
   }
 
+  /**
+   * The property name where {@link org.whispersystems.websocket.auth.WebsocketAuthValueFactoryProvider} can find an
+   * {@link ReusableAuth} object that lives for the lifetime of the websocket
+   */
+  public static final String REUSABLE_AUTH_PROPERTY = WebSocketResourceProvider.class.getName() + ".reusableAuth";
+
+  /**
+   * The property name where {@link org.whispersystems.websocket.auth.WebsocketAuthValueFactoryProvider} can install a
+   * {@link org.whispersystems.websocket.ReusableAuth.MutableRef} for us to close when the request is finished
+   */
+  public static final String RESOLVED_PRINCIPAL_PROPERTY = WebSocketResourceProvider.class.getName() + ".resolvedPrincipal";
   private void handleRequest(WebSocketRequestMessage requestMessage) {
     ContainerRequest containerRequest = new ContainerRequest(null, URI.create(requestMessage.getPath()),
         requestMessage.getVerb(), new WebSocketSecurityContext(new ContextPrincipal(context)),
@@ -173,30 +184,43 @@ public class WebSocketResourceProvider<T extends Principal> implements WebSocket
     }
 
     containerRequest.setProperty(remoteAddressPropertyName, remoteAddress);
+    containerRequest.setProperty(REUSABLE_AUTH_PROPERTY, reusableAuth);
 
     ByteArrayOutputStream responseBody = new ByteArrayOutputStream();
     CompletableFuture<ContainerResponse> responseFuture = (CompletableFuture<ContainerResponse>) jerseyHandler.apply(
         containerRequest, responseBody);
 
-    responseFuture.thenAccept(response -> {
-      try {
-        sendResponse(requestMessage, response, responseBody);
-      } catch (IOException e) {
-        throw new RuntimeException(e);
-      }
-      requestLog.log(remoteAddress, containerRequest, response);
-    }).exceptionally(exception -> {
-      logger.warn("Websocket Error: " + requestMessage.getVerb() + " " + requestMessage.getPath() + "\n"
-          + requestMessage.getBody(), exception);
-      try {
-        sendErrorResponse(requestMessage, Response.status(500).build());
-      } catch (IOException e) {
-        logger.warn("Failed to send error response", e);
-      }
-      requestLog.log(remoteAddress, containerRequest,
-          new ContainerResponse(containerRequest, Response.status(500).build()));
-      return null;
-    });
+    responseFuture
+        .whenComplete((ignoredResponse, ignoredError) -> {
+          // If the request ended up being one that mutates our principal, we have to close it to indicate we're done
+          // with the mutation operation
+          final Object resolvedPrincipal = containerRequest.getProperty(RESOLVED_PRINCIPAL_PROPERTY);
+          if (resolvedPrincipal instanceof ReusableAuth.MutableRef ref) {
+              ref.close();
+          } else if (resolvedPrincipal != null) {
+            logger.warn("unexpected resolved principal type {} : {}", resolvedPrincipal.getClass(), resolvedPrincipal);
+          }
+        })
+        .thenAccept(response -> {
+          try {
+            sendResponse(requestMessage, response, responseBody);
+          } catch (IOException e) {
+            throw new RuntimeException(e);
+          }
+          requestLog.log(remoteAddress, containerRequest, response);
+        })
+        .exceptionally(exception -> {
+          logger.warn("Websocket Error: " + requestMessage.getVerb() + " " + requestMessage.getPath() + "\n"
+              + requestMessage.getBody(), exception);
+          try {
+            sendErrorResponse(requestMessage, Response.status(500).build());
+          } catch (IOException e) {
+            logger.warn("Failed to send error response", e);
+          }
+          requestLog.log(remoteAddress, containerRequest,
+              new ContainerResponse(containerRequest, Response.status(500).build()));
+          return null;
+        });
   }
 
   @VisibleForTesting

websocket-resources/src/main/java/org/whispersystems/websocket/WebSocketResourceProviderFactory.java

@@ -22,7 +22,6 @@ import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.whispersystems.websocket.auth.AuthenticationException;
 import org.whispersystems.websocket.auth.WebSocketAuthenticator;
-import org.whispersystems.websocket.auth.WebSocketAuthenticator.AuthenticationResult;
 import org.whispersystems.websocket.auth.WebsocketAuthValueFactoryProvider;
 import org.whispersystems.websocket.configuration.WebSocketConfiguration;
 import org.whispersystems.websocket.session.WebSocketSessionContextValueFactoryProvider;
@@ -57,17 +56,17 @@ public class WebSocketResourceProviderFactory<T extends Principal> extends Jetty
   public Object createWebSocket(final JettyServerUpgradeRequest request, final JettyServerUpgradeResponse response) {
     try {
       Optional<WebSocketAuthenticator<T>> authenticator = Optional.ofNullable(environment.getAuthenticator());
-      T authenticated = null;
 
+      final ReusableAuth<T> authenticated;
       if (authenticator.isPresent()) {
-        AuthenticationResult<T> authenticationResult = authenticator.get().authenticate(request);
+        authenticated = authenticator.get().authenticate(request);
 
-        if (authenticationResult.getUser().isEmpty() && authenticationResult.credentialsPresented()) {
+        if (authenticated.invalidCredentialsProvided()) {
           response.sendForbidden("Unauthorized");
           return null;
-        } else {
-          authenticated = authenticationResult.getUser().orElse(null);
         }
+      } else {
+        authenticated = ReusableAuth.anonymous();
       }
 
       return new WebSocketResourceProvider<>(getRemoteAddress(request),

websocket-resources/src/main/java/org/whispersystems/websocket/auth/Mutable.java

@@ -0,0 +1,26 @@
+/*
+ * Copyright 2024 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+package org.whispersystems.websocket.auth;
+
+import io.dropwizard.auth.Auth;
+
+import java.lang.annotation.ElementType;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+import java.lang.annotation.Target;
+
+/**
+ * An @{@link Auth} object annotated with {@link Mutable} indicates that the consumer of the object
+ * will modify the object or its underlying canonical source.
+ *
+ * Note: An {@link Auth} object that does not specify @{@link ReadOnly} will be assumed to be @Mutable
+ *
+ * @see org.whispersystems.websocket.ReusableAuth
+ */
+@Retention(RetentionPolicy.RUNTIME)
+@Target({ElementType.FIELD, ElementType.PARAMETER})
+public @interface Mutable {
+}
+

websocket-resources/src/main/java/org/whispersystems/websocket/auth/PrincipalSupplier.java

@@ -0,0 +1,58 @@
+/*
+ * Copyright 2024 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+package org.whispersystems.websocket.auth;
+
+/**
+ * Teach {@link org.whispersystems.websocket.ReusableAuth} how to make a deep copy of a principal (that is safe to
+ * concurrently modify while the original principal is being read), and how to refresh a principal after it has been
+ * potentially modified.
+ *
+ * @param <T> The underlying principal type
+ */
+public interface PrincipalSupplier<T> {
+
+  /**
+   * Re-fresh the principal after it has been modified.
+   * <p>
+   * If the principal is populated from a backing store, refresh should re-read it.
+   *
+   * @param t the potentially stale principal to refresh
+   * @return The up-to-date principal
+   */
+  T refresh(T t);
+
+  /**
+   * Create a deep, in-memory copy of the principal. This should be identical to the original principal, but should
+   * share no mutable state with the original. It should be safe for two threads to independently write and read from
+   * two independent deep copies.
+   *
+   * @param t the principal to copy
+   * @return An in-memory copy of the principal
+   */
+  T deepCopy(T t);
+
+  class ImmutablePrincipalSupplier<T> implements PrincipalSupplier<T> {
+    @SuppressWarnings({"rawtypes"})
+    private static final PrincipalSupplier INSTANCE = new ImmutablePrincipalSupplier();
+
+    @Override
+    public T refresh(final T t) {
+      return t;
+    }
+
+    @Override
+    public T deepCopy(final T t) {
+      return t;
+    }
+  }
+
+  /**
+   * @return A principal supplier that can be used if the principal type does not support modification.
+   */
+  static <T> PrincipalSupplier<T> forImmutablePrincipal() {
+    //noinspection unchecked
+    return (PrincipalSupplier<T>) ImmutablePrincipalSupplier.INSTANCE;
+  }
+}

websocket-resources/src/main/java/org/whispersystems/websocket/auth/ReadOnly.java

@@ -0,0 +1,25 @@
+/*
+ * Copyright 2024 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+package org.whispersystems.websocket.auth;
+
+import java.lang.annotation.ElementType;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+import java.lang.annotation.Target;
+
+/**
+ * An @{@link io.dropwizard.auth.Auth} object annotated with {@link ReadOnly} indicates that the consumer of the object
+ * will never modify the object, nor its underlying canonical source.
+ * <p>
+ * For example, a consumer of a @ReadOnly AuthenticatedAccount promises to never modify the in-memory
+ * AuthenticatedAccount and to never modify the underlying Account database for the account.
+ *
+ * @see org.whispersystems.websocket.ReusableAuth
+ */
+@Retention(RetentionPolicy.RUNTIME)
+@Target({ElementType.FIELD, ElementType.PARAMETER})
+public @interface ReadOnly {
+}
+

websocket-resources/src/main/java/org/whispersystems/websocket/auth/WebSocketAuthenticator.java

@@ -7,27 +7,8 @@ package org.whispersystems.websocket.auth;
 import java.security.Principal;
 import java.util.Optional;
 import org.eclipse.jetty.websocket.api.UpgradeRequest;
+import org.whispersystems.websocket.ReusableAuth;
 
 public interface WebSocketAuthenticator<T extends Principal> {
-
-  AuthenticationResult<T> authenticate(UpgradeRequest request) throws AuthenticationException;
-
-  @SuppressWarnings("OptionalUsedAsFieldOrParameterType")
-  class AuthenticationResult<T> {
-    private final Optional<T> user;
-    private final boolean credentialsPresented;
-
-    public AuthenticationResult(final Optional<T> user, final boolean credentialsPresented) {
-      this.user = user;
-      this.credentialsPresented = credentialsPresented;
-    }
-
-    public Optional<T> getUser() {
-      return user;
-    }
-
-    public boolean credentialsPresented() {
-      return credentialsPresented;
-    }
-  }
+  ReusableAuth<T> authenticate(UpgradeRequest request) throws AuthenticationException;
 }

websocket-resources/src/main/java/org/whispersystems/websocket/auth/WebsocketAuthValueFactoryProvider.java

@@ -5,24 +5,28 @@
 package org.whispersystems.websocket.auth;
 
 import io.dropwizard.auth.Auth;
+import java.lang.reflect.ParameterizedType;
+import java.security.Principal;
+import java.util.Optional;
+import java.util.function.Function;
+import javax.annotation.Nullable;
+import javax.inject.Inject;
+import javax.inject.Singleton;
+import javax.ws.rs.WebApplicationException;
 import org.glassfish.jersey.internal.inject.AbstractBinder;
 import org.glassfish.jersey.server.ContainerRequest;
 import org.glassfish.jersey.server.internal.inject.AbstractValueParamProvider;
 import org.glassfish.jersey.server.internal.inject.MultivaluedParameterExtractorProvider;
 import org.glassfish.jersey.server.model.Parameter;
 import org.glassfish.jersey.server.spi.internal.ValueParamProvider;
-
-import javax.annotation.Nullable;
-import javax.inject.Inject;
-import javax.inject.Singleton;
-import javax.ws.rs.WebApplicationException;
-import java.lang.reflect.ParameterizedType;
-import java.security.Principal;
-import java.util.Optional;
-import java.util.function.Function;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import org.whispersystems.websocket.ReusableAuth;
+import org.whispersystems.websocket.WebSocketResourceProvider;
 
 @Singleton
 public class WebsocketAuthValueFactoryProvider<T extends Principal> extends AbstractValueParamProvider  {
+  private static final Logger logger = LoggerFactory.getLogger(WebsocketAuthValueFactoryProvider.class);
 
   private final Class<T> principalClass;
 
@@ -39,18 +43,38 @@ public class WebsocketAuthValueFactoryProvider<T extends Principal> extends Abst
       return null;
     }
 
-    if (parameter.getRawType() == Optional.class                                 &&
-        ParameterizedType.class.isAssignableFrom(parameter.getType().getClass()) &&
-        principalClass == ((ParameterizedType)parameter.getType()).getActualTypeArguments()[0])
-    {
-      return request -> new OptionalContainerRequestValueFactory(request).provide();
+    final boolean readOnly = parameter.isAnnotationPresent(ReadOnly.class);
+
+    if (parameter.getRawType() == Optional.class
+        && ParameterizedType.class.isAssignableFrom(parameter.getType().getClass())
+        && principalClass == ((ParameterizedType) parameter.getType()).getActualTypeArguments()[0]) {
+      return containerRequest -> createPrincipal(containerRequest, readOnly);
     } else if (principalClass.equals(parameter.getRawType())) {
-      return request -> new StandardContainerRequestValueFactory(request).provide();
+      return containerRequest ->
+          createPrincipal(containerRequest, readOnly)
+              .orElseThrow(() -> new WebApplicationException("Authenticated resource", 401));
     } else {
       throw new IllegalStateException("Can't inject unassignable principal: " + principalClass + " for parameter: " + parameter);
     }
   }
 
+  private Optional<? extends Principal> createPrincipal(final ContainerRequest request, final boolean readOnly) {
+    final Object obj = request.getProperty(WebSocketResourceProvider.REUSABLE_AUTH_PROPERTY);
+    if (!(obj instanceof ReusableAuth<?>)) {
+      logger.warn("Unexpected reusable auth property type {} : {}", obj.getClass(), obj);
+      return Optional.empty();
+    }
+    @SuppressWarnings("unchecked") final ReusableAuth<T> reusableAuth = (ReusableAuth<T>) obj;
+    if (readOnly) {
+      return reusableAuth.ref();
+    } else {
+      return reusableAuth.mutableRef().map(writeRef -> {
+        request.setProperty(WebSocketResourceProvider.RESOLVED_PRINCIPAL_PROPERTY, writeRef);
+        return writeRef.ref();
+      });
+    }
+  }
+
   @Singleton
   static class WebsocketPrincipalClassProvider<T extends Principal> {
 
@@ -80,38 +104,4 @@ public class WebsocketAuthValueFactoryProvider<T extends Principal> extends Abst
       bind(WebsocketAuthValueFactoryProvider.class).to(ValueParamProvider.class).in(Singleton.class);
     }
   }
-
-  private static class StandardContainerRequestValueFactory {
-
-    private final ContainerRequest request;
-
-    public StandardContainerRequestValueFactory(ContainerRequest request) {
-      this.request = request;
-    }
-
-    public Principal provide() {
-      final Principal principal = request.getSecurityContext().getUserPrincipal();
-
-      if (principal == null) {
-        throw new WebApplicationException("Authenticated resource", 401);
-      }
-
-      return principal;
-    }
-
-  }
-
-  private static class OptionalContainerRequestValueFactory {
-
-    private final ContainerRequest request;
-
-    public OptionalContainerRequestValueFactory(ContainerRequest request) {
-      this.request = request;
-    }
-
-    public Optional<Principal> provide() {
-      return Optional.ofNullable(request.getSecurityContext().getUserPrincipal());
-    }
-  }
-
 }

websocket-resources/src/test/java/org/whispersystems/websocket/WebSocketResourceProviderFactoryTest.java

@@ -7,6 +7,7 @@ package org.whispersystems.websocket;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
 import static org.junit.jupiter.api.Assertions.assertNull;
+import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.verify;
@@ -27,6 +28,7 @@ import org.glassfish.jersey.server.ResourceConfig;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.whispersystems.websocket.auth.AuthenticationException;
+import org.whispersystems.websocket.auth.PrincipalSupplier;
 import org.whispersystems.websocket.auth.WebSocketAuthenticator;
 import org.whispersystems.websocket.configuration.WebSocketConfiguration;
 import org.whispersystems.websocket.setup.WebSocketEnvironment;
@@ -56,8 +58,7 @@ public class WebSocketResourceProviderFactoryTest {
   @Test
   void testUnauthorized() throws AuthenticationException, IOException {
     when(environment.getAuthenticator()).thenReturn(authenticator);
-    when(authenticator.authenticate(eq(request))).thenReturn(
-        new WebSocketAuthenticator.AuthenticationResult<>(Optional.empty(), true));
+    when(authenticator.authenticate(eq(request))).thenReturn(ReusableAuth.invalid());
     when(environment.jersey()).thenReturn(jerseyEnvironment);
 
     WebSocketResourceProviderFactory<?> factory = new WebSocketResourceProviderFactory<>(environment, Account.class,
@@ -74,8 +75,8 @@ public class WebSocketResourceProviderFactoryTest {
     Account account = new Account();
 
     when(environment.getAuthenticator()).thenReturn(authenticator);
-    when(authenticator.authenticate(eq(request))).thenReturn(
-        new WebSocketAuthenticator.AuthenticationResult<>(Optional.of(account), true));
+    when(authenticator.authenticate(eq(request)))
+        .thenReturn(ReusableAuth.authenticated(account, PrincipalSupplier.forImmutablePrincipal()));
     when(environment.jersey()).thenReturn(jerseyEnvironment);
     final HttpServletRequest httpServletRequest = mock(HttpServletRequest.class);
     when(httpServletRequest.getAttribute(REMOTE_ADDRESS_PROPERTY_NAME)).thenReturn("127.0.0.1");
@@ -137,6 +138,7 @@ public class WebSocketResourceProviderFactoryTest {
     public boolean implies(Subject subject) {
       return false;
     }
+
   }
 
 

websocket-resources/src/test/java/org/whispersystems/websocket/WebSocketResourceProviderTest.java

@@ -59,6 +59,7 @@ import org.glassfish.jersey.server.ResourceConfig;
 import org.junit.jupiter.api.Test;
 import org.mockito.ArgumentCaptor;
 import org.mockito.stubbing.Answer;
+import org.whispersystems.websocket.auth.PrincipalSupplier;
 import org.whispersystems.websocket.auth.WebsocketAuthValueFactoryProvider;
 import org.whispersystems.websocket.logging.WebsocketRequestLog;
 import org.whispersystems.websocket.messages.protobuf.ProtobufWebSocketMessageFactory;
@@ -80,7 +81,7 @@ class WebSocketResourceProviderTest {
     WebSocketResourceProvider<TestPrincipal> provider = new WebSocketResourceProvider<>("127.0.0.1",
         REMOTE_ADDRESS_PROPERTY_NAME,
         applicationHandler, requestLog,
-        new TestPrincipal("fooz"),
+        immutableTestPrincipal("fooz"),
         new ProtobufWebSocketMessageFactory(),
         Optional.of(connectListener),
         Duration.ofMillis(30000));
@@ -108,7 +109,7 @@ class WebSocketResourceProviderTest {
     ApplicationHandler applicationHandler = mock(ApplicationHandler.class);
     WebsocketRequestLog requestLog = mock(WebsocketRequestLog.class);
     WebSocketResourceProvider<TestPrincipal> provider = new WebSocketResourceProvider<>("127.0.0.1",
-        REMOTE_ADDRESS_PROPERTY_NAME, applicationHandler, requestLog, new TestPrincipal("foo"),
+        REMOTE_ADDRESS_PROPERTY_NAME, applicationHandler, requestLog, immutableTestPrincipal("foo"),
         new ProtobufWebSocketMessageFactory(), Optional.empty(), Duration.ofMillis(30000));
 
     Session session = mock(Session.class);
@@ -184,7 +185,7 @@ class WebSocketResourceProviderTest {
     ApplicationHandler applicationHandler = mock(ApplicationHandler.class);
     WebsocketRequestLog requestLog = mock(WebsocketRequestLog.class);
     WebSocketResourceProvider<TestPrincipal> provider = new WebSocketResourceProvider<>("127.0.0.1",
-        REMOTE_ADDRESS_PROPERTY_NAME, applicationHandler, requestLog, new TestPrincipal("foo"),
+        REMOTE_ADDRESS_PROPERTY_NAME, applicationHandler, requestLog, immutableTestPrincipal("foo"),
         new ProtobufWebSocketMessageFactory(), Optional.empty(), Duration.ofMillis(30000));
 
     Session session = mock(Session.class);
@@ -240,7 +241,7 @@ class WebSocketResourceProviderTest {
     ApplicationHandler applicationHandler = new ApplicationHandler(resourceConfig);
     WebsocketRequestLog requestLog = mock(WebsocketRequestLog.class);
     WebSocketResourceProvider<TestPrincipal> provider = new WebSocketResourceProvider<>("127.0.0.1",
-        REMOTE_ADDRESS_PROPERTY_NAME, applicationHandler, requestLog, new TestPrincipal("foo"),
+        REMOTE_ADDRESS_PROPERTY_NAME, applicationHandler, requestLog, immutableTestPrincipal("foo"),
         new ProtobufWebSocketMessageFactory(), Optional.empty(), Duration.ofMillis(30000));
 
     Session session = mock(Session.class);
@@ -280,7 +281,7 @@ class WebSocketResourceProviderTest {
     ApplicationHandler applicationHandler = new ApplicationHandler(resourceConfig);
     WebsocketRequestLog requestLog = mock(WebsocketRequestLog.class);
     WebSocketResourceProvider<TestPrincipal> provider = new WebSocketResourceProvider<>("127.0.0.1",
-        REMOTE_ADDRESS_PROPERTY_NAME, applicationHandler, requestLog, new TestPrincipal("foo"),
+        REMOTE_ADDRESS_PROPERTY_NAME, applicationHandler, requestLog, immutableTestPrincipal("foo"),
         new ProtobufWebSocketMessageFactory(), Optional.empty(), Duration.ofMillis(30000));
 
     Session session = mock(Session.class);
@@ -320,7 +321,7 @@ class WebSocketResourceProviderTest {
     ApplicationHandler applicationHandler = new ApplicationHandler(resourceConfig);
     WebsocketRequestLog requestLog = mock(WebsocketRequestLog.class);
     WebSocketResourceProvider<TestPrincipal> provider = new WebSocketResourceProvider<>("127.0.0.1",
-        REMOTE_ADDRESS_PROPERTY_NAME, applicationHandler, requestLog, new TestPrincipal("authorizedUserName"),
+        REMOTE_ADDRESS_PROPERTY_NAME, applicationHandler, requestLog, immutableTestPrincipal("authorizedUserName"),
         new ProtobufWebSocketMessageFactory(), Optional.empty(), Duration.ofMillis(30000));
 
     Session session = mock(Session.class);
@@ -360,8 +361,8 @@ class WebSocketResourceProviderTest {
     ApplicationHandler applicationHandler = new ApplicationHandler(resourceConfig);
     WebsocketRequestLog requestLog = mock(WebsocketRequestLog.class);
     WebSocketResourceProvider<TestPrincipal> provider = new WebSocketResourceProvider<>("127.0.0.1",
-        REMOTE_ADDRESS_PROPERTY_NAME, applicationHandler, requestLog, null, new ProtobufWebSocketMessageFactory(),
-        Optional.empty(), Duration.ofMillis(30000));
+        REMOTE_ADDRESS_PROPERTY_NAME, applicationHandler, requestLog, ReusableAuth.anonymous(),
+        new ProtobufWebSocketMessageFactory(), Optional.empty(), Duration.ofMillis(30000));
 
     Session session = mock(Session.class);
     RemoteEndpoint remoteEndpoint = mock(RemoteEndpoint.class);
@@ -399,7 +400,7 @@ class WebSocketResourceProviderTest {
     ApplicationHandler applicationHandler = new ApplicationHandler(resourceConfig);
     WebsocketRequestLog requestLog = mock(WebsocketRequestLog.class);
     WebSocketResourceProvider<TestPrincipal> provider = new WebSocketResourceProvider<>("127.0.0.1",
-        REMOTE_ADDRESS_PROPERTY_NAME, applicationHandler, requestLog, new TestPrincipal("something"),
+        REMOTE_ADDRESS_PROPERTY_NAME, applicationHandler, requestLog, immutableTestPrincipal("something"),
         new ProtobufWebSocketMessageFactory(), Optional.empty(), Duration.ofMillis(30000));
 
     Session session = mock(Session.class);
@@ -439,8 +440,8 @@ class WebSocketResourceProviderTest {
     ApplicationHandler applicationHandler = new ApplicationHandler(resourceConfig);
     WebsocketRequestLog requestLog = mock(WebsocketRequestLog.class);
     WebSocketResourceProvider<TestPrincipal> provider = new WebSocketResourceProvider<>("127.0.0.1",
-        REMOTE_ADDRESS_PROPERTY_NAME, applicationHandler, requestLog, null, new ProtobufWebSocketMessageFactory(),
-        Optional.empty(), Duration.ofMillis(30000));
+        REMOTE_ADDRESS_PROPERTY_NAME, applicationHandler, requestLog, ReusableAuth.anonymous(),
+        new ProtobufWebSocketMessageFactory(), Optional.empty(), Duration.ofMillis(30000));
 
     Session session = mock(Session.class);
     RemoteEndpoint remoteEndpoint = mock(RemoteEndpoint.class);
@@ -479,7 +480,7 @@ class WebSocketResourceProviderTest {
     ApplicationHandler applicationHandler = new ApplicationHandler(resourceConfig);
     WebsocketRequestLog requestLog = mock(WebsocketRequestLog.class);
     WebSocketResourceProvider<TestPrincipal> provider = new WebSocketResourceProvider<>("127.0.0.1",
-        REMOTE_ADDRESS_PROPERTY_NAME, applicationHandler, requestLog, new TestPrincipal("gooduser"),
+        REMOTE_ADDRESS_PROPERTY_NAME, applicationHandler, requestLog, immutableTestPrincipal("gooduser"),
         new ProtobufWebSocketMessageFactory(), Optional.empty(), Duration.ofMillis(30000));
 
     Session session = mock(Session.class);
@@ -520,7 +521,7 @@ class WebSocketResourceProviderTest {
     ApplicationHandler applicationHandler = new ApplicationHandler(resourceConfig);
     WebsocketRequestLog requestLog = mock(WebsocketRequestLog.class);
     WebSocketResourceProvider<TestPrincipal> provider = new WebSocketResourceProvider<>("127.0.0.1",
-        REMOTE_ADDRESS_PROPERTY_NAME, applicationHandler, requestLog, new TestPrincipal("gooduser"),
+        REMOTE_ADDRESS_PROPERTY_NAME, applicationHandler, requestLog, immutableTestPrincipal("gooduser"),
         new ProtobufWebSocketMessageFactory(), Optional.empty(), Duration.ofMillis(30000));
 
     Session session = mock(Session.class);
@@ -562,7 +563,7 @@ class WebSocketResourceProviderTest {
     ApplicationHandler applicationHandler = new ApplicationHandler(resourceConfig);
     WebsocketRequestLog requestLog = mock(WebsocketRequestLog.class);
     WebSocketResourceProvider<TestPrincipal> provider = new WebSocketResourceProvider<>("127.0.0.1",
-        REMOTE_ADDRESS_PROPERTY_NAME, applicationHandler, requestLog, new TestPrincipal("gooduser"),
+        REMOTE_ADDRESS_PROPERTY_NAME, applicationHandler, requestLog, immutableTestPrincipal("gooduser"),
         new ProtobufWebSocketMessageFactory(), Optional.empty(), Duration.ofMillis(30000));
 
     Session session = mock(Session.class);
@@ -602,7 +603,7 @@ class WebSocketResourceProviderTest {
     ApplicationHandler applicationHandler = new ApplicationHandler(resourceConfig);
     WebsocketRequestLog requestLog = mock(WebsocketRequestLog.class);
     WebSocketResourceProvider<TestPrincipal> provider = new WebSocketResourceProvider<>("127.0.0.1",
-        REMOTE_ADDRESS_PROPERTY_NAME, applicationHandler, requestLog, new TestPrincipal("gooduser"),
+        REMOTE_ADDRESS_PROPERTY_NAME, applicationHandler, requestLog, immutableTestPrincipal("gooduser"),
         new ProtobufWebSocketMessageFactory(), Optional.empty(), Duration.ofMillis(30000));
 
     Session session = mock(Session.class);
@@ -727,6 +728,10 @@ class WebSocketResourceProviderTest {
     }
   }
 
+  public static ReusableAuth<TestPrincipal> immutableTestPrincipal(final String name) {
+    return ReusableAuth.authenticated(new TestPrincipal(name), PrincipalSupplier.forImmutablePrincipal());
+  }
+
   public static class TestException extends Exception {
 
     public TestException(String message) {