update utp encoding

Author: unit-adm

batch.go

@@ -5,7 +5,7 @@ import (
 	"sync"
 	"time"
 
-	"github.com/unit-io/unitdb-go/internal/utp"
+	"github.com/unit-io/unitdb/server/utp"
 )
 
 const (

client.go

@@ -11,10 +11,11 @@ import (
 	"time"
 
 	"github.com/golang/protobuf/proto"
+	lp "github.com/unit-io/unitdb-go/internal/net"
 	"github.com/unit-io/unitdb-go/internal/store"
-	"github.com/unit-io/unitdb-go/internal/utp"
 	"github.com/unit-io/unitdb/server/common"
 	pbx "github.com/unit-io/unitdb/server/proto"
+	"github.com/unit-io/unitdb/server/utp"
 	"google.golang.org/grpc"
 
 	// Database store
@@ -58,7 +59,7 @@ type client struct {
 	epoch      uint32   // The session ID of the connection.
 	conn       net.Conn // the network connection
 	send       chan *MessageAndResult
-	recv       chan utp.Message
+	recv       chan lp.MessagePack
 	pub        chan *utp.Publish
 	callbacks  map[uint64]MessageHandler
 
@@ -84,7 +85,7 @@ func NewClient(target, clientID string, opts ...Options) (Client, error) {
 		cancel:     cancel,
 		messageIds: messageIds{index: make(map[MID]Result)},
 		send:       make(chan *MessageAndResult, 1), // buffered
-		recv:       make(chan utp.Message),
+		recv:       make(chan lp.MessagePack),
 		pub:        make(chan *utp.Publish),
 		callbacks:  make(map[uint64]MessageHandler),
 		// close
@@ -543,12 +544,12 @@ func TimeNow() time.Time {
 	return time.Now().UTC().Round(time.Millisecond)
 }
 
-func (c *client) inboundID(id int32) MID {
-	return MID(c.connID - id)
+func (c *client) inboundID(id uint16) MID {
+	return MID(c.connID - int32(id))
 }
 
-func (c *client) outboundID(mid MID) (id int32) {
-	return c.connID - (int32(mid))
+func (c *client) outboundID(mid MID) (id uint16) {
+	return uint16(c.connID - (int32(mid)))
 }
 
 func (c *client) updateLastAction() {
@@ -561,11 +562,11 @@ func (c *client) updateLastTouched() {
 	c.lastTouched.Store(TimeNow())
 }
 
-func (c *client) storeInbound(m utp.Message) {
+func (c *client) storeInbound(m lp.MessagePack) {
 	store.Log.PersistInbound(uint32(c.sessID), m)
 }
 
-func (c *client) storeOutbound(m utp.Message) {
+func (c *client) storeOutbound(m lp.MessagePack) {
 	store.Log.PersistOutbound(uint32(c.sessID), m)
 }
 

hdl_conn.go

@@ -9,14 +9,15 @@ import (
 	"net"
 	"time"
 
-	"github.com/unit-io/unitdb-go/internal/utp"
+	lp "github.com/unit-io/unitdb-go/internal/net"
+	"github.com/unit-io/unitdb/server/utp"
 )
 
 // Connect takes a connected net.Conn and performs the initial handshake. Paramaters are:
 // conn - Connected net.Conn
 // cm - Connect Message
-func Connect(conn net.Conn, cm *utp.Connect) (rc int32, epoch int32, cid int32, err error) {
-	m, err := utp.Encode(cm)
+func Connect(conn net.Conn, cm *utp.Connect) (rc uint8, epoch int32, cid int32, err error) {
+	m, err := lp.Encode(cm)
 	if err != nil {
 		fmt.Println(err)
 		return utp.ErrRefusedServerUnavailable, 0, 0, err
@@ -31,21 +32,24 @@ func Connect(conn net.Conn, cm *utp.Connect) (rc int32, epoch int32, cid int32,
 // when the connection is first started.
 // This prevents receiving incoming data while resume
 // is in progress if clean session is false.
-func verifyCONNACK(conn net.Conn) (int32, int32, int32, error) {
-	ca, err := utp.Read(conn)
+func verifyCONNACK(conn net.Conn) (uint8, int32, int32, error) {
+	ca, err := lp.Read(conn)
 	if err != nil {
 		return utp.ErrRefusedServerUnavailable, 0, 0, err
 	}
 	if ca == nil {
-		return utp.ErrRefusedServerUnavailable, 0, 0, errors.New("nil Connect Acknowledge Message")
+		return utp.ErrRefusedServerUnavailable, 0, 0, errors.New("nil connect acknowledge message")
 	}
 
-	msg, ok := ca.(*utp.ConnectAcknowledge)
+	pack, ok := ca.(*utp.ControlMessage)
 	if !ok {
 		return utp.ErrRefusedServerUnavailable, 0, 0, errors.New("first message must be connect acknowledge message")
 	}
 
-	return msg.ReturnCode, msg.Epoch, msg.ConnID, nil
+	connack := &utp.ConnectAcknowledge{}
+	connack.FromBinary(utp.FixedHeader{MessageType: utp.CONNECT, FlowControl: utp.ACKNOWLEDGE}, pack.Message)
+
+	return connack.ReturnCode, connack.Epoch, connack.ConnID, nil
 }
 
 // Handle handles incoming messages
@@ -67,7 +71,7 @@ func (c *client) readLoop(ctx context.Context) error {
 			return nil
 		default:
 			// Unpack an incoming Message
-			msg, err := utp.Read(reader)
+			msg, err := lp.Read(reader)
 			if err != nil {
 				return err
 			}
@@ -84,7 +88,7 @@ func (c *client) readLoop(ctx context.Context) error {
 }
 
 // handle handles inbound messages.
-func (c *client) handler(inMsg utp.Message) error {
+func (c *client) handler(inMsg lp.MessagePack) error {
 	c.updateLastAction()
 
 	switch inMsg.Type() {
@@ -143,7 +147,7 @@ func (c *client) writeLoop(ctx context.Context) {
 				mId := c.inboundID(msg.MessageID)
 				c.freeID(mId)
 			}
-			m, err := utp.Encode(outMsg.m)
+			m, err := lp.Encode(outMsg.m)
 			if err != nil {
 				fmt.Println(err)
 				// return
@@ -183,7 +187,7 @@ func (c *client) dispatcher(ctx context.Context) {
 // keepalive - Send ping when connection unused for set period
 // connection passed in to avoid race condition on shutdown
 func (c *client) keepalive(ctx context.Context) {
-	var pingInterval int64
+	var pingInterval int32
 	var pingSent time.Time
 
 	if c.opts.keepAlive > 10 {
@@ -192,7 +196,7 @@ func (c *client) keepalive(ctx context.Context) {
 		pingInterval = c.opts.keepAlive / 2
 	}
 
-	pingTicker := time.NewTicker(time.Duration(pingInterval * int64(time.Second)))
+	pingTicker := time.NewTicker(time.Duration(pingInterval * int32(time.Second)))
 	defer func() {
 		pingTicker.Stop()
 	}()
@@ -206,12 +210,12 @@ func (c *client) keepalive(ctx context.Context) {
 		case <-pingTicker.C:
 			lastAction := c.lastAction.Load().(time.Time)
 			lastTouched := c.lastTouched.Load().(time.Time)
-			live := TimeNow().Add(-time.Duration(c.opts.keepAlive * int64(time.Second)))
+			live := TimeNow().Add(-time.Duration(c.opts.keepAlive * int32(time.Second)))
 			timeout := TimeNow().Add(-c.opts.pingTimeout)
 
 			if lastAction.After(live) || lastTouched.After(live) {
 				ping := &utp.Pingreq{}
-				m, err := utp.Encode(ping)
+				m, err := lp.Encode(ping)
 				if err != nil {
 					fmt.Println(err)
 				}

internal/net/message.go

@@ -0,0 +1,72 @@
+package net
+
+import (
+	"bytes"
+	"fmt"
+	"io"
+
+	"github.com/unit-io/unitdb/server/utp"
+)
+
+// MessagePack is the interface for all Messages
+type MessagePack interface {
+	Type() utp.MessageType
+	Info() utp.Info
+}
+
+// Read unpacks the packet from the provided reader.
+func Read(r io.Reader) (MessagePack, error) {
+	var fh utp.FixedHeader
+	fh.FromBinary(r)
+
+	// Check for empty packets
+	switch fh.MessageType {
+	case utp.DISCONNECT:
+		return &utp.Disconnect{}, nil
+	}
+
+	rawMsg := make([]byte, fh.MessageLength)
+	_, err := io.ReadFull(r, rawMsg)
+	if err != nil {
+		return nil, err
+	}
+
+	// unpack the body
+	if fh.FlowControl != utp.NONE {
+		pack := &utp.ControlMessage{}
+		pack.FromBinary(fh, rawMsg)
+		return pack, nil
+	}
+	switch fh.MessageType {
+	case utp.PUBLISH:
+		pack := &utp.Publish{}
+		pack.FromBinary(fh, rawMsg)
+		return pack, nil
+	default:
+		return nil, fmt.Errorf("message::Read: Invalid zero-length packet type %d", fh.MessageType)
+	}
+}
+
+// Encode encodes the message into binary data
+func Encode(pack MessagePack) (bytes.Buffer, error) {
+	switch pack.Type() {
+	case utp.PINGREQ:
+		return pack.(*utp.Pingreq).ToBinary()
+	case utp.CONNECT:
+		return pack.(*utp.Connect).ToBinary()
+	case utp.DISCONNECT:
+		return pack.(*utp.Disconnect).ToBinary()
+	case utp.RELAY:
+		return pack.(*utp.Relay).ToBinary()
+	case utp.SUBSCRIBE:
+		return pack.(*utp.Subscribe).ToBinary()
+	case utp.UNSUBSCRIBE:
+		return pack.(*utp.Unsubscribe).ToBinary()
+	case utp.PUBLISH:
+		return pack.(*utp.Publish).ToBinary()
+	case utp.FLOWCONTROL:
+		return pack.(*utp.ControlMessage).ToBinary()
+	default:
+		return bytes.Buffer{}, fmt.Errorf("message::Encode: Invalid zero-length packet type %d", pack.Type())
+	}
+}

internal/store/store.go

@@ -6,7 +6,8 @@ import (
 	"fmt"
 
 	adapter "github.com/unit-io/unitdb-go/internal/db"
-	"github.com/unit-io/unitdb-go/internal/utp"
+	lp "github.com/unit-io/unitdb-go/internal/net"
+	"github.com/unit-io/unitdb/server/utp"
 )
 
 var adp adapter.Adapter
@@ -95,13 +96,13 @@ type MessageLog struct{}
 var Log MessageLog
 
 // handle which outgoing messages are stored
-func (l *MessageLog) PersistOutbound(blockID uint32, outMsg utp.Message) {
+func (l *MessageLog) PersistOutbound(blockID uint32, outMsg lp.MessagePack) {
 	switch outMsg.(type) {
 	case *utp.Publish, *utp.Subscribe, *utp.Unsubscribe:
 		// Received a publish. store it in ibound
 		// until ACKNOWLEDGE or COMPLETE is sent
 		okey := uint64(blockID)<<32 + uint64(outMsg.Info().MessageID)
-		m, err := utp.Encode(outMsg)
+		m, err := lp.Encode(outMsg)
 		if err != nil {
 			fmt.Println(err)
 			return
@@ -115,7 +116,7 @@ func (l *MessageLog) PersistOutbound(blockID uint32, outMsg utp.Message) {
 			// Received a RECEIPT control message. store it in obound
 			// until COMPLETE is received.
 			okey := uint64(blockID)<<32 + uint64(outMsg.Info().MessageID)
-			m, err := utp.Encode(outMsg)
+			m, err := lp.Encode(outMsg)
 			if err != nil {
 				fmt.Println(err)
 				return
@@ -126,13 +127,13 @@ func (l *MessageLog) PersistOutbound(blockID uint32, outMsg utp.Message) {
 }
 
 // handle which incoming messages are stored
-func (l *MessageLog) PersistInbound(blockID uint32, inMsg utp.Message) {
+func (l *MessageLog) PersistInbound(blockID uint32, inMsg lp.MessagePack) {
 	switch inMsg.(type) {
 	case *utp.Publish:
 		// Received a publish. store it in ibound
 		// until COMPLETE sent
 		ikey := uint64(blockID)<<32 + uint64(inMsg.Info().MessageID)
-		m, err := utp.Encode(inMsg)
+		m, err := lp.Encode(inMsg)
 		if err != nil {
 			fmt.Println(err)
 			return
@@ -151,7 +152,7 @@ func (l *MessageLog) PersistInbound(blockID uint32, inMsg utp.Message) {
 			// Sending RECEIPT. store in obound
 			// until COMPLETE received
 			ikey := uint64(blockID)<<32 + uint64(inMsg.Info().MessageID)
-			m, err := utp.Encode(inMsg)
+			m, err := lp.Encode(inMsg)
 			if err != nil {
 				fmt.Println(err)
 				return
@@ -162,10 +163,10 @@ func (l *MessageLog) PersistInbound(blockID uint32, inMsg utp.Message) {
 }
 
 // Get performs a query and attempts to fetch message for the given key
-func (l *MessageLog) Get(key uint64) utp.Message {
+func (l *MessageLog) Get(key uint64) lp.MessagePack {
 	if raw, err := adp.GetMessage(key); raw != nil && err == nil {
 		r := bytes.NewReader(raw)
-		if msg, err := utp.Read(r); err == nil {
+		if msg, err := lp.Read(r); err == nil {
 			return msg
 		}
 	}