Respect the hostname of current repository in queries

api/client.go

@@ -11,6 +11,7 @@ import (
 	"regexp"
 	"strings"
 
+	"github.com/cli/cli/internal/ghinstance"
 	"github.com/henvic/httpretty"
 	"github.com/shurcooL/graphql"
 )
@@ -43,25 +44,21 @@ func NewClientFromHTTP(httpClient *http.Client) *Client {
 func AddHeader(name, value string) ClientOption {
 	return func(tr http.RoundTripper) http.RoundTripper {
 		return &funcTripper{roundTrip: func(req *http.Request) (*http.Response, error) {
-			// prevent the token from leaking to non-GitHub hosts
-			// TODO: GHE support
-			if !strings.EqualFold(name, "Authorization") || strings.HasSuffix(req.URL.Hostname(), ".github.com") {
-				req.Header.Add(name, value)
-			}
+			req.Header.Add(name, value)
 			return tr.RoundTrip(req)
 		}}
 	}
 }
 
 // AddHeaderFunc is an AddHeader that gets the string value from a function
-func AddHeaderFunc(name string, value func() string) ClientOption {
+func AddHeaderFunc(name string, getValue func(*http.Request) (string, error)) ClientOption {
 	return func(tr http.RoundTripper) http.RoundTripper {
 		return &funcTripper{roundTrip: func(req *http.Request) (*http.Response, error) {
-			// prevent the token from leaking to non-GitHub hosts
-			// TODO: GHE support
-			if !strings.EqualFold(name, "Authorization") || strings.HasSuffix(req.URL.Hostname(), ".github.com") {
-				req.Header.Add(name, value())
+			value, err := getValue(req)
+			if err != nil {
+				return nil, err
 			}
+			req.Header.Add(name, value)
 			return tr.RoundTrip(req)
 		}}
 	}
@@ -238,14 +235,13 @@ func (c Client) HasScopes(wantedScopes ...string) (bool, string, error) {
 }
 
 // GraphQL performs a GraphQL request and parses the response
-func (c Client) GraphQL(query string, variables map[string]interface{}, data interface{}) error {
-	url := "https://api.github.com/graphql"
+func (c Client) GraphQL(hostname string, query string, variables map[string]interface{}, data interface{}) error {
 	reqBody, err := json.Marshal(map[string]interface{}{"query": query, "variables": variables})
 	if err != nil {
 		return err
 	}
 
-	req, err := http.NewRequest("POST", url, bytes.NewBuffer(reqBody))
+	req, err := http.NewRequest("POST", ghinstance.GraphQLEndpoint(hostname), bytes.NewBuffer(reqBody))
 	if err != nil {
 		return err
 	}
@@ -261,13 +257,13 @@ func (c Client) GraphQL(query string, variables map[string]interface{}, data int
 	return handleResponse(resp, data)
 }
 
-func graphQLClient(h *http.Client) *graphql.Client {
-	return graphql.NewClient("https://api.github.com/graphql", h)
+func graphQLClient(h *http.Client, hostname string) *graphql.Client {
+	return graphql.NewClient(ghinstance.GraphQLEndpoint(hostname), h)
 }
 
 // REST performs a REST request and parses the response.
-func (c Client) REST(method string, p string, body io.Reader, data interface{}) error {
-	url := "https://api.github.com/" + p
+func (c Client) REST(hostname string, method string, p string, body io.Reader, data interface{}) error {
+	url := ghinstance.RESTPrefix(hostname) + p
 	req, err := http.NewRequest(method, url, body)
 	if err != nil {
 		return err

api/client_test.go

@@ -33,7 +33,7 @@ func TestGraphQL(t *testing.T) {
 	}{}
 
 	http.StubResponse(200, bytes.NewBufferString(`{"data":{"viewer":{"login":"hubot"}}}`))
-	err := client.GraphQL("QUERY", vars, &response)
+	err := client.GraphQL("github.com", "QUERY", vars, &response)
 	eq(t, err, nil)
 	eq(t, response.Viewer.Login, "hubot")
 
@@ -55,7 +55,7 @@ func TestGraphQLError(t *testing.T) {
 	  ]
 	}`))
 
-	err := client.GraphQL("", nil, &response)
+	err := client.GraphQL("github.com", "", nil, &response)
 	if err == nil || err.Error() != "GraphQL error: OH NO\nthis is fine" {
 		t.Fatalf("got %q", err.Error())
 	}
@@ -71,7 +71,7 @@ func TestRESTGetDelete(t *testing.T) {
 	http.StubResponse(204, bytes.NewBuffer([]byte{}))
 
 	r := bytes.NewReader([]byte(`{}`))
-	err := client.REST("DELETE", "applications/CLIENTID/grant", r, nil)
+	err := client.REST("github.com", "DELETE", "applications/CLIENTID/grant", r, nil)
 	eq(t, err, nil)
 }
 
@@ -82,7 +82,7 @@ func TestRESTError(t *testing.T) {
 	http.StubResponse(422, bytes.NewBufferString(`{"message": "OH NO"}`))
 
 	var httpErr HTTPError
-	err := client.REST("DELETE", "repos/branch", nil, nil)
+	err := client.REST("github.com", "DELETE", "repos/branch", nil, nil)
 	if err == nil || !errors.As(err, &httpErr) {
 		t.Fatalf("got %v", err)
 	}

api/queries_issue.go

@@ -112,7 +112,7 @@ func IssueCreate(client *Client, repo *Repository, params map[string]interface{}
 		}
 	}{}
 
-	err := client.GraphQL(query, variables, &result)
+	err := client.GraphQL(repo.RepoHost(), query, variables, &result)
 	if err != nil {
 		return nil, err
 	}
@@ -171,7 +171,7 @@ func IssueStatus(client *Client, repo ghrepo.Interface, currentUsername string)
 	}
 
 	var resp response
-	err := client.GraphQL(query, variables, &resp)
+	err := client.GraphQL(repo.RepoHost(), query, variables, &resp)
 	if err != nil {
 		return nil, err
 	}
@@ -270,7 +270,7 @@ func IssueList(client *Client, repo ghrepo.Interface, state string, labels []str
 loop:
 	for {
 		variables["limit"] = pageLimit
-		err := client.GraphQL(query, variables, &response)
+		err := client.GraphQL(repo.RepoHost(), query, variables, &response)
 		if err != nil {
 			return nil, err
 		}
@@ -361,7 +361,7 @@ func IssueByNumber(client *Client, repo ghrepo.Interface, number int) (*Issue, e
 	}
 
 	var resp response
-	err := client.GraphQL(query, variables, &resp)
+	err := client.GraphQL(repo.RepoHost(), query, variables, &resp)
 	if err != nil {
 		return nil, err
 	}
@@ -389,7 +389,7 @@ func IssueClose(client *Client, repo ghrepo.Interface, issue Issue) error {
 		},
 	}
 
-	gql := graphQLClient(client.http)
+	gql := graphQLClient(client.http, repo.RepoHost())
 	err := gql.MutateNamed(context.Background(), "IssueClose", &mutation, variables)
 
 	if err != nil {
@@ -414,7 +414,7 @@ func IssueReopen(client *Client, repo ghrepo.Interface, issue Issue) error {
 		},
 	}
 
-	gql := graphQLClient(client.http)
+	gql := graphQLClient(client.http, repo.RepoHost())
 	err := gql.MutateNamed(context.Background(), "IssueReopen", &mutation, variables)
 
 	return err

api/queries_org.go

@@ -4,6 +4,8 @@ import (
 	"context"
 	"fmt"
 
+	"github.com/cli/cli/internal/ghinstance"
+	"github.com/cli/cli/internal/ghrepo"
 	"github.com/shurcooL/githubv4"
 )
 
@@ -12,7 +14,8 @@ func resolveOrganization(client *Client, orgName string) (string, error) {
 	var response struct {
 		NodeID string `json:"node_id"`
 	}
-	err := client.REST("GET", fmt.Sprintf("users/%s", orgName), nil, &response)
+	// TODO: GHE support
+	err := client.REST(ghinstance.Default(), "GET", fmt.Sprintf("users/%s", orgName), nil, &response)
 	return response.NodeID, err
 }
 
@@ -24,12 +27,13 @@ func resolveOrganizationTeam(client *Client, orgName, teamSlug string) (string,
 			NodeID string `json:"node_id"`
 		}
 	}
-	err := client.REST("GET", fmt.Sprintf("orgs/%s/teams/%s", orgName, teamSlug), nil, &response)
+	// TODO: GHE support
+	err := client.REST(ghinstance.Default(), "GET", fmt.Sprintf("orgs/%s/teams/%s", orgName, teamSlug), nil, &response)
 	return response.Organization.NodeID, response.NodeID, err
 }
 
 // OrganizationProjects fetches all open projects for an organization
-func OrganizationProjects(client *Client, owner string) ([]RepoProject, error) {
+func OrganizationProjects(client *Client, repo ghrepo.Interface) ([]RepoProject, error) {
 	var query struct {
 		Organization struct {
 			Projects struct {
@@ -43,11 +47,11 @@ func OrganizationProjects(client *Client, owner string) ([]RepoProject, error) {
 	}
 
 	variables := map[string]interface{}{
-		"owner":     githubv4.String(owner),
+		"owner":     githubv4.String(repo.RepoOwner()),
 		"endCursor": (*githubv4.String)(nil),
 	}
 
-	gql := graphQLClient(client.http)
+	gql := graphQLClient(client.http, repo.RepoHost())
 
 	var projects []RepoProject
 	for {
@@ -72,7 +76,7 @@ type OrgTeam struct {
 }
 
 // OrganizationTeams fetches all the teams in an organization
-func OrganizationTeams(client *Client, owner string) ([]OrgTeam, error) {
+func OrganizationTeams(client *Client, repo ghrepo.Interface) ([]OrgTeam, error) {
 	var query struct {
 		Organization struct {
 			Teams struct {
@@ -86,11 +90,11 @@ func OrganizationTeams(client *Client, owner string) ([]OrgTeam, error) {
 	}
 
 	variables := map[string]interface{}{
-		"owner":     githubv4.String(owner),
+		"owner":     githubv4.String(repo.RepoOwner()),
 		"endCursor": (*githubv4.String)(nil),
 	}
 
-	gql := graphQLClient(client.http)
+	gql := graphQLClient(client.http, repo.RepoHost())
 
 	var teams []OrgTeam
 	for {