Skip to content

Commit 159086b

Browse files
JanCizmarJanCizmar
committed
feat: Allow websocket subscription with PAK/PAT
1 parent ebc8acc commit 159086b

File tree

8 files changed

+419
-28
lines changed

8 files changed

+419
-28
lines changed

‎backend/api/src/main/kotlin/io/tolgee/websocket/WebSocketConfig.kt‎

Lines changed: 37 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,7 @@
11
package io.tolgee.websocket
22

3-
import io.tolgee.dtos.cacheable.UserAccountDto
3+
import io.tolgee.dtos.cacheable.ApiKeyDto
44
import io.tolgee.model.enums.Scope
5-
import io.tolgee.security.authentication.JwtService
65
import io.tolgee.security.authentication.TolgeeAuthentication
76
import io.tolgee.service.security.SecurityService
87
import org.springframework.context.annotation.Configuration
@@ -23,10 +22,10 @@ import org.springframework.web.socket.config.annotation.WebSocketMessageBrokerCo
2322
@Configuration
2423
@EnableWebSocketMessageBroker
2524
class WebSocketConfig(
26-
@Lazy
27-
private val jwtService: JwtService,
2825
@Lazy
2926
private val securityService: SecurityService,
27+
@Lazy
28+
private val websocketAuthenticationResolver: WebsocketAuthenticationResolver,
3029
) : WebSocketMessageBrokerConfigurer {
3130
override fun configureMessageBroker(config: MessageBrokerRegistry) {
3231
config.enableSimpleBroker("/")
@@ -46,15 +45,17 @@ class WebSocketConfig(
4645
val accessor = MessageHeaderAccessor.getAccessor(message, StompHeaderAccessor::class.java)
4746

4847
if (accessor?.command == StompCommand.CONNECT) {
49-
val tokenString = accessor.getNativeHeader("jwtToken")?.firstOrNull()
50-
accessor.user = if (tokenString == null) null else jwtService.validateToken(tokenString)
48+
val authorization = accessor.getNativeHeader("authorization")?.firstOrNull()
49+
val xApiKey = accessor.getNativeHeader("x-api-key")?.firstOrNull()
50+
val legacyJwt = accessor.getNativeHeader("jwtToken")?.firstOrNull()
51+
accessor.user = websocketAuthenticationResolver.resolve(authorization, xApiKey, legacyJwt)
5152
}
5253

53-
val user = (accessor?.user as? TolgeeAuthentication)?.principal
54+
val authentication = accessor?.user as? TolgeeAuthentication
5455

5556
if (accessor?.command == StompCommand.SUBSCRIBE) {
56-
checkProjectPathPermissions(user, accessor.destination)
57-
checkUserPathPermissions(user, accessor.destination)
57+
checkProjectPathPermissionsAuth(authentication, accessor.destination)
58+
checkUserPathPermissionsAuth(authentication, accessor.destination)
5859
}
5960

6061
return message
@@ -63,8 +64,8 @@ class WebSocketConfig(
6364
)
6465
}
6566

66-
fun checkProjectPathPermissions(
67-
user: UserAccountDto?,
67+
fun checkProjectPathPermissionsAuth(
68+
authentication: TolgeeAuthentication?,
6869
destination: String?,
6970
) {
7071
val projectId =
@@ -73,19 +74,30 @@ class WebSocketConfig(
7374
?.getOrNull(1)?.toLong()
7475
} ?: return
7576

76-
if (user == null) {
77+
if (authentication == null) {
7778
throw MessagingException("Unauthenticated")
7879
}
7980

81+
val creds = authentication.credentials
82+
if (creds is ApiKeyDto) {
83+
val matchesProject = creds.projectId == projectId
84+
val hasScope = creds.scopes.contains(Scope.KEYS_VIEW)
85+
if (!matchesProject || !hasScope) {
86+
throw MessagingException("Forbidden")
87+
}
88+
return
89+
}
90+
91+
val user = authentication.principal
8092
try {
8193
securityService.checkProjectPermissionNoApiKey(projectId = projectId, Scope.KEYS_VIEW, user)
8294
} catch (e: Exception) {
8395
throw MessagingException("Forbidden")
8496
}
8597
}
8698

87-
fun checkUserPathPermissions(
88-
user: UserAccountDto?,
99+
fun checkUserPathPermissionsAuth(
100+
authentication: TolgeeAuthentication?,
89101
destination: String?,
90102
) {
91103
val userId =
@@ -94,6 +106,17 @@ class WebSocketConfig(
94106
?.getOrNull(1)?.toLong()
95107
} ?: return
96108

109+
if (authentication == null) {
110+
throw MessagingException("Forbidden")
111+
}
112+
113+
val creds = authentication.credentials
114+
if (creds is ApiKeyDto) {
115+
// API keys must not subscribe to user topics
116+
throw MessagingException("Forbidden")
117+
}
118+
119+
val user = (authentication as? TolgeeAuthentication)?.principal
97120
if (user?.id != userId) {
98121
throw MessagingException("Forbidden")
99122
}

‎backend/api/src/main/kotlin/io/tolgee/websocket/WebsocketAuthenticationResolver.kt‎

Lines changed: 104 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,104 @@
1+
package io.tolgee.websocket
2+
3+
import io.tolgee.constants.Message
4+
import io.tolgee.dtos.cacheable.UserAccountDto
5+
import io.tolgee.exceptions.AuthenticationException
6+
import io.tolgee.security.PAT_PREFIX
7+
import io.tolgee.security.authentication.JwtService
8+
import io.tolgee.security.authentication.TolgeeAuthentication
9+
import io.tolgee.security.authentication.TolgeeAuthenticationDetails
10+
import io.tolgee.service.security.ApiKeyService
11+
import io.tolgee.service.security.PatService
12+
import io.tolgee.service.security.UserAccountService
13+
import org.springframework.context.annotation.Lazy
14+
import org.springframework.stereotype.Component
15+
16+
@Component
17+
class WebsocketAuthenticationResolver(
18+
@Lazy private val jwtService: JwtService,
19+
@Lazy private val apiKeyService: ApiKeyService,
20+
@Lazy private val patService: PatService,
21+
@Lazy private val userAccountService: UserAccountService,
22+
) {
23+
/**
24+
* Resolves STOMP CONNECT headers into TolgeeAuthentication.
25+
* Supports:
26+
* - Authorization: Bearer <jwt>
27+
* - X-API-Key: tgpat_<token> (PAT) or tgpak_<...> (PAK, incl. legacy/raw)
28+
* - jwtToken: <jwt> (legacy header)
29+
*/
30+
fun resolve(
31+
authorizationHeader: String?,
32+
xApiKeyHeader: String?,
33+
legacyJwtHeader: String?,
34+
): TolgeeAuthentication? {
35+
// Authorization: Bearer <jwt>
36+
val bearer = extractBearer(authorizationHeader)
37+
if (bearer != null) {
38+
return runCatching { jwtService.validateToken(bearer) }.getOrNull()
39+
}
40+
41+
// X-API-Key: PAT / PAK
42+
val xApiKey = xApiKeyHeader
43+
if (!xApiKey.isNullOrBlank()) {
44+
return when {
45+
xApiKey.startsWith(PAT_PREFIX) -> runCatching { patAuth(xApiKey) }.getOrNull()
46+
else -> runCatching { pakAuth(xApiKey) }.getOrNull()
47+
}
48+
}
49+
50+
// Legacy jwtToken header
51+
if (!legacyJwtHeader.isNullOrBlank()) {
52+
return runCatching { jwtService.validateToken(legacyJwtHeader) }.getOrNull()
53+
}
54+
55+
return null
56+
}
57+
58+
private fun extractBearer(value: String?): String? {
59+
if (value == null) return null
60+
val prefix = "Bearer "
61+
return if (value.startsWith(prefix, ignoreCase = true)) value.substring(prefix.length).trim() else null
62+
}
63+
64+
private fun pakAuth(key: String): TolgeeAuthentication {
65+
val parsed = apiKeyService.parseApiKey(key) ?: throw AuthenticationException(Message.INVALID_PROJECT_API_KEY)
66+
val hash = apiKeyService.hashKey(parsed)
67+
val pak = apiKeyService.findDto(hash) ?: throw AuthenticationException(Message.INVALID_PROJECT_API_KEY)
68+
69+
if (pak.expiresAt?.before(java.util.Date()) == true) {
70+
throw AuthenticationException(Message.PROJECT_API_KEY_EXPIRED)
71+
}
72+
73+
val userAccount: UserAccountDto =
74+
userAccountService.findDto(pak.userAccountId) ?: throw AuthenticationException(Message.USER_NOT_FOUND)
75+
76+
apiKeyService.updateLastUsedAsync(pak.id)
77+
78+
return TolgeeAuthentication(
79+
pak,
80+
userAccount,
81+
TolgeeAuthenticationDetails(false),
82+
)
83+
}
84+
85+
private fun patAuth(key: String): TolgeeAuthentication {
86+
val hash = patService.hashToken(key.substring(PAT_PREFIX.length))
87+
val pat = patService.findDto(hash) ?: throw AuthenticationException(Message.INVALID_PAT)
88+
89+
if (pat.expiresAt?.before(java.util.Date()) == true) {
90+
throw AuthenticationException(Message.PAT_EXPIRED)
91+
}
92+
93+
val userAccount: UserAccountDto =
94+
userAccountService.findDto(pat.userAccountId) ?: throw AuthenticationException(Message.USER_NOT_FOUND)
95+
96+
patService.updateLastUsedAsync(pat.id)
97+
98+
return TolgeeAuthentication(
99+
pat,
100+
userAccount,
101+
TolgeeAuthenticationDetails(false),
102+
)
103+
}
104+
}

‎backend/app/src/test/kotlin/io/tolgee/batch/BatchJobTestUtil.kt‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -261,7 +261,7 @@ class BatchJobTestUtil(
261261
websocketHelper =
262262
WebsocketTestHelper(
263263
port,
264-
jwtService.emitToken(testData.user.id),
264+
WebsocketTestHelper.Auth(jwtToken = jwtService.emitToken(testData.user.id)),
265265
testData.projectBuilder.self.id,
266266
testData.user.id,
267267
)

‎backend/app/src/test/kotlin/io/tolgee/websocket/AbstractWebsocketTest.kt‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -45,14 +45,14 @@ abstract class AbstractWebsocketTest : ProjectAuthControllerTest("/v2/projects/"
4545
currentUserWebsocket =
4646
WebsocketTestHelper(
4747
port,
48-
jwtService.emitToken(testData.user.id),
48+
WebsocketTestHelper.Auth(jwtToken = jwtService.emitToken(testData.user.id)),
4949
testData.projectBuilder.self.id,
5050
testData.user.id,
5151
)
5252
anotherUserWebsocket =
5353
WebsocketTestHelper(
5454
port,
55-
jwtService.emitToken(anotherUser.id),
55+
WebsocketTestHelper.Auth(jwtToken = jwtService.emitToken(anotherUser.id)),
5656
testData.projectBuilder.self.id,
5757
anotherUser.id,
5858
)
@@ -239,12 +239,13 @@ abstract class AbstractWebsocketTest : ProjectAuthControllerTest("/v2/projects/"
239239
val spyingUserWebsocket =
240240
WebsocketTestHelper(
241241
port,
242-
jwtService.emitToken(anotherUser.id),
242+
WebsocketTestHelper.Auth(jwtToken = jwtService.emitToken(anotherUser.id)),
243243
testData.projectBuilder.self.id,
244244
// anotherUser trying to spy on other user's websocket
245245
testData.user.id,
246246
)
247247
spyingUserWebsocket.listenForNotificationsChanged()
248+
spyingUserWebsocket.waitForForbidden()
248249
saveNotificationForCurrentUser()
249250

250251
assertCurrentUserReceivedMessage()

0 commit comments

Comments
 (0)