<commit_before>// Copyright 2020 Google LLC
//
// 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.

// Implements dart/sdk buildpack.
// The sdk buildpack installs the Dart SDK.
package main

import (
	"os"

	"github.com/GoogleCloudPlatform/buildpacks/pkg/env"
	gcp "github.com/GoogleCloudPlatform/buildpacks/pkg/gcpbuildpack"
	"github.com/GoogleCloudPlatform/buildpacks/pkg/runtime"
	"github.com/buildpacks/libcnb"
)

const (
	dartLayer      = "dart"
	defaultVersion = "2.15.1"
	dartEnabledEnv = "GOOGLE_DART_ENABLED"
)

func main() {
	gcp.Main(detectFn, buildFn)
}

func detectFn(ctx *gcp.Context) (gcp.DetectResult, error) {
	if !isDartEnabled() {
		return gcp.OptOutEnvNotSet(dartEnabledEnv), nil
	}
	if result := runtime.CheckOverride(ctx, "dart"); result != nil {
		return result, nil
	}
	if ctx.FileExists("pubspec.yaml") {
		return gcp.OptInFileFound("pubspec.yaml"), nil
	}
	if len(ctx.Glob("*.dart")) > 0 {
		return gcp.OptIn("found .dart files"), nil
	}

	return gcp.OptOut("neither pubspec.yaml nor any .dart files found"), nil
}

func buildFn(ctx *gcp.Context) error {
	version := defaultVersion
	if envVersion := os.Getenv(env.RuntimeVersion); envVersion != "" {
		ctx.Logf("Using runtime version from %s: %s", env.RuntimeVersion, envVersion)
		version = envVersion
	}

	// The Dart SDK is only required at compile time. It is not included in the run image.
	drl := ctx.Layer(dartLayer, gcp.BuildLayer, gcp.CacheLayer)
	ctx.AddBOMEntry(libcnb.BOMEntry{
		Name:     dartLayer,
		Metadata: map[string]interface{}{"version": version},
		Build:    true,
	})

	if runtime.IsCached(ctx, drl, version) {
		ctx.CacheHit(dartLayer)
		ctx.Logf("Runtime cache hit, skipping installation.")
		return nil
	}
	ctx.CacheMiss(dartLayer)

	return runtime.InstallDartSDK(ctx, drl, version)
}

// isDartEnabled returns true if we should enable the experimental Dart buildpacks.
func isDartEnabled() bool {
	res, err := env.IsPresentAndTrue(dartEnabledEnv)
	return err == nil && res
}
<commit_msg>Bump Dart SDK 2.15.1 -> 2.16.0<commit_after>// Copyright 2020 Google LLC
//
// 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.

// Implements dart/sdk buildpack.
// The sdk buildpack installs the Dart SDK.
package main

import (
	"os"

	"github.com/GoogleCloudPlatform/buildpacks/pkg/env"
	gcp "github.com/GoogleCloudPlatform/buildpacks/pkg/gcpbuildpack"
	"github.com/GoogleCloudPlatform/buildpacks/pkg/runtime"
	"github.com/buildpacks/libcnb"
)

const (
	dartLayer      = "dart"
	defaultVersion = "2.16.0"
	dartEnabledEnv = "GOOGLE_DART_ENABLED"
)

func main() {
	gcp.Main(detectFn, buildFn)
}

func detectFn(ctx *gcp.Context) (gcp.DetectResult, error) {
	if !isDartEnabled() {
		return gcp.OptOutEnvNotSet(dartEnabledEnv), nil
	}
	if result := runtime.CheckOverride(ctx, "dart"); result != nil {
		return result, nil
	}
	if ctx.FileExists("pubspec.yaml") {
		return gcp.OptInFileFound("pubspec.yaml"), nil
	}
	if len(ctx.Glob("*.dart")) > 0 {
		return gcp.OptIn("found .dart files"), nil
	}

	return gcp.OptOut("neither pubspec.yaml nor any .dart files found"), nil
}

func buildFn(ctx *gcp.Context) error {
	version := defaultVersion
	if envVersion := os.Getenv(env.RuntimeVersion); envVersion != "" {
		ctx.Logf("Using runtime version from %s: %s", env.RuntimeVersion, envVersion)
		version = envVersion
	}

	// The Dart SDK is only required at compile time. It is not included in the run image.
	drl := ctx.Layer(dartLayer, gcp.BuildLayer, gcp.CacheLayer)
	ctx.AddBOMEntry(libcnb.BOMEntry{
		Name:     dartLayer,
		Metadata: map[string]interface{}{"version": version},
		Build:    true,
	})

	if runtime.IsCached(ctx, drl, version) {
		ctx.CacheHit(dartLayer)
		ctx.Logf("Runtime cache hit, skipping installation.")
		return nil
	}
	ctx.CacheMiss(dartLayer)

	return runtime.InstallDartSDK(ctx, drl, version)
}

// isDartEnabled returns true if we should enable the experimental Dart buildpacks.
func isDartEnabled() bool {
	res, err := env.IsPresentAndTrue(dartEnabledEnv)
	return err == nil && res
}
<|endoftext|>


<commit_before>package main

import (
	"crypto/rand"
	"crypto/tls"
	"crypto/x509"
	"expvar"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	_ "net/http/pprof"
	"os"
	"os/signal"
	"path/filepath"
	"runtime"
	"runtime/pprof"
	"syscall"
	"time"

	"github.com/ekanite/ekanite"
	"github.com/ekanite/ekanite/input"
	"github.com/ekanite/ekanite/status"
)

var (
	stats = expvar.NewMap("ekanite")
)

// Program parameters
var datadir string
var tcpIface string
var udpIface string
var caPemPath string
var caKeyPath string
var queryIface string
var batchSize int
var batchTimeout int
var indexMaxPending int
var gomaxprocs int
var numShards int
var retentionPeriod string
var cpuProfile string
var memProfile string
var inputFormat string

// Flag set
var fs *flag.FlagSet

// Types
const (
	DefaultDataDir         = "/var/opt/ekanite"
	DefaultBatchSize       = 300
	DefaultBatchTimeout    = 1000
	DefaultIndexMaxPending = 1000
	DefaultNumShards       = 4
	DefaultRetentionPeriod = "168h"
	DefaultQueryAddr       = "localhost:9950"
	DefaultHTTPQueryAddr   = "localhost:8080"
	DefaultDiagsIface      = "localhost:9951"
	DefaultTCPServer       = "localhost:5514"
	DefaultInputFormat     = "syslog"
)

func main() {
	fs = flag.NewFlagSet("", flag.ExitOnError)
	var (
		datadir         = fs.String("datadir", DefaultDataDir, "Set data directory")
		batchSize       = fs.Int("batchsize", DefaultBatchSize, "Indexing batch size")
		batchTimeout    = fs.Int("batchtime", DefaultBatchTimeout, "Indexing batch timeout, in milliseconds")
		indexMaxPending = fs.Int("maxpending", DefaultIndexMaxPending, "Maximum pending index events")
		tcpIface        = fs.String("tcp", DefaultTCPServer, "Syslog server TCP bind address in the form host:port. To disable set to empty string")
		udpIface        = fs.String("udp", "", "Syslog server UDP bind address in the form host:port. If not set, not started")
		diagIface       = fs.String("diag", DefaultDiagsIface, "expvar and pprof bind address in the form host:port. If not set, not started")
		caPemPath       = fs.String("tlspem", "", "path to CA PEM file for TLS-enabled TCP server. If not set, TLS not activated")
		caKeyPath       = fs.String("tlskey", "", "path to CA key file for TLS-enabled TCP server. If not set, TLS not activated")
		queryIface      = fs.String("query", DefaultQueryAddr, "TCP Bind address for query server in the form host:port. To disable set to empty string")
		queryIfaceHttp  = fs.String("queryhttp", DefaultHTTPQueryAddr, "TCP Bind address for http query server in the form host:port. To disable set to empty string")
		numShards       = fs.Int("numshards", DefaultNumShards, "Set number of shards per index")
		retentionPeriod = fs.String("retention", DefaultRetentionPeriod, "Data retention period. Minimum is 24 hours")
		cpuProfile      = fs.String("cpuprof", "", "Where to write CPU profiling data. Not written if not set")
		memProfile      = fs.String("memprof", "", "Where to write memory profiling data. Not written if not set")
		inputFormat     = fs.String("input", DefaultInputFormat, "Message format of input (only syslog supported)")
	)
	fs.Usage = printHelp
	fs.Parse(os.Args[1:])

	absDataDir, err := filepath.Abs(*datadir)
	if err != nil {
		log.Fatalf("failed to get absolute data path for '%s': %s", *datadir, err.Error())
	}

	// Get the retention period.
	retention, err := time.ParseDuration(*retentionPeriod)
	if err != nil {
		log.Fatalf("failed to parse retention period '%s'", *retentionPeriod)
	}

	log.SetFlags(log.LstdFlags)
	log.SetPrefix("[ekanite] ")
	log.Printf("ekanite started using %s for index storage", absDataDir)

	runtime.GOMAXPROCS(runtime.NumCPU())
	log.Println("GOMAXPROCS set to", runtime.GOMAXPROCS(0))

	// Start the expvar handler if requested.
	if *diagIface != "" {
		diagServer := status.NewService(*diagIface)

		if diagServer.Start(); err != nil {
			log.Fatalf("failed to start status server: %s", err.Error())
		}
	}

	// Create and open the Engine.
	engine := ekanite.NewEngine(absDataDir)
	if engine == nil {
		log.Fatalf("failed to create indexing engine at %s", absDataDir)
	}
	engine.NumShards = *numShards
	engine.RetentionPeriod = retention

	if err := engine.Open(); err != nil {
		log.Fatalf("failed to open engine: %s", err.Error())
	}
	log.Printf("engine opened with shard number of %d, retention period of %s",
		engine.NumShards, engine.RetentionPeriod)

	// Start the simple query server if requested.
	if *queryIface != "" {
		server := ekanite.NewServer(*queryIface, engine)
		if server == nil {
			log.Fatal("failed to create query server")
		}
		if err := server.Start(); err != nil {
			log.Fatalf("failed to start query server: %s", err.Error())
		}
		log.Printf("query server listening to %s", *queryIface)
	}

	// Start the http query server if requested.
	if *queryIfaceHttp != "" {
		server := ekanite.NewHTTPServer(*queryIfaceHttp, engine)
		if server == nil {
			log.Fatal("failed to create HTTP query server")
		}
		if err := server.Start(); err != nil {
			log.Fatalf("failed to start HTTP query server: %s", err.Error())
		}
		log.Printf("HTTP query server listening to %s", *queryIfaceHttp)
	}

	// Create and start the batcher.
	batcherTimeout := time.Duration(*batchTimeout) * time.Millisecond
	batcher := ekanite.NewBatcher(engine, *batchSize, batcherTimeout, *indexMaxPending)
	if batcher == nil {
		log.Fatal("failed to create indexing batcher")
	}

	errChan := make(chan error)
	if err := batcher.Start(errChan); err != nil {
		log.Fatalf("failed to start indexing batcher: %s", err.Error())
	}
	log.Printf("batching configured with size %d, timeout %s, max pending %d",
		*batchSize, batcherTimeout, *indexMaxPending)

	// Start draining batcher errors.
	go func() {
		for {
			select {
			case err := <-errChan:
				if err != nil {
					log.Printf("error indexing batch: %s", err.Error())
				}
			}
		}
	}()

	// Start TCP collector if requested.
	if *tcpIface != "" {
		var tlsConfig *tls.Config
		if *caPemPath != "" && *caKeyPath != "" {
			tlsConfig, err = newTLSConfig(*caPemPath, *caKeyPath)
			if err != nil {
				log.Fatalf("failed to configure TLS: %s", err.Error())
			}
			log.Printf("TLS successfully configured")
		}

		collector, err := input.NewCollector("tcp", *tcpIface, *inputFormat, tlsConfig)
		if err != nil {
			log.Fatalf("failed to create TCP collector: %s", err.Error())
		}
		if err := collector.Start(batcher.C()); err != nil {
			log.Fatalf("failed to start TCP collector: %s", err.Error())
		}
		log.Printf("TCP collector listening to %s", *tcpIface)
	}

	// Start UDP collector if requested.
	if *udpIface != "" {
		collector, err := input.NewCollector("udp", *udpIface, *inputFormat, nil)
		if collector == nil {
			log.Fatalf("failed to create UDP collector: %s", err.Error())
		}
		if err := collector.Start(batcher.C()); err != nil {
			log.Fatalf("failed to start UDP collector: %s", err.Error())
		}
		log.Printf("UDP collector listening to %s", *udpIface)
	}

	// Start profiling.
	startProfile(*cpuProfile, *memProfile)

	stats.Set("launch", time.Now().UTC())

	// Set up signal handling.
	signalCh := make(chan os.Signal, 1)
	signal.Notify(signalCh, os.Interrupt, syscall.SIGTERM)

	// Block until one of the signals above is received
	select {
	case <-signalCh:
		log.Println("signal received, shutting down...")
	}

	stopProfile()
}

