Fix review comments

Signed-off-by: Vladislav Byrgazov <[email protected]>

main.go

@@ -34,6 +34,7 @@ import (
 	"os"
 	"os/signal"
 	"path/filepath"
+	"strings"
 	"sync/atomic"
 	"syscall"
 	"time"
@@ -92,7 +93,7 @@ type Config struct {
 	ConnectTo              url.URL           `default:"unix:///var/lib/networkservicemesh/nsm.io.sock" desc:"url to connect to" split_words:"true"`
 	MaxTokenLifetime       time.Duration     `default:"10m" desc:"maximum lifetime of tokens" split_words:"true"`
 	RegistryClientPolicies []string          `default:"etc/nsm/opa/common/.*.rego,etc/nsm/opa/registry/.*.rego,etc/nsm/opa/client/.*.rego" desc:"paths to files and directories that contain registry client policies" split_words:"true"`
-	IPAMPolicy             string            `default:"polite" desc:"defines NSE's IPAM Policy. Possible values: polite, strict. Polite policy accepts any addresses sent by client. Strict policy resets ip_context if any of the client's addresses doesn't match endpoint's CIDR." split_words:"true"`
+	IPAMPolicy             ipamPolicyFunc    `default:"polite" desc:"defines NSE's IPAM Policy. Possible values: polite, strict. Polite policy accepts any addresses sent by client. Strict policy resets ip_context if any of the client's addresses doesn't match endpoint's CIDR." split_words:"true"`
 	ServiceNames           []string          `default:"icmp-responder" desc:"Name of provided services" split_words:"true"`
 	Payload                string            `default:"ETHERNET" desc:"Name of provided service payload" split_words:"true"`
 	Labels                 map[string]string `default:"" desc:"Endpoint labels"`
@@ -109,6 +110,32 @@ type Config struct {
 	PprofListenOn          string            `default:"localhost:6060" desc:"pprof URL to ListenAndServe" split_words:"true"`
 }
 
+type ipamPolicyFunc func([][]*net.IPNet) networkservice.NetworkServiceServer
+
+// Decode takes a string IPAM Policy and returns the IPAM Policy func
+func (f *ipamPolicyFunc) Decode(policy string) error {
+	switch strings.ToLower(policy) {
+	case "strict":
+		*f = func(cidrPrefix [][]*net.IPNet) networkservice.NetworkServiceServer {
+			var ipnetList []*net.IPNet
+			for _, group := range cidrPrefix {
+				ipnetList = append(ipnetList, group...)
+			}
+			return strictipam.NewServer(func(i ...*net.IPNet) networkservice.NetworkServiceServer {
+				return groupipam.NewServer(cidrPrefix)
+			}, ipnetList...)
+		}
+		return nil
+	case "polite":
+		*f = func(cidrPrefix [][]*net.IPNet) networkservice.NetworkServiceServer {
+			return groupipam.NewServer(cidrPrefix)
+		}
+		return nil
+	default:
+		return errors.Errorf("not a valid IPAM Policy: %s", policy)
+	}
+}
+
 // Process prints and processes env to config
 func (c *Config) Process() error {
 	if err := envconfig.Usage("nsm", c); err != nil {
@@ -220,30 +247,13 @@ func main() {
 
 	tokenServer := getSriovTokenServerChainElement(ctx)
 
-	ipnetList := []*net.IPNet{}
-	for _, group := range config.CidrPrefix {
-		ipnetList = append(ipnetList, group...)
-	}
-	var IPAMServer networkservice.NetworkServiceServer
-
-	switch config.IPAMPolicy {
-	case "strict":
-		IPAMServer = strictipam.NewServer(func(i ...*net.IPNet) networkservice.NetworkServiceServer {
-			return groupipam.NewServer(config.CidrPrefix)
-		}, ipnetList...)
-	case "polite":
-		IPAMServer = groupipam.NewServer(config.CidrPrefix)
-	default:
-		logrus.Fatalf("not a valid IPAM Policy: %s", config.IPAMPolicy)
-	}
-
 	responderEndpoint := endpoint.NewServer(ctx,
 		spiffejwt.TokenGeneratorFunc(source, config.MaxTokenLifetime),
 		endpoint.WithName(config.Name),
 		endpoint.WithAuthorizeServer(authorize.NewServer()),
 		endpoint.WithAdditionalFunctionality(
 			onidle.NewServer(ctx, cancel, config.IdleTimeout),
-			IPAMServer,
+			config.IPAMPolicy(config.CidrPrefix),
 			policyroute.NewServer(newPolicyRoutesGetter(ctx, config.PBRConfigPath).Get),
 			mechanisms.NewServer(map[string]networkservice.NetworkServiceServer{
 				kernelmech.MECHANISM: kernel.NewServer(),