Team LBAC: Limit access to data source (#78376)

* Team LBAC: Limit access to data source

* Fix tests

pkg/api/datasources.go

@@ -376,9 +376,12 @@ func validateTeamHTTPHeaderJSON(jsonData *simplejson.Json) error {
 		datasourcesLogger.Error("Unable to marshal TeamHTTPHeaders")
 		return errors.New("validation error, invalid format of TeamHTTPHeaders")
 	}
+	if teamHTTPHeadersJSON == nil {
+		return nil
+	}
 	// whitelisting ValidHeaders
 	// each teams headers
-	for _, teamheaders := range teamHTTPHeadersJSON {
+	for _, teamheaders := range teamHTTPHeadersJSON.Headers {
 		for _, header := range teamheaders {
 			if !slices.ContainsFunc(validHeaders, func(v string) bool {
 				return http.CanonicalHeaderKey(v) == http.CanonicalHeaderKey(header.Header)

pkg/api/datasources_test.go

@@ -233,47 +233,52 @@ func TestUpdateDataSourceTeamHTTPHeaders_InvalidJSONData(t *testing.T) {
 	}{
 		{
 			desc: "We should only allow for headers being X-Prom-Label-Policy",
-			data: datasources.TeamHTTPHeaders{tenantID: []datasources.TeamHTTPHeader{
-				{
-					Header: "Authorization",
-					Value:  "foo!=bar",
-				},
-			},
-			},
+			data: datasources.TeamHTTPHeaders{
+				Headers: datasources.TeamHeaders{
+					tenantID: []datasources.TeamHTTPHeader{
+						{
+							Header: "Authorization",
+							Value:  "foo!=bar",
+						},
+					},
+				}},
 			want: 400,
 		},
 		{
 			desc: "Allowed header but no team id",
-			data: datasources.TeamHTTPHeaders{"": []datasources.TeamHTTPHeader{
-				{
-					Header: "X-Prom-Label-Policy",
-					Value:  "foo=bar",
+			data: datasources.TeamHTTPHeaders{
+				Headers: datasources.TeamHeaders{"": []datasources.TeamHTTPHeader{
+					{
+						Header: "X-Prom-Label-Policy",
+						Value:  "foo=bar",
+					},
 				},
-			},
-			},
+				}},
 			want: 400,
 		},
 		{
 			desc: "Allowed team id and header name with invalid header values ",
-			data: datasources.TeamHTTPHeaders{tenantID: []datasources.TeamHTTPHeader{
-				{
-					Header: "X-Prom-Label-Policy",
-					Value:  "Bad value",
+			data: datasources.TeamHTTPHeaders{
+				Headers: datasources.TeamHeaders{tenantID: []datasources.TeamHTTPHeader{
+					{
+						Header: "X-Prom-Label-Policy",
+						Value:  "Bad value",
+					},
 				},
-			},
-			},
+				}},
 			want: 400,
 		},
 		// Complete valid case, with team id, header name and header value
 		{
 			desc: "Allowed header and header values ",
-			data: datasources.TeamHTTPHeaders{tenantID: []datasources.TeamHTTPHeader{
-				{
-					Header: "X-Prom-Label-Policy",
-					Value:  `1234:{ name!="value",foo!~"bar" }`,
+			data: datasources.TeamHTTPHeaders{
+				Headers: datasources.TeamHeaders{tenantID: []datasources.TeamHTTPHeader{
+					{
+						Header: "X-Prom-Label-Policy",
+						Value:  `1234:{ name!="value",foo!~"bar" }`,
+					},
 				},
-			},
-			},
+				}},
 			want: 200,
 		},
 	}

pkg/services/datasources/models.go

@@ -71,7 +71,12 @@ type TeamHTTPHeadersJSONData struct {
 	TeamHTTPHeaders TeamHTTPHeaders `json:"teamHttpHeaders"`
 }
 
-type TeamHTTPHeaders map[string][]TeamHTTPHeader
+type TeamHTTPHeaders struct {
+	Headers        TeamHeaders `json:"headers"`
+	RestrictAccess bool        `json:"restrictAccess"`
+}
+
+type TeamHeaders map[string][]TeamHTTPHeader
 
 type TeamHTTPHeader struct {
 	Header string `json:"header"`
@@ -80,38 +85,43 @@ type TeamHTTPHeader struct {
 
 const DefaultTeamHTTPHeader = "default"
 
-func (ds DataSource) TeamHTTPHeaders() (TeamHTTPHeaders, error) {
+func (ds DataSource) TeamHTTPHeaders() (*TeamHTTPHeaders, error) {
 	return GetTeamHTTPHeaders(ds.JsonData)
 }
 
-func GetTeamHTTPHeaders(jsonData *simplejson.Json) (TeamHTTPHeaders, error) {
-	teamHTTPHeadersJSON := TeamHTTPHeaders{}
-	if jsonData != nil && jsonData.Get("teamHttpHeaders") != nil {
-		jsonData, err := jsonData.Get("teamHttpHeaders").MarshalJSON()
-		if err != nil {
-			return nil, err
-		}
-		err = json.Unmarshal(jsonData, &teamHTTPHeadersJSON)
-		if err != nil {
-			return nil, err
+func GetTeamHTTPHeaders(jsonData *simplejson.Json) (*TeamHTTPHeaders, error) {
+	teamHTTPHeaders := &TeamHTTPHeaders{}
+	if jsonData == nil {
+		return nil, nil
+	}
+	if _, ok := jsonData.CheckGet("teamHttpHeaders"); !ok {
+		return nil, nil
+	}
+
+	teamHTTPHeadersJSON, err := jsonData.Get("teamHttpHeaders").MarshalJSON()
+	if err != nil {
+		return nil, err
+	}
+	err = json.Unmarshal(teamHTTPHeadersJSON, teamHTTPHeaders)
+	if err != nil {
+		return nil, err
+	}
+	for teamID, headers := range teamHTTPHeaders.Headers {
+		if teamID == "" {
+			return nil, errors.New("teamID is missing or empty in teamHttpHeaders")
 		}
-		for teamID, headers := range teamHTTPHeadersJSON {
-			if teamID == "" {
-				return nil, errors.New("teamID is missing or empty in teamHttpHeaders")
-			}
 
-			for _, header := range headers {
-				if header.Header == "" {
-					return nil, errors.New("header name is missing or empty")
-				}
-				if header.Value == "" {
-					return nil, errors.New("header value is missing or empty")
-				}
+		for _, header := range headers {
+			if header.Header == "" {
+				return nil, errors.New("header name is missing or empty")
+			}
+			if header.Value == "" {
+				return nil, errors.New("header value is missing or empty")
 			}
 		}
 	}
 
-	return teamHTTPHeadersJSON, nil
+	return teamHTTPHeaders, nil
 }
 
 // AllowedCookies parses the jsondata.keepCookies and returns a list of

pkg/services/datasources/models_test.go

@@ -63,14 +63,16 @@ func TestTeamHTTPHeaders(t *testing.T) {
 	testCases := []struct {
 		desc  string
 		given string
-		want  TeamHTTPHeaders
+		want  *TeamHTTPHeaders
 	}{
 		{
 			desc:  "Usual json data with teamHttpHeaders",
-			given: `{"teamHttpHeaders": {"101": [{"header": "X-CUSTOM-HEADER", "value": "foo"}]}}`,
-			want: TeamHTTPHeaders{
-				"101": {
-					{Header: "X-CUSTOM-HEADER", Value: "foo"},
+			given: `{"teamHttpHeaders": {"headers": {"101": [{"header": "X-CUSTOM-HEADER", "value": "foo"}]}}}`,
+			want: &TeamHTTPHeaders{
+				Headers: TeamHeaders{
+					"101": {
+						{Header: "X-CUSTOM-HEADER", Value: "foo"},
+					},
 				},
 			},
 		},