diff --git a/java/ql/lib/semmle/code/java/security/InsufficientKeySizeQuery.qll b/java/ql/lib/semmle/code/java/security/InsufficientKeySizeQuery.qll
index 4cf274f3263..c7cb3a10331 100644
--- a/java/ql/lib/semmle/code/java/security/InsufficientKeySizeQuery.qll
+++ b/java/ql/lib/semmle/code/java/security/InsufficientKeySizeQuery.qll
@@ -2,6 +2,11 @@ import semmle.code.java.security.Encryption
 import semmle.code.java.dataflow.TaintTracking
 import semmle.code.java.dataflow.DataFlow
 
+// TODO:
+// todo #1: make representation of source that can be shared across the configs
+// todo #2: make representation of sink that can be shared across the configs
+// todo #3: finish adding tracking for algo type/name... need flow/taint-tracking for across methods??
+// todo #3a: make list of algo names more easily reusable (either as constant-type variable at top of file, or model as own class to share, etc.)
 // ******* DATAFLOW BELOW *************************************************************************
 /**
  * Asymmetric (RSA, DSA, DH) key length data flow tracking configuration.
@@ -10,13 +15,27 @@ class AsymmetricKeyTrackingConfiguration extends DataFlow::Configuration {
   AsymmetricKeyTrackingConfiguration() { this = "AsymmetricKeyTrackingConfiguration" }
 
   override predicate isSource(DataFlow::Node source) {
-    source.asExpr() instanceof IntegerLiteral and // ! this works with current test cases, but reconsider IntegerLiteral when variables are used
-    source.toString().toInt() < 2048
+    exists(ClassInstanceExpr rsaGenParamSpec |
+      rsaGenParamSpec.getConstructedType() instanceof RSAGenParameterSpec and // ! double-check if should just use getType() instead
+      rsaGenParamSpec.getArgument(0).(IntegerLiteral).getIntValue() < 2048 and
+      source.asExpr() = rsaGenParamSpec
+    )
+    or
+    source.asExpr().(IntegerLiteral).getIntValue() < 2048
   }
 
   override predicate isSink(DataFlow::Node sink) {
-    exists(MethodAccess ma |
+    exists(MethodAccess ma, VarAccess va |
       ma.getMethod() instanceof KeyPairGeneratorInitMethod and
+      va.getVariable()
+          .getAnAssignedValue()
+          .(JavaSecurityKeyPairGenerator)
+          .getAlgoSpec()
+          .(StringLiteral)
+          .getValue()
+          .toUpperCase()
+          .matches(["RSA", "DSA", "DH"]) and
+      ma.getQualifier() = va and
       sink.asExpr() = ma.getArgument(0)
     )
   }
@@ -30,15 +49,26 @@ class AsymmetricECCKeyTrackingConfiguration extends DataFlow::Configuration {
 
   override predicate isSource(DataFlow::Node source) {
     exists(ClassInstanceExpr ecGenParamSpec |
-      getECKeySize(ecGenParamSpec.getArgument(0).(StringLiteral).getValue()) < 256 and
+      getECKeySize(ecGenParamSpec.getArgument(0).(StringLiteral).getValue()) < 256 and // ! can generate EC with just the keysize and not the curve apparently... (based on netty/netty FP example)
       source.asExpr() = ecGenParamSpec
     )
+    or
+    source.asExpr().(IntegerLiteral).getIntValue() < 256
   }
 
   override predicate isSink(DataFlow::Node sink) {
-    exists(MethodAccess ma |
+    exists(MethodAccess ma, VarAccess va |
       ma.getMethod() instanceof KeyPairGeneratorInitMethod and
-      ma.getArgument(0).getType() instanceof ECGenParameterSpec and
+      //ma.getArgument(0).getType() instanceof ECGenParameterSpec and // ! can generate EC with just the keysize and not the curve apparently... (based on netty/netty FP example)
+      va.getVariable()
+          .getAnAssignedValue()
+          .(JavaSecurityKeyPairGenerator)
+          .getAlgoSpec()
+          .(StringLiteral)
+          .getValue()
+          .toUpperCase()
+          .matches(["EC%"]) and
+      ma.getQualifier() = va and
       sink.asExpr() = ma.getArgument(0)
     )
   }
@@ -51,13 +81,21 @@ class SymmetricKeyTrackingConfiguration extends DataFlow::Configuration {
   SymmetricKeyTrackingConfiguration() { this = "SymmetricKeyTrackingConfiguration2" }
 
   override predicate isSource(DataFlow::Node source) {
-    source.asExpr() instanceof IntegerLiteral and // ! this works with current test cases, but reconsider IntegerLiteral when variables are used
-    source.toString().toInt() < 128
+    source.asExpr().(IntegerLiteral).getIntValue() < 128
   }
 
   override predicate isSink(DataFlow::Node sink) {
-    exists(MethodAccess ma |
+    exists(MethodAccess ma, VarAccess va |
       ma.getMethod() instanceof KeyGeneratorInitMethod and
+      va.getVariable()
+          .getAnAssignedValue()
+          .(JavaxCryptoKeyGenerator)
+          .getAlgoSpec()
+          .(StringLiteral)
+          .getValue()
+          .toUpperCase()
+          .matches(["AES"]) and
+      ma.getQualifier() = va and
       sink.asExpr() = ma.getArgument(0)
     )
   }
@@ -77,6 +115,11 @@ private class ECGenParameterSpec extends RefType {
   ECGenParameterSpec() { this.hasQualifiedName("java.security.spec", "ECGenParameterSpec") }
 }
 
+/** The Java class `java.security.spec.ECGenParameterSpec`. */
+private class RSAGenParameterSpec extends RefType {
+  RSAGenParameterSpec() { this.hasQualifiedName("java.security.spec", "RSAKeyGenParameterSpec") }
+}
+
 // ! move to Encryption.qll?
 /** Returns the key size in the EC algorithm string */
 bindingset[algorithm]
