<commit_before>// Copyright 2016 ePoxy 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.
//////////////////////////////////////////////////////////////////////////////

// The ePoxy boot server is the first point of contact for managed machines as
// they boot. The boot server serves all client connections over TLS. And, the
// boot server restricts all state-changing requests to administrative users
// (any machine) and managed machines (only itself).
//
// Managed machines progress through three boot stages:
//   stage1) local boot media like an iPXE ROM, or an immutable CD image
//   stage2) a minimal, linux-based network boot environment
//   stage3) the final system image.
//
// Managed machines are treated as stateless. So, the ePoxy boot server acts as
// an external state manager that mediates the transition of successive boot
// stages. Managed machines positively acknowlege every stage transition using
// session ids generated on the first request and known only to the ePoxy boot
// server and the remote machine.
//
// So, if a managed machine acknowleges the final stage successfully, then we
// know that this machine is the same one that first contacted the ePoxy boot
// server.
package main

import (
	"context"
	"fmt"
	"log"
	"net/http"
	"os"

	"github.com/m-lab/go/httpx"

	"cloud.google.com/go/datastore"
	"github.com/gorilla/mux"
	"github.com/m-lab/epoxy/handler"
	"github.com/m-lab/epoxy/metrics"
	"github.com/m-lab/epoxy/storage"
	"github.com/m-lab/go/rtx"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promhttp"
	"golang.org/x/crypto/acme/autocert"
)

var (
	// Environment variables are preferred to flags for deployments in
	// AppEngine and Docker containers. Using environment variables is encouraged
	// for twelve-factor apps -- https://12factor.net/config

	// projectID must be set using the GCLOUD_PROJECT environment variable.
	projectID = os.Getenv("GCLOUD_PROJECT")

	// publicHostname must be set if *not* running in AppEngine. When running in
	// AppEngine publicHostname is set automatically from a combination of
	// GCLOUD_PROJECT and GAE_SERVICE environment variables. When not running in
	// AppEngine, or to override the default in AppEngine, the PUBLIC_HOSTNAME
	// environment variable must be set instead.
	publicHostname = os.Getenv("PUBLIC_HOSTNAME")

	// bindAddress may be set using the LISTEN environment variable. By default,
	// ePoxy listens on all available interfaces.
	bindAddress = os.Getenv("LISTEN")

	// bindPort may be set using the PORT environment variable.
	bindPort = "8080"

	// allowForwardedRequests controls how the ePoxy server evaluates and applies
	// the Host IP whitelist to incoming requests.
	// DEPRECATED.
	allowForwardedRequests = false

	// serverCert and serverKey are the filenames for the iPXE server certificate.
	serverCert = os.Getenv("IPXE_CERT_FILE")
	serverKey  = os.Getenv("IPXE_KEY_FILE")

	// enableLetsEncrypt enables allocation of let's encrypt certs. Should be
	// disabled for testing.
	enableLetsEncrypt = os.Getenv("ENABLE_LETSENCRYPT")

	// Create a unified context and a cancel method for main().
	ctx, cancelCtx = context.WithCancel(context.Background())
)

// init checks the environment for configuration values.
func init() {
	// Only use the automatic public address if PUBLIC_HOSTNAME is not already set.
	if service := os.Getenv("GAE_SERVICE"); service != "" && projectID != "" && publicHostname == "" {
		publicHostname = fmt.Sprintf("%s-dot-%s.appspot.com", service, projectID)
	}
	if port := os.Getenv("PORT"); port != "" {
		bindPort = port
	}
	if os.Getenv("ALLOW_FORWARDED_REQUESTS") == "true" {
		allowForwardedRequests = true
	}
}

// addRoute adds a new handler for a pattern-based URL target to a Gorilla mux.Router.
func addRoute(router *mux.Router, method, pattern string, handler http.Handler) {
	router.Methods(method).Path(pattern).Handler(handler)
}

// newRouter creates and initializes all routes for the ePoxy boot server.
func newRouter(env *handler.Env) *mux.Router {
	router := mux.NewRouter()

	// A health checker for running in Docker or AppEngine.
	addRoute(router, "GET", "/_ah/health", http.HandlerFunc(checkHealth))

	///////////////////////////////////////////////////////////////////////////
	// Boot stage targets.
	//
	// Immediately after boot, a machine unconditionally requests a stage1 target.
	// After that, the machine should sequentially request the stage2, stage3, and
	// report targets in order. As each is requested, the session ID for the
	// previous is invalidated.

	// Stage1 scripts are always the first script fetched by a booting machine.
	// "stage1.ipxe" is the target for ROM-based iPXE clients.
	addRoute(router, "POST", "/v1/boot/{hostname}/stage1.ipxe",
		promhttp.InstrumentHandlerDuration(metrics.RequestDuration,
			http.HandlerFunc(env.GenerateStage1IPXE)))

	// "stage1.json" is the target for native ePoxy clients.
	addRoute(router, "POST", "/v1/boot/{hostname}/stage1.json",
		http.HandlerFunc(env.GenerateStage1JSON))

	// TODO: make the names stage2 and stage3 arbitrary when we need to support
	// the case where not every machine has the same stage2 or stage3.

	// TODO: consider placing stage sequence names under their own subpath, e.g.
	//   /v1/boot/{hostname}/{sessionID}/stage/{stage}"

	// Stage2, stage3, and report targets load after stage1 runs successfully. Stage2
	// and stage3 targets return an epoxy action. The report target returns no content.
	addRoute(router, "POST", "/v1/boot/{hostname}/{sessionID}/stage2",
		http.HandlerFunc(env.GenerateJSONConfig))
	addRoute(router, "POST", "/v1/boot/{hostname}/{sessionID}/stage3",
		http.HandlerFunc(env.GenerateJSONConfig))
	addRoute(router, "POST", "/v1/boot/{hostname}/{sessionID}/report",
		http.HandlerFunc(env.ReceiveReport))

	///////////////////////////////////////////////////////////////////////////
	// Extension targets.
	//
	// Extension operations may be requested at any time during boot. The session
	// is revoked after successful use. Extensions may return any content type
	// supported by the extension service.
	addRoute(router, "POST", "/v1/boot/{hostname}/{sessionID}/extension/{operation}",
		http.HandlerFunc(env.HandleExtension))

	return router
}

// checkHealth reports whether the server is healthy. checkHealth will
// typically be registered as the http.Handler for the path "/_ah/health" when
// running in Docker or AppEngine.
func checkHealth(rw http.ResponseWriter, req *http.Request) {
	fmt.Fprint(rw, "ok")
}

func setupMetricsHandler(dsCfg *storage.DatastoreConfig) *http.ServeMux {
	// Note: we use custom collectors to read directly from datastore rather than
	// instrumenting http handlers because we want to guarantee that metrics are
	// always available, even after an appengine server restart. These metrics will
	// be critical for defining alerts on boot failures.
	prometheus.Register(metrics.NewCollector("epoxy_last_boot", dsCfg))
	prometheus.Register(metrics.NewCollector("epoxy_last_success", dsCfg))
	// Define a custom serve mux for prometheus to listen on a separate port.
	// We listen on a separate port so we can forward this port on the host VM.
	// We cannot forward port 8080 because it is used by AppEngine.
	mux := http.NewServeMux()
	// Assign the default prometheus handler to the standard exporter path.
	mux.Handle("/metrics", promhttp.Handler())
	return mux
}

func setupPXEServer(addr string, r *mux.Router) *http.Server {
	return &http.Server{
		Addr:    addr,
		Handler: r,
	}
}

func setupLetsEncryptServer(addr string, r *mux.Router, hostname string) *http.Server {
	// We will listen on standard TLS port using LetsEncrypt certificates.
	m := &autocert.Manager{
		// Certificates are cached to a local directory.
		Cache: autocert.DirCache("autocert.cache"),
		// The "Let's Encrypt Terms of Service" are accepted automatically.
		Prompt: autocert.AcceptTOS,
		// The ePoxy server will only accept TLS host requests from given hostname.
		HostPolicy: autocert.HostWhitelist(hostname),
	}
	// Server with custom TLS config.
	return &http.Server{
		Addr:      addr,
		Handler:   r,
		TLSConfig: m.TLSConfig(),
	}
}

func startAppEngineServerAsync(addr string, router *mux.Router) {
	ipxeServer := setupPXEServer(addr, router)
	httpx.ListenAndServeAsync(ipxeServer)
}

func startTLSServerAsync(addr string, router *mux.Router, hostname string) {
	tlsServer := setupLetsEncryptServer(addr, router, hostname)
	// Certificates are already configured in the server.TLSConfig.
	httpx.ListenAndServeTLSAsync(tlsServer, "", "")

	// Because we're running LetsEncrypt certificates on the standard TLS port,
	// run the iPXE server on a higher port.
	ipxeServer := setupPXEServer(addr+"0", router)
	if serverCert == "" || serverKey == "" {
		log.Fatalln("WARNING: IPXE_CERT_FILE and IPXE_KEY_FILE were not specified.")
	}
	httpx.ListenAndServeTLSAsync(ipxeServer, serverCert, serverKey)
}

func main() {
	defer cancelCtx()

	if projectID == "" {
		log.Fatalf("Environment variable GCLOUD_PROJECT must specify a project ID for Datastore.")
	}
	if publicHostname == "" {
		log.Fatalf("Environment variable PUBLIC_HOSTNAME must specify a public service name.")
	}

	tlsAddr := fmt.Sprintf("%s:%s", bindAddress, bindPort)
	pxeAddr := tlsAddr
	// addr := fmt.Sprintf("%s:%s", bindAddress, bindPort)
	client, err := datastore.NewClient(ctx, projectID)
	rtx.Must(err, "Failed to create new datastore client")

	dsCfg := storage.NewDatastoreConfig(client)
	env := &handler.Env{
		Config:                 dsCfg,
		ServerAddr:             publicHostname,
		AllowForwardedRequests: allowForwardedRequests,
	}

	go func() {
		mux := setupMetricsHandler(dsCfg)
		log.Fatal(http.ListenAndServe(":9090", mux))
	}()

	router := newRouter(env)
	if service := os.Getenv("GAE_SERVICE"); service != "" {
		// Assume we are running in AppEngine.
		startAppEngineServerAsync(pxeAddr, router)
	} else {
		startTLSServerAsync(tlsAddr, router, publicHostname)
	}

	<-ctx.Done()

}
<commit_msg>Add startMetricsServerAsync and clean up dead logic<commit_after>// Copyright 2016 ePoxy 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.
//////////////////////////////////////////////////////////////////////////////

// The ePoxy boot server is the first point of contact for managed machines as
// they boot. The boot server serves all client connections over TLS. And, the
// boot server restricts all state-changing requests to administrative users
// (any machine) and managed machines (only itself).
//
// Managed machines progress through three boot stages:
//   stage1) local boot media like an iPXE ROM, or an immutable CD image
//   stage2) a minimal, linux-based network boot environment
//   stage3) the final system image.
//
// Managed machines are treated as stateless. So, the ePoxy boot server acts as
// an external state manager that mediates the transition of successive boot
// stages. Managed machines positively acknowlege every stage transition using
// session ids generated on the first request and known only to the ePoxy boot
// server and the remote machine.
//
// So, if a managed machine acknowleges the final stage successfully, then we
// know that this machine is the same one that first contacted the ePoxy boot
// server.
package main

import (
	"context"
	"fmt"
	"log"
	"net/http"
	"os"

	"github.com/m-lab/go/httpx"

	"cloud.google.com/go/datastore"
	"github.com/gorilla/mux"
	"github.com/m-lab/epoxy/handler"
	"github.com/m-lab/epoxy/metrics"
	"github.com/m-lab/epoxy/storage"
	"github.com/m-lab/go/rtx"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promhttp"
	"golang.org/x/crypto/acme/autocert"
)

var (
	// Environment variables are preferred to flags for deployments in
	// AppEngine and Docker containers. Using environment variables is encouraged
	// for twelve-factor apps -- https://12factor.net/config

	// projectID must be set using the GCLOUD_PROJECT environment variable.
	projectID = os.Getenv("GCLOUD_PROJECT")

	// publicHostname must be set if *not* running in AppEngine. When running in
	// AppEngine publicHostname is set automatically from a combination of
	// GCLOUD_PROJECT and GAE_SERVICE environment variables. When not running in
	// AppEngine, or to override the default in AppEngine, the PUBLIC_HOSTNAME
	// environment variable must be set instead.
	publicHostname = os.Getenv("PUBLIC_HOSTNAME")

	// bindAddress may be set using the LISTEN environment variable. By default,
	// ePoxy listens on all available interfaces.
	bindAddress = os.Getenv("LISTEN")

	// bindPort may be set using the PORT environment variable.
	bindPort = "8080"

	// allowForwardedRequests controls how the ePoxy server evaluates and applies
	// the Host IP whitelist to incoming requests.
	// DEPRECATED.
	allowForwardedRequests = false

	// serverCert and serverKey are the filenames for the iPXE server certificate.
	serverCert = os.Getenv("IPXE_CERT_FILE")
	serverKey  = os.Getenv("IPXE_KEY_FILE")

	// Create a unified context and a cancel method for main().
	ctx, cancelCtx = context.WithCancel(context.Background())
)

