<commit_before>package proxy

import (
	"bufio"
	"bytes"
	"crypto/tls"
	"encoding/base64"
	"fmt"
	"io"
	"log"
	"net"
	"net/http"
	"sync"
	"time"
)

var (
	ConnectTimeout     time.Duration = 30 * time.Second
	RetryListenTimeout time.Duration = 5 * time.Second
)

type (
	proxy struct {
		config    *Config
		frontends map[int]net.Listener
		backends  map[string]*backend

		configChannel          chan *Config
		incomingRequestChannel chan *incomingRequest
		tlsConfigChannel       chan *tls.Config
	}
	backend struct {
		endpoints []string
		count     uint64
	}
)

func newProxy() *proxy {
	return &proxy{
		config:                 &Config{},
		frontends:              make(map[int]net.Listener),
		backends:               make(map[string]*backend),
		configChannel:          make(chan *Config),
		incomingRequestChannel: make(chan *incomingRequest),
		tlsConfigChannel:       make(chan *tls.Config),
	}
}

func (this *proxy) reload(filename string) {
	cfg, err := GetConfig(filename)
	if err != nil {
		log.Println("error loading config:", err)
		return
	}
	this.configChannel <- cfg
}

func (this *proxy) start() {
	for {
		select {
		case cfg := <-this.configChannel:
			log.Println("load config")

			err := this.listen(cfg.Port)
			if err != nil {
				log.Println("error listening on", cfg.Port, ":", err)
				time.AfterFunc(RetryListenTimeout, func() {
					this.configChannel <- cfg
				})
				break
			}

			err = this.listenTLS(cfg.TLSPort)
			if err != nil {
				log.Println("error listening on", cfg.TLSPort, ":", err)
				time.AfterFunc(RetryListenTimeout, func() {
					this.configChannel <- cfg
				})
				break
			}

			this.backends = map[string]*backend{}
			for host, route := range cfg.Routes {
				be := &backend{route.Endpoints, 0}
				this.backends[host] = be
				log.Println(host, "-->", be.endpoints)
			}

			this.tlsConfigChannel <- this.getTLSConfig(cfg)

			this.config = cfg
		case ir := <-this.incomingRequestChannel:
			if ir.err != nil {
				ir.writeError(ir.err.Error(), 500)
			} else if be, ok := this.backends[ir.request.Host]; ok {
				endpoint := be.endpoints[be.count%uint64(len(be.endpoints))]
				go this.join(ir, endpoint)
				be.count++
			} else {
				ir.writeError(fmt.Sprintf("unknown host '%v'", ir.request.Host), 404)
			}
		}
	}
}

func (this *proxy) getTLSConfig(cfg *Config) *tlsConfig {
	tlsConfig := &tls.Config{
		Certificates: []tls.Certificate{},
	}
	for host, route := range cfg.Routes {
		if route.TLS.Certificate != "" && route.TLS.Key != "" {
			certBs, err := base64.StdEncoding.DecodeString(route.TLS.Certificate)
			if err != nil {
				log.Println("invalid certificate for", host)
				continue
			}
			keyBs, err := base64.StdEncoding.DecodeString(route.TLS.Key)
			if err != nil {
				log.Println("invalid key for", host)
				continue
			}
			cert, err := tls.X509KeyPair(certBs, keyBs)
			if err != nil {
				log.Println("invalid tls for", host)
				continue
			}
			tlsConfig.Certificates = append(tlsConfig.Certificates, cert)
		}
	}
	tlsConfig.BuildNameToCertificate()
	return tlsConfig
}

func (this *proxy) listen(port int) error {
	if port == this.config.Port {
		return nil
	}

	l, ok := this.frontends[this.config.Port]
	if ok {
		l.Close()
		delete(this.frontends, this.config.Port)
	}

	l, err := net.Listen("tcp", fmt.Sprint(":", port))
	if err != nil {
		return err
	}

	go this.accept(l)
	this.frontends[port] = l
	this.config.Port = port

	return nil
}

func (this *proxy) listenTLS(port int) error {
	if port == this.config.TLSPort {
		return nil
	}

	l, ok := this.frontends[this.config.TLSPort]
	if ok {
		l.Close()
		delete(this.frontends, this.config.TLSPort)
	}

	l, err := net.Listen("tcp", fmt.Sprint(":", port))
	if err != nil {
		return err
	}

	go this.accept(l)
	this.frontends[port] = l
	this.config.TLSPort = port

	return nil
}

func (this *proxy) accept(listener net.Listener) {
	for {
		conn, err := listener.Accept()
		if err != nil {
			log.Println("error accepting from", listener, ":", err)
			return
		}
		go this.handleConn(conn)
	}
}

func (this *proxy) acceptTLS(listener net.Listener) {
	conns := make(chan net.Conn)
	go func() {
		for {
			conn, err := listener.Accept()
			if err != nil {
				log.Println("error accepting from", listener, ":", err)
				break
			}
			conns <- conn
		}
		close(conns)
	}()

	config := <-this.tlsConfigChannel

	for {
		select {
		case cfg := <-this.tlsConfigChannel:
			config = cfg
		case conn, ok := <-conns:
			if !ok {
				return
			}
			conn = tls.Server(conn, config)
			go this.handleConn(conn)
		}
	}
}

func (this *proxy) handleConn(conn net.Conn) {
	ir := &incomingRequest{
		Conn:   conn,
		buffer: bytes.NewBuffer(nil),
	}
	rdr := io.TeeReader(ir.Conn, ir.buffer)
	ir.reader = bufio.NewReader(rdr)
	ir.request, ir.err = http.ReadRequest(ir.reader)
	this.incomingRequestChannel <- ir
	if ir.request != nil {
		log.Println(ir.request.Method, ir.request.URL)
	}
}

func (this *proxy) join(ir *incomingRequest, endpoint string) {
	conn, err := net.Dial("tcp", endpoint)
	if err != nil {
		log.Println("error connecting to", endpoint, ":", err)

		if ir.opened.Add(-ConnectTimeout).After(time.Now()) {
			ir.writeError("timeout", 504)
		} else {
			time.AfterFunc(time.Second, func() {
				this.incomingRequestChannel <- ir
			})
		}

		return
	}

	defer ir.Conn.Close()
	defer conn.Close()

	_, err = io.Copy(conn, ir.buffer)
	if err != nil && err != io.EOF {
		log.Println("error copying buffer to", conn.RemoteAddr(), ":", err)
		return
	}

	var wg sync.WaitGroup
	halfJoin := func(dst net.Conn, src net.Conn) {
		defer wg.Done()
		_, err := io.Copy(dst, src)
		if err != nil {
			log.Println("error copying from", src.RemoteAddr(), "to", dst.RemoteAddr(), ":", err)
		}
	}

	wg.Add(2)
	go halfJoin(ir.Conn, conn)
	go halfJoin(conn, ir.Conn)
	wg.Wait()
}
<commit_msg>fix typo<commit_after>package proxy

import (
	"bufio"
	"bytes"
	"crypto/tls"
	"encoding/base64"
	"fmt"
	"io"
	"log"
	"net"
	"net/http"
	"sync"
	"time"
)

var (
	ConnectTimeout     time.Duration = 30 * time.Second
	RetryListenTimeout time.Duration = 5 * time.Second
)

type (
	proxy struct {
		config    *Config
		frontends map[int]net.Listener
		backends  map[string]*backend

		configChannel          chan *Config
		incomingRequestChannel chan *incomingRequest
		tlsConfigChannel       chan *tls.Config
	}
	backend struct {
		endpoints []string
		count     uint64
	}
)

func newProxy() *proxy {
	return &proxy{
		config:                 &Config{},
		frontends:              make(map[int]net.Listener),
		backends:               make(map[string]*backend),
		configChannel:          make(chan *Config),
		incomingRequestChannel: make(chan *incomingRequest),
		tlsConfigChannel:       make(chan *tls.Config),
	}
}

func (this *proxy) reload(filename string) {
	cfg, err := GetConfig(filename)
	if err != nil {
		log.Println("error loading config:", err)
		return
	}
	this.configChannel <- cfg
}

func (this *proxy) start() {
	for {
		select {
		case cfg := <-this.configChannel:
			log.Println("load config")

			err := this.listen(cfg.Port)
			if err != nil {
				log.Println("error listening on", cfg.Port, ":", err)
				time.AfterFunc(RetryListenTimeout, func() {
					this.configChannel <- cfg
				})
				break
			}

			err = this.listenTLS(cfg.TLSPort)
			if err != nil {
				log.Println("error listening on", cfg.TLSPort, ":", err)
				time.AfterFunc(RetryListenTimeout, func() {
					this.configChannel <- cfg
				})
				break
			}

			this.backends = map[string]*backend{}
			for host, route := range cfg.Routes {
				be := &backend{route.Endpoints, 0}
				this.backends[host] = be
				log.Println(host, "-->", be.endpoints)
			}

			this.tlsConfigChannel <- this.getTLSConfig(cfg)

			this.config = cfg
		case ir := <-this.incomingRequestChannel:
			if ir.err != nil {
				ir.writeError(ir.err.Error(), 500)
			} else if be, ok := this.backends[ir.request.Host]; ok {
				endpoint := be.endpoints[be.count%uint64(len(be.endpoints))]
				go this.join(ir, endpoint)
				be.count++
			} else {
				ir.writeError(fmt.Sprintf("unknown host '%v'", ir.request.Host), 404)
			}
		}
	}
}

func (this *proxy) getTLSConfig(cfg *Config) *tls.Config {
	tlsConfig := &tls.Config{
		Certificates: []tls.Certificate{},
	}
	for host, route := range cfg.Routes {
		if route.TLS.Certificate != "" && route.TLS.Key != "" {
			certBs, err := base64.StdEncoding.DecodeString(route.TLS.Certificate)
			if err != nil {
				log.Println("invalid certificate for", host)
				continue
			}
			keyBs, err := base64.StdEncoding.DecodeString(route.TLS.Key)
			if err != nil {
				log.Println("invalid key for", host)
				continue
			}
			cert, err := tls.X509KeyPair(certBs, keyBs)
			if err != nil {
				log.Println("invalid tls for", host)
				continue
			}
			tlsConfig.Certificates = append(tlsConfig.Certificates, cert)
		}
	}
	tlsConfig.BuildNameToCertificate()
	return tlsConfig
}

func (this *proxy) listen(port int) error {
	if port == this.config.Port {
		return nil
	}

	l, ok := this.frontends[this.config.Port]
	if ok {
		l.Close()
		delete(this.frontends, this.config.Port)
	}

	l, err := net.Listen("tcp", fmt.Sprint(":", port))
	if err != nil {
		return err
	}

	go this.accept(l)
	this.frontends[port] = l
	this.config.Port = port

	return nil
}

func (this *proxy) listenTLS(port int) error {
	if port == this.config.TLSPort {
		return nil
	}

	l, ok := this.frontends[this.config.TLSPort]
	if ok {
		l.Close()
		delete(this.frontends, this.config.TLSPort)
	}

	l, err := net.Listen("tcp", fmt.Sprint(":", port))
	if err != nil {
		return err
	}

	go this.accept(l)
	this.frontends[port] = l
	this.config.TLSPort = port

	return nil
}

func (this *proxy) accept(listener net.Listener) {
	for {
		conn, err := listener.Accept()
		if err != nil {
			log.Println("error accepting from", listener, ":", err)
			return
		}
		go this.handleConn(conn)
	}
}

func (this *proxy) acceptTLS(listener net.Listener) {
	conns := make(chan net.Conn)
	go func() {
		for {
			conn, err := listener.Accept()
			if err != nil {
				log.Println("error accepting from", listener, ":", err)
				break
			}
			conns <- conn
		}
		close(conns)
	}()

	config := <-this.tlsConfigChannel

	for {
		select {
		case cfg := <-this.tlsConfigChannel:
			config = cfg
		case conn, ok := <-conns:
			if !ok {
				return
			}
			conn = tls.Server(conn, config)
			go this.handleConn(conn)
		}
	}
}

func (this *proxy) handleConn(conn net.Conn) {
	ir := &incomingRequest{
		Conn:   conn,
		buffer: bytes.NewBuffer(nil),
	}
	rdr := io.TeeReader(ir.Conn, ir.buffer)
	ir.reader = bufio.NewReader(rdr)
	ir.request, ir.err = http.ReadRequest(ir.reader)
	this.incomingRequestChannel <- ir
	if ir.request != nil {
		log.Println(ir.request.Method, ir.request.URL)
	}
}

func (this *proxy) join(ir *incomingRequest, endpoint string) {
	conn, err := net.Dial("tcp", endpoint)
	if err != nil {
		log.Println("error connecting to", endpoint, ":", err)

		if ir.opened.Add(-ConnectTimeout).After(time.Now()) {
			ir.writeError("timeout", 504)
		} else {
			time.AfterFunc(time.Second, func() {
				this.incomingRequestChannel <- ir
			})
		}

		return
	}

	defer ir.Conn.Close()
	defer conn.Close()

	_, err = io.Copy(conn, ir.buffer)
	if err != nil && err != io.EOF {
		log.Println("error copying buffer to", conn.RemoteAddr(), ":", err)
		return
	}

	var wg sync.WaitGroup
	halfJoin := func(dst net.Conn, src net.Conn) {
		defer wg.Done()
		_, err := io.Copy(dst, src)
		if err != nil {
			log.Println("error copying from", src.RemoteAddr(), "to", dst.RemoteAddr(), ":", err)
		}
	}

	wg.Add(2)
	go halfJoin(ir.Conn, conn)
	go halfJoin(conn, ir.Conn)
	wg.Wait()
}
<|endoftext|>


<commit_before>package profile

import (
	"bufio"
	"bytes"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"testing"
)

type checkFn func(t *testing.T, stdout, stderr []byte, err error)

