go-libp2p/itest/echo.go

package itest

import (
	"context"
	"fmt"
	"io"
	"sync"
	"time"

	"github.com/libp2p/go-libp2p-core/host"
	"github.com/libp2p/go-libp2p-core/network"
	"github.com/libp2p/go-libp2p-core/peer"

	logging "github.com/ipfs/go-log/v2"
)

const (
	EchoService = "test.echo"
	EchoProtoID = "/test/echo"
)

var (
	echoLog = logging.Logger("echo")
)

type Echo struct {
	Host host.Host

	mx     sync.Mutex
	status EchoStatus

	beforeReserve, beforeRead, beforeWrite, beforeDone func() error
	done                                               func()
}

type EchoStatus struct {
	StreamsIn                 int
	EchosIn, EchosOut         int
	IOErrors                  int
	ResourceServiceErrors     int
	ResourceReservationErrors int
}

func NewEcho(h host.Host) *Echo {
	e := &Echo{Host: h}
	h.SetStreamHandler(EchoProtoID, e.handleStream)
	return e
}

func (e *Echo) Status() EchoStatus {
	e.mx.Lock()
	defer e.mx.Unlock()

	return e.status
}

func (e *Echo) BeforeReserve(f func() error) {
	e.mx.Lock()
	defer e.mx.Unlock()

	e.beforeReserve = f
}

func (e *Echo) BeforeRead(f func() error) {
	e.mx.Lock()
	defer e.mx.Unlock()

	e.beforeRead = f
}

func (e *Echo) BeforeWrite(f func() error) {
	e.mx.Lock()
	defer e.mx.Unlock()

	e.beforeWrite = f
}

func (e *Echo) BeforeDone(f func() error) {
	e.mx.Lock()
	defer e.mx.Unlock()

	e.beforeDone = f
}

func (e *Echo) Done(f func()) {
	e.mx.Lock()
	defer e.mx.Unlock()

	e.done = f
}

func (e *Echo) getBeforeReserve() func() error {
	e.mx.Lock()
	defer e.mx.Unlock()

	return e.beforeReserve
}

func (e *Echo) getBeforeRead() func() error {
	e.mx.Lock()
	defer e.mx.Unlock()

	return e.beforeRead
}

func (e *Echo) getBeforeWrite() func() error {
	e.mx.Lock()
	defer e.mx.Unlock()

	return e.beforeWrite
}

func (e *Echo) getBeforeDone() func() error {
	e.mx.Lock()
	defer e.mx.Unlock()

	return e.beforeDone
}

func (e *Echo) getDone() func() {
	e.mx.Lock()
	defer e.mx.Unlock()

	return e.done
}

func (e *Echo) handleStream(s network.Stream) {
	defer s.Close()

	if done := e.getDone(); done != nil {
		defer done()
	}

	e.mx.Lock()
	e.status.StreamsIn++
	e.mx.Unlock()

	if beforeReserve := e.getBeforeReserve(); beforeReserve != nil {
		if err := beforeReserve(); err != nil {
			echoLog.Debugf("error syncing before reserve: %s", err)

			s.Reset()
			return
		}
	}

	if err := s.Scope().SetService(EchoService); err != nil {
		echoLog.Debugf("error attaching stream to echo service: %s", err)

		e.mx.Lock()
		e.status.ResourceServiceErrors++
		e.mx.Unlock()

		s.Reset()
		return
	}

	if err := s.Scope().ReserveMemory(4096, network.ReservationPriorityAlways); err != nil {
		echoLog.Debugf("error reserving memory: %s", err)

		e.mx.Lock()
		e.status.ResourceReservationErrors++
		e.mx.Unlock()

		s.Reset()
		return
	}

	if beforeRead := e.getBeforeRead(); beforeRead != nil {
		if err := beforeRead(); err != nil {
			echoLog.Debugf("error syncing before read: %s", err)

			s.Reset()
			return
		}
	}

	buf := make([]byte, 4096)

	s.SetReadDeadline(time.Now().Add(5 * time.Second))
	n, err := s.Read(buf)
	switch {
	case err == io.EOF:
		if n == 0 {
			return
		}

	case err != nil:
		echoLog.Debugf("I/O error : %s", err)

		e.mx.Lock()
		e.status.IOErrors++
		e.mx.Unlock()

		s.Reset()
		return
	}

	e.mx.Lock()
	e.status.EchosIn++
	e.mx.Unlock()

	if beforeWrite := e.getBeforeWrite(); beforeWrite != nil {
		if err := beforeWrite(); err != nil {
			echoLog.Debugf("error syncing before write: %s", err)

			s.Reset()
			return
		}
	}

	s.SetWriteDeadline(time.Now().Add(5 * time.Second))
	_, err = s.Write(buf[:n])
	if err != nil {
		echoLog.Debugf("I/O error: %s", err)

		e.mx.Lock()
		e.status.IOErrors++
		e.mx.Unlock()

		s.Reset()
		return
	}

	e.mx.Lock()
	e.status.EchosOut++
	e.mx.Unlock()

	s.CloseWrite()

	if beforeDone := e.getBeforeDone(); beforeDone != nil {
		if err := beforeDone(); err != nil {
			echoLog.Debugf("error syncing before done: %s", err)

			s.Reset()
		}
	}
}

func (e *Echo) Echo(p peer.ID, what string) error {
	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
	defer cancel()

	s, err := e.Host.NewStream(ctx, p, EchoProtoID)
	if err != nil {
		return err
	}
	defer s.Close()

	if err := s.Scope().SetService(EchoService); err != nil {
		echoLog.Debugf("error attaching stream to echo service: %s", err)

		s.Reset()
		return err
	}

	if err := s.Scope().ReserveMemory(4096, network.ReservationPriorityAlways); err != nil {
		echoLog.Debugf("error reserving memory: %s", err)

		s.Reset()
		return err
	}

	s.SetWriteDeadline(time.Now().Add(5 * time.Second))
	_, err = s.Write([]byte(what))
	if err != nil {
		return err
	}
	s.CloseWrite()

	buf := make([]byte, 4096)

	s.SetReadDeadline(time.Now().Add(5 * time.Second))
	n, err := s.Read(buf)
	switch {
	case err == io.EOF:
		if n == 0 {
			return err
		}

	case err != nil:
		echoLog.Debugf("I/O error : %s", err)

		s.Reset()
		return err
	}

	if what != string(buf[:n]) {
		return fmt.Errorf("echo output doesn't match input")
	}

	return nil
}