Refactor SymmetricKeyEncSessionPacket

Author: peterdettman

pg/src/main/java/org/bouncycastle/bcpg/SymmetricKeyEncSessionPacket.java

@@ -29,15 +29,13 @@ public class SymmetricKeyEncSessionPacket
      */
     public static final int VERSION_6 = 6;
 
-    private int version;          // V4, V5, V6
-    private int encAlgorithm;     // V4, V5, V6
-    private S2K s2k;              // V4,
-    // array for exposing raw S2K parameters. Useful for forwards compat.
-    private byte[] s2kBytes;         // Makes only sense for v6, since there we have a counter
-    private byte[] secKeyData;       // V4, V5, V6
-    private int aeadAlgorithm;    // V5, V6
-    private byte[] iv;               // V5, V6
-    private byte[] authTag;          // V5, V6
+    private final int version;          // V4, V5, V6
+    private final int encAlgorithm;     // V4, V5, V6
+    private final int aeadAlgorithm;    // V5, V6
+    private S2K s2k;                    // V4, V5, V6
+    private byte[] secKeyData;          // V4, V5, V6
+    private byte[] iv;                  // V5, V6
+    private byte[] authTag;             // V5, V6
 
     public SymmetricKeyEncSessionPacket(
         BCPGInputStream in)
@@ -57,6 +55,7 @@ public SymmetricKeyEncSessionPacket(
         if (version == VERSION_4)
         {
             encAlgorithm = in.read();
+            aeadAlgorithm = 0;
 
             s2k = new S2K(in);
 
@@ -207,6 +206,7 @@ public SymmetricKeyEncSessionPacket(
 
         this.version = VERSION_4;
         this.encAlgorithm = encAlgorithm;
+        this.aeadAlgorithm = 0;
         this.s2k = s2k;
         this.secKeyData = secKeyData;
     }
@@ -339,51 +339,63 @@ public void encode(
         BCPGOutputStream out)
         throws IOException
     {
+        PacketFormat packetFormat = version > 4 ? PacketFormat.CURRENT : PacketFormat.ROUNDTRIP;
+
         ByteArrayOutputStream bOut = new ByteArrayOutputStream();
-        BCPGOutputStream pOut;
-        if (version == 4)
-        {
-            pOut = new BCPGOutputStream(bOut);
-        }
-        else
-        {
-            pOut = new BCPGOutputStream(bOut, true);
-        }
+        BCPGOutputStream pOut = new BCPGOutputStream(bOut, packetFormat);
 
         pOut.write(version);
-        if (version == VERSION_4)
+
+        switch (version)
+        {
+        case VERSION_4:
         {
             pOut.write(encAlgorithm);
-            pOut.writeObject(s2k);
+            s2k.encode(pOut);
 
             if (secKeyData != null && secKeyData.length > 0)
             {
                 pOut.write(secKeyData);
             }
+            break;
         }
-        else
+        case VERSION_5:
         {
-            int s2kLen = 0;
-            if (version == VERSION_6)
-            {
-                s2kLen = s2k.getEncoded().length;
-                int count = 1 + 1 + 1 + s2kLen + iv.length;
-                pOut.write(count); // len of 5 following fields
-            }
             pOut.write(encAlgorithm);
             pOut.write(aeadAlgorithm);
-            if (version == VERSION_6)
+            s2k.encode(pOut);
+            pOut.write(iv);
+
+            if (secKeyData != null && secKeyData.length > 0)
             {
-                pOut.write(s2kLen);
+                pOut.write(secKeyData);
             }
-            pOut.writeObject(s2k);
+
+            pOut.write(authTag);
+            break;
+        }
+        case VERSION_6:
+        {
+            byte[] s2kEncoded = s2k.getEncoded();
+            int count = 1 + 1 + 1 + s2kEncoded.length + iv.length; // len of 5 following fields
+
+            pOut.write(count);
+            pOut.write(encAlgorithm);
+            pOut.write(aeadAlgorithm);
+            pOut.write(s2kEncoded.length);
+            pOut.write(s2kEncoded);
             pOut.write(iv);
 
             if (secKeyData != null && secKeyData.length > 0)
             {
                 pOut.write(secKeyData);
             }
+
             pOut.write(authTag);
+            break;
+        }
+        default:
+            throw new IllegalStateException();
         }
 
         pOut.close();