2020-10-26 23:22:13 +03:00
|
|
|
// Copyright 2020 Security Scorecard Authors
|
|
|
|
//
|
|
|
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
|
|
|
// you may not use this file except in compliance with the License.
|
|
|
|
// You may obtain a copy of the License at
|
|
|
|
//
|
|
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
|
|
|
//
|
|
|
|
// Unless required by applicable law or agreed to in writing, software
|
|
|
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
|
|
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
|
|
// See the License for the specific language governing permissions and
|
|
|
|
// limitations under the License.
|
|
|
|
|
2020-10-09 17:47:59 +03:00
|
|
|
package roundtripper
|
|
|
|
|
|
|
|
import (
|
|
|
|
"context"
|
2020-11-13 20:06:46 +03:00
|
|
|
"log"
|
2020-10-09 17:47:59 +03:00
|
|
|
"net/http"
|
|
|
|
"os"
|
|
|
|
"strconv"
|
2020-12-02 16:59:43 +03:00
|
|
|
"strings"
|
|
|
|
"sync/atomic"
|
2020-10-09 17:47:59 +03:00
|
|
|
"time"
|
|
|
|
|
2020-11-13 20:06:46 +03:00
|
|
|
"github.com/bradleyfalzon/ghinstallation"
|
2021-02-22 20:18:28 +03:00
|
|
|
cache "github.com/naveensrinivasan/httpcache"
|
|
|
|
"github.com/naveensrinivasan/httpcache/diskcache"
|
2021-02-23 01:19:26 +03:00
|
|
|
"github.com/peterbourgon/diskv"
|
2021-02-22 20:18:28 +03:00
|
|
|
"github.com/pkg/errors"
|
2020-10-13 19:29:29 +03:00
|
|
|
"go.uber.org/zap"
|
2020-10-09 17:47:59 +03:00
|
|
|
"golang.org/x/oauth2"
|
|
|
|
)
|
|
|
|
|
2020-11-13 20:06:46 +03:00
|
|
|
const (
|
2021-02-22 20:18:28 +03:00
|
|
|
GithubAuthToken = "GITHUB_AUTH_TOKEN" // #nosec G101
|
|
|
|
GithubAppKeyPath = "GITHUB_APP_KEY_PATH"
|
|
|
|
GithubAppID = "GITHUB_APP_ID"
|
|
|
|
GithubAppInstallationID = "GITHUB_APP_INSTALLATION_ID"
|
|
|
|
UseDiskCache = "USE_DISK_CACHE"
|
|
|
|
DiskCachePath = "DISK_CACHE_PATH"
|
2021-02-23 18:01:49 +03:00
|
|
|
UseBlobCache = "USE_BLOB_CACHE"
|
|
|
|
BucketURL = "BLOB_URL"
|
2020-11-13 20:06:46 +03:00
|
|
|
)
|
2020-10-09 17:47:59 +03:00
|
|
|
|
2021-03-18 00:41:29 +03:00
|
|
|
var counter int64
|
|
|
|
|
2020-10-09 17:47:59 +03:00
|
|
|
// RateLimitRoundTripper is a rate-limit aware http.Transport for Github.
|
|
|
|
type RateLimitRoundTripper struct {
|
2020-10-13 19:29:29 +03:00
|
|
|
Logger *zap.SugaredLogger
|
2020-10-09 17:47:59 +03:00
|
|
|
InnerTransport http.RoundTripper
|
|
|
|
}
|
|
|
|
|
2020-12-02 16:59:43 +03:00
|
|
|
type RoundRobinTokenSource struct {
|
|
|
|
AccessTokens []string
|
2021-03-18 00:41:29 +03:00
|
|
|
log *zap.SugaredLogger
|
2020-12-02 16:59:43 +03:00
|
|
|
}
|
|
|
|
|
|
|
|
func (r *RoundRobinTokenSource) Token() (*oauth2.Token, error) {
|
2021-03-18 00:41:29 +03:00
|
|
|
c := atomic.AddInt64(&counter, 1)
|
|
|
|
// not locking it because it is never modified
|
|
|
|
l := len(r.AccessTokens)
|
|
|
|
index := c % int64(l)
|
2020-12-02 16:59:43 +03:00
|
|
|
return &oauth2.Token{
|
|
|
|
AccessToken: r.AccessTokens[index],
|
|
|
|
}, nil
|
|
|
|
}
|
|
|
|
|
2021-02-22 20:18:28 +03:00
|
|
|
// NewTransport returns a configured http.Transport for use with GitHub.
|
2020-10-13 19:29:29 +03:00
|
|
|
func NewTransport(ctx context.Context, logger *zap.SugaredLogger) http.RoundTripper {
|
2020-10-09 17:47:59 +03:00
|
|
|
// Start with oauth
|
|
|
|
transport := http.DefaultTransport
|
2021-02-22 20:18:28 +03:00
|
|
|
if token := os.Getenv(GithubAuthToken); token != "" {
|
2020-12-02 16:59:43 +03:00
|
|
|
ts := &RoundRobinTokenSource{
|
|
|
|
AccessTokens: strings.Split(token, ","),
|
2021-03-18 00:41:29 +03:00
|
|
|
log: logger,
|
2020-12-02 16:59:43 +03:00
|
|
|
}
|
2020-10-09 17:47:59 +03:00
|
|
|
transport = oauth2.NewClient(ctx, ts).Transport
|
2021-02-22 20:18:28 +03:00
|
|
|
} else if keyPath := os.Getenv(GithubAppKeyPath); keyPath != "" { // Also try a GITHUB_APP
|
|
|
|
appID, err := strconv.Atoi(os.Getenv(GithubAppID))
|
2020-11-13 20:06:46 +03:00
|
|
|
if err != nil {
|
|
|
|
log.Panic(err)
|
|
|
|
}
|
2021-02-22 20:18:28 +03:00
|
|
|
installationID, err := strconv.Atoi(os.Getenv(GithubAppInstallationID))
|
2020-11-13 20:06:46 +03:00
|
|
|
if err != nil {
|
|
|
|
log.Panic(err)
|
|
|
|
}
|
2021-02-22 20:18:28 +03:00
|
|
|
transport, err = ghinstallation.NewKeyFromFile(transport, int64(appID), int64(installationID), keyPath)
|
2020-11-13 20:06:46 +03:00
|
|
|
if err != nil {
|
|
|
|
log.Panic(err)
|
|
|
|
}
|
2020-10-09 17:47:59 +03:00
|
|
|
}
|
|
|
|
|
|
|
|
// Wrap that with the rate limiter
|
|
|
|
rateLimit := &RateLimitRoundTripper{
|
2020-10-13 19:29:29 +03:00
|
|
|
Logger: logger,
|
2020-10-09 17:47:59 +03:00
|
|
|
InnerTransport: transport,
|
|
|
|
}
|
2021-02-23 18:01:49 +03:00
|
|
|
|
|
|
|
// uses blob cache like GCS,S3.
|
|
|
|
if cachePath, useBlob := shouldUseBlobCache(); useBlob {
|
|
|
|
b, e := New(context.Background(), cachePath)
|
|
|
|
if e != nil {
|
|
|
|
log.Panic(e)
|
|
|
|
}
|
|
|
|
|
|
|
|
c := cache.NewTransport(b)
|
|
|
|
c.Transport = rateLimit
|
|
|
|
return c
|
|
|
|
}
|
|
|
|
|
2021-02-22 20:18:28 +03:00
|
|
|
// uses the disk cache
|
|
|
|
if cachePath, useDisk := shouldUseDiskCache(); useDisk {
|
2021-02-23 18:01:49 +03:00
|
|
|
const cacheSize uint64 = 10000 * 1024 * 1024 // 10gb
|
2021-02-23 01:19:26 +03:00
|
|
|
c := cache.NewTransport(diskcache.NewWithDiskv(
|
|
|
|
diskv.New(diskv.Options{BasePath: cachePath, CacheSizeMax: cacheSize})))
|
2021-02-22 20:18:28 +03:00
|
|
|
c.Transport = rateLimit
|
|
|
|
return c
|
2020-10-09 17:47:59 +03:00
|
|
|
}
|
|
|
|
|
2021-02-22 20:18:28 +03:00
|
|
|
// uses memory cache
|
|
|
|
c := cache.NewTransport(cache.NewMemoryCache())
|
|
|
|
c.Transport = rateLimit
|
|
|
|
return c
|
|
|
|
}
|
|
|
|
|
|
|
|
// shouldUseDiskCache checks the env variables USE_DISK_CACHE and DISK_CACHE_PATH to determine if
|
|
|
|
// disk should be used for caching.
|
|
|
|
func shouldUseDiskCache() (string, bool) {
|
|
|
|
if isDiskCache := os.Getenv(UseDiskCache); isDiskCache != "" {
|
|
|
|
if result, err := strconv.ParseBool(isDiskCache); err == nil && result {
|
|
|
|
if cachePath := os.Getenv(DiskCachePath); cachePath != "" {
|
|
|
|
return cachePath, true
|
|
|
|
}
|
|
|
|
}
|
|
|
|
}
|
|
|
|
return "", false
|
2020-10-09 17:47:59 +03:00
|
|
|
}
|
|
|
|
|
2021-02-23 18:01:49 +03:00
|
|
|
// shouldUseBlobCache checks the env variables USE_BLOB_CACHE and BLOB_URL to determine if
|
|
|
|
// blob should be used for caching.
|
|
|
|
func shouldUseBlobCache() (string, bool) {
|
|
|
|
if result, err := strconv.ParseBool(os.Getenv(UseBlobCache)); err == nil && result {
|
|
|
|
if cachePath := os.Getenv(BucketURL); cachePath != "" {
|
|
|
|
return cachePath, true
|
|
|
|
}
|
|
|
|
}
|
|
|
|
return "", false
|
|
|
|
}
|
|
|
|
|
2020-10-09 17:47:59 +03:00
|
|
|
// Roundtrip handles caching and ratelimiting of responses from GitHub.
|
|
|
|
func (gh *RateLimitRoundTripper) RoundTrip(r *http.Request) (*http.Response, error) {
|
|
|
|
resp, err := gh.InnerTransport.RoundTrip(r)
|
|
|
|
if err != nil {
|
2021-02-22 20:18:28 +03:00
|
|
|
return nil, errors.Wrap(err, "error in round trip")
|
2020-10-09 17:47:59 +03:00
|
|
|
}
|
|
|
|
|
|
|
|
rateLimit := resp.Header.Get("X-RateLimit-Remaining")
|
|
|
|
remaining, err := strconv.Atoi(rateLimit)
|
|
|
|
if err != nil {
|
|
|
|
return resp, nil
|
|
|
|
}
|
|
|
|
|
|
|
|
if remaining <= 0 {
|
|
|
|
reset, err := strconv.Atoi(resp.Header.Get("X-RateLimit-Reset"))
|
|
|
|
if err != nil {
|
|
|
|
return resp, nil
|
|
|
|
}
|
|
|
|
|
|
|
|
duration := time.Until(time.Unix(int64(reset), 0))
|
2020-10-13 22:43:12 +03:00
|
|
|
gh.Logger.Warnf("Rate limit exceeded. Waiting %s to retry...", duration)
|
2020-10-09 17:47:59 +03:00
|
|
|
|
|
|
|
// Retry
|
|
|
|
time.Sleep(duration)
|
2020-10-13 19:29:29 +03:00
|
|
|
gh.Logger.Warnf("Rate limit exceeded. Retrying...")
|
2020-10-09 17:47:59 +03:00
|
|
|
return gh.RoundTrip(r)
|
|
|
|
}
|
2021-03-18 00:41:29 +03:00
|
|
|
|
2021-02-22 20:18:28 +03:00
|
|
|
return resp, nil
|
2020-10-09 17:47:59 +03:00
|
|
|
}
|