var profileTests = []struct {
	name   string
	code   string
	checks []checkFn
}{{
	name: "default profile (cpu)",
	code: `
package main

import "github.com/pkg/profile"

func main() {
	defer profile.Start().Stop()
}
`,
	checks: []checkFn{
		NoStdout,
		Stderr("profile: cpu profiling enabled"),
		NoErr,
	},
}, {
	name: "memory profile",
	code: `
package main

import "github.com/pkg/profile"

func main() {
	defer profile.Start(profile.MemProfile).Stop()
}
`,
	checks: []checkFn{
		NoStdout,
		Stderr("profile: memory profiling enabled"),
		NoErr,
	},
}, {
	name: "memory profile (rate 2048)",
	code: `
package main

import "github.com/pkg/profile"

func main() {
	defer profile.Start(profile.MemProfileRate(2048)).Stop()
}
`,
	checks: []checkFn{
		NoStdout,
		Stderr("profile: memory profiling enabled (rate 2048)"),
		NoErr,
	},
}, {
	name: "double start",
	code: `
package main

import "github.com/pkg/profile"

func main() {
	defer profile.Start().Stop()
	defer profile.Start().Stop()
}
`,
	checks: []checkFn{
		NoStdout,
		Stderr("cpu profiling enabled", "profile: Start() already called"),
		Err,
	},
}, {
	name: "block profile",
	code: `
package main

import "github.com/pkg/profile"

func main() {
	defer profile.Start(profile.BlockProfile).Stop()
}
`,
	checks: []checkFn{
		NoStdout,
		Stderr("profile: block profiling enabled"),
		NoErr,
	},
}, {
	name: "profile path",
	code: `
package main

import "github.com/pkg/profile"

func main() {
	defer profile.Start(profile.ProfilePath(".")).Stop()
}
`,
	checks: []checkFn{
		NoStdout,
		Stderr("profile: cpu profiling enabled, cpu.pprof"),
		NoErr,
	},
}, {
	name: "profile path error",
	code: `
package main

import "github.com/pkg/profile"

func main() {
		defer profile.Start(profile.ProfilePath("README.md")).Stop()
}
`,
	checks: []checkFn{
		NoStdout,
		Stderr("could not create initial output"),
		Err,
	},
}, {
	name: "profile quiet",
	code: `
package main

import "github.com/pkg/profile"

func main() {
        defer profile.Start(profile.Quiet).Stop()
}
`,
	checks: []checkFn{NoStdout, NoStderr, NoErr},
}}

func TestProfile(t *testing.T) {
	for _, tt := range profileTests {
		t.Log(tt.name)
		stdout, stderr, err := runTest(t, tt.code)
		for _, f := range tt.checks {
			f(t, stdout, stderr, err)
		}
	}
}

// Stdout verifies that the given lines match the output from stdout
func Stdout(lines ...string) checkFn {
	return func(t *testing.T, stdout, _ []byte, _ error) {
		buf := bytes.NewBuffer(stdout)
		if !validateOutput(buf, lines) {
			t.Errorf("stderr: wanted '%s', got '%s'", lines, stdout)
		}

	}
}

// NoStdout checks that stdout was blank.
func NoStdout(t *testing.T, stdout, _ []byte, _ error) {
	if len := len(stdout); len > 0 {
		t.Errorf("stdout: wanted 0 bytes, got %d", len)
	}
}

// Stderr verifies that the given lines match the output from stderr
func Stderr(lines ...string) checkFn {
	return func(t *testing.T, stdout, stderr []byte, _ error) {
		buf := bytes.NewBuffer(stderr)
		if !validateOutput(buf, lines) {
			t.Errorf("stderr: wanted '%s', got '%s'", lines, stderr)
		}
	}
}

// NoStderr checks that stderr was blank.
func NoStderr(t *testing.T, _, stderr []byte, _ error) {
	if len := len(stderr); len > 0 {
		t.Errorf("stderr: wanted 0 bytes, got %d", len)
	}
}

// Err checks that there was an error returned
func Err(t *testing.T, _, _ []byte, err error) {
	if err == nil {
		t.Errorf("expected error")
	}
}

// NoErr checks that err was nil
func NoErr(t *testing.T, _, _ []byte, err error) {
	if err != nil {
		t.Errorf("error: expected nil, got %v", err)
	}
}

func validateOutput(buf *bytes.Buffer, lines []string) bool {
	s := bufio.NewScanner(buf)
	for _, l := range lines {
		if !s.Scan() {
			return false
		}
		if !strings.Contains(s.Text(), l) {
			return false
		}
	}
	return true
}

// runTest executes the go program supplied and returns the contents of stdout,
// stderr, and an error which may contain status information about the result
// of the program.
func runTest(t *testing.T, code string) ([]byte, []byte, error) {
	chk := func(err error) {
		if err != nil {
			t.Fatal(err)
		}
	}
	gopath, err := ioutil.TempDir("", "profile-gopath")
	chk(err)
	defer os.RemoveAll(gopath)

	srcdir := filepath.Join(gopath, "src")
	err = os.Mkdir(srcdir, 0755)
	chk(err)
	src := filepath.Join(srcdir, "main.go")
	err = ioutil.WriteFile(src, []byte(code), 0644)
	chk(err)

	cmd := exec.Command("go", "run", src)

	var stdout, stderr bytes.Buffer
	cmd.Stdout = &stdout
	cmd.Stderr = &stderr
	err = cmd.Run()
	return stdout.Bytes(), stderr.Bytes(), err
}
<commit_msg>comment<commit_after>package profile

import (
	"bufio"
	"bytes"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"testing"
)

type checkFn func(t *testing.T, stdout, stderr []byte, err error)

var profileTests = []struct {
	name   string
	code   string
	checks []checkFn
}{{
	name: "default profile (cpu)",
	code: `
package main

import "github.com/pkg/profile"

func main() {
	defer profile.Start().Stop()
}
`,
	checks: []checkFn{
		NoStdout,
		Stderr("profile: cpu profiling enabled"),
		NoErr,
	},
}, {
	name: "memory profile",
	code: `
package main

import "github.com/pkg/profile"

func main() {
	defer profile.Start(profile.MemProfile).Stop()
}
`,
	checks: []checkFn{
		NoStdout,
		Stderr("profile: memory profiling enabled"),
		NoErr,
	},
}, {
	name: "memory profile (rate 2048)",
	code: `
package main

import "github.com/pkg/profile"

func main() {
	defer profile.Start(profile.MemProfileRate(2048)).Stop()
}
`,
	checks: []checkFn{
		NoStdout,
		Stderr("profile: memory profiling enabled (rate 2048)"),
		NoErr,
	},
}, {
	name: "double start",
	code: `
package main

import "github.com/pkg/profile"

func main() {
	defer profile.Start().Stop()
	defer profile.Start().Stop()
}
`,
	checks: []checkFn{
		NoStdout,
		Stderr("cpu profiling enabled", "profile: Start() already called"),
		Err,
	},
}, {
	name: "block profile",
	code: `
package main

import "github.com/pkg/profile"

func main() {
	defer profile.Start(profile.BlockProfile).Stop()
}
`,
	checks: []checkFn{
		NoStdout,
		Stderr("profile: block profiling enabled"),
		NoErr,
	},
}, {
	name: "profile path",
	code: `
package main

import "github.com/pkg/profile"

func main() {
	defer profile.Start(profile.ProfilePath(".")).Stop()
}
`,
	checks: []checkFn{
		NoStdout,
		Stderr("profile: cpu profiling enabled, cpu.pprof"),
		NoErr,
	},
}, {
	name: "profile path error",
	code: `
package main

import "github.com/pkg/profile"

func main() {
		defer profile.Start(profile.ProfilePath("README.md")).Stop()
}
`,
	checks: []checkFn{
		NoStdout,
		Stderr("could not create initial output"),
		Err,
	},
}, {
	name: "profile quiet",
	code: `
package main

import "github.com/pkg/profile"

func main() {
        defer profile.Start(profile.Quiet).Stop()
}
`,
	checks: []checkFn{NoStdout, NoStderr, NoErr},
}}

func TestProfile(t *testing.T) {
	for _, tt := range profileTests {
		t.Log(tt.name)
		stdout, stderr, err := runTest(t, tt.code)
		for _, f := range tt.checks {
			f(t, stdout, stderr, err)
		}
	}
}

// Stdout verifies that the given lines match the output from stdout
func Stdout(lines ...string) checkFn {
	return func(t *testing.T, stdout, _ []byte, _ error) {
		buf := bytes.NewBuffer(stdout)
		if !validateOutput(buf, lines) {
			t.Errorf("stderr: wanted '%s', got '%s'", lines, stdout)
		}

	}
}

// NoStdout checks that stdout was blank.
func NoStdout(t *testing.T, stdout, _ []byte, _ error) {
	if len := len(stdout); len > 0 {
		t.Errorf("stdout: wanted 0 bytes, got %d", len)
	}
}

// Stderr verifies that the given lines match the output from stderr
func Stderr(lines ...string) checkFn {
	return func(t *testing.T, stdout, stderr []byte, _ error) {
		buf := bytes.NewBuffer(stderr)
		if !validateOutput(buf, lines) {
			t.Errorf("stderr: wanted '%s', got '%s'", lines, stderr)
		}
	}
}

// NoStderr checks that stderr was blank.
func NoStderr(t *testing.T, _, stderr []byte, _ error) {
	if len := len(stderr); len > 0 {
		t.Errorf("stderr: wanted 0 bytes, got %d", len)
	}
}

// Err checks that there was an error returned
func Err(t *testing.T, _, _ []byte, err error) {
	if err == nil {
		t.Errorf("expected error")
	}
}

// NoErr checks that err was nil
func NoErr(t *testing.T, _, _ []byte, err error) {
	if err != nil {
		t.Errorf("error: expected nil, got %v", err)
	}
}

// validatedOutput validates the given slice of lines against a scanned buffer
func validateOutput(buf *bytes.Buffer, lines []string) bool {
	s := bufio.NewScanner(buf)
	for _, l := range lines {
		if !s.Scan() {
			return false
		}
		if !strings.Contains(s.Text(), l) {
			return false
		}
	}
	return true
}

// runTest executes the go program supplied and returns the contents of stdout,
// stderr, and an error which may contain status information about the result
// of the program.
func runTest(t *testing.T, code string) ([]byte, []byte, error) {
	chk := func(err error) {
		if err != nil {
			t.Fatal(err)
		}
	}
	gopath, err := ioutil.TempDir("", "profile-gopath")
	chk(err)
	defer os.RemoveAll(gopath)

	srcdir := filepath.Join(gopath, "src")
	err = os.Mkdir(srcdir, 0755)
	chk(err)
	src := filepath.Join(srcdir, "main.go")
	err = ioutil.WriteFile(src, []byte(code), 0644)
	chk(err)

	cmd := exec.Command("go", "run", src)

	var stdout, stderr bytes.Buffer
	cmd.Stdout = &stdout
	cmd.Stderr = &stderr
	err = cmd.Run()
	return stdout.Bytes(), stderr.Bytes(), err
}
<|endoftext|>


<commit_before>package seq

import (
	"fmt"
	"math"
	"reflect"
	"testing"
)

var pf = fmt.Printf

var shortAlpha = NewAlphabet('A', 'B', 'C')

func o(freq, coltotal, nullfreq, nulltot int) Prob {
	if freq == 0 || nullfreq == 0 {
		return MinProb
	}
	num := float64(freq) / float64(coltotal)
	den := float64(nullfreq) / float64(nulltot)
	return -Prob(math.Log(num / den))
}

var tests = []struct {
	pexpected *Profile
	fexpected *FrequencyProfile
	null      *FrequencyProfile
	seqs      []Sequence
}{
	{
		pexpected: &Profile{
			Alphabet: shortAlpha,
			Emissions: []EProbs{
				{'A': o(3, 3, 3, 9), 'B': o(0, 3, 3, 9), 'C': o(0, 3, 3, 9)},
				{'A': o(0, 3, 3, 9), 'B': o(3, 3, 3, 9), 'C': o(0, 3, 3, 9)},
				{'A': o(0, 3, 3, 9), 'B': o(0, 3, 3, 9), 'C': o(3, 3, 3, 9)},
			},
		},
		fexpected: &FrequencyProfile{
			Alphabet: shortAlpha,
			Freqs: []map[Residue]int{
				{'A': 3, 'B': 0, 'C': 0},
				{'A': 0, 'B': 3, 'C': 0},
				{'A': 0, 'B': 0, 'C': 3},
			},
		},
		null: &FrequencyProfile{
			Alphabet: shortAlpha,
			Freqs:    []map[Residue]int{{'A': 3, 'B': 3, 'C': 3}},
		},
		seqs: []Sequence{
			{"1", strr("ABC")},
			{"2", strr("ABC")},
			{"3", strr("ABC")},
		},
	},
	{
		pexpected: &Profile{
			Alphabet: shortAlpha,
			Emissions: []EProbs{
				{'A': o(1, 3, 3, 9), 'B': o(2, 3, 3, 9), 'C': o(0, 3, 3, 9)},
				{'A': o(0, 3, 3, 9), 'B': o(1, 3, 3, 9), 'C': o(2, 3, 3, 9)},
				{'A': o(2, 3, 3, 9), 'B': o(0, 3, 3, 9), 'C': o(1, 3, 3, 9)},
			},
		},
		fexpected: &FrequencyProfile{
			Alphabet: shortAlpha,
			Freqs: []map[Residue]int{
				{'A': 1, 'B': 2, 'C': 0},
				{'A': 0, 'B': 1, 'C': 2},
				{'A': 2, 'B': 0, 'C': 1},
			},
		},
		null: &FrequencyProfile{
			Alphabet: shortAlpha,
			Freqs:    []map[Residue]int{{'A': 3, 'B': 3, 'C': 3}},
		},
		seqs: []Sequence{
			{"1", strr("ABC")},
			{"2", strr("BCA")},
			{"3", strr("BCA")},
		},
	},
	{
		pexpected: &Profile{
			Alphabet: shortAlpha,
			Emissions: []EProbs{
				{'A': o(0, 3, 1, 9), 'B': o(2, 3, 2, 9), 'C': o(1, 3, 6, 9)},
				{'A': o(0, 3, 1, 9), 'B': o(0, 3, 2, 9), 'C': o(3, 3, 6, 9)},
				{'A': o(1, 3, 1, 9), 'B': o(0, 3, 2, 9), 'C': o(2, 3, 6, 9)},
			},
		},
		fexpected: &FrequencyProfile{
			Alphabet: shortAlpha,
			Freqs: []map[Residue]int{
				{'A': 0, 'B': 2, 'C': 1},
				{'A': 0, 'B': 0, 'C': 3},
				{'A': 1, 'B': 0, 'C': 2},
			},
		},
		null: &FrequencyProfile{
			Alphabet: shortAlpha,
			Freqs:    []map[Residue]int{{'A': 1, 'B': 2, 'C': 6}},
		},
		seqs: []Sequence{
			{"1", strr("BCC")},
			{"2", strr("BCC")},
			{"3", strr("CCA")},
		},
	},
}

func TestProfile(t *testing.T) {
	for i, test := range tests {
		columns := len(test.seqs[0].Residues)
		fprof := NewFrequencyProfileAlphabet(columns, shortAlpha)
		for _, seq := range test.seqs {
			fprof.Add(seq)
		}
		prof := fprof.Profile(test.null)
		if !reflect.DeepEqual(test.pexpected, prof) {
			t.Logf("Profile test %d failed.", i)
			t.Logf("Expected\n%v\nbut got\n%v\n", test.pexpected, prof)
			t.Fail()
		}
	}
}

