<commit_before>// +build windows

package main

import (
	"fmt"
	"net/http"
	"os"
	"sort"
	"strconv"
	"strings"
	"sync"
	"time"

	"golang.org/x/sys/windows/svc"

	"github.com/StackExchange/wmi"
	"github.com/martinlindhe/wmi_exporter/collector"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promhttp"
	"github.com/prometheus/common/log"
	"github.com/prometheus/common/version"
	"gopkg.in/alecthomas/kingpin.v2"
)

// WmiCollector implements the prometheus.Collector interface.
type WmiCollector struct {
	maxScrapeDuration time.Duration
	collectors        map[string]collector.Collector
}

const (
	defaultCollectors            = "cpu,cs,logical_disk,net,os,service,system,textfile"
	defaultCollectorsPlaceholder = "[defaults]"
	serviceName                  = "wmi_exporter"
)

var (
	scrapeDurationDesc = prometheus.NewDesc(
		prometheus.BuildFQName(collector.Namespace, "exporter", "collector_duration_seconds"),
		"wmi_exporter: Duration of a collection.",
		[]string{"collector"},
		nil,
	)
	scrapeSuccessDesc = prometheus.NewDesc(
		prometheus.BuildFQName(collector.Namespace, "exporter", "collector_success"),
		"wmi_exporter: Whether the collector was successful.",
		[]string{"collector"},
		nil,
	)
	scrapeTimeoutDesc = prometheus.NewDesc(
		prometheus.BuildFQName(collector.Namespace, "exporter", "collector_timeout"),
		"wmi_exporter: Whether the collector timed out.",
		[]string{"collector"},
		nil,
	)
	snapshotDuration = prometheus.NewDesc(
		prometheus.BuildFQName(collector.Namespace, "exporter", "perflib_snapshot_duration_seconds"),
		"Duration of perflib snapshot capture",
		nil,
		nil,
	)

	// This can be removed when client_golang exposes this on Windows
	// (See https://github.com/prometheus/client_golang/issues/376)
	startTime     = float64(time.Now().Unix())
	startTimeDesc = prometheus.NewDesc(
		"process_start_time_seconds",
		"Start time of the process since unix epoch in seconds.",
		nil,
		nil,
	)
)

// Describe sends all the descriptors of the collectors included to
// the provided channel.
func (coll WmiCollector) Describe(ch chan<- *prometheus.Desc) {
	ch <- scrapeDurationDesc
	ch <- scrapeSuccessDesc
}

// Collect sends the collected metrics from each of the collectors to
// prometheus.
func (coll WmiCollector) Collect(ch chan<- prometheus.Metric) {
	t := time.Now()
	scrapeContext, err := collector.PrepareScrapeContext()
	ch <- prometheus.MustNewConstMetric(
		snapshotDuration,
		prometheus.GaugeValue,
		time.Since(t).Seconds(),
	)
	if err != nil {
		ch <- prometheus.NewInvalidMetric(scrapeSuccessDesc, fmt.Errorf("failed to prepare scrape: %v", err))
		return
	}

	remainingCollectors := make(map[string]bool)
	for name := range coll.collectors {
		remainingCollectors[name] = true
	}

	metricsBuffer := make(chan prometheus.Metric)
	allDone := make(chan struct{})
	stopped := false
	go func() {
		for {
			select {
			case m, ok := <-metricsBuffer:
				if ok && !stopped {
					ch <- m
				}
			case <-allDone:
				return
			}
		}
	}()

	wg := sync.WaitGroup{}
	wg.Add(len(coll.collectors))
	go func() {
		wg.Wait()
		close(allDone)
		close(metricsBuffer)
	}()

	for name, c := range coll.collectors {
		go func(name string, c collector.Collector) {
			defer wg.Done()
			execute(name, c, scrapeContext, metricsBuffer)
			delete(remainingCollectors, name)
		}(name, c)
	}

	ch <- prometheus.MustNewConstMetric(
		startTimeDesc,
		prometheus.CounterValue,
		startTime,
	)

	select {
	case <-allDone:
		stopped = true
		return
	case <-time.After(coll.maxScrapeDuration):
		stopped = true
		remainingCollectorNames := make([]string, 0, len(remainingCollectors))
		for rc := range remainingCollectors {
			remainingCollectorNames = append(remainingCollectorNames, rc)
		}
		log.Warn("Collection timed out, still waiting for ", remainingCollectorNames)
		for name := range remainingCollectors {
			ch <- prometheus.MustNewConstMetric(
				scrapeSuccessDesc,
				prometheus.GaugeValue,
				0.0,
				name,
			)
			ch <- prometheus.MustNewConstMetric(
				scrapeTimeoutDesc,
				prometheus.GaugeValue,
				1.0,
				name,
			)
		}
	}
}

func filterAvailableCollectors(collectors string) string {
	var availableCollectors []string
	for _, c := range strings.Split(collectors, ",") {
		_, ok := collector.Factories[c]
		if ok {
			availableCollectors = append(availableCollectors, c)
		}
	}
	return strings.Join(availableCollectors, ",")
}

func execute(name string, c collector.Collector, ctx *collector.ScrapeContext, ch chan<- prometheus.Metric) {
	begin := time.Now()
	err := c.Collect(ctx, ch)
	duration := time.Since(begin)
	var success float64

	if err != nil {
		log.Errorf("collector %s failed after %fs: %s", name, duration.Seconds(), err)
		success = 0
	} else {
		log.Debugf("collector %s succeeded after %fs.", name, duration.Seconds())
		success = 1
	}
	ch <- prometheus.MustNewConstMetric(
		scrapeDurationDesc,
		prometheus.GaugeValue,
		duration.Seconds(),
		name,
	)
	ch <- prometheus.MustNewConstMetric(
		scrapeSuccessDesc,
		prometheus.GaugeValue,
		success,
		name,
	)
	ch <- prometheus.MustNewConstMetric(
		scrapeTimeoutDesc,
		prometheus.GaugeValue,
		0.0,
		name,
	)
}

func expandEnabledCollectors(enabled string) []string {
	expanded := strings.Replace(enabled, defaultCollectorsPlaceholder, defaultCollectors, -1)
	separated := strings.Split(expanded, ",")
	unique := map[string]bool{}
	for _, s := range separated {
		if s != "" {
			unique[s] = true
		}
	}
	result := make([]string, 0, len(unique))
	for s := range unique {
		result = append(result, s)
	}
	return result
}

func loadCollectors(list string) (map[string]collector.Collector, error) {
	collectors := map[string]collector.Collector{}
	enabled := expandEnabledCollectors(list)

	for _, name := range enabled {
		fn, ok := collector.Factories[name]
		if !ok {
			return nil, fmt.Errorf("collector '%s' not available", name)
		}
		c, err := fn()
		if err != nil {
			return nil, err
		}
		collectors[name] = c
	}
	return collectors, nil
}

func initWbem() {
	// This initialization prevents a memory leak on WMF 5+. See
	// https://github.com/martinlindhe/wmi_exporter/issues/77 and linked issues
	// for details.
	log.Debugf("Initializing SWbemServices")
	s, err := wmi.InitializeSWbemServices(wmi.DefaultClient)
	if err != nil {
		log.Fatal(err)
	}
	wmi.DefaultClient.AllowMissingFields = true
	wmi.DefaultClient.SWbemServicesClient = s
}

func main() {
	var (
		listenAddress = kingpin.Flag(
			"telemetry.addr",
			"host:port for WMI exporter.",
		).Default(":9182").String()
		metricsPath = kingpin.Flag(
			"telemetry.path",
			"URL path for surfacing collected metrics.",
		).Default("/metrics").String()
		enabledCollectors = kingpin.Flag(
			"collectors.enabled",
			"Comma-separated list of collectors to use. Use '[defaults]' as a placeholder for all the collectors enabled by default.").
			Default(filterAvailableCollectors(defaultCollectors)).String()
		printCollectors = kingpin.Flag(
			"collectors.print",
			"If true, print available collectors and exit.",
		).Bool()
		timeoutMargin = kingpin.Flag(
			"scrape.timeout-margin",
			"Seconds to subtract from the timeout allowed by the client. Tune to allow for overhead or high loads.",
		).Default("0.5").Float64()
	)

	log.AddFlags(kingpin.CommandLine)
	kingpin.Version(version.Print("wmi_exporter"))
	kingpin.HelpFlag.Short('h')
	kingpin.Parse()

	if *printCollectors {
		collectorNames := make(sort.StringSlice, 0, len(collector.Factories))
		for n := range collector.Factories {
			collectorNames = append(collectorNames, n)
		}
		collectorNames.Sort()
		fmt.Printf("Available collectors:\n")
		for _, n := range collectorNames {
			fmt.Printf(" - %s\n", n)
		}
		return
	}

	initWbem()

	isInteractive, err := svc.IsAnInteractiveSession()
	if err != nil {
		log.Fatal(err)
	}

	stopCh := make(chan bool)
	if !isInteractive {
		go func() {
			err = svc.Run(serviceName, &wmiExporterService{stopCh: stopCh})
			if err != nil {
				log.Errorf("Failed to start service: %v", err)
			}
		}()
	}

	collectors, err := loadCollectors(*enabledCollectors)
	if err != nil {
		log.Fatalf("Couldn't load collectors: %s", err)
	}

	log.Infof("Enabled collectors: %v", strings.Join(keys(collectors), ", "))

	h := &metricsHandler{
		timeoutMargin: *timeoutMargin,
		collectorFactory: func(timeout time.Duration) *WmiCollector {
			return &WmiCollector{
				collectors:        collectors,
				maxScrapeDuration: timeout,
			}
		},
	}

	http.Handle(*metricsPath, h)
	http.HandleFunc("/health", healthCheck)
	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		http.Redirect(w, r, *metricsPath, http.StatusMovedPermanently)
	})

	log.Infoln("Starting WMI exporter", version.Info())
	log.Infoln("Build context", version.BuildContext())

	go func() {
		log.Infoln("Starting server on", *listenAddress)
		log.Fatalf("cannot start WMI exporter: %s", http.ListenAndServe(*listenAddress, nil))
	}()

	for {
		if <-stopCh {
			log.Info("Shutting down WMI exporter")
			break
		}
	}
}

func healthCheck(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "application/json")
	_, err := fmt.Fprintln(w, `{"status":"ok"}`)
	if err != nil {
		log.Debugf("Failed to write to stream: %v", err)
	}
}

func keys(m map[string]collector.Collector) []string {
	ret := make([]string, 0, len(m))
	for key := range m {
		ret = append(ret, key)
	}
	return ret
}

type wmiExporterService struct {
	stopCh chan<- bool
}

func (s *wmiExporterService) Execute(args []string, r <-chan svc.ChangeRequest, changes chan<- svc.Status) (ssec bool, errno uint32) {
	const cmdsAccepted = svc.AcceptStop | svc.AcceptShutdown
	changes <- svc.Status{State: svc.StartPending}
	changes <- svc.Status{State: svc.Running, Accepts: cmdsAccepted}
loop:
	for {
		select {
		case c := <-r:
			switch c.Cmd {
			case svc.Interrogate:
				changes <- c.CurrentStatus
			case svc.Stop, svc.Shutdown:
				s.stopCh <- true
				break loop
			default:
				log.Error(fmt.Sprintf("unexpected control request #%d", c))
			}
		}
	}
	changes <- svc.Status{State: svc.StopPending}
	return
}

type metricsHandler struct {
	timeoutMargin    float64
	collectorFactory func(timeout time.Duration) *WmiCollector
}

func (mh *metricsHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	const defaultTimeout = 10.0

	var timeoutSeconds float64
	if v := r.Header.Get("X-Prometheus-Scrape-Timeout-Seconds"); v != "" {
		var err error
		timeoutSeconds, err = strconv.ParseFloat(v, 64)
		if err != nil {
			log.Warnf("Couldn't parse X-Prometheus-Scrape-Timeout-Seconds: %q. Defaulting timeout to %d", v, defaultTimeout)
		}
	}
	if timeoutSeconds == 0 {
		timeoutSeconds = defaultTimeout
	}
	timeoutSeconds = timeoutSeconds - mh.timeoutMargin

	reg := prometheus.NewRegistry()
	reg.MustRegister(mh.collectorFactory(time.Duration(timeoutSeconds * float64(time.Second))))
	reg.MustRegister(
		prometheus.NewProcessCollector(os.Getpid(), ""),
		prometheus.NewGoCollector(),
		version.NewCollector("wmi_exporter"),
	)

	h := promhttp.HandlerFor(reg, promhttp.HandlerOpts{})
	h.ServeHTTP(w, r)
}
<commit_msg>Fix float-format<commit_after>// +build windows

package main

import (
	"fmt"
	"net/http"
	"os"
	"sort"
	"strconv"
	"strings"
	"sync"
	"time"

	"golang.org/x/sys/windows/svc"

	"github.com/StackExchange/wmi"
	"github.com/martinlindhe/wmi_exporter/collector"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promhttp"
	"github.com/prometheus/common/log"
	"github.com/prometheus/common/version"
	"gopkg.in/alecthomas/kingpin.v2"
)

// WmiCollector implements the prometheus.Collector interface.
type WmiCollector struct {
	maxScrapeDuration time.Duration
	collectors        map[string]collector.Collector
}

const (
	defaultCollectors            = "cpu,cs,logical_disk,net,os,service,system,textfile"
	defaultCollectorsPlaceholder = "[defaults]"
	serviceName                  = "wmi_exporter"
)