// init checks the environment for configuration values.
func init() {
	// Only use the automatic public address if PUBLIC_HOSTNAME is not already set.
	if service := os.Getenv("GAE_SERVICE"); service != "" && projectID != "" && publicHostname == "" {
		publicHostname = fmt.Sprintf("%s-dot-%s.appspot.com", service, projectID)
	}
	if port := os.Getenv("PORT"); port != "" {
		bindPort = port
	}
	if os.Getenv("ALLOW_FORWARDED_REQUESTS") == "true" {
		allowForwardedRequests = true
	}
}

// addRoute adds a new handler for a pattern-based URL target to a Gorilla mux.Router.
func addRoute(router *mux.Router, method, pattern string, handler http.Handler) {
	router.Methods(method).Path(pattern).Handler(handler)
}

// newRouter creates and initializes all routes for the ePoxy boot server.
func newRouter(env *handler.Env) *mux.Router {
	router := mux.NewRouter()

	// A health checker for running in Docker or AppEngine.
	addRoute(router, "GET", "/_ah/health", http.HandlerFunc(checkHealth))

	///////////////////////////////////////////////////////////////////////////
	// Boot stage targets.
	//
	// Immediately after boot, a machine unconditionally requests a stage1 target.
	// After that, the machine should sequentially request the stage2, stage3, and
	// report targets in order. As each is requested, the session ID for the
	// previous is invalidated.

	// Stage1 scripts are always the first script fetched by a booting machine.
	// "stage1.ipxe" is the target for ROM-based iPXE clients.
	addRoute(router, "POST", "/v1/boot/{hostname}/stage1.ipxe",
		promhttp.InstrumentHandlerDuration(metrics.RequestDuration,
			http.HandlerFunc(env.GenerateStage1IPXE)))

	// "stage1.json" is the target for native ePoxy clients.
	addRoute(router, "POST", "/v1/boot/{hostname}/stage1.json",
		http.HandlerFunc(env.GenerateStage1JSON))

	// TODO: make the names stage2 and stage3 arbitrary when we need to support
	// the case where not every machine has the same stage2 or stage3.

	// TODO: consider placing stage sequence names under their own subpath, e.g.
	//   /v1/boot/{hostname}/{sessionID}/stage/{stage}"

	// Stage2, stage3, and report targets load after stage1 runs successfully. Stage2
	// and stage3 targets return an epoxy action. The report target returns no content.
	addRoute(router, "POST", "/v1/boot/{hostname}/{sessionID}/stage2",
		http.HandlerFunc(env.GenerateJSONConfig))
	addRoute(router, "POST", "/v1/boot/{hostname}/{sessionID}/stage3",
		http.HandlerFunc(env.GenerateJSONConfig))
	addRoute(router, "POST", "/v1/boot/{hostname}/{sessionID}/report",
		http.HandlerFunc(env.ReceiveReport))

	///////////////////////////////////////////////////////////////////////////
	// Extension targets.
	//
	// Extension operations may be requested at any time during boot. The session
	// is revoked after successful use. Extensions may return any content type
	// supported by the extension service.
	addRoute(router, "POST", "/v1/boot/{hostname}/{sessionID}/extension/{operation}",
		http.HandlerFunc(env.HandleExtension))

	return router
}

// checkHealth reports whether the server is healthy. checkHealth will
// typically be registered as the http.Handler for the path "/_ah/health" when
// running in Docker or AppEngine.
func checkHealth(rw http.ResponseWriter, req *http.Request) {
	fmt.Fprint(rw, "ok")
}

func setupMetricsHandler(dsCfg *storage.DatastoreConfig) *http.ServeMux {
	// Note: we use custom collectors to read directly from datastore rather than
	// instrumenting http handlers because we want to guarantee that metrics are
	// always available, even after an appengine server restart. These metrics will
	// be critical for defining alerts on boot failures.
	prometheus.Register(metrics.NewCollector("epoxy_last_boot", dsCfg))
	prometheus.Register(metrics.NewCollector("epoxy_last_success", dsCfg))
	// Define a custom serve mux for prometheus to listen on a separate port.
	// We listen on a separate port so we can forward this port on the host VM.
	// We cannot forward port 8080 because it is used by AppEngine.
	mux := http.NewServeMux()
	// Assign the default prometheus handler to the standard exporter path.
	mux.Handle("/metrics", promhttp.Handler())
	return mux
}

func setupPXEServer(addr string, r *mux.Router) *http.Server {
	return &http.Server{
		Addr:    addr,
		Handler: r,
	}
}

func setupLetsEncryptServer(addr string, r *mux.Router, hostname string) *http.Server {
	// We will listen on standard TLS port using LetsEncrypt certificates.
	m := &autocert.Manager{
		// Certificates are cached to a local directory.
		Cache: autocert.DirCache("autocert.cache"),
		// The "Let's Encrypt Terms of Service" are accepted automatically.
		Prompt: autocert.AcceptTOS,
		// The ePoxy server will only accept TLS host requests from given hostname.
		HostPolicy: autocert.HostWhitelist(hostname),
	}
	// Server with custom TLS config.
	return &http.Server{
		Addr:      addr,
		Handler:   r,
		TLSConfig: m.TLSConfig(),
	}
}

func startMetricsServerAsync(dsCfg *storage.DatastoreConfig) {
	mux := setupMetricsHandler(dsCfg)
	s := &http.Server{
		Addr:    ":9000",
		Handler: mux,
	}
	httpx.ListenAndServeAsync(s)
}

func startAppEngineServerAsync(addr string, router *mux.Router) {
	// Start the standard PXE server with the default address.
	ipxeServer := setupPXEServer(addr, router)
	httpx.ListenAndServeAsync(ipxeServer)
}

func startTLSServerAsync(addr string, router *mux.Router, hostname string) {
	// Allocate and use LetsEncrypt certificates on given port.
	tlsServer := setupLetsEncryptServer(addr, router, hostname)
	// Certificates are already configured in the server.TLSConfig.
	httpx.ListenAndServeTLSAsync(tlsServer, "", "")

	// Because we're running LetsEncrypt certificates on the given port,
	// run the iPXE server on a higher port.
	ipxeServer := setupPXEServer(addr+"0", router)
	if serverCert == "" || serverKey == "" {
		log.Fatalln("WARNING: IPXE_CERT_FILE and IPXE_KEY_FILE were not specified.")
	}
	httpx.ListenAndServeTLSAsync(ipxeServer, serverCert, serverKey)
}

func main() {
	defer cancelCtx()

	if projectID == "" {
		log.Fatalf("Environment variable GCLOUD_PROJECT must specify a project ID for Datastore.")
	}
	if publicHostname == "" {
		log.Fatalf("Environment variable PUBLIC_HOSTNAME must specify a public service name.")
	}

	tlsAddr := fmt.Sprintf("%s:%s", bindAddress, bindPort)
	pxeAddr := tlsAddr

	client, err := datastore.NewClient(ctx, projectID)
	rtx.Must(err, "Failed to create new datastore client")

	dsCfg := storage.NewDatastoreConfig(client)
	env := &handler.Env{
		Config:                 dsCfg,
		ServerAddr:             publicHostname,
		AllowForwardedRequests: allowForwardedRequests,
	}

	startMetricsServerAsync(dsCfg)
	router := newRouter(env)
	if service := os.Getenv("GAE_SERVICE"); service != "" {
		startAppEngineServerAsync(pxeAddr, router)
	} else {
		startTLSServerAsync(tlsAddr, router, publicHostname)
	}

	<-ctx.Done()

}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"

	"github.com/Sirupsen/logrus"
	"github.com/moul/as-a-service"
)

type Action func([]string) (interface{}, error)
type Actions map[string]Action

func GetManfredTouron(args []string) (interface{}, error) {
	return moul.GetManfredTouron(), nil
}

func main() {
	actions := map[string]Action{
		"manfred-touron": GetManfredTouron,
	}

	action := "manfred-touron"
	ret, err := actions[action](nil)
	if err != nil {
		logrus.Fatalf("Failed to execute %q: %v", action, err)
	}

	out, err := json.MarshalIndent(ret, "", "  ")
	if err != nil {
		logrus.Fatalf("Failed to marshal json: %v", err)
	}
	fmt.Printf("%s\n", out)
}
<commit_msg>Using codegangsta/cli for parsing<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"os"

	"github.com/Sirupsen/logrus"
	"github.com/codegangsta/cli"
	"github.com/moul/as-a-service"
)

type Action func([]string) (interface{}, error)

func GetManfredTouron(args []string) (interface{}, error) {
	return moul.GetManfredTouron(), nil
}

var Actions map[string]Action

func init() {
	Actions = make(map[string]Action)
	Actions["manfred-touron"] = GetManfredTouron
}

func main() {
	app := cli.NewApp()
	app.Name = "moul-as-a-service"
	app.Usage = "moul, but as a service"
	app.Commands = []cli.Command{}

	for action := range Actions {
		command := cli.Command{
			Name:   action,
			Action: CliActionCallback,
		}
		app.Commands = append(app.Commands, command)
	}

	app.Run(os.Args)
}

func CliActionCallback(c *cli.Context) {
	action := c.Command.Name
	ret, err := Actions[action](c.Args())
	if err != nil {
		logrus.Fatalf("Failed to execute %q: %v", action, err)
	}

	out, err := json.MarshalIndent(ret, "", "  ")
	if err != nil {
		logrus.Fatalf("Failed to marshal json: %v", err)
	}
	fmt.Printf("%s\n", out)
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"context"
	"fmt"
	"os"
	"sort"
	"strings"
	"time"

	imagev1 "github.com/openshift/client-go/image/clientset/versioned/typed/image/v1"
	"github.com/openshift/library-go/pkg/image/reference"
	"github.com/openshift/origin/pkg/monitor/monitorapi"
	"github.com/openshift/origin/pkg/test/ginkgo"
	"github.com/openshift/origin/test/extended/util/image"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/util/sets"
	"k8s.io/client-go/rest"
	"k8s.io/klog/v2"
	e2e "k8s.io/kubernetes/test/e2e/framework"
	k8simage "k8s.io/kubernetes/test/utils/image"
)

// OpenShift tests consume images from the cluster, from a number of vetted community locations,
// and from the upstream Kubernetes test suite which may reference images produced by a variety
// of build systems. To better organize and consolidate these images, the utility code within
// Kubernetes and OpenShift that is consulted to find image streams is made remappable - so that
// all images used by the test code can be located in one place. During normal operation the test
// images are read from the mirror, and administrators can choose to copy those images into
// restricted environment image registries and then run the tests against that subset. This also
// allows us to control the process whereby new images are introduced and review those in one spot.
//
// Test code utilizes helpers to get image strings throughout the code base, or is expected to use
// one or more of the images every OpenShift distribution includes in the 'openshift' namespace.
//
// See test/extended/util/image/README.md for a description of the process of adding a new image.

// defaultTestImageMirrorLocation is where all Kube test inputs are sourced.
const defaultTestImageMirrorLocation = "quay.io/openshift/community-e2e-images"

// createImageMirrorForInternalImages returns a list of 'oc image mirror' mappings from source to
// target or returns an error. If mirrored is true the images are assumed to have already been copied
// from their upstream location into our official mirror, in the REPO:TAG format where TAG is a hash
// of the original internal name and the index of the image in the array. Otherwise the mappings will
// be set to mirror the location as defined in the test code into our official mirror, where the target
// TAG is the hash described above.
func createImageMirrorForInternalImages(prefix string, ref reference.DockerImageReference, mirrored bool) ([]string, error) {
	source := ref.Exact()

	initialDefaults := k8simage.GetOriginalImageConfigs()
	exceptions := image.Exceptions.List()
	defaults := map[int]k8simage.Config{}
	for i, config := range initialDefaults {
		for _, exception := range exceptions {
			if strings.Contains(config.GetE2EImage(), exception) {
				continue
			}
			defaults[i] = config
		}
	}
	updated := k8simage.GetMappedImageConfigs(defaults, ref.Exact())

	openshiftDefaults := image.OriginalImages()
	openshiftUpdated := image.GetMappedImages(openshiftDefaults, defaultTestImageMirrorLocation)

	// if we've mirrored, then the source is going to be our repo, not upstream's
	if mirrored {
		baseRef, err := reference.Parse(defaultTestImageMirrorLocation)
		if err != nil {
			return nil, fmt.Errorf("invalid default mirror location: %v", err)
		}

		// calculate the mapping of upstream images by setting defaults to baseRef
		covered := sets.NewString()
		for i, config := range updated {
			defaultConfig := defaults[i]
			pullSpec := config.GetE2EImage()
			if pullSpec == defaultConfig.GetE2EImage() {
				continue
			}
			if covered.Has(pullSpec) {
				continue
			}
			covered.Insert(pullSpec)
			e2eRef, err := reference.Parse(pullSpec)
			if err != nil {
				return nil, fmt.Errorf("invalid test image: %s: %v", pullSpec, err)
			}
			if len(e2eRef.Tag) == 0 {
				return nil, fmt.Errorf("invalid test image: %s: no tag", pullSpec)
			}
			config.SetRegistry(baseRef.Registry)
			config.SetName(baseRef.RepositoryName())
			config.SetVersion(e2eRef.Tag)
			defaults[i] = config
		}

		// calculate the mapping for openshift images by populating openshiftUpdated
		openshiftUpdated = make(map[string]string)
		sourceMappings := image.GetMappedImages(openshiftDefaults, defaultTestImageMirrorLocation)
		targetMappings := image.GetMappedImages(openshiftDefaults, source)

		for from, to := range targetMappings {
			if from == to {
				continue
			}
			if covered.Has(to) {
				continue
			}
			covered.Insert(to)
			from := sourceMappings[from]
			openshiftUpdated[from] = to
		}
	}

	covered := sets.NewString()
	var lines []string
	for i := range updated {
		a, b := defaults[i], updated[i]
		from, to := a.GetE2EImage(), b.GetE2EImage()
		if from == to {
			continue
		}
		if covered.Has(from) {
			continue
		}
		covered.Insert(from)
		lines = append(lines, fmt.Sprintf("%s %s%s", from, prefix, to))
	}

	for from, to := range openshiftUpdated {
		if from == to {
			continue
		}
		if covered.Has(from) {
			continue
		}
		covered.Insert(from)
		lines = append(lines, fmt.Sprintf("%s %s%s", from, prefix, to))
	}

	sort.Strings(lines)
	return lines, nil
}

