frostfs-node/cmd/frostfs-node/qos.go

package main

import (
	"bytes"
	"context"

	qosconfig "git.frostfs.info/TrueCloudLab/frostfs-node/cmd/frostfs-node/config/qos"
	"git.frostfs.info/TrueCloudLab/frostfs-node/internal/logs"
	"git.frostfs.info/TrueCloudLab/frostfs-node/internal/qos"
	"git.frostfs.info/TrueCloudLab/frostfs-node/pkg/core/netmap"
	"git.frostfs.info/TrueCloudLab/frostfs-node/pkg/util/logger"
	qosTagging "git.frostfs.info/TrueCloudLab/frostfs-qos/tagging"
	"go.uber.org/zap"
)

type cfgQoSService struct {
	netmapSource        netmap.Source
	logger              *logger.Logger
	allowedCriticalPubs [][]byte
	allowedInternalPubs [][]byte
}

func initQoSService(c *cfg) {
	criticalPubs := qosconfig.CriticalAuthorizedKeys(c.appCfg)
	internalPubs := qosconfig.InternalAuthorizedKeys(c.appCfg)
	rawCriticalPubs := make([][]byte, 0, len(criticalPubs))
	rawInternalPubs := make([][]byte, 0, len(internalPubs))
	for i := range criticalPubs {
		rawCriticalPubs = append(rawCriticalPubs, criticalPubs[i].Bytes())
	}
	for i := range internalPubs {
		rawInternalPubs = append(rawInternalPubs, internalPubs[i].Bytes())
	}

	c.cfgQoSService = cfgQoSService{
		netmapSource:        c.netMapSource,
		logger:              c.log,
		allowedCriticalPubs: rawCriticalPubs,
		allowedInternalPubs: rawInternalPubs,
	}
}

func (s *cfgQoSService) AdjustIncomingTag(ctx context.Context, requestSignPublicKey []byte) context.Context {
	rawTag, defined := qosTagging.IOTagFromContext(ctx)
	if !defined {
		if s.isInternalIOTagPublicKey(ctx, requestSignPublicKey) {
			return qosTagging.ContextWithIOTag(ctx, qos.IOTagInternal.String())
		}
		return qosTagging.ContextWithIOTag(ctx, qos.IOTagClient.String())
	}
	ioTag, err := qos.FromRawString(rawTag)
	if err != nil {
		s.logger.Debug(ctx, logs.FailedToParseIncomingIOTag, zap.Error(err))
		return qosTagging.ContextWithIOTag(ctx, qos.IOTagClient.String())
	}

	switch ioTag {
	case qos.IOTagClient:
		return ctx
	case qos.IOTagCritical:
		for _, pk := range s.allowedCriticalPubs {
			if bytes.Equal(pk, requestSignPublicKey) {
				return ctx
			}
		}
		nm, err := s.netmapSource.GetNetMap(ctx, 0)
		if err != nil {
			s.logger.Debug(ctx, logs.FailedToGetNetmapToAdjustIOTag, zap.Error(err))
			return qosTagging.ContextWithIOTag(ctx, qos.IOTagClient.String())
		}
		for _, node := range nm.Nodes() {
			if bytes.Equal(node.PublicKey(), requestSignPublicKey) {
				return ctx
			}
		}
		s.logger.Debug(ctx, logs.FailedToValidateIncomingIOTag)
		return qosTagging.ContextWithIOTag(ctx, qos.IOTagClient.String())
	case qos.IOTagInternal:
		if s.isInternalIOTagPublicKey(ctx, requestSignPublicKey) {
			return ctx
		}
		s.logger.Debug(ctx, logs.FailedToValidateIncomingIOTag)
		return qosTagging.ContextWithIOTag(ctx, qos.IOTagClient.String())
	default:
		s.logger.Debug(ctx, logs.NotSupportedIncomingIOTagReplacedWithClient, zap.Stringer("io_tag", ioTag))
		return qosTagging.ContextWithIOTag(ctx, qos.IOTagClient.String())
	}
}

func (s *cfgQoSService) isInternalIOTagPublicKey(ctx context.Context, publicKey []byte) bool {
	for _, pk := range s.allowedInternalPubs {
		if bytes.Equal(pk, publicKey) {
			return true
		}
	}
	nm, err := s.netmapSource.GetNetMap(ctx, 0)
	if err != nil {
		s.logger.Debug(ctx, logs.FailedToGetNetmapToAdjustIOTag, zap.Error(err))
		return false
	}
	for _, node := range nm.Nodes() {
		if bytes.Equal(node.PublicKey(), publicKey) {
			return true
		}
	}

	return false
}