package servlet import ( "crypto/rsa" "fmt" "io/ioutil" "sync" logging "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" uuid "github.com/satori/go.uuid" "github.com/valyala/fasthttp" ) var ( verifyKey *rsa.PublicKey signKey *rsa.PrivateKey ) func init() { signBytes, err := ioutil.ReadFile("overFlow-private.key") if nil != err { panic(err) } signKey, err = jwt.ParseRSAPrivateKeyFromPEM(signBytes) if nil != err { panic(err) } verifyBytes, err := ioutil.ReadFile("overFlow-public.pem") if nil != err { panic(err) } verifyKey, err = jwt.ParseRSAPublicKeyFromPEM(verifyBytes) if nil != err { panic(err) } } type WebappServlet interface { ogrs.RPCServlet } type WebappServlets struct { ogrs.RPCServlets 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 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, ogs.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) }