<commit_before>package requester

import (
	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	"sync"
	"time"

	"github.com/integralist/go-requester/aggregator"
	"github.com/integralist/go-requester/config"
	"gopkg.in/yaml.v2"
)

type component struct {
	ID        string `yaml:"id"`
	URL       string `yaml:"url"`
	Mandatory bool   `yaml:"mandatory"`
}

type componentsList struct {
	Components []component `yaml:"components"`
}

func checkError(err error) int {
	if e, ok := err.(net.Error); ok && e.Timeout() {
		return 408
	}
	return 500
}

func getSummary(status int) string {
	if status == 200 || status == 304 {
		return "success"
	}
	return "failure"
}

func getComponent(wg *sync.WaitGroup, client *http.Client, i int, v component, ch chan aggregator.ComponentResponse) {
	defer wg.Done()

	resp, err := client.Get(v.URL)

	if err != nil {
		fmt.Printf("Problem getting the response: %s\n\n", err)
		status := checkError(err)

		ch <- aggregator.ComponentResponse{
			ID:        v.ID,
			Status:    status,
			Body:      err.Error(),
			Summary:   getSummary(status),
			Mandatory: v.Mandatory,
		}
	} else {
		defer resp.Body.Close()
		contents, err := ioutil.ReadAll(resp.Body)
		if err != nil {
			fmt.Printf("Problem reading the body for %s -> %s\n", v.ID, err)
		}

		ch <- aggregator.ComponentResponse{
			ID:        v.ID,
			Status:    resp.StatusCode,
			Body:      string(contents),
			Summary:   getSummary(resp.StatusCode),
			Mandatory: v.Mandatory,
		}
	}
}

// Process function kick starts the concurrent requesting of components
func Process(w http.ResponseWriter, r *http.Request, configPath string) {
	var cr []aggregator.ComponentResponse
	var components componentsList

	ch := make(chan aggregator.ComponentResponse)

	config := config.Parse(configPath)
	yaml.Unmarshal(config, &components)

	timeout := time.Duration(1 * time.Second)
	client := http.Client{
		Timeout: timeout,
	}

	var wg sync.WaitGroup
	for i, v := range components.Components {
		wg.Add(1)
		go getComponent(&wg, &client, i, v, ch)
		cr = append(cr, <-ch)
	}
	wg.Wait()

	response, err := aggregator.Process(cr)
	if err != nil {
		fmt.Printf("There was an error aggregating a response: %s", err.Error())
		return
	}

	w.Header().Set("Content-Type", "application/json")
	w.Write(response)
}
<commit_msg>Refactor to improve performance<commit_after>package requester

import (
	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	"sync"
	"time"

	"github.com/integralist/go-requester/aggregator"
	"github.com/integralist/go-requester/config"
	"gopkg.in/yaml.v2"
)

var wg sync.WaitGroup
var ch chan aggregator.ComponentResponse

type component struct {
	ID        string `yaml:"id"`
	URL       string `yaml:"url"`
	Mandatory bool   `yaml:"mandatory"`
}

type componentsList struct {
	Components []component `yaml:"components"`
}

func checkError(err error) int {
	if e, ok := err.(net.Error); ok && e.Timeout() {
		return 408
	}
	return 500
}

func getSummary(status int) string {
	if status == 200 || status == 304 {
		return "success"
	}
	return "failure"
}

func getComponent(i int, v component) {
	defer wg.Done()

	timeout := time.Duration(1 * time.Second)
	client := http.Client{
		Timeout: timeout,
	}
	resp, err := client.Get(v.URL)

	if err != nil {
		fmt.Printf("Problem getting the response: %s\n\n", err)
		status := checkError(err)

		ch <- aggregator.ComponentResponse{
			ID:        v.ID,
			Status:    status,
			Body:      err.Error(),
			Summary:   getSummary(status),
			Mandatory: v.Mandatory,
		}
	} else {
		defer resp.Body.Close()

		contents, err := ioutil.ReadAll(resp.Body)
		if err != nil {
			fmt.Printf("Problem reading the body for %s -> %s\n", v.ID, err)
		}

		ch <- aggregator.ComponentResponse{
			ID:        v.ID,
			Status:    resp.StatusCode,
			Body:      string(contents),
			Summary:   getSummary(resp.StatusCode),
			Mandatory: v.Mandatory,
		}
	}
}

// Process function kick starts the concurrent requesting of components
func Process(w http.ResponseWriter, r *http.Request, configPath string) {
	var cr []aggregator.ComponentResponse
	var components componentsList

	config := config.Parse(configPath)
	yaml.Unmarshal(config, &components)

	ch = make(chan aggregator.ComponentResponse, len(components.Components))

	for i, v := range components.Components {
		wg.Add(1)
		go getComponent(i, v)
	}
	wg.Wait()
	close(ch)

	for component := range ch {
		cr = append(cr, component)
	}

	response, err := aggregator.Process(cr)
	if err != nil {
		fmt.Printf("There was an error aggregating a response: %s", err.Error())
		return
	}

	w.Header().Set("Content-Type", "application/json")
	w.Write(response)
}
<|endoftext|>


<commit_before>package hostdb

import (
	"testing"
	"time"

	"github.com/NebulousLabs/Sia/modules"
	"github.com/NebulousLabs/Sia/types"
)

// TestInsertHost tests the insertHost method, which also depends on the
// scanHostEntry method.
func TestInsertHost(t *testing.T) {
	hdb := bareHostDB()

	// invalid host should not be scanned
	var dbe modules.HostDBEntry
	dbe.NetAddress = "foo"
	hdb.insertHost(dbe)
	select {
	case <-hdb.scanPool:
		t.Error("invalid host was added to scan pool")
	case <-time.After(100 * time.Millisecond):
	}

	// valid host should be scanned
	dbe.NetAddress = "foo.com:1234"
	hdb.insertHost(dbe)
	select {
	case <-hdb.scanPool:
	case <-time.After(time.Second):
		t.Error("host was not scanned")
	}

	// duplicate host should not be scanned
	hdb.insertHost(dbe)
	select {
	case <-hdb.scanPool:
		t.Error("duplicate host was added to scan pool")
	case <-time.After(100 * time.Millisecond):
	}
}

// TestActiveHosts tests the ActiveHosts method.
func TestActiveHosts(t *testing.T) {
	hdb := bareHostDB()

	// empty
	if hosts := hdb.ActiveHosts(); len(hosts) != 0 {
		t.Errorf("wrong number of hosts: expected %v, got %v", 0, len(hosts))
	}

	// with one host
	h1 := new(hostEntry)
	h1.NetAddress = "foo"
	h1.Weight = types.NewCurrency64(1)
	h1.AcceptingContracts = true
	hdb.insertNode(h1)
	if hosts := hdb.ActiveHosts(); len(hosts) != 1 {
		t.Errorf("wrong number of hosts: expected %v, got %v", 1, len(hosts))
	} else if hosts[0].NetAddress != h1.NetAddress {
		t.Errorf("ActiveHosts returned wrong host: expected %v, got %v", h1.NetAddress, hosts[0].NetAddress)
	}

	// with multiple hosts
	h2 := new(hostEntry)
	h2.NetAddress = "bar"
	h2.Weight = types.NewCurrency64(1)
	h2.AcceptingContracts = true
	hdb.insertNode(h2)
	if hosts := hdb.ActiveHosts(); len(hosts) != 2 {
		t.Errorf("wrong number of hosts: expected %v, got %v", 2, len(hosts))
	} else if hosts[0].NetAddress != h1.NetAddress && hosts[1].NetAddress != h1.NetAddress {
		t.Errorf("ActiveHosts did not contain an inserted host: %v (missing %v)", hosts, h1.NetAddress)
	} else if hosts[0].NetAddress != h2.NetAddress && hosts[1].NetAddress != h2.NetAddress {
		t.Errorf("ActiveHosts did not contain an inserted host: %v (missing %v)", hosts, h2.NetAddress)
	}
}

// TestAverageContractPrice tests the AverageContractPrice method, which also depends on the
// randomHosts method.
func TestAverageContractPrice(t *testing.T) {
	hdb := bareHostDB()

	// empty
	if avg := hdb.AverageContractPrice(); !avg.IsZero() {
		t.Error("average of empty hostdb should be zero:", avg)
	}

	// with one host
	h1 := new(hostEntry)
	h1.NetAddress = "foo"
	h1.ContractPrice = types.NewCurrency64(100)
	h1.Weight = baseWeight
	h1.AcceptingContracts = true
	hdb.insertNode(h1)
	if avg := hdb.AverageContractPrice(); avg.Cmp(h1.ContractPrice) != 0 {
		t.Error("average of one host should be that host's price:", avg)
	}

	// with two hosts
	h2 := new(hostEntry)
	h2.NetAddress = "bar"
	h2.ContractPrice = types.NewCurrency64(300)
	h2.Weight = baseWeight
	h2.AcceptingContracts = true
	hdb.insertNode(h2)
	if len(hdb.activeHosts) != 2 {
		t.Error("host was not added:", hdb.activeHosts)
	}
	if avg := hdb.AverageContractPrice(); avg.Cmp64(200) != 0 {
		t.Error("average of two hosts should be their sum/2:", avg)
	}
}

// TestIsOffline tests the IsOffline method.
func TestIsOffline(t *testing.T) {
	now := time.Now()
	hdb := &HostDB{
		allHosts: map[modules.NetAddress]*hostEntry{
			"foo.com:1234": {LastScanned: now, LastSeen: now},
			"bar.com:1234": {LastScanned: now, LastSeen: now.Add(-uptimeThreshold)},
			"baz.com:1234": {LastScanned: now, LastSeen: now.Add(-uptimeThreshold / 2)},
		},
		activeHosts: map[modules.NetAddress]*hostNode{
			"foo.com:1234": nil,
		},
		scanPool: make(chan *hostEntry),
	}

	tests := []struct {
		addr    modules.NetAddress
		offline bool
	}{
		{"foo.com:1234", false},
		{"bar.com:1234", true},
		{"baz.com:1234", false},
		{"quux.com:1234", false},
	}
	for _, test := range tests {
		if offline := hdb.IsOffline(test.addr); offline != test.offline {
			t.Errorf("IsOffline(%v) = %v, expected %v", test.addr, offline, test.offline)
		}
	}
}
<commit_msg>pad offline test time<commit_after>package hostdb

import (
	"testing"
	"time"

	"github.com/NebulousLabs/Sia/modules"
	"github.com/NebulousLabs/Sia/types"
)

// TestInsertHost tests the insertHost method, which also depends on the
// scanHostEntry method.
func TestInsertHost(t *testing.T) {
	hdb := bareHostDB()

	// invalid host should not be scanned
	var dbe modules.HostDBEntry
	dbe.NetAddress = "foo"
	hdb.insertHost(dbe)
	select {
	case <-hdb.scanPool:
		t.Error("invalid host was added to scan pool")
	case <-time.After(100 * time.Millisecond):
	}

	// valid host should be scanned
	dbe.NetAddress = "foo.com:1234"
	hdb.insertHost(dbe)
	select {
	case <-hdb.scanPool:
	case <-time.After(time.Second):
		t.Error("host was not scanned")
	}

	// duplicate host should not be scanned
	hdb.insertHost(dbe)
	select {
	case <-hdb.scanPool:
		t.Error("duplicate host was added to scan pool")
	case <-time.After(100 * time.Millisecond):
	}
}

// TestActiveHosts tests the ActiveHosts method.
func TestActiveHosts(t *testing.T) {
	hdb := bareHostDB()

	// empty
	if hosts := hdb.ActiveHosts(); len(hosts) != 0 {
		t.Errorf("wrong number of hosts: expected %v, got %v", 0, len(hosts))
	}

	// with one host
	h1 := new(hostEntry)
	h1.NetAddress = "foo"
	h1.Weight = types.NewCurrency64(1)
	h1.AcceptingContracts = true
	hdb.insertNode(h1)
	if hosts := hdb.ActiveHosts(); len(hosts) != 1 {
		t.Errorf("wrong number of hosts: expected %v, got %v", 1, len(hosts))
	} else if hosts[0].NetAddress != h1.NetAddress {
		t.Errorf("ActiveHosts returned wrong host: expected %v, got %v", h1.NetAddress, hosts[0].NetAddress)
	}

	// with multiple hosts
	h2 := new(hostEntry)
	h2.NetAddress = "bar"
	h2.Weight = types.NewCurrency64(1)
	h2.AcceptingContracts = true
	hdb.insertNode(h2)
	if hosts := hdb.ActiveHosts(); len(hosts) != 2 {
		t.Errorf("wrong number of hosts: expected %v, got %v", 2, len(hosts))
	} else if hosts[0].NetAddress != h1.NetAddress && hosts[1].NetAddress != h1.NetAddress {
		t.Errorf("ActiveHosts did not contain an inserted host: %v (missing %v)", hosts, h1.NetAddress)
	} else if hosts[0].NetAddress != h2.NetAddress && hosts[1].NetAddress != h2.NetAddress {
		t.Errorf("ActiveHosts did not contain an inserted host: %v (missing %v)", hosts, h2.NetAddress)
	}
}

// TestAverageContractPrice tests the AverageContractPrice method, which also depends on the
// randomHosts method.
func TestAverageContractPrice(t *testing.T) {
	hdb := bareHostDB()

	// empty
	if avg := hdb.AverageContractPrice(); !avg.IsZero() {
		t.Error("average of empty hostdb should be zero:", avg)
	}

	// with one host
	h1 := new(hostEntry)
	h1.NetAddress = "foo"
	h1.ContractPrice = types.NewCurrency64(100)
	h1.Weight = baseWeight
	h1.AcceptingContracts = true
	hdb.insertNode(h1)
	if avg := hdb.AverageContractPrice(); avg.Cmp(h1.ContractPrice) != 0 {
		t.Error("average of one host should be that host's price:", avg)
	}

	// with two hosts
	h2 := new(hostEntry)
	h2.NetAddress = "bar"
	h2.ContractPrice = types.NewCurrency64(300)
	h2.Weight = baseWeight
	h2.AcceptingContracts = true
	hdb.insertNode(h2)
	if len(hdb.activeHosts) != 2 {
		t.Error("host was not added:", hdb.activeHosts)
	}
	if avg := hdb.AverageContractPrice(); avg.Cmp64(200) != 0 {
		t.Error("average of two hosts should be their sum/2:", avg)
	}
}