func TestFrequencyProfile(t *testing.T) {
	for i, test := range tests {
		columns := len(test.seqs[0].Residues)
		prof := NewFrequencyProfileAlphabet(columns, shortAlpha)
		for _, seq := range test.seqs {
			prof.Add(seq)
		}
		if !reflect.DeepEqual(test.fexpected, prof) {
			t.Logf("Freq test %d failed.", i)
			t.Logf("Expected\n%v\nbut got\n%v\n", test.fexpected, prof)
			t.Fail()
		}
	}
}

func strr(s string) []Residue {
	bs := []byte(s)
	rs := make([]Residue, len(bs))
	for i := range bs {
		rs[i] = Residue(bs[i])
	}
	return rs
}
<commit_msg>Update frequency profile tests to work with new EProbs representation.<commit_after>package seq

import (
	"fmt"
	"math"
	"reflect"
	"testing"
)

var pf = fmt.Printf

var shortAlpha = NewAlphabet('A', 'B', 'C')

func o(freq, coltotal, nullfreq, nulltot int) Prob {
	if freq == 0 || nullfreq == 0 {
		return MinProb
	}
	num := float64(freq) / float64(coltotal)
	den := float64(nullfreq) / float64(nulltot)
	return -Prob(math.Log(num / den))
}

func newEProbs(alpha Alphabet, probs []map[Residue]Prob) []EProbs {
	eps := make([]EProbs, 0, len(probs))
	for _, m := range probs {
		ep := NewEProbs(alpha)
		for r, p := range m {
			ep.Set(r, p)
		}
		eps = append(eps, ep)
	}
	return eps
}

var tests = []struct {
	pexpected *Profile
	fexpected *FrequencyProfile
	null      *FrequencyProfile
	seqs      []Sequence
}{
	{
		pexpected: &Profile{
			Alphabet: shortAlpha,
			Emissions: newEProbs(shortAlpha, []map[Residue]Prob{
				{'A': o(3, 3, 3, 9), 'B': o(0, 3, 3, 9), 'C': o(0, 3, 3, 9)},
				{'A': o(0, 3, 3, 9), 'B': o(3, 3, 3, 9), 'C': o(0, 3, 3, 9)},
				{'A': o(0, 3, 3, 9), 'B': o(0, 3, 3, 9), 'C': o(3, 3, 3, 9)},
			}),
		},
		fexpected: &FrequencyProfile{
			Alphabet: shortAlpha,
			Freqs: []map[Residue]int{
				{'A': 3, 'B': 0, 'C': 0},
				{'A': 0, 'B': 3, 'C': 0},
				{'A': 0, 'B': 0, 'C': 3},
			},
		},
		null: &FrequencyProfile{
			Alphabet: shortAlpha,
			Freqs:    []map[Residue]int{{'A': 3, 'B': 3, 'C': 3}},
		},
		seqs: []Sequence{
			{"1", strr("ABC")},
			{"2", strr("ABC")},
			{"3", strr("ABC")},
		},
	},
	{
		pexpected: &Profile{
			Alphabet: shortAlpha,
			Emissions: newEProbs(shortAlpha, []map[Residue]Prob{
				{'A': o(1, 3, 3, 9), 'B': o(2, 3, 3, 9), 'C': o(0, 3, 3, 9)},
				{'A': o(0, 3, 3, 9), 'B': o(1, 3, 3, 9), 'C': o(2, 3, 3, 9)},
				{'A': o(2, 3, 3, 9), 'B': o(0, 3, 3, 9), 'C': o(1, 3, 3, 9)},
			}),
		},
		fexpected: &FrequencyProfile{
			Alphabet: shortAlpha,
			Freqs: []map[Residue]int{
				{'A': 1, 'B': 2, 'C': 0},
				{'A': 0, 'B': 1, 'C': 2},
				{'A': 2, 'B': 0, 'C': 1},
			},
		},
		null: &FrequencyProfile{
			Alphabet: shortAlpha,
			Freqs:    []map[Residue]int{{'A': 3, 'B': 3, 'C': 3}},
		},
		seqs: []Sequence{
			{"1", strr("ABC")},
			{"2", strr("BCA")},
			{"3", strr("BCA")},
		},
	},
	{
		pexpected: &Profile{
			Alphabet: shortAlpha,
			Emissions: newEProbs(shortAlpha, []map[Residue]Prob{
				{'A': o(0, 3, 1, 9), 'B': o(2, 3, 2, 9), 'C': o(1, 3, 6, 9)},
				{'A': o(0, 3, 1, 9), 'B': o(0, 3, 2, 9), 'C': o(3, 3, 6, 9)},
				{'A': o(1, 3, 1, 9), 'B': o(0, 3, 2, 9), 'C': o(2, 3, 6, 9)},
			}),
		},
		fexpected: &FrequencyProfile{
			Alphabet: shortAlpha,
			Freqs: []map[Residue]int{
				{'A': 0, 'B': 2, 'C': 1},
				{'A': 0, 'B': 0, 'C': 3},
				{'A': 1, 'B': 0, 'C': 2},
			},
		},
		null: &FrequencyProfile{
			Alphabet: shortAlpha,
			Freqs:    []map[Residue]int{{'A': 1, 'B': 2, 'C': 6}},
		},
		seqs: []Sequence{
			{"1", strr("BCC")},
			{"2", strr("BCC")},
			{"3", strr("CCA")},
		},
	},
}

func TestProfile(t *testing.T) {
	for i, test := range tests {
		columns := len(test.seqs[0].Residues)
		fprof := NewFrequencyProfileAlphabet(columns, shortAlpha)
		for _, seq := range test.seqs {
			fprof.Add(seq)
		}
		prof := fprof.Profile(test.null)
		if !reflect.DeepEqual(test.pexpected, prof) {
			t.Logf("Profile test %d failed.", i)
			t.Logf("Expected\n%v\nbut got\n%v\n", test.pexpected, prof)
			t.Fail()
		}
	}
}

func TestFrequencyProfile(t *testing.T) {
	for i, test := range tests {
		columns := len(test.seqs[0].Residues)
		prof := NewFrequencyProfileAlphabet(columns, shortAlpha)
		for _, seq := range test.seqs {
			prof.Add(seq)
		}
		if !reflect.DeepEqual(test.fexpected, prof) {
			t.Logf("Freq test %d failed.", i)
			t.Logf("Expected\n%v\nbut got\n%v\n", test.fexpected, prof)
			t.Fail()
		}
	}
}

func strr(s string) []Residue {
	bs := []byte(s)
	rs := make([]Residue, len(bs))
	for i := range bs {
		rs[i] = Residue(bs[i])
	}
	return rs
}
<|endoftext|>


<commit_before>package purse

import (
	"io/ioutil"
	"os"
	"path/filepath"
	"sync"
)

const (
	ext string = ".sql"
)

// Purse is a key/value collection of loaded SQL files by name : content.
// It is safe for concurrent use by multiple goroutines.
type Purse struct {
	mu    sync.RWMutex
	files map[string]string
}

// New loads all SQL files in the specified directory dir.
//
// A loaded file's contents can be accessed via Get().
//
// Returns an error if the directory does not exist or is not a directory.
func New(dir string) (*Purse, error) {
	f, err := os.Open(dir)
	if err != nil {
		return nil, err
	}
	defer f.Close()

	fis, err := f.Readdir(-1)
	if err != nil {
		return nil, err
	}

	s := &Purse{files: make(map[string]string, 0)}

	for _, fi := range fis {
		if !fi.IsDir() && filepath.Ext(fi.Name()) == ext {
			f, err := os.Open(filepath.Join(dir, fi.Name()))
			if err != nil {
				return nil, err
			}
			defer f.Close()

			fd, err := ioutil.ReadAll(f)
			if err != nil {
				return nil, err
			}
			s.files[fi.Name()] = string(fd)
		}
	}
	return s, nil
}

// Get returns a loaded file's contents and existence of the file by filename.
func (p *Purse) Get(filename string) (v string, ok bool) {
	p.mux.RLock()
	v, ok = p.files[filename]
	p.mux.RUnlock()
	return
}
<commit_msg>changed confusing variable<commit_after>package purse

import (
	"io/ioutil"
	"os"
	"path/filepath"
	"sync"
)

const (
	ext string = ".sql"
)

// Purse is a key/value collection of loaded SQL files by name : content.
// It is safe for concurrent use by multiple goroutines.
type Purse struct {
	mu    sync.RWMutex
	files map[string]string
}

// New loads all SQL files in the specified directory dir.
//
// A loaded file's contents can be accessed via Get().
//
// Returns an error if the directory does not exist or is not a directory.
func New(dir string) (*Purse, error) {
	f, err := os.Open(dir)
	if err != nil {
		return nil, err
	}
	defer f.Close()

	fis, err := f.Readdir(-1)
	if err != nil {
		return nil, err
	}

	s := &Purse{files: make(map[string]string, 0)}

	for _, fi := range fis {
		if !fi.IsDir() && filepath.Ext(fi.Name()) == ext {
			f, err := os.Open(filepath.Join(dir, fi.Name()))
			if err != nil {
				return nil, err
			}
			defer f.Close()

			fd, err := ioutil.ReadAll(f)
			if err != nil {
				return nil, err
			}
			s.files[fi.Name()] = string(fd)
		}
	}
	return s, nil
}

// Get returns a loaded file's contents and existence of the file by filename.
func (p *Purse) Get(filename string) (v string, ok bool) {
	p.mu.RLock()
	v, ok = p.files[filename]
	p.mu.RUnlock()
	return
}
<|endoftext|>


<commit_before>package resource

import (
	"fmt"
	"io"
	"os"
	"os/user"
	"path/filepath"

	"github.com/dnaeon/gru/utils"
	"github.com/hashicorp/hcl"
	"github.com/hashicorp/hcl/hcl/ast"
	"github.com/imdario/mergo"
)

// Name and description of the resource
const fileResourceType = "file"
const fileResourceDesc = "manages files"

// BaseFileResource is the base resource for managing files
type BaseFileResource struct {
	// Path to the file
	Path string `hcl:"path"`

	// Permission bits to set on the file
	Mode int `hcl:"mode"`

	// Owner of the file
	Owner string `hcl:"owner"`

	// Group of the file
	Group string `hcl:"group"`

	// Source file to use when creating/updating the file
	Source string `hcl:"source"`
}

// FileResource is a resource which manages files
type FileResource struct {
	BaseResource     `hcl:",squash"`
	BaseFileResource `hcl:",squash"`
}

// NewFileResource creates a new resource for managing files
func NewFileResource(title string, obj *ast.ObjectItem) (Resource, error) {
	// Defaults for owner and group
	currentUser, err := user.Current()
	if err != nil {
		return nil, err
	}

	currentGroup, err := user.LookupGroupId(currentUser.Gid)
	if err != nil {
		return nil, err
	}

	defaultOwner := currentUser.Username
	defaultGroup := currentGroup.Name

	// Resource defaults
	defaults := FileResource{
		BaseResource: BaseResource{
			Title: title,
			Type:  fileResourceType,
			State: StatePresent,
		},
		BaseFileResource: BaseFileResource{
			Path:  title,
			Mode:  0644,
			Owner: defaultOwner,
			Group: defaultGroup,
		},
	}

	var fr FileResource
	err = hcl.DecodeObject(&fr, obj)
	if err != nil {
		return nil, err
	}

	// Merge the decoded object with the resource defaults
	err = mergo.Merge(&fr, defaults)

	return &fr, err
}

// Evaluate evaluates the file resource
func (fr *FileResource) Evaluate(w io.Writer, opts *Options) (State, error) {
	rs := State{
		Current: StateUnknown,
		Want:    fr.State,
		Update:  false,
	}

	// The file we manage
	dst := utils.NewFileUtil(fr.Path)

	// File does not exist
	fi, err := os.Stat(fr.Path)
	if os.IsNotExist(err) {
		rs.Current = StateAbsent
		return rs, nil
	} else {
		rs.Current = StatePresent
	}

	// Ensure that the file we manage is a regular file
	if !fi.Mode().IsRegular() {
		return rs, fmt.Errorf("%s is not a regular file", fr.Path)
	}

	// Check file content
	if fr.Source != "" {
		srcPath := filepath.Join(opts.SiteDir, "data", fr.Source)
		same, err := dst.SameContentWith(srcPath)
		if err != nil {
			return rs, err
		}
		if !same {
			fr.Printf(w, "content is out of date\n")
			rs.Update = true
		}
	}

	// Check file permissions
	mode, err := dst.Mode()
	if err != nil {
		return rs, err
	}

	if mode.Perm() != os.FileMode(fr.Mode) {
		fr.Printf(w, "permissions are out of date\n")
		rs.Update = true
	}

	// Check ownership
	owner, err := dst.Owner()
	if err != nil {
		return rs, err
	}

	if fr.Owner != owner.User.Username || fr.Group != owner.Group.Name {
		fr.Printf(w, "owner is out of date\n")
		rs.Update = true
	}

	return rs, nil
}

// Create creates the file managed by the resource
func (fr *FileResource) Create(w io.Writer, opts *Options) error {
	dst := utils.NewFileUtil(fr.Path)
	fr.Printf(w, "creating file\n")

	// Set file content
	switch {
	case fr.Source == "" && dst.Exists():
		// Do nothing
		break
	case fr.Source == "" && !dst.Exists():
		// Create an empty file
		if _, err := os.Create(fr.Path); err != nil {
			return err
		}
	case fr.Source != "" && dst.Exists():
		// File exists and we have a source file
		srcPath := filepath.Join(opts.SiteDir, "data", fr.Source)
		if err := dst.CopyFrom(srcPath); err != nil {
			return err
		}
	}

	// Set file owner
	if err := dst.SetOwner(fr.Owner, fr.Group); err != nil {
		return err
	}

	// Set file permissions
	return dst.Chmod(os.FileMode(fr.Mode))
}

// Delete deletes the file
func (fr *FileResource) Delete(w io.Writer, opts *Options) error {
	fr.Printf(w, "removing file\n")
	dst := utils.NewFileUtil(fr.Path)

	return dst.Remove()
}

