VerseVDI-Data-Plane/gateway/clipboard.go

package gateway

import (
	"crypto/rand"
	"crypto/sha256"
	"encoding/base64"
	"errors"
	"sync"
	"time"
	"unicode/utf8"

	protocol "git.sechmachine.io.vn/sechmachine/VerseVDI-Protocol/gen/go/protocol"
)

var (
	ErrClipboardDenied = errors.New("clipboard policy denied")
	ErrClipboardRate   = errors.New("clipboard rate limited")
)

const clipboardRetention = time.Minute

type clipboardRecord struct {
	digest    [sha256.Size]byte
	direction string
	at        time.Time
}

type clipboardGate struct {
	policy  protocol.ClipboardPolicy
	now     func() time.Time
	mu      sync.Mutex
	updates []time.Time
	seen    map[string]clipboardRecord
}

func newClipboardGate(policy protocol.ClipboardPolicy, now func() time.Time) (*clipboardGate, error) {
	if err := policy.Validate(); err != nil || now == nil {
		return nil, ErrProviderMalformed
	}
	return &clipboardGate{policy: policy, now: now, seen: make(map[string]clipboardRecord, policy.MaxUpdatesPerMinute)}, nil
}

// ValidateGatewayClipboard applies the authenticated Server policy before any
// clipboard value can reach a provider or Verse client.
func ValidateGatewayClipboard(policy protocol.ClipboardPolicy, value protocol.GatewayClipboardText) error {
	if err := policy.Validate(); err != nil || value.Validate() != nil || !utf8.ValidString(value.Text) {
		return ErrProviderMalformed
	}
	if len(value.Text) > int(policy.MaxTextBytes) || len(value.LoopToken) < 16 || len(value.LoopToken) > 128 {
		return ErrProviderMalformed
	}
	if _, err := base64.RawURLEncoding.DecodeString(value.LoopToken); err != nil {
		return ErrProviderMalformed
	}
	switch value.Direction {
	case "client_to_provider":
		if !policy.ClientToProviderEnabled {
			return ErrClipboardDenied
		}
	case "provider_to_client":
		if !policy.ProviderToClientEnabled {
			return ErrClipboardDenied
		}
	default:
		return ErrProviderMalformed
	}
	return nil
}

func (g *clipboardGate) fromClient(value protocol.GatewayClipboardText) (bool, error) {
	if g == nil {
		return false, ErrClipboardDenied
	}
	if err := ValidateGatewayClipboard(g.policy, value); err != nil {
		return false, err
	}
	now := g.now()
	digest := sha256.Sum256([]byte(value.Text))
	g.mu.Lock()
	defer g.mu.Unlock()
	g.pruneLocked(now)
	if record, ok := g.seen[value.LoopToken]; ok {
		if record.direction == "provider_to_client" && record.digest == digest {
			return true, nil
		}
		return false, ErrClipboardDenied
	}
	if !g.allowUpdateLocked(now) {
		return false, ErrClipboardRate
	}
	g.seen[value.LoopToken] = clipboardRecord{digest: digest, direction: value.Direction, at: now}
	return false, nil
}

func (g *clipboardGate) fromProvider(text string) (protocol.GatewayClipboardText, bool, error) {
	if g == nil || !g.policy.ProviderToClientEnabled {
		return protocol.GatewayClipboardText{}, false, ErrClipboardDenied
	}
	if !utf8.ValidString(text) || len(text) > int(g.policy.MaxTextBytes) {
		return protocol.GatewayClipboardText{}, false, ErrProviderMalformed
	}
	now := g.now()
	digest := sha256.Sum256([]byte(text))
	g.mu.Lock()
	defer g.mu.Unlock()
	g.pruneLocked(now)
	for _, record := range g.seen {
		if record.digest == digest {
			return protocol.GatewayClipboardText{}, true, nil
		}
	}
	if !g.allowUpdateLocked(now) {
		return protocol.GatewayClipboardText{}, false, ErrClipboardRate
	}
	for attempts := 0; attempts < 3; attempts++ {
		var raw [24]byte
		if _, err := rand.Read(raw[:]); err != nil {
			return protocol.GatewayClipboardText{}, false, err
		}
		token := base64.RawURLEncoding.EncodeToString(raw[:])
		if _, exists := g.seen[token]; exists {
			continue
		}
		value := protocol.GatewayClipboardText{Direction: "provider_to_client", Text: text, Encoding: "utf-8", LoopToken: token}
		g.seen[token] = clipboardRecord{digest: digest, direction: value.Direction, at: now}
		return value, false, nil
	}
	return protocol.GatewayClipboardText{}, false, ErrClipboardDenied
}

func (g *clipboardGate) retractClient(value protocol.GatewayClipboardText) {
	if g == nil {
		return
	}
	digest := sha256.Sum256([]byte(value.Text))
	g.mu.Lock()
	defer g.mu.Unlock()
	if record, ok := g.seen[value.LoopToken]; ok && record.direction == "client_to_provider" && record.digest == digest {
		delete(g.seen, value.LoopToken)
	}
}

func (g *clipboardGate) pruneLocked(now time.Time) {
	minimum := now.Add(-clipboardRetention)
	index := 0
	for _, update := range g.updates {
		if update.After(minimum) {
			g.updates[index] = update
			index++
		}
	}
	g.updates = g.updates[:index]
	for token, record := range g.seen {
		if !record.at.After(minimum) {
			delete(g.seen, token)
		}
	}
}

func (g *clipboardGate) allowUpdateLocked(now time.Time) bool {
	if len(g.updates) >= int(g.policy.MaxUpdatesPerMinute) {
		return false
	}
	g.updates = append(g.updates, now)
	return true
}