<commit_before>package influxdb_output

import (
	"context"
	"crypto/tls"
	"encoding/json"
	"io/ioutil"
	"log"
	"math"
	"time"

	influxdb2 "github.com/influxdata/influxdb-client-go/v2"
	"github.com/karimra/gnmic/formatters"
	"github.com/karimra/gnmic/outputs"
	"github.com/openconfig/gnmi/proto/gnmi"
	"github.com/prometheus/client_golang/prometheus"
	"google.golang.org/protobuf/proto"
)

const (
	defaultURL               = "http://localhost:8086"
	defaultBatchSize         = 1000
	defaultFlushTimer        = 10 * time.Second
	defaultHealthCheckPeriod = 30 * time.Second

	numWorkers    = 1
	loggingPrefix = "[influxdb_output] "
)

func init() {
	outputs.Register("influxdb", func() outputs.Output {
		return &InfluxDBOutput{
			Cfg:       &Config{},
			eventChan: make(chan *formatters.EventMsg),
			reset:     make(chan struct{}),
			startSig:  make(chan struct{}),
			logger:    log.New(ioutil.Discard, loggingPrefix, log.LstdFlags|log.Lmicroseconds),
		}
	})
}

type InfluxDBOutput struct {
	Cfg       *Config
	client    influxdb2.Client
	logger    *log.Logger
	cancelFn  context.CancelFunc
	eventChan chan *formatters.EventMsg
	reset     chan struct{}
	startSig  chan struct{}
	wasup     bool
	evps      []formatters.EventProcessor
}
type Config struct {
	URL                string        `mapstructure:"url,omitempty"`
	Org                string        `mapstructure:"org,omitempty"`
	Bucket             string        `mapstructure:"bucket,omitempty"`
	Token              string        `mapstructure:"token,omitempty"`
	BatchSize          uint          `mapstructure:"batch-size,omitempty"`
	FlushTimer         time.Duration `mapstructure:"flush-timer,omitempty"`
	UseGzip            bool          `mapstructure:"use-gzip,omitempty"`
	EnableTLS          bool          `mapstructure:"enable-tls,omitempty"`
	HealthCheckPeriod  time.Duration `mapstructure:"health-check-period,omitempty"`
	Debug              bool          `mapstructure:"debug,omitempty"`
	EventProcessors    []string      `mapstructure:"event-processors,omitempty"`
	EnableMetrics      bool          `mapstructure:"enable-metrics,omitempty"`
	OverrideTimestamps bool          `mapstructure:"override-timestamps,omitempty"`
}

func (k *InfluxDBOutput) String() string {
	b, err := json.Marshal(k)
	if err != nil {
		return ""
	}
	return string(b)
}

func (i *InfluxDBOutput) SetLogger(logger *log.Logger) {
	if logger != nil && i.logger != nil {
		i.logger.SetOutput(logger.Writer())
		i.logger.SetFlags(logger.Flags())
	}
}

func (i *InfluxDBOutput) SetEventProcessors(ps map[string]map[string]interface{}, logger *log.Logger, tcs map[string]interface{}) {
	for _, epName := range i.Cfg.EventProcessors {
		if epCfg, ok := ps[epName]; ok {
			epType := ""
			for k := range epCfg {
				epType = k
				break
			}
			if in, ok := formatters.EventProcessors[epType]; ok {
				ep := in()
				err := ep.Init(epCfg[epType], formatters.WithLogger(logger), formatters.WithTargets(tcs))
				if err != nil {
					i.logger.Printf("failed initializing event processor '%s' of type='%s': %v", epName, epType, err)
					continue
				}
				i.evps = append(i.evps, ep)
				i.logger.Printf("added event processor '%s' of type=%s to influxdb output", epName, epType)
				continue
			}
			i.logger.Printf("%q event processor has an unknown type=%q", epName, epType)
			continue
		}
		i.logger.Printf("%q event processor not found!", epName)
	}
}

func (i *InfluxDBOutput) Init(ctx context.Context, name string, cfg map[string]interface{}, opts ...outputs.Option) error {
	err := outputs.DecodeConfig(cfg, i.Cfg)
	if err != nil {
		return err
	}
	for _, opt := range opts {
		opt(i)
	}
	if i.Cfg.URL == "" {
		i.Cfg.URL = defaultURL
	}
	if i.Cfg.BatchSize == 0 {
		i.Cfg.BatchSize = defaultBatchSize
	}
	if i.Cfg.FlushTimer == 0 {
		i.Cfg.FlushTimer = defaultFlushTimer
	}
	if i.Cfg.HealthCheckPeriod == 0 {
		i.Cfg.HealthCheckPeriod = defaultHealthCheckPeriod
	}

	iopts := influxdb2.DefaultOptions().
		SetUseGZip(i.Cfg.UseGzip).
		SetBatchSize(i.Cfg.BatchSize).
		SetFlushInterval(uint(i.Cfg.FlushTimer.Milliseconds()))
	if i.Cfg.EnableTLS {
		iopts.SetTLSConfig(&tls.Config{
			InsecureSkipVerify: true,
		})
	}
	if i.Cfg.Debug {
		iopts.SetLogLevel(3)
	}
	ctx, i.cancelFn = context.WithCancel(ctx)
CRCLIENT:
	i.client = influxdb2.NewClientWithOptions(i.Cfg.URL, i.Cfg.Token, iopts)
	// start influx health check
	err = i.health(ctx)
	if err != nil {
		i.logger.Printf("failed to check influxdb health: %v", err)
		time.Sleep(10 * time.Second)
		goto CRCLIENT
	}
	i.wasup = true
	go i.healthCheck(ctx)
	i.logger.Printf("initialized influxdb client: %s", i.String())

	for k := 0; k < numWorkers; k++ {
		go i.worker(ctx, k)
	}
	go func() {
		<-ctx.Done()
		i.Close()
	}()
	return nil
}

func (i *InfluxDBOutput) Write(ctx context.Context, rsp proto.Message, meta outputs.Meta) {
	if rsp == nil {
		return
	}
	switch rsp := rsp.(type) {
	case *gnmi.SubscribeResponse:
		measName := "default"
		if subName, ok := meta["subscription-name"]; ok {
			measName = subName
		}
		events, err := formatters.ResponseToEventMsgs(measName, rsp, meta, i.evps...)
		if err != nil {
			i.logger.Printf("failed to convert message to event: %v", err)
			return
		}
		for _, ev := range events {
			select {
			case <-ctx.Done():
				return
			case <-i.reset:
				return
			case i.eventChan <- ev:
			}
		}
	}
}

func (i *InfluxDBOutput) WriteEvent(ctx context.Context, ev *formatters.EventMsg) {}

func (i *InfluxDBOutput) Close() error {
	i.logger.Printf("closing client...")
	i.cancelFn()
	i.logger.Printf("closed.")
	return nil
}
func (i *InfluxDBOutput) RegisterMetrics(reg *prometheus.Registry) {}

func (i *InfluxDBOutput) healthCheck(ctx context.Context) {
	ticker := time.NewTicker(i.Cfg.HealthCheckPeriod)
	for {
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
			i.health(ctx)
		}
	}
}

func (i *InfluxDBOutput) health(ctx context.Context) error {
	res, err := i.client.Health(ctx)
	if err != nil {
		i.logger.Printf("failed health check: %v", err)
		if i.wasup {
			close(i.reset)
			i.reset = make(chan struct{})
		}
		return err
	}
	if res != nil {
		b, err := json.Marshal(res)
		if err != nil {
			i.logger.Printf("failed to marshal health check result: %v", err)
			i.logger.Printf("health check result: %+v", res)
			if i.wasup {
				close(i.reset)
				i.reset = make(chan struct{})
			}
			return err
		}
		i.wasup = true
		close(i.startSig)
		i.startSig = make(chan struct{})
		i.logger.Printf("health check result: %s", string(b))
		return nil
	}
	i.wasup = true
	close(i.startSig)
	i.startSig = make(chan struct{})
	i.logger.Print("health check result is nil")
	return nil
}

func (i *InfluxDBOutput) worker(ctx context.Context, idx int) {
	firstStart := true
START:
	if !firstStart {
		i.logger.Printf("worker-%d waiting for client recovery", idx)
		<-i.startSig
	}
	i.logger.Printf("starting worker-%d", idx)
	writer := i.client.WriteAPI(i.Cfg.Org, i.Cfg.Bucket)
	//defer writer.Flush()
	for {
		select {
		case <-ctx.Done():
			if ctx.Err() != nil {
				i.logger.Printf("worker-%d err=%v", idx, ctx.Err())
			}
			i.logger.Printf("worker-%d terminating...", idx)
			return
		case ev := <-i.eventChan:
			for n, v := range ev.Values {
				switch v := v.(type) {
				case *gnmi.Decimal64:
					ev.Values[n] = float64(v.Digits) / math.Pow10(int(v.Precision))
				}
			}
			if ev.Timestamp == 0 || i.Cfg.OverrideTimestamps {
				ev.Timestamp = time.Now().UnixNano()
			}
			writer.WritePoint(influxdb2.NewPoint(ev.Name, ev.Tags, ev.Values, time.Unix(0, ev.Timestamp)))
		case <-i.reset:
			firstStart = false
			i.logger.Printf("resetting worker-%d...", idx)
			goto START
		case err := <-writer.Errors():
			i.logger.Printf("worker-%d write error: %v", idx, err)
		}
	}
}

func (i *InfluxDBOutput) SetName(name string)        {}
func (i *InfluxDBOutput) SetClusterName(name string) {}
<commit_msg>convert uints to int if infxludb version is 1.8.x<commit_after>package influxdb_output

import (
	"context"
	"crypto/tls"
	"encoding/json"
	"io/ioutil"
	"log"
	"math"
	"strings"
	"time"

	influxdb2 "github.com/influxdata/influxdb-client-go/v2"
	"github.com/karimra/gnmic/formatters"
	"github.com/karimra/gnmic/outputs"
	"github.com/openconfig/gnmi/proto/gnmi"
	"github.com/prometheus/client_golang/prometheus"
	"google.golang.org/protobuf/proto"
)

const (
	defaultURL               = "http://localhost:8086"
	defaultBatchSize         = 1000
	defaultFlushTimer        = 10 * time.Second
	defaultHealthCheckPeriod = 30 * time.Second

	numWorkers    = 1
	loggingPrefix = "[influxdb_output] "
)

func init() {
	outputs.Register("influxdb", func() outputs.Output {
		return &InfluxDBOutput{
			Cfg:       &Config{},
			eventChan: make(chan *formatters.EventMsg),
			reset:     make(chan struct{}),
			startSig:  make(chan struct{}),
			logger:    log.New(ioutil.Discard, loggingPrefix, log.LstdFlags|log.Lmicroseconds),
		}
	})
}

type InfluxDBOutput struct {
	Cfg       *Config
	client    influxdb2.Client
	logger    *log.Logger
	cancelFn  context.CancelFunc
	eventChan chan *formatters.EventMsg
	reset     chan struct{}
	startSig  chan struct{}
	wasUP     bool
	evps      []formatters.EventProcessor
	dbVersion string
}
type Config struct {
	URL                string        `mapstructure:"url,omitempty"`
	Org                string        `mapstructure:"org,omitempty"`
	Bucket             string        `mapstructure:"bucket,omitempty"`
	Token              string        `mapstructure:"token,omitempty"`
	BatchSize          uint          `mapstructure:"batch-size,omitempty"`
	FlushTimer         time.Duration `mapstructure:"flush-timer,omitempty"`
	UseGzip            bool          `mapstructure:"use-gzip,omitempty"`
	EnableTLS          bool          `mapstructure:"enable-tls,omitempty"`
	HealthCheckPeriod  time.Duration `mapstructure:"health-check-period,omitempty"`
	Debug              bool          `mapstructure:"debug,omitempty"`
	EventProcessors    []string      `mapstructure:"event-processors,omitempty"`
	EnableMetrics      bool          `mapstructure:"enable-metrics,omitempty"`
	OverrideTimestamps bool          `mapstructure:"override-timestamps,omitempty"`
}

func (k *InfluxDBOutput) String() string {
	b, err := json.Marshal(k)
	if err != nil {
		return ""
	}
	return string(b)
}

func (i *InfluxDBOutput) SetLogger(logger *log.Logger) {
	if logger != nil && i.logger != nil {
		i.logger.SetOutput(logger.Writer())
		i.logger.SetFlags(logger.Flags())
	}
}

func (i *InfluxDBOutput) SetEventProcessors(ps map[string]map[string]interface{}, logger *log.Logger, tcs map[string]interface{}) {
	for _, epName := range i.Cfg.EventProcessors {
		if epCfg, ok := ps[epName]; ok {
			epType := ""
			for k := range epCfg {
				epType = k
				break
			}
			if in, ok := formatters.EventProcessors[epType]; ok {
				ep := in()
				err := ep.Init(epCfg[epType], formatters.WithLogger(logger), formatters.WithTargets(tcs))
				if err != nil {
					i.logger.Printf("failed initializing event processor '%s' of type='%s': %v", epName, epType, err)
					continue
				}
				i.evps = append(i.evps, ep)
				i.logger.Printf("added event processor '%s' of type=%s to influxdb output", epName, epType)
				continue
			}
			i.logger.Printf("%q event processor has an unknown type=%q", epName, epType)
			continue
		}
		i.logger.Printf("%q event processor not found!", epName)
	}
}

