Improved Security Store to support race conditions during registration

@@ -75,13 +75,11 @@ public class X509LwM2MIntegrationTest extends AbstractLwM2MIntegrationTest {
         return device;
     }
 
-    //TODO: use different endpoints to isolate tests.
-    @Ignore()
     @Test
     public void testConnectAndObserveTelemetry() throws Exception {
         createDeviceProfile(TRANSPORT_CONFIGURATION);
         X509ClientCredentials credentials = new X509ClientCredentials();
-        credentials.setEndpoint(endpoint+1);
+        credentials.setEndpoint(endpoint);
         Device device = createDevice(credentials);
 
         SingleEntityFilter sef = new SingleEntityFilter();
@@ -99,7 +97,7 @@ public class X509LwM2MIntegrationTest extends AbstractLwM2MIntegrationTest {
         wsClient.waitForReply();
 
         wsClient.registerWaitForUpdate();
-        LwM2MTestClient client = new LwM2MTestClient(executor, endpoint+1);
+        LwM2MTestClient client = new LwM2MTestClient(executor, endpoint);
         Security security = x509(serverUri, 123, clientX509Cert.getEncoded(), clientPrivateKeyFromCert.getEncoded(), serverX509Cert.getEncoded());
         client.init(security, coapConfig);
         String msg = wsClient.waitForUpdate();

@@ -42,8 +42,8 @@ import org.thingsboard.server.common.transport.util.SslUtil;
 import org.thingsboard.server.queue.util.TbLwM2mTransportComponent;
 import org.thingsboard.server.transport.lwm2m.config.LwM2MTransportServerConfig;
 import org.thingsboard.server.transport.lwm2m.secure.credentials.LwM2MCredentials;
-import org.thingsboard.server.transport.lwm2m.server.store.TbEditableSecurityStore;
 import org.thingsboard.server.transport.lwm2m.server.store.TbLwM2MDtlsSessionStore;
+import org.thingsboard.server.transport.lwm2m.server.store.TbMainSecurityStore;
 
 import javax.annotation.PostConstruct;
 import javax.security.auth.x500.X500Principal;
@@ -67,7 +67,7 @@ public class TbLwM2MDtlsCertificateVerifier implements NewAdvancedCertificateVer
     private final TbLwM2MDtlsSessionStore sessionStorage;
     private final LwM2MTransportServerConfig config;
     private final LwM2mCredentialsSecurityInfoValidator securityInfoValidator;
-    private final TbEditableSecurityStore securityStore;
+    private final TbMainSecurityStore securityStore;
 
     @SuppressWarnings("deprecation")
     private StaticCertificateVerifier staticCertificateVerifier;
@@ -134,7 +134,7 @@ public class TbLwM2MDtlsCertificateVerifier implements NewAdvancedCertificateVer
                                 if (msg.hasDeviceInfo() && deviceProfile != null) {
                                     sessionStorage.put(endpoint, new TbX509DtlsSessionInfo(cert.getSubjectX500Principal().getName(), msg));
                                     try {
-                                        securityStore.put(securityInfo);
+                                        securityStore.putX509(securityInfo);
                                     } catch (NonUniqueSecurityInfoException e) {
                                         log.trace("Failed to add security info: {}", securityInfo, e);
                                     }

@@ -17,6 +17,7 @@ package org.thingsboard.server.transport.lwm2m.server.client;
 
 import lombok.RequiredArgsConstructor;
 import lombok.extern.slf4j.Slf4j;
+import org.eclipse.leshan.core.SecurityMode;
 import org.eclipse.leshan.core.model.ResourceModel;
 import org.eclipse.leshan.core.node.LwM2mPath;
 import org.eclipse.leshan.server.registration.Registration;
@@ -30,7 +31,7 @@ import org.thingsboard.server.transport.lwm2m.config.LwM2MTransportServerConfig;
 import org.thingsboard.server.transport.lwm2m.secure.TbLwM2MSecurityInfo;
 import org.thingsboard.server.transport.lwm2m.server.LwM2mTransportContext;
 import org.thingsboard.server.transport.lwm2m.server.LwM2mTransportUtil;
-import org.thingsboard.server.transport.lwm2m.server.store.TbEditableSecurityStore;
+import org.thingsboard.server.transport.lwm2m.server.store.TbMainSecurityStore;
 
 import java.util.Arrays;
 import java.util.Collection;
@@ -54,7 +55,7 @@ public class LwM2mClientContextImpl implements LwM2mClientContext {
 
     private final LwM2mTransportContext context;
     private final LwM2MTransportServerConfig config;
-    private final TbEditableSecurityStore securityStore;
+    private final TbMainSecurityStore securityStore;
     private final Map<String, LwM2mClient> lwM2mClientsByEndpoint = new ConcurrentHashMap<>();
     private final Map<String, LwM2mClient> lwM2mClientsByRegistrationId = new ConcurrentHashMap<>();
     private final Map<UUID, Lwm2mDeviceProfileTransportConfiguration> profiles = new ConcurrentHashMap<>();
@@ -75,6 +76,9 @@ public class LwM2mClientContextImpl implements LwM2mClientContext {
             oldSession = lwM2MClient.getSession();
             TbLwM2MSecurityInfo securityInfo = securityStore.getTbLwM2MSecurityInfoByEndpoint(lwM2MClient.getEndpoint());
             if (securityInfo.getSecurityMode() != null) {
+                if (SecurityMode.X509.equals(securityInfo.getSecurityMode())) {
+                    securityStore.registerX509(registration.getEndpoint(), registration.getId());
+                }
                 if (securityInfo.getDeviceProfile() != null) {
                     profileUpdate(securityInfo.getDeviceProfile());
                     if (securityInfo.getSecurityInfo() != null) {
@@ -124,7 +128,7 @@ public class LwM2mClientContextImpl implements LwM2mClientContext {
             if (currentRegistration.getId().equals(registration.getId())) {
                 lwM2MClient.setState(LwM2MClientState.UNREGISTERED);
                 lwM2mClientsByEndpoint.remove(lwM2MClient.getEndpoint());
-                this.securityStore.remove(lwM2MClient.getEndpoint());
+                this.securityStore.remove(lwM2MClient.getEndpoint(), registration.getId());
                 UUID profileId = lwM2MClient.getProfileId();
                 if (profileId != null) {
                     Optional<LwM2mClient> otherClients = lwM2mClientsByRegistrationId.values().stream().filter(e -> e.getProfileId().equals(profileId)).findFirst();

@@ -22,13 +22,22 @@ import org.jetbrains.annotations.Nullable;
 import org.thingsboard.server.transport.lwm2m.secure.LwM2mCredentialsSecurityInfoValidator;
 import org.thingsboard.server.transport.lwm2m.secure.TbLwM2MSecurityInfo;
 
+import java.util.HashSet;
+import java.util.Map;
+import java.util.Set;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ConcurrentMap;
+import java.util.concurrent.locks.Lock;
+import java.util.concurrent.locks.ReentrantLock;
+
 import static org.thingsboard.server.transport.lwm2m.server.uplink.LwM2mTypeServer.CLIENT;
 
 @Slf4j
-public class TbLwM2mSecurityStore implements TbEditableSecurityStore {
+public class TbLwM2mSecurityStore implements TbMainSecurityStore {
 
     private final TbEditableSecurityStore securityStore;
     private final LwM2mCredentialsSecurityInfoValidator validator;
+    private final ConcurrentMap<String, Set<String>> endpointRegistrations = new ConcurrentHashMap<>();
 
     public TbLwM2mSecurityStore(TbEditableSecurityStore securityStore, LwM2mCredentialsSecurityInfoValidator validator) {
         this.securityStore = securityStore;
@@ -61,24 +70,42 @@ public class TbLwM2mSecurityStore implements TbEditableSecurityStore {
     @Nullable
     public SecurityInfo fetchAndPutSecurityInfo(String credentialsId) {
         TbLwM2MSecurityInfo securityInfo = validator.getEndpointSecurityInfoByCredentialsId(credentialsId, CLIENT);
-        try {
-            if (securityInfo != null) {
+        doPut(securityInfo);
+        return securityInfo != null ? securityInfo.getSecurityInfo() : null;
+    }
+
+    private void doPut(TbLwM2MSecurityInfo securityInfo) {
+        if (securityInfo != null) {
+            try {
                 securityStore.put(securityInfo);
+            } catch (NonUniqueSecurityInfoException e) {
+                log.trace("Failed to add security info: {}", securityInfo, e);
             }
-        } catch (NonUniqueSecurityInfoException e) {
-            log.trace("Failed to add security info: {}", securityInfo, e);
         }
-        return securityInfo != null ? securityInfo.getSecurityInfo() : null;
     }
 
     @Override
-    public void put(TbLwM2MSecurityInfo tbSecurityInfo) throws NonUniqueSecurityInfoException {
-        securityStore.put(tbSecurityInfo);
+    public void putX509(TbLwM2MSecurityInfo securityInfo) throws NonUniqueSecurityInfoException {
+        securityStore.put(securityInfo);
     }
 
     @Override
-    public void remove(String endpoint) {
-        //TODO: Make sure we delay removal of security store from endpoint due to reg/unreg race condition.
-//        securityStore.remove(endpoint);
+    public void registerX509(String endpoint, String registrationId) {
+        endpointRegistrations.computeIfAbsent(endpoint, ep -> new HashSet<>()).add(registrationId);
+    }
+
+    @Override
+    public void remove(String endpoint, String registrationId) {
+        Set<String> epRegistrationIds = endpointRegistrations.get(endpoint);
+        boolean shouldRemove;
+        if (epRegistrationIds == null) {
+            shouldRemove = true;
+        } else {
+            epRegistrationIds.remove(registrationId);
+            shouldRemove = epRegistrationIds.isEmpty();
+        }
+        if (shouldRemove) {
+            securityStore.remove(endpoint);
+        }
     }
 }

@@ -51,7 +51,7 @@ public class TbLwM2mStoreFactory {
     }
 
     @Bean
-    private TbEditableSecurityStore securityStore() {
+    private TbMainSecurityStore securityStore() {
         return new TbLwM2mSecurityStore(redisConfiguration.isPresent() && useRedis ?
                 new TbLwM2mRedisSecurityStore(redisConfiguration.get().redisConnectionFactory()) : new TbInMemorySecurityStore(), validator);
     }

+/**
+ * Copyright © 2016-2021 The Thingsboard Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.thingsboard.server.transport.lwm2m.server.store;
+
+import org.eclipse.leshan.server.security.NonUniqueSecurityInfoException;
+import org.thingsboard.server.transport.lwm2m.secure.TbLwM2MSecurityInfo;
+
+public interface TbMainSecurityStore extends TbSecurityStore {
+
+    void putX509(TbLwM2MSecurityInfo tbSecurityInfo) throws NonUniqueSecurityInfoException;
+
+    void registerX509(String endpoint, String registrationId);
+
+    void remove(String endpoint, String registrationId);
+
+}