support tagChunkSize replacements at arbitrary depths

Dwolla · Jan 30, 2024 · df9f1a3 · df9f1a3
1 parent e74ff4c
commit df9f1a3
Show file tree

Hide file tree

Showing 3 changed files with 46 additions and 4 deletions.
diff --git a/scalafix/input/src/main/scala/example/SeveralDifferentUses.scala b/scalafix/input/src/main/scala/example/SeveralDifferentUses.scala
@@ -4,6 +4,7 @@ rule = com.dwolla.security.crypto.V04to05
 */
 import cats.effect.*
 import com.dwolla.security.crypto.*
+import eu.timepit.refined.auto.*
 import org.bouncycastle.openpgp.PGPPublicKey
 import org.typelevel.log4cats.LoggerFactory
 
@@ -16,13 +17,19 @@ object SeveralDifferentUses {
 
     // all arguments set
     alg.encrypt(key, chunkSize, Option("filename"), Encryption.Aes256, Compression.Bzip2, PgpLiteralDataPacketFormat.Utf8)
+    alg.encrypt(key, tagChunkSize(42), Option("filename"), Encryption.Aes256, Compression.Bzip2, PgpLiteralDataPacketFormat.Utf8)
 
     // all arguments set with named parameters
     alg.encrypt(key = key, chunkSize = chunkSize, fileName = Option("filename"), encryption = Encryption.Aes256, compression = Compression.Bzip2, packetFormat = PgpLiteralDataPacketFormat.Utf8)
+    alg.encrypt(key = key, chunkSize = {
+      tagChunkSize(42)
+    }, fileName = Option("filename"), encryption = Encryption.Aes256, compression = Compression.Bzip2, packetFormat = PgpLiteralDataPacketFormat.Utf8)
     alg.encrypt(key, chunkSize = chunkSize, fileName = Option("filename"), encryption = Encryption.Aes256, compression = Compression.Bzip2, packetFormat = PgpLiteralDataPacketFormat.Utf8)
 
     // random assortments
-    alg.encrypt(key, chunkSize, Option("filename"))
+    alg.encrypt(key, {
+      tagChunkSize(42)
+    }, Option("filename"))
     alg.encrypt(key, fileName = Option("filename"), chunkSize = chunkSize)
     alg.encrypt(key = key, encryption = Encryption.Aes256, packetFormat = PgpLiteralDataPacketFormat.Utf8)
     alg.encrypt(key, fileName = Option("filename"), compression = Compression.Bzip2)

diff --git a/scalafix/output/src/main/scala/example/SeveralDifferentUses.scala b/scalafix/output/src/main/scala/example/SeveralDifferentUses.scala
@@ -2,6 +2,7 @@ package example
 
 import cats.effect.*
 import com.dwolla.security.crypto.*
+import eu.timepit.refined.auto.*
 import org.bouncycastle.openpgp.PGPPublicKey
 import org.typelevel.log4cats.LoggerFactory
 
@@ -14,13 +15,19 @@ object SeveralDifferentUses {
 
     // all arguments set
     alg.encrypt(EncryptionConfig().withChunkSize(chunkSize).withFileName(Option("filename")).withEncryption(Encryption.Aes256).withCompression(Compression.Bzip2).withPacketFormat(PgpLiteralDataPacketFormat.Utf8), key)
+    alg.encrypt(EncryptionConfig().withChunkSize(ChunkSize(42)).withFileName(Option("filename")).withEncryption(Encryption.Aes256).withCompression(Compression.Bzip2).withPacketFormat(PgpLiteralDataPacketFormat.Utf8), key)
 
     // all arguments set with named parameters
     alg.encrypt(EncryptionConfig().withChunkSize(chunkSize).withFileName(Option("filename")).withEncryption(Encryption.Aes256).withCompression(Compression.Bzip2).withPacketFormat(PgpLiteralDataPacketFormat.Utf8), key)
+    alg.encrypt(EncryptionConfig().withChunkSize({
+      ChunkSize(42)
+    }).withFileName(Option("filename")).withEncryption(Encryption.Aes256).withCompression(Compression.Bzip2).withPacketFormat(PgpLiteralDataPacketFormat.Utf8), key)
     alg.encrypt(EncryptionConfig().withChunkSize(chunkSize).withFileName(Option("filename")).withEncryption(Encryption.Aes256).withCompression(Compression.Bzip2).withPacketFormat(PgpLiteralDataPacketFormat.Utf8), key)
 
     // random assortments
-    alg.encrypt(EncryptionConfig().withChunkSize(chunkSize).withFileName(Option("filename")), key)
+    alg.encrypt(EncryptionConfig().withChunkSize({
+      ChunkSize(42)
+    }).withFileName(Option("filename")), key)
     alg.encrypt(EncryptionConfig().withFileName(Option("filename")).withChunkSize(chunkSize), key)
     alg.encrypt(EncryptionConfig().withEncryption(Encryption.Aes256).withPacketFormat(PgpLiteralDataPacketFormat.Utf8), key)
     alg.encrypt(EncryptionConfig().withFileName(Option("filename")).withCompression(Compression.Bzip2), key)

diff --git a/scalafix/rules/src/main/scala/com/dwolla/security/crypto/V04to05.scala b/scalafix/rules/src/main/scala/com/dwolla/security/crypto/V04to05.scala
@@ -50,6 +50,7 @@ class V04to05 extends SemanticRule("com.dwolla.security.crypto.V04to05") {
                 val parameterName = parameter.displayName
 
                 s :+ (parameterName.capitalize -> value)
+              case _ => throw new RuntimeException(s"Unexpected method signature ${info.signature} (how did the code being modified ever compile?)")
             }
           case _ => s
         }
@@ -58,11 +59,38 @@ class V04to05 extends SemanticRule("com.dwolla.security.crypto.V04to05") {
 
     Patch.replaceTree(t, sortKeyToEnd(map).foldLeft(s"(EncryptionConfig()") {
       case (s, ("key", term)) => s + s", $term)"
-      case (s, (argName, Term.Apply.After_4_6_0(Term.Name("tagChunkSize"), Term.ArgClause(List(term), _)))) => s + s".with$argName(ChunkSize($term))"
-      case (s, (argName, term)) => s + s".with$argName($term)"
+      case (s, (argName, term)) =>
+        tagChunkSizeTermReplacement(term)
+          .map(updatedChunkSize => s + s".with$argName($updatedChunkSize)")
+          .getOrElse(s + s".with$argName($term)")
     })
   }
 
+  private def tagChunkSizeTermReplacement(term: Term): Option[String] = {
+    val chunkSizeTerms = term.collect {
+      case t@Term.Apply.After_4_6_0(Term.Name("tagChunkSize"), Term.ArgClause(List(_), _)) => t.pos
+    }
+
+    val updatedChunkSize =
+      chunkSizeTerms
+        .foldLeft(new StringBuilder) { case (acc, t) =>
+          val tokens = term.tokens.toString()
+
+          val replacementLength = "tagChunkSize".length
+          val start = t.start - term.pos.start + replacementLength
+          val end = tokens.length
+
+          acc
+            .append(tokens.slice(0, t.start - term.pos.start).mkString)
+            .append("ChunkSize")
+            .append(tokens.slice(start, end).mkString)
+        }
+        .toString()
+
+    if (chunkSizeTerms.isEmpty) None
+    else Some(updatedChunkSize)
+  }
+
   private def sortKeyToEnd(terms: List[(String, Term)]): List[(String, Term)] = {
     val key = terms.toMap.apply("key")