// TestIsOffline tests the IsOffline method.
func TestIsOffline(t *testing.T) {
	now := time.Now()
	hdb := &HostDB{
		allHosts: map[modules.NetAddress]*hostEntry{
			"foo.com:1234": {LastScanned: now, LastSeen: now},
			"bar.com:1234": {LastScanned: now, LastSeen: now.Add(-uptimeThreshold * 2)},
			"baz.com:1234": {LastScanned: now, LastSeen: now.Add(-uptimeThreshold / 2)},
		},
		activeHosts: map[modules.NetAddress]*hostNode{
			"foo.com:1234": nil,
		},
		scanPool: make(chan *hostEntry),
	}

	tests := []struct {
		addr    modules.NetAddress
		offline bool
	}{
		{"foo.com:1234", false},
		{"bar.com:1234", true},
		{"baz.com:1234", false},
		{"quux.com:1234", false},
	}
	for _, test := range tests {
		if offline := hdb.IsOffline(test.addr); offline != test.offline {
			t.Errorf("IsOffline(%v) = %v, expected %v", test.addr, offline, test.offline)
		}
	}
}
<|endoftext|>


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

// root-ls dumps the content of a ROOT file
package main // import "go-hep.org/x/hep/rootio/cmd/root-ls"

import (
	"bytes"
	"flag"
	"fmt"
	"io"
	"log"
	"os"
	"runtime/pprof"
	"text/tabwriter"

	"go-hep.org/x/hep/rootio"
)

var (
	doProf = flag.String("profile", "", "filename of cpuprofile")
	doSI   = flag.Bool("sinfos", false, "print StreamerInfos")
	doTree = flag.Bool("t", false, "print Tree(s) (recursively)")
)

func main() {
	flag.Parse()

	if *doProf != "" {
		f, err := os.Create(*doProf)
		if err != nil {
			log.Fatal(err)
		}
		pprof.StartCPUProfile(f)
		defer pprof.StopCPUProfile()
	}

	if flag.NArg() <= 0 {
		fmt.Fprintf(os.Stderr, "**error** you need to give a ROOT file\n")
		flag.Usage()
		os.Exit(1)
	}

	for ii, fname := range flag.Args() {

		if ii > 0 {
			fmt.Printf("\n")
		}

		fmt.Printf("=== [%s] ===\n", fname)
		f, err := rootio.Open(fname)
		if err != nil {
			fmt.Fprintf(os.Stderr, "rootio: failed to open [%s]: %v\n", fname, err)
			os.Exit(1)
		}
		fmt.Printf("version: %v\n", f.Version())
		if *doSI {
			fmt.Printf("streamer-infos:\n")
			sinfos := f.StreamerInfo()
			for _, v := range sinfos {
				name := v.Name()
				fmt.Printf(" StreamerInfo for %q version=%d title=%q\n", name, v.ClassVersion(), v.Title())
				w := tabwriter.NewWriter(os.Stdout, 8, 4, 1, ' ', 0)
				for _, elm := range v.Elements() {
					fmt.Fprintf(w, "  %s\t%s\toffset=%3d\ttype=%3d\tsize=%3d\t %s\n", elm.TypeName(), elm.Name(), elm.Offset(), elm.Type(), elm.Size(), elm.Title())
				}
				w.Flush()
			}
		}

		w := tabwriter.NewWriter(os.Stdout, 8, 4, 1, ' ', 0)
		for _, k := range f.Keys() {
			if *doTree {
				tree, ok := k.Value().(rootio.Tree)
				if ok {
					w := tabwriter.NewWriter(w, 8, 4, 1, ' ', 0)
					fmt.Fprintf(w, "%s\t%s\t%s\t(entries=%d)\n", k.Class(), k.Name(), k.Title(), tree.Entries())
					displayBranches(w, tree, 2)
					w.Flush()
					continue
				}
			}
			fmt.Fprintf(w, "%s\t%s\t%s\t(cycle=%d)\n", k.Class(), k.Name(), k.Title(), k.Cycle())
		}
		w.Flush()
	}
}

type windent struct {
	hdr []byte
	w   io.Writer
}

func newWindent(n int, w io.Writer) *windent {
	return &windent{
		hdr: bytes.Repeat([]byte(" "), n),
		w:   w,
	}
}

func (w *windent) Write(data []byte) (int, error) {
	return w.w.Write(append(w.hdr, data...))
}

func (w *windent) Flush() error {
	ww, ok := w.w.(flusher)
	if !ok {
		return nil
	}
	return ww.Flush()
}

type flusher interface {
	Flush() error
}

type brancher interface {
	Branches() []rootio.Branch
}

func displayBranches(w io.Writer, bres brancher, indent int) {
	branches := bres.Branches()
	if len(branches) <= 0 {
		return
	}
	ww := newWindent(indent, w)
	for _, b := range branches {
		fmt.Fprintf(ww, "%s\t%q\t%v\n", b.Name(), b.Title(), b.Class())
		displayBranches(ww, b, 2)
	}
	ww.Flush()
}
<commit_msg>rootio/cmd/root-ls: properly handle subdirs<commit_after>// Copyright 2017 The go-hep Authors.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// root-ls dumps the content of a ROOT file
package main // import "go-hep.org/x/hep/rootio/cmd/root-ls"

import (
	"bytes"
	"flag"
	"fmt"
	"io"
	"log"
	"os"
	"runtime/pprof"
	"text/tabwriter"

	"go-hep.org/x/hep/rootio"
)

var (
	doProf = flag.String("profile", "", "filename of cpuprofile")
	doSI   = flag.Bool("sinfos", false, "print StreamerInfos")
	doTree = flag.Bool("t", false, "print Tree(s) (recursively)")
)

func main() {
	flag.Parse()

	if *doProf != "" {
		f, err := os.Create(*doProf)
		if err != nil {
			log.Fatal(err)
		}
		pprof.StartCPUProfile(f)
		defer pprof.StopCPUProfile()
	}

	if flag.NArg() <= 0 {
		fmt.Fprintf(os.Stderr, "**error** you need to give a ROOT file\n")
		flag.Usage()
		os.Exit(1)
	}

	for ii, fname := range flag.Args() {

		if ii > 0 {
			fmt.Printf("\n")
		}

		fmt.Printf("=== [%s] ===\n", fname)
		f, err := rootio.Open(fname)
		if err != nil {
			fmt.Fprintf(os.Stderr, "rootio: failed to open [%s]: %v\n", fname, err)
			os.Exit(1)
		}
		fmt.Printf("version: %v\n", f.Version())
		if *doSI {
			fmt.Printf("streamer-infos:\n")
			sinfos := f.StreamerInfo()
			for _, v := range sinfos {
				name := v.Name()
				fmt.Printf(" StreamerInfo for %q version=%d title=%q\n", name, v.ClassVersion(), v.Title())
				w := tabwriter.NewWriter(os.Stdout, 8, 4, 1, ' ', 0)
				for _, elm := range v.Elements() {
					fmt.Fprintf(w, "  %s\t%s\toffset=%3d\ttype=%3d\tsize=%3d\t %s\n", elm.TypeName(), elm.Name(), elm.Offset(), elm.Type(), elm.Size(), elm.Title())
				}
				w.Flush()
			}
		}

		w := tabwriter.NewWriter(os.Stdout, 8, 4, 1, ' ', 0)
		for _, k := range f.Keys() {
			walk(w, k)
		}
		w.Flush()
	}
}

func walk(w io.Writer, k rootio.Key) {
	obj := k.Value()
	if *doTree {
		tree, ok := obj.(rootio.Tree)
		if ok {
			w := tabwriter.NewWriter(w, 8, 4, 1, ' ', 0)
			fmt.Fprintf(w, "%s\t%s\t%s\t(entries=%d)\n", k.Class(), k.Name(), k.Title(), tree.Entries())
			displayBranches(w, tree, 2)
			w.Flush()
			return
		}
	}
	fmt.Fprintf(w, "%s\t%s\t%s\t(cycle=%d)\n", k.Class(), k.Name(), k.Title(), k.Cycle())
	if dir, ok := obj.(rootio.Directory); ok {
		w := newWindent(2, w)
		for _, k := range dir.Keys() {
			walk(w, k)
		}
		w.Flush()
	}
}

type windent struct {
	hdr []byte
	w   io.Writer
}

func newWindent(n int, w io.Writer) *windent {
	return &windent{
		hdr: bytes.Repeat([]byte(" "), n),
		w:   w,
	}
}

func (w *windent) Write(data []byte) (int, error) {
	return w.w.Write(append(w.hdr, data...))
}

func (w *windent) Flush() error {
	ww, ok := w.w.(flusher)
	if !ok {
		return nil
	}
	return ww.Flush()
}

type flusher interface {
	Flush() error
}

type brancher interface {
	Branches() []rootio.Branch
}

func displayBranches(w io.Writer, bres brancher, indent int) {
	branches := bres.Branches()
	if len(branches) <= 0 {
		return
	}
	ww := newWindent(indent, w)
	for _, b := range branches {
		fmt.Fprintf(ww, "%s\t%q\t%v\n", b.Name(), b.Title(), b.Class())
		displayBranches(ww, b, 2)
	}
	ww.Flush()
}
<|endoftext|>


<commit_before>package unitconfig

/*
UnitConfig is a struct representation of what is expected to be inside a
Builderfile for a single build/tag/push sequence.
*/
type UnitConfig struct {
	Version          int                 `toml:"version" json:"version" yaml:"version"`
	Docker           Docker              `toml:"docker" json:"docker" yaml:"docker"`
	ContainerArr     []*ContainerSection `toml:"container" json:"container" yaml:"container"`
	ContainerGlobals *ContainerSection   `toml:"container_globals" json:"container_globals" yaml:"container_globals"`
}

/*
Docker is a struct representation of the "docker" section of a Builderfile.
*/
type Docker struct {
	BuildOpts []string `toml:"build_opts" json:"build_opts" yaml:"build_opts"`
	TagOpts   []string `toml:"tag_opts" json:"tag_opts" yaml:"tag_opts"`
}

/*
ContainerSection is a struct representation of an individual member of the  "containers"
section of a Builderfile. Each of these sections defines a docker container to
be built and other related options.
*/
type ContainerSection struct {
	Name       string   `toml:"name" json:"name" yaml:"name"`
	Dockerfile string   `toml:"Dockerfile" json:"Dockerfile" yaml:"Dockerfile"`
	Registry   string   `toml:"registry" json:"registry" yaml:"registry"`
	Project    string   `toml:"project" json:"project" yaml:"project"`
	Tags       []string `toml:"tags" json:"tags" yaml:"tags"`
	SkipPush   bool     `toml:"skip_push" json:"skip_push" yaml:"skip_push"`
	CfgUn      string   `toml:"dockercfg_un" json:"dockercfg_un" yaml:"dockercfg_un"`
	CfgPass    string   `toml:"dockercfg_pass" json:"dockercfg_pass" yaml:"dockercfg_pass"`
	CfgEmail   string   `toml:"dockercfg_email" json:"dockercfg_email" yaml:"dockercfg_email"`
}

// SkipPush sets whether or not pushes should be skipped for the given unit config
func (config *UnitConfig) SkipPush(shouldSkip bool) {
	if config.ContainerGlobals == nil {
		config.ContainerGlobals = &ContainerSection{}
	}
	config.ContainerGlobals.SkipPush = shouldSkip
}
<commit_msg>Adding UnitConfig SetGlobals for setting global values<commit_after>package unitconfig

/*
UnitConfig is a struct representation of what is expected to be inside a
Builderfile for a single build/tag/push sequence.
*/
type UnitConfig struct {
	Version          int                 `toml:"version" json:"version" yaml:"version"`
	Docker           Docker              `toml:"docker" json:"docker" yaml:"docker"`
	ContainerArr     []*ContainerSection `toml:"container" json:"container" yaml:"container"`
	ContainerGlobals *ContainerSection   `toml:"container_globals" json:"container_globals" yaml:"container_globals"`
}

/*
Docker is a struct representation of the "docker" section of a Builderfile.
*/
type Docker struct {
	BuildOpts []string `toml:"build_opts" json:"build_opts" yaml:"build_opts"`
	TagOpts   []string `toml:"tag_opts" json:"tag_opts" yaml:"tag_opts"`
}

/*
ContainerSection is a struct representation of an individual member of the  "containers"
section of a Builderfile. Each of these sections defines a docker container to
be built and other related options.
*/
type ContainerSection struct {
	Name       string   `toml:"name" json:"name" yaml:"name"`
	Dockerfile string   `toml:"Dockerfile" json:"Dockerfile" yaml:"Dockerfile"`
	Registry   string   `toml:"registry" json:"registry" yaml:"registry"`
	Project    string   `toml:"project" json:"project" yaml:"project"`
	Tags       []string `toml:"tags" json:"tags" yaml:"tags"`
	SkipPush   bool     `toml:"skip_push" json:"skip_push" yaml:"skip_push"`
	CfgUn      string   `toml:"dockercfg_un" json:"dockercfg_un" yaml:"dockercfg_un"`
	CfgPass    string   `toml:"dockercfg_pass" json:"dockercfg_pass" yaml:"dockercfg_pass"`
	CfgEmail   string   `toml:"dockercfg_email" json:"dockercfg_email" yaml:"dockercfg_email"`
}

type UnitConfigGlobals struct {
	SkipPush bool
	CfgUn    string
	CfgPass  string
	CfgEmail string
}

// SkipPush sets whether or not pushes should be skipped for the given unit config
func (config *UnitConfig) SetGlobals(globals UnitConfigGlobals) {
	if config.ContainerGlobals == nil {
		config.ContainerGlobals = &ContainerSection{}
	}

	if !config.ContainerGlobals.SkipPush {
		config.ContainerGlobals.SkipPush = globals.SkipPush
	}
	if config.ContainerGlobals.CfgUn == "" {
		config.ContainerGlobals.CfgUn = globals.CfgUn
	}
	if config.ContainerGlobals.CfgPass == "" {
		config.ContainerGlobals.CfgPass = globals.CfgPass
	}
	if config.ContainerGlobals.CfgEmail == "" {
		config.ContainerGlobals.CfgEmail = globals.CfgEmail
	}

}
<|endoftext|>


<commit_before>package irmaserver

import (
	"testing"
	"time"

	irma "github.com/privacybydesign/irmago"
	"github.com/privacybydesign/irmago/server"
	"github.com/sirupsen/logrus"
	"github.com/stretchr/testify/require"
)

var logger = logrus.New()

func init() {
	irma.SetLogger(logger)
	logger.Level = logrus.FatalLevel
}

func TestSessionHandlerInvokedOnCancel(t *testing.T) {
	s, err := New(&server.Configuration{Logger: logger})
	require.NoError(t, err)
	defer s.Stop()

	request := irma.NewDisclosureRequest(irma.NewAttributeTypeIdentifier("irma-demo.RU.studentCard.studentID"))

	var handlerInvoked bool
	_, token, _, err := s.StartSession(request, func(result *server.SessionResult) {
		handlerInvoked = true
	})
	require.NoError(t, err)

	require.NoError(t, s.CancelSession(token))
	time.Sleep(100 * time.Millisecond) // give session handler time to run
	require.True(t, handlerInvoked)
}

