<commit_before>package main

import (
	"archive/tar"
	"bytes"
	"compress/gzip"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"path"
	"path/filepath"
	"strconv"
	"strings"
	"text/template"
	"time"

	"github.com/dynport/gocloud/aws/s3"
	"github.com/dynport/gossh"
)

var logger = log.New(os.Stderr, "", 0)

const sshExample = "ubuntu@127.0.0.1"

func main() {
	dir := flag.String("dir", "", "Dir to build. Default: current directory")
	host := flag.String("host", os.Getenv("DEV_HOST"), "Host to build on. Example: "+sshExample)
	deploy := flag.String("deploy", "", "Deploy to host after building. Example: "+sshExample)
	bucket := flag.String("bucket", "", "Upload binary to s3 bucket after building")
	public := flag.Bool("public", false, "Upload to s3 and make public")
	verbose := flag.Bool("verbose", false, "Build using -v flag")
	goVersion := flag.String("go-version", "1.3.3", "Go version")
	flag.Parse()
	logger.Printf("running with host=%q go_version=%q", *host, *goVersion)
	b := &build{Host: *host, Dir: *dir, DeployTo: *deploy, Bucket: *bucket, verbose: *verbose, Public: *public, GoVersion: *goVersion}
	e := b.Run()
	if e != nil {
		logger.Fatalf("ERROR: %s", e)
	}
}

type build struct {
	Host      string
	Dir       string
	Bucket    string
	Public    bool
	DeployTo  string
	verbose   bool
	GoVersion string
}

func benchmark(message string) func() {
	started := time.Now()
	return func() {
		logger.Printf("finished %s in %.06f", message, time.Since(started).Seconds())
	}
}

func (r *build) deps() ([]string, error) {
	s, e := r.exec("go", "list", "-f", `{{ join .Deps " " }}`)
	if e != nil {
		return nil, e
	}
	return strings.Fields(s), nil
}

func (r *build) exec(cmd string, vals ...string) (string, error) {
	c := exec.Command(cmd, vals...)
	c.Dir = r.Dir
	out, e := c.CombinedOutput()
	if e != nil {
		return "", e
	}
	return string(out), nil
}

func (r *build) currentPackage() (string, error) {
	s, e := r.exec("go", "list")
	if e != nil {
		return "", e
	}
	return strings.TrimSpace(s), nil
}

func (r *build) filesMap() (map[string]os.FileInfo, error) {
	cp, e := r.currentPackage()
	if e != nil {
		return nil, e
	}
	pkgs, e := r.deps()
	if e != nil {
		return nil, e
	}
	pkgs = append(pkgs, cp)
	files := map[string]os.FileInfo{}
	sum := int64(0)
	for _, p := range pkgs {
		if !strings.Contains(p, ".") {
			continue
		}
		prefix := os.ExpandEnv("$GOPATH/src")
		dbg.Printf("walking %q", p)
		e := filepath.Walk(os.ExpandEnv(prefix+"/"+p+"/"), func(p string, info os.FileInfo, e error) error {
			skip := func() bool {
				for _, s := range []string{".git", ".bzr", ".hg"} {
					if strings.Contains(p, "/"+s+"/") {
						return true
					}
				}
				return false
			}()
			if skip {
				return nil
			}
			if _, ok := files[p]; !ok {
				sum += info.Size()
				files[p] = info
			}
			return nil
		})
		if e != nil {
			return nil, e
		}
	}
	return files, nil
}

func (b *build) createArchive() (string, error) {
	defer benchmark("create archive")()
	var name string
	e := func() error {
		f, e := ioutil.TempFile("/tmp", "gobuild-archive-")
		if e != nil {
			return e
		}
		name = f.Name()
		defer f.Close()
		files, e := b.filesMap()
		if e != nil {
			return e
		}

		gz := gzip.NewWriter(f)
		sum := int64(0)
		defer gz.Close()
		t := tar.NewWriter(gz)
		defer t.Close()

		for p, info := range files {
			name := strings.TrimPrefix(p, os.ExpandEnv("$GOPATH/src/"))
			if info.IsDir() {
				continue
			}
			dbg.Printf("adding %q", p)
			h := &tar.Header{ModTime: info.ModTime(), Size: info.Size(), Mode: int64(info.Mode()), Name: name}
			e = t.WriteHeader(h)
			if e != nil {
				return e
			}
			e := func() error {
				f, e := os.Open(p)
				if e != nil {
					return e
				}
				defer f.Close()
				i, e := io.Copy(t, f)
				sum += i
				return e
			}()
			if e != nil {
				return e
			}
		}
		dbg.Printf("%s", sizePretty(sum))
		return nil
	}()
	return name, e
}

type buildConfig struct {
	Current string
	Sudo    bool
	Verbose bool
	Version string
}

func (b *buildConfig) Goroot() string {
	return "{{ .BuildHome }}/.go/go-{{ .Version }}/go"
}

func (b *buildConfig) Gopath() string {
	return "{{ .BuildHome }}/{{ .Current }}"
}

func (b *buildConfig) BuildHome() string {
	return "$HOME/.gobuild"
}

func (b *buildConfig) BinName() string {
	return path.Base(b.Current)
}

func (b *build) Run() error {
	defer benchmark("build")()
	currentPkg, e := b.currentPackage()
	if e != nil {
		return e
	}

	cfg, e := parseConfig(b.Host)
	if e != nil {
		return e
	}
	dbg.Printf("using config %#v", cfg)
	con, e := cfg.Connection()
	if e != nil {
		return e
	}
	defer con.Close()

	name, e := b.createArchive()
	if e != nil {
		return e
	}
	dbg.Printf("created archive at %q", name)
	defer os.RemoveAll(name)
	f, e := os.Open(name)
	if e != nil {
		return e
	}
	defer f.Close()

	ses, e := con.NewSession()
	if e != nil {
		return e
	}
	ses.Stdin = f
	ses.Stdout = os.Stdout
	ses.Stderr = os.Stderr

	buildCfg := &buildConfig{
		Current: currentPkg,
		Sudo:    cfg.User != "root",
		Verbose: b.verbose,
		Version: b.GoVersion,
	}
	if buildCfg.Version == "" {
		buildCfg.Version = "1.3.3"
	}

	cmd := renderRecursive(buildCmd, buildCfg)
	e = ses.Run(cmd)
	if e != nil {
		return e
	}

	name = path.Base(currentPkg)
	var binPath string

	if b.Bucket != "" || b.DeployTo != "" {
		defer os.RemoveAll(binPath)
		e = func() error {
			ses, e := con.NewSession()
			if e != nil {
				return e
			}
			defer ses.Close()

			f, e := ioutil.TempFile("/tmp", "gobuild-bin-")
			if e != nil {
				return e
			}
			binPath = f.Name()
			defer f.Close()
			ses.Stdout = f
			ses.Stderr = os.Stderr

			cmd := renderRecursive("cat {{ .Gopath }}/bin/{{ .BinName }}", buildCfg)
			e = ses.Run(cmd)
			if e != nil {
				return e
			}
			return nil
		}()
		if e != nil {
			return e
		}
	}
	if name != "" {
		defer os.RemoveAll(binPath)
	}
	if b.Bucket != "" {
		e = func() error {
			f, e := os.Open(binPath)
			if e != nil {
				return e
			}
			defer f.Close()
			client := s3.NewFromEnv()
			client.CustomEndpointHost = "s3-eu-west-1.amazonaws.com"
			bucket, key := bucketAndKey(b.Bucket, name)
			logger.Printf("uploading to bucket=%q key=%q", bucket, key)

			opts := &s3.PutOptions{}
			if b.Public {
				opts.AmzAcl = "public-read"
			}
			return client.PutStream(bucket, key, f, opts)
		}()
		if e != nil {
			return e
		}
		logger.Printf("uploaded to bucket %q", b.Bucket)
	}

	if b.DeployTo != "" {
		e := func() error {
			cfg, e := parseConfig(b.DeployTo)
			if e != nil {
				return e
			}
			con, e := cfg.Connection()
			if e != nil {
				return e
			}
			defer con.Close()
			ses, e := con.NewSession()
			if e != nil {
				return e
			}
			defer ses.Close()
			f, e := os.Open(binPath)
			if e != nil {
				return e
			}
			defer f.Close()
			ses.Stdin = f
			ses.Stdout = os.Stdout
			ses.Stderr = os.Stderr
			s := struct {
				Name string
				Sudo bool
			}{
				Name: name, Sudo: cfg.User != "root",
			}
			cmd := renderRecursive("cd /usr/local/bin && cat - | {{ if .Sudo }}sudo {{ end}}tee {{ .Name }}.tmp > /dev/null && {{ if .Sudo }}sudo {{ end }}chmod 0755 {{ .Name }}.tmp && {{ if .Sudo }}sudo {{ end }}mv {{ .Name }}.tmp {{ .Name }}", s)
			dbg.Printf("%s", cmd)
			return ses.Run(cmd)
		}()
		if e != nil {
			return e
		}
	}
	return nil
}

func bucketAndKey(bucketWithPrefix, name string) (bucket, key string) {
	parts := strings.Split(bucketWithPrefix, "/")
	bucket = parts[0]
	key = name
	if len(parts) > 1 {
		key = strings.TrimSuffix(strings.Join(parts[1:], "/"), "/") + "/" + key
	}
	return bucket, key
}

func parseConfig(s string) (*gossh.Config, error) {
	cfg := &gossh.Config{}
	parts := strings.Split(s, "@")
	hostAndPort := ""
	switch len(parts) {
	case 0:
		return nil, fmt.Errorf("Host must be set")
	case 1:
		hostAndPort = parts[0]
	case 2:
		cfg.User = parts[0]
		hostAndPort = parts[1]
	case 3:
		return nil, fmt.Errorf("format of host %q not understood", s)
	}
	parts = strings.Split(hostAndPort, ":")
	cfg.Host = parts[0]
	if len(parts) == 2 {
		var e error
		cfg.Port, e = strconv.Atoi(parts[1])
		if e != nil {
			return nil, e
		}
	}
	return cfg, nil
}

const buildCmd = `#!/bin/bash
export BUILD_HOME={{ .BuildHome }}
export GOPATH={{ .Gopath }}
export GOROOT={{ .Goroot }}
export PATH=$GOROOT/bin:$PATH

if [ ! -f $GOROOT/bin/go ]; then
  echo "installing go {{ .Version }}"
  tmp=$(dirname $GOROOT)
  mkdir -p $tmp
  cd $tmp
  curl -sL "https://storage.googleapis.com/golang/go{{ .Version }}.linux-amd64.tar.gz" | tar xfz -
fi

set -xe
rm -Rf $GOPATH
mkdir -p $GOPATH/src
cd $GOPATH/src
tar xfz -
cd {{ .Current }}
go get {{ if .Verbose }}-v{{ end }} .
{{ with .Sudo }}sudo {{ end }}cp $GOPATH/bin/{{ .BinName }} /usr/local/bin/{{ .BinName }}.tmp
{{ with .Sudo }}sudo {{ end }}mv /usr/local/bin/{{ .BinName }}.tmp /usr/local/bin/{{ .BinName }}
`

func debugStream() io.Writer {
	if os.Getenv("DEBUG") == "true" {
		return os.Stderr
	}
	return ioutil.Discard
}

var dbg = log.New(debugStream(), "[DEBUG] ", log.Lshortfile)

func renderRecursive(tpl string, i interface{}) string {
	s := tpl
	for j := 0; j < 10; j++ {
		rendered := mustRender([]byte(s), i)
		if rendered == s {
			return rendered
		}
		s = rendered
	}
	logger.Fatal("rendering loop, rendered 10 times")
	return ""
}

func mustRender(raw []byte, i interface{}) string {
	out, e := render(raw, i)
	if e != nil {
		logger.Fatal(e)
	}
	return out
}

func render(raw []byte, i interface{}) (string, error) {
	tpl, e := template.New(string(raw)).Parse(string(raw))
	if e != nil {
		return "", e
	}
	buf := &bytes.Buffer{}
	e = tpl.Execute(buf, i)
	if e != nil {
		return "", e
	}
	return buf.String(), nil
}

var (
	oneKb = 1024.0
	oneMb = oneKb * 1024.0
	oneGb = oneMb * 1024.0
)

func sizePretty(raw int64) string {
	f := float64(raw)
	if f < oneKb {
		return fmt.Sprintf("%.0f", f)
	} else if f < oneMb {
		return fmt.Sprintf("%.2fKB", f/oneKb)
	} else if f < oneGb {
		return fmt.Sprintf("%.2fMB", f/oneMb)
	} else {
		return fmt.Sprintf("%.2fGB", f/oneGb)
	}
}
<commit_msg>add support for custom deploy dir<commit_after>package main

import (
	"archive/tar"
	"bytes"
	"compress/gzip"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"path"
	"path/filepath"
	"strconv"
	"strings"
	"text/template"
	"time"

	"github.com/dynport/gocloud/aws/s3"
	"github.com/dynport/gossh"
)

var logger = log.New(os.Stderr, "", 0)

const sshExample = "ubuntu@127.0.0.1"

func main() {
	dir := flag.String("dir", "", "Dir to build. Default: current directory")
	host := flag.String("host", os.Getenv("DEV_HOST"), "Host to build on. Example: "+sshExample)
	deploy := flag.String("deploy", "", "Deploy to host after building. Example: "+sshExample)
	deployDir := flag.String("deploy-dir", "/usr/local/bin", "Deploy to this directory for -deploy")
	bucket := flag.String("bucket", "", "Upload binary to s3 bucket after building")
	public := flag.Bool("public", false, "Upload to s3 and make public")
	verbose := flag.Bool("verbose", false, "Build using -v flag")
	goVersion := flag.String("go-version", "1.3.3", "Go version")
	flag.Parse()
	logger.Printf("running with host=%q go_version=%q", *host, *goVersion)
	b := &build{DeployDir: *deployDir, Host: *host, Dir: *dir, DeployTo: *deploy, Bucket: *bucket, verbose: *verbose, Public: *public, GoVersion: *goVersion}
	e := b.Run()
	if e != nil {
		logger.Fatalf("ERROR: %s", e)
	}
}