func verifyImages() error {
	if len(os.Getenv("KUBE_TEST_REPO")) > 0 {
		return fmt.Errorf("KUBE_TEST_REPO may not be specified when this command is run")
	}
	return verifyImagesWithoutEnv()
}

func verifyImagesWithoutEnv() error {
	defaults := k8simage.GetOriginalImageConfigs()

	for originalPullSpec, index := range image.OriginalImages() {
		if index == -1 {
			continue
		}
		existing, ok := defaults[index]
		if !ok {
			return fmt.Errorf("image %q not found in upstream images, must be moved to test/extended/util/image", originalPullSpec)
		}
		if existing.GetE2EImage() != originalPullSpec {
			return fmt.Errorf("image %q defines index %d but is defined upstream as %q, must be fixed in test/extended/util/image", originalPullSpec, index, existing.GetE2EImage())
		}
		mirror := image.LocationFor(originalPullSpec)
		upstreamMirror := k8simage.GetE2EImage(index)
		if mirror != upstreamMirror {
			return fmt.Errorf("image %q defines index %d and mirror %q but is mirrored upstream as %q, must be fixed in test/extended/util/image", originalPullSpec, index, mirror, upstreamMirror)
		}
	}

	return nil
}

// pulledInvalidImages returns a function that checks whether the cluster pulled an image that is
// outside the allowed list of images. The list is defined as a set of static test case images, the
// local cluster registry, any repository referenced by the image streams in the cluster's 'openshift'
// namespace, or the location that input images are cloned from. Only namespaces prefixed with 'e2e-'
// are checked.
func pulledInvalidImages(fromRepository string) ginkgo.JUnitForEventsFunc {
	// static allowed images
	allowedImages := sets.NewString("image/webserver:404")
	allowedPrefixes := sets.NewString(
		"image-registry.openshift-image-registry.svc",
		"gcr.io/k8s-authenticated-test/",
		"gcr.io/authenticated-image-pulling/",
		"invalid.com/",

		// installed alongside OLM and managed externally
		"registry.redhat.io/redhat/community-operator-index",
		"registry.redhat.io/redhat/certified-operator-index",
		"registry.redhat.io/redhat/redhat-marketplace-index",
		"registry.redhat.io/redhat/redhat-operator-index",

		// used by OLM tests
		"registry.redhat.io/amq7/amq-streams-rhel7-operator",
		"registry.redhat.io/amq7/amqstreams-rhel7-operator-metadata",

		// used to test pull secrets against an authenticated registry
		// TODO: will not work for a disconnected test environment and should be emulated by launching
		//   an authenticated registry in a pod on cluster
		"registry.redhat.io/ubi8/nodejs-14:latest",
	)
	if len(fromRepository) > 0 {
		allowedPrefixes.Insert(fromRepository)
	}

	// any image not in the allowed prefixes is considered a failure, as the user
	// may have added a new test image without calling the appropriate helpers
	return func(events monitorapi.Intervals, _ time.Duration, _ *rest.Config) []*ginkgo.JUnitTestCase {
		imageStreamPrefixes, err := imagePrefixesFromNamespaceImageStreams("openshift")
		if err != nil {
			klog.Errorf("Unable to identify image prefixes from the openshift namespace: %v", err)
		}
		allowedPrefixes.Insert(imageStreamPrefixes.UnsortedList()...)

		allowedPrefixes := allowedPrefixes.List()

		var tests []*ginkgo.JUnitTestCase

		pulls := make(map[string]sets.String)
		for _, event := range events {
			// only messages that include a Pulled reason
			if !strings.Contains(" "+event.Message, " reason/Pulled ") {
				continue
			}
			// only look at pull events from an e2e-* namespace
			if !strings.Contains(" "+event.Locator, " ns/e2e-") {
				continue
			}

			parts := strings.Split(event.Message, " ")
			if len(parts) == 0 {
				continue
			}
			image := strings.TrimPrefix(parts[len(parts)-1], "image/")
			if hasAnyStringPrefix(image, allowedPrefixes) || allowedImages.Has(image) {
				continue
			}
			byImage, ok := pulls[image]
			if !ok {
				byImage = sets.NewString()
				pulls[image] = byImage
			}
			byImage.Insert(event.Locator)
		}
		if len(pulls) > 0 {
			images := make([]string, 0, len(pulls))
			for image := range pulls {
				images = append(images, image)
			}
			sort.Strings(images)
			buf := &bytes.Buffer{}
			for _, image := range images {
				fmt.Fprintf(buf, "%s from pods:\n", image)
				for _, locator := range pulls[image].List() {
					fmt.Fprintf(buf, "  %s\n", locator)
				}
			}
			tests = append(tests, &ginkgo.JUnitTestCase{
				Name:      "[sig-arch] Only known images used by tests",
				SystemOut: buf.String(),
				FailureOutput: &ginkgo.FailureOutput{
					Output: fmt.Sprintf("Cluster accessed images that were not mirrored to the testing repository or already part of the cluster, see test/extended/util/image/README.md in the openshift/origin repo:\n\n%s", buf.String()),
				},
			})
		}

		return tests
	}
}

func hasAnyStringPrefix(s string, prefixes []string) bool {
	for _, prefix := range prefixes {
		if strings.HasPrefix(s, prefix) {
			return true
		}
	}
	return false
}

// imagePrefixesFromNamespaceImageStreams identifies all image repositories referenced by
// image streams in the provided namespace and returns them as a set (for both tags and
// digests). This set of prefixes can be used to verify that image references are coming
// from a location the cluster knows about.
func imagePrefixesFromNamespaceImageStreams(ns string) (sets.String, error) {
	clientConfig, err := e2e.LoadConfig(true)
	if err != nil {
		return nil, err
	}
	client, err := imagev1.NewForConfig(clientConfig)
	if err != nil {
		return nil, err
	}
	streams, err := client.ImageStreams(ns).List(context.TODO(), metav1.ListOptions{})
	if err != nil {
		return nil, err
	}

	allowedPrefixes := sets.NewString()
	for _, stream := range streams.Items {
		for _, tag := range stream.Spec.Tags {
			if tag.From == nil || tag.From.Kind != "DockerImage" {
				continue
			}
			ref, err := reference.Parse(tag.From.Name)
			if err != nil {
				continue
			}
			repo := ref.AsRepository().Exact()
			allowedPrefixes.Insert(repo + ":")
			allowedPrefixes.Insert(repo + "@")
		}
		for _, tag := range stream.Status.Tags {
			for _, event := range tag.Items {
				if len(event.DockerImageReference) == 0 {
					continue
				}
				ref, err := reference.Parse(event.DockerImageReference)
				if err != nil {
					continue
				}
				repo := ref.AsRepository().Exact()
				allowedPrefixes.Insert(repo + ":")
				allowedPrefixes.Insert(repo + "@")
			}
		}
	}
	return allowedPrefixes, nil
}
<commit_msg>record only known images test pass so we can get a percentage of passes<commit_after>package main

import (
	"bytes"
	"context"
	"fmt"
	"os"
	"sort"
	"strings"
	"time"

	imagev1 "github.com/openshift/client-go/image/clientset/versioned/typed/image/v1"
	"github.com/openshift/library-go/pkg/image/reference"
	"github.com/openshift/origin/pkg/monitor/monitorapi"
	"github.com/openshift/origin/pkg/test/ginkgo"
	"github.com/openshift/origin/test/extended/util/image"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/util/sets"
	"k8s.io/client-go/rest"
	"k8s.io/klog/v2"
	e2e "k8s.io/kubernetes/test/e2e/framework"
	k8simage "k8s.io/kubernetes/test/utils/image"
)

// OpenShift tests consume images from the cluster, from a number of vetted community locations,
// and from the upstream Kubernetes test suite which may reference images produced by a variety
// of build systems. To better organize and consolidate these images, the utility code within
// Kubernetes and OpenShift that is consulted to find image streams is made remappable - so that
// all images used by the test code can be located in one place. During normal operation the test
// images are read from the mirror, and administrators can choose to copy those images into
// restricted environment image registries and then run the tests against that subset. This also
// allows us to control the process whereby new images are introduced and review those in one spot.
//
// Test code utilizes helpers to get image strings throughout the code base, or is expected to use
// one or more of the images every OpenShift distribution includes in the 'openshift' namespace.
//
// See test/extended/util/image/README.md for a description of the process of adding a new image.

// defaultTestImageMirrorLocation is where all Kube test inputs are sourced.
const defaultTestImageMirrorLocation = "quay.io/openshift/community-e2e-images"

// createImageMirrorForInternalImages returns a list of 'oc image mirror' mappings from source to
// target or returns an error. If mirrored is true the images are assumed to have already been copied
// from their upstream location into our official mirror, in the REPO:TAG format where TAG is a hash
// of the original internal name and the index of the image in the array. Otherwise the mappings will
// be set to mirror the location as defined in the test code into our official mirror, where the target
// TAG is the hash described above.
func createImageMirrorForInternalImages(prefix string, ref reference.DockerImageReference, mirrored bool) ([]string, error) {
	source := ref.Exact()

	initialDefaults := k8simage.GetOriginalImageConfigs()
	exceptions := image.Exceptions.List()
	defaults := map[int]k8simage.Config{}
	for i, config := range initialDefaults {
		for _, exception := range exceptions {
			if strings.Contains(config.GetE2EImage(), exception) {
				continue
			}
			defaults[i] = config
		}
	}
	updated := k8simage.GetMappedImageConfigs(defaults, ref.Exact())

	openshiftDefaults := image.OriginalImages()
	openshiftUpdated := image.GetMappedImages(openshiftDefaults, defaultTestImageMirrorLocation)

	// if we've mirrored, then the source is going to be our repo, not upstream's
	if mirrored {
		baseRef, err := reference.Parse(defaultTestImageMirrorLocation)
		if err != nil {
			return nil, fmt.Errorf("invalid default mirror location: %v", err)
		}

		// calculate the mapping of upstream images by setting defaults to baseRef
		covered := sets.NewString()
		for i, config := range updated {
			defaultConfig := defaults[i]
			pullSpec := config.GetE2EImage()
			if pullSpec == defaultConfig.GetE2EImage() {
				continue
			}
			if covered.Has(pullSpec) {
				continue
			}
			covered.Insert(pullSpec)
			e2eRef, err := reference.Parse(pullSpec)
			if err != nil {
				return nil, fmt.Errorf("invalid test image: %s: %v", pullSpec, err)
			}
			if len(e2eRef.Tag) == 0 {
				return nil, fmt.Errorf("invalid test image: %s: no tag", pullSpec)
			}
			config.SetRegistry(baseRef.Registry)
			config.SetName(baseRef.RepositoryName())
			config.SetVersion(e2eRef.Tag)
			defaults[i] = config
		}

		// calculate the mapping for openshift images by populating openshiftUpdated
		openshiftUpdated = make(map[string]string)
		sourceMappings := image.GetMappedImages(openshiftDefaults, defaultTestImageMirrorLocation)
		targetMappings := image.GetMappedImages(openshiftDefaults, source)

		for from, to := range targetMappings {
			if from == to {
				continue
			}
			if covered.Has(to) {
				continue
			}
			covered.Insert(to)
			from := sourceMappings[from]
			openshiftUpdated[from] = to
		}
	}

	covered := sets.NewString()
	var lines []string
	for i := range updated {
		a, b := defaults[i], updated[i]
		from, to := a.GetE2EImage(), b.GetE2EImage()
		if from == to {
			continue
		}
		if covered.Has(from) {
			continue
		}
		covered.Insert(from)
		lines = append(lines, fmt.Sprintf("%s %s%s", from, prefix, to))
	}

	for from, to := range openshiftUpdated {
		if from == to {
			continue
		}
		if covered.Has(from) {
			continue
		}
		covered.Insert(from)
		lines = append(lines, fmt.Sprintf("%s %s%s", from, prefix, to))
	}

	sort.Strings(lines)
	return lines, nil
}

