cloud-services/services/gateway/websocket/session_mobile.go

package websocket

import (
	"context"
	"encoding/json"
	"net/http"
	"time"

	"fiskerinc.com/modules/common"
	"fiskerinc.com/modules/jwt"
	"fiskerinc.com/modules/kafka"
	"fiskerinc.com/modules/logger"
	"fiskerinc.com/modules/validator"
	"google.golang.org/protobuf/proto"

	"github.com/gobwas/ws"
	"github.com/pkg/errors"
	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/tracer"
)

// NewMobileSession serves as the constructor for HMI sessions
func NewMobileSession(w http.ResponseWriter, r *http.Request, version string) (SessionInterface, error) {
	var s SessionInterface

	conn, _, _, err := ws.UpgradeHTTP(r, w)
	if err != nil {
		return s, errors.WithStack(err)
	}

	return &SessionMobile{
		Session: &Session{
			Websocket: conn,
			Type:      common.Mobile,
			Version:   version,
			epoch:     time.Now().UnixNano(),
		},
	}, nil
}

// SessionMobile contains websocket info
type SessionMobile struct {
	*Session
}

// Authenticate returns id if proper authentication, else returns error
func (s *SessionMobile) Authenticate() error {
	err := s.authenticate()

	data, _ := json.Marshal(AuthResponse{
		Handler: "verify",
		Data: AuthResponseData{
			Authenticated: err == nil,
		},
	})
	s.SendMsgToClient(data)
	if err != nil {
		return errors.WithStack(err)
	}

	return nil
}

func (s *SessionMobile) authenticate() error {
	msg, _, err := s.Receive()
	if err != nil {
		return err
	}

	var m common.MobileSessionMessage
	err = json.Unmarshal(msg, &m)
	if err != nil {
		return errors.WithStack(err)
	} else if m.Handler != "verify" {
		return errors.Errorf("incorrect auth handler specified %v", m.Handler)
	}

	err = validator.ValidateStruct(m)
	if err != nil {
		return errors.WithStack(err)
	}

	valid := jwt.NewJWTValidator("")
	payload, err := valid.ValidateToken(m.Data.Token)
	if err != nil {
		return err
	}

	id, ok := payload["username"].(string)
	if !ok {
		return ErrInvalidToken
	}

	s.ID = id
	return nil
}

// Listen to websocket session and use handler upon message received
func (s *SessionMobile) Listen(ctx context.Context, producer kafka.ProducerInterface) error {
	span, _ := tracer.StartSpanFromContext(ctx, "listen")
	defer span.Finish()

	key := s.Key()
	for {
		msg, op, err := s.Receive()
		if op == ws.OpClose {
			logger.At(logger.Info(), key, "conn").Msg("OpClose")
			return nil
		} else if err != nil {
			logger.At(logger.Info(), key, "conn").Err(err).Send()
			return err
		}

		err = s.Route(producer, msg)
		if err != nil {
			logger.At(logger.Warn(), key, "route").
				Err(err).Send()
		}
	}
}

// Route Mobile messages
func (s *SessionMobile) Route(producer kafka.ProducerInterface, data []byte) error {
	var m common.MessageRawJSON
	err := m.Unmarshal(data)
	if err != nil {
		return err
	}

	key := s.Key()
	logger.At(logger.Debug(), key, "route").
		Str("handler", m.Handler).
		Int("size", len(data)).
		Msgf("received from %v", key)

	topic, ok := kafka.MobileHandlerTopics[m.Handler]
	if !ok {
		return ErrInvalidHandler(m.Handler)
	}

	switch topic {
	case kafka.AttendantServiceGRPCKafka:
		valetData := common.AttendantRouteMobileGRPC(m)
		grpcData, err := proto.Marshal(valetData)
		if err != nil {
			logger.At(logger.Error(), key, "route").
				Err(errors.Wrap(err, "unable to marshal trexlogs GRPC "+topic)).Send()
			return errors.WithStack(err)
		}
		err = producer.ProduceBinary(kafka.AttendantServiceGRPCKafka, key, grpcData, nil)
		grpcData = nil
		if err != nil {
			logger.At(logger.Error(), key, "route").
				Err(errors.Wrap(err, "Producer for trex logs failed")).Send()
			return err
		}
	case kafka.DepotServiceGRPCKafka:
		valetData := common.DepotRouteMobileToGRPC(m)
		grpcData, err := proto.Marshal(valetData)
		if err != nil {
			logger.At(logger.Error(), key, "route").
				Err(errors.Wrap(err, "unable to marshal trexlogs GRPC "+topic)).Send()
			return errors.WithStack(err)
		}
		err = producer.ProduceBinary(kafka.DepotServiceGRPCKafka, key, grpcData, nil)
		grpcData = nil
		if err != nil {
			logger.At(logger.Error(), key, "route").
				Err(errors.Wrap(err, "Producer for trex logs failed")).Send()
			return err
		}
	case kafka.ValetServiceGRPCKafka:
		valetData := common.ValetRouteMobilePayloadGRPC(m)
		grpcData, err := proto.Marshal(valetData)
		if err != nil {
			logger.At(logger.Error(), key, "route").
				Err(errors.Wrap(err, "unable to marshal trexlogs GRPC "+topic)).Send()
			return errors.WithStack(err)
		}
		err = producer.ProduceBinary(kafka.ValetServiceGRPCKafka, key, grpcData, nil)
		grpcData = nil
		if err != nil {
			logger.At(logger.Error(), key, "route").
				Err(errors.Wrap(err, "Producer for trex logs failed")).Send()
			return err
		}
	default:
		err = producer.Produce(topic, key, m, nil)
	}
	if err != nil {
		return err
	}

	return nil
}