Initial commit
This commit is contained in:
Yaser
331
pkg/routing/client.go
Normal file
331
pkg/routing/client.go
Normal file
@@ -0,0 +1,331 @@
|
||||
package routing
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"os"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/gorilla/websocket"
|
||||
msg "github.com/openware/rango/pkg/message"
|
||||
"github.com/openware/rango/pkg/metrics"
|
||||
"github.com/rs/zerolog/log"
|
||||
)
|
||||
|
||||
const (
|
||||
// Time allowed to write a message to the peer.
|
||||
writeWait = 10 * time.Second
|
||||
|
||||
// Time allowed to read the next pong message from the peer.
|
||||
pongWait = 60 * time.Second
|
||||
|
||||
// Send pings to peer with this period. Must be less than pongWait.
|
||||
pingPeriod = (pongWait * 9) / 10
|
||||
|
||||
// Maximum message size allowed from peer.
|
||||
maxMessageSize = 512
|
||||
)
|
||||
|
||||
var (
|
||||
newline = []byte{'\n'}
|
||||
space = []byte{' '}
|
||||
)
|
||||
|
||||
var upgrader = websocket.Upgrader{
|
||||
ReadBufferSize: 1024,
|
||||
WriteBufferSize: 1024,
|
||||
CheckOrigin: checkSameOrigin(os.Getenv("API_CORS_ORIGINS")),
|
||||
}
|
||||
|
||||
var maxBufferedMessages = 256
|
||||
|
||||
type Auth struct {
|
||||
UID string
|
||||
Role string
|
||||
}
|
||||
|
||||
// FIXME: IClient looks very wrong.
|
||||
type IClient interface {
|
||||
Send(string)
|
||||
Close()
|
||||
GetAuth() Auth
|
||||
GetSubscriptions() []string
|
||||
SubscribePublic(string)
|
||||
SubscribePrivate(string)
|
||||
UnsubscribePublic(string)
|
||||
UnsubscribePrivate(string)
|
||||
}
|
||||
|
||||
// Client is a middleman between the websocket connection and the hub.
|
||||
type Client struct {
|
||||
hub *Hub
|
||||
|
||||
// User ID if authorized
|
||||
Auth Auth
|
||||
|
||||
pubSub []string
|
||||
privSub []string
|
||||
|
||||
// The websocket connection.
|
||||
conn *websocket.Conn
|
||||
|
||||
// Buffered channel of outbound messages.
|
||||
send chan []byte
|
||||
}
|
||||
|
||||
func checkSameOrigin(origins string) func(r *http.Request) bool {
|
||||
if origins == "" {
|
||||
return func(r *http.Request) bool {
|
||||
origin := r.Header["Origin"]
|
||||
if len(origin) == 0 {
|
||||
return true
|
||||
}
|
||||
u, err := url.Parse(origin[0])
|
||||
if err != nil {
|
||||
return false
|
||||
}
|
||||
return strings.EqualFold(u.Host, r.Host)
|
||||
}
|
||||
}
|
||||
|
||||
hosts := []string{}
|
||||
|
||||
for _, o := range strings.Split(origins, ",") {
|
||||
o = strings.TrimSpace(o)
|
||||
if strings.HasPrefix(o, "http://") || strings.HasPrefix(o, "https://") {
|
||||
u, err := url.Parse(o)
|
||||
if err != nil || u.Host == "" {
|
||||
panic("Failed to parse url in API_CORS_ORIGINS: " + o)
|
||||
}
|
||||
hosts = append(hosts, u.Host)
|
||||
} else {
|
||||
hosts = append(hosts, o)
|
||||
}
|
||||
}
|
||||
|
||||
return func(r *http.Request) bool {
|
||||
origin := r.Header["Origin"]
|
||||
if len(origin) == 0 {
|
||||
return true
|
||||
}
|
||||
u, err := url.Parse(origin[0])
|
||||
if err != nil {
|
||||
return false
|
||||
}
|
||||
|
||||
for _, host := range hosts {
|
||||
if strings.EqualFold(u.Host, host) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
// NewClient handles websocket requests from the peer.
|
||||
func NewClient(hub *Hub, w http.ResponseWriter, r *http.Request) {
|
||||
conn, err := upgrader.Upgrade(w, r, nil)
|
||||
if err != nil {
|
||||
log.Error().Msg("Websocket upgrade failed: " + err.Error())
|
||||
return
|
||||
}
|
||||
client := &Client{
|
||||
hub: hub,
|
||||
conn: conn,
|
||||
send: make(chan []byte, maxBufferedMessages),
|
||||
Auth: Auth{
|
||||
UID: r.Header.Get("JwtUID"),
|
||||
Role: r.Header.Get("JwtRole"),
|
||||
},
|
||||
pubSub: []string{},
|
||||
privSub: []string{},
|
||||
}
|
||||
|
||||
if client.Auth.UID == "" {
|
||||
log.Info().Msgf("New anonymous connection")
|
||||
} else {
|
||||
log.Info().Msgf("New authenticated connection: %s", client.Auth.UID)
|
||||
}
|
||||
|
||||
hub.handleSubscribe(&Request{
|
||||
client: client,
|
||||
Request: msg.Request{
|
||||
Streams: parseStreamsFromURI(r.RequestURI),
|
||||
},
|
||||
})
|
||||
|
||||
metrics.RecordHubClientNew()
|
||||
|
||||
// Allow collection of memory referenced by the caller by doing all work in
|
||||
// new goroutines.
|
||||
go client.write()
|
||||
go client.read()
|
||||
}
|
||||
|
||||
func (c *Client) Send(s string) {
|
||||
if len(c.send) == maxBufferedMessages {
|
||||
log.Warn().Msg("Closing slow websocket connection")
|
||||
c.conn.Close()
|
||||
} else {
|
||||
c.send <- []byte(s)
|
||||
}
|
||||
}
|
||||
|
||||
func (c *Client) Close() {
|
||||
close(c.send)
|
||||
}
|
||||
|
||||
func (c *Client) GetAuth() Auth {
|
||||
return c.Auth
|
||||
}
|
||||
|
||||
func (c *Client) GetSubscriptions() []string {
|
||||
return append(c.pubSub, c.privSub...)
|
||||
}
|
||||
|
||||
func (c *Client) SubscribePublic(s string) {
|
||||
if !contains(c.pubSub, s) {
|
||||
c.pubSub = append(c.pubSub, s)
|
||||
}
|
||||
}
|
||||
|
||||
func (c *Client) SubscribePrivate(s string) {
|
||||
if !contains(c.privSub, s) {
|
||||
c.privSub = append(c.privSub, s)
|
||||
}
|
||||
}
|
||||
|
||||
func (c *Client) UnsubscribePublic(s string) {
|
||||
l := make([]string, len(c.pubSub)-1)
|
||||
i := 0
|
||||
for _, el := range c.pubSub {
|
||||
if s != el {
|
||||
l[i] = el
|
||||
i++
|
||||
}
|
||||
}
|
||||
c.pubSub = l
|
||||
}
|
||||
|
||||
func (c *Client) UnsubscribePrivate(s string) {
|
||||
l := make([]string, len(c.privSub)-1)
|
||||
i := 0
|
||||
for _, el := range c.privSub {
|
||||
if s != el {
|
||||
l[i] = el
|
||||
i++
|
||||
}
|
||||
}
|
||||
c.privSub = l
|
||||
}
|
||||
|
||||
func parseStreamsFromURI(uri string) []string {
|
||||
streams := make([]string, 0)
|
||||
path := strings.Split(uri, "?")
|
||||
if len(path) != 2 {
|
||||
return streams
|
||||
}
|
||||
for _, up := range strings.Split(path[1], "&") {
|
||||
p := strings.Split(up, "=")
|
||||
if len(p) != 2 || p[0] != "stream" {
|
||||
continue
|
||||
}
|
||||
streams = append(streams, strings.Split(p[1], ",")...)
|
||||
|
||||
}
|
||||
return streams
|
||||
}
|
||||
|
||||
// read pumps messages from the websocket connection to the hub.
|
||||
//
|
||||
// The application runs read in a per-connection goroutine. The application
|
||||
// ensures that there is at most one reader on a connection by executing all
|
||||
// reads from this goroutine.
|
||||
func (c *Client) read() {
|
||||
defer func() {
|
||||
log.Debug().Msgf("Closing client read (%s)", c.GetAuth().UID)
|
||||
c.hub.Unregister <- c
|
||||
metrics.RecordHubClientClose()
|
||||
c.conn.Close()
|
||||
}()
|
||||
|
||||
c.conn.SetReadLimit(maxMessageSize)
|
||||
c.conn.SetReadDeadline(time.Now().Add(pongWait))
|
||||
c.conn.SetPongHandler(func(string) error {
|
||||
c.conn.SetReadDeadline(time.Now().Add(pongWait))
|
||||
return nil
|
||||
})
|
||||
|
||||
for {
|
||||
_, message, err := c.conn.ReadMessage()
|
||||
if err != nil {
|
||||
if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
|
||||
log.Info().Msgf("error: %v", err)
|
||||
}
|
||||
break
|
||||
}
|
||||
message = bytes.TrimSpace(bytes.Replace(message, newline, space, -1))
|
||||
if len(message) == 0 {
|
||||
continue
|
||||
}
|
||||
if isDebug() {
|
||||
log.Debug().Msgf("Received message %s", message)
|
||||
}
|
||||
|
||||
// handle ping
|
||||
if string(message) == "ping" {
|
||||
c.send <- []byte("pong")
|
||||
continue
|
||||
}
|
||||
|
||||
req, err := msg.ParseRequest(message)
|
||||
if err != nil {
|
||||
c.send <- []byte(responseMust(err, nil))
|
||||
continue
|
||||
}
|
||||
|
||||
c.hub.Requests <- Request{c, req}
|
||||
}
|
||||
}
|
||||
|
||||
// write pumps messages from the hub to the websocket connection.
|
||||
//
|
||||
// A goroutine running write is started for each connection. The
|
||||
// application ensures that there is at most one writer to a connection by
|
||||
// executing all writes from this goroutine.
|
||||
func (c *Client) write() {
|
||||
ticker := time.NewTicker(pingPeriod)
|
||||
defer func() {
|
||||
log.Debug().Msgf("Closing client write (%s)", c.GetAuth().UID)
|
||||
ticker.Stop()
|
||||
c.conn.Close()
|
||||
}()
|
||||
|
||||
for {
|
||||
select {
|
||||
case message, ok := <-c.send:
|
||||
c.conn.SetWriteDeadline(time.Now().Add(writeWait))
|
||||
if !ok {
|
||||
// The hub closed the channel.
|
||||
c.conn.WriteMessage(websocket.CloseMessage, []byte{})
|
||||
return
|
||||
}
|
||||
|
||||
w, err := c.conn.NextWriter(websocket.TextMessage)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
w.Write(message)
|
||||
if err := w.Close(); err != nil {
|
||||
return
|
||||
}
|
||||
case <-ticker.C:
|
||||
c.conn.SetWriteDeadline(time.Now().Add(writeWait))
|
||||
if err := c.conn.WriteMessage(websocket.PingMessage, nil); err != nil {
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
139
pkg/routing/client_test.go
Normal file
139
pkg/routing/client_test.go
Normal file
@@ -0,0 +1,139 @@
|
||||
package routing
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
)
|
||||
|
||||
func TestClient(t *testing.T) {
|
||||
hub := NewHub(nil)
|
||||
client := &Client{
|
||||
hub: hub,
|
||||
send: make(chan []byte, 256),
|
||||
Auth: Auth{UID: "UIDABC001", Role: "admin"},
|
||||
pubSub: []string{},
|
||||
privSub: []string{},
|
||||
}
|
||||
|
||||
assert.Equal(t, "UIDABC001", client.GetAuth().UID)
|
||||
assert.Equal(t, []string{}, client.GetSubscriptions())
|
||||
|
||||
client.SubscribePublic("a.x")
|
||||
assert.Equal(t, []string{"a.x"}, client.GetSubscriptions())
|
||||
assert.Equal(t, []string{"a.x"}, client.pubSub)
|
||||
assert.Equal(t, []string{}, client.privSub)
|
||||
|
||||
client.SubscribePublic("a.y")
|
||||
assert.Equal(t, []string{"a.x", "a.y"}, client.GetSubscriptions())
|
||||
assert.Equal(t, []string{"a.x", "a.y"}, client.pubSub)
|
||||
assert.Equal(t, []string{}, client.privSub)
|
||||
|
||||
client.UnsubscribePublic("a.y")
|
||||
assert.Equal(t, []string{"a.x"}, client.GetSubscriptions())
|
||||
assert.Equal(t, []string{"a.x"}, client.pubSub)
|
||||
assert.Equal(t, []string{}, client.privSub)
|
||||
|
||||
client.SubscribePrivate("b")
|
||||
assert.Equal(t, []string{"a.x", "b"}, client.GetSubscriptions())
|
||||
assert.Equal(t, []string{"a.x"}, client.pubSub)
|
||||
assert.Equal(t, []string{"b"}, client.privSub)
|
||||
|
||||
client.SubscribePrivate("c")
|
||||
assert.Equal(t, []string{"a.x", "b", "c"}, client.GetSubscriptions())
|
||||
assert.Equal(t, []string{"a.x"}, client.pubSub)
|
||||
assert.Equal(t, []string{"b", "c"}, client.privSub)
|
||||
|
||||
client.UnsubscribePrivate("b")
|
||||
assert.Equal(t, []string{"a.x", "c"}, client.GetSubscriptions())
|
||||
assert.Equal(t, []string{"a.x"}, client.pubSub)
|
||||
assert.Equal(t, []string{"c"}, client.privSub)
|
||||
|
||||
client.UnsubscribePrivate("c")
|
||||
assert.Equal(t, []string{"a.x"}, client.GetSubscriptions())
|
||||
assert.Equal(t, []string{"a.x"}, client.pubSub)
|
||||
assert.Equal(t, []string{}, client.privSub)
|
||||
|
||||
client.UnsubscribePublic("a.x")
|
||||
assert.Equal(t, []string{}, client.GetSubscriptions())
|
||||
assert.Equal(t, []string{}, client.pubSub)
|
||||
assert.Equal(t, []string{}, client.privSub)
|
||||
}
|
||||
|
||||
func TestParseStreamsFromURI(t *testing.T) {
|
||||
assert.Equal(t, []string{}, parseStreamsFromURI("/?"))
|
||||
assert.Equal(t, []string{}, parseStreamsFromURI(""))
|
||||
assert.Equal(t, []string{"aaa", "bbb"}, parseStreamsFromURI("/?stream=aaa&stream=bbb"))
|
||||
assert.Equal(t, []string{"aaa", "bbb"}, parseStreamsFromURI("/?stream=aaa,bbb"))
|
||||
assert.Equal(t, []string{"aaa", "bbb"}, parseStreamsFromURI("/public/?stream=aaa,bbb"))
|
||||
}
|
||||
|
||||
func TestCheckSameOriginEmpty(t *testing.T) {
|
||||
var checkSameOriginTests = []struct {
|
||||
ok bool
|
||||
r *http.Request
|
||||
}{
|
||||
{false, &http.Request{Host: "example.org", Header: map[string][]string{"Origin": {"https://other.org"}}}},
|
||||
{true, &http.Request{Host: "example.org", Header: map[string][]string{"Origin": {"https://example.org"}}}},
|
||||
{true, &http.Request{Host: "Example.org", Header: map[string][]string{"Origin": {"https://example.org"}}}},
|
||||
}
|
||||
|
||||
for _, tt := range checkSameOriginTests {
|
||||
ok := checkSameOrigin("")(tt.r)
|
||||
if tt.ok != ok {
|
||||
t.Errorf("checkSameOrigin(%+v) returned %v, want %v", tt.r, ok, tt.ok)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestCheckSameOriginDomainsSetup(t *testing.T) {
|
||||
var checkSameOriginTests = []struct {
|
||||
ok bool
|
||||
r *http.Request
|
||||
}{
|
||||
{false, &http.Request{Host: "example.org", Header: map[string][]string{"Origin": {"https://other.org"}}}},
|
||||
{true, &http.Request{Host: "whatever.org", Header: map[string][]string{"Origin": {"https://example.org"}}}},
|
||||
{true, &http.Request{Host: "whatever.org", Header: map[string][]string{"Origin": {"https://Example.org"}}}},
|
||||
{true, &http.Request{Host: "whatever.org", Header: map[string][]string{"Origin": {"https://example.com"}}}},
|
||||
{true, &http.Request{Host: "whatever.org", Header: map[string][]string{"Origin": {"https://Example.com"}}}},
|
||||
{true, &http.Request{Host: "whatever.org", Header: map[string][]string{"Origin": {}}}},
|
||||
}
|
||||
|
||||
checker := checkSameOrigin("example.org,example.com")
|
||||
for _, tt := range checkSameOriginTests {
|
||||
ok := checker(tt.r)
|
||||
if tt.ok != ok {
|
||||
t.Errorf("checkSameOrigin(%+v) returned %v, want %v", tt.r, ok, tt.ok)
|
||||
}
|
||||
}
|
||||
|
||||
checker = checkSameOrigin("example.org, example.com")
|
||||
for _, tt := range checkSameOriginTests {
|
||||
ok := checker(tt.r)
|
||||
if tt.ok != ok {
|
||||
t.Errorf("checkSameOrigin(%+v) returned %v, want %v", tt.r, ok, tt.ok)
|
||||
}
|
||||
}
|
||||
|
||||
checker = checkSameOrigin("https://example.org,https://example.com")
|
||||
for _, tt := range checkSameOriginTests {
|
||||
ok := checker(tt.r)
|
||||
if tt.ok != ok {
|
||||
t.Errorf("checkSameOrigin(%+v) returned %v, want %v", tt.r, ok, tt.ok)
|
||||
}
|
||||
}
|
||||
|
||||
checker = checkSameOrigin("https://example.org, https://example.com")
|
||||
for _, tt := range checkSameOriginTests {
|
||||
ok := checker(tt.r)
|
||||
if tt.ok != ok {
|
||||
t.Errorf("checkSameOrigin(%+v) returned %v, want %v", tt.r, ok, tt.ok)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestCheckSameOriginBadConfiguration(t *testing.T) {
|
||||
assert.Panics(t, func() { checkSameOrigin("https://ex ample.org") })
|
||||
assert.Panics(t, func() { checkSameOrigin("https://ex:ample.org") })
|
||||
}
|
||||
572
pkg/routing/hub.go
Normal file
572
pkg/routing/hub.go
Normal file
@@ -0,0 +1,572 @@
|
||||
package routing
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"strings"
|
||||
"sync"
|
||||
|
||||
msg "github.com/openware/rango/pkg/message"
|
||||
"github.com/openware/rango/pkg/metrics"
|
||||
"github.com/rs/zerolog"
|
||||
"github.com/rs/zerolog/log"
|
||||
"github.com/streadway/amqp"
|
||||
)
|
||||
|
||||
type Request struct {
|
||||
client IClient
|
||||
msg.Request
|
||||
}
|
||||
|
||||
// Hub maintains the set of active clients and broadcasts messages to the
|
||||
// clients.
|
||||
type Hub struct {
|
||||
// Register Requests from the clients.
|
||||
Requests chan Request
|
||||
|
||||
// Unregister requests from clients.
|
||||
Unregister chan IClient
|
||||
|
||||
// List of clients registered to public topics
|
||||
PublicTopics map[string]*Topic
|
||||
|
||||
// List of clients registered to private topics
|
||||
PrivateTopics map[string]map[string]*Topic
|
||||
|
||||
// map[prefix -> map[topic -> *Topic]]
|
||||
PrefixedTopics map[string]map[string]*Topic
|
||||
|
||||
// Storage for incremental objects
|
||||
IncrementalObjects map[string]*IncrementalObject
|
||||
|
||||
// map[prefix -> allowed roles]
|
||||
RBAC map[string][]string
|
||||
|
||||
mutex sync.Mutex
|
||||
}
|
||||
|
||||
type Event struct {
|
||||
Scope string // global, public, private
|
||||
Stream string // channel routing key
|
||||
Type string // event type
|
||||
Topic string // topic routing key (stream.type)
|
||||
Body interface{} // event json body
|
||||
}
|
||||
|
||||
type IncrementalObject struct {
|
||||
Snapshot string
|
||||
Increments []string
|
||||
}
|
||||
|
||||
func NewHub(rbac map[string][]string) *Hub {
|
||||
return &Hub{
|
||||
Requests: make(chan Request),
|
||||
Unregister: make(chan IClient),
|
||||
PublicTopics: make(map[string]*Topic, 100),
|
||||
PrivateTopics: make(map[string]map[string]*Topic, 1000),
|
||||
PrefixedTopics: make(map[string]map[string]*Topic, 100),
|
||||
IncrementalObjects: make(map[string]*IncrementalObject, 5),
|
||||
RBAC: rbac,
|
||||
}
|
||||
}
|
||||
|
||||
func isIncrementObject(s string) bool {
|
||||
return strings.HasSuffix(s, "-inc")
|
||||
}
|
||||
|
||||
func isSnapshotObject(s string) bool {
|
||||
return strings.HasSuffix(s, "-snap")
|
||||
}
|
||||
|
||||
func isDebug() bool {
|
||||
return log.Logger.GetLevel() <= zerolog.DebugLevel
|
||||
}
|
||||
|
||||
func isTrace() bool {
|
||||
return log.Logger.GetLevel() <= zerolog.TraceLevel
|
||||
}
|
||||
|
||||
func getTopic(scope, stream, typ string) string {
|
||||
if isSnapshotObject(typ) {
|
||||
typ = strings.Replace(typ, "-snap", "-inc", 1)
|
||||
}
|
||||
if scope == "private" {
|
||||
return typ
|
||||
}
|
||||
return stream + "." + typ
|
||||
}
|
||||
|
||||
func (h *Hub) ListenWebsocketEvents() {
|
||||
for {
|
||||
select {
|
||||
case req := <-h.Requests:
|
||||
h.handleRequest(&req)
|
||||
|
||||
case client := <-h.Unregister:
|
||||
log.Info().Msgf("Unregistering client (%s)", client.GetAuth().UID)
|
||||
h.unsubscribeAll(client)
|
||||
client.Close()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ReceiveMsg handles AMQP messages
|
||||
func (h *Hub) ReceiveMsg(delivery amqp.Delivery) {
|
||||
if isTrace() {
|
||||
log.Trace().Msgf("AMQP msg received: %s -> %s", delivery.RoutingKey, delivery.Body)
|
||||
}
|
||||
s := strings.Split(delivery.RoutingKey, ".")
|
||||
|
||||
var o interface{}
|
||||
err := json.Unmarshal(delivery.Body, &o)
|
||||
|
||||
if err != nil {
|
||||
log.Error().Msgf("JSON parse error: %s, msg: %s", err.Error(), delivery.Body)
|
||||
return
|
||||
}
|
||||
|
||||
switch len(s) {
|
||||
case 2:
|
||||
msg := Event{
|
||||
Scope: s[0],
|
||||
Stream: "",
|
||||
Type: s[1],
|
||||
Topic: getTopic(s[0], s[0], s[1]),
|
||||
Body: o,
|
||||
}
|
||||
|
||||
h.routeMessage(&msg)
|
||||
|
||||
case 3:
|
||||
msg := Event{
|
||||
Scope: s[0],
|
||||
Stream: s[1],
|
||||
Type: s[2],
|
||||
Topic: getTopic(s[0], s[1], s[2]),
|
||||
Body: o,
|
||||
}
|
||||
|
||||
h.routeMessage(&msg)
|
||||
|
||||
default:
|
||||
log.Error().Msgf("Bad routing key: %s", delivery.RoutingKey)
|
||||
}
|
||||
}
|
||||
|
||||
func (h *Hub) SkipPrivateMsg(delivery amqp.Delivery) {
|
||||
if strings.HasPrefix(delivery.RoutingKey, "private") {
|
||||
return
|
||||
}
|
||||
|
||||
h.ReceiveMsg(delivery)
|
||||
}
|
||||
|
||||
func (h *Hub) handleSnapshot(msg *Event) (string, error) {
|
||||
topic := msg.Stream + "." + msg.Type
|
||||
body, err := json.Marshal(map[string]interface{}{
|
||||
topic: msg.Body,
|
||||
})
|
||||
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
o, ok := h.IncrementalObjects[msg.Topic]
|
||||
if !ok {
|
||||
o = &IncrementalObject{}
|
||||
h.IncrementalObjects[msg.Topic] = o
|
||||
}
|
||||
o.Snapshot = string(body)
|
||||
o.Increments = []string{}
|
||||
|
||||
return string(body), nil
|
||||
}
|
||||
|
||||
func (h *Hub) handleIncrement(msg *Event) (string, error) {
|
||||
body, err := json.Marshal(map[string]interface{}{
|
||||
msg.Topic: msg.Body,
|
||||
})
|
||||
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
o, ok := h.IncrementalObjects[msg.Topic]
|
||||
if !ok {
|
||||
return "", fmt.Errorf("No snapshot received before the increment for topic %s, ignoring", msg.Topic)
|
||||
}
|
||||
o.Increments = append(o.Increments, string(body))
|
||||
return string(body), nil
|
||||
|
||||
}
|
||||
|
||||
func (h *Hub) handleMessage(topic *Topic, ok bool, msg *Event) {
|
||||
switch {
|
||||
case isIncrementObject(msg.Type):
|
||||
rm, err := h.handleIncrement(msg)
|
||||
if err != nil {
|
||||
log.Error().Msgf("handleIncrement failed: %s", err.Error())
|
||||
return
|
||||
}
|
||||
if ok {
|
||||
topic.broadcastRaw(rm)
|
||||
}
|
||||
|
||||
case isSnapshotObject(msg.Type):
|
||||
_, err := h.handleSnapshot(msg)
|
||||
if err != nil {
|
||||
log.Error().Msgf("handleSnapshot failed: %s", err.Error())
|
||||
return
|
||||
}
|
||||
|
||||
default:
|
||||
if ok {
|
||||
topic.broadcast(msg)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (h *Hub) routeMessage(msg *Event) {
|
||||
if isTrace() {
|
||||
log.Trace().Msgf("Routing message %v", msg)
|
||||
}
|
||||
h.mutex.Lock()
|
||||
defer h.mutex.Unlock()
|
||||
|
||||
switch msg.Scope {
|
||||
case "public", "global":
|
||||
topic, ok := h.PublicTopics[msg.Topic]
|
||||
h.handleMessage(topic, ok, msg)
|
||||
|
||||
if !ok {
|
||||
if isTrace() {
|
||||
log.Trace().Msgf("No public registration to %s", msg.Topic)
|
||||
log.Trace().Msgf("Public topics: %v", h.PublicTopics)
|
||||
}
|
||||
}
|
||||
|
||||
case "private":
|
||||
uid := msg.Stream
|
||||
uTopic, ok := h.PrivateTopics[uid]
|
||||
if ok {
|
||||
topic, ok := uTopic[msg.Topic]
|
||||
if ok {
|
||||
topic.broadcast(msg)
|
||||
break
|
||||
}
|
||||
}
|
||||
if isTrace() {
|
||||
log.Trace().Msgf("No private registration to %s", msg.Topic)
|
||||
log.Trace().Msgf("Private topics: %v", h.PrivateTopics)
|
||||
}
|
||||
|
||||
default:
|
||||
scope, ok := h.PrefixedTopics[msg.Scope]
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
|
||||
topic, ok := scope[msg.Topic]
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
|
||||
topic.broadcast(msg)
|
||||
|
||||
log.Trace().Msgf("Broadcasted message scope %s", msg.Scope)
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
func (h *Hub) unsubscribeAll(client IClient) {
|
||||
h.mutex.Lock()
|
||||
defer h.mutex.Unlock()
|
||||
|
||||
for t, topic := range h.PublicTopics {
|
||||
if topic.unsubscribe(client) {
|
||||
metrics.RecordHubUnsubscription("public", t)
|
||||
}
|
||||
if topic.len() == 0 {
|
||||
delete(h.PublicTopics, t)
|
||||
}
|
||||
}
|
||||
|
||||
for k, scope := range h.PrefixedTopics {
|
||||
for t, topic := range scope {
|
||||
if topic.unsubscribe(client) {
|
||||
metrics.RecordHubUnsubscription("prefixed", t)
|
||||
}
|
||||
|
||||
if topic.len() == 0 {
|
||||
delete(scope, t)
|
||||
}
|
||||
}
|
||||
|
||||
if len(scope) == 0 {
|
||||
delete(h.PrefixedTopics, k)
|
||||
}
|
||||
}
|
||||
|
||||
uid := client.GetAuth().UID
|
||||
topics, ok := h.PrivateTopics[uid]
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
|
||||
for t, topic := range topics {
|
||||
if topic.unsubscribe(client) {
|
||||
metrics.RecordHubUnsubscription("private", t)
|
||||
}
|
||||
if topic.len() == 0 {
|
||||
delete(topics, t)
|
||||
}
|
||||
}
|
||||
|
||||
if len(topics) == 0 {
|
||||
delete(h.PrivateTopics, uid)
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
func responseMust(e error, r interface{}) string {
|
||||
res, err := msg.PackOutgoingResponse(e, r)
|
||||
if err != nil {
|
||||
log.Panic().Msg("responseMust failed:" + err.Error())
|
||||
panic(err.Error())
|
||||
}
|
||||
|
||||
return string(res)
|
||||
}
|
||||
|
||||
func isPrivateStream(s string) bool {
|
||||
return strings.Count(s, ".") == 0
|
||||
}
|
||||
func isPrefixedStream(s string) bool {
|
||||
return strings.Count(s, ".") == 2
|
||||
}
|
||||
|
||||
func (h *Hub) handleRequest(req *Request) {
|
||||
switch req.Method {
|
||||
case "subscribe":
|
||||
h.handleSubscribe(req)
|
||||
case "unsubscribe":
|
||||
h.handleUnsubscribe(req)
|
||||
default:
|
||||
req.client.Send(responseMust(errors.New("unsupported method"), nil))
|
||||
}
|
||||
}
|
||||
|
||||
func (h *Hub) subscribePrivate(t string, req *Request) {
|
||||
uid := req.client.GetAuth().UID
|
||||
if uid == "" {
|
||||
log.Error().Msgf("Anonymous user tried to subscribe to private stream %s", t)
|
||||
return
|
||||
}
|
||||
|
||||
uTopics, ok := h.PrivateTopics[uid]
|
||||
if !ok {
|
||||
uTopics = make(map[string]*Topic, 3)
|
||||
h.PrivateTopics[uid] = uTopics
|
||||
}
|
||||
|
||||
topic, ok := uTopics[t]
|
||||
if !ok {
|
||||
topic = NewTopic(h)
|
||||
uTopics[t] = topic
|
||||
}
|
||||
|
||||
if topic.subscribe(req.client) {
|
||||
metrics.RecordHubSubscription("private", t)
|
||||
req.client.SubscribePrivate(t)
|
||||
}
|
||||
}
|
||||
|
||||
func (h *Hub) subscribePublic(t string, req *Request) {
|
||||
topic, ok := h.PublicTopics[t]
|
||||
if !ok {
|
||||
topic = NewTopic(h)
|
||||
h.PublicTopics[t] = topic
|
||||
}
|
||||
|
||||
// Replay snapshot and buffered increments before joining the topic, so live
|
||||
// broadcasts cannot interleave ahead of the initial ob-snap on market switch.
|
||||
if isIncrementObject(t) {
|
||||
o, ok := h.IncrementalObjects[t]
|
||||
if ok && o.Snapshot != "" {
|
||||
req.client.Send(o.Snapshot)
|
||||
for _, inc := range o.Increments {
|
||||
req.client.Send(inc)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if topic.subscribe(req.client) {
|
||||
metrics.RecordHubSubscription("public", t)
|
||||
req.client.SubscribePublic(t)
|
||||
}
|
||||
}
|
||||
|
||||
func (h *Hub) premittedRBAC(prefix string, auth Auth) bool {
|
||||
rbac := h.RBAC[prefix]
|
||||
|
||||
for _, role := range rbac {
|
||||
if role == auth.Role {
|
||||
return true
|
||||
}
|
||||
}
|
||||
|
||||
return false
|
||||
}
|
||||
|
||||
func splitPrefixedTopic(prefixed string) (string, string) {
|
||||
spl := strings.Split(prefixed, ".")
|
||||
prefix := spl[0]
|
||||
t := strings.TrimPrefix(prefixed, prefix+".")
|
||||
|
||||
return prefix, t
|
||||
}
|
||||
|
||||
func (h *Hub) subscribePrefixed(prefixed string, req *Request) {
|
||||
prefix, t := splitPrefixedTopic(prefixed)
|
||||
|
||||
if !h.premittedRBAC(prefix, req.client.GetAuth()) {
|
||||
req.client.Send(responseMust(nil, map[string]interface{}{
|
||||
"message": "cannot subscribe to " + prefixed,
|
||||
}))
|
||||
|
||||
return
|
||||
}
|
||||
|
||||
topics, ok := h.PrefixedTopics[prefix]
|
||||
if !ok {
|
||||
topics := make(map[string]*Topic, 0)
|
||||
h.PrefixedTopics[prefix] = topics
|
||||
}
|
||||
|
||||
topic, ok := topics[t]
|
||||
if !ok {
|
||||
topic = NewTopic(h)
|
||||
h.PrefixedTopics[prefix][t] = topic
|
||||
}
|
||||
|
||||
if isIncrementObject(t) {
|
||||
o, ok := h.IncrementalObjects[t]
|
||||
if ok && o.Snapshot != "" {
|
||||
req.client.Send(o.Snapshot)
|
||||
for _, inc := range o.Increments {
|
||||
req.client.Send(inc)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if topic.subscribe(req.client) {
|
||||
metrics.RecordHubSubscription("prefixed", prefixed)
|
||||
req.client.SubscribePublic(prefixed)
|
||||
}
|
||||
}
|
||||
|
||||
func (h *Hub) handleSubscribe(req *Request) {
|
||||
h.mutex.Lock()
|
||||
defer h.mutex.Unlock()
|
||||
|
||||
for _, t := range req.Streams {
|
||||
switch {
|
||||
case isPrivateStream(t):
|
||||
h.subscribePrivate(t, req)
|
||||
case isPrefixedStream(t):
|
||||
h.subscribePrefixed(t, req)
|
||||
default:
|
||||
h.subscribePublic(t, req)
|
||||
}
|
||||
}
|
||||
|
||||
req.client.Send(responseMust(nil, map[string]interface{}{
|
||||
"message": "subscribed",
|
||||
"streams": req.client.GetSubscriptions(),
|
||||
}))
|
||||
}
|
||||
|
||||
func (h *Hub) unsubscribePrivate(t string, req *Request) {
|
||||
uid := req.client.GetAuth().UID
|
||||
if uid == "" {
|
||||
return
|
||||
}
|
||||
uTopics, ok := h.PrivateTopics[uid]
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
|
||||
topic, ok := uTopics[t]
|
||||
if ok {
|
||||
if topic.unsubscribe(req.client) {
|
||||
metrics.RecordHubUnsubscription("private", t)
|
||||
req.client.UnsubscribePrivate(t)
|
||||
}
|
||||
|
||||
if topic.len() == 0 {
|
||||
delete(uTopics, t)
|
||||
}
|
||||
}
|
||||
|
||||
uTopics, ok = h.PrivateTopics[uid]
|
||||
if ok && len(uTopics) == 0 {
|
||||
delete(h.PrivateTopics, uid)
|
||||
}
|
||||
}
|
||||
|
||||
func (h *Hub) unsubscribePrefixed(prefixed string, req *Request) {
|
||||
scope, t := splitPrefixedTopic(prefixed)
|
||||
topics, ok := h.PrefixedTopics[scope]
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
|
||||
topic, ok := topics[t]
|
||||
if ok {
|
||||
if topic.unsubscribe(req.client) {
|
||||
metrics.RecordHubUnsubscription("prefixed", t)
|
||||
req.client.UnsubscribePublic(t)
|
||||
}
|
||||
|
||||
if topic.len() == 0 {
|
||||
delete(topics, t)
|
||||
h.PrefixedTopics[scope] = topics
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (h *Hub) unsubscribePublic(t string, req *Request) {
|
||||
topic, ok := h.PublicTopics[t]
|
||||
if ok {
|
||||
if topic.unsubscribe(req.client) {
|
||||
metrics.RecordHubUnsubscription("public", t)
|
||||
req.client.UnsubscribePublic(t)
|
||||
}
|
||||
|
||||
if topic.len() == 0 {
|
||||
delete(h.PublicTopics, t)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (h *Hub) handleUnsubscribe(req *Request) {
|
||||
h.mutex.Lock()
|
||||
defer h.mutex.Unlock()
|
||||
|
||||
for _, t := range req.Streams {
|
||||
switch {
|
||||
case isPrivateStream(t):
|
||||
h.unsubscribePrivate(t, req)
|
||||
case isPrefixedStream(t):
|
||||
h.unsubscribePrefixed(t, req)
|
||||
default:
|
||||
h.unsubscribePublic(t, req)
|
||||
}
|
||||
}
|
||||
|
||||
req.client.Send(responseMust(nil, map[string]interface{}{
|
||||
"message": "unsubscribed",
|
||||
"streams": req.client.GetSubscriptions(),
|
||||
}))
|
||||
}
|
||||
355
pkg/routing/hub_test.go
Normal file
355
pkg/routing/hub_test.go
Normal file
@@ -0,0 +1,355 @@
|
||||
package routing
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/openware/rango/pkg/message"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/mock"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
type MockedClient struct {
|
||||
mock.Mock
|
||||
}
|
||||
|
||||
func (c *MockedClient) Send(m string) {
|
||||
c.Called(m)
|
||||
}
|
||||
|
||||
func (c *MockedClient) Close() {
|
||||
}
|
||||
|
||||
func (c *MockedClient) GetAuth() Auth {
|
||||
args := c.Called()
|
||||
return args.Get(0).(Auth)
|
||||
}
|
||||
|
||||
func (c *MockedClient) GetSubscriptions() []string {
|
||||
args := c.Called()
|
||||
return args.Get(0).([]string)
|
||||
}
|
||||
|
||||
func (c *MockedClient) SubscribePublic(s string) {
|
||||
c.Called(s)
|
||||
}
|
||||
|
||||
func (c *MockedClient) SubscribePrivate(s string) {
|
||||
c.Called(s)
|
||||
}
|
||||
|
||||
func (c *MockedClient) UnsubscribePublic(s string) {
|
||||
c.Called(s)
|
||||
}
|
||||
|
||||
func (c *MockedClient) UnsubscribePrivate(s string) {
|
||||
c.Called(s)
|
||||
}
|
||||
|
||||
func setup(c *MockedClient, streams []string) *Hub {
|
||||
h := NewHub(nil)
|
||||
h.handleSubscribe(&Request{
|
||||
client: c,
|
||||
Request: message.Request{
|
||||
Streams: streams,
|
||||
},
|
||||
})
|
||||
return h
|
||||
}
|
||||
|
||||
func teardown(h *Hub, c *MockedClient, streams []string) {
|
||||
h.handleUnsubscribe(&Request{
|
||||
client: c,
|
||||
Request: message.Request{
|
||||
Streams: streams,
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
func TestAnonymous(t *testing.T) {
|
||||
t.Run("subscribe to a public single stream", func(t *testing.T) {
|
||||
c := &MockedClient{}
|
||||
|
||||
streams := []string{
|
||||
"eurusd.trades",
|
||||
}
|
||||
|
||||
c.On("GetAuth").Return(Auth{})
|
||||
c.On("GetSubscriptions").Return(streams).Once()
|
||||
c.On("SubscribePublic", streams[0]).Return().Once()
|
||||
c.On("Send", `{"success":{"message":"subscribed","streams":["`+streams[0]+`"]}}`).Return()
|
||||
|
||||
h := setup(c, streams)
|
||||
assert.Equal(t, 1, len(h.PublicTopics))
|
||||
assert.Equal(t, 0, len(h.PrivateTopics))
|
||||
|
||||
c.On("UnsubscribePublic", streams[0]).Return()
|
||||
c.On("GetSubscriptions").Return([]string{}).Once()
|
||||
c.On("Send", `{"success":{"message":"unsubscribed","streams":[]}}`).Return()
|
||||
|
||||
teardown(h, c, streams)
|
||||
assert.Equal(t, 0, len(h.PublicTopics))
|
||||
assert.Equal(t, 0, len(h.PrivateTopics))
|
||||
})
|
||||
|
||||
t.Run("subscribe to multiple public streams", func(t *testing.T) {
|
||||
c := &MockedClient{}
|
||||
streams := []string{
|
||||
"eurusd.trades",
|
||||
"eurusd.updates",
|
||||
}
|
||||
|
||||
c.On("GetAuth").Return(Auth{})
|
||||
c.On("GetSubscriptions").Return(streams).Once()
|
||||
c.On("SubscribePublic", "eurusd.trades").Return()
|
||||
c.On("SubscribePublic", "eurusd.updates").Return()
|
||||
c.On("Send", `{"success":{"message":"subscribed","streams":["eurusd.trades","eurusd.updates"]}}`).Return()
|
||||
|
||||
h := setup(c, []string{
|
||||
"eurusd.trades",
|
||||
"eurusd.updates",
|
||||
})
|
||||
|
||||
assert.Equal(t, 2, len(h.PublicTopics))
|
||||
assert.Equal(t, 0, len(h.PrivateTopics))
|
||||
|
||||
c.On("UnsubscribePublic", streams[0]).Return().Once()
|
||||
c.On("UnsubscribePublic", streams[1]).Return().Once()
|
||||
c.On("GetSubscriptions").Return([]string{}).Once()
|
||||
c.On("Send", `{"success":{"message":"unsubscribed","streams":[]}}`).Return()
|
||||
|
||||
teardown(h, c, streams)
|
||||
assert.Equal(t, 0, len(h.PublicTopics))
|
||||
assert.Equal(t, 0, len(h.PrivateTopics))
|
||||
|
||||
})
|
||||
|
||||
t.Run("subscribe to a private single stream", func(t *testing.T) {
|
||||
c := MockedClient{}
|
||||
|
||||
c.On("GetAuth").Return(Auth{})
|
||||
c.On("GetSubscriptions").Return([]string{})
|
||||
c.On("SubscribePrivate", "trades").Return()
|
||||
c.On("Send", `{"success":{"message":"subscribed","streams":[]}}`).Return()
|
||||
|
||||
h := setup(&c, []string{
|
||||
"trades",
|
||||
})
|
||||
|
||||
assert.Equal(t, 0, len(h.PublicTopics))
|
||||
assert.Equal(t, 0, len(h.PrivateTopics))
|
||||
})
|
||||
}
|
||||
func TestAuthenticated(t *testing.T) {
|
||||
t.Run("subscribe to a private single stream", func(t *testing.T) {
|
||||
c := &MockedClient{}
|
||||
|
||||
c.On("GetAuth").Return(Auth{UID: "UIDABC00001"})
|
||||
c.On("GetSubscriptions").Return([]string{"trades"}).Once()
|
||||
c.On("SubscribePrivate", "trades").Return()
|
||||
c.On("Send", `{"success":{"message":"subscribed","streams":["trades"]}}`).Return()
|
||||
|
||||
h := setup(c, []string{
|
||||
"trades",
|
||||
})
|
||||
assert.Equal(t, 0, len(h.PublicTopics))
|
||||
assert.Equal(t, 1, len(h.PrivateTopics))
|
||||
|
||||
c.On("UnsubscribePrivate", "trades").Return().Once()
|
||||
c.On("GetSubscriptions").Return([]string{}).Once()
|
||||
c.On("Send", `{"success":{"message":"unsubscribed","streams":[]}}`).Return()
|
||||
|
||||
teardown(h, c, []string{"trades"})
|
||||
assert.Equal(t, 0, len(h.PublicTopics))
|
||||
assert.Equal(t, 0, len(h.PrivateTopics))
|
||||
})
|
||||
|
||||
t.Run("subscribe to multiple private streams", func(t *testing.T) {
|
||||
c := &MockedClient{}
|
||||
|
||||
c.On("GetSubscriptions").Return([]string{"trades", "orders"}).Once()
|
||||
c.On("GetAuth").Return(Auth{UID: "UIDABC00001"})
|
||||
c.On("SubscribePrivate", "trades").Return()
|
||||
c.On("SubscribePrivate", "orders").Return()
|
||||
c.On("Send", `{"success":{"message":"subscribed","streams":["trades","orders"]}}`).Return()
|
||||
|
||||
h := setup(c, []string{"trades", "orders"})
|
||||
assert.Equal(t, 0, len(h.PublicTopics))
|
||||
assert.Equal(t, 1, len(h.PrivateTopics))
|
||||
|
||||
uTopics, ok := h.PrivateTopics["UIDABC00001"]
|
||||
require.True(t, ok)
|
||||
assert.Equal(t, 2, len(uTopics))
|
||||
|
||||
c.On("UnsubscribePrivate", "trades").Return().Once()
|
||||
c.On("UnsubscribePrivate", "orders").Return().Once()
|
||||
c.On("GetSubscriptions").Return([]string{}).Once()
|
||||
c.On("Send", `{"success":{"message":"unsubscribed","streams":[]}}`).Return()
|
||||
|
||||
teardown(h, c, []string{"trades", "orders"})
|
||||
assert.Equal(t, 0, len(h.PublicTopics))
|
||||
assert.Equal(t, 0, len(h.PrivateTopics))
|
||||
|
||||
})
|
||||
|
||||
t.Run("subscribe to multiple private and public streams", func(t *testing.T) {
|
||||
c := &MockedClient{}
|
||||
|
||||
c.On("GetSubscriptions").Return([]string{"trades", "orders", "eurusd.updates"}).Once()
|
||||
c.On("GetAuth").Return(Auth{UID: "UIDABC00001"})
|
||||
c.On("SubscribePrivate", "trades").Return()
|
||||
c.On("SubscribePrivate", "orders").Return()
|
||||
c.On("SubscribePublic", "eurusd.updates").Return()
|
||||
c.On("Send", `{"success":{"message":"subscribed","streams":["trades","orders","eurusd.updates"]}}`).Return()
|
||||
|
||||
h := setup(c, []string{"trades", "orders", "eurusd.updates"})
|
||||
assert.Equal(t, 1, len(h.PublicTopics))
|
||||
assert.Equal(t, 1, len(h.PrivateTopics))
|
||||
|
||||
uTopics, ok := h.PrivateTopics["UIDABC00001"]
|
||||
require.True(t, ok)
|
||||
assert.Equal(t, 2, len(uTopics))
|
||||
|
||||
c.On("UnsubscribePrivate", "trades").Return().Once()
|
||||
c.On("UnsubscribePrivate", "orders").Return().Once()
|
||||
c.On("UnsubscribePublic", "eurusd.updates").Return().Once()
|
||||
c.On("GetSubscriptions").Return([]string{}).Once()
|
||||
c.On("Send", `{"success":{"message":"unsubscribed","streams":[]}}`).Return()
|
||||
|
||||
teardown(h, c, []string{"trades", "orders", "eurusd.updates"})
|
||||
assert.Equal(t, 0, len(h.PublicTopics))
|
||||
assert.Equal(t, 0, len(h.PrivateTopics))
|
||||
})
|
||||
}
|
||||
|
||||
func TestIsIncremental(t *testing.T) {
|
||||
assert.True(t, isIncrementObject("public.eurusd.ob-inc"))
|
||||
assert.False(t, isIncrementObject("public.eurusd.ob-snap"))
|
||||
assert.False(t, isIncrementObject("public.eurusd.ob"))
|
||||
|
||||
assert.True(t, isSnapshotObject("public.eurusd.ob-snap"))
|
||||
assert.False(t, isSnapshotObject("public.eurusd.ob-inc"))
|
||||
assert.False(t, isSnapshotObject("public.eurusd.ob"))
|
||||
}
|
||||
|
||||
func TestGetTopic(t *testing.T) {
|
||||
assert.Equal(t, "abc.count", getTopic("public", "abc", "count"))
|
||||
assert.Equal(t, "count", getTopic("private", "abc", "count"))
|
||||
assert.Equal(t, "abc.count-inc", getTopic("public", "abc", "count-inc"))
|
||||
assert.Equal(t, "abc.count-inc", getTopic("public", "abc", "count-snap"))
|
||||
}
|
||||
|
||||
func TestHandleMessage(t *testing.T) {
|
||||
h := NewHub(nil)
|
||||
c := &MockedClient{}
|
||||
c.On("SubscribePublic", "abc.ticker").Return()
|
||||
c.On("Send", "{\"abc.ticker\":{\"some\":\"data\"}}").Return()
|
||||
|
||||
h.subscribePublic("abc.ticker", &Request{
|
||||
client: c,
|
||||
})
|
||||
|
||||
h.routeMessage(&Event{
|
||||
Scope: "public",
|
||||
Stream: "abc",
|
||||
Type: "ticker",
|
||||
Topic: "abc.ticker",
|
||||
Body: map[string]interface{}{
|
||||
"some": "data",
|
||||
},
|
||||
})
|
||||
|
||||
c.AssertExpectations(t)
|
||||
}
|
||||
|
||||
func TestIncrementalObjectStorage(t *testing.T) {
|
||||
h := NewHub(nil)
|
||||
|
||||
// Increments before the first snapshot must be ignored
|
||||
h.routeMessage(&Event{
|
||||
Scope: "public",
|
||||
Stream: "abc",
|
||||
Type: "count-inc",
|
||||
Topic: "abc.count-inc",
|
||||
Body: map[string]interface{}{
|
||||
"data": 1,
|
||||
"sequence": 11,
|
||||
},
|
||||
})
|
||||
|
||||
require.Equal(t, 0, len(h.IncrementalObjects))
|
||||
|
||||
// Initial snapshot
|
||||
h.routeMessage(&Event{
|
||||
Scope: "public",
|
||||
Stream: "abc",
|
||||
Type: "count-snap",
|
||||
Topic: "abc.count-inc",
|
||||
Body: map[string]interface{}{
|
||||
"data": []int{2, 3, 4},
|
||||
"sequence": 12,
|
||||
},
|
||||
})
|
||||
|
||||
require.Equal(t, 1, len(h.IncrementalObjects))
|
||||
|
||||
o, ok := h.IncrementalObjects["abc.count-inc"]
|
||||
require.True(t, ok)
|
||||
require.Equal(t, 0, len(o.Increments))
|
||||
require.Equal(t, `{"abc.count-snap":{"data":[2,3,4],"sequence":12}}`, o.Snapshot)
|
||||
|
||||
// First Increment
|
||||
h.routeMessage(&Event{
|
||||
Scope: "public",
|
||||
Stream: "abc",
|
||||
Type: "count-inc",
|
||||
Topic: "abc.count-inc",
|
||||
Body: map[string]interface{}{
|
||||
"data": 5,
|
||||
"sequence": 13,
|
||||
},
|
||||
})
|
||||
require.Equal(t, 1, len(h.IncrementalObjects))
|
||||
o, ok = h.IncrementalObjects["abc.count-inc"]
|
||||
require.True(t, ok)
|
||||
require.Equal(t, 1, len(o.Increments))
|
||||
require.Equal(t, `{"abc.count-snap":{"data":[2,3,4],"sequence":12}}`, o.Snapshot)
|
||||
require.Equal(t, `{"abc.count-inc":{"data":5,"sequence":13}}`, o.Increments[0])
|
||||
|
||||
// Second Increment
|
||||
h.routeMessage(&Event{
|
||||
Scope: "public",
|
||||
Stream: "abc",
|
||||
Type: "count-inc",
|
||||
Topic: "abc.count-inc",
|
||||
Body: map[string]interface{}{
|
||||
"data": 6,
|
||||
"sequence": 14,
|
||||
},
|
||||
})
|
||||
require.Equal(t, 1, len(h.IncrementalObjects))
|
||||
o, ok = h.IncrementalObjects["abc.count-inc"]
|
||||
require.True(t, ok)
|
||||
require.Equal(t, 2, len(o.Increments))
|
||||
require.Equal(t, `{"abc.count-snap":{"data":[2,3,4],"sequence":12}}`, o.Snapshot)
|
||||
require.Equal(t, `{"abc.count-inc":{"data":5,"sequence":13}}`, o.Increments[0])
|
||||
require.Equal(t, `{"abc.count-inc":{"data":6,"sequence":14}}`, o.Increments[1])
|
||||
|
||||
// Second snapshot
|
||||
h.routeMessage(&Event{
|
||||
Scope: "public",
|
||||
Stream: "abc",
|
||||
Type: "count-snap",
|
||||
Topic: "abc.count-inc",
|
||||
Body: map[string]interface{}{
|
||||
"data": []int{2, 3, 4, 5, 6},
|
||||
"sequence": 14,
|
||||
},
|
||||
})
|
||||
|
||||
require.Equal(t, 1, len(h.IncrementalObjects))
|
||||
o, ok = h.IncrementalObjects["abc.count-inc"]
|
||||
require.True(t, ok)
|
||||
require.Equal(t, 0, len(o.Increments))
|
||||
require.Equal(t, `{"abc.count-snap":{"data":[2,3,4,5,6],"sequence":14}}`, o.Snapshot)
|
||||
}
|
||||
79
pkg/routing/topic.go
Normal file
79
pkg/routing/topic.go
Normal file
@@ -0,0 +1,79 @@
|
||||
package routing
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
|
||||
msg "github.com/openware/rango/pkg/message"
|
||||
"github.com/rs/zerolog/log"
|
||||
)
|
||||
|
||||
type Topic struct {
|
||||
hub *Hub
|
||||
clients map[IClient]struct{}
|
||||
}
|
||||
|
||||
func NewTopic(h *Hub) *Topic {
|
||||
return &Topic{
|
||||
clients: make(map[IClient]struct{}),
|
||||
hub: h,
|
||||
}
|
||||
}
|
||||
|
||||
func eventMust(method string, data interface{}) []byte {
|
||||
ev, err := msg.PackOutgoingEvent(method, data)
|
||||
if err != nil {
|
||||
log.Panic().Msg(err.Error())
|
||||
}
|
||||
|
||||
return ev
|
||||
}
|
||||
|
||||
func contains(list []string, el string) bool {
|
||||
for _, l := range list {
|
||||
if l == el {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func (t *Topic) len() int {
|
||||
return len(t.clients)
|
||||
}
|
||||
|
||||
func (t *Topic) broadcast(message *Event) {
|
||||
body, err := json.Marshal(map[string]interface{}{
|
||||
message.Topic: message.Body,
|
||||
})
|
||||
|
||||
if err != nil {
|
||||
log.Error().Msgf("Fail to JSON marshal: %s", err.Error())
|
||||
return
|
||||
}
|
||||
|
||||
for client := range t.clients {
|
||||
client.Send(string(body))
|
||||
}
|
||||
}
|
||||
|
||||
func (t *Topic) broadcastRaw(msgBody string) {
|
||||
for client := range t.clients {
|
||||
client.Send(msgBody)
|
||||
}
|
||||
}
|
||||
|
||||
func (t *Topic) subscribe(c IClient) bool {
|
||||
if _, ok := t.clients[c]; ok {
|
||||
return false
|
||||
}
|
||||
t.clients[c] = struct{}{}
|
||||
|
||||
return true
|
||||
}
|
||||
|
||||
func (t *Topic) unsubscribe(c IClient) bool {
|
||||
_, ok := t.clients[c]
|
||||
delete(t.clients, c)
|
||||
|
||||
return ok
|
||||
}
|
||||
Reference in New Issue
Block a user