func verifyImages() error {
	if len(os.Getenv("KUBE_TEST_REPO")) > 0 {
		return fmt.Errorf("KUBE_TEST_REPO may not be specified when this command is run")
	}
	return verifyImagesWithoutEnv()
}

func verifyImagesWithoutEnv() error {
	defaults := k8simage.GetOriginalImageConfigs()

	for originalPullSpec, index := range image.OriginalImages() {
		if index == -1 {
			continue
		}
		existing, ok := defaults[index]
		if !ok {
			return fmt.Errorf("image %q not found in upstream images, must be moved to test/extended/util/image", originalPullSpec)
		}
		if existing.GetE2EImage() != originalPullSpec {
			return fmt.Errorf("image %q defines index %d but is defined upstream as %q, must be fixed in test/extended/util/image", originalPullSpec, index, existing.GetE2EImage())
		}
		mirror := image.LocationFor(originalPullSpec)
		upstreamMirror := k8simage.GetE2EImage(index)
		if mirror != upstreamMirror {
			return fmt.Errorf("image %q defines index %d and mirror %q but is mirrored upstream as %q, must be fixed in test/extended/util/image", originalPullSpec, index, mirror, upstreamMirror)
		}
	}

	return nil
}

// pulledInvalidImages returns a function that checks whether the cluster pulled an image that is
// outside the allowed list of images. The list is defined as a set of static test case images, the
// local cluster registry, any repository referenced by the image streams in the cluster's 'openshift'
// namespace, or the location that input images are cloned from. Only namespaces prefixed with 'e2e-'
// are checked.
func pulledInvalidImages(fromRepository string) ginkgo.JUnitForEventsFunc {
	// static allowed images
	allowedImages := sets.NewString("image/webserver:404")
	allowedPrefixes := sets.NewString(
		"image-registry.openshift-image-registry.svc",
		"gcr.io/k8s-authenticated-test/",
		"gcr.io/authenticated-image-pulling/",
		"invalid.com/",

		// installed alongside OLM and managed externally
		"registry.redhat.io/redhat/community-operator-index",
		"registry.redhat.io/redhat/certified-operator-index",
		"registry.redhat.io/redhat/redhat-marketplace-index",
		"registry.redhat.io/redhat/redhat-operator-index",

		// used by OLM tests
		"registry.redhat.io/amq7/amq-streams-rhel7-operator",
		"registry.redhat.io/amq7/amqstreams-rhel7-operator-metadata",

		// used to test pull secrets against an authenticated registry
		// TODO: will not work for a disconnected test environment and should be emulated by launching
		//   an authenticated registry in a pod on cluster
		"registry.redhat.io/ubi8/nodejs-14:latest",
	)
	if len(fromRepository) > 0 {
		allowedPrefixes.Insert(fromRepository)
	}

	// any image not in the allowed prefixes is considered a failure, as the user
	// may have added a new test image without calling the appropriate helpers
	return func(events monitorapi.Intervals, _ time.Duration, _ *rest.Config) []*ginkgo.JUnitTestCase {
		imageStreamPrefixes, err := imagePrefixesFromNamespaceImageStreams("openshift")
		if err != nil {
			klog.Errorf("Unable to identify image prefixes from the openshift namespace: %v", err)
		}
		allowedPrefixes.Insert(imageStreamPrefixes.UnsortedList()...)

		allowedPrefixes := allowedPrefixes.List()

		var tests []*ginkgo.JUnitTestCase

		pulls := make(map[string]sets.String)
		for _, event := range events {
			// only messages that include a Pulled reason
			if !strings.Contains(" "+event.Message, " reason/Pulled ") {
				continue
			}
			// only look at pull events from an e2e-* namespace
			if !strings.Contains(" "+event.Locator, " ns/e2e-") {
				continue
			}

			parts := strings.Split(event.Message, " ")
			if len(parts) == 0 {
				continue
			}
			image := strings.TrimPrefix(parts[len(parts)-1], "image/")
			if hasAnyStringPrefix(image, allowedPrefixes) || allowedImages.Has(image) {
				continue
			}
			byImage, ok := pulls[image]
			if !ok {
				byImage = sets.NewString()
				pulls[image] = byImage
			}
			byImage.Insert(event.Locator)
		}
		if len(pulls) > 0 {
			images := make([]string, 0, len(pulls))
			for image := range pulls {
				images = append(images, image)
			}
			sort.Strings(images)
			buf := &bytes.Buffer{}
			for _, image := range images {
				fmt.Fprintf(buf, "%s from pods:\n", image)
				for _, locator := range pulls[image].List() {
					fmt.Fprintf(buf, "  %s\n", locator)
				}
			}
			tests = append(tests, &ginkgo.JUnitTestCase{
				Name:      "[sig-arch] Only known images used by tests",
				SystemOut: buf.String(),
				FailureOutput: &ginkgo.FailureOutput{
					Output: fmt.Sprintf("Cluster accessed images that were not mirrored to the testing repository or already part of the cluster, see test/extended/util/image/README.md in the openshift/origin repo:\n\n%s", buf.String()),
				},
			})

		} else {
			// if the test passed, indicate that too.
			tests = append(tests, &ginkgo.JUnitTestCase{
				Name: "[sig-arch] Only known images used by tests",
			})
		}

		return tests
	}
}

func hasAnyStringPrefix(s string, prefixes []string) bool {
	for _, prefix := range prefixes {
		if strings.HasPrefix(s, prefix) {
			return true
		}
	}
	return false
}

// imagePrefixesFromNamespaceImageStreams identifies all image repositories referenced by
// image streams in the provided namespace and returns them as a set (for both tags and
// digests). This set of prefixes can be used to verify that image references are coming
// from a location the cluster knows about.
func imagePrefixesFromNamespaceImageStreams(ns string) (sets.String, error) {
	clientConfig, err := e2e.LoadConfig(true)
	if err != nil {
		return nil, err
	}
	client, err := imagev1.NewForConfig(clientConfig)
	if err != nil {
		return nil, err
	}
	streams, err := client.ImageStreams(ns).List(context.TODO(), metav1.ListOptions{})
	if err != nil {
		return nil, err
	}

	allowedPrefixes := sets.NewString()
	for _, stream := range streams.Items {
		for _, tag := range stream.Spec.Tags {
			if tag.From == nil || tag.From.Kind != "DockerImage" {
				continue
			}
			ref, err := reference.Parse(tag.From.Name)
			if err != nil {
				continue
			}
			repo := ref.AsRepository().Exact()
			allowedPrefixes.Insert(repo + ":")
			allowedPrefixes.Insert(repo + "@")
		}
		for _, tag := range stream.Status.Tags {
			for _, event := range tag.Items {
				if len(event.DockerImageReference) == 0 {
					continue
				}
				ref, err := reference.Parse(event.DockerImageReference)
				if err != nil {
					continue
				}
				repo := ref.AsRepository().Exact()
				allowedPrefixes.Insert(repo + ":")
				allowedPrefixes.Insert(repo + "@")
			}
		}
	}
	return allowedPrefixes, nil
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"context"
	"fmt"
	"os"
	"sort"
	"strings"
	"time"

	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/util/sets"
	"k8s.io/klog/v2"

	imagev1 "github.com/openshift/client-go/image/clientset/versioned/typed/image/v1"
	"github.com/openshift/library-go/pkg/image/reference"

	"github.com/openshift/origin/pkg/monitor"
	"github.com/openshift/origin/pkg/test/ginkgo"
	"github.com/openshift/origin/test/extended/util/image"

	e2e "k8s.io/kubernetes/test/e2e/framework"
	k8simage "k8s.io/kubernetes/test/utils/image"
)

// OpenShift tests consume images from the cluster, from a number of vetted community locations,
// and from the upstream Kubernetes test suite which may reference images produced by a variety
// of build systems. To better organize and consolidate these images, the utility code within
// Kubernetes and OpenShift that is consulted to find image streams is made remappable - so that
// all images used by the test code can be located in one place. During normal operation the test
// images are read from the mirror, and administrators can choose to copy those images into
// restricted environment image registries and then run the tests against that subset. This also
// allows us to control the process whereby new images are introduced and review those in one spot.
//
// Test code utilizes helpers to get image strings throughout the code base, or is expected to use
// one or more of the images every OpenShift distribution includes in the 'openshift' namespace.
//
// See test/extended/util/image/README.md for a description of the process of adding a new image.

// defaultTestImageMirrorLocation is where all Kube test inputs are sourced.
const defaultTestImageMirrorLocation = "quay.io/openshift/community-e2e-images"

// createImageMirrorForInternalImages returns a list of 'oc image mirror' mappings from source to
// target or returns an error. If mirrored is true the images are assumed to have already been copied
// from their upstream location into our official mirror, in the REPO:TAG format where TAG is a hash
// of the original internal name and the index of the image in the array. Otherwise the mappings will
// be set to mirror the location as defined in the test code into our official mirror, where the target
// TAG is the hash described above.
func createImageMirrorForInternalImages(prefix string, ref reference.DockerImageReference, mirrored bool) ([]string, error) {
	source := ref.Exact()

	defaults := k8simage.GetOriginalImageConfigs()
	updated := k8simage.GetMappedImageConfigs(defaults, ref.Exact())

	openshiftDefaults := image.OriginalImages()
	openshiftUpdated := image.GetMappedImages(openshiftDefaults, defaultTestImageMirrorLocation)

	// if we've mirrored, then the source is going to be our repo, not upstream's
	if mirrored {
		baseRef, err := reference.Parse(defaultTestImageMirrorLocation)
		if err != nil {
			return nil, fmt.Errorf("invalid default mirror location: %v", err)
		}

		// calculate the mapping of upstream images by setting defaults to baseRef
		covered := sets.NewString()
		for i, config := range updated {
			defaultConfig := defaults[i]
			pullSpec := config.GetE2EImage()
			if pullSpec == defaultConfig.GetE2EImage() {
				continue
			}
			if covered.Has(pullSpec) {
				continue
			}
			covered.Insert(pullSpec)
			e2eRef, err := reference.Parse(pullSpec)
			if err != nil {
				return nil, fmt.Errorf("invalid test image: %s: %v", pullSpec, err)
			}
			if len(e2eRef.Tag) == 0 {
				return nil, fmt.Errorf("invalid test image: %s: no tag", pullSpec)
			}
			config.SetRegistry(baseRef.Registry)
			config.SetName(baseRef.RepositoryName())
			config.SetVersion(e2eRef.Tag)
			defaults[i] = config
		}

		// calculate the mapping for openshift images by populating openshiftUpdated
		openshiftUpdated = make(map[string]string)
		sourceMappings := image.GetMappedImages(openshiftDefaults, defaultTestImageMirrorLocation)
		targetMappings := image.GetMappedImages(openshiftDefaults, source)

		for from, to := range targetMappings {
			if from == to {
				continue
			}
			if covered.Has(to) {
				continue
			}
			covered.Insert(to)
			from := sourceMappings[from]
			openshiftUpdated[from] = to
		}
	}

	covered := sets.NewString()
	var lines []string
	for i := range updated {
		a, b := defaults[i], updated[i]
		from, to := a.GetE2EImage(), b.GetE2EImage()
		if from == to {
			continue
		}
		if covered.Has(from) {
			continue
		}
		covered.Insert(from)
		lines = append(lines, fmt.Sprintf("%s %s%s", from, prefix, to))
	}

	for from, to := range openshiftUpdated {
		if from == to {
			continue
		}
		if covered.Has(from) {
			continue
		}
		covered.Insert(from)
		lines = append(lines, fmt.Sprintf("%s %s%s", from, prefix, to))
	}

	sort.Strings(lines)
	return lines, nil
}

func verifyImages() error {
	if len(os.Getenv("KUBE_TEST_REPO")) > 0 {
		return fmt.Errorf("KUBE_TEST_REPO may not be specified when this command is run")
	}
	return verifyImagesWithoutEnv()
}

func verifyImagesWithoutEnv() error {
	defaults := k8simage.GetOriginalImageConfigs()

	for originalPullSpec, index := range image.OriginalImages() {
		if index == -1 {
			continue
		}
		existing, ok := defaults[index]
		if !ok {
			return fmt.Errorf("image %q not found in upstream images, must be moved to test/extended/util/image", originalPullSpec)
		}
		if existing.GetE2EImage() != originalPullSpec {
			return fmt.Errorf("image %q defines index %d but is defined upstream as %q, must be fixed in test/extended/util/image", originalPullSpec, index, existing.GetE2EImage())
		}
		mirror := image.LocationFor(originalPullSpec)
		upstreamMirror := k8simage.GetE2EImage(index)
		if mirror != upstreamMirror {
			return fmt.Errorf("image %q defines index %d and mirror %q but is mirrored upstream as %q, must be fixed in test/extended/util/image", originalPullSpec, index, mirror, upstreamMirror)
		}
	}

	return nil
}

