add AsyncSocketSecurity

AsyncSocket/src/main/kotlin/cn/tursom/socket/AsyncSocket.kt

@@ -50,7 +50,6 @@ interface AsyncSocket : Closeable {
   /**
    * 在有数据读取的时候自动由内存池分配内存
    */
-  @Throws(SocketException::class)
   suspend fun read(pool: MemoryPool, timeout: Long = 0L): ByteBuffer = read(timeout) {
     val buffer = pool.get()
     if (channel.read(buffer) < 0) throw SocketException()

AsyncSocket/src/main/kotlin/cn/tursom/socket/security/AsyncSocketSecurityUtil.kt

@@ -0,0 +1,37 @@
+package cn.tursom.socket.security
+
+import cn.tursom.core.buffer.impl.HeapByteBuffer
+import cn.tursom.core.encrypt.AES
+import cn.tursom.core.encrypt.RSA
+import cn.tursom.core.pool.HeapMemoryPool
+import cn.tursom.socket.AsyncSocket
+import cn.tursom.socket.enhance.EnhanceSocket
+import cn.tursom.socket.enhance.impl.ByteArrayWriter
+import cn.tursom.socket.enhance.impl.SocketReaderImpl
+import cn.tursom.socket.enhance.impl.SocketWriterImpl
+
+object AsyncSocketSecurityUtil {
+  private val memoryPool = HeapMemoryPool(4096)
+
+  suspend fun initAESByRSAServer(socket: AsyncSocket, rsa: RSA): EnhanceSocket<ByteArray, ByteArray> {
+    // 发送RSA公钥
+    socket.write(HeapByteBuffer(rsa.publicKeyEncoded))
+    // 接受AES密钥
+    val aesKey = socket.read(memoryPool)
+    val aes = AES(aesKey.getBytes())
+    return SecurityEnhanceSocket(SocketReaderImpl(socket), ByteArrayWriter(SocketWriterImpl(socket)), aes)
+  }
+
+  suspend fun AsyncSocket.initAESByRSAClient(socket: AsyncSocket): EnhanceSocket<ByteArray, ByteArray> {
+    // 接受RSA公钥
+    val rsaKey = socket.read(memoryPool)
+    val rsaInstance = RSA(rsaKey.getBytes())
+    // 生成AES对象
+    val aes = AES()
+    // 对AES密钥使用RSA公钥进行加密
+    val aesKey = HeapByteBuffer(rsaInstance.encrypt(aes.secKey.encoded))
+    socket.write(aesKey)
+    return SecurityEnhanceSocket(SocketReaderImpl(socket), ByteArrayWriter(SocketWriterImpl(socket)), aes)
+  }
+}
+

AsyncSocket/src/main/kotlin/cn/tursom/socket/security/SecurityEnhanceSocket.kt

@@ -0,0 +1,30 @@
+package cn.tursom.socket.security
+
+import cn.tursom.core.buffer.ByteBuffer
+import cn.tursom.core.encrypt.AES
+import cn.tursom.core.pool.MemoryPool
+import cn.tursom.socket.enhance.EnhanceSocket
+import cn.tursom.socket.enhance.SocketReader
+import cn.tursom.socket.enhance.SocketWriter
+
+class SecurityEnhanceSocket(
+  val reader: SocketReader<ByteBuffer>,
+  val writer: SocketWriter<ByteArray>,
+  val aes: AES
+) : EnhanceSocket<ByteArray, ByteArray> {
+  override fun close() {
+    reader.close()
+    writer.close()
+  }
+
+  override suspend fun read(pool: MemoryPool, timeout: Long): ByteArray {
+    val buffer = reader.read(pool, timeout)
+    return aes.decrypt(buffer)
+  }
+
+  override suspend fun write(value: ByteArray) {
+    writer.write(aes.encrypt(value))
+  }
+
+  override suspend fun flush(timeout: Long): Long = writer.flush()
+}

AsyncSocket/src/main/kotlin/cn/tursom/socket/security/SecurityNioServer.kt

@@ -0,0 +1,26 @@
+package cn.tursom.socket.security
+
+import cn.tursom.core.encrypt.RSA
+import cn.tursom.socket.AsyncSocket
+import cn.tursom.socket.server.NioServer
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.GlobalScope
+
+class SecurityNioServer(
+  port: Int,
+  backlog: Int = 50,
+  coroutineScope: CoroutineScope = GlobalScope,
+  @Suppress("MemberVisibilityCanBePrivate") val rsa: RSA = RSA(),
+  val handler: suspend AsyncSocket.() -> Unit
+) : NioServer(port, backlog, coroutineScope, {
+  AsyncSocketSecurityUtil.initAESByRSAServer(this, rsa)
+  handler()
+}) {
+  constructor(
+    port: Int,
+    backlog: Int = 50,
+    coroutineScope: CoroutineScope = GlobalScope,
+    keySize: Int,
+    handler: suspend AsyncSocket.() -> Unit
+  ) : this(port, backlog, coroutineScope, RSA(keySize), handler)
+}

AsyncSocket/src/main/kotlin/cn/tursom/socket/server/NioServer.kt

@@ -1,9 +1,9 @@
 package cn.tursom.socket.server
 
-import cn.tursom.socket.AsyncSocket
-import cn.tursom.socket.NioSocket
-import cn.tursom.socket.NioProtocol
 import cn.tursom.niothread.NioThread
+import cn.tursom.socket.AsyncSocket
+import cn.tursom.socket.NioProtocol
+import cn.tursom.socket.NioSocket
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.GlobalScope
 import kotlinx.coroutines.launch
@@ -17,7 +17,7 @@ open class NioServer(
   override val port: Int,
   backlog: Int = 50,
   coroutineScope: CoroutineScope = GlobalScope,
-  val handler: suspend AsyncSocket.() -> Unit
+  private val handler: suspend AsyncSocket.() -> Unit
 ) : SocketServer by NioLoopServer(port, object : NioProtocol by NioSocket.nioSocketProtocol {
   override fun handleConnect(key: SelectionKey, nioThread: NioThread) {
     coroutineScope.launch {

src/main/kotlin/cn/tursom/core/encrypt/Encrypt.kt

@@ -1,8 +1,26 @@
 package cn.tursom.core.encrypt
 
+import cn.tursom.core.buffer.ByteBuffer
+
 interface Encrypt {
   fun encrypt(data: ByteArray, offset: Int = 0, size: Int = data.size - offset): ByteArray
   fun decrypt(data: ByteArray, offset: Int = 0, size: Int = data.size - offset): ByteArray
   fun encrypt(data: ByteArray, buffer: ByteArray, bufferOffset: Int = 0, offset: Int = 0, size: Int = data.size - offset): Int
   fun decrypt(data: ByteArray, buffer: ByteArray, bufferOffset: Int = 0, offset: Int = 0, size: Int = data.size - offset): Int
+
+  fun encrypt(data: ByteBuffer): ByteArray {
+    return if (data.hasArray) {
+      encrypt(data.array, data.readOffset, data.writeOffset)
+    } else {
+      encrypt(data.getBytes())
+    }
+  }
+
+  fun decrypt(data: ByteBuffer): ByteArray {
+    return if (data.hasArray) {
+      decrypt(data.array, data.readOffset, data.writeOffset)
+    } else {
+      decrypt(data.getBytes())
+    }
+  }
 }