var (
	scrapeDurationDesc = prometheus.NewDesc(
		prometheus.BuildFQName(collector.Namespace, "exporter", "collector_duration_seconds"),
		"wmi_exporter: Duration of a collection.",
		[]string{"collector"},
		nil,
	)
	scrapeSuccessDesc = prometheus.NewDesc(
		prometheus.BuildFQName(collector.Namespace, "exporter", "collector_success"),
		"wmi_exporter: Whether the collector was successful.",
		[]string{"collector"},
		nil,
	)
	scrapeTimeoutDesc = prometheus.NewDesc(
		prometheus.BuildFQName(collector.Namespace, "exporter", "collector_timeout"),
		"wmi_exporter: Whether the collector timed out.",
		[]string{"collector"},
		nil,
	)
	snapshotDuration = prometheus.NewDesc(
		prometheus.BuildFQName(collector.Namespace, "exporter", "perflib_snapshot_duration_seconds"),
		"Duration of perflib snapshot capture",
		nil,
		nil,
	)

	// This can be removed when client_golang exposes this on Windows
	// (See https://github.com/prometheus/client_golang/issues/376)
	startTime     = float64(time.Now().Unix())
	startTimeDesc = prometheus.NewDesc(
		"process_start_time_seconds",
		"Start time of the process since unix epoch in seconds.",
		nil,
		nil,
	)
)

// Describe sends all the descriptors of the collectors included to
// the provided channel.
func (coll WmiCollector) Describe(ch chan<- *prometheus.Desc) {
	ch <- scrapeDurationDesc
	ch <- scrapeSuccessDesc
}

// Collect sends the collected metrics from each of the collectors to
// prometheus.
func (coll WmiCollector) Collect(ch chan<- prometheus.Metric) {
	t := time.Now()
	scrapeContext, err := collector.PrepareScrapeContext()
	ch <- prometheus.MustNewConstMetric(
		snapshotDuration,
		prometheus.GaugeValue,
		time.Since(t).Seconds(),
	)
	if err != nil {
		ch <- prometheus.NewInvalidMetric(scrapeSuccessDesc, fmt.Errorf("failed to prepare scrape: %v", err))
		return
	}

	remainingCollectors := make(map[string]bool)
	for name := range coll.collectors {
		remainingCollectors[name] = true
	}

	metricsBuffer := make(chan prometheus.Metric)
	allDone := make(chan struct{})
	stopped := false
	go func() {
		for {
			select {
			case m, ok := <-metricsBuffer:
				if ok && !stopped {
					ch <- m
				}
			case <-allDone:
				return
			}
		}
	}()

	wg := sync.WaitGroup{}
	wg.Add(len(coll.collectors))
	go func() {
		wg.Wait()
		close(allDone)
		close(metricsBuffer)
	}()

	for name, c := range coll.collectors {
		go func(name string, c collector.Collector) {
			defer wg.Done()
			execute(name, c, scrapeContext, metricsBuffer)
			delete(remainingCollectors, name)
		}(name, c)
	}

	ch <- prometheus.MustNewConstMetric(
		startTimeDesc,
		prometheus.CounterValue,
		startTime,
	)

	select {
	case <-allDone:
		stopped = true
		return
	case <-time.After(coll.maxScrapeDuration):
		stopped = true
		remainingCollectorNames := make([]string, 0, len(remainingCollectors))
		for rc := range remainingCollectors {
			remainingCollectorNames = append(remainingCollectorNames, rc)
		}
		log.Warn("Collection timed out, still waiting for ", remainingCollectorNames)
		for name := range remainingCollectors {
			ch <- prometheus.MustNewConstMetric(
				scrapeSuccessDesc,
				prometheus.GaugeValue,
				0.0,
				name,
			)
			ch <- prometheus.MustNewConstMetric(
				scrapeTimeoutDesc,
				prometheus.GaugeValue,
				1.0,
				name,
			)
		}
	}
}

func filterAvailableCollectors(collectors string) string {
	var availableCollectors []string
	for _, c := range strings.Split(collectors, ",") {
		_, ok := collector.Factories[c]
		if ok {
			availableCollectors = append(availableCollectors, c)
		}
	}
	return strings.Join(availableCollectors, ",")
}

func execute(name string, c collector.Collector, ctx *collector.ScrapeContext, ch chan<- prometheus.Metric) {
	begin := time.Now()
	err := c.Collect(ctx, ch)
	duration := time.Since(begin)
	var success float64

	if err != nil {
		log.Errorf("collector %s failed after %fs: %s", name, duration.Seconds(), err)
		success = 0
	} else {
		log.Debugf("collector %s succeeded after %fs.", name, duration.Seconds())
		success = 1
	}
	ch <- prometheus.MustNewConstMetric(
		scrapeDurationDesc,
		prometheus.GaugeValue,
		duration.Seconds(),
		name,
	)
	ch <- prometheus.MustNewConstMetric(
		scrapeSuccessDesc,
		prometheus.GaugeValue,
		success,
		name,
	)
	ch <- prometheus.MustNewConstMetric(
		scrapeTimeoutDesc,
		prometheus.GaugeValue,
		0.0,
		name,
	)
}

func expandEnabledCollectors(enabled string) []string {
	expanded := strings.Replace(enabled, defaultCollectorsPlaceholder, defaultCollectors, -1)
	separated := strings.Split(expanded, ",")
	unique := map[string]bool{}
	for _, s := range separated {
		if s != "" {
			unique[s] = true
		}
	}
	result := make([]string, 0, len(unique))
	for s := range unique {
		result = append(result, s)
	}
	return result
}

func loadCollectors(list string) (map[string]collector.Collector, error) {
	collectors := map[string]collector.Collector{}
	enabled := expandEnabledCollectors(list)

	for _, name := range enabled {
		fn, ok := collector.Factories[name]
		if !ok {
			return nil, fmt.Errorf("collector '%s' not available", name)
		}
		c, err := fn()
		if err != nil {
			return nil, err
		}
		collectors[name] = c
	}
	return collectors, nil
}

func initWbem() {
	// This initialization prevents a memory leak on WMF 5+. See
	// https://github.com/martinlindhe/wmi_exporter/issues/77 and linked issues
	// for details.
	log.Debugf("Initializing SWbemServices")
	s, err := wmi.InitializeSWbemServices(wmi.DefaultClient)
	if err != nil {
		log.Fatal(err)
	}
	wmi.DefaultClient.AllowMissingFields = true
	wmi.DefaultClient.SWbemServicesClient = s
}

func main() {
	var (
		listenAddress = kingpin.Flag(
			"telemetry.addr",
			"host:port for WMI exporter.",
		).Default(":9182").String()
		metricsPath = kingpin.Flag(
			"telemetry.path",
			"URL path for surfacing collected metrics.",
		).Default("/metrics").String()
		enabledCollectors = kingpin.Flag(
			"collectors.enabled",
			"Comma-separated list of collectors to use. Use '[defaults]' as a placeholder for all the collectors enabled by default.").
			Default(filterAvailableCollectors(defaultCollectors)).String()
		printCollectors = kingpin.Flag(
			"collectors.print",
			"If true, print available collectors and exit.",
		).Bool()
		timeoutMargin = kingpin.Flag(
			"scrape.timeout-margin",
			"Seconds to subtract from the timeout allowed by the client. Tune to allow for overhead or high loads.",
		).Default("0.5").Float64()
	)

	log.AddFlags(kingpin.CommandLine)
	kingpin.Version(version.Print("wmi_exporter"))
	kingpin.HelpFlag.Short('h')
	kingpin.Parse()

	if *printCollectors {
		collectorNames := make(sort.StringSlice, 0, len(collector.Factories))
		for n := range collector.Factories {
			collectorNames = append(collectorNames, n)
		}
		collectorNames.Sort()
		fmt.Printf("Available collectors:\n")
		for _, n := range collectorNames {
			fmt.Printf(" - %s\n", n)
		}
		return
	}

	initWbem()

	isInteractive, err := svc.IsAnInteractiveSession()
	if err != nil {
		log.Fatal(err)
	}

	stopCh := make(chan bool)
	if !isInteractive {
		go func() {
			err = svc.Run(serviceName, &wmiExporterService{stopCh: stopCh})
			if err != nil {
				log.Errorf("Failed to start service: %v", err)
			}
		}()
	}

	collectors, err := loadCollectors(*enabledCollectors)
	if err != nil {
		log.Fatalf("Couldn't load collectors: %s", err)
	}

	log.Infof("Enabled collectors: %v", strings.Join(keys(collectors), ", "))

	h := &metricsHandler{
		timeoutMargin: *timeoutMargin,
		collectorFactory: func(timeout time.Duration) *WmiCollector {
			return &WmiCollector{
				collectors:        collectors,
				maxScrapeDuration: timeout,
			}
		},
	}

	http.Handle(*metricsPath, h)
	http.HandleFunc("/health", healthCheck)
	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		http.Redirect(w, r, *metricsPath, http.StatusMovedPermanently)
	})

	log.Infoln("Starting WMI exporter", version.Info())
	log.Infoln("Build context", version.BuildContext())

	go func() {
		log.Infoln("Starting server on", *listenAddress)
		log.Fatalf("cannot start WMI exporter: %s", http.ListenAndServe(*listenAddress, nil))
	}()

	for {
		if <-stopCh {
			log.Info("Shutting down WMI exporter")
			break
		}
	}
}

func healthCheck(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "application/json")
	_, err := fmt.Fprintln(w, `{"status":"ok"}`)
	if err != nil {
		log.Debugf("Failed to write to stream: %v", err)
	}
}

func keys(m map[string]collector.Collector) []string {
	ret := make([]string, 0, len(m))
	for key := range m {
		ret = append(ret, key)
	}
	return ret
}

type wmiExporterService struct {
	stopCh chan<- bool
}

func (s *wmiExporterService) Execute(args []string, r <-chan svc.ChangeRequest, changes chan<- svc.Status) (ssec bool, errno uint32) {
	const cmdsAccepted = svc.AcceptStop | svc.AcceptShutdown
	changes <- svc.Status{State: svc.StartPending}
	changes <- svc.Status{State: svc.Running, Accepts: cmdsAccepted}
loop:
	for {
		select {
		case c := <-r:
			switch c.Cmd {
			case svc.Interrogate:
				changes <- c.CurrentStatus
			case svc.Stop, svc.Shutdown:
				s.stopCh <- true
				break loop
			default:
				log.Error(fmt.Sprintf("unexpected control request #%d", c))
			}
		}
	}
	changes <- svc.Status{State: svc.StopPending}
	return
}

type metricsHandler struct {
	timeoutMargin    float64
	collectorFactory func(timeout time.Duration) *WmiCollector
}

func (mh *metricsHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	const defaultTimeout = 10.0

	var timeoutSeconds float64
	if v := r.Header.Get("X-Prometheus-Scrape-Timeout-Seconds"); v != "" {
		var err error
		timeoutSeconds, err = strconv.ParseFloat(v, 64)
		if err != nil {
			log.Warnf("Couldn't parse X-Prometheus-Scrape-Timeout-Seconds: %q. Defaulting timeout to %f", v, defaultTimeout)
		}
	}
	if timeoutSeconds == 0 {
		timeoutSeconds = defaultTimeout
	}
	timeoutSeconds = timeoutSeconds - mh.timeoutMargin

	reg := prometheus.NewRegistry()
	reg.MustRegister(mh.collectorFactory(time.Duration(timeoutSeconds * float64(time.Second))))
	reg.MustRegister(
		prometheus.NewProcessCollector(os.Getpid(), ""),
		prometheus.NewGoCollector(),
		version.NewCollector("wmi_exporter"),
	)

	h := promhttp.HandlerFor(reg, promhttp.HandlerOpts{})
	h.ServeHTTP(w, r)
}
<|endoftext|>


<commit_before>// Copyright (c) 2015 The btcsuite developers
// Use of this source code is governed by an ISC
// license that can be found in the LICENSE file.

package txscript

import (
	"sync"

	"github.com/btcsuite/btcd/btcec"
	"github.com/btcsuite/btcd/wire"
)

// sigCacheEntry represents an entry in the SigCache. Entries within the
// SigCache are keyed according to the sigHash of the signature. In the
// scenario of a cache-hit (according to the sigHash), an additional comparison
// of the signature, and public key will be executed in order to ensure a complete
// match. In the occasion that two sigHashes collide, the newer sigHash will
// simply overwrite the existing entry.
type sigCacheEntry struct {
	sig    *btcec.Signature
	pubKey *btcec.PublicKey
}

// SigCache implements an ECDSA signature verification cache with a randomized
// entry eviction policy. Only valid signatures will be added to the cache. The
// benefits of SigCache are two fold. Firstly, usage of SigCache mitigates a DoS
// attack wherein an attack causes a victim's client to hang due to worst-case
// behavior triggered while processing attacker crafted invalid transactions. A
// detailed description of the mitigated DoS attack can be found here:
// https://bitslog.wordpress.com/2013/01/23/fixed-bitcoin-vulnerability-explanation-why-the-signature-cache-is-a-dos-protection/.
// Secondly, usage of the SigCache introduces a signature verification
// optimization which speeds up the validation of transactions within a block,
// if they've already been seen and verified within the mempool.
type SigCache struct {
	sync.RWMutex
	validSigs  map[wire.ShaHash]sigCacheEntry
	maxEntries uint
}

// NewSigCache creates and initializes a new instance of SigCache. Its sole
// parameter 'maxEntries' represents the maximum number of entries allowed to
// exist in the SigCache at any particular moment. Random entries are evicted
// to make room for new entries that would cause the number of entries in the
// cache to exceed the max.
func NewSigCache(maxEntries uint) *SigCache {
	return &SigCache{
		validSigs:  make(map[wire.ShaHash]sigCacheEntry, maxEntries),
		maxEntries: maxEntries,
	}
}

// Exists returns true if an existing entry of 'sig' over 'sigHash' for public
// key 'pubKey' is found within the SigCache. Otherwise, false is returned.
//
// NOTE: This function is safe for concurrent access. Readers won't be blocked
// unless there exists a writer, adding an entry to the SigCache.
func (s *SigCache) Exists(sigHash wire.ShaHash, sig *btcec.Signature, pubKey *btcec.PublicKey) bool {
	s.RLock()
	defer s.RUnlock()

	if entry, ok := s.validSigs[sigHash]; ok {
		return entry.pubKey.IsEqual(pubKey) && entry.sig.IsEqual(sig)
	}

	return false
}

