woop/p2p/helper.go

package p2p

import (
	"bufio"
	"bytes"
	"encoding/binary"
	"io"
	"log"
	"net"
	"time"
)

/*
P2p Message data structure:

---- message start -----
1 byte            - p2p type
                    0x00: unicast (no need to forward)
                    0x01: broadcast (may need to forward to other neighbors)
4 bytes           - message content size n in number of bytes
content (n bytes) - actual message content
----  message end  -----


*/

// Read the message type and content size, and return the actual content.
func ReadMessageContent(conn net.Conn) ([]byte, error) {
	var (
		contentBuf = bytes.NewBuffer([]byte{})
		r          = bufio.NewReader(conn)
	)

	timeoutDuration := 1 * time.Second
	conn.SetReadDeadline(time.Now().Add(timeoutDuration))

	//// Read 1 byte for message type
	_, err := r.ReadByte()
	switch err {
	case io.EOF:
		log.Printf("Error reading the p2p message type field: %s", err)
		return contentBuf.Bytes(), err
	case nil:
		//log.Printf("Received p2p message type: %x\n", msgType)
	default:
		log.Printf("Error reading the p2p message type field: %s", err)
		return contentBuf.Bytes(), err
	}
	// TODO: check on msgType and take actions accordingly

	//// Read 4 bytes for message size
	fourBytes := make([]byte, 4)
	n, err := r.Read(fourBytes)
	if err != nil {
		log.Printf("Error reading the p2p message size field")
		return contentBuf.Bytes(), err
	} else if n < len(fourBytes) {
		log.Printf("Failed reading the p2p message size field: only read %d bytes", n)
		return contentBuf.Bytes(), err
	}

	//log.Print(fourBytes)
	// Number of bytes for the message content
	bytesToRead := binary.BigEndian.Uint32(fourBytes)
	//log.Printf("The content size is %d bytes.", bytesToRead)

	//// Read the content in chunk of 1024 bytes
	tmpBuf := make([]byte, 1024)
ILOOP:
	for {
		timeoutDuration := 1 * time.Second
		conn.SetReadDeadline(time.Now().Add(timeoutDuration))
		if bytesToRead < 1024 {
			// Read the last number of bytes less than 1024
			tmpBuf = make([]byte, bytesToRead)
		}
		n, err := r.Read(tmpBuf)
		contentBuf.Write(tmpBuf[:n])

		switch err {
		case io.EOF:
			// TODO: should we return error here, or just ignore it?
			log.Printf("EOF reached while reading p2p message")
			break ILOOP
		case nil:
			bytesToRead -= uint32(n) // TODO: think about avoid the casting in every loop
			if bytesToRead <= 0 {
				break ILOOP
			}
		default:
			log.Printf("Error reading p2p message")
			return []byte{}, err
		}
	}
	return contentBuf.Bytes(), nil
}

func CreateMessage(msgType byte, data []byte) []byte {
	buffer := bytes.NewBuffer([]byte{})

	buffer.WriteByte(msgType)

	fourBytes := make([]byte, 4)
	binary.BigEndian.PutUint32(fourBytes, uint32(len(data)))
	buffer.Write(fourBytes)

	buffer.Write(data)
	return buffer.Bytes()
}

func SendMessageContent(conn net.Conn, data []byte) {
	msgToSend := CreateMessage(byte(1), data)
	w := bufio.NewWriter(conn)
	w.Write(msgToSend)
	w.Flush()
}