// pulledInvalidImages returns a function that checks whether the cluster pulled an image that is
// outside the allowed list of images. The list is defined as a set of static test case images, the
// local cluster registry, any repository referenced by the image streams in the cluster's 'openshift'
// namespace, or the location that input images are cloned from. Only namespaces prefixed with 'e2e-'
// are checked.
func pulledInvalidImages(fromRepository string) ginkgo.JUnitForEventsFunc {
	// static allowed images
	allowedImages := sets.NewString("image/webserver:404")
	allowedPrefixes := sets.NewString(
		"image-registry.openshift-image-registry.svc",
		"gcr.io/k8s-authenticated-test/",
		"gcr.io/authenticated-image-pulling/",
		"invalid.com/",

		// installed alongside OLM and managed externally
		"registry.redhat.io/redhat/community-operator-index",
		"registry.redhat.io/redhat/certified-operator-index",
		"registry.redhat.io/redhat/redhat-marketplace-index",
		"registry.redhat.io/redhat/redhat-operator-index",

		// used by OLM tests
		"registry.redhat.io/amq7/amq-streams-rhel7-operator",
		"registry.redhat.io/amq7/amqstreams-rhel7-operator-metadata",

		// used to test pull secrets against an authenticated registry
		// TODO: will not work for a disconnected test environment and should be emulated by launching
		//   an authenticated registry in a pod on cluster
		"registry.redhat.io/rhscl/nodejs-10-rhel7:latest",
	)
	if len(fromRepository) > 0 {
		allowedPrefixes.Insert(fromRepository)
	}

	// any image not in the allowed prefixes is considered a failure, as the user
	// may have added a new test image without calling the appropriate helpers
	return func(events monitor.EventIntervals, _ time.Duration) ([]*ginkgo.JUnitTestCase, bool) {
		imageStreamPrefixes, err := imagePrefixesFromNamespaceImageStreams("openshift")
		if err != nil {
			klog.Errorf("Unable to identify image prefixes from the openshift namespace: %v", err)
		}
		allowedPrefixes.Insert(imageStreamPrefixes.UnsortedList()...)

		allowedPrefixes := allowedPrefixes.List()

		passed := true
		var tests []*ginkgo.JUnitTestCase

		pulls := make(map[string]sets.String)
		for _, event := range events {
			// only messages that include a Pulled reason
			if !strings.Contains(event.Message, " reason/Pulled ") {
				continue
			}
			// only look at pull events from an e2e-* namespace
			if !strings.Contains(event.Locator, " ns/e2e-") {
				continue
			}

			parts := strings.Split(event.Message, " ")
			if len(parts) == 0 {
				continue
			}
			image := strings.TrimPrefix(parts[len(parts)-1], "image/")
			if hasAnyStringPrefix(image, allowedPrefixes) || allowedImages.Has(image) {
				continue
			}
			byImage, ok := pulls[image]
			if !ok {
				byImage = sets.NewString()
				pulls[image] = byImage
			}
			byImage.Insert(event.Locator)
		}
		if len(pulls) > 0 {
			images := make([]string, 0, len(pulls))
			for image := range pulls {
				images = append(images, image)
			}
			sort.Strings(images)
			buf := &bytes.Buffer{}
			for _, image := range images {
				fmt.Fprintf(buf, "%s from pods:\n", image)
				for _, locator := range pulls[image].List() {
					fmt.Fprintf(buf, "  %s\n", locator)
				}
			}
			tests = append(tests, &ginkgo.JUnitTestCase{
				Name:      "[sig-arch] Only known images used by tests",
				SystemOut: buf.String(),
				FailureOutput: &ginkgo.FailureOutput{
					Output: fmt.Sprintf("Cluster accessed images that were not mirrored to the testing repository or already part of the cluster, see test/extended/util/image/README.md in the openshift/origin repo:\n\n%s", buf.String()),
				},
			})
			passed = false
		}

		return tests, passed
	}
}

func hasAnyStringPrefix(s string, prefixes []string) bool {
	for _, prefix := range prefixes {
		if strings.HasPrefix(s, prefix) {
			return true
		}
	}
	return false
}

// imagePrefixesFromNamespaceImageStreams identifies all image repositories referenced by
// image streams in the provided namespace and returns them as a set (for both tags and
// digests). This set of prefixes can be used to verify that image references are coming
// from a location the cluster knows about.
func imagePrefixesFromNamespaceImageStreams(ns string) (sets.String, error) {
	clientConfig, err := e2e.LoadConfig(true)
	if err != nil {
		return nil, err
	}
	client, err := imagev1.NewForConfig(clientConfig)
	if err != nil {
		return nil, err
	}
	streams, err := client.ImageStreams(ns).List(context.TODO(), metav1.ListOptions{})
	if err != nil {
		return nil, err
	}

	allowedPrefixes := sets.NewString()
	for _, stream := range streams.Items {
		for _, tag := range stream.Spec.Tags {
			if tag.From == nil || tag.From.Kind != "DockerImage" {
				continue
			}
			ref, err := reference.Parse(tag.From.Name)
			if err != nil {
				continue
			}
			repo := ref.AsRepository().Exact()
			allowedPrefixes.Insert(repo + ":")
			allowedPrefixes.Insert(repo + "@")
		}
		for _, tag := range stream.Status.Tags {
			for _, event := range tag.Items {
				if len(event.DockerImageReference) == 0 {
					continue
				}
				ref, err := reference.Parse(event.DockerImageReference)
				if err != nil {
					continue
				}
				repo := ref.AsRepository().Exact()
				allowedPrefixes.Insert(repo + ":")
				allowedPrefixes.Insert(repo + "@")
			}
		}
	}
	return allowedPrefixes, nil
}
<commit_msg>test: Correctly detect non-mirrored images<commit_after>package main

import (
	"bytes"
	"context"
	"fmt"
	"os"
	"sort"
	"strings"
	"time"

	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/util/sets"
	"k8s.io/klog/v2"

	imagev1 "github.com/openshift/client-go/image/clientset/versioned/typed/image/v1"
	"github.com/openshift/library-go/pkg/image/reference"

	"github.com/openshift/origin/pkg/monitor"
	"github.com/openshift/origin/pkg/test/ginkgo"
	"github.com/openshift/origin/test/extended/util/image"

	e2e "k8s.io/kubernetes/test/e2e/framework"
	k8simage "k8s.io/kubernetes/test/utils/image"
)

// OpenShift tests consume images from the cluster, from a number of vetted community locations,
// and from the upstream Kubernetes test suite which may reference images produced by a variety
// of build systems. To better organize and consolidate these images, the utility code within
// Kubernetes and OpenShift that is consulted to find image streams is made remappable - so that
// all images used by the test code can be located in one place. During normal operation the test
// images are read from the mirror, and administrators can choose to copy those images into
// restricted environment image registries and then run the tests against that subset. This also
// allows us to control the process whereby new images are introduced and review those in one spot.
//
// Test code utilizes helpers to get image strings throughout the code base, or is expected to use
// one or more of the images every OpenShift distribution includes in the 'openshift' namespace.
//
// See test/extended/util/image/README.md for a description of the process of adding a new image.

// defaultTestImageMirrorLocation is where all Kube test inputs are sourced.
const defaultTestImageMirrorLocation = "quay.io/openshift/community-e2e-images"

// createImageMirrorForInternalImages returns a list of 'oc image mirror' mappings from source to
// target or returns an error. If mirrored is true the images are assumed to have already been copied
// from their upstream location into our official mirror, in the REPO:TAG format where TAG is a hash
// of the original internal name and the index of the image in the array. Otherwise the mappings will
// be set to mirror the location as defined in the test code into our official mirror, where the target
// TAG is the hash described above.
func createImageMirrorForInternalImages(prefix string, ref reference.DockerImageReference, mirrored bool) ([]string, error) {
	source := ref.Exact()

	defaults := k8simage.GetOriginalImageConfigs()
	updated := k8simage.GetMappedImageConfigs(defaults, ref.Exact())

	openshiftDefaults := image.OriginalImages()
	openshiftUpdated := image.GetMappedImages(openshiftDefaults, defaultTestImageMirrorLocation)

	// if we've mirrored, then the source is going to be our repo, not upstream's
	if mirrored {
		baseRef, err := reference.Parse(defaultTestImageMirrorLocation)
		if err != nil {
			return nil, fmt.Errorf("invalid default mirror location: %v", err)
		}

		// calculate the mapping of upstream images by setting defaults to baseRef
		covered := sets.NewString()
		for i, config := range updated {
			defaultConfig := defaults[i]
			pullSpec := config.GetE2EImage()
			if pullSpec == defaultConfig.GetE2EImage() {
				continue
			}
			if covered.Has(pullSpec) {
				continue
			}
			covered.Insert(pullSpec)
			e2eRef, err := reference.Parse(pullSpec)
			if err != nil {
				return nil, fmt.Errorf("invalid test image: %s: %v", pullSpec, err)
			}
			if len(e2eRef.Tag) == 0 {
				return nil, fmt.Errorf("invalid test image: %s: no tag", pullSpec)
			}
			config.SetRegistry(baseRef.Registry)
			config.SetName(baseRef.RepositoryName())
			config.SetVersion(e2eRef.Tag)
			defaults[i] = config
		}

		// calculate the mapping for openshift images by populating openshiftUpdated
		openshiftUpdated = make(map[string]string)
		sourceMappings := image.GetMappedImages(openshiftDefaults, defaultTestImageMirrorLocation)
		targetMappings := image.GetMappedImages(openshiftDefaults, source)

		for from, to := range targetMappings {
			if from == to {
				continue
			}
			if covered.Has(to) {
				continue
			}
			covered.Insert(to)
			from := sourceMappings[from]
			openshiftUpdated[from] = to
		}
	}

	covered := sets.NewString()
	var lines []string
	for i := range updated {
		a, b := defaults[i], updated[i]
		from, to := a.GetE2EImage(), b.GetE2EImage()
		if from == to {
			continue
		}
		if covered.Has(from) {
			continue
		}
		covered.Insert(from)
		lines = append(lines, fmt.Sprintf("%s %s%s", from, prefix, to))
	}

	for from, to := range openshiftUpdated {
		if from == to {
			continue
		}
		if covered.Has(from) {
			continue
		}
		covered.Insert(from)
		lines = append(lines, fmt.Sprintf("%s %s%s", from, prefix, to))
	}

	sort.Strings(lines)
	return lines, nil
}

func verifyImages() error {
	if len(os.Getenv("KUBE_TEST_REPO")) > 0 {
		return fmt.Errorf("KUBE_TEST_REPO may not be specified when this command is run")
	}
	return verifyImagesWithoutEnv()
}

func verifyImagesWithoutEnv() error {
	defaults := k8simage.GetOriginalImageConfigs()

	for originalPullSpec, index := range image.OriginalImages() {
		if index == -1 {
			continue
		}
		existing, ok := defaults[index]
		if !ok {
			return fmt.Errorf("image %q not found in upstream images, must be moved to test/extended/util/image", originalPullSpec)
		}
		if existing.GetE2EImage() != originalPullSpec {
			return fmt.Errorf("image %q defines index %d but is defined upstream as %q, must be fixed in test/extended/util/image", originalPullSpec, index, existing.GetE2EImage())
		}
		mirror := image.LocationFor(originalPullSpec)
		upstreamMirror := k8simage.GetE2EImage(index)
		if mirror != upstreamMirror {
			return fmt.Errorf("image %q defines index %d and mirror %q but is mirrored upstream as %q, must be fixed in test/extended/util/image", originalPullSpec, index, mirror, upstreamMirror)
		}
	}

	return nil
}

// pulledInvalidImages returns a function that checks whether the cluster pulled an image that is
// outside the allowed list of images. The list is defined as a set of static test case images, the
// local cluster registry, any repository referenced by the image streams in the cluster's 'openshift'
// namespace, or the location that input images are cloned from. Only namespaces prefixed with 'e2e-'
// are checked.
func pulledInvalidImages(fromRepository string) ginkgo.JUnitForEventsFunc {
	// static allowed images
	allowedImages := sets.NewString("image/webserver:404")
	allowedPrefixes := sets.NewString(
		"image-registry.openshift-image-registry.svc",
		"gcr.io/k8s-authenticated-test/",
		"gcr.io/authenticated-image-pulling/",
		"invalid.com/",

		// installed alongside OLM and managed externally
		"registry.redhat.io/redhat/community-operator-index",
		"registry.redhat.io/redhat/certified-operator-index",
		"registry.redhat.io/redhat/redhat-marketplace-index",
		"registry.redhat.io/redhat/redhat-operator-index",

		// used by OLM tests
		"registry.redhat.io/amq7/amq-streams-rhel7-operator",
		"registry.redhat.io/amq7/amqstreams-rhel7-operator-metadata",

		// used to test pull secrets against an authenticated registry
		// TODO: will not work for a disconnected test environment and should be emulated by launching
		//   an authenticated registry in a pod on cluster
		"registry.redhat.io/rhscl/nodejs-10-rhel7:latest",
	)
	if len(fromRepository) > 0 {
		allowedPrefixes.Insert(fromRepository)
	}

	// any image not in the allowed prefixes is considered a failure, as the user
	// may have added a new test image without calling the appropriate helpers
	return func(events monitor.EventIntervals, _ time.Duration) ([]*ginkgo.JUnitTestCase, bool) {
		imageStreamPrefixes, err := imagePrefixesFromNamespaceImageStreams("openshift")
		if err != nil {
			klog.Errorf("Unable to identify image prefixes from the openshift namespace: %v", err)
		}
		allowedPrefixes.Insert(imageStreamPrefixes.UnsortedList()...)

		allowedPrefixes := allowedPrefixes.List()

		passed := true
		var tests []*ginkgo.JUnitTestCase

		pulls := make(map[string]sets.String)
		for _, event := range events {
			// only messages that include a Pulled reason
			if !strings.Contains(" "+event.Message, " reason/Pulled ") {
				continue
			}
			// only look at pull events from an e2e-* namespace
			if !strings.Contains(" "+event.Locator, " ns/e2e-") {
				continue
			}

			parts := strings.Split(event.Message, " ")
			if len(parts) == 0 {
				continue
			}
			image := strings.TrimPrefix(parts[len(parts)-1], "image/")
			if hasAnyStringPrefix(image, allowedPrefixes) || allowedImages.Has(image) {
				continue
			}
			byImage, ok := pulls[image]
			if !ok {
				byImage = sets.NewString()
				pulls[image] = byImage
			}
			byImage.Insert(event.Locator)
		}
		if len(pulls) > 0 {
			images := make([]string, 0, len(pulls))
			for image := range pulls {
				images = append(images, image)
			}
			sort.Strings(images)
			buf := &bytes.Buffer{}
			for _, image := range images {
				fmt.Fprintf(buf, "%s from pods:\n", image)
				for _, locator := range pulls[image].List() {
					fmt.Fprintf(buf, "  %s\n", locator)
				}
			}
			tests = append(tests, &ginkgo.JUnitTestCase{
				Name:      "[sig-arch] Only known images used by tests",
				SystemOut: buf.String(),
				FailureOutput: &ginkgo.FailureOutput{
					Output: fmt.Sprintf("Cluster accessed images that were not mirrored to the testing repository or already part of the cluster, see test/extended/util/image/README.md in the openshift/origin repo:\n\n%s", buf.String()),
				},
			})
			passed = false
		}

		return tests, passed
	}
}