// Add adds an entry for a signature over 'sigHash' under public key 'pubKey'
// to the signature cache. In the event that the SigCache is 'full', an
// existing entry is randomly chosen to be evicted in order to make space for
// the new entry.
//
// NOTE: This function is safe for concurrent access. Writers will block
// simultaneous readers until function execution has concluded.
func (s *SigCache) Add(sigHash wire.ShaHash, sig *btcec.Signature, pubKey *btcec.PublicKey) {
	s.Lock()
	defer s.Unlock()

	if s.maxEntries <= 0 {
		return
	}

	// If adding this new entry will put us over the max number of allowed
	// entries, then evict an entry.
	if uint(len(s.validSigs)+1) > s.maxEntries {
		// Remove a random entry from the map relaying on the random
		// starting point of Go's map iteration. It's worth noting that
		// the random iteration starting point is not 100% guaranteed
		// by the spec, however most Go compilers support it.
		// Ultimately, the iteration order isn't important here because
		// in order to manipulate which items are evicted, an adversary
		// would need to be able to execute preimage attacks on the
		// hashing function in order to start eviction at a specific
		// entry.
		for sigEntry := range s.validSigs {
			delete(s.validSigs, sigEntry)
			break
		}
	}
	s.validSigs[sigHash] = sigCacheEntry{sig, pubKey}
}
<commit_msg>txscript: Fix typo. (#700)<commit_after>// Copyright (c) 2015-2016 The btcsuite developers
// Use of this source code is governed by an ISC
// license that can be found in the LICENSE file.

package txscript

import (
	"sync"

	"github.com/btcsuite/btcd/btcec"
	"github.com/btcsuite/btcd/wire"
)

// sigCacheEntry represents an entry in the SigCache. Entries within the
// SigCache are keyed according to the sigHash of the signature. In the
// scenario of a cache-hit (according to the sigHash), an additional comparison
// of the signature, and public key will be executed in order to ensure a complete
// match. In the occasion that two sigHashes collide, the newer sigHash will
// simply overwrite the existing entry.
type sigCacheEntry struct {
	sig    *btcec.Signature
	pubKey *btcec.PublicKey
}

// SigCache implements an ECDSA signature verification cache with a randomized
// entry eviction policy. Only valid signatures will be added to the cache. The
// benefits of SigCache are two fold. Firstly, usage of SigCache mitigates a DoS
// attack wherein an attack causes a victim's client to hang due to worst-case
// behavior triggered while processing attacker crafted invalid transactions. A
// detailed description of the mitigated DoS attack can be found here:
// https://bitslog.wordpress.com/2013/01/23/fixed-bitcoin-vulnerability-explanation-why-the-signature-cache-is-a-dos-protection/.
// Secondly, usage of the SigCache introduces a signature verification
// optimization which speeds up the validation of transactions within a block,
// if they've already been seen and verified within the mempool.
type SigCache struct {
	sync.RWMutex
	validSigs  map[wire.ShaHash]sigCacheEntry
	maxEntries uint
}

// NewSigCache creates and initializes a new instance of SigCache. Its sole
// parameter 'maxEntries' represents the maximum number of entries allowed to
// exist in the SigCache at any particular moment. Random entries are evicted
// to make room for new entries that would cause the number of entries in the
// cache to exceed the max.
func NewSigCache(maxEntries uint) *SigCache {
	return &SigCache{
		validSigs:  make(map[wire.ShaHash]sigCacheEntry, maxEntries),
		maxEntries: maxEntries,
	}
}

// Exists returns true if an existing entry of 'sig' over 'sigHash' for public
// key 'pubKey' is found within the SigCache. Otherwise, false is returned.
//
// NOTE: This function is safe for concurrent access. Readers won't be blocked
// unless there exists a writer, adding an entry to the SigCache.
func (s *SigCache) Exists(sigHash wire.ShaHash, sig *btcec.Signature, pubKey *btcec.PublicKey) bool {
	s.RLock()
	defer s.RUnlock()

	if entry, ok := s.validSigs[sigHash]; ok {
		return entry.pubKey.IsEqual(pubKey) && entry.sig.IsEqual(sig)
	}

	return false
}

// Add adds an entry for a signature over 'sigHash' under public key 'pubKey'
// to the signature cache. In the event that the SigCache is 'full', an
// existing entry is randomly chosen to be evicted in order to make space for
// the new entry.
//
// NOTE: This function is safe for concurrent access. Writers will block
// simultaneous readers until function execution has concluded.
func (s *SigCache) Add(sigHash wire.ShaHash, sig *btcec.Signature, pubKey *btcec.PublicKey) {
	s.Lock()
	defer s.Unlock()

	if s.maxEntries <= 0 {
		return
	}

	// If adding this new entry will put us over the max number of allowed
	// entries, then evict an entry.
	if uint(len(s.validSigs)+1) > s.maxEntries {
		// Remove a random entry from the map. Relying on the random
		// starting point of Go's map iteration. It's worth noting that
		// the random iteration starting point is not 100% guaranteed
		// by the spec, however most Go compilers support it.
		// Ultimately, the iteration order isn't important here because
		// in order to manipulate which items are evicted, an adversary
		// would need to be able to execute preimage attacks on the
		// hashing function in order to start eviction at a specific
		// entry.
		for sigEntry := range s.validSigs {
			delete(s.validSigs, sigEntry)
			break
		}
	}
	s.validSigs[sigHash] = sigCacheEntry{sig, pubKey}
}
<|endoftext|>


<commit_before>package nvim

import (
	"log"
	"time"
)

// Profile measurement of the time it took to any func and output log file.
// Usage: defer nvim.Profile(time.Now(), "func name")
func Profile(start time.Time, name string) {
	elapsed := time.Since(start)
	log.Printf("%s: %s", name, elapsed)
}
<commit_msg>Profile: Add 2 newline for visibility of the log<commit_after>package nvim

import (
	"log"
	"time"
)

// Profile measurement of the time it took to any func and output log file.
// Usage: defer nvim.Profile(time.Now(), "func name")
func Profile(start time.Time, name string) {
	elapsed := time.Since(start)
	log.Printf("%s: %s\n\n", name, elapsed)
}
<|endoftext|>


<commit_before>//  Copyright (c) 2017 Couchbase, Inc.
//
// 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.

package zap

import (
	"bufio"
	"bytes"
	"encoding/binary"
	"math"
	"os"
	"sort"

	"github.com/Smerity/govarint"
	"github.com/blevesearch/bleve/index/scorch/segment/mem"
	"github.com/couchbase/vellum"
	"github.com/golang/snappy"
)

const version uint32 = 4

const fieldNotUninverted = math.MaxUint64

// PersistSegmentBase persists SegmentBase in the zap file format.
func PersistSegmentBase(sb *SegmentBase, path string) error {
	flag := os.O_RDWR | os.O_CREATE

	f, err := os.OpenFile(path, flag, 0600)
	if err != nil {
		return err
	}

	cleanup := func() {
		_ = f.Close()
		_ = os.Remove(path)
	}

	br := bufio.NewWriter(f)

	_, err = br.Write(sb.mem)
	if err != nil {
		cleanup()
		return err
	}

	err = persistFooter(sb.numDocs, sb.storedIndexOffset, sb.fieldsIndexOffset, sb.docValueOffset,
		sb.chunkFactor, sb.memCRC, br)
	if err != nil {
		cleanup()
		return err
	}

	err = br.Flush()
	if err != nil {
		cleanup()
		return err
	}

	err = f.Sync()
	if err != nil {
		cleanup()
		return err
	}

	err = f.Close()
	if err != nil {
		cleanup()
		return err
	}

	return nil
}

// PersistSegment takes the in-memory segment and persists it to
// the specified path in the zap file format.
func PersistSegment(memSegment *mem.Segment, path string, chunkFactor uint32) error {
	flag := os.O_RDWR | os.O_CREATE

	f, err := os.OpenFile(path, flag, 0600)
	if err != nil {
		return err
	}

	cleanup := func() {
		_ = f.Close()
		_ = os.Remove(path)
	}

	// buffer the output
	br := bufio.NewWriter(f)

	// wrap it for counting (tracking offsets)
	cr := NewCountHashWriter(br)

	numDocs, storedIndexOffset, fieldsIndexOffset, docValueOffset, _, err :=
		persistBase(memSegment, cr, chunkFactor)
	if err != nil {
		cleanup()
		return err
	}

	err = persistFooter(numDocs, storedIndexOffset, fieldsIndexOffset, docValueOffset,
		chunkFactor, cr.Sum32(), cr)
	if err != nil {
		cleanup()
		return err
	}

	err = br.Flush()
	if err != nil {
		cleanup()
		return err
	}

	err = f.Sync()
	if err != nil {
		cleanup()
		return err
	}

	err = f.Close()
	if err != nil {
		cleanup()
		return err
	}

	return nil
}

func persistBase(memSegment *mem.Segment, cr *CountHashWriter, chunkFactor uint32) (
	numDocs, storedIndexOffset, fieldsIndexOffset, docValueOffset uint64,
	dictLocs []uint64, err error) {
	docValueOffset = uint64(fieldNotUninverted)

	if len(memSegment.Stored) > 0 {
		storedIndexOffset, err = persistStored(memSegment, cr)
		if err != nil {
			return 0, 0, 0, 0, nil, err
		}

		freqOffsets, locOffsets, err := persistPostingDetails(memSegment, cr, chunkFactor)
		if err != nil {
			return 0, 0, 0, 0, nil, err
		}

		postingsListLocs, err := persistPostingsLocs(memSegment, cr)
		if err != nil {
			return 0, 0, 0, 0, nil, err
		}

		postingsLocs, err := persistPostingsLists(memSegment, cr, postingsListLocs, freqOffsets, locOffsets)
		if err != nil {
			return 0, 0, 0, 0, nil, err
		}

		dictLocs, err = persistDictionary(memSegment, cr, postingsLocs)
		if err != nil {
			return 0, 0, 0, 0, nil, err
		}

		docValueOffset, err = persistFieldDocValues(memSegment, cr, chunkFactor)
		if err != nil {
			return 0, 0, 0, 0, nil, err
		}
	} else {
		dictLocs = make([]uint64, len(memSegment.FieldsInv))
	}

	fieldsIndexOffset, err = persistFields(memSegment.FieldsInv, cr, dictLocs)
	if err != nil {
		return 0, 0, 0, 0, nil, err
	}

	return uint64(len(memSegment.Stored)), storedIndexOffset, fieldsIndexOffset, docValueOffset,
		dictLocs, nil
}

func persistStored(memSegment *mem.Segment, w *CountHashWriter) (uint64, error) {
	var curr int
	var metaBuf bytes.Buffer
	var data, compressed []byte

	metaEncoder := govarint.NewU64Base128Encoder(&metaBuf)

	docNumOffsets := make(map[int]uint64, len(memSegment.Stored))

	for docNum, storedValues := range memSegment.Stored {
		if docNum != 0 {
			// reset buffer if necessary
			curr = 0
			metaBuf.Reset()
			data = data[:0]
			compressed = compressed[:0]
		}

		st := memSegment.StoredTypes[docNum]
		sp := memSegment.StoredPos[docNum]

		// encode fields in order
		for fieldID := range memSegment.FieldsInv {
			if storedFieldValues, ok := storedValues[uint16(fieldID)]; ok {
				stf := st[uint16(fieldID)]
				spf := sp[uint16(fieldID)]

				var err2 error
				curr, data, err2 = persistStoredFieldValues(fieldID,
					storedFieldValues, stf, spf, curr, metaEncoder, data)
				if err2 != nil {
					return 0, err2
				}
			}
		}

		metaEncoder.Close()
		metaBytes := metaBuf.Bytes()

		// compress the data
		compressed = snappy.Encode(compressed, data)

		// record where we're about to start writing
		docNumOffsets[docNum] = uint64(w.Count())

		// write out the meta len and compressed data len
		_, err := writeUvarints(w, uint64(len(metaBytes)), uint64(len(compressed)))
		if err != nil {
			return 0, err
		}

		// now write the meta
		_, err = w.Write(metaBytes)
		if err != nil {
			return 0, err
		}
		// now write the compressed data
		_, err = w.Write(compressed)
		if err != nil {
			return 0, err
		}
	}

	// return value is the start of the stored index
	rv := uint64(w.Count())
	// now write out the stored doc index
	for docNum := range memSegment.Stored {
		err := binary.Write(w, binary.BigEndian, docNumOffsets[docNum])
		if err != nil {
			return 0, err
		}
	}

	return rv, nil
}

func persistStoredFieldValues(fieldID int,
	storedFieldValues [][]byte, stf []byte, spf [][]uint64,
	curr int, metaEncoder *govarint.Base128Encoder, data []byte) (
	int, []byte, error) {
	for i := 0; i < len(storedFieldValues); i++ {
		// encode field
		_, err := metaEncoder.PutU64(uint64(fieldID))
		if err != nil {
			return 0, nil, err
		}
		// encode type
		_, err = metaEncoder.PutU64(uint64(stf[i]))
		if err != nil {
			return 0, nil, err
		}
		// encode start offset
		_, err = metaEncoder.PutU64(uint64(curr))
		if err != nil {
			return 0, nil, err
		}
		// end len
		_, err = metaEncoder.PutU64(uint64(len(storedFieldValues[i])))
		if err != nil {
			return 0, nil, err
		}
		// encode number of array pos
		_, err = metaEncoder.PutU64(uint64(len(spf[i])))
		if err != nil {
			return 0, nil, err
		}
		// encode all array positions
		for _, pos := range spf[i] {
			_, err = metaEncoder.PutU64(pos)
			if err != nil {
				return 0, nil, err
			}
		}

		data = append(data, storedFieldValues[i]...)
		curr += len(storedFieldValues[i])
	}

	return curr, data, nil
}

func persistPostingDetails(memSegment *mem.Segment, w *CountHashWriter, chunkFactor uint32) ([]uint64, []uint64, error) {
	freqOffsets := make([]uint64, 0, len(memSegment.Postings))
	tfEncoder := newChunkedIntCoder(uint64(chunkFactor), uint64(len(memSegment.Stored)-1))
	for postingID := range memSegment.Postings {
		if postingID != 0 {
			tfEncoder.Reset()
		}
		freqs := memSegment.Freqs[postingID]
		norms := memSegment.Norms[postingID]
		postingsListItr := memSegment.Postings[postingID].Iterator()
		var offset int
		for postingsListItr.HasNext() {
			docNum := uint64(postingsListItr.Next())

			// put freq & norm
			err := tfEncoder.Add(docNum, freqs[offset], uint64(math.Float32bits(norms[offset])))
			if err != nil {
				return nil, nil, err
			}

			offset++
		}

		// record where this postings freq info starts
		freqOffsets = append(freqOffsets, uint64(w.Count()))

		tfEncoder.Close()
		_, err := tfEncoder.Write(w)
		if err != nil {
			return nil, nil, err
		}
	}

	// now do it again for the locations
	locOffsets := make([]uint64, 0, len(memSegment.Postings))
	locEncoder := newChunkedIntCoder(uint64(chunkFactor), uint64(len(memSegment.Stored)-1))
	for postingID := range memSegment.Postings {
		if postingID != 0 {
			locEncoder.Reset()
		}
		freqs := memSegment.Freqs[postingID]
		locfields := memSegment.Locfields[postingID]
		locpos := memSegment.Locpos[postingID]
		locstarts := memSegment.Locstarts[postingID]
		locends := memSegment.Locends[postingID]
		locarraypos := memSegment.Locarraypos[postingID]
		postingsListItr := memSegment.Postings[postingID].Iterator()
		var offset int
		var locOffset int
		for postingsListItr.HasNext() {
			docNum := uint64(postingsListItr.Next())
			n := int(freqs[offset])
			for i := 0; i < n; i++ {
				if len(locfields) > 0 {
					err := locEncoder.Add(docNum, uint64(locfields[locOffset]),
						locpos[locOffset], locstarts[locOffset], locends[locOffset],
						uint64(len(locarraypos[locOffset])))
					if err != nil {
						return nil, nil, err
					}

					// put each array position
					err = locEncoder.Add(docNum, locarraypos[locOffset]...)
					if err != nil {
						return nil, nil, err
					}
				}
				locOffset++
			}
			offset++
		}

		// record where this postings loc info starts
		locOffsets = append(locOffsets, uint64(w.Count()))

		locEncoder.Close()
		_, err := locEncoder.Write(w)
		if err != nil {
			return nil, nil, err
		}
	}

	return freqOffsets, locOffsets, nil
}

func persistPostingsLocs(memSegment *mem.Segment, w *CountHashWriter) (rv []uint64, err error) {
	rv = make([]uint64, 0, len(memSegment.PostingsLocs))
	reuseBufVarint := make([]byte, binary.MaxVarintLen64)
	for postingID := range memSegment.PostingsLocs {
		// record where we start this posting loc
		rv = append(rv, uint64(w.Count()))
		// write out the length and bitmap
		_, err = writeRoaringWithLen(memSegment.PostingsLocs[postingID], w, reuseBufVarint)
		if err != nil {
			return nil, err
		}
	}
	return rv, nil
}

func persistPostingsLists(memSegment *mem.Segment, w *CountHashWriter,
	postingsListLocs, freqOffsets, locOffsets []uint64) (rv []uint64, err error) {
	rv = make([]uint64, 0, len(memSegment.Postings))
	reuseBufVarint := make([]byte, binary.MaxVarintLen64)
	for postingID := range memSegment.Postings {
		// record where we start this posting list
		rv = append(rv, uint64(w.Count()))

		// write out the term info, loc info, and loc posting list offset
		_, err = writeUvarints(w, freqOffsets[postingID],
			locOffsets[postingID], postingsListLocs[postingID])
		if err != nil {
			return nil, err
		}

		// write out the length and bitmap
		_, err = writeRoaringWithLen(memSegment.Postings[postingID], w, reuseBufVarint)
		if err != nil {
			return nil, err
		}
	}
	return rv, nil
}

func persistDictionary(memSegment *mem.Segment, w *CountHashWriter, postingsLocs []uint64) ([]uint64, error) {
	rv := make([]uint64, 0, len(memSegment.DictKeys))

	varintBuf := make([]byte, binary.MaxVarintLen64)

	var buffer bytes.Buffer
	builder, err := vellum.New(&buffer, nil)
	if err != nil {
		return nil, err
	}
	for fieldID, fieldTerms := range memSegment.DictKeys {

		dict := memSegment.Dicts[fieldID]
		// now walk the dictionary in order of fieldTerms (already sorted)
		for _, fieldTerm := range fieldTerms {
			postingID := dict[fieldTerm] - 1
			postingsAddr := postingsLocs[postingID]
			err = builder.Insert([]byte(fieldTerm), postingsAddr)
			if err != nil {
				return nil, err
			}
		}
		err = builder.Close()
		if err != nil {
			return nil, err
		}

		// record where this dictionary starts
		rv = append(rv, uint64(w.Count()))

		vellumData := buffer.Bytes()

		// write out the length of the vellum data
		n := binary.PutUvarint(varintBuf, uint64(len(vellumData)))
		_, err = w.Write(varintBuf[:n])
		if err != nil {
			return nil, err
		}

		// write this vellum to disk
		_, err = w.Write(vellumData)
		if err != nil {
			return nil, err
		}

		// reset buffer and vellum builder
		buffer.Reset()
		err = builder.Reset(&buffer)
		if err != nil {
			return nil, err
		}
	}

	return rv, nil
}

type docIDRange []uint64

func (a docIDRange) Len() int           { return len(a) }
func (a docIDRange) Swap(i, j int)      { a[i], a[j] = a[j], a[i] }
func (a docIDRange) Less(i, j int) bool { return a[i] < a[j] }

func persistDocValues(memSegment *mem.Segment, w *CountHashWriter,
	chunkFactor uint32) (map[uint16]uint64, error) {
	fieldChunkOffsets := make(map[uint16]uint64, len(memSegment.FieldsInv))
	fdvEncoder := newChunkedContentCoder(uint64(chunkFactor), uint64(len(memSegment.Stored)-1))

	var postings *mem.PostingsList
	var postingsItr *mem.PostingsIterator

	for fieldID := range memSegment.DocValueFields {
		field := memSegment.FieldsInv[fieldID]
		docTermMap := make(map[uint64][]byte, 0)
		dict, err := memSegment.Dictionary(field)
		if err != nil {
			return nil, err
		}

		dictItr := dict.Iterator()
		next, err := dictItr.Next()
		for err == nil && next != nil {
			var err1 error
			postings, err1 = dict.(*mem.Dictionary).InitPostingsList(next.Term, nil, postings)
			if err1 != nil {
				return nil, err1
			}

			postingsItr = postings.InitIterator(postingsItr)
			nextPosting, err2 := postingsItr.Next()
			for err2 == nil && nextPosting != nil {
				docNum := nextPosting.Number()
				docTermMap[docNum] = append(append(docTermMap[docNum], []byte(next.Term)...), termSeparator)
				nextPosting, err2 = postingsItr.Next()
			}
			if err2 != nil {
				return nil, err2
			}

			next, err = dictItr.Next()
		}
		if err != nil {
			return nil, err
		}

		// sort wrt to docIDs
		docNumbers := make(docIDRange, 0, len(docTermMap))
		for k := range docTermMap {
			docNumbers = append(docNumbers, k)
		}
		sort.Sort(docNumbers)

		for _, docNum := range docNumbers {
			err = fdvEncoder.Add(docNum, docTermMap[docNum])
			if err != nil {
				return nil, err
			}
		}

		fieldChunkOffsets[fieldID] = uint64(w.Count())
		err = fdvEncoder.Close()
		if err != nil {
			return nil, err
		}
		// persist the doc value details for this field
		_, err = fdvEncoder.Write(w)
		if err != nil {
			return nil, err
		}
		// reseting encoder for the next field
		fdvEncoder.Reset()
	}

	return fieldChunkOffsets, nil
}

func persistFieldDocValues(memSegment *mem.Segment, w *CountHashWriter,
	chunkFactor uint32) (uint64, error) {
	fieldDvOffsets, err := persistDocValues(memSegment, w, chunkFactor)
	if err != nil {
		return 0, err
	}

	fieldDocValuesOffset := uint64(w.Count())
	buf := make([]byte, binary.MaxVarintLen64)
	offset := uint64(0)
	ok := true
	for fieldID := range memSegment.FieldsInv {
		// if the field isn't configured for docValue, then mark
		// the offset accordingly
		if offset, ok = fieldDvOffsets[uint16(fieldID)]; !ok {
			offset = fieldNotUninverted
		}
		n := binary.PutUvarint(buf, uint64(offset))
		_, err := w.Write(buf[:n])
		if err != nil {
			return 0, err
		}
	}

	return fieldDocValuesOffset, nil
}

func NewSegmentBase(memSegment *mem.Segment, chunkFactor uint32) (*SegmentBase, error) {
	var br bytes.Buffer

	cr := NewCountHashWriter(&br)

	numDocs, storedIndexOffset, fieldsIndexOffset, docValueOffset, dictLocs, err :=
		persistBase(memSegment, cr, chunkFactor)
	if err != nil {
		return nil, err
	}

	return InitSegmentBase(br.Bytes(), cr.Sum32(), chunkFactor,
		memSegment.FieldsMap, memSegment.FieldsInv, numDocs,
		storedIndexOffset, fieldsIndexOffset, docValueOffset, dictLocs)
}

func InitSegmentBase(mem []byte, memCRC uint32, chunkFactor uint32,
	fieldsMap map[string]uint16, fieldsInv []string, numDocs uint64,
	storedIndexOffset uint64, fieldsIndexOffset uint64, docValueOffset uint64,
	dictLocs []uint64) (*SegmentBase, error) {
	sb := &SegmentBase{
		mem:               mem,
		memCRC:            memCRC,
		chunkFactor:       chunkFactor,
		fieldsMap:         fieldsMap,
		fieldsInv:         fieldsInv,
		numDocs:           numDocs,
		storedIndexOffset: storedIndexOffset,
		fieldsIndexOffset: fieldsIndexOffset,
		docValueOffset:    docValueOffset,
		dictLocs:          dictLocs,
		fieldDvIterMap:    make(map[uint16]*docValueIterator),
	}

	err := sb.loadDvIterators()
	if err != nil {
		return nil, err
	}

	return sb, nil
}
<commit_msg>scorch zap remove mem.Segment usage from persist / build.go<commit_after>//  Copyright (c) 2017 Couchbase, Inc.
//
// 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.

package zap

import (
	"bufio"
	"math"
	"os"

	"github.com/Smerity/govarint"
)

const version uint32 = 4

const fieldNotUninverted = math.MaxUint64

// PersistSegmentBase persists SegmentBase in the zap file format.
func PersistSegmentBase(sb *SegmentBase, path string) error {
	flag := os.O_RDWR | os.O_CREATE

	f, err := os.OpenFile(path, flag, 0600)
	if err != nil {
		return err
	}

	cleanup := func() {
		_ = f.Close()
		_ = os.Remove(path)
	}

	br := bufio.NewWriter(f)

	_, err = br.Write(sb.mem)
	if err != nil {
		cleanup()
		return err
	}

	err = persistFooter(sb.numDocs, sb.storedIndexOffset, sb.fieldsIndexOffset, sb.docValueOffset,
		sb.chunkFactor, sb.memCRC, br)
	if err != nil {
		cleanup()
		return err
	}

	err = br.Flush()
	if err != nil {
		cleanup()
		return err
	}

	err = f.Sync()
	if err != nil {
		cleanup()
		return err
	}

	err = f.Close()
	if err != nil {
		cleanup()
		return err
	}

	return nil
}

func persistStoredFieldValues(fieldID int,
	storedFieldValues [][]byte, stf []byte, spf [][]uint64,
	curr int, metaEncoder *govarint.Base128Encoder, data []byte) (
	int, []byte, error) {
	for i := 0; i < len(storedFieldValues); i++ {
		// encode field
		_, err := metaEncoder.PutU64(uint64(fieldID))
		if err != nil {
			return 0, nil, err
		}
		// encode type
		_, err = metaEncoder.PutU64(uint64(stf[i]))
		if err != nil {
			return 0, nil, err
		}
		// encode start offset
		_, err = metaEncoder.PutU64(uint64(curr))
		if err != nil {
			return 0, nil, err
		}
		// end len
		_, err = metaEncoder.PutU64(uint64(len(storedFieldValues[i])))
		if err != nil {
			return 0, nil, err
		}
		// encode number of array pos
		_, err = metaEncoder.PutU64(uint64(len(spf[i])))
		if err != nil {
			return 0, nil, err
		}
		// encode all array positions
		for _, pos := range spf[i] {
			_, err = metaEncoder.PutU64(pos)
			if err != nil {
				return 0, nil, err
			}
		}

		data = append(data, storedFieldValues[i]...)
		curr += len(storedFieldValues[i])
	}

	return curr, data, nil
}

func InitSegmentBase(mem []byte, memCRC uint32, chunkFactor uint32,
	fieldsMap map[string]uint16, fieldsInv []string, numDocs uint64,
	storedIndexOffset uint64, fieldsIndexOffset uint64, docValueOffset uint64,
	dictLocs []uint64) (*SegmentBase, error) {
	sb := &SegmentBase{
		mem:               mem,
		memCRC:            memCRC,
		chunkFactor:       chunkFactor,
		fieldsMap:         fieldsMap,
		fieldsInv:         fieldsInv,
		numDocs:           numDocs,
		storedIndexOffset: storedIndexOffset,
		fieldsIndexOffset: fieldsIndexOffset,
		docValueOffset:    docValueOffset,
		dictLocs:          dictLocs,
		fieldDvIterMap:    make(map[uint16]*docValueIterator),
	}

	err := sb.loadDvIterators()
	if err != nil {
		return nil, err
	}

	return sb, nil
}
<|endoftext|>


<commit_before>package turnpike

import (
	"fmt"
	"sync"
	"time"

	logrus "github.com/sirupsen/logrus"
)

const (
	defaultAuthTimeout = 2 * time.Minute
)

// A Realm is a WAMP routing and administrative domain.
//
// Clients that have connected to a WAMP router are joined to a realm and all
// message delivery is handled by the realm.
type Realm struct {
	_                string
	URI              URI
	Broker           Broker
	Dealer           Dealer
	Authorizer       Authorizer
	Interceptor      Interceptor
	CRAuthenticators map[string]CRAuthenticator
	Authenticators   map[string]Authenticator
	// DefaultAuth      func(details map[string]interface{}) (map[string]interface{}, error)
	AuthTimeout time.Duration
	clients     cmap.ConcurrentMap
	localClient *localClient

	lock sync.RWMutex
}

type localClient struct {
	*Client
	sync.Mutex
}

func (r *Realm) getPeer(details map[string]interface{}) (Peer, error) {
	peerA, peerB := localPipe()
	sess := &Session{Peer: peerA, Id: NewID(), Details: details, kill: make(chan URI, 1)}
	if details == nil {
		details = make(map[string]interface{})
	}
	go func() {
		r.handleSession(sess)
		sess.Close()
	}()
	log.WithField("session_id", sess.Id).Info("established internal session")
	return peerB, nil
}

// Close disconnects all clients after sending a goodbye message
func (r Realm) Close() {
	iter := r.clients.Iter()
	for client := range iter {
		sess, isSession := client.Val.(*Session)
		if !isSession {
			continue
		}
		sess.kill <- ErrSystemShutdown
	}
}

func (r *Realm) init() {
	r.lock.Lock()
	defer r.lock.Unlock()

	r.clients = cmap.New()

	if r.localClient == nil {
		p, _ := r.getPeer(nil)
		client := NewClient(p)
		r.localClient = new(localClient)
		r.localClient.Client = client
	}

	if r.Broker == nil {
		r.Broker = NewDefaultBroker()
	}
	if r.Dealer == nil {
		r.Dealer = NewDefaultDealer()
	}
	if r.Authorizer == nil {
		r.Authorizer = NewDefaultAuthorizer()
	}
	if r.Interceptor == nil {
		r.Interceptor = NewDefaultInterceptor()
	}
	if r.AuthTimeout == 0 {
		r.AuthTimeout = defaultAuthTimeout
	}
}

func (l *localClient) onJoin(details map[string]interface{}) {
	l.Publish("wamp.session.on_join", nil, []interface{}{details}, nil)
}

func (l *localClient) onLeave(session ID) {
	l.Publish("wamp.session.on_leave", nil, []interface{}{session}, nil)
}

func (r *Realm) doOne(c <-chan Message, sess *Session) bool {
	r.lock.RLock()
	defer r.lock.RUnlock()

	var msg Message
	var open bool
	select {
	case msg, open = <-c:
		if !open {
			log.WithField("session_id", sess.Id).Error("lost session")
			return false
		}
	case reason := <-sess.kill:
		logErr(sess.Send(&Goodbye{Reason: reason, Details: make(map[string]interface{})}))
		log.Printf("kill session %s: %v", sess, reason)
		// TODO: wait for client Goodbye?
		return false
	}

	redactedMsg := redactMessage(msg)

	log.WithFields(logrus.Fields{
		"session_id":   sess.Id,
		"message_type": msg.MessageType().String(),
		"message":      redactedMsg,
	}).Debug("new message")

	if isAuthz, err := r.Authorizer.Authorize(sess, msg); !isAuthz {
		errMsg := &Error{Type: msg.MessageType()}
		if err != nil {
			errMsg.Error = ErrAuthorizationFailed
			log.WithFields(logrus.Fields{
				"session_id":   sess.Id,
				"message_type": msg.MessageType().String(),
				"message":      redactedMsg,
				"err":          err,
			}).Error("authorization failed")
		} else {
			errMsg.Error = ErrNotAuthorized
			log.WithFields(logrus.Fields{
				"session_id":   sess.Id,
				"message_type": msg.MessageType().String(),
				"message":      redactedMsg,
				"err":          err,
			}).Error("UNAUTHORIZED")
		}
		logErr(sess.Send(errMsg))
		return true
	}

	r.Interceptor.Intercept(sess, &msg)

	switch msg := msg.(type) {
	case *Goodbye:
		logErr(sess.Send(&Goodbye{Reason: ErrGoodbyeAndOut, Details: make(map[string]interface{})}))
		log.WithFields(logrus.Fields{
			"session_id": sess.Id,
			"reason":     msg.Reason,
		}).Warning("leaving")
		return false

	// Broker messages
	case *Publish:
		r.Broker.Publish(sess, msg)
	case *Subscribe:
		r.Broker.Subscribe(sess, msg)
	case *Unsubscribe:
		r.Broker.Unsubscribe(sess, msg)

	// Dealer messages
	case *Register:
		r.Dealer.Register(sess, msg)
	case *Unregister:
		r.Dealer.Unregister(sess, msg)
	case *Call:
		r.Dealer.Call(sess, msg)
	case *Yield:
		r.Dealer.Yield(sess, msg)

	// Error messages
	case *Error:
		if msg.Type == INVOCATION {
			// the only type of ERROR message the router should receive
			r.Dealer.Error(sess, msg)
		} else {
			log.Infof("invalid ERROR message received: %v", msg)
		}

	default:
		log.Warningf("Unhandled message:", msg.MessageType())
	}
	return true
}

func redactMessage(msg Message) Message {
	switch msg := msg.(type) {
	case *Call:
		var redacted Call
		redacted.Request = msg.Request
		redacted.Arguments = append(redacted.Arguments, msg.Arguments...)
		redacted.ArgumentsKw = make(map[string]interface{})
		for k, v := range msg.ArgumentsKw {
			if k == "token" {
				v = "redacted"
			}
			redacted.ArgumentsKw[k] = v
		}
		redacted.Options = make(map[string]interface{})
		for k, v := range msg.Options {
			redacted.Options[k] = v
		}
		return &redacted
	}
	return msg
}

func (r *Realm) handleSession(sess *Session) {
	r.lock.RLock()
	r.clients.Set(string(sess.Id), sess)
	r.localClient.onJoin(sess.Details)
	r.lock.RUnlock()

	defer func() {
		r.lock.RLock()
		defer r.lock.RUnlock()

		r.clients.Remove(string(sess.Id))
		r.Broker.RemoveSession(sess)
		r.Dealer.RemoveSession(sess)
		r.localClient.onLeave(sess.Id)
	}()
	c := sess.Receive()
	// TODO: what happens if the realm is closed?

	for r.doOne(c, sess) {
	}
}

func (r *Realm) handleAuth(client Peer, details map[string]interface{}) (*Welcome, error) {
	msg, err := r.authenticate(details)
	if err != nil {
		return nil, err
	}
	// we should never get anything besides WELCOME and CHALLENGE
	if msg.MessageType() == WELCOME {
		return msg.(*Welcome), nil
	}
	// Challenge response
	challenge := msg.(*Challenge)
	if err := client.Send(challenge); err != nil {
		return nil, err
	}

	msg, err = GetMessageTimeout(client, r.AuthTimeout)
	if err != nil {
		return nil, err
	}
	log.Printf("%s: %+v", msg.MessageType(), msg)
	if authenticate, ok := msg.(*Authenticate); !ok {
		return nil, fmt.Errorf("unexpected %s message received", msg.MessageType())
	} else {
		return r.checkResponse(challenge, authenticate)
	}
}

// Authenticate either authenticates a client or returns a challenge message if
// challenge/response authentication is to be used.
func (r Realm) authenticate(details map[string]interface{}) (Message, error) {
	log.Println("details:", details)
	if len(r.Authenticators) == 0 && len(r.CRAuthenticators) == 0 {
		return &Welcome{}, nil
	}
	// TODO: this might not always be a []interface{}. Using the JSON unmarshaller it will be,
	// but we may have serializations that preserve more of the original type.
	// For now, the tests just explicitly send a []interface{}
	_authmethods, ok := details["authmethods"].([]interface{})
	if !ok {
		return nil, fmt.Errorf("No authentication supplied")
	}
	authmethods := []string{}
	for _, method := range _authmethods {
		if m, ok := method.(string); ok {
			authmethods = append(authmethods, m)
		} else {
			log.Printf("invalid authmethod value: %v", method)
		}
	}
	for _, method := range authmethods {
		if auth, ok := r.CRAuthenticators[method]; ok {
			if challenge, err := auth.Challenge(details); err != nil {
				return nil, err
			} else {
				return &Challenge{AuthMethod: method, Extra: challenge}, nil
			}
		}
		if auth, ok := r.Authenticators[method]; ok {
			if authDetails, err := auth.Authenticate(details); err != nil {
				return nil, err
			} else {
				return &Welcome{Details: addAuthMethod(authDetails, method)}, nil
			}
		}
	}
	// TODO: check default auth (special '*' auth?)
	return nil, fmt.Errorf("could not authenticate with any method")
}

// checkResponse determines whether the response to the challenge is sufficient to gain access to the Realm.
func (r Realm) checkResponse(chal *Challenge, auth *Authenticate) (*Welcome, error) {
	authenticator, ok := r.CRAuthenticators[chal.AuthMethod]
	if !ok {
		return nil, fmt.Errorf("authentication method has been removed")
	}
	if details, err := authenticator.Authenticate(chal.Extra, auth.Signature); err != nil {
		return nil, err
	} else {
		return &Welcome{Details: addAuthMethod(details, chal.AuthMethod)}, nil
	}
}

func addAuthMethod(details map[string]interface{}, method string) map[string]interface{} {
	if details == nil {
		details = make(map[string]interface{})
	}
	details["authmethod"] = method
	return details
}

// r := Realm{
// 	Authenticators: map[string]gowamp.Authenticator{
// 		"wampcra": gowamp.NewCRAAuthenticatorFactoryFactory(mySecret),
// 		"ticket": gowamp.NewTicketAuthenticator(myTicket),
// 		"asdfasdf": myAsdfAuthenticator,
// 	},
// 	BasicAuthenticators: map[string]turnpike.BasicAuthenticator{
// 		"anonymous": nil,
// 	},
// }
<commit_msg>Add back in erroneously removed concurrent-map import<commit_after>package turnpike

import (
	"fmt"
	"sync"
	"time"

	logrus "github.com/sirupsen/logrus"
	"github.com/streamrail/concurrent-map"
)

const (
	defaultAuthTimeout = 2 * time.Minute
)

// A Realm is a WAMP routing and administrative domain.
//
// Clients that have connected to a WAMP router are joined to a realm and all
// message delivery is handled by the realm.
type Realm struct {
	_                string
	URI              URI
	Broker           Broker
	Dealer           Dealer
	Authorizer       Authorizer
	Interceptor      Interceptor
	CRAuthenticators map[string]CRAuthenticator
	Authenticators   map[string]Authenticator
	// DefaultAuth      func(details map[string]interface{}) (map[string]interface{}, error)
	AuthTimeout time.Duration
	clients     cmap.ConcurrentMap
	localClient *localClient

	lock sync.RWMutex
}

type localClient struct {
	*Client
	sync.Mutex
}

func (r *Realm) getPeer(details map[string]interface{}) (Peer, error) {
	peerA, peerB := localPipe()
	sess := &Session{Peer: peerA, Id: NewID(), Details: details, kill: make(chan URI, 1)}
	if details == nil {
		details = make(map[string]interface{})
	}
	go func() {
		r.handleSession(sess)
		sess.Close()
	}()
	log.WithField("session_id", sess.Id).Info("established internal session")
	return peerB, nil
}

// Close disconnects all clients after sending a goodbye message
func (r Realm) Close() {
	iter := r.clients.Iter()
	for client := range iter {
		sess, isSession := client.Val.(*Session)
		if !isSession {
			continue
		}
		sess.kill <- ErrSystemShutdown
	}
}

func (r *Realm) init() {
	r.lock.Lock()
	defer r.lock.Unlock()

	r.clients = cmap.New()

	if r.localClient == nil {
		p, _ := r.getPeer(nil)
		client := NewClient(p)
		r.localClient = new(localClient)
		r.localClient.Client = client
	}

	if r.Broker == nil {
		r.Broker = NewDefaultBroker()
	}
	if r.Dealer == nil {
		r.Dealer = NewDefaultDealer()
	}
	if r.Authorizer == nil {
		r.Authorizer = NewDefaultAuthorizer()
	}
	if r.Interceptor == nil {
		r.Interceptor = NewDefaultInterceptor()
	}
	if r.AuthTimeout == 0 {
		r.AuthTimeout = defaultAuthTimeout
	}
}

func (l *localClient) onJoin(details map[string]interface{}) {
	l.Publish("wamp.session.on_join", nil, []interface{}{details}, nil)
}

func (l *localClient) onLeave(session ID) {
	l.Publish("wamp.session.on_leave", nil, []interface{}{session}, nil)
}

func (r *Realm) doOne(c <-chan Message, sess *Session) bool {
	r.lock.RLock()
	defer r.lock.RUnlock()

	var msg Message
	var open bool
	select {
	case msg, open = <-c:
		if !open {
			log.WithField("session_id", sess.Id).Error("lost session")
			return false
		}
	case reason := <-sess.kill:
		logErr(sess.Send(&Goodbye{Reason: reason, Details: make(map[string]interface{})}))
		log.Printf("kill session %s: %v", sess, reason)
		// TODO: wait for client Goodbye?
		return false
	}

	redactedMsg := redactMessage(msg)

	log.WithFields(logrus.Fields{
		"session_id":   sess.Id,
		"message_type": msg.MessageType().String(),
		"message":      redactedMsg,
	}).Debug("new message")

	if isAuthz, err := r.Authorizer.Authorize(sess, msg); !isAuthz {
		errMsg := &Error{Type: msg.MessageType()}
		if err != nil {
			errMsg.Error = ErrAuthorizationFailed
			log.WithFields(logrus.Fields{
				"session_id":   sess.Id,
				"message_type": msg.MessageType().String(),
				"message":      redactedMsg,
				"err":          err,
			}).Error("authorization failed")
		} else {
			errMsg.Error = ErrNotAuthorized
			log.WithFields(logrus.Fields{
				"session_id":   sess.Id,
				"message_type": msg.MessageType().String(),
				"message":      redactedMsg,
				"err":          err,
			}).Error("UNAUTHORIZED")
		}
		logErr(sess.Send(errMsg))
		return true
	}

	r.Interceptor.Intercept(sess, &msg)

	switch msg := msg.(type) {
	case *Goodbye:
		logErr(sess.Send(&Goodbye{Reason: ErrGoodbyeAndOut, Details: make(map[string]interface{})}))
		log.WithFields(logrus.Fields{
			"session_id": sess.Id,
			"reason":     msg.Reason,
		}).Warning("leaving")
		return false

	// Broker messages
	case *Publish:
		r.Broker.Publish(sess, msg)
	case *Subscribe:
		r.Broker.Subscribe(sess, msg)
	case *Unsubscribe:
		r.Broker.Unsubscribe(sess, msg)

	// Dealer messages
	case *Register:
		r.Dealer.Register(sess, msg)
	case *Unregister:
		r.Dealer.Unregister(sess, msg)
	case *Call:
		r.Dealer.Call(sess, msg)
	case *Yield:
		r.Dealer.Yield(sess, msg)

	// Error messages
	case *Error:
		if msg.Type == INVOCATION {
			// the only type of ERROR message the router should receive
			r.Dealer.Error(sess, msg)
		} else {
			log.Infof("invalid ERROR message received: %v", msg)
		}

	default:
		log.Warningf("Unhandled message:", msg.MessageType())
	}
	return true
}

func redactMessage(msg Message) Message {
	switch msg := msg.(type) {
	case *Call:
		var redacted Call
		redacted.Request = msg.Request
		redacted.Arguments = append(redacted.Arguments, msg.Arguments...)
		redacted.ArgumentsKw = make(map[string]interface{})
		for k, v := range msg.ArgumentsKw {
			if k == "token" {
				v = "redacted"
			}
			redacted.ArgumentsKw[k] = v
		}
		redacted.Options = make(map[string]interface{})
		for k, v := range msg.Options {
			redacted.Options[k] = v
		}
		return &redacted
	}
	return msg
}

func (r *Realm) handleSession(sess *Session) {
	r.lock.RLock()
	r.clients.Set(string(sess.Id), sess)
	r.localClient.onJoin(sess.Details)
	r.lock.RUnlock()

	defer func() {
		r.lock.RLock()
		defer r.lock.RUnlock()

		r.clients.Remove(string(sess.Id))
		r.Broker.RemoveSession(sess)
		r.Dealer.RemoveSession(sess)
		r.localClient.onLeave(sess.Id)
	}()
	c := sess.Receive()
	// TODO: what happens if the realm is closed?

	for r.doOne(c, sess) {
	}
}

func (r *Realm) handleAuth(client Peer, details map[string]interface{}) (*Welcome, error) {
	msg, err := r.authenticate(details)
	if err != nil {
		return nil, err
	}
	// we should never get anything besides WELCOME and CHALLENGE
	if msg.MessageType() == WELCOME {
		return msg.(*Welcome), nil
	}
	// Challenge response
	challenge := msg.(*Challenge)
	if err := client.Send(challenge); err != nil {
		return nil, err
	}

	msg, err = GetMessageTimeout(client, r.AuthTimeout)
	if err != nil {
		return nil, err
	}
	log.Printf("%s: %+v", msg.MessageType(), msg)
	if authenticate, ok := msg.(*Authenticate); !ok {
		return nil, fmt.Errorf("unexpected %s message received", msg.MessageType())
	} else {
		return r.checkResponse(challenge, authenticate)
	}
}

// Authenticate either authenticates a client or returns a challenge message if
// challenge/response authentication is to be used.
func (r Realm) authenticate(details map[string]interface{}) (Message, error) {
	log.Println("details:", details)
	if len(r.Authenticators) == 0 && len(r.CRAuthenticators) == 0 {
		return &Welcome{}, nil
	}
	// TODO: this might not always be a []interface{}. Using the JSON unmarshaller it will be,
	// but we may have serializations that preserve more of the original type.
	// For now, the tests just explicitly send a []interface{}
	_authmethods, ok := details["authmethods"].([]interface{})
	if !ok {
		return nil, fmt.Errorf("No authentication supplied")
	}
	authmethods := []string{}
	for _, method := range _authmethods {
		if m, ok := method.(string); ok {
			authmethods = append(authmethods, m)
		} else {
			log.Printf("invalid authmethod value: %v", method)
		}
	}
	for _, method := range authmethods {
		if auth, ok := r.CRAuthenticators[method]; ok {
			if challenge, err := auth.Challenge(details); err != nil {
				return nil, err
			} else {
				return &Challenge{AuthMethod: method, Extra: challenge}, nil
			}
		}
		if auth, ok := r.Authenticators[method]; ok {
			if authDetails, err := auth.Authenticate(details); err != nil {
				return nil, err
			} else {
				return &Welcome{Details: addAuthMethod(authDetails, method)}, nil
			}
		}
	}
	// TODO: check default auth (special '*' auth?)
	return nil, fmt.Errorf("could not authenticate with any method")
}

// checkResponse determines whether the response to the challenge is sufficient to gain access to the Realm.
func (r Realm) checkResponse(chal *Challenge, auth *Authenticate) (*Welcome, error) {
	authenticator, ok := r.CRAuthenticators[chal.AuthMethod]
	if !ok {
		return nil, fmt.Errorf("authentication method has been removed")
	}
	if details, err := authenticator.Authenticate(chal.Extra, auth.Signature); err != nil {
		return nil, err
	} else {
		return &Welcome{Details: addAuthMethod(details, chal.AuthMethod)}, nil
	}
}

func addAuthMethod(details map[string]interface{}, method string) map[string]interface{} {
	if details == nil {
		details = make(map[string]interface{})
	}
	details["authmethod"] = method
	return details
}

// r := Realm{
// 	Authenticators: map[string]gowamp.Authenticator{
// 		"wampcra": gowamp.NewCRAAuthenticatorFactoryFactory(mySecret),
// 		"ticket": gowamp.NewTicketAuthenticator(myTicket),
// 		"asdfasdf": myAsdfAuthenticator,
// 	},
// 	BasicAuthenticators: map[string]turnpike.BasicAuthenticator{
// 		"anonymous": nil,
// 	},
// }
<|endoftext|>


<commit_before>package imap_test

import (
	"testing"

	"github.com/emersion/go-imap"
)

func TestRespHandle_Accept(t *testing.T) {
	ch := make(chan bool, 1)
	hdlr := &imap.RespHandle{
		Accepts: ch,
	}

	hdlr.Accept()

	v := <-ch
	if v != true {
		t.Error("Invalid return value:", v)
	}
}

func TestRespHandle_Reject(t *testing.T) {
	ch := make(chan bool, 1)
	hdlr := &imap.RespHandle{
		Accepts: ch,
	}

	hdlr.Reject()

	v := <-ch
	if v != false {
		t.Error("Invalid return value:", v)
	}
}

func TestRespHandle_AcceptNamedResp_Matching(t *testing.T) {
	ch := make(chan bool, 1)
	hdlr := &imap.RespHandle{
		Resp: &imap.Resp{
			Tag:    "*",
			Fields: []interface{}{"SEARCH", "42"},
		},
		Accepts: ch,
	}

	fields, ok := hdlr.AcceptNamedResp("SEARCH")
	if ok != true {
		t.Error("Matching response not accepted")
	}
	if len(fields) != 1 {
		t.Error("Invalid fields")
	}
	if f, ok := fields[0].(string); !ok || f != "42" {
		t.Error("Invalid first field")
	}

	v := <-ch
	if v != true {
		t.Error("Invalid return value:", v)
	}
}

func TestRespHandle_AcceptNamedResp_NotMatching(t *testing.T) {
	ch := make(chan bool, 1)
	hdlr := &imap.RespHandle{
		Resp: &imap.Resp{
			Tag:    "*",
			Fields: []interface{}{"26", "EXISTS"},
		},
		Accepts: ch,
	}

	_, ok := hdlr.AcceptNamedResp("SEARCH")
	if ok != false {
		t.Error("Response not matching has been accepted")
	}

	v := <-ch
	if v != false {
		t.Error("Invalid return value:", v)
	}
}

func MultiRespHandler(t *testing.T) {
	mh := imap.NewMultiRespHandler()

	h1 := make(imap.RespHandler)
	mh.Add(h1)
	go func() {
		(<-h1).Accept()
		(<-h1).Reject()
		mh.Del(h1)
	}()

	h2 := make(imap.RespHandler)
	mh.Add(h2)
	go func() {
		(<-h2).Reject()
		(<-h2).Reject()
		mh.Del(h2)
	}()

	// Should not add it, or will block forever
	var h3 imap.RespHandler
	mh.Add(h3)

	rh1 := &imap.RespHandle{Accepts: make(chan bool, 1)}
	rh2 := &imap.RespHandle{Accepts: make(chan bool, 1)}

	h := make(imap.RespHandler, 2)
	h <- rh1
	h <- rh2
	close(h)

	if err := mh.HandleFrom(h); err != nil {
		t.Fatal("Expected no error while handling response, got:", err)
	}
	if accepted := <-rh1.Accepts; !accepted {
		t.Error("First response was not accepted")
	}
	if accepted := <-rh2.Accepts; accepted {
		t.Error("First response was not rejected")
	}
}
<commit_msg>imap: fixes test function naming<commit_after>package imap_test

import (
	"testing"

	"github.com/emersion/go-imap"
)

func TestRespHandle_Accept(t *testing.T) {
	ch := make(chan bool, 1)
	hdlr := &imap.RespHandle{
		Accepts: ch,
	}

	hdlr.Accept()

	v := <-ch
	if v != true {
		t.Error("Invalid return value:", v)
	}
}

func TestRespHandle_Reject(t *testing.T) {
	ch := make(chan bool, 1)
	hdlr := &imap.RespHandle{
		Accepts: ch,
	}

	hdlr.Reject()

	v := <-ch
	if v != false {
		t.Error("Invalid return value:", v)
	}
}

func TestRespHandle_AcceptNamedResp_Matching(t *testing.T) {
	ch := make(chan bool, 1)
	hdlr := &imap.RespHandle{
		Resp: &imap.Resp{
			Tag:    "*",
			Fields: []interface{}{"SEARCH", "42"},
		},
		Accepts: ch,
	}

	fields, ok := hdlr.AcceptNamedResp("SEARCH")
	if ok != true {
		t.Error("Matching response not accepted")
	}
	if len(fields) != 1 {
		t.Error("Invalid fields")
	}
	if f, ok := fields[0].(string); !ok || f != "42" {
		t.Error("Invalid first field")
	}

	v := <-ch
	if v != true {
		t.Error("Invalid return value:", v)
	}
}

func TestRespHandle_AcceptNamedResp_NotMatching(t *testing.T) {
	ch := make(chan bool, 1)
	hdlr := &imap.RespHandle{
		Resp: &imap.Resp{
			Tag:    "*",
			Fields: []interface{}{"26", "EXISTS"},
		},
		Accepts: ch,
	}

	_, ok := hdlr.AcceptNamedResp("SEARCH")
	if ok != false {
		t.Error("Response not matching has been accepted")
	}

	v := <-ch
	if v != false {
		t.Error("Invalid return value:", v)
	}
}

func TestMultiRespHandler(t *testing.T) {
	mh := imap.NewMultiRespHandler()

	h1 := make(imap.RespHandler)
	mh.Add(h1)
	go func() {
		(<-h1).Accept()
		(<-h1).Reject()
		mh.Del(h1)
	}()

	h2 := make(imap.RespHandler)
	mh.Add(h2)
	go func() {
		(<-h2).Reject()
		(<-h2).Reject()
		mh.Del(h2)
	}()

	// Should not add it, or will block forever
	var h3 imap.RespHandler
	mh.Add(h3)

	rh1 := &imap.RespHandle{Accepts: make(chan bool, 1)}
	rh2 := &imap.RespHandle{Accepts: make(chan bool, 1)}

	h := make(imap.RespHandler, 2)
	h <- rh1
	h <- rh2
	close(h)

	if err := mh.HandleFrom(h); err != nil {
		t.Fatal("Expected no error while handling response, got:", err)
	}
	if accepted := <-rh1.Accepts; !accepted {
		t.Error("First response was not accepted")
	}
	if accepted := <-rh2.Accepts; accepted {
		t.Error("First response was not rejected")
	}
}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"fmt"
	"net/url"
	"regexp"
	"strings"
)

