From 0d5d63c236db6f3a3a6f32775bd12f3aed56dc1d Mon Sep 17 00:00:00 2001
From: wangyan <wangyan@vmware.com>
Date: Thu, 15 Jun 2017 04:05:07 -0700
Subject: [PATCH] notary inteceptor

update etst

update

update

update

update

update

update

update

update

update

update

update

update

update

update

update

update

update
---
 src/ui/proxy/interceptor_test.go | 31 +++++++++++++++++++++++++++++++
 src/ui/proxy/interceptors.go     | 30 +++++++++++++++++++++++++++++-
 2 files changed, 60 insertions(+), 1 deletion(-)

diff --git a/src/ui/proxy/interceptor_test.go b/src/ui/proxy/interceptor_test.go
index d1abcb8df..72206aa9a 100644
--- a/src/ui/proxy/interceptor_test.go
+++ b/src/ui/proxy/interceptor_test.go
@@ -2,10 +2,13 @@ package proxy
 
 import (
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 	"github.com/vmware/harbor/src/common"
+	"github.com/vmware/harbor/src/common/models"
 	notarytest "github.com/vmware/harbor/src/common/utils/notary/test"
 	utilstest "github.com/vmware/harbor/src/common/utils/test"
 	"github.com/vmware/harbor/src/ui/config"
+	"github.com/vmware/harbor/src/ui/projectmanager/pms"
 
 	"net/http"
 	"net/http/httptest"
@@ -17,6 +20,9 @@ var endpoint = "10.117.4.142"
 var notaryServer *httptest.Server
 var adminServer *httptest.Server
 
+var admiralEndpoint = "http://127.0.0.1:8282"
+var token = ""
+
 func TestMain(m *testing.M) {
 	notaryServer = notarytest.NewNotaryServer(endpoint)
 	defer notaryServer.Close()
@@ -95,6 +101,31 @@ func TestEnvPolicyChecker(t *testing.T) {
 	assert.False(vulFlag)
 }
 
+func TestPMSPolicyChecker(t *testing.T) {
+	pm := pms.NewProjectManager(admiralEndpoint, token)
+	name := "project_for_test_get_true"
+	id, err := pm.Create(&models.Project{
+		Name:               name,
+		EnableContentTrust: true,
+	})
+	require.Nil(t, err)
+	defer func(id int64) {
+		if err := pm.Delete(id); err != nil {
+			require.Nil(t, err)
+		}
+	}(id)
+	project, err := pm.Get(id)
+	assert.Nil(t, err)
+	assert.Equal(t, id, project.ProjectID)
+	server, err2 := utilstest.NewAdminserver(nil)
+	if err2 != nil {
+		t.Fatalf("failed to create a mock admin server: %v", err2)
+	}
+	defer server.Close()
+	contentTrustFlag := getPolicyChecker().contentTrustEnabled("project_for_test_get_true")
+	assert.True(t, contentTrustFlag)
+}
+
 func TestMatchNotaryDigest(t *testing.T) {
 	assert := assert.New(t)
 	//The data from common/utils/notary/helper_test.go
diff --git a/src/ui/proxy/interceptors.go b/src/ui/proxy/interceptors.go
index c68e86ee3..536be11f5 100644
--- a/src/ui/proxy/interceptors.go
+++ b/src/ui/proxy/interceptors.go
@@ -5,6 +5,8 @@ import (
 	"github.com/vmware/harbor/src/common/utils/log"
 	"github.com/vmware/harbor/src/common/utils/notary"
 	"github.com/vmware/harbor/src/ui/config"
+	"github.com/vmware/harbor/src/ui/projectmanager"
+	"github.com/vmware/harbor/src/ui/projectmanager/pms"
 
 	"context"
 	"fmt"
@@ -64,8 +66,34 @@ func (ec envPolicyChecker) vulnerableEnabled(name string) bool {
 	return os.Getenv("PROJECT_VULNERABBLE") == "1"
 }
 
-//TODO: integrate with PMS to get project policies
+type pmsPolicyChecker struct {
+	pm projectmanager.ProjectManager
+}
+
+func (pc pmsPolicyChecker) contentTrustEnabled(name string) bool {
+	project, err := pc.pm.Get(name)
+	if err != nil {
+		log.Errorf("Unexpected error when getting the project, error: %v", err)
+		return true
+	}
+	return project.EnableContentTrust
+}
+func (pc pmsPolicyChecker) vulnerableEnabled(name string) bool {
+	return true
+}
+
+// newPMSPolicyChecker returns an instance of an pmsPolicyChecker
+func newPMSPolicyChecker(pm projectmanager.ProjectManager) policyChecker {
+	return &pmsPolicyChecker{
+		pm: pm,
+	}
+}
+
+// TODO: Get project manager with PM factory.
 func getPolicyChecker() policyChecker {
+	if config.WithAdmiral() {
+		return newPMSPolicyChecker(pms.NewProjectManager(config.AdmiralEndpoint(), ""))
+	}
 	return EnvChecker
 }