diff --git a/java/ql/src/Security/CWE/CWE-326/InsufficientKeySize.ql b/java/ql/src/Security/CWE/CWE-326/InsufficientKeySize.ql
index 93411a101b2..c6dce740588 100644
--- a/java/ql/src/Security/CWE/CWE-326/InsufficientKeySize.ql
+++ b/java/ql/src/Security/CWE/CWE-326/InsufficientKeySize.ql
@@ -28,7 +28,6 @@ import DataFlow::PathGraph
 //   sink.getNode(), "size"
 from DataFlow::PathNode source, DataFlow::PathNode sink
 where
-  //hasInsufficientKeySize2(source, sink)
   exists(AsymmetricKeyTrackingConfiguration config1 | config1.hasFlowPath(source, sink)) or
   exists(AsymmetricECCKeyTrackingConfiguration config2 | config2.hasFlowPath(source, sink)) or
   exists(SymmetricKeyTrackingConfiguration config3 | config3.hasFlowPath(source, sink))
diff --git a/java/ql/test/query-tests/security/CWE-326/InsufficientKeySizeTest.java b/java/ql/test/query-tests/security/CWE-326/InsufficientKeySizeTest.java
index 50dc7b4637f..b383892361b 100644
--- a/java/ql/test/query-tests/security/CWE-326/InsufficientKeySizeTest.java
+++ b/java/ql/test/query-tests/security/CWE-326/InsufficientKeySizeTest.java
@@ -1,5 +1,6 @@
 import java.security.KeyPairGenerator;
 import java.security.spec.ECGenParameterSpec;
+import java.security.spec.RSAKeyGenParameterSpec;
 import javax.crypto.KeyGenerator;
 
 public class InsufficientKeySizeTest {
@@ -27,6 +28,16 @@ public class InsufficientKeySizeTest {
             // GOOD: Key size is no less than 2048
             KeyPairGenerator keyPairGen2 = KeyPairGenerator.getInstance("RSA");
             keyPairGen2.initialize(2048); // Safe
+
+            // test with spec
+            // BAD: Key size is less than 2048
+            KeyPairGenerator keyPairGen3 = KeyPairGenerator.getInstance("RSA");
+            RSAKeyGenParameterSpec rsaSpec = new RSAKeyGenParameterSpec(1024, null);
+            keyPairGen3.initialize(rsaSpec); // $ hasInsufficientKeySize
+
+            // BAD: Key size is less than 2048
+            KeyPairGenerator keyPairGen4 = KeyPairGenerator.getInstance("RSA");
+            keyPairGen4.initialize(new RSAKeyGenParameterSpec(1024, null)); // $ hasInsufficientKeySize
         }
 
         // DSA (Asymmetric)
@@ -145,7 +156,7 @@ public class InsufficientKeySizeTest {
             int size = 64; // test integer variable
             KeyGenerator keyGen = KeyGenerator.getInstance("AES"); // test KeyGenerator variable
             testSymmetric(size, keyGen); // test with variable as key size
-            testSymmetric2(64); // test with int constant as key size
+            testSymmetric2(64); // test with int literal as key size
          }
 
 
@@ -153,15 +164,16 @@ public class InsufficientKeySizeTest {
          {
             int size = 1024; // test integer variable
             KeyPairGenerator keyPairGen21 = KeyPairGenerator.getInstance("RSA"); // test KeyPairGenerator variable
-            testAsymmetricNonEC(size, keyPairGen21);
+            testAsymmetricNonEC(size, keyPairGen21); // test with variable as key size
+            testAsymmetricNonEC2(1024); // test with int literal as key size
          }
 
          // Test variable passed to other method(s) - Asymmetric, EC
          {
             ECGenParameterSpec ecSpec = new ECGenParameterSpec("secp112r1"); // test ECGenParameterSpec variable
             KeyPairGenerator keyPairGen22 = KeyPairGenerator.getInstance("EC"); // test KeyPairGenerator variable
-            testAsymmetricEC(ecSpec, keyPairGen22); // test with variable as key size
-            testAsymmetricNonEC2(1024); // test with int constant as key size
+            testAsymmetricEC(ecSpec, keyPairGen22);
+
          }
 
     }