From bc8200e48128365eb99484649919c702c9bbc694 Mon Sep 17 00:00:00 2001
From: d-gol <dejan.golubovic.27@gmail.com>
Date: Sat, 3 Dec 2022 14:38:26 +0100
Subject: [PATCH] Add auth + use constants for labels + cleanup

---
 pkg/new-ui/v1beta1/backend.go | 57 ++++++++++++++++++++---------------
 1 file changed, 33 insertions(+), 24 deletions(-)

diff --git a/pkg/new-ui/v1beta1/backend.go b/pkg/new-ui/v1beta1/backend.go
index 3301a16a00b..3ef4ed44230 100644
--- a/pkg/new-ui/v1beta1/backend.go
+++ b/pkg/new-ui/v1beta1/backend.go
@@ -24,6 +24,7 @@ import (
 	"log"
 	"net/http"
 	"path/filepath"
+	"io"
 
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/credentials/insecure"
@@ -40,8 +41,10 @@ import (
 	apiv1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/types"
-	corev1 "k8s.io/client-go/kubernetes/typed/core/v1"
 	"sigs.k8s.io/controller-runtime/pkg/client/config"
+	common "github.com/kubeflow/katib/pkg/apis/controller/common/v1beta1"
+	mccommon "github.com/kubeflow/katib/pkg/metricscollector/v1beta1/common"
+	"k8s.io/client-go/kubernetes"
 )
 
 func NewKatibUIHandler(dbManagerAddr string) *KatibUIHandler {
@@ -586,11 +589,19 @@ func (k *KatibUIHandler) FetchTrial(w http.ResponseWriter, r *http.Request) {
 
 // FetchTrialLogs fetches logs for a trial in specific namespace.
 func (k *KatibUIHandler) FetchTrialLogs(w http.ResponseWriter, r *http.Request) {
-	log.Printf("Requesting logs")
-
 	trialName := r.URL.Query()["trialName"][0]
 	namespace := r.URL.Query()["namespace"][0]
-	log.Printf("Requesting logs")
+
+	user, err := IsAuthorized(consts.ActionTypeGet, namespace, consts.PluralTrial, "", trialName, trialsv1beta1.SchemeGroupVersion, k.katibClient.GetClient(), r)
+	if user == "" && err != nil {
+		log.Printf("No user provided in kubeflow-userid header.")
+		http.Error(w, err.Error(), http.StatusUnauthorized)
+		return
+	} else if err != nil {
+		log.Printf("The user: %s is not authorized to get trial: %s in namespace: %s \n", user, trialName, namespace)
+		http.Error(w, err.Error(), http.StatusForbidden)
+		return
+	}
 
 	logs, err := getTrialLogs(k, trialName, namespace)
 	if err != nil {
@@ -613,28 +624,26 @@ func (k *KatibUIHandler) FetchTrialLogs(w http.ResponseWriter, r *http.Request)
 
 // GetTrialLogs returns logs of a master Pod for the given job name and namespace
 func getTrialLogs(k *KatibUIHandler, trialName string, namespace string) (string, error) {
-	cfg, err := config.GetConfig()
-	if err != nil {
-		return "", err
-	}
-
-	clientset, err := corev1.NewForConfig(cfg)
-	if err != nil {
-		return "", err
-	}
-
 	trial := &trialsv1beta1.Trial{}
 	if err := k.katibClient.GetClient().Get(context.TODO(), types.NamespacedName{Name: trialName, Namespace: namespace}, trial); err != nil {
 		return "", err
 	}
 
-	selectionLabel := "training.kubeflow.org/job-name=" + trialName
-
+	selectionLabel := consts.LabelTrialName + "=" + trialName
 	for primaryKey, primaryValue := range trial.Spec.PrimaryPodLabels {
 		selectionLabel = selectionLabel + "," + primaryKey + "=" + primaryValue
 	}
 
-	podList, err := clientset.Pods(namespace).List(context.Background(), metav1.ListOptions{LabelSelector: selectionLabel})
+	cfg, err := config.GetConfig()
+	if err != nil {
+		return "", err
+	}
+	clientset, err := kubernetes.NewForConfig(cfg)
+	if err != nil {
+		return "", err
+	}
+
+	podList, err := clientset.CoreV1().Pods(namespace).List(context.Background(), metav1.ListOptions{LabelSelector: selectionLabel})
 	if err != nil {
 		return "", err
 	}
@@ -642,21 +651,21 @@ func getTrialLogs(k *KatibUIHandler, trialName string, namespace string) (string
 	if len(podList.Items) == 0 {
 		message := `Logs for the trial could not be found.
 Was 'retain: true' specified in the Experiment definition?
-An example can be found here: https://github.com/kubeflow/katib/blob/master/examples/v1beta1/argo/argo-workflow.yaml#L33`
+An example can be found here: https://github.com/kubeflow/katib/blob/7bf39225f7235ee4ba6cf285ecc2c455c6471234/examples/v1beta1/argo/argo-workflow.yaml#L33`
 
 		return message, nil
 	}
+	if len(podList.Items) > 1 {
+		return "", errors.New("More than one master replica found")
+	}
 
 	podLogOpts := apiv1.PodLogOptions{}
 	podLogOpts.Container = trial.Spec.PrimaryContainerName
-	for container := range podList.Items[0].Spec.Containers {
-		if podList.Items[0].Spec.Containers[container].Name == "metrics-logger-and-collector" {
-			podLogOpts.Container = "metrics-logger-and-collector"
-			break
-		}
+	if trial.Spec.MetricsCollector.Collector.Kind == common.StdOutCollector {
+		podLogOpts.Container = mccommon.MetricLoggerCollectorContainerName
 	}
 
-	req := clientset.Pods(namespace).GetLogs(podList.Items[0].Name, &podLogOpts)
+	req := clientset.CoreV1().Pods(namespace).GetLogs(podList.Items[0].Name, &podLogOpts)
 	podLogs, err := req.Stream(context.Background())
 	if err != nil {
 		return "", err