From 2aa0daf4d58ffc930fde611e7efe6ae3c9515130 Mon Sep 17 00:00:00 2001
From: Mathias Hall-Andersen <mathias@hall-andersen.dk>
Date: Sun, 2 Jul 2017 15:28:38 +0200
Subject: [PATCH] Fixed transport header problem

---
 src/helper_test.go    |  4 ++--
 src/keypair.go        |  3 ++-
 src/noise_protocol.go | 13 ++++++++-----
 src/receive.go        |  5 +++--
 src/send.go           | 27 +++++++++++++++++----------
 5 files changed, 32 insertions(+), 20 deletions(-)

diff --git a/src/helper_test.go b/src/helper_test.go
index 464292f..6d85771 100644
--- a/src/helper_test.go
+++ b/src/helper_test.go
@@ -10,7 +10,7 @@ import (
 
 type DummyTUN struct {
 	name    string
-	mtu     uint
+	mtu     int
 	packets chan []byte
 }
 
@@ -18,7 +18,7 @@ func (tun *DummyTUN) Name() string {
 	return tun.name
 }
 
-func (tun *DummyTUN) MTU() uint {
+func (tun *DummyTUN) MTU() int {
 	return tun.mtu
 }
 
diff --git a/src/keypair.go b/src/keypair.go
index 0e845f7..0fac5cb 100644
--- a/src/keypair.go
+++ b/src/keypair.go
@@ -13,7 +13,8 @@ type KeyPair struct {
 	sendNonce   uint64
 	isInitiator bool
 	created     time.Time
-	id          uint32
+	localIndex  uint32
+	remoteIndex uint32
 }
 
 type KeyPairs struct {
diff --git a/src/noise_protocol.go b/src/noise_protocol.go
index adb00ec..5a62901 100644
--- a/src/noise_protocol.go
+++ b/src/noise_protocol.go
@@ -32,10 +32,11 @@ const (
 )
 
 const (
-	MessageInitiationSize  = 148
-	MessageResponseSize    = 92
-	MessageCookieReplySize = 64
-	MessageTransportSize   = 16 + poly1305.TagSize // size of empty transport
+	MessageInitiationSize      = 148
+	MessageResponseSize        = 92
+	MessageCookieReplySize     = 64
+	MessageTransportHeaderSize = 16
+	MessageTransportSize       = MessageTransportHeaderSize + poly1305.TagSize // size of empty transport
 )
 
 const (
@@ -449,6 +450,8 @@ func (peer *Peer) NewKeyPair() *KeyPair {
 	keyPair.sendNonce = 0
 	keyPair.recvNonce = 0
 	keyPair.created = time.Now()
+	keyPair.localIndex = peer.handshake.localIndex
+	keyPair.remoteIndex = peer.handshake.remoteIndex
 
 	// remap index
 
@@ -471,7 +474,7 @@ func (peer *Peer) NewKeyPair() *KeyPair {
 			if kp.previous != nil {
 				kp.previous.send = nil
 				kp.previous.recv = nil
-				peer.device.indices.Delete(kp.previous.id)
+				peer.device.indices.Delete(kp.previous.localIndex)
 			}
 			kp.previous = kp.current
 			kp.current = keyPair
diff --git a/src/receive.go b/src/receive.go
index ab28944..5afbf7f 100644
--- a/src/receive.go
+++ b/src/receive.go
@@ -75,6 +75,7 @@ func (device *Device) RoutineReceiveIncomming() {
 		// handle packet
 
 		packet = packet[:size]
+		debugLog.Println("GOT:", packet)
 		msgType := binary.LittleEndian.Uint32(packet[:4])
 
 		func() {
@@ -371,6 +372,8 @@ func (peer *Peer) RoutineSequentialReceiver() {
 
 		// check for replay
 
+		// strip padding
+
 		// check for keep-alive
 
 		if len(elem.packet) == 0 {
@@ -393,8 +396,6 @@ func (device *Device) RoutineWriteToTUN(tun TUNDevice) {
 		case packet = <-device.queue.inbound:
 		}
 
-		device.log.Debug.Println("GOT:", packet)
-
 		size, err := tun.Write(packet)
 		device.log.Debug.Println("DEBUG:", size, err)
 		if err != nil {
diff --git a/src/send.go b/src/send.go
index 7a10560..3fe4733 100644
--- a/src/send.go
+++ b/src/send.go
@@ -171,8 +171,6 @@ func (peer *Peer) RoutineNonce() {
 				}
 			}
 
-			logger.Println("PACKET:", packet)
-
 			// wait for key pair
 
 			for {
@@ -221,8 +219,6 @@ func (peer *Peer) RoutineNonce() {
 				work.peer = peer
 				work.mutex.Lock()
 
-				logger.Println("WORK:", work)
-
 				packet = nil
 
 				// drop packets until there is space
@@ -263,7 +259,7 @@ func (device *Device) RoutineEncryption() {
 
 		// pad packet
 
-		padding := device.mtu - len(work.packet)
+		padding := device.mtu - len(work.packet) - MessageTransportSize
 		if padding < 0 {
 			work.Drop()
 			continue
@@ -272,19 +268,30 @@ func (device *Device) RoutineEncryption() {
 		for n := 0; n < padding; n += 1 {
 			work.packet = append(work.packet, 0)
 		}
-		device.log.Debug.Println(work.packet)
+		content := work.packet[MessageTransportHeaderSize:]
+		copy(content, work.packet)
 
-		// encrypt
+		// prepare header
+
+		binary.LittleEndian.PutUint32(work.packet[:4], MessageTransportType)
+		binary.LittleEndian.PutUint32(work.packet[4:8], work.keyPair.remoteIndex)
+		binary.LittleEndian.PutUint64(work.packet[8:16], work.nonce)
+
+		device.log.Debug.Println(work.packet, work.nonce)
+
+		// encrypt content
 
 		binary.LittleEndian.PutUint64(nonce[4:], work.nonce)
-		work.packet = work.keyPair.send.Seal(
-			work.packet[:0],
+		work.keyPair.send.Seal(
+			content[:0],
 			nonce[:],
-			work.packet,
+			content,
 			nil,
 		)
 		work.mutex.Unlock()
 
+		device.log.Debug.Println(work.packet, work.nonce)
+
 		// initiate new handshake
 
 		work.peer.KeepKeyFreshSending()