Merge pull request #1337 from keep-network/unicast-local

Unicast support - move local implementation from keep-tecdsa to keep-core

Author: pdyraga

pkg/net/libp2p/channel.go

@@ -148,10 +148,6 @@ func (c *channel) RegisterUnmarshaler(unmarshaler func() net.TaggedUnmarshaler)
 	c.unmarshalersMutex.Lock()
 	defer c.unmarshalersMutex.Unlock()
 
-	if _, exists := c.unmarshalersByType[tpe]; exists {
-		return fmt.Errorf("type %s already has an associated unmarshaler", tpe)
-	}
-
 	c.unmarshalersByType[tpe] = unmarshaler
 	return nil
 }

pkg/net/libp2p/libp2p.go

@@ -2,6 +2,7 @@ package libp2p
 
 import (
 	"context"
+	"crypto/ecdsa"
 	"fmt"
 	"sync"
 	"time"
@@ -133,6 +134,14 @@ func (p *provider) ConnectionManager() net.ConnectionManager {
 	return p.connectionManager
 }
 
+func (p *provider) CreateTransportIdentifier(publicKey ecdsa.PublicKey) (
+	net.TransportIdentifier,
+	error,
+) {
+	networkPublicKey := key.NetworkPublic(publicKey)
+	return peer.IDFromPublicKey(&networkPublicKey)
+}
+
 type connectionManager struct {
 	host.Host
 }

pkg/net/local/broadcast_channel.go

@@ -0,0 +1,160 @@
+package local
+
+import (
+	"context"
+	"fmt"
+	"sync"
+	"sync/atomic"
+
+	"github.com/keep-network/keep-core/pkg/net"
+	"github.com/keep-network/keep-core/pkg/net/internal"
+	"github.com/keep-network/keep-core/pkg/net/key"
+	"github.com/keep-network/keep-core/pkg/net/retransmission"
+)
+
+const messageHandlerThrottle = 256
+
+type messageHandler struct {
+	ctx     context.Context
+	channel chan net.Message
+}
+
+type localChannel struct {
+	counter              uint64
+	name                 string
+	identifier           net.TransportIdentifier
+	staticKey            *key.NetworkPublic
+	messageHandlersMutex sync.Mutex
+	messageHandlers      []*messageHandler
+	unmarshalersMutex    sync.Mutex
+	unmarshalersByType   map[string]func() net.TaggedUnmarshaler
+	retransmissionTicker *retransmission.Ticker
+}
+
+func (lc *localChannel) nextSeqno() uint64 {
+	return atomic.AddUint64(&lc.counter, 1)
+}
+
+func (lc *localChannel) Name() string {
+	return lc.name
+}
+
+func (lc *localChannel) Send(ctx context.Context, message net.TaggedMarshaler) error {
+	bytes, err := message.Marshal()
+	if err != nil {
+		return err
+	}
+
+	unmarshaler, found := lc.unmarshalersByType[string(message.Type())]
+	if !found {
+		return fmt.Errorf("couldn't find unmarshaler for type %s", string(message.Type()))
+	}
+
+	unmarshaled := unmarshaler()
+	err = unmarshaled.Unmarshal(bytes)
+	if err != nil {
+		return err
+	}
+
+	netMessage := internal.BasicMessage(
+		lc.identifier,
+		unmarshaled,
+		"local",
+		key.Marshal(lc.staticKey),
+		lc.nextSeqno(),
+	)
+
+	retransmission.ScheduleRetransmissions(
+		ctx,
+		lc.retransmissionTicker,
+		func() error {
+			return broadcastMessage(lc.name, netMessage)
+		},
+	)
+
+	return broadcastMessage(lc.name, netMessage)
+}
+
+func (lc *localChannel) deliver(message net.Message) {
+	lc.messageHandlersMutex.Lock()
+	snapshot := make([]*messageHandler, len(lc.messageHandlers))
+	copy(snapshot, lc.messageHandlers)
+	lc.messageHandlersMutex.Unlock()
+
+	for _, handler := range snapshot {
+		select {
+		case handler.channel <- message:
+		default:
+			logger.Warningf("handler too slow, dropping message")
+		}
+	}
+}
+
+func (lc *localChannel) Recv(ctx context.Context, handler func(m net.Message)) {
+	messageHandler := &messageHandler{
+		ctx:     ctx,
+		channel: make(chan net.Message, messageHandlerThrottle),
+	}
+
+	lc.messageHandlersMutex.Lock()
+	lc.messageHandlers = append(lc.messageHandlers, messageHandler)
+	lc.messageHandlersMutex.Unlock()
+
+	handleWithRetransmissions := retransmission.WithRetransmissionSupport(handler)
+
+	go func() {
+		for {
+			select {
+			case <-ctx.Done():
+				logger.Debug("context is done, removing handler")
+				lc.removeHandler(messageHandler)
+				return
+
+			case msg := <-messageHandler.channel:
+				// Go language specification says that if one or more of the
+				// communications in the select statement can proceed, a single
+				// one that will proceed is chosen via a uniform pseudo-random
+				// selection.
+				// Thus, it can happen this communication is called when ctx is
+				// already done. Since we guarantee in the network channel API
+				// that handler is not called after ctx is done (client code
+				// could e.g. perform come cleanup), we need to double-check
+				// the context state here.
+				if messageHandler.ctx.Err() != nil {
+					continue
+				}
+
+				handleWithRetransmissions(msg)
+			}
+		}
+	}()
+}
+
+func (lc *localChannel) removeHandler(handler *messageHandler) {
+	lc.messageHandlersMutex.Lock()
+	defer lc.messageHandlersMutex.Unlock()
+
+	for i, h := range lc.messageHandlers {
+		if h.channel == handler.channel {
+			lc.messageHandlers[i] = lc.messageHandlers[len(lc.messageHandlers)-1]
+			lc.messageHandlers = lc.messageHandlers[:len(lc.messageHandlers)-1]
+			break
+		}
+	}
+}
+
+func (lc *localChannel) RegisterUnmarshaler(
+	unmarshaler func() net.TaggedUnmarshaler,
+) (err error) {
+	tpe := unmarshaler().Type()
+
+	lc.unmarshalersMutex.Lock()
+	defer lc.unmarshalersMutex.Unlock()
+
+	lc.unmarshalersByType[tpe] = unmarshaler
+	return nil
+}
+
+func (lc *localChannel) SetFilter(filter net.BroadcastChannelFilter) error {
+	return nil // no-op
+}

pkg/net/local/broadcast_channel_manager.go

@@ -0,0 +1,62 @@
+package local
+
+import (
+	"context"
+	"sync"
+	"time"
+
+	"github.com/keep-network/keep-core/pkg/net"
+	"github.com/keep-network/keep-core/pkg/net/key"
+	"github.com/keep-network/keep-core/pkg/net/retransmission"
+)
+
+var broadcastChannelsMutex sync.Mutex
+var broadcastChannels map[string][]*localChannel
+
+// getBroadcastChannel returns a BroadcastChannel designed to mediate between local
+// participants. It delivers all messages sent to the channel through its
+// receive channels. RecvChan on a LocalChannel creates a new receive channel
+// that is returned to the caller, so that all receive channels can receive
+// the message.
+func getBroadcastChannel(name string, staticKey *key.NetworkPublic) net.BroadcastChannel {
+	broadcastChannelsMutex.Lock()
+	defer broadcastChannelsMutex.Unlock()
+	if broadcastChannels == nil {
+		broadcastChannels = make(map[string][]*localChannel)
+	}
+
+	localChannels, exists := broadcastChannels[name]
+	if !exists {
+		localChannels = make([]*localChannel, 0)
+		broadcastChannels[name] = localChannels
+	}
+
+	identifier := randomLocalIdentifier()
+	channel := &localChannel{
+		name:                 name,
+		identifier:           &identifier,
+		staticKey:            staticKey,
+		messageHandlersMutex: sync.Mutex{},
+		messageHandlers:      make([]*messageHandler, 0),
+		unmarshalersMutex:    sync.Mutex{},
+		unmarshalersByType:   make(map[string]func() net.TaggedUnmarshaler, 0),
+		retransmissionTicker: retransmission.NewTimeTicker(
+			context.Background(), 50*time.Millisecond,
+		),
+	}
+	broadcastChannels[name] = append(broadcastChannels[name], channel)
+
+	return channel
+}
+
+func broadcastMessage(name string, message net.Message) error {
+	broadcastChannelsMutex.Lock()
+	targetChannels := broadcastChannels[name]
+	broadcastChannelsMutex.Unlock()
+
+	for _, targetChannel := range targetChannels {
+		targetChannel.deliver(message)
+	}
+
+	return nil
+}

pkg/net/local/local_test.go pkg/net/local/broadcast_channel_test.gopkg/net/local/local_test.go renamed to pkg/net/local/broadcast_channel_test.go

@@ -0,0 +1,32 @@
+package local
+
+import (
+	"encoding/hex"
+	"math/rand"
+
+	"github.com/keep-network/keep-core/pkg/net/key"
+)
+
+var letterRunes = [52]rune{'a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j',
+	'k', 'l', 'm', 'n', 'o', 'p', 'q', 'r', 's', 't', 'u', 'v', 'w', 'x', 'y',
+	'z', 'A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J', 'K', 'L', 'M', 'N',
+	'O', 'P', 'Q', 'R', 'S', 'T', 'U', 'V', 'W', 'X', 'Y', 'Z'}
+
+type localIdentifier string
+
+func (li localIdentifier) String() string {
+	return string(li)
+}
+
+func randomLocalIdentifier() localIdentifier {
+	runes := make([]rune, 32)
+	for i := range runes {
+		runes[i] = letterRunes[rand.Intn(len(letterRunes))]
+	}
+
+	return localIdentifier(runes)
+}
+
+func createLocalIdentifier(staticKey *key.NetworkPublic) localIdentifier {
+	return localIdentifier(hex.EncodeToString(key.Marshal(staticKey)))
+}