// Update updates the file managed by the resource
func (fr *FileResource) Update(w io.Writer, opts *Options) error {
	dst := utils.NewFileUtil(fr.Path)

	// Update file content if needed
	if fr.Source != "" {
		srcPath := filepath.Join(opts.SiteDir, "data", fr.Source)
		same, err := dst.SameContentWith(srcPath)
		if err != nil {
			return err
		}

		if !same {
			srcFile := utils.NewFileUtil(srcPath)
			srcMd5, err := srcFile.Md5()
			if err != nil {
				return err
			}

			fr.Printf(w, "updating content to md5:%s\n", srcMd5)
			if err := dst.CopyFrom(srcPath); err != nil {
				return err
			}
		}
	}

	// Fix permissions if needed
	mode, err := dst.Mode()
	if err != nil {
		return err
	}

	if mode.Perm() != os.FileMode(fr.Mode) {
		fr.Printf(w, "setting permissions to %#o\n", fr.Mode)
		dst.Chmod(os.FileMode(fr.Mode))
	}

	// Fix ownership if needed
	owner, err := dst.Owner()
	if err != nil {
		return err
	}

	if fr.Owner != owner.User.Username || fr.Group != owner.Group.Name {
		fr.Printf(w, "setting owner to %s:%s\n", fr.Owner, fr.Group)
		if err := dst.SetOwner(fr.Owner, fr.Group); err != nil {
			return err
		}
	}

	return nil
}

func init() {
	item := RegistryItem{
		Name:        fileResourceType,
		Description: fileResourceDesc,
		Provider:    NewFileResource,
	}

	Register(item)
}
<commit_msg>Implemenet basic support for different file types in resource.FileResource<commit_after>package resource

import (
	"fmt"
	"io"
	"os"
	"os/user"
	"path/filepath"

	"github.com/dnaeon/gru/utils"
	"github.com/hashicorp/hcl"
	"github.com/hashicorp/hcl/hcl/ast"
	"github.com/imdario/mergo"
)

// Name and description of the resource
const fileResourceType = "file"
const fileResourceDesc = "manages files"

// File types
const fileTypeRegular = "file"
const fileTypeDirectory = "directory"

// BaseFileResource is the base resource for managing files
type BaseFileResource struct {
	// Path to the file
	Path string `hcl:"path"`

	// Permission bits to set on the file
	Mode int `hcl:"mode"`

	// Owner of the file
	Owner string `hcl:"owner"`

	// Group of the file
	Group string `hcl:"group"`

	// Source file to use when creating/updating the file
	Source string `hcl:"source"`

	// The file type we manage
	FileType string `hcl:"type"`
}

// FileResource is a resource which manages files
type FileResource struct {
	BaseResource     `hcl:",squash"`
	BaseFileResource `hcl:",squash"`
}

// NewFileResource creates a new resource for managing files
func NewFileResource(title string, obj *ast.ObjectItem) (Resource, error) {
	// Defaults for owner and group
	currentUser, err := user.Current()
	if err != nil {
		return nil, err
	}

	currentGroup, err := user.LookupGroupId(currentUser.Gid)
	if err != nil {
		return nil, err
	}

	defaultOwner := currentUser.Username
	defaultGroup := currentGroup.Name

	// Resource defaults
	defaults := FileResource{
		BaseResource: BaseResource{
			Title: title,
			Type:  fileResourceType,
			State: StatePresent,
		},
		BaseFileResource: BaseFileResource{
			Path:     title,
			Mode:     0644,
			Owner:    defaultOwner,
			Group:    defaultGroup,
			FileType: fileTypeRegular,
		},
	}

	var fr FileResource
	err = hcl.DecodeObject(&fr, obj)
	if err != nil {
		return nil, err
	}

	// Merge the decoded object with the resource defaults
	err = mergo.Merge(&fr, defaults)
	if err != nil {
		return nil, err
	}

	if fr.FileType != fileTypeRegular || fr.FileType != fileTypeDirectory {
		return nil, fmt.Errorf("Unknown file type '%s'", fr.FileType)
	}

	return &fr, nil
}

// Evaluate evaluates the file resource
func (fr *FileResource) Evaluate(w io.Writer, opts *Options) (State, error) {
	rs := State{
		Current: StateUnknown,
		Want:    fr.State,
		Update:  false,
	}

	// The file we manage
	dst := utils.NewFileUtil(fr.Path)

	// File does not exist
	if !dst.Exists() {
		rs.Current = StateAbsent
		return rs, nil
	} else {
		rs.Current = StatePresent
	}

	fi, err := os.Stat(fr.Path)
	if err != nil {
		return rs, err
	}

	// Check the target file we manage
	switch fr.FileType {
	case fileTypeRegular:
		if !fi.Mode().IsRegular() {
			return rs, fmt.Errorf("%s exists, but is not a regular file", fr.Path)
		}
	case fileTypeDirectory:
		if !fi.IsDir() {
			return rs, fmt.Errorf("%s exists, but is not a directory", fr.Path)
		}
	}

	// Check file content
	if fr.Source != "" {
		srcPath := filepath.Join(opts.SiteDir, "data", fr.Source)
		same, err := dst.SameContentWith(srcPath)
		if err != nil {
			return rs, err
		}
		if !same {
			fr.Printf(w, "content is out of date\n")
			rs.Update = true
		}
	}

	// Check file permissions
	mode, err := dst.Mode()
	if err != nil {
		return rs, err
	}

	if mode.Perm() != os.FileMode(fr.Mode) {
		fr.Printf(w, "permissions are out of date\n")
		rs.Update = true
	}

	// Check ownership
	owner, err := dst.Owner()
	if err != nil {
		return rs, err
	}

	if fr.Owner != owner.User.Username || fr.Group != owner.Group.Name {
		fr.Printf(w, "owner is out of date\n")
		rs.Update = true
	}

	return rs, nil
}

// createRegularFile creates the file and content managed by the resource
func (fr *FileResource) createRegularFile(opts *Options) error {
	dst := utils.NewFileUtil(fr.Path)

	switch {
	case fr.Source == "" && dst.Exists():
		// We have no source, do nothing
		break
	case fr.Source == "" && !dst.Exists():
		// Create an empty file
		if _, err := os.Create(fr.Path); err != nil {
			return err
		}
	case fr.Source != "" && dst.Exists():
		// File exists and we have a source file
		srcPath := filepath.Join(opts.SiteDir, "data", fr.Source)
		if err := dst.CopyFrom(srcPath); err != nil {
			return err
		}
	}

	return nil
}

// createDirectory creates the directory and content managed by the resource
func (fr *FileResource) createDirectory(opts *Options) error {
	if err := os.Mkdir(fr.Path, os.FileMode(fr.Mode)); err != nil {
		return err
	}

	return nil
}

// Create creates the file managed by the resource
func (fr *FileResource) Create(w io.Writer, opts *Options) error {
	dst := utils.NewFileUtil(fr.Path)
	fr.Printf(w, "creating resource\n")

	// Set content
	switch fr.FileType {
	case fileTypeRegular:
		if err := fr.createRegularFile(opts); err != nil {
			return err
		}
	case fileTypeDirectory:
		if err := fr.createDirectory(opts); err != nil {
			return err
		}
	}

	// Set file owner
	if err := dst.SetOwner(fr.Owner, fr.Group); err != nil {
		return err
	}

	// Set file permissions
	return dst.Chmod(os.FileMode(fr.Mode))
}

// Delete deletes the file
func (fr *FileResource) Delete(w io.Writer, opts *Options) error {
	fr.Printf(w, "removing file\n")
	dst := utils.NewFileUtil(fr.Path)

	return dst.Remove()
}

// Update updates the file managed by the resource
func (fr *FileResource) Update(w io.Writer, opts *Options) error {
	dst := utils.NewFileUtil(fr.Path)

	// Update file content if needed
	if fr.Source != "" {
		srcPath := filepath.Join(opts.SiteDir, "data", fr.Source)
		same, err := dst.SameContentWith(srcPath)
		if err != nil {
			return err
		}

		if !same {
			srcFile := utils.NewFileUtil(srcPath)
			srcMd5, err := srcFile.Md5()
			if err != nil {
				return err
			}

			fr.Printf(w, "updating content to md5:%s\n", srcMd5)
			if err := dst.CopyFrom(srcPath); err != nil {
				return err
			}
		}
	}

	// Fix permissions if needed
	mode, err := dst.Mode()
	if err != nil {
		return err
	}

	if mode.Perm() != os.FileMode(fr.Mode) {
		fr.Printf(w, "setting permissions to %#o\n", fr.Mode)
		dst.Chmod(os.FileMode(fr.Mode))
	}

	// Fix ownership if needed
	owner, err := dst.Owner()
	if err != nil {
		return err
	}

	if fr.Owner != owner.User.Username || fr.Group != owner.Group.Name {
		fr.Printf(w, "setting owner to %s:%s\n", fr.Owner, fr.Group)
		if err := dst.SetOwner(fr.Owner, fr.Group); err != nil {
			return err
		}
	}

	return nil
}

func init() {
	item := RegistryItem{
		Name:        fileResourceType,
		Description: fileResourceDesc,
		Provider:    NewFileResource,
	}

	Register(item)
}
<|endoftext|>


<commit_before>package main

import (
	"github.com/AsynkronIT/gam/actor"
	"github.com/AsynkronIT/gam/remoting"

	"log"
	"sync"

	"time"
)

// import "runtime/pprof"
import "github.com/AsynkronIT/gam/examples/remotebenchmark/messages"
import "runtime"

type localActor struct {
	count        int
	wgStop       *sync.WaitGroup
	start        time.Time
	messageCount int
}

func (state *localActor) Receive(context actor.Context) {
	switch context.Message().(type) {
	case *messages.Pong:
		state.count++
		if state.count%50000 == 0 {
			log.Println(state.count)
		}
		if state.count == state.messageCount {
			state.wgStop.Done()
		}
	}
}

type FakeMessage struct {
	pid *actor.PID
}

func newLocalActor(stop *sync.WaitGroup, messageCount int) actor.Producer {
	return func() actor.Actor {
		return &localActor{
			wgStop:       stop,
			messageCount: messageCount,
		}
	}
}

func main() {
	runtime.GOMAXPROCS(runtime.NumCPU())

	var wg sync.WaitGroup

	messageCount := 1000000

	remoting.Start("127.0.0.1:8081")

	props := actor.
		FromProducer(newLocalActor(&wg, messageCount)).
		WithMailbox(actor.NewBoundedMailbox(1000, 10000))

	pid := actor.Spawn(props)

	remote := actor.NewPID("127.0.0.1:8080", "remote")
	res, _ := remote.AskFuture(&messages.StartRemote{})
	defer res.Stop()
	res.Wait()
	wg.Add(1)

	start := time.Now()
	log.Println("Starting to send")

	message := &messages.Ping{}
	for i := 0; i < messageCount; i++ {
		remote.Ask(message, pid)
	}

	wg.Wait()
	elapsed := time.Since(start)
	log.Printf("Elapsed %s", elapsed)

	x := int(float32(messageCount*2) / (float32(elapsed) / float32(time.Second)))
	log.Printf("Msg per sec %v", x)

	// f, err := os.Create("memprofile")
	// if err != nil {
	// 	log.Fatal(err)
	// }
	// pprof.WriteHeapProfile(f)
	// f.Close()
}
<commit_msg>optimize remote settings<commit_after>package main

import (
	"github.com/AsynkronIT/gam/actor"
	"github.com/AsynkronIT/gam/remoting"

	"log"
	"sync"

	"time"
)

// import "runtime/pprof"
import "github.com/AsynkronIT/gam/examples/remotebenchmark/messages"
import "runtime"

type localActor struct {
	count        int
	wgStop       *sync.WaitGroup
	messageCount int
}

func (state *localActor) Receive(context actor.Context) {
	switch context.Message().(type) {
	case *messages.Pong:
		state.count++
		if state.count%50000 == 0 {
			log.Println(state.count)
		}
		if state.count == state.messageCount {
			state.wgStop.Done()
		}
	}
}

type FakeMessage struct {
	pid *actor.PID
}

func newLocalActor(stop *sync.WaitGroup, messageCount int) actor.Producer {
	return func() actor.Actor {
		return &localActor{
			wgStop:       stop,
			messageCount: messageCount,
		}
	}
}

func main() {
	runtime.GOMAXPROCS(runtime.NumCPU())

	var wg sync.WaitGroup

	messageCount := 1000000

	remoting.Start("127.0.0.1:8081")

	props := actor.
		FromProducer(newLocalActor(&wg, messageCount)).
		WithMailbox(actor.NewBoundedMailbox(1000, 1000))

	pid := actor.Spawn(props)

	remote := actor.NewPID("127.0.0.1:8080", "remote")
	res, _ := remote.AskFuture(&messages.StartRemote{})
	defer res.Stop()
	res.Wait()
	wg.Add(1)

	start := time.Now()
	log.Println("Starting to send")

	message := &messages.Ping{}
	for i := 0; i < messageCount; i++ {
		remote.Ask(message, pid)
	}

	wg.Wait()
	elapsed := time.Since(start)
	log.Printf("Elapsed %s", elapsed)

	x := int(float32(messageCount*2) / (float32(elapsed) / float32(time.Second)))
	log.Printf("Msg per sec %v", x)

	// f, err := os.Create("memprofile")
	// if err != nil {
	// 	log.Fatal(err)
	// }
	// pprof.WriteHeapProfile(f)
	// f.Close()
}
<|endoftext|>


<commit_before>package grep

import (
	"github.com/monochromegane/the_platinum_searcher/search/file"
	"github.com/monochromegane/the_platinum_searcher/search/option"
	"github.com/monochromegane/the_platinum_searcher/search/print"
	"testing"
)

type Assert struct {
	path, pattern, fileType, match string
}

var Asserts = []Assert{
	Assert{"ascii.txt", "go", file.ASCII, "go test"},
	Assert{"ja/euc-jp.txt", "go", file.EUCJP, "go テスト"},
	Assert{"ja/shift_jis.txt", "go", file.SHIFTJIS, "go テスト"},
	Assert{"ja/utf8.txt", "go", file.UTF8, "go テスト"},
	Assert{"ja/broken_euc-jp.txt", "go", file.EUCJP, "go テスト"},
	Assert{"ja/broken_shift_jis.txt", "go", file.SHIFTJIS, "go テスト"},
	Assert{"ja/broken_utf8.txt", "go", file.UTF8, "go テスト"},
}

func TestGrep(t *testing.T) {

	for _, g := range Asserts {
		in := make(chan *Params)
		out := make(chan *print.Params)
		grepper := Grepper{in, out, &option.Option{Proc: 1}}

                sem := make(chan bool, 1)
                sem <- true
		go grepper.Grep("../../files/" + g.path, g.fileType, g.pattern, sem)
		o := <-out
		if o.Path != "../../files/"+g.path {
			t.Errorf("It should be equal ../../files/%s.", g.path)
		}
		if o.Matches[0].Match != g.match {
			t.Errorf("%s should be equal %s", g.path, g.match)
		}
	}
}

