add muxing to service handlers

Author: faheelsattar

.gitignore

@@ -33,3 +33,4 @@
 .env*
 !.env.example
 
+examples

config/mux_config.go

@@ -1,6 +1,7 @@
 package config
 
 import (
+	"errors"
 	"fmt"
 	"io"
 	"os"
@@ -9,6 +10,21 @@ import (
 	"gopkg.in/yaml.v3"
 )
 
+var (
+	ErrNoFilePathSpecified      = errors.New("no file path specified")
+	ErrNoPolicyDefined          = errors.New("at least one policy must be defined")
+	ErrPolicyNameEmpty          = errors.New("policy name cannot be empty")
+	ErrDuplicatePolicyName      = errors.New("duplicate policy name")
+	ErrPolicyNoRelayers         = errors.New("policy must have at least one relayer")
+	ErrRelayerNameEmpty         = errors.New("relayer name cannot be empty")
+	ErrRelayerURLEmpty          = errors.New("relayer URL cannot be empty")
+	ErrMappingNameEmpty         = errors.New("mapping name cannot be empty")
+	ErrMappingNoPolicySpecified = errors.New("mapping must specify a policy")
+	ErrMappingUnknownPolicy     = errors.New("mapping references unknown policy")
+	ErrMappingNoPublicKeyFilter = errors.New("mapping must specify at least one public key filter")
+	ErrPolicyNotFound           = errors.New("policy not found")
+)
+
 type MuxConfig struct {
 	Policies []Policy  `yaml:"policies"`
 	Mappings []Mapping `yaml:"mappings"`
@@ -38,7 +54,7 @@ type Filters struct {
 // LoadMuxConfig loads the muxing configuration from a yaml file
 func LoadMuxConfig(configPath string) (*MuxConfig, error) {
 	if configPath == "" {
-		return nil, nil
+		return nil, ErrNoFilePathSpecified
 	}
 
 	file, err := os.Open(configPath)
@@ -66,31 +82,31 @@ func LoadMuxConfig(configPath string) (*MuxConfig, error) {
 
 func (c *MuxConfig) validate() error {
 	if len(c.Policies) == 0 {
-		return fmt.Errorf("atleast one policy must be defined")
+		return ErrNoPolicyDefined
 	}
 
 	policyNames := make(map[string]bool)
 	// check if policies are valid
 	for _, policy := range c.Policies {
 		if policy.Name == "" {
-			return fmt.Errorf("policy name cant be empty")
+			return ErrPolicyNameEmpty
 		}
 		if policyNames[policy.Name] {
-			return fmt.Errorf("duplicate policy name: %s", policy.Name)
+			return fmt.Errorf("%w: %s", ErrDuplicatePolicyName, policy.Name)
 		}
 		policyNames[policy.Name] = true
 
 		if len(policy.Relayers) == 0 {
-			return fmt.Errorf("policy %s must have atleast one relayer", policy.Name)
+			return fmt.Errorf("policy %s: %w", policy.Name, ErrPolicyNoRelayers)
 		}
 
 		// check for the relayers if valid
 		for _, relayer := range policy.Relayers {
 			if relayer.Name == "" {
-				return fmt.Errorf("relayer name cant be empty in policy %s", policy.Name)
+				return fmt.Errorf("policy %s: %w", policy.Name, ErrRelayerNameEmpty)
 			}
 			if relayer.URL == "" {
-				return fmt.Errorf("relayer url cant be empty for %s in policy %s", relayer.Name, policy.Name)
+				return fmt.Errorf("policy %s, relayer %s: %w", policy.Name, relayer.Name, ErrRelayerURLEmpty)
 			}
 			if _, err := types.NewRelayEntry(relayer.URL); err != nil {
 				return err
@@ -102,16 +118,16 @@ func (c *MuxConfig) validate() error {
 	// also check if they reference the correct policies
 	for _, mapping := range c.Mappings {
 		if mapping.Name == "" {
-			return fmt.Errorf("mapping name cant be empty")
+			return ErrMappingNameEmpty
 		}
 		if mapping.Policy == "" {
-			return fmt.Errorf("mapping %s must specify a policy", mapping.Name)
+			return fmt.Errorf("mapping %s: %w", mapping.Name, ErrMappingNoPolicySpecified)
 		}
 		if !policyNames[mapping.Policy] {
-			return fmt.Errorf("mapping %s references unknown policy: %s", mapping.Name, mapping.Policy)
+			return fmt.Errorf("mapping %s: %w: %s", mapping.Name, ErrMappingUnknownPolicy, mapping.Policy)
 		}
 		if len(mapping.Filters.PublicKeys) == 0 {
-			return fmt.Errorf("mapping %s must specify atleast one public key filter", mapping.Name)
+			return fmt.Errorf("mapping %s: %w", mapping.Name, ErrMappingNoPublicKeyFilter)
 		}
 	}
 
@@ -145,13 +161,13 @@ func (c *MuxConfig) GetRelaysForPolicy(policyName string) ([]types.RelayEntry, e
 			return relays, nil
 		}
 	}
-	return nil, fmt.Errorf("policy not found: %s", policyName)
+	return nil, fmt.Errorf("%w: %s", ErrPolicyNotFound, policyName)
 }
 
 func (c *MuxConfig) GetAllPolicies() []string {
 	policies := make([]string, len(c.Policies))
-	for _, policy := range c.Policies {
-		policies = append(policies, policy.Name)
+	for i, policy := range c.Policies {
+		policies[i] = policy.Name
 	}
 	return policies
 }

go.mod

@@ -58,5 +58,5 @@ require (
 	golang.org/x/crypto v0.37.0 // indirect
 	golang.org/x/sys v0.32.0 // indirect
 	gopkg.in/yaml.v2 v2.4.0 // indirect
-	gopkg.in/yaml.v3 v3.0.1 // indirect
+	gopkg.in/yaml.v3 v3.0.1
 )

server/get_header.go

@@ -68,8 +68,10 @@ func (m *BoostService) getHeader(log *logrus.Entry, slot phase0.Slot, pubkey, pa
 		relays = make(map[BlockHashHex][]types.RelayEntry)
 	)
 
+	relaysForValidator := m.getRelaysForValidator(pubkey)
+
 	// Request a bid from each relay
-	for _, relay := range m.relays {
+	for _, relay := range relaysForValidator {
 		wg.Add(1)
 		go func(relay types.RelayEntry) {
 			defer wg.Done()

server/get_payload.go

@@ -141,8 +141,15 @@ func (m *BoostService) innerGetPayload(log *logrus.Entry, signedBlindedBeaconBlo
 		log.Warn("bid found but no associated relays")
 	}
 
+	relays := m.relays
+	if len(originalBid.relays) > 0 {
+		// substitute to use originalBid relays since they are def going to contain policy based relays
+		// which are the only ones we want to request the payload from and not from all the relays.
+		relays = originalBid.relays
+	}
+
 	// Prepare for requests
-	resultCh := make(chan payloadResult, len(m.relays))
+	resultCh := make(chan payloadResult, len(relays))
 	var received atomic.Bool
 	go func() {
 		// Make sure we receive a response within the timeout
@@ -154,7 +161,7 @@ func (m *BoostService) innerGetPayload(log *logrus.Entry, signedBlindedBeaconBlo
 	requestCtx, requestCtxCancel := context.WithTimeout(context.Background(), m.httpClientGetPayload.Timeout)
 	defer requestCtxCancel()
 
-	for _, relay := range m.relays {
+	for _, relay := range relays {
 		go func(relay types.RelayEntry) {
 			var url string
 			if version == GetPayloadV1 {

server/register_validator.go

@@ -3,25 +3,49 @@ package server
 import (
 	"bytes"
 	"context"
+	"encoding/json"
+	"errors"
 	"fmt"
+	"mime"
 	"net/http"
 
+	builderApiV1 "github.com/attestantio/go-builder-client/api/v1"
 	"github.com/flashbots/mev-boost/server/params"
 	"github.com/flashbots/mev-boost/server/types"
 	"github.com/sirupsen/logrus"
 )
 
+var ErrUnsupportedContentType = errors.New("unsupported content typee")
+
 func (m *BoostService) registerValidator(log *logrus.Entry, regBytes []byte, header http.Header) error {
-	respErrCh := make(chan error, len(m.relays))
+	if m.muxConfig == nil {
+		return m.sendRegistrationsToRelays(log, regBytes, header, m.relays)
+	}
+	pubkeys, err := m.extractValidatorPubkeys(regBytes, header)
+	if err != nil {
+		return m.sendRegistrationsToRelays(log, regBytes, header, m.relays)
+	}
+	relaysToUse := m.getRelaysForValidators(pubkeys)
+
+	log.WithFields(logrus.Fields{
+		"numValidators": len(pubkeys),
+		"numRelays":     len(relaysToUse),
+	}).Debug("sending validator registrations to relevant relays")
+
+	return m.sendRegistrationsToRelays(log, regBytes, header, relaysToUse)
+}
+
+func (m *BoostService) sendRegistrationsToRelays(log *logrus.Entry, regBytes []byte, header http.Header, relays []types.RelayEntry) error {
+	respErrCh := make(chan error, len(relays))
 
 	log.WithFields(logrus.Fields{
 		"timeout":   m.httpClientRegVal.Timeout,
-		"numRelays": len(m.relays),
+		"numRelays": len(relays),
 		"regBytes":  len(regBytes),
 	}).Info("calling registerValidator on relays")
 
 	// Forward request to each relay
-	for _, relay := range m.relays {
+	for _, relay := range relays {
 		go func(relay types.RelayEntry) {
 			// Get the URL for this relay
 			requestURL := relay.GetURI(params.PathRegisterValidator)
@@ -67,7 +91,7 @@ func (m *BoostService) registerValidator(log *logrus.Entry, regBytes []byte, hea
 	}
 
 	// Return OK if any relay responds OK
-	for range m.relays {
+	for range relays {
 		respErr := <-respErrCh
 		if respErr == nil {
 			// Goroutines are independent, so if there are a lot of configured
@@ -82,3 +106,51 @@ func (m *BoostService) registerValidator(log *logrus.Entry, regBytes []byte, hea
 	log.Debug("no relays accepted the registrations")
 	return errNoSuccessfulRelayResponse
 }
+
+func (m *BoostService) extractValidatorPubkeys(regBytes []byte, header http.Header) ([]string, error) {
+	contentType := header.Get("Content-Type")
+	parsedContentType, _, err := mime.ParseMediaType(contentType)
+	if err != nil {
+		parsedContentType = "application/json"
+	}
+
+	var registrations []*builderApiV1.SignedValidatorRegistration
+
+	switch parsedContentType {
+	case "application/json":
+		if err := json.Unmarshal(regBytes, &registrations); err != nil {
+			return nil, err
+		}
+	case "application/octet-stream":
+		var sszRegistrations builderApiV1.SignedValidatorRegistrations
+		if err := sszRegistrations.UnmarshalSSZ(regBytes); err != nil {
+			return nil, err
+		}
+		registrations = sszRegistrations.Registrations
+	default:
+		return nil, ErrUnsupportedContentType
+	}
+	pubkeys := make([]string, len(registrations))
+	for i, reg := range registrations {
+		pubkeys[i] = reg.Message.Pubkey.String()
+	}
+
+	return pubkeys, nil
+}
+
+func (m *BoostService) getRelaysForValidators(pubkeys []string) []types.RelayEntry {
+	relayMap := make(map[string]types.RelayEntry)
+
+	for _, pubkey := range pubkeys {
+		validatorRelays := m.getRelaysForValidator(pubkey)
+		for _, relay := range validatorRelays {
+			relayMap[relay.URL.String()] = relay
+		}
+	}
+	relays := make([]types.RelayEntry, 0, len(relayMap))
+	for _, relay := range relayMap {
+		relays = append(relays, relay)
+	}
+
+	return relays
+}

server/service.go

@@ -54,6 +54,7 @@ type BoostServiceOpts struct {
 	Log                   *logrus.Entry
 	ListenAddr            string
 	Relays                []types.RelayEntry
+	MuxConfig             *config.MuxConfig
 	GenesisForkVersionHex string
 	GenesisTime           uint64
 	RelayCheck            bool
@@ -69,6 +70,7 @@ type BoostServiceOpts struct {
 type BoostService struct {
 	listenAddr  string
 	relays      []types.RelayEntry
+	muxConfig   *config.MuxConfig
 	log         *logrus.Entry
 	srv         *http.Server
 	relayCheck  bool
@@ -102,6 +104,7 @@ func NewBoostService(opts BoostServiceOpts) (*BoostService, error) {
 	return &BoostService{
 		listenAddr:  opts.ListenAddr,
 		relays:      opts.Relays,
+		muxConfig:   opts.MuxConfig,
 		log:         opts.Log,
 		relayCheck:  opts.RelayCheck,
 		relayMinBid: opts.RelayMinBid,
@@ -482,3 +485,20 @@ func (m *BoostService) CheckRelays() int {
 	wg.Wait()
 	return int(numSuccessRequestsToRelay)
 }
+
+func (m *BoostService) getRelaysForValidator(pubkey string) []types.RelayEntry {
+	if m.muxConfig == nil {
+		return m.relays
+	}
+
+	policyName := m.muxConfig.GetPolicyForValidator(pubkey)
+	if policyName == "" {
+		return m.relays
+	}
+
+	policyRelays, err := m.muxConfig.GetRelaysForPolicy(policyName)
+	if err != nil {
+		return m.relays
+	}
+	return policyRelays
+}