type build struct {
	Host      string
	Dir       string
	Bucket    string
	Public    bool
	DeployTo  string
	DeployDir string
	verbose   bool
	GoVersion string
}

func benchmark(message string) func() {
	started := time.Now()
	return func() {
		logger.Printf("finished %s in %.06f", message, time.Since(started).Seconds())
	}
}

func (r *build) deps() ([]string, error) {
	s, e := r.exec("go", "list", "-f", `{{ join .Deps " " }}`)
	if e != nil {
		return nil, e
	}
	return strings.Fields(s), nil
}

func (r *build) exec(cmd string, vals ...string) (string, error) {
	c := exec.Command(cmd, vals...)
	c.Dir = r.Dir
	out, e := c.CombinedOutput()
	if e != nil {
		return "", e
	}
	return string(out), nil
}

func (r *build) currentPackage() (string, error) {
	s, e := r.exec("go", "list")
	if e != nil {
		return "", e
	}
	return strings.TrimSpace(s), nil
}

func (r *build) filesMap() (map[string]os.FileInfo, error) {
	cp, e := r.currentPackage()
	if e != nil {
		return nil, e
	}
	pkgs, e := r.deps()
	if e != nil {
		return nil, e
	}
	pkgs = append(pkgs, cp)
	files := map[string]os.FileInfo{}
	sum := int64(0)
	for _, p := range pkgs {
		if !strings.Contains(p, ".") {
			continue
		}
		prefix := os.ExpandEnv("$GOPATH/src")
		dbg.Printf("walking %q", p)
		e := filepath.Walk(os.ExpandEnv(prefix+"/"+p+"/"), func(p string, info os.FileInfo, e error) error {
			skip := func() bool {
				for _, s := range []string{".git", ".bzr", ".hg"} {
					if strings.Contains(p, "/"+s+"/") {
						return true
					}
				}
				return false
			}()
			if skip {
				return nil
			}
			if _, ok := files[p]; !ok {
				sum += info.Size()
				files[p] = info
			}
			return nil
		})
		if e != nil {
			return nil, e
		}
	}
	return files, nil
}

func (b *build) createArchive() (string, error) {
	defer benchmark("create archive")()
	var name string
	e := func() error {
		f, e := ioutil.TempFile("/tmp", "gobuild-archive-")
		if e != nil {
			return e
		}
		name = f.Name()
		defer f.Close()
		files, e := b.filesMap()
		if e != nil {
			return e
		}

		gz := gzip.NewWriter(f)
		sum := int64(0)
		defer gz.Close()
		t := tar.NewWriter(gz)
		defer t.Close()

		for p, info := range files {
			name := strings.TrimPrefix(p, os.ExpandEnv("$GOPATH/src/"))
			if info.IsDir() {
				continue
			}
			dbg.Printf("adding %q", p)
			h := &tar.Header{ModTime: info.ModTime(), Size: info.Size(), Mode: int64(info.Mode()), Name: name}
			e = t.WriteHeader(h)
			if e != nil {
				return e
			}
			e := func() error {
				f, e := os.Open(p)
				if e != nil {
					return e
				}
				defer f.Close()
				i, e := io.Copy(t, f)
				sum += i
				return e
			}()
			if e != nil {
				return e
			}
		}
		dbg.Printf("%s", sizePretty(sum))
		return nil
	}()
	return name, e
}

type buildConfig struct {
	Current string
	Sudo    bool
	Verbose bool
	Version string
}

func (b *buildConfig) Goroot() string {
	return "{{ .BuildHome }}/.go/go-{{ .Version }}/go"
}

func (b *buildConfig) Gopath() string {
	return "{{ .BuildHome }}/{{ .Current }}"
}

func (b *buildConfig) BuildHome() string {
	return "$HOME/.gobuild"
}

func (b *buildConfig) BinName() string {
	return path.Base(b.Current)
}

func (b *build) Run() error {
	defer benchmark("build")()
	currentPkg, e := b.currentPackage()
	if e != nil {
		return e
	}

	cfg, e := parseConfig(b.Host)
	if e != nil {
		return e
	}
	dbg.Printf("using config %#v", cfg)
	con, e := cfg.Connection()
	if e != nil {
		return e
	}
	defer con.Close()

	name, e := b.createArchive()
	if e != nil {
		return e
	}
	dbg.Printf("created archive at %q", name)
	defer os.RemoveAll(name)
	f, e := os.Open(name)
	if e != nil {
		return e
	}
	defer f.Close()

	ses, e := con.NewSession()
	if e != nil {
		return e
	}
	ses.Stdin = f
	ses.Stdout = os.Stdout
	ses.Stderr = os.Stderr

	buildCfg := &buildConfig{
		Current: currentPkg,
		Sudo:    cfg.User != "root",
		Verbose: b.verbose,
		Version: b.GoVersion,
	}
	if buildCfg.Version == "" {
		buildCfg.Version = "1.3.3"
	}

	cmd := renderRecursive(buildCmd, buildCfg)
	e = ses.Run(cmd)
	if e != nil {
		return e
	}

	name = path.Base(currentPkg)
	var binPath string

	if b.Bucket != "" || b.DeployTo != "" {
		defer os.RemoveAll(binPath)
		e = func() error {
			ses, e := con.NewSession()
			if e != nil {
				return e
			}
			defer ses.Close()

			f, e := ioutil.TempFile("/tmp", "gobuild-bin-")
			if e != nil {
				return e
			}
			binPath = f.Name()
			defer f.Close()
			ses.Stdout = f
			ses.Stderr = os.Stderr

			cmd := renderRecursive("cat {{ .Gopath }}/bin/{{ .BinName }}", buildCfg)
			e = ses.Run(cmd)
			if e != nil {
				return e
			}
			return nil
		}()
		if e != nil {
			return e
		}
	}
	if name != "" {
		defer os.RemoveAll(binPath)
	}
	if b.Bucket != "" {
		e = func() error {
			f, e := os.Open(binPath)
			if e != nil {
				return e
			}
			defer f.Close()
			client := s3.NewFromEnv()
			client.CustomEndpointHost = "s3-eu-west-1.amazonaws.com"
			bucket, key := bucketAndKey(b.Bucket, name)
			logger.Printf("uploading to bucket=%q key=%q", bucket, key)

			opts := &s3.PutOptions{}
			if b.Public {
				opts.AmzAcl = "public-read"
			}
			return client.PutStream(bucket, key, f, opts)
		}()
		if e != nil {
			return e
		}
		logger.Printf("uploaded to bucket %q", b.Bucket)
	}

	if b.DeployTo != "" {
		e := func() error {
			cfg, e := parseConfig(b.DeployTo)
			if e != nil {
				return e
			}
			con, e := cfg.Connection()
			if e != nil {
				return e
			}
			defer con.Close()
			ses, e := con.NewSession()
			if e != nil {
				return e
			}
			defer ses.Close()
			f, e := os.Open(binPath)
			if e != nil {
				return e
			}
			defer f.Close()
			ses.Stdin = f
			ses.Stdout = os.Stdout
			ses.Stderr = os.Stderr
			s := struct {
				Name      string
				Sudo      bool
				DeployDir string
			}{
				Name:      name,
				Sudo:      cfg.User != "root",
				DeployDir: b.DeployDir,
			}
			cmd := renderRecursive("cd {{ .DeployDir }} && cat - | {{ if .Sudo }}sudo {{ end}}tee {{ .Name }}.tmp > /dev/null && {{ if .Sudo }}sudo {{ end }}chmod 0755 {{ .Name }}.tmp && {{ if .Sudo }}sudo {{ end }}mv {{ .Name }}.tmp {{ .Name }}", s)
			dbg.Printf("%s", cmd)
			return ses.Run(cmd)
		}()
		if e != nil {
			return e
		}
	}
	return nil
}

func bucketAndKey(bucketWithPrefix, name string) (bucket, key string) {
	parts := strings.Split(bucketWithPrefix, "/")
	bucket = parts[0]
	key = name
	if len(parts) > 1 {
		key = strings.TrimSuffix(strings.Join(parts[1:], "/"), "/") + "/" + key
	}
	return bucket, key
}

func parseConfig(s string) (*gossh.Config, error) {
	cfg := &gossh.Config{}
	parts := strings.Split(s, "@")
	hostAndPort := ""
	switch len(parts) {
	case 0:
		return nil, fmt.Errorf("Host must be set")
	case 1:
		hostAndPort = parts[0]
	case 2:
		cfg.User = parts[0]
		hostAndPort = parts[1]
	case 3:
		return nil, fmt.Errorf("format of host %q not understood", s)
	}
	parts = strings.Split(hostAndPort, ":")
	cfg.Host = parts[0]
	if len(parts) == 2 {
		var e error
		cfg.Port, e = strconv.Atoi(parts[1])
		if e != nil {
			return nil, e
		}
	}
	return cfg, nil
}

const buildCmd = `#!/bin/bash
export BUILD_HOME={{ .BuildHome }}
export GOPATH={{ .Gopath }}
export GOROOT={{ .Goroot }}
export PATH=$GOROOT/bin:$PATH

if [ ! -f $GOROOT/bin/go ]; then
  echo "installing go {{ .Version }}"
  tmp=$(dirname $GOROOT)
  mkdir -p $tmp
  cd $tmp
  curl -sL "https://storage.googleapis.com/golang/go{{ .Version }}.linux-amd64.tar.gz" | tar xfz -
fi

set -xe
rm -Rf $GOPATH
mkdir -p $GOPATH/src
cd $GOPATH/src
tar xfz -
cd {{ .Current }}
go get {{ if .Verbose }}-v{{ end }} .
{{ with .Sudo }}sudo {{ end }}cp $GOPATH/bin/{{ .BinName }} /usr/local/bin/{{ .BinName }}.tmp
{{ with .Sudo }}sudo {{ end }}mv /usr/local/bin/{{ .BinName }}.tmp /usr/local/bin/{{ .BinName }}
`

func debugStream() io.Writer {
	if os.Getenv("DEBUG") == "true" {
		return os.Stderr
	}
	return ioutil.Discard
}

var dbg = log.New(debugStream(), "[DEBUG] ", log.Lshortfile)

func renderRecursive(tpl string, i interface{}) string {
	s := tpl
	for j := 0; j < 10; j++ {
		rendered := mustRender([]byte(s), i)
		if rendered == s {
			return rendered
		}
		s = rendered
	}
	logger.Fatal("rendering loop, rendered 10 times")
	return ""
}

func mustRender(raw []byte, i interface{}) string {
	out, e := render(raw, i)
	if e != nil {
		logger.Fatal(e)
	}
	return out
}

func render(raw []byte, i interface{}) (string, error) {
	tpl, e := template.New(string(raw)).Parse(string(raw))
	if e != nil {
		return "", e
	}
	buf := &bytes.Buffer{}
	e = tpl.Execute(buf, i)
	if e != nil {
		return "", e
	}
	return buf.String(), nil
}

var (
	oneKb = 1024.0
	oneMb = oneKb * 1024.0
	oneGb = oneMb * 1024.0
)

func sizePretty(raw int64) string {
	f := float64(raw)
	if f < oneKb {
		return fmt.Sprintf("%.0f", f)
	} else if f < oneMb {
		return fmt.Sprintf("%.2fKB", f/oneKb)
	} else if f < oneGb {
		return fmt.Sprintf("%.2fMB", f/oneMb)
	} else {
		return fmt.Sprintf("%.2fGB", f/oneGb)
	}
}
<|endoftext|>


<commit_before>package client

import (
	"bytes"
	"errors"
	"fmt"
	"strconv"

	ct "github.com/google/certificate-transparency/go"
	"golang.org/x/net/context"
)

// GetRawEntries exposes the /ct/v1/get-entries result with only the JSON parsing done.
func (c *LogClient) GetRawEntries(ctx context.Context, start, end int64) (*ct.GetEntriesResponse, error) {
	if end < 0 {
		return nil, errors.New("end should be >= 0")
	}
	if end < start {
		return nil, errors.New("start should be <= end")
	}

	params := map[string]string{
		"start": strconv.FormatInt(start, 10),
		"end":   strconv.FormatInt(end, 10),
	}
	if ctx == nil {
		ctx = context.TODO()
	}

	var resp ct.GetEntriesResponse
	_, err := c.GetAndParse(ctx, ct.GetEntriesPath, params, &resp)
	if err != nil {
		return nil, err
	}

	return &resp, nil
}

// GetEntries attempts to retrieve the entries in the sequence [|start|, |end|] from the CT log server. (see section 4.6.)
// Returns a slice of LeafInputs or a non-nil error.
func (c *LogClient) GetEntries(start, end int64) ([]ct.LogEntry, error) {
	resp, err := c.GetRawEntries(context.TODO(), start, end)
	if err != nil {
		return nil, err
	}
	entries := make([]ct.LogEntry, len(resp.Entries))
	for index, entry := range resp.Entries {
		leaf, err := ct.ReadMerkleTreeLeaf(bytes.NewBuffer(entry.LeafInput))
		if err != nil {
			return nil, err
		}
		entries[index].Leaf = *leaf

		var chain []ct.ASN1Cert
		switch leaf.TimestampedEntry.EntryType {
		case ct.X509LogEntryType:
			chain, err = ct.UnmarshalX509ChainArray(entry.ExtraData)

		case ct.PrecertLogEntryType:
			chain, err = ct.UnmarshalPrecertChainArray(entry.ExtraData)

		default:
			return nil, fmt.Errorf("saw unknown entry type: %v", leaf.TimestampedEntry.EntryType)
		}
		if err != nil {
			return nil, err
		}
		entries[index].Chain = chain
		entries[index].Index = start + int64(index)
	}
	return entries, nil
}
<commit_msg>go/client: use tls library to parse entries<commit_after>package client

