ct_hammer: add options for copy-chain generation

trillian/integration/ct_hammer/main.go

@@ -18,32 +18,51 @@ package main
 import (
 	"bytes"
 	"compress/gzip"
+	"context"
+	"crypto/tls"
 	"encoding/base64"
 	"flag"
 	"fmt"
 	"io"
+	"io/ioutil"
 	"math/rand"
 	"net/http"
 	"os"
+	"strings"
 	"sync"
 	"time"
 
 	"github.com/golang/glog"
 	"github.com/golang/protobuf/ptypes"
+	"github.com/google/certificate-transparency-go/client"
 	"github.com/google/certificate-transparency-go/fixchain/ratelimiter"
+	"github.com/google/certificate-transparency-go/jsonclient"
+	"github.com/google/certificate-transparency-go/loglist"
 	"github.com/google/certificate-transparency-go/trillian/ctfe"
 	"github.com/google/certificate-transparency-go/trillian/ctfe/configpb"
 	"github.com/google/certificate-transparency-go/trillian/integration"
+	"github.com/google/certificate-transparency-go/x509util"
 	"github.com/google/trillian/monitoring"
 	"github.com/google/trillian/monitoring/prometheus"
 	"github.com/prometheus/client_golang/prometheus/promhttp"
 )
 
 var (
-	banner              = flag.Bool("banner", true, "Display intro")
-	httpServers         = flag.String("ct_http_servers", "localhost:8092", "Comma-separated list of (assumed interchangeable) servers, each as address:port")
-	testDir             = flag.String("testdata_dir", "testdata", "Name of directory with test data")
-	leafNotAfter        = flag.String("leaf_not_after", "", "Not-After date to use for leaf certs, RFC3339/ISO-8601 format (e.g. 2017-11-26T12:29:19Z)")
+	banner      = flag.Bool("banner", true, "Display intro")
+	httpServers = flag.String("ct_http_servers", "localhost:8092", "Comma-separated list of (assumed interchangeable) servers, each as address:port")
+
+	// Options for synthetic cert generation.
+	testDir      = flag.String("testdata_dir", "testdata", "Name of directory with test data")
+	leafNotAfter = flag.String("leaf_not_after", "", "Not-After date to use for leaf certs, RFC3339/ISO-8601 format (e.g. 2017-11-26T12:29:19Z)")
+	// Options for copied-cert generation.
+	srcLogURI       = flag.String("src_log_uri", "", "URI for source log to copy certificates from")
+	srcPubKey       = flag.String("src_pub_key", "", "Name of file containing source log's public key")
+	srcLogName      = flag.String("src_log_name", "", "Name of source log to copy certificate from  (from --log_list)")
+	logList         = flag.String("log_list", loglist.AllLogListURL, "Location of master log list (URL or filename)")
+	skipHTTPSVerify = flag.Bool("skip_https_verify", false, "Skip verification of HTTPS transport connection to source log")
+	chainBufSize    = flag.Int("buffered_chains", 100, "Number of buffered certificate chains to hold")
+	startIndex      = flag.Int("start_index", 0, "Index of start point in source log to scan from")
+
 	metricsEndpoint     = flag.String("metrics_endpoint", "", "Endpoint for serving metrics; if left empty, metrics will not be exposed")
 	seed                = flag.Int64("seed", -1, "Seed for random number generation")
 	logConfig           = flag.String("log_config", "", "File holding log config in JSON")
@@ -91,10 +110,76 @@ func main() {
 	if err != nil {
 		glog.Exitf("Failed to read log config: %v", err)
 	}
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
 
 	var generatorFactory func(c *configpb.LogConfig) (integration.ChainGenerator, error) // Almost Java-esque...
-	if *testDir != "" {
-		// Retrieve the test data.
+	if len(*srcLogURI) > 0 || len(*srcLogName) > 0 {
+		// Test cert chains will be generated by copying from a source log.
+		var tlsCfg *tls.Config
+		if *skipHTTPSVerify {
+			glog.Warning("Skipping HTTPS connection verification")
+			tlsCfg = &tls.Config{InsecureSkipVerify: *skipHTTPSVerify}
+		}
+		httpClient := &http.Client{
+			Timeout: 10 * time.Second,
+			Transport: &http.Transport{
+				TLSHandshakeTimeout:   30 * time.Second,
+				ResponseHeaderTimeout: 30 * time.Second,
+				MaxIdleConnsPerHost:   10,
+				DisableKeepAlives:     false,
+				MaxIdleConns:          100,
+				IdleConnTimeout:       90 * time.Second,
+				ExpectContinueTimeout: 1 * time.Second,
+				TLSClientConfig:       tlsCfg,
+			},
+		}
+		uri := *srcLogURI
+		var opts jsonclient.Options
+		if *srcPubKey != "" {
+			pubkey, err := ioutil.ReadFile(*srcPubKey)
+			if err != nil {
+				glog.Exit(err)
+			}
+			opts.PublicKey = string(pubkey)
+		}
+		if len(*srcLogName) > 0 {
+			llData, err := x509util.ReadFileOrURL(*logList, httpClient)
+			if err != nil {
+				glog.Exitf("Failed to read log list: %v", err)
+			}
+			ll, err := loglist.NewFromJSON(llData)
+			if err != nil {
+				glog.Exitf("Failed to build log list: %v", err)
+			}
+
+			logs := ll.FindLogByName(*srcLogName)
+			if len(logs) == 0 {
+				glog.Exitf("No log with name like %q found in loglist %q", *srcLogName, *logList)
+			}
+			if len(logs) > 1 {
+				logNames := make([]string, len(logs))
+				for i, log := range logs {
+					logNames[i] = fmt.Sprintf("%q", log.Description)
+				}
+				glog.Exitf("Multiple logs with name like %q found in loglist: %s", *srcLogName, strings.Join(logNames, ","))
+			}
+			uri = "https://" + logs[0].URL
+			if opts.PublicKey == "" {
+				opts.PublicKeyDER = logs[0].Key
+			}
+		}
+		logClient, err := client.New(uri, httpClient, opts)
+		if err != nil {
+			glog.Exitf("Failed to create client for %q: %v", uri, err)
+		}
+		glog.Infof("Testing with certs copied from log at %s starting at index %d", uri, *startIndex)
+		generatorFactory = func(c *configpb.LogConfig) (integration.ChainGenerator, error) {
+			return integration.NewCopyChainGenerator(ctx, logClient, c, *startIndex, *chainBufSize)
+		}
+	} else if *testDir != "" {
+		// Test cert chains will be generated as synthetic certs from a template.
+		// Retrieve the test data holding the template and key.
 		leafChain, err := integration.GetChain(*testDir, "leaf01.chain")
 		if err != nil {
 			glog.Exitf("failed to load certificate: %v", err)
@@ -111,6 +196,7 @@ func main() {
 			}
 		}
 		// Build a synthetic generator for each target log.
+		glog.Infof("Testing with synthetic certs based on data from %s", *testDir)
 		generatorFactory = func(c *configpb.LogConfig) (integration.ChainGenerator, error) {
 			notAfter := notAfterOverride
 			if notAfter.IsZero() {