github.com/noisysockets/noisysockets@v0.21.2-0.20240515114641-7f467e651c90/internal/transport/indextable.go (about)

     1  // SPDX-License-Identifier: MPL-2.0
     2  /*
     3   * Copyright (C) 2024 The Noisy Sockets Authors.
     4   *
     5   * This Source Code Form is subject to the terms of the Mozilla Public
     6   * License, v. 2.0. If a copy of the MPL was not distributed with this
     7   * file, You can obtain one at http://mozilla.org/MPL/2.0/.
     8   *
     9   * Portions of this file are based on code originally from wireguard-go,
    10   *
    11   * Copyright (C) 2017-2023 WireGuard LLC. All Rights Reserved.
    12   *
    13   * Permission is hereby granted, free of charge, to any person obtaining a copy of
    14   * this software and associated documentation files (the "Software"), to deal in
    15   * the Software without restriction, including without limitation the rights to
    16   * use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies
    17   * of the Software, and to permit persons to whom the Software is furnished to do
    18   * so, subject to the following conditions:
    19   *
    20   * The above copyright notice and this permission notice shall be included in all
    21   * copies or substantial portions of the Software.
    22   *
    23   * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
    24   * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
    25   * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
    26   * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
    27   * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
    28   * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
    29   * SOFTWARE.
    30   */
    31  
    32  package transport
    33  
    34  import (
    35  	"crypto/rand"
    36  	"encoding/binary"
    37  	"sync"
    38  )
    39  
    40  type IndexTableEntry struct {
    41  	peer      *Peer
    42  	handshake *Handshake
    43  	keypair   *Keypair
    44  }
    45  
    46  type IndexTable struct {
    47  	sync.RWMutex
    48  	table map[uint32]IndexTableEntry
    49  }
    50  
    51  func randUint32() (uint32, error) {
    52  	var integer [4]byte
    53  	_, err := rand.Read(integer[:])
    54  	// Arbitrary endianness; both are intrinsified by the Go compiler.
    55  	return binary.LittleEndian.Uint32(integer[:]), err
    56  }
    57  
    58  func (table *IndexTable) Init() {
    59  	table.Lock()
    60  	defer table.Unlock()
    61  	table.table = make(map[uint32]IndexTableEntry)
    62  }
    63  
    64  func (table *IndexTable) Delete(index uint32) {
    65  	table.Lock()
    66  	defer table.Unlock()
    67  	delete(table.table, index)
    68  }
    69  
    70  func (table *IndexTable) SwapIndexForKeypair(index uint32, keypair *Keypair) {
    71  	table.Lock()
    72  	defer table.Unlock()
    73  	entry, ok := table.table[index]
    74  	if !ok {
    75  		return
    76  	}
    77  	table.table[index] = IndexTableEntry{
    78  		peer:      entry.peer,
    79  		keypair:   keypair,
    80  		handshake: nil,
    81  	}
    82  }
    83  
    84  func (table *IndexTable) NewIndexForHandshake(peer *Peer, handshake *Handshake) (uint32, error) {
    85  	for {
    86  		// generate random index
    87  
    88  		index, err := randUint32()
    89  		if err != nil {
    90  			return index, err
    91  		}
    92  
    93  		// check if index used
    94  
    95  		table.RLock()
    96  		_, ok := table.table[index]
    97  		table.RUnlock()
    98  		if ok {
    99  			continue
   100  		}
   101  
   102  		// check again while locked
   103  
   104  		table.Lock()
   105  		_, found := table.table[index]
   106  		if found {
   107  			table.Unlock()
   108  			continue
   109  		}
   110  		table.table[index] = IndexTableEntry{
   111  			peer:      peer,
   112  			handshake: handshake,
   113  			keypair:   nil,
   114  		}
   115  		table.Unlock()
   116  		return index, nil
   117  	}
   118  }
   119  
   120  func (table *IndexTable) Lookup(id uint32) IndexTableEntry {
   121  	table.RLock()
   122  	defer table.RUnlock()
   123  	return table.table[id]
   124  }