import (
	"errors"
	"fmt"
	"strconv"

	ct "github.com/google/certificate-transparency/go"
	"github.com/google/certificate-transparency/go/tls"
	"golang.org/x/net/context"
)

// GetRawEntries exposes the /ct/v1/get-entries result with only the JSON parsing done.
func (c *LogClient) GetRawEntries(ctx context.Context, start, end int64) (*ct.GetEntriesResponse, error) {
	if end < 0 {
		return nil, errors.New("end should be >= 0")
	}
	if end < start {
		return nil, errors.New("start should be <= end")
	}

	params := map[string]string{
		"start": strconv.FormatInt(start, 10),
		"end":   strconv.FormatInt(end, 10),
	}
	if ctx == nil {
		ctx = context.TODO()
	}

	var resp ct.GetEntriesResponse
	_, err := c.GetAndParse(ctx, ct.GetEntriesPath, params, &resp)
	if err != nil {
		return nil, err
	}

	return &resp, nil
}

// GetEntries attempts to retrieve the entries in the sequence [|start|, |end|] from the CT log server. (see section 4.6.)
// Returns a slice of LeafInputs or a non-nil error.
func (c *LogClient) GetEntries(start, end int64) ([]ct.LogEntry, error) {
	resp, err := c.GetRawEntries(context.TODO(), start, end)
	if err != nil {
		return nil, err
	}
	entries := make([]ct.LogEntry, len(resp.Entries))
	for index, entry := range resp.Entries {
		var leaf ct.MerkleTreeLeaf
		if rest, err := tls.Unmarshal(entry.LeafInput, &leaf); err != nil {
			return nil, err
		} else if len(rest) > 0 {
			return nil, fmt.Errorf("trailing data (%d bytes) after MerkleTreeLeaf", len(rest))
		}
		entries[index].Leaf = leaf

		var chain []ct.ASN1Cert
		switch leaf.TimestampedEntry.EntryType {
		case ct.X509LogEntryType:
			var certChain ct.CertificateChain
			if rest, err := tls.Unmarshal(entry.ExtraData, &certChain); err != nil {
				return nil, err
			} else if len(rest) > 0 {
				return nil, fmt.Errorf("trailing data (%d bytes) after CertificateChain", len(rest))
			}
			chain = certChain.Entries

		case ct.PrecertLogEntryType:
			var precertChain ct.PrecertChainEntry
			if rest, err := tls.Unmarshal(entry.ExtraData, &precertChain); err != nil {
				return nil, err
			} else if len(rest) > 0 {
				return nil, fmt.Errorf("trailing data (%d bytes) after PrecertChainEntry", len(rest))
			}
			chain = append(chain, precertChain.PreCertificate)
			chain = append(chain, precertChain.CertificateChain...)

		default:
			return nil, fmt.Errorf("saw unknown entry type: %v", leaf.TimestampedEntry.EntryType)
		}
		entries[index].Chain = chain
		entries[index].Index = start + int64(index)
	}
	return entries, nil
}
<|endoftext|>


<commit_before>package model

// Customer is a model in the "customers" table.
type Customer struct {
	ID   int     `json:"id"`
	Name *string `json:"name" gorm:"not null"`
}

// Order is a model in the "orders" table.
type Order struct {
	ID       int     `json:"id"`
	Subtotal float64 `json:"subtotal" gorm:"type:decimal(18,2)"`

	Customer   Customer `json:"customer" gorm:"ForeignKey:CustomerID"`
	CustomerID int      `json:"-"`

	Products []Product `json:"products" gorm:"many2many:order_products"`
}

// Product is a model in the "products" table.
type Product struct {
	ID    int     `json:"id"`
	Name  *string `json:"name"  gorm:"not null;unique"`
	Price float64 `json:"price" gorm:"type:decimal(18,2)"`
}
<commit_msg>testing: Add omitempty to id fields<commit_after>package model

// Customer is a model in the "customers" table.
type Customer struct {
	ID   int     `json:"id,omitempty"`
	Name *string `json:"name" gorm:"not null"`
}

// Order is a model in the "orders" table.
type Order struct {
	ID       int     `json:"id,omitempty"`
	Subtotal float64 `json:"subtotal" gorm:"type:decimal(18,2)"`

	Customer   Customer `json:"customer" gorm:"ForeignKey:CustomerID"`
	CustomerID int      `json:"-"`

	Products []Product `json:"products" gorm:"many2many:order_products"`
}

// Product is a model in the "products" table.
type Product struct {
	ID    int     `json:"id,omitempty"`
	Name  *string `json:"name"  gorm:"not null;unique"`
	Price float64 `json:"price" gorm:"type:decimal(18,2)"`
}
<|endoftext|>


<commit_before>package goat

import (
	"fmt"
	"log"
	"net"
	"net/http"
	"strconv"
	"strings"
	"sync/atomic"
)

// Handle incoming HTTP connections and serve
func handleHTTP(l net.Listener, httpDoneChan chan bool) {
	// Create shutdown function
	go func(l net.Listener, httpDoneChan chan bool) {
		// Wait for done signal
		<-static.ShutdownChan

		// Close listener
		l.Close()
		log.Println("HTTP listener stopped")
		httpDoneChan <- true
	}(l, httpDoneChan)

	// Log API configuration
	if static.Config.API {
		log.Println("API functionality enabled")
	}

	// Set up HTTP routes for handling functions
	http.HandleFunc("/", parseHTTP)

	// Serve HTTP requests
	http.Serve(l, nil)
}

// Parse incoming HTTP connections before making tracker calls
func parseHTTP(w http.ResponseWriter, r *http.Request) {
	// Count incoming connections
	atomic.AddInt64(&static.HTTP.Current, 1)
	atomic.AddInt64(&static.HTTP.Total, 1)

	// Add header to identify goat
	w.Header().Add("Server", fmt.Sprintf("%s/%s", App, Version))

	// Parse querystring into a Values map
	query := r.URL.Query()

	// Check if IP was previously set
	if query.Get("ip") == "" {
		// If no IP set, detect and store it in query map
		query.Set("ip", strings.Split(r.RemoteAddr, ":")[0])
	}

	// Store current URL path
	url := r.URL.Path

	// Split URL into segments
	urlArr := strings.Split(url, "/")

	// If configured, Detect if client is making an API call
	url = urlArr[1]
	if url == "api" {
		// Output JSON
		w.Header().Add("Content-Type", "application/json")

		// API enabled
		if static.Config.API {
			// RATE LIMITER
			if !apiRateLimit(strings.Split(r.RemoteAddr, ":")[0]) {
				http.Error(w, string(apiErrorResponse("Rate limit exceeded")), 429)
				return
			}

			// API authentication
			auth := new(basicAPIAuthenticator).Auth(r)
			if !auth {
				http.Error(w, string(apiErrorResponse("Authentication failed")), 401)
				return
			}

			// Handle API calls, output JSON
			apiRouter(w, r)
			return
		} else {
			http.Error(w, string(apiErrorResponse("API is currently disabled")), 503)
			return
		}
	}

	// Detect if passkey present in URL
	var passkey string
	if len(urlArr) == 3 {
		passkey = urlArr[1]
		url = urlArr[2]
	}

	// Make sure URL is valid torrent function
	if url != "announce" && url != "scrape" {
		w.Write(httpTrackerError("Malformed announce"))
		return
	}

	// Verify that torrent client is advertising its User-Agent, so we can use a whitelist
	if r.Header.Get("User-Agent") == "" {
		w.Write(httpTrackerError("Your client is not identifying itself"))
		return
	}

	client := r.Header.Get("User-Agent")

	// If configured, verify that torrent client is on whitelist
	if static.Config.Whitelist {
		whitelist := new(whitelistRecord).Load(client, "client")
		if whitelist == (whitelistRecord{}) || !whitelist.Approved {
			w.Write(httpTrackerError("Your client is not whitelisted"))

			// Block things like browsers and web crawlers, because they will just clutter up the table
			if strings.Contains(client, "Mozilla") || strings.Contains(client, "Opera") {
				return
			}

			// Insert unknown clients into list for later approval
			if whitelist == (whitelistRecord{}) {
				whitelist.Client = client
				whitelist.Approved = false

				log.Printf("whitelist: detected new client '%s', awaiting manual approval", client)

				go whitelist.Save()
			}

			return
		}
	}

	// Put client in query map
	query.Set("client", client)

	// Check if server is configured for passkey announce
	if static.Config.Passkey && passkey == "" {
		w.Write(httpTrackerError("No passkey found in announce URL"))
		return
	}

	// Validate passkey if needed
	user := new(userRecord).Load(passkey, "passkey")
	if static.Config.Passkey && user == (userRecord{}) {
		w.Write(httpTrackerError("Invalid passkey"))
		return
	}

	// Put passkey in query map
	query.Set("passkey", user.Passkey)

	// Mark client as HTTP
	query.Set("udp", "0")

	// Get user's total number of active torrents
	seeding := user.Seeding()
	leeching := user.Leeching()
	if seeding == -1 || leeching == -1 {
		w.Write(httpTrackerError("Failed to calculate active torrents"))
		return
	}

	// Verify that client has not exceeded this user's torrent limit
	activeSum := seeding + leeching
	if user.TorrentLimit < activeSum {
		w.Write(httpTrackerError(fmt.Sprintf("Exceeded active torrent limit: %d > %d", activeSum, user.TorrentLimit)))
		return
	}

	// Create channel to return response to client
	resChan := make(chan []byte)

	// Tracker announce
	if url == "announce" {
		// Validate required parameter input
		required := []string{"info_hash", "ip", "port", "uploaded", "downloaded", "left"}
		// Validate required integer input
		reqInt := []string{"port", "uploaded", "downloaded", "left"}

		// Check for required parameters
		for _, r := range required {
			if query.Get(r) == "" {
				w.Write(httpTrackerError("Missing required parameter: " + r))
				close(resChan)
				return
			}
		}

		// Check for all valid integers
		for _, r := range reqInt {
			if query.Get(r) != "" {
				_, err := strconv.Atoi(query.Get(r))
				if err != nil {
					w.Write(httpTrackerError("Invalid integer parameter: " + r))
					close(resChan)
					return
				}
			}
		}

		// Only allow compact announce
		if query.Get("compact") == "" || query.Get("compact") != "1" {
			w.Write(httpTrackerError("Your client does not support compact announce"))
			close(resChan)
			return
		}

		// Perform tracker announce
		go trackerAnnounce(user, query, nil, resChan)
	// Tracker scrape
	} else if url == "scrape" {
		go trackerScrape(user, query, resChan)
	}

	// Wait for response, and send it when ready
	w.Write(<-resChan)
	close(resChan)
	return
}
<commit_msg>Remove accidental commit of rate limiter entry point<commit_after>package goat

import (
	"fmt"
	"log"
	"net"
	"net/http"
	"strconv"
	"strings"
	"sync/atomic"
)

// Handle incoming HTTP connections and serve
func handleHTTP(l net.Listener, httpDoneChan chan bool) {
	// Create shutdown function
	go func(l net.Listener, httpDoneChan chan bool) {
		// Wait for done signal
		<-static.ShutdownChan

		// Close listener
		l.Close()
		log.Println("HTTP listener stopped")
		httpDoneChan <- true
	}(l, httpDoneChan)

	// Log API configuration
	if static.Config.API {
		log.Println("API functionality enabled")
	}

	// Set up HTTP routes for handling functions
	http.HandleFunc("/", parseHTTP)

	// Serve HTTP requests
	http.Serve(l, nil)
}

