DHT-kademlia-P2Psec/pkg/packets.go

package kademlia

import (
	"crypto/sha256"
	"encoding/binary"
	"errors"
	"log"
	"net"
)

const (
	PUT           = 650
	GET           = 651
	SUCCESS       = 652
	FAILURE       = 653
	STORE         = 660
	PING          = 661
	FINDNODE      = 662
	FINDVALUE     = 663
	PONG          = 671
	FOUNDNODE     = 672
	FOUNDVALUE    = 673
	NOTFOUNDVALUE = 674
)

type DHTPutPacket struct {
	TTL         uint16
	Replication byte
	Key         [32]byte
	Value       []byte
}

type DHTGetPacket struct {
	Key [32]byte
}

type DHTSuccessPacket struct {
	Key   [32]byte
	Value []byte
}

type DHTFailurePacket struct {
	Key [32]byte
}

func SerialiseSuccessPacket(packet DHTSuccessPacket) []byte {
	//Placeholder for size
	buf := make([]byte, 2)
	// Message type
	buf = binary.BigEndian.AppendUint16(buf, uint16(SUCCESS))
	// Key
	buf = append(buf, packet.Key[:]...)
	// Value
	buf = append(buf, packet.Value...)
	// Size
	binary.BigEndian.PutUint16(buf, uint16(len(buf)))
	return buf
}

func ParseSucceessPacket(buf []byte) (*DHTSuccessPacket, error) {
	if len(buf) < 36 {
		return nil, errors.New("missing fields")
	}

	return &DHTSuccessPacket{
		Key:   [32]byte(buf[4:36]),
		Value: buf[36:],
	}, nil
}

func SerialiseFailurePacket(packet DHTFailurePacket) []byte {
	//Placeholder for size
	buf := make([]byte, 2)
	// Message type
	buf = binary.BigEndian.AppendUint16(buf, uint16(FAILURE))
	// Key
	buf = append(buf, packet.Key[:]...)
	// Size
	binary.BigEndian.PutUint16(buf, uint16(len(buf)))
	return buf
}

func ParsePutPacket(buf []byte) (*DHTPutPacket, error) {
	if len(buf) <= 40 {
		return nil, errors.New("missing fields")
	}
	return &DHTPutPacket{
		TTL:         binary.BigEndian.Uint16(buf[4:6]),
		Replication: buf[6],
		Key:         [32]byte(buf[8:40]),
		Value:       buf[40:],
	}, nil
}

func SerialisePutPacket(packet *DHTPutPacket) []byte {
	//Placeholder for size
	buf := make([]byte, 2)
	// Message type
	buf = binary.BigEndian.AppendUint16(buf, uint16(PUT))
	// Time to live
	buf = binary.BigEndian.AppendUint16(buf, packet.TTL)
	// Replication and reserved
	buf = append(buf, packet.Replication, 0)
	// Key
	buf = append(buf, packet.Key[:]...)
	// Value
	buf = append(buf, packet.Value...)
	// Size
	binary.BigEndian.PutUint16(buf, uint16(len(buf)))
	return buf
}

func ParseGetPacket(buf []byte) (*DHTGetPacket, error) {
	if len(buf) < 36 {
		return nil, errors.New("missing fields")
	}

	return &DHTGetPacket{
		Key: [32]byte(buf[4:36]),
	}, nil
}

func SerialiseGetPacket(packet *DHTGetPacket) []byte {
	//Placeholder for size
	buf := make([]byte, 2)
	// Message type
	buf = binary.BigEndian.AppendUint16(buf, uint16(GET))
	// Key
	buf = append(buf, packet.Key[:]...)
	// Size
	binary.BigEndian.PutUint16(buf, uint16(len(buf)))
	return buf
}

type P2PStorePacket struct {
	SourceIP    net.IP
	SourcePort  uint16
	TTL         uint16
	RepCount    uint16
	Certificate []byte
	PoW         [32]byte
	Key         [32]byte
	Value       []byte
}

type P2PPingPacket struct {
	SourceIP    net.IP
	SourcePort  uint16
	Certificate []byte
}

type P2PFindNodePacket struct {
	SourceIP    net.IP
	SourcePort  uint16
	Certificate []byte
	Key         [32]byte
	Amount      uint16
}

type P2PFindValuePacket struct {
	SourceIP    net.IP
	SourcePort  uint16
	Certificate []byte
	Key         [32]byte
}

type P2PFoundNodePacket struct {
	Nodes []P2PNode
}

type P2PFoundValuePacket struct {
	Value []byte
}

type P2PPongPacket struct {
}

type P2PNotFoundPValuePacket struct {
}

func ProofWorkForStore(key [32]byte, value []byte) [32]byte {
	payload := append(key[:], value...)
	var hash [32]byte
	log.Printf("Starting pow")
	for {
		res := sha256.Sum256(append(hash[:], payload...))
		if getBitsAt(res, 0, ProofOfWorkBits) == 0 {
			break
		}
		hash = res
	}
	log.Printf("Ending pow")
	return hash
}

