<commit_before>package sia

import (
	"errors"
	"net"
	"reflect"
	"strconv"
	"time"
)

const (
	timeout   = time.Second * 5
	maxMsgLen = 1 << 16
)

// A NetAddress contains the information needed to contact a peer over TCP.
type NetAddress struct {
	Host string
	Port uint16
}

// String returns the NetAddress as a string, concatentating the hostname and
// port number.
func (na *NetAddress) String() string {
	return net.JoinHostPort(na.Host, strconv.Itoa(int(na.Port)))
}

// Call establishes a TCP connection to the NetAddress, calls the provided
// function on it, and closes the connection.
func (na *NetAddress) Call(fn func(net.Conn) error) error {
	conn, err := net.DialTimeout("tcp", na.String(), timeout)
	if err != nil {
		return err
	}
	defer conn.Close()
	return fn(conn)
}

// SendVal returns a closure that can be used in conjuction with Call to send
// a value to a NetAddress. It prefixes the encoded data with a header,
// containing the message's type and length
func SendVal(t byte, val interface{}) func(net.Conn) error {
	encVal := Marshal(val)
	encLen := EncUint64(uint64(len(encVal)))
	msg := append([]byte{t},
		append(encLen[:4], encVal...)...)

	return func(conn net.Conn) error {
		_, err := conn.Write(msg)
		return err
	}
}

// TBD
var BootstrapPeers = []NetAddress{}

// A TCPServer sends and receives messages. It also maintains an address book
// of peers to broadcast to and make requests of.
type TCPServer struct {
	net.Listener
	myAddr      NetAddress
	addressbook map[NetAddress]struct{}
	handlerMap  map[byte]func(net.Conn, []byte) error
}

// Broadcast calls the specified function on each peer in the address book.
func (tcps *TCPServer) Broadcast(fn func(net.Conn) error) {
	for addr := range tcps.addressbook {
		addr.Call(fn)
	}
}

// RegisterHandler registers a message type with a message handler. The
// existing handler for that type will be overwritten.
func (tcps *TCPServer) RegisterHandler(t byte, fn func(net.Conn, []byte) error) {
	tcps.handlerMap[t] = fn
}

// RegisterRPC is for simple handlers. A simple handler decodes the message
// data and passes it to fn. fn must have the type signature:
//   func(Type) error
// i.e. a 1-adic function that returns an error.
func (tcps *TCPServer) RegisterRPC(t byte, fn interface{}) error {
	// if fn not correct type, panic
	val, typ := reflect.ValueOf(fn), reflect.TypeOf(fn)
	if typ.Kind() != reflect.Func || typ.NumIn() != 1 ||
		typ.NumOut() != 1 || typ.Out(0) != reflect.TypeOf((*error)(nil)).Elem() {
		return errors.New("registered function has wrong type signature")
	}

	// create function:
	sfn := func(_ net.Conn, b []byte) error {
		v := reflect.New(typ.In(0))
		if err := Unmarshal(b, v.Interface()); err != nil {
			return err
		}
		if err := val.Call([]reflect.Value{v.Elem()})[0].Interface(); err != nil {
			return err.(error)
		}
		return nil
	}

	tcps.RegisterHandler(t, sfn)
	return nil
}

// NewTCPServer creates a TCPServer that listens on the specified port.
func NewTCPServer(port uint16) (tcps *TCPServer, err error) {
	tcpServ, err := net.Listen("tcp", ":"+strconv.Itoa(int(port)))
	if err != nil {
		return
	}
	tcps = &TCPServer{
		Listener:    tcpServ,
		addressbook: make(map[NetAddress]struct{}),
		// default handlers
		handlerMap: map[byte]func(net.Conn, []byte) error{
			'H': sendHostname,
			'P': tcps.sharePeers,
			'A': tcps.addPeer,
		},
	}

	// spawn listener
	go tcps.listen()
	return
}

// listen runs in the background, accepting incoming connections and serving
// them. listen will return after TCPServer.Close() is called, because the
// Accept() call will fail.
func (tcps *TCPServer) listen() {
	for {
		conn, err := tcps.Accept()
		if err != nil {
			return
		}
		// it is the handler's responsibility to close the connection
		go tcps.handleConn(conn)
	}
}