func TestSessionHandlerInvokedOnTimeout(t *testing.T) {
	s, err := New(&server.Configuration{Logger: logger})
	require.NoError(t, err)
	defer s.Stop()

	request := &irma.ServiceProviderRequest{
		RequestorBaseRequest: irma.RequestorBaseRequest{
			ClientTimeout: 1,
		},
		Request: irma.NewDisclosureRequest(irma.NewAttributeTypeIdentifier("irma-demo.RU.studentCard.studentID")),
	}

	var handlerInvoked bool
	_, _, _, err = s.StartSession(request, func(result *server.SessionResult) {
		handlerInvoked = true
	})
	require.NoError(t, err)

	time.Sleep(2 * time.Second)
	s.sessions.deleteExpired()
	time.Sleep(100 * time.Millisecond) // give session handler time to run

	require.True(t, handlerInvoked)
}

func TestMemoryStoreNoDeadlock(t *testing.T) {
	s, err := New(&server.Configuration{Logger: logger})
	require.NoError(t, err)
	defer s.Stop()

	req, err := server.ParseSessionRequest(`{"request":{"@context":"https://irma.app/ld/request/disclosure/v2","context":"AQ==","nonce":"MtILupG0g0J23GNR1YtupQ==","devMode":true,"disclose":[[[{"type":"test.test.email.email","value":"example@example.com"}]]]}}`)
	require.NoError(t, err)
	session, err := s.newSession(irma.ActionDisclosing, req, nil, "")
	require.NoError(t, err)

	session.Lock()
	deletingCompleted := false
	addingCompleted := false
	// Make sure the deleting continues on completion such that the test itself will not hang.
	defer func() {
		session.Unlock()
		time.Sleep(100 * time.Millisecond)
		require.True(t, deletingCompleted)
	}()

	go func() {
		s.sessions.(*memorySessionStore).deleteExpired()
		deletingCompleted = true
	}()

	// Make sure the goroutine above is running
	time.Sleep(100 * time.Millisecond)

	// Make a new session; this involves adding it to the memory session store.
	go func() {
		_, _ = s.newSession(irma.ActionDisclosing, req, nil, "")
		addingCompleted = true
	}()

	// Check whether the IRMA server doesn't hang
	time.Sleep(100 * time.Millisecond)
	require.True(t, addingCompleted)
	require.False(t, deletingCompleted)
}
<commit_msg>fix failing test after merging (type assertion needed).<commit_after>package irmaserver

import (
	"testing"
	"time"

	irma "github.com/privacybydesign/irmago"
	"github.com/privacybydesign/irmago/server"
	"github.com/sirupsen/logrus"
	"github.com/stretchr/testify/require"
)

var logger = logrus.New()

func init() {
	irma.SetLogger(logger)
	logger.Level = logrus.FatalLevel
}

func TestSessionHandlerInvokedOnCancel(t *testing.T) {
	s, err := New(&server.Configuration{Logger: logger})
	require.NoError(t, err)
	defer s.Stop()

	request := irma.NewDisclosureRequest(irma.NewAttributeTypeIdentifier("irma-demo.RU.studentCard.studentID"))

	var handlerInvoked bool
	_, token, _, err := s.StartSession(request, func(result *server.SessionResult) {
		handlerInvoked = true
	})
	require.NoError(t, err)

	require.NoError(t, s.CancelSession(token))
	time.Sleep(100 * time.Millisecond) // give session handler time to run
	require.True(t, handlerInvoked)
}

func TestSessionHandlerInvokedOnTimeout(t *testing.T) {
	s, err := New(&server.Configuration{Logger: logger})
	require.NoError(t, err)
	defer s.Stop()

	request := &irma.ServiceProviderRequest{
		RequestorBaseRequest: irma.RequestorBaseRequest{
			ClientTimeout: 1,
		},
		Request: irma.NewDisclosureRequest(irma.NewAttributeTypeIdentifier("irma-demo.RU.studentCard.studentID")),
	}

	var handlerInvoked bool
	_, _, _, err = s.StartSession(request, func(result *server.SessionResult) {
		handlerInvoked = true
	})
	require.NoError(t, err)

	time.Sleep(2 * time.Second)
	s.sessions.(*memorySessionStore).deleteExpired()
	time.Sleep(100 * time.Millisecond) // give session handler time to run

	require.True(t, handlerInvoked)
}

func TestMemoryStoreNoDeadlock(t *testing.T) {
	s, err := New(&server.Configuration{Logger: logger})
	require.NoError(t, err)
	defer s.Stop()

	req, err := server.ParseSessionRequest(`{"request":{"@context":"https://irma.app/ld/request/disclosure/v2","context":"AQ==","nonce":"MtILupG0g0J23GNR1YtupQ==","devMode":true,"disclose":[[[{"type":"test.test.email.email","value":"example@example.com"}]]]}}`)
	require.NoError(t, err)
	session, err := s.newSession(irma.ActionDisclosing, req, nil, "")
	require.NoError(t, err)

	session.Lock()
	deletingCompleted := false
	addingCompleted := false
	// Make sure the deleting continues on completion such that the test itself will not hang.
	defer func() {
		session.Unlock()
		time.Sleep(100 * time.Millisecond)
		require.True(t, deletingCompleted)
	}()

	go func() {
		s.sessions.(*memorySessionStore).deleteExpired()
		deletingCompleted = true
	}()

	// Make sure the goroutine above is running
	time.Sleep(100 * time.Millisecond)

	// Make a new session; this involves adding it to the memory session store.
	go func() {
		_, _ = s.newSession(irma.ActionDisclosing, req, nil, "")
		addingCompleted = true
	}()

	// Check whether the IRMA server doesn't hang
	time.Sleep(100 * time.Millisecond)
	require.True(t, addingCompleted)
	require.False(t, deletingCompleted)
}
<|endoftext|>


<commit_before>package terrain

import (
	log "github.com/Sirupsen/logrus"
	"testing"
)

func init() {
	log.SetLevel(log.DebugLevel)
}

func TestTerrainMaxSize(t *testing.T) {
	ter := New(10, 65535)

	if ter.Max != 1025 {
		t.Error("Expected 1025, got", ter.Max)
	}
}

func TestTerrainLayout(t *testing.T) {
	ter := New(10, 65535)

	if len(ter.layout) != 1025 {
		t.Error("Expected length 1025 but got", len(ter.layout))
	}

	for i := range ter.layout {
		if len(ter.layout[i]) != 1025 {
			t.Error("Expected length 1025 but got", len(ter.layout[i]), "on iteration", i)
		}
	}
}

func TestGetWhenInRange(t *testing.T) {
	ter := New(10, 65535)
	ter.layout[22][22] = 3252

	num, err := ter.GetHeight(22, 22)
	if err != nil {
		t.Errorf("Should not have produced an error")
	}

	if num != 3252 {
		t.Errorf("Returned incorrect number")
	}
}

func TestGetWhenXAboveBounds(t *testing.T) {
	ter := New(10, 65535)

	_, err := ter.GetHeight(22222, 22)
	if err == nil {
		t.Errorf("Should have produced an error")
	}
}

func TestGetWhenYAboveBounds(t *testing.T) {
	ter := New(10, 65535)

	_, err := ter.GetHeight(22, 22222)
	if err == nil {
		t.Errorf("Should have produced an error")
	}
}

func TestGetWhenYAtBounds(t *testing.T) {
	ter := New(10, 65535)

	_, err := ter.GetHeight(22, 1024)
	if err != nil {
		t.Errorf("Should not have produced an error")
	}
}
<commit_msg>Fixing tests to adhere to new single layout<commit_after>package terrain

import (
	log "github.com/Sirupsen/logrus"
	"testing"
)

func init() {
	log.SetLevel(log.DebugLevel)
}

func TestTerrainMaxSize(t *testing.T) {
	ter := New(10, 65535)

	if ter.Max != 1025 {
		t.Error("Expected 1025, got", ter.Max)
	}
}

func TestTerrainLayout(t *testing.T) {
	ter := New(10, 65535)

	if len(ter.layout) != 1025*1025 {
		t.Error("Expected length 1025 but got", len(ter.layout))
	}
}

func TestGetWhenInRange(t *testing.T) {
	ter := New(10, 65535)
	ter.SetHeight(22, 22, 3252)
	num, err := ter.GetHeight(22, 22)
	if err != nil {
		t.Errorf("Should not have produced an error")
	}

	if num != 3252 {
		t.Errorf("Returned incorrect number")
	}
}

func TestGetWhenXAboveBounds(t *testing.T) {
	ter := New(10, 65535)

	_, err := ter.GetHeight(22222, 22)
	if err == nil {
		t.Errorf("Should have produced an error")
	}
}

func TestGetWhenYAboveBounds(t *testing.T) {
	ter := New(10, 65535)

	_, err := ter.GetHeight(22, 22222)
	if err == nil {
		t.Errorf("Should have produced an error")
	}
}

func TestGetWhenYAtBounds(t *testing.T) {
	ter := New(10, 65535)

	_, err := ter.GetHeight(22, 1024)
	if err != nil {
		t.Errorf("Should not have produced an error")
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"net/http"
	"net/http/httptest"
	"net/url"
	"testing"
	"time"

	. "github.com/aandryashin/matchers"
	. "github.com/aandryashin/matchers/httpresp"
)

var (
	srv *httptest.Server
)

func hostport(u string) string {
	uri, _ := url.Parse(u)
	return uri.Host
}

func root(p string) string {
	return fmt.Sprintf("%s%s", srv.URL, p)
}

func peek() (host string) {
	select {
	case host = <-hosts:
	default:
	}
	return host
}

func init() {
	srv = httptest.NewServer(handlers())
	listen = hostport(srv.URL)
}

func TestNewSessionWithGet(t *testing.T) {
	rsp, err := http.Get(root("/wd/hub/session"))
	AssertThat(t, err, Is{nil})
	AssertThat(t, rsp, Code{http.StatusMethodNotAllowed})
}

func TestNewSessionNotFound(t *testing.T) {
	queue(stringSlice{":4444"})
	rsp, err := http.Post(root("/wd/hub/session/123"), "", nil)
	AssertThat(t, err, Is{nil})
	AssertThat(t, rsp, Code{http.StatusNotFound})
}

func TestNewSessionHostDown(t *testing.T) {
	queue(stringSlice{":4444"})
	rsp, err := http.Post(root("/wd/hub/session"), "", nil)
	AssertThat(t, err, Is{nil})
	AssertThat(t, rsp, Code{http.StatusInternalServerError})

	host := peek()
	AssertThat(t, host, EqualTo{":4444"})
}

func TestNewSessionBadGateway(t *testing.T) {
	mux := http.NewServeMux()
	mux.HandleFunc("/wd/hub/session", func(w http.ResponseWriter, r *http.Request) {
		http.Error(w, "", http.StatusBadRequest)
	})
	driver := httptest.NewServer(mux)
	defer driver.Close()

	queue(stringSlice{hostport(driver.URL)})
	rsp, err := http.Post(root("/wd/hub/session"), "", nil)
	AssertThat(t, err, Is{nil})
	AssertThat(t, rsp, Code{http.StatusBadRequest})

	host := peek()
	AssertThat(t, host, EqualTo{hostport(driver.URL)})
}

func TestNewSessionCreated(t *testing.T) {
	mux := http.NewServeMux()
	mux.HandleFunc("/wd/hub/session", func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte(`{"sessionId":"123"}`))
	})
	driver := httptest.NewServer(mux)
	defer driver.Close()

	queue(stringSlice{hostport(driver.URL)})
	rsp, err := http.Post(root("/wd/hub/session"), "", nil)
	AssertThat(t, err, Is{nil})
	AssertThat(t, rsp, Code{http.StatusOK})

	host := peek()
	AssertThat(t, host, EqualTo{""})
}

func TestNewSessionTimeout(t *testing.T) {
	mux := http.NewServeMux()
	mux.HandleFunc("/wd/hub/session", func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte(`{"sessionId":"123"}`))
	})
	driver := httptest.NewServer(mux)
	defer driver.Close()

	timeout = 30 * time.Millisecond
	queue(stringSlice{hostport(driver.URL)})
	rsp, err := http.Post(root("/wd/hub/session"), "", nil)
	AssertThat(t, err, Is{nil})
	AssertThat(t, rsp, Code{http.StatusOK})

	host := peek()
	AssertThat(t, host, EqualTo{""})

	<- time.After(50 * time.Millisecond)
	host = peek()
	AssertThat(t, host, EqualTo{hostport(driver.URL)})
}

func TestProxySessionTimeout(t *testing.T) {
	mux := http.NewServeMux()
	mux.HandleFunc("/wd/hub/session", func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte(`{"sessionId":"123"}`))
	})
		mux.HandleFunc("/wd/hub/session/123", func(w http.ResponseWriter, r *http.Request) {
	})
	driver := httptest.NewServer(mux)
	defer driver.Close()

	timeout = 30 * time.Millisecond
	queue(stringSlice{hostport(driver.URL)})
	rsp, err := http.Post(root("/wd/hub/session"), "", nil)
	AssertThat(t, err, Is{nil})
	AssertThat(t, rsp, Code{http.StatusOK})

	host := peek()
	AssertThat(t, host, EqualTo{""})

	<-time.After(10 * time.Millisecond)
	rsp, err = http.Post(root("/wd/hub/session/123"), "", nil)
	AssertThat(t, err, Is{nil})
	AssertThat(t, rsp, Code{http.StatusOK})

	<- time.After(50 * time.Millisecond)
	host = peek()
	AssertThat(t, host, EqualTo{hostport(driver.URL)})
}

func TestUseSession(t *testing.T) {
	mux := http.NewServeMux()
	mux.HandleFunc("/wd/hub/session", func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte(`{"sessionId":"123"}`))
	})
	mux.HandleFunc("/wd/hub/session/123", func(w http.ResponseWriter, r *http.Request) {
	})
	driver := httptest.NewServer(mux)
	defer driver.Close()

	queue(stringSlice{hostport(driver.URL)})
	rsp, err := http.Post(root("/wd/hub/session"), "", nil)
	AssertThat(t, err, Is{nil})
	AssertThat(t, rsp, Code{http.StatusOK})

	rsp, err = http.Post(root("/wd/hub/session/123"), "", nil)
	AssertThat(t, err, Is{nil})
	AssertThat(t, rsp, Code{http.StatusOK})
}