// Parse incoming HTTP connections before making tracker calls
func parseHTTP(w http.ResponseWriter, r *http.Request) {
	// Count incoming connections
	atomic.AddInt64(&static.HTTP.Current, 1)
	atomic.AddInt64(&static.HTTP.Total, 1)

	// Add header to identify goat
	w.Header().Add("Server", fmt.Sprintf("%s/%s", App, Version))

	// Parse querystring into a Values map
	query := r.URL.Query()

	// Check if IP was previously set
	if query.Get("ip") == "" {
		// If no IP set, detect and store it in query map
		query.Set("ip", strings.Split(r.RemoteAddr, ":")[0])
	}

	// Store current URL path
	url := r.URL.Path

	// Split URL into segments
	urlArr := strings.Split(url, "/")

	// If configured, Detect if client is making an API call
	url = urlArr[1]
	if url == "api" {
		// Output JSON
		w.Header().Add("Content-Type", "application/json")

		// API enabled
		if static.Config.API {
			/*
			// RATE LIMITER
			if !apiRateLimit(strings.Split(r.RemoteAddr, ":")[0]) {
				http.Error(w, string(apiErrorResponse("Rate limit exceeded")), 429)
				return
			}
			*/

			// API authentication
			auth := new(basicAPIAuthenticator).Auth(r)
			if !auth {
				http.Error(w, string(apiErrorResponse("Authentication failed")), 401)
				return
			}

			// Handle API calls, output JSON
			apiRouter(w, r)
			return
		} else {
			http.Error(w, string(apiErrorResponse("API is currently disabled")), 503)
			return
		}
	}

	// Detect if passkey present in URL
	var passkey string
	if len(urlArr) == 3 {
		passkey = urlArr[1]
		url = urlArr[2]
	}

	// Make sure URL is valid torrent function
	if url != "announce" && url != "scrape" {
		w.Write(httpTrackerError("Malformed announce"))
		return
	}

	// Verify that torrent client is advertising its User-Agent, so we can use a whitelist
	if r.Header.Get("User-Agent") == "" {
		w.Write(httpTrackerError("Your client is not identifying itself"))
		return
	}

	client := r.Header.Get("User-Agent")

	// If configured, verify that torrent client is on whitelist
	if static.Config.Whitelist {
		whitelist := new(whitelistRecord).Load(client, "client")
		if whitelist == (whitelistRecord{}) || !whitelist.Approved {
			w.Write(httpTrackerError("Your client is not whitelisted"))

			// Block things like browsers and web crawlers, because they will just clutter up the table
			if strings.Contains(client, "Mozilla") || strings.Contains(client, "Opera") {
				return
			}

			// Insert unknown clients into list for later approval
			if whitelist == (whitelistRecord{}) {
				whitelist.Client = client
				whitelist.Approved = false

				log.Printf("whitelist: detected new client '%s', awaiting manual approval", client)

				go whitelist.Save()
			}

			return
		}
	}

	// Put client in query map
	query.Set("client", client)

	// Check if server is configured for passkey announce
	if static.Config.Passkey && passkey == "" {
		w.Write(httpTrackerError("No passkey found in announce URL"))
		return
	}

	// Validate passkey if needed
	user := new(userRecord).Load(passkey, "passkey")
	if static.Config.Passkey && user == (userRecord{}) {
		w.Write(httpTrackerError("Invalid passkey"))
		return
	}

	// Put passkey in query map
	query.Set("passkey", user.Passkey)

	// Mark client as HTTP
	query.Set("udp", "0")

	// Get user's total number of active torrents
	seeding := user.Seeding()
	leeching := user.Leeching()
	if seeding == -1 || leeching == -1 {
		w.Write(httpTrackerError("Failed to calculate active torrents"))
		return
	}

	// Verify that client has not exceeded this user's torrent limit
	activeSum := seeding + leeching
	if user.TorrentLimit < activeSum {
		w.Write(httpTrackerError(fmt.Sprintf("Exceeded active torrent limit: %d > %d", activeSum, user.TorrentLimit)))
		return
	}

	// Create channel to return response to client
	resChan := make(chan []byte)

	// Tracker announce
	if url == "announce" {
		// Validate required parameter input
		required := []string{"info_hash", "ip", "port", "uploaded", "downloaded", "left"}
		// Validate required integer input
		reqInt := []string{"port", "uploaded", "downloaded", "left"}

		// Check for required parameters
		for _, r := range required {
			if query.Get(r) == "" {
				w.Write(httpTrackerError("Missing required parameter: " + r))
				close(resChan)
				return
			}
		}

		// Check for all valid integers
		for _, r := range reqInt {
			if query.Get(r) != "" {
				_, err := strconv.Atoi(query.Get(r))
				if err != nil {
					w.Write(httpTrackerError("Invalid integer parameter: " + r))
					close(resChan)
					return
				}
			}
		}

		// Only allow compact announce
		if query.Get("compact") == "" || query.Get("compact") != "1" {
			w.Write(httpTrackerError("Your client does not support compact announce"))
			close(resChan)
			return
		}

		// Perform tracker announce
		go trackerAnnounce(user, query, nil, resChan)
	// Tracker scrape
	} else if url == "scrape" {
		go trackerScrape(user, query, resChan)
	}

	// Wait for response, and send it when ready
	w.Write(<-resChan)
	close(resChan)
	return
}
<|endoftext|>


<commit_before>package handlers

import (
	"fmt"
	"net/http"
	"strings"

	"appengine"

	"github.com/icco/natnatnat/models"
	"github.com/pilu/traffic"
)

type TagData struct {
	Posts interface{}
	Tag   string
}

func TagHandler(w traffic.ResponseWriter, r *traffic.Request) {
	c := appengine.NewContext(r.Request)
	tag := r.Param("id")

	if tag == "" {
		http.Redirect(w, r.Request, "/tags", 301)
	}

	if tag != strings.ToLower(tag) {
		http.Redirect(w, r.Request, fmt.Sprintf("/tags/%s", strings.ToLower(tag)), 301)
	}

	entries, err := models.PostsWithTag(c, tag)
	if err != nil {
		http.Error(w, err.Error(), 500)
		return
	}
	data := &TagData{Posts: entries, Tag: tag}
	w.Render("tag", data)
}

type TagsData struct {
	Tags map[string]int
}

func TagsHandler(w traffic.ResponseWriter, r *traffic.Request) {
	c := appengine.NewContext(r.Request)
	w.Render("tags", &TagsData{Tags: models.AllTags(c)})
}

func TagAliasGetHandler(w traffic.ResponseWriter, r *traffic.Request) {
	c := appengine.NewContext(r.Request)
	u := user.Current(c)
	if u == nil {
		url, _ := user.LoginURL(c, "/post/new")
		http.Redirect(w, r.Request, url, 302)
		return
	} else {
		c.Infof("Logged in as: %s", u.String())
	}

	if u != nil && !user.IsAdmin(c) {
		http.Error(w, errors.New("Not a valid user.").Error(), 403)
		return
	} else {
		w.Render("aliases", &TagsData{Tags: models.AllTags(c)})
		return
	}
}

func TagAliasPostHandler(w traffic.ResponseWriter, r *traffic.Request) {
	c := appengine.NewContext(r.Request)
	u := user.Current(c)
	if u == nil {
		url, _ := user.LoginURL(c, "/post/new")
		http.Redirect(w, r.Request, url, 302)
		return
	} else {
		c.Infof("Logged in as: %s", u.String())
	}

	if u != nil && !user.IsAdmin(c) {
		http.Error(w, errors.New("Not a valid user.").Error(), 403)
		return
	} else {
		err := r.ParseForm()
		if err != nil {
			c.Warningf("Couldn't parse form: %v", r)
		}
		xsrf := r.Request.FormValue("xsrf")

		if xsrftoken.Valid(xsrf, models.GetFlagLogError(c, "SESSION_KEY"), u.String(), r.Request.Url.Path) {
			c.Infof("Valid Token!")
		} else {
			c.Infof("Invalid Token...")
			http.Error(w, errors.New("Invalid Token").Error(), 403)
			return
		}

		http.Redirect(w, r.Request, "/aliases", 302)
		return
	}
}
<commit_msg>tweak<commit_after>package handlers

import (
	"errors"
	"fmt"
	"net/http"
	"strings"

	"appengine"
	"appengine/user"

	"code.google.com/p/xsrftoken"

	"github.com/icco/natnatnat/models"
	"github.com/pilu/traffic"
)

type TagData struct {
	Posts interface{}
	Tag   string
}

func TagHandler(w traffic.ResponseWriter, r *traffic.Request) {
	c := appengine.NewContext(r.Request)
	tag := r.Param("id")

	if tag == "" {
		http.Redirect(w, r.Request, "/tags", 301)
	}

	if tag != strings.ToLower(tag) {
		http.Redirect(w, r.Request, fmt.Sprintf("/tags/%s", strings.ToLower(tag)), 301)
	}

	entries, err := models.PostsWithTag(c, tag)
	if err != nil {
		http.Error(w, err.Error(), 500)
		return
	}
	data := &TagData{Posts: entries, Tag: tag}
	w.Render("tag", data)
}

type TagsData struct {
	Tags map[string]int
}

func TagsHandler(w traffic.ResponseWriter, r *traffic.Request) {
	c := appengine.NewContext(r.Request)
	w.Render("tags", &TagsData{Tags: models.AllTags(c)})
}

func TagAliasGetHandler(w traffic.ResponseWriter, r *traffic.Request) {
	c := appengine.NewContext(r.Request)
	u := user.Current(c)
	if u == nil {
		url, _ := user.LoginURL(c, "/post/new")
		http.Redirect(w, r.Request, url, 302)
		return
	} else {
		c.Infof("Logged in as: %s", u.String())
	}

	if u != nil && !user.IsAdmin(c) {
		http.Error(w, errors.New("Not a valid user.").Error(), 403)
		return
	} else {
		w.Render("aliases", &TagsData{Tags: models.AllTags(c)})
		return
	}
}

func TagAliasPostHandler(w traffic.ResponseWriter, r *traffic.Request) {
	c := appengine.NewContext(r.Request)
	u := user.Current(c)
	if u == nil {
		url, _ := user.LoginURL(c, "/post/new")
		http.Redirect(w, r.Request, url, 302)
		return
	} else {
		c.Infof("Logged in as: %s", u.String())
	}

	if u != nil && !user.IsAdmin(c) {
		http.Error(w, errors.New("Not a valid user.").Error(), 403)
		return
	} else {
		err := r.ParseForm()
		if err != nil {
			c.Warningf("Couldn't parse form: %v", r)
		}
		xsrf := r.Request.FormValue("xsrf")

		if xsrftoken.Valid(xsrf, models.GetFlagLogError(c, "SESSION_KEY"), u.String(), r.Request.URL.Path) {
			c.Infof("Valid Token!")
		} else {
			c.Infof("Invalid Token...")
			http.Error(w, errors.New("Invalid Token").Error(), 403)
			return
		}

		http.Redirect(w, r.Request, "/aliases", 302)
		return
	}
}
<|endoftext|>


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

import (
	"context"
	"encoding/base64"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"strings"
	"time"

	"github.com/gorilla/websocket"

	huproxy "github.com/google/huproxy/lib"
)

var (
	writeTimeout = flag.Duration("write_timeout", 10*time.Second, "Write timeout")
	basicAuth    = flag.String("auth", "", "HTTP Basic Auth in @<filename> or <username>:<password> format.")
	verbose      = flag.Bool("verbose", false, "Verbose.")
)

func secretString(s string) (string, error) {
	if strings.HasPrefix(s, "@") {
		fn := s[1:]
		st, err := os.Stat(fn)
		if err != nil {
			return "", err
		}
		p := st.Mode() & os.ModePerm
		if p&0177 > 0 {
			return "", fmt.Errorf("valid permissions for %q is %0o, was %0o", fn, 0600, p)
		}
		b, err := ioutil.ReadFile(fn)
		return strings.TrimSpace(string(b)), err
	}
	return s, nil
}

func dialError(url string, resp *http.Response, err error) {
	if resp != nil {
		extra := ""
		if *verbose {
			b, err := ioutil.ReadAll(resp.Body)
			if err != nil {
				log.Printf("Failed to read HTTP body: %v", err)
			}
			extra = "Body:\n" + string(b)
		}
		log.Fatalf("%s: HTTP error: %d %s\n%s", err, resp.StatusCode, resp.Status, extra)

	}
	log.Fatalf("Dial to %q fail: %v", url, err)
}

func main() {
	flag.Parse()

	if flag.NArg() != 1 {
		log.Fatalf("Want exactly one arg")
	}
	url := flag.Arg(0)

	if *verbose {
		log.Printf("huproxyclient %s", huproxy.Version)
	}

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	dialer := websocket.Dialer{}
	head := map[string][]string{}

	// Add basic auth.
	if *basicAuth != "" {
		ss, err := secretString(*basicAuth)
		if err != nil {
			log.Fatalf("Error reading secret string %q: %v", *basicAuth, err)
		}
		a := base64.StdEncoding.EncodeToString([]byte(ss))
		head["Authorization"] = []string{
			"Basic " + a,
		}
	}

	conn, resp, err := dialer.Dial(url, head)
	if err != nil {
		dialError(url, resp, err)
	}
	defer conn.Close()

	// websocket -> stdout
	go func() {
		for {
			mt, r, err := conn.NextReader()
			if websocket.IsCloseError(err, websocket.CloseNormalClosure) {
				return
			}
			if err != nil {
				log.Fatal(err)
			}
			if mt != websocket.BinaryMessage {
				log.Fatal("blah")
			}
			if _, err := io.Copy(os.Stdout, r); err != nil {
				log.Printf("Reading from websocket: %v", err)
				cancel()
			}
		}
	}()

	// stdin -> websocket
	// TODO: NextWriter() seems to be broken.
	if err := huproxy.File2WS(ctx, cancel, os.Stdin, conn); err == io.EOF {
		if err := conn.WriteControl(websocket.CloseMessage,
			websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""),
			time.Now().Add(*writeTimeout)); err == websocket.ErrCloseSent {
		} else if err != nil {
			log.Printf("Error sending close message: %v", err)
		}
	} else if err != nil {
		log.Printf("reading from stdin: %v", err)
		cancel()
	}

	if ctx.Err() != nil {
		os.Exit(1)
	}
}
<commit_msg>Add insecure wss connection<commit_after>// Copyright 2017 Google Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package main

import (
	"context"
	"crypto/tls"
	"encoding/base64"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"strings"
	"time"

	"github.com/gorilla/websocket"

	huproxy "github.com/google/huproxy/lib"
)

var (
	writeTimeout = flag.Duration("write_timeout", 10*time.Second, "Write timeout")
	basicAuth    = flag.String("auth", "", "HTTP Basic Auth in @<filename> or <username>:<password> format.")
	verbose      = flag.Bool("verbose", false, "Verbose.")
	insecure     = flag.Bool("insecure_conn", false, "Skip certificate validation")
)

func secretString(s string) (string, error) {
	if strings.HasPrefix(s, "@") {
		fn := s[1:]
		st, err := os.Stat(fn)
		if err != nil {
			return "", err
		}
		p := st.Mode() & os.ModePerm
		if p&0177 > 0 {
			return "", fmt.Errorf("valid permissions for %q is %0o, was %0o", fn, 0600, p)
		}
		b, err := ioutil.ReadFile(fn)
		return strings.TrimSpace(string(b)), err
	}
	return s, nil
}

