Merge pull request #24 from Ilhasoft/feat/session-token-sending

Send token feature for session verification

api/main.go

@@ -31,15 +31,15 @@ func init() {
 	})
 
 	if config.Get.SentryDSN != "" {
-	hook, err := logrus_sentry.NewSentryHook(config.Get.SentryDSN, []log.Level{log.PanicLevel, log.FatalLevel, log.ErrorLevel})
-	hook.Timeout = 0
-	hook.StacktraceConfiguration.Enable = true
-	hook.StacktraceConfiguration.Skip = 4
-	hook.StacktraceConfiguration.Context = 5
-	if err != nil {
-		log.Fatalf("invalid sentry DSN: '%s': %s", config.Get.SentryDSN, err)
-	}
-	log.StandardLogger().Hooks.Add(hook)
+		hook, err := logrus_sentry.NewSentryHook(config.Get.SentryDSN, []log.Level{log.PanicLevel, log.FatalLevel, log.ErrorLevel})
+		hook.Timeout = 0
+		hook.StacktraceConfiguration.Enable = true
+		hook.StacktraceConfiguration.Skip = 4
+		hook.StacktraceConfiguration.Context = 5
+		if err != nil {
+			log.Fatalf("invalid sentry DSN: '%s': %s", config.Get.SentryDSN, err)
+		}
+		log.StandardLogger().Hooks.Add(hook)
 	}
 }
 

go.mod

@@ -6,6 +6,7 @@ require (
 	github.com/adjust/rmq/v4 v4.0.1
 	github.com/aws/aws-sdk-go v1.38.8
 	github.com/certifi/gocertifi v0.0.0-20210507211836-431795d63e8d // indirect
+	github.com/dchest/uniuri v0.0.0-20200228104902-7aecb25e1fe5 // indirect
 	github.com/evalphobia/logrus_sentry v0.8.2 // indirect
 	github.com/getsentry/raven-go v0.2.0 // indirect
 	github.com/go-playground/universal-translator v0.17.0 // indirect

go.sum

@@ -20,6 +20,8 @@ github.com/cespare/xxhash/v2 v2.1.2/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XL
 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 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/dchest/uniuri v0.0.0-20200228104902-7aecb25e1fe5 h1:RAV05c0xOkJ3dZGS0JFybxFKZ2WMLabgx3uXnd7rpGs=
+github.com/dchest/uniuri v0.0.0-20200228104902-7aecb25e1fe5/go.mod h1:GgB8SF9nRG+GqaDtLcwJZsQFhcogVCJ79j4EdT0c2V4=
 github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f h1:lO4WD4F/rVNCu3HqELle0jiPLLBs70cWOduZpkS1E78=
 github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f/go.mod h1:cuUVRXasLTGF7a8hSLbxyZXjz+1KgoB3wDUb6vlszIc=
 github.com/evalphobia/logrus_sentry v0.8.2 h1:dotxHq+YLZsT1Bb45bB5UQbfCh3gM/nFFetyN46VoDQ=

pkg/websocket/client.go

@@ -16,6 +16,7 @@ import (
 	"github.com/ilhasoft/wwcs/pkg/metric"
 	"github.com/ilhasoft/wwcs/pkg/queue"
 	log "github.com/sirupsen/logrus"
+	uni "github.com/dchest/uniuri"
 )
 
 // Client errors
@@ -36,6 +37,7 @@ type Client struct {
 	Origin          string
 	Channel         string
 	Host            string
+	AuthToken       string
 }
 
 func (c *Client) Read(app *App) {
@@ -88,11 +90,35 @@ func (c *Client) ParsePayload(app *App, payload OutgoingPayload, to postJSON) er
 		return c.Redirect(payload, to, app)
 	case "ping":
 		return c.Redirect(payload, to, app)
+	case "close_session":
+		return CloseSession(payload, app)	
 	}
 
 	return ErrorInvalidPayloadType
 }
 
+func CloseSession(payload OutgoingPayload, app *App) error{
+
+	client := app.Pool.Clients[payload.From]
+	if client != nil {
+		if client.AuthToken == payload.Token {
+			errorPayload := IncomingPayload{
+				Type:  "warning",
+				Warning: "Connection closed by request",
+			}
+			err := client.Send(errorPayload)
+			if err != nil {
+				log.Error(err)
+			}
+			client.Conn.Close()
+			return nil
+		} else {
+			return ErrorInvalidToken
+		}
+	}
+	return ErrorInvalidClient
+}
+
 // Register register an user
 func (c *Client) Register(payload OutgoingPayload, triggerTo postJSON, app *App) error {
 	start := time.Now()
@@ -101,12 +127,21 @@ func (c *Client) Register(payload OutgoingPayload, triggerTo postJSON, app *App)
 		return err
 	}
 
-	if _, found := app.Pool.Clients[payload.From]; found {
+	if client, found := app.Pool.Clients[payload.From]; found {
+		tokenPayload := IncomingPayload{
+			Type: "token",
+			Token: client.AuthToken,
+		}
+		err = c.Send(tokenPayload)
+		if err != nil {
+			return err
+		}
 		return ErrorIDAlreadyExists
 	}
 
 	c.ID = payload.From
 	c.Callback = payload.Callback
+	c.AuthToken = uni.NewLen(32)
 	c.setupClientQueue(app.RDB)
 
 	u, err := url.Parse(payload.Callback)
@@ -162,6 +197,16 @@ func (c *Client) Register(payload OutgoingPayload, triggerTo postJSON, app *App)
 		app.Metrics.SaveSocketRegistration(socketRegistrationMetrics)
 	}
 
+	// token sending
+	tokenPayload := IncomingPayload {
+		Type: "token",
+		Token: c.AuthToken,
+	}
+	err = c.Send(tokenPayload)
+	if err != nil {
+		return err
+	}
+
 	return nil
 }
 

