I think I might have finally figured it out?

That, or the changes I just made should at least be a small optimization, we'll see.
This commit is contained in:
NovaFox161
2023-10-01 18:09:53 -05:00
parent 09735b305d
commit ce3fee274e
6 changed files with 63 additions and 21 deletions

View File

@@ -33,7 +33,6 @@ class GoogleAuth(
private val calendarService: CalendarService,
private val objectMapper: ObjectMapper,
) {
private final val aes: AESEncryption = AESEncryption(Config.SECRET_GOOGLE_CREDENTIAL_KEY.getString())
suspend fun requestNewAccessToken(calendar: Calendar): CredentialData? {
val aes = AESEncryption(calendar.secrets.privateKey)
@@ -57,16 +56,13 @@ class GoogleAuth(
suspend fun requestNewAccessToken(credentialId: Int): CredentialData {
val credential = credentialService.getCredential(credentialId) ?: throw NotFoundException()
if (!credential.expiresAt.isExpiredTtl()) {
val accessToken = aes.decrypt(credential.encryptedAccessToken).awaitSingle()
return CredentialData(accessToken, credential.expiresAt)
}
if (!credential.expiresAt.isExpiredTtl()) return CredentialData(credential.accessToken, credential.expiresAt)
LOGGER.debug("Refreshing access token | credentialId:$credentialId")
val refreshToken = aes.decrypt(credential.encryptedRefreshToken).awaitSingle()
val refreshedCredentialData = doAccessTokenRequest(refreshToken) ?: throw EmptyNotAllowedException()
credential.encryptedAccessToken = aes.encrypt(refreshedCredentialData.accessToken).awaitSingle()
val refreshedCredentialData = doAccessTokenRequest(credential.refreshToken) ?: throw EmptyNotAllowedException()
credential.accessToken = refreshedCredentialData.accessToken
credential.expiresAt = refreshedCredentialData.validUntil.minusSeconds(60) // Add a minute of wiggle room
credentialService.updateCredential(credential)

View File

@@ -204,12 +204,14 @@ class AnnouncementService(
private fun getEvents(guild: Guild, announcement: Announcement): Flux<Event> {
val cached = getCached(announcement.guildId)
if (cached.events.contains(announcement.calendarNumber))
return Flux.fromIterable(cached.events[announcement.calendarNumber]!!)
return if (!cached.events.contains(announcement.calendarNumber)) {
getCalendar(guild, announcement).flatMapMany {
it.getUpcomingEvents(20).cache()
}
} else cached.events[announcement.calendarNumber]!!
return getCalendar(guild, announcement).flatMapMany {
it.getUpcomingEvents(20)
}.collectList()
.doOnNext { cached.events[announcement.calendarNumber] = it }
.flatMapIterable { it }
}
private fun getCached(guildId: Snowflake): AnnouncementCache {

View File

@@ -26,6 +26,7 @@ class AESEncryption(privateKey: String) {
}
}
@Deprecated("Use #decryptFixed(string) instead")
fun encrypt(data: String): Mono<String> {
return Mono.fromCallable {
this.cipher?.init(Cipher.ENCRYPT_MODE, this.secretKeySpec, this.ivParameterSpec)
@@ -39,6 +40,20 @@ class AESEncryption(privateKey: String) {
}.subscribeOn(Schedulers.single()).switchIfEmpty(Mono.error(EmptyNotAllowedException()))
}
fun encryptFixed(data: String): String {
return try {
this.cipher?.init(Cipher.ENCRYPT_MODE, this.secretKeySpec, this.ivParameterSpec)
val encrypted = this.cipher?.doFinal(data.toByteArray(StandardCharsets.UTF_8))
Base64.encodeBase64String(encrypted)
} catch (ex: Exception) {
LOGGER.error("Encrypt failure", ex)
throw IllegalStateException("Encrypt Failure", ex)
}
}
@Deprecated("Use #decryptFixed(string) instead")
fun decrypt(data: String): Mono<String> {
return Mono.fromCallable {
this.cipher?.init(Cipher.DECRYPT_MODE, this.secretKeySpec, this.ivParameterSpec)
@@ -51,4 +66,16 @@ class AESEncryption(privateKey: String) {
Mono.error(IllegalStateException("Decrypt Failure", it))
}.subscribeOn(Schedulers.single()).switchIfEmpty(Mono.error(EmptyNotAllowedException()))
}
fun decryptFixed(data: String): String {
return try {
this.cipher?.init(Cipher.DECRYPT_MODE, this.secretKeySpec, this.ivParameterSpec)
val decrypted = this.cipher?.doFinal(Base64.decodeBase64(data))
String(decrypted!!, StandardCharsets.UTF_8)
} catch (ex: Exception) {
LOGGER.error("Decrypt failure", ex)
throw IllegalStateException("Decrypt Failure", ex)
}
}
}

View File

@@ -3,11 +3,10 @@ package org.dreamexposure.discal.core.`object`.announcement
import discord4j.common.util.Snowflake
import org.dreamexposure.discal.core.entities.Calendar
import org.dreamexposure.discal.core.entities.Event
import reactor.core.publisher.Flux
import java.util.concurrent.ConcurrentHashMap
data class AnnouncementCache(
val id: Snowflake,
val calendars: ConcurrentHashMap<Int, Calendar> = ConcurrentHashMap(),
val events: ConcurrentHashMap<Int, Flux<Event>> = ConcurrentHashMap(),
val events: ConcurrentHashMap<Int, List<Event>> = ConcurrentHashMap(),
)

View File

@@ -35,8 +35,8 @@ data class Calendar(
data class Secrets(
val credentialId: Int,
val privateKey: String,
val encryptedRefreshToken: String,
var encryptedAccessToken: String,
val encryptedRefreshToken: String, // TODO: Secrets should be unencrypted immediately before/after Db write/read respectively
var encryptedAccessToken: String, // TODO: Secrets should be unencrypted immediately before/after Db write/read respectively
var expiresAt: Instant,
)
}

View File

@@ -1,19 +1,37 @@
package org.dreamexposure.discal.core.`object`.new
import org.dreamexposure.discal.core.config.Config
import org.dreamexposure.discal.core.crypto.AESEncryption
import org.dreamexposure.discal.core.database.CredentialData
import org.dreamexposure.discal.core.extensions.asInstantMilli
import java.time.Instant
data class Credential(
val credentialNumber: Int,
var encryptedRefreshToken: String,
var encryptedAccessToken: String,
var expiresAt: Instant,
var refreshToken: String,
var accessToken: String,
) {
/**
* Returns the encrypted refresh token, note that this will run the encryption on every access.
*/
val encryptedRefreshToken: String
get() = aes.encryptFixed(refreshToken)
/**
* Returns the encrypted access token, note that this will run the encryption on every access.
*/
val encryptedAccessToken: String
get() = aes.encryptFixed(accessToken)
constructor(data: CredentialData) : this(
credentialNumber = data.credentialNumber,
encryptedRefreshToken = data.refreshToken,
encryptedAccessToken = data.accessToken,
expiresAt = data.expiresAt.asInstantMilli(),
refreshToken = aes.decryptFixed(data.refreshToken),
accessToken = aes.decryptFixed(data.accessToken),
)
companion object {
private val aes = AESEncryption(Config.SECRET_GOOGLE_CREDENTIAL_KEY.getString())
}
}