func dialError(url string, resp *http.Response, err error) {
	if resp != nil {
		extra := ""
		if *verbose {
			b, err := ioutil.ReadAll(resp.Body)
			if err != nil {
				log.Printf("Failed to read HTTP body: %v", err)
			}
			extra = "Body:\n" + string(b)
		}
		log.Fatalf("%s: HTTP error: %d %s\n%s", err, resp.StatusCode, resp.Status, extra)

	}
	log.Fatalf("Dial to %q fail: %v", url, err)
}

func main() {
	flag.Parse()

	if flag.NArg() != 1 {
		log.Fatalf("Want exactly one arg")
	}
	url := flag.Arg(0)

	if *verbose {
		log.Printf("huproxyclient %s", huproxy.Version)
	}

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	var tlsConfig *tls.Config
	if *insecure {
		tlsConfig = &tls.Config{InsecureSkipVerify: true}
	} else {
		tlsConfig = nil
	}
	dialer := websocket.Dialer{
		TLSClientConfig: tlsConfig,
	}
	head := map[string][]string{}

	// Add basic auth.
	if *basicAuth != "" {
		ss, err := secretString(*basicAuth)
		if err != nil {
			log.Fatalf("Error reading secret string %q: %v", *basicAuth, err)
		}
		a := base64.StdEncoding.EncodeToString([]byte(ss))
		head["Authorization"] = []string{
			"Basic " + a,
		}
	}

	conn, resp, err := dialer.Dial(url, head)
	if err != nil {
		dialError(url, resp, err)
	}
	defer conn.Close()

	// websocket -> stdout
	go func() {
		for {
			mt, r, err := conn.NextReader()
			if websocket.IsCloseError(err, websocket.CloseNormalClosure) {
				return
			}
			if err != nil {
				log.Fatal(err)
			}
			if mt != websocket.BinaryMessage {
				log.Fatal("blah")
			}
			if _, err := io.Copy(os.Stdout, r); err != nil {
				log.Printf("Reading from websocket: %v", err)
				cancel()
			}
		}
	}()

	// stdin -> websocket
	// TODO: NextWriter() seems to be broken.
	if err := huproxy.File2WS(ctx, cancel, os.Stdin, conn); err == io.EOF {
		if err := conn.WriteControl(websocket.CloseMessage,
			websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""),
			time.Now().Add(*writeTimeout)); err == websocket.ErrCloseSent {
		} else if err != nil {
			log.Printf("Error sending close message: %v", err)
		}
	} else if err != nil {
		log.Printf("reading from stdin: %v", err)
		cancel()
	}

	if ctx.Err() != nil {
		os.Exit(1)
	}
}
<|endoftext|>


<commit_before>// +build linux

package fs2

import (
	"bufio"
	"os"
	"strconv"
	"strings"

	"github.com/opencontainers/runc/libcontainer/cgroups"
	"github.com/opencontainers/runc/libcontainer/cgroups/fscommon"
	"github.com/opencontainers/runc/libcontainer/configs"
)

func isIoSet(r *configs.Resources) bool {
	return r.BlkioWeight != 0 ||
		len(r.BlkioThrottleReadBpsDevice) > 0 ||
		len(r.BlkioThrottleWriteBpsDevice) > 0 ||
		len(r.BlkioThrottleReadIOPSDevice) > 0 ||
		len(r.BlkioThrottleWriteIOPSDevice) > 0
}

func setIo(dirPath string, r *configs.Resources) error {
	if !isIoSet(r) {
		return nil
	}

	if r.BlkioWeight != 0 {
		filename := "io.bfq.weight"
		if err := fscommon.WriteFile(dirPath, filename,
			strconv.FormatUint(uint64(r.BlkioWeight), 10)); err != nil {
			// if io.bfq.weight does not exist, then bfq module is not loaded.
			// Fallback to use io.weight with a conversion scheme
			if !os.IsNotExist(err) {
				return err
			}
			v := cgroups.ConvertBlkIOToIOWeightValue(r.BlkioWeight)
			if err := fscommon.WriteFile(dirPath, "io.weight", strconv.FormatUint(v, 10)); err != nil {
				return err
			}
		}
	}
	for _, td := range r.BlkioThrottleReadBpsDevice {
		if err := fscommon.WriteFile(dirPath, "io.max", td.StringName("rbps")); err != nil {
			return err
		}
	}
	for _, td := range r.BlkioThrottleWriteBpsDevice {
		if err := fscommon.WriteFile(dirPath, "io.max", td.StringName("wbps")); err != nil {
			return err
		}
	}
	for _, td := range r.BlkioThrottleReadIOPSDevice {
		if err := fscommon.WriteFile(dirPath, "io.max", td.StringName("riops")); err != nil {
			return err
		}
	}
	for _, td := range r.BlkioThrottleWriteIOPSDevice {
		if err := fscommon.WriteFile(dirPath, "io.max", td.StringName("wiops")); err != nil {
			return err
		}
	}

	return nil
}

func readCgroup2MapFile(dirPath string, name string) (map[string][]string, error) {
	ret := map[string][]string{}
	f, err := fscommon.OpenFile(dirPath, name, os.O_RDONLY)
	if err != nil {
		return nil, err
	}
	defer f.Close()
	scanner := bufio.NewScanner(f)
	for scanner.Scan() {
		line := scanner.Text()
		parts := strings.Fields(line)
		if len(parts) < 2 {
			continue
		}
		ret[parts[0]] = parts[1:]
	}
	if err := scanner.Err(); err != nil {
		return nil, err
	}
	return ret, nil
}

func statIo(dirPath string, stats *cgroups.Stats) error {
	// more details on the io.stat file format: https://www.kernel.org/doc/Documentation/cgroup-v2.txt
	var ioServiceBytesRecursive []cgroups.BlkioStatEntry
	values, err := readCgroup2MapFile(dirPath, "io.stat")
	if err != nil {
		return err
	}
	for k, v := range values {
		d := strings.Split(k, ":")
		if len(d) != 2 {
			continue
		}
		major, err := strconv.ParseUint(d[0], 10, 0)
		if err != nil {
			return err
		}
		minor, err := strconv.ParseUint(d[1], 10, 0)
		if err != nil {
			return err
		}

		for _, item := range v {
			d := strings.Split(item, "=")
			if len(d) != 2 {
				continue
			}
			op := d[0]

			// Accommodate the cgroup v1 naming
			switch op {
			case "rbytes":
				op = "Read"
			case "wbytes":
				op = "Write"
			}

			value, err := strconv.ParseUint(d[1], 10, 0)
			if err != nil {
				return err
			}

			entry := cgroups.BlkioStatEntry{
				Op:    op,
				Major: major,
				Minor: minor,
				Value: value,
			}
			ioServiceBytesRecursive = append(ioServiceBytesRecursive, entry)
		}
	}
	stats.BlkioStats = cgroups.BlkioStats{IoServiceBytesRecursive: ioServiceBytesRecursive}
	return nil
}
<commit_msg>cgroup2: io: map io.stats to v1 blkio.stats correctly<commit_after>// +build linux

package fs2

import (
	"bufio"
	"os"
	"strconv"
	"strings"

	"github.com/sirupsen/logrus"

	"github.com/opencontainers/runc/libcontainer/cgroups"
	"github.com/opencontainers/runc/libcontainer/cgroups/fscommon"
	"github.com/opencontainers/runc/libcontainer/configs"
)

func isIoSet(r *configs.Resources) bool {
	return r.BlkioWeight != 0 ||
		len(r.BlkioThrottleReadBpsDevice) > 0 ||
		len(r.BlkioThrottleWriteBpsDevice) > 0 ||
		len(r.BlkioThrottleReadIOPSDevice) > 0 ||
		len(r.BlkioThrottleWriteIOPSDevice) > 0
}

func setIo(dirPath string, r *configs.Resources) error {
	if !isIoSet(r) {
		return nil
	}

	if r.BlkioWeight != 0 {
		filename := "io.bfq.weight"
		if err := fscommon.WriteFile(dirPath, filename,
			strconv.FormatUint(uint64(r.BlkioWeight), 10)); err != nil {
			// if io.bfq.weight does not exist, then bfq module is not loaded.
			// Fallback to use io.weight with a conversion scheme
			if !os.IsNotExist(err) {
				return err
			}
			v := cgroups.ConvertBlkIOToIOWeightValue(r.BlkioWeight)
			if err := fscommon.WriteFile(dirPath, "io.weight", strconv.FormatUint(v, 10)); err != nil {
				return err
			}
		}
	}
	for _, td := range r.BlkioThrottleReadBpsDevice {
		if err := fscommon.WriteFile(dirPath, "io.max", td.StringName("rbps")); err != nil {
			return err
		}
	}
	for _, td := range r.BlkioThrottleWriteBpsDevice {
		if err := fscommon.WriteFile(dirPath, "io.max", td.StringName("wbps")); err != nil {
			return err
		}
	}
	for _, td := range r.BlkioThrottleReadIOPSDevice {
		if err := fscommon.WriteFile(dirPath, "io.max", td.StringName("riops")); err != nil {
			return err
		}
	}
	for _, td := range r.BlkioThrottleWriteIOPSDevice {
		if err := fscommon.WriteFile(dirPath, "io.max", td.StringName("wiops")); err != nil {
			return err
		}
	}

	return nil
}

func readCgroup2MapFile(dirPath string, name string) (map[string][]string, error) {
	ret := map[string][]string{}
	f, err := fscommon.OpenFile(dirPath, name, os.O_RDONLY)
	if err != nil {
		return nil, err
	}
	defer f.Close()
	scanner := bufio.NewScanner(f)
	for scanner.Scan() {
		line := scanner.Text()
		parts := strings.Fields(line)
		if len(parts) < 2 {
			continue
		}
		ret[parts[0]] = parts[1:]
	}
	if err := scanner.Err(); err != nil {
		return nil, err
	}
	return ret, nil
}

func statIo(dirPath string, stats *cgroups.Stats) error {
	values, err := readCgroup2MapFile(dirPath, "io.stat")
	if err != nil {
		return err
	}
	// more details on the io.stat file format: https://www.kernel.org/doc/Documentation/cgroup-v2.txt
	var parsedStats cgroups.BlkioStats
	for k, v := range values {
		d := strings.Split(k, ":")
		if len(d) != 2 {
			continue
		}
		major, err := strconv.ParseUint(d[0], 10, 0)
		if err != nil {
			return err
		}
		minor, err := strconv.ParseUint(d[1], 10, 0)
		if err != nil {
			return err
		}

		for _, item := range v {
			d := strings.Split(item, "=")
			if len(d) != 2 {
				continue
			}
			op := d[0]

			// Map to the cgroupv1 naming and layout (in separate tables).
			var targetTable *[]cgroups.BlkioStatEntry
			switch op {
			// Equivalent to cgroupv1's blkio.io_service_bytes.
			case "rbytes":
				op = "Read"
				targetTable = &parsedStats.IoServiceBytesRecursive
			case "wbytes":
				op = "Write"
				targetTable = &parsedStats.IoServiceBytesRecursive
			// Equivalent to cgroupv1's blkio.io_serviced.
			case "rios":
				op = "Read"
				targetTable = &parsedStats.IoServicedRecursive
			case "wios":
				op = "Write"
				targetTable = &parsedStats.IoServicedRecursive
			default:
				// Skip over entries we cannot map to cgroupv1 stats for now.
				// In the future we should expand the stats struct to include
				// them.
				logrus.Debugf("cgroupv2 io stats: skipping over unmappable %s entry", item)
				continue
			}

			value, err := strconv.ParseUint(d[1], 10, 0)
			if err != nil {
				return err
			}

			entry := cgroups.BlkioStatEntry{
				Op:    op,
				Major: major,
				Minor: minor,
				Value: value,
			}
			*targetTable = append(*targetTable, entry)
		}
	}
	stats.BlkioStats = parsedStats
	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"fmt"
	"github.com/holys/goredis"
	"strconv"
	"strings"
)

// Read-Eval-Print Loop
func repl(server RedisServer, commandLine string) string {
	cmds, err := parseEditorCommand(commandLine)
	if err != nil {
		return err.Error()
	}

	if len(cmds) == 0 {
		return ""
	}

	cmd := strings.ToLower(cmds[0])
	if cmd == "help" || cmd == "?" {
		return ""
	} else if cmd == "quit" || cmd == "exit" {
		return ""
	} else if cmd == "clear" {
		return ""
	} else if cmd == "connect" {
		return ""
	} else {
		client := cliConnect(server)
		defer client.Close()
		return cliSendCommand(client, cmds)
	}
}

// Returns the executable path and arguments
func parseEditorCommand(editorCmd string) ([]string, error) {
	var args []string
	state := "start"
	current := ""
	quote := "\""
	for i := 0; i < len(editorCmd); i++ {
		c := editorCmd[i]

		if state == "quotes" {
			if string(c) != quote {
				if c == '\\' {
					i++
					if i >= len(editorCmd) {
						return []string{}, errors.New(fmt.Sprintf("nothing escape in command line: %s", editorCmd))
					}
					current += editorCmd[i-1 : i+1]
				} else {
					current += string(c)
				}
			} else {
				current = "\"" + current + "\""
				unquoted, _ := strconv.Unquote(current)
				args = append(args, unquoted)
				current = ""
				state = "start"
			}
			continue
		}

		if c == '"' || c == '\'' {
			state = "quotes"
			quote = string(c)
			continue
		}

		if state == "arg" {
			if c == ' ' || c == '\t' {
				args = append(args, current)
				current = ""
				state = "start"
			} else {
				current += string(c)
			}
			continue
		}

		if c != ' ' && c != '\t' {
			state = "arg"
			current += string(c)
		}
	}

	if state == "quotes" {
		return []string{}, errors.New(fmt.Sprintf("Unclosed quote in command line: %s", editorCmd))
	}

	if current != "" {
		args = append(args, current)
	}

	if len(args) <= 0 {
		return []string{}, errors.New("Empty command line")
	}

	return args, nil
}

