indextable.go raw

   1  /* SPDX-License-Identifier: MIT
   2   *
   3   * Copyright (C) 2017-2025 WireGuard LLC. All Rights Reserved.
   4   */
   5  
   6  package device
   7  
   8  import (
   9  	"crypto/rand"
  10  	"encoding/binary"
  11  	"sync"
  12  )
  13  
  14  type IndexTableEntry struct {
  15  	peer      *Peer
  16  	handshake *Handshake
  17  	keypair   *Keypair
  18  }
  19  
  20  type IndexTable struct {
  21  	sync.RWMutex
  22  	table map[uint32]IndexTableEntry
  23  }
  24  
  25  func randUint32() (uint32, error) {
  26  	var integer [4]byte
  27  	_, err := rand.Read(integer[:])
  28  	// Arbitrary endianness; both are intrinsified by the Go compiler.
  29  	return binary.LittleEndian.Uint32(integer[:]), err
  30  }
  31  
  32  func (table *IndexTable) Init() {
  33  	table.Lock()
  34  	defer table.Unlock()
  35  	table.table = make(map[uint32]IndexTableEntry)
  36  }
  37  
  38  func (table *IndexTable) Delete(index uint32) {
  39  	table.Lock()
  40  	defer table.Unlock()
  41  	delete(table.table, index)
  42  }
  43  
  44  func (table *IndexTable) SwapIndexForKeypair(index uint32, keypair *Keypair) {
  45  	table.Lock()
  46  	defer table.Unlock()
  47  	entry, ok := table.table[index]
  48  	if !ok {
  49  		return
  50  	}
  51  	table.table[index] = IndexTableEntry{
  52  		peer:      entry.peer,
  53  		keypair:   keypair,
  54  		handshake: nil,
  55  	}
  56  }
  57  
  58  func (table *IndexTable) NewIndexForHandshake(peer *Peer, handshake *Handshake) (uint32, error) {
  59  	for {
  60  		// generate random index
  61  
  62  		index, err := randUint32()
  63  		if err != nil {
  64  			return index, err
  65  		}
  66  
  67  		// check if index used
  68  
  69  		table.RLock()
  70  		_, ok := table.table[index]
  71  		table.RUnlock()
  72  		if ok {
  73  			continue
  74  		}
  75  
  76  		// check again while locked
  77  
  78  		table.Lock()
  79  		_, found := table.table[index]
  80  		if found {
  81  			table.Unlock()
  82  			continue
  83  		}
  84  		table.table[index] = IndexTableEntry{
  85  			peer:      peer,
  86  			handshake: handshake,
  87  			keypair:   nil,
  88  		}
  89  		table.Unlock()
  90  		return index, nil
  91  	}
  92  }
  93  
  94  func (table *IndexTable) Lookup(id uint32) IndexTableEntry {
  95  	table.RLock()
  96  	defer table.RUnlock()
  97  	return table.table[id]
  98  }
  99