func (i *InfluxDBOutput) Init(ctx context.Context, name string, cfg map[string]interface{}, opts ...outputs.Option) error {
	err := outputs.DecodeConfig(cfg, i.Cfg)
	if err != nil {
		return err
	}
	for _, opt := range opts {
		opt(i)
	}
	if i.Cfg.URL == "" {
		i.Cfg.URL = defaultURL
	}
	if i.Cfg.BatchSize == 0 {
		i.Cfg.BatchSize = defaultBatchSize
	}
	if i.Cfg.FlushTimer == 0 {
		i.Cfg.FlushTimer = defaultFlushTimer
	}
	if i.Cfg.HealthCheckPeriod == 0 {
		i.Cfg.HealthCheckPeriod = defaultHealthCheckPeriod
	}

	iopts := influxdb2.DefaultOptions().
		SetUseGZip(i.Cfg.UseGzip).
		SetBatchSize(i.Cfg.BatchSize).
		SetFlushInterval(uint(i.Cfg.FlushTimer.Milliseconds()))
	if i.Cfg.EnableTLS {
		iopts.SetTLSConfig(&tls.Config{
			InsecureSkipVerify: true,
		})
	}
	if i.Cfg.Debug {
		iopts.SetLogLevel(3)
	}
	ctx, i.cancelFn = context.WithCancel(ctx)
CRCLIENT:
	i.client = influxdb2.NewClientWithOptions(i.Cfg.URL, i.Cfg.Token, iopts)
	// start influx health check
	err = i.health(ctx)
	if err != nil {
		i.logger.Printf("failed to check influxdb health: %v", err)
		time.Sleep(10 * time.Second)
		goto CRCLIENT
	}
	i.wasUP = true
	go i.healthCheck(ctx)
	i.logger.Printf("initialized influxdb client: %s", i.String())

	for k := 0; k < numWorkers; k++ {
		go i.worker(ctx, k)
	}
	go func() {
		<-ctx.Done()
		i.Close()
	}()
	return nil
}

func (i *InfluxDBOutput) Write(ctx context.Context, rsp proto.Message, meta outputs.Meta) {
	if rsp == nil {
		return
	}
	switch rsp := rsp.(type) {
	case *gnmi.SubscribeResponse:
		measName := "default"
		if subName, ok := meta["subscription-name"]; ok {
			measName = subName
		}
		events, err := formatters.ResponseToEventMsgs(measName, rsp, meta, i.evps...)
		if err != nil {
			i.logger.Printf("failed to convert message to event: %v", err)
			return
		}
		for _, ev := range events {
			select {
			case <-ctx.Done():
				return
			case <-i.reset:
				return
			case i.eventChan <- ev:
			}
		}
	}
}

func (i *InfluxDBOutput) WriteEvent(ctx context.Context, ev *formatters.EventMsg) {}

func (i *InfluxDBOutput) Close() error {
	i.logger.Printf("closing client...")
	i.cancelFn()
	i.logger.Printf("closed.")
	return nil
}
func (i *InfluxDBOutput) RegisterMetrics(reg *prometheus.Registry) {}

func (i *InfluxDBOutput) healthCheck(ctx context.Context) {
	ticker := time.NewTicker(i.Cfg.HealthCheckPeriod)
	for {
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
			i.health(ctx)
		}
	}
}

func (i *InfluxDBOutput) health(ctx context.Context) error {
	res, err := i.client.Health(ctx)
	if err != nil {
		i.logger.Printf("failed health check: %v", err)
		if i.wasUP {
			close(i.reset)
			i.reset = make(chan struct{})
		}
		return err
	}
	if res != nil {
		if res.Version != nil {
			i.dbVersion = *res.Version
		}
		b, err := json.Marshal(res)
		if err != nil {
			i.logger.Printf("failed to marshal health check result: %v", err)
			i.logger.Printf("health check result: %+v", res)
			if i.wasUP {
				close(i.reset)
				i.reset = make(chan struct{})
			}
			return err
		}
		i.wasUP = true
		close(i.startSig)
		i.startSig = make(chan struct{})
		i.logger.Printf("health check result: %s", string(b))
		return nil
	}
	i.wasUP = true
	close(i.startSig)
	i.startSig = make(chan struct{})
	i.logger.Print("health check result is nil")
	return nil
}

func (i *InfluxDBOutput) worker(ctx context.Context, idx int) {
	firstStart := true
START:
	if !firstStart {
		i.logger.Printf("worker-%d waiting for client recovery", idx)
		<-i.startSig
	}
	i.logger.Printf("starting worker-%d", idx)
	writer := i.client.WriteAPI(i.Cfg.Org, i.Cfg.Bucket)
	//defer writer.Flush()
	for {
		select {
		case <-ctx.Done():
			if ctx.Err() != nil {
				i.logger.Printf("worker-%d err=%v", idx, ctx.Err())
			}
			i.logger.Printf("worker-%d terminating...", idx)
			return
		case ev := <-i.eventChan:
			for n, v := range ev.Values {
				switch v := v.(type) {
				case *gnmi.Decimal64:
					ev.Values[n] = float64(v.Digits) / math.Pow10(int(v.Precision))
				}
			}
			if ev.Timestamp == 0 || i.Cfg.OverrideTimestamps {
				ev.Timestamp = time.Now().UnixNano()
			}
			i.convertUints(ev)
			writer.WritePoint(influxdb2.NewPoint(ev.Name, ev.Tags, ev.Values, time.Unix(0, ev.Timestamp)))
		case <-i.reset:
			firstStart = false
			i.logger.Printf("resetting worker-%d...", idx)
			goto START
		case err := <-writer.Errors():
			i.logger.Printf("worker-%d write error: %v", idx, err)
		}
	}
}

func (i *InfluxDBOutput) SetName(name string)        {}
func (i *InfluxDBOutput) SetClusterName(name string) {}

func (i *InfluxDBOutput) convertUints(ev *formatters.EventMsg) {
	if !strings.HasPrefix(i.dbVersion, "1.8") {
		return
	}
	for k, v := range ev.Values {
		switch v := v.(type) {
		case uint:
			ev.Values[k] = int(v)
		}
	}
}
<|endoftext|>


<commit_before>// Copyright 2017 Google LLC. All Rights Reserved.
//
// 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 merkle

import (
	"bytes"
	"errors"
	"fmt"
	"math/bits"
)

// RootMismatchError occurs when an inclusion proof fails.
type RootMismatchError struct {
	ExpectedRoot   []byte
	CalculatedRoot []byte
}

func (e RootMismatchError) Error() string {
	return fmt.Sprintf("calculated root:\n%v\n does not match expected root:\n%v", e.CalculatedRoot, e.ExpectedRoot)
}

// LogHasher provides the hash functions needed to compute dense merkle trees.
type LogHasher interface {
	// EmptyRoot supports returning a special case for the root of an empty tree.
	EmptyRoot() []byte
	// HashLeaf computes the hash of a leaf that exists.
	HashLeaf(leaf []byte) []byte
	// HashChildren computes interior nodes.
	HashChildren(l, r []byte) []byte
}

// LogVerifier verifies inclusion and consistency proofs for append only logs.
type LogVerifier struct {
	hasher LogHasher
}

// New returns a new LogVerifier for a tree.
func New(hasher LogHasher) LogVerifier {
	return LogVerifier{hasher}
}

// VerifyInclusionProof verifies the correctness of the proof given the passed
// in information about the tree and leaf.
func (v LogVerifier) VerifyInclusionProof(leafIndex, treeSize int64, proof [][]byte, root []byte, leafHash []byte) error {
	calcRoot, err := v.RootFromInclusionProof(leafIndex, treeSize, proof, leafHash)
	if err != nil {
		return err
	}
	if !bytes.Equal(calcRoot, root) {
		return RootMismatchError{
			CalculatedRoot: calcRoot,
			ExpectedRoot:   root,
		}
	}
	return nil
}

// RootFromInclusionProof calculates the expected tree root given the proof and leaf.
// leafIndex starts at 0.  treeSize is the number of nodes in the tree.
// proof is an array of neighbor nodes from the bottom to the root.
func (v LogVerifier) RootFromInclusionProof(leafIndex, treeSize int64, proof [][]byte, leafHash []byte) ([]byte, error) {
	switch {
	case leafIndex < 0:
		return nil, fmt.Errorf("leafIndex %d < 0", leafIndex)
	case treeSize < 0:
		return nil, fmt.Errorf("treeSize %d < 0", treeSize)
	case leafIndex >= treeSize:
		return nil, fmt.Errorf("leafIndex is beyond treeSize: %d >= %d", leafIndex, treeSize)
	}

	inner, border := decompInclProof(leafIndex, treeSize)
	if got, want := len(proof), inner+border; got != want {
		return nil, fmt.Errorf("wrong proof size %d, want %d", got, want)
	}

	ch := hashChainer(v)
	res := ch.chainInner(leafHash, proof[:inner], leafIndex)
	res = ch.chainBorderRight(res, proof[inner:])
	return res, nil
}

// VerifyConsistencyProof checks that the passed in consistency proof is valid
// between the passed in tree snapshots. Snapshots are the respective tree
// sizes. Accepts shapshot2 >= snapshot1 >= 0.
func (v LogVerifier) VerifyConsistencyProof(snapshot1, snapshot2 int64, root1, root2 []byte, proof [][]byte) error {
	switch {
	case snapshot1 < 0:
		return fmt.Errorf("snapshot1 (%d) < 0 ", snapshot1)
	case snapshot2 < snapshot1:
		return fmt.Errorf("snapshot2 (%d) < snapshot1 (%d)", snapshot1, snapshot2)
	case snapshot1 == snapshot2:
		if !bytes.Equal(root1, root2) {
			return RootMismatchError{
				CalculatedRoot: root1,
				ExpectedRoot:   root2,
			}
		} else if len(proof) > 0 {
			return errors.New("root1 and root2 match, but proof is non-empty")
		}
		return nil // Proof OK.
	case snapshot1 == 0:
		// Any snapshot greater than 0 is consistent with snapshot 0.
		if len(proof) > 0 {
			return fmt.Errorf("expected empty proof, but got %d components", len(proof))
		}
		return nil // Proof OK.
	case len(proof) == 0:
		return errors.New("empty proof")
	}

	inner, border := decompInclProof(snapshot1-1, snapshot2)
	shift := bits.TrailingZeros64(uint64(snapshot1))
	inner -= shift // Note: shift < inner if snapshot1 < snapshot2.

	// The proof includes the root hash for the sub-tree of size 2^shift.
	seed, start := proof[0], 1
	if snapshot1 == 1<<uint(shift) { // Unless snapshot1 is that very 2^shift.
		seed, start = root1, 0
	}
	if got, want := len(proof), start+inner+border; got != want {
		return fmt.Errorf("wrong proof size %d, want %d", got, want)
	}
	proof = proof[start:]
	// Now len(proof) == inner+border, and proof is effectively a suffix of
	// inclusion proof for entry |snapshot1-1| in a tree of size |snapshot2|.

	// Verify the first root.
	ch := hashChainer(v)
	mask := (snapshot1 - 1) >> uint(shift) // Start chaining from level |shift|.
	hash1 := ch.chainInnerRight(seed, proof[:inner], mask)
	hash1 = ch.chainBorderRight(hash1, proof[inner:])
	if !bytes.Equal(hash1, root1) {
		return RootMismatchError{
			CalculatedRoot: hash1,
			ExpectedRoot:   root1,
		}
	}

	// Verify the second root.
	hash2 := ch.chainInner(seed, proof[:inner], mask)
	hash2 = ch.chainBorderRight(hash2, proof[inner:])
	if !bytes.Equal(hash2, root2) {
		return RootMismatchError{
			CalculatedRoot: hash2,
			ExpectedRoot:   root2,
		}
	}

	return nil // Proof OK.
}

// VerifiedPrefixHashFromInclusionProof calculates a root hash over leaves
// [0..subSize), based on the inclusion |proof| and |leafHash| for a leaf at
// index |subSize-1| in a tree of the specified |size| with the passed in
// |root| hash.
// Returns an error if the |proof| verification fails. The resulting smaller
// tree's root hash is trusted iff the bigger tree's |root| hash is trusted.
func (v LogVerifier) VerifiedPrefixHashFromInclusionProof(
	subSize, size int64,
	proof [][]byte, root []byte, leafHash []byte,
) ([]byte, error) {
	if subSize <= 0 {
		return nil, fmt.Errorf("subtree size is %d, want > 0", subSize)
	}
	leaf := subSize - 1
	if err := v.VerifyInclusionProof(leaf, size, proof, root, leafHash); err != nil {
		return nil, err
	}

	inner := innerProofSize(leaf, size)
	ch := hashChainer(v)
	res := ch.chainInnerRight(leafHash, proof[:inner], leaf)
	res = ch.chainBorderRight(res, proof[inner:])
	return res, nil
}

// decompInclProof breaks down inclusion proof for a leaf at the specified
// |index| in a tree of the specified |size| into 2 components. The splitting
// point between them is where paths to leaves |index| and |size-1| diverge.
// Returns lengths of the bottom and upper proof parts correspondingly. The sum
// of the two determines the correct length of the inclusion proof.
func decompInclProof(index, size int64) (int, int) {
	inner := innerProofSize(index, size)
	border := bits.OnesCount64(uint64(index) >> uint(inner))
	return inner, border
}

func innerProofSize(index, size int64) int {
	return bits.Len64(uint64(index ^ (size - 1)))
}
<commit_msg>New -> NewLogVerifier<commit_after>// Copyright 2017 Google LLC. All Rights Reserved.
//
// 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 merkle

import (
	"bytes"
	"errors"
	"fmt"
	"math/bits"
)

// RootMismatchError occurs when an inclusion proof fails.
type RootMismatchError struct {
	ExpectedRoot   []byte
	CalculatedRoot []byte
}

func (e RootMismatchError) Error() string {
	return fmt.Sprintf("calculated root:\n%v\n does not match expected root:\n%v", e.CalculatedRoot, e.ExpectedRoot)
}

// LogHasher provides the hash functions needed to compute dense merkle trees.
type LogHasher interface {
	// EmptyRoot supports returning a special case for the root of an empty tree.
	EmptyRoot() []byte
	// HashLeaf computes the hash of a leaf that exists.
	HashLeaf(leaf []byte) []byte
	// HashChildren computes interior nodes.
	HashChildren(l, r []byte) []byte
}

// LogVerifier verifies inclusion and consistency proofs for append only logs.
type LogVerifier struct {
	hasher LogHasher
}

// NewLogVerifier returns a new LogVerifier for a tree.
func NewLogVerifier(hasher LogHasher) LogVerifier {
	return LogVerifier{hasher}
}

