feat(webflux): add data masking function to load aggregate state (#1055)

Ahoo-Wang · Dec 18, 2024 · 5ad7127 · 5ad7127
1 parent 6f4e6f6
commit 5ad7127
Show file tree

Hide file tree

Showing 5 changed files with 35 additions and 9 deletions.
diff --git a/wow-query/src/main/kotlin/me/ahoo/wow/query/mask/DataMasking.kt b/wow-query/src/main/kotlin/me/ahoo/wow/query/mask/DataMasking.kt
@@ -20,6 +20,14 @@ interface DataMasking<SOURCE : DataMasking<SOURCE>> : DataMasker {
     fun mask(): SOURCE
 }
 
+fun <S : Any> S.tryMask(): S {
+    if (this is DataMasking<*>) {
+        @Suppress("UNCHECKED_CAST")
+        return this.mask() as S
+    }
+    return this
+}
+
 fun <S : Any> MaterializedSnapshot<S>.tryMask(): MaterializedSnapshot<S> {
     val state = this.state
     if (state is DataMasking<*>) {

diff --git a/wow-query/src/test/kotlin/me/ahoo/wow/query/mask/DataMaskingKtTest.kt b/wow-query/src/test/kotlin/me/ahoo/wow/query/mask/DataMaskingKtTest.kt
@@ -15,8 +15,9 @@ package me.ahoo.wow.query.mask
 
 import me.ahoo.wow.api.query.MaterializedSnapshot
 import me.ahoo.wow.api.query.PagedList
-import org.hamcrest.CoreMatchers
-import org.hamcrest.MatcherAssert
+import org.hamcrest.CoreMatchers.equalTo
+import org.hamcrest.CoreMatchers.sameInstance
+import org.hamcrest.MatcherAssert.*
 import org.junit.jupiter.api.Test
 
 class DataMaskingKtTest {
@@ -39,7 +40,7 @@ class DataMaskingKtTest {
             deleted = false
         )
         val masked = snapshot.tryMask()
-        MatcherAssert.assertThat(masked, CoreMatchers.sameInstance(snapshot))
+        assertThat(masked, sameInstance(snapshot))
     }
 
     @Test
@@ -60,7 +61,7 @@ class DataMaskingKtTest {
             deleted = false
         )
         val masked = snapshot.tryMask()
-        MatcherAssert.assertThat(masked.state.pwd, CoreMatchers.equalTo("******"))
+        assertThat(masked.state.pwd, equalTo("******"))
     }
 
     @Test
@@ -82,7 +83,7 @@ class DataMaskingKtTest {
         )
         val pagedList = PagedList(1, listOf(snapshot))
         val masked = pagedList.tryMask()
-        MatcherAssert.assertThat(masked, CoreMatchers.sameInstance(pagedList))
+        assertThat(masked, sameInstance(pagedList))
     }
 
     @Test
@@ -104,7 +105,21 @@ class DataMaskingKtTest {
         )
         val pagedList = PagedList(1, listOf(snapshot))
         val masked = pagedList.tryMask()
-        MatcherAssert.assertThat(masked.list.first().state.pwd, CoreMatchers.equalTo("******"))
+        assertThat(masked.list.first().state.pwd, equalTo("******"))
+    }
+
+    @Test
+    fun tryMaskState() {
+        val state = MockMaskingData("pwd")
+        val masked = state.tryMask()
+        assertThat(masked.pwd, equalTo("******"))
+    }
+
+    @Test
+    fun tryMaskAnyState() {
+        val source = Any()
+        val masked = source.tryMask()
+        assertThat(masked, sameInstance(source))
     }
 
     data class MockMaskingData(val pwd: String) : DataMasking<MockMaskingData> {

diff --git a/...x/src/main/kotlin/me/ahoo/wow/webflux/route/state/AbstractLoadAggregateHandlerFunction.kt b/...x/src/main/kotlin/me/ahoo/wow/webflux/route/state/AbstractLoadAggregateHandlerFunction.kt
@@ -19,6 +19,7 @@ import me.ahoo.wow.modeling.matedata.AggregateMetadata
 import me.ahoo.wow.modeling.state.StateAggregate
 import me.ahoo.wow.modeling.state.StateAggregateRepository
 import me.ahoo.wow.openapi.RoutePaths
+import me.ahoo.wow.query.mask.tryMask
 import me.ahoo.wow.webflux.exception.RequestExceptionHandler
 import me.ahoo.wow.webflux.exception.toServerResponse
 import me.ahoo.wow.webflux.route.command.CommandParser.getTenantIdOrDefault
@@ -48,7 +49,7 @@ abstract class AbstractLoadAggregateHandlerFunction(
             }
             .map {
                 checkVersion(version, it)
-                it.state
+                it.state.tryMask()
             }
             .throwNotFoundIfEmpty()
             .toServerResponse(request, exceptionHandler)

diff --git a/...bflux/src/main/kotlin/me/ahoo/wow/webflux/route/state/IdsQueryAggregateHandlerFunction.kt b/...bflux/src/main/kotlin/me/ahoo/wow/webflux/route/state/IdsQueryAggregateHandlerFunction.kt
@@ -17,6 +17,7 @@ import me.ahoo.wow.modeling.aggregateId
 import me.ahoo.wow.modeling.matedata.AggregateMetadata
 import me.ahoo.wow.modeling.state.StateAggregateRepository
 import me.ahoo.wow.openapi.state.IdsQueryAggregateRouteSpec
+import me.ahoo.wow.query.mask.tryMask
 import me.ahoo.wow.webflux.exception.RequestExceptionHandler
 import me.ahoo.wow.webflux.exception.toServerResponse
 import me.ahoo.wow.webflux.route.RouteHandlerFunctionFactory
@@ -47,7 +48,7 @@ class IdsQueryAggregateHandlerFunction(
             }.filter {
                 it.initialized && !it.deleted
             }
-            .map { it.state }
+            .map { it.state.tryMask() }
             .collectList()
             .toServerResponse(request, exceptionHandler)
     }

diff --git a/wow-webflux/src/main/kotlin/me/ahoo/wow/webflux/route/state/ScanAggregateHandlerFunction.kt b/wow-webflux/src/main/kotlin/me/ahoo/wow/webflux/route/state/ScanAggregateHandlerFunction.kt
@@ -18,6 +18,7 @@ import me.ahoo.wow.modeling.matedata.AggregateMetadata
 import me.ahoo.wow.modeling.state.StateAggregateRepository
 import me.ahoo.wow.openapi.RoutePaths
 import me.ahoo.wow.openapi.state.ScanAggregateRouteSpec
+import me.ahoo.wow.query.mask.tryMask
 import me.ahoo.wow.webflux.exception.RequestExceptionHandler
 import me.ahoo.wow.webflux.exception.toServerResponse
 import me.ahoo.wow.webflux.route.RouteHandlerFunctionFactory
@@ -44,7 +45,7 @@ class ScanAggregateHandlerFunction(
             stateAggregateRepository.load<Any>(it)
         }.filter {
             it.initialized && !it.deleted
-        }.map { it.state }
+        }.map { it.state.tryMask() }
             .collectList()
             .toServerResponse(request, exceptionHandler)
     }