1 files changed, 18 insertions, 0 deletions
diff --git a/src/device.go b/src/device.go
index d03057d..9f1daa6 100644
--- a/src/device.go
+++ b/src/device.go
@@ -1,12 +1,17 @@
 package main
 
 import (
+	"math/rand"
 	"sync"
 )
 
+/* TODO: Locking may be a little broad here
+ */
+
 type Device struct {
 	mutex        sync.RWMutex
 	peers        map[NoisePublicKey]*Peer
+	sessions     map[uint32]*Handshake
 	privateKey   NoisePrivateKey
 	publicKey    NoisePublicKey
 	fwMark       uint32
@@ -14,6 +19,19 @@ type Device struct {
 	routingTable RoutingTable
 }
 
+func (dev *Device) NewID(h *Handshake) uint32 {
+	dev.mutex.Lock()
+	defer dev.mutex.Unlock()
+	for {
+		id := rand.Uint32()
+		_, ok := dev.sessions[id]
+		if !ok {
+			dev.sessions[id] = h
+			return id
+		}
+	}
+}
+
 func (dev *Device) RemovePeer(key NoisePublicKey) {
 	dev.mutex.Lock()
 	defer dev.mutex.Unlock()