// VerifyInclusionProof verifies the correctness of the proof given the passed
// in information about the tree and leaf.
func (v LogVerifier) VerifyInclusionProof(leafIndex, treeSize int64, proof [][]byte, root []byte, leafHash []byte) error {
	calcRoot, err := v.RootFromInclusionProof(leafIndex, treeSize, proof, leafHash)
	if err != nil {
		return err
	}
	if !bytes.Equal(calcRoot, root) {
		return RootMismatchError{
			CalculatedRoot: calcRoot,
			ExpectedRoot:   root,
		}
	}
	return nil
}

// RootFromInclusionProof calculates the expected tree root given the proof and leaf.
// leafIndex starts at 0.  treeSize is the number of nodes in the tree.
// proof is an array of neighbor nodes from the bottom to the root.
func (v LogVerifier) RootFromInclusionProof(leafIndex, treeSize int64, proof [][]byte, leafHash []byte) ([]byte, error) {
	switch {
	case leafIndex < 0:
		return nil, fmt.Errorf("leafIndex %d < 0", leafIndex)
	case treeSize < 0:
		return nil, fmt.Errorf("treeSize %d < 0", treeSize)
	case leafIndex >= treeSize:
		return nil, fmt.Errorf("leafIndex is beyond treeSize: %d >= %d", leafIndex, treeSize)
	}

	inner, border := decompInclProof(leafIndex, treeSize)
	if got, want := len(proof), inner+border; got != want {
		return nil, fmt.Errorf("wrong proof size %d, want %d", got, want)
	}

	ch := hashChainer(v)
	res := ch.chainInner(leafHash, proof[:inner], leafIndex)
	res = ch.chainBorderRight(res, proof[inner:])
	return res, nil
}

// VerifyConsistencyProof checks that the passed in consistency proof is valid
// between the passed in tree snapshots. Snapshots are the respective tree
// sizes. Accepts shapshot2 >= snapshot1 >= 0.
func (v LogVerifier) VerifyConsistencyProof(snapshot1, snapshot2 int64, root1, root2 []byte, proof [][]byte) error {
	switch {
	case snapshot1 < 0:
		return fmt.Errorf("snapshot1 (%d) < 0 ", snapshot1)
	case snapshot2 < snapshot1:
		return fmt.Errorf("snapshot2 (%d) < snapshot1 (%d)", snapshot1, snapshot2)
	case snapshot1 == snapshot2:
		if !bytes.Equal(root1, root2) {
			return RootMismatchError{
				CalculatedRoot: root1,
				ExpectedRoot:   root2,
			}
		} else if len(proof) > 0 {
			return errors.New("root1 and root2 match, but proof is non-empty")
		}
		return nil // Proof OK.
	case snapshot1 == 0:
		// Any snapshot greater than 0 is consistent with snapshot 0.
		if len(proof) > 0 {
			return fmt.Errorf("expected empty proof, but got %d components", len(proof))
		}
		return nil // Proof OK.
	case len(proof) == 0:
		return errors.New("empty proof")
	}

	inner, border := decompInclProof(snapshot1-1, snapshot2)
	shift := bits.TrailingZeros64(uint64(snapshot1))
	inner -= shift // Note: shift < inner if snapshot1 < snapshot2.

	// The proof includes the root hash for the sub-tree of size 2^shift.
	seed, start := proof[0], 1
	if snapshot1 == 1<<uint(shift) { // Unless snapshot1 is that very 2^shift.
		seed, start = root1, 0
	}
	if got, want := len(proof), start+inner+border; got != want {
		return fmt.Errorf("wrong proof size %d, want %d", got, want)
	}
	proof = proof[start:]
	// Now len(proof) == inner+border, and proof is effectively a suffix of
	// inclusion proof for entry |snapshot1-1| in a tree of size |snapshot2|.

	// Verify the first root.
	ch := hashChainer(v)
	mask := (snapshot1 - 1) >> uint(shift) // Start chaining from level |shift|.
	hash1 := ch.chainInnerRight(seed, proof[:inner], mask)
	hash1 = ch.chainBorderRight(hash1, proof[inner:])
	if !bytes.Equal(hash1, root1) {
		return RootMismatchError{
			CalculatedRoot: hash1,
			ExpectedRoot:   root1,
		}
	}

	// Verify the second root.
	hash2 := ch.chainInner(seed, proof[:inner], mask)
	hash2 = ch.chainBorderRight(hash2, proof[inner:])
	if !bytes.Equal(hash2, root2) {
		return RootMismatchError{
			CalculatedRoot: hash2,
			ExpectedRoot:   root2,
		}
	}

	return nil // Proof OK.
}

// VerifiedPrefixHashFromInclusionProof calculates a root hash over leaves
// [0..subSize), based on the inclusion |proof| and |leafHash| for a leaf at
// index |subSize-1| in a tree of the specified |size| with the passed in
// |root| hash.
// Returns an error if the |proof| verification fails. The resulting smaller
// tree's root hash is trusted iff the bigger tree's |root| hash is trusted.
func (v LogVerifier) VerifiedPrefixHashFromInclusionProof(
	subSize, size int64,
	proof [][]byte, root []byte, leafHash []byte,
) ([]byte, error) {
	if subSize <= 0 {
		return nil, fmt.Errorf("subtree size is %d, want > 0", subSize)
	}
	leaf := subSize - 1
	if err := v.VerifyInclusionProof(leaf, size, proof, root, leafHash); err != nil {
		return nil, err
	}

	inner := innerProofSize(leaf, size)
	ch := hashChainer(v)
	res := ch.chainInnerRight(leafHash, proof[:inner], leaf)
	res = ch.chainBorderRight(res, proof[inner:])
	return res, nil
}

// decompInclProof breaks down inclusion proof for a leaf at the specified
// |index| in a tree of the specified |size| into 2 components. The splitting
// point between them is where paths to leaves |index| and |size-1| diverge.
// Returns lengths of the bottom and upper proof parts correspondingly. The sum
// of the two determines the correct length of the inclusion proof.
func decompInclProof(index, size int64) (int, int) {
	inner := innerProofSize(index, size)
	border := bits.OnesCount64(uint64(index) >> uint(inner))
	return inner, border
}

func innerProofSize(index, size int64) int {
	return bits.Len64(uint64(index ^ (size - 1)))
}
<|endoftext|>


<commit_before>package logging

import (
	"os"
	"strings"
	"testing"
)

// These are provided to ease testing code that uses the logging pkg

type mockWriter struct {
	written []byte
}

func (w *mockWriter) Write(p []byte) (n int, err os.Error) {
	w.written = append(w.written, p...)
	return len(p), nil
}

func (w *mockWriter) reset() {
	w.written = w.written[:0]
}

type writerMap map[LogLevel]*mockWriter

// This doesn't create a mock Logger but a Logger that writes to mock outputs
func NewMock() (*logger, writerMap) {
	wMap := writerMap{
		Debug: &mockWriter{make([]byte, 0)},
		Info:  &mockWriter{make([]byte, 0)},
		Warn:  &mockWriter{make([]byte, 0)},
		Error: &mockWriter{make([]byte, 0)},
		Fatal: &mockWriter{make([]byte, 0)},
	}
	logMap := make(LogMap)
	for lv, w := range wMap {
		logMap[lv] = makeLogger(w)
	}
	return New(logMap, Error, false), wMap
}

func (m writerMap) CheckNothingWritten(t *testing.T) {
	for lv, w := range m {
		if len(w.written) > 0 {
			t.Errorf("%d bytes logged at level %s, expected none:",
				len(w.written), logString[lv])
			t.Errorf("\t%s", w.written)
			w.reset()
		}
	}
}

func (m writerMap) CheckWrittenAtLevel(t *testing.T, lv LogLevel, exp string) {
	var w *mockWriter
	if _, ok := m[lv]; !ok {
		w = m[Debug]
	} else {
		w = m[lv]
	}
	// 32 bytes covers the date, time and filename up to the colon in
	// 2011/10/22 10:22:57 log_test.go:<line no>: <level> <log message>
	if len(w.written) <= 32 {
		t.Errorf("Not enough bytes logged at level %s:", LogString(lv))
		t.Errorf("\tgot: %s", string(w.written))
	}
	s := string(w.written[32:])
	// 2 covers the : itself and the extra space
	idx := strings.Index(s, ":") + 2
	// s will end in "\n", so -1 to chop that off too
	s = s[idx:len(s)-1]
	// expected won't contain the log level prefix, so prepend that
	exp = LogString(lv) + " " + exp
	if s != exp {
		t.Errorf("Log message at level %s differed.", LogString(lv))
		t.Errorf("\texp: %s\n\tgot: %s", exp, s)
	}
	w.reset()
	// Calling checkNothingWritten here both tests that w.reset() works
	// and verifies that nothing was written at any other levels.
	m.CheckNothingWritten(t)
}

<commit_msg>Hmm, slice out of range error, wtf x3!<commit_after>package logging

import (
	"os"
	"strings"
	"testing"
)

// These are provided to ease testing code that uses the logging pkg

type mockWriter struct {
	written []byte
}

func (w *mockWriter) Write(p []byte) (n int, err os.Error) {
	w.written = append(w.written, p...)
	return len(p), nil
}

func (w *mockWriter) reset() {
	w.written = w.written[:0]
}

type writerMap map[LogLevel]*mockWriter

// This doesn't create a mock Logger but a Logger that writes to mock outputs
func NewMock() (*logger, writerMap) {
	wMap := writerMap{
		Debug: &mockWriter{make([]byte, 0)},
		Info:  &mockWriter{make([]byte, 0)},
		Warn:  &mockWriter{make([]byte, 0)},
		Error: &mockWriter{make([]byte, 0)},
		Fatal: &mockWriter{make([]byte, 0)},
	}
	logMap := make(LogMap)
	for lv, w := range wMap {
		logMap[lv] = makeLogger(w)
	}
	return New(logMap, Error, false), wMap
}

func (m writerMap) CheckNothingWritten(t *testing.T) {
	for lv, w := range m {
		if len(w.written) > 0 {
			t.Errorf("%d bytes logged at level %s, expected none:",
				len(w.written), logString[lv])
			t.Errorf("\t%s", w.written)
			w.reset()
		}
	}
}

func (m writerMap) CheckWrittenAtLevel(t *testing.T, lv LogLevel, exp string) {
	var w *mockWriter
	if _, ok := m[lv]; !ok {
		w = m[Debug]
	} else {
		w = m[lv]
	}
	// 32 bytes covers the date, time and filename up to the colon in
	// 2011/10/22 10:22:57 log_test.go:<line no>: <level> <log message>
	if len(w.written) <= 33 {
		t.Errorf("Not enough bytes logged at level %s:", LogString(lv))
		t.Errorf("\tgot: %s", string(w.written))
	}
	s := string(w.written[32:])
	// 2 covers the : itself and the extra space
	idx := strings.Index(s, ":") + 2
	// s will end in "\n", so -1 to chop that off too
	s = s[idx:len(s)-1]
	// expected won't contain the log level prefix, so prepend that
	exp = LogString(lv) + " " + exp
	if s != exp {
		t.Errorf("Log message at level %s differed.", LogString(lv))
		t.Errorf("\texp: %s\n\tgot: %s", exp, s)
	}
	w.reset()
	// Calling checkNothingWritten here both tests that w.reset() works
	// and verifies that nothing was written at any other levels.
	m.CheckNothingWritten(t)
}

<|endoftext|>


<commit_before>package main

import (
	"log"
	"os"
	"os/signal"
	"sync"
	"syscall"
	"time"
)

// globals are eeeeevil
var paused bool
var signalLock = sync.RWMutex{}

// we wrap access to `paused` in a RLock so that if we're in the middle of
// marking services for maintenance we don't get stale reads
func inMaintenanceMode() bool {
	signalLock.RLock()
	defer signalLock.RUnlock()
	return paused
}

func toggleMaintenanceMode(config *Config) {
	signalLock.Lock()
	defer signalLock.Unlock()
	paused = !paused
	if paused {
		forAllServices(config, func(service *ServiceConfig) {
			log.Printf("Marking for maintenance: %s\n", service.Name)
			service.MarkForMaintenance()
		})
	}
}

func terminate(config *Config) {
	signalLock.Lock()
	defer signalLock.Unlock()
	stopPolling(config)
	forAllServices(config, func(service *ServiceConfig) {
		log.Printf("Deregistering service: %s\n", service.Name)
		service.Deregister()
	})

	// Run and wait for preStop command to exit
	run(config.preStopCmd)

	cmd := config.Command
	if cmd == nil || cmd.Process == nil {
		// Not managing the process, so don't do anything
		return
	}
	if config.StopTimeout > 0 {
		if err := cmd.Process.Signal(syscall.SIGTERM); err != nil {
			log.Printf("Error sending SIGTERM to application: %s\n", err)
		} else {
			time.AfterFunc(time.Duration(config.StopTimeout)*time.Second, func() {
				log.Printf("Killing Process %#v\n", cmd.Process)
				cmd.Process.Kill()
			})
			return
		}
	}
	log.Printf("Killing Process %#v\n", config.Command.Process)
	cmd.Process.Kill()
}

func reloadConfig(config *Config) *Config {
	signalLock.Lock()
	defer signalLock.Unlock()
	log.Printf("Reloading configuration.\n")
	newConfig, err := loadConfig()
	if err != nil {
		log.Printf("Could not reload config: %v\n", err)
		return nil
	}
	// stop advertising the existing services so that we can
	// make sure we update them if ports, etc. change.
	stopPolling(config)
	forAllServices(config, func(service *ServiceConfig) {
		log.Printf("Deregistering service: %s\n", service.Name)
		service.Deregister()
	})

	signal.Reset()
	handleSignals(newConfig)
	handlePolling(newConfig)

	return newConfig // return for debuggability
}

func stopPolling(config *Config) {
	for _, quit := range config.QuitChannels {
		quit <- true
	}
}

type serviceFunc func(service *ServiceConfig)

func forAllServices(config *Config, fn serviceFunc) {
	for _, service := range config.Services {
		fn(service)
	}
}

