From d991decb5135e22999025cb132f6c3613433d3cd Mon Sep 17 00:00:00 2001
From: Sanskarzz <sanskar.gur@gmail.com>
Date: Thu, 23 May 2024 16:18:55 +0530
Subject: [PATCH] feat: some minor changes added address flag and liveness
 handler

Signed-off-by: Sanskarzz <sanskar.gur@gmail.com>
---
 main.go              |  4 +++-
 pkg/server/server.go | 37 ++++++++++++++++++++++++++++++++++---
 2 files changed, 37 insertions(+), 4 deletions(-)

diff --git a/main.go b/main.go
index 500e5c5..9b6f1fc 100644
--- a/main.go
+++ b/main.go
@@ -9,16 +9,18 @@ import (
 )
 
 var policies []string
+var address string
 
 func init() {
 	serveCmd.Flags().StringSliceVar(&policies, "policy", nil, "Path to kyverno-json policies")
+	serveCmd.Flags().StringVar(&address, "address", ":9000", "Address to listen on")
 }
 
 var serveCmd = &cobra.Command{
 	Use:   "serve",
 	Short: "Start the kyverno-envoy-plugin server",
 	Run: func(cmd *cobra.Command, args []string) {
-		srv := server.NewServers(policies)
+		srv := server.NewServers(policies, address)
 		server.StartServers(srv)
 	},
 }
diff --git a/pkg/server/server.go b/pkg/server/server.go
index 047b225..5aa080c 100644
--- a/pkg/server/server.go
+++ b/pkg/server/server.go
@@ -2,12 +2,14 @@ package server
 
 import (
 	"context"
+	"strings"
 
 	"fmt"
 	"io"
 	"log"
 	"net"
 	"net/http"
+	"net/url"
 	"os"
 	"os/signal"
 	"syscall"
@@ -27,6 +29,7 @@ import (
 
 type extAuthzServerV3 struct {
 	policies []string
+	address  string
 }
 
 type Servers struct {
@@ -35,10 +38,11 @@ type Servers struct {
 	grpcV3     *extAuthzServerV3
 }
 
-func NewServers(policies []string) *Servers {
+func NewServers(policies []string, address string) *Servers {
 	return &Servers{
 		grpcV3: &extAuthzServerV3{
 			policies: policies,
+			address:  address,
 		},
 	}
 }
@@ -89,12 +93,39 @@ func handler(w http.ResponseWriter, r *http.Request) {
 }
 
 func (s *Servers) startGRPCServer(ctx context.Context) {
-	lis, err := net.Listen("tcp", ":9000")
+
+	address := s.grpcV3.address
+	if !strings.Contains(address, "://") {
+		address = "grpc://" + address
+	}
+
+	parsedURL, err := url.Parse(address)
+	if err != nil {
+		log.Fatalf("failed to parse address url: %v", err)
+	}
+
+	var lis net.Listener
+
+	switch parsedURL.Scheme {
+	case "unix":
+		socketPath := parsedURL.Host + parsedURL.Path
+		if strings.HasPrefix(parsedURL.String(), parsedURL.Scheme+"://@") {
+			socketPath = "@" + socketPath
+		} else {
+			os.Remove(socketPath)
+		}
+		lis, err = net.Listen("unix", socketPath)
+	case "grpc":
+		lis, err = net.Listen("tcp", parsedURL.Host)
+	default:
+		err = fmt.Errorf("invalid url schema %q", parsedURL.Scheme)
+	}
+
 	if err != nil {
 		log.Fatalf("failed to listen: %v", err)
 	}
 	s.grpcServer = grpc.NewServer()
-	fmt.Println("Starting GRPC server on Port 9000")
+	log.Printf("Starting GRPC server on %s", s.grpcV3.address)
 
 	authv3.RegisterAuthorizationServer(s.grpcServer, s.grpcV3)