func hasAnyStringPrefix(s string, prefixes []string) bool {
	for _, prefix := range prefixes {
		if strings.HasPrefix(s, prefix) {
			return true
		}
	}
	return false
}

// imagePrefixesFromNamespaceImageStreams identifies all image repositories referenced by
// image streams in the provided namespace and returns them as a set (for both tags and
// digests). This set of prefixes can be used to verify that image references are coming
// from a location the cluster knows about.
func imagePrefixesFromNamespaceImageStreams(ns string) (sets.String, error) {
	clientConfig, err := e2e.LoadConfig(true)
	if err != nil {
		return nil, err
	}
	client, err := imagev1.NewForConfig(clientConfig)
	if err != nil {
		return nil, err
	}
	streams, err := client.ImageStreams(ns).List(context.TODO(), metav1.ListOptions{})
	if err != nil {
		return nil, err
	}

	allowedPrefixes := sets.NewString()
	for _, stream := range streams.Items {
		for _, tag := range stream.Spec.Tags {
			if tag.From == nil || tag.From.Kind != "DockerImage" {
				continue
			}
			ref, err := reference.Parse(tag.From.Name)
			if err != nil {
				continue
			}
			repo := ref.AsRepository().Exact()
			allowedPrefixes.Insert(repo + ":")
			allowedPrefixes.Insert(repo + "@")
		}
		for _, tag := range stream.Status.Tags {
			for _, event := range tag.Items {
				if len(event.DockerImageReference) == 0 {
					continue
				}
				ref, err := reference.Parse(event.DockerImageReference)
				if err != nil {
					continue
				}
				repo := ref.AsRepository().Exact()
				allowedPrefixes.Insert(repo + ":")
				allowedPrefixes.Insert(repo + "@")
			}
		}
	}
	return allowedPrefixes, nil
}
<|endoftext|>


<commit_before>// Copyright 2020 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// The errors below are produced by the parser, but we check
// them here for consistency with the types2 tests.

package p

import ; // ERROR missing import path
import
var /* ERROR missing import path */ _ int
import .; //  ERROR missing import path
import 'x' // ERROR import path must be a string
var _ int
import /* ERROR imports must appear before other declarations */ _ "math"

// Don't repeat previous error for each immediately following import ...
import ()
import (.) // ERROR missing import path
import (
	"fmt"
	.
) // ERROR missing import path

// ... but remind with error again if we start a new import section after
// other declarations
var _ = fmt.Println
import /* ERROR imports must appear before other declarations */ _ "math"
import _ "math"
<commit_msg>internal/types: add a test for empty string import path<commit_after>// Copyright 2020 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// The errors below are produced by the parser, but we check
// them here for consistency with the types2 tests.

package p

import ; // ERROR missing import path
import "" // ERROR invalid import path \(empty string\)
import
var /* ERROR missing import path */ _ int
import .; //  ERROR missing import path
import 'x' // ERROR import path must be a string
var _ int
import /* ERROR imports must appear before other declarations */ _ "math"

// Don't repeat previous error for each immediately following import ...
import ()
import (.) // ERROR missing import path
import (
	"fmt"
	.
) // ERROR missing import path

// ... but remind with error again if we start a new import section after
// other declarations
var _ = fmt.Println
import /* ERROR imports must appear before other declarations */ _ "math"
import _ "math"
<|endoftext|>


<commit_before>package prometheus_output

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"hash/fnv"
	"log"
	"math"
	"net"
	"net/http"
	"os"
	"path/filepath"
	"sort"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/karimra/gnmic/collector"
	"github.com/karimra/gnmic/outputs"
	"github.com/mitchellh/mapstructure"
	"github.com/openconfig/gnmi/proto/gnmi"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promhttp"
	dto "github.com/prometheus/client_model/go"
	"google.golang.org/protobuf/proto"
)

const (
	defaultListen     = ":9804"
	defaultPath       = "/metrics"
	defaultExpiration = time.Minute
	defaultMetricHelp = "gNMIc generated metric"
)

type labelPair struct {
	Name  string
	Value string
}
type promMetric struct {
	name   string
	labels []*labelPair
	time   time.Time
	value  float64
}

func init() {
	outputs.Register("prometheus", func() outputs.Output {
		return &PrometheusOutput{
			Cfg:       &Config{},
			eventChan: make(chan *collector.EventMsg),
			wg:        new(sync.WaitGroup),
			entries:   make(map[uint64]*promMetric),
		}
	})
}

type PrometheusOutput struct {
	Cfg       *Config
	metrics   []prometheus.Collector
	logger    *log.Logger
	cancelFn  context.CancelFunc
	eventChan chan *collector.EventMsg

	wg     *sync.WaitGroup
	server *http.Server
	sync.Mutex
	entries map[uint64]*promMetric

	replacer *strings.Replacer
}
type Config struct {
	Listen     string        `mapstructure:"listen,omitempty"`
	Path       string        `mapstructure:"path,omitempty"`
	Expiration time.Duration `mapstructure:"expiration,omitempty"`
	Debug      bool          `mapstructure:"debug,omitempty"`
}

func (p *PrometheusOutput) String() string {
	b, err := json.Marshal(p)
	if err != nil {
		return ""
	}
	return string(b)
}
func (p *PrometheusOutput) Init(ctx context.Context, cfg map[string]interface{}, logger *log.Logger) error {
	ctx, p.cancelFn = context.WithCancel(ctx)
	decoder, err := mapstructure.NewDecoder(
		&mapstructure.DecoderConfig{
			DecodeHook: mapstructure.StringToTimeDurationHookFunc(),
			Result:     p.Cfg,
		},
	)
	if err != nil {
		return err
	}
	err = decoder.Decode(cfg)
	if err != nil {
		return err
	}
	if p.Cfg.Listen == "" {
		p.Cfg.Listen = defaultListen
	}
	if p.Cfg.Path == "" {
		p.Cfg.Path = defaultPath
	}
	if p.Cfg.Expiration == 0 {
		p.Cfg.Expiration = defaultExpiration
	}
	p.logger = log.New(os.Stderr, "prometheus_output ", log.LstdFlags|log.Lmicroseconds)
	if logger != nil {
		p.logger.SetOutput(logger.Writer())
		p.logger.SetFlags(logger.Flags())
	}
	p.replacer = strings.NewReplacer("-", "_", ":", "_", "/", "_")
	// create prometheus registery
	registry := prometheus.NewRegistry()

	err = registry.Register(p)
	if err != nil {
		return err
	}
	// create http server
	promHandler := promhttp.HandlerFor(registry, promhttp.HandlerOpts{ErrorHandling: promhttp.ContinueOnError})

	mux := http.NewServeMux()
	mux.Handle(p.Cfg.Path, promHandler)

	p.server = &http.Server{
		Addr:    p.Cfg.Listen,
		Handler: mux,
	}

	// create tcp listener
	listener, err := net.Listen("tcp", p.Cfg.Listen)
	if err != nil {
		return err
	}
	// start worker
	p.wg.Add(2)
	go p.worker(ctx)
	go func() {
		defer p.wg.Done()
		err = p.server.Serve(listener)
		if err != nil && err != http.ErrServerClosed {
			p.logger.Printf("prometheus server error: %v", err)
		}
	}()
	p.logger.Printf("initialized prometheus output: %s", p.String())
	return nil
}
func (p *PrometheusOutput) Write(ctx context.Context, rsp proto.Message, meta outputs.Meta) {
	if rsp == nil {
		return
	}
	switch rsp := rsp.(type) {
	case *gnmi.SubscribeResponse:
		measName := "default"
		if subName, ok := meta["subscription-name"]; ok {
			measName = subName
		}
		events, err := collector.ResponseToEventMsgs(measName, rsp, meta)
		if err != nil {
			p.logger.Printf("failed to convert message to event: %v", err)
			return
		}
		for _, ev := range events {
			select {
			case <-ctx.Done():
				return
			case p.eventChan <- ev:
			}
		}
	}
}
func (p *PrometheusOutput) Close() error {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	err := p.server.Shutdown(ctx)
	if err != nil {
		p.logger.Printf("failed to shutdown http server: %v", err)
	}
	p.cancelFn()
	p.wg.Wait()
	return nil
}
func (p *PrometheusOutput) Metrics() []prometheus.Collector { return p.metrics }

///
func (p *PrometheusOutput) Describe(ch chan<- *prometheus.Desc) {}

func (p *PrometheusOutput) Collect(ch chan<- prometheus.Metric) {
	p.Lock()
	defer p.Unlock()

	for _, entry := range p.entries {
		ch <- entry
	}
}

func (p *PrometheusOutput) getLabels(ev *collector.EventMsg) []*labelPair {
	labels := make([]*labelPair, 0, len(ev.Tags))
	addedLabels := make(map[string]struct{})
	for k, v := range ev.Tags {
		labelName := p.replacer.Replace(filepath.Base(k))
		if _, ok := addedLabels[labelName]; ok {
			continue
		}
		labels = append(labels, &labelPair{Name: labelName, Value: v})
		addedLabels[labelName] = struct{}{}
	}
	return labels
}

func (p *PrometheusOutput) worker(ctx context.Context) {
	defer p.wg.Done()
	for {
		select {
		case <-ctx.Done():
			return
		case ev := <-p.eventChan:
			if p.Cfg.Debug {
				p.logger.Printf("got event to store: %+v", ev)
			}
			p.Lock()
			labels := p.getLabels(ev)
			for vName, val := range ev.Values {
				v, err := getFloat(val)
				if err != nil {
					continue
				}
				metricName := strings.TrimRight(p.replacer.Replace(ev.Name), "_")
				metricName += "_" + strings.TrimLeft(p.replacer.Replace(vName), "_")
				pm := &promMetric{
					name:   metricName,
					labels: labels,
					time:   time.Unix(0, ev.Timestamp),
					value:  v,
				}
				key := pm.calculateKey()
				if e, ok := p.entries[key]; ok {
					if e.time.Before(pm.time) {
						p.entries[key] = pm
					}
				} else {
					p.entries[key] = pm
				}
				if p.Cfg.Debug {
					p.logger.Printf("saved key=%d, metric: %+v", key, pm)
				}
			}
			// expire entries
			expiry := time.Now().Add(-p.Cfg.Expiration)
			for k, e := range p.entries {
				if e.time.Before(expiry) {
					delete(p.entries, k)
				}
			}
			p.Unlock()
		}
	}
}

// Metric
func (p *promMetric) calculateKey() uint64 {
	h := fnv.New64a()
	h.Write([]byte(p.name))
	if len(p.labels) > 0 {
		h.Write([]byte(":"))
		sort.Slice(p.labels, func(i, j int) bool {
			return p.labels[i].Name < p.labels[j].Name
		})
		for _, label := range p.labels {
			h.Write([]byte(label.Name))
			h.Write([]byte(":"))
			h.Write([]byte(label.Value))
			h.Write([]byte(":"))
		}
	}
	return h.Sum64()
}
func (p *promMetric) String() string {
	if p == nil {
		return ""
	}
	sb := strings.Builder{}
	sb.WriteString("name=")
	sb.WriteString(p.name)
	sb.WriteString(",")
	numLabels := len(p.labels)
	if numLabels > 0 {
		sb.WriteString("labels=[")
		for i, lb := range p.labels {
			sb.WriteString(lb.Name)
			sb.WriteString("=")
			sb.WriteString(lb.Value)
			if i < numLabels-1 {
				sb.WriteString(",")
			}
		}
		sb.WriteString("],")
	}
	sb.WriteString(fmt.Sprintf("value=%f,", p.value))
	sb.WriteString("time=")
	sb.WriteString(p.time.String())
	return sb.String()
}
func (p *promMetric) Desc() *prometheus.Desc {
	labelNames := make([]string, 0, len(p.labels))
	for _, label := range p.labels {
		labelNames = append(labelNames, label.Name)
	}

	return prometheus.NewDesc(p.name, defaultMetricHelp, labelNames, nil)
}
func (p *promMetric) Write(out *dto.Metric) error {
	out.Untyped = &dto.Untyped{
		Value: &p.value,
	}
	out.Label = make([]*dto.LabelPair, 0, len(p.labels))
	for _, lb := range p.labels {
		out.Label = append(out.Label, &dto.LabelPair{Name: &lb.Name, Value: &lb.Value})
	}
	timestamp := p.time.UnixNano() / 1000000
	out.TimestampMs = &timestamp
	return nil
}

