1 files changed, 103 insertions, 0 deletions
diff --git a/java/com/tigervnc/rdr/AESOutStream.java b/java/com/tigervnc/rdr/AESOutStream.java
new file mode 100644
index 00000000..e0014615
--- /dev/null
+++ b/java/com/tigervnc/rdr/AESOutStream.java
@@ -0,0 +1,103 @@
+/* Copyright (C) 2022 Dinglan Peng
+ *
+ * This is free software; you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation; either version 2 of the License, or
+ * (at your option) any later version.
+ *
+ * This software is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this software; if not, write to the Free Software
+ * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301,
+ * USA.
+ */
+
+package com.tigervnc.rdr;
+
+import java.nio.ByteBuffer;
+
+public class AESOutStream extends OutStream {
+
+  static final int maxMessageSize = 8192;
+
+  public AESOutStream(OutStream _out, byte[] key)
+  {
+    out = _out;
+    bufSize = maxMessageSize;
+    b = new byte[bufSize];
+    buffer = new byte[bufSize + 16 + 2];
+    ptr = offset = start = 0;
+    end = start + bufSize;
+    cipher = new AESEAXCipher(key);
+    counter = new byte[16];
+  }
+
+  public int length()
+  {
+    return offset + ptr - start;
+  }
+
+  public void flush()
+  {
+    int sentUpTo = start;
+    while (sentUpTo < ptr) {
+      int n = writeMessage(b, sentUpTo, ptr - sentUpTo);
+      sentUpTo += n;
+      offset += n;
+    }
+
+    ptr = start;
+  }
+
+  protected int overrun(int itemSize, int nItems)
+  {
+    if (itemSize > bufSize)
+      throw new Exception("AESOutStream overrun: max itemSize exceeded");
+
+    flush();
+
+    int nAvail;
+    nAvail = (end - ptr) / itemSize;
+    if (nAvail < nItems)
+      return nAvail;
+
+    return nItems;
+  }
+
+  protected int writeMessage(byte[] data, int dataPtr, int length)
+  {
+    if (length == 0)
+      return 0;
+    buffer[0] = (byte)((length & 0xff00) >> 8);
+    buffer[1] = (byte)(length & 0xff);
+    cipher.encrypt(data, dataPtr, length,
+                   buffer, 0, 2,
+                   counter,
+                   buffer, 2,
+                   buffer, 2 + length);
+    out.writeBytes(ByteBuffer.wrap(buffer, 0, length + 16 + 2),
+                   length + 16 + 2);
+    out.flush();
+    // Update nonce by incrementing the counter as a
+    // 128bit little endian unsigned integer
+    for (int i = 0; i < 16; ++i) {
+      // increment until there is no carry
+      if (++counter[i] != 0) {
+        break;
+      }
+    }
+    return length;
+  }
+
+  private AESEAXCipher cipher;
+  private int start;
+  private int offset;
+  private int bufSize;
+  private byte[] buffer;
+  private byte[] counter;
+  private OutStream out;
+}