// Listen for and capture signals from the OS
func handleSignals(config *Config) {
	sig := make(chan os.Signal, 1)
	signal.Notify(sig, syscall.SIGUSR1, syscall.SIGTERM, syscall.SIGHUP)
	go func() {
		for signal := range sig {
			switch signal {
			case syscall.SIGUSR1:
				toggleMaintenanceMode(config)
			case syscall.SIGTERM:
				terminate(config)
			case syscall.SIGHUP:
				reloadConfig(config)
			}
		}
	}()
}
<commit_msg>Split maintenance mode lock from signal handler lock.<commit_after>package main

import (
	"log"
	"os"
	"os/signal"
	"sync"
	"syscall"
	"time"
)

// globals are eeeeevil
var paused bool
var signalLock = sync.RWMutex{}
var maintModeLock = sync.RWMutex{}

// we wrap access to `paused` in a RLock so that if we're in the middle of
// marking services for maintenance we don't get stale reads
func inMaintenanceMode() bool {
	maintModeLock.RLock()
	defer maintModeLock.RUnlock()
	return paused
}

func toggleMaintenanceMode(config *Config) {
	maintModeLock.RLock()
	signalLock.Lock()
	defer signalLock.Unlock()
	defer maintModeLock.RUnlock()
	paused = !paused
	if paused {
		forAllServices(config, func(service *ServiceConfig) {
			log.Printf("Marking for maintenance: %s\n", service.Name)
			service.MarkForMaintenance()
		})
	}
}

func terminate(config *Config) {
	signalLock.Lock()
	defer signalLock.Unlock()
	stopPolling(config)
	forAllServices(config, func(service *ServiceConfig) {
		log.Printf("Deregistering service: %s\n", service.Name)
		service.Deregister()
	})

	// Run and wait for preStop command to exit
	run(config.preStopCmd)

	cmd := config.Command
	if cmd == nil || cmd.Process == nil {
		// Not managing the process, so don't do anything
		return
	}
	if config.StopTimeout > 0 {
		if err := cmd.Process.Signal(syscall.SIGTERM); err != nil {
			log.Printf("Error sending SIGTERM to application: %s\n", err)
		} else {
			time.AfterFunc(time.Duration(config.StopTimeout)*time.Second, func() {
				log.Printf("Killing Process %#v\n", cmd.Process)
				cmd.Process.Kill()
			})
			return
		}
	}
	log.Printf("Killing Process %#v\n", config.Command.Process)
	cmd.Process.Kill()
}

func reloadConfig(config *Config) *Config {
	signalLock.Lock()
	defer signalLock.Unlock()
	log.Printf("Reloading configuration.\n")
	newConfig, err := loadConfig()
	if err != nil {
		log.Printf("Could not reload config: %v\n", err)
		return nil
	}
	// stop advertising the existing services so that we can
	// make sure we update them if ports, etc. change.
	stopPolling(config)
	forAllServices(config, func(service *ServiceConfig) {
		log.Printf("Deregistering service: %s\n", service.Name)
		service.Deregister()
	})

	signal.Reset()
	handleSignals(newConfig)
	handlePolling(newConfig)

	return newConfig // return for debuggability
}

func stopPolling(config *Config) {
	for _, quit := range config.QuitChannels {
		quit <- true
	}
}

type serviceFunc func(service *ServiceConfig)

func forAllServices(config *Config, fn serviceFunc) {
	for _, service := range config.Services {
		fn(service)
	}
}

// Listen for and capture signals from the OS
func handleSignals(config *Config) {
	sig := make(chan os.Signal, 1)
	signal.Notify(sig, syscall.SIGUSR1, syscall.SIGTERM, syscall.SIGHUP)
	go func() {
		for signal := range sig {
			switch signal {
			case syscall.SIGUSR1:
				toggleMaintenanceMode(config)
			case syscall.SIGTERM:
				terminate(config)
			case syscall.SIGHUP:
				reloadConfig(config)
			}
		}
	}()
}
<|endoftext|>


<commit_before>// Package resources provides unfancy resources embedding with Go.
package resources

import (
	"bytes"
	"fmt"
	"io"
	"log"
	"os"
	"text/template"
)

// File mimicks the os.File and http.File interface.
type File interface {
	io.Reader
	Stat() (os.FileInfo, error)
}

// New creates a new Package.
func New() *Package {
	return &Package{
		Config: Config{
			Pkg:     "resources",
			Var:     "FS",
			Declare: true,
		},
		Files: make(map[string]File),
	}
}

// Config defines some details about the output Go file.
type Config struct {
	Pkg     string // Package name
	Var     string // Variable name to assign the file system to.
	Tag     string // Build tag, leave empty for no tag.
	Declare bool   // Dictates whatever there should be a defintion Variable
}

// Package describes...
type Package struct {
	Config
	Files map[string]File
}

//Add a file to the package at the give path.
func (p *Package) Add(path string, file File) error {
	p.Files[path] = file
	return nil
}

// AddFile is a helper function that adds the files from the path into the package under the path file.
func (p *Package) AddFile(path string, file string) error {
	f, err := os.Open(file)
	if err != nil {
		return err
	}
	return p.Add(path, f)
}

// Build the package
func (p *Package) Build(out io.Writer) error {
	return pkg.Execute(out, p)
}

// Write the build to a file, you don't need to call Build.
func (p *Package) Write(path string) error {
	f, err := os.Create(path)
	if err != nil {
		return err
	}
	defer func() {
		err := f.Close()
		if err != nil {
			log.Panicf("Failed to close file: %s", err)
		}
	}()

	return p.Build(f)
}

// Template
var pkg *template.Template

func reader(input io.Reader) (string, error) {

	var (
		buff       bytes.Buffer
		err        error
		blockwidth = 12
		curblock   = 0
	)

	b := make([]byte, blockwidth)

	for n, err := input.Read(b); err == nil; n, err = input.Read(b) {
		for i := 0; i < n; i++ {
			_, err = fmt.Fprintf(&buff, "0x%02x,", b[i])
			if err != nil {
				break
			}
			curblock++
			if curblock < blockwidth {
				continue
			}
			buff.WriteByte('\n')
			buff.Write([]byte{'\t', '\t'})
			curblock = 0
		}
	}

	return buff.String(), err
}

func init() {

	pkg = template.Must(template.New("file").Funcs(template.FuncMap{"reader": reader}).Parse(` File{
	  data: []byte{
	{{ reader . }}
  },
  fi: FileInfo {
	name:    "{{ .Stat.Name }}",
    size:    {{ .Stat.Size }},
	modTime: time.Unix(0, {{ .Stat.ModTime.UnixNano }}),
    isDir:   {{ .Stat.IsDir }},
  },
}`))

	pkg = template.Must(pkg.New("pkg").Parse(`{{ if .Tag }}// +build {{ .Tag }}

{{ end }}// Package {{.Pkg }} is generated by github.com/omeid/go-resources
package {{ .Pkg }}

import (
	"bytes"
	"errors"
	"net/http"
	"os"
	"path/filepath"
	"strings"
	"time"
)

// FileSystem is an http.FileSystem implementation.
type FileSystem struct {
	files map[string]File
}

// String returns the content of the file as string.
func (fs *FileSystem) String(name string) (string, bool) {
	if filepath.Separator != '/' && strings.IndexRune(name, filepath.Separator) >= 0 ||
		strings.Contains(name, "\x00") {
		return "", false
	}

	file, ok := fs.files[name]

	if !ok {
		return "", false
	}

	return string(file.data), true
}

// Open implements http.FileSystem.Open
func (fs *FileSystem) Open(name string) (http.File, error) {
	if filepath.Separator != '/' && strings.IndexRune(name, filepath.Separator) >= 0 ||
		strings.Contains(name, "\x00") {
		return nil, errors.New("http: invalid character in file path")
	}
	file, ok := fs.files[name]
	if !ok {
		files := []os.FileInfo{}
		for path, file := range fs.files {
			if strings.HasPrefix(path, name) {
				fi := file.fi
				files = append(files, &fi)
			}
		}

		if len(files) == 0 {
			return nil, os.ErrNotExist
		}

		//We have a directory.
		return &File{
			fi: FileInfo{
				isDir: true,
				files: files,
			}}, nil
	}
	file.Reader = bytes.NewReader(file.data)
	return &file, nil
}

// File implements http.File
type File struct {
	*bytes.Reader
	data []byte
	fi   FileInfo
}

// Close is a noop-closer.
func (f *File) Close() error {
	return nil
}

// Readdir implements http.File.Readdir
func (f *File) Readdir(count int) ([]os.FileInfo, error) {
	return nil, os.ErrNotExist
}

// Stat implements http.Stat.Readdir
func (f *File) Stat() (os.FileInfo, error) {
	return &f.fi, nil
}

// FileInfo implements the os.FileInfo interface.
type FileInfo struct {
	name    string
	size    int64
	mode    os.FileMode
	modTime time.Time
	isDir   bool
	sys     interface{}

	files []os.FileInfo
}

// Name implements os.FileInfo.Name
func (f *FileInfo) Name() string {
	return f.name
}

// Size implements os.FileInfo.Size
func (f *FileInfo) Size() int64 {
	return f.size
}

// Mode implements os.FileInfo.Mode
func (f *FileInfo) Mode() os.FileMode {
	return f.mode
}

// ModTime implements os.FileInfo.ModTime
func (f *FileInfo) ModTime() time.Time {
	return f.modTime
}

// IsDir implements os.FileInfo.IsDir
func (f *FileInfo) IsDir() bool {
	return f.isDir
}

// Readdir implements os.FileInfo.Readdir
func (f *FileInfo) Readdir(count int) ([]os.FileInfo, error) {
	return f.files, nil
}

// Sys returns the underlying value.
func (f *FileInfo) Sys() interface{} {
	return f.sys
}

func init() {
  {{ .Var }} = &FileSystem{
		files: map[string]File{
		  {{range $path, $file := .Files }} "/{{ $path }}": {{ template "file" $file }}, {{ end }}
		},
	  }
}
`))
}
<commit_msg>gofmt resources.go<commit_after>// Package resources provides unfancy resources embedding with Go.
package resources

import (
	"bytes"
	"fmt"
	"io"
	"log"
	"os"
	"text/template"
)

// File mimicks the os.File and http.File interface.
type File interface {
	io.Reader
	Stat() (os.FileInfo, error)
}

// New creates a new Package.
func New() *Package {
	return &Package{
		Config: Config{
			Pkg:     "resources",
			Var:     "FS",
			Declare: true,
		},
		Files: make(map[string]File),
	}
}

// Config defines some details about the output Go file.
type Config struct {
	Pkg     string // Package name
	Var     string // Variable name to assign the file system to.
	Tag     string // Build tag, leave empty for no tag.
	Declare bool   // Dictates whatever there should be a defintion Variable
}

// Package describes...
type Package struct {
	Config
	Files map[string]File
}

//Add a file to the package at the give path.
func (p *Package) Add(path string, file File) error {
	p.Files[path] = file
	return nil
}

// AddFile is a helper function that adds the files from the path into the package under the path file.
func (p *Package) AddFile(path string, file string) error {
	f, err := os.Open(file)
	if err != nil {
		return err
	}
	return p.Add(path, f)
}

// Build the package
func (p *Package) Build(out io.Writer) error {
	return pkg.Execute(out, p)
}

// Write the build to a file, you don't need to call Build.
func (p *Package) Write(path string) error {
	f, err := os.Create(path)
	if err != nil {
		return err
	}
	defer func() {
		err := f.Close()
		if err != nil {
			log.Panicf("Failed to close file: %s", err)
		}
	}()

	return p.Build(f)
}

// Template
var pkg *template.Template

func reader(input io.Reader) (string, error) {

	var (
		buff       bytes.Buffer
		err        error
		blockwidth = 12
		curblock   = 0
	)

	b := make([]byte, blockwidth)

	for n, err := input.Read(b); err == nil; n, err = input.Read(b) {
		for i := 0; i < n; i++ {
			_, err = fmt.Fprintf(&buff, "0x%02x,", b[i])
			if err != nil {
				break
			}
			curblock++
			if curblock < blockwidth {
				continue
			}
			buff.WriteByte('\n')
			buff.Write([]byte{'\t', '\t'})
			curblock = 0
		}
	}

	return buff.String(), err
}

func init() {

	pkg = template.Must(template.New("file").Funcs(template.FuncMap{"reader": reader}).Parse(` File{
	  data: []byte{
	{{ reader . }}
  },
  fi: FileInfo {
	name:    "{{ .Stat.Name }}",
    size:    {{ .Stat.Size }},
	modTime: time.Unix(0, {{ .Stat.ModTime.UnixNano }}),
    isDir:   {{ .Stat.IsDir }},
  },
}`))

	pkg = template.Must(pkg.New("pkg").Parse(`{{ if .Tag }}// +build {{ .Tag }}

{{ end }}// Package {{.Pkg }} is generated by github.com/omeid/go-resources
package {{ .Pkg }}

import (
	"bytes"
	"errors"
	"net/http"
	"os"
	"path/filepath"
	"strings"
	"time"
)

// FileSystem is an http.FileSystem implementation.
type FileSystem struct {
	files map[string]File
}

// String returns the content of the file as string.
func (fs *FileSystem) String(name string) (string, bool) {
	if filepath.Separator != '/' && strings.IndexRune(name, filepath.Separator) >= 0 ||
		strings.Contains(name, "\x00") {
		return "", false
	}

	file, ok := fs.files[name]

	if !ok {
		return "", false
	}

	return string(file.data), true
}

// Open implements http.FileSystem.Open
func (fs *FileSystem) Open(name string) (http.File, error) {
	if filepath.Separator != '/' && strings.IndexRune(name, filepath.Separator) >= 0 ||
		strings.Contains(name, "\x00") {
		return nil, errors.New("http: invalid character in file path")
	}
	file, ok := fs.files[name]
	if !ok {
		files := []os.FileInfo{}
		for path, file := range fs.files {
			if strings.HasPrefix(path, name) {
				fi := file.fi
				files = append(files, &fi)
			}
		}

		if len(files) == 0 {
			return nil, os.ErrNotExist
		}

		//We have a directory.
		return &File{
			fi: FileInfo{
				isDir: true,
				files: files,
			}}, nil
	}
	file.Reader = bytes.NewReader(file.data)
	return &file, nil
}

// File implements http.File
type File struct {
	*bytes.Reader
	data []byte
	fi   FileInfo
}

// Close is a noop-closer.
func (f *File) Close() error {
	return nil
}

// Readdir implements http.File.Readdir
func (f *File) Readdir(count int) ([]os.FileInfo, error) {
	return nil, os.ErrNotExist
}

// Stat implements http.Stat.Readdir
func (f *File) Stat() (os.FileInfo, error) {
	return &f.fi, nil
}

// FileInfo implements the os.FileInfo interface.
type FileInfo struct {
	name    string
	size    int64
	mode    os.FileMode
	modTime time.Time
	isDir   bool
	sys     interface{}

	files []os.FileInfo
}

// Name implements os.FileInfo.Name
func (f *FileInfo) Name() string {
	return f.name
}

// Size implements os.FileInfo.Size
func (f *FileInfo) Size() int64 {
	return f.size
}

// Mode implements os.FileInfo.Mode
func (f *FileInfo) Mode() os.FileMode {
	return f.mode
}

// ModTime implements os.FileInfo.ModTime
func (f *FileInfo) ModTime() time.Time {
	return f.modTime
}

// IsDir implements os.FileInfo.IsDir
func (f *FileInfo) IsDir() bool {
	return f.isDir
}

// Readdir implements os.FileInfo.Readdir
func (f *FileInfo) Readdir(count int) ([]os.FileInfo, error) {
	return f.files, nil
}

// Sys returns the underlying value.
func (f *FileInfo) Sys() interface{} {
	return f.sys
}

func init() {
  {{ .Var }} = &FileSystem{
		files: map[string]File{
		  {{range $path, $file := .Files }} "/{{ $path }}": {{ template "file" $file }}, {{ end }}
		},
	  }
}
`))
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"net/url"
	"os"
	"path/filepath"
	"time"

	"github.com/mohae/autofact/conf"
	"github.com/uber-go/zap"
)

