From 1698f0017b2646106db7c1651ad8d1f90c6b1c98 Mon Sep 17 00:00:00 2001
From: Justin Li <jli.justinli@gmail.com>
Date: Thu, 19 Feb 2015 22:25:42 -0500
Subject: [PATCH] udp: Add AES-based connection ID generation

---
 udp/connection.go      | 50 ++++++++++++++++++++++++++++++++++++++++++
 udp/connection_test.go | 44 +++++++++++++++++++++++++++++++++++++
 2 files changed, 94 insertions(+)
 create mode 100644 udp/connection.go
 create mode 100644 udp/connection_test.go

diff --git a/udp/connection.go b/udp/connection.go
new file mode 100644
index 0000000..c805210
--- /dev/null
+++ b/udp/connection.go
@@ -0,0 +1,50 @@
+// Copyright 2015 The Chihaya Authors. All rights reserved.
+// Use of this source code is governed by the BSD 2-Clause license,
+// which can be found in the LICENSE file.
+
+package udp
+
+import (
+	"crypto/aes"
+	"crypto/cipher"
+	"crypto/rand"
+)
+
+var connectionKey, connectionIV []byte
+
+func InitConnectionIDEncryption() error {
+	connectionKey = make([]byte, 16)
+	_, err := rand.Read(connectionKey)
+	if err != nil {
+		return err
+	}
+
+	connectionIV = make([]byte, 16)
+	_, err = rand.Read(connectionIV)
+	return err
+}
+
+func GenerateConnectionID(ip []byte) []byte {
+	block, err := aes.NewCipher(connectionKey)
+	if err != nil {
+		panic(err)
+	}
+
+	if len(ip) > 16 {
+		panic("IP larger than 16 bytes")
+	}
+
+	for len(ip) < 8 {
+		ip = append(ip, ip...) // Not enough bits in output.
+	}
+
+	ct := make([]byte, 16)
+	stream := cipher.NewCFBDecrypter(block, connectionIV)
+	stream.XORKeyStream(ct, ip)
+
+	for i := len(ip) - 1; i >= 8; i-- {
+		ct[i-8] ^= ct[i]
+	}
+
+	return ct[:8]
+}
diff --git a/udp/connection_test.go b/udp/connection_test.go
new file mode 100644
index 0000000..7c4aa39
--- /dev/null
+++ b/udp/connection_test.go
@@ -0,0 +1,44 @@
+// Copyright 2015 The Chihaya Authors. All rights reserved.
+// Use of this source code is governed by the BSD 2-Clause license,
+// which can be found in the LICENSE file.
+
+package udp
+
+import (
+	"bytes"
+	"net"
+	"testing"
+)
+
+func TestInitReturnsNoError(t *testing.T) {
+	if err := InitConnectionIDEncryption(); err != nil {
+		t.Error("InitConnectionIDEncryption returned", err)
+	}
+}
+
+func testGenerateConnectionID(t *testing.T, ip net.IP) {
+	InitConnectionIDEncryption()
+
+	id1 := GenerateConnectionID(ip)
+	id2 := GenerateConnectionID(ip)
+
+	if !bytes.Equal(id1, id2) {
+		t.Errorf("Connection ID mismatch: %x != %x", id1, id2)
+	}
+
+	if len(id1) != 8 {
+		t.Errorf("Connection ID had length: %d != 8", len(id1))
+	}
+
+	if bytes.Count(id1, []byte{0}) == 8 {
+		t.Errorf("Connection ID was 0")
+	}
+}
+
+func TestGenerateConnectionIDIPv4(t *testing.T) {
+	testGenerateConnectionID(t, net.ParseIP("192.168.1.123").To4())
+}
+
+func TestGenerateConnectionIDIPv6(t *testing.T) {
+	testGenerateConnectionID(t, net.ParseIP("1:2:3:4::5:6"))
+}