From ee8c0a22045c8b36b563a499c0a51ec80c286e6e Mon Sep 17 00:00:00 2001
From: Erik Brakkee <erik@brakkee.org>
Date: Fri, 17 Jan 2025 21:16:41 +0100
Subject: [PATCH] inferring application ports in case not configure so that
 they can be used for linkerd authorization.

---
 cmd/policygen/cluster.go         | 40 ++++++++++++++++++++++++++------
 cmd/policygen/config.go          | 13 +++++++++--
 cmd/policygen/configvalidator.go | 16 +++++++------
 3 files changed, 53 insertions(+), 16 deletions(-)

diff --git a/cmd/policygen/cluster.go b/cmd/policygen/cluster.go
index 2866076..fcfd3bb 100644
--- a/cmd/policygen/cluster.go
+++ b/cmd/policygen/cluster.go
@@ -2,12 +2,14 @@ package main
 
 import (
 	"context"
+	"fmt"
 	v1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/labels"
 	"k8s.io/client-go/kubernetes"
 	"log"
 	"slices"
+	"strconv"
 )
 
 type Cluster struct {
@@ -93,11 +95,35 @@ func (c *Cluster) IsLinkerdEnabled(application *Application) bool {
 	return ns.Annotations["linkerd.io/inject"] == "enabled"
 }
 
-func (c *Cluster) Ports(application *Application, nameBased bool) []Port {
-	// gather unique ports based on name
-	// or based on number.
-	// Warning: same name different ports
-	//          same port different names.
-	// Can occur if the selector matches multiple replicasets.
-	return nil
+func (c *Cluster) PortNumbers(application *Application) []Port {
+	if !c.IsLinkerdEnabled(application) {
+		return nil
+	}
+	tcpPorts := make(map[int]Port)
+	udpPorts := make(map[int]Port)
+	for _, pod := range c.Pods(application) {
+		for _, container := range pod.Spec.Containers {
+			for _, port := range container.Ports {
+				switch port.Protocol {
+				case "TCP":
+					tcpPorts[int(port.ContainerPort)] = Port{
+						Port:     strconv.Itoa(int(port.ContainerPort)),
+						Protocol: string(port.Protocol),
+					}
+				case "UDP":
+					udpPorts[int(port.ContainerPort)] = Port{
+						Port:     strconv.Itoa(int(port.ContainerPort)),
+						Protocol: string(port.Protocol),
+					}
+
+				default:
+					panic(fmt.Sprintf("Unknown port type for pod %s/%s: %s",
+						pod.Namespace, pod.Name, port.Protocol))
+				}
+			}
+		}
+	}
+	res := MapValues(tcpPorts)
+	res = append(res, MapValues(udpPorts)...)
+	return res
 }
diff --git a/cmd/policygen/config.go b/cmd/policygen/config.go
index 7da91a6..b8c682d 100644
--- a/cmd/policygen/config.go
+++ b/cmd/policygen/config.go
@@ -6,9 +6,11 @@ import (
 	"fmt"
 	"github.com/goccy/go-yaml"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"log"
 	"net"
 	"os"
 	"slices"
+	"strings"
 )
 
 var PREDEFINED_APPS = []string{"apiserver"}
@@ -180,11 +182,18 @@ func (c Config) GetApplication(name string) (*Application, *Network, string) {
 	return nil, nil, ""
 }
 
-func (c *Config) Infer(resolver func(application *Application) []string) {
+func (c *Config) Infer(resolver Resolver) {
 	for _, ns := range c.Namespaces {
 		for _, app := range ns.Applications {
 			if len(app.ServiceAccounts) == 0 {
-				app.ServiceAccounts = resolver(app)
+				app.ServiceAccounts = resolver.ServiceAccounts(app)
+				log.Printf("Inferred service accounts: %s/%s: %v", app.Namespace.Name, app.Name,
+					app.ServiceAccounts)
+			}
+			if len(app.Ports) == 0 && !strings.HasPrefix(ns.Name, "linkerd") {
+				app.Ports = resolver.PortNumbers(app)
+				log.Printf("Inferred ports: %s/%s: %v", app.Namespace.Name, app.Name,
+					app.Ports)
 			}
 		}
 	}
diff --git a/cmd/policygen/configvalidator.go b/cmd/policygen/configvalidator.go
index bdf43d4..c8733a2 100644
--- a/cmd/policygen/configvalidator.go
+++ b/cmd/policygen/configvalidator.go
@@ -6,7 +6,6 @@ import (
 	"iter"
 	"k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
-	"log"
 	"maps"
 	"os"
 	"slices"
@@ -21,6 +20,11 @@ const (
 	Error
 )
 
+type Resolver interface {
+	ServiceAccounts(application *Application) []string
+	PortNumbers(application *Application) []Port
+}
+
 func LogValidationMsg(level ValidationLevel, msg string, v ...any) {
 	fmt.Fprintf(os.Stderr, "NOTICE: "+msg+"\n", v...)
 }
@@ -36,6 +40,9 @@ func IterToSlice[K any](i iter.Seq[K]) []K {
 func MapKeys[K comparable, V any](m map[K]V) []K {
 	return IterToSlice(maps.Keys(m))
 }
+func MapValues[K comparable, V any](m map[K]V) []V {
+	return IterToSlice(maps.Values(m))
+}
 
 func validate(files []string, options *Options) error {
 	clientset, _ := GetKubernetesConnection()
@@ -49,12 +56,7 @@ func validate(files []string, options *Options) error {
 		return err
 	}
 
-	config.Infer(func(application *Application) []string {
-		res := cluster.ServiceAccounts(application)
-		log.Printf("Inferred service accounts: %s/%s: %v", application.Namespace.Name, application.Name,
-			res)
-		return res
-	})
+	config.Infer(cluster)
 
 	// map  applname1 -> appname2 where appname1 is in an open namespace and app2 is in a closed namespace.
 	// Exclusing when 'from' side is a CIDR.