func getFloat(v interface{}) (float64, error) {
	switch i := v.(type) {
	case float64:
		return float64(i), nil
	case float32:
		return float64(i), nil
	case int64:
		return float64(i), nil
	case int32:
		return float64(i), nil
	case int16:
		return float64(i), nil
	case int8:
		return float64(i), nil
	case uint64:
		return float64(i), nil
	case uint32:
		return float64(i), nil
	case uint16:
		return float64(i), nil
	case uint8:
		return float64(i), nil
	case int:
		return float64(i), nil
	case uint:
		return float64(i), nil
	case string:
		f, err := strconv.ParseFloat(i, 64)
		if err != nil {
			return math.NaN(), err
		}
		return f, err
	default:
		return math.NaN(), errors.New("getFloat: unknown value is of incompatible type")
	}
}
<commit_msg>separate metric name generation<commit_after>package prometheus_output

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"hash/fnv"
	"log"
	"math"
	"net"
	"net/http"
	"os"
	"path/filepath"
	"sort"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/karimra/gnmic/collector"
	"github.com/karimra/gnmic/outputs"
	"github.com/mitchellh/mapstructure"
	"github.com/openconfig/gnmi/proto/gnmi"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promhttp"
	dto "github.com/prometheus/client_model/go"
	"google.golang.org/protobuf/proto"
)

const (
	defaultListen     = ":9804"
	defaultPath       = "/metrics"
	defaultExpiration = time.Minute
	defaultMetricHelp = "gNMIc generated metric"
)

type labelPair struct {
	Name  string
	Value string
}
type promMetric struct {
	name   string
	labels []*labelPair
	time   time.Time
	value  float64
}

func init() {
	outputs.Register("prometheus", func() outputs.Output {
		return &PrometheusOutput{
			Cfg:       &Config{},
			eventChan: make(chan *collector.EventMsg),
			wg:        new(sync.WaitGroup),
			entries:   make(map[uint64]*promMetric),
		}
	})
}

type PrometheusOutput struct {
	Cfg       *Config
	metrics   []prometheus.Collector
	logger    *log.Logger
	cancelFn  context.CancelFunc
	eventChan chan *collector.EventMsg

	wg     *sync.WaitGroup
	server *http.Server
	sync.Mutex
	entries map[uint64]*promMetric

	replacer *strings.Replacer
}
type Config struct {
	Listen     string        `mapstructure:"listen,omitempty"`
	Path       string        `mapstructure:"path,omitempty"`
	Expiration time.Duration `mapstructure:"expiration,omitempty"`
	Debug      bool          `mapstructure:"debug,omitempty"`
}

func (p *PrometheusOutput) String() string {
	b, err := json.Marshal(p)
	if err != nil {
		return ""
	}
	return string(b)
}
func (p *PrometheusOutput) Init(ctx context.Context, cfg map[string]interface{}, logger *log.Logger) error {
	ctx, p.cancelFn = context.WithCancel(ctx)
	decoder, err := mapstructure.NewDecoder(
		&mapstructure.DecoderConfig{
			DecodeHook: mapstructure.StringToTimeDurationHookFunc(),
			Result:     p.Cfg,
		},
	)
	if err != nil {
		return err
	}
	err = decoder.Decode(cfg)
	if err != nil {
		return err
	}
	if p.Cfg.Listen == "" {
		p.Cfg.Listen = defaultListen
	}
	if p.Cfg.Path == "" {
		p.Cfg.Path = defaultPath
	}
	if p.Cfg.Expiration == 0 {
		p.Cfg.Expiration = defaultExpiration
	}
	p.logger = log.New(os.Stderr, "prometheus_output ", log.LstdFlags|log.Lmicroseconds)
	if logger != nil {
		p.logger.SetOutput(logger.Writer())
		p.logger.SetFlags(logger.Flags())
	}
	p.replacer = strings.NewReplacer("-", "_", ":", "_", "/", "_")
	// create prometheus registery
	registry := prometheus.NewRegistry()

	err = registry.Register(p)
	if err != nil {
		return err
	}
	// create http server
	promHandler := promhttp.HandlerFor(registry, promhttp.HandlerOpts{ErrorHandling: promhttp.ContinueOnError})

	mux := http.NewServeMux()
	mux.Handle(p.Cfg.Path, promHandler)

	p.server = &http.Server{
		Addr:    p.Cfg.Listen,
		Handler: mux,
	}

	// create tcp listener
	listener, err := net.Listen("tcp", p.Cfg.Listen)
	if err != nil {
		return err
	}
	// start worker
	p.wg.Add(2)
	go p.worker(ctx)
	go func() {
		defer p.wg.Done()
		err = p.server.Serve(listener)
		if err != nil && err != http.ErrServerClosed {
			p.logger.Printf("prometheus server error: %v", err)
		}
	}()
	p.logger.Printf("initialized prometheus output: %s", p.String())
	return nil
}
func (p *PrometheusOutput) Write(ctx context.Context, rsp proto.Message, meta outputs.Meta) {
	if rsp == nil {
		return
	}
	switch rsp := rsp.(type) {
	case *gnmi.SubscribeResponse:
		measName := "default"
		if subName, ok := meta["subscription-name"]; ok {
			measName = subName
		}
		events, err := collector.ResponseToEventMsgs(measName, rsp, meta)
		if err != nil {
			p.logger.Printf("failed to convert message to event: %v", err)
			return
		}
		for _, ev := range events {
			select {
			case <-ctx.Done():
				return
			case p.eventChan <- ev:
			}
		}
	}
}
func (p *PrometheusOutput) Close() error {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	err := p.server.Shutdown(ctx)
	if err != nil {
		p.logger.Printf("failed to shutdown http server: %v", err)
	}
	p.cancelFn()
	p.wg.Wait()
	return nil
}
func (p *PrometheusOutput) Metrics() []prometheus.Collector { return p.metrics }

///
func (p *PrometheusOutput) Describe(ch chan<- *prometheus.Desc) {}

func (p *PrometheusOutput) Collect(ch chan<- prometheus.Metric) {
	p.Lock()
	defer p.Unlock()

	for _, entry := range p.entries {
		ch <- entry
	}
}

func (p *PrometheusOutput) getLabels(ev *collector.EventMsg) []*labelPair {
	labels := make([]*labelPair, 0, len(ev.Tags))
	addedLabels := make(map[string]struct{})
	for k, v := range ev.Tags {
		labelName := p.replacer.Replace(filepath.Base(k))
		if _, ok := addedLabels[labelName]; ok {
			continue
		}
		labels = append(labels, &labelPair{Name: labelName, Value: v})
		addedLabels[labelName] = struct{}{}
	}
	return labels
}

func (p *PrometheusOutput) worker(ctx context.Context) {
	defer p.wg.Done()
	for {
		select {
		case <-ctx.Done():
			return
		case ev := <-p.eventChan:
			if p.Cfg.Debug {
				p.logger.Printf("got event to store: %+v", ev)
			}
			p.Lock()
			labels := p.getLabels(ev)
			for vName, val := range ev.Values {
				v, err := getFloat(val)
				if err != nil {
					continue
				}
				pm := &promMetric{
					name:   p.metricName(ev.Name, vName),
					labels: labels,
					time:   time.Unix(0, ev.Timestamp),
					value:  v,
				}
				key := pm.calculateKey()
				if e, ok := p.entries[key]; ok {
					if e.time.Before(pm.time) {
						p.entries[key] = pm
					}
				} else {
					p.entries[key] = pm
				}
				if p.Cfg.Debug {
					p.logger.Printf("saved key=%d, metric: %+v", key, pm)
				}
			}
			// expire entries
			expiry := time.Now().Add(-p.Cfg.Expiration)
			for k, e := range p.entries {
				if e.time.Before(expiry) {
					delete(p.entries, k)
				}
			}
			p.Unlock()
		}
	}
}

// Metric
func (p *promMetric) calculateKey() uint64 {
	h := fnv.New64a()
	h.Write([]byte(p.name))
	if len(p.labels) > 0 {
		h.Write([]byte(":"))
		sort.Slice(p.labels, func(i, j int) bool {
			return p.labels[i].Name < p.labels[j].Name
		})
		for _, label := range p.labels {
			h.Write([]byte(label.Name))
			h.Write([]byte(":"))
			h.Write([]byte(label.Value))
			h.Write([]byte(":"))
		}
	}
	return h.Sum64()
}
func (p *promMetric) String() string {
	if p == nil {
		return ""
	}
	sb := strings.Builder{}
	sb.WriteString("name=")
	sb.WriteString(p.name)
	sb.WriteString(",")
	numLabels := len(p.labels)
	if numLabels > 0 {
		sb.WriteString("labels=[")
		for i, lb := range p.labels {
			sb.WriteString(lb.Name)
			sb.WriteString("=")
			sb.WriteString(lb.Value)
			if i < numLabels-1 {
				sb.WriteString(",")
			}
		}
		sb.WriteString("],")
	}
	sb.WriteString(fmt.Sprintf("value=%f,", p.value))
	sb.WriteString("time=")
	sb.WriteString(p.time.String())
	return sb.String()
}
func (p *promMetric) Desc() *prometheus.Desc {
	labelNames := make([]string, 0, len(p.labels))
	for _, label := range p.labels {
		labelNames = append(labelNames, label.Name)
	}

	return prometheus.NewDesc(p.name, defaultMetricHelp, labelNames, nil)
}
func (p *promMetric) Write(out *dto.Metric) error {
	out.Untyped = &dto.Untyped{
		Value: &p.value,
	}
	out.Label = make([]*dto.LabelPair, 0, len(p.labels))
	for _, lb := range p.labels {
		out.Label = append(out.Label, &dto.LabelPair{Name: &lb.Name, Value: &lb.Value})
	}
	timestamp := p.time.UnixNano() / 1000000
	out.TimestampMs = &timestamp
	return nil
}

func getFloat(v interface{}) (float64, error) {
	switch i := v.(type) {
	case float64:
		return float64(i), nil
	case float32:
		return float64(i), nil
	case int64:
		return float64(i), nil
	case int32:
		return float64(i), nil
	case int16:
		return float64(i), nil
	case int8:
		return float64(i), nil
	case uint64:
		return float64(i), nil
	case uint32:
		return float64(i), nil
	case uint16:
		return float64(i), nil
	case uint8:
		return float64(i), nil
	case int:
		return float64(i), nil
	case uint:
		return float64(i), nil
	case string:
		f, err := strconv.ParseFloat(i, 64)
		if err != nil {
			return math.NaN(), err
		}
		return f, err
	default:
		return math.NaN(), errors.New("getFloat: unknown value is of incompatible type")
	}
}

func (p *PrometheusOutput) metricName(measName, valueName string) string {
	sb := strings.Builder{}
	sb.WriteString(strings.TrimRight(p.replacer.Replace(measName), "_"))
	sb.WriteString("_")
	sb.WriteString(strings.TrimLeft(p.replacer.Replace(valueName), "_"))
	return sb.String()
}
<|endoftext|>


<commit_before>// Copyright (c) 2013 Conformal Systems LLC.
// Copyright (c) 2014 Filippo Valsorda
// Use of this source code is governed by an ISC
// license that can be found in the LICENSE file.

package main

import (
	"errors"
	"github.com/conformal/btcchain"
	"github.com/conformal/btcdb"
	_ "github.com/conformal/btcdb/ldb"
	"github.com/conformal/btcec"
	"github.com/conformal/btclog"
	"github.com/conformal/btcutil"
	"github.com/conformal/btcwire"
	"github.com/conformal/go-flags"
	"github.com/davecgh/go-spew/spew"
	// "math/big"
	"os"
	"os/signal"
	"path/filepath"
	"time"
)

type ShaHash btcwire.ShaHash

type config struct {
	DataDir  string `short:"b" long:"datadir" description:"Directory to store data"`
	DbType   string `long:"dbtype" description:"Database backend"`
	TestNet3 bool   `long:"testnet" description:"Use the test network"`
	// Height   int64  `short:"b" description:"Block height to process" required:"true"`
}

var (
	btcdHomeDir    = btcutil.AppDataDir("btcd", false)
	defaultDataDir = filepath.Join(btcdHomeDir, "data")
	log            btclog.Logger
)

const (
	ArgSha = iota
	ArgHeight
)

var rValuesMap = make(map[string]int64)
var duplicates = make(map[string][]int64)

