ws-handler.go

/*
 * Copyright (c) 2018. Abstrium SAS <team (at) pydio.com>
 * This file is part of Pydio Cells.
 *
 * Pydio Cells is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * Pydio Cells is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with Pydio Cells.  If not, see <http://www.gnu.org/licenses/>.
 *
 * The latest code can be found at <https://pydio.com>.
 */

package websocket

import (
	"context"
	"strings"
	"sync"
	"time"

	json "github.com/pydio/cells/x/jsonx"

	context2 "github.com/pydio/cells/common/utils/context"

	servicecontext "github.com/pydio/cells/common/service/context"

	"github.com/micro/go-micro/metadata"
	"github.com/micro/protobuf/jsonpb"
	"github.com/pydio/melody"
	"go.uber.org/zap"
	"golang.org/x/time/rate"

	"github.com/pydio/cells/common"
	"github.com/pydio/cells/common/auth"
	"github.com/pydio/cells/common/auth/claim"
	"github.com/pydio/cells/common/log"
	"github.com/pydio/cells/common/proto/activity"
	"github.com/pydio/cells/common/proto/idm"
	"github.com/pydio/cells/common/proto/jobs"
	"github.com/pydio/cells/common/proto/tree"
	"github.com/pydio/cells/common/utils/permissions"
	"github.com/pydio/cells/common/views"
)

type WebsocketHandler struct {
	Websocket   *melody.Melody
	EventRouter *views.RouterEventFilter

	batcherLock   *sync.Mutex
	batchers      map[string]*NodeEventsBatcher
	dispatcher    chan *NodeChangeEventWithInfo
	done          chan string
	silentDropper *rate.Limiter
}

func NewWebSocketHandler(serviceCtx context.Context) *WebsocketHandler {
	w := &WebsocketHandler{
		batchers:      make(map[string]*NodeEventsBatcher),
		dispatcher:    make(chan *NodeChangeEventWithInfo),
		done:          make(chan string),
		batcherLock:   &sync.Mutex{},
		silentDropper: rate.NewLimiter(20, 10),
	}
	w.InitHandlers(serviceCtx)
	go func() {
		for {
			select {
			case e := <-w.dispatcher:
				w.BroadcastNodeChangeEvent(context.Background(), e)
			case finished := <-w.done:
				w.batcherLock.Lock()
				delete(w.batchers, finished)
				w.batcherLock.Unlock()
			}
		}
	}()
	return w
}

func (w *WebsocketHandler) InitHandlers(serviceCtx context.Context) {

	w.Websocket = melody.New()
	w.Websocket.Config.MaxMessageSize = 2048

	w.Websocket.HandleError(func(session *melody.Session, i error) {
		if !strings.Contains(i.Error(), "close 1000 (normal)") {
			log.Logger(serviceCtx).Debug("HandleError", zap.Error(i))
		}
		ClearSession(session)
	})

	w.Websocket.HandleClose(func(session *melody.Session, i int, i2 string) error {
		ClearSession(session)
		return nil
	})

	w.Websocket.HandleMessage(func(session *melody.Session, bytes []byte) {

		msg := &Message{}
		e := json.Unmarshal(bytes, msg)
		if e != nil {
			session.CloseWithMsg(NewErrorMessage(e))
			return
		}
		switch msg.Type {
		case MsgSubscribe:

			if msg.JWT == "" {
				session.CloseWithMsg(NewErrorMessageString("empty jwt"))
				log.Logger(serviceCtx).Debug("empty jwt")
				return
			}
			ctx := context.Background()
			verifier := auth.DefaultJWTVerifier()
			_, claims, e := verifier.Verify(ctx, msg.JWT)
			if e != nil {
				log.Logger(serviceCtx).Error("invalid jwt received from websocket connection")
				session.CloseWithMsg(NewErrorMessage(e))
				return
			}
			UpdateSessionFromClaims(session, claims, w.EventRouter.GetClientsPool())

		case MsgUnsubscribe:

			ClearSession(session)

		default:
			return
		}

	})

}

func (w *WebsocketHandler) getBatcherForUuid(uuid string) *NodeEventsBatcher {
	var batcher *NodeEventsBatcher
	w.batcherLock.Lock()
	if b, ok := w.batchers[uuid]; ok && !b.closed {
		batcher = b
	} else {
		batcher = NewEventsBatcher(1*time.Second, uuid, w.dispatcher, w.done)
		w.batchers[uuid] = batcher
	}
	w.batcherLock.Unlock()
	return batcher
}

