aead udp support now,

current aead cipher test pass tcp and udp.
TongxiJi · Oct 25, 2018 · f83d0bf · f83d0bf
1 parent cf8355a
commit f83d0bf
Show file tree

Hide file tree

Showing 3 changed files with 56 additions and 27 deletions.
diff --git a/README.md b/README.md
@@ -2,7 +2,7 @@
 A  implementation of Shadowsocks in Java base on netty4 framework.
 
 # Features
-- [x] AEAD Ciphers support(current tcp test pass in game)
+- [x] AEAD Ciphers support
 - [x] TCP & UDP full support
 - [x] DNS proxy optimization
 
@@ -29,7 +29,6 @@ shadowsocks-netty-server.bat
 2. maven package
 
 ## TODO
-* [ ] AEAD Ciphers support udp test
 * [ ] ss-local implementation
 * [ ] ssr obfs features implementation(maybe no use,but for fun)
 * [ ] performance optimization

diff --git a/src/main/java/cn/wowspeeder/encryption/CryptAeadBase.java b/src/main/java/cn/wowspeeder/encryption/CryptAeadBase.java
@@ -28,7 +28,7 @@ public abstract class CryptAeadBase implements ICrypt {
 
     private static byte[] info = "ss-subkey".getBytes();
 
-    private static byte[] ZERO_NONCE = null;
+    private static byte[] ZERO_NONCE = new byte[getNonceLength()];
 
 
     protected final String _name;
@@ -73,10 +73,6 @@ public void isForUdp(boolean isForUdp) {
                 encNonce = new byte[getNonceLength()];
                 decNonce = new byte[getNonceLength()];
             }
-        } else {
-            if (ZERO_NONCE == null) {
-                ZERO_NONCE = new byte[getTagLength()];
-            }
         }
     }
 
@@ -124,7 +120,11 @@ public void encrypt(byte[] data, ByteArrayOutputStream stream) throws Exception
                 encCipher = getCipher(true);
                 _encryptSaltSet = true;
             }
-            _encrypt(data, stream);
+            if (!isForUdp) {
+                _tcpEncrypt(data, stream);
+            } else {
+                _udpEncrypt(data, stream);
+            }
         }
     }
 
@@ -152,7 +152,11 @@ public void decrypt(byte[] data, ByteArrayOutputStream stream) throws Exception
             } else {
                 temp = data;
             }
-            _decrypt(temp, stream);
+            if (!isForUdp) {
+                _tcpDecrypt(temp, stream);
+            } else {
+                _udpDecrypt(temp, stream);
+            }
         }
     }
 
@@ -163,24 +167,34 @@ public void decrypt(byte[] data, int length, ByteArrayOutputStream stream) throw
         decrypt(d, stream);
     }
 
-    private byte[] randomBytes(int size) {
+    private static byte[] randomBytes(int size) {
         byte[] bytes = new byte[size];
         new SecureRandom().nextBytes(bytes);
         return bytes;
     }
 
+    private static int getNonceLength() {
+        return 12;
+    }
+
+    protected static int getTagLength() {
+        return 16;
+    }
+
     protected abstract AEADBlockCipher getCipher(boolean isEncrypted)
             throws GeneralSecurityException;
 
-    protected abstract void _encrypt(byte[] data, ByteArrayOutputStream stream) throws Exception;
+    protected abstract void _tcpEncrypt(byte[] data, ByteArrayOutputStream stream) throws Exception;
+
+    protected abstract void _tcpDecrypt(byte[] data, ByteArrayOutputStream stream) throws Exception;
+
+    protected abstract void _udpEncrypt(byte[] data, ByteArrayOutputStream stream) throws Exception;
 
-    protected abstract void _decrypt(byte[] data, ByteArrayOutputStream stream) throws Exception;
+    protected abstract void _udpDecrypt(byte[] data, ByteArrayOutputStream stream) throws Exception;
 
     protected abstract int getKeyLength();
 
     protected abstract int getSaltLength();
 
-    protected abstract int getNonceLength();
 
-    protected abstract int getTagLength();
 }
diff --git a/src/main/java/cn/wowspeeder/encryption/impl/AesGcmCrypt.java b/src/main/java/cn/wowspeeder/encryption/impl/AesGcmCrypt.java
@@ -77,16 +77,6 @@ public int getSaltLength() {
         return 0;
     }
 
-    @Override
-    protected int getNonceLength() {
-        return 12;
-    }
-
-    @Override
-    protected int getTagLength() {
-        return 16;
-    }
-
     /**
      * TCP:[encrypted payload length][length tag][encrypted payload][payload tag]
      * UDP:[salt][encrypted payload][tag]
@@ -98,7 +88,7 @@ protected int getTagLength() {
      * @throws IOException
      */
     @Override
-    protected void _encrypt(byte[] data, ByteArrayOutputStream stream) throws GeneralSecurityException, IOException, InvalidCipherTextException {
+    protected void _tcpEncrypt(byte[] data, ByteArrayOutputStream stream) throws GeneralSecurityException, IOException, InvalidCipherTextException {
 //        byte[] buffer = new byte[data.length];
 //        int noBytesProcessed = encCipher.processBytes(data, 0, data.length, buffer, 0);
 //        stream.write(buffer, 0, noBytesProcessed);
@@ -133,11 +123,11 @@ protected void _encrypt(byte[] data, ByteArrayOutputStream stream) throws Genera
      * @throws InvalidCipherTextException
      */
     @Override
-    protected void _decrypt(byte[] data, ByteArrayOutputStream stream) throws InvalidCipherTextException {
+    protected void _tcpDecrypt(byte[] data, ByteArrayOutputStream stream) throws InvalidCipherTextException {
 //        byte[] buffer = new byte[data.length];
 //        int noBytesProcessed = decCipher.processBytes(data, 0, data.length, buffer,
 //                0);
-//        logger.debug("remaining _decrypt");
+//        logger.debug("remaining _tcpDecrypt");
 //        stream.write(buffer, 0, noBytesProcessed);
 //        logger.debug("ciphertext len:{}", data.length);
         ByteBuffer buffer = ByteBuffer.wrap(data);
@@ -195,4 +185,30 @@ protected void _decrypt(byte[] data, ByteArrayOutputStream stream) throws Invali
 //            logger.debug("cipher text decode finish");
         }
     }
+
+    @Override
+    protected void _udpEncrypt(byte[] data, ByteArrayOutputStream stream) throws Exception {
+        ByteBuffer buffer = ByteBuffer.wrap(data);
+        int remaining = buffer.remaining();
+        buffer.get(encBuffer, 0, remaining);
+        encCipher.init(true, getCipherParameters(true));
+        encCipher.doFinal(
+                encBuffer,
+                encCipher.processBytes(encBuffer, 0, remaining, encBuffer, 0)
+        );
+        stream.write(encBuffer, 0, remaining + getTagLength());
+    }
+
+    @Override
+    protected void _udpDecrypt(byte[] data, ByteArrayOutputStream stream) throws Exception {
+        ByteBuffer buffer = ByteBuffer.wrap(data);
+        int remaining = buffer.remaining();
+        buffer.get(decBuffer, 0, remaining);
+        decCipher.init(false, getCipherParameters(false));
+        decCipher.doFinal(
+                decBuffer,
+                decCipher.processBytes(decBuffer, 0, remaining, decBuffer, 0)
+        );
+        stream.write(decBuffer, 0, remaining - getTagLength());
+    }
 }