func SerialiseStorePacket(packet *P2PStorePacket) []byte {
	//Placeholder for size
	buf := make([]byte, 2)
	// Message type
	buf = binary.BigEndian.AppendUint16(buf, uint16(STORE))
	// IP
	buf = append(buf, packet.SourceIP.To16()...)
	// Port
	buf = binary.BigEndian.AppendUint16(buf, packet.SourcePort)
	// Certificate length
	buf = binary.BigEndian.AppendUint16(buf, uint16(len(packet.Certificate)))
	// Cert
	buf = append(buf, packet.Certificate...)
	// Rep Count
	buf = binary.BigEndian.AppendUint16(buf, packet.RepCount)
	// TTL
	buf = binary.BigEndian.AppendUint16(buf, packet.TTL)
	// ProofOfWork
	buf = append(buf, packet.PoW[:]...)
	// Key
	buf = append(buf, packet.Key[:]...)
	// Value
	buf = append(buf, packet.Value...)
	// Size
	binary.BigEndian.PutUint16(buf, uint16(len(buf)))
	return buf
}

func ParseStorePacket(buf []byte) (*P2PStorePacket, error) {
	if len(buf) <= 92 {
		return nil, errors.New("packet length too short")
	}

	certLen := binary.BigEndian.Uint16(buf[22:24])

	packet := &P2PStorePacket{
		SourceIP:    net.IP(buf[4:20]),
		SourcePort:  binary.BigEndian.Uint16(buf[20:22]),
		Certificate: buf[24 : 24+certLen],
		TTL:         binary.BigEndian.Uint16(buf[24+certLen : 26+certLen]),
		RepCount:    binary.BigEndian.Uint16(buf[26+certLen : 28+certLen]),
		PoW:         [32]byte(buf[28+certLen : 60+certLen]),
		Key:         [32]byte(buf[60+certLen : 92+certLen]),
		Value:       buf[92+certLen:],
	}
	hash := sha256.Sum256(buf[28+certLen:])
	if getBitsAt(hash, 0, ProofOfWorkBits) != 0 {
		return nil, errors.New("invalid proof of work")
	}
	return packet, nil
}

func SerialisePingPacket(packet *P2PPingPacket) []byte {
	//Placeholder for size
	buf := make([]byte, 2)
	// Message type
	buf = binary.BigEndian.AppendUint16(buf, uint16(PING))
	// IP
	buf = append(buf, packet.SourceIP.To16()...)
	// Port
	buf = binary.BigEndian.AppendUint16(buf, uint16(packet.SourcePort))
	// Certificate length
	buf = binary.BigEndian.AppendUint16(buf, uint16(len(packet.Certificate)))
	// Cert
	buf = append(buf, packet.Certificate...)
	// Size
	binary.BigEndian.PutUint16(buf, uint16(len(buf)))
	return buf
}

func ParsePingPacket(buf []byte) (*P2PPingPacket, error) {
	if len(buf) < 22 {
		return nil, errors.New("packet length too short")
	}

	certLen := binary.BigEndian.Uint16(buf[22:24])

	packet := &P2PPingPacket{
		SourceIP:    net.IP(buf[4:20]),
		SourcePort:  binary.BigEndian.Uint16(buf[20:22]),
		Certificate: buf[24 : 24+certLen],
	}
	return packet, nil
}

func SerialiseFindNodePacket(packet *P2PFindNodePacket) []byte {
	//Placeholder for size
	buf := make([]byte, 2)
	// Message type
	buf = binary.BigEndian.AppendUint16(buf, uint16(FINDNODE))
	// IP
	buf = append(buf, packet.SourceIP.To16()...)
	// Port
	buf = binary.BigEndian.AppendUint16(buf, packet.SourcePort)
	// Certificate length
	buf = binary.BigEndian.AppendUint16(buf, uint16(len(packet.Certificate)))
	// Cert
	buf = append(buf, packet.Certificate...)
	// Key
	buf = append(buf, packet.Key[:]...)
	// Amount
	buf = binary.BigEndian.AppendUint16(buf, packet.Amount)
	// Size
	binary.BigEndian.PutUint16(buf, uint16(len(buf)))
	return buf
}

func ParseFindNodePacket(buf []byte) (*P2PFindNodePacket, error) {
	if len(buf) < 56 {
		return nil, errors.New("packet length too short")
	}

	certLen := binary.BigEndian.Uint16(buf[22:24])

	packet := &P2PFindNodePacket{
		SourceIP:    net.IP(buf[4:20]),
		SourcePort:  binary.BigEndian.Uint16(buf[20:22]),
		Certificate: buf[24 : 24+certLen],
		Key:         [32]byte(buf[24+certLen : 56+certLen]),
		Amount:      binary.BigEndian.Uint16(buf[56+certLen : 58+certLen]),
	}

	return packet, nil
}