func newTLSConfig(caPemPath, caKeyPath string) (*tls.Config, error) {
	var config *tls.Config

	caPem, err := ioutil.ReadFile(caPemPath)
	if err != nil {
		return nil, err
	}
	ca, err := x509.ParseCertificate(caPem)
	if err != nil {
		return nil, err
	}

	caKey, err := ioutil.ReadFile(caKeyPath)
	if err != nil {
		return nil, err
	}
	key, err := x509.ParsePKCS1PrivateKey(caKey)
	if err != nil {
		return nil, err
	}
	pool := x509.NewCertPool()
	pool.AddCert(ca)

	cert := tls.Certificate{
		Certificate: [][]byte{caPem},
		PrivateKey:  key,
	}

	config = &tls.Config{
		ClientAuth:   tls.RequireAndVerifyClientCert,
		Certificates: []tls.Certificate{cert},
		ClientCAs:    pool,
	}

	config.Rand = rand.Reader

	return config, nil
}

// prof stores the file locations of active profiles.
var prof struct {
	cpu *os.File
	mem *os.File
}

// StartProfile initializes the cpu and memory profile, if specified.
func startProfile(cpuprofile, memprofile string) {
	if cpuprofile != "" {
		f, err := os.Create(cpuprofile)
		if err != nil {
			log.Fatalf("cpuprofile: %v", err)
		}
		log.Printf("writing CPU profile to: %s\n", cpuprofile)
		prof.cpu = f
		pprof.StartCPUProfile(prof.cpu)
	}

	if memprofile != "" {
		f, err := os.Create(memprofile)
		if err != nil {
			log.Fatalf("memprofile: %v", err)
		}
		log.Printf("writing memory profile to: %s\n", memprofile)
		prof.mem = f
		runtime.MemProfileRate = 4096
	}

}

// StopProfile closes the cpu and memory profiles if they are running.
func stopProfile() {
	if prof.cpu != nil {
		pprof.StopCPUProfile()
		prof.cpu.Close()
		log.Println("CPU profile stopped")
	}
	if prof.mem != nil {
		pprof.Lookup("heap").WriteTo(prof.mem, 0)
		prof.mem.Close()
		log.Println("memory profile stopped")
	}
}

func printHelp() {
	fmt.Println("ekanited [options]")
	fs.PrintDefaults()
}
<commit_msg>Reduce cyclo complexity of main<commit_after>package main

import (
	"crypto/rand"
	"crypto/tls"
	"crypto/x509"
	"expvar"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	_ "net/http/pprof"
	"os"
	"os/signal"
	"path/filepath"
	"runtime"
	"runtime/pprof"
	"syscall"
	"time"

	"github.com/ekanite/ekanite"
	"github.com/ekanite/ekanite/input"
	"github.com/ekanite/ekanite/status"
)

var (
	stats = expvar.NewMap("ekanite")
)

// Program parameters
var datadir string
var tcpIface string
var udpIface string
var caPemPath string
var caKeyPath string
var queryIface string
var batchSize int
var batchTimeout int
var indexMaxPending int
var gomaxprocs int
var numShards int
var retentionPeriod string
var cpuProfile string
var memProfile string
var inputFormat string

// Flag set
var fs *flag.FlagSet

// Types
const (
	DefaultDataDir         = "/var/opt/ekanite"
	DefaultBatchSize       = 300
	DefaultBatchTimeout    = 1000
	DefaultIndexMaxPending = 1000
	DefaultNumShards       = 4
	DefaultRetentionPeriod = "168h"
	DefaultQueryAddr       = "localhost:9950"
	DefaultHTTPQueryAddr   = "localhost:8080"
	DefaultDiagsIface      = "localhost:9951"
	DefaultTCPServer       = "localhost:5514"
	DefaultInputFormat     = "syslog"
)

func main() {
	fs = flag.NewFlagSet("", flag.ExitOnError)
	var (
		datadir         = fs.String("datadir", DefaultDataDir, "Set data directory")
		batchSize       = fs.Int("batchsize", DefaultBatchSize, "Indexing batch size")
		batchTimeout    = fs.Int("batchtime", DefaultBatchTimeout, "Indexing batch timeout, in milliseconds")
		indexMaxPending = fs.Int("maxpending", DefaultIndexMaxPending, "Maximum pending index events")
		tcpIface        = fs.String("tcp", DefaultTCPServer, "Syslog server TCP bind address in the form host:port. To disable set to empty string")
		udpIface        = fs.String("udp", "", "Syslog server UDP bind address in the form host:port. If not set, not started")
		diagIface       = fs.String("diag", DefaultDiagsIface, "expvar and pprof bind address in the form host:port. If not set, not started")
		caPemPath       = fs.String("tlspem", "", "path to CA PEM file for TLS-enabled TCP server. If not set, TLS not activated")
		caKeyPath       = fs.String("tlskey", "", "path to CA key file for TLS-enabled TCP server. If not set, TLS not activated")
		queryIface      = fs.String("query", DefaultQueryAddr, "TCP Bind address for query server in the form host:port. To disable set to empty string")
		queryIfaceHttp  = fs.String("queryhttp", DefaultHTTPQueryAddr, "TCP Bind address for http query server in the form host:port. To disable set to empty string")
		numShards       = fs.Int("numshards", DefaultNumShards, "Set number of shards per index")
		retentionPeriod = fs.String("retention", DefaultRetentionPeriod, "Data retention period. Minimum is 24 hours")
		cpuProfile      = fs.String("cpuprof", "", "Where to write CPU profiling data. Not written if not set")
		memProfile      = fs.String("memprof", "", "Where to write memory profiling data. Not written if not set")
		inputFormat     = fs.String("input", DefaultInputFormat, "Message format of input (only syslog supported)")
	)
	fs.Usage = printHelp
	fs.Parse(os.Args[1:])

	absDataDir, err := filepath.Abs(*datadir)
	if err != nil {
		log.Fatalf("failed to get absolute data path for '%s': %s", *datadir, err.Error())
	}

	// Get the retention period.
	retention, err := time.ParseDuration(*retentionPeriod)
	if err != nil {
		log.Fatalf("failed to parse retention period '%s'", *retentionPeriod)
	}

	log.SetFlags(log.LstdFlags)
	log.SetPrefix("[ekanite] ")
	log.Printf("ekanite started using %s for index storage", absDataDir)

	runtime.GOMAXPROCS(runtime.NumCPU())
	log.Println("GOMAXPROCS set to", runtime.GOMAXPROCS(0))

	// Start the expvar handler if requested.
	if *diagIface != "" {
		startDiagServer(*diagIface)
	}

	// Create and open the Engine.
	engine := ekanite.NewEngine(absDataDir)
	if engine == nil {
		log.Fatalf("failed to create indexing engine at %s", absDataDir)
	}
	engine.NumShards = *numShards
	engine.RetentionPeriod = retention

	if err := engine.Open(); err != nil {
		log.Fatalf("failed to open engine: %s", err.Error())
	}
	log.Printf("engine opened with shard number of %d, retention period of %s",
		engine.NumShards, engine.RetentionPeriod)

	// Start the simple query server if requested.
	if *queryIface != "" {
		startQueryServer(*queryIface, engine)
	}

	// Start the http query server if requested.
	if *queryIfaceHttp != "" {
		startHTTPQueryServer(*queryIfaceHttp, engine)
	}

	// Create and start the batcher.
	batcherTimeout := time.Duration(*batchTimeout) * time.Millisecond
	batcher := ekanite.NewBatcher(engine, *batchSize, batcherTimeout, *indexMaxPending)
	if batcher == nil {
		log.Fatal("failed to create indexing batcher")
	}

	errChan := make(chan error)
	if err := batcher.Start(errChan); err != nil {
		log.Fatalf("failed to start indexing batcher: %s", err.Error())
	}
	log.Printf("batching configured with size %d, timeout %s, max pending %d",
		*batchSize, batcherTimeout, *indexMaxPending)

	// Start draining batcher errors.
	go func() {
		for {
			select {
			case err := <-errChan:
				if err != nil {
					log.Printf("error indexing batch: %s", err.Error())
				}
			}
		}
	}()

	// Start TCP collector if requested.
	if *tcpIface != "" {
		var tlsConfig *tls.Config
		if *caPemPath != "" && *caKeyPath != "" {
			tlsConfig, err = newTLSConfig(*caPemPath, *caKeyPath)
			if err != nil {
				log.Fatalf("failed to configure TLS: %s", err.Error())
			}
			log.Printf("TLS successfully configured")
		}

		collector, err := input.NewCollector("tcp", *tcpIface, *inputFormat, tlsConfig)
		if err != nil {
			log.Fatalf("failed to create TCP collector: %s", err.Error())
		}
		if err := collector.Start(batcher.C()); err != nil {
			log.Fatalf("failed to start TCP collector: %s", err.Error())
		}
		log.Printf("TCP collector listening to %s", *tcpIface)
	}

	// Start UDP collector if requested.
	if *udpIface != "" {
		collector, err := input.NewCollector("udp", *udpIface, *inputFormat, nil)
		if collector == nil {
			log.Fatalf("failed to create UDP collector: %s", err.Error())
		}
		if err := collector.Start(batcher.C()); err != nil {
			log.Fatalf("failed to start UDP collector: %s", err.Error())
		}
		log.Printf("UDP collector listening to %s", *udpIface)
	}

	// Start profiling.
	startProfile(*cpuProfile, *memProfile)

	stats.Set("launch", time.Now().UTC())

	// Set up signal handling.
	signalCh := make(chan os.Signal, 1)
	signal.Notify(signalCh, os.Interrupt, syscall.SIGTERM)

	// Block until one of the signals above is received
	select {
	case <-signalCh:
		log.Println("signal received, shutting down...")
	}

	stopProfile()
}

func startQueryServer(iface string, engine *ekanite.Engine) {
	server := ekanite.NewServer(iface, engine)
	if server == nil {
		log.Fatal("failed to create query server")
	}
	if err := server.Start(); err != nil {
		log.Fatalf("failed to start query server: %s", err.Error())
	}
	log.Printf("query server listening to %s", iface)
}

func startHTTPQueryServer(iface string, engine *ekanite.Engine) {
	server := ekanite.NewHTTPServer(iface, engine)
	if server == nil {
		log.Fatal("failed to create HTTP query server")
	}
	if err := server.Start(); err != nil {
		log.Fatalf("failed to start HTTP query server: %s", err.Error())
	}
	log.Printf("HTTP query server listening to %s", iface)
}

func startDiagServer(iface string) {
	diagServer := status.NewService(iface)
	if err := diagServer.Start(); err != nil {
		log.Fatalf("failed to start status server on %s: %s", iface, err.Error())
	}
}

func newTLSConfig(caPemPath, caKeyPath string) (*tls.Config, error) {
	var config *tls.Config

	caPem, err := ioutil.ReadFile(caPemPath)
	if err != nil {
		return nil, err
	}
	ca, err := x509.ParseCertificate(caPem)
	if err != nil {
		return nil, err
	}

	caKey, err := ioutil.ReadFile(caKeyPath)
	if err != nil {
		return nil, err
	}
	key, err := x509.ParsePKCS1PrivateKey(caKey)
	if err != nil {
		return nil, err
	}
	pool := x509.NewCertPool()
	pool.AddCert(ca)

	cert := tls.Certificate{
		Certificate: [][]byte{caPem},
		PrivateKey:  key,
	}

	config = &tls.Config{
		ClientAuth:   tls.RequireAndVerifyClientCert,
		Certificates: []tls.Certificate{cert},
		ClientCAs:    pool,
	}

	config.Rand = rand.Reader

	return config, nil
}

// prof stores the file locations of active profiles.
var prof struct {
	cpu *os.File
	mem *os.File
}

// StartProfile initializes the cpu and memory profile, if specified.
func startProfile(cpuprofile, memprofile string) {
	if cpuprofile != "" {
		f, err := os.Create(cpuprofile)
		if err != nil {
			log.Fatalf("cpuprofile: %v", err)
		}
		log.Printf("writing CPU profile to: %s\n", cpuprofile)
		prof.cpu = f
		pprof.StartCPUProfile(prof.cpu)
	}

	if memprofile != "" {
		f, err := os.Create(memprofile)
		if err != nil {
			log.Fatalf("memprofile: %v", err)
		}
		log.Printf("writing memory profile to: %s\n", memprofile)
		prof.mem = f
		runtime.MemProfileRate = 4096
	}

}

// StopProfile closes the cpu and memory profiles if they are running.
func stopProfile() {
	if prof.cpu != nil {
		pprof.StopCPUProfile()
		prof.cpu.Close()
		log.Println("CPU profile stopped")
	}
	if prof.mem != nil {
		pprof.Lookup("heap").WriteTo(prof.mem, 0)
		prof.mem.Close()
		log.Println("memory profile stopped")
	}
}

func printHelp() {
	fmt.Println("ekanited [options]")
	fs.PrintDefaults()
}
<|endoftext|>


<commit_before>// Copyright 2015 SpectoLabs. 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.

// hoverfly is an HTTP/s proxy configurable via flags/environment variables/admin HTTP API
//
// this proxy can be dynamically configured through HTTP calls when it's running, to change modes,
// export and import requests.