type AliasCmd struct{}

func init() {
	h := NewHandler()
	h.CommandName = "alias"
	h.CommandPattern = "(alias)( )(.*)"
	h.Usage = "alias (create host port indexName aliasName) | (remove host port indexName aliasName) | (move host port fromIndex toIndex aliasName)"
	h.CommandParser = func(cmd *Command) (string, bool) {
		pattFn := map[*regexp.Regexp]func([]string) (string, bool){
			// List all aliases on host
			regexp.MustCompile(`^alias$`): func(s []string) (string, bool) {
				d := Resource{
					Endpoint: "_aliases",
					Scheme:   "http",
					Host:     server.host,
					Port:     server.port,
				}
				c := AliasCmd{}
				r, ok := c.GetAll(d)
				return r, ok
			},
			// Alias help
			regexp.MustCompile(`^alias /\?$`): func(s []string) (string, bool) {
				return "", false
			},
			// Create alias
			regexp.MustCompile(`^alias create ([a-zA-Z0-9\.\-]+) ([0-9]{1,5}) ([a-zA-Z0-9\.\-]+) ([a-zA-Z0-9\.\-]+)$`): func(s []string) (string, bool) {
				d := Resource{
					Endpoint: "_aliases",
					Index:    s[3],
					Alias:    s[4],
				}
				c := AliasCmd{}
				r, ok := c.Create(d)
				return r, ok
			},
			// Remove alias
			regexp.MustCompile(`^alias remove ([a-zA-Z0-9\.\-]+) ([0-9]{1,5}) ([a-zA-Z0-9\.\-]+) ([a-zA-Z0-9\.\-]+)$`): func(s []string) (string, bool) {
				d := Resource{
					Endpoint: "_aliases",
					Index:    s[3],
					Alias:    s[4],
				}
				c := AliasCmd{}
				r, ok := c.Remove(d)
				return r, ok
			},
			// Move alias
			regexp.MustCompile(`^alias move ([a-zA-Z0-9\.\-]+) ([0-9]{1,5}) ([a-zA-Z0-9\.\-]+) ([a-zA-Z0-9\.\-]+) ([a-zA-Z0-9\.\-]+)$`): func(s []string) (string, bool) {
				dFrom := Resource{
					Endpoint: "_aliases",
					Index:    s[3],
					Alias:    s[5],
				}
				dTarget := Resource{
					Index: s[4],
					Alias: s[5],
				}
				c := AliasCmd{}
				r, ok := c.Move(dFrom, dTarget)
				return r, ok
			},
		}
		r, ok := h.Tokenize(strings.TrimSpace(cmd.Instruction), pattFn)
		return r, ok
	}
	h.HandlerFunc = func(cmd *Command) string {
		r, ok := h.CommandParser(cmd)
		if !ok {
			if r != "" {
				r += "\n\n"
			}
			return r + usageMessage(h.Usage)
		}
		return r
	}
	HandlerRegistry[h.CommandName] = h
}

