Repository: tomarrell/lbadd
Pull request #112 (2020-05-05 14:39:40)

Add a network layer with tcp implementation
Added a network layer, which is just an abstraction over a `net.Listener` and `net.Conn`s. However, this abstraction adds framing to messages, meaning that it writes the length of a message before sending the actual message, and also reading a number before reading the message. This makes it very different from, but easier to use than an `io.Reader` or `io.Writer`.



Patch:
diff --git a/go.mod b/go.mod
index 2508086..a9c0edd 100644
--- a/go.mod
+++ b/go.mod
@@ -4,13 +4,15 @@ go 1.13
 
 require (
 	github.com/awnumar/memguard v0.22.1
-	github.com/dvyukov/go-fuzz v0.0.0-20200318091601-be3528f3a813 // indirect
 	github.com/google/go-cmp v0.4.0
 	github.com/kr/pretty v0.2.0 // indirect
+	github.com/oklog/ulid v1.3.1
 	github.com/rs/zerolog v1.18.0
 	github.com/spf13/afero v1.1.2
 	github.com/spf13/cobra v0.0.6
 	github.com/stretchr/testify v1.4.0
+	golang.org/x/net v0.0.0-20200226121028-0de0cce0169b
+	golang.org/x/sync v0.0.0-20190911185100-cd5d95a43a6e
 	golang.org/x/text v0.3.2
 	golang.org/x/tools v0.0.0-20200328031815-3db5fc6bac03
 	gopkg.in/check.v1 v1.0.0-20190902080502-41f04d3bba15 // indirect
diff --git a/go.sum b/go.sum
index 97914cc..ef992df 100644
--- a/go.sum
+++ b/go.sum
@@ -23,8 +23,6 @@ github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/dgrijalva/jwt-go v3.2.0+incompatible/go.mod h1:E3ru+11k8xSBh+hMPgOLZmtrrCbhqsmaPHjLKYnJCaQ=
 github.com/dgryski/go-sip13 v0.0.0-20181026042036-e10d5fee7954/go.mod h1:vAd38F8PWV+bWy6jNmig1y/TA+kYO4g3RSRF0IAv0no=
-github.com/dvyukov/go-fuzz v0.0.0-20200318091601-be3528f3a813 h1:NgO45/5mBLRVfiXerEFzH6ikcZ7DNRPS639xFg3ENzU=
-github.com/dvyukov/go-fuzz v0.0.0-20200318091601-be3528f3a813/go.mod h1:11Gm+ccJnvAhCNLlf5+cS9KjtbaD5I5zaZpFMsTHWTw=
 github.com/fsnotify/fsnotify v1.4.7 h1:IXs+QLmnXW2CcXuY+8Mzv/fWEsPGWxqefPtCP5CnV9I=
 github.com/fsnotify/fsnotify v1.4.7/go.mod h1:jwhsz4b93w/PPRr/qN1Yymfu8t87LnFCMoQvtojpjFo=
 github.com/ghodss/yaml v1.0.0/go.mod h1:4dBDuWmgqj2HViK6kFavaiC9ZROes6MMH2rRYeMEF04=
@@ -71,6 +69,7 @@ github.com/mitchellh/go-homedir v1.1.0/go.mod h1:SfyaCUpYCn1Vlf4IUYiD9fPX4A5wJrk
 github.com/mitchellh/mapstructure v1.1.2 h1:fmNYVwqnSfB9mZU6OS2O6GsXM+wcskZDuKQzvN1EDeE=
 github.com/mitchellh/mapstructure v1.1.2/go.mod h1:FVVH3fgwuzCH5S8UJGiWEs2h04kUh9fWfEaFds41c1Y=
 github.com/mwitkow/go-conntrack v0.0.0-20161129095857-cc309e4a2223/go.mod h1:qRWi+5nqEBWmkhHvq77mSJWrCKwh8bxhgT7d/eI7P4U=
+github.com/oklog/ulid v1.3.1 h1:EGfNDEx6MqHz8B3uNV6QAib1UR2Lm97sHi3ocA6ESJ4=
 github.com/oklog/ulid v1.3.1/go.mod h1:CirwcVhetQ6Lv90oh/F+FBtV6XMibvdAFo93nm5qn4U=
 github.com/pelletier/go-toml v1.2.0 h1:T5zMGML61Wp+FlcbWjRDT7yAxhJNAiPPLOFECq181zc=
 github.com/pelletier/go-toml v1.2.0/go.mod h1:5z9KED0ma1S8pY6P1sdut58dfprrGBbd/94hg7ilaic=
@@ -100,6 +99,8 @@ github.com/soheilhy/cmux v0.1.4/go.mod h1:IM3LyeVVIOuxMH7sFAkER9+bJ4dT7Ms6E4xg4k
 github.com/spaolacci/murmur3 v0.0.0-20180118202830-f09979ecbc72/go.mod h1:JwIasOWyU6f++ZhiEuf87xNszmSA2myDM2Kzu9HwQUA=
 github.com/spf13/afero v1.1.2 h1:m8/z1t7/fwjysjQRYbP0RD+bUIF/8tJwPdEZsI83ACI=
 github.com/spf13/afero v1.1.2/go.mod h1:j4pytiNVoe2o6bmDsKpLACNPDBIoEAkihy7loJ1B0CQ=
+github.com/spf13/afero v1.2.2 h1:5jhuqJyZCZf2JRofRvN/nIFgIWNzPa3/Vz8mYylgbWc=
+github.com/spf13/afero v1.2.2/go.mod h1:9ZxEEn6pIJ8Rxe320qSDBk6AsU0r9pR7Q4OcevTdifk=
 github.com/spf13/cast v1.3.0 h1:oget//CVOEoFewqQxwr0Ej5yjygnqGkvggSE/gB35Q8=
 github.com/spf13/cast v1.3.0/go.mod h1:Qx5cxh0v+4UWYiBimWS+eyWzqEqokIECu5etghLkUJE=
 github.com/spf13/cobra v0.0.6 h1:breEStsVwemnKh2/s6gMvSdMEkwW0sK8vGStnlVBMCs=
@@ -144,12 +145,14 @@ golang.org/x/net v0.0.0-20190311183353-d8887717615a/go.mod h1:t9HGtf8HONx5eT2rtn
 golang.org/x/net v0.0.0-20190404232315-eb5bcb51f2a3/go.mod h1:t9HGtf8HONx5eT2rtn7q6eTqICYqUVnKs3thJo3Qplg=
 golang.org/x/net v0.0.0-20190522155817-f3200d17e092/go.mod h1:HSz+uSET+XFnRR8LxR5pz3Of3rY3CfYBVs4xY44aLks=
 golang.org/x/net v0.0.0-20190620200207-3b0461eec859/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s=
+golang.org/x/net v0.0.0-20200226121028-0de0cce0169b h1:0mm1VjtFUOIlE1SbDlwjYaDxZVDP2S5ou6y0gSgXHu8=
 golang.org/x/net v0.0.0-20200226121028-0de0cce0169b/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s=
 golang.org/x/oauth2 v0.0.0-20180821212333-d2e6202438be/go.mod h1:N/0e6XlmueqKjAGxoOufVs8QHGRruUQn6yWY3a++T0U=
 golang.org/x/sync v0.0.0-20180314180146-1d60e4601c6f/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sync v0.0.0-20181108010431-42b317875d0f/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sync v0.0.0-20181221193216-37e7f081c4d4/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
+golang.org/x/sync v0.0.0-20190911185100-cd5d95a43a6e h1:vcxGaoTs7kV8m5Np9uUNQin4BrLOthgV7252N8V+FwY=
 golang.org/x/sync v0.0.0-20190911185100-cd5d95a43a6e/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sys v0.0.0-20180830151530-49385e6e1522/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
 golang.org/x/sys v0.0.0-20180905080454-ebe1bf3edb33/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
diff --git a/internal/network/doc.go b/internal/network/doc.go
new file mode 100644
index 0000000..1c72fcc
--- /dev/null
+++ b/internal/network/doc.go
@@ -0,0 +1,3 @@
+// Package network implements a communication layer with a server and a client.
+// Clients can connect to servers and receive and send messages.
+package network
diff --git a/internal/network/errors.go b/internal/network/errors.go
new file mode 100644
index 0000000..789b3fb
--- /dev/null
+++ b/internal/network/errors.go
@@ -0,0 +1,15 @@
+package network
+
+// Error is a helper type for creating constant errors.
+type Error string
+
+func (e Error) Error() string { return string(e) }
+
+const (
+	// ErrOpen indicates, that the component was already opened, and it is
+	// unable to be opened another time.
+	ErrOpen Error = "already open"
+	// ErrClosed indicates, that the component is already closed, and it cannot
+	// be used anymore.
+	ErrClosed Error = "already closed"
+)
diff --git a/internal/network/id.go b/internal/network/id.go
new file mode 100644
index 0000000..fc227f5
--- /dev/null
+++ b/internal/network/id.go
@@ -0,0 +1,42 @@
+package network
+
+import (
+	"fmt"
+	"log"
+	"math/rand"
+	"sync"
+	"time"
+
+	"github.com/oklog/ulid"
+)
+
+var _ ID = (*id)(nil)
+
+type id ulid.ULID
+
+var (
+	lock       sync.Mutex
+	randSource = rand.New(rand.NewSource(time.Now().UnixNano()))
+	entropy    = ulid.Monotonic(randSource, 0)
+)
+
+func createID() ID {
+	lock.Lock()
+	defer lock.Unlock()
+
+	id, err := ulid.New(ulid.Now(), entropy)
+	if err != nil {
+		// For this to happen, the random module would have to fail. Since we
+		// use Go's pseudo RNG, which just jumps around a few numbers, instead
+		// of using crypto/rand, and we also made this function safe for
+		// concurrent use, this is nearly impossible to happen. However, with
+		// the current version of oklog/ulid v1.3.1, this will also break after
+		// 2121-04-11 11:53:25.01172576 UTC.
+		log.Fatal(fmt.Errorf("new ulid: %w", err))
+	}
+	return ID(id)
+}
+
+func (id id) String() string {
+	return ulid.ULID(id).String()
+}
diff --git a/internal/network/server.go b/internal/network/server.go
new file mode 100644
index 0000000..d957d05
--- /dev/null
+++ b/internal/network/server.go
@@ -0,0 +1,57 @@
+package network
+
+import (
+	"fmt"
+	"io"
+	"net"
+)
+
+// ConnHandler is a handler function for handling new connections. It will be
+// called with a fully initialized Conn, and is used as a callback in the
+// server.
+type ConnHandler func(Conn)
+
+// Server describes a server component, that listens for connecting clients.
+// Before opening, it is recommended that one sets a connect handler with
+// Server.OnConnect. A server can only be opened once. Closing a server must not
+// close the accepted connections, but must only stop accepting new connections
+// and release the allocated address.
+type Server interface {
+	io.Closer
+
+	// Open opens the server on the given address. To choose the server a random
+	// free port for you, specify a port ":0".
+	Open(string) error
+	// Addr returns the address that this server is listening to.
+	Addr() net.Addr
+
+	// OnConnect sets a callback that will be executed whenever a new connection
+	// connects to this server.
+	OnConnect(ConnHandler)
+}
+
+// Conn describes a network connection. One can send a message with Conn.Send,
+// and receive one with Conn.Receive. Unlike an io.Writer, the data that is
+// passed into Send is guaranteed to be returned in a single Receive call on the
+// other end, meaning that you don't have to worry about where your messages
+// end. Maximum message length is 2GiB.
+type Conn interface {
+	io.Closer
+
+	// ID returns the ID of this connection. It can be used to uniquely identify
+	// this connection globally.
+	ID() ID
+	// Send sends the given payload to the remote part of this connection. The
+	// message will not be chunked, and can be read with a single call to
+	// Conn.Receive.
+	Send([]byte) error
+	// Receive reads a whole message and returns it in a byte slice. A message
+	// is a byte slice that was sent with a single call to Conn.Send.
+	Receive() ([]byte, error)
+}
+
+// ID describes an identifier that is used for connections. An ID has to be
+// unique application-wide. IDs must not be re-used.
+type ID interface {
+	fmt.Stringer
+}
diff --git a/internal/network/tcp_conn.go b/internal/network/tcp_conn.go
new file mode 100644
index 0000000..5aa2c5b
--- /dev/null
+++ b/internal/network/tcp_conn.go
@@ -0,0 +1,107 @@
+package network
+
+import (
+	"encoding/binary"
+	"fmt"
+	"net"
+
+	"golang.org/x/net/context"
+	"golang.org/x/sync/errgroup"
+)
+
+const (
+	frameSizeBytes int = 4
+)
+
+var (
+	byteOrder = binary.BigEndian
+)
+
+var _ Conn = (*tcpConn)(nil)
+
+type tcpConn struct {
+	id         ID
+	underlying net.Conn
+	closed     bool
+}
+
+// DialTCP dials to the given address, assuming a TCP network. The returned Conn
+// is ready to use.
+func DialTCP(addr string) (Conn, error) {
+	conn, err := net.Dial("tcp", addr)
+	if err != nil {
+		return nil, fmt.Errorf("dial tcp: %w", err)
+	}
+	return newTCPConn(conn), nil
+}
+
+func newTCPConn(underlying net.Conn) *tcpConn {
+	id := createID()
+	conn := &tcpConn{
+		id:         id,
+		underlying: underlying,
+	}
+	return conn
+}
+
+func (c *tcpConn) ID() ID {
+	return c.id
+}
+
+func (c *tcpConn) Send(payload []byte) error {
+	if c.closed {
+		return ErrClosed
+	}
+
+	var frameSize [frameSizeBytes]byte
+	byteOrder.PutUint32(frameSize[:], uint32(len(payload)))
+
+	n, err := c.underlying.Write(frameSize[:])
+	if err != nil {
+		return fmt.Errorf("write size: %w", err)
+	}
+	if n != frameSizeBytes {
+		return fmt.Errorf("write bytes: written %v of %v size bytes", n, len(payload))
+	}
+
+	n, err = c.underlying.Write(payload)
+	if err != nil {
+		return fmt.Errorf("write payload: %w", err)
+	}
+	if n != len(payload) {
+		return fmt.Errorf("write bytes: written %v of %v payload bytes", n, len(payload))
+	}
+	return nil
+}
+
+func (c *tcpConn) Receive() ([]byte, error) {
+	var frameSizeB [frameSizeBytes]byte
+	n, err := c.underlying.Read(frameSizeB[:])
+	if err != nil {
+		return nil, fmt.Errorf("read frame size: %w", err)
+	}
+	if n != frameSizeBytes {
+		return nil, fmt.Errorf("read only %v frame size bytes of %v expected", n, frameSizeBytes)
+	}
+
+	frameSize := byteOrder.Uint32(frameSizeB[:])
+	frameData := make([]byte, frameSize)
+	n, err = c.underlying.Read(frameData)
+	if err != nil {
+		return nil, fmt.Errorf("read frame payload: %w", err)
+	}
+	if n != int(frameSize) {
+		return nil, fmt.Errorf("read only %v frame payload bytes of %v expected", n, frameSize)
+	}
+	return frameData, nil
+}
+
+func (c *tcpConn) Close() error {
+	c.closed = true
+
+	// release all resources
+	ctx := context.Background()
+	errs, _ := errgroup.WithContext(ctx)
+	errs.Go(c.underlying.Close)
+	return errs.Wait()
+}
diff --git a/internal/network/tcp_server.go b/internal/network/tcp_server.go
new file mode 100644
index 0000000..e23ee83
--- /dev/null
+++ b/internal/network/tcp_server.go
@@ -0,0 +1,97 @@
+package network
+
+import (
+	"fmt"
+	"net"
+
+	"github.com/rs/zerolog"
+	"golang.org/x/net/context"
+	"golang.org/x/sync/errgroup"
+)
+
+var _ Server = (*tcpServer)(nil)
+
+type tcpServer struct {
+	log zerolog.Logger
+
+	open bool
+	lis  net.Listener
+
+	onConnect ConnHandler
+}
+
+// NewTCPServer creates a new ready to use TCP server that uses the given
+// logger.
+func NewTCPServer(log zerolog.Logger) Server {
+	return &tcpServer{
+		log: log,
+	}
+}
+
+func (s *tcpServer) Open(addr string) error {
+	if s.open {
+		return ErrOpen
+	}
+
+	lis, err := net.Listen("tcp", addr)
+	if err != nil {
+		return fmt.Errorf("listen: %w", err)
+	}
+
+	s.open = true
+	s.lis = lis
+
+	s.log.Debug().
+		Str("addr", lis.Addr().String()).
+		Msg("tcp open")
+
+	s.handleIncomingConnections()
+	return nil
+}
+
+func (s *tcpServer) Addr() net.Addr {
+	if s.lis == nil {
+		return nil
+	}
+	return s.lis.Addr()
+}
+
+func (s *tcpServer) OnConnect(h ConnHandler) {
+	s.onConnect = h
+}
+
+func (s *tcpServer) Close() error {
+	s.open = false
+
+	// release all resources
+	ctx := context.Background()
+	errs, _ := errgroup.WithContext(ctx)
+	errs.Go(s.lis.Close)
+	return errs.Wait()
+}
+
+func (s *tcpServer) handleIncomingConnections() {
+	for {
+		conn, err := s.lis.Accept()
+		if err != nil {
+			if !s.open {
+				// server was already closed, we can discard the error, but we
+				// also need to stop accepting further connections
+				break
+			}
+
+			// otherwise, an error occurred while accepting a connection, but
+			// since we can't that let us stop from accepting further
+			// connections, we just log it
+			s.log.Error().
+				Err(err).
+				Msg("accept")
+		}
+
+		tcpConn := newTCPConn(conn)
+
+		if s.onConnect != nil {
+			go s.onConnect(tcpConn)
+		}
+	}
+}


Test patch:
diff --git a/internal/network/id_test.go b/internal/network/id_test.go
new file mode 100644
index 0000000..8510026
--- /dev/null
+++ b/internal/network/id_test.go
@@ -0,0 +1,13 @@
+package network
+
+import "testing"
+
+func TestIDThreadSafe(t *testing.T) {
+	// This is sufficient for the race detector to detect a race if createID is
+	// not safe for concurrent use.
+	for i := 0; i < 5; i++ {
+		go func() {
+			_ = createID()
+		}()
+	}
+}
diff --git a/internal/network/tcp_conn_test.go b/internal/network/tcp_conn_test.go
new file mode 100644
index 0000000..2c97b8d
--- /dev/null
+++ b/internal/network/tcp_conn_test.go
@@ -0,0 +1,59 @@
+package network
+
+import (
+	"net"
+	"strconv"
+	"sync"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestTCPConnSendReceive(t *testing.T) {
+	assert := assert.New(t)
+	conn1, conn2 := net.Pipe()
+	tcpConn1, tcpConn2 := newTCPConn(conn1), newTCPConn(conn2)
+
+	payload := []byte("Hello, World!")
+	recv := make([]byte, len(payload))
+	wg := &sync.WaitGroup{}
+	wg.Add(1)
+	go func() {
+		var err error
+		recv, err = tcpConn2.Receive()
+		assert.NoError(err)
+		wg.Done()
+	}()
+
+	err := tcpConn1.Send(payload)
+	assert.NoError(err)
+
+	wg.Wait()
+	assert.Equal(payload, recv)
+}
+
+func TestDialTCP(t *testing.T) {
+	assert := assert.New(t)
+	payload := []byte("Hello, World!")
+
+	lis, err := net.Listen("tcp", ":0")
+	assert.NoError(err)
+
+	go func() {
+		conn, err := lis.Accept()
+		assert.NoError(err)
+
+		tcpConn := newTCPConn(conn)
+		assert.NoError(tcpConn.Send(payload))
+	}()
+
+	port := lis.Addr().(*net.TCPAddr).Port
+
+	conn, err := DialTCP(":" + strconv.Itoa(port))
+	assert.NoError(err)
+	defer func() { assert.NoError(conn.Close()) }()
+
+	recv, err := conn.Receive()
+	assert.NoError(err)
+	assert.Equal(payload, recv)
+}