DHT-kademlia-P2Psec/pkg/dht_utils.go
2024-09-13 14:04:40 +09:00

163 lines
3.6 KiB
Go

package kademlia
import (
"crypto/sha256"
"crypto/tls"
"crypto/x509"
"encoding/binary"
"errors"
"net"
"slices"
"sort"
"strconv"
"time"
)
func GetPeerIdFromAddr(addr net.Addr) [32]byte {
addrTCP := addr.(*net.TCPAddr)
ipPortConcat := binary.BigEndian.AppendUint16(addrTCP.IP.To16(), uint16(addrTCP.Port))
peerId := sha256.Sum256(ipPortConcat)
return peerId
}
func GetPeerId(ip net.IP, port uint16) [32]byte {
val := binary.BigEndian.AppendUint16(ip.To16(), port)
res := sha256.Sum256(val)
return res
}
func CreateNode(ip string, port uint64, cert []byte) P2PNode {
ipA := net.ParseIP(ip)
return P2PNode{
NodeID: GetPeerId(ipA, uint16(port)),
IPAddress: ipA,
Port: port,
Certificate: cert,
}
}
func NodeFromIpPort(ip net.IP, port uint64, cert []byte) P2PNode {
return P2PNode{
NodeID: GetPeerId(ip, uint16(port)),
IPAddress: ip,
Port: port,
Certificate: cert,
}
}
func ContainsID(nodes []P2PNode, target [32]byte) bool {
for _, node := range nodes {
if node.NodeID == target {
return true
}
}
return false
}
func SortNodesByDistanceToKey(key [32]byte, nodes []P2PNode) []P2PNode {
sort.SliceStable(nodes, func(i, j int) bool {
return IsCloserToKey(nodes[i].NodeID, nodes[j].NodeID, key)
})
return nodes
}
func RemoveDuplicates(nodes []P2PNode) []P2PNode {
return slices.CompactFunc(nodes, func(node P2PNode, node2 P2PNode) bool {
return node.NodeID == node2.NodeID
})
}
func NodeEquals(node P2PNode, node2 P2PNode) bool {
return node.NodeID == node2.NodeID
}
// IsCloserToKey returns whether the first argument is closer to the key than the second one when using the XOR metric
// everything is closer to a key than nil.
func IsCloserToKey(a, b, key [32]byte) bool {
for i, _ := range key {
if a[i] == b[i] {
continue
} else {
return a[i]^key[i] < b[i]^key[i]
}
}
return false
}
/*
ConnectToNode attempts to establish a connection to the provided node and puts a 30 second timeout on it.
*/
func ConnectToNode(node P2PNode) (net.Conn, error) {
certPool := x509.NewCertPool()
if ok := certPool.AppendCertsFromPEM(node.Certificate); !ok {
return nil, errors.New("[Utils]: Unable to parse cert")
}
tlsConfig := &tls.Config{RootCAs: certPool}
addr := node.IPAddress.String() + ":" + strconv.Itoa(int(node.Port))
conn, err := tls.Dial("tcp", addr, tlsConfig)
if err == nil {
// Set lenient timeout in case of unresponsive other peer
_ = conn.SetDeadline(time.Now().Add(time.Second * 30))
}
return conn, err
}
func ContainsNode(bucket []P2PNode, node P2PNode) bool {
for _, n := range bucket {
if n.NodeID == node.NodeID {
return true
}
}
return false
}
/*
getBitsAt retrieves a specified number of bits from a [32]byte. Beginning with the bit at index startBit, it
retrieves numBits bits. numBits cannot be greater than 64 and has to stay within bounds.
*/
func getBitsAt(id [32]byte, startBit uint64, numBits uint64) uint64 {
if numBits <= 0 || startBit < 0 || startBit+numBits > 256 || numBits > 64 {
return uint64(0)
}
var result uint64
var bitPos uint64
for i := uint64(0); i < numBits; i++ {
byteIndex := (startBit + i) / 8
bitIndex := (startBit + i) % 8
bit := (id[byteIndex] >> (7 - bitIndex)) & 1
result = (result << 1) | uint64(bit)
bitPos++
}
return result
}
/*
setBitAt sets the bit of a [32]byte at the specified index to the specified value
*/
func setBitAt(id [32]byte, index int, value bool) [32]byte {
if index >= 256 {
return id
}
byteIndex := index / 8
bitIndex := index % 8
if value {
// Set the bit to 1
id[byteIndex] |= 1 << uint(7-bitIndex)
} else {
// Set the bit to 0
id[byteIndex] &^= 1 << uint(7-bitIndex)
}
return id
}