const (
	addressVar = "address"
	aVar       = "a"
	portVar    = "port"
	pVar       = "p"
)

var (
	connFile = "autofact.json"
	// This is the default directory for autofact-client app data.
	autofactPath    = "$HOME/.autofact"
	autofactEnvName = "AUTOFACT_PATH"
	// default
	connConf   conf.Conn
	serverless bool // if the client is being run without a server
)

// TODO determine loglevel mapping to actual usage:
// Proposed:
//  DebugLevel == not used
//	InfoLevel == Gathered data
//  WarnLevel == Connection info an non-error messages: status type
//  ErrorLevel == Errors
//  PanicLevel == Panic: shouldn't be used
//  FatalLevel == Unrecoverable error that results in app shutdown
// TODO: implement data logging
var (
	log      zap.Logger
	loglevel = zap.LevelFlag("loglevel", zap.WarnLevel, "log level")
	logfile  string

	// These are globals so that they can be closed
	f        *os.File
	isStdErr bool
)

// TODO: reconcile these flags with config file usage.  Probably add contour
// to handle this after the next refactor of contour.
// TODO: make connectInterval/period handling consistent, e.g. should they be
// flags, what is precedence in relation to Conn?
func init() {
	flag.StringVar(&connConf.ServerAddress, addressVar, "127.0.0.1", "the server address")
	flag.StringVar(&connConf.ServerAddress, aVar, "127.0.0.1", "the server address (short)")
	flag.StringVar(&connConf.ServerPort, portVar, "8675", "the connection port")
	flag.StringVar(&connConf.ServerPort, pVar, "8675", "the connection port (short)")
	flag.StringVar(&logfile, "logfile", "autofact.log", "application log file; if empty stderr will be used")
	flag.StringVar(&logfile, "l", "autofact.log", "application log file; if empty stderr will be used")
	flag.BoolVar(&serverless, "serverless", false, "serverless: the client will run standalone and write the collected data to the log")
	connConf.ConnectInterval.Duration = 5 * time.Second
	connConf.ConnectPeriod.Duration = 15 * time.Minute
}

func main() {
	// Load the AUTOPATH value
	tmp := os.Getenv(autofactEnvName)
	if tmp != "" {
		autofactPath = tmp
	}
	autofactPath = os.ExpandEnv(autofactPath)

	// make sure the autofact path exists (create if it doesn't)
	err := os.MkdirAll(autofactPath, 0760)
	if err != nil {
		fmt.Fprintf(os.Stderr, "unable to create AUTOFACT_PATH: %s\n", err)
		fmt.Fprintln(os.Stderr, "startup error: exiting")
		os.Exit(1)
	}

	// finalize the paths
	connFile = filepath.Join(autofactPath, connFile)

	// process the settings
	var connMsg string
	err = connConf.Load(connFile)
	if err != nil {
		// capture the error for logging once it is setup and continue.  An error
		// is not a show stopper as the file may not exist if this is the first
		// time autofact has run on this node.
		connMsg = fmt.Sprintf("using default settings")
	}

	// Parse the flags.
	flag.Parse()

	// now that everything is parsed; set up logging
	SetLogging()
	defer CloseLog()
	// if there was an error reading the connection configuration and this isn't
	// being run serverless, log it
	if connMsg != "" && !serverless {
		log.Warn(
			err.Error(),
			zap.String("conf", connMsg),
		)
	}

	// TODO add env var support

	// get a client
	c := NewClient(connConf)
	c.AutoPath = autofactPath

	// doneCh is used to signal that the connection has been closed
	doneCh := make(chan struct{})

	if !serverless {
		// connect to the Server
		c.ServerURL = url.URL{Scheme: "ws", Host: fmt.Sprintf("%s:%s", c.ServerAddress, c.ServerPort), Path: "/client"}

		// must have a connection before doing anything
		for i := 0; i < 3; i++ {
			connected := c.Connect()
			if connected {
				break
			}
			// retry on fail until retry attempts have been exceeded
		}
		if !c.IsConnected() {
			CloseLog() // defer doesn't run on fatal
			log.Fatal(
				"unable to connect",
				zap.String("server", c.ServerURL.String()),
			)
		}
	}

	// start the go routines first
	go c.Listen(doneCh)
	go c.MemInfo(doneCh)
	go c.CPUUtilization(doneCh)
	go c.NetUsage(doneCh)
	// start the connection handler
	go c.MessageWriter(doneCh)

	if !serverless {
		// if connected, save the conf: this will also save the ClientID
		err = c.Conn.Save()
		if err != nil {
			log.Error(
				err.Error(),
				zap.String("op", "save conn"),
				zap.String("file", c.Filename),
			)
		}
	}
	<-doneCh
}

func SetLogging() {
	// if logfile is empty, use Stderr
	var err error
	if logfile == "" {
		f = os.Stderr
		isStdErr = true
	} else {
		f, err = os.OpenFile(logfile, os.O_CREATE|os.O_RDWR|os.O_APPEND, 0664)
		if err != nil {
			panic(err)
		}
	}
	log = zap.New(
		zap.NewJSONEncoder(
			zap.RFC3339Formatter("ts"),
		),
		zap.Output(f),
	)
	log.SetLevel(*loglevel)
}

// CloseLog closes the log file before exiting.
func CloseLog() {
	if !isStdErr && f != nil {
		f.Close()
	}
}
<commit_msg>add interrupt signal handling; send a close connection requrest to the server when applicable<commit_after>package main

import (
	"flag"
	"fmt"
	"net/url"
	"os"
	"os/signal"
	"path/filepath"
	"time"

	"github.com/gorilla/websocket"
	"github.com/mohae/autofact/conf"
	"github.com/uber-go/zap"
)

const (
	addressVar = "address"
	aVar       = "a"
	portVar    = "port"
	pVar       = "p"
)

var (
	connFile = "autofact.json"
	// This is the default directory for autofact-client app data.
	autofactPath    = "$HOME/.autofact"
	autofactEnvName = "AUTOFACT_PATH"
	// default
	connConf   conf.Conn
	serverless bool // if the client is being run without a server
)

// TODO determine loglevel mapping to actual usage:
// Proposed:
//  DebugLevel == not used
//	InfoLevel == Gathered data
//  WarnLevel == Connection info an non-error messages: status type
//  ErrorLevel == Errors
//  PanicLevel == Panic: shouldn't be used
//  FatalLevel == Unrecoverable error that results in app shutdown
// TODO: implement data logging
var (
	log      zap.Logger
	loglevel = zap.LevelFlag("loglevel", zap.WarnLevel, "log level")
	logfile  string

	// These are globals so that they can be closed
	f        *os.File
	isStdErr bool
)

// TODO: reconcile these flags with config file usage.  Probably add contour
// to handle this after the next refactor of contour.
// TODO: make connectInterval/period handling consistent, e.g. should they be
// flags, what is precedence in relation to Conn?
func init() {
	flag.StringVar(&connConf.ServerAddress, addressVar, "127.0.0.1", "the server address")
	flag.StringVar(&connConf.ServerAddress, aVar, "127.0.0.1", "the server address (short)")
	flag.StringVar(&connConf.ServerPort, portVar, "8675", "the connection port")
	flag.StringVar(&connConf.ServerPort, pVar, "8675", "the connection port (short)")
	flag.StringVar(&logfile, "logfile", "autofact.log", "application log file; if empty stderr will be used")
	flag.StringVar(&logfile, "l", "autofact.log", "application log file; if empty stderr will be used")
	flag.BoolVar(&serverless, "serverless", false, "serverless: the client will run standalone and write the collected data to the log")
	connConf.ConnectInterval.Duration = 5 * time.Second
	connConf.ConnectPeriod.Duration = 15 * time.Minute
}

func main() {
	// Load the AUTOPATH value
	tmp := os.Getenv(autofactEnvName)
	if tmp != "" {
		autofactPath = tmp
	}
	autofactPath = os.ExpandEnv(autofactPath)

	// make sure the autofact path exists (create if it doesn't)
	err := os.MkdirAll(autofactPath, 0760)
	if err != nil {
		fmt.Fprintf(os.Stderr, "unable to create AUTOFACT_PATH: %s\n", err)
		fmt.Fprintln(os.Stderr, "startup error: exiting")
		os.Exit(1)
	}

	// finalize the paths
	connFile = filepath.Join(autofactPath, connFile)

	// process the settings
	var connMsg string
	err = connConf.Load(connFile)
	if err != nil {
		// capture the error for logging once it is setup and continue.  An error
		// is not a show stopper as the file may not exist if this is the first
		// time autofact has run on this node.
		connMsg = fmt.Sprintf("using default settings")
	}

	// Parse the flags.
	flag.Parse()

	// now that everything is parsed; set up logging
	SetLogging()
	defer CloseLog()
	// if there was an error reading the connection configuration and this isn't
	// being run serverless, log it
	if connMsg != "" && !serverless {
		log.Warn(
			err.Error(),
			zap.String("conf", connMsg),
		)
	}

	// TODO add env var support

	// get a client
	c := NewClient(connConf)
	c.AutoPath = autofactPath

	// handle signals
	go handleSignals(c)

	// doneCh is used to signal that the connection has been closed
	doneCh := make(chan struct{})

	if !serverless {
		// connect to the Server
		c.ServerURL = url.URL{Scheme: "ws", Host: fmt.Sprintf("%s:%s", c.ServerAddress, c.ServerPort), Path: "/client"}

		// must have a connection before doing anything
		for i := 0; i < 3; i++ {
			connected := c.Connect()
			if connected {
				break
			}
			// retry on fail until retry attempts have been exceeded
		}
		if !c.IsConnected() {
			CloseLog() // defer doesn't run on fatal
			log.Fatal(
				"unable to connect",
				zap.String("server", c.ServerURL.String()),
			)
		}
	}

	// start the go routines for socket communications
	go c.Listen(doneCh)
	go c.MemInfo(doneCh)
	go c.CPUUtilization(doneCh)
	go c.NetUsage(doneCh)
	// start the connection handler
	go c.MessageWriter(doneCh)

	if !serverless {
		// if connected, save the conf: this will also save the ClientID
		err = c.Conn.Save()
		if err != nil {
			log.Error(
				err.Error(),
				zap.String("op", "save conn"),
				zap.String("file", c.Filename),
			)
		}
	}
	<-doneCh
}

func SetLogging() {
	// if logfile is empty, use Stderr
	var err error
	if logfile == "" {
		f = os.Stderr
		isStdErr = true
	} else {
		f, err = os.OpenFile(logfile, os.O_CREATE|os.O_RDWR|os.O_APPEND, 0664)
		if err != nil {
			panic(err)
		}
	}
	log = zap.New(
		zap.NewJSONEncoder(
			zap.RFC3339Formatter("ts"),
		),
		zap.Output(f),
	)
	log.SetLevel(*loglevel)
}

// CloseLog closes the log file before exiting.
func CloseLog() {
	if !isStdErr && f != nil {
		f.Close()
	}
}

func handleSignals(c *Client) {
	ch := make(chan os.Signal, 1)
	signal.Notify(ch, os.Interrupt)
	v := <-ch
	log.Info(
		"os signal received: shutting down autofact",
		zap.Object("signal", v.String()),
	)
	// If there's a connection send a close signal
	if c.IsConnected() {
		log.Debug(
			"closing connection",
			zap.String("op", "shutdown"),
		)
		c.WS.WriteMessage(websocket.CloseMessage, []byte(string(c.Conn.ID)+" shutting down"))
	}
	CloseLog()
	os.Exit(1)
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"fmt"
	"go/format"
	"go/token"
	"io"
	"io/ioutil"
	"strings"
	"text/template"
	"unicode"

	"github.com/cilium/ebpf"
)

const ebpfModule = "github.com/cilium/ebpf"