var blocksCounter int64 = 0
var sigCounter int64 = 0

func main() {
	cfg := config{
		DbType:  "leveldb",
		DataDir: defaultDataDir,
	}
	parser := flags.NewParser(&cfg, flags.Default)
	_, err := parser.Parse()
	if err != nil {
		if e, ok := err.(*flags.Error); !ok || e.Type != flags.ErrHelp {
			parser.WriteHelp(os.Stderr)
		}
		return
	}

	backendLogger := btclog.NewDefaultBackendLogger()
	defer backendLogger.Flush()
	log = btclog.NewSubsystemLogger(backendLogger, "")
	btcdb.UseLogger(log)

	var testnet string
	if cfg.TestNet3 {
		testnet = "testnet"
	} else {
		testnet = "mainnet"
	}

	cfg.DataDir = filepath.Join(cfg.DataDir, testnet)

	blockDbNamePrefix := "blocks"
	dbName := blockDbNamePrefix + "_" + cfg.DbType
	if cfg.DbType == "sqlite" {
		dbName = dbName + ".db"
	}
	dbPath := filepath.Join(cfg.DataDir, dbName)

	log.Infof("loading db %v", cfg.DbType)
	db, err := btcdb.OpenDB(cfg.DbType, dbPath)
	if err != nil {
		log.Warnf("db open failed: %v", err)
		return
	}
	defer db.Close()
	log.Infof("db load complete")

	c := make(chan os.Signal, 1)
	signal.Notify(c, os.Interrupt, os.Kill)

	go func() {
		_ = <-c
		log.Info(spew.Sdump(duplicates))
		// this will have to be reworked while parallelizing
		backendLogger.Flush()
		db.Close()
		os.Exit(2)
	}()

	_, max_heigth, err := db.NewestSha()
	if err != nil {
		log.Warnf("db NewestSha failed: %v", err)
		return
	}
	log.Infof("max_heigth: %v", max_heigth)

	last_time := time.Now()
	for h := int64(0); h < max_heigth; h++ {
		// TODO: parallelize
		err = DumpBlock(db, h)
		if err != nil {
			log.Warnf("Failed to dump block %v, err %v", h, err)
		}
		if blocksCounter%10000 == 0 {
			log.Infof("%v blocks processed in %v, %v signatures stored",
				blocksCounter, time.Since(last_time), sigCounter)
			last_time = time.Now()
		}
	}

	log.Info(spew.Sdump(duplicates))
}

var notDataError = errors.New("the first opcode is not a data push")

func parseData(SignatureScript []uint8) ([]uint8, error) {
	opcode := SignatureScript[0]

	if opcode >= 1 && opcode <= 75 {
		sigStr := SignatureScript[1 : opcode+1]
		return sigStr, nil
	}

	// TODO: OP_PUSHDATA1 OP_PUSHDATA2 OP_PUSHDATA3

	return nil, notDataError
}

func DumpBlock(db btcdb.Db, height int64) error {
	sha, err := db.FetchBlockShaByHeight(height)
	if err != nil {
		return err
	}
	blk, err := db.FetchBlockBySha(sha)
	if err != nil {
		return err
	}
	rblk, err := blk.Bytes()
	if err != nil {
		return err
	}
	blkid := blk.Height()
	if blkid != height {
		return errors.New("WHAT!?")
	}

	log.Debugf("Block %v depth %v", sha, blkid)

	log.Debugf("Block %v depth %v %v", sha, blkid, spew.Sdump(rblk))
	mblk := blk.MsgBlock()
	log.Debugf("Block %v depth %v %v", sha, blkid, spew.Sdump(mblk))

	log.Debugf("Num transactions %v", len(mblk.Transactions))
	for i, tx := range mblk.Transactions {

		txsha, err := tx.TxSha()
		if err != nil {
			log.Warnf("Block %v (%v)", blkid, sha)
			log.Warnf("tx %v (%v)", i, &txsha)
			log.Warnf("Error: %v", err)
			continue
		}

		log.Debugf("tx %v: %v", i, &txsha)

		if btcchain.IsCoinBase(btcutil.NewTx(tx)) {
			log.Debugf("tx %v: skipping (coinbase)", i)
			continue
		}

		for t, txin := range tx.TxIn {
			log.Debugf("tx %v: TxIn %v: SignatureScript: %v", i, t, spew.Sdump(txin.SignatureScript))

			sigStr, err := parseData(txin.SignatureScript)
			if err != nil {
				log.Warnf("Block %v (%v)", blkid, sha)
				log.Warnf("tx %v (%v)", i, &txsha)
				log.Warnf("txin %v (parseData)", t)
				log.Warnf("Error: %v", err)
				log.Info(spew.Sdump(txin.SignatureScript))
				continue
			}
			log.Debugf("tx %v: TxIn %v: sigStr: %v", i, t, spew.Sdump(sigStr))

			signature, err := btcec.ParseSignature(sigStr, btcec.S256())
			if err != nil {
				log.Warnf("Block %v (%v)", blkid, sha)
				log.Warnf("tx %v (%v)", i, &txsha)
				log.Warnf("txin %v (ParseSignature)", t)
				log.Warnf("Error: %v", err)
				log.Info(spew.Sdump(sigStr))
				continue

			}
			log.Debugf("tx %v: TxIn %v: signature: %v", i, t, spew.Sdump(signature))

			signatureString := signature.R.String()
			if rValuesMap[signatureString] != int64(0) {
				log.Infof("%v:%v:%v DUPLICATE FOUND: %v",
					blkid, txsha.String()[:5], t, rValuesMap[signatureString])
				if len(duplicates[signatureString]) == 0 {
					duplicates[signatureString] = append(duplicates[signatureString], rValuesMap[signatureString])
				}
				duplicates[signatureString] = append(duplicates[signatureString], blkid)
			} else {
				rValuesMap[signatureString] = blkid
				sigCounter++
			}

		}

	}

	blocksCounter++

	return nil
}
<commit_msg>minor<commit_after>// Copyright (c) 2013 Conformal Systems LLC.
// Copyright (c) 2014 Filippo Valsorda
// Use of this source code is governed by an ISC
// license that can be found in the LICENSE file.

package main

import (
	"fmt"
	"github.com/conformal/btcchain"
	"github.com/conformal/btcdb"
	_ "github.com/conformal/btcdb/ldb"
	"github.com/conformal/btcec"
	"github.com/conformal/btclog"
	"github.com/conformal/btcutil"
	"github.com/conformal/btcwire"
	"github.com/conformal/go-flags"
	"github.com/davecgh/go-spew/spew"
	// "math/big"
	"os"
	"os/signal"
	"path/filepath"
	"time"
)

// TODO: monky-patch OP_CHECKSIG

type ShaHash btcwire.ShaHash

type config struct {
	DataDir  string `short:"b" long:"datadir" description:"Directory to store data"`
	DbType   string `long:"dbtype" description:"Database backend"`
	TestNet3 bool   `long:"testnet" description:"Use the test network"`
	// Height   int64  `short:"b" description:"Block height to process" required:"true"`
}

var (
	btcdHomeDir    = btcutil.AppDataDir("btcd", false)
	defaultDataDir = filepath.Join(btcdHomeDir, "data")
	log            btclog.Logger
)

const (
	ArgSha = iota
	ArgHeight
)

var rValuesMap = make(map[string]int64)
var duplicates = make(map[string][]int64)

var blocksCounter int64 = 0
var sigCounter int64 = 0

func main() {
	cfg := config{
		DbType:  "leveldb",
		DataDir: defaultDataDir,
	}
	parser := flags.NewParser(&cfg, flags.Default)
	_, err := parser.Parse()
	if err != nil {
		if e, ok := err.(*flags.Error); !ok || e.Type != flags.ErrHelp {
			parser.WriteHelp(os.Stderr)
		}
		return
	}

	backendLogger := btclog.NewDefaultBackendLogger()
	defer backendLogger.Flush()
	log = btclog.NewSubsystemLogger(backendLogger, "")
	btcdb.UseLogger(log)

	var testnet string
	if cfg.TestNet3 {
		testnet = "testnet"
	} else {
		testnet = "mainnet"
	}

	cfg.DataDir = filepath.Join(cfg.DataDir, testnet)

	blockDbNamePrefix := "blocks"
	dbName := blockDbNamePrefix + "_" + cfg.DbType
	if cfg.DbType == "sqlite" {
		dbName = dbName + ".db"
	}
	dbPath := filepath.Join(cfg.DataDir, dbName)

	log.Infof("loading db %v", cfg.DbType)
	db, err := btcdb.OpenDB(cfg.DbType, dbPath)
	if err != nil {
		log.Warnf("db open failed: %v", err)
		return
	}
	defer db.Close()
	log.Infof("db load complete")

	signal_chan := make(chan os.Signal, 1)
	signal.Notify(signal_chan, os.Interrupt, os.Kill)

	go func() {
		_ = <-signal_chan
		log.Info(spew.Sdump(duplicates))
		// this will have to be reworked while parallelizing
		backendLogger.Flush()
		db.Close()
		os.Exit(2)
	}()

	_, max_heigth, err := db.NewestSha()
	if err != nil {
		log.Warnf("db NewestSha failed: %v", err)
		return
	}
	log.Infof("max_heigth: %v", max_heigth)

	last_time := time.Now()
	for h := int64(0); h < max_heigth; h++ {
		// TODO: parallelize
		err = DumpBlock(db, h)
		if err != nil {
			log.Warnf("Failed to dump block %v, err %v", h, err)
		}
		if blocksCounter%10000 == 0 {
			log.Infof("%v blocks processed in %v, %v signatures stored",
				blocksCounter, time.Since(last_time), sigCounter)
			last_time = time.Now()
		}
	}

	log.Info(spew.Sdump(duplicates))
}

func parseData(SignatureScript []uint8) ([]uint8, error) {
	opcode := SignatureScript[0]

	if opcode >= 1 && opcode <= 75 {
		sigStr := SignatureScript[1 : opcode+1]
		return sigStr, nil
	}

	// TODO: OP_PUSHDATA1 OP_PUSHDATA2 OP_PUSHDATA3

	return nil, fmt.Errorf("the first opcode (%x) is not a data push", opcode)
}

func DumpBlock(db btcdb.Db, height int64) error {
	sha, err := db.FetchBlockShaByHeight(height)
	if err != nil {
		return err
	}
	blk, err := db.FetchBlockBySha(sha)
	if err != nil {
		return err
	}
	rblk, err := blk.Bytes()
	if err != nil {
		return err
	}
	blkid := blk.Height()
	if blkid != height {
		return fmt.Errorf("WHAT!?")
	}

	log.Debugf("Block %v depth %v", sha, blkid)

	log.Debugf("Block %v depth %v %v", sha, blkid, spew.Sdump(rblk))
	mblk := blk.MsgBlock()
	log.Debugf("Block %v depth %v %v", sha, blkid, spew.Sdump(mblk))

	log.Debugf("Num transactions %v", len(mblk.Transactions))
	for i, tx := range mblk.Transactions {

		txsha, err := tx.TxSha()
		if err != nil {
			log.Warnf("Block %v (%v)", blkid, sha)
			log.Warnf("tx %v (%v)", i, &txsha)
			log.Warnf("Error: %v", err)
			continue
		}

		log.Debugf("tx %v: %v", i, &txsha)

		if btcchain.IsCoinBase(btcutil.NewTx(tx)) {
			log.Debugf("tx %v: skipping (coinbase)", i)
			continue
		}

		for t, txin := range tx.TxIn {
			log.Debugf("tx %v: TxIn %v: SignatureScript: %v", i, t, spew.Sdump(txin.SignatureScript))

			sigStr, err := parseData(txin.SignatureScript)
			if err != nil {
				log.Warnf("Block %v (%v)", blkid, sha)
				log.Warnf("tx %v (%v)", i, &txsha)
				log.Warnf("txin %v (parseData)", t)
				log.Warnf("Error: %v", err)
				log.Info(spew.Sdump(txin.SignatureScript))
				continue
			}
			log.Debugf("tx %v: TxIn %v: sigStr: %v", i, t, spew.Sdump(sigStr))

			signature, err := btcec.ParseSignature(sigStr, btcec.S256())
			if err != nil {
				log.Warnf("Block %v (%v)", blkid, sha)
				log.Warnf("tx %v (%v)", i, &txsha)
				log.Warnf("txin %v (ParseSignature)", t)
				log.Warnf("Error: %v", err)
				log.Info(spew.Sdump(sigStr))
				continue

			}
			log.Debugf("tx %v: TxIn %v: signature: %v", i, t, spew.Sdump(signature))

			signatureString := signature.R.String()
			if rValuesMap[signatureString] != int64(0) {
				log.Infof("%v:%v:%v DUPLICATE FOUND: %v",
					blkid, txsha.String()[:5], t, rValuesMap[signatureString])
				if len(duplicates[signatureString]) == 0 {
					duplicates[signatureString] = append(duplicates[signatureString], rValuesMap[signatureString])
				}
				duplicates[signatureString] = append(duplicates[signatureString], blkid)
			} else {
				rValuesMap[signatureString] = blkid
				sigCounter++
			}

		}

	}

	blocksCounter++

	return nil
}
<|endoftext|>