func receive(in chan *Params, params *Params) {
	in <- params
	close(in)
}
<commit_msg>go fmt<commit_after>package grep

import (
	"github.com/monochromegane/the_platinum_searcher/search/file"
	"github.com/monochromegane/the_platinum_searcher/search/option"
	"github.com/monochromegane/the_platinum_searcher/search/print"
	"testing"
)

type Assert struct {
	path, pattern, fileType, match string
}

var Asserts = []Assert{
	Assert{"ascii.txt", "go", file.ASCII, "go test"},
	Assert{"ja/euc-jp.txt", "go", file.EUCJP, "go テスト"},
	Assert{"ja/shift_jis.txt", "go", file.SHIFTJIS, "go テスト"},
	Assert{"ja/utf8.txt", "go", file.UTF8, "go テスト"},
	Assert{"ja/broken_euc-jp.txt", "go", file.EUCJP, "go テスト"},
	Assert{"ja/broken_shift_jis.txt", "go", file.SHIFTJIS, "go テスト"},
	Assert{"ja/broken_utf8.txt", "go", file.UTF8, "go テスト"},
}

func TestGrep(t *testing.T) {

	for _, g := range Asserts {
		in := make(chan *Params)
		out := make(chan *print.Params)
		grepper := Grepper{in, out, &option.Option{Proc: 1}}

		sem := make(chan bool, 1)
		sem <- true
		go grepper.Grep("../../files/"+g.path, g.fileType, g.pattern, sem)
		o := <-out
		if o.Path != "../../files/"+g.path {
			t.Errorf("It should be equal ../../files/%s.", g.path)
		}
		if o.Matches[0].Match != g.match {
			t.Errorf("%s should be equal %s", g.path, g.match)
		}
	}
}

func receive(in chan *Params, params *Params) {
	in <- params
	close(in)
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"fmt"
	"go/build"
	"go/format"
	"io"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"strings"
	"time"
)

func operationEmbedGo(pkg *build.Package) {

	boxMap := findBoxes(pkg)

	// notify user when no calls to rice.FindBox are made (is this an error and therefore os.Exit(1) ?
	if len(boxMap) == 0 {
		fmt.Println("no calls to rice.FindBox() found")
		return
	}

	verbosef("\n")

	for boxname := range boxMap {
		// find path and filename for this box
		boxPath := filepath.Join(pkg.Dir, boxname)
		boxFilename := strings.Replace(boxname, "/", "-", -1)
		boxFilename = strings.Replace(boxFilename, "..", "back", -1)
		boxFilename = boxFilename + `.rice-box.go`

		// verbose info
		verbosef("embedding box '%s'\n", boxname)
		verbosef("\tto file %s\n", boxFilename)

		// create box datastructure (used by template)
		box := &boxDataType{
			Package: pkg.Name,
			BoxName: boxname,
			UnixNow: time.Now().Unix(),
			Files:   make([]*fileDataType, 0),
			Dirs:    make(map[string]*dirDataType),
		}

		// fill box datastructure with file data
		filepath.Walk(boxPath, func(path string, info os.FileInfo, err error) error {
			if err != nil {
				fmt.Printf("error walking box: %s\n", err)
				os.Exit(1)
			}

			filename := strings.TrimPrefix(path, boxPath)
			filename = strings.Replace(filename, "\\", "/", -1)
			filename = strings.TrimPrefix(filename, "/")
			if info.IsDir() {
				dirData := &dirDataType{
					Identifier: "dir" + nextIdentifier(),
					FileName:   filename,
					ModTime:    info.ModTime().Unix(),
					ChildFiles: make([]*fileDataType, 0),
					ChildDirs:  make([]*dirDataType, 0),
				}
				verbosef("\tincludes dir: '%s'\n", dirData.FileName)
				box.Dirs[dirData.FileName] = dirData

				// add tree entry (skip for root, it'll create a recursion)
				if dirData.FileName != "" {
					pathParts := strings.Split(dirData.FileName, "/")
					parentDir := box.Dirs[strings.Join(pathParts[:len(pathParts)-1], "/")]
					parentDir.ChildDirs = append(parentDir.ChildDirs, dirData)
				}
			} else {
				fileData := &fileDataType{
					Identifier: "file" + nextIdentifier(),
					FileName:   filename,
					ModTime:    info.ModTime().Unix(),
				}
				verbosef("\tincludes file: '%s'\n", fileData.FileName)
				fileData.Content, err = ioutil.ReadFile(path)
				if err != nil {
					fmt.Printf("error reading file content while walking box: %s\n", err)
					os.Exit(1)
				}
				box.Files = append(box.Files, fileData)

				// add tree entry
				pathParts := strings.Split(fileData.FileName, "/")
				parentDir := box.Dirs[strings.Join(pathParts[:len(pathParts)-1], "/")]
				parentDir.ChildFiles = append(parentDir.ChildFiles, fileData)
			}
			return nil
		})

		embedSourceUnformated := bytes.NewBuffer(make([]byte, 0))

		// execute template to buffer
		err := tmplEmbeddedBox.Execute(embedSourceUnformated, box)
		if err != nil {
			log.Printf("error writing embedded box to file (template execute): %s\n", err)
			os.Exit(1)
		}

		// format the source code
		embedSource, err := format.Source(embedSourceUnformated.Bytes())
		if err != nil {
			log.Printf("error formatting embedSource: %s\n", err)
			os.Exit(1)
		}

		// create go file for box
		boxFile, err := os.Create(filepath.Join(pkg.Dir, boxFilename))
		if err != nil {
			log.Printf("error creating embedded box file: %s\n", err)
			os.Exit(1)
		}
		defer boxFile.Close()

		// write source to file
		_, err = io.Copy(boxFile, bytes.NewBuffer(embedSource))
		if err != nil {
			log.Printf("error writing embedSource to file: %s\n", err)
			os.Exit(1)
		}
	}
}
<commit_msg>Handles case where a box is a symbolic link<commit_after>package main

import (
	"bytes"
	"fmt"
	"go/build"
	"go/format"
	"io"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"strings"
	"time"
)