const commonRaw = `// Code generated by bpf2go; DO NOT EDIT.
{{- range .Tags }}
// +build {{ . }}
{{- end }}

package {{ .Package }}

import (
	"bytes"
	"fmt"
	"io"

	"{{ .Module }}"
)

// {{ .Name.Load }} returns the embedded CollectionSpec for {{ .Name }}.
func {{ .Name.Load }}() (*ebpf.CollectionSpec, error) {
	reader := bytes.NewReader({{ .Name.Bytes }})
	spec, err := ebpf.LoadCollectionSpecFromReader(reader)
	if err != nil {
		return nil, fmt.Errorf("can't load {{ .Name }}: %w", err)
	}

	return spec, err
}

// {{ .Name.LoadObjects }} loads {{ .Name }} and converts it into a struct.
//
// The following types are suitable as obj argument:
//
//     *{{ .Name.Objects }}
//     *{{ .Name.Programs }}
//     *{{ .Name.Maps }}
//
// See ebpf.CollectionSpec.LoadAndAssign documentation for details.
func {{ .Name.LoadObjects }}(obj interface{}, opts *ebpf.CollectionOptions) (error) {
	spec, err := {{ .Name.Load }}()
	if err != nil {
		return err
	}

	return spec.LoadAndAssign(obj, opts)
}

// {{ .Name.Specs }} contains maps and programs before they are loaded into the kernel.
//
// It can be passed ebpf.CollectionSpec.Assign.
type {{ .Name.Specs }} struct {
	{{ .Name.ProgramSpecs }}
	{{ .Name.MapSpecs }}
}

// {{ .Name.Specs }} contains programs before they are loaded into the kernel.
//
// It can be passed ebpf.CollectionSpec.Assign.
type {{ .Name.ProgramSpecs }} struct {
{{- range $name, $id := .Programs }}
	{{ $id }} *ebpf.ProgramSpec {{ tag $name }}
{{- end }}
}

// {{ .Name.MapSpecs }} contains maps before they are loaded into the kernel.
//
// It can be passed ebpf.CollectionSpec.Assign.
type {{ .Name.MapSpecs }} struct {
{{- range $name, $id := .Maps }}
	{{ $id }} *ebpf.MapSpec {{ tag $name }}
{{- end }}
}

// {{ .Name.Objects }} contains all objects after they have been loaded into the kernel.
//
// It can be passed to {{ .Name.LoadObjects }} or ebpf.CollectionSpec.LoadAndAssign.
type {{ .Name.Objects }} struct {
	{{ .Name.Programs }}
	{{ .Name.Maps }}
}

func (o *{{ .Name.Objects }}) Close() error {
	return {{ .Name.CloseHelper }}(
		&o.{{ .Name.Programs }},
		&o.{{ .Name.Maps }},
	)
}

// {{ .Name.Maps }} contains all maps after they have been loaded into the kernel.
//
// It can be passed to {{ .Name.LoadObjects }} or ebpf.CollectionSpec.LoadAndAssign.
type {{ .Name.Maps }} struct {
{{- range $name, $id := .Maps }}
	{{ $id }} *ebpf.Map {{ tag $name }}
{{- end }}
}

func (m *{{ .Name.Maps }}) Close() error {
	return {{ .Name.CloseHelper }}(
{{- range $id := .Maps }}
		m.{{ $id }},
{{- end }}
	)
}

// {{ .Name.Programs }} contains all programs after they have been loaded into the kernel.
//
// It can be passed to {{ .Name.LoadObjects }} or ebpf.CollectionSpec.LoadAndAssign.
type {{ .Name.Programs }} struct {
{{- range $name, $id := .Programs }}
	{{ $id }} *ebpf.Program {{ tag $name }}
{{- end }}
}

func (p *{{ .Name.Programs }}) Close() error {
	return {{ .Name.CloseHelper }}(
{{- range $id := .Programs }}
		p.{{ $id }},
{{- end }}
	)
}

func {{ .Name.CloseHelper }}(closers ...io.Closer) error {
	for _, closer := range closers {
		if err := closer.Close(); err != nil {
			return err
		}
	}
	return nil
}

// Do not access this directly.
var {{ .Name.Bytes }} = []byte("{{ .Bytes }}")

`

var (
	tplFuncs = map[string]interface{}{
		"tag": tag,
	}
	commonTemplate = template.Must(template.New("common").Funcs(tplFuncs).Parse(commonRaw))
)

type templateName string

func (n templateName) maybeExport(str string) string {
	if token.IsExported(string(n)) {
		return toUpperFirst(str)
	}

	return str
}

func (n templateName) Bytes() string {
	return "_" + toUpperFirst(string(n)) + "Bytes"
}

func (n templateName) Specs() string {
	return n.maybeExport(string(n) + "Specs")
}

func (n templateName) ProgramSpecs() string {
	return n.maybeExport(string(n) + "ProgramSpecs")
}

func (n templateName) MapSpecs() string {
	return n.maybeExport(string(n) + "MapSpecs")
}

func (n templateName) Load() string {
	return n.maybeExport("load" + toUpperFirst(string(n)))
}

func (n templateName) LoadObjects() string {
	return n.maybeExport("load" + toUpperFirst(string(n)) + "Objects")
}

func (n templateName) Objects() string {
	return n.maybeExport(string(n) + "Objects")
}

func (n templateName) Maps() string {
	return n.maybeExport(string(n) + "Maps")
}

func (n templateName) Programs() string {
	return n.maybeExport(string(n) + "Programs")
}

func (n templateName) CloseHelper() string {
	return "_" + toUpperFirst(string(n)) + "Close"
}

type writeArgs struct {
	pkg   string
	ident string
	tags  []string
	obj   io.Reader
	out   io.Writer
}

func writeCommon(args writeArgs) error {
	obj, err := ioutil.ReadAll(args.obj)
	if err != nil {
		return fmt.Errorf("read object file contents: %s", err)
	}

	spec, err := ebpf.LoadCollectionSpecFromReader(bytes.NewReader(obj))
	if err != nil {
		return fmt.Errorf("can't load BPF from ELF: %s", err)
	}

	maps := make(map[string]string)
	for name := range spec.Maps {
		if strings.HasPrefix(name, ".") {
			// Skip .rodata, .data, .bss, etc. sections
			continue
		}

		maps[name] = identifier(name)
	}

	programs := make(map[string]string)
	for name := range spec.Programs {
		programs[name] = identifier(name)
	}

	ctx := struct {
		Module   string
		Package  string
		Tags     []string
		Name     templateName
		Maps     map[string]string
		Programs map[string]string
		Bytes    string
	}{
		ebpfModule,
		args.pkg,
		args.tags,
		templateName(args.ident),
		maps,
		programs,
		binaryString(obj),
	}

	var buf bytes.Buffer
	if err := commonTemplate.Execute(&buf, &ctx); err != nil {
		return fmt.Errorf("can't generate types: %s", err)
	}

	return writeFormatted(buf.Bytes(), args.out)
}

func binaryString(buf []byte) string {
	var builder strings.Builder
	for _, b := range buf {
		builder.WriteString(`\x`)
		builder.WriteString(fmt.Sprintf("%02x", b))
	}
	return builder.String()
}

func writeFormatted(src []byte, out io.Writer) error {
	formatted, err := format.Source(src)
	if err != nil {
		return fmt.Errorf("can't format source: %s", err)
	}

	_, err = out.Write(formatted)
	return err
}

func identifier(str string) string {
	prev := rune(-1)
	return strings.Map(func(r rune) rune {
		// See https://golang.org/ref/spec#Identifiers
		switch {
		case unicode.IsLetter(r):
			if prev == -1 {
				r = unicode.ToUpper(r)
			}

		case r == '_':
			switch {
			// The previous rune was deleted, or we are at the
			// beginning of the string.
			case prev == -1:
				fallthrough

			// The previous rune is a lower case letter or a digit.
			case unicode.IsDigit(prev) || (unicode.IsLetter(prev) && unicode.IsLower(prev)):
				// delete the current rune, and force the
				// next character to be uppercased.
				r = -1
			}

		case unicode.IsDigit(r):

		default:
			// Delete the current rune. prev is unchanged.
			return -1
		}

		prev = r
		return r
	}, str)
}

func tag(str string) string {
	return "`ebpf:\"" + str + "\"`"
}
<commit_msg>cmd/bpf2go: make formatting errors more descriptive<commit_after>package main

import (
	"bytes"
	"errors"
	"fmt"
	"go/format"
	"go/scanner"
	"go/token"
	"io"
	"io/ioutil"
	"strings"
	"text/template"
	"unicode"

	"github.com/cilium/ebpf"
)

const ebpfModule = "github.com/cilium/ebpf"

const commonRaw = `// Code generated by bpf2go; DO NOT EDIT.
{{- range .Tags }}
// +build {{ . }}
{{- end }}

package {{ .Package }}

import (
	"bytes"
	"fmt"
	"io"

	"{{ .Module }}"
)

// {{ .Name.Load }} returns the embedded CollectionSpec for {{ .Name }}.
func {{ .Name.Load }}() (*ebpf.CollectionSpec, error) {
	reader := bytes.NewReader({{ .Name.Bytes }})
	spec, err := ebpf.LoadCollectionSpecFromReader(reader)
	if err != nil {
		return nil, fmt.Errorf("can't load {{ .Name }}: %w", err)
	}

	return spec, err
}

// {{ .Name.LoadObjects }} loads {{ .Name }} and converts it into a struct.
//
// The following types are suitable as obj argument:
//
//     *{{ .Name.Objects }}
//     *{{ .Name.Programs }}
//     *{{ .Name.Maps }}
//
// See ebpf.CollectionSpec.LoadAndAssign documentation for details.
func {{ .Name.LoadObjects }}(obj interface{}, opts *ebpf.CollectionOptions) (error) {
	spec, err := {{ .Name.Load }}()
	if err != nil {
		return err
	}

	return spec.LoadAndAssign(obj, opts)
}

// {{ .Name.Specs }} contains maps and programs before they are loaded into the kernel.
//
// It can be passed ebpf.CollectionSpec.Assign.
type {{ .Name.Specs }} struct {
	{{ .Name.ProgramSpecs }}
	{{ .Name.MapSpecs }}
}

// {{ .Name.Specs }} contains programs before they are loaded into the kernel.
//
// It can be passed ebpf.CollectionSpec.Assign.
type {{ .Name.ProgramSpecs }} struct {
{{- range $name, $id := .Programs }}
	{{ $id }} *ebpf.ProgramSpec {{ tag $name }}
{{- end }}
}

// {{ .Name.MapSpecs }} contains maps before they are loaded into the kernel.
//
// It can be passed ebpf.CollectionSpec.Assign.
type {{ .Name.MapSpecs }} struct {
{{- range $name, $id := .Maps }}
	{{ $id }} *ebpf.MapSpec {{ tag $name }}
{{- end }}
}

// {{ .Name.Objects }} contains all objects after they have been loaded into the kernel.
//
// It can be passed to {{ .Name.LoadObjects }} or ebpf.CollectionSpec.LoadAndAssign.
type {{ .Name.Objects }} struct {
	{{ .Name.Programs }}
	{{ .Name.Maps }}
}

func (o *{{ .Name.Objects }}) Close() error {
	return {{ .Name.CloseHelper }}(
		&o.{{ .Name.Programs }},
		&o.{{ .Name.Maps }},
	)
}

// {{ .Name.Maps }} contains all maps after they have been loaded into the kernel.
//
// It can be passed to {{ .Name.LoadObjects }} or ebpf.CollectionSpec.LoadAndAssign.
type {{ .Name.Maps }} struct {
{{- range $name, $id := .Maps }}
	{{ $id }} *ebpf.Map {{ tag $name }}
{{- end }}
}

func (m *{{ .Name.Maps }}) Close() error {
	return {{ .Name.CloseHelper }}(
{{- range $id := .Maps }}
		m.{{ $id }},
{{- end }}
	)
}

// {{ .Name.Programs }} contains all programs after they have been loaded into the kernel.
//
// It can be passed to {{ .Name.LoadObjects }} or ebpf.CollectionSpec.LoadAndAssign.
type {{ .Name.Programs }} struct {
{{- range $name, $id := .Programs }}
	{{ $id }} *ebpf.Program {{ tag $name }}
{{- end }}
}

func (p *{{ .Name.Programs }}) Close() error {
	return {{ .Name.CloseHelper }}(
{{- range $id := .Programs }}
		p.{{ $id }},
{{- end }}
	)
}

func {{ .Name.CloseHelper }}(closers ...io.Closer) error {
	for _, closer := range closers {
		if err := closer.Close(); err != nil {
			return err
		}
	}
	return nil
}

// Do not access this directly.
var {{ .Name.Bytes }} = []byte("{{ .Bytes }}")

`

var (
	tplFuncs = map[string]interface{}{
		"tag": tag,
	}
	commonTemplate = template.Must(template.New("common").Funcs(tplFuncs).Parse(commonRaw))
)

type templateName string

func (n templateName) maybeExport(str string) string {
	if token.IsExported(string(n)) {
		return toUpperFirst(str)
	}

	return str
}

func (n templateName) Bytes() string {
	return "_" + toUpperFirst(string(n)) + "Bytes"
}

func (n templateName) Specs() string {
	return n.maybeExport(string(n) + "Specs")
}

func (n templateName) ProgramSpecs() string {
	return n.maybeExport(string(n) + "ProgramSpecs")
}

func (n templateName) MapSpecs() string {
	return n.maybeExport(string(n) + "MapSpecs")
}

func (n templateName) Load() string {
	return n.maybeExport("load" + toUpperFirst(string(n)))
}

func (n templateName) LoadObjects() string {
	return n.maybeExport("load" + toUpperFirst(string(n)) + "Objects")
}

func (n templateName) Objects() string {
	return n.maybeExport(string(n) + "Objects")
}

func (n templateName) Maps() string {
	return n.maybeExport(string(n) + "Maps")
}

func (n templateName) Programs() string {
	return n.maybeExport(string(n) + "Programs")
}

func (n templateName) CloseHelper() string {
	return "_" + toUpperFirst(string(n)) + "Close"
}

type writeArgs struct {
	pkg   string
	ident string
	tags  []string
	obj   io.Reader
	out   io.Writer
}

