1 files changed, 102 insertions, 7 deletions
diff --git a/src/macs_device.go b/src/macs.go
index 68181e6..841ef31 100644
--- a/src/macs_device.go
+++ b/src/macs.go
@@ -1,10 +1,9 @@
 package main
 
 import (
-	"crypto/cipher"
 	"crypto/hmac"
 	"crypto/rand"
-	"github.com/aead/chacha20poly1305" // Needed for XChaCha20Poly1305, TODO:
+	"errors"
 	"golang.org/x/crypto/blake2s"
 	"net"
 	"sync"
@@ -17,9 +16,22 @@ type MACStateDevice struct {
 	secret    [blake2s.Size]byte
 	keyMAC1   [blake2s.Size]byte
 	keyMAC2   [blake2s.Size]byte
-	xaead     cipher.AEAD
 }
 
+type MACStatePeer struct {
+	mutex     sync.RWMutex
+	cookieSet time.Time
+	cookie    [blake2s.Size128]byte
+	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()
@@ -38,7 +50,6 @@ func (state *MACStateDevice) Init(pk NoisePublicKey) {
 		hsh.Sum(state.keyMAC2[:0])
 	}()
 
-	state.xaead, _ = chacha20poly1305.NewXCipher(state.keyMAC2[:])
 	state.refreshed = time.Time{}
 }
 
@@ -90,7 +101,10 @@ func (state *MACStateDevice) CheckMAC2(msg []byte, addr *net.UDPAddr) bool {
 	return hmac.Equal(mac2[:], msg[start:])
 }
 
-func (device *Device) CreateMessageCookieReply(msg []byte, receiver uint32, addr *net.UDPAddr) (*MessageCookieReply, error) {
+func (device *Device) CreateMessageCookieReply(
+	msg []byte, receiver uint32, addr *net.UDPAddr,
+) (*MessageCookieReply, error) {
+
 	state := &device.mac
 	state.mutex.RLock()
 
@@ -137,7 +151,15 @@ func (device *Device) CreateMessageCookieReply(msg []byte, receiver uint32, addr
 		state.mutex.RUnlock()
 		return nil, err
 	}
-	state.xaead.Seal(reply.Cookie[:0], reply.Nonce[:], cookie[:], mac1)
+
+	XChaCha20Poly1305Encrypt(
+		reply.Cookie[:0],
+		&reply.Nonce,
+		cookie[:],
+		mac1,
+		&state.keyMAC2,
+	)
+
 	state.mutex.RUnlock()
 	return reply, nil
 }
@@ -163,7 +185,14 @@ func (device *Device) ConsumeMessageCookieReply(msg *MessageCookieReply) bool {
 	state.mutex.Lock()
 	defer state.mutex.Unlock()
 
-	_, err := state.xaead.Open(cookie[:0], msg.Nonce[:], msg.Cookie[:], state.lastMAC1[:])
+	_, err := XChaCha20Poly1305Decrypt(
+		cookie[:0],
+		&msg.Nonce,
+		msg.Cookie[:],
+		state.lastMAC1[:],
+		&state.keyMAC2,
+	)
+
 	if err != nil {
 		return false
 	}
@@ -171,3 +200,69 @@ func (device *Device) ConsumeMessageCookieReply(msg *MessageCookieReply) bool {
 	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)
+
+	if size < blake2s.Size128*2 {
+		panic(errors.New("bug: message too short"))
+	}
+
+	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)
+
+	// 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])
+	}()
+}