implement Cipher#auth_data= and support decryption with GCM mode

Author: kares

src/main/java/org/jruby/ext/openssl/Cipher.java

@@ -1090,26 +1090,23 @@ private String getCipherAlgorithm() {
     public IRubyObject update(final ThreadContext context, final IRubyObject arg) {
         final Ruby runtime = context.runtime;
 
-        if ( auth_tag != null ) throw newAuthTagPresentError(runtime);
-
         if ( isDebug(runtime) ) dumpVars( runtime.getOut(), "update()" );
 
         checkCipherNotNull(runtime);
+        checkAuthTag(runtime);
 
         final ByteList data = arg.asString().getByteList();
         final int length = data.length();
         if ( length == 0 ) {
             throw runtime.newArgumentError("data must not be empty");
         }
 
-        if ( ! cipherInited ) {
-            //if ( debug ) runtime.getOut().println("BEFORE INITING");
-            doInitCipher(runtime);
-            //if ( debug ) runtime.getOut().println("AFTER INITING");
-        }
+        if ( ! cipherInited ) doInitCipher(runtime);
 
         final ByteList str;
         try {
+            updateAuthData(runtime); // if any
+
             final byte[] in = data.getUnsafeBytes();
             final int offset = data.begin();
             final byte[] out = cipher.update(in, offset, length);
@@ -1143,43 +1140,32 @@ public IRubyObject update_deprecated(final ThreadContext context, final IRubyObj
     public IRubyObject do_final(final ThreadContext context) {
         final Ruby runtime = context.runtime;
 
-        if ( auth_tag != null ) throw newAuthTagPresentError(runtime);
-
         checkCipherNotNull(runtime);
+        checkAuthTag(runtime);
+
         if ( ! cipherInited ) doInitCipher(runtime);
         // trying to allow update after final like cruby-openssl. Bad idea.
         if ( "RC4".equalsIgnoreCase(cryptoBase) ) return runtime.newString("");
 
-        final ByteList str; boolean shared = false;
+        final ByteList str;
         try {
-            final byte[] out = cipher.doFinal();
-            if ( out != null ) {
-                if ( isAuthDataMode() ) { // TODO only implemented encryption !
-
-                    // if GCM/CCM is being used, the authentication tag is appended
-                    // in the case of encryption, or verified in the case of decryption.
-                    // The result is stored in a new buffer.
-                    final int len = getAuthTagLength(); int strLen;
-                    if ( ( strLen = out.length - len ) > 0 ) {
-                        str = new ByteList(out, 0, strLen, false); shared = true;
-                    }
-                    else {
-                        str = new ByteList(ByteList.NULL_ARRAY); strLen = 0;
-                    }
-                    auth_tag = new ByteList(out, strLen, out.length - strLen);
-                }
-                // TODO: Modifying this line appears to fix the issue, but I do
-                // not have a good reason for why. Best I can tell, lastIv needs
-                // to be set regardless of encryptMode, so we'll go with this
-                // for now. JRUBY-3335.
-                //if ( realIV != null && encryptMode ) ...
-                else {
+            if ( isAuthDataMode() ) {
+                str = do_final_with_auth(runtime);
+            }
+            else {
+                final byte[] out = cipher.doFinal();
+                if ( out != null ) {
+                    // TODO: Modifying this line appears to fix the issue, but I do
+                    // not have a good reason for why. Best I can tell, lastIv needs
+                    // to be set regardless of encryptMode, so we'll go with this
+                    // for now. JRUBY-3335.
+                    //if ( realIV != null && encryptMode ) ...
                     str = new ByteList(out, false);
                     if ( realIV != null ) setLastIVIfNeeded(out);
                 }
-            }
-            else {
-                str = new ByteList(ByteList.NULL_ARRAY);
+                else {
+                    str = new ByteList(ByteList.NULL_ARRAY);
+                }
             }
 
             //if ( ! isStreamCipher() ) {
@@ -1202,18 +1188,46 @@ public IRubyObject do_final(final ThreadContext context) {
             debugStackTrace(runtime, e);
             throw newCipherError(runtime, e);
         }
-        return shared ? RubyString.newStringShared(runtime, str) : RubyString.newString(runtime, str);
+        return RubyString.newString(runtime, str);
     }
 
-    private RaiseException newAuthTagPresentError(final Ruby runtime) {
-        final String error;
-        //if ( encryptMode ) {
-            error = "authentication tag already generated by cipher";
-        //}
-        //else {
-            //error = "authentication tag already consumed by cipher";
-        //}
-        return newCipherError(runtime, error);
+    private ByteList do_final_with_auth(final Ruby runtime) throws GeneralSecurityException {
+        updateAuthData(runtime); // if any
+
+        final ByteList str;
+        // if GCM/CCM is being used, the authentication tag is appended
+        // in the case of encryption, or verified in the case of decryption.
+        // The result is stored in a new buffer.
+        if ( encryptMode ) {
+            final byte[] out = cipher.doFinal();
+
+            final int len = getAuthTagLength(); int strLen;
+            if ( ( strLen = out.length - len ) > 0 ) {
+                str = new ByteList(out, 0, strLen, false);
+            }
+            else {
+                str = new ByteList(ByteList.NULL_ARRAY); strLen = 0;
+            }
+            auth_tag = new ByteList(out, strLen, out.length - strLen);
+            return str;
+        }
+        else {
+            final byte[] out;
+            if ( auth_tag != null ) {
+                final byte[] tag = auth_tag.getUnsafeBytes();
+                out = cipher.doFinal(tag, auth_tag.begin(), auth_tag.length());
+            }
+            else {
+                out = cipher.doFinal();
+            }
+            return new ByteList(out, false);
+        }
+    }
+
+    private void checkAuthTag(final Ruby runtime) {
+        if ( auth_tag != null && encryptMode ) {
+            throw newCipherError(runtime, "authentication tag already generated by cipher");
+        }
     }
 
     private void setLastIVIfNeeded(final byte[] tmpIV) {
@@ -1247,6 +1261,17 @@ public IRubyObject auth_tag(final ThreadContext context) {
         return context.nil;
     }
 
+    @JRubyMethod(name = "auth_tag=")
+    public IRubyObject set_auth_tag(final ThreadContext context, final IRubyObject tag) {
+        if ( ! isAuthDataMode() ) {
+            throw newCipherError(context.runtime, "authentication tag not supported by this cipher");
+        }
+        final RubyString auth_tag = tag.asString();
+        this.auth_tag = auth_tag.getByteList();
+        auth_tag.setByteListShared();
+        return auth_tag;
+    }
+
     private boolean isAuthDataMode() { // Authenticated Encryption with Associated Data (AEAD)
         return "GCM".equalsIgnoreCase(cryptoMode) || "CCM".equalsIgnoreCase(cryptoMode);
     }
@@ -1257,6 +1282,33 @@ private int getAuthTagLength() {
         return Math.min(MAX_AUTH_TAG_LENGTH, this.key.length); // in bytes
     }
 
+    private transient ByteList auth_data;
+
+    @JRubyMethod(name = "auth_data=")
+    public IRubyObject set_auth_data(final ThreadContext context, final IRubyObject data) {
+        if ( ! isAuthDataMode() ) {
+            throw newCipherError(context.runtime, "authentication data not supported by this cipher");
+        }
+        final RubyString auth_data = data.asString();
+        this.auth_data = auth_data.getByteList();
+        auth_data.setByteListShared();
+        return auth_data;
+    }
+
+    private boolean updateAuthData(final Ruby runtime) {
+        if ( auth_data == null ) return false; // only to be set if auth-mode
+        //try {
+            final byte[] data = auth_data.getUnsafeBytes();
+            cipher.updateAAD(data, auth_data.begin(), auth_data.length());
+        //}
+        //catch (RuntimeException e) {
+        //    debugStackTrace( runtime, e );
+        //    throw newCipherError(runtime, e);
+        //}
+        auth_data = null;
+        return true;
+    }
+
     @JRubyMethod(name = "authenticated?")
     public RubyBoolean authenticated_p(final ThreadContext context) {
         return context.runtime.newBoolean( isAuthDataMode() );

src/test/ruby/test_cipher.rb

@@ -363,6 +363,42 @@ def test_aes_128_gcm
     #assert_equal "", cipher.final
   end
 
+  def test_aes_gcm
+    ['aes-128-gcm', 'aes-192-gcm', 'aes-256-gcm'].each do |algo|
+      pt = "You should all use Authenticated Encryption!"
+      cipher, key, iv = new_encryptor(algo)
+
+      cipher.auth_data = "aad"
+      ct  = cipher.update(pt) + cipher.final
+      tag = cipher.auth_tag
+      assert_equal(16, tag.size)
+
+      decipher = new_decryptor(algo, key, iv)
+      decipher.auth_tag = tag
+      decipher.auth_data = "aad"
+
+      assert_equal(pt, decipher.update(ct) + decipher.final)
+    end
+  end
+
+  def new_encryptor(algo)
+    cipher = OpenSSL::Cipher.new(algo)
+    cipher.encrypt
+    key = cipher.random_key
+    iv = cipher.random_iv
+    [cipher, key, iv]
+  end
+  private :new_encryptor
+
+  def new_decryptor(algo, key, iv)
+    OpenSSL::Cipher.new(algo).tap do |cipher|
+      cipher.decrypt
+      cipher.key = key
+      cipher.iv = iv
+    end
+  end
+  private :new_decryptor
+
   def test_aes_128_gcm_with_auth_tag
     cipher = OpenSSL::Cipher.new('aes-128-gcm')
     cipher.encrypt