func (c *AliasCmd) GetAll(d Resource) (string, bool) {
	if server.host == "" || server.port == "" {
		return "Missing host or port environment config.", false
	}
	u := new(url.URL)
	u.Scheme = d.Scheme
	u.Host = d.Host + ":" + d.Port
	u.Path = "/" + d.Endpoint
	q := u.Query()
	q.Add("pretty", "true")
	u.RawQuery = q.Encode()
	fmt.Println("Request:", u)
	err := errors.New("")
	res, err := getHttpResource(u.String())
	if err != nil {
		return err.Error(), false
	}
	return res, true
}

func (c *AliasCmd) Create(d Resource) (string, bool) {
	if server.host == "" || server.port == "" {
		return "Missing host or port environment config.", false
	}
	//curl -XPOST "http://10.1.1.12:9200/_aliases" -d '{ "actions": [ { "add": { "index": "podcasts-2014-07-29-001", "alias": "podcastsupdater" } } ] }'
	//post _alias?pretty { "actions": [ { "add": { "index": "podcasts-2014-07-29-001", "alias": "podcastsupdater" } } ] }
	urlString := "post " + d.Endpoint + " " + "{\"actions\": [ { \"add\": { \"index\": \"" + d.Index + "\", \"alias\": \"" + d.Alias + "\" } } ] }"
	cmdParser := NewCommandParser()
	newCmd, err := cmdParser.Parse(urlString)
	if err != nil {
		return err.Error(), false
	}
	dispatcher := NewDispatcher()
	res := dispatcher.Dispatch(newCmd)
	return res, true
}