// HandleNodeChangeEvent listens to NodeChangeEvents and either broadcast them directly, or use NodeEventsBatcher
// to buffer them and flatten them into one.
func (w *WebsocketHandler) HandleNodeChangeEvent(ctx context.Context, event *tree.NodeChangeEvent) error {

	defer func() {
		if e := recover(); e != nil {
			log.Logger(ctx).Info("recovered a panic in WebSocket handler", zap.Any("e", e))
		}
	}()

	switch event.Type {
	case tree.NodeChangeEvent_UPDATE_META, tree.NodeChangeEvent_CREATE, tree.NodeChangeEvent_UPDATE_CONTENT:
		if event.Target != nil {
			batcher := w.getBatcherForUuid(event.Target.Uuid)
			batcher.in <- event
			return nil
		} else {
			e := &NodeChangeEventWithInfo{}
			e.NodeChangeEvent = *event
			return w.BroadcastNodeChangeEvent(ctx, e)
		}
	case tree.NodeChangeEvent_DELETE, tree.NodeChangeEvent_UPDATE_PATH:
		e := &NodeChangeEventWithInfo{}
		e.NodeChangeEvent = *event
		e.refreshTarget = true
		return w.BroadcastNodeChangeEvent(ctx, e)
	case tree.NodeChangeEvent_READ:
		// Ignore READ events
		return nil
	default:
		return nil
	}

}

// BroadcastNodeChangeEvent will browse the currently registered websocket sessions and decide whether to broadcast
// the event or not.
func (w *WebsocketHandler) BroadcastNodeChangeEvent(ctx context.Context, event *NodeChangeEventWithInfo) error {

	jsonMarshaler := &jsonpb.Marshaler{}

	if event.Silent && !w.silentDropper.Allow() {
		//log.Logger(ctx).Warn("Dropping Silent Event")
		return nil
	}

	return w.Websocket.BroadcastFilter([]byte(`"dump"`), func(session *melody.Session) bool {

		var workspaces map[string]*idm.Workspace
		var accessList *permissions.AccessList

		if value, ok := session.Get(SessionWorkspacesKey); !ok || value == nil {
			return false
		} else {
			workspaces = value.(map[string]*idm.Workspace)
		}

		if value, ok := session.Get(SessionAccessListKey); !ok || value == nil {
			return false
		} else {
			accessList = value.(*permissions.AccessList)
		}

		// Rate-limit events (let Optimistic events always go through)
		if lim, ok := session.Get(SessionLimiterKey); ok && !event.Optimistic {
			limiter := lim.(*rate.Limiter)
			if err := limiter.Wait(ctx); err != nil {
				log.Logger(ctx).Warn("WebSocket: some events were dropped (session rate limiter)")
				return false
			}
		}

		var (
			hasData bool
		)

		claims, _ := session.Get(SessionClaimsKey)
		uName, _ := session.Get(SessionUsernameKey)
		metaCtx := context.Background()
		metaCtx = metadata.NewContext(metaCtx, map[string]string{
			common.PydioContextUserKey: uName.(string),
		})
		if md, o := session.Get(SessionMetaContext); o {
			metaCtx = context2.WithAdditionalMetadata(metaCtx, md.(metadata.Metadata))
		}
		metaCtx = servicecontext.WithServiceName(metaCtx, common.ServiceGatewayNamespace_+common.ServiceWebSocket)
		metaCtx = auth.ToMetadata(metaCtx, claims.(claim.Claims))

		if event.refreshTarget && event.Target != nil {
			if respNode, err := w.EventRouter.GetClientsPool().GetTreeClient().ReadNode(metaCtx, &tree.ReadNodeRequest{Node: event.Target}); err == nil {
				event.Target = respNode.Node
			}
		}

		for wsId, workspace := range workspaces {
			nTarget, t1 := w.EventRouter.WorkspaceCanSeeNode(metaCtx, accessList, workspace, event.Target)
			nSource, t2 := w.EventRouter.WorkspaceCanSeeNode(metaCtx, nil, workspace, event.Source) // Do not deep-check acl on source nodes (deleted!)
			// log.Logger(ctx).Info("Ws can see", zap.String("eType", event.Type.String()), zap.Bool("source", t2), event.Source.ZapPath(), zap.Bool("target", t1), event.Target.ZapPath())
			// Depending on node, broadcast now
			if t1 || t2 {
				eType := event.Type
				if nTarget != nil {
					nTarget.SetMeta("EventWorkspaceId", workspace.UUID)
					nTarget = nTarget.WithoutReservedMetas()
					log.Logger(ctx).Debug("Broadcasting event to this session for workspace", zap.Any("type", event.Type), zap.String("wsId", wsId), zap.Any("path", event.Target.Path))
				}
				if nSource != nil {
					nSource.SetMeta("EventWorkspaceId", workspace.UUID)
					nSource = nSource.WithoutReservedMetas()
				}
				// Eventually update event type if one node is out of scope
				if eType == tree.NodeChangeEvent_UPDATE_PATH {
					if nSource == nil {
						eType = tree.NodeChangeEvent_CREATE
					} else if nTarget == nil {
						eType = tree.NodeChangeEvent_DELETE
					}
				}

				s, _ := jsonMarshaler.MarshalToString(&tree.NodeChangeEvent{
					Type:   eType,
					Target: nTarget,
					Source: nSource,
				})
				data := []byte(s)

				session.Write(data)
				hasData = true
			}
		}

		return hasData
	})

}