pkg/websocket/client_test.go

@@ -45,11 +45,78 @@ var ttParsePayload = []struct {
 func TestParsePayload(t *testing.T) {
 	rdb := redis.NewClient(&redis.Options{Addr: "localhost:6379", DB: 3})
 	app := NewApp(NewPool(), nil, rdb, nil)
+	client, ws, s := newTestClient(t)
+	defer client.Conn.Close()
+	defer ws.Close()
+	defer s.Close()
+
+	for _, tt := range ttParsePayload {
+		t.Run(tt.TestName, func(t *testing.T) {
+			client.ID = tt.Payload.From
+			client.Callback = tt.Payload.Callback
+
+			err := client.ParsePayload(app, tt.Payload, toTest)
+			if err != tt.Err {
+				t.Errorf("got %v, want %v", err, tt.Err)
+			}
+		})
+	}
+}
+
+var ttCloseSession = []struct {
+	TestName string
+	Payload  OutgoingPayload
+	Err      error
+}{
+	{
+		TestName: "Close Session",
+		Payload: OutgoingPayload{
+			Type:     "close_session",
+			Callback: "https://foo.bar",
+			From:     "00005",
+			Token:    "abcde",
+		},
+		Err: nil,
+	},
+	{
+		TestName: "Invalid Token",
+		Payload: OutgoingPayload{
+			Type:     "close_session",
+			Callback: "https://foo.bar",
+			From:     "00005",
+			Token:    "abce",
+		},
+		Err: ErrorInvalidToken,
+	},
+	{
+		TestName: "Invalid Client",
+		Payload: OutgoingPayload{
+			Type:     "close_session",
+			Callback: "https://foo.bar",
+			From:     "00000",
+			Token:    "abcde",
+		},
+		Err: ErrorInvalidClient,
+	},
+}
+
+func TestCloseSession(t *testing.T) {
+	rdb := redis.NewClient(&redis.Options{Addr: "localhost:6379", DB: 3})
+	app := NewApp(NewPool(), nil, rdb, nil)
+	conn := NewOpenConnection(t)
+
 	client := &Client{
-		Conn: nil,
+		ID:   "00005",
+		Conn: conn,
+		AuthToken: "abcde",
 	}
 
-	for _, tt := range ttParsePayload {
+	defer client.Conn.Close()
+
+	// Register client that will have the session closed
+	app.Pool.Clients[client.ID] = client
+
+	for _, tt := range ttCloseSession {
 		t.Run(tt.TestName, func(t *testing.T) {
 			client.ID = tt.Payload.From
 			client.Callback = tt.Payload.Callback
@@ -127,9 +194,11 @@ func TestClientRegister(t *testing.T) {
 	rdb := redis.NewClient(&redis.Options{Addr: "localhost:6379", DB: 3})
 	app := NewApp(NewPool(), nil, rdb, nil)
 	var poolSize int
-	client := &Client{
-		Conn: nil,
-	}
+
+	client, ws, s := newTestClient(t)
+	defer client.Conn.Close()
+	defer ws.Close()
+	defer s.Close()
 
 	for _, tt := range ttClientRegister {
 		t.Run(tt.TestName, func(t *testing.T) {
@@ -384,6 +453,15 @@ var ttSend = []struct {
 		Want: fmt.Sprintln(`{"type":"pong","to":"","from":"","message":{"type":"","timestamp":""}}`),
 		Err:  nil,
 	},
+	{
+		TestName: "Token Message",
+		Payload: IncomingPayload{
+			Type:  "token",
+			Token: "aaaaaa",
+		},
+		Want: fmt.Sprintln(`{"type":"token","to":"","from":"","message":{"type":"","timestamp":""},"token":"aaaaaa"}`),
+		Err:  nil,
+	},
 }
 
 func TestSend(t *testing.T) {
@@ -428,3 +506,10 @@ func newTestClient(t *testing.T) (*Client, *websocket.Conn, *httptest.Server) {
 
 	return client, ws, server
 }
+
+func NewOpenConnection(t  *testing.T) (*websocket.Conn){
+	t.Helper()
+	_, _, conn := newTestServer(t)
+
+	return conn
+}

pkg/websocket/helper.go

@@ -28,7 +28,9 @@ var (
 	ErrorInvalidMessageType = fmt.Errorf("%s invalid message type", errorPrefix)
 	ErrorDecodingMedia      = fmt.Errorf("%s could not decode media", errorPrefix)
 	ErrorUploadingToS3      = fmt.Errorf("%s can not upload image to s3", errorPrefix)
-	// register
+	// close_session
+	ErrorInvalidToken  = fmt.Errorf("token does not match that of the client")
+	ErrorInvalidClient = fmt.Errorf("Client not found")
 )
 
 func formatOutgoingPayload(payload OutgoingPayload) (OutgoingPayload, error) {

pkg/websocket/payload.go

@@ -6,7 +6,9 @@ type IncomingPayload struct {
 	To      string  `json:"to" validate:"required"`
 	From    string  `json:"from" validate:"required"`
 	Error   string  `json:"error,omitempty"`
-	Message Message `json:"message"`
+	Message Message `json:"message,omitempty"`
+	Token   string  `json:"token,omitempty"`
+	Warning  string  `json:"warning,omitempty"`
 }
 
 // OutgoingPayload data (outgoing messages)
@@ -16,6 +18,7 @@ type OutgoingPayload struct {
 	Callback string  `json:"callback,omitempty"`
 	Trigger  string  `json:"trigger,omitempty"`
 	Message  Message `json:"message,omitempty"`
+	Token    string  `json:"token,omitempty"`
 }
 
 // Message data