func (c *AliasCmd) Remove(d Resource) (string, bool) {
	if server.host == "" || server.port == "" {
		return "Missing host or port environment config.", false
	}
	//curl -XPOST "http://10.1.1.12:9200/_aliases" -d '{ "actions": [ { "remove": { "index": "podcasts-2014-05-07-0103", "alias": "podcastsupdater" } } ] }'
	//post _alias?pretty { "actions": [ { "remove": { "index": "podcasts-2014-05-07-0103", "alias": "podcastsupdater" } } ] }
	urlString := "post " + d.Endpoint + " " + "{\"actions\": [ { \"remove\": { \"index\": \"" + d.Index + "\", \"alias\": \"" + d.Alias + "\" } } ] }"
	cmdParser := NewCommandParser()
	newCmd, err := cmdParser.Parse(urlString)
	if err != nil {
		return err.Error(), false
	}
	dispatcher := NewDispatcher()
	res := dispatcher.Dispatch(newCmd)
	return res, true
}

func (c *AliasCmd) Move(dFrom Resource, dTarget Resource) (string, bool) {
	if server.host == "" || server.port == "" {
		return "Missing host or port environment config.", false
	}
	postData := "{ \"actions\": [ { \"remove\": { \"alias\": \"" + dFrom.Alias + "\", \"index\": \"" + dFrom.Index + "\" }}, { \"add\": { \"alias\": \"" + dTarget.Alias + "\", \"index\": \"" + dTarget.Index + "\" } } ] }"
	urlString := "post " + dFrom.Endpoint + " " + postData
	cmdParser := NewCommandParser()
	newCmd, err := cmdParser.Parse(urlString)
	if err != nil {
		return err.Error(), false
	}
	dispatcher := NewDispatcher()
	res := dispatcher.Dispatch(newCmd)
	return res, true
}
<commit_msg>Cleaned.<commit_after>package main