// BroadcastTaskChangeEvent listens to tasks events and broadcast them to sessions with the adequate user.
func (w *WebsocketHandler) BroadcastTaskChangeEvent(ctx context.Context, event *jobs.TaskChangeEvent) error {

	if w.Websocket == nil {
		return nil
	}

	taskOwner := event.TaskUpdated.TriggerOwner
	marshaller := jsonpb.Marshaler{}
	message, _ := marshaller.MarshalToString(event)
	return w.Websocket.BroadcastFilter([]byte(message), func(session *melody.Session) bool {
		var isAdmin, o bool
		var v interface{}
		if v, o = session.Get(SessionProfileKey); o && v == common.PydioProfileAdmin {
			isAdmin = true
		}
		value, ok := session.Get(SessionUsernameKey)
		if !ok || value == nil {
			return false
		}
		isOwner := value.(string) == taskOwner || (taskOwner == common.PydioSystemUsername && isAdmin)
		if isOwner {
			log.Logger(ctx).Debug("Should Broadcast Task Event : ", zap.Any("task", event.TaskUpdated), zap.Any("job", event.Job))
		} else {
			log.Logger(ctx).Debug("Owner was " + taskOwner + " while session user was " + value.(string))
		}
		return isOwner
	})

}

// BroadcastIDMChangeEvent listens to ACL events and broadcast them to sessions if the Role, User, or Workspace is concerned
// This triggers a registry reload in the UX (and eventually a change of permissions)
func (w *WebsocketHandler) BroadcastIDMChangeEvent(ctx context.Context, event *idm.ChangeEvent) error {

	marshaller := jsonpb.Marshaler{}
	event.JsonType = "idm"
	message, _ := marshaller.MarshalToString(event)
	return w.Websocket.BroadcastFilter([]byte(message), func(session *melody.Session) bool {

		var checkRoleId string
		var checkUserId string
		var checkWorkspaceId string
		if event.Acl != nil && event.Acl.RoleID != "" && !strings.HasPrefix(event.Acl.Action.Name, "parameter:") && !strings.HasPrefix(event.Acl.Action.Name, "action:") {
			checkRoleId = event.Acl.RoleID
		} else if event.Role != nil {
			checkRoleId = event.Role.Uuid
		} else if event.User != nil {
			checkUserId = event.User.Uuid
		} else if event.Workspace != nil {
			checkWorkspaceId = event.Workspace.UUID
		}

		if checkUserId != "" {
			if val, ok := session.Get(SessionUsernameKey); ok && val != nil {
				return checkUserId == val.(string)
			}
		}

		if checkRoleId != "" {
			if value, ok := session.Get(SessionRolesKey); ok && value != nil {
				roles := value.([]*idm.Role)
				for _, r := range roles {
					if r.Uuid == checkRoleId {
						return true
					}
				}
			}
		}

		if checkWorkspaceId != "" {
			if value, ok := session.Get(SessionWorkspacesKey); ok && value != nil {
				if _, has := value.(map[string]*idm.Workspace)[checkWorkspaceId]; has {
					return true
				}
			}
		}

		return false
	})

}

// BroadcastActivityEvent listens to activities and broadcast them to sessions with the adequate user.
func (w *WebsocketHandler) BroadcastActivityEvent(ctx context.Context, event *activity.PostActivityEvent) error {

	// Only handle "users inbox" events for now
	if event.BoxName != "inbox" && event.OwnerType != activity.OwnerType_USER {
		return nil
	}
	marshaller := jsonpb.Marshaler{}
	event.JsonType = "activity"
	message, _ := marshaller.MarshalToString(event)
	return w.Websocket.BroadcastFilter([]byte(message), func(session *melody.Session) bool {
		if val, ok := session.Get(SessionUsernameKey); ok && val != nil {
			return event.OwnerId == val.(string) && event.Activity.Actor.Id != val.(string)
		}
		return false
	})

}