Improved cookie/mac computation code

10 files changed, 523 insertions, 442 deletions

diff --git a/src/constants.go b/src/constants.go
index d366610..5ee8646 100644
--- a/src/constants.go
+++ b/src/constants.go
@@ -7,15 +7,16 @@ import (
 /* Specification constants */
 
 const (
-	RekeyAfterMessages  = (1 << 64) - (1 << 16) - 1
-	RejectAfterMessages = (1 << 64) - (1 << 4) - 1
-	RekeyAfterTime      = time.Second * 120
-	RekeyAttemptTime    = time.Second * 90
-	RekeyTimeout        = time.Second * 5
-	RejectAfterTime     = time.Second * 180
-	KeepaliveTimeout    = time.Second * 10
-	CookieRefreshTime   = time.Second * 120
-	PaddingMultiple     = 16
+	RekeyAfterMessages     = (1 << 64) - (1 << 16) - 1
+	RejectAfterMessages    = (1 << 64) - (1 << 4) - 1
+	RekeyAfterTime         = time.Second * 120
+	RekeyAttemptTime       = time.Second * 90
+	RekeyTimeout           = time.Second * 5
+	RejectAfterTime        = time.Second * 180
+	KeepaliveTimeout       = time.Second * 10
+	CookieRefreshTime      = time.Second * 120
+	HandshakeInitationRate = time.Second / 20
+	PaddingMultiple        = 16
 )
 
 const (
diff --git a/src/cookie.go b/src/cookie.go
new file mode 100644
index 0000000..a81819b
--- /dev/null
+++ b/src/cookie.go
@@ -0,0 +1,256 @@
+package main
+
+import (
+	"crypto/hmac"
+	"crypto/rand"
+	"golang.org/x/crypto/blake2s"
+	"golang.org/x/crypto/chacha20poly1305"
+	"net"
+	"sync"
+	"time"
+	"unsafe"
+)
+
+type CookieChecker struct {
+	mutex sync.RWMutex
+	mac1  struct {
+		key [blake2s.Size]byte
+	}
+	mac2 struct {
+		secret        [blake2s.Size]byte
+		secretSet     time.Time
+		encryptionKey [chacha20poly1305.KeySize]byte
+	}
+}
+
+type CookieGenerator struct {
+	mutex sync.RWMutex
+	mac1  struct {
+		key [blake2s.Size]byte
+	}
+	mac2 struct {
+		cookie        [blake2s.Size128]byte
+		cookieSet     time.Time
+		hasLastMAC1   bool
+		lastMAC1      [blake2s.Size128]byte
+		encryptionKey [chacha20poly1305.KeySize]byte
+	}
+}
+
+func (st *CookieChecker) Init(pk NoisePublicKey) {
+	st.mutex.Lock()
+	defer st.mutex.Unlock()
+
+	// mac1 state
+
+	func() {
+		hsh, _ := blake2s.New256(nil)
+		hsh.Write([]byte(WGLabelMAC1))
+		hsh.Write(pk[:])
+		hsh.Sum(st.mac1.key[:0])
+	}()
+
+	// mac2 state
+
+	func() {
+		hsh, _ := blake2s.New256(nil)
+		hsh.Write([]byte(WGLabelCookie))
+		hsh.Write(pk[:])
+		hsh.Sum(st.mac2.encryptionKey[:0])
+	}()
+
+	st.mac2.secretSet = time.Time{}
+}
+
+func (st *CookieChecker) CheckMAC1(msg []byte) bool {
+	size := len(msg)
+	smac2 := size - blake2s.Size128
+	smac1 := smac2 - blake2s.Size128
+
+	var mac1 [blake2s.Size128]byte
+
+	mac, _ := blake2s.New128(st.mac1.key[:])
+	mac.Write(msg[:smac1])
+	mac.Sum(mac1[:0])
+
+	return hmac.Equal(mac1[:], msg[smac1:smac2])
+}
+
+func (st *CookieChecker) CheckMAC2(msg []byte, src *net.UDPAddr) bool {
+	st.mutex.RLock()
+	defer st.mutex.RUnlock()
+
+	if time.Now().Sub(st.mac2.secretSet) > CookieRefreshTime {
+		return false
+	}
+
+	// derive cookie key
+
+	var cookie [blake2s.Size128]byte
+	func() {
+		mac, _ := blake2s.New128(st.mac2.secret[:])
+		mac.Write(src.IP)
+		mac.Write((*[unsafe.Sizeof(src.Port)]byte)(unsafe.Pointer(&src.Port))[:])
+		mac.Sum(cookie[:0])
+	}()
+
+	// calculate mac of packet (including mac1)
+
+	smac2 := len(msg) - blake2s.Size128
+
+	var mac2 [blake2s.Size128]byte
+	func() {
+		mac, _ := blake2s.New128(cookie[:])
+		mac.Write(msg[:smac2])
+		mac.Sum(mac2[:0])
+	}()
+
+	return hmac.Equal(mac2[:], msg[smac2:])
+}
+
+func (st *CookieChecker) CreateReply(
+	msg []byte,
+	recv uint32,
+	src *net.UDPAddr,
+) (*MessageCookieReply, error) {
+
+	st.mutex.RLock()
+
+	// refresh cookie secret
+
+	if time.Now().Sub(st.mac2.secretSet) > CookieRefreshTime {
+		st.mutex.RUnlock()
+		st.mutex.Lock()
+		_, err := rand.Read(st.mac2.secret[:])
+		if err != nil {
+			st.mutex.Unlock()
+			return nil, err
+		}
+		st.mac2.secretSet = time.Now()
+		st.mutex.Unlock()
+		st.mutex.RLock()
+	}
+
+	// derive cookie
+
+	var cookie [blake2s.Size128]byte
+	func() {
+		mac, _ := blake2s.New128(st.mac2.secret[:])
+		mac.Write(src.IP)
+		mac.Write((*[unsafe.Sizeof(src.Port)]byte)(unsafe.Pointer(&src.Port))[:])
+		mac.Sum(cookie[:0])
+	}()
+
+	// encrypt cookie
+
+	size := len(msg)
+
+	smac2 := size - blake2s.Size128
+	smac1 := smac2 - blake2s.Size128
+
+	reply := new(MessageCookieReply)
+	reply.Type = MessageCookieReplyType
+	reply.Receiver = recv
+
+	_, err := rand.Read(reply.Nonce[:])
+	if err != nil {
+		st.mutex.RUnlock()
+		return nil, err
+	}
+
+	XChaCha20Poly1305Encrypt(
+		reply.Cookie[:0],
+		&reply.Nonce,
+		cookie[:],
+		msg[smac1:smac2],
+		&st.mac2.encryptionKey,
+	)
+
+	st.mutex.RUnlock()
+
+	return reply, nil
+}
+
+func (st *CookieGenerator) Init(pk NoisePublicKey) {
+	st.mutex.Lock()
+	defer st.mutex.Unlock()
+
+	func() {
+		hsh, _ := blake2s.New256(nil)
+		hsh.Write([]byte(WGLabelMAC1))
+		hsh.Write(pk[:])
+		hsh.Sum(st.mac1.key[:0])
+	}()
+
+	func() {
+		hsh, _ := blake2s.New256(nil)
+		hsh.Write([]byte(WGLabelCookie))
+		hsh.Write(pk[:])
+		hsh.Sum(st.mac2.encryptionKey[:0])
+	}()
+
+	st.mac2.cookieSet = time.Time{}
+}
+
+func (st *CookieGenerator) ConsumeReply(msg *MessageCookieReply) bool {
+	st.mutex.Lock()
+	defer st.mutex.Unlock()
+
+	if !st.mac2.hasLastMAC1 {
+		return false
+	}
+
+	var cookie [blake2s.Size128]byte
+
+	_, err := XChaCha20Poly1305Decrypt(
+		cookie[:0],
+		&msg.Nonce,
+		msg.Cookie[:],
+		st.mac2.lastMAC1[:],
+		&st.mac2.encryptionKey,
+	)
+
+	if err != nil {
+		return false
+	}
+
+	st.mac2.cookieSet = time.Now()
+	st.mac2.cookie = cookie
+	return true
+}
+
+func (st *CookieGenerator) AddMacs(msg []byte) {
+
+	size := len(msg)
+
+	smac2 := size - blake2s.Size128
+	smac1 := smac2 - blake2s.Size128
+
+	mac1 := msg[smac1:smac2]
+	mac2 := msg[smac2:]
+
+	st.mutex.Lock()
+	defer st.mutex.Unlock()
+
+	// set mac1
+
+	func() {
+		mac, _ := blake2s.New128(st.mac1.key[:])
+		mac.Write(msg[:smac1])
+		mac.Sum(mac1[:0])
+	}()
+	copy(st.mac2.lastMAC1[:], mac1)
+	st.mac2.hasLastMAC1 = true
+
+	// set mac2
+
+	if time.Now().Sub(st.mac2.cookieSet) > CookieRefreshTime {
+		return
+	}
+
+	func() {
+		mac, _ := blake2s.New128(st.mac2.cookie[:])
+		mac.Write(msg[:smac2])
+		mac.Sum(mac2[:0])
+	}()
+}
diff --git a/src/cookie_test.go b/src/cookie_test.go
new file mode 100644
index 0000000..193a76e
--- /dev/null
+++ b/src/cookie_test.go
@@ -0,0 +1,187 @@
+package main
+
+import (
+	"net"
+	"testing"
+)
+
+func TestCookieMAC1(t *testing.T) {
+
+	// setup generator / checker
+
+	var (
+		generator CookieGenerator
+		checker   CookieChecker
+	)
+
+	sk, err := newPrivateKey()
+	if err != nil {
+		t.Fatal(err)
+	}
+	pk := sk.publicKey()
+
+	generator.Init(pk)
+	checker.Init(pk)
+
+	// check mac1
+
+	src, _ := net.ResolveUDPAddr("udp", "192.168.13.37:4000")
+
+	checkMAC1 := func(msg []byte) {
+		generator.AddMacs(msg)
+		if !checker.CheckMAC1(msg) {
+			t.Fatal("MAC1 generation/verification failed")
+		}
+		if checker.CheckMAC2(msg, src) {
+			t.Fatal("MAC2 generation/verification failed")
+		}
+	}
+
+	checkMAC1([]byte{
+		0x99, 0xbb, 0xa5, 0xfc, 0x99, 0xaa, 0x83, 0xbd,
+		0x7b, 0x00, 0xc5, 0x9a, 0x4c, 0xb9, 0xcf, 0x62,
+		0x40, 0x23, 0xf3, 0x8e, 0xd8, 0xd0, 0x62, 0x64,
+		0x5d, 0xb2, 0x80, 0x13, 0xda, 0xce, 0xc6, 0x91,
+		0x61, 0xd6, 0x30, 0xf1, 0x32, 0xb3, 0xa2, 0xf4,
+		0x7b, 0x43, 0xb5, 0xa7, 0xe2, 0xb1, 0xf5, 0x6c,
+		0x74, 0x6b, 0xb0, 0xcd, 0x1f, 0x94, 0x86, 0x7b,
+		0xc8, 0xfb, 0x92, 0xed, 0x54, 0x9b, 0x44, 0xf5,
+		0xc8, 0x7d, 0xb7, 0x8e, 0xff, 0x49, 0xc4, 0xe8,
+		0x39, 0x7c, 0x19, 0xe0, 0x60, 0x19, 0x51, 0xf8,
+		0xe4, 0x8e, 0x02, 0xf1, 0x7f, 0x1d, 0xcc, 0x8e,
+		0xb0, 0x07, 0xff, 0xf8, 0xaf, 0x7f, 0x66, 0x82,
+		0x83, 0xcc, 0x7c, 0xfa, 0x80, 0xdb, 0x81, 0x53,
+		0xad, 0xf7, 0xd8, 0x0c, 0x10, 0xe0, 0x20, 0xfd,
+		0xe8, 0x0b, 0x3f, 0x90, 0x15, 0xcd, 0x93, 0xad,
+		0x0b, 0xd5, 0x0c, 0xcc, 0x88, 0x56, 0xe4, 0x3f,
+	})
+
+	checkMAC1([]byte{
+		0x33, 0xe7, 0x2a, 0x84, 0x9f, 0xff, 0x57, 0x6c,
+		0x2d, 0xc3, 0x2d, 0xe1, 0xf5, 0x5c, 0x97, 0x56,
+		0xb8, 0x93, 0xc2, 0x7d, 0xd4, 0x41, 0xdd, 0x7a,
+		0x4a, 0x59, 0x3b, 0x50, 0xdd, 0x7a, 0x7a, 0x8c,
+		0x9b, 0x96, 0xaf, 0x55, 0x3c, 0xeb, 0x6d, 0x0b,
+		0x13, 0x0b, 0x97, 0x98, 0xb3, 0x40, 0xc3, 0xcc,
+		0xb8, 0x57, 0x33, 0x45, 0x6e, 0x8b, 0x09, 0x2b,
+		0x81, 0x2e, 0xd2, 0xb9, 0x66, 0x0b, 0x93, 0x05,
+	})
+
+	checkMAC1([]byte{
+		0x9b, 0x96, 0xaf, 0x55, 0x3c, 0xeb, 0x6d, 0x0b,
+		0x13, 0x0b, 0x97, 0x98, 0xb3, 0x40, 0xc3, 0xcc,
+		0xb8, 0x57, 0x33, 0x45, 0x6e, 0x8b, 0x09, 0x2b,
+		0x81, 0x2e, 0xd2, 0xb9, 0x66, 0x0b, 0x93, 0x05,
+	})
+
+	// exchange cookie reply
+
+	func() {
+		msg := []byte{
+			0x6d, 0xd7, 0xc3, 0x2e, 0xb0, 0x76, 0xd8, 0xdf,
+			0x30, 0x65, 0x7d, 0x62, 0x3e, 0xf8, 0x9a, 0xe8,
+			0xe7, 0x3c, 0x64, 0xa3, 0x78, 0x48, 0xda, 0xf5,
+			0x25, 0x61, 0x28, 0x53, 0x79, 0x32, 0x86, 0x9f,
+			0xa0, 0x27, 0x95, 0x69, 0xb6, 0xba, 0xd0, 0xa2,
+			0xf8, 0x68, 0xea, 0xa8, 0x62, 0xf2, 0xfd, 0x1b,
+			0xe0, 0xb4, 0x80, 0xe5, 0x6b, 0x3a, 0x16, 0x9e,
+			0x35, 0xf6, 0xa8, 0xf2, 0x4f, 0x9a, 0x7b, 0xe9,
+			0x77, 0x0b, 0xc2, 0xb4, 0xed, 0xba, 0xf9, 0x22,
+			0xc3, 0x03, 0x97, 0x42, 0x9f, 0x79, 0x74, 0x27,
+			0xfe, 0xf9, 0x06, 0x6e, 0x97, 0x3a, 0xa6, 0x8f,
+			0xc9, 0x57, 0x0a, 0x54, 0x4c, 0x64, 0x4a, 0xe2,
+			0x4f, 0xa1, 0xce, 0x95, 0x9b, 0x23, 0xa9, 0x2b,
+			0x85, 0x93, 0x42, 0xb0, 0xa5, 0x53, 0xed, 0xeb,
+			0x63, 0x2a, 0xf1, 0x6d, 0x46, 0xcb, 0x2f, 0x61,
+			0x8c, 0xe1, 0xe8, 0xfa, 0x67, 0x20, 0x80, 0x6d,
+		}
+		generator.AddMacs(msg)
+		reply, err := checker.CreateReply(msg, 1377, src)
+		if err != nil {
+			t.Fatal("Failed to create cookie reply:", err)
+		}
+		if !generator.ConsumeReply(reply) {
+			t.Fatal("Failed to consume cookie reply")
+		}
+	}()
+
+	// check mac2
+
+	checkMAC2 := func(msg []byte) {
+		generator.AddMacs(msg)
+
+		if !checker.CheckMAC1(msg) {
+			t.Fatal("MAC1 generation/verification failed")
+		}
+		if !checker.CheckMAC2(msg, src) {
+			t.Fatal("MAC2 generation/verification failed")
+		}
+
+		msg[5] ^= 0x20
+
+		if checker.CheckMAC1(msg) {
+			t.Fatal("MAC1 generation/verification failed")
+		}
+		if checker.CheckMAC2(msg, src) {
+			t.Fatal("MAC2 generation/verification failed")
+		}
+
+		msg[5] ^= 0x20
+
+		srcBad1, _ := net.ResolveUDPAddr("udp", "192.168.13.37:4001")
+		if checker.CheckMAC2(msg, srcBad1) {
+			t.Fatal("MAC2 generation/verification failed")
+		}
+
+		srcBad2, _ := net.ResolveUDPAddr("udp", "192.168.13.38:4000")
+		if checker.CheckMAC2(msg, srcBad2) {
+			t.Fatal("MAC2 generation/verification failed")
+		}
+	}
+
+	checkMAC2([]byte{
+		0x03, 0x31, 0xb9, 0x9e, 0xb0, 0x2a, 0x54, 0xa3,
+		0xc1, 0x3f, 0xb4, 0x96, 0x16, 0xb9, 0x25, 0x15,
+		0x3d, 0x3a, 0x82, 0xf9, 0x58, 0x36, 0x86, 0x3f,
+		0x13, 0x2f, 0xfe, 0xb2, 0x53, 0x20, 0x8c, 0x3f,
+		0xba, 0xeb, 0xfb, 0x4b, 0x1b, 0x22, 0x02, 0x69,
+		0x2c, 0x90, 0xbc, 0xdc, 0xcf, 0xcf, 0x85, 0xeb,
+		0x62, 0x66, 0x6f, 0xe8, 0xe1, 0xa6, 0xa8, 0x4c,
+		0xa0, 0x04, 0x23, 0x15, 0x42, 0xac, 0xfa, 0x38,
+	})
+
+	checkMAC2([]byte{
+		0x0e, 0x2f, 0x0e, 0xa9, 0x29, 0x03, 0xe1, 0xf3,
+		0x24, 0x01, 0x75, 0xad, 0x16, 0xa5, 0x66, 0x85,
+		0xca, 0x66, 0xe0, 0xbd, 0xc6, 0x34, 0xd8, 0x84,
+		0x09, 0x9a, 0x58, 0x14, 0xfb, 0x05, 0xda, 0xf5,
+		0x90, 0xf5, 0x0c, 0x4e, 0x22, 0x10, 0xc9, 0x85,
+		0x0f, 0xe3, 0x77, 0x35, 0xe9, 0x6b, 0xc2, 0x55,
+		0x32, 0x46, 0xae, 0x25, 0xe0, 0xe3, 0x37, 0x7a,
+		0x4b, 0x71, 0xcc, 0xfc, 0x91, 0xdf, 0xd6, 0xca,
+		0xfe, 0xee, 0xce, 0x3f, 0x77, 0xa2, 0xfd, 0x59,
+		0x8e, 0x73, 0x0a, 0x8d, 0x5c, 0x24, 0x14, 0xca,
+		0x38, 0x91, 0xb8, 0x2c, 0x8c, 0xa2, 0x65, 0x7b,
+		0xbc, 0x49, 0xbc, 0xb5, 0x58, 0xfc, 0xe3, 0xd7,
+		0x02, 0xcf, 0xf7, 0x4c, 0x60, 0x91, 0xed, 0x55,
+		0xe9, 0xf9, 0xfe, 0xd1, 0x44, 0x2c, 0x75, 0xf2,
+		0xb3, 0x5d, 0x7b, 0x27, 0x56, 0xc0, 0x48, 0x4f,
+		0xb0, 0xba, 0xe4, 0x7d, 0xd0, 0xaa, 0xcd, 0x3d,
+		0xe3, 0x50, 0xd2, 0xcf, 0xb9, 0xfa, 0x4b, 0x2d,
+		0xc6, 0xdf, 0x3b, 0x32, 0x98, 0x45, 0xe6, 0x8f,
+		0x1c, 0x5c, 0xa2, 0x20, 0x7d, 0x1c, 0x28, 0xc2,
+		0xd4, 0xa1, 0xe0, 0x21, 0x52, 0x8f, 0x1c, 0xd0,
+		0x62, 0x97, 0x48, 0xbb, 0xf4, 0xa9, 0xcb, 0x35,
+		0xf2, 0x07, 0xd3, 0x50, 0xd8, 0xa9, 0xc5, 0x9a,
+		0x0f, 0xbd, 0x37, 0xaf, 0xe1, 0x45, 0x19, 0xee,
+		0x41, 0xf3, 0xf7, 0xe5, 0xe0, 0x30, 0x3f, 0xbe,
+		0x3d, 0x39, 0x64, 0x00, 0x7a, 0x1a, 0x51, 0x5e,
+		0xe1, 0x70, 0x0b, 0xb9, 0x77, 0x5a, 0xf0, 0xc4,
+		0x8a, 0xa1, 0x3a, 0x77, 0x1a, 0xe0, 0xc2, 0x06,
+		0x91, 0xd5, 0xe9, 0x1c, 0xd3, 0xfe, 0xab, 0x93,
+		0x1a, 0x0a, 0x4c, 0xbb, 0xf0, 0xff, 0xdc, 0xaa,
+		0x61, 0x73, 0xcb, 0x03, 0x4b, 0x71, 0x68, 0x64,
+		0x3d, 0x82, 0x31, 0x41, 0xd7, 0x8b, 0x22, 0x7b,
+		0x7d, 0xa1, 0xd5, 0x85, 0x6d, 0xf0, 0x1b, 0xaa,
+	})
+}
diff --git a/src/device.go b/src/device.go
index 781b525..dfd2f35 100644
--- a/src/device.go
+++ b/src/device.go
@@ -42,7 +42,7 @@ type Device struct {
 	underLoadUntil atomic.Value
 	ratelimiter    Ratelimiter
 	peers          map[NoisePublicKey]*Peer
-	mac            MACStateDevice
+	mac            CookieChecker
 }
 
 /* Warning:
diff --git a/src/keypair.go b/src/keypair.go
index b5f46df..ba9c437 100644
--- a/src/keypair.go
+++ b/src/keypair.go
@@ -29,3 +29,9 @@ func (kp *KeyPairs) Current() *KeyPair {
 	defer kp.mutex.RUnlock()
 	return kp.current
 }
+
+func (device *Device) DeleteKeyPair(key *KeyPair) {
+	key.send = nil
+	key.receive = nil
+	device.indices.Delete(key.localIndex)
+}
diff --git a/src/macs.go b/src/macs.go
deleted file mode 100644
index d55e18f..0000000
--- a/src/macs.go
+++ /dev/null
@@ -1,270 +0,0 @@
-package main
-
-import (
-	"crypto/hmac"
-	"crypto/rand"
-	"golang.org/x/crypto/blake2s"
-	"net"
-	"sync"
-	"time"
-)
-
-type MACStateDevice struct {
-	mutex     sync.RWMutex
-	refreshed time.Time
-	secret    [blake2s.Size]byte
-	keyMAC1   [blake2s.Size]byte
-	keyMAC2   [blake2s.Size]byte // TODO: Change to more descriptive size constant, rename to something.
-}
-
-type MACStatePeer struct {
-	mutex       sync.RWMutex
-	cookieSet   time.Time
-	cookie      [blake2s.Size128]byte
-	lastMAC1Set bool
-	lastMAC1    [blake2s.Size128]byte
-	keyMAC1     [blake2s.Size]byte
-	keyMAC2     [blake2s.Size]byte
-}
-
-/* Methods for verifing MAC fields
- * and creating/consuming cookies replies
- * (per device)
- */
-
-func (state *MACStateDevice) Init(pk NoisePublicKey) {
-	state.mutex.Lock()
-	defer state.mutex.Unlock()
-
-	func() {
-		hsh, _ := blake2s.New256(nil)
-		hsh.Write([]byte(WGLabelMAC1))
-		hsh.Write(pk[:])
-		hsh.Sum(state.keyMAC1[:0])
-	}()
-
-	func() {
-		hsh, _ := blake2s.New256(nil)
-		hsh.Write([]byte(WGLabelCookie))
-		hsh.Write(pk[:])
-		hsh.Sum(state.keyMAC2[:0])
-	}()
-
-	state.refreshed = time.Time{}
-}
-
-func (state *MACStateDevice) CheckMAC1(msg []byte) bool {
-	size := len(msg)
-	startMac1 := size - (blake2s.Size128 * 2)
-	startMac2 := size - blake2s.Size128
-
-	var mac1 [blake2s.Size128]byte
-	func() {
-		mac, _ := blake2s.New128(state.keyMAC1[:])
-		mac.Write(msg[:startMac1])
-		mac.Sum(mac1[:0])
-	}()
-
-	return hmac.Equal(mac1[:], msg[startMac1:startMac2])
-}
-
-func (state *MACStateDevice) CheckMAC2(msg []byte, addr *net.UDPAddr) bool {
-	state.mutex.RLock()
-	defer state.mutex.RUnlock()
-
-	if time.Now().Sub(state.refreshed) > CookieRefreshTime {
-		return false
-	}
-
-	// derive cookie key
-
-	var cookie [blake2s.Size128]byte
-	func() {
-		port := [2]byte{byte(addr.Port >> 8), byte(addr.Port)}
-		mac, _ := blake2s.New128(state.secret[:])
-		mac.Write(addr.IP)
-		mac.Write(port[:]) // TODO: Be faster and more platform dependent?
-		mac.Sum(cookie[:0])
-	}()
-
-	// calculate mac of packet
-
-	start := len(msg) - blake2s.Size128
-
-	var mac2 [blake2s.Size128]byte
-	func() {
-		mac, _ := blake2s.New128(cookie[:])
-		mac.Write(msg[:start])
-		mac.Sum(mac2[:0])
-	}()
-
-	return hmac.Equal(mac2[:], msg[start:])
-}
-
-func (device *Device) CreateMessageCookieReply(
-	msg []byte, receiver uint32, addr *net.UDPAddr,
-) (*MessageCookieReply, error) {
-
-	state := &device.mac
-	state.mutex.RLock()
-
-	// refresh cookie secret
-
-	if time.Now().Sub(state.refreshed) > CookieRefreshTime {
-		state.mutex.RUnlock()
-		state.mutex.Lock()
-		_, err := rand.Read(state.secret[:])
-		if err != nil {
-			state.mutex.Unlock()
-			return nil, err
-		}
-		state.refreshed = time.Now()
-		state.mutex.Unlock()
-		state.mutex.RLock()
-	}
-
-	// derive cookie key
-
-	var cookie [blake2s.Size128]byte
-	func() {
-		port := [2]byte{byte(addr.Port >> 8), byte(addr.Port)}
-		mac, _ := blake2s.New128(state.secret[:])
-		mac.Write(addr.IP)
-		mac.Write(port[:]) // TODO: Do whatever we did above
-		mac.Sum(cookie[:0])
-	}()
-
-	// encrypt cookie
-
-	size := len(msg)
-
-	startMac1 := size - (blake2s.Size128 * 2)
-	startMac2 := size - blake2s.Size128
-
-	mac1 := msg[startMac1:startMac2]
-
-	reply := new(MessageCookieReply)
-	reply.Type = MessageCookieReplyType
-	reply.Receiver = receiver
-	_, err := rand.Read(reply.Nonce[:])
-	if err != nil {
-		state.mutex.RUnlock()
-		return nil, err
-	}
-
-	XChaCha20Poly1305Encrypt(
-		reply.Cookie[:0],
-		&reply.Nonce,
-		cookie[:],
-		mac1,
-		&state.keyMAC2,
-	)
-
-	state.mutex.RUnlock()
-	return reply, nil
-}
-
-func (device *Device) ConsumeMessageCookieReply(msg *MessageCookieReply) bool {
-
-	if msg.Type != MessageCookieReplyType {
-		return false
-	}
-
-	// lookup peer
-
-	lookup := device.indices.Lookup(msg.Receiver)
-	if lookup.handshake == nil {
-		return false
-	}
-
-	// decrypt and store cookie
-
-	var cookie [blake2s.Size128]byte
-	state := &lookup.peer.mac
-
-	state.mutex.Lock()
-	defer state.mutex.Unlock()
-
-	if !state.lastMAC1Set {
-		return false
-	}
-
-	_, err := XChaCha20Poly1305Decrypt(
-		cookie[:0],
-		&msg.Nonce,
-		msg.Cookie[:],
-		state.lastMAC1[:],
-		&state.keyMAC2,
-	)
-
-	if err != nil {
-		return false
-	}
-
-	state.cookieSet = time.Now()
-	state.cookie = cookie
-	return true
-}
-
-/* Methods for generating the MAC fields
- * (per peer)
- */
-
-func (state *MACStatePeer) Init(pk NoisePublicKey) {
-	state.mutex.Lock()
-	defer state.mutex.Unlock()
-
-	func() {
-		hsh, _ := blake2s.New256(nil)
-		hsh.Write([]byte(WGLabelMAC1))
-		hsh.Write(pk[:])
-		hsh.Sum(state.keyMAC1[:0])
-	}()
-
-	func() {
-		hsh, _ := blake2s.New256(nil)
-		hsh.Write([]byte(WGLabelCookie))
-		hsh.Write(pk[:])
-		hsh.Sum(state.keyMAC2[:0])
-	}()
-
-	state.cookieSet = time.Time{} // never
-}
-
-func (state *MACStatePeer) AddMacs(msg []byte) {
-	size := len(msg)
-
-	startMac1 := size - (blake2s.Size128 * 2)
-	startMac2 := size - blake2s.Size128
-
-	mac1 := msg[startMac1 : startMac1+blake2s.Size128]
-	mac2 := msg[startMac2 : startMac2+blake2s.Size128]
-
-	state.mutex.Lock()
-	defer state.mutex.Unlock()
-
-	// set mac1
-
-	func() {
-		mac, _ := blake2s.New128(state.keyMAC1[:])
-		mac.Write(msg[:startMac1])
-		mac.Sum(mac1[:0])
-	}()
-	copy(state.lastMAC1[:], mac1)
-	state.lastMAC1Set = true
-
-	// set mac2
-
-	if state.cookieSet.IsZero() {
-		return
-	}
-	if time.Now().Sub(state.cookieSet) > CookieRefreshTime {
-		state.cookieSet = time.Time{}
-		return
-	}
-	func() {
-		mac, _ := blake2s.New128(state.cookie[:])
-		mac.Write(msg[:startMac2])
-		mac.Sum(mac2[:0])
-	}()
-}
diff --git a/src/macs_test.go b/src/macs_test.go
deleted file mode 100644
index 3575ccb..0000000
--- a/src/macs_test.go
+++ /dev/null
@@ -1,116 +0,0 @@
-package main
-
-import (
-	"net"
-	"testing"
-	"testing/quick"
-)
-
-func TestMAC1(t *testing.T) {
-	dev1 := randDevice(t)
-	dev2 := randDevice(t)
-
-	defer dev1.Close()
-	defer dev2.Close()
-
-	peer1, _ := dev2.NewPeer(dev1.privateKey.publicKey())
-	peer2, _ := dev1.NewPeer(dev2.privateKey.publicKey())
-
-	assertEqual(t, peer1.mac.keyMAC1[:], dev1.mac.keyMAC1[:])
-	assertEqual(t, peer2.mac.keyMAC1[:], dev2.mac.keyMAC1[:])
-
-	msg1 := make([]byte, 256)
-	copy(msg1, []byte("some content"))
-	peer1.mac.AddMacs(msg1)
-	if dev1.mac.CheckMAC1(msg1) == false {
-		t.Fatal("failed to verify mac1")
-	}
-}
-
-func TestMACs(t *testing.T) {
-	assertion := func(
-		addr net.UDPAddr,
-		addrInvalid net.UDPAddr,
-		sk1 NoisePrivateKey,
-		sk2 NoisePrivateKey,
-		msg []byte,
-		receiver uint32,
-	) bool {
-		device1 := randDevice(t)
-		device1.SetPrivateKey(sk1)
-
-		device2 := randDevice(t)
-		device2.SetPrivateKey(sk2)
-
-		defer device1.Close()
-		defer device2.Close()
-
-		peer1, _ := device2.NewPeer(device1.privateKey.publicKey())
-		peer2, _ := device1.NewPeer(device2.privateKey.publicKey())
-
-		if addr.Port < 0 {
-			return true
-		}
-
-		addr.Port &= 0xffff
-
-		if len(msg) < 32 {
-			return true
-		}
-
-		assertEqual(t, peer1.mac.keyMAC1[:], device1.mac.keyMAC1[:])
-		assertEqual(t, peer2.mac.keyMAC1[:], device2.mac.keyMAC1[:])
-
-		device2.indices.Insert(receiver, IndexTableEntry{
-			peer:      peer1,
-			handshake: &peer1.handshake,
-		})
-
-		// test just MAC1
-
-		peer1.mac.AddMacs(msg)
-		if device1.mac.CheckMAC1(msg) == false {
-			return false
-		}
-
-		// exchange cookie reply
-
-		cr, err := device1.CreateMessageCookieReply(msg, receiver, &addr)
-		if err != nil {
-			return false
-		}
-
-		if !device2.ConsumeMessageCookieReply(cr) {
-			return false
-		}
-
-		// test MAC1 + MAC2
-
-		peer1.mac.AddMacs(msg)
-		if !device1.mac.CheckMAC1(msg) {
-			return false
-		}
-		if !device1.mac.CheckMAC2(msg, &addr) {
-			return false
-		}
-
-		// test invalid
-
-		if device1.mac.CheckMAC2(msg, &addrInvalid) {
-			return false
-		}
-		msg[5] ^= 1
-		if device1.mac.CheckMAC1(msg) {
-			return false
-		}
-
-		t.Log("Passed")
-
-		return true
-	}
-
-	err := quick.Check(assertion, nil)
-	if err != nil {
-		t.Error(err)
-	}
-}
diff --git a/src/noise_protocol.go b/src/noise_protocol.go
index 5c776a8..0d78c84 100644
--- a/src/noise_protocol.go
+++ b/src/noise_protocol.go
@@ -87,18 +87,19 @@ type MessageCookieReply struct {
 }
 
 type Handshake struct {
-	state                   int
-	mutex                   sync.RWMutex
-	hash                    [blake2s.Size]byte       // hash value
-	chainKey                [blake2s.Size]byte       // chain key
-	presharedKey            NoiseSymmetricKey        // psk
-	localEphemeral          NoisePrivateKey          // ephemeral secret key
-	localIndex              uint32                   // used to clear hash-table
-	remoteIndex             uint32                   // index for sending
-	remoteStatic            NoisePublicKey           // long term key
-	remoteEphemeral         NoisePublicKey           // ephemeral public key
-	precomputedStaticStatic [NoisePublicKeySize]byte // precomputed shared secret
-	lastTimestamp           TAI64N
+	state                     int
+	mutex                     sync.RWMutex
+	hash                      [blake2s.Size]byte       // hash value
+	chainKey                  [blake2s.Size]byte       // chain key
+	presharedKey              NoiseSymmetricKey        // psk
+	localEphemeral            NoisePrivateKey          // ephemeral secret key
+	localIndex                uint32                   // used to clear hash-table
+	remoteIndex               uint32                   // index for sending
+	remoteStatic              NoisePublicKey           // long term key
+	remoteEphemeral           NoisePublicKey           // ephemeral public key
+	precomputedStaticStatic   [NoisePublicKeySize]byte // precomputed shared secret
+	lastTimestamp             TAI64N
+	lastInitiationConsumption time.Time
 }
 
 var (
@@ -239,34 +240,27 @@ func (device *Device) ConsumeMessageInitiation(msg *MessageInitiation) *Peer {
 	// verify identity
 
 	var timestamp TAI64N
-	ok := func() bool {
-
-		// read lock handshake
-
-		handshake.mutex.RLock()
-		defer handshake.mutex.RUnlock()
-
-		// decrypt timestamp
-
-		func() {
-			var key [chacha20poly1305.KeySize]byte
-			chainKey, key = KDF2(
-				chainKey[:],
-				handshake.precomputedStaticStatic[:],
-			)
-			aead, _ := chacha20poly1305.New(key[:])
-			_, err = aead.Open(timestamp[:0], ZeroNonce[:], msg.Timestamp[:], hash[:])
-		}()
-		if err != nil {
-			return false
-		}
-		hash = mixHash(hash, msg.Timestamp[:])
+	var key [chacha20poly1305.KeySize]byte
 
-		// check for replay attack
+	handshake.mutex.RLock()
+	chainKey, key = KDF2(
+		chainKey[:],
+		handshake.precomputedStaticStatic[:],
+	)
+	aead, _ := chacha20poly1305.New(key[:])
+	_, err = aead.Open(timestamp[:0], ZeroNonce[:], msg.Timestamp[:], hash[:])
+	if err != nil {
+		handshake.mutex.RUnlock()
+		return nil
+	}
+	hash = mixHash(hash, msg.Timestamp[:])
 
-		return timestamp.After(handshake.lastTimestamp)
-	}()
+	// protect against replay & flood
 
+	var ok bool
+	ok = timestamp.After(handshake.lastTimestamp)
+	ok = ok && time.Now().Sub(handshake.lastInitiationConsumption) > HandshakeInitationRate
+	handshake.mutex.RUnlock()
 	if !ok {
 		return nil
 	}
@@ -280,6 +274,7 @@ func (device *Device) ConsumeMessageInitiation(msg *MessageInitiation) *Peer {
 	handshake.remoteIndex = msg.Sender
 	handshake.remoteEphemeral = msg.Ephemeral
 	handshake.lastTimestamp = timestamp
+	handshake.lastInitiationConsumption = time.Now()
 	handshake.state = HandshakeInitiationConsumed
 
 	handshake.mutex.Unlock()
@@ -483,15 +478,21 @@ func (peer *Peer) NewKeyPair() *KeyPair {
 		// TODO: Adapt kernel behavior noise.c:161
 		if isInitiator {
 			if kp.previous != nil {
-				kp.previous.send = nil
-				kp.previous.receive = nil
 				indices.Delete(kp.previous.localIndex)
 			}
-			kp.previous = kp.current
-			kp.current = keyPair
-			signalSend(peer.signal.newKeyPair)
+
+			if kp.next != nil {
+				kp.previous = kp.next
+				kp.next = keyPair
+			} else {
+				kp.previous = kp.current
+				kp.current = keyPair
+				signalSend(peer.signal.newKeyPair) // TODO: This more places (after confirming the key)
+			}
+
 		} else {
 			kp.next = keyPair
+			kp.previous = nil // TODO: Discuss why
 		}
 	}()
 
diff --git a/src/peer.go b/src/peer.go
index 02aac3b..a4feb2f 100644
--- a/src/peer.go
+++ b/src/peer.go
@@ -56,7 +56,7 @@ type Peer struct {
 		outbound chan *QueueOutboundElement // sequential ordering of work
 		inbound  chan *QueueInboundElement  // sequential ordering of work
 	}
-	mac MACStatePeer
+	mac CookieGenerator
 }
 
 func (device *Device) NewPeer(pk NoisePublicKey) (*Peer, error) {
diff --git a/src/receive.go b/src/receive.go
index c47d93c..4c76bbf 100644
--- a/src/receive.go
+++ b/src/receive.go
@@ -272,7 +272,9 @@ func (device *Device) RoutineHandshake() {
 
 		case MessageCookieReplyType:
 
-			// verify and update peer cookie state
+			// unmarshal packet
+
+			logDebug.Println("Process cookie reply from:", elem.source.String())
 
 			var reply MessageCookieReply
 			reader := bytes.NewReader(elem.packet)
@@ -281,7 +283,14 @@ func (device *Device) RoutineHandshake() {
 				logDebug.Println("Failed to decode cookie reply")
 				return
 			}
-			device.ConsumeMessageCookieReply(&reply)
+
+			// lookup peer and consume response
+
+			entry := device.indices.Lookup(reply.Receiver)
+			if entry.peer == nil {
+				return
+			}
+			entry.peer.mac.ConsumeReply(&reply)
 			continue
 
 		case MessageInitiationType, MessageResponseType:
@@ -298,12 +307,17 @@ func (device *Device) RoutineHandshake() {
 
 					// construct cookie reply
 
+					logDebug.Println("Sending cookie reply to:", elem.source.String())
+
 					sender := binary.LittleEndian.Uint32(elem.packet[4:8]) // "sender" always follows "type"
-					reply, err := device.CreateMessageCookieReply(elem.packet, sender, elem.source)
+					reply, err := device.mac.CreateReply(elem.packet, sender, elem.source)
 					if err != nil {
 						logError.Println("Failed to create cookie reply:", err)
 						return
 					}
+
+					// marshal and send reply
+
 					writer := bytes.NewBuffer(temp[:0])
 					binary.Write(writer, binary.LittleEndian, reply)
 					_, err = device.net.conn.WriteToUDP(
@@ -392,6 +406,8 @@ func (device *Device) RoutineHandshake() {
 
 		case MessageResponseType:
 
+			logDebug.Println("Process response")
+
 			// unmarshal
 
 			var msg MessageResponse