func TestDeleteSession(t *testing.T) {
	mux := http.NewServeMux()
	mux.HandleFunc("/wd/hub/session", func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte(`{"sessionId":"123"}`))
	})
	mux.HandleFunc("/wd/hub/session/123", func(w http.ResponseWriter, r *http.Request) {
	})

	driver := httptest.NewServer(mux)
	defer driver.Close()

	queue(stringSlice{hostport(driver.URL)})
	rsp, err := http.Post(root("/wd/hub/session"), "", nil)
	AssertThat(t, err, Is{nil})
	AssertThat(t, rsp, Code{http.StatusOK})

	host := peek()
	AssertThat(t, host, EqualTo{""})

	req, err := http.NewRequest(http.MethodDelete, fmt.Sprintf("http://%s/wd/hub/session/123", listen), nil)
	http.DefaultClient.Do(req)
	AssertThat(t, err, Is{nil})
	AssertThat(t, rsp, Code{http.StatusOK})
}

<commit_msg>go fmt<commit_after>package main

import (
	"fmt"
	"net/http"
	"net/http/httptest"
	"net/url"
	"testing"
	"time"

	. "github.com/aandryashin/matchers"
	. "github.com/aandryashin/matchers/httpresp"
)

var (
	srv *httptest.Server
)

func hostport(u string) string {
	uri, _ := url.Parse(u)
	return uri.Host
}

func root(p string) string {
	return fmt.Sprintf("%s%s", srv.URL, p)
}

func peek() (host string) {
	select {
	case host = <-hosts:
	default:
	}
	return host
}

func init() {
	srv = httptest.NewServer(handlers())
	listen = hostport(srv.URL)
}

func TestNewSessionWithGet(t *testing.T) {
	rsp, err := http.Get(root("/wd/hub/session"))
	AssertThat(t, err, Is{nil})
	AssertThat(t, rsp, Code{http.StatusMethodNotAllowed})
}

func TestNewSessionNotFound(t *testing.T) {
	queue(stringSlice{":4444"})
	rsp, err := http.Post(root("/wd/hub/session/123"), "", nil)
	AssertThat(t, err, Is{nil})
	AssertThat(t, rsp, Code{http.StatusNotFound})
}

func TestNewSessionHostDown(t *testing.T) {
	queue(stringSlice{":4444"})
	rsp, err := http.Post(root("/wd/hub/session"), "", nil)
	AssertThat(t, err, Is{nil})
	AssertThat(t, rsp, Code{http.StatusInternalServerError})

	host := peek()
	AssertThat(t, host, EqualTo{":4444"})
}

func TestNewSessionBadGateway(t *testing.T) {
	mux := http.NewServeMux()
	mux.HandleFunc("/wd/hub/session", func(w http.ResponseWriter, r *http.Request) {
		http.Error(w, "", http.StatusBadRequest)
	})
	driver := httptest.NewServer(mux)
	defer driver.Close()

	queue(stringSlice{hostport(driver.URL)})
	rsp, err := http.Post(root("/wd/hub/session"), "", nil)
	AssertThat(t, err, Is{nil})
	AssertThat(t, rsp, Code{http.StatusBadRequest})

	host := peek()
	AssertThat(t, host, EqualTo{hostport(driver.URL)})
}

func TestNewSessionCreated(t *testing.T) {
	mux := http.NewServeMux()
	mux.HandleFunc("/wd/hub/session", func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte(`{"sessionId":"123"}`))
	})
	driver := httptest.NewServer(mux)
	defer driver.Close()

	queue(stringSlice{hostport(driver.URL)})
	rsp, err := http.Post(root("/wd/hub/session"), "", nil)
	AssertThat(t, err, Is{nil})
	AssertThat(t, rsp, Code{http.StatusOK})

	host := peek()
	AssertThat(t, host, EqualTo{""})
}

func TestNewSessionTimeout(t *testing.T) {
	mux := http.NewServeMux()
	mux.HandleFunc("/wd/hub/session", func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte(`{"sessionId":"123"}`))
	})
	driver := httptest.NewServer(mux)
	defer driver.Close()

	timeout = 30 * time.Millisecond
	queue(stringSlice{hostport(driver.URL)})
	rsp, err := http.Post(root("/wd/hub/session"), "", nil)
	AssertThat(t, err, Is{nil})
	AssertThat(t, rsp, Code{http.StatusOK})

	host := peek()
	AssertThat(t, host, EqualTo{""})

	<-time.After(50 * time.Millisecond)
	host = peek()
	AssertThat(t, host, EqualTo{hostport(driver.URL)})
}

func TestProxySessionTimeout(t *testing.T) {
	mux := http.NewServeMux()
	mux.HandleFunc("/wd/hub/session", func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte(`{"sessionId":"123"}`))
	})
	mux.HandleFunc("/wd/hub/session/123", func(w http.ResponseWriter, r *http.Request) {
	})
	driver := httptest.NewServer(mux)
	defer driver.Close()

	timeout = 30 * time.Millisecond
	queue(stringSlice{hostport(driver.URL)})
	rsp, err := http.Post(root("/wd/hub/session"), "", nil)
	AssertThat(t, err, Is{nil})
	AssertThat(t, rsp, Code{http.StatusOK})

	host := peek()
	AssertThat(t, host, EqualTo{""})

	<-time.After(10 * time.Millisecond)
	rsp, err = http.Post(root("/wd/hub/session/123"), "", nil)
	AssertThat(t, err, Is{nil})
	AssertThat(t, rsp, Code{http.StatusOK})

	<-time.After(50 * time.Millisecond)
	host = peek()
	AssertThat(t, host, EqualTo{hostport(driver.URL)})
}

func TestUseSession(t *testing.T) {
	mux := http.NewServeMux()
	mux.HandleFunc("/wd/hub/session", func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte(`{"sessionId":"123"}`))
	})
	mux.HandleFunc("/wd/hub/session/123", func(w http.ResponseWriter, r *http.Request) {
	})
	driver := httptest.NewServer(mux)
	defer driver.Close()

	queue(stringSlice{hostport(driver.URL)})
	rsp, err := http.Post(root("/wd/hub/session"), "", nil)
	AssertThat(t, err, Is{nil})
	AssertThat(t, rsp, Code{http.StatusOK})

	rsp, err = http.Post(root("/wd/hub/session/123"), "", nil)
	AssertThat(t, err, Is{nil})
	AssertThat(t, rsp, Code{http.StatusOK})
}

func TestDeleteSession(t *testing.T) {
	mux := http.NewServeMux()
	mux.HandleFunc("/wd/hub/session", func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte(`{"sessionId":"123"}`))
	})
	mux.HandleFunc("/wd/hub/session/123", func(w http.ResponseWriter, r *http.Request) {
	})

	driver := httptest.NewServer(mux)
	defer driver.Close()

	queue(stringSlice{hostport(driver.URL)})
	rsp, err := http.Post(root("/wd/hub/session"), "", nil)
	AssertThat(t, err, Is{nil})
	AssertThat(t, rsp, Code{http.StatusOK})

	host := peek()
	AssertThat(t, host, EqualTo{""})

	req, err := http.NewRequest(http.MethodDelete, fmt.Sprintf("http://%s/wd/hub/session/123", listen), nil)
	http.DefaultClient.Do(req)
	AssertThat(t, err, Is{nil})
	AssertThat(t, rsp, Code{http.StatusOK})
}
<|endoftext|>


<commit_before>// Copyright 2017 Authors of Cilium
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package helpers

import (
	"fmt"
	"io"
	"os"
	"os/exec"
	"strings"

	"github.com/cilium/cilium/test/config"
	ginkgoext "github.com/cilium/cilium/test/ginkgo-ext"

	"github.com/onsi/ginkgo"
	"github.com/sirupsen/logrus"
)

// CreateVM creates a new vagrant server.Receives a scope which indicates the
// target server that needs to be created. In case of any error on vagrant
// [provision|up|ssh-config] error will be returned.
func CreateVM(scope string) error {
	createCMD := "vagrant up %s --provision"

	for _, v := range Status(scope) {
		switch v {
		case "running":
			createCMD = "vagrant provision %s"
		case "not_created":
			createCMD = "vagrant up %s --provision"
		default:
			// Sometimes servers are stoped and not destroyed. Destroy VM just in case
			DestroyVM(scope)
		}
	}
	createCMD = fmt.Sprintf(createCMD, scope)
	log.Infof("Vagrant:Create: running '%s'", createCMD)
	cmd := getCmd(createCMD)
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return fmt.Errorf("error getting stdout: %s", err)
	}
	stderr, err := cmd.StderrPipe()
	if err != nil {
		return fmt.Errorf("error getting stderr: %s", err)
	}

	go io.Copy(ginkgoext.GinkgoWriter, stderr)
	go io.Copy(ginkgoext.GinkgoWriter, stdout)

	go io.Copy(ginkgo.GinkgoWriter, stderr)
	go io.Copy(ginkgo.GinkgoWriter, stdout)

	if err := cmd.Start(); err != nil {
		log.WithFields(logrus.Fields{
			"command": createCMD,
			"err":     err,
		}).Fatalf("Create error on start")
		return err
	}
	return cmd.Wait()
}

// GetVagrantSSHMetadata returns a string containing the output of `vagrant ssh-config`
// for the provided Vagrant of name vmName. Returns an error if
// `vagrant ssh-config` fails to execute.
func GetVagrantSSHMetadata(vmName string) ([]byte, error) {
	cmd := getCmd(fmt.Sprintf("vagrant ssh-config %s", vmName))
	if config.CiliumTestConfig.SSHConfig != "" {
		cmd = getCmd(config.CiliumTestConfig.SSHConfig)
	}
	result, err := cmd.Output()
	if err != nil {
		return nil, err
	}
	return result, nil
}

//DestroyVM destroys all running Vagrant VMs in the provided scope. It returns an
//error if deletion of either the VMs fails
func DestroyVM(scope string) error {
	command := fmt.Sprintf("vagrant destroy -f %s ", scope)
	cmd := getCmd(command)
	_, err := cmd.CombinedOutput()
	if err != nil {
		return err
	}
	return nil
}

func getCmd(vmCommand string) *exec.Cmd {
	log.Infof("Vagrant: running command '%v'", vmCommand)
	cmd := exec.Command(getPath("bash"), "-c", vmCommand)
	cmd.Dir = getDir()
	return cmd
}

func getDir() string {
	dir, err := os.Getwd()
	if err != nil {
		return "/tmp/"
	}
	return fmt.Sprintf("%s/", dir)
}

func getPath(prog string) string {
	path, err := exec.LookPath(prog)
	if err != nil {
		return ""
	}
	return path
}

//Status returns a mapping of Vagrant VM name to its status
func Status(key string) map[string]string {
	result := map[string]string{}

	cmd := getCmd(fmt.Sprintf("vagrant status %s --machine-readable", key))
	data, err := cmd.CombinedOutput()
	if err != nil {
		return result
	}
	for _, line := range strings.Split(string(data), "\n") {
		val := strings.Split(line, ",")
		if len(val) > 2 && val[2] == "state" {
			result[val[1]] = val[3]
		}
	}
	return result
}
<commit_msg>Test: Dump Vagrant output to Jenkins console<commit_after>// Copyright 2017 Authors of Cilium
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package helpers

import (
	"fmt"
	"io"
	"os"
	"os/exec"
	"strings"

	"github.com/cilium/cilium/test/config"
	ginkgoext "github.com/cilium/cilium/test/ginkgo-ext"

	"github.com/onsi/ginkgo"
	"github.com/sirupsen/logrus"
)

// CreateVM creates a new vagrant server.Receives a scope which indicates the
// target server that needs to be created. In case of any error on vagrant
// [provision|up|ssh-config] error will be returned.
func CreateVM(scope string) error {
	createCMD := "vagrant up %s --provision"

	for _, v := range Status(scope) {
		switch v {
		case "running":
			createCMD = "vagrant provision %s"
		case "not_created":
			createCMD = "vagrant up %s --provision"
		default:
			// Sometimes servers are stoped and not destroyed. Destroy VM just in case
			DestroyVM(scope)
		}
	}
	createCMD = fmt.Sprintf(createCMD, scope)
	log.Infof("Vagrant:Create: running '%s'", createCMD)
	cmd := getCmd(createCMD)
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return fmt.Errorf("error getting stdout: %s", err)
	}
	stderr, err := cmd.StderrPipe()
	if err != nil {
		return fmt.Errorf("error getting stderr: %s", err)
	}

	globalWriter := ginkgoext.NewWriter(ginkgo.GinkgoWriter)

	go io.Copy(globalWriter, stderr)
	go io.Copy(globalWriter, stdout)

	if err := cmd.Start(); err != nil {
		log.WithFields(logrus.Fields{
			"command": createCMD,
			"err":     err,
		}).Fatalf("Create error on start")
		return err
	}
	result := cmd.Wait()
	io.Copy(ginkgoext.GinkgoWriter, globalWriter.Buffer)
	return result
}

// GetVagrantSSHMetadata returns a string containing the output of `vagrant ssh-config`
// for the provided Vagrant of name vmName. Returns an error if
// `vagrant ssh-config` fails to execute.
func GetVagrantSSHMetadata(vmName string) ([]byte, error) {
	cmd := getCmd(fmt.Sprintf("vagrant ssh-config %s", vmName))
	if config.CiliumTestConfig.SSHConfig != "" {
		cmd = getCmd(config.CiliumTestConfig.SSHConfig)
	}
	result, err := cmd.Output()
	if err != nil {
		return nil, err
	}
	return result, nil
}

//DestroyVM destroys all running Vagrant VMs in the provided scope. It returns an
//error if deletion of either the VMs fails
func DestroyVM(scope string) error {
	command := fmt.Sprintf("vagrant destroy -f %s ", scope)
	cmd := getCmd(command)
	_, err := cmd.CombinedOutput()
	if err != nil {
		return err
	}
	return nil
}

func getCmd(vmCommand string) *exec.Cmd {
	log.Infof("Vagrant: running command '%v'", vmCommand)
	cmd := exec.Command(getPath("bash"), "-c", vmCommand)
	cmd.Dir = getDir()
	return cmd
}

func getDir() string {
	dir, err := os.Getwd()
	if err != nil {
		return "/tmp/"
	}
	return fmt.Sprintf("%s/", dir)
}

func getPath(prog string) string {
	path, err := exec.LookPath(prog)
	if err != nil {
		return ""
	}
	return path
}

//Status returns a mapping of Vagrant VM name to its status
func Status(key string) map[string]string {
	result := map[string]string{}

	cmd := getCmd(fmt.Sprintf("vagrant status %s --machine-readable", key))
	data, err := cmd.CombinedOutput()
	if err != nil {
		return result
	}
	for _, line := range strings.Split(string(data), "\n") {
		val := strings.Split(line, ",")
		if len(val) > 2 && val[2] == "state" {
			result[val[1]] = val[3]
		}
	}
	return result
}
<|endoftext|>