package main

import (
	log "github.com/Sirupsen/logrus"
	hv "github.com/SpectoLabs/hoverfly"
	"github.com/SpectoLabs/hoverfly/authentication/backends"

	"flag"
	"fmt"
	"net/http"
)

func main() {
	log.SetFormatter(&log.JSONFormatter{})

	// getting proxy configuration
	verbose := flag.Bool("v", false, "should every proxy request be logged to stdout")
	// modes
	capture := flag.Bool("capture", false, "should proxy capture requests")
	synthesize := flag.Bool("synthesize", false, "should proxy capture requests")
	modify := flag.Bool("modify", false, "should proxy only modify requests")

	destination := flag.String("destination", ".", "destination URI to catch")
	middleware := flag.String("middleware", "", "should proxy use middleware")

	// proxy port
	proxyPort := flag.String("pp", "", "proxy port - run proxy on another port (i.e. '-pp 9999' to run proxy on port 9999)")
	// admin port
	adminPort := flag.String("ap", "", "admin port - run admin interface on another port (i.e. '-ap 1234' to run admin UI on port 1234)")

	// database location
	database := flag.String("db", "", "database location - supply it if you want to provide specific to database (will be created there if it doesn't exist)")

	// delete current database on startup
	wipeDb := flag.Bool("wipedb", false, "supply -wipedb flag to delete all records from given database on startup")

	// metrics
	metrics := flag.Bool("metrics", false, "supply -metrics flag to enable metrics logging to stdout")

	// development
	dev := flag.Bool("dev", false, "supply -dev flag to serve directly from ./static/dist instead from statik binary")

	// import flag
	imp := flag.String("import", "", "import from file or from URL (i.e. '-import my_service.json' or '-import http://mypage.com/service_x.json'")

	// adding new user
	addNew := flag.Bool("add", false, "add new user '-add -username hfadmin -password hfpass'")
	addUser := flag.String("username", "", "username for new user")
	addPassword := flag.String("password", "", "password for new user")
	isAdmin := flag.Bool("admin", true, "supply '-admin false' to make this non admin user (defaults to 'true') ")

	// TODO: this should be enabled by default when UI and documentation is ready
	authEnabled := flag.Bool("auth", false, "enable authentication, currently it is disabled by default")

	flag.Parse()

	// getting settings
	cfg := hv.InitSettings()

	if *verbose {
		// Only log the warning severity or above.
		log.SetLevel(log.DebugLevel)
	}
	cfg.Verbose = *verbose

	if *dev {
		// making text pretty
		log.SetFormatter(&log.TextFormatter{})
	}

	// overriding environment variables (proxy and admin ports)
	if *proxyPort != "" {
		cfg.ProxyPort = *proxyPort
	}
	if *adminPort != "" {
		cfg.AdminPort = *adminPort
	}

	// development settings
	cfg.Development = *dev

	// overriding database location
	if *database != "" {
		cfg.DatabaseName = *database
	}

	if *wipeDb {
		os.Remove(cfg.DatabaseName)
	}

	// overriding default middleware setting
	cfg.Middleware = *middleware

	// setting default mode
	mode := hv.VirtualizeMode

	if *capture {
		mode = hv.CaptureMode
		// checking whether user supplied other modes
		if *synthesize == true || *modify == true {
			log.Fatal("Two or more modes supplied, check your flags")
		}
	} else if *synthesize {
		mode = hv.SynthesizeMode

		if cfg.Middleware == "" {
			log.Fatal("Synthesize mode chosen although middleware not supplied")
		}

		if *capture == true || *modify == true {
			log.Fatal("Two or more modes supplied, check your flags")
		}
	} else if *modify {
		mode = hv.ModifyMode

		if cfg.Middleware == "" {
			log.Fatal("Modify mode chosen although middleware not supplied")
		}

		if *capture == true || *synthesize == true {
			log.Fatal("Two or more modes supplied, check your flags")
		}
	}

	// overriding default settings
	cfg.Mode = mode

	// enabling authentication if flag or env variable is set to 'true'
	if cfg.AuthEnabled || *authEnabled {
		cfg.AuthEnabled = true
	}

	// overriding destination
	cfg.Destination = *destination

	// getting boltDB
	db := hv.GetDB(cfg.DatabaseName)
	cache := hv.NewBoltDBCache(db, []byte(hv.RequestsBucketName))
	defer cache.CloseDB()

	proxy, dbClient := hv.GetNewHoverfly(cfg, cache)

	ab := backends.NewBoltDBAuthBackend(db, []byte(backends.TokenBucketName), []byte(backends.UserBucketName))

	// assigning auth backend
	dbClient.AB = ab

	// if add new user supplied - adding it to database
	if *addNew {
		err := ab.AddUser([]byte(*addUser), []byte(*addPassword), *isAdmin)
		if err != nil {
			log.WithFields(log.Fields{
				"error":    err.Error(),
				"username": *addUser,
			}).Fatal("failed to add new user")
		} else {
			log.WithFields(log.Fields{
				"username": *addUser,
			}).Info("user added successfuly")
		}
		return
	}

	// importing stuff
	if *imp != "" {
		err := dbClient.Import(*imp)
		if err != nil {
			log.WithFields(log.Fields{
				"error":  err.Error(),
				"import": *imp,
			}).Fatal("Failed to import given resource")
		}
	}

	// starting admin interface
	dbClient.StartAdminInterface()

	// start metrics registry flush
	if *metrics {
		dbClient.Counter.Init()
	}

	log.Warn(http.ListenAndServe(fmt.Sprintf(":%s", cfg.ProxyPort), proxy))
}
<commit_msg>updated mode descriptions<commit_after>// Copyright 2015 SpectoLabs. 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.

// hoverfly is an HTTP/s proxy configurable via flags/environment variables/admin HTTP API
//
// this proxy can be dynamically configured through HTTP calls when it's running, to change modes,
// export and import requests.

package main

import (
	log "github.com/Sirupsen/logrus"
	hv "github.com/SpectoLabs/hoverfly"
	"github.com/SpectoLabs/hoverfly/authentication/backends"

	"flag"
	"fmt"
	"net/http"
)

func main() {
	log.SetFormatter(&log.JSONFormatter{})

	// getting proxy configuration
	verbose := flag.Bool("v", false, "should every proxy request be logged to stdout")
	// modes
	capture := flag.Bool("capture", false, "start Hoverfly in capture mode - transparently intercepts and saves requests/response")
	synthesize := flag.Bool("synthesize", false, "start Hoverfly in synthesize mode (middleware is required)")
	modify := flag.Bool("modify", false, "start Hoverfly in modify mode - applies middleware (required) to both outgoing and incomming HTTP traffic")

	destination := flag.String("destination", ".", "destination URI to catch")
	middleware := flag.String("middleware", "", "should proxy use middleware")

	// proxy port
	proxyPort := flag.String("pp", "", "proxy port - run proxy on another port (i.e. '-pp 9999' to run proxy on port 9999)")
	// admin port
	adminPort := flag.String("ap", "", "admin port - run admin interface on another port (i.e. '-ap 1234' to run admin UI on port 1234)")

	// database location
	database := flag.String("db", "", "database location - supply it if you want to provide specific to database (will be created there if it doesn't exist)")

	// delete current database on startup
	wipeDb := flag.Bool("wipedb", false, "supply -wipedb flag to delete all records from given database on startup")

	// metrics
	metrics := flag.Bool("metrics", false, "supply -metrics flag to enable metrics logging to stdout")

	// development
	dev := flag.Bool("dev", false, "supply -dev flag to serve directly from ./static/dist instead from statik binary")

	// import flag
	imp := flag.String("import", "", "import from file or from URL (i.e. '-import my_service.json' or '-import http://mypage.com/service_x.json'")

	// adding new user
	addNew := flag.Bool("add", false, "add new user '-add -username hfadmin -password hfpass'")
	addUser := flag.String("username", "", "username for new user")
	addPassword := flag.String("password", "", "password for new user")
	isAdmin := flag.Bool("admin", true, "supply '-admin false' to make this non admin user (defaults to 'true') ")

	// TODO: this should be enabled by default when UI and documentation is ready
	authEnabled := flag.Bool("auth", false, "enable authentication, currently it is disabled by default")

	flag.Parse()

	// getting settings
	cfg := hv.InitSettings()

	if *verbose {
		// Only log the warning severity or above.
		log.SetLevel(log.DebugLevel)
	}
	cfg.Verbose = *verbose

	if *dev {
		// making text pretty
		log.SetFormatter(&log.TextFormatter{})
	}

	// overriding environment variables (proxy and admin ports)
	if *proxyPort != "" {
		cfg.ProxyPort = *proxyPort
	}
	if *adminPort != "" {
		cfg.AdminPort = *adminPort
	}

	// development settings
	cfg.Development = *dev

	// overriding database location
	if *database != "" {
		cfg.DatabaseName = *database
	}

	if *wipeDb {
		os.Remove(cfg.DatabaseName)
	}

	// overriding default middleware setting
	cfg.Middleware = *middleware

	// setting default mode
	mode := hv.VirtualizeMode

	if *capture {
		mode = hv.CaptureMode
		// checking whether user supplied other modes
		if *synthesize == true || *modify == true {
			log.Fatal("Two or more modes supplied, check your flags")
		}
	} else if *synthesize {
		mode = hv.SynthesizeMode

		if cfg.Middleware == "" {
			log.Fatal("Synthesize mode chosen although middleware not supplied")
		}

		if *capture == true || *modify == true {
			log.Fatal("Two or more modes supplied, check your flags")
		}
	} else if *modify {
		mode = hv.ModifyMode

		if cfg.Middleware == "" {
			log.Fatal("Modify mode chosen although middleware not supplied")
		}

		if *capture == true || *synthesize == true {
			log.Fatal("Two or more modes supplied, check your flags")
		}
	}

	// overriding default settings
	cfg.Mode = mode

	// enabling authentication if flag or env variable is set to 'true'
	if cfg.AuthEnabled || *authEnabled {
		cfg.AuthEnabled = true
	}

	// overriding destination
	cfg.Destination = *destination

	// getting boltDB
	db := hv.GetDB(cfg.DatabaseName)
	cache := hv.NewBoltDBCache(db, []byte(hv.RequestsBucketName))
	defer cache.CloseDB()

	proxy, dbClient := hv.GetNewHoverfly(cfg, cache)

	ab := backends.NewBoltDBAuthBackend(db, []byte(backends.TokenBucketName), []byte(backends.UserBucketName))

	// assigning auth backend
	dbClient.AB = ab

	// if add new user supplied - adding it to database
	if *addNew {
		err := ab.AddUser([]byte(*addUser), []byte(*addPassword), *isAdmin)
		if err != nil {
			log.WithFields(log.Fields{
				"error":    err.Error(),
				"username": *addUser,
			}).Fatal("failed to add new user")
		} else {
			log.WithFields(log.Fields{
				"username": *addUser,
			}).Info("user added successfuly")
		}
		return
	}

	// importing stuff
	if *imp != "" {
		err := dbClient.Import(*imp)
		if err != nil {
			log.WithFields(log.Fields{
				"error":  err.Error(),
				"import": *imp,
			}).Fatal("Failed to import given resource")
		}
	}

	// starting admin interface
	dbClient.StartAdminInterface()

	// start metrics registry flush
	if *metrics {
		dbClient.Counter.Init()
	}

	log.Warn(http.ListenAndServe(fmt.Sprintf(":%s", cfg.ProxyPort), proxy))
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"log"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/nelhage/taktician/playtak"
	"github.com/nelhage/taktician/ptn"
	"github.com/nelhage/taktician/tak"
)

type Game struct {
	id       string
	opponent string
	color    tak.Color
	size     int
	time     time.Duration

	positions []*tak.Position
	moves     []tak.Move
}

type Bot interface {
	NewGame(g *Game)
	GameOver()
	GetMove(p *tak.Position, mine, theirs time.Duration) tak.Move
	AcceptUndo() bool
	HandleChat(who, msg string)
}

type Client interface {
	Recv() <-chan string
	SendCommand(...string)
}

