1 files changed, 47 insertions, 28 deletions

diff --git a/src/index.go b/src/index.go
index 81f71e9..9178510 100644
--- a/src/index.go
+++ b/src/index.go
@@ -11,10 +11,15 @@ import (
  *
  */
 
+type IndexTableEntry struct {
+	peer      *Peer
+	handshake *Handshake
+	keyPair   *KeyPair
+}
+
 type IndexTable struct {
-	mutex      sync.RWMutex
-	keypairs   map[uint32]*KeyPair
-	handshakes map[uint32]*Peer
+	mutex sync.RWMutex
+	table map[uint32]IndexTableEntry
 }
 
 func randUint32() (uint32, error) {
@@ -32,52 +37,66 @@ func randUint32() (uint32, error) {
 
 func (table *IndexTable) Init() {
 	table.mutex.Lock()
-	defer table.mutex.Unlock()
-	table.keypairs = make(map[uint32]*KeyPair)
-	table.handshakes = make(map[uint32]*Peer)
+	table.table = make(map[uint32]IndexTableEntry)
+	table.mutex.Unlock()
 }
 
-func (table *IndexTable) NewIndex(peer *Peer) (uint32, error) {
+func (table *IndexTable) ClearIndex(index uint32) {
+	if index == 0 {
+		return
+	}
+	table.mutex.Lock()
+	delete(table.table, index)
+	table.mutex.Unlock()
+}
+
+func (table *IndexTable) Insert(key uint32, value IndexTableEntry) {
 	table.mutex.Lock()
-	defer table.mutex.Unlock()
+	table.table[key] = value
+	table.mutex.Unlock()
+}
+
+func (table *IndexTable) NewIndex(peer *Peer) (uint32, error) {
 	for {
 		// generate random index
 
-		id, err := randUint32()
+		index, err := randUint32()
 		if err != nil {
-			return id, err
+			return index, err
 		}
-		if id == 0 {
+		if index == 0 {
 			continue
 		}
 
 		// check if index used
 
-		_, ok := table.keypairs[id]
-		if ok {
-			continue
-		}
-		_, ok = table.handshakes[id]
+		table.mutex.RLock()
+		_, ok := table.table[index]
 		if ok {
 			continue
 		}
+		table.mutex.RUnlock()
 
-		// clean old index
+		// replace index
 
-		delete(table.handshakes, peer.handshake.localIndex)
-		table.handshakes[id] = peer
-		return id, nil
+		table.mutex.Lock()
+		_, found := table.table[index]
+		if found {
+			table.mutex.Unlock()
+			continue
+		}
+		table.table[index] = IndexTableEntry{
+			peer:      peer,
+			handshake: &peer.handshake,
+			keyPair:   nil,
+		}
+		table.mutex.Unlock()
+		return index, nil
 	}
 }
 
-func (table *IndexTable) LookupKeyPair(id uint32) *KeyPair {
-	table.mutex.RLock()
-	defer table.mutex.RUnlock()
-	return table.keypairs[id]
-}
-
-func (table *IndexTable) LookupHandshake(id uint32) *Peer {
+func (table *IndexTable) Lookup(id uint32) IndexTableEntry {
 	table.mutex.RLock()
 	defer table.mutex.RUnlock()
-	return table.handshakes[id]
+	return table.table[id]
 }