<commit_before>package seqio

import (
	"bufio"
	"bytes"
	"errors"
	"fmt"
	"io"
	"os"
	"regexp"

	"github.com/shenwei356/bio/seq"
)

// FastaRecord struct
type FastaRecord struct {
	ID  []byte
	Seq *seq.Seq
}

// NewFastaRecord is constructor of type *FastaRecord*
func NewFastaRecord(t *seq.Alphabet, id, str []byte) (*FastaRecord, error) {
	sequence, err := seq.NewSeq(t, str)
	if err != nil {
		return nil, fmt.Errorf("%s: %s", err, id)
	}
	return &FastaRecord{id, sequence}, nil
}

// FormatSeq formats output
func (record *FastaRecord) FormatSeq(width int) []byte {
	return seq.FormatSeq(record.Seq.Seq, width)
}

// FastaWriter struct
type FastaWriter struct {
	filename  string
	lineWidth int
}

// NewFastaWriter constructor of type *FastaWriter*
func NewFastaWriter(filename string, lineWidth int) *FastaWriter {
	return &FastaWriter{filename, lineWidth}
}

// Write list of FastaRecords to file
func (writer FastaWriter) Write(records []*FastaRecord) (int, error) {
	fout, err := os.Create(writer.filename)
	defer fout.Close()
	if err != nil {
		return 0, err
	}

	w := bufio.NewWriter(fout)
	n := 0
	for _, record := range records {
		fmt.Fprintf(w, ">%s\n%s", record.ID, record.FormatSeq(writer.lineWidth))
		n++
	}
	w.Flush()
	return n, nil
}

/*FastaReader Usage:

  // Sequence type must be specified.
  fasta, err := seqio.NewFastaReader(seq.RNAredundant, "hairpin.fa")

  // check err, usually caused by a wrong file path.
  if err != nil {
      fmt.Println(err)
      return
  }

  // read and check if more record existed
  for fasta.HasNext() {
       record, err := fasta.NextSeq()
       // check err, the record may contain invalid sequence !!!
       if err != nil {
           fmt.Println(err)
           continue
       }

       // deal with the fasta record
       // format output
       fmt.Printf(">%s\n%s", record.Id, record.FormatSeq(70))
  }

*/
type FastaReader struct {
	t        *seq.Alphabet
	filename string
	nextseq  *FastaRecord

	fh                *os.File
	reader            *bufio.Reader
	buffer            bytes.Buffer
	secondLastHead    []byte // report when NextSeq() return a invalid record
	lastHead          []byte
	hasSeq            bool
	fileHandlerClosed bool
}

// NewFastaReader is constructor of FastaReader. Sequence Type (Alphabet) must be given,
// to validate the sequence.
func NewFastaReader(t *seq.Alphabet, filename string) (*FastaReader, error) {
	fh, err := os.Open(filename)
	if err != nil {
		return nil, errors.New("failed to open file: " + filename)
	}

	fasta := new(FastaReader)
	fasta.t = t
	fasta.filename = filename
	fasta.fh = fh

	fasta.reader = bufio.NewReader(fh)
	fasta.hasSeq = false
	fasta.fileHandlerClosed = false

	return fasta, nil
}

// NextSeq returns the next record
func (fasta *FastaReader) NextSeq() (*FastaRecord, error) {
	if fasta.nextseq == nil {
		err := fmt.Errorf("HasNext() should be called firstly. Or invalid %s sequence: %s",
			fasta.t.Type(), string(fasta.secondLastHead))
		return nil, err
	}
	return fasta.nextseq, nil
}

var reTrimLeftSpace = regexp.MustCompile(`^\s+`)
var reTrimRightSpace = regexp.MustCompile(`[\r\n]+$`)
var reTrimSpace = regexp.MustCompile(`[\r\n\s]+`)

// HasNext parses fasta file to check if more record existed,
// if existed, store it.
func (fasta *FastaReader) HasNext() bool {
	if fasta.fileHandlerClosed {
		return false
	}
	for {
		str, err := fasta.reader.ReadBytes('\n')
		str = reTrimLeftSpace.ReplaceAll(str, []byte("")) // spaces before > are allowed

		if err == io.EOF {
			// do not forget the last line,
			// even which does not ends with "\n"
			fasta.buffer.Write(reTrimSpace.ReplaceAll(str, []byte("")))
			sequence := fasta.buffer.Bytes()
			// fmt.Printf("1 [%p] %s\n", sequence, sequence)
			fasta.buffer.Reset()

			fasta.secondLastHead = fasta.lastHead
			head := fasta.lastHead

			fasta.fh.Close()
			fasta.fileHandlerClosed = true

			// I have to do this to solve the problem that different FastaRecord
			// may point to the same address of records.Seq.Seq, this happened when
			// calling Iterator(). Because sequence initially points the ADDRESS
			// of []byte.
			sequence = []byte(string(sequence))

			fasta.nextseq, err = NewFastaRecord(fasta.t, head, sequence)
			if err != nil {
				fmt.Fprintf(os.Stderr, "error when reading %s: %s\n", head, err)
				os.Exit(1)
				return false
			}
			return true
		}

		if bytes.HasPrefix(str, []byte(">")) {
			fasta.hasSeq = true

			thisHead := reTrimRightSpace.ReplaceAll(str[1:], []byte(""))
			if fasta.buffer.Len() > 0 { // no-first seq head
				sequence := fasta.buffer.Bytes()
				// fmt.Printf("2 [%p] %s\n", sequence, sequence)
				fasta.buffer.Reset()
				fasta.secondLastHead = fasta.lastHead
				head := fasta.lastHead
				fasta.lastHead = thisHead

				sequence = []byte(string(sequence))
				fasta.nextseq, err = NewFastaRecord(fasta.t, head, sequence)
				if err != nil {
					fmt.Fprintf(os.Stderr, "error when reading %s: %s\n", head, err)
					os.Exit(1)
					return false
				}
				return true
			}
			// first sequence head
			fasta.secondLastHead = thisHead
			fasta.lastHead = thisHead
		} else if fasta.hasSeq { // append sequence
			fasta.buffer.Write(reTrimSpace.ReplaceAll(str, []byte("")))
		} else {
			// some line before the first ">"
		}
	}
}

// Iterator returns an fasta record iterator.
// It's Go-ish by using channel!
//
// The arguments buffersize specifies the buffer size of channel.
// If buffersize == 0, it reads one fasta record and waits the
// record being deliveried.
// When buffersize > 0, records are well prepared in the backgound.
// It's useful when reading large sequences (human chromosome) and
// processing of sequence is also time-consuming.
//
// Usage:
//
//		fasta, err := NewFastaReader(seq.Unlimit, "test.fa")
//		if err != nil {
//			t.Error(err)
//			return
//		}
//
//		for record = range fasta.Iterator(10) {
//			fmt.Printf(">%s\n%s", record.Id, record.FormatSeq(70))
//		}
//
//
func (fasta *FastaReader) Iterator(buffersize int) chan *FastaRecord {
	if buffersize < 0 {
		buffersize = 0
	}

	ch := make(chan *FastaRecord, buffersize)
	go func() {
		for fasta.HasNext() {
			record, err := fasta.NextSeq()
			if err != nil {
				fmt.Fprintf(os.Stderr, "read fasta error: %s\n", err)
				close(ch)
				os.Exit(1)
				return
			}
			// fmt.Printf(">>%s\n%s", record.ID, record.FormatSeq(70))
			ch <- record
		}
		close(ch)
	}()
	return ch
}
<commit_msg>fix a bug when no sequences in file<commit_after>package seqio

import (
	"bufio"
	"bytes"
	"errors"
	"fmt"
	"io"
	"os"
	"regexp"

	"github.com/shenwei356/bio/seq"
)

// FastaRecord struct
type FastaRecord struct {
	ID  []byte
	Seq *seq.Seq
}

// NewFastaRecord is constructor of type *FastaRecord*
func NewFastaRecord(t *seq.Alphabet, id, str []byte) (*FastaRecord, error) {
	sequence, err := seq.NewSeq(t, str)
	if err != nil {
		return nil, fmt.Errorf("%s: %s", err, id)
	}
	return &FastaRecord{id, sequence}, nil
}

// FormatSeq formats output
func (record *FastaRecord) FormatSeq(width int) []byte {
	return seq.FormatSeq(record.Seq.Seq, width)
}

// FastaWriter struct
type FastaWriter struct {
	filename  string
	lineWidth int
}

// NewFastaWriter constructor of type *FastaWriter*
func NewFastaWriter(filename string, lineWidth int) *FastaWriter {
	return &FastaWriter{filename, lineWidth}
}

// Write list of FastaRecords to file
func (writer FastaWriter) Write(records []*FastaRecord) (int, error) {
	fout, err := os.Create(writer.filename)
	defer fout.Close()
	if err != nil {
		return 0, err
	}

	w := bufio.NewWriter(fout)
	n := 0
	for _, record := range records {
		fmt.Fprintf(w, ">%s\n%s", record.ID, record.FormatSeq(writer.lineWidth))
		n++
	}
	w.Flush()
	return n, nil
}

/*FastaReader Usage:

  // Sequence type must be specified.
  fasta, err := seqio.NewFastaReader(seq.RNAredundant, "hairpin.fa")

  // check err, usually caused by a wrong file path.
  if err != nil {
      fmt.Println(err)
      return
  }

  // read and check if more record existed
  for fasta.HasNext() {
       record, err := fasta.NextSeq()
       // check err, the record may contain invalid sequence !!!
       if err != nil {
           fmt.Println(err)
           continue
       }

       // deal with the fasta record
       // format output
       fmt.Printf(">%s\n%s", record.Id, record.FormatSeq(70))
  }

*/
type FastaReader struct {
	t        *seq.Alphabet
	filename string
	nextseq  *FastaRecord

	fh                *os.File
	reader            *bufio.Reader
	buffer            bytes.Buffer
	secondLastHead    []byte // report when NextSeq() return a invalid record
	lastHead          []byte
	hasSeq            bool
	fileHandlerClosed bool
}

// NewFastaReader is constructor of FastaReader. Sequence Type (Alphabet) must be given,
// to validate the sequence.
func NewFastaReader(t *seq.Alphabet, filename string) (*FastaReader, error) {
	fh, err := os.Open(filename)
	if err != nil {
		return nil, errors.New("failed to open file: " + filename)
	}

	fasta := new(FastaReader)
	fasta.t = t
	fasta.filename = filename
	fasta.fh = fh

	fasta.reader = bufio.NewReader(fh)
	fasta.hasSeq = false
	fasta.fileHandlerClosed = false

	return fasta, nil
}

// NextSeq returns the next record
func (fasta *FastaReader) NextSeq() (*FastaRecord, error) {
	if fasta.nextseq == nil {
		err := fmt.Errorf("HasNext() should be called firstly. Or invalid %s sequence: %s",
			fasta.t.Type(), string(fasta.secondLastHead))
		return nil, err
	}
	return fasta.nextseq, nil
}

var reTrimLeftSpace = regexp.MustCompile(`^\s+`)
var reTrimRightSpace = regexp.MustCompile(`[\r\n]+$`)
var reTrimSpace = regexp.MustCompile(`[\r\n\s]+`)

// HasNext parses fasta file to check if more record existed,
// if existed, store it.
func (fasta *FastaReader) HasNext() bool {
	if fasta.fileHandlerClosed {
		return false
	}
	for {
		str, err := fasta.reader.ReadBytes('\n')
		str = reTrimLeftSpace.ReplaceAll(str, []byte("")) // spaces before > are allowed

		if err == io.EOF {
			// do not forget the last line,
			// even which does not ends with "\n"
			fasta.buffer.Write(reTrimSpace.ReplaceAll(str, []byte("")))
			sequence := fasta.buffer.Bytes()
			// fmt.Printf("1 [%p] %s\n", sequence, sequence)
			fasta.buffer.Reset()

			fasta.secondLastHead = fasta.lastHead
			head := fasta.lastHead

			fasta.fh.Close()
			fasta.fileHandlerClosed = true

			if len(head) == 0 && len(sequence) == 0 {
				return false
			}

			// I have to do this to solve the problem that different FastaRecord
			// may point to the same address of records.Seq.Seq, this happened when
			// calling Iterator(). Because sequence initially points the ADDRESS
			// of []byte.
			sequence = []byte(string(sequence))

			fasta.nextseq, err = NewFastaRecord(fasta.t, head, sequence)
			if err != nil {
				fmt.Fprintf(os.Stderr, "error when reading %s: %s\n", head, err)
				os.Exit(1)
				return false
			}
			return true
		}

		if bytes.HasPrefix(str, []byte(">")) {
			fasta.hasSeq = true

			thisHead := reTrimRightSpace.ReplaceAll(str[1:], []byte(""))
			if fasta.buffer.Len() > 0 { // no-first seq head
				sequence := fasta.buffer.Bytes()
				// fmt.Printf("2 [%p] %s\n", sequence, sequence)
				fasta.buffer.Reset()
				fasta.secondLastHead = fasta.lastHead
				head := fasta.lastHead
				fasta.lastHead = thisHead

				sequence = []byte(string(sequence))
				fasta.nextseq, err = NewFastaRecord(fasta.t, head, sequence)
				if err != nil {
					fmt.Fprintf(os.Stderr, "error when reading %s: %s\n", head, err)
					os.Exit(1)
					return false
				}
				return true
			}
			// first sequence head
			fasta.secondLastHead = thisHead
			fasta.lastHead = thisHead
		} else if fasta.hasSeq { // append sequence
			fasta.buffer.Write(reTrimSpace.ReplaceAll(str, []byte("")))
		} else {
			// some line before the first ">"
		}
	}
}

// Iterator returns an fasta record iterator.
// It's Go-ish by using channel!
//
// The arguments buffersize specifies the buffer size of channel.
// If buffersize == 0, it reads one fasta record and waits the
// record being deliveried.
// When buffersize > 0, records are well prepared in the backgound.
// It's useful when reading large sequences (human chromosome) and
// processing of sequence is also time-consuming.
//
// Usage:
//
//		fasta, err := NewFastaReader(seq.Unlimit, "test.fa")
//		if err != nil {
//			t.Error(err)
//			return
//		}
//
//		for record = range fasta.Iterator(10) {
//			fmt.Printf(">%s\n%s", record.Id, record.FormatSeq(70))
//		}
//
//
func (fasta *FastaReader) Iterator(buffersize int) chan *FastaRecord {
	if buffersize < 0 {
		buffersize = 0
	}

	ch := make(chan *FastaRecord, buffersize)
	go func() {
		for fasta.HasNext() {
			record, err := fasta.NextSeq()
			if err != nil {
				fmt.Fprintf(os.Stderr, "read fasta error: %s\n", err)
				close(ch)
				os.Exit(1)
				return
			}
			// fmt.Printf(">>%s\n%s", record.ID, record.FormatSeq(70))
			ch <- record
		}
		close(ch)
	}()
	return ch
}
<|endoftext|>