func playGame(c Client, b Bot, line string) {
	var g Game
	bits := strings.Split(line, " ")
	g.size, _ = strconv.Atoi(bits[3])
	g.id = bits[2]
	switch bits[7] {
	case "white":
		g.color = tak.White
		g.opponent = bits[6]
	case "black":
		g.color = tak.Black
		g.opponent = bits[4]
	default:
		panic(fmt.Sprintf("bad color: %s", bits[7]))
	}

	secs, _ := strconv.Atoi(bits[8])
	g.time = time.Duration(secs) * time.Second

	gameStr := fmt.Sprintf("Game#%s", g.id)
	p := tak.New(tak.Config{Size: g.size})
	g.positions = append(g.positions, p)
	b.NewGame(&g)
	defer b.GameOver()

	log.Printf("new game game-id=%q size=%d opponent=%q color=%q time=%q",
		g.id, g.size, g.opponent, g.color, g.time)

	var times struct {
		mine, theirs time.Duration
	}
	times.mine = g.time
	times.theirs = g.time

	var moves chan tak.Move
	var moveLock sync.Mutex

	for {
		over, _ := p.GameOver()
		if g.color == p.ToMove() && !over {
			moves = make(chan tak.Move, 1)
			go func(mc chan<- tak.Move) {
				moveLock.Lock()
				defer moveLock.Unlock()
				mc <- b.GetMove(p, times.mine, times.theirs)
			}(moves)
		}

		var timeout <-chan time.Time
	eventLoop:
		for {
			var line string
			var ok bool
			select {
			case line, ok = <-c.Recv():
				if !ok {
					return
				}
			case move := <-moves:
				next, err := p.Move(&move)
				if err != nil {
					log.Printf("ai returned bad move: %s: %s",
						ptn.FormatMove(&move), err)
					break eventLoop
				}
				c.SendCommand(gameStr, playtak.FormatServer(&move))
				log.Printf("my-move game-id=%s ply=%d ptn=%d.%s move=%q",
					g.id,
					p.MoveNumber(),
					p.MoveNumber()/2+1,
					strings.ToUpper(p.ToMove().String()[:1]),
					ptn.FormatMove(&move))
				p = next
				g.positions = append(g.positions, p)
				g.moves = append(g.moves, move)
				continue eventLoop
			case <-timeout:
				break eventLoop
			}

			bits = strings.Split(line, " ")
			switch bits[0] {
			case gameStr:
			case "Shout":
				who, msg := playtak.ParseShout(line)
				if who != "" {
					b.HandleChat(who, msg)
				}
				fallthrough
			default:
				continue eventLoop
			}
			switch bits[1] {
			case "P", "M":
				move, err := playtak.ParseServer(strings.Join(bits[1:], " "))
				if err != nil {
					panic(err)
				}
				next, err := p.Move(&move)
				if err != nil {
					panic(err)
				}
				log.Printf("their-move game-id=%s ply=%d ptn=%d.%s move=%q",
					g.id,
					p.MoveNumber(),
					p.MoveNumber()/2+1,
					strings.ToUpper(p.ToMove().String()[:1]),
					ptn.FormatMove(&move))
				p = next
				g.positions = append(g.positions, p)
				g.moves = append(g.moves, move)
				timeout = time.NewTimer(500 * time.Millisecond).C
			case "Abandoned.":
				log.Printf("game-over game-id=%s opponent=%s ply=%d result=abandoned",
					g.id, g.opponent, p.MoveNumber())
				return
			case "Over":
				log.Printf("game-over game-id=%s opponent=%s ply=%d result=%q",
					g.id, g.opponent, p.MoveNumber(), bits[2])
				return
			case "Time":
				w, _ := strconv.Atoi(bits[2])
				b, _ := strconv.Atoi(bits[3])
				if g.color == tak.White {
					times.mine = time.Duration(w) * time.Second
					times.theirs = time.Duration(b) * time.Second
				} else {
					times.theirs = time.Duration(w) * time.Second
					times.mine = time.Duration(b) * time.Second
				}
				break eventLoop
			case "RequestUndo":
				if b.AcceptUndo() {
					c.SendCommand(gameStr, "RequestUndo")
				}
				moves = nil
			case "Undo":
				g.positions = g.positions[:len(g.positions)-1]
				g.moves = g.moves[:len(g.moves)-1]
				p = g.positions[len(g.positions)-1]
				break eventLoop
			}
		}
	}
}
<commit_msg>log undos<commit_after>package main

import (
	"fmt"
	"log"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/nelhage/taktician/playtak"
	"github.com/nelhage/taktician/ptn"
	"github.com/nelhage/taktician/tak"
)

type Game struct {
	id       string
	opponent string
	color    tak.Color
	size     int
	time     time.Duration

	positions []*tak.Position
	moves     []tak.Move
}

type Bot interface {
	NewGame(g *Game)
	GameOver()
	GetMove(p *tak.Position, mine, theirs time.Duration) tak.Move
	AcceptUndo() bool
	HandleChat(who, msg string)
}

type Client interface {
	Recv() <-chan string
	SendCommand(...string)
}

func playGame(c Client, b Bot, line string) {
	var g Game
	bits := strings.Split(line, " ")
	g.size, _ = strconv.Atoi(bits[3])
	g.id = bits[2]
	switch bits[7] {
	case "white":
		g.color = tak.White
		g.opponent = bits[6]
	case "black":
		g.color = tak.Black
		g.opponent = bits[4]
	default:
		panic(fmt.Sprintf("bad color: %s", bits[7]))
	}

	secs, _ := strconv.Atoi(bits[8])
	g.time = time.Duration(secs) * time.Second

	gameStr := fmt.Sprintf("Game#%s", g.id)
	p := tak.New(tak.Config{Size: g.size})
	g.positions = append(g.positions, p)
	b.NewGame(&g)
	defer b.GameOver()

	log.Printf("new game game-id=%q size=%d opponent=%q color=%q time=%q",
		g.id, g.size, g.opponent, g.color, g.time)

	var times struct {
		mine, theirs time.Duration
	}
	times.mine = g.time
	times.theirs = g.time

	var moves chan tak.Move
	var moveLock sync.Mutex

	for {
		over, _ := p.GameOver()
		if g.color == p.ToMove() && !over {
			moves = make(chan tak.Move, 1)
			go func(mc chan<- tak.Move) {
				moveLock.Lock()
				defer moveLock.Unlock()
				mc <- b.GetMove(p, times.mine, times.theirs)
			}(moves)
		}

		var timeout <-chan time.Time
	eventLoop:
		for {
			var line string
			var ok bool
			select {
			case line, ok = <-c.Recv():
				if !ok {
					return
				}
			case move := <-moves:
				next, err := p.Move(&move)
				if err != nil {
					log.Printf("ai returned bad move: %s: %s",
						ptn.FormatMove(&move), err)
					break eventLoop
				}
				c.SendCommand(gameStr, playtak.FormatServer(&move))
				log.Printf("my-move game-id=%s ply=%d ptn=%d.%s move=%q",
					g.id,
					p.MoveNumber(),
					p.MoveNumber()/2+1,
					strings.ToUpper(p.ToMove().String()[:1]),
					ptn.FormatMove(&move))
				p = next
				g.positions = append(g.positions, p)
				g.moves = append(g.moves, move)
				continue eventLoop
			case <-timeout:
				break eventLoop
			}

			bits = strings.Split(line, " ")
			switch bits[0] {
			case gameStr:
			case "Shout":
				who, msg := playtak.ParseShout(line)
				if who != "" {
					b.HandleChat(who, msg)
				}
				fallthrough
			default:
				continue eventLoop
			}
			switch bits[1] {
			case "P", "M":
				move, err := playtak.ParseServer(strings.Join(bits[1:], " "))
				if err != nil {
					panic(err)
				}
				next, err := p.Move(&move)
				if err != nil {
					panic(err)
				}
				log.Printf("their-move game-id=%s ply=%d ptn=%d.%s move=%q",
					g.id,
					p.MoveNumber(),
					p.MoveNumber()/2+1,
					strings.ToUpper(p.ToMove().String()[:1]),
					ptn.FormatMove(&move))
				p = next
				g.positions = append(g.positions, p)
				g.moves = append(g.moves, move)
				timeout = time.NewTimer(500 * time.Millisecond).C
			case "Abandoned.":
				log.Printf("game-over game-id=%s opponent=%s ply=%d result=abandoned",
					g.id, g.opponent, p.MoveNumber())
				return
			case "Over":
				log.Printf("game-over game-id=%s opponent=%s ply=%d result=%q",
					g.id, g.opponent, p.MoveNumber(), bits[2])
				return
			case "Time":
				w, _ := strconv.Atoi(bits[2])
				b, _ := strconv.Atoi(bits[3])
				if g.color == tak.White {
					times.mine = time.Duration(w) * time.Second
					times.theirs = time.Duration(b) * time.Second
				} else {
					times.theirs = time.Duration(w) * time.Second
					times.mine = time.Duration(b) * time.Second
				}
				break eventLoop
			case "RequestUndo":
				if b.AcceptUndo() {
					c.SendCommand(gameStr, "RequestUndo")
				}
				moves = nil
			case "Undo":
				log.Printf("undo game-id=%s ply=%d", g.id, p.MoveNumber())
				g.positions = g.positions[:len(g.positions)-1]
				g.moves = g.moves[:len(g.moves)-1]
				p = g.positions[len(g.positions)-1]
				break eventLoop
			}
		}
	}
}
<|endoftext|>


<commit_before>// Copyright 2015 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 upload command writes a file to Google Cloud Storage. It's used
// exclusively by the Makefiles in the Go project repos. Think of it
// as a very light version of gsutil or gcloud, but with some
// Go-specific configuration knowledge baked in.
package main

import (
	"bytes"
	"crypto/md5"
	"flag"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"os/exec"
	"runtime"
	"strings"
	"time"

	"golang.org/x/build/auth"
	"golang.org/x/build/envutil"
	"golang.org/x/net/context"
	"golang.org/x/oauth2"
	"google.golang.org/cloud"
	"google.golang.org/cloud/storage"
)

var (
	public    = flag.Bool("public", false, "object should be world-readable")
	cacheable = flag.Bool("cacheable", true, "object should be cacheable")
	file      = flag.String("file", "-", "Filename to read object from, or '-' for stdin. If it begins with 'go:' then the rest is considered to be a Go target to install first, and then upload.")
	verbose   = flag.Bool("verbose", false, "verbose logging")
	osarch    = flag.String("osarch", "", "Optional 'GOOS-GOARCH' value to cross-compile; used only if --file begins with 'go:'. As a special case, if the value contains a '.' byte, anything up to and including that period is discarded.")
	project   = flag.String("project", "", "GCE Project. If blank, it's automatically inferred from the bucket name for the common Go buckets.")
	tags      = flag.String("tags", "", "tags to pass to go list, go install, etc. Only applicable if the --file value begins with 'go:'")
)

func main() {
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage: upload [--public] [--file=...] <bucket/object>\n")
		flag.PrintDefaults()
	}
	flag.Parse()
	if flag.NArg() != 1 {
		flag.Usage()
		os.Exit(1)
	}
	args := strings.SplitN(flag.Arg(0), "/", 2)
	if len(args) != 2 {
		flag.Usage()
		os.Exit(1)
	}
	if strings.HasPrefix(*file, "go:") {
		buildGoTarget()
	}
	bucket, object := args[0], args[1]

	proj := *project
	if proj == "" {
		proj, _ = bucketProject[bucket]
		if proj == "" {
			log.Fatalf("bucket %q doesn't have an associated project in upload.go", bucket)
		}
	}

	ts, err := auth.ProjectTokenSource(proj, storage.ScopeReadWrite)
	if err != nil {
		log.Fatalf("Failed to get an OAuth2 token source: %v", err)
	}
	httpClient := oauth2.NewClient(oauth2.NoContext, ts)
	ctx := cloud.NewContext(proj, httpClient)

	if alreadyUploaded(ctx, bucket, object) {
		if *verbose {
			log.Printf("Already uploaded.")
		}
		return
	}

	w := storage.NewWriter(ctx, bucket, object)
	// If you don't give the owners access, the web UI seems to
	// have a bug and doesn't have access to see that it's public, so
	// won't render the "Shared Publicly" link. So we do that, even
	// though it's dumb and unnecessary otherwise:
	w.ACL = append(w.ACL, storage.ACLRule{Entity: storage.ACLEntity("project-owners-" + proj), Role: storage.RoleOwner})
	if *public {
		w.ACL = append(w.ACL, storage.ACLRule{Entity: storage.AllUsers, Role: storage.RoleReader})
		if !*cacheable {
			w.CacheControl = "no-cache"
		}
	}
	var content io.Reader
	if *file == "-" {
		content = os.Stdin
	} else {
		content, err = os.Open(*file)
		if err != nil {
			log.Fatal(err)
		}
	}

	const maxSlurp = 1 << 20
	var buf bytes.Buffer
	n, err := io.CopyN(&buf, content, maxSlurp)
	if err != nil && err != io.EOF {
		log.Fatalf("Error reading from stdin: %v, %v", n, err)
	}
	w.ContentType = http.DetectContentType(buf.Bytes())

	_, err = io.Copy(w, io.MultiReader(&buf, content))
	if cerr := w.Close(); cerr != nil && err == nil {
		err = cerr
	}
	if err != nil {
		log.Fatalf("Write error: %v", err)
	}
	if *verbose {
		log.Printf("Wrote %v", object)
	}
	os.Exit(0)
}

var bucketProject = map[string]string{
	"dev-gccgo-builder-data": "gccgo-dashboard-dev",
	"dev-go-builder-data":    "go-dashboard-dev",
	"gccgo-builder-data":     "gccgo-dashboard-builders",
	"go-builder-data":        "symbolic-datum-552",
	"go-build-log":           "symbolic-datum-552",
	"http2-demo-server-tls":  "symbolic-datum-552",
	"winstrap":               "999119582588",
	"gobuilder":              "999119582588", // deprecated
}