func cliSendCommand(client *goredis.Client, cmds []string) string {
	args := make([]interface{}, len(cmds[1:]))
	for i := range args {
		args[i] = string(cmds[1+i])
	}

	cmd := strings.ToLower(cmds[0])

	r, err := client.Do(cmd, args...)
	if err != nil {
		return fmt.Sprintf("(error) %s", err.Error())
	}

	if cmd == "info" {
		return printInfo(r)
	} else {
		return printReply(0, r)
	}

}

func cliConnect(server RedisServer) *goredis.Client {
	client := goredis.NewClient(server.Addr, server.Password)
	client.SetMaxIdleConns(1)
	return client
}

func printInfo(reply interface{}) string {
	switch reply := reply.(type) {
	case []byte:
		return fmt.Sprintf("%s", reply)
		//some redis proxies don't support this command.
	case goredis.Error:
		return fmt.Sprintf("(error) %s", string(reply))
	}

	return "unknown reply"
}

func printReply(level int, reply interface{}) string {
	switch reply := reply.(type) {
	case int64:
		return fmt.Sprintf("(integer) %d", reply)
	case string:
		return fmt.Sprintf("%s", reply)
	case []byte:
		return fmt.Sprintf("%q", reply)
	case nil:
		return fmt.Sprintf("(nil)")
	case goredis.Error:
		return fmt.Sprintf("(error) %s", string(reply))
	case []interface{}:
		resp := ""
		for i, v := range reply {
			if i != 0 {
				resp += fmt.Sprintf("%s", strings.Repeat(" ", level*4))
			}

			s := fmt.Sprintf("%d) ", i+1)
			resp += fmt.Sprintf("%-4s", s)

			resp += printReply(level+1, v)
			if i != len(reply)-1 {
				resp += "\n"
			}
		}
		return resp
	default:
		return fmt.Sprintf("Unknown reply type: %+v", reply)
	}
}
<commit_msg>Update redis_cli.go<commit_after>package main

import (
	"errors"
	"fmt"
	"github.com/holys/goredis"
	"strconv"
	"strings"
)

// Read-Eval-Print Loop
func repl(server RedisServer, commandLine string) string {
	cmds, err := parseEditorCommand(commandLine)
	if err != nil {
		return err.Error()
	}

	if len(cmds) == 0 {
		return ""
	}

	cmd := strings.ToLower(cmds[0])
	if cmd == "help" || cmd == "?" {
		return ""
	} else if cmd == "quit" || cmd == "exit" {
		return ""
	} else if cmd == "clear" {
		return ""
	} else if cmd == "connect" {
		return ""
	} else {
		client := cliConnect(server)
		defer client.Close()
		return cliSendCommand(client, cmds)
	}
}

// Returns the executable path and arguments
func parseEditorCommand(editorCmd string) ([]string, error) {
	var args []string
	state := "start"
	current := ""
	quote := "\""
	for i := 0; i < len(editorCmd); i++ {
		c := editorCmd[i]

		if state == "quotes" {
			if string(c) != quote {
				if c == '\\' {
					i++
					if i >= len(editorCmd) {
						return []string{}, errors.New(fmt.Sprintf("nothing escape in command line: %s", editorCmd))
					}
					current += editorCmd[i-1 : i+1]
				} else {
					current += string(c)
				}
			} else {
				current = "\"" + current + "\""
				unquoted, _ := strconv.Unquote(current)
				args = append(args, unquoted)
				current = ""
				state = "start"
			}
			continue
		}

		if c == '"' || c == '\'' {
			state = "quotes"
			quote = string(c)
			continue
		}

		if state == "arg" {
			if c == ' ' || c == '\t' {
				args = append(args, current)
				current = ""
				state = "start"
			} else {
				current += string(c)
			}
			continue
		}

		if c != ' ' && c != '\t' {
			state = "arg"
			current += string(c)
		}
	}

	if state == "quotes" {
		return []string{}, errors.New(fmt.Sprintf("Unclosed quote in command line: %s", editorCmd))
	}

	if current != "" {
		args = append(args, current)
	}

	if len(args) <= 0 {
		return []string{}, errors.New("Empty command line")
	}

	return args, nil
}

func cliSendCommand(client *goredis.Client, cmds []string) string {
	args := make([]interface{}, len(cmds[1:]))
	for i := range args {
		args[i] = string(cmds[1+i])
	}

	cmd := strings.ToLower(cmds[0])

	r, err := client.Do(cmd, args...)
	if err != nil {
		return fmt.Sprintf("(error) %s", err.Error())
	}

	if cmd == "info" {
		return printInfo(r)
	} else {
		return printReply(0, r)
	}

}

func cliConnect(server RedisServer) *goredis.Client {
	client := goredis.NewClient(server.Addr, server.Password)
	client.SetMaxIdleConns(1)
	_, _ = client.Do("SELECT", server.DefaultDb)
	return client
}

func printInfo(reply interface{}) string {
	switch reply := reply.(type) {
	case []byte:
		return fmt.Sprintf("%s", reply)
		//some redis proxies don't support this command.
	case goredis.Error:
		return fmt.Sprintf("(error) %s", string(reply))
	}

	return "unknown reply"
}

func printReply(level int, reply interface{}) string {
	switch reply := reply.(type) {
	case int64:
		return fmt.Sprintf("(integer) %d", reply)
	case string:
		return fmt.Sprintf("%s", reply)
	case []byte:
		return fmt.Sprintf("%q", reply)
	case nil:
		return fmt.Sprintf("(nil)")
	case goredis.Error:
		return fmt.Sprintf("(error) %s", string(reply))
	case []interface{}:
		resp := ""
		for i, v := range reply {
			if i != 0 {
				resp += fmt.Sprintf("%s", strings.Repeat(" ", level*4))
			}

			s := fmt.Sprintf("%d) ", i+1)
			resp += fmt.Sprintf("%-4s", s)

			resp += printReply(level+1, v)
			if i != len(reply)-1 {
				resp += "\n"
			}
		}
		return resp
	default:
		return fmt.Sprintf("Unknown reply type: %+v", reply)
	}
}
<|endoftext|>


<commit_before>// Package redisence provides simple user presence system
package redisence

import (
	"fmt"
	"strconv"
	"time"

	gredis "github.com/garyburd/redigo/redis"
	"github.com/koding/redis"
)

// Status defines what is the current status of a user
// in presence system
type Status int

const (
	Offline Status = iota
	Online
	Closed
)

// Event is the data type for
// occuring events in the system
type Event struct {
	// Id is the given key by the application
	Id string

	// Status holds the changing type of event
	Status Status
}

// Prefix for redisence package
const RedisencePrefix = "redisence"

// Session holds the required connection data for redis
type Session struct {
	// main redis connection
	redis *redis.RedisSession

	// inactiveDuration specifies no-probe allowance time
	inactiveDuration time.Duration

	// receiving offline events pattern
	becameOfflinePattern string

	// receiving online events pattern
	becameOnlinePattern string
}

// New creates a session for any broker system that is architected to use,
// communicate, forward events to the presence system
func New(server string, db int, inactiveDuration time.Duration) (*Session, error) {
	redis, err := redis.NewRedisSession(&redis.RedisConf{Server: server, DB: db})
	if err != nil {
		return nil, err
	}
	redis.SetPrefix(RedisencePrefix)

	return &Session{
		redis:                redis,
		becameOfflinePattern: fmt.Sprintf("__keyevent@%d__:expired", db),
		becameOnlinePattern:  fmt.Sprintf("__keyevent@%d__:set", db),
		inactiveDuration:     inactiveDuration,
	}, nil
}

// Ping resets the expiration time for any given key
// if key doesnt exists, it means user is now online and should be set as online
// Whenever application gets any prob from a client
// should call this function
func (s *Session) Online(ids ...string) error {
	if len(ids) == 1 {
		// if member exits increase ttl
		if s.redis.Expire(ids[0], s.inactiveDuration) == nil {
			return nil
		}

		// if member doesnt exist set it
		return s.redis.Setex(ids[0], s.inactiveDuration, ids[0])
	}

	existance, err := s.sendMultiExpire(ids)
	if err != nil {
		return err
	}

	return s.sendMultiSetIfRequired(ids, existance)
}

// Offline sets given ids as offline, ignores any error
// since not exist keys returned as nilErr
func (s *Session) Offline(ids ...string) error {
	if len(ids) == 1 {
		if s.redis.Expire(ids[0], time.Second*0) == nil {
			return nil
		}
	}

	s.sendMultiExpire(ids, "0")
	return nil
}

// sendMultiSetIfRequired accepts set of ids and their existtance status
// traverse over them and any key is not exists in db, set them in a multi/exec
// request
func (s *Session) sendMultiSetIfRequired(ids []string, existance []int) error {
	if len(ids) != len(existance) {
		return fmt.Errorf("Length is not same Ids: %d Existance: %d", len(ids), len(existance))
	}

	// cache inactive duration as string
	seconds := strconv.Itoa(int(s.inactiveDuration.Seconds()))

	// get one connection from pool
	c := s.redis.Pool().Get()

	// item count for non-existent members
	notExistsCount := 0

	for i, exists := range existance {
		// `0` means, member doesnt exists in presence system
		if exists != 0 {
			continue
		}

		// init multi command lazily
		if notExistsCount == 0 {
			c.Send("MULTI")
		}

		notExistsCount++
		c.Send("SETEX", s.redis.AddPrefix(ids[i]), seconds, ids[i])
	}

	// execute multi command if only we flushed some to connection
	if notExistsCount != 0 {
		// ignore values
		if _, err := c.Do("EXEC"); err != nil {
			return err
		}
	}

	// do not forget to close the connection
	if err := c.Close(); err != nil {
		return err
	}

	return nil
}

// sendMultiExpire if the system tries to update more than one key at a time
// inorder to leverage rtt, send multi expire
func (s *Session) sendMultiExpire(ids []string, duration string) ([]int, error) {
	// get one connection from pool
	c := s.redis.Pool().Get()

	// init multi command
	c.Send("MULTI")

	// send expire command for all members
	for _, id := range ids {
		c.Send("EXPIRE", s.redis.AddPrefix(id), seconds)
	}

	// execute command
	r, err := c.Do("EXEC")
	if err != nil {
		return make([]int, 0), err
	}

	// close connection
	if err := c.Close(); err != nil {
		return make([]int, 0), err
	}

	values, err := s.redis.Values(r)
	if err != nil {
		return make([]int, 0), err
	}

	res := make([]int, len(values))
	for i, value := range values {
		res[i], err = s.redis.Int(value)
		if err != nil {
			// what about returning half-generated slice?
			// instead of an empty one
			return make([]int, 0), err
		}

	}

	return res, nil
}

// MultipleStatus returns the current status multiple keys from system
func (s *Session) MultipleStatus(ids []string) ([]Event, error) {
	// get one connection from pool
	c := s.redis.Pool().Get()

	// init multi command
	c.Send("MULTI")

	// send expire command for all members
	for _, id := range ids {
		c.Send("EXISTS", s.redis.AddPrefix(id))
	}

	// execute command
	r, err := c.Do("EXEC")
	if err != nil {
		return make([]Event, 0), err
	}

	// close connection
	if err := c.Close(); err != nil {
		return make([]Event, 0), err
	}

	values, err := s.redis.Values(r)
	if err != nil {
		return make([]Event, 0), err
	}

	res := make([]Event, len(values))
	for i, value := range values {
		status, err := s.redis.Int(value)
		if err != nil {
			return make([]Event, 0), err
		}

		res[i] = Event{
			Id: ids[i],
			// cast redis response to Status
			Status: Status(status),
		}
	}

	return res, nil
}

// Status returns the current status a key from system
func (s *Session) Status(id string) (Event, error) {
	res := Event{
		Id:     id,
		Status: Offline,
	}

	if s.redis.Exists(id) {
		res.Status = Online
	}

	return res, nil
}
// createEvent Creates the event with the required properties
func (s *Session) createEvent(n gredis.PMessage) Event {
	e := Event{}

	switch n.Pattern {
	case s.becameOfflinePattern:
		e.Id = string(n.Data[len(RedisencePrefix)+1:])
		e.Status = Offline
	case s.becameOnlinePattern:
		e.Id = string(n.Data[len(RedisencePrefix)+1:])
		e.Status = Online
	default:
		//ignore other events
	}

	return e
}

// ListenStatusChanges pubscribes to the redis and
// gets online and offline status changes from it
func (s *Session) ListenStatusChanges(events chan Event) {
	psc := s.redis.CreatePubSubConn()

	psc.PSubscribe(s.becameOnlinePattern, s.becameOfflinePattern)

	var wg sync.WaitGroup
	wg.Add(1)

	go func() {
		defer wg.Done()
		for {
			switch n := psc.Receive().(type) {
			case gredis.PMessage:
				events <- s.createEvent(n)
			case error:
				fmt.Printf("error: %v\n", n)
				return
			}
		}
	}()

	// avoid lock
	go func() {
		wg.Wait()
		psc.PUnsubscribe(s.becameOfflinePattern, s.becameOnlinePattern)
		psc.Close()
		events <- Event{Status: Closed}
	}()
}
<commit_msg>Redisence: refactor event listening system<commit_after>// Package redisence provides simple user presence system
package redisence

