transport.go

package main

import (
	"context"
	"errors"
	"fmt"

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

	tpt "github.com/libp2p/go-libp2p/core/transport"

	ma "github.com/multiformats/go-multiaddr"

	mafmt "github.com/multiformats/go-multiaddr-fmt"

	"github.com/joomcode/errorx"
)

type CustomTransport struct {
	// Used to upgrade unsecure TCP connections to secure multiplexed and
	// authenticate connections.
	upgrader tpt.Upgrader

	myCustomNode *CustomProtocol
}

type Option func(*CustomTransport) error

func WithCustomNode(h *CustomProtocol) Option {
	return func(t *CustomTransport) error {
		t.myCustomNode = h
		return nil
	}
}

func NewCustomTransport(u tpt.Upgrader, rcmgr network.ResourceManager, opts ...Option) (tpt.Transport, error) {
	t := &CustomTransport{
		upgrader: u,
	}

	for _, opt := range opts {
		if err := opt(t); err != nil {
			return nil, err
		}
	}

	if t.myCustomNode == nil {
		return nil, errors.New("custom node required")
	}

	return t, nil
}

func (*CustomTransport) Proxy() bool {
	return false
}

var supportedProtos = []int{P_CUSTOM}

func (t *CustomTransport) Protocols() []int {
	return supportedProtos
}

var matcher = mafmt.And(
	mafmt.IP,
	mafmt.Base(ma.P_TCP),
	mafmt.Base(P_CUSTOM),
)

func (t *CustomTransport) CanDial(maddr ma.Multiaddr) bool {
	return matcher.Matches(maddr)
}

func (t *CustomTransport) Close() {
	// TODO
}

type transportListener struct {
	tpt.Listener
}

func (l *transportListener) Accept() (tpt.CapableConn, error) {
	conn, err := l.Listener.Accept()
	if err != nil {
		return nil, err
	}
	return &capableConn{CapableConn: conn}, nil
}

func (t *CustomTransport) Listen(laddr ma.Multiaddr) (tpt.Listener, error) {
	malist := NewCustomListener(t.myCustomNode)
	return &transportListener{Listener: t.upgrader.UpgradeListener(t, malist)}, nil
}

func (t *CustomTransport) Dial(ctx context.Context, raddr ma.Multiaddr, p peer.ID) (tpt.CapableConn, error) {
	if !t.CanDial(raddr) {
		return nil, errorx.IllegalArgument.New(fmt.Sprintf("Can't dial \"%s\".", raddr))
	}

	conn, err := t.dialAndUpgrade(ctx, raddr, p)
	if err != nil {
		return nil, err
	}

	return conn, nil
}

func (c *CustomTransport) dialAndUpgrade(ctx context.Context, a ma.Multiaddr, p peer.ID) (tpt.CapableConn, error) {
	ctx, cancel := context.WithCancel(ctx)

	conn := &CustomConn{
		myCustomNode: c.myCustomNode,
		remotePeer:   EncapsulatePeerID(p, a),
		ctx:          ctx,
		cancel:       cancel,
	}
	cc, err := c.upgrader.Upgrade(ctx, c, conn, network.DirOutbound, p, &network.NullScope{})
	if err != nil {
		return nil, err
	}
	fmt.Println("> UPGRADE DONE")
	return capableConn{cc}, nil
}