Feat/optimised file writing (#15)

* Vault token renewal

Adds a new parameter for the Vault docs to drive the auth token renewal.

* Add file headers

* Complete the Vault ticket renewal.

* Avoid file to be created over and over again

For the AES decoder there was a strategy to create a file (using random UUID) every time the "key" was returned. The new code is changing the approach to make sure we only create the file once. This means though for the AES encoding the encoding format should contain an id for a file:${encoding}_${id}. The id has to be unique as well.

Small code improvement around decoding

* Fix compilation

Author: stheppi

src/main/scala/io/lenses/connect/secrets/io/FileWriter.scala

@@ -0,0 +1,52 @@
+/*
+ *
+ *  * Copyright 2017-2020 Lenses.io Ltd
+ *
+ */
+package io.lenses.connect.secrets.io
+
+import java.io.BufferedOutputStream
+import java.io.FileOutputStream
+import java.nio.file.Path
+import java.nio.file.Paths
+
+import com.typesafe.scalalogging.StrictLogging
+import io.lenses.connect.secrets.utils.WithRetry
+
+import scala.concurrent.duration._
+import scala.util.Try
+
+trait FileWriter {
+  def write(fileName: String, content: Array[Byte], key: String): Path
+}
+
+class FileWriterOnce(rootPath: Path)
+    extends FileWriter
+    with WithRetry
+    with StrictLogging {
+  rootPath.toFile.mkdirs()
+  def write(fileName: String, content: Array[Byte], key: String): Path = {
+    val fullPath = Paths.get(rootPath.toString, fileName)
+    val file = fullPath.toFile
+    if (file.exists()) fullPath
+    else {
+      val tempFile = Paths.get(rootPath.toString, fileName + ".bak").toFile
+      withRetry(10, Some(500.milliseconds)) {
+        if (tempFile.exists()) tempFile.delete()
+        tempFile.createNewFile()
+        val fos = new BufferedOutputStream(new FileOutputStream(tempFile))
+        try {
+          fos.write(content)
+          fos.flush()
+          logger.info(
+            s"Payload written to [${file.getAbsolutePath}] for key [$key]"
+          )
+        } finally {
+          Try(fos.close())
+        }
+        Try(tempFile.renameTo(file))
+      }
+      fullPath
+    }
+  }
+}

src/main/scala/io/lenses/connect/secrets/package.scala

@@ -6,23 +6,26 @@
 
 package io.lenses.connect.secrets
 
-import java.io.{File, FileOutputStream}
+import java.io.File
+import java.io.FileOutputStream
 import java.time.OffsetDateTime
 import java.util.Base64
 
 import com.typesafe.scalalogging.StrictLogging
 import org.apache.kafka.common.config.ConfigData
 import org.apache.kafka.connect.errors.ConnectException
 
-import scala.util.{Failure, Success, Try}
 import scala.collection.JavaConverters._
 import scala.collection.mutable
+import scala.util.Failure
+import scala.util.Success
+import scala.util.Try
 
 package object connect extends StrictLogging {
 
-  val FILE_ENCODING = "file-encoding"
-  val FILE_DIR = "file.dir"
-  val FILE_DIR_DESC =
+  val FILE_ENCODING: String = "file-encoding"
+  val FILE_DIR: String = "file.dir"
+  val FILE_DIR_DESC: String =
     """
       | Location to write any files for any secrets that need to
       | be written to disk. For example java keystores.
@@ -77,32 +80,19 @@ package object connect extends StrictLogging {
   }
 
   // decode a key bases on the prefix encoding
-  def decodeKey(key: String, value: String, fileName: String): String = {
-
-    key.toLowerCase match {
-      case k
-          if k.startsWith(
-            s"${Encoding.BASE64_FILE.toString.toLowerCase}_"
-          ) =>
+  def decodeKey(
+      encoding: Option[Encoding.Value],
+      key: String,
+      value: String,
+      writeFileFn: Array[Byte] => String
+  ): String = {
+    encoding.fold(value) {
+      case Encoding.BASE64 => decode(key, value)
+      case Encoding.BASE64_FILE =>
         val decoded = decodeToBytes(key, value)
-        fileWriter(fileName, decoded, key)
-        fileName
-
-      case k
-          if k.startsWith(
-            s"${Encoding.BASE64.toString.toLowerCase}_"
-          ) =>
-        decode(key, value)
-
-      case k
-          if k.startsWith(
-            s"${Encoding.UTF8_FILE.toString.toLowerCase}_"
-          ) =>
-        fileWriter(fileName, value.getBytes(), key)
-        fileName
-
-      case _ =>
-        value
+        writeFileFn(decoded)
+      case Encoding.UTF8      => value
+      case Encoding.UTF8_FILE => writeFileFn(value.getBytes())
     }
   }
 
@@ -154,8 +144,8 @@ package object connect extends StrictLogging {
 
   //calculate the min expiry for secrets and return the configData and expiry
   def getSecretsAndExpiry(
-      secrets: Map[String, (String, Option[OffsetDateTime])])
-    : (Option[OffsetDateTime], ConfigData) = {
+      secrets: Map[String, (String, Option[OffsetDateTime])]
+  ): (Option[OffsetDateTime], ConfigData) = {
     var expiryList = mutable.ListBuffer.empty[OffsetDateTime]
 
     val data = secrets
@@ -176,6 +166,11 @@ package object connect extends StrictLogging {
     }
   }
 
-  def getFileName(rootDir: String, path: String, key: String, separator: String): String =
+  def getFileName(
+      rootDir: String,
+      path: String,
+      key: String,
+      separator: String
+  ): String =
     s"${rootDir.stripSuffix(separator)}$separator$path$separator${key.toLowerCase}"
 }

src/main/scala/io/lenses/connect/secrets/providers/AWSHelper.scala

@@ -7,6 +7,7 @@
 package io.lenses.connect.secrets.providers
 
 import java.nio.file.FileSystems
+import java.nio.file.Paths
 import java.time.OffsetDateTime
 import java.util.Calendar
 
@@ -16,7 +17,10 @@ import com.amazonaws.services.secretsmanager.{AWSSecretsManager, AWSSecretsManag
 import com.fasterxml.jackson.databind.ObjectMapper
 import com.typesafe.scalalogging.StrictLogging
 import io.lenses.connect.secrets.config.AWSProviderSettings
-import io.lenses.connect.secrets.connect.{AuthMode, decodeKey, getFileName}
+import io.lenses.connect.secrets.connect.{decodeKey, getFileName, AuthMode}
+import io.lenses.connect.secrets.io.FileWriter
+import io.lenses.connect.secrets.io.FileWriterOnce
+import io.lenses.connect.secrets.utils.EncodingAndId
 import org.apache.kafka.connect.errors.ConnectException
 
 import scala.collection.JavaConverters._
@@ -108,12 +112,17 @@ trait AWSHelper extends StrictLogging {
               )
             )
 
+        val fileWriter:FileWriter = new FileWriterOnce(Paths.get(rootDir, secretId))
         // decode the value
+        val encodingAndId = EncodingAndId.from(key)
         (
           decodeKey(
             key = key,
             value = value,
-            fileName = getFileName(rootDir, secretId, key.toLowerCase, separator)
+            encoding = encodingAndId.encoding,
+            writeFileFn = content=>{
+              fileWriter.write(key.toLowerCase, content, key).toString
+            }
           ),
           getTTL(client, secretId)
         )

src/main/scala/io/lenses/connect/secrets/providers/Aes256DecodingHelper.scala

@@ -1,16 +1,14 @@
 package io.lenses.connect.secrets.providers
 
+import java.security.SecureRandom
+import java.util.Base64
+
 import javax.crypto.Cipher
-import javax.crypto.SecretKey
 import javax.crypto.spec.IvParameterSpec
 import javax.crypto.spec.SecretKeySpec
 
-import scala.util.Try
-
-import java.security.SecureRandom
-import java.util.Base64
 import scala.util.Failure
-import scala.util.Success
+import scala.util.Try
 
 private[providers] object Aes256DecodingHelper {
 
@@ -35,13 +33,14 @@ private[providers] object Aes256DecodingHelper {
     }
 }
 
-private[providers] class Aes256DecodingHelper private (
-    key: String,
-    ivSeparator: String
-) {
-  import B64._
+private[providers] class Aes256DecodingHelper private(
+                                                       key: String,
+                                                       ivSeparator: String
+                                                     ) {
+
   import Aes256DecodingHelper.CHARSET
-  import InitializationVector._
+  import B64._
+
   private val secret = new SecretKeySpec(key.getBytes(CHARSET), "AES")
 
   def decrypt(s: String): Try[String] =
@@ -52,9 +51,9 @@ private[providers] class Aes256DecodingHelper private (
     } yield new String(decrypted, CHARSET)
 
   private def decryptBytes(
-      iv: InitializationVector,
-      bytes: Array[Byte]
-  ): Try[Array[Byte]] =
+                            iv: InitializationVector,
+                            bytes: Array[Byte]
+                          ): Try[Array[Byte]] =
     for {
       cipher <- getCipher(Cipher.DECRYPT_MODE, iv)
       encrypted <- Try(cipher.doFinal(bytes))
@@ -69,10 +68,12 @@ private[providers] class Aes256DecodingHelper private (
     }
 }
 
-private case class InitializationVector private (bytes: Array[Byte])
+private case class InitializationVector private(bytes: Array[Byte])
 
 private object InitializationVector {
+
   import B64._
+
   private val random = new SecureRandom()
   private val length = 16
 
@@ -84,9 +85,9 @@ private object InitializationVector {
   }
 
   def extractInitialisationVector(
-      s: String,
-      ivSeparator: String
-  ): Try[(InitializationVector, String)] =
+                                   s: String,
+                                   ivSeparator: String
+                                 ): Try[(InitializationVector, String)] =
     s.indexOf(ivSeparator) match {
       case -1 =>
         Failure(

src/main/scala/io/lenses/connect/secrets/providers/Aes256DecodingProvider.scala

@@ -1,55 +1,64 @@
 package io.lenses.connect.secrets.providers
 
+import java.nio.file.Paths
+import java.util
+
+import io.lenses.connect.secrets.config.Aes256ProviderConfig
+import io.lenses.connect.secrets.connect.decodeKey
+import io.lenses.connect.secrets.connect.Encoding
+import io.lenses.connect.secrets.io.FileWriter
+import io.lenses.connect.secrets.io.FileWriterOnce
+import io.lenses.connect.secrets.utils.EncodingAndId
 import org.apache.kafka.common.config.ConfigData
 import org.apache.kafka.common.config.provider.ConfigProvider
-import org.apache.kafka.connect.errors.ConnectException
 import org.apache.kafka.common.config.ConfigException
-import java.util
-import javax.crypto.SecretKey
-import javax.crypto.SecretKeyFactory
-import javax.crypto.spec.PBEKeySpec
-import javax.crypto.spec.SecretKeySpec
-import io.lenses.connect.secrets.config.Aes256ProviderConfig
-import io.lenses.connect.secrets.connect.{decodeKey, getFileName}
+import org.apache.kafka.connect.errors.ConnectException
+
 import scala.collection.JavaConverters._
-import scala.util.Try
-import io.lenses.connect.secrets.connect.Encoding
-import java.nio.file.FileSystems
-import java.util.UUID.randomUUID
+
 
 class Aes256DecodingProvider extends ConfigProvider {
 
   var decoder: Option[Aes256DecodingHelper] = None
   var writeDir: String = ""
-  
+
+  private var fileWriter:FileWriter = _
   override def configure(configs: util.Map[String, _]): Unit = {
     val aes256Cfg = Aes256ProviderConfig(configs)
     val aes256Key = aes256Cfg.aes256Key
     decoder = Option(aes256Key)
       .map(Aes256DecodingHelper.init)
       .map(_.fold(e => throw new ConfigException(e), identity))
     writeDir = aes256Cfg.writeDirectory
+    fileWriter = new FileWriterOnce(Paths.get(writeDir, "secrets"))
   }
 
   override def get(path: String): ConfigData = new ConfigData(Map.empty[String, String].asJava)
-  
-  override def get(path: String, keys: util.Set[String]): ConfigData =
+
+  override def get(path: String, keys: util.Set[String]): ConfigData = {
+    val encodingAndId = EncodingAndId.from(path)
     decoder match {
       case Some(d) =>
         def decrypt(key: String): String = {
-          val decrypted = d.decrypt(key).fold(e => throw new ConnectException(e.getMessage(), e), identity)
-          val keyPrefixedWithEncoding = if (path.nonEmpty) s"${path.toLowerCase}_$key" else key
+          val decrypted = d.decrypt(key).fold(e => throw new ConnectException("Failed to decrypt the secret.", e), identity)
           decodeKey(
-            key = keyPrefixedWithEncoding,
+            key = key,
             value = decrypted,
-            fileName = getFileName(writeDir, "secrets", randomUUID().toString, FileSystems.getDefault.getSeparator)
+            encoding = encodingAndId.encoding,
+            writeFileFn = { content =>
+              encodingAndId.id match {
+                case Some(value) => fileWriter.write(value, content, key).toString
+                case None => throw new ConnectException(s"Invalid argument received for key:$key. Expecting a file identifier.")
+              }
+            }
           )
         }
-          
+
         new ConfigData(keys.asScala.map(k => k -> decrypt(k)).toMap.asJava)
       case None =>
         throw new ConnectException("decoder is not configured.")
     }
+  }
 
   override def close(): Unit = {}
 }