func buildGoTarget() {
	target := strings.TrimPrefix(*file, "go:")
	var goos, goarch string
	if *osarch != "" {
		*osarch = (*osarch)[strings.LastIndex(*osarch, ".")+1:]
		v := strings.Split(*osarch, "-")
		if len(v) != 2 || v[0] == "" || v[1] == "" {
			log.Fatalf("invalid -osarch value %q", *osarch)
		}
		goos, goarch = v[0], v[1]
	}

	env := envutil.Dedup(runtime.GOOS == "windows", append(os.Environ(), "GOOS="+goos, "GOARCH="+goarch))
	cmd := exec.Command("go", "list", "--tags="+*tags, "-f", "{{.Target}}", target)
	cmd.Env = env
	out, err := cmd.Output()
	if err != nil {
		log.Fatalf("go list: %v", err)
	}
	outFile := string(bytes.TrimSpace(out))
	fi0, err := os.Stat(outFile)
	if os.IsNotExist(err) {
		if *verbose {
			log.Printf("File %s doesn't exist; building...", outFile)
		}
	}

	version := os.Getenv("USER") + "-" + time.Now().Format(time.RFC3339)
	cmd = exec.Command("go", "install", "--tags="+*tags, "-x", "--ldflags=-X main.Version="+version, target)
	var stderr bytes.Buffer
	cmd.Stderr = &stderr
	if *verbose {
		cmd.Stdout = os.Stdout
	}
	cmd.Env = env
	if err := cmd.Run(); err != nil {
		log.Fatalf("go install %s: %v, %s", target, err, stderr.Bytes())
	}

	fi1, err := os.Stat(outFile)
	if err != nil {
		log.Fatalf("Expected output file %s stat failure after go install %v: %v", outFile, target, err)
	}
	if !os.SameFile(fi0, fi1) {
		if *verbose {
			log.Printf("File %s rebuilt.", outFile)
		}
	}
	*file = outFile
}

// alreadyUploaded reports whether *file has already been uploaded and the correct contents
// are on cloud storage already.
func alreadyUploaded(ctx context.Context, bucket, object string) bool {
	if *file == "-" {
		return false // don't know.
	}
	o, err := storage.StatObject(ctx, bucket, object)
	if err == storage.ErrObjectNotExist {
		return false
	}
	if err != nil {
		log.Printf("Warning: stat failure: %v", err)
		return false
	}
	m5 := md5.New()
	fi, err := os.Stat(*file)
	if err != nil {
		log.Fatal(err)
	}
	if fi.Size() != o.Size {
		return false
	}
	f, err := os.Open(*file)
	if err != nil {
		log.Fatal(err)
	}
	defer f.Close()
	n, err := io.Copy(m5, f)
	if err != nil {
		log.Fatal(err)
	}
	if n != fi.Size() {
		log.Printf("Warning: file size of %v changed", *file)
	}
	return bytes.Equal(m5.Sum(nil), o.MD5)
}
<commit_msg>cmd/upload: update to new Google Cloud Storage API<commit_after>// Copyright 2015 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 upload command writes a file to Google Cloud Storage. It's used
// exclusively by the Makefiles in the Go project repos. Think of it
// as a very light version of gsutil or gcloud, but with some
// Go-specific configuration knowledge baked in.
package main

import (
	"bytes"
	"crypto/md5"
	"flag"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"os/exec"
	"runtime"
	"strings"
	"time"

	"golang.org/x/build/auth"
	"golang.org/x/build/envutil"
	"golang.org/x/net/context"
	"google.golang.org/cloud"
	"google.golang.org/cloud/storage"
)

var (
	public    = flag.Bool("public", false, "object should be world-readable")
	cacheable = flag.Bool("cacheable", true, "object should be cacheable")
	file      = flag.String("file", "-", "Filename to read object from, or '-' for stdin. If it begins with 'go:' then the rest is considered to be a Go target to install first, and then upload.")
	verbose   = flag.Bool("verbose", false, "verbose logging")
	osarch    = flag.String("osarch", "", "Optional 'GOOS-GOARCH' value to cross-compile; used only if --file begins with 'go:'. As a special case, if the value contains a '.' byte, anything up to and including that period is discarded.")
	project   = flag.String("project", "", "GCE Project. If blank, it's automatically inferred from the bucket name for the common Go buckets.")
	tags      = flag.String("tags", "", "tags to pass to go list, go install, etc. Only applicable if the --file value begins with 'go:'")
)

func main() {
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage: upload [--public] [--file=...] <bucket/object>\n")
		flag.PrintDefaults()
	}
	flag.Parse()
	if flag.NArg() != 1 {
		flag.Usage()
		os.Exit(1)
	}
	args := strings.SplitN(flag.Arg(0), "/", 2)
	if len(args) != 2 {
		flag.Usage()
		os.Exit(1)
	}
	if strings.HasPrefix(*file, "go:") {
		buildGoTarget()
	}
	bucket, object := args[0], args[1]

	proj := *project
	if proj == "" {
		proj, _ = bucketProject[bucket]
		if proj == "" {
			log.Fatalf("bucket %q doesn't have an associated project in upload.go", bucket)
		}
	}

	ts, err := auth.ProjectTokenSource(proj, storage.ScopeReadWrite)
	if err != nil {
		log.Fatalf("Failed to get an OAuth2 token source: %v", err)
	}

	ctx := context.Background()
	storageClient, err := storage.NewClient(ctx, cloud.WithTokenSource(ts))
	if err != nil {
		log.Fatalf("storage.NewClient: %v", err)
	}

	if alreadyUploaded(storageClient, bucket, object) {
		if *verbose {
			log.Printf("Already uploaded.")
		}
		return
	}

	w := storageClient.Bucket(bucket).Object(object).NewWriter(ctx)
	// If you don't give the owners access, the web UI seems to
	// have a bug and doesn't have access to see that it's public, so
	// won't render the "Shared Publicly" link. So we do that, even
	// though it's dumb and unnecessary otherwise:
	w.ACL = append(w.ACL, storage.ACLRule{Entity: storage.ACLEntity("project-owners-" + proj), Role: storage.RoleOwner})
	if *public {
		w.ACL = append(w.ACL, storage.ACLRule{Entity: storage.AllUsers, Role: storage.RoleReader})
		if !*cacheable {
			w.CacheControl = "no-cache"
		}
	}
	var content io.Reader
	if *file == "-" {
		content = os.Stdin
	} else {
		content, err = os.Open(*file)
		if err != nil {
			log.Fatal(err)
		}
	}

	const maxSlurp = 1 << 20
	var buf bytes.Buffer
	n, err := io.CopyN(&buf, content, maxSlurp)
	if err != nil && err != io.EOF {
		log.Fatalf("Error reading from stdin: %v, %v", n, err)
	}
	w.ContentType = http.DetectContentType(buf.Bytes())

	_, err = io.Copy(w, io.MultiReader(&buf, content))
	if cerr := w.Close(); cerr != nil && err == nil {
		err = cerr
	}
	if err != nil {
		log.Fatalf("Write error: %v", err)
	}
	if *verbose {
		log.Printf("Wrote %v", object)
	}
	os.Exit(0)
}

var bucketProject = map[string]string{
	"dev-gccgo-builder-data": "gccgo-dashboard-dev",
	"dev-go-builder-data":    "go-dashboard-dev",
	"gccgo-builder-data":     "gccgo-dashboard-builders",
	"go-builder-data":        "symbolic-datum-552",
	"go-build-log":           "symbolic-datum-552",
	"http2-demo-server-tls":  "symbolic-datum-552",
	"winstrap":               "999119582588",
	"gobuilder":              "999119582588", // deprecated
}

func buildGoTarget() {
	target := strings.TrimPrefix(*file, "go:")
	var goos, goarch string
	if *osarch != "" {
		*osarch = (*osarch)[strings.LastIndex(*osarch, ".")+1:]
		v := strings.Split(*osarch, "-")
		if len(v) != 2 || v[0] == "" || v[1] == "" {
			log.Fatalf("invalid -osarch value %q", *osarch)
		}
		goos, goarch = v[0], v[1]
	}

	env := envutil.Dedup(runtime.GOOS == "windows", append(os.Environ(), "GOOS="+goos, "GOARCH="+goarch))
	cmd := exec.Command("go", "list", "--tags="+*tags, "-f", "{{.Target}}", target)
	cmd.Env = env
	out, err := cmd.Output()
	if err != nil {
		log.Fatalf("go list: %v", err)
	}
	outFile := string(bytes.TrimSpace(out))
	fi0, err := os.Stat(outFile)
	if os.IsNotExist(err) {
		if *verbose {
			log.Printf("File %s doesn't exist; building...", outFile)
		}
	}

	version := os.Getenv("USER") + "-" + time.Now().Format(time.RFC3339)
	cmd = exec.Command("go", "install", "--tags="+*tags, "-x", "--ldflags=-X main.Version="+version, target)
	var stderr bytes.Buffer
	cmd.Stderr = &stderr
	if *verbose {
		cmd.Stdout = os.Stdout
	}
	cmd.Env = env
	if err := cmd.Run(); err != nil {
		log.Fatalf("go install %s: %v, %s", target, err, stderr.Bytes())
	}

	fi1, err := os.Stat(outFile)
	if err != nil {
		log.Fatalf("Expected output file %s stat failure after go install %v: %v", outFile, target, err)
	}
	if !os.SameFile(fi0, fi1) {
		if *verbose {
			log.Printf("File %s rebuilt.", outFile)
		}
	}
	*file = outFile
}

// alreadyUploaded reports whether *file has already been uploaded and the correct contents
// are on cloud storage already.
func alreadyUploaded(storageClient *storage.Client, bucket, object string) bool {
	if *file == "-" {
		return false // don't know.
	}
	o, err := storageClient.Bucket(bucket).Object(object).Attrs(context.Background())
	if err == storage.ErrObjectNotExist {
		return false
	}
	if err != nil {
		log.Printf("Warning: stat failure: %v", err)
		return false
	}
	m5 := md5.New()
	fi, err := os.Stat(*file)
	if err != nil {
		log.Fatal(err)
	}
	if fi.Size() != o.Size {
		return false
	}
	f, err := os.Open(*file)
	if err != nil {
		log.Fatal(err)
	}
	defer f.Close()
	n, err := io.Copy(m5, f)
	if err != nil {
		log.Fatal(err)
	}
	if n != fi.Size() {
		log.Printf("Warning: file size of %v changed", *file)
	}
	return bytes.Equal(m5.Sum(nil), o.MD5)
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"github.com/codegangsta/cli"
	"github.com/keybase/go-jsonw"
	"github.com/keybase/go-libkb"
	"io"
	"os"
	"regexp"
	"sort"
	"strings"
)

type TrackList []*libkb.TrackChainLink

func (tl TrackList) Len() int {
	return len(tl)
}

func (tl TrackList) Swap(i, j int) {
	tl[i], tl[j] = tl[j], tl[i]
}

func (tl TrackList) Less(i, j int) bool {
	return strings.ToLower(tl[i].ToDisplayString()) < strings.ToLower(tl[j].ToDisplayString())
}

type CmdListTracking struct {
	filter  string
	json    bool
	verbose bool
	headers bool
	user    *libkb.User
	tracks  TrackList
}

func (s *CmdListTracking) ParseArgv(ctx *cli.Context) error {
	nargs := len(ctx.Args())
	var err error

	s.json = ctx.Bool("json")
	s.verbose = ctx.Bool("verbose")
	s.headers = ctx.Bool("headers")
	s.filter = ctx.String("filter")

	if nargs > 0 {
		err = fmt.Errorf("list-tracking takes no args")
	}

	return err
}

func (s *CmdListTracking) filterTracks(f func(libkb.TrackChainLink) bool) {
	tracks := make([]*libkb.TrackChainLink, 0, 0)
	for _, link := range s.tracks {
		if f(*link) {
			tracks = append(tracks, link)
		}
	}
	s.tracks = tracks
}

func (s *CmdListTracking) FilterRxx() error {
	if len(s.filter) == 0 {
		return nil
	}
	rxx, err := regexp.Compile(s.filter)
	if err != nil {
		return err
	}
	s.filterTracks(func(l libkb.TrackChainLink) bool {
		if rxx.MatchString(l.ToDisplayString()) {
			return true
		}
		for _, sb := range l.ToServiceBlocks() {
			_, v := sb.ToKeyValuePair()
			if rxx.MatchString(v) {
				return true
			}
		}
		return false
	})
	return nil
}

func (s *CmdListTracking) ProcessTracks() (err error) {
	if err = s.FilterRxx(); err != nil {
		return
	}
	return
}

func (s *CmdListTracking) skipLink(link libkb.TypedChainLink) bool {
	return link.IsRevoked() || link.IsRevocationIsh() || !s.IsActiveKey(link)
}

func (s *CmdListTracking) IsActiveKey(link libkb.TypedChainLink) bool {
	fp1, _ := s.user.GetActivePgpFingerprint()
	if fp1 == nil {
		return false
	}
	fp2 := link.GetPgpFingerprint()
	return fp1.Eq(fp2)
}