func writeCommon(args writeArgs) error {
	obj, err := ioutil.ReadAll(args.obj)
	if err != nil {
		return fmt.Errorf("read object file contents: %s", err)
	}

	spec, err := ebpf.LoadCollectionSpecFromReader(bytes.NewReader(obj))
	if err != nil {
		return fmt.Errorf("can't load BPF from ELF: %s", err)
	}

	maps := make(map[string]string)
	for name := range spec.Maps {
		if strings.HasPrefix(name, ".") {
			// Skip .rodata, .data, .bss, etc. sections
			continue
		}

		maps[name] = identifier(name)
	}

	programs := make(map[string]string)
	for name := range spec.Programs {
		programs[name] = identifier(name)
	}

	ctx := struct {
		Module   string
		Package  string
		Tags     []string
		Name     templateName
		Maps     map[string]string
		Programs map[string]string
		Bytes    string
	}{
		ebpfModule,
		args.pkg,
		args.tags,
		templateName(args.ident),
		maps,
		programs,
		binaryString(obj),
	}

	var buf bytes.Buffer
	if err := commonTemplate.Execute(&buf, &ctx); err != nil {
		return fmt.Errorf("can't generate types: %s", err)
	}

	return writeFormatted(buf.Bytes(), args.out)
}

func binaryString(buf []byte) string {
	var builder strings.Builder
	for _, b := range buf {
		builder.WriteString(`\x`)
		builder.WriteString(fmt.Sprintf("%02x", b))
	}
	return builder.String()
}

func writeFormatted(src []byte, out io.Writer) error {
	formatted, err := format.Source(src)
	if err == nil {
		_, err = out.Write(formatted)
		return err
	}

	var el scanner.ErrorList
	if !errors.As(err, &el) {
		return err
	}

	var nel scanner.ErrorList
	for _, err := range el {
		if !err.Pos.IsValid() {
			nel = append(nel, err)
			continue
		}

		buf := src[err.Pos.Offset:]
		nl := bytes.IndexRune(buf, '\n')
		if nl == -1 {
			nel = append(nel, err)
			continue
		}

		err.Msg += ": " + string(buf[:nl])
		nel = append(nel, err)
	}

	return nel
}

func identifier(str string) string {
	prev := rune(-1)
	return strings.Map(func(r rune) rune {
		// See https://golang.org/ref/spec#Identifiers
		switch {
		case unicode.IsLetter(r):
			if prev == -1 {
				r = unicode.ToUpper(r)
			}

		case r == '_':
			switch {
			// The previous rune was deleted, or we are at the
			// beginning of the string.
			case prev == -1:
				fallthrough

			// The previous rune is a lower case letter or a digit.
			case unicode.IsDigit(prev) || (unicode.IsLetter(prev) && unicode.IsLower(prev)):
				// delete the current rune, and force the
				// next character to be uppercased.
				r = -1
			}

		case unicode.IsDigit(r):

		default:
			// Delete the current rune. prev is unchanged.
			return -1
		}

		prev = r
		return r
	}, str)
}

func tag(str string) string {
	return "`ebpf:\"" + str + "\"`"
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"log"
	"time"

	"github.com/notnil/chess"
	"github.com/srom/chessbot/play"
	"path/filepath"
)

const DEPTH = 3

func main() {
	modelPath, err := filepath.Abs("../../model/chessbot.pb")
	if err != nil {
		log.Fatalf("Error reading path: %v", err)
	}
	model, err := play.LoadModel(modelPath)
	if err != nil {
		log.Fatalf("Error parsing model: %v", err)
	}
	defer model.Close()

	game := chess.NewGame()

	player := int64(1)
	moveUnit := 0
	for game.Outcome() == chess.NoOutcome {
		moveUnit += 1
		start := time.Now()
		//player = -1 * player
		moveNodes, _ := play.Negamax(model, game, DEPTH, -1 * play.MAX_SCORE, play.MAX_SCORE, player)
		randomMoveNode := play.PickRandomMove(moveNodes)
		game.Move(randomMoveNode.Move)

		moves := []string{}
		for _, moveNode := range moveNodes {
			moves = append(moves, moveNode.Move.String())
		}

		log.Printf("Moves: %d %v %v", len(moveNodes), moves, randomMoveNode.Move.String())
		if moveUnit % 2 == 0 {
			log.Printf("Round %d", moveUnit / 2)
			fmt.Print(game.Position().Board().Draw())
			log.Printf("Elapsed: %v", time.Since(start))
			fmt.Print("\n")
		}
	}

	log.Printf("Outcome: %v", game.Outcome())
	fmt.Print(game.Position().Board().Draw())
	fmt.Print("\n")
}
<commit_msg>Revert player inversion<commit_after>package main

import (
	"fmt"
	"log"
	"time"

	"github.com/notnil/chess"
	"github.com/srom/chessbot/play"
	"path/filepath"
)

const DEPTH = 3

func main() {
	modelPath, err := filepath.Abs("../../model/chessbot.pb")
	if err != nil {
		log.Fatalf("Error reading path: %v", err)
	}
	model, err := play.LoadModel(modelPath)
	if err != nil {
		log.Fatalf("Error parsing model: %v", err)
	}
	defer model.Close()

	game := chess.NewGame()

	player := int64(-11)
	moveUnit := 0
	for game.Outcome() == chess.NoOutcome {
		moveUnit += 1
		start := time.Now()
		player = -1 * player
		moveNodes, _ := play.Negamax(model, game, DEPTH, -1 * play.MAX_SCORE, play.MAX_SCORE, player)
		randomMoveNode := play.PickRandomMove(moveNodes)
		game.Move(randomMoveNode.Move)

		moves := []string{}
		for _, moveNode := range moveNodes {
			moves = append(moves, moveNode.Move.String())
		}

		log.Printf("Moves: %d %v %v", len(moveNodes), moves, randomMoveNode.Move.String())
		if moveUnit % 2 == 0 {
			log.Printf("Round %d", moveUnit / 2)
			fmt.Print(game.Position().Board().Draw())
			log.Printf("Elapsed: %v", time.Since(start))
			fmt.Print("\n")
		}
	}

	log.Printf("Outcome: %v", game.Outcome())
	fmt.Print(game.Position().Board().Draw())
	fmt.Print("\n")
}
<|endoftext|>


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

// The frontend runs a service to serve user-facing traffic.
package main

import (
	"context"
	"flag"
	"net/http"
	"os"
	"path/filepath"
	"time"

	"cloud.google.com/go/profiler"
	"contrib.go.opencensus.io/integrations/ocsql"
	"github.com/go-redis/redis/v7"
	"github.com/google/safehtml/template"
	"golang.org/x/pkgsite/cmd/internal/cmdconfig"
	"golang.org/x/pkgsite/internal"
	"golang.org/x/pkgsite/internal/config"
	"golang.org/x/pkgsite/internal/database"
	"golang.org/x/pkgsite/internal/dcensus"
	"golang.org/x/pkgsite/internal/derrors"
	"golang.org/x/pkgsite/internal/frontend"
	"golang.org/x/pkgsite/internal/localdatasource"
	"golang.org/x/pkgsite/internal/log"
	"golang.org/x/pkgsite/internal/middleware"
	"golang.org/x/pkgsite/internal/postgres"
	"golang.org/x/pkgsite/internal/proxy"
	"golang.org/x/pkgsite/internal/proxydatasource"
	"golang.org/x/pkgsite/internal/queue"
	"golang.org/x/pkgsite/internal/source"
)

var (
	queueName      = config.GetEnv("GO_DISCOVERY_FRONTEND_TASK_QUEUE", "")
	workers        = flag.Int("workers", 10, "number of concurrent requests to the fetch service, when running locally")
	_              = flag.String("static", "content/static", "path to folder containing static files served")
	thirdPartyPath = flag.String("third_party", "third_party", "path to folder containing third-party libraries")
	devMode        = flag.Bool("dev", false, "enable developer mode (reload templates on each page load, serve non-minified JS/CSS, etc.)")
	disableCSP     = flag.Bool("nocsp", false, "disable Content Security Policy")
	proxyURL       = flag.String("proxy_url", "https://proxy.golang.org", "Uses the module proxy referred to by this URL "+
		"for direct proxy mode and frontend fetches")
	directProxy = flag.Bool("direct_proxy", false, "if set to true, uses the module proxy referred to by this URL "+
		"as a direct backend, bypassing the database")
	localPaths         = flag.String("local", "", "run locally, accepts a GOPATH-like collection of local paths for modules to load to memory")
	gopathMode         = flag.Bool("gopath_mode", false, "assume that local modules' paths are relative to GOPATH/src, used only with -local")
	bypassLicenseCheck = flag.Bool("bypass_license_check", false, "display all information, even for non-redistributable paths")
)

func main() {
	flag.Parse()
	ctx := context.Background()
	cfg, err := config.Init(ctx)
	if err != nil {
		log.Fatal(ctx, err)
	}
	cfg.Dump(os.Stderr)
	if cfg.UseProfiler {
		if err := profiler.Start(profiler.Config{}); err != nil {
			log.Fatalf(ctx, "profiler.Start: %v", err)
		}
	}

	log.SetLevel(cfg.LogLevel)

	var (
		dsg        func(context.Context) internal.DataSource
		fetchQueue queue.Queue
	)
	if *bypassLicenseCheck {
		log.Info(ctx, "BYPASSING LICENSE CHECKING: DISPLAYING NON-REDISTRIBUTABLE INFORMATION")
	}

	log.Infof(ctx, "cmd/frontend: initializing cmdconfig.ExperimentGetter")
	expg := cmdconfig.ExperimentGetter(ctx, cfg)
	log.Infof(ctx, "cmd/frontend: initialized cmdconfig.ExperimentGetter")

	if *localPaths != "" {
		lds := localdatasource.New()
		dsg = func(context.Context) internal.DataSource { return lds }
	} else {
		proxyClient, err := proxy.New(*proxyURL)
		if err != nil {
			log.Fatal(ctx, err)
		}

		if *directProxy {
			var pds *proxydatasource.DataSource
			if *bypassLicenseCheck {
				pds = proxydatasource.NewBypassingLicenseCheck(proxyClient)
			} else {
				pds = proxydatasource.New(proxyClient)
			}
			dsg = func(context.Context) internal.DataSource { return pds }
		} else {
			// Wrap the postgres driver with OpenCensus instrumentation.
			ocDriver, err := ocsql.Register("postgres", ocsql.WithAllTraceOptions())
			if err != nil {
				log.Fatalf(ctx, "unable to register the ocsql driver: %v\n", err)
			}

			log.Infof(ctx, "cmd/frontend: openDB start")
			ddb, err := openDB(ctx, cfg, ocDriver)
			if err != nil {
				log.Fatal(ctx, err)
			}
			log.Infof(ctx, "cmd/frontend: openDB finished")

			var db *postgres.DB
			if *bypassLicenseCheck {
				db = postgres.NewBypassingLicenseCheck(ddb)
			} else {
				db = postgres.New(ddb)
			}
			defer db.Close()
			dsg = func(context.Context) internal.DataSource { return db }
			sourceClient := source.NewClient(config.SourceTimeout)
			// The closure passed to queue.New is only used for testing and local
			// execution, not in production. So it's okay that it doesn't use a
			// per-request connection.
			fetchQueue, err = queue.New(ctx, cfg, queueName, *workers, expg,
				func(ctx context.Context, modulePath, version string) (int, error) {
					return frontend.FetchAndUpdateState(ctx, modulePath, version, proxyClient, sourceClient, db)
				})
			if err != nil {
				log.Fatalf(ctx, "queue.New: %v", err)
			}
		}
	}

	var haClient *redis.Client
	if cfg.RedisHAHost != "" {
		haClient = redis.NewClient(&redis.Options{
			Addr: cfg.RedisHAHost + ":" + cfg.RedisHAPort,
		})
	}
	server, err := frontend.NewServer(frontend.ServerConfig{
		DataSourceGetter:     dsg,
		Queue:                fetchQueue,
		CompletionClient:     haClient,
		TaskIDChangeInterval: config.TaskIDChangeIntervalFrontend,
		StaticPath:           template.TrustedSourceFromFlag(flag.Lookup("static").Value),
		ThirdPartyPath:       *thirdPartyPath,
		DevMode:              *devMode,
		AppVersionLabel:      cfg.AppVersionLabel(),
		GoogleTagManagerID:   cfg.GoogleTagManagerID,
		ServeStats:           cfg.ServeStats,
	})
	if err != nil {
		log.Fatalf(ctx, "frontend.NewServer: %v", err)
	}

	if *localPaths != "" {
		lds, ok := dsg(ctx).(*localdatasource.DataSource)
		if ok {
			load(ctx, lds, *localPaths)
		}
	}

	router := dcensus.NewRouter(frontend.TagRoute)
	var cacheClient *redis.Client
	if cfg.RedisCacheHost != "" {
		cacheClient = redis.NewClient(&redis.Options{
			Addr: cfg.RedisCacheHost + ":" + cfg.RedisCachePort,
		})
	}
	server.Install(router.Handle, cacheClient, cfg.AuthValues)
	views := append(dcensus.ServerViews,
		postgres.SearchLatencyDistribution,
		postgres.SearchResponseCount,
		frontend.FetchLatencyDistribution,
		frontend.FetchResponseCount,
		frontend.PlaygroundShareRequestCount,
		frontend.VersionTypeCount,
		middleware.CacheResultCount,
		middleware.CacheErrorCount,
		middleware.CacheLatency,
		middleware.QuotaResultCount,
	)
	if err := dcensus.Init(cfg, views...); err != nil {
		log.Fatal(ctx, err)
	}
	// We are not currently forwarding any ports on AppEngine, so serving debug
	// information is broken.
	if !cfg.OnAppEngine() {
		dcensusServer, err := dcensus.NewServer()
		if err != nil {
			log.Fatal(ctx, err)
		}
		go http.ListenAndServe(cfg.DebugAddr("localhost:8081"), dcensusServer)
	}
	panicHandler, err := server.PanicHandler()
	if err != nil {
		log.Fatal(ctx, err)
	}
	rc := cmdconfig.ReportingClient(ctx, cfg)
	log.Infof(ctx, "cmd/frontend: initializing cmdconfig.Experimenter")
	experimenter := cmdconfig.Experimenter(ctx, cfg, expg, rc)
	log.Infof(ctx, "cmd/frontend: initialized cmdconfig.Experimenter")

	ermw := middleware.Identity()
	if rc != nil {
		ermw = middleware.ErrorReporting(rc.Report)
	}
	mw := middleware.Chain(
		middleware.RequestLog(cmdconfig.Logger(ctx, cfg, "frontend-log")),
		middleware.AcceptRequests(http.MethodGet, http.MethodPost), // accept only GETs and POSTs
		middleware.Quota(cfg.Quota),
		middleware.GodocURL(),                                                                 // potentially redirects so should be early in chain
		middleware.SecureHeaders(!*disableCSP),                                                // must come before any caching for nonces to work
		middleware.LatestVersions(server.GetLatestMinorVersion, server.GetLatestMajorVersion), // must come before caching for version badge to work
		middleware.Panic(panicHandler),
		ermw,
		middleware.Timeout(54*time.Second),
		middleware.Experiment(experimenter),
	)
	addr := cfg.HostAddr("localhost:8080")
	log.Infof(ctx, "Listening on addr %s", addr)
	log.Fatal(ctx, http.ListenAndServe(addr, mw(router)))
}