<commit_before>// Copyright 2012 Apcera Inc. All rights reserved.

package server

import (
	"bufio"
	"encoding/json"
	"fmt"
	"math/rand"
	"net"
	"sync"
	"time"

	"github.com/apcera/gnatsd/hashmap"
	"github.com/apcera/gnatsd/sublist"
)

// The size of the bufio reader/writer on top of the socket.
const defaultBufSize = 32768

type client struct {
	mu   sync.Mutex
	cid  uint64
	opts clientOpts
	conn net.Conn
	bw   *bufio.Writer
	srv  *Server
	subs *hashmap.HashMap
	pcd  map[*client]struct{}
	atmr *time.Timer
	cstats
	parseState
}

func (c *client) String() string {
	return fmt.Sprintf("cid:%d", c.cid)
}

type cstats struct {
	nr int
	nb int
	nm int
}

type subscription struct {
	client  *client
	subject []byte
	queue   []byte
	sid     []byte
	nm      int64
	max     int64
}

type clientOpts struct {
	Verbose       bool   `json:"verbose"`
	Pedantic      bool   `json:"pedantic"`
	SslRequired   bool   `json:"ssl_required"`
	Authorization string `json:"auth_token"`
	Username      string `json:"user"`
	Password      string `json:"pass"`
	Name          string `json:"name"`
}

var defaultOpts = clientOpts{Verbose: true, Pedantic: true}

func init() {
	rand.Seed(time.Now().UnixNano())
}

func (c *client) readLoop() {
	b := make([]byte, defaultBufSize)
	for {
		n, err := c.conn.Read(b)
		if err != nil {
			c.closeConnection()
			return
		}
		if err := c.parse(b[:n]); err != nil {
			Log(err.Error(), clientConnStr(c.conn), c.cid)
			c.closeConnection()
			return
		}
		// Check pending clients for flush.
		for cp, _ := range c.pcd {
			// Flush those in the set
			cp.mu.Lock()
			err := cp.bw.Flush()
			cp.mu.Unlock()
			if err != nil {
				// FIXME, close connection?
				Logf("Error flushing client connection: %v\n", err)
			}
			delete(c.pcd, cp)
		}
	}
}

func (c *client) traceMsg(msg []byte) {
	pm := fmt.Sprintf("Processing msg: %d", c.nm)
	opa := []interface{}{pm, string(c.pa.subject), string(c.pa.reply), string(msg)}
	Trace(logStr(opa), fmt.Sprintf("c: %d", c.cid))
}

func (c *client) traceOp(op string, arg []byte) {
	if !trace {
		return
	}
	opa := []interface{}{fmt.Sprintf("%s OP", op)}
	if arg != nil {
		opa = append(opa, fmt.Sprintf("%s %s", op, string(arg)))
	}
	Trace(logStr(opa), fmt.Sprintf("c: %d", c.cid))
}

func (c *client) processConnect(arg []byte) error {
	c.traceOp("CONNECT", arg)

	// This will be resolved regardless before we exit this func,
	// so we can just clear it here.
	if c.atmr != nil {
		c.atmr.Stop()
		c.atmr = nil
	}

	// FIXME, check err
	if err := json.Unmarshal(arg, &c.opts); err != nil {
		return err
	}
	// Check for Auth
	if c.srv != nil {
		if ok := c.srv.checkAuth(c); !ok {
			c.sendErr("Authorization is Required")
			return fmt.Errorf("Authorization Error")
		}
	}
	if c.opts.Verbose {
		c.sendOK()
	}
	return nil
}

func (c *client) authViolation() {
	c.sendErr("Authorization is Required")
	fmt.Printf("AUTH TIMER EXPIRED!!\n")
	c.closeConnection()
}

func (c *client) sendErr(err string) {
	c.mu.Lock()
	c.bw.WriteString(fmt.Sprintf("-ERR '%s'\r\n", err))
	c.pcd[c] = needFlush
	c.mu.Unlock()
}

func (c *client) sendOK() {
	c.mu.Lock()
	c.bw.WriteString("+OK\r\n")
	c.pcd[c] = needFlush
	c.mu.Unlock()
}

func (c *client) processPing() {
	c.traceOp("PING", nil)
	if c.conn == nil {
		return
	}
	c.mu.Lock()
	c.bw.WriteString("PONG\r\n")
	err := c.bw.Flush()
	c.mu.Unlock()
	if err != nil {
		// FIXME, close connection?
		Logf("Error flushing client connection [PING]: %v\n", err)
	}
}

const argsLenMax = 3

func (c *client) processPub(arg []byte) error {
	if trace {
		c.traceOp("PUB", arg)
	}

	// Unroll splitArgs to avoid runtime/heap issues
	a := [argsLenMax][]byte{}
	args := a[:0]
	start := -1
	for i, b := range arg {
		switch b {
		case ' ', '\t', '\r', '\n':
			if start >= 0 {
				args = append(args, arg[start:i])
				start = -1
			}
		default:
			if start < 0 {
				start = i
			}
		}
	}
	if start >= 0 {
		args = append(args, arg[start:])
	}

	switch len(args) {
	case 2:
		c.pa.subject = args[0]
		c.pa.reply = nil
		c.pa.size = parseSize(args[1])
		c.pa.szb = args[1]
	case 3:
		c.pa.subject = args[0]
		c.pa.reply = args[1]
		c.pa.size = parseSize(args[2])
		c.pa.szb = args[2]
	default:
		return fmt.Errorf("processPub Parse Error: '%s'", arg)
	}
	if c.pa.size < 0 {
		return fmt.Errorf("processPub Bad or Missing Size: '%s'", arg)
	}
	if c.opts.Pedantic && !sublist.IsValidLiteralSubject(c.pa.subject) {
		c.sendErr("Invalid Subject")
	}
	return nil
}

func splitArg(arg []byte) [][]byte {
	a := [argsLenMax][]byte{}
	args := a[:0]
	start := -1
	for i, b := range arg {
		switch b {
		case ' ', '\t', '\r', '\n':
			if start >= 0 {
				args = append(args, arg[start:i])
				start = -1
			}
		default:
			if start < 0 {
				start = i
			}
		}
	}
	if start >= 0 {
		args = append(args, arg[start:])
	}
	return args
}

func (c *client) processSub(argo []byte) (err error) {
	c.traceOp("SUB", argo)
	// Copy so we do not reference a potentially large buffer
	arg := make([]byte, len(argo))
	copy(arg, argo)
	args := splitArg(arg)
	sub := &subscription{client: c}
	switch len(args) {
	case 2:
		sub.subject = args[0]
		sub.queue = nil
		sub.sid = args[1]
	case 3:
		sub.subject = args[0]
		sub.queue = args[1]
		sub.sid = args[2]
	default:
		return fmt.Errorf("processSub Parse Error: '%s'", arg)
	}

	c.mu.Lock()
	c.subs.Set(sub.sid, sub)
	if c.srv != nil {
		err = c.srv.sl.Insert(sub.subject, sub)
	}
	c.mu.Unlock()
	if err != nil {
		c.sendErr("Invalid Subject")
	} else if c.opts.Verbose {
		c.sendOK()
	}
	return nil
}

func (c *client) unsubscribe(sub *subscription) {
	c.mu.Lock()
	defer c.mu.Unlock()
	if sub.max > 0 && sub.nm <= sub.max {
		return
	}
	c.traceOp("DELSUB", sub.sid)
	c.subs.Remove(sub.sid)
	if c.srv != nil {
		c.srv.sl.Remove(sub.subject, sub)
	}
}

func (c *client) processUnsub(arg []byte) error {
	c.traceOp("UNSUB", arg)
	args := splitArg(arg)
	var sid []byte
	max := -1

	switch len(args) {
	case 1:
		sid = args[0]
	case 2:
		sid = args[0]
		max = parseSize(args[1])
	default:
		return fmt.Errorf("processUnsub Parse Error: '%s'", arg)
	}
	if sub, ok := (c.subs.Get(sid)).(*subscription); ok {
		if max > 0 {
			sub.max = int64(max)
		}
		c.unsubscribe(sub)
	}
	if c.opts.Verbose {
		c.sendOK()
	}
	return nil
}

func (c *client) msgHeader(mh []byte, sub *subscription) []byte {
	mh = append(mh, sub.sid...)
	mh = append(mh, ' ')
	if c.pa.reply != nil {
		mh = append(mh, c.pa.reply...)
		mh = append(mh, ' ')
	}
	mh = append(mh, c.pa.szb...)
	mh = append(mh, "\r\n"...)
	return mh
}

// Used to treat map as efficient set
type empty struct{}

var needFlush = empty{}

func (c *client) deliverMsg(sub *subscription, mh, msg []byte) {
	if sub.client == nil || sub.client.conn == nil {
		return
	}
	client := sub.client
	client.mu.Lock()
	sub.nm++
	if sub.max > 0 && sub.nm > sub.max {
		client.mu.Unlock()
		client.unsubscribe(sub)
		return
	}
	// Deliver to the client.
	_, err := client.bw.Write(mh)
	if err != nil {
		Logf("Error writing msg header: %v\n", err)
	}
	_, err = client.bw.Write(msg)
	if err != nil {
		Logf("Error writing msg: %v\n", err)
	}
	// FIXME, this is already attached to original message
	_, err = client.bw.WriteString(CR_LF)
	if err != nil {
		Logf("Error writing CRLF: %v\n", err)
	}
	client.mu.Unlock()
	c.pcd[sub.client] = needFlush
}

func (c *client) processMsg(msg []byte) {
	c.nm++
	if trace {
		c.traceMsg(msg)
	}
	if c.srv == nil {
		return
	}
	if c.opts.Verbose {
		c.sendOK()
	}

	scratch := [512]byte{}
	msgh := scratch[:0]

	r := c.srv.sl.Match(c.pa.subject)
	if len(r) <= 0 {
		return
	}

	// msg header
	// FIXME, put MSG into initializer
	msgh = append(msgh, "MSG "...)
	msgh = append(msgh, c.pa.subject...)
	msgh = append(msgh, ' ')
	si := len(msgh)

	var qmap map[string][]*subscription
	var qsubs []*subscription

	for _, v := range r {
		sub := v.(*subscription)
		if sub.queue != nil {
			// FIXME, this can be more efficient
			if qmap == nil {
				qmap = make(map[string][]*subscription)
			}
			//qname := *(*string)(unsafe.Pointer(&sub.queue))
			qname := string(sub.queue)
			qsubs = qmap[qname]
			if qsubs == nil {
				qsubs = make([]*subscription, 0, 4)
			}
			qsubs = append(qsubs, sub)
			qmap[qname] = qsubs
			continue
		}
		mh := c.msgHeader(msgh[:si], sub)
		c.deliverMsg(sub, mh, msg)
	}
	if qmap != nil {
		for _, qsubs := range qmap {
			index := rand.Int() % len(qsubs)
			sub := qsubs[index]
			mh := c.msgHeader(msgh[:si], sub)
			c.deliverMsg(sub, mh, msg)
		}
	}
}

// Lock should be held
func (c *client) clearAuthTimer() {
	if c.atmr == nil {
		return
	}
	c.atmr.Stop()
	c.atmr = nil
}

// Lock should be held
func (c *client) clearConnection() {
	if c.conn == nil {
		return
	}
	c.bw.Flush()
	c.conn.Close()
	c.conn = nil
}

func (c *client) closeConnection() {
	if c.conn == nil {
		return
	}
	Debug("Client connection closed", clientConnStr(c.conn), c.cid)

	c.mu.Lock()
	c.clearAuthTimer()
	c.clearConnection()
	subs := c.subs.All()
	c.mu.Unlock()

	if c.srv != nil {
		for _, s := range subs {
			if sub, ok := s.(*subscription); ok {
				c.srv.sl.Remove(sub.subject, sub)
			}
		}
	}

	/*
		log.Printf("Sublist Stats: %+v\n", c.srv.sl.Stats())
		if c.nr > 0 {
			log.Printf("stats: %d %d %d\n", c.nr, c.nb, c.nm)
			log.Printf("bytes per read: %d\n", c.nb/c.nr)
			log.Printf("msgs per read: %d\n", c.nm/c.nr)
		}
	*/
}
<commit_msg>Check for dummy client with no bw/bufio<commit_after>// Copyright 2012 Apcera Inc. All rights reserved.

package server

import (
	"bufio"
	"encoding/json"
	"fmt"
	"math/rand"
	"net"
	"sync"
	"time"

	"github.com/apcera/gnatsd/hashmap"
	"github.com/apcera/gnatsd/sublist"
)

// The size of the bufio reader/writer on top of the socket.
const defaultBufSize = 32768

type client struct {
	mu   sync.Mutex
	cid  uint64
	opts clientOpts
	conn net.Conn
	bw   *bufio.Writer
	srv  *Server
	subs *hashmap.HashMap
	pcd  map[*client]struct{}
	atmr *time.Timer
	cstats
	parseState
}

func (c *client) String() string {
	return fmt.Sprintf("cid:%d", c.cid)
}

type cstats struct {
	nr int
	nb int
	nm int
}

type subscription struct {
	client  *client
	subject []byte
	queue   []byte
	sid     []byte
	nm      int64
	max     int64
}

type clientOpts struct {
	Verbose       bool   `json:"verbose"`
	Pedantic      bool   `json:"pedantic"`
	SslRequired   bool   `json:"ssl_required"`
	Authorization string `json:"auth_token"`
	Username      string `json:"user"`
	Password      string `json:"pass"`
	Name          string `json:"name"`
}

var defaultOpts = clientOpts{Verbose: true, Pedantic: true}

func init() {
	rand.Seed(time.Now().UnixNano())
}

func (c *client) readLoop() {
	b := make([]byte, defaultBufSize)
	for {
		n, err := c.conn.Read(b)
		if err != nil {
			c.closeConnection()
			return
		}
		if err := c.parse(b[:n]); err != nil {
			Log(err.Error(), clientConnStr(c.conn), c.cid)
			c.closeConnection()
			return
		}
		// Check pending clients for flush.
		for cp, _ := range c.pcd {
			// Flush those in the set
			cp.mu.Lock()
			err := cp.bw.Flush()
			cp.mu.Unlock()
			if err != nil {
				// FIXME, close connection?
				Logf("Error flushing client connection: %v\n", err)
			}
			delete(c.pcd, cp)
		}
	}
}

