member_gateway_rpc/servlet/webapp-servlet.go

package servlet

import (
	"crypto/rsa"
	"fmt"
	"sync"

	"git.loafle.net/commons/logging-go"
	"git.loafle.net/commons/server-go"
	"git.loafle.net/commons/server-go/socket"
	og "git.loafle.net/overflow/gateway"
	ogs "git.loafle.net/overflow/gateway/subscribe"
	ogrs "git.loafle.net/overflow/gateway_rpc/servlet"
	"git.loafle.net/overflow/member_gateway_rpc/subscribe"
	"github.com/dgrijalva/jwt-go"
	"github.com/satori/go.uuid"
	"github.com/valyala/fasthttp"
)

type WebappServlet interface {
	ogrs.RPCServlet
}

type WebappServlets struct {
	ogrs.RPCServlets

	VerifyKey *rsa.PublicKey
	SignKey   *rsa.PrivateKey

	connections sync.Map
}

func (s *WebappServlets) Init(serverCtx server.ServerCtx) error {
	if err := s.RPCServlets.Init(serverCtx); nil != err {
		return err
	}

	return nil
}

func (s *WebappServlets) OnStart(serverCtx server.ServerCtx) error {
	if err := s.RPCServlets.OnStart(serverCtx); nil != err {
		return err
	}

	subscribeChan, err := subscribe.Subscriber.Subscribe("/webapp")
	if nil != err {
		return err
	}
	go s.handleSubscribe(serverCtx, subscribeChan)

	return nil
}

func (s *WebappServlets) OnStop(serverCtx server.ServerCtx) {
	if err := subscribe.Subscriber.Unsubscribe("/webapp"); nil != err {
		logging.Logger().Warn(err)
	}

	s.RPCServlets.OnStop(serverCtx)
}

func (s *WebappServlets) Destroy(serverCtx server.ServerCtx) {

	s.RPCServlets.Destroy(serverCtx)
}

func (s *WebappServlets) Handshake(servletCtx server.ServletCtx, ctx *fasthttp.RequestCtx) (*fasthttp.ResponseHeader, error) {
	var ok bool

	tokenString := string(ctx.QueryArgs().Peek("authToken"))
	token, err := jwt.Parse(tokenString, func(token *jwt.Token) (interface{}, error) {
		// Don't forget to validate the alg is what you expect:
		if _, ok := token.Method.(*jwt.SigningMethodRSA); !ok {
			return nil, fmt.Errorf("Unexpected signing method: %v", token.Header["alg"])
		}
		// hmacSampleSecret is a []byte containing your secret, e.g. []byte("my_secret_key")
		return s.VerifyKey, nil
	})

	if nil != err {
		return nil, err
	}

	var claims jwt.MapClaims
	if claims, ok = token.Claims.(jwt.MapClaims); !ok || !token.Valid {
		return nil, fmt.Errorf("Token is not valid %v", token)
	}

	userEmail := claims["sub"].(string)
	sessionID := uuid.NewV4().String()

	servletCtx.SetAttribute(og.SessionIDKey, sessionID)
	servletCtx.SetAttribute(og.SessionClientTypeKey, og.MEMBER)
	servletCtx.SetAttribute(og.SessionTargetIDKey, userEmail)

	return nil, nil
}

func (s *WebappServlets) OnConnect(servletCtx server.ServletCtx, conn socket.Conn) {
	s.RPCServlets.OnConnect(servletCtx, conn)

	sessionID := servletCtx.GetAttribute(og.SessionIDKey)
	targetID := servletCtx.GetAttribute(og.SessionTargetIDKey)
	if nil != sessionID && nil != targetID {
		s.connections.Store(sessionID.(string), retainConnection(targetID.(string), servletCtx))
	}
}

func (s *WebappServlets) OnDisconnect(servletCtx server.ServletCtx) {
	s.RPCServlets.OnDisconnect(servletCtx)

	sessionID := servletCtx.GetAttribute(og.SessionIDKey)
	if nil != sessionID {
		s.connections.Delete(sessionID.(string))
	}
}

func (s *WebappServlets) handleSubscribe(serverCtx server.ServerCtx, subscribeChan <-chan *ogs.Message) {
	for {
		select {
		case msg, ok := <-subscribeChan:
			if !ok {
				return
			}
			switch msg.TargetType {
			case ogs.MEMBER:
				for _, targetID := range msg.Targets {
					_connections := s.getMemberConnections(targetID)
					if nil == _connections || 0 == len(_connections) {
						break
					}

					for _, _connection := range _connections {
						_writeChan := _connection.servletCtx.GetAttribute(og.SessionWriteChanKey)
						if nil != _writeChan {
							writeChan := _writeChan.(chan<- []byte)
							writeChan <- msg.Message
						}
					}
				}

			case ogs.MEMBER_SESSION:
				for _, sessionID := range msg.Targets {
					__connection, ok := s.connections.Load(sessionID)
					if !ok {
						logging.Logger().Debugf("Client[%s] is not exist", sessionID)
						break
					}
					_connection := __connection.(*connection)
					_writeChan := _connection.servletCtx.GetAttribute(og.SessionWriteChanKey)
					if nil != _writeChan {
						writeChan := _writeChan.(chan<- []byte)
						writeChan <- msg.Message
					}
				}
			default:
				logging.Logger().Warnf("Subscriber: Unknown TargetType %s", msg.TargetType)
			}
		}
	}
}

func (s *WebappServlets) getMemberConnections(targetID string) []*connection {
	var connections []*connection

	s.connections.Range(func(k, v interface{}) bool {
		_connection := v.(*connection)
		if _connection.targetID == targetID {
			connections = append(connections, _connection)
		}
		return true
	})

	return connections
}

type connection struct {
	targetID   string
	servletCtx server.ServletCtx
}

var connectionPool sync.Pool

func retainConnection(targetID string, servletCtx server.ServletCtx) *connection {
	v := connectionPool.Get()
	var _connection *connection
	if v == nil {
		_connection = &connection{}
	} else {
		_connection = v.(*connection)
	}

	_connection.targetID = targetID
	_connection.servletCtx = servletCtx

	return _connection
}

func releaseConnection(_connection *connection) {
	_connection.targetID = ""
	_connection.servletCtx = nil

	connectionPool.Put(_connection)
}