import (
	"errors"
	"fmt"
	"net/url"
	"regexp"
	"strings"
)

type AliasCmd struct{}

func init() {
	h := NewHandler()
	h.CommandName = "alias"
	h.CommandPattern = "(alias)( )(.*)"
	h.Usage = "alias\n" +
		"alias host port\n" +
		"alias create host port indexName aliasName\n" +
		"alias remove host port indexName aliasName\n" +
		"alias move host port fromIndex toIndex aliasName"
	h.CommandParser = func(cmd *Command) (string, bool) {
		pattFn := map[*regexp.Regexp]func([]string) (string, bool){
			// List all aliases on host
			regexp.MustCompile(`^alias$`): func(s []string) (string, bool) {
				d := Resource{
					Endpoint: "_aliases",
					Scheme:   "http",
					Host:     server.host,
					Port:     server.port,
				}
				c := AliasCmd{}
				r, ok := c.GetAll(d)
				return r, ok
			},
			regexp.MustCompile(`^alias ([a-zA-Z0-9\.\-]+) ([0-9]{1,5})$`): func(s []string) (string, bool) {
				d := Resource{
					Endpoint: "_aliases",
					Scheme:   "http",
					Host:     s[1],
					Port:     s[2],
				}
				c := AliasCmd{}
				r, ok := c.GetAll(d)
				return r, ok
			},
			// Alias help
			regexp.MustCompile(`^alias /\?$`): func(s []string) (string, bool) {
				return "", false
			},
			// Create alias
			regexp.MustCompile(`^alias create ([a-zA-Z0-9\.\-]+) ([0-9]{1,5}) ([a-zA-Z0-9\.\-]+) ([a-zA-Z0-9\.\-]+)$`): func(s []string) (string, bool) {
				d := Resource{
					Endpoint: "_aliases",
					Host:     s[1],
					Port:     s[2],
					Index:    s[3],
					Alias:    s[4],
				}
				c := AliasCmd{}
				r, ok := c.Create(d)
				return r, ok
			},
			// Remove alias
			regexp.MustCompile(`^alias remove ([a-zA-Z0-9\.\-]+) ([0-9]{1,5}) ([a-zA-Z0-9\.\-]+) ([a-zA-Z0-9\.\-]+)$`): func(s []string) (string, bool) {
				d := Resource{
					Endpoint: "_aliases",
					Host:     s[1],
					Port:     s[2],
					Index:    s[3],
					Alias:    s[4],
				}
				c := AliasCmd{}
				r, ok := c.Remove(d)
				return r, ok
			},
			// Move alias
			regexp.MustCompile(`^alias move ([a-zA-Z0-9\.\-]+) ([0-9]{1,5}) ([a-zA-Z0-9\.\-]+) ([a-zA-Z0-9\.\-]+) ([a-zA-Z0-9\.\-]+)$`): func(s []string) (string, bool) {
				dFrom := Resource{
					Endpoint: "_aliases",
					Host:     s[1],
					Port:     s[2],
					Index:    s[3],
					Alias:    s[5],
				}
				dTarget := Resource{
					Index: s[4],
					Alias: s[5],
				}
				c := AliasCmd{}
				r, ok := c.Move(dFrom, dTarget)
				return r, ok
			},
		}
		r, ok := h.Tokenize(strings.TrimSpace(cmd.Instruction), pattFn)
		return r, ok
	}
	h.HandlerFunc = func(cmd *Command) string {
		r, ok := h.CommandParser(cmd)
		if !ok {
			if r != "" {
				r += "\n\n"
			}
			return r + usageMessage(h.Usage)
		}
		return r
	}
	HandlerRegistry[h.CommandName] = h
}

func (c *AliasCmd) GetAll(d Resource) (string, bool) {
	if d.Host == "" || d.Port == "" {
		return "Missing host or port.", false
	}
	u := new(url.URL)
	u.Scheme = d.Scheme
	u.Host = d.Host + ":" + d.Port
	u.Path = "/" + d.Endpoint
	q := u.Query()
	q.Add("pretty", "true")
	u.RawQuery = q.Encode()
	fmt.Println("Request:", u)
	err := errors.New("")
	res, err := getHttpResource(u.String())
	if err != nil {
		return err.Error(), false
	}
	return res, true
}

func (c *AliasCmd) Create(d Resource) (string, bool) {
	if d.Host == "" || d.Port == "" {
		return "Missing host or port.", false
	}
	postData := "{\"actions\": [ { \"add\": { \"index\": \"" + d.Index + "\", \"alias\": \"" + d.Alias + "\" } } ] }"
	urlString := "post " + d.Host + " " + d.Port + " " + d.Endpoint + " " + postData
	cmdParser := NewCommandParser()
	newCmd, err := cmdParser.Parse(urlString)
	if err != nil {
		return err.Error(), false
	}
	dispatcher := NewDispatcher()
	res := dispatcher.Dispatch(newCmd)
	return res, true
}

func (c *AliasCmd) Remove(d Resource) (string, bool) {
	if d.Host == "" || d.Port == "" {
		return "Missing host or port.", false
	}
	postData := "{\"actions\": [ { \"remove\": { \"index\": \"" + d.Index + "\", \"alias\": \"" + d.Alias + "\" } } ] }"
	urlString := "post " + d.Host + " " + d.Port + " " + d.Endpoint + " " + postData
	cmdParser := NewCommandParser()
	newCmd, err := cmdParser.Parse(urlString)
	if err != nil {
		return err.Error(), false
	}
	dispatcher := NewDispatcher()
	res := dispatcher.Dispatch(newCmd)
	return res, true
}

func (c *AliasCmd) Move(dFrom Resource, dTarget Resource) (string, bool) {
	if dFrom.Host == "" || dFrom.Port == "" {
		return "Missing host or port.", false
	}
	postData := "{ \"actions\": [ { \"remove\": { \"alias\": \"" + dFrom.Alias + "\", \"index\": \"" + dFrom.Index + "\" }}, { \"add\": { \"alias\": \"" + dTarget.Alias + "\", \"index\": \"" + dTarget.Index + "\" } } ] }"
	urlString := "post " + dFrom.Host + " " + dFrom.Port + " " + dFrom.Endpoint + " " + postData
	cmdParser := NewCommandParser()
	newCmd, err := cmdParser.Parse(urlString)
	if err != nil {
		return err.Error(), false
	}
	dispatcher := NewDispatcher()
	res := dispatcher.Dispatch(newCmd)
	return res, true
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"os"
	"strings"
)

type ARIN_OrgNets struct {
	Nets struct {
		InaccuracyReportURL string `json:"@inaccuracyReportUrl"`
		TermsOfUse          string `json:"@termsOfUse"`
		LimitExceeded       struct {
			Limit string `json:"@limit"`
			Value string `json:"$"`
		} `json:"limitExceeded"`
		NetRef []struct {
			EndAddress   string `json:"@endAddress"`
			StartAddress string `json:"@startAddress"`
			Handle       string `json:"@handle"`
			Name         string `json:"@name"`
			Value        string `json:"$"`
		} `json:"netRef"`
	} `json:"nets"`
}

type ARIN_OrgNet struct {
	Nets struct {
		InaccuracyReportURL string `json:"@inaccuracyReportUrl"`
		TermsOfUse          string `json:"@termsOfUse"`
		LimitExceeded       struct {
			Limit string `json:"@limit"`
			Value string `json:"$"`
		} `json:"limitExceeded"`
		NetRef struct {
			EndAddress   string `json:"@endAddress"`
			StartAddress string `json:"@startAddress"`
			Handle       string `json:"@handle"`
			Name         string `json:"@name"`
			Value        string `json:"$"`
		} `json:"netRef"`
	} `json:"nets"`
}

type ARIN_Nets struct {
	Net struct {
		InaccuracyReportURL string `json:"@inaccuracyReportUrl"`
		TermsOfUse          string `json:"@termsOfUse"`
		RegistrationDate    struct {
			Value string `json:"$"`
		} `json:"registrationDate"`
		Ref struct {
			Value string `json:"$"`
		} `json:"ref"`
		EndAddress struct {
			Value string `json:"$"`
		} `json:"endAddress"`
		Handle struct {
			Value string `json:"$"`
		} `json:"handle"`
		Name struct {
			Value string `json:"$"`
		} `json:"name"`
		NetBlocks struct {
			NetBlock []struct {
				CidrLength struct {
					Value string `json:"$"`
				} `json:"cidrLength"`
				EndAddress struct {
					Value string `json:"$"`
				} `json:"endAddress"`
				Description struct {
					Value string `json:"$"`
				} `json:"description"`
				Type struct {
					Value string `json:"$"`
				} `json:"type"`
				StartAddress struct {
					Value string `json:"$"`
				} `json:"startAddress"`
			} `json:"netBlock"`
		} `json:"netBlocks"`
		Resources struct {
			InaccuracyReportURL string `json:"@inaccuracyReportUrl"`
			TermsOfUse          string `json:"@termsOfUse"`
			LimitExceeded       struct {
				Limit string `json:"@limit"`
				Value string `json:"$"`
			} `json:"limitExceeded"`
		} `json:"resources"`
		OrgRef struct {
			Handle string `json:"@handle"`
			Name   string `json:"@name"`
			Value  string `json:"$"`
		} `json:"orgRef"`
		ParentNetRef struct {
			Handle string `json:"@handle"`
			Name   string `json:"@name"`
			Value  string `json:"$"`
		} `json:"parentNetRef"`
		StartAddress struct {
			Value string `json:"$"`
		} `json:"startAddress"`
		UpdateDate struct {
			Value string `json:"$"`
		} `json:"updateDate"`
		Version struct {
			Value string `json:"$"`
		} `json:"version"`
	} `json:"net"`
}