func (s *CmdListTracking) DisplayTable() (err error) {

	var cols []string

	if s.headers && s.verbose {
		cols = []string{
			"Username",
			"Sig ID",
			"PGP fingerprint",
			"When Tracked",
			"Proofs",
		}
	}

	i := 0
	idtab := s.tracks

	rowfunc := func() []string {
		var row []string

		for ; i < len(idtab) && row == nil; i++ {
			link := idtab[i]

			if s.skipLink(link) {
				continue
			}

			if !s.verbose {
				row = []string{link.ToDisplayString()}
				continue
			}

			fp, err := link.GetTrackedPgpFingerprint()
			if err != nil {
				G.Log.Warning("Bad track of %s: %s", link.ToDisplayString(), err.Error())
				continue
			}

			row = []string{
				link.ToDisplayString(),
				link.GetSigId().ToDisplayString(true),
				strings.ToUpper(fp.ToString()),
				libkb.FormatTime(link.GetCTime()),
			}
			for _, sb := range link.ToServiceBlocks() {
				row = append(row, sb.ToIdString())
			}
		}
		return row
	}

	libkb.Tablify(os.Stdout, cols, rowfunc)

	return
}

func (s *CmdListTracking) DisplayJson() (err error) {
	ret := jsonw.NewArray(0)
	_, err = io.WriteString(os.Stdout, ret.MarshalPretty()+"\n")
	return
}

func (s *CmdListTracking) Display() (err error) {
	if s.json {
		err = s.DisplayJson()
	} else {
		err = s.DisplayTable()
	}
	return
}

func (s *CmdListTracking) Run() (err error) {

	arg := libkb.LoadUserArg{Self: true}
	s.user, err = libkb.LoadUser(arg)

	if err != nil {
		return
	}

	s.tracks = s.user.IdTable.GetTrackList()

	if err = s.ProcessTracks(); err != nil {
		return
	}

	sort.Sort(s.tracks)

	if err = s.Display(); err != nil {
		return
	}

	return
}

func NewCmdListTracking(cl *CommandLine) cli.Command {
	return cli.Command{
		Name:        "list-tracking",
		Usage:       "keybase list-tracking",
		Description: "list who you're tracking",
		Action: func(c *cli.Context) {
			cl.ChooseCommand(&CmdListTracking{}, "list-tracking", c)
		},
		Flags: []cli.Flag{
			cli.BoolFlag{
				Name:  "j, json",
				Usage: "output in json format; default is text",
			},
			cli.BoolFlag{
				Name:  "v, verbose",
				Usage: "a full dump, with more gory detail",
			},
			cli.BoolFlag{
				Name:  "H, headers",
				Usage: "show column headers",
			},
			cli.StringFlag{
				Name:  "f, filter",
				Usage: "provide a regex filter",
			},
		},
	}
}

func (v *CmdListTracking) GetUsage() libkb.Usage {
	return libkb.Usage{
		Config: true,
		API:    true,
	}
}
<commit_msg>JSON feature for list_tracking<commit_after>package main

import (
	"fmt"
	"github.com/codegangsta/cli"
	"github.com/keybase/go-jsonw"
	"github.com/keybase/go-libkb"
	"io"
	"os"
	"regexp"
	"sort"
	"strings"
)

type TrackList []*libkb.TrackChainLink

func (tl TrackList) Len() int {
	return len(tl)
}

func (tl TrackList) Swap(i, j int) {
	tl[i], tl[j] = tl[j], tl[i]
}

func (tl TrackList) Less(i, j int) bool {
	return strings.ToLower(tl[i].ToDisplayString()) < strings.ToLower(tl[j].ToDisplayString())
}

type CmdListTracking struct {
	filter  string
	json    bool
	verbose bool
	headers bool
	user    *libkb.User
	tracks  TrackList
}

func (s *CmdListTracking) ParseArgv(ctx *cli.Context) error {
	nargs := len(ctx.Args())
	var err error

	s.json = ctx.Bool("json")
	s.verbose = ctx.Bool("verbose")
	s.headers = ctx.Bool("headers")
	s.filter = ctx.String("filter")

	if nargs > 0 {
		err = fmt.Errorf("list-tracking takes no args")
	}

	return err
}

func (s *CmdListTracking) filterTracks(f func(libkb.TrackChainLink) bool) {
	tracks := make([]*libkb.TrackChainLink, 0, 0)
	for _, link := range s.tracks {
		if f(*link) {
			tracks = append(tracks, link)
		}
	}
	s.tracks = tracks
}

func (s *CmdListTracking) FilterRxx() error {
	if len(s.filter) == 0 {
		return nil
	}
	rxx, err := regexp.Compile(s.filter)
	if err != nil {
		return err
	}
	s.filterTracks(func(l libkb.TrackChainLink) bool {
		if rxx.MatchString(l.ToDisplayString()) {
			return true
		}
		for _, sb := range l.ToServiceBlocks() {
			_, v := sb.ToKeyValuePair()
			if rxx.MatchString(v) {
				return true
			}
		}
		return false
	})
	return nil
}

func (s *CmdListTracking) ProcessTracks() (err error) {
	if err = s.FilterRxx(); err != nil {
		return
	}
	return
}

func (s *CmdListTracking) skipLink(link libkb.TypedChainLink) bool {
	return link.IsRevoked() || link.IsRevocationIsh() || !s.IsActiveKey(link)
}

func (s *CmdListTracking) IsActiveKey(link libkb.TypedChainLink) bool {
	fp1, _ := s.user.GetActivePgpFingerprint()
	if fp1 == nil {
		return false
	}
	fp2 := link.GetPgpFingerprint()
	return fp1.Eq(fp2)
}

func (s *CmdListTracking) CondenseRecord(l *libkb.TrackChainLink) (out *jsonw.Wrapper, err error) {
	var uid *libkb.UID
	var fp *libkb.PgpFingerprint
	var un string
	out = jsonw.NewDictionary()
	rp := l.RemoteKeyProofs()

	if uid, err = l.GetTrackedUid(); err != nil {
		return
	}

	if fp, err = l.GetTrackedPgpFingerprint(); err != nil {
		return
	}

	if un, err = l.GetTrackedUsername(); err != nil {
		return
	}

	out.SetKey("uid", jsonw.NewString(uid.ToString()))
	out.SetKey("key", jsonw.NewString(strings.ToUpper(fp.ToString())))
	out.SetKey("ctime", jsonw.NewInt64(l.GetCTime().Unix()))
	out.SetKey("username", jsonw.NewString(un))
	out.SetKey("proofs", rp)

	return
}

func (s *CmdListTracking) DisplayTable() (err error) {

	var cols []string

	if s.headers && s.verbose {
		cols = []string{
			"Username",
			"Sig ID",
			"PGP fingerprint",
			"When Tracked",
			"Proofs",
		}
	}

	i := 0
	idtab := s.tracks

	rowfunc := func() []string {
		var row []string

		for ; i < len(idtab) && row == nil; i++ {
			link := idtab[i]

			if s.skipLink(link) {
				continue
			}

			if !s.verbose {
				row = []string{link.ToDisplayString()}
				continue
			}

			fp, err := link.GetTrackedPgpFingerprint()
			if err != nil {
				G.Log.Warning("Bad track of %s: %s", link.ToDisplayString(), err.Error())
				continue
			}

			row = []string{
				link.ToDisplayString(),
				link.GetSigId().ToDisplayString(true),
				strings.ToUpper(fp.ToString()),
				libkb.FormatTime(link.GetCTime()),
			}
			for _, sb := range link.ToServiceBlocks() {
				row = append(row, sb.ToIdString())
			}
		}
		return row
	}

	libkb.Tablify(os.Stdout, cols, rowfunc)

	return
}

func (s *CmdListTracking) DisplayJson() (err error) {
	tmp := make([]*jsonw.Wrapper, 0, 1)
	for _, e := range s.tracks {
		var rec *jsonw.Wrapper
		var e2 error
		if s.verbose {
			rec = e.GetPayloadJson()
		} else if rec, e2 = s.CondenseRecord(e); e2 != nil {
			G.Log.Warning("In conversion to JSON: %s", e2.Error())
		}
		if e2 == nil {
			tmp = append(tmp, rec)
		}
	}

	ret := jsonw.NewArray(len(tmp))
	for i, r := range tmp {
		ret.SetIndex(i, r)
	}

	_, err = io.WriteString(os.Stdout, ret.MarshalPretty()+"\n")
	return
}

func (s *CmdListTracking) Display() (err error) {
	if s.json {
		err = s.DisplayJson()
	} else {
		err = s.DisplayTable()
	}
	return
}

func (s *CmdListTracking) Run() (err error) {

	arg := libkb.LoadUserArg{Self: true}
	s.user, err = libkb.LoadUser(arg)

	if err != nil {
		return
	}

	s.tracks = s.user.IdTable.GetTrackList()

	if err = s.ProcessTracks(); err != nil {
		return
	}

	sort.Sort(s.tracks)

	if err = s.Display(); err != nil {
		return
	}

	return
}

func NewCmdListTracking(cl *CommandLine) cli.Command {
	return cli.Command{
		Name:        "list-tracking",
		Usage:       "keybase list-tracking",
		Description: "list who you're tracking",
		Action: func(c *cli.Context) {
			cl.ChooseCommand(&CmdListTracking{}, "list-tracking", c)
		},
		Flags: []cli.Flag{
			cli.BoolFlag{
				Name:  "j, json",
				Usage: "output in json format; default is text",
			},
			cli.BoolFlag{
				Name:  "v, verbose",
				Usage: "a full dump, with more gory detail",
			},
			cli.BoolFlag{
				Name:  "H, headers",
				Usage: "show column headers",
			},
			cli.StringFlag{
				Name:  "f, filter",
				Usage: "provide a regex filter",
			},
		},
	}
}

func (v *CmdListTracking) GetUsage() libkb.Usage {
	return libkb.Usage{
		Config: true,
		API:    true,
	}
}
<|endoftext|>


<commit_before>package cmd

import (
	"fmt"
	"github.com/spf13/cobra/cobra/tpl"
	"os"
	"text/template"
)

// Project contains name, license and paths to projects.
type Project struct {
	// v2
	PkgName      string
	Copyright    string
	AbsolutePath string
	Legal        License
	Viper        bool
	AppName      string

	//absPath string
	//cmdPath string
	//srcPath string
	//license License
	//name    string
}

type Command struct {
	CmdName   string
	CmdParent string
	*Project
}

func (p *Project) Create() error {

	// check if AbsolutePath exists
	if _, err := os.Stat(p.AbsolutePath); os.IsNotExist(err) {
		// create directory
		if err := os.Mkdir(p.AbsolutePath, 0754); err != nil {
			return err
		}
	}

	// create main.go
	mainFile, err := os.Create(fmt.Sprintf("%s/main.go", p.AbsolutePath))
	if err != nil {
		return err
	}
	defer mainFile.Close()

	mainTemplate := template.Must(template.New("main").Parse(string(tpl.MainTemplate())))
	err = mainTemplate.Execute(mainFile, p)
	if err != nil {
		return err
	}

	// create cmd/root.go
	if _, err = os.Stat(fmt.Sprintf("%s/cmd", p.AbsolutePath)); os.IsNotExist(err) {
		os.Mkdir(fmt.Sprintf("%s/cmd", p.AbsolutePath), 0751)
	}
	rootFile, err := os.Create(fmt.Sprintf("%s/cmd/root.go", p.AbsolutePath))
	if err != nil {
		return err
	}
	defer rootFile.Close()

	rootTemplate := template.Must(template.New("root").Parse(string(tpl.RootTemplate())))
	err = rootTemplate.Execute(rootFile, p)
	if err != nil {
		return err
	}

	// create license
	return p.createLicenseFile()
}

func (p *Project) createLicenseFile() error {
	data := map[string]interface{}{
		"copyright": copyrightLine(),
	}
	licenseFile, err := os.Create(fmt.Sprintf("%s/LICENSE", p.AbsolutePath))
	if err != nil {
		return err
	}

	licenseTemplate := template.Must(template.New("license").Parse(p.Legal.Text))
	return licenseTemplate.Execute(licenseFile, data)
}

func (c *Command) Create() error {
	cmdFile, err := os.Create(fmt.Sprintf("%s/cmd/%s.go", c.AbsolutePath, c.CmdName))
	if err != nil {
		return err
	}
	defer cmdFile.Close()

	commandTemplate := template.Must(template.New("sub").Parse(string(tpl.AddCommandTemplate())))
	err = commandTemplate.Execute(cmdFile, c)
	if err != nil {
		return err
	}
	return nil
}
<commit_msg>remove unused struct fields<commit_after>package cmd

import (
	"fmt"
	"github.com/spf13/cobra/cobra/tpl"
	"os"
	"text/template"
)

// Project contains name, license and paths to projects.
type Project struct {
	// v2
	PkgName      string
	Copyright    string
	AbsolutePath string
	Legal        License
	Viper        bool
	AppName      string
}

type Command struct {
	CmdName   string
	CmdParent string
	*Project
}

