github.com/liloew/wireguard-go@v0.0.0-20220224014633-9cd745e6f114/device/indextable.go (about)

     1  /* SPDX-License-Identifier: MIT
     2   *
     3   * Copyright (C) 2017-2021 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  }