From a1036c38470a6141c5d7c0678dc13037fd87ce06 Mon Sep 17 00:00:00 2001
From: Fadi Abbud <fadi.abbud@intevation.de>
Date: Mon, 30 May 2022 13:38:29 +0200
Subject: [PATCH] Add 'Rate' config option for download  throttling (Checker)

---
 cmd/csaf_checker/main.go      | 13 +++++-----
 cmd/csaf_checker/processor.go | 47 ++++++++++++++++++++++++++++-------
 2 files changed, 45 insertions(+), 15 deletions(-)

diff --git a/cmd/csaf_checker/main.go b/cmd/csaf_checker/main.go
index a4ae759..41203df 100644
--- a/cmd/csaf_checker/main.go
+++ b/cmd/csaf_checker/main.go
@@ -26,12 +26,13 @@ import (
 var reportHTML string
 
 type options struct {
-	Output     string  `short:"o" long:"output" description:"File name of the generated report" value-name:"REPORT-FILE"`
-	Format     string  `short:"f" long:"format" choice:"json" choice:"html" description:"Format of report" default:"json"`
-	Insecure   bool    `long:"insecure" description:"Do not check TLS certificates from provider"`
-	ClientCert *string `long:"client-cert" description:"TLS client certificate file (PEM encoded data)" value-name:"CERT-FILE"`
-	ClientKey  *string `long:"client-key" description:"TLS client private key file (PEM encoded data)" value-name:"KEY-FILE"`
-	Version    bool    `long:"version" description:"Display version of the binary"`
+	Output     string   `short:"o" long:"output" description:"File name of the generated report" value-name:"REPORT-FILE"`
+	Format     string   `short:"f" long:"format" choice:"json" choice:"html" description:"Format of report" default:"json"`
+	Insecure   bool     `long:"insecure" description:"Do not check TLS certificates from provider"`
+	ClientCert *string  `long:"client-cert" description:"TLS client certificate file (PEM encoded data)" value-name:"CERT-FILE"`
+	ClientKey  *string  `long:"client-key" description:"TLS client private key file (PEM encoded data)" value-name:"KEY-FILE"`
+	Version    bool     `long:"version" description:"Display version of the binary"`
+	Rate       *float64 `long:"rate" short:"t"`
 }
 
 func errCheck(err error) {
diff --git a/cmd/csaf_checker/processor.go b/cmd/csaf_checker/processor.go
index e3b9083..59e9549 100644
--- a/cmd/csaf_checker/processor.go
+++ b/cmd/csaf_checker/processor.go
@@ -11,6 +11,7 @@ package main
 import (
 	"bufio"
 	"bytes"
+	"context"
 	"crypto/sha256"
 	"crypto/sha512"
 	"crypto/tls"
@@ -30,6 +31,7 @@ import (
 	"time"
 
 	"github.com/ProtonMail/gopenpgp/v2/crypto"
+	"golang.org/x/time/rate"
 
 	"github.com/csaf-poc/csaf_distribution/csaf"
 	"github.com/csaf-poc/csaf_distribution/util"
@@ -38,9 +40,22 @@ import (
 // topicMessages stores the collected topicMessages for a specific topic.
 type topicMessages []string
 
+type client interface {
+	Get(url string) (*http.Response, error)
+}
+type limitingClient struct {
+	client
+	limiter *rate.Limiter
+}
+
+func (lc *limitingClient) Get(url string) (*http.Response, error) {
+	lc.limiter.Wait(context.Background())
+	return lc.client.Get(url)
+}
+
 type processor struct {
 	opts   *options
-	client *http.Client
+	client client
 
 	redirects      map[string]string
 	noneTLS        map[string]struct{}
@@ -263,15 +278,14 @@ func (p *processor) checkRedirect(r *http.Request, via []*http.Request) error {
 	return nil
 }
 
-func (p *processor) httpClient() *http.Client {
+func (p *processor) httpClient() client {
 
 	if p.client != nil {
 		return p.client
 	}
 
-	p.client = &http.Client{
-		CheckRedirect: p.checkRedirect,
-	}
+	client := http.Client{}
+	client.CheckRedirect = p.checkRedirect
 	var tlsConfig tls.Config
 	if p.opts.Insecure {
 		tlsConfig.InsecureSkipVerify = true
@@ -283,10 +297,26 @@ func (p *processor) httpClient() *http.Client {
 		}
 		tlsConfig.Certificates = []tls.Certificate{cert}
 	}
-	p.client.Transport = &http.Transport{
+	client.Transport = &http.Transport{
 		TLSClientConfig: &tlsConfig,
 	}
+	p.client = &client
+
+	if p.opts.Rate == nil {
+		return &client
+	}
+
+	var r float64
+	if p.opts.Rate != nil {
+		r = *p.opts.Rate
+
+	}
+	p.client = &limitingClient{
+		client:  &client,
+		limiter: rate.NewLimiter(rate.Limit(r), 1),
+	}
 	return p.client
+
 }
 
 var yearFromURL = regexp.MustCompile(`.*/(\d{4})/[^/]+$`)
@@ -458,7 +488,6 @@ func (p *processor) integrity(
 }
 
 func (p *processor) processROLIEFeed(feed string) error {
-
 	client := p.httpClient()
 	res, err := client.Get(feed)
 	if err != nil {
@@ -531,6 +560,7 @@ func (p *processor) processROLIEFeed(feed string) error {
 // It returns error if fetching/reading the file(s) fails, otherwise nil.
 func (p *processor) checkIndex(base string, mask whereType) error {
 	client := p.httpClient()
+
 	index := base + "/index.txt"
 	p.checkTLS(index)
 
@@ -795,10 +825,10 @@ func (p *processor) locateProviderMetadata(
 ) error {
 
 	client := p.httpClient()
-
 	tryURL := func(url string) (bool, error) {
 		log.Printf("Trying: %v\n", url)
 		res, err := client.Get(url)
+
 		if err != nil || res.StatusCode != http.StatusOK ||
 			res.Header.Get("Content-Type") != "application/json" {
 			// ignore this as it is expected.
@@ -943,7 +973,6 @@ func (p *processor) checkProviderMetadata(domain string) error {
 func (p *processor) checkSecurity(domain string) error {
 
 	client := p.httpClient()
-
 	p.badSecurity.use()
 
 	path := "https://" + domain + "/.well-known/security.txt"