func (p *Project) Create() error {

	// check if AbsolutePath exists
	if _, err := os.Stat(p.AbsolutePath); os.IsNotExist(err) {
		// create directory
		if err := os.Mkdir(p.AbsolutePath, 0754); err != nil {
			return err
		}
	}

	// create main.go
	mainFile, err := os.Create(fmt.Sprintf("%s/main.go", p.AbsolutePath))
	if err != nil {
		return err
	}
	defer mainFile.Close()

	mainTemplate := template.Must(template.New("main").Parse(string(tpl.MainTemplate())))
	err = mainTemplate.Execute(mainFile, p)
	if err != nil {
		return err
	}

	// create cmd/root.go
	if _, err = os.Stat(fmt.Sprintf("%s/cmd", p.AbsolutePath)); os.IsNotExist(err) {
		os.Mkdir(fmt.Sprintf("%s/cmd", p.AbsolutePath), 0751)
	}
	rootFile, err := os.Create(fmt.Sprintf("%s/cmd/root.go", p.AbsolutePath))
	if err != nil {
		return err
	}
	defer rootFile.Close()

	rootTemplate := template.Must(template.New("root").Parse(string(tpl.RootTemplate())))
	err = rootTemplate.Execute(rootFile, p)
	if err != nil {
		return err
	}

	// create license
	return p.createLicenseFile()
}

func (p *Project) createLicenseFile() error {
	data := map[string]interface{}{
		"copyright": copyrightLine(),
	}
	licenseFile, err := os.Create(fmt.Sprintf("%s/LICENSE", p.AbsolutePath))
	if err != nil {
		return err
	}

	licenseTemplate := template.Must(template.New("license").Parse(p.Legal.Text))
	return licenseTemplate.Execute(licenseFile, data)
}

func (c *Command) Create() error {
	cmdFile, err := os.Create(fmt.Sprintf("%s/cmd/%s.go", c.AbsolutePath, c.CmdName))
	if err != nil {
		return err
	}
	defer cmdFile.Close()

	commandTemplate := template.Must(template.New("sub").Parse(string(tpl.AddCommandTemplate())))
	err = commandTemplate.Execute(cmdFile, c)
	if err != nil {
		return err
	}
	return nil
}
<|endoftext|>


<commit_before>package middleware

import (
	"bufio"
	"compress/flate"
	"compress/gzip"
	"errors"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"strings"
	"sync"
)

var defaultCompressibleContentTypes = []string{
	"text/html",
	"text/css",
	"text/plain",
	"text/javascript",
	"application/javascript",
	"application/x-javascript",
	"application/json",
	"application/atom+xml",
	"application/rss+xml",
	"image/svg+xml",
}

// A default compressor that allows for the old API to use the new code.
// DEPRECATED
var defaultCompressor *Compressor

// Compressor represents a set of encoding configurations.
type Compressor struct {
	level int // The compression level.
	// The mapping of encoder names to encoder functions.
	encoders map[string]EncoderFunc
	// The mapping of pooled encoders to pools.
	pooledEncoders map[string]*sync.Pool
	// The set of content types allowed to be compressed.
	allowedTypes map[string]bool
	// The list of encoders in order of decreasing precedence.
	encodingPrecedence []string
}

// NewCompressor creates a new Compressor that will handle encoding responses.
//
// The level should be one of the ones defined in the flate package.
// The types are the content types that are allowed to be compressed.
func NewCompressor(level int, types ...string) *Compressor {
	// If types are provided, set those as the allowed types. If none are
	// provided, use the default list.
	allowedTypes := make(map[string]bool)
	if len(types) > 0 {
		for _, t := range types {
			allowedTypes[t] = true
		}
	} else {
		for _, t := range defaultCompressibleContentTypes {
			allowedTypes[t] = true
		}
	}

	c := &Compressor{
		level:          level,
		encoders:       make(map[string]EncoderFunc),
		pooledEncoders: make(map[string]*sync.Pool),
		allowedTypes:   allowedTypes,
	}
	// Set the default encoders.  The precedence order uses the reverse
	// ordering that the encoders were added. This means adding new encoders
	// will move them to the front of the order.
	//
	// TODO:
	// lzma: Opera.
	// sdch: Chrome, Android. Gzip output + dictionary header.
	// br:   Brotli, see https://github.com/go-chi/chi/pull/326

	// HTTP 1.1 "deflate" (RFC 2616) stands for DEFLATE data (RFC 1951)
	// wrapped with zlib (RFC 1950). The zlib wrapper uses Adler-32
	// checksum compared to CRC-32 used in "gzip" and thus is faster.
	//
	// But.. some old browsers (MSIE, Safari 5.1) incorrectly expect
	// raw DEFLATE data only, without the mentioned zlib wrapper.
	// Because of this major confusion, most modern browsers try it
	// both ways, first looking for zlib headers.
	// Quote by Mark Adler: http://stackoverflow.com/a/9186091/385548
	//
	// The list of browsers having problems is quite big, see:
	// http://zoompf.com/blog/2012/02/lose-the-wait-http-compression
	// https://web.archive.org/web/20120321182910/http://www.vervestudios.co/projects/compression-tests/results
	//
	// That's why we prefer gzip over deflate. It's just more reliable
	// and not significantly slower than gzip.
	c.SetEncoder("deflate", encoderDeflate)

	// TODO: Exception for old MSIE browsers that can't handle non-HTML?
	// https://zoompf.com/blog/2012/02/lose-the-wait-http-compression
	c.SetEncoder("gzip", encoderGzip)

	// NOTE: Not implemented, intentionally:
	// case "compress": // LZW. Deprecated.
	// case "bzip2":    // Too slow on-the-fly.
	// case "zopfli":   // Too slow on-the-fly.
	// case "xz":       // Too slow on-the-fly.
	return c
}

// SetEncoder can be used to set the implementation of a compression algorithm.
//
// The encoding should be a standardised identifier. See:
// https://developer.mozilla.org/en-US/docs/Web/HTTP/Headers/Accept-Encoding
//
// For example, add the Brotli algortithm:
//
//  import brotli_enc "gopkg.in/kothar/brotli-go.v0/enc"
//
//  compressor := middleware.NewCompressor(5, "text/html")
//  compressor.SetEncoder("br", func(w http.ResponseWriter, level int) io.Writer {
//    params := brotli_enc.NewBrotliParams()
//    params.SetQuality(level)
//    return brotli_enc.NewBrotliWriter(params, w)
//  })
func (c *Compressor) SetEncoder(encoding string, fn EncoderFunc) {
	encoding = strings.ToLower(encoding)
	if encoding == "" {
		panic("the encoding can not be empty")
	}
	if fn == nil {
		panic("attempted to set a nil encoder function")
	}

	// If we are adding a new encoder that is already registered, we have to
	// clear that one out first.
	if _, ok := c.pooledEncoders[encoding]; ok {
		delete(c.pooledEncoders, encoding)
	}
	if _, ok := c.encoders[encoding]; ok {
		delete(c.encoders, encoding)
	}

	// If the encoder supports Resetting (IoReseterWriter), then it can be pooled.
	encoder := fn(ioutil.Discard, c.level)
	if encoder != nil {
		if _, ok := encoder.(ioResetterWriter); ok {
			pool := &sync.Pool{
				New: func() interface{} {
					return fn(ioutil.Discard, c.level)
				},
			}
			c.pooledEncoders[encoding] = pool
		}
	}
	// If the encoder is not in the pooledEncoders, add it to the normal encoders.
	if _, ok := c.pooledEncoders[encoding]; !ok {
		c.encoders[encoding] = fn
	}

	for i, v := range c.encodingPrecedence {
		if v == encoding {
			c.encodingPrecedence = append(c.encodingPrecedence[:i], c.encodingPrecedence[i+1:]...)
		}
	}

	c.encodingPrecedence = append([]string{encoding}, c.encodingPrecedence...)
}

// Handler returns a new middleware that will compress the response based on the
// current Compressor.
func (c *Compressor) Handler() func(next http.Handler) http.Handler {
	return func(next http.Handler) http.Handler {
		fn := func(w http.ResponseWriter, r *http.Request) {
			encoder, encoding, cleanup := c.selectEncoder(r.Header, w)

			cw := &compressResponseWriter{
				ResponseWriter: w,
				w:              w,
				contentTypes:   c.allowedTypes,
				encoding:       encoding,
			}
			if encoder != nil {
				cw.w = encoder
			}
			// Re-add the encoder to the pool if applicable.
			defer cleanup()
			defer cw.Close()

			next.ServeHTTP(cw, r)
		}

		return http.HandlerFunc(fn)
	}

}

// selectEncoder returns the encoder, the name of the encoder, and a closer function.
func (c *Compressor) selectEncoder(h http.Header, w io.Writer) (io.Writer, string, func()) {
	header := h.Get("Accept-Encoding")

	// Parse the names of all accepted algorithms from the header.
	accepted := strings.Split(strings.ToLower(header), ",")

	// Find supported encoder by accepted list by precedence
	for _, name := range c.encodingPrecedence {
		if matchAcceptEncoding(accepted, name) {
			if pool, ok := c.pooledEncoders[name]; ok {
				encoder := pool.Get().(ioResetterWriter)
				cleanup := func() {
					pool.Put(encoder)
				}
				encoder.Reset(w)
				return encoder, name, cleanup

			}
			if fn, ok := c.encoders[name]; ok {
				return fn(w, c.level), name, func() {}
			}
		}

	}

	// No encoder found to match the accepted encoding
	return nil, "", func() {}
}

func matchAcceptEncoding(accepted []string, encoding string) bool {
	for _, v := range accepted {
		if strings.Contains(v, encoding) {
			return true
		}
	}
	return false
}

// An EncoderFunc is a function that wraps the provided io.Writer with a
// streaming compression algorithm and returns it.
//
// In case of failure, the function should return nil.
type EncoderFunc func(w io.Writer, level int) io.Writer

// Interface for types that allow resetting io.Writers.
type ioResetterWriter interface {
	io.Writer
	Reset(w io.Writer)
}

// SetEncoder can be used to set the implementation of a compression algorithm.
//
// The encoding should be a standardised identifier. See:
// https://developer.mozilla.org/en-US/docs/Web/HTTP/Headers/Accept-Encoding
//
// For example, add the Brotli algortithm:
//
//  import brotli_enc "gopkg.in/kothar/brotli-go.v0/enc"
//
//  middleware.SetEncoder("br", func(w http.ResponseWriter, level int) io.Writer {
//    params := brotli_enc.NewBrotliParams()
//    params.SetQuality(level)
//    return brotli_enc.NewBrotliWriter(params, w)
//  })
//
//  DEPRECATED
func SetEncoder(encoding string, fn EncoderFunc) {
	if defaultCompressor == nil {
		panic("no compressor to set encoders on. Call Compress() first")
	}
	defaultCompressor.SetEncoder(encoding, fn)
}

// DefaultCompress is a middleware that compresses response
// body of predefined content types to a data format based
// on Accept-Encoding request header. It uses a default
// compression level.
// DEPRECATED
func DefaultCompress(next http.Handler) http.Handler {
	return Compress(flate.DefaultCompression)(next)
}

// Compress is a middleware that compresses response
// body of a given content types to a data format based
// on Accept-Encoding request header. It uses a given
// compression level.
//
// NOTE: make sure to set the Content-Type header on your response
// otherwise this middleware will not compress the response body. For ex, in
// your handler you should set w.Header().Set("Content-Type", http.DetectContentType(yourBody))
// or set it manually.
//
// DEPRECATED
func Compress(level int, types ...string) func(next http.Handler) http.Handler {
	defaultCompressor = NewCompressor(level, types...)
	return defaultCompressor.Handler()
}

type compressResponseWriter struct {
	http.ResponseWriter
	// The streaming encoder writer to be used if there is one. Otherwise,
	// this is just the normal writer.
	w            io.Writer
	encoding     string
	contentTypes map[string]bool
	wroteHeader  bool
}

func (cw *compressResponseWriter) WriteHeader(code int) {
	if cw.wroteHeader {
		return
	}
	cw.wroteHeader = true
	defer cw.ResponseWriter.WriteHeader(code)

	// Already compressed data?
	if cw.Header().Get("Content-Encoding") != "" {
		return
	}

	// Parse the first part of the Content-Type response header.
	contentType := cw.Header().Get("Content-Type")
	if idx := strings.Index(contentType, ";"); idx >= 0 {
		contentType = contentType[0:idx]
	}

	// Is the content type compressable?
	if _, ok := cw.contentTypes[contentType]; !ok {
		return
	}

	if cw.encoding != "" {
		cw.Header().Set("Content-Encoding", cw.encoding)
		cw.Header().Set("Vary", "Accept-Encoding")

		// The content-length after compression is unknown
		cw.Header().Del("Content-Length")
	}
}

func (cw *compressResponseWriter) Write(p []byte) (int, error) {
	if !cw.wroteHeader {
		cw.WriteHeader(http.StatusOK)
	}

	return cw.w.Write(p)
}

func (cw *compressResponseWriter) Flush() {
	if f, ok := cw.w.(http.Flusher); ok {
		f.Flush()
	}
}

func (cw *compressResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
	if hj, ok := cw.w.(http.Hijacker); ok {
		return hj.Hijack()
	}
	return nil, nil, errors.New("chi/middleware: http.Hijacker is unavailable on the writer")
}