func operationEmbedGo(pkg *build.Package) {

	boxMap := findBoxes(pkg)

	// notify user when no calls to rice.FindBox are made (is this an error and therefore os.Exit(1) ?
	if len(boxMap) == 0 {
		fmt.Println("no calls to rice.FindBox() found")
		return
	}

	verbosef("\n")

	for boxname := range boxMap {
		// find path and filename for this box
		boxPath := filepath.Join(pkg.Dir, boxname)

		// Check to see if the path for the box is a symbolic link.  If so, simply
		// box what the symbolic link points to.  Note: the filepath.Walk function
		// will NOT follow any nested symbolic links.  This only handles the case
		// where the root of the box is a symbolic link.
		symPath, serr := os.Readlink(boxPath)
		if serr == nil {
			boxPath = symPath
		}

		boxFilename := strings.Replace(boxname, "/", "-", -1)
		boxFilename = strings.Replace(boxFilename, "..", "back", -1)
		boxFilename = boxFilename + `.rice-box.go`

		// verbose info
		verbosef("embedding box '%s'\n", boxname)
		verbosef("\tto file %s\n", boxFilename)

		// create box datastructure (used by template)
		box := &boxDataType{
			Package: pkg.Name,
			BoxName: boxname,
			UnixNow: time.Now().Unix(),
			Files:   make([]*fileDataType, 0),
			Dirs:    make(map[string]*dirDataType),
		}

		// fill box datastructure with file data
		filepath.Walk(boxPath, func(path string, info os.FileInfo, err error) error {
			if err != nil {
				fmt.Printf("error walking box: %s\n", err)
				os.Exit(1)
			}

			filename := strings.TrimPrefix(path, boxPath)
			filename = strings.Replace(filename, "\\", "/", -1)
			filename = strings.TrimPrefix(filename, "/")
			if info.IsDir() {
				dirData := &dirDataType{
					Identifier: "dir" + nextIdentifier(),
					FileName:   filename,
					ModTime:    info.ModTime().Unix(),
					ChildFiles: make([]*fileDataType, 0),
					ChildDirs:  make([]*dirDataType, 0),
				}
				verbosef("\tincludes dir: '%s'\n", dirData.FileName)
				box.Dirs[dirData.FileName] = dirData

				// add tree entry (skip for root, it'll create a recursion)
				if dirData.FileName != "" {
					pathParts := strings.Split(dirData.FileName, "/")
					parentDir := box.Dirs[strings.Join(pathParts[:len(pathParts)-1], "/")]
					parentDir.ChildDirs = append(parentDir.ChildDirs, dirData)
				}
			} else {
				fileData := &fileDataType{
					Identifier: "file" + nextIdentifier(),
					FileName:   filename,
					ModTime:    info.ModTime().Unix(),
				}
				verbosef("\tincludes file: '%s'\n", fileData.FileName)
				fileData.Content, err = ioutil.ReadFile(path)
				if err != nil {
					fmt.Printf("error reading file content while walking box: %s\n", err)
					os.Exit(1)
				}
				box.Files = append(box.Files, fileData)

				// add tree entry
				pathParts := strings.Split(fileData.FileName, "/")
				parentDir := box.Dirs[strings.Join(pathParts[:len(pathParts)-1], "/")]
				parentDir.ChildFiles = append(parentDir.ChildFiles, fileData)
			}
			return nil
		})

		embedSourceUnformated := bytes.NewBuffer(make([]byte, 0))

		// execute template to buffer
		err := tmplEmbeddedBox.Execute(embedSourceUnformated, box)
		if err != nil {
			log.Printf("error writing embedded box to file (template execute): %s\n", err)
			os.Exit(1)
		}

		// format the source code
		embedSource, err := format.Source(embedSourceUnformated.Bytes())
		if err != nil {
			log.Printf("error formatting embedSource: %s\n", err)
			os.Exit(1)
		}

		// create go file for box
		boxFile, err := os.Create(filepath.Join(pkg.Dir, boxFilename))
		if err != nil {
			log.Printf("error creating embedded box file: %s\n", err)
			os.Exit(1)
		}
		defer boxFile.Close()

		// write source to file
		_, err = io.Copy(boxFile, bytes.NewBuffer(embedSource))
		if err != nil {
			log.Printf("error writing embedSource to file: %s\n", err)
			os.Exit(1)
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"github.com/ansible-semaphore/semaphore/db"
	"github.com/ansible-semaphore/semaphore/db/sql"
	"github.com/ansible-semaphore/semaphore/util"
	"github.com/snikch/goodman/transaction"
	"math/rand"
	"os"
	"time"
)

var tablesShouldBeTruncated = [...]string{
	"access_key",
	"event",
	"user__token",
	"project",
	"task__output",
	"task",
	"session",
	"project__environment",
	"project__inventory",
	"project__repository",
	"project__template",
	"project__schedule",
	"project__user",
	"user",
	"project__view",
}

// Test Runner User
func addTestRunnerUser() {
	uid := getUUID()
	testRunnerUser = &db.User{
		Username: "ITU-" + uid,
		Name:     "ITU-" + uid,
		Email:    uid + "@semaphore.test",
		Created:  db.GetParsedTime(time.Now()),
		Admin:    true,
	}

	dbConnect()
	defer store.Sql().Db.Close()

	truncateAll()

	newUser, err := store.CreateUserWithoutPassword(*testRunnerUser)

	if err != nil {
		panic(err)
	}

	testRunnerUser.ID = newUser.ID

	addToken(adminToken, testRunnerUser.ID)
}

func truncateAll() {
	tx, err := store.Sql().Begin()
	if err != nil {
		panic(err)
	}

	_, err = tx.Exec("SET FOREIGN_KEY_CHECKS = 0")
	if err == nil {
		for _, tableName := range tablesShouldBeTruncated {
			tx.Exec("TRUNCATE TABLE " + tableName)
		}
		tx.Exec("SET FOREIGN_KEY_CHECKS = 1")
	}

	if err := tx.Commit(); err != nil {
		panic(err)
	}
}

func removeTestRunnerUser(transactions []*transaction.Transaction) {
	dbConnect()
	defer store.Sql().Db.Close()
	deleteToken(adminToken, testRunnerUser.ID)
	deleteObject(testRunnerUser)
}

// Parameter Substitution
func setupObjectsAndPaths(t *transaction.Transaction) {
	alterRequestPath(t)
	alterRequestBody(t)
}

// Object Lifecycle
func addUserProjectRelation(pid int, user int) {
	_, err := store.CreateProjectUser(db.ProjectUser{
		ProjectID: pid,
		UserID:    user,
		Admin:     true,
	})
	if err != nil {
		panic(err)
	}
}

func deleteUserProjectRelation(pid int, user int) {
	err := store.DeleteProjectUser(pid, user)
	//_, err := store.Sql().Exec("delete from project__user where project_id=? and user_id=?", strconv.Itoa(pid), strconv.Itoa(user))
	if err != nil {
		panic(err)
	}
}

func addAccessKey(pid *int) *db.AccessKey {
	uid := getUUID()
	secret := "5up3r53cr3t\n"
	key := db.AccessKey{
		Name:      "ITK-" + uid,
		Type:      "ssh",
		Secret:    &secret,
		ProjectID: pid,
	}
	if err := store.Sql().Insert(&key); err != nil {
		panic(err)
	}
	return &key
}

func addProject() *db.Project {
	uid := getUUID()
	project := db.Project{
		Name:    "ITP-" + uid,
		Created: time.Now(),
	}
	project, err := store.CreateProject(project)
	if err != nil {
		panic(err)
	}
	return &project
}

func addUser() *db.User {
	uid := getUUID()
	user := db.User{
		Created:  time.Now(),
		Username: "ITU-" + uid,
		Email:    "test@semaphore." + uid,
	}

	user, err := store.CreateUserWithoutPassword(user)

	if err != nil {
		panic(err)
	}
	return &user
}

func addView() *db.View {
	view, err := store.CreateView(db.View{
		ProjectID: userProject.ID,
		Title:     "Test",
		Position:  1,
	})

	if err != nil {
		panic(err)
	}

	return &view
}

func addSchedule() *db.Schedule {
	schedule, err := store.CreateSchedule(db.Schedule{
		TemplateID: int(templateID),
		CronFormat: "* * * 1 *",
		ProjectID:  userProject.ID,
	})

	if err != nil {
		panic(err)
	}

	return &schedule
}

func addTask() *db.Task {
	t := db.Task{
		ProjectID:  userProject.ID,
		TemplateID: int(templateID),
		Status:     "testing",
		UserID:     &userPathTestUser.ID,
		Created:    db.GetParsedTime(time.Now()),
	}
	if err := store.Sql().Insert(&t); err != nil {
		fmt.Println("error during insertion of task:")
		if j, err := json.Marshal(t); err == nil {
			fmt.Println(string(j))
		} else {
			fmt.Println("can not stringify task object")
		}
		panic(err)
	}
	return &t
}

func deleteObject(i interface{}) {
	_, err := store.Sql().Delete(i)
	if err != nil {
		panic(err)
	}
}

// Token Handling
func addToken(tok string, user int) {
	token := db.APIToken{
		ID:      tok,
		Created: time.Now(),
		UserID:  user,
		Expired: false,
	}
	if err := store.Sql().Insert(&token); err != nil {
		panic(err)
	}
}

func deleteToken(tok string, user int) {
	token := db.APIToken{
		ID:     tok,
		UserID: user,
	}
	deleteObject(&token)
}

// HELPERS
var r *rand.Rand
var randSetup = false

func getUUID() string {
	if !randSetup {
		r = rand.New(rand.NewSource(time.Now().UnixNano()))
		randSetup = true
	}
	return randomString(8)
}
func randomString(strlen int) string {
	const chars = "abcdefghijklmnopqrstuvwxyz0123456789"
	result := ""
	for i := 0; i < strlen; i++ {
		index := r.Intn(len(chars))
		result += chars[index : index+1]
	}
	return result
}

func loadConfig() {
	cwd, _ := os.Getwd()
	file, _ := os.Open(cwd + "/.dredd/config.json")
	if err := json.NewDecoder(file).Decode(&util.Config); err != nil {
		fmt.Println("Could not decode configuration!")
		panic(err)
	}
}

var store sql.SqlDb

func dbConnect() {
	store = sql.SqlDb{}

	if err := store.Connect(); err != nil {
		panic(err)
	}
}

func stringInSlice(a string, list []string) (int, bool) {
	for k, b := range list {
		if b == a {
			return k, true
		}
	}
	return 0, false
}

func printError(err error) {
	if err != nil {
		//fmt.Println(err)
		panic(err)
	}
}
<commit_msg>test: create access key via store<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"github.com/ansible-semaphore/semaphore/db"
	"github.com/ansible-semaphore/semaphore/db/sql"
	"github.com/ansible-semaphore/semaphore/util"
	"github.com/snikch/goodman/transaction"
	"math/rand"
	"os"
	"time"
)

var tablesShouldBeTruncated = [...]string{
	"access_key",
	"event",
	"user__token",
	"project",
	"task__output",
	"task",
	"session",
	"project__environment",
	"project__inventory",
	"project__repository",
	"project__template",
	"project__schedule",
	"project__user",
	"user",
	"project__view",
}

// Test Runner User
func addTestRunnerUser() {
	uid := getUUID()
	testRunnerUser = &db.User{
		Username: "ITU-" + uid,
		Name:     "ITU-" + uid,
		Email:    uid + "@semaphore.test",
		Created:  db.GetParsedTime(time.Now()),
		Admin:    true,
	}

	dbConnect()
	defer store.Sql().Db.Close()

	truncateAll()

	newUser, err := store.CreateUserWithoutPassword(*testRunnerUser)

	if err != nil {
		panic(err)
	}

	testRunnerUser.ID = newUser.ID

	addToken(adminToken, testRunnerUser.ID)
}

func truncateAll() {
	tx, err := store.Sql().Begin()
	if err != nil {
		panic(err)
	}

	_, err = tx.Exec("SET FOREIGN_KEY_CHECKS = 0")
	if err == nil {
		for _, tableName := range tablesShouldBeTruncated {
			tx.Exec("TRUNCATE TABLE " + tableName)
		}
		tx.Exec("SET FOREIGN_KEY_CHECKS = 1")
	}

	if err := tx.Commit(); err != nil {
		panic(err)
	}
}

func removeTestRunnerUser(transactions []*transaction.Transaction) {
	dbConnect()
	defer store.Sql().Db.Close()
	deleteToken(adminToken, testRunnerUser.ID)
	deleteObject(testRunnerUser)
}

// Parameter Substitution
func setupObjectsAndPaths(t *transaction.Transaction) {
	alterRequestPath(t)
	alterRequestBody(t)
}

// Object Lifecycle
func addUserProjectRelation(pid int, user int) {
	_, err := store.CreateProjectUser(db.ProjectUser{
		ProjectID: pid,
		UserID:    user,
		Admin:     true,
	})
	if err != nil {
		panic(err)
	}
}

func deleteUserProjectRelation(pid int, user int) {
	err := store.DeleteProjectUser(pid, user)
	if err != nil {
		panic(err)
	}
}

func addAccessKey(pid *int) *db.AccessKey {
	uid := getUUID()
	secret := "5up3r53cr3t\n"

	key, err := store.CreateAccessKey(db.AccessKey{
		Name:      "ITK-" + uid,
		Type:      "ssh",
		Secret:    &secret,
		ProjectID: pid,
	})

	if err != nil {
		panic(err)
	}
	return &key
}

func addProject() *db.Project {
	uid := getUUID()
	project := db.Project{
		Name:    "ITP-" + uid,
		Created: time.Now(),
	}
	project, err := store.CreateProject(project)
	if err != nil {
		panic(err)
	}
	return &project
}

func addUser() *db.User {
	uid := getUUID()
	user := db.User{
		Created:  time.Now(),
		Username: "ITU-" + uid,
		Email:    "test@semaphore." + uid,
	}

	user, err := store.CreateUserWithoutPassword(user)

	if err != nil {
		panic(err)
	}
	return &user
}

func addView() *db.View {
	view, err := store.CreateView(db.View{
		ProjectID: userProject.ID,
		Title:     "Test",
		Position:  1,
	})

	if err != nil {
		panic(err)
	}

	return &view
}

func addSchedule() *db.Schedule {
	schedule, err := store.CreateSchedule(db.Schedule{
		TemplateID: int(templateID),
		CronFormat: "* * * 1 *",
		ProjectID:  userProject.ID,
	})

	if err != nil {
		panic(err)
	}

	return &schedule
}

func addTask() *db.Task {
	t := db.Task{
		ProjectID:  userProject.ID,
		TemplateID: int(templateID),
		Status:     "testing",
		UserID:     &userPathTestUser.ID,
		Created:    db.GetParsedTime(time.Now()),
	}
	if err := store.Sql().Insert(&t); err != nil {
		fmt.Println("error during insertion of task:")
		if j, err := json.Marshal(t); err == nil {
			fmt.Println(string(j))
		} else {
			fmt.Println("can not stringify task object")
		}
		panic(err)
	}
	return &t
}

func deleteObject(i interface{}) {
	_, err := store.Sql().Delete(i)
	if err != nil {
		panic(err)
	}
}

// Token Handling
func addToken(tok string, user int) {
	token := db.APIToken{
		ID:      tok,
		Created: time.Now(),
		UserID:  user,
		Expired: false,
	}
	if err := store.Sql().Insert(&token); err != nil {
		panic(err)
	}
}

func deleteToken(tok string, user int) {
	token := db.APIToken{
		ID:     tok,
		UserID: user,
	}
	deleteObject(&token)
}

// HELPERS
var r *rand.Rand
var randSetup = false

func getUUID() string {
	if !randSetup {
		r = rand.New(rand.NewSource(time.Now().UnixNano()))
		randSetup = true
	}
	return randomString(8)
}
func randomString(strlen int) string {
	const chars = "abcdefghijklmnopqrstuvwxyz0123456789"
	result := ""
	for i := 0; i < strlen; i++ {
		index := r.Intn(len(chars))
		result += chars[index : index+1]
	}
	return result
}

func loadConfig() {
	cwd, _ := os.Getwd()
	file, _ := os.Open(cwd + "/.dredd/config.json")
	if err := json.NewDecoder(file).Decode(&util.Config); err != nil {
		fmt.Println("Could not decode configuration!")
		panic(err)
	}
}

var store sql.SqlDb

func dbConnect() {
	store = sql.SqlDb{}

	if err := store.Connect(); err != nil {
		panic(err)
	}
}

func stringInSlice(a string, list []string) (int, bool) {
	for k, b := range list {
		if b == a {
			return k, true
		}
	}
	return 0, false
}

func printError(err error) {
	if err != nil {
		//fmt.Println(err)
		panic(err)
	}
}
<|endoftext|>


<commit_before>// Copyright 2016-2017 The psh Authors. All rights reserved.
package psh
<commit_msg>Test hostname<commit_after>// Copyright 2016-2017 The psh Authors. All rights reserved.
package psh

import "testing"

func TestSegmentHostnameCompile(t *testing.T) {
	expected := "\\h"
	segment := NewSegmentHostname()
	segment.Compile()
	if string(segment.Data) != expected {
		t.Fatalf("Compiled data expected to be %s but got %s", expected, segment.Data)
	}
}
<|endoftext|>


<commit_before>// Update protocol:
//
//   GET hk.heroku.com/hk/linux-amd64.json
//
//   200 ok
//   {
//       "Version": "2",
//       "Sha256": "..." // base64
//   }
//
// then
//
//   GET hkpatch.s3.amazonaws.com/hk/1/2/linux-amd64
//
//   200 ok
//   [bsdiff data]
//
// or
//
//   GET hkdist.s3.amazonaws.com/hk/2/linux-amd64.gz
//
//   200 ok
//   [gzipped executable data]
//
//
package selfupdate

import (
	"bytes"
	"compress/gzip"
	"crypto/sha256"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"math/rand"
	"net/http"
	"os"
	"path/filepath"
	"runtime"
	"time"

	"github.com/inconshreveable/go-update"
	"github.com/kardianos/osext"
	"github.com/kr/binarydist"
)

const (
	upcktimePath = "cktime"
	plat         = runtime.GOOS + "-" + runtime.GOARCH
)

const devValidTime = 7 * 24 * time.Hour

var ErrHashMismatch = errors.New("new file hash mismatch after patch")
var up = update.New()

// Updater is the configuration and runtime data for doing an update.
//
// Note that ApiURL, BinURL and DiffURL should have the same value if all files are available at the same location.
//
// Example:
//
//  updater := &selfupdate.Updater{
//  	CurrentVersion: version,
//  	ApiURL:         "http://updates.yourdomain.com/",
//  	BinURL:         "http://updates.yourdownmain.com/",
//  	DiffURL:        "http://updates.yourdomain.com/",
//  	Dir:            "update/",
//  	CmdName:        "myapp", // app name
//  }
//  if updater != nil {
//  	go updater.BackgroundRun()
//  }
type Updater struct {
	CurrentVersion string // Currently running version.
	ApiURL         string // Base URL for API requests (json files).
	CmdName        string // Command name is appended to the ApiURL like http://apiurl/CmdName/. This represents one binary.
	BinURL         string // Base URL for full binary downloads.
	DiffURL        string // Base URL for diff downloads.
	Dir            string // Directory to store selfupdate state.
	Info           struct {
		Version string
		Sha256  []byte
	}
}

func (u *Updater) getExecRelativeDir(dir string) string {
	filename, _ := osext.Executable()
	path := filepath.Join(filepath.Dir(filename), dir)
	return path
}

// BackgroundRun starts the update check and apply cycle.
func (u *Updater) BackgroundRun() error {
	os.MkdirAll(u.getExecRelativeDir(u.Dir), 0777)
	if u.wantUpdate() {
		if err := up.CanUpdate(); err != nil {
			// fail
			return err
		}
		//self, err := osext.Executable()
		//if err != nil {
		// fail update, couldn't figure out path to self
		//return
		//}
		// TODO(bgentry): logger isn't on Windows. Replace w/ proper error reports.
		if err := u.Update(); err != nil {
			return err
		}
	}
	return nil
}

func (u *Updater) wantUpdate() bool {
	path := u.getExecRelativeDir(u.Dir + upcktimePath)
	if u.CurrentVersion == "dev" || readTime(path).After(time.Now()) {
		return false
	}
	wait := 24*time.Hour + randDuration(24*time.Hour)
	return writeTime(path, time.Now().Add(wait))
}

func (u *Updater) Update() error {
	path, err := osext.Executable()
	if err != nil {
		return err
	}
	old, err := os.Open(path)
	if err != nil {
		return err
	}
	defer old.Close()

	err = u.Check()
	if err != nil {
		return err
	}
	if u.Info.Version == u.CurrentVersion {
		return nil
	}
	// TODO(sqs): reenable fetching just the patch
	/*bin, err := u.fetchAndVerifyPatch(old)
	if err != nil {
		if err == ErrHashMismatch {
			log.Println("update: hash mismatch from patched binary, downloading full binary")
		} else {
			if u.DiffURL != "" {
				log.Println("update: downloading full binary because patching binary failed:", err)
			}
		}*/

	bin, err := u.fetchAndVerifyFullBin()
	if err != nil {
		if err == ErrHashMismatch {
			log.Println("update: hash mismatch from full binary")
		} else {
			log.Println("update: fetching full binary failed:", err)
		}
		return err
	}
	//}

	// close the old binary before installing because on windows
	// it can't be renamed if a handle to the file is still open
	old.Close()

	err, errRecover := up.FromStream(bytes.NewBuffer(bin))
	if errRecover != nil {
		return fmt.Errorf("update and recovery errors: %q %q", err, errRecover)
	}
	if err != nil {
		return err
	}
	return nil
}

func (u *Updater) Check() error {
	r, err := fetch(u.ApiURL + u.CmdName + "/" + plat + "/" + u.CmdName + ".json")
	if err != nil {
		return err
	}
	defer r.Close()
	err = json.NewDecoder(r).Decode(&u.Info)
	if err != nil {
		return err
	}
	if len(u.Info.Sha256) != sha256.Size {
		return errors.New("bad cmd hash in info")
	}
	return nil
}

func (u *Updater) fetchAndVerifyPatch(old io.Reader) ([]byte, error) {
	bin, err := u.fetchAndApplyPatch(old)
	if err != nil {
		return nil, err
	}
	if !verifySha(bin, u.Info.Sha256) {
		return nil, ErrHashMismatch
	}
	return bin, nil
}

func (u *Updater) fetchAndApplyPatch(old io.Reader) ([]byte, error) {
	r, err := fetch(u.DiffURL + u.CmdName + "/" + u.CurrentVersion + "/" + u.Info.Version + "/" + plat + "/" + u.CmdName)
	if err != nil {
		return nil, err
	}
	defer r.Close()
	var buf bytes.Buffer
	err = binarydist.Patch(old, &buf, r)
	return buf.Bytes(), err
}

func (u *Updater) fetchAndVerifyFullBin() ([]byte, error) {
	bin, err := u.fetchBin()
	if err != nil {
		return nil, err
	}
	verified := verifySha(bin, u.Info.Sha256)
	if !verified {
		return nil, ErrHashMismatch
	}
	return bin, nil
}

func (u *Updater) fetchBin() ([]byte, error) {
	r, err := fetch(u.BinURL + u.CmdName + "/" + u.Info.Version + "/" + plat + "/" + u.CmdName + ".gz")
	if err != nil {
		return nil, err
	}
	defer r.Close()
	buf := new(bytes.Buffer)
	gz, err := gzip.NewReader(r)
	if err != nil {
		return nil, err
	}
	if _, err = io.Copy(buf, gz); err != nil {
		return nil, err
	}

	return buf.Bytes(), nil
}

// returns a random duration in [0,n).
func randDuration(n time.Duration) time.Duration {
	return time.Duration(rand.Int63n(int64(n)))
}

func fetch(url string) (io.ReadCloser, error) {
	resp, err := http.Get(url)
	if err != nil {
		return nil, err
	}
	if resp.StatusCode != 200 {
		return nil, fmt.Errorf("bad http status from %s: %v", url, resp.Status)
	}
	return resp.Body, nil
}

func readTime(path string) time.Time {
	p, err := ioutil.ReadFile(path)
	if os.IsNotExist(err) {
		return time.Time{}
	}
	if err != nil {
		return time.Now().Add(1000 * time.Hour)
	}
	t, err := time.Parse(time.RFC3339, string(p))
	if err != nil {
		return time.Now().Add(1000 * time.Hour)
	}
	return t
}

func verifySha(bin []byte, sha []byte) bool {
	h := sha256.New()
	h.Write(bin)
	return bytes.Equal(h.Sum(nil), sha)
}

func writeTime(path string, t time.Time) bool {
	return ioutil.WriteFile(path, []byte(t.Format(time.RFC3339)), 0644) == nil
}
<commit_msg>use go-update.v0<commit_after>// Update protocol:
//
//   GET hk.heroku.com/hk/linux-amd64.json
//
//   200 ok
//   {
//       "Version": "2",
//       "Sha256": "..." // base64
//   }
//
// then
//
//   GET hkpatch.s3.amazonaws.com/hk/1/2/linux-amd64
//
//   200 ok
//   [bsdiff data]
//
// or
//
//   GET hkdist.s3.amazonaws.com/hk/2/linux-amd64.gz
//
//   200 ok
//   [gzipped executable data]
//
//
package selfupdate

import (
	"bytes"
	"compress/gzip"
	"crypto/sha256"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"math/rand"
	"net/http"
	"os"
	"path/filepath"
	"runtime"
	"time"

	"github.com/kardianos/osext"
	"github.com/kr/binarydist"
	"gopkg.in/inconshreveable/go-update.v0"
)

const (
	upcktimePath = "cktime"
	plat         = runtime.GOOS + "-" + runtime.GOARCH
)

const devValidTime = 7 * 24 * time.Hour

var ErrHashMismatch = errors.New("new file hash mismatch after patch")
var up = update.New()

// Updater is the configuration and runtime data for doing an update.
//
// Note that ApiURL, BinURL and DiffURL should have the same value if all files are available at the same location.
//
// Example:
//
//  updater := &selfupdate.Updater{
//  	CurrentVersion: version,
//  	ApiURL:         "http://updates.yourdomain.com/",
//  	BinURL:         "http://updates.yourdownmain.com/",
//  	DiffURL:        "http://updates.yourdomain.com/",
//  	Dir:            "update/",
//  	CmdName:        "myapp", // app name
//  }
//  if updater != nil {
//  	go updater.BackgroundRun()
//  }
type Updater struct {
	CurrentVersion string // Currently running version.
	ApiURL         string // Base URL for API requests (json files).
	CmdName        string // Command name is appended to the ApiURL like http://apiurl/CmdName/. This represents one binary.
	BinURL         string // Base URL for full binary downloads.
	DiffURL        string // Base URL for diff downloads.
	Dir            string // Directory to store selfupdate state.
	Info           struct {
		Version string
		Sha256  []byte
	}
}

func (u *Updater) getExecRelativeDir(dir string) string {
	filename, _ := osext.Executable()
	path := filepath.Join(filepath.Dir(filename), dir)
	return path
}

// BackgroundRun starts the update check and apply cycle.
func (u *Updater) BackgroundRun() error {
	os.MkdirAll(u.getExecRelativeDir(u.Dir), 0777)
	if u.wantUpdate() {
		if err := up.CanUpdate(); err != nil {
			// fail
			return err
		}
		//self, err := osext.Executable()
		//if err != nil {
		// fail update, couldn't figure out path to self
		//return
		//}
		// TODO(bgentry): logger isn't on Windows. Replace w/ proper error reports.
		if err := u.Update(); err != nil {
			return err
		}
	}
	return nil
}

func (u *Updater) wantUpdate() bool {
	path := u.getExecRelativeDir(u.Dir + upcktimePath)
	if u.CurrentVersion == "dev" || readTime(path).After(time.Now()) {
		return false
	}
	wait := 24*time.Hour + randDuration(24*time.Hour)
	return writeTime(path, time.Now().Add(wait))
}

func (u *Updater) Update() error {
	path, err := osext.Executable()
	if err != nil {
		return err
	}
	old, err := os.Open(path)
	if err != nil {
		return err
	}
	defer old.Close()

	err = u.Check()
	if err != nil {
		return err
	}
	if u.Info.Version == u.CurrentVersion {
		return nil
	}
	// TODO(sqs): reenable fetching just the patch
	/*bin, err := u.fetchAndVerifyPatch(old)
	if err != nil {
		if err == ErrHashMismatch {
			log.Println("update: hash mismatch from patched binary, downloading full binary")
		} else {
			if u.DiffURL != "" {
				log.Println("update: downloading full binary because patching binary failed:", err)
			}
		}*/

	bin, err := u.fetchAndVerifyFullBin()
	if err != nil {
		if err == ErrHashMismatch {
			log.Println("update: hash mismatch from full binary")
		} else {
			log.Println("update: fetching full binary failed:", err)
		}
		return err
	}
	//}

	// close the old binary before installing because on windows
	// it can't be renamed if a handle to the file is still open
	old.Close()

	err, errRecover := up.FromStream(bytes.NewBuffer(bin))
	if errRecover != nil {
		return fmt.Errorf("update and recovery errors: %q %q", err, errRecover)
	}
	if err != nil {
		return err
	}
	return nil
}

func (u *Updater) Check() error {
	r, err := fetch(u.ApiURL + u.CmdName + "/" + plat + "/" + u.CmdName + ".json")
	if err != nil {
		return err
	}
	defer r.Close()
	err = json.NewDecoder(r).Decode(&u.Info)
	if err != nil {
		return err
	}
	if len(u.Info.Sha256) != sha256.Size {
		return errors.New("bad cmd hash in info")
	}
	return nil
}

func (u *Updater) fetchAndVerifyPatch(old io.Reader) ([]byte, error) {
	bin, err := u.fetchAndApplyPatch(old)
	if err != nil {
		return nil, err
	}
	if !verifySha(bin, u.Info.Sha256) {
		return nil, ErrHashMismatch
	}
	return bin, nil
}

func (u *Updater) fetchAndApplyPatch(old io.Reader) ([]byte, error) {
	r, err := fetch(u.DiffURL + u.CmdName + "/" + u.CurrentVersion + "/" + u.Info.Version + "/" + plat + "/" + u.CmdName)
	if err != nil {
		return nil, err
	}
	defer r.Close()
	var buf bytes.Buffer
	err = binarydist.Patch(old, &buf, r)
	return buf.Bytes(), err
}

func (u *Updater) fetchAndVerifyFullBin() ([]byte, error) {
	bin, err := u.fetchBin()
	if err != nil {
		return nil, err
	}
	verified := verifySha(bin, u.Info.Sha256)
	if !verified {
		return nil, ErrHashMismatch
	}
	return bin, nil
}

func (u *Updater) fetchBin() ([]byte, error) {
	r, err := fetch(u.BinURL + u.CmdName + "/" + u.Info.Version + "/" + plat + "/" + u.CmdName + ".gz")
	if err != nil {
		return nil, err
	}
	defer r.Close()
	buf := new(bytes.Buffer)
	gz, err := gzip.NewReader(r)
	if err != nil {
		return nil, err
	}
	if _, err = io.Copy(buf, gz); err != nil {
		return nil, err
	}

	return buf.Bytes(), nil
}

// returns a random duration in [0,n).
func randDuration(n time.Duration) time.Duration {
	return time.Duration(rand.Int63n(int64(n)))
}

func fetch(url string) (io.ReadCloser, error) {
	resp, err := http.Get(url)
	if err != nil {
		return nil, err
	}
	if resp.StatusCode != 200 {
		return nil, fmt.Errorf("bad http status from %s: %v", url, resp.Status)
	}
	return resp.Body, nil
}

func readTime(path string) time.Time {
	p, err := ioutil.ReadFile(path)
	if os.IsNotExist(err) {
		return time.Time{}
	}
	if err != nil {
		return time.Now().Add(1000 * time.Hour)
	}
	t, err := time.Parse(time.RFC3339, string(p))
	if err != nil {
		return time.Now().Add(1000 * time.Hour)
	}
	return t
}

func verifySha(bin []byte, sha []byte) bool {
	h := sha256.New()
	h.Write(bin)
	return bytes.Equal(h.Sum(nil), sha)
}

func writeTime(path string, t time.Time) bool {
	return ioutil.WriteFile(path, []byte(t.Format(time.RFC3339)), 0644) == nil
}
<|endoftext|>


<commit_before>// Package doctransport observes API registrations and traffic and generates
// Swagger. Ideally you could use this with unit tests which simulate API
// traffic.
package doctransport

import (
	"context"
	"fmt"
	"net/http"
	"reflect"
	"sync"

	"github.com/go-openapi/spec"
	"github.com/pkg/errors"

	"github.com/paultyng/resttransport"
	"github.com/paultyng/resttransport/routename"
)

type docTransport struct {
	sync.Mutex
	inner            resttransport.Transport
	spec             *spec.Swagger
	referenceStructs map[reflect.Type]bool
	namer            routename.Namer
}

type docRequestResponse struct {
	*docTransport
	inner resttransport.RequestResponse
	op    *spec.Operation
}

// SwaggerTransport is the interface for a Transport that has Swagger spec generation.
type SwaggerTransport interface {
	resttransport.Transport
	Generate() (*spec.Swagger, error)
}

// New returns an resttransport middleware that records interactions for documenting.
func New(inner resttransport.Transport) SwaggerTransport {
	return &docTransport{
		inner: inner,
		spec: &spec.Swagger{
			SwaggerProps: spec.SwaggerProps{
				// TODO: swagger 3?
				Swagger: "2.0",
				// TODO: other types?
				Consumes: []string{"application/json"},
				Produces: []string{"application/json"},
			},
		},
		referenceStructs: map[reflect.Type]bool{},
		namer:            routename.New(),
	}
}

func (t *docTransport) Generate() (*spec.Swagger, error) {
	t.Lock()
	swagger := *t.spec
	t.Unlock()

	// TODO: support additional auth types
	swagger.SecurityDefinitions = spec.SecurityDefinitions{
		bearerTokenAuthorizationName: &spec.SecurityScheme{
			SecuritySchemeProps: spec.SecuritySchemeProps{
				Description: "Requires a bearer token",
				In:          "header",
				Name:        "Authorization",
				Type:        "apiKey",
			},
		},
	}

	if swagger.Definitions == nil {
		swagger.Definitions = spec.Definitions{}
	}

	for s := range t.referenceStructs {
		sch, err := schema(s, false)
		if err != nil {
			return nil, errors.Wrapf(err, "unable to create schema for %v", s)
		}
		swagger.Definitions[s.Name()] = sch
	}

	return &swagger, nil
}

func addStructs(structs map[reflect.Type]bool, t reflect.Type) {
	switch t.Kind() {
	case reflect.Ptr,
		reflect.Slice,
		reflect.Array:
		addStructs(structs, t.Elem())
		return
	case reflect.Struct:
		if !structs[t] && t != wkTime {
			structs[t] = true
		}
	default:
		return
	}

	for i := 0; i < t.NumField(); i++ {
		f := t.Field(i)
		// assumes body only, so only json tags
		tag := f.Tag.Get("json")
		if isExportedField(f, tag) && !f.Anonymous {
			addStructs(structs, f.Type)
		}
	}
}

func (t *docTransport) addReferenceStruct(ref reflect.Type) {
	t.Lock()
	defer t.Unlock()

	addStructs(t.referenceStructs, ref)
}

func getOperation(pi spec.PathItem, httpMethod string) *spec.Operation {
	switch httpMethod {
	case "DELETE":
		return pi.Delete
	case "GET":
		return pi.Get
	case "HEAD":
		return pi.Head
	case "OPTIONS":
		return pi.Options
	case "PATCH":
		return pi.Patch
	case "POST":
		return pi.Post
	case "PUT":
		return pi.Put
	default:
		panic("unexpected http method " + httpMethod)
	}
}

func setOperation(pi *spec.PathItem, httpMethod string, op *spec.Operation) {
	switch httpMethod {
	case "DELETE":
		pi.Delete = op
	case "GET":
		pi.Get = op
	case "HEAD":
		pi.Head = op
	case "OPTIONS":
		pi.Options = op
	case "PATCH":
		pi.Patch = op
	case "POST":
		pi.Post = op
	case "PUT":
		pi.Put = op
	default:
		panic("unexpected http method " + httpMethod)
	}
}

func (t *docTransport) wrapHandler(auth bool, httpMethod, path string, inner resttransport.Handler) resttransport.Handler {
	t.Lock()
	defer t.Unlock()

	if t.spec.Paths == nil {
		t.spec.Paths = &spec.Paths{
			Paths: map[string]spec.PathItem{},
		}
	}

	id := t.namer.Name(httpMethod, path)
	pi := t.spec.Paths.Paths[path]
	op := getOperation(pi, httpMethod)

	if op == nil {
		op = &spec.Operation{
			OperationProps: spec.OperationProps{
				ID:          id,
				Description: "",
				Parameters:  []spec.Parameter{},
				Responses: &spec.Responses{
					ResponsesProps: spec.ResponsesProps{
						StatusCodeResponses: map[int]spec.Response{},
					},
				},
			},
		}
		if auth {
			op.Security = append(op.OperationProps.Security, map[string][]string{"Bearer": []string{}})
			op.Responses.StatusCodeResponses[http.StatusUnauthorized] = spec.Response{
				ResponseProps: spec.ResponseProps{
					Description: http.StatusText(http.StatusUnauthorized),
				},
			}
		}
		setOperation(&pi, httpMethod, op)
	}

	t.spec.Paths.Paths[path] = pi

	return func(ctx context.Context, reqres resttransport.RequestResponse) error {
		wrapper := &docRequestResponse{
			inner:        reqres,
			op:           op,
			docTransport: t,
		}

		return inner(ctx, wrapper)
	}
}

func (t *docTransport) RegisterHandler(httpMethod, path string, h resttransport.Handler) error {
	return t.inner.RegisterHandler(httpMethod, path, t.wrapHandler(false, httpMethod, path, h))
}

func (t *docTransport) RegisterAuthenticatedHandler(httpMethod, path string, h resttransport.Handler) error {
	return t.inner.RegisterAuthenticatedHandler(httpMethod, path, t.wrapHandler(true, httpMethod, path, h))
}

func (reqres *docRequestResponse) hasBodyParameter() bool {
	reqres.Lock()
	defer reqres.Unlock()
	for _, p := range reqres.op.Parameters {
		if p.In == "body" {
			return true
		}
	}
	return false
}

func (reqres *docRequestResponse) appendBodyParameter(v interface{}) error {
	const in = "body"

	if reqres.hasBodyParameter() {
		//short circuit if body param exists
		return nil
	}

	t := reflect.TypeOf(v)
	reqres.addReferenceStruct(t)

	typeSchema, err := schema(t, true)
	if err != nil {
		return errors.Wrap(err, "unable to map type for schema")
	}

	reqres.Lock()
	defer reqres.Unlock()
	reqres.op.Parameters = append(reqres.op.Parameters, spec.Parameter{
		ParamProps: spec.ParamProps{
			In:          in,
			Name:        in,
			Description: "",
			Required:    true,
			Schema:      &typeSchema,
		},
	})
	return nil
}

func (reqres *docRequestResponse) RequestHeader() http.Header {
	return reqres.inner.RequestHeader()
}

func (reqres *docRequestResponse) BindBody(v interface{}) error {
	err := reqres.appendBodyParameter(v)
	if err != nil {
		return errors.Wrap(err, "unable to append body parameter")
	}
	return reqres.inner.BindBody(v)
}

func (reqres *docRequestResponse) BindQuery(v interface{}) error {
	const in = "query"
	err := reqres.appendSimpleSchemaParameters(in, v)
	if err != nil {
		return err
	}
	return reqres.inner.BindQuery(v)
}

func (reqres *docRequestResponse) BindPath(v interface{}) error {
	const in = "path"
	err := reqres.appendSimpleSchemaParameters(in, v)
	if err != nil {
		return err
	}
	return reqres.inner.BindPath(v)
}

func (reqres *docRequestResponse) appendSimpleSchemaParameters(in string, v interface{}) error {
	reqres.Lock()
	defer reqres.Unlock()

	t := reflect.TypeOf(v)
	if t.Kind() == reflect.Ptr {
		t = t.Elem()
	}
	if t.Kind() != reflect.Struct {
		return fmt.Errorf("type %s in %s is not a struct", t.Name(), in)
	}
	return eachStructField(t, in, func(name string, required bool, s spec.Schema) error {
		st := ""
		if len(s.Type) > 0 {
			st = s.Type[0]
		}
		for _, p := range reqres.op.Parameters {
			if p.Name == name {
				//already exists
				return nil
			}
		}
		reqres.op.Parameters = append(reqres.op.Parameters, spec.Parameter{
			SimpleSchema: spec.SimpleSchema{
				Type:    st,
				Format:  s.Format,
				Default: s.Default,
			},
			ParamProps: spec.ParamProps{
				In:          in,
				Name:        name,
				Description: s.Description,
				Required:    required,
			},
		})

		return nil
	})
}

func (reqres *docRequestResponse) User() interface{} {
	return reqres.inner.User()
}

func (reqres *docRequestResponse) Body(status int, v interface{}) error {
	t := reflect.TypeOf(v)

	reqres.addReferenceStruct(t)

	typeSchema, err := schema(t, true)
	if err != nil {
		return errors.Wrap(err, "unable to map type for operation request")
	}

	resp := spec.Response{
		ResponseProps: spec.ResponseProps{
			Description: http.StatusText(status),
			Schema:      &typeSchema,
		},
	}

	reqres.Lock()
	defer reqres.Unlock()
	reqres.op.Responses.StatusCodeResponses[status] = resp

	return reqres.inner.Body(status, v)
}

func (reqres *docRequestResponse) NoBody(status int) error {
	resp := spec.Response{
		ResponseProps: spec.ResponseProps{
			Description: http.StatusText(status),
		},
	}

	reqres.Lock()
	defer reqres.Unlock()
	reqres.op.Responses.StatusCodeResponses[status] = resp

	return reqres.inner.NoBody(status)
}
<commit_msg>Lock higher up in doc middleware<commit_after>// Package doctransport observes API registrations and traffic and generates
// Swagger. Ideally you could use this with unit tests which simulate API
// traffic.
package doctransport

import (
	"context"
	"fmt"
	"net/http"
	"reflect"
	"sync"

	"github.com/go-openapi/spec"
	"github.com/pkg/errors"

	"github.com/paultyng/resttransport"
	"github.com/paultyng/resttransport/routename"
)

type docTransport struct {
	sync.Mutex
	inner            resttransport.Transport
	spec             *spec.Swagger
	referenceStructs map[reflect.Type]bool
	namer            routename.Namer
}

type docRequestResponse struct {
	*docTransport
	inner resttransport.RequestResponse
	op    *spec.Operation
}

// SwaggerTransport is the interface for a Transport that has Swagger spec generation.
type SwaggerTransport interface {
	resttransport.Transport
	Generate() (*spec.Swagger, error)
}

// New returns an resttransport middleware that records interactions for documenting.
func New(inner resttransport.Transport) SwaggerTransport {
	return &docTransport{
		inner: inner,
		spec: &spec.Swagger{
			SwaggerProps: spec.SwaggerProps{
				// TODO: swagger 3?
				Swagger: "2.0",
				// TODO: other types?
				Consumes: []string{"application/json"},
				Produces: []string{"application/json"},
			},
		},
		referenceStructs: map[reflect.Type]bool{},
		namer:            routename.New(),
	}
}

func (t *docTransport) Generate() (*spec.Swagger, error) {
	t.Lock()
	defer t.Unlock()
	swagger := *t.spec

	// TODO: support additional auth types
	swagger.SecurityDefinitions = spec.SecurityDefinitions{
		bearerTokenAuthorizationName: &spec.SecurityScheme{
			SecuritySchemeProps: spec.SecuritySchemeProps{
				Description: "Requires a bearer token",
				In:          "header",
				Name:        "Authorization",
				Type:        "apiKey",
			},
		},
	}

	if swagger.Definitions == nil {
		swagger.Definitions = spec.Definitions{}
	}

	for s := range t.referenceStructs {
		sch, err := schema(s, false)
		if err != nil {
			return nil, errors.Wrapf(err, "unable to create schema for %v", s)
		}
		swagger.Definitions[s.Name()] = sch
	}

	return &swagger, nil
}

func addStructs(structs map[reflect.Type]bool, t reflect.Type) {
	switch t.Kind() {
	case reflect.Ptr,
		reflect.Slice,
		reflect.Array:
		addStructs(structs, t.Elem())
		return
	case reflect.Struct:
		if !structs[t] && t != wkTime {
			structs[t] = true
		}
	default:
		return
	}

	for i := 0; i < t.NumField(); i++ {
		f := t.Field(i)
		// assumes body only, so only json tags
		tag := f.Tag.Get("json")
		if isExportedField(f, tag) && !f.Anonymous {
			addStructs(structs, f.Type)
		}
	}
}

func (t *docTransport) addReferenceStruct(ref reflect.Type) {
	addStructs(t.referenceStructs, ref)
}

func getOperation(pi spec.PathItem, httpMethod string) *spec.Operation {
	switch httpMethod {
	case "DELETE":
		return pi.Delete
	case "GET":
		return pi.Get
	case "HEAD":
		return pi.Head
	case "OPTIONS":
		return pi.Options
	case "PATCH":
		return pi.Patch
	case "POST":
		return pi.Post
	case "PUT":
		return pi.Put
	default:
		panic("unexpected http method " + httpMethod)
	}
}

func setOperation(pi *spec.PathItem, httpMethod string, op *spec.Operation) {
	switch httpMethod {
	case "DELETE":
		pi.Delete = op
	case "GET":
		pi.Get = op
	case "HEAD":
		pi.Head = op
	case "OPTIONS":
		pi.Options = op
	case "PATCH":
		pi.Patch = op
	case "POST":
		pi.Post = op
	case "PUT":
		pi.Put = op
	default:
		panic("unexpected http method " + httpMethod)
	}
}

func (t *docTransport) wrapHandler(auth bool, httpMethod, path string, inner resttransport.Handler) resttransport.Handler {
	if t.spec.Paths == nil {
		t.spec.Paths = &spec.Paths{
			Paths: map[string]spec.PathItem{},
		}
	}

	id := t.namer.Name(httpMethod, path)
	pi := t.spec.Paths.Paths[path]
	op := getOperation(pi, httpMethod)

	if op == nil {
		op = &spec.Operation{
			OperationProps: spec.OperationProps{
				ID:          id,
				Description: "",
				Parameters:  []spec.Parameter{},
				Responses: &spec.Responses{
					ResponsesProps: spec.ResponsesProps{
						StatusCodeResponses: map[int]spec.Response{},
					},
				},
			},
		}
		if auth {
			op.Security = append(op.OperationProps.Security, map[string][]string{"Bearer": []string{}})
			op.Responses.StatusCodeResponses[http.StatusUnauthorized] = spec.Response{
				ResponseProps: spec.ResponseProps{
					Description: http.StatusText(http.StatusUnauthorized),
				},
			}
		}
		setOperation(&pi, httpMethod, op)
	}

	t.spec.Paths.Paths[path] = pi

	return func(ctx context.Context, reqres resttransport.RequestResponse) error {
		wrapper := &docRequestResponse{
			inner:        reqres,
			op:           op,
			docTransport: t,
		}

		return inner(ctx, wrapper)
	}
}

func (t *docTransport) RegisterHandler(httpMethod, path string, h resttransport.Handler) error {
	t.Lock()
	defer t.Unlock()
	return t.inner.RegisterHandler(httpMethod, path, t.wrapHandler(false, httpMethod, path, h))
}

func (t *docTransport) RegisterAuthenticatedHandler(httpMethod, path string, h resttransport.Handler) error {
	t.Lock()
	defer t.Unlock()
	return t.inner.RegisterAuthenticatedHandler(httpMethod, path, t.wrapHandler(true, httpMethod, path, h))
}

func (reqres *docRequestResponse) hasBodyParameter() bool {
	for _, p := range reqres.op.Parameters {
		if p.In == "body" {
			return true
		}
	}
	return false
}

func (reqres *docRequestResponse) appendBodyParameter(v interface{}) error {
	const in = "body"

	if reqres.hasBodyParameter() {
		//short circuit if body param exists
		return nil
	}

	t := reflect.TypeOf(v)
	reqres.addReferenceStruct(t)

	typeSchema, err := schema(t, true)
	if err != nil {
		return errors.Wrap(err, "unable to map type for schema")
	}

	reqres.op.Parameters = append(reqres.op.Parameters, spec.Parameter{
		ParamProps: spec.ParamProps{
			In:          in,
			Name:        in,
			Description: "",
			Required:    true,
			Schema:      &typeSchema,
		},
	})
	return nil
}

func (reqres *docRequestResponse) RequestHeader() http.Header {
	return reqres.inner.RequestHeader()
}

func (reqres *docRequestResponse) BindBody(v interface{}) error {
	err := func() error {
		reqres.Lock()
		defer reqres.Unlock()
		return reqres.appendBodyParameter(v)
	}()
	if err != nil {
		return errors.Wrap(err, "unable to append body parameter")
	}
	return reqres.inner.BindBody(v)
}

func (reqres *docRequestResponse) BindQuery(v interface{}) error {
	const in = "query"
	err := func() error {
		reqres.Lock()
		defer reqres.Unlock()
		return reqres.appendSimpleSchemaParameters(in, v)
	}()
	if err != nil {
		return err
	}
	return reqres.inner.BindQuery(v)
}

func (reqres *docRequestResponse) BindPath(v interface{}) error {
	const in = "path"
	err := func() error {
		reqres.Lock()
		defer reqres.Unlock()
		return reqres.appendSimpleSchemaParameters(in, v)
	}()
	if err != nil {
		return err
	}
	return reqres.inner.BindPath(v)
}

func (reqres *docRequestResponse) appendSimpleSchemaParameters(in string, v interface{}) error {
	t := reflect.TypeOf(v)
	if t.Kind() == reflect.Ptr {
		t = t.Elem()
	}
	if t.Kind() != reflect.Struct {
		return fmt.Errorf("type %s in %s is not a struct", t.Name(), in)
	}
	return eachStructField(t, in, func(name string, required bool, s spec.Schema) error {
		st := ""
		if len(s.Type) > 0 {
			st = s.Type[0]
		}
		for _, p := range reqres.op.Parameters {
			if p.Name == name {
				//already exists
				return nil
			}
		}
		reqres.op.Parameters = append(reqres.op.Parameters, spec.Parameter{
			SimpleSchema: spec.SimpleSchema{
				Type:    st,
				Format:  s.Format,
				Default: s.Default,
			},
			ParamProps: spec.ParamProps{
				In:          in,
				Name:        name,
				Description: s.Description,
				Required:    required,
			},
		})

		return nil
	})
}

func (reqres *docRequestResponse) User() interface{} {
	return reqres.inner.User()
}

func (reqres *docRequestResponse) Body(status int, v interface{}) error {
	t := reflect.TypeOf(v)

	err := func() error {
		reqres.Lock()
		defer reqres.Unlock()

		reqres.addReferenceStruct(t)

		typeSchema, err := schema(t, true)
		if err != nil {
			return errors.Wrap(err, "unable to map type for operation request")
		}

		resp := spec.Response{
			ResponseProps: spec.ResponseProps{
				Description: http.StatusText(status),
				Schema:      &typeSchema,
			},
		}

		reqres.op.Responses.StatusCodeResponses[status] = resp
		return nil
	}()
	if err != nil {
		return err
	}

	return reqres.inner.Body(status, v)
}

func (reqres *docRequestResponse) NoBody(status int) error {
	resp := spec.Response{
		ResponseProps: spec.ResponseProps{
			Description: http.StatusText(status),
		},
	}

	func() {
		reqres.Lock()
		defer reqres.Unlock()
		reqres.op.Responses.StatusCodeResponses[status] = resp
	}()

	return reqres.inner.NoBody(status)
}
<|endoftext|>