func SerialiseFindValuePacket(packet *P2PFindValuePacket) []byte {
	//Placeholder for size
	buf := make([]byte, 2)
	// Message type
	buf = binary.BigEndian.AppendUint16(buf, uint16(FINDVALUE))
	// IP
	buf = append(buf, packet.SourceIP.To16()...)
	// Port
	buf = binary.BigEndian.AppendUint16(buf, packet.SourcePort)
	// Certificate length
	buf = binary.BigEndian.AppendUint16(buf, uint16(len(packet.Certificate)))
	// Cert
	buf = append(buf, packet.Certificate...)
	// Key
	buf = append(buf, packet.Key[:]...)
	// Size
	binary.BigEndian.PutUint16(buf, uint16(len(buf)))
	return buf
}

func ParseFindValuePacket(buf []byte) (*P2PFindValuePacket, error) {
	if len(buf) < 56 {
		return nil, errors.New("packet length too short")
	}

	certLen := binary.BigEndian.Uint16(buf[22:24])

	packet := &P2PFindValuePacket{
		SourceIP:    net.IP(buf[4:20]),
		SourcePort:  binary.BigEndian.Uint16(buf[20:22]),
		Certificate: buf[24 : 24+certLen],
		Key:         [32]byte(buf[24+certLen : 56+certLen]),
	}

	return packet, nil
}

func generateNodesList(nodes []P2PNode) []byte {
	nodeList := make([]byte, 0)

	for _, node := range nodes {
		nodeList = append(nodeList, net.IP.To16(node.IPAddress)...)
		nodeList = binary.BigEndian.AppendUint16(nodeList, uint16(node.Port))
		nodeList = binary.BigEndian.AppendUint16(nodeList, uint16(len(node.Certificate)))
		nodeList = append(nodeList, node.Certificate...)
	}
	return nodeList
}

func SerialiseFoundNodePacket(packet *P2PFoundNodePacket) []byte {
	buf := make([]byte, 2)
	// Type
	buf = binary.BigEndian.AppendUint16(buf, FOUNDNODE)
	// reserved
	buf = binary.BigEndian.AppendUint16(buf, 0)
	// number of nodes
	buf = binary.BigEndian.AppendUint16(buf, uint16(len(packet.Nodes)))
	// list of nodes
	buf = append(buf, generateNodesList(packet.Nodes)...)
	// Size
	binary.BigEndian.PutUint16(buf, uint16(len(buf)))
	return buf
}

func ParseFoundNodesPacket(packetToParse []byte) (*P2PFoundNodePacket, error) {
	if len(packetToParse) < 8 {
		return nil, errors.New("Invalid FoundNodes header")
	}
	numberOfNodes := binary.BigEndian.Uint16(packetToParse[6:8])

	nodes := make([]P2PNode, numberOfNodes)

	currentIndex := 8
	for i := 0; i < int(numberOfNodes); i++ {
		if len(packetToParse) < currentIndex+20 {
			return nil, errors.New("received invalid packet, fields not present")
		}
		ip := packetToParse[currentIndex : currentIndex+16]
		port := uint64(binary.BigEndian.Uint16(packetToParse[currentIndex+16 : currentIndex+18]))
		certLen := int(binary.BigEndian.Uint16(packetToParse[currentIndex+18 : currentIndex+20]))
		currentIndex += 20
		if len(packetToParse) < currentIndex+certLen {
			return nil, errors.New("received invalid packet, certificate too short")
		}
		cert := packetToParse[currentIndex : currentIndex+certLen]
		currentIndex += certLen
		nodes[i] = P2PNode{
			NodeID:      GetPeerId(ip, uint16(port)),
			IPAddress:   ip,
			Port:        port,
			Certificate: cert,
		}
	}
	return &P2PFoundNodePacket{Nodes: nodes}, nil
}

func SerialiseFoundValuePacket(packet *P2PFoundValuePacket) []byte {
	//Placeholder for size
	buf := make([]byte, 2)
	// Message type
	buf = binary.BigEndian.AppendUint16(buf, uint16(FOUNDVALUE))
	// Value
	buf = append(buf, packet.Value...)
	// Size
	binary.BigEndian.PutUint16(buf, uint16(len(buf)))
	return buf
}

func ParseFoundValuePacket(buf []byte) *P2PFoundValuePacket {
	return &P2PFoundValuePacket{Value: buf[4:]}
}

func SerialiseNotFoundValuePacket(packet *P2PNotFoundPValuePacket) []byte {
	buf := make([]byte, 2)
	buf = binary.BigEndian.AppendUint16(buf, NOTFOUNDVALUE)
	binary.BigEndian.PutUint16(buf, uint16(len(buf)))
	return buf
}

func SerialisePongPacket(packet *P2PPongPacket) []byte {
	buf := make([]byte, 2)
	buf = binary.BigEndian.AppendUint16(buf, PONG)
	binary.BigEndian.PutUint16(buf, uint16(len(buf)))
	return buf
}