func (cw *compressResponseWriter) Push(target string, opts *http.PushOptions) error {
	if ps, ok := cw.w.(http.Pusher); ok {
		return ps.Push(target, opts)
	}
	return errors.New("chi/middleware: http.Pusher is unavailable on the writer")
}

func (cw *compressResponseWriter) Close() error {
	if c, ok := cw.w.(io.WriteCloser); ok {
		return c.Close()
	}
	return errors.New("chi/middleware: io.WriteCloser is unavailable on the writer")
}

func encoderGzip(w io.Writer, level int) io.Writer {
	gw, err := gzip.NewWriterLevel(w, level)
	if err != nil {
		return nil
	}
	return gw
}

func encoderDeflate(w io.Writer, level int) io.Writer {
	dw, err := flate.NewWriter(w, level)
	if err != nil {
		return nil
	}
	return dw
}
<commit_msg>middleware.Compress: only use encoder writer for allowed content-types, properly<commit_after>package middleware

import (
	"bufio"
	"compress/flate"
	"compress/gzip"
	"errors"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"strings"
	"sync"
)

var defaultCompressibleContentTypes = []string{
	"text/html",
	"text/css",
	"text/plain",
	"text/javascript",
	"application/javascript",
	"application/x-javascript",
	"application/json",
	"application/atom+xml",
	"application/rss+xml",
	"image/svg+xml",
}

// A default compressor that allows for the old API to use the new code.
// DEPRECATED
var defaultCompressor *Compressor

// Compressor represents a set of encoding configurations.
type Compressor struct {
	level int // The compression level.
	// The mapping of encoder names to encoder functions.
	encoders map[string]EncoderFunc
	// The mapping of pooled encoders to pools.
	pooledEncoders map[string]*sync.Pool
	// The set of content types allowed to be compressed.
	allowedTypes map[string]bool
	// The list of encoders in order of decreasing precedence.
	encodingPrecedence []string
}

// NewCompressor creates a new Compressor that will handle encoding responses.
//
// The level should be one of the ones defined in the flate package.
// The types are the content types that are allowed to be compressed.
func NewCompressor(level int, types ...string) *Compressor {
	// If types are provided, set those as the allowed types. If none are
	// provided, use the default list.
	allowedTypes := make(map[string]bool)
	if len(types) > 0 {
		for _, t := range types {
			allowedTypes[t] = true
		}
	} else {
		for _, t := range defaultCompressibleContentTypes {
			allowedTypes[t] = true
		}
	}

	c := &Compressor{
		level:          level,
		encoders:       make(map[string]EncoderFunc),
		pooledEncoders: make(map[string]*sync.Pool),
		allowedTypes:   allowedTypes,
	}

	// Set the default encoders.  The precedence order uses the reverse
	// ordering that the encoders were added. This means adding new encoders
	// will move them to the front of the order.
	//
	// TODO:
	// lzma: Opera.
	// sdch: Chrome, Android. Gzip output + dictionary header.
	// br:   Brotli, see https://github.com/go-chi/chi/pull/326

	// HTTP 1.1 "deflate" (RFC 2616) stands for DEFLATE data (RFC 1951)
	// wrapped with zlib (RFC 1950). The zlib wrapper uses Adler-32
	// checksum compared to CRC-32 used in "gzip" and thus is faster.
	//
	// But.. some old browsers (MSIE, Safari 5.1) incorrectly expect
	// raw DEFLATE data only, without the mentioned zlib wrapper.
	// Because of this major confusion, most modern browsers try it
	// both ways, first looking for zlib headers.
	// Quote by Mark Adler: http://stackoverflow.com/a/9186091/385548
	//
	// The list of browsers having problems is quite big, see:
	// http://zoompf.com/blog/2012/02/lose-the-wait-http-compression
	// https://web.archive.org/web/20120321182910/http://www.vervestudios.co/projects/compression-tests/results
	//
	// That's why we prefer gzip over deflate. It's just more reliable
	// and not significantly slower than gzip.
	c.SetEncoder("deflate", encoderDeflate)

	// TODO: Exception for old MSIE browsers that can't handle non-HTML?
	// https://zoompf.com/blog/2012/02/lose-the-wait-http-compression
	c.SetEncoder("gzip", encoderGzip)

	// NOTE: Not implemented, intentionally:
	// case "compress": // LZW. Deprecated.
	// case "bzip2":    // Too slow on-the-fly.
	// case "zopfli":   // Too slow on-the-fly.
	// case "xz":       // Too slow on-the-fly.
	return c
}

// SetEncoder can be used to set the implementation of a compression algorithm.
//
// The encoding should be a standardised identifier. See:
// https://developer.mozilla.org/en-US/docs/Web/HTTP/Headers/Accept-Encoding
//
// For example, add the Brotli algortithm:
//
//  import brotli_enc "gopkg.in/kothar/brotli-go.v0/enc"
//
//  compressor := middleware.NewCompressor(5, "text/html")
//  compressor.SetEncoder("br", func(w http.ResponseWriter, level int) io.Writer {
//    params := brotli_enc.NewBrotliParams()
//    params.SetQuality(level)
//    return brotli_enc.NewBrotliWriter(params, w)
//  })
func (c *Compressor) SetEncoder(encoding string, fn EncoderFunc) {
	encoding = strings.ToLower(encoding)
	if encoding == "" {
		panic("the encoding can not be empty")
	}
	if fn == nil {
		panic("attempted to set a nil encoder function")
	}

	// If we are adding a new encoder that is already registered, we have to
	// clear that one out first.
	if _, ok := c.pooledEncoders[encoding]; ok {
		delete(c.pooledEncoders, encoding)
	}
	if _, ok := c.encoders[encoding]; ok {
		delete(c.encoders, encoding)
	}

	// If the encoder supports Resetting (IoReseterWriter), then it can be pooled.
	encoder := fn(ioutil.Discard, c.level)
	if encoder != nil {
		if _, ok := encoder.(ioResetterWriter); ok {
			pool := &sync.Pool{
				New: func() interface{} {
					return fn(ioutil.Discard, c.level)
				},
			}
			c.pooledEncoders[encoding] = pool
		}
	}
	// If the encoder is not in the pooledEncoders, add it to the normal encoders.
	if _, ok := c.pooledEncoders[encoding]; !ok {
		c.encoders[encoding] = fn
	}

	for i, v := range c.encodingPrecedence {
		if v == encoding {
			c.encodingPrecedence = append(c.encodingPrecedence[:i], c.encodingPrecedence[i+1:]...)
		}
	}

	c.encodingPrecedence = append([]string{encoding}, c.encodingPrecedence...)
}

// Handler returns a new middleware that will compress the response based on the
// current Compressor.
func (c *Compressor) Handler() func(next http.Handler) http.Handler {
	return func(next http.Handler) http.Handler {
		fn := func(w http.ResponseWriter, r *http.Request) {
			encoder, encoding, cleanup := c.selectEncoder(r.Header, w)

			cw := &compressResponseWriter{
				ResponseWriter: w,
				w:              w,
				contentTypes:   c.allowedTypes,
				encoding:       encoding,
				compressable:   false, // determined in post-handler
			}
			if encoder != nil {
				cw.w = encoder
			}
			// Re-add the encoder to the pool if applicable.
			defer cleanup()
			defer cw.Close()

			next.ServeHTTP(cw, r)
		}

		return http.HandlerFunc(fn)
	}

}

// selectEncoder returns the encoder, the name of the encoder, and a closer function.
func (c *Compressor) selectEncoder(h http.Header, w io.Writer) (io.Writer, string, func()) {
	header := h.Get("Accept-Encoding")

	// Parse the names of all accepted algorithms from the header.
	accepted := strings.Split(strings.ToLower(header), ",")

	// Find supported encoder by accepted list by precedence
	for _, name := range c.encodingPrecedence {
		if matchAcceptEncoding(accepted, name) {
			if pool, ok := c.pooledEncoders[name]; ok {
				encoder := pool.Get().(ioResetterWriter)
				cleanup := func() {
					pool.Put(encoder)
				}
				encoder.Reset(w)
				return encoder, name, cleanup

			}
			if fn, ok := c.encoders[name]; ok {
				return fn(w, c.level), name, func() {}
			}
		}

	}

	// No encoder found to match the accepted encoding
	return nil, "", func() {}
}

func matchAcceptEncoding(accepted []string, encoding string) bool {
	for _, v := range accepted {
		if strings.Contains(v, encoding) {
			return true
		}
	}
	return false
}

// An EncoderFunc is a function that wraps the provided io.Writer with a
// streaming compression algorithm and returns it.
//
// In case of failure, the function should return nil.
type EncoderFunc func(w io.Writer, level int) io.Writer

// Interface for types that allow resetting io.Writers.
type ioResetterWriter interface {
	io.Writer
	Reset(w io.Writer)
}

// SetEncoder can be used to set the implementation of a compression algorithm.
//
// The encoding should be a standardised identifier. See:
// https://developer.mozilla.org/en-US/docs/Web/HTTP/Headers/Accept-Encoding
//
// For example, add the Brotli algortithm:
//
//  import brotli_enc "gopkg.in/kothar/brotli-go.v0/enc"
//
//  middleware.SetEncoder("br", func(w http.ResponseWriter, level int) io.Writer {
//    params := brotli_enc.NewBrotliParams()
//    params.SetQuality(level)
//    return brotli_enc.NewBrotliWriter(params, w)
//  })
//
//  DEPRECATED
func SetEncoder(encoding string, fn EncoderFunc) {
	if defaultCompressor == nil {
		panic("no compressor to set encoders on. Call Compress() first")
	}
	defaultCompressor.SetEncoder(encoding, fn)
}

// DefaultCompress is a middleware that compresses response
// body of predefined content types to a data format based
// on Accept-Encoding request header. It uses a default
// compression level.
// DEPRECATED
func DefaultCompress(next http.Handler) http.Handler {
	return Compress(flate.DefaultCompression)(next)
}

// Compress is a middleware that compresses response
// body of a given content types to a data format based
// on Accept-Encoding request header. It uses a given
// compression level.
//
// NOTE: make sure to set the Content-Type header on your response
// otherwise this middleware will not compress the response body. For ex, in
// your handler you should set w.Header().Set("Content-Type", http.DetectContentType(yourBody))
// or set it manually.
//
// DEPRECATED
func Compress(level int, types ...string) func(next http.Handler) http.Handler {
	defaultCompressor = NewCompressor(level, types...)
	return defaultCompressor.Handler()
}

type compressResponseWriter struct {
	http.ResponseWriter

	// The streaming encoder writer to be used if there is one. Otherwise,
	// this is just the normal writer.
	w            io.Writer
	encoding     string
	contentTypes map[string]bool
	wroteHeader  bool
	compressable bool
}

func (cw *compressResponseWriter) WriteHeader(code int) {
	if cw.wroteHeader {
		return
	}
	cw.wroteHeader = true
	defer cw.ResponseWriter.WriteHeader(code)

	// Already compressed data?
	if cw.Header().Get("Content-Encoding") != "" {
		return
	}

	// Parse the first part of the Content-Type response header.
	contentType := cw.Header().Get("Content-Type")
	if idx := strings.Index(contentType, ";"); idx >= 0 {
		contentType = contentType[0:idx]
	}

	// Is the content type compressable?
	if _, ok := cw.contentTypes[contentType]; !ok {
		cw.compressable = false
		return
	}

	if cw.encoding != "" {
		cw.compressable = true
		cw.Header().Set("Content-Encoding", cw.encoding)
		cw.Header().Set("Vary", "Accept-Encoding")

		// The content-length after compression is unknown
		cw.Header().Del("Content-Length")
	}
}

func (cw *compressResponseWriter) Write(p []byte) (int, error) {
	if !cw.wroteHeader {
		cw.WriteHeader(http.StatusOK)
	}

	return cw.writer().Write(p)
}

func (cw *compressResponseWriter) writer() io.Writer {
	if cw.compressable {
		return cw.w
	} else {
		return cw.ResponseWriter
	}
}

func (cw *compressResponseWriter) Flush() {
	if f, ok := cw.writer().(http.Flusher); ok {
		f.Flush()
	}
}

func (cw *compressResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
	if hj, ok := cw.writer().(http.Hijacker); ok {
		return hj.Hijack()
	}
	return nil, nil, errors.New("chi/middleware: http.Hijacker is unavailable on the writer")
}

func (cw *compressResponseWriter) Push(target string, opts *http.PushOptions) error {
	if ps, ok := cw.writer().(http.Pusher); ok {
		return ps.Push(target, opts)
	}
	return errors.New("chi/middleware: http.Pusher is unavailable on the writer")
}

func (cw *compressResponseWriter) Close() error {
	if c, ok := cw.writer().(io.WriteCloser); ok {
		return c.Close()
	}
	return errors.New("chi/middleware: io.WriteCloser is unavailable on the writer")
}

func encoderGzip(w io.Writer, level int) io.Writer {
	gw, err := gzip.NewWriterLevel(w, level)
	if err != nil {
		return nil
	}
	return gw
}

func encoderDeflate(w io.Writer, level int) io.Writer {
	dw, err := flate.NewWriter(w, level)
	if err != nil {
		return nil
	}
	return dw
}
<|endoftext|>