type ARIN_Net struct {
	Net struct {
		InaccuracyReportURL string `json:"@inaccuracyReportUrl"`
		TermsOfUse          string `json:"@termsOfUse"`
		RegistrationDate    struct {
			Value string `json:"$"`
		} `json:"registrationDate"`
		Ref struct {
			Value string `json:"$"`
		} `json:"ref"`
		EndAddress struct {
			Value string `json:"$"`
		} `json:"endAddress"`
		Handle struct {
			Value string `json:"$"`
		} `json:"handle"`
		Name struct {
			Value string `json:"$"`
		} `json:"name"`
		NetBlocks struct {
			NetBlock struct {
				CidrLength struct {
					Value string `json:"$"`
				} `json:"cidrLength"`
				EndAddress struct {
					Value string `json:"$"`
				} `json:"endAddress"`
				Description struct {
					Value string `json:"$"`
				} `json:"description"`
				Type struct {
					Value string `json:"$"`
				} `json:"type"`
				StartAddress struct {
					Value string `json:"$"`
				} `json:"startAddress"`
			} `json:"netBlock"`
		} `json:"netBlocks"`
		Resources struct {
			InaccuracyReportURL string `json:"@inaccuracyReportUrl"`
			TermsOfUse          string `json:"@termsOfUse"`
			LimitExceeded       struct {
				Limit string `json:"@limit"`
				Value string `json:"$"`
			} `json:"limitExceeded"`
		} `json:"resources"`
		OrgRef struct {
			Handle string `json:"@handle"`
			Name   string `json:"@name"`
			Value  string `json:"$"`
		} `json:"orgRef"`
		ParentNetRef struct {
			Handle string `json:"@handle"`
			Name   string `json:"@name"`
			Value  string `json:"$"`
		} `json:"parentNetRef"`
		StartAddress struct {
			Value string `json:"$"`
		} `json:"startAddress"`
		UpdateDate struct {
			Value string `json:"$"`
		} `json:"updateDate"`
		Version struct {
			Value string `json:"$"`
		} `json:"version"`
	} `json:"net"`
}

func LookupOrgNets(org string) ([]string, error) {
	handles := []string{}

	safe_org := url.QueryEscape(org)
	u := fmt.Sprintf("http://whois.arin.net/rest/org/%s/nets", safe_org)

	req, err := http.NewRequest("GET", u, nil)
	if err != nil {
		return handles, err
	}
	req.Header.Set("Accept", "application/json")

	client := &http.Client{}
	resp, err := client.Do(req)
	if err != nil {
		return handles, err
	}

	defer resp.Body.Close()

	content, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return handles, err
	}

	// No network handles associated with this organization
	if strings.Contains(string(content), "No related resources were found for the handle provided") {
		return handles, nil
	}

	var nets ARIN_OrgNets

	if err := json.Unmarshal(content, &nets); err == nil {
		for i := range nets.Nets.NetRef {
			handles = append(handles, nets.Nets.NetRef[i].Handle)
		}
	} else {
		// Try to decode as a single-net organization
		var net ARIN_OrgNet
		if err := json.Unmarshal(content, &net); err != nil {
			return handles, err
		}

		handles = append(handles, net.Nets.NetRef.Handle)
	}

	return handles, nil
}

func LookupNetCidrs(handle string) ([]string, error) {
	cidrs := []string{}

	safe_handle := url.QueryEscape(handle)
	u := fmt.Sprintf("http://whois.arin.net/rest/net/%s", safe_handle)

	req, err := http.NewRequest("GET", u, nil)
	if err != nil {
		return cidrs, err
	}
	req.Header.Set("Accept", "application/json")

	client := &http.Client{}
	resp, err := client.Do(req)
	if err != nil {
		return cidrs, err
	}

	defer resp.Body.Close()

	content, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return cidrs, err
	}

	var nets ARIN_Nets

	if err := json.Unmarshal(content, &nets); err == nil {
		for i := range nets.Net.NetBlocks.NetBlock {
			cidrs = append(cidrs, fmt.Sprintf("%s/%s", nets.Net.NetBlocks.NetBlock[i].StartAddress.Value, nets.Net.NetBlocks.NetBlock[i].CidrLength.Value))
		}
	} else {
		// Try to decode as a single-block network
		var net ARIN_Net
		if err := json.Unmarshal(content, &net); err != nil {
			return cidrs, err
		}
		cidrs = append(cidrs, fmt.Sprintf("%s/%s", net.Net.NetBlocks.NetBlock.StartAddress.Value, net.Net.NetBlocks.NetBlock.CidrLength.Value))
	}

	return cidrs, nil
}

func main() {

	if len(os.Args) != 2 {
		fmt.Println("Usage: inetdata-arin-org2nets <org-handle>\n")
		os.Exit(1)
	}

	org := os.Args[1]

	handles, e := LookupOrgNets(org)
	if e != nil {
		fmt.Fprintf(os.Stderr, "Could not list network handles: %s", e.Error())
		os.Exit(1)
	}

	for i := range handles {
		cidrs, e := LookupNetCidrs(handles[i])
		if e != nil {
			fmt.Fprintf(os.Stderr, "Could not list CIDRs for %s: %s", handles[i], e.Error())
			continue
		}
		fmt.Println(strings.Join(cidrs, "\n"))
	}
}
<commit_msg>Handle Customer records as well<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"os"
	"regexp"
	"strings"
)

var IsCustomerHandle = regexp.MustCompile(`^C[A-F0-9]{8}$`)

type ARIN_OrgNets struct {
	Nets struct {
		InaccuracyReportURL string `json:"@inaccuracyReportUrl"`
		TermsOfUse          string `json:"@termsOfUse"`
		LimitExceeded       struct {
			Limit string `json:"@limit"`
			Value string `json:"$"`
		} `json:"limitExceeded"`
		NetRef []struct {
			EndAddress   string `json:"@endAddress"`
			StartAddress string `json:"@startAddress"`
			Handle       string `json:"@handle"`
			Name         string `json:"@name"`
			Value        string `json:"$"`
		} `json:"netRef"`
	} `json:"nets"`
}

type ARIN_OrgNet struct {
	Nets struct {
		InaccuracyReportURL string `json:"@inaccuracyReportUrl"`
		TermsOfUse          string `json:"@termsOfUse"`
		LimitExceeded       struct {
			Limit string `json:"@limit"`
			Value string `json:"$"`
		} `json:"limitExceeded"`
		NetRef struct {
			EndAddress   string `json:"@endAddress"`
			StartAddress string `json:"@startAddress"`
			Handle       string `json:"@handle"`
			Name         string `json:"@name"`
			Value        string `json:"$"`
		} `json:"netRef"`
	} `json:"nets"`
}

type ARIN_Nets struct {
	Net struct {
		InaccuracyReportURL string `json:"@inaccuracyReportUrl"`
		TermsOfUse          string `json:"@termsOfUse"`
		RegistrationDate    struct {
			Value string `json:"$"`
		} `json:"registrationDate"`
		Ref struct {
			Value string `json:"$"`
		} `json:"ref"`
		EndAddress struct {
			Value string `json:"$"`
		} `json:"endAddress"`
		Handle struct {
			Value string `json:"$"`
		} `json:"handle"`
		Name struct {
			Value string `json:"$"`
		} `json:"name"`
		NetBlocks struct {
			NetBlock []struct {
				CidrLength struct {
					Value string `json:"$"`
				} `json:"cidrLength"`
				EndAddress struct {
					Value string `json:"$"`
				} `json:"endAddress"`
				Description struct {
					Value string `json:"$"`
				} `json:"description"`
				Type struct {
					Value string `json:"$"`
				} `json:"type"`
				StartAddress struct {
					Value string `json:"$"`
				} `json:"startAddress"`
			} `json:"netBlock"`
		} `json:"netBlocks"`
		Resources struct {
			InaccuracyReportURL string `json:"@inaccuracyReportUrl"`
			TermsOfUse          string `json:"@termsOfUse"`
			LimitExceeded       struct {
				Limit string `json:"@limit"`
				Value string `json:"$"`
			} `json:"limitExceeded"`
		} `json:"resources"`
		OrgRef struct {
			Handle string `json:"@handle"`
			Name   string `json:"@name"`
			Value  string `json:"$"`
		} `json:"orgRef"`
		ParentNetRef struct {
			Handle string `json:"@handle"`
			Name   string `json:"@name"`
			Value  string `json:"$"`
		} `json:"parentNetRef"`
		StartAddress struct {
			Value string `json:"$"`
		} `json:"startAddress"`
		UpdateDate struct {
			Value string `json:"$"`
		} `json:"updateDate"`
		Version struct {
			Value string `json:"$"`
		} `json:"version"`
	} `json:"net"`
}

type ARIN_Net struct {
	Net struct {
		InaccuracyReportURL string `json:"@inaccuracyReportUrl"`
		TermsOfUse          string `json:"@termsOfUse"`
		RegistrationDate    struct {
			Value string `json:"$"`
		} `json:"registrationDate"`
		Ref struct {
			Value string `json:"$"`
		} `json:"ref"`
		EndAddress struct {
			Value string `json:"$"`
		} `json:"endAddress"`
		Handle struct {
			Value string `json:"$"`
		} `json:"handle"`
		Name struct {
			Value string `json:"$"`
		} `json:"name"`
		NetBlocks struct {
			NetBlock struct {
				CidrLength struct {
					Value string `json:"$"`
				} `json:"cidrLength"`
				EndAddress struct {
					Value string `json:"$"`
				} `json:"endAddress"`
				Description struct {
					Value string `json:"$"`
				} `json:"description"`
				Type struct {
					Value string `json:"$"`
				} `json:"type"`
				StartAddress struct {
					Value string `json:"$"`
				} `json:"startAddress"`
			} `json:"netBlock"`
		} `json:"netBlocks"`
		Resources struct {
			InaccuracyReportURL string `json:"@inaccuracyReportUrl"`
			TermsOfUse          string `json:"@termsOfUse"`
			LimitExceeded       struct {
				Limit string `json:"@limit"`
				Value string `json:"$"`
			} `json:"limitExceeded"`
		} `json:"resources"`
		OrgRef struct {
			Handle string `json:"@handle"`
			Name   string `json:"@name"`
			Value  string `json:"$"`
		} `json:"orgRef"`
		ParentNetRef struct {
			Handle string `json:"@handle"`
			Name   string `json:"@name"`
			Value  string `json:"$"`
		} `json:"parentNetRef"`
		StartAddress struct {
			Value string `json:"$"`
		} `json:"startAddress"`
		UpdateDate struct {
			Value string `json:"$"`
		} `json:"updateDate"`
		Version struct {
			Value string `json:"$"`
		} `json:"version"`
	} `json:"net"`
}

func LookupOrgNets(org string) ([]string, error) {
	handles := []string{}

	safe_org := url.QueryEscape(org)

	u := ""

	// Organizations are split into Customers and Non-Customers, which
	// determines which API endpoint to use. Fortunately we can tell
	// which one is what based on the naming convention.
	if IsCustomerHandle.Match([]byte(org)) {
		u = fmt.Sprintf("http://whois.arin.net/rest/customer/%s/nets", safe_org)
	} else {
		u = fmt.Sprintf("http://whois.arin.net/rest/org/%s/nets", safe_org)
	}

	req, err := http.NewRequest("GET", u, nil)
	if err != nil {
		return handles, err
	}
	req.Header.Set("Accept", "application/json")

	client := &http.Client{}
	resp, err := client.Do(req)
	if err != nil {
		return handles, err
	}

	defer resp.Body.Close()

	content, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return handles, err
	}

	// No network handles associated with this organization
	if strings.Contains(string(content), "No related resources were found for the handle provided") {
		return handles, nil
	}

	var nets ARIN_OrgNets

	if err := json.Unmarshal(content, &nets); err == nil {
		for i := range nets.Nets.NetRef {
			handles = append(handles, nets.Nets.NetRef[i].Handle)
		}
	} else {
		// Try to decode as a single-net organization
		var net ARIN_OrgNet
		if err := json.Unmarshal(content, &net); err != nil {
			return handles, err
		}

		handles = append(handles, net.Nets.NetRef.Handle)
	}

	return handles, nil
}

func LookupNetCidrs(handle string) ([]string, error) {
	cidrs := []string{}

	safe_handle := url.QueryEscape(handle)
	u := fmt.Sprintf("http://whois.arin.net/rest/net/%s", safe_handle)

	req, err := http.NewRequest("GET", u, nil)
	if err != nil {
		return cidrs, err
	}
	req.Header.Set("Accept", "application/json")

	client := &http.Client{}
	resp, err := client.Do(req)
	if err != nil {
		return cidrs, err
	}

	defer resp.Body.Close()

	content, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return cidrs, err
	}

	var nets ARIN_Nets

	if err := json.Unmarshal(content, &nets); err == nil {
		for i := range nets.Net.NetBlocks.NetBlock {
			cidrs = append(cidrs, fmt.Sprintf("%s/%s", nets.Net.NetBlocks.NetBlock[i].StartAddress.Value, nets.Net.NetBlocks.NetBlock[i].CidrLength.Value))
		}
	} else {
		// Try to decode as a single-block network
		var net ARIN_Net
		if err := json.Unmarshal(content, &net); err != nil {
			return cidrs, err
		}
		cidrs = append(cidrs, fmt.Sprintf("%s/%s", net.Net.NetBlocks.NetBlock.StartAddress.Value, net.Net.NetBlocks.NetBlock.CidrLength.Value))
	}

	return cidrs, nil
}

func main() {

	if len(os.Args) != 2 {
		fmt.Println("Usage: inetdata-arin-org2nets <org-handle>\n")
		os.Exit(1)
	}

	org := os.Args[1]

	handles, e := LookupOrgNets(org)
	if e != nil {
		fmt.Fprintf(os.Stderr, "Could not list network handles: %s", e.Error())
		os.Exit(1)
	}

	for i := range handles {
		cidrs, e := LookupNetCidrs(handles[i])
		if e != nil {
			fmt.Fprintf(os.Stderr, "Could not list CIDRs for %s: %s", handles[i], e.Error())
			continue
		}
		fmt.Println(strings.Join(cidrs, "\n"))
	}
}
<|endoftext|>