// handleConn reads header data from a connection, unmarshals the data
// structures it contains, and routes the data to other functions for
// processing.
// TODO: set deadlines?
func (tcps *TCPServer) handleConn(conn net.Conn) {
	defer conn.Close()
	var (
		msgHead []byte = make([]byte, 5)
		msgData []byte // length determined by msgHead
	)
	if n, err := conn.Read(msgHead); err != nil || n != 5 {
		// TODO: log error
		return
	}
	msgLen := DecUint64(msgHead[1:])
	if msgLen > maxMsgLen {
		// TODO: log error
		return
	}
	if msgLen > 0 {
		msgData = make([]byte, msgLen)
		if n, err := conn.Read(msgData); err != nil || uint64(n) != msgLen {
			// TODO: log error
			return
		}
	}

	// call registered handler for this message type
	if fn, ok := tcps.handlerMap[msgHead[0]]; ok {
		fn(conn, msgData)
		// TODO: log error
		// no wait, send the error?
	}
	return
}

// sendHostname replies to the send with the sender's external IP.
func sendHostname(conn net.Conn, _ []byte) error {
	_, err := conn.Write([]byte(conn.RemoteAddr().String()))
	return err
}

// sharePeers transmits at most 'num' peers over the connection.
// TODO: choose random peers?
func (tcps *TCPServer) sharePeers(conn net.Conn, msgData []byte) error {
	if len(msgData) != 1 {
		return errors.New("invalid number of peers")
	}
	num := msgData[0]
	var addrs []NetAddress
	for addr := range tcps.addressbook {
		if num == 0 {<commit_msg>add ReadPrefix helper<commit_after>package sia

import (
	"errors"
	"net"
	"reflect"
	"strconv"
	"time"
)

const (
	timeout   = time.Second * 5
	maxMsgLen = 1 << 16
)

// A NetAddress contains the information needed to contact a peer over TCP.
type NetAddress struct {
	Host string
	Port uint16
}

// String returns the NetAddress as a string, concatentating the hostname and
// port number.
func (na *NetAddress) String() string {
	return net.JoinHostPort(na.Host, strconv.Itoa(int(na.Port)))
}

// Call establishes a TCP connection to the NetAddress, calls the provided
// function on it, and closes the connection.
func (na *NetAddress) Call(fn func(net.Conn) error) error {
	conn, err := net.DialTimeout("tcp", na.String(), timeout)
	if err != nil {
		return err
	}
	defer conn.Close()
	return fn(conn)
}

func ReadPrefix(conn net.Conn) ([]byte, error) {
	prefix := make([]byte, 4)
	if n, err := conn.Read(prefix); err != nil || n != len(prefix) {
		return nil, errors.New("could not read length prefix")
	}
	msgLen := DecUint64(prefix)
	if msgLen > maxMsgLen {
		return nil, errors.New("bad message length")
	}
	msgData := make([]byte, msgLen)
	if msgLen > 0 {
		if n, err := conn.Read(msgData); err != nil || uint64(n) != msgLen {
			return nil, errors.New("could not read message content")
		}
	}
	return msgData, nil
}

// SendVal returns a closure that can be used in conjuction with Call to send
// a value to a NetAddress. It prefixes the encoded data with a header,
// containing the message's type and length
func SendVal(t byte, val interface{}) func(net.Conn) error {
	encVal := Marshal(val)
	encLen := EncUint64(uint64(len(encVal)))
	msg := append([]byte{t},
		append(encLen[:4], encVal...)...)

	return func(conn net.Conn) error {
		_, err := conn.Write(msg)
		return err
	}
}

// TBD
var BootstrapPeers = []NetAddress{}

// A TCPServer sends and receives messages. It also maintains an address book
// of peers to broadcast to and make requests of.
type TCPServer struct {
	net.Listener
	myAddr      NetAddress
	addressbook map[NetAddress]struct{}
	handlerMap  map[byte]func(net.Conn, []byte) error
}

// Broadcast calls the specified function on each peer in the address book.
func (tcps *TCPServer) Broadcast(fn func(net.Conn) error) {
	for addr := range tcps.addressbook {
		addr.Call(fn)
	}
}

// RegisterHandler registers a message type with a message handler. The
// existing handler for that type will be overwritten.
func (tcps *TCPServer) RegisterHandler(t byte, fn func(net.Conn, []byte) error) {
	tcps.handlerMap[t] = fn
}

// RegisterRPC is for simple handlers. A simple handler decodes the message
// data and passes it to fn. fn must have the type signature:
//   func(Type) error
// i.e. a 1-adic function that returns an error.
func (tcps *TCPServer) RegisterRPC(t byte, fn interface{}) error {
	// if fn not correct type, panic
	val, typ := reflect.ValueOf(fn), reflect.TypeOf(fn)
	if typ.Kind() != reflect.Func || typ.NumIn() != 1 ||
		typ.NumOut() != 1 || typ.Out(0) != reflect.TypeOf((*error)(nil)).Elem() {
		return errors.New("registered function has wrong type signature")
	}

	// create function:
	sfn := func(_ net.Conn, b []byte) error {
		v := reflect.New(typ.In(0))
		if err := Unmarshal(b, v.Interface()); err != nil {
			return err
		}
		if err := val.Call([]reflect.Value{v.Elem()})[0].Interface(); err != nil {
			return err.(error)
		}
		return nil
	}

	tcps.RegisterHandler(t, sfn)
	return nil
}

// NewTCPServer creates a TCPServer that listens on the specified port.
func NewTCPServer(port uint16) (tcps *TCPServer, err error) {
	tcpServ, err := net.Listen("tcp", ":"+strconv.Itoa(int(port)))
	if err != nil {
		return
	}
	tcps = &TCPServer{
		Listener:    tcpServ,
		addressbook: make(map[NetAddress]struct{}),
		// default handlers
		handlerMap: map[byte]func(net.Conn, []byte) error{
			'H': sendHostname,
			'P': tcps.sharePeers,
			'A': tcps.addPeer,
		},
	}

	// spawn listener
	go tcps.listen()
	return
}

// listen runs in the background, accepting incoming connections and serving
// them. listen will return after TCPServer.Close() is called, because the
// Accept() call will fail.
func (tcps *TCPServer) listen() {
	for {
		conn, err := tcps.Accept()
		if err != nil {
			return
		}
		// it is the handler's responsibility to close the connection
		go tcps.handleConn(conn)
	}
}

// handleConn reads header data from a connection, unmarshals the data
// structures it contains, and routes the data to other functions for
// processing.
// TODO: set deadlines?
func (tcps *TCPServer) handleConn(conn net.Conn) {
	defer conn.Close()
	msgType := make([]byte, 1)
	if n, err := conn.Read(msgType); err != nil || n != 1 {
		// TODO: log error
		return
	}
	msgData, err := ReadPrefix(conn)
	if err != nil {
		// TODO: log error
		return
	}
	// call registered handler for this message type
	if fn, ok := tcps.handlerMap[msgType[0]]; ok {
		fn(conn, msgData)
		// TODO: log error
		// no wait, send the error?
	}
	return
}

// sendHostname replies to the send with the sender's external IP.
func sendHostname(conn net.Conn, _ []byte) error {
	_, err := conn.Write([]byte(conn.RemoteAddr().String()))
	return err
}

// sharePeers transmits at most 'num' peers over the connection.
// TODO: choose random peers?
func (tcps *TCPServer) sharePeers(conn net.Conn, msgData []byte) error {
	if len(msgData) != 1 {
		return errors.New("invalid number of peers")
	}
	num := msgData[0]
	var addrs []NetAddress
	for addr := range tcps.addressbook {
		if num == 0 {
{'fraction_non_alphanumeric': 0.13288603112477868, 'fraction_numerical': 0.00577765352716429, 'mean_word_length': 3.6600086843247936, 'pattern_counts': {'":': 2, '<': 7, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5652215', 'n_tokens_mistral': 3701, 'n_tokens_neox': 3559, 'n_words': 1228}