func (c *client) traceMsg(msg []byte) {
	pm := fmt.Sprintf("Processing msg: %d", c.nm)
	opa := []interface{}{pm, string(c.pa.subject), string(c.pa.reply), string(msg)}
	Trace(logStr(opa), fmt.Sprintf("c: %d", c.cid))
}

func (c *client) traceOp(op string, arg []byte) {
	if !trace {
		return
	}
	opa := []interface{}{fmt.Sprintf("%s OP", op)}
	if arg != nil {
		opa = append(opa, fmt.Sprintf("%s %s", op, string(arg)))
	}
	Trace(logStr(opa), fmt.Sprintf("c: %d", c.cid))
}

func (c *client) processConnect(arg []byte) error {
	c.traceOp("CONNECT", arg)

	// This will be resolved regardless before we exit this func,
	// so we can just clear it here.
	if c.atmr != nil {
		c.atmr.Stop()
		c.atmr = nil
	}

	// FIXME, check err
	if err := json.Unmarshal(arg, &c.opts); err != nil {
		return err
	}
	// Check for Auth
	if c.srv != nil {
		if ok := c.srv.checkAuth(c); !ok {
			c.sendErr("Authorization is Required")
			return fmt.Errorf("Authorization Error")
		}
	}
	if c.opts.Verbose {
		c.sendOK()
	}
	return nil
}

func (c *client) authViolation() {
	c.sendErr("Authorization is Required")
	fmt.Printf("AUTH TIMER EXPIRED!!\n")
	c.closeConnection()
}

func (c *client) sendErr(err string) {
	c.mu.Lock()
	if c.bw != nil {
		c.bw.WriteString(fmt.Sprintf("-ERR '%s'\r\n", err))
		c.pcd[c] = needFlush
	}
	c.mu.Unlock()
}

func (c *client) sendOK() {
	c.mu.Lock()
	c.bw.WriteString("+OK\r\n")
	c.pcd[c] = needFlush
	c.mu.Unlock()
}

func (c *client) processPing() {
	c.traceOp("PING", nil)
	if c.conn == nil {
		return
	}
	c.mu.Lock()
	c.bw.WriteString("PONG\r\n")
	err := c.bw.Flush()
	c.mu.Unlock()
	if err != nil {
		// FIXME, close connection?
		Logf("Error flushing client connection [PING]: %v\n", err)
	}
}

const argsLenMax = 3

func (c *client) processPub(arg []byte) error {
	if trace {
		c.traceOp("PUB", arg)
	}

	// Unroll splitArgs to avoid runtime/heap issues
	a := [argsLenMax][]byte{}
	args := a[:0]
	start := -1
	for i, b := range arg {
		switch b {
		case ' ', '\t', '\r', '\n':
			if start >= 0 {
				args = append(args, arg[start:i])
				start = -1
			}
		default:
			if start < 0 {
				start = i
			}
		}
	}
	if start >= 0 {
		args = append(args, arg[start:])
	}

	switch len(args) {
	case 2:
		c.pa.subject = args[0]
		c.pa.reply = nil
		c.pa.size = parseSize(args[1])
		c.pa.szb = args[1]
	case 3:
		c.pa.subject = args[0]
		c.pa.reply = args[1]
		c.pa.size = parseSize(args[2])
		c.pa.szb = args[2]
	default:
		return fmt.Errorf("processPub Parse Error: '%s'", arg)
	}
	if c.pa.size < 0 {
		return fmt.Errorf("processPub Bad or Missing Size: '%s'", arg)
	}
	if c.opts.Pedantic && !sublist.IsValidLiteralSubject(c.pa.subject) {
		c.sendErr("Invalid Subject")
	}
	return nil
}

func splitArg(arg []byte) [][]byte {
	a := [argsLenMax][]byte{}
	args := a[:0]
	start := -1
	for i, b := range arg {
		switch b {
		case ' ', '\t', '\r', '\n':
			if start >= 0 {
				args = append(args, arg[start:i])
				start = -1
			}
		default:
			if start < 0 {
				start = i
			}
		}
	}
	if start >= 0 {
		args = append(args, arg[start:])
	}
	return args
}

func (c *client) processSub(argo []byte) (err error) {
	c.traceOp("SUB", argo)
	// Copy so we do not reference a potentially large buffer
	arg := make([]byte, len(argo))
	copy(arg, argo)
	args := splitArg(arg)
	sub := &subscription{client: c}
	switch len(args) {
	case 2:
		sub.subject = args[0]
		sub.queue = nil
		sub.sid = args[1]
	case 3:
		sub.subject = args[0]
		sub.queue = args[1]
		sub.sid = args[2]
	default:
		return fmt.Errorf("processSub Parse Error: '%s'", arg)
	}

	c.mu.Lock()
	c.subs.Set(sub.sid, sub)
	if c.srv != nil {
		err = c.srv.sl.Insert(sub.subject, sub)
	}
	c.mu.Unlock()
	if err != nil {
		c.sendErr("Invalid Subject")
	} else if c.opts.Verbose {
		c.sendOK()
	}
	return nil
}

func (c *client) unsubscribe(sub *subscription) {
	c.mu.Lock()
	defer c.mu.Unlock()
	if sub.max > 0 && sub.nm <= sub.max {
		return
	}
	c.traceOp("DELSUB", sub.sid)
	c.subs.Remove(sub.sid)
	if c.srv != nil {
		c.srv.sl.Remove(sub.subject, sub)
	}
}

func (c *client) processUnsub(arg []byte) error {
	c.traceOp("UNSUB", arg)
	args := splitArg(arg)
	var sid []byte
	max := -1

	switch len(args) {
	case 1:
		sid = args[0]
	case 2:
		sid = args[0]
		max = parseSize(args[1])
	default:
		return fmt.Errorf("processUnsub Parse Error: '%s'", arg)
	}
	if sub, ok := (c.subs.Get(sid)).(*subscription); ok {
		if max > 0 {
			sub.max = int64(max)
		}
		c.unsubscribe(sub)
	}
	if c.opts.Verbose {
		c.sendOK()
	}
	return nil
}

func (c *client) msgHeader(mh []byte, sub *subscription) []byte {
	mh = append(mh, sub.sid...)
	mh = append(mh, ' ')
	if c.pa.reply != nil {
		mh = append(mh, c.pa.reply...)
		mh = append(mh, ' ')
	}
	mh = append(mh, c.pa.szb...)
	mh = append(mh, "\r\n"...)
	return mh
}

// Used to treat map as efficient set
type empty struct{}

var needFlush = empty{}

func (c *client) deliverMsg(sub *subscription, mh, msg []byte) {
	if sub.client == nil || sub.client.conn == nil {
		return
	}
	client := sub.client
	client.mu.Lock()
	sub.nm++
	if sub.max > 0 && sub.nm > sub.max {
		client.mu.Unlock()
		client.unsubscribe(sub)
		return
	}
	// Deliver to the client.
	_, err := client.bw.Write(mh)
	if err != nil {
		Logf("Error writing msg header: %v\n", err)
	}
	_, err = client.bw.Write(msg)
	if err != nil {
		Logf("Error writing msg: %v\n", err)
	}
	// FIXME, this is already attached to original message
	_, err = client.bw.WriteString(CR_LF)
	if err != nil {
		Logf("Error writing CRLF: %v\n", err)
	}
	client.mu.Unlock()
	c.pcd[sub.client] = needFlush
}

func (c *client) processMsg(msg []byte) {
	c.nm++
	if trace {
		c.traceMsg(msg)
	}
	if c.srv == nil {
		return
	}
	if c.opts.Verbose {
		c.sendOK()
	}

	scratch := [512]byte{}
	msgh := scratch[:0]

	r := c.srv.sl.Match(c.pa.subject)
	if len(r) <= 0 {
		return
	}

	// msg header
	// FIXME, put MSG into initializer
	msgh = append(msgh, "MSG "...)
	msgh = append(msgh, c.pa.subject...)
	msgh = append(msgh, ' ')
	si := len(msgh)

	var qmap map[string][]*subscription
	var qsubs []*subscription

	for _, v := range r {
		sub := v.(*subscription)
		if sub.queue != nil {
			// FIXME, this can be more efficient
			if qmap == nil {
				qmap = make(map[string][]*subscription)
			}
			//qname := *(*string)(unsafe.Pointer(&sub.queue))
			qname := string(sub.queue)
			qsubs = qmap[qname]
			if qsubs == nil {
				qsubs = make([]*subscription, 0, 4)
			}
			qsubs = append(qsubs, sub)
			qmap[qname] = qsubs
			continue
		}
		mh := c.msgHeader(msgh[:si], sub)
		c.deliverMsg(sub, mh, msg)
	}
	if qmap != nil {
		for _, qsubs := range qmap {
			index := rand.Int() % len(qsubs)
			sub := qsubs[index]
			mh := c.msgHeader(msgh[:si], sub)
			c.deliverMsg(sub, mh, msg)
		}
	}
}

// Lock should be held
func (c *client) clearAuthTimer() {
	if c.atmr == nil {
		return
	}
	c.atmr.Stop()
	c.atmr = nil
}

// Lock should be held
func (c *client) clearConnection() {
	if c.conn == nil {
		return
	}
	c.bw.Flush()
	c.conn.Close()
	c.conn = nil
}

func (c *client) closeConnection() {
	if c.conn == nil {
		return
	}
	Debug("Client connection closed", clientConnStr(c.conn), c.cid)

	c.mu.Lock()
	c.clearAuthTimer()
	c.clearConnection()
	subs := c.subs.All()
	c.mu.Unlock()

	if c.srv != nil {
		for _, s := range subs {
			if sub, ok := s.(*subscription); ok {
				c.srv.sl.Remove(sub.subject, sub)
			}
		}
	}

	/*
		log.Printf("Sublist Stats: %+v\n", c.srv.sl.Stats())
		if c.nr > 0 {
			log.Printf("stats: %d %d %d\n", c.nr, c.nb, c.nm)
			log.Printf("bytes per read: %d\n", c.nb/c.nr)
			log.Printf("msgs per read: %d\n", c.nm/c.nr)
		}
	*/
}
<|endoftext|>


<commit_before>package server

import (
	"errors"
	"fmt"
	"html/template"
	"net/url"
	"os"
	"path/filepath"
	texttemplate "text/template"
	"time"

	"github.com/coreos/go-oidc/key"
	"github.com/coreos/pkg/health"

	"github.com/coreos/dex/client"
	"github.com/coreos/dex/connector"
	"github.com/coreos/dex/db"
	"github.com/coreos/dex/email"
	"github.com/coreos/dex/refresh"
	"github.com/coreos/dex/repo"
	"github.com/coreos/dex/session"
	"github.com/coreos/dex/user"
	useremail "github.com/coreos/dex/user/email"
)

type ServerConfig struct {
	IssuerURL          string
	IssuerName         string
	IssuerLogoURL      string
	TemplateDir        string
	EmailTemplateDirs  []string
	EmailFromAddress   string
	EmailerConfigFile  string
	StateConfig        StateConfigurer
	EnableRegistration bool
}

type StateConfigurer interface {
	Configure(*Server) error
}

type SingleServerConfig struct {
	ClientsFile    string
	ConnectorsFile string
	UsersFile      string
}

type MultiServerConfig struct {
	KeySecrets     [][]byte
	DatabaseConfig db.Config
	UseOldFormat   bool
}

func (cfg *ServerConfig) Server() (*Server, error) {
	iu, err := url.Parse(cfg.IssuerURL)
	if err != nil {
		return nil, err
	}

	tpl, err := getTemplates(cfg.IssuerName, cfg.IssuerLogoURL, cfg.EnableRegistration, cfg.TemplateDir)
	if err != nil {
		return nil, err
	}

	km := key.NewPrivateKeyManager()
	srv := Server{
		IssuerURL:  *iu,
		KeyManager: km,
		Templates:  tpl,

		HealthChecks: []health.Checkable{km},
		Connectors:   []connector.Connector{},

		EnableRegistration: cfg.EnableRegistration,
	}

	err = cfg.StateConfig.Configure(&srv)
	if err != nil {
		return nil, err
	}

	err = setTemplates(&srv, tpl)
	if err != nil {
		return nil, err
	}

	err = setEmailer(&srv, cfg.EmailFromAddress, cfg.EmailerConfigFile, cfg.EmailTemplateDirs)
	if err != nil {
		return nil, err
	}
	return &srv, nil
}

func (cfg *SingleServerConfig) Configure(srv *Server) error {
	k, err := key.GeneratePrivateKey()
	if err != nil {
		return err
	}

	ks := key.NewPrivateKeySet([]*key.PrivateKey{k}, time.Now().Add(24*time.Hour))
	kRepo := key.NewPrivateKeySetRepo()
	if err = kRepo.Set(ks); err != nil {
		return err
	}

	cf, err := os.Open(cfg.ClientsFile)
	if err != nil {
		return fmt.Errorf("unable to read clients from file %s: %v", cfg.ClientsFile, err)
	}
	defer cf.Close()
	ciRepo, err := client.NewClientIdentityRepoFromReader(cf)
	if err != nil {
		return fmt.Errorf("unable to read client identities from file %s: %v", cfg.ClientsFile, err)
	}

	cfgRepo, err := connector.NewConnectorConfigRepoFromFile(cfg.ConnectorsFile)
	if err != nil {
		return fmt.Errorf("unable to create ConnectorConfigRepo: %v", err)
	}

	sRepo := session.NewSessionRepo()
	skRepo := session.NewSessionKeyRepo()
	sm := session.NewSessionManager(sRepo, skRepo)

	userRepo, err := user.NewUserRepoFromFile(cfg.UsersFile)
	if err != nil {
		return fmt.Errorf("unable to read users from file: %v", err)
	}

	pwiRepo := user.NewPasswordInfoRepo()

	refTokRepo := refresh.NewRefreshTokenRepo()

	txnFactory := repo.InMemTransactionFactory
	userManager := user.NewManager(userRepo, pwiRepo, txnFactory, user.ManagerOptions{})
	srv.ClientIdentityRepo = ciRepo
	srv.KeySetRepo = kRepo
	srv.ConnectorConfigRepo = cfgRepo
	srv.UserRepo = userRepo
	srv.UserManager = userManager
	srv.PasswordInfoRepo = pwiRepo
	srv.SessionManager = sm
	srv.RefreshTokenRepo = refTokRepo
	return nil

}

