Skip to content

Commit

Permalink
WIP - Still working on this refactor, I like it I think
Browse files Browse the repository at this point in the history
  • Loading branch information
@NovaFox161
NovaFox161 committed Aug 31, 2023
1 parent 0343ad2 commit 7f160ca
Show file tree
Hide file tree
Showing 13 changed files with 270 additions and 84 deletions.
11 changes: 5 additions & 6 deletions cam/src/main/kotlin/org/dreamexposure/discal/cam/endpoints/v1/GetEndpoint.kt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@ import discord4j.common.util.Snowflake
import kotlinx.coroutines.reactor.awaitSingle
import org.dreamexposure.discal.cam.google.GoogleAuth
import org.dreamexposure.discal.core.annotations.Authentication
import org.dreamexposure.discal.core.database.DatabaseManager
import org.dreamexposure.discal.core.business.CalendarService
import org.dreamexposure.discal.core.enums.calendar.CalendarHost
import org.dreamexposure.discal.core.`object`.network.discal.CredentialData
import org.springframework.web.bind.annotation.GetMapping
Expand All @@ -15,22 +15,21 @@ import org.springframework.web.bind.annotation.RestController
@RestController
@RequestMapping("/v1/")
class GetEndpoint(
private val calendarService: CalendarService,
private val googleAuth: GoogleAuth,
) {
@Authentication(access = Authentication.AccessLevel.ADMIN)
@GetMapping("token", produces = ["application/json"])
suspend fun get(@RequestParam host: CalendarHost, @RequestParam id: Int, @RequestParam guild: Snowflake?): CredentialData {
suspend fun get(@RequestParam host: CalendarHost, @RequestParam id: Int, @RequestParam guild: Snowflake?): CredentialData? {
return when (host) {
CalendarHost.GOOGLE -> {
if (guild == null) {
// Internal (owned by DisCal, should never go bad)
googleAuth.requestNewAccessToken(id).awaitSingle()
} else {
// External (owned by user)
// TODO: Replace this db manager call
DatabaseManager.getCalendar(guild, id)
.flatMap(googleAuth::requestNewAccessToken)
.awaitSingle()
val calendar = calendarService.getCalendar(guild, id) ?: return null
googleAuth.requestNewAccessToken(calendar)
}
}
}
Expand Down
140 changes: 71 additions & 69 deletions cam/src/main/kotlin/org/dreamexposure/discal/cam/google/GoogleAuth.kt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,22 +2,24 @@ package org.dreamexposure.discal.cam.google

import com.google.api.client.http.HttpStatusCodes.STATUS_CODE_BAD_REQUEST
import com.google.api.client.http.HttpStatusCodes.STATUS_CODE_OK
import kotlinx.coroutines.reactor.awaitSingle
import kotlinx.coroutines.reactor.mono
import okhttp3.FormBody
import okhttp3.Request
import org.dreamexposure.discal.cam.json.google.ErrorData
import org.dreamexposure.discal.cam.json.google.RefreshData
import org.dreamexposure.discal.core.business.CalendarService
import org.dreamexposure.discal.core.business.CredentialService
import org.dreamexposure.discal.core.config.Config
import org.dreamexposure.discal.core.crypto.AESEncryption
import org.dreamexposure.discal.core.database.DatabaseManager
import org.dreamexposure.discal.core.entities.google.DisCalGoogleCredential
import org.dreamexposure.discal.core.exceptions.AccessRevokedException
import org.dreamexposure.discal.core.exceptions.EmptyNotAllowedException
import org.dreamexposure.discal.core.exceptions.NotFoundException
import org.dreamexposure.discal.core.extensions.isExpiredTtl
import org.dreamexposure.discal.core.logger.LOGGER
import org.dreamexposure.discal.core.`object`.calendar.CalendarData
import org.dreamexposure.discal.core.`object`.network.discal.CredentialData
import org.dreamexposure.discal.core.`object`.new.Calendar
import org.dreamexposure.discal.core.utils.GlobalVal.DEFAULT
import org.dreamexposure.discal.core.utils.GlobalVal.HTTP_CLIENT
import org.dreamexposure.discal.core.utils.GlobalVal.JSON_FORMAT
Expand All @@ -31,8 +33,9 @@ import kotlin.system.exitProcess
@Component
class GoogleAuth(
private val credentialService: CredentialService,
private val calendarService: CalendarService,
) {
private final val CREDENTIALS: Flux<DisCalGoogleCredential>
private final val CREDENTIALS: Flux<DisCalGoogleCredential> // TODO: Refactor as kc

init {
val credCount = Config.SECRET_GOOGLE_CREDENTIAL_COUNT.getInt()
Expand All @@ -44,27 +47,26 @@ class GoogleAuth(
.cache()
}

fun requestNewAccessToken(calendarData: CalendarData): Mono<CredentialData> {
return Mono.just(AESEncryption(calendarData.privateKey)).flatMap { aes ->
if (!calendarData.expired()) {
return@flatMap aes.decrypt(calendarData.encryptedAccessToken)
.map { CredentialData(it, calendarData.expiresAt) }
}
suspend fun requestNewAccessToken(calendar: Calendar): CredentialData? {
val aes = AESEncryption(calendar.secrets.privateKey)
if (!calendar.secrets.expiresAt.isExpiredTtl()) {
return aes.decrypt(calendar.secrets.encryptedAccessToken)
.map { CredentialData(it, calendar.secrets.expiresAt) }
.awaitSingle()
}

aes.decrypt(calendarData.encryptedRefreshToken)
.flatMap(this::doAccessTokenRequest)
.flatMap { data ->
//calendarData.encryptedAccessToken = aes.encrypt(data.accessToken)
calendarData.expiresAt = data.validUntil.minusSeconds(60) // Add a minute of wiggle room
val refreshToken = aes.decrypt(calendar.secrets.encryptedRefreshToken).awaitSingle()
val refreshedCredential = doAccessTokenRequest(refreshToken) ?: return null

aes.encrypt(data.accessToken)
.doOnNext { calendarData.encryptedAccessToken = it }
.then(DatabaseManager.updateCalendar(calendarData).thenReturn(data))//TODO: Replace this
}
}
calendar.secrets.expiresAt = refreshedCredential.validUntil.minusSeconds(60) // Add a minute of wiggle room
calendar.secrets.encryptedAccessToken = aes.encrypt(refreshedCredential.accessToken).awaitSingle()

calendarService.updateCalendar(calendar)

return refreshedCredential
}

fun requestNewAccessToken(credentialId: Int): Mono<CredentialData> {
fun requestNewAccessToken(credentialId: Int): Mono<CredentialData> { // TODO: Refactor as kc
return CREDENTIALS
.filter { it.credential.credentialNumber == credentialId }
.next()
Expand All @@ -76,63 +78,63 @@ class GoogleAuth(
}

credential.getRefreshToken()
.flatMap(this::doAccessTokenRequest)
.flatMap { credential.setAccessToken(it.accessToken).thenReturn(it) }
.flatMap { mono { doAccessTokenRequest(it) } }
.flatMap { credential.setAccessToken(it.accessToken).thenReturn(it) }
.doOnNext { credential.credential.expiresAt = it.validUntil }
.flatMap(mono { credentialService.updateCredential(credential.credential) }::thenReturn)
}.switchIfEmpty(Mono.error(EmptyNotAllowedException()))

}

private fun doAccessTokenRequest(refreshToken: String): Mono<CredentialData> {
return Mono.fromCallable {
val body = FormBody.Builder()
.addEncoded("client_id", Config.SECRET_GOOGLE_CLIENT_ID.getString())
.addEncoded("client_secret", Config.SECRET_GOOGLE_CLIENT_SECRET.getString())
.addEncoded("refresh_token", refreshToken)
.addEncoded("grant_type", "refresh_token")
.build()

val request = Request.Builder()
.url("https://www.googleapis.com/oauth2/v4/token")
.post(body)
.header("Content-Type", "application/x-www-form-urlencoded")
.build()

HTTP_CLIENT.newCall(request).execute()
}.subscribeOn(Schedulers.boundedElastic()).flatMap { response ->
when (response.code) {
STATUS_CODE_OK -> {
val body = JSON_FORMAT.decodeFromString(RefreshData.serializer(), response.body!!.string())
response.body?.close()
response.close()

Mono.just(CredentialData(body.accessToken, Instant.now().plusSeconds(body.expiresIn.toLong())))
}
STATUS_CODE_BAD_REQUEST -> {
val body = JSON_FORMAT.decodeFromString(ErrorData.serializer(), response.body!!.string())
response.body?.close()
response.close()

LOGGER.error("[Google] Access Token Request: $body")

if (body.error == "invalid_grant") {
LOGGER.debug(DEFAULT, "[Google] Access to resource has been revoked")
Mono.error(AccessRevokedException())
} else {
LOGGER.debug(DEFAULT, "[Google] Error requesting new access token | ${response.code} | ${response.message} | $body")
Mono.empty()
}
}
else -> {
// Failed to get OK. Send debug info
LOGGER.debug(DEFAULT, "[Google] Error requesting new access token | ${response.code} " +
"| ${response.message} | ${response.body?.string()}")
response.body?.close()
response.close()
Mono.empty()
private suspend fun doAccessTokenRequest(refreshToken: String): CredentialData? {
val requestFormBody = FormBody.Builder()
.addEncoded("client_id", Config.SECRET_GOOGLE_CLIENT_ID.getString())
.addEncoded("client_secret", Config.SECRET_GOOGLE_CLIENT_SECRET.getString())
.addEncoded("refresh_token", refreshToken)
.addEncoded("grant_type", "refresh_token")
.build()
val request = Request.Builder()
.url("https://www.googleapis.com/oauth2/v4/token")
.post(requestFormBody)
.header("Content-Type", "application/x-www-form-urlencoded")
.build()


val response = Mono.fromCallable(HTTP_CLIENT.newCall(request)::execute)
.subscribeOn(Schedulers.boundedElastic())
.awaitSingle()

return when (response.code) {
STATUS_CODE_OK -> {
val body = JSON_FORMAT.decodeFromString(RefreshData.serializer(), response.body!!.string())
response.body?.close()
response.close()

CredentialData(body.accessToken, Instant.now().plusSeconds(body.expiresIn.toLong()))
}
STATUS_CODE_BAD_REQUEST -> {
val body = JSON_FORMAT.decodeFromString(ErrorData.serializer(), response.body!!.string())
response.body?.close()
response.close()

LOGGER.error("[Google] Access Token Request: $body")

if (body.error == "invalid_grant") {
LOGGER.debug(DEFAULT, "[Google] Access to resource has been revoked")
throw AccessRevokedException()
} else {
LOGGER.error(DEFAULT, "[Google] Error requesting new access token | ${response.code} | ${response.message} | $body")
return null
}
}
else -> {
// Failed to get OK. Send error info
LOGGER.error(DEFAULT, "[Google] Error requesting new access token | ${response.code} ${response.message} | ${response.body?.string()}")
response.body?.close()
response.close()

null
}
}
}
}
65 changes: 65 additions & 0 deletions core/src/main/kotlin/org/dreamexposure/discal/core/business/CalendarService.kt
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,65 @@
package org.dreamexposure.discal.core.business

import discord4j.common.util.Snowflake
import kotlinx.coroutines.reactor.awaitSingle
import org.dreamexposure.discal.CalendarCache
import org.dreamexposure.discal.core.database.CalendarRepository
import org.dreamexposure.discal.core.`object`.new.Calendar
import org.springframework.stereotype.Component

@Component
class DefaultCalendarService(
private val calendarRepository: CalendarRepository,
private val calendarCache: CalendarCache,
) : CalendarService {
override suspend fun getAllCalendars(guildId: Snowflake): List<Calendar> {
var calendars = calendarCache.get(guildId.asLong())?.toList()
if (calendars != null) return calendars

calendars = calendarRepository.findAllByGuildId(guildId.asLong())
.map(::Calendar)
.collectList()
.awaitSingle()

calendarCache.put(guildId.asLong(), calendars.toTypedArray())
return calendars
}

override suspend fun getCalendar(guildId: Snowflake, number: Int): Calendar? {
return getAllCalendars(guildId).first { it.number == number }
}

override suspend fun updateCalendar(calendar: Calendar) {
calendarRepository.updateCalendarByGuildIdAndCalendarNumber(
guildId = calendar.guildId.asLong(),
calendarNumber = calendar.number,
host = calendar.host.name,
calendarId = calendar.id,
calendarAddress = calendar.address,
external = calendar.external,
credentialId = calendar.secrets.credentialId,
privateKey = calendar.secrets.privateKey,
accessToken = calendar.secrets.encryptedAccessToken,
refreshToken = calendar.secrets.encryptedRefreshToken,
expiresAt = calendar.secrets.expiresAt.toEpochMilli(),
).awaitSingle()

val cached = calendarCache.get(calendar.guildId.asLong())
if (cached != null) {
val newList = cached.toMutableList()
newList.removeIf { it.number == calendar.number }
calendarCache.put(calendar.guildId.asLong(), (newList + calendar).toTypedArray())
}
}

}

interface CalendarService {
// TODO: Need a function to invalidate cache because bot and API are using Db Manager

suspend fun getAllCalendars(guildId: Snowflake): List<Calendar>

suspend fun getCalendar(guildId: Snowflake, number: Int): Calendar?

suspend fun updateCalendar(calendar: Calendar)
}
3 changes: 2 additions & 1 deletion core/src/main/kotlin/org/dreamexposure/discal/core/cache/JdkCacheRepository.kt
View file
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
package org.dreamexposure.discal.core.cache

import org.dreamexposure.discal.core.extensions.isExpiredTtl
import reactor.core.publisher.Flux
import java.time.Duration
import java.time.Instant
Expand Down Expand Up @@ -31,7 +32,7 @@ class JdkCacheRepository<K : Any, V>(override val ttl: Duration) : CacheReposito
val cached = cache[key] ?: return null
evict(key)

return if (Instant.now().isAfter(cached.first)) null else cached.second
return if (cached.first.isExpiredTtl()) null else cached.second

}

Expand Down
24 changes: 18 additions & 6 deletions core/src/main/kotlin/org/dreamexposure/discal/core/config/CacheConfig.kt
View file
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
package org.dreamexposure.discal.core.config

import com.fasterxml.jackson.databind.ObjectMapper
import org.dreamexposure.discal.CalendarCache
import org.dreamexposure.discal.CredentialsCache
import org.dreamexposure.discal.OauthStateCache
import org.dreamexposure.discal.core.cache.JdkCacheRepository
Expand All @@ -21,10 +22,12 @@ class CacheConfig {
private val settingsCacheName = "$prefix.settingsCache"
private val credentialsCacheName = "$prefix.credentialsCache"
private val oauthStateCacheName = "$prefix.oauthStateCache"
private val calendarCacheName = "$prefix.calendarCache"

private val settingsTtl = Config.CACHE_TTL_SETTINGS_MINUTES.getLong().asMinutes()
private val credentialsTll = Config.CACHE_TTL_CREDENTIALS_MINUTES.getLong().asMinutes()
private val oauthStateTtl = Config.CACHE_TTL_OAUTH_STATE_MINUTES.getLong().asMinutes()
private val calendarTtl = Config.CACHE_TTL_CALENDAR_MINUTES.getLong().asMinutes()


// Redis caching
Expand All @@ -34,13 +37,13 @@ class CacheConfig {
return RedisCacheManager.builder(connection)
.withCacheConfiguration(settingsCacheName,
RedisCacheConfiguration.defaultCacheConfig().entryTtl(settingsTtl)
)
.withCacheConfiguration(credentialsCacheName,
).withCacheConfiguration(credentialsCacheName,
RedisCacheConfiguration.defaultCacheConfig().entryTtl(credentialsTll)
)
.withCacheConfiguration(oauthStateCacheName,
RedisCacheConfiguration.defaultCacheConfig().entryTtl(oauthStateTtl))
.build()
).withCacheConfiguration(oauthStateCacheName,
RedisCacheConfiguration.defaultCacheConfig().entryTtl(oauthStateTtl)
).withCacheConfiguration(calendarCacheName,
RedisCacheConfiguration.defaultCacheConfig().entryTtl(calendarTtl)
).build()
}

@Bean
Expand All @@ -55,11 +58,20 @@ class CacheConfig {
fun oauthStateRedisCache(cacheManager: RedisCacheManager, objectMapper: ObjectMapper): OauthStateCache =
RedisCacheRepository(cacheManager, objectMapper, oauthStateCacheName)

@Bean
@Primary
@ConditionalOnProperty("bot.cache.redis", havingValue = "true")
fun calendarRedisCache(cacheManager: RedisCacheManager, objectMapper: ObjectMapper): CalendarCache =
RedisCacheRepository(cacheManager, objectMapper, calendarCacheName)


// In-memory fallback caching
@Bean
fun credentialsFallbackCache(): CredentialsCache = JdkCacheRepository(settingsTtl)

@Bean
fun oauthStateFallbackCache(): OauthStateCache = JdkCacheRepository(settingsTtl)

@Bean
fun calendarFallbackCache(): CalendarCache = JdkCacheRepository(calendarTtl)
}
1 change: 1 addition & 0 deletions core/src/main/kotlin/org/dreamexposure/discal/core/config/Config.kt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@ enum class Config(private val key: String, private var value: Any? = null) {
CACHE_TTL_CREDENTIALS_MINUTES("bot.cache.ttl-minutes.credentials", 120),
CACHE_TTL_ACCOUNTS_MINUTES("bot.cache.ttl-minutes.accounts", 60),
CACHE_TTL_OAUTH_STATE_MINUTES("bot.cache.ttl-minutes.oauth.state", 5),
CACHE_TTL_CALENDAR_MINUTES("bots.cache.ttl-minutes.calendar", 120),

// Security configuration

Expand Down
Loading

0 comments on commit 7f160ca

Please sign in to comment.