probe_gateway_rpc/servlet/auth-servlet.go

package servlet

import (
	"context"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"sync"

	"github.com/valyala/fasthttp"

	"git.loafle.net/commons/logging-go"
	"git.loafle.net/commons/server-go"
	"git.loafle.net/commons/server-go/socket"
	ocnc "git.loafle.net/overflow/commons-go/noauthprobe/constants"
	ocnm "git.loafle.net/overflow/commons-go/noauthprobe/model"
	og "git.loafle.net/overflow/gateway"
	"git.loafle.net/overflow/gateway/external/grpc"
	ogs "git.loafle.net/overflow/gateway/subscribe"
	ogrs "git.loafle.net/overflow/gateway_rpc/servlet"
	"git.loafle.net/overflow/member_gateway_rpc/subscribe"
)

type AuthServlet interface {
	ogrs.RPCServlet
}

type AuthServlets struct {
	ogrs.RPCServlets
	sessions sync.Map
}

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

	return nil
}

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

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

	return nil
}

func (s *AuthServlets) OnStop(serverCtx server.ServerCtx) {
	if err := subscribe.Subscriber.Unsubscribe(ocnc.HTTPEntry_Auth); nil != err {
		logging.Logger().Warn(err)
	}

	s.RPCServlets.OnStop(serverCtx)
}

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

	s.RPCServlets.Destroy(serverCtx)
}

func (s *AuthServlets) Handshake(servletCtx server.ServletCtx, ctx *fasthttp.RequestCtx) (*fasthttp.ResponseHeader, error) {
	bMethod := ctx.Request.Header.Peek(ocnc.HTTPRequestHeaderKey_NoAuthProbe_Method)

	if nil == bMethod {
		return nil, fmt.Errorf("Unexpected noauth probe method: %v", bMethod)
	}

	method := string(bMethod)

	switch method {
	case ocnc.HTTPRequestHeaderValue_NoAuthProbe_Method_Regist:
		bInfo := ctx.Request.Header.Peek(ocnc.HTTPRequestHeaderKey_NoAuthProbe_Info)
		if nil == bInfo {
			return nil, fmt.Errorf("Unexpected Request Header Key NoAuthProbe Info")
		}
		rJSON := make([]byte, base64.StdEncoding.DecodedLen(len(bInfo)))
		_, err := base64.StdEncoding.Decode(rJSON, bInfo)
		if nil != err {
			return nil, fmt.Errorf("Base64 Encoding Error: %s", err.Error())
		}
		grpcCTX := context.Background()
		r, err := grpc.Exec(grpcCTX, "NoAuthProbeService.regist", string(rJSON))

		if nil != err {
			return nil, fmt.Errorf("grpc call Error: %s", err.Error())
		}
		nap := &ocnm.NoAuthProbe{}
		err = json.Unmarshal([]byte(r), nap)
		if nil != err {
			return nil, fmt.Errorf("grpc result unMarshal Error: %s", err.Error())
		}

		extHeader := &fasthttp.ResponseHeader{}
		extHeader.Add(ocnc.HTTPResponseHeaderKey_NoAuthProbe_SetTempProbeKey, nap.TempProbeKey)

		servletCtx.SetAttribute(og.SessionIDKey, nap.TempProbeKey)
		servletCtx.SetAttribute(og.SessionClientTypeKey, og.PROBE)
		servletCtx.SetAttribute(og.SessionTargetIDKey, nap.TempProbeKey)

		return extHeader, nil
	case ocnc.HTTPRequestHeaderValue_NoAuthProbe_Method_Connect:
		bTempProbeKey := ctx.Request.Header.Peek(ocnc.HTTPRequestHeaderKey_NoAuthProbe_TempProbeKey)
		if nil == bTempProbeKey {
			return nil, fmt.Errorf("Noauth probe temp key is not an existing key")
		}
		tempProbeKey := string(bTempProbeKey)

		grpcCTX := context.Background()
		_, err := grpc.Exec(grpcCTX, "NoAuthProbeService.readByTempKey", tempProbeKey)

		if nil != err {
			return nil, fmt.Errorf("grpc result error: %s", err.Error())
		}

		servletCtx.SetAttribute(og.SessionIDKey, tempProbeKey)
		servletCtx.SetAttribute(og.SessionClientTypeKey, og.PROBE)
		servletCtx.SetAttribute(og.SessionTargetIDKey, tempProbeKey)

		return nil, nil

	default:
		return nil, fmt.Errorf("Unexpected noauth probe method: %s", method)
	}
}

func (s *AuthServlets) 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.sessions.Store(sessionID.(string), ogrs.RetainSession(targetID.(string), servletCtx))
	}
}

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

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

func (s *AuthServlets) handleSubscribe(serverCtx server.ServerCtx, subscribeChan <-chan *ogs.Message) {
LOOP:
	for {
		select {
		case msg, ok := <-subscribeChan:
			if !ok {
				return
			}

			switch msg.TargetType {
			case ogs.PROBE:
				sessions := s.getAuthSessions(msg.Targets)
				if nil == sessions || 0 == len(sessions) {
					continue LOOP
				}

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

func (s *AuthServlets) getAuthSessions(targetIDs []string) []*ogrs.Session {
	var sessions []*ogrs.Session

	if nil == targetIDs || 0 == len(targetIDs) {
		return sessions
	}

	for _, targetID := range targetIDs {
		session, ok := s.sessions.Load(targetID)
		if ok {
			sessions = append(sessions, session.(*ogrs.Session))
		}
	}

	return sessions
}