func (cfg *MultiServerConfig) Configure(srv *Server) error {
	if len(cfg.KeySecrets) == 0 {
		return errors.New("missing key secret")
	}

	if cfg.DatabaseConfig.DSN == "" {
		return errors.New("missing database connection string")
	}

	dbc, err := db.NewConnection(cfg.DatabaseConfig)
	if err != nil {
		return fmt.Errorf("unable to initialize database connection: %v", err)
	}

	kRepo, err := db.NewPrivateKeySetRepo(dbc, cfg.UseOldFormat, cfg.KeySecrets...)
	if err != nil {
		return fmt.Errorf("unable to create PrivateKeySetRepo: %v", err)
	}

	ciRepo := db.NewClientIdentityRepo(dbc)
	sRepo := db.NewSessionRepo(dbc)
	skRepo := db.NewSessionKeyRepo(dbc)
	cfgRepo := db.NewConnectorConfigRepo(dbc)
	userRepo := db.NewUserRepo(dbc)
	pwiRepo := db.NewPasswordInfoRepo(dbc)
	userManager := user.NewManager(userRepo, pwiRepo, db.TransactionFactory(dbc), user.ManagerOptions{})
	refreshTokenRepo := db.NewRefreshTokenRepo(dbc)

	sm := session.NewSessionManager(sRepo, skRepo)

	srv.ClientIdentityRepo = ciRepo
	srv.KeySetRepo = kRepo
	srv.ConnectorConfigRepo = cfgRepo
	srv.UserRepo = userRepo
	srv.UserManager = userManager
	srv.PasswordInfoRepo = pwiRepo
	srv.SessionManager = sm
	srv.RefreshTokenRepo = refreshTokenRepo
	return nil
}

func getTemplates(issuerName, issuerLogoURL string,
	enableRegister bool, dir string) (*template.Template, error) {
	tpl := template.New("").Funcs(map[string]interface{}{
		"issuerName": func() string {
			return issuerName
		},
		"issuerLogoURL": func() string {
			return issuerLogoURL
		},
		"enableRegister": func() bool {
			return enableRegister
		},
	})

	return tpl.ParseGlob(dir + "/*.html")
}

func setTemplates(srv *Server, tpls *template.Template) error {
	ltpl, err := findTemplate(LoginPageTemplateName, tpls)
	if err != nil {
		return err
	}
	srv.LoginTemplate = ltpl

	rtpl, err := findTemplate(RegisterTemplateName, tpls)
	if err != nil {
		return err
	}
	srv.RegisterTemplate = rtpl

	vtpl, err := findTemplate(VerifyEmailTemplateName, tpls)
	if err != nil {
		return err
	}
	srv.VerifyEmailTemplate = vtpl

	srtpl, err := findTemplate(SendResetPasswordEmailTemplateName, tpls)
	if err != nil {
		return err
	}
	srv.SendResetPasswordEmailTemplate = srtpl

	rpwtpl, err := findTemplate(ResetPasswordTemplateName, tpls)
	if err != nil {
		return err
	}
	srv.ResetPasswordTemplate = rpwtpl

	return nil
}

func setEmailer(srv *Server, fromAddress, emailerConfigFile string, emailTemplateDirs []string) error {

	cfg, err := email.NewEmailerConfigFromFile(emailerConfigFile)
	if err != nil {
		return err
	}

	emailer, err := cfg.Emailer()
	if err != nil {
		return err
	}

	getFileNames := func(dir, ext string) ([]string, error) {
		fns, err := filepath.Glob(dir + "/*." + ext)
		if err != nil {
			return nil, err
		}
		return fns, nil
	}
	getTextFiles := func(dir string) ([]string, error) {
		return getFileNames(dir, "txt")
	}
	getHTMLFiles := func(dir string) ([]string, error) {
		return getFileNames(dir, "html")
	}

	textTemplates := texttemplate.New("textTemplates")
	htmlTemplates := template.New("htmlTemplates")
	for _, dir := range emailTemplateDirs {
		textFileNames, err := getTextFiles(dir)
		if err != nil {
			return err
		}
		if len(textFileNames) != 0 {
			textTemplates, err = textTemplates.ParseFiles(textFileNames...)
		}
		if err != nil {
			return err
		}

		htmlFileNames, err := getHTMLFiles(dir)
		if err != nil {
			return err
		}
		if len(htmlFileNames) != 0 {
			htmlTemplates, err = htmlTemplates.ParseFiles(htmlFileNames...)
		}
		if err != nil {
			return err
		}
	}
	tMailer := email.NewTemplatizedEmailerFromTemplates(textTemplates, htmlTemplates, emailer)

	ue := useremail.NewUserEmailer(srv.UserRepo,
		srv.PasswordInfoRepo,
		srv.KeyManager.Signer,
		srv.SessionManager.ValidityWindow,
		srv.IssuerURL,
		tMailer,
		fromAddress,
		srv.absURL(httpPathResetPassword),
		srv.absURL(httpPathEmailVerify))

	srv.UserEmailer = ue
	return nil
}

func findTemplate(name string, tpls *template.Template) (*template.Template, error) {
	tpl := tpls.Lookup(name)
	if tpl == nil {
		return nil, fmt.Errorf("unable to find template: %q", name)
	}
	return tpl, nil
}
<commit_msg>server: pass issuer name to emailer<commit_after>package server

import (
	"errors"
	"fmt"
	"html/template"
	"net/url"
	"os"
	"path/filepath"
	texttemplate "text/template"
	"time"

	"github.com/coreos/go-oidc/key"
	"github.com/coreos/pkg/health"

	"github.com/coreos/dex/client"
	"github.com/coreos/dex/connector"
	"github.com/coreos/dex/db"
	"github.com/coreos/dex/email"
	"github.com/coreos/dex/refresh"
	"github.com/coreos/dex/repo"
	"github.com/coreos/dex/session"
	"github.com/coreos/dex/user"
	useremail "github.com/coreos/dex/user/email"
)

type ServerConfig struct {
	IssuerURL          string
	IssuerName         string
	IssuerLogoURL      string
	TemplateDir        string
	EmailTemplateDirs  []string
	EmailFromAddress   string
	EmailerConfigFile  string
	StateConfig        StateConfigurer
	EnableRegistration bool
}

type StateConfigurer interface {
	Configure(*Server) error
}

type SingleServerConfig struct {
	ClientsFile    string
	ConnectorsFile string
	UsersFile      string
}

type MultiServerConfig struct {
	KeySecrets     [][]byte
	DatabaseConfig db.Config
	UseOldFormat   bool
}

func (cfg *ServerConfig) Server() (*Server, error) {
	iu, err := url.Parse(cfg.IssuerURL)
	if err != nil {
		return nil, err
	}

	tpl, err := getTemplates(cfg.IssuerName, cfg.IssuerLogoURL, cfg.EnableRegistration, cfg.TemplateDir)
	if err != nil {
		return nil, err
	}

	km := key.NewPrivateKeyManager()
	srv := Server{
		IssuerURL:  *iu,
		KeyManager: km,
		Templates:  tpl,

		HealthChecks: []health.Checkable{km},
		Connectors:   []connector.Connector{},

		EnableRegistration: cfg.EnableRegistration,
	}

	err = cfg.StateConfig.Configure(&srv)
	if err != nil {
		return nil, err
	}

	err = setTemplates(&srv, tpl)
	if err != nil {
		return nil, err
	}

	err = setEmailer(&srv, cfg.IssuerName, cfg.EmailFromAddress, cfg.EmailerConfigFile, cfg.EmailTemplateDirs)
	if err != nil {
		return nil, err
	}
	return &srv, nil
}

func (cfg *SingleServerConfig) Configure(srv *Server) error {
	k, err := key.GeneratePrivateKey()
	if err != nil {
		return err
	}

	ks := key.NewPrivateKeySet([]*key.PrivateKey{k}, time.Now().Add(24*time.Hour))
	kRepo := key.NewPrivateKeySetRepo()
	if err = kRepo.Set(ks); err != nil {
		return err
	}

	cf, err := os.Open(cfg.ClientsFile)
	if err != nil {
		return fmt.Errorf("unable to read clients from file %s: %v", cfg.ClientsFile, err)
	}
	defer cf.Close()
	ciRepo, err := client.NewClientIdentityRepoFromReader(cf)
	if err != nil {
		return fmt.Errorf("unable to read client identities from file %s: %v", cfg.ClientsFile, err)
	}

	cfgRepo, err := connector.NewConnectorConfigRepoFromFile(cfg.ConnectorsFile)
	if err != nil {
		return fmt.Errorf("unable to create ConnectorConfigRepo: %v", err)
	}

	sRepo := session.NewSessionRepo()
	skRepo := session.NewSessionKeyRepo()
	sm := session.NewSessionManager(sRepo, skRepo)

	userRepo, err := user.NewUserRepoFromFile(cfg.UsersFile)
	if err != nil {
		return fmt.Errorf("unable to read users from file: %v", err)
	}

	pwiRepo := user.NewPasswordInfoRepo()

	refTokRepo := refresh.NewRefreshTokenRepo()

	txnFactory := repo.InMemTransactionFactory
	userManager := user.NewManager(userRepo, pwiRepo, txnFactory, user.ManagerOptions{})
	srv.ClientIdentityRepo = ciRepo
	srv.KeySetRepo = kRepo
	srv.ConnectorConfigRepo = cfgRepo
	srv.UserRepo = userRepo
	srv.UserManager = userManager
	srv.PasswordInfoRepo = pwiRepo
	srv.SessionManager = sm
	srv.RefreshTokenRepo = refTokRepo
	return nil

}

func (cfg *MultiServerConfig) Configure(srv *Server) error {
	if len(cfg.KeySecrets) == 0 {
		return errors.New("missing key secret")
	}

	if cfg.DatabaseConfig.DSN == "" {
		return errors.New("missing database connection string")
	}

	dbc, err := db.NewConnection(cfg.DatabaseConfig)
	if err != nil {
		return fmt.Errorf("unable to initialize database connection: %v", err)
	}

	kRepo, err := db.NewPrivateKeySetRepo(dbc, cfg.UseOldFormat, cfg.KeySecrets...)
	if err != nil {
		return fmt.Errorf("unable to create PrivateKeySetRepo: %v", err)
	}

	ciRepo := db.NewClientIdentityRepo(dbc)
	sRepo := db.NewSessionRepo(dbc)
	skRepo := db.NewSessionKeyRepo(dbc)
	cfgRepo := db.NewConnectorConfigRepo(dbc)
	userRepo := db.NewUserRepo(dbc)
	pwiRepo := db.NewPasswordInfoRepo(dbc)
	userManager := user.NewManager(userRepo, pwiRepo, db.TransactionFactory(dbc), user.ManagerOptions{})
	refreshTokenRepo := db.NewRefreshTokenRepo(dbc)

	sm := session.NewSessionManager(sRepo, skRepo)

	srv.ClientIdentityRepo = ciRepo
	srv.KeySetRepo = kRepo
	srv.ConnectorConfigRepo = cfgRepo
	srv.UserRepo = userRepo
	srv.UserManager = userManager
	srv.PasswordInfoRepo = pwiRepo
	srv.SessionManager = sm
	srv.RefreshTokenRepo = refreshTokenRepo
	return nil
}

func getTemplates(issuerName, issuerLogoURL string,
	enableRegister bool, dir string) (*template.Template, error) {
	tpl := template.New("").Funcs(map[string]interface{}{
		"issuerName": func() string {
			return issuerName
		},
		"issuerLogoURL": func() string {
			return issuerLogoURL
		},
		"enableRegister": func() bool {
			return enableRegister
		},
	})

	return tpl.ParseGlob(dir + "/*.html")
}

func setTemplates(srv *Server, tpls *template.Template) error {
	ltpl, err := findTemplate(LoginPageTemplateName, tpls)
	if err != nil {
		return err
	}
	srv.LoginTemplate = ltpl

	rtpl, err := findTemplate(RegisterTemplateName, tpls)
	if err != nil {
		return err
	}
	srv.RegisterTemplate = rtpl

	vtpl, err := findTemplate(VerifyEmailTemplateName, tpls)
	if err != nil {
		return err
	}
	srv.VerifyEmailTemplate = vtpl

	srtpl, err := findTemplate(SendResetPasswordEmailTemplateName, tpls)
	if err != nil {
		return err
	}
	srv.SendResetPasswordEmailTemplate = srtpl

	rpwtpl, err := findTemplate(ResetPasswordTemplateName, tpls)
	if err != nil {
		return err
	}
	srv.ResetPasswordTemplate = rpwtpl

	return nil
}

func setEmailer(srv *Server, issuerName, fromAddress, emailerConfigFile string, emailTemplateDirs []string) error {

	cfg, err := email.NewEmailerConfigFromFile(emailerConfigFile)
	if err != nil {
		return err
	}

	emailer, err := cfg.Emailer()
	if err != nil {
		return err
	}

	getFileNames := func(dir, ext string) ([]string, error) {
		fns, err := filepath.Glob(dir + "/*." + ext)
		if err != nil {
			return nil, err
		}
		return fns, nil
	}
	getTextFiles := func(dir string) ([]string, error) {
		return getFileNames(dir, "txt")
	}
	getHTMLFiles := func(dir string) ([]string, error) {
		return getFileNames(dir, "html")
	}

	textTemplates := texttemplate.New("textTemplates")
	htmlTemplates := template.New("htmlTemplates")
	for _, dir := range emailTemplateDirs {
		textFileNames, err := getTextFiles(dir)
		if err != nil {
			return err
		}
		if len(textFileNames) != 0 {
			textTemplates, err = textTemplates.ParseFiles(textFileNames...)
		}
		if err != nil {
			return err
		}

		htmlFileNames, err := getHTMLFiles(dir)
		if err != nil {
			return err
		}
		if len(htmlFileNames) != 0 {
			htmlTemplates, err = htmlTemplates.ParseFiles(htmlFileNames...)
		}
		if err != nil {
			return err
		}
	}
	tMailer := email.NewTemplatizedEmailerFromTemplates(textTemplates, htmlTemplates, emailer)
	tMailer.SetGlobalContext(map[string]interface{}{
		"issuer_name": issuerName,
	})

	ue := useremail.NewUserEmailer(srv.UserRepo,
		srv.PasswordInfoRepo,
		srv.KeyManager.Signer,
		srv.SessionManager.ValidityWindow,
		srv.IssuerURL,
		tMailer,
		fromAddress,
		srv.absURL(httpPathResetPassword),
		srv.absURL(httpPathEmailVerify))

	srv.UserEmailer = ue
	return nil
}

func findTemplate(name string, tpls *template.Template) (*template.Template, error) {
	tpl := tpls.Lookup(name)
	if tpl == nil {
		return nil, fmt.Errorf("unable to find template: %q", name)
	}
	return tpl, nil
}
<|endoftext|>