// TODO(https://github.com/golang/go/issues/40097): factor out to reduce
// duplication with cmd/worker/main.go.

// openDB opens a connection to a database with the given driver, using connection info from
// the given config.
// It first tries the main connection info (DBConnInfo), and if that fails, it uses backup
// connection info it if exists (DBSecondaryConnInfo).
func openDB(ctx context.Context, cfg *config.Config, driver string) (_ *database.DB, err error) {
	defer derrors.Wrap(&err, "openDB(ctx, cfg, %q)", driver)
	log.Infof(ctx, "opening database on host %s", cfg.DBHost)
	ddb, err := database.Open(driver, cfg.DBConnInfo(), cfg.InstanceID)
	if err == nil {
		log.Infof(ctx, "connected to primary host: %s", cfg.DBHost)
		return ddb, nil
	}

	ci := cfg.DBSecondaryConnInfo()
	if ci == "" {
		log.Infof(ctx, "no secondary DB host")
		return nil, err
	}
	log.Errorf(ctx, "database.Open for primary host %s failed with %v; trying secondary host %s ",
		cfg.DBHost, err, cfg.DBSecondaryHost)
	db, err := database.Open(driver, ci, cfg.InstanceID)
	if err != nil {
		return nil, err
	}
	log.Infof(ctx, "connected to secondary host %s", cfg.DBSecondaryHost)
	return db, nil
}

// load loads local modules from pathList.
func load(ctx context.Context, ds *localdatasource.DataSource, pathList string) {
	paths := filepath.SplitList(pathList)
	for _, path := range paths {
		var err error
		if *gopathMode {
			err = ds.LoadFromGOPATH(ctx, path)
		} else {
			err = ds.Load(ctx, path)
		}
		if err != nil {
			log.Error(ctx, err)
		}
	}
}
<commit_msg>cmd/frontend: exit if loading fails in local mode<commit_after>// Copyright 2019 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// The frontend runs a service to serve user-facing traffic.
package main

import (
	"context"
	"flag"
	"net/http"
	"os"
	"path/filepath"
	"time"

	"cloud.google.com/go/profiler"
	"contrib.go.opencensus.io/integrations/ocsql"
	"github.com/go-redis/redis/v7"
	"github.com/google/safehtml/template"
	"golang.org/x/pkgsite/cmd/internal/cmdconfig"
	"golang.org/x/pkgsite/internal"
	"golang.org/x/pkgsite/internal/config"
	"golang.org/x/pkgsite/internal/database"
	"golang.org/x/pkgsite/internal/dcensus"
	"golang.org/x/pkgsite/internal/derrors"
	"golang.org/x/pkgsite/internal/frontend"
	"golang.org/x/pkgsite/internal/localdatasource"
	"golang.org/x/pkgsite/internal/log"
	"golang.org/x/pkgsite/internal/middleware"
	"golang.org/x/pkgsite/internal/postgres"
	"golang.org/x/pkgsite/internal/proxy"
	"golang.org/x/pkgsite/internal/proxydatasource"
	"golang.org/x/pkgsite/internal/queue"
	"golang.org/x/pkgsite/internal/source"
)

var (
	queueName      = config.GetEnv("GO_DISCOVERY_FRONTEND_TASK_QUEUE", "")
	workers        = flag.Int("workers", 10, "number of concurrent requests to the fetch service, when running locally")
	_              = flag.String("static", "content/static", "path to folder containing static files served")
	thirdPartyPath = flag.String("third_party", "third_party", "path to folder containing third-party libraries")
	devMode        = flag.Bool("dev", false, "enable developer mode (reload templates on each page load, serve non-minified JS/CSS, etc.)")
	disableCSP     = flag.Bool("nocsp", false, "disable Content Security Policy")
	proxyURL       = flag.String("proxy_url", "https://proxy.golang.org", "Uses the module proxy referred to by this URL "+
		"for direct proxy mode and frontend fetches")
	directProxy = flag.Bool("direct_proxy", false, "if set to true, uses the module proxy referred to by this URL "+
		"as a direct backend, bypassing the database")
	localPaths         = flag.String("local", "", "run locally, accepts a GOPATH-like collection of local paths for modules to load to memory")
	gopathMode         = flag.Bool("gopath_mode", false, "assume that local modules' paths are relative to GOPATH/src, used only with -local")
	bypassLicenseCheck = flag.Bool("bypass_license_check", false, "display all information, even for non-redistributable paths")
)

func main() {
	flag.Parse()
	ctx := context.Background()
	cfg, err := config.Init(ctx)
	if err != nil {
		log.Fatal(ctx, err)
	}
	cfg.Dump(os.Stderr)
	if cfg.UseProfiler {
		if err := profiler.Start(profiler.Config{}); err != nil {
			log.Fatalf(ctx, "profiler.Start: %v", err)
		}
	}

	log.SetLevel(cfg.LogLevel)

	var (
		dsg        func(context.Context) internal.DataSource
		fetchQueue queue.Queue
	)
	if *bypassLicenseCheck {
		log.Info(ctx, "BYPASSING LICENSE CHECKING: DISPLAYING NON-REDISTRIBUTABLE INFORMATION")
	}

	log.Infof(ctx, "cmd/frontend: initializing cmdconfig.ExperimentGetter")
	expg := cmdconfig.ExperimentGetter(ctx, cfg)
	log.Infof(ctx, "cmd/frontend: initialized cmdconfig.ExperimentGetter")

	if *localPaths != "" {
		lds := localdatasource.New()
		dsg = func(context.Context) internal.DataSource { return lds }
	} else {
		proxyClient, err := proxy.New(*proxyURL)
		if err != nil {
			log.Fatal(ctx, err)
		}

		if *directProxy {
			var pds *proxydatasource.DataSource
			if *bypassLicenseCheck {
				pds = proxydatasource.NewBypassingLicenseCheck(proxyClient)
			} else {
				pds = proxydatasource.New(proxyClient)
			}
			dsg = func(context.Context) internal.DataSource { return pds }
		} else {
			// Wrap the postgres driver with OpenCensus instrumentation.
			ocDriver, err := ocsql.Register("postgres", ocsql.WithAllTraceOptions())
			if err != nil {
				log.Fatalf(ctx, "unable to register the ocsql driver: %v\n", err)
			}

			log.Infof(ctx, "cmd/frontend: openDB start")
			ddb, err := openDB(ctx, cfg, ocDriver)
			if err != nil {
				log.Fatal(ctx, err)
			}
			log.Infof(ctx, "cmd/frontend: openDB finished")

			var db *postgres.DB
			if *bypassLicenseCheck {
				db = postgres.NewBypassingLicenseCheck(ddb)
			} else {
				db = postgres.New(ddb)
			}
			defer db.Close()
			dsg = func(context.Context) internal.DataSource { return db }
			sourceClient := source.NewClient(config.SourceTimeout)
			// The closure passed to queue.New is only used for testing and local
			// execution, not in production. So it's okay that it doesn't use a
			// per-request connection.
			fetchQueue, err = queue.New(ctx, cfg, queueName, *workers, expg,
				func(ctx context.Context, modulePath, version string) (int, error) {
					return frontend.FetchAndUpdateState(ctx, modulePath, version, proxyClient, sourceClient, db)
				})
			if err != nil {
				log.Fatalf(ctx, "queue.New: %v", err)
			}
		}
	}

	var haClient *redis.Client
	if cfg.RedisHAHost != "" {
		haClient = redis.NewClient(&redis.Options{
			Addr: cfg.RedisHAHost + ":" + cfg.RedisHAPort,
		})
	}
	server, err := frontend.NewServer(frontend.ServerConfig{
		DataSourceGetter:     dsg,
		Queue:                fetchQueue,
		CompletionClient:     haClient,
		TaskIDChangeInterval: config.TaskIDChangeIntervalFrontend,
		StaticPath:           template.TrustedSourceFromFlag(flag.Lookup("static").Value),
		ThirdPartyPath:       *thirdPartyPath,
		DevMode:              *devMode,
		AppVersionLabel:      cfg.AppVersionLabel(),
		GoogleTagManagerID:   cfg.GoogleTagManagerID,
		ServeStats:           cfg.ServeStats,
	})
	if err != nil {
		log.Fatalf(ctx, "frontend.NewServer: %v", err)
	}

	if *localPaths != "" {
		lds, ok := dsg(ctx).(*localdatasource.DataSource)
		if ok {
			load(ctx, lds, *localPaths)
		}
	}

	router := dcensus.NewRouter(frontend.TagRoute)
	var cacheClient *redis.Client
	if cfg.RedisCacheHost != "" {
		cacheClient = redis.NewClient(&redis.Options{
			Addr: cfg.RedisCacheHost + ":" + cfg.RedisCachePort,
		})
	}
	server.Install(router.Handle, cacheClient, cfg.AuthValues)
	views := append(dcensus.ServerViews,
		postgres.SearchLatencyDistribution,
		postgres.SearchResponseCount,
		frontend.FetchLatencyDistribution,
		frontend.FetchResponseCount,
		frontend.PlaygroundShareRequestCount,
		frontend.VersionTypeCount,
		middleware.CacheResultCount,
		middleware.CacheErrorCount,
		middleware.CacheLatency,
		middleware.QuotaResultCount,
	)
	if err := dcensus.Init(cfg, views...); err != nil {
		log.Fatal(ctx, err)
	}
	// We are not currently forwarding any ports on AppEngine, so serving debug
	// information is broken.
	if !cfg.OnAppEngine() {
		dcensusServer, err := dcensus.NewServer()
		if err != nil {
			log.Fatal(ctx, err)
		}
		go http.ListenAndServe(cfg.DebugAddr("localhost:8081"), dcensusServer)
	}
	panicHandler, err := server.PanicHandler()
	if err != nil {
		log.Fatal(ctx, err)
	}
	rc := cmdconfig.ReportingClient(ctx, cfg)
	log.Infof(ctx, "cmd/frontend: initializing cmdconfig.Experimenter")
	experimenter := cmdconfig.Experimenter(ctx, cfg, expg, rc)
	log.Infof(ctx, "cmd/frontend: initialized cmdconfig.Experimenter")

	ermw := middleware.Identity()
	if rc != nil {
		ermw = middleware.ErrorReporting(rc.Report)
	}
	mw := middleware.Chain(
		middleware.RequestLog(cmdconfig.Logger(ctx, cfg, "frontend-log")),
		middleware.AcceptRequests(http.MethodGet, http.MethodPost), // accept only GETs and POSTs
		middleware.Quota(cfg.Quota),
		middleware.GodocURL(),                                                                 // potentially redirects so should be early in chain
		middleware.SecureHeaders(!*disableCSP),                                                // must come before any caching for nonces to work
		middleware.LatestVersions(server.GetLatestMinorVersion, server.GetLatestMajorVersion), // must come before caching for version badge to work
		middleware.Panic(panicHandler),
		ermw,
		middleware.Timeout(54*time.Second),
		middleware.Experiment(experimenter),
	)
	addr := cfg.HostAddr("localhost:8080")
	log.Infof(ctx, "Listening on addr %s", addr)
	log.Fatal(ctx, http.ListenAndServe(addr, mw(router)))
}

// TODO(https://github.com/golang/go/issues/40097): factor out to reduce
// duplication with cmd/worker/main.go.

// openDB opens a connection to a database with the given driver, using connection info from
// the given config.
// It first tries the main connection info (DBConnInfo), and if that fails, it uses backup
// connection info it if exists (DBSecondaryConnInfo).
func openDB(ctx context.Context, cfg *config.Config, driver string) (_ *database.DB, err error) {
	defer derrors.Wrap(&err, "openDB(ctx, cfg, %q)", driver)
	log.Infof(ctx, "opening database on host %s", cfg.DBHost)
	ddb, err := database.Open(driver, cfg.DBConnInfo(), cfg.InstanceID)
	if err == nil {
		log.Infof(ctx, "connected to primary host: %s", cfg.DBHost)
		return ddb, nil
	}

	ci := cfg.DBSecondaryConnInfo()
	if ci == "" {
		log.Infof(ctx, "no secondary DB host")
		return nil, err
	}
	log.Errorf(ctx, "database.Open for primary host %s failed with %v; trying secondary host %s ",
		cfg.DBHost, err, cfg.DBSecondaryHost)
	db, err := database.Open(driver, ci, cfg.InstanceID)
	if err != nil {
		return nil, err
	}
	log.Infof(ctx, "connected to secondary host %s", cfg.DBSecondaryHost)
	return db, nil
}

// load loads local modules from pathList.
func load(ctx context.Context, ds *localdatasource.DataSource, pathList string) {
	paths := filepath.SplitList(pathList)
	loaded := len(paths)
	for _, path := range paths {
		var err error
		if *gopathMode {
			err = ds.LoadFromGOPATH(ctx, path)
		} else {
			err = ds.Load(ctx, path)
		}
		if err != nil {
			log.Error(ctx, err)
			loaded--
		}
	}

	if loaded == 0 {
		log.Fatalf(ctx, "failed to load module(s) at %s", pathList)
	}
}
<|endoftext|>