import (
	"fmt"
	"strconv"
	"time"

	gredis "github.com/garyburd/redigo/redis"
	"github.com/koding/redis"
)

// Status defines what is the current status of a user
// in presence system
type Status int

const (
	Offline Status = iota
	Online
	Closed
)

// Event is the data type for
// occuring events in the system
type Event struct {
	// Id is the given key by the application
	Id string

	// Status holds the changing type of event
	Status Status
}

// Prefix for redisence package
const RedisencePrefix = "redisence"

// Session holds the required connection data for redis
type Session struct {
	// main redis connection
	redis *redis.RedisSession

	// inactiveDuration specifies no-probe allowance time
	inactiveDuration time.Duration

	// receiving offline events pattern
	becameOfflinePattern string

	// receiving online events pattern
	becameOnlinePattern string

	// inactiveDurationAsString holds the expiration duration as string
	inactiveDurationAsString string

	// closeChan is used for giving close signal
	closeChan chan bool

	// errChan pipe all errors  the this channel
	errChan chan error

	// closed holds the status of connection
	closed bool

	//psc holds the pubsub channel if opened
	psc *gredis.PubSubConn
}

// New creates a session for any broker system that is architected to use,
// communicate, forward events to the presence system
func New(server string, db int, inactiveDuration time.Duration) (*Session, error) {
	redis, err := redis.NewRedisSession(&redis.RedisConf{Server: server, DB: db})
	if err != nil {
		return nil, err
	}
	redis.SetPrefix(RedisencePrefix)

	return &Session{
		redis:                redis,
		becameOfflinePattern: fmt.Sprintf("__keyevent@%d__:expired", db),
		becameOnlinePattern:  fmt.Sprintf("__keyevent@%d__:set", db),
		inactiveDuration:     inactiveDuration,
		// cache inactive duration as string
		inactiveDurationAsString: strconv.Itoa(int(inactiveDuration.Seconds())),
		closeChan:                make(chan bool, 1),
		errChan:                  make(chan error, 1),
	}, nil
}

// Close closes the redis connection gracefully
func (s *Session) Close() error {
	s.closeChan <- true
	return s.close()
}
// Ping resets the expiration time for any given key
// if key doesnt exists, it means user is now online and should be set as online
// Whenever application gets any prob from a client
// should call this function
func (s *Session) Online(ids ...string) error {
	if len(ids) == 1 {
		// if member exits increase ttl
		if s.redis.Expire(ids[0], s.inactiveDuration) == nil {
			return nil
		}

		// if member doesnt exist set it
		return s.redis.Setex(ids[0], s.inactiveDuration, ids[0])
	}

	existance, err := s.sendMultiExpire(ids)
	if err != nil {
		return err
	}

	return s.sendMultiSetIfRequired(ids, existance)
}

// Offline sets given ids as offline, ignores any error
// since not exist keys returned as nilErr
func (s *Session) Offline(ids ...string) error {
	if len(ids) == 1 {
		if s.redis.Expire(ids[0], time.Second*0) == nil {
			return nil
		}
	}

	s.sendMultiExpire(ids, "0")
	return nil
}

// sendMultiSetIfRequired accepts set of ids and their existtance status
// traverse over them and any key is not exists in db, set them in a multi/exec
// request
func (s *Session) sendMultiSetIfRequired(ids []string, existance []int) error {
	if len(ids) != len(existance) {
		return fmt.Errorf("Length is not same Ids: %d Existance: %d", len(ids), len(existance))
	}

	// cache inactive duration as string
	seconds := strconv.Itoa(int(s.inactiveDuration.Seconds()))

	// get one connection from pool
	c := s.redis.Pool().Get()

	// item count for non-existent members
	notExistsCount := 0

	for i, exists := range existance {
		// `0` means, member doesnt exists in presence system
		if exists != 0 {
			continue
		}

		// init multi command lazily
		if notExistsCount == 0 {
			c.Send("MULTI")
		}

		notExistsCount++
		c.Send("SETEX", s.redis.AddPrefix(ids[i]), seconds, ids[i])
	}

	// execute multi command if only we flushed some to connection
	if notExistsCount != 0 {
		// ignore values
		if _, err := c.Do("EXEC"); err != nil {
			return err
		}
	}

	// do not forget to close the connection
	if err := c.Close(); err != nil {
		return err
	}

	return nil
}

// sendMultiExpire if the system tries to update more than one key at a time
// inorder to leverage rtt, send multi expire
func (s *Session) sendMultiExpire(ids []string, duration string) ([]int, error) {
	// get one connection from pool
	c := s.redis.Pool().Get()

	// init multi command
	c.Send("MULTI")

	// send expire command for all members
	for _, id := range ids {
		c.Send("EXPIRE", s.redis.AddPrefix(id), seconds)
	}

	// execute command
	r, err := c.Do("EXEC")
	if err != nil {
		return make([]int, 0), err
	}

	// close connection
	if err := c.Close(); err != nil {
		return make([]int, 0), err
	}

	values, err := s.redis.Values(r)
	if err != nil {
		return make([]int, 0), err
	}

	res := make([]int, len(values))
	for i, value := range values {
		res[i], err = s.redis.Int(value)
		if err != nil {
			// what about returning half-generated slice?
			// instead of an empty one
			return make([]int, 0), err
		}

	}

	return res, nil
}

// MultipleStatus returns the current status multiple keys from system
func (s *Session) MultipleStatus(ids []string) ([]Event, error) {
	// get one connection from pool
	c := s.redis.Pool().Get()

	// init multi command
	c.Send("MULTI")

	// send expire command for all members
	for _, id := range ids {
		c.Send("EXISTS", s.redis.AddPrefix(id))
	}

	// execute command
	r, err := c.Do("EXEC")
	if err != nil {
		return make([]Event, 0), err
	}

	// close connection
	if err := c.Close(); err != nil {
		return make([]Event, 0), err
	}

	values, err := s.redis.Values(r)
	if err != nil {
		return make([]Event, 0), err
	}

	res := make([]Event, len(values))
	for i, value := range values {
		status, err := s.redis.Int(value)
		if err != nil {
			return make([]Event, 0), err
		}

		res[i] = Event{
			Id: ids[i],
			// cast redis response to Status
			Status: Status(status),
		}
	}

	return res, nil
}

// Status returns the current status a key from system
func (s *Session) Status(id string) (Event, error) {
	res := Event{
		Id:     id,
		Status: Offline,
	}

	if s.redis.Exists(id) {
		res.Status = Online
	}

	return res, nil
}

// ListenStatusChanges subscribes with a pattern to the redis and
// gets online and offline status changes from it
func (s *Session) ListenStatusChanges(events chan Event) {
	s.psc = s.redis.CreatePubSubConn()
	s.psc.PSubscribe(s.becameOnlinePattern, s.becameOfflinePattern)

	go s.listenEvents(events)

	<-s.closeChan
	events <- Event{Status: Closed}
}

func (s *Session) close() error {
	s.closed = true
	if s.psc != nil {
		s.psc.PUnsubscribe()
	}

	return s.redis.Close()
}

// createEvent Creates the event with the required properties
func (s *Session) listenEvents(events chan Event) {
	for {
		if s.closed {
			break
		}
		switch n := s.psc.Receive().(type) {
		case gredis.PMessage:
			events <- s.createEvent(n)
		case error:
			s.errChan <- n
			return
		}
	}
}

// createEvent Creates the event with the required properties
func (s *Session) createEvent(n gredis.PMessage) Event {
	e := Event{}

	switch n.Pattern {
	case s.becameOfflinePattern:
		e.Id = string(n.Data[len(RedisencePrefix)+1:])
		e.Status = Offline
	case s.becameOnlinePattern:
		e.Id = string(n.Data[len(RedisencePrefix)+1:])
		e.Status = Online
	default:
		//ignore other events
	}

	return e
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"regexp"

	"github.com/gocolly/colly"
)

func main() {
	// Instantiate default collector
	c := colly.NewCollector(
		// Visit only root url and urls which start with "e" or "h" on httpbin.org
		colly.URLFilters(
			regexp.MustCompile("http://httpbin\\.org/(|e.+)$"),
			regexp.MustCompile("http://httpbin\\.org/h.+"),
		),
	)

	// On every a element which has href attribute call callback
	c.OnHTML("a[href]", func(e *colly.HTMLElement) {
		link := e.Attr("href")
		// Print link
		fmt.Printf("Link found: %q -> %s\n", e.Text, link)
		// Visit link found on page
		// Only those links are visited which are matched by  any of the URLFilter regexps
		c.Visit(e.Request.AbsoluteURL(link))
	})

	// Before making a request print "Visiting ..."
	c.OnRequest(func(r *colly.Request) {
		fmt.Println("Visiting", r.URL.String())
	})

	// Start scraping on https://hackerspaces.org
	c.Visit("http://httpbin.org/")
}
<commit_msg>Fix link in comment<commit_after>package main

import (
	"fmt"
	"regexp"

	"github.com/gocolly/colly"
)

func main() {
	// Instantiate default collector
	c := colly.NewCollector(
		// Visit only root url and urls which start with "e" or "h" on httpbin.org
		colly.URLFilters(
			regexp.MustCompile("http://httpbin\\.org/(|e.+)$"),
			regexp.MustCompile("http://httpbin\\.org/h.+"),
		),
	)

	// On every a element which has href attribute call callback
	c.OnHTML("a[href]", func(e *colly.HTMLElement) {
		link := e.Attr("href")
		// Print link
		fmt.Printf("Link found: %q -> %s\n", e.Text, link)
		// Visit link found on page
		// Only those links are visited which are matched by  any of the URLFilter regexps
		c.Visit(e.Request.AbsoluteURL(link))
	})

	// Before making a request print "Visiting ..."
	c.OnRequest(func(r *colly.Request) {
		fmt.Println("Visiting", r.URL.String())
	})

	// Start scraping on http://httpbin.org
	c.Visit("http://httpbin.org/")
}
<|endoftext|>


<commit_before>// Copyright 2017 The OPA Authors.  All rights reserved.
// Use of this source code is governed by an Apache2
// license that can be found in the LICENSE file.

// Package rego exposes high level APIs for evaluating Rego policies.
package rego

import (
	"context"
	"fmt"
	"strings"

	"github.com/open-policy-agent/opa/ast"
	"github.com/open-policy-agent/opa/storage"
	"github.com/open-policy-agent/opa/topdown"
)

// Result defines the output of Rego evaluation.
type Result struct {
	Expressions []*ExpressionValue `json:"expressions"`
	Bindings    Vars               `json:"bindings,omitempty"`
}

func newResult() Result {
	return Result{
		Bindings: Vars{},
	}
}

// Location defines a position in a Rego query or module.
type Location struct {
	Row int `json:"row"`
	Col int `json:"col"`
}

// ExpressionValue defines the value of an expression in a Rego query.
type ExpressionValue struct {
	Value    interface{} `json:"value"`
	Text     string      `json:"text"`
	Location *Location   `json:"location"`
}

func newExpressionValue(expr *ast.Expr, value interface{}) *ExpressionValue {
	return &ExpressionValue{
		Value: value,
		Text:  string(expr.Location.Text),
		Location: &Location{
			Row: expr.Location.Row,
			Col: expr.Location.Col,
		},
	}
}

// ResultSet represents a collection of output from Rego evaluation. An empty
// result set represents an undefined query.
type ResultSet []Result

// Vars represents a collection of variable bindings. The keys are the variable
// names and the values are the binding values.
type Vars map[string]interface{}

// Errors represents a collection of errors returned when evaluating Rego.
type Errors []error

func (errs Errors) Error() string {
	if len(errs) == 0 {
		return "no error"
	}
	if len(errs) == 1 {
		return fmt.Sprintf("1 error occurred: %v", errs[0].Error())
	}
	buf := []string{fmt.Sprintf("%v errors occurred", len(errs))}
	for _, err := range errs {
		buf = append(buf, err.Error())
	}
	return strings.Join(buf, "\n")
}

// Rego constructs a query and can be evaluated to obtain results.
type Rego struct {
	query     string
	pkg       string
	imports   []string
	rawInput  *interface{}
	input     ast.Value
	modules   []rawModule
	compiler  *ast.Compiler
	storage   *storage.Storage
	termVarID int
}

// Query returns an argument that sets the Rego query.
func Query(q string) func(r *Rego) {
	return func(r *Rego) {
		r.query = q
	}
}

// Package returns an argument that sets the Rego package on the query's
// context.
func Package(p string) func(r *Rego) {
	return func(r *Rego) {
		r.pkg = p
	}
}

// Imports returns an argument that adds a Rego import to the query's context.
func Imports(p []string) func(r *Rego) {
	return func(r *Rego) {
		r.imports = append(r.imports, p...)
	}
}

// Input returns an argument that sets the Rego input document.
func Input(x interface{}) func(r *Rego) {
	return func(r *Rego) {
		r.rawInput = &x
	}
}

// Module returns an argument that adds a Rego module.
func Module(filename, input string) func(r *Rego) {
	return func(r *Rego) {
		r.modules = append(r.modules, rawModule{
			filename: filename,
			module:   input,
		})
	}
}

// Compiler returns an argument that sets the Rego compiler.
func Compiler(c *ast.Compiler) func(r *Rego) {
	return func(r *Rego) {
		r.compiler = c
	}
}

// Storage returns an argument that sets the policy engine's data storage layer.
func Storage(s *storage.Storage) func(r *Rego) {
	return func(r *Rego) {
		r.storage = s
	}
}

// New returns a new Rego object.
func New(options ...func(*Rego)) *Rego {
	r := &Rego{}

	for _, option := range options {
		option(r)
	}

	if r.compiler == nil {
		r.compiler = ast.NewCompiler()
	}

	if r.storage == nil {
		r.storage = storage.New(storage.InMemoryConfig())
	}

	return r
}

// Eval evaluates this Rego object and returns a ResultSet.
func (r *Rego) Eval(ctx context.Context) (ResultSet, error) {

	if len(r.query) == 0 {
		return nil, fmt.Errorf("cannot evaluate empty query")
	}

	// Parse inputs
	parsed, query, err := r.parse()
	if err != nil {
		return nil, err
	}

	// If the query contains expressions that consist of a single term, rewrite
	// those expressions so that we capture the value of the term in a variable
	// that can be included in the result.
	for i := range query {
		if !query[i].Negated {
			if term, ok := query[i].Terms.(*ast.Term); ok {
				query[i].Terms = ast.Equality.Expr(term, r.generateTermVar()).Terms
			}
		}
	}

	// Compile inputs
	compiled, err := r.compile(parsed, query)
	if err != nil {
		return nil, err
	}

	// Prepare storage layer. Transaction could be an argument in the future.
	txn, err := r.storage.NewTransaction(ctx)
	if err != nil {
		return nil, err
	}

	// Evaluate query
	return r.eval(ctx, compiled, txn)
}

func (r *Rego) parse() (map[string]*ast.Module, ast.Body, error) {
	var errs Errors
	parsed := map[string]*ast.Module{}

	for _, module := range r.modules {
		p, err := module.Parse()
		if err != nil {
			errs = append(errs, err)
		}
		parsed[module.filename] = p
	}

	query, err := ast.ParseBody(r.query)
	if err != nil {
		errs = append(errs, err)
	}

	if len(errs) > 0 {
		return nil, nil, errs
	}

	return parsed, query, nil
}

func (r *Rego) compile(modules map[string]*ast.Module, query ast.Body) (ast.Body, error) {

	if len(modules) > 0 {
		r.compiler.Compile(modules)

		if r.compiler.Failed() {
			var errs Errors
			for _, err := range r.compiler.Errors {
				errs = append(errs, err)
			}
			return nil, errs
		}
	}

	var qctx *ast.QueryContext

	if r.pkg != "" {
		pkg, err := ast.ParsePackage(fmt.Sprintf("package %v", r.pkg))
		if err != nil {
			return nil, err
		}
		qctx = qctx.WithPackage(pkg)
	}

	if len(r.imports) > 0 {
		s := make([]string, len(r.imports))
		for i := range r.imports {
			s[i] = fmt.Sprintf("import %v", r.imports[i])
		}
		imports, err := ast.ParseImports(strings.Join(s, "\n"))
		if err != nil {
			return nil, err
		}
		qctx = qctx.WithImports(imports)
	}

	if r.rawInput != nil {
		val, err := ast.InterfaceToValue(*r.rawInput)
		if err != nil {
			return nil, err
		}
		qctx = qctx.WithInput(val)
		r.input = val
	}

	return r.compiler.QueryCompiler().WithContext(qctx).Compile(query)
}

func (r *Rego) eval(ctx context.Context, compiled ast.Body, txn storage.Transaction) (rs ResultSet, err error) {

	t := topdown.New(ctx, compiled, r.compiler, r.storage, txn)

	if r.input != nil {
		t.Input = r.input
	}

	exprs := map[*ast.Expr]struct{}{}

	err = topdown.Eval(t, func(t *topdown.Topdown) error {
		result := newResult()
		for key, value := range t.Vars() {
			val, err := topdown.ValueToInterface(value, t)
			if err != nil {
				return err
			}
			if !isTermVar(key) {
				result.Bindings[string(key)] = val
			} else if expr := findExprForTermVar(compiled, key); expr != nil {
				result.Expressions = append(result.Expressions, newExpressionValue(expr, val))
				exprs[expr] = struct{}{}
			}
		}
		for _, expr := range compiled {
			if _, ok := exprs[expr]; !ok {
				result.Expressions = append(result.Expressions, newExpressionValue(expr, true))
			}
		}
		rs = append(rs, result)
		return nil
	})

	if err != nil {
		return nil, err
	}

	if len(rs) == 0 {
		return nil, nil
	}

	return rs, nil
}

func (r *Rego) generateTermVar() *ast.Term {
	r.termVarID++
	return ast.VarTerm(ast.WildcardPrefix + fmt.Sprintf("term%v", r.termVarID))
}

func isTermVar(v ast.Var) bool {
	return strings.HasPrefix(string(v), ast.WildcardPrefix+"term")
}

func findExprForTermVar(query ast.Body, v ast.Var) *ast.Expr {
	for i := range query {
		vis := ast.NewVarVisitor()
		ast.Walk(vis, query[i])
		if vis.Vars().Contains(v) {
			return query[i]
		}
	}
	return nil
}

type rawModule struct {
	filename string
	module   string
}

func (m rawModule) Parse() (*ast.Module, error) {
	return ast.ParseModule(m.filename, m.module)
}
<commit_msg>Fix rego.Eval to close transactions<commit_after>// Copyright 2017 The OPA Authors.  All rights reserved.
// Use of this source code is governed by an Apache2
// license that can be found in the LICENSE file.

// Package rego exposes high level APIs for evaluating Rego policies.
package rego

import (
	"context"
	"fmt"
	"strings"

	"github.com/open-policy-agent/opa/ast"
	"github.com/open-policy-agent/opa/storage"
	"github.com/open-policy-agent/opa/topdown"
)

// Result defines the output of Rego evaluation.
type Result struct {
	Expressions []*ExpressionValue `json:"expressions"`
	Bindings    Vars               `json:"bindings,omitempty"`
}

func newResult() Result {
	return Result{
		Bindings: Vars{},
	}
}

// Location defines a position in a Rego query or module.
type Location struct {
	Row int `json:"row"`
	Col int `json:"col"`
}

// ExpressionValue defines the value of an expression in a Rego query.
type ExpressionValue struct {
	Value    interface{} `json:"value"`
	Text     string      `json:"text"`
	Location *Location   `json:"location"`
}

func newExpressionValue(expr *ast.Expr, value interface{}) *ExpressionValue {
	return &ExpressionValue{
		Value: value,
		Text:  string(expr.Location.Text),
		Location: &Location{
			Row: expr.Location.Row,
			Col: expr.Location.Col,
		},
	}
}

// ResultSet represents a collection of output from Rego evaluation. An empty
// result set represents an undefined query.
type ResultSet []Result

// Vars represents a collection of variable bindings. The keys are the variable
// names and the values are the binding values.
type Vars map[string]interface{}

// Errors represents a collection of errors returned when evaluating Rego.
type Errors []error

func (errs Errors) Error() string {
	if len(errs) == 0 {
		return "no error"
	}
	if len(errs) == 1 {
		return fmt.Sprintf("1 error occurred: %v", errs[0].Error())
	}
	buf := []string{fmt.Sprintf("%v errors occurred", len(errs))}
	for _, err := range errs {
		buf = append(buf, err.Error())
	}
	return strings.Join(buf, "\n")
}

// Rego constructs a query and can be evaluated to obtain results.
type Rego struct {
	query     string
	pkg       string
	imports   []string
	rawInput  *interface{}
	input     ast.Value
	modules   []rawModule
	compiler  *ast.Compiler
	storage   *storage.Storage
	termVarID int
}

// Query returns an argument that sets the Rego query.
func Query(q string) func(r *Rego) {
	return func(r *Rego) {
		r.query = q
	}
}

// Package returns an argument that sets the Rego package on the query's
// context.
func Package(p string) func(r *Rego) {
	return func(r *Rego) {
		r.pkg = p
	}
}

// Imports returns an argument that adds a Rego import to the query's context.
func Imports(p []string) func(r *Rego) {
	return func(r *Rego) {
		r.imports = append(r.imports, p...)
	}
}

// Input returns an argument that sets the Rego input document. Input should be
// a native Go value representing the input document.
func Input(x interface{}) func(r *Rego) {
	return func(r *Rego) {
		r.rawInput = &x
	}
}

// Module returns an argument that adds a Rego module.
func Module(filename, input string) func(r *Rego) {
	return func(r *Rego) {
		r.modules = append(r.modules, rawModule{
			filename: filename,
			module:   input,
		})
	}
}

// Compiler returns an argument that sets the Rego compiler.
func Compiler(c *ast.Compiler) func(r *Rego) {
	return func(r *Rego) {
		r.compiler = c
	}
}

// Storage returns an argument that sets the policy engine's data storage layer.
func Storage(s *storage.Storage) func(r *Rego) {
	return func(r *Rego) {
		r.storage = s
	}
}

// New returns a new Rego object.
func New(options ...func(*Rego)) *Rego {
	r := &Rego{}

	for _, option := range options {
		option(r)
	}

	if r.compiler == nil {
		r.compiler = ast.NewCompiler()
	}

	if r.storage == nil {
		r.storage = storage.New(storage.InMemoryConfig())
	}

	return r
}

// Eval evaluates this Rego object and returns a ResultSet.
func (r *Rego) Eval(ctx context.Context) (ResultSet, error) {

	if len(r.query) == 0 {
		return nil, fmt.Errorf("cannot evaluate empty query")
	}

	// Parse inputs
	parsed, query, err := r.parse()
	if err != nil {
		return nil, err
	}

	// If the query contains expressions that consist of a single term, rewrite
	// those expressions so that we capture the value of the term in a variable
	// that can be included in the result.
	for i := range query {
		if !query[i].Negated {
			if term, ok := query[i].Terms.(*ast.Term); ok {
				query[i].Terms = ast.Equality.Expr(term, r.generateTermVar()).Terms
			}
		}
	}

	// Compile inputs
	compiled, err := r.compile(parsed, query)
	if err != nil {
		return nil, err
	}

	// Prepare storage layer. Transaction could be an argument in the future.
	txn, err := r.storage.NewTransaction(ctx)
	if err != nil {
		return nil, err
	}

	defer r.storage.Close(ctx, txn)

	// Evaluate query
	return r.eval(ctx, compiled, txn)
}

func (r *Rego) parse() (map[string]*ast.Module, ast.Body, error) {
	var errs Errors
	parsed := map[string]*ast.Module{}

	for _, module := range r.modules {
		p, err := module.Parse()
		if err != nil {
			errs = append(errs, err)
		}
		parsed[module.filename] = p
	}

	query, err := ast.ParseBody(r.query)
	if err != nil {
		errs = append(errs, err)
	}

	if len(errs) > 0 {
		return nil, nil, errs
	}

	return parsed, query, nil
}

func (r *Rego) compile(modules map[string]*ast.Module, query ast.Body) (ast.Body, error) {

	if len(modules) > 0 {
		r.compiler.Compile(modules)

		if r.compiler.Failed() {
			var errs Errors
			for _, err := range r.compiler.Errors {
				errs = append(errs, err)
			}
			return nil, errs
		}
	}

	var qctx *ast.QueryContext

	if r.pkg != "" {
		pkg, err := ast.ParsePackage(fmt.Sprintf("package %v", r.pkg))
		if err != nil {
			return nil, err
		}
		qctx = qctx.WithPackage(pkg)
	}

	if len(r.imports) > 0 {
		s := make([]string, len(r.imports))
		for i := range r.imports {
			s[i] = fmt.Sprintf("import %v", r.imports[i])
		}
		imports, err := ast.ParseImports(strings.Join(s, "\n"))
		if err != nil {
			return nil, err
		}
		qctx = qctx.WithImports(imports)
	}

	if r.rawInput != nil {
		val, err := ast.InterfaceToValue(*r.rawInput)
		if err != nil {
			return nil, err
		}
		qctx = qctx.WithInput(val)
		r.input = val
	}

	return r.compiler.QueryCompiler().WithContext(qctx).Compile(query)
}

func (r *Rego) eval(ctx context.Context, compiled ast.Body, txn storage.Transaction) (rs ResultSet, err error) {

	t := topdown.New(ctx, compiled, r.compiler, r.storage, txn)

	if r.input != nil {
		t.Input = r.input
	}

	exprs := map[*ast.Expr]struct{}{}

	err = topdown.Eval(t, func(t *topdown.Topdown) error {
		result := newResult()
		for key, value := range t.Vars() {
			val, err := topdown.ValueToInterface(value, t)
			if err != nil {
				return err
			}
			if !isTermVar(key) {
				result.Bindings[string(key)] = val
			} else if expr := findExprForTermVar(compiled, key); expr != nil {
				result.Expressions = append(result.Expressions, newExpressionValue(expr, val))
				exprs[expr] = struct{}{}
			}
		}
		for _, expr := range compiled {
			if _, ok := exprs[expr]; !ok {
				result.Expressions = append(result.Expressions, newExpressionValue(expr, true))
			}
		}
		rs = append(rs, result)
		return nil
	})

	if err != nil {
		return nil, err
	}

	if len(rs) == 0 {
		return nil, nil
	}

	return rs, nil
}

func (r *Rego) generateTermVar() *ast.Term {
	r.termVarID++
	return ast.VarTerm(ast.WildcardPrefix + fmt.Sprintf("term%v", r.termVarID))
}

func isTermVar(v ast.Var) bool {
	return strings.HasPrefix(string(v), ast.WildcardPrefix+"term")
}

func findExprForTermVar(query ast.Body, v ast.Var) *ast.Expr {
	for i := range query {
		vis := ast.NewVarVisitor()
		ast.Walk(vis, query[i])
		if vis.Vars().Contains(v) {
			return query[i]
		}
	}
	return nil
}

type rawModule struct {
	filename string
	module   string
}

func (m rawModule) Parse() (*ast.Module, error) {
	return ast.ParseModule(m.filename, m.module)
}
<|endoftext|>