enhanced prepareServer for better readabilty

Author: ShohamBit

pkg/cmd/flags/server/server.go

@@ -10,13 +10,13 @@ import (
 )
 
 const (
-	HTTPServer          = "http"
-	GRPCServer          = "grpc"
-	MetricsEndpointFlag = "metrics"
-	HealthzEndpointFlag = "healthz"
-	PProfEndpointFlag   = "pprof"
-	ListenEndpointFlag  = "address"
-	PyroscopeAgentFlag  = "pyroscope"
+	HTTPServer                 = "http"
+	GRPCServer                 = "grpc"
+	MetricsEndpointFlag        = "metrics"
+	HealthzEndpointFlag        = "healthz"
+	PProfEndpointFlag          = "pprof"
+	ListenEndpointFlag         = "address"
+	PyroscopeAgentEndpointFlag = "pyroscope"
 )
 
 // TODO: this should be extract to be under 'pkg/cmd/flags' once we remove the binary tracee-rules.
@@ -39,60 +39,75 @@ func PrepareServer(serverSlice []string) (*Server, error) {
 		enablePProf     = false
 		enablePyroscope = false
 	)
-	grpcAddr := ""
 	for _, endpoint := range serverSlice {
-		if strings.Contains(endpoint, HTTPServer) {
-			if strings.Contains(endpoint, MetricsEndpointFlag) {
-				if strings.Contains(endpoint, "true") {
+		// split flag http.address or grpc.address for example
+		serverParts := strings.SplitN(endpoint, ".", 2)
+		switch serverParts[0] {
+		//flag http.Xxx
+		case HTTPServer:
+			httpParts := strings.SplitN(serverParts[1], "=", 1)
+			switch httpParts[0] {
+			case ListenEndpointFlag:
+				server.HTTPServer = http.New(httpParts[1])
+			case MetricsEndpointFlag:
+				if strings.Compare(httpParts[1], "true") == 0 {
 					enableMetrics = true
 				}
-			} else if strings.Contains(endpoint, HealthzEndpointFlag) {
-				if strings.Contains(endpoint, "true") {
+			case HealthzEndpointFlag:
+				if strings.Compare(httpParts[1], "true") == 0 {
 					enableHealthz = true
 				}
-			} else if strings.Contains(endpoint, PProfEndpointFlag) {
-				if strings.Contains(endpoint, "true") {
+			case PProfEndpointFlag:
+				if strings.Compare(httpParts[1], "true") == 0 {
 					enablePProf = true
 				}
-			} else if strings.Contains(endpoint, PyroscopeAgentFlag) {
-				if strings.Contains(endpoint, "true") {
+			case PyroscopeAgentEndpointFlag:
+				if strings.Compare(endpoint, "true") == 0 {
 					enablePyroscope = true
 				}
-			} else if strings.Contains(endpoint, ListenEndpointFlag) {
-				server.HTTPServer = http.New(endpoint[len(HTTPServer)+1+len(ListenEndpointFlag):])
-			} else {
+			default:
 				server.HTTPServer = http.New("")
 			}
-		} else if strings.Contains(endpoint, GRPCServer) {
-			if strings.Contains(endpoint, ListenEndpointFlag) {
-				grpcAddr = endpoint[len(GRPCServer)+1+len(ListenEndpointFlag):]
-
-				addrParts := strings.SplitN(grpcAddr, ":", 2)
-				protocol := addrParts[0]
-
-				if protocol != "tcp" && protocol != "unix" {
-					return nil, errfmt.Errorf("grpc supported protocols are tcp or unix. eg: tcp:4466, unix:/tmp/tracee.sock")
-				}
-
-				if len(addrParts) == 1 {
-					if protocol == "tcp" {
-						grpcAddr += ":4466"
-					} else { // protocol == "unix"
-						grpcAddr += ":/var/run/tracee.sock"
+		//flag grpc.Xxx
+		case GRPCServer:
+			grpcParts := strings.SplitN(serverParts[1], "=", 1)
+			switch grpcParts[0] {
+			case ListenEndpointFlag:
+				addressParts := strings.SplitN(grpcParts[1], ":", 2)
+				switch addressParts[0] {
+				case "unix":
+					if len(addressParts) == 1 {
+						addressParts = append(addressParts, "/var/run/tracee.sock")
 					}
-				}
-				// cleanup listen address if needed (unix socket), for example if a panic happened
-				if protocol == "unix" {
-					path := strings.SplitN(grpcAddr, ":", 2)[1]
-					if _, err = os.Stat(path); err == nil {
-						err = os.Remove(path)
+					if _, err = os.Stat(addressParts[1]); err == nil {
+						err = os.Remove(addressParts[1])
 						if err != nil {
-							return nil, errfmt.Errorf("failed to cleanup gRPC listening address (%s): %v", path, err)
+							return nil, errfmt.Errorf("failed to cleanup gRPC listening address (%s): %v", addressParts[1], err)
 						}
 					}
+				case "tcp":
+					if len(addressParts) == 1 {
+						addressParts = append(addressParts, "4466")
+					}
+				default:
+					return nil, errfmt.Errorf("grpc supported protocols are tcp or unix. eg: tcp:4466, unix:/tmp/tracee.sock")
+				}
+				server.GRPCServer, err = grpc.New(addressParts[0], addressParts[1])
+				if err != nil {
+					return nil, err
 				}
-				server.GRPCServer, err = grpc.New("dsa", "dsa") //protocol, grpcAddr)
 
+			default:
+				if _, err = os.Stat("/var/run/tracee.sock"); err == nil {
+					err = os.Remove("/var/run/tracee.sock")
+					if err != nil {
+						return nil, errfmt.Errorf("failed to cleanup gRPC listening address (%s): %v", "/var/run/tracee.sock", err)
+					}
+				}
+				server.GRPCServer, err = grpc.New("unix", "/var/run/tracee.sock")
+				if err != nil {
+					return nil, err
+				}
 			}
 		}
 	}