<commit_before>// Copyright 2015 Keybase, Inc. All rights reserved. Use of
// this source code is governed by the included BSD license.

package s3

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"log"
	"net/url"
	"os"
	"path/filepath"
	"strings"
	"time"

	"github.com/alecthomas/template"
	"github.com/goamz/goamz/aws"
	"github.com/goamz/goamz/s3"
	"github.com/keybase/release/version"
)

type Section struct {
	Header   string
	Releases []Release
}

type Release struct {
	Name    string
	URL     string
	Version string
	Date    string
	Commit  string
}

func NewClient() (client *s3.S3, err error) {
	auth, err := aws.EnvAuth()
	if err != nil {
		return
	}
	client = s3.New(auth, aws.USEast)
	return
}

func WriteHTML(path string, bucketName string, prefixes string, suffix string) error {
	client, err := NewClient()
	if err != nil {
		return err
	}
	bucket := client.Bucket(bucketName)
	if bucket == nil {
		return fmt.Errorf("Bucket %s not found", bucketName)
	}

	var sections []Section
	for _, prefix := range strings.Split(prefixes, ",") {
		resp, err := bucket.List(prefix, "", "", 0)
		if err != nil {
			return err
		}
		keys := reverseKey(resp.Contents, 20)

		var releases []Release
		for _, k := range keys {
			if strings.HasSuffix(k.Key, suffix) {
				key := k.Key
				name := key[len(prefix):]
				urlString := fmt.Sprintf("https://s3.amazonaws.com/%s/%s%s", bucketName, prefix, url.QueryEscape(name))
				version, date, commit, err := version.Parse(name)
				if err != nil {
					log.Printf("Couldn't get version from name: %s\n", name)
				}

				// Convert to Eastern
				locationNewYork, err := time.LoadLocation("America/New_York")
				if err != nil {
					log.Printf("Couldn't load location: %s", err)
				}
				date = date.In(locationNewYork)

				releases = append(releases,
					Release{
						Name:    name,
						URL:     urlString,
						Version: version,
						Date:    date.Format("Mon Jan _2 15:04:05 MST 2006"),
						Commit:  commit,
					})
			}
		}
		sections = append(sections, Section{
			Header:   prefix,
			Releases: releases,
		})
	}

	return WriteHTMLForLinks(path, bucketName, sections)
}

var htmlTemplate = `
<!doctype html>

  <title>{{ .Title }}</title>

	{{ range $index, $sec := .Sections }}

{{ $sec.Header }}

		{{ range $index2, $rel := $sec.Releases }}

{{ $rel.Name }} {{ $rel.Version }} {{ $rel.Date }} {{ $rel.Commit }}

		{{ end }}

	{{ end }}

`

func WriteHTMLForLinks(path string, title string, sections []Section) error {
	vars := map[string]interface{}{
		"Title":    title,
		"Sections": sections,
	}

	t, err := template.New("t").Parse(htmlTemplate)
	if err != nil {
		return err
	}

	if path != "" {
		var data bytes.Buffer
		err = t.Execute(&data, vars)
		if err != nil {
			return err
		}
		err := makeParentDirs(path)
		if err != nil {
			return err
		}
		return ioutil.WriteFile(path, data.Bytes(), 0644)
	}
	return nil
}

type Link struct {
	Prefix string
	Suffix string
	Name   string
}

func CopyLatest(bucketName string) error {
	client, err := NewClient()
	if err != nil {
		return err
	}
	bucket := client.Bucket(bucketName)

	linksForPrefix := []Link{
		Link{Prefix: "darwin/", Name: "Keybase.dmg"},
		Link{Prefix: "linux_binaries/deb/", Suffix: "_amd64.deb", Name: "keybase_amd64.deb"},
		Link{Prefix: "linux_binaries/rpm/", Suffix: ".x86_64.rpm", Name: "keybase_amd64.rpm"},
	}

	for _, link := range linksForPrefix {
		resp, err := bucket.List(link.Prefix, "", "", 0)
		if err != nil {
			return err
		}
		keys := reverseKey(resp.Contents, 0)
		for _, k := range keys {
			if !strings.HasSuffix(k.Key, link.Suffix) {
				continue
			}

			url := urlString(k, bucketName, link.Prefix)
			// Instead of linking, we're making copies. S3 linking has some issues.
			// headers := map[string][]string{
			// 	"x-amz-website-redirect-location": []string{url},
			// }
			//err = bucket.PutHeader(name, []byte{}, headers, s3.PublicRead)
			log.Printf("Copying %s from %s (latest)\n", link.Name, k.Key)
			_, err = bucket.PutCopy(link.Name, s3.PublicRead, s3.CopyOptions{}, url)
			if err != nil {
				return err
			}
			break
		}
	}
	return nil
}

func urlString(k s3.Key, bucketName string, prefix string) string {
	key := k.Key
	name := key[len(prefix):]
	return fmt.Sprintf("https://s3.amazonaws.com/%s/%s%s", bucketName, prefix, url.QueryEscape(name))
}

func reverseKey(a []s3.Key, truncate int) []s3.Key {
	for left, right := 0, len(a)-1; left < right; left, right = left+1, right-1 {
		a[left], a[right] = a[right], a[left]
	}
	if truncate > 0 && len(a) > truncate {
		a = a[0:truncate]
	}
	return a
}

func makeParentDirs(filename string) error {
	dir, _ := filepath.Split(filename)
	exists, err := fileExists(dir)
	if err != nil {
		return err
	}

	if !exists {
		err = os.MkdirAll(dir, 0755)
		if err != nil {
			return err
		}
	}
	return nil
}

func fileExists(path string) (bool, error) {
	_, err := os.Stat(path)
	if err == nil {
		return true, nil
	}
	if os.IsNotExist(err) {
		return false, nil
	}
	return false, err
}
<commit_msg>Fixing release ordering<commit_after>// Copyright 2015 Keybase, Inc. All rights reserved. Use of
// this source code is governed by the included BSD license.

package s3

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"log"
	"net/url"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"time"

	"github.com/alecthomas/template"
	"github.com/goamz/goamz/aws"
	"github.com/goamz/goamz/s3"
	"github.com/keybase/release/version"
)

type Section struct {
	Header   string
	Releases []Release
}

type Release struct {
	Name       string
	Key        s3.Key
	URL        string
	Version    string
	DateString string
	Date       time.Time
	Commit     string
}

type ByRelease []Release

func (s ByRelease) Len() int {
	return len(s)
}

func (s ByRelease) Swap(i, j int) {
	s[i], s[j] = s[j], s[i]
}

func (s ByRelease) Less(i, j int) bool {
	// Reverse date order
	return s[j].Date.Before(s[i].Date)
}

func NewClient() (client *s3.S3, err error) {
	auth, err := aws.EnvAuth()
	if err != nil {
		return
	}
	client = s3.New(auth, aws.USEast)
	return
}

func loadReleases(keys []s3.Key, bucketName string, prefix string, suffix string, truncate int) []Release {
	var releases []Release
	for _, k := range keys {
		if strings.HasSuffix(k.Key, suffix) {
			key := k.Key
			name := key[len(prefix):]
			urlString := fmt.Sprintf("https://s3.amazonaws.com/%s/%s%s", bucketName, prefix, url.QueryEscape(name))
			version, date, commit, err := version.Parse(name)
			if err != nil {
				log.Printf("Couldn't get version from name: %s\n", name)
			}

			// Convert to Eastern
			locationNewYork, err := time.LoadLocation("America/New_York")
			if err != nil {
				log.Printf("Couldn't load location: %s", err)
			}
			date = date.In(locationNewYork)

			releases = append(releases,
				Release{
					Name:       name,
					Key:        k,
					URL:        urlString,
					Version:    version,
					Date:       date,
					DateString: date.Format("Mon Jan _2 15:04:05 MST 2006"),
					Commit:     commit,
				})
		}
	}
	// TODO: Should also sanity check that version sort is same as time sort
	// otherwise something got messed up
	sort.Sort(ByRelease(releases))
	if truncate > 0 && len(releases) > truncate {
		releases = releases[0:truncate]
	}
	return releases
}

func WriteHTML(path string, bucketName string, prefixes string, suffix string) error {
	client, err := NewClient()
	if err != nil {
		return err
	}
	bucket := client.Bucket(bucketName)
	if bucket == nil {
		return fmt.Errorf("Bucket %s not found", bucketName)
	}

	var sections []Section
	for _, prefix := range strings.Split(prefixes, ",") {
		resp, err := bucket.List(prefix, "", "", 0)
		if err != nil {
			return err
		}

		releases := loadReleases(resp.Contents, bucketName, prefix, suffix, 20)
		if len(releases) > 0 {
			log.Printf("Found %d release(s) at %s\n", len(releases), prefix)
			for _, release := range releases {
				log.Printf(" %s %s %s\n", release.Name, release.Version, release.DateString)
			}
		}
		sections = append(sections, Section{
			Header:   prefix,
			Releases: releases,
		})
	}

	return WriteHTMLForLinks(path, bucketName, sections)
}

var htmlTemplate = `
<!doctype html>

  <title>{{ .Title }}</title>

	{{ range $index, $sec := .Sections }}

{{ $sec.Header }}

		{{ range $index2, $rel := $sec.Releases }}

{{ $rel.Name }} {{ $rel.Version }} {{ $rel.Date }} {{ $rel.Commit }}

		{{ end }}

	{{ end }}

`

func WriteHTMLForLinks(path string, title string, sections []Section) error {
	vars := map[string]interface{}{
		"Title":    title,
		"Sections": sections,
	}

	t, err := template.New("t").Parse(htmlTemplate)
	if err != nil {
		return err
	}

	if path != "" {
		var data bytes.Buffer
		err = t.Execute(&data, vars)
		if err != nil {
			return err
		}
		err := makeParentDirs(path)
		if err != nil {
			return err
		}
		return ioutil.WriteFile(path, data.Bytes(), 0644)
	}
	return nil
}

type Link struct {
	Prefix string
	Suffix string
	Name   string
}

func CopyLatest(bucketName string) error {
	client, err := NewClient()
	if err != nil {
		return err
	}
	bucket := client.Bucket(bucketName)

	linksForPrefix := []Link{
		Link{Prefix: "darwin/", Name: "Keybase.dmg"},
		Link{Prefix: "linux_binaries/deb/", Suffix: "_amd64.deb", Name: "keybase_amd64.deb"},
		Link{Prefix: "linux_binaries/rpm/", Suffix: ".x86_64.rpm", Name: "keybase_amd64.rpm"},
	}

	for _, link := range linksForPrefix {
		resp, err := bucket.List(link.Prefix, "", "", 0)
		if err != nil {
			return err
		}
		releases := loadReleases(resp.Contents, bucketName, link.Prefix, link.Suffix, 0)
		for _, release := range releases {
			k := release.Key
			if !strings.HasSuffix(k.Key, link.Suffix) {
				continue
			}

			url := urlString(k, bucketName, link.Prefix)
			// Instead of linking, we're making copies. S3 linking has some issues.
			// headers := map[string][]string{
			// 	"x-amz-website-redirect-location": []string{url},
			// }
			//err = bucket.PutHeader(name, []byte{}, headers, s3.PublicRead)
			log.Printf("Copying %s from %s (latest)\n", link.Name, k.Key)
			_, err = bucket.PutCopy(link.Name, s3.PublicRead, s3.CopyOptions{}, url)
			if err != nil {
				return err
			}
			break
		}
	}
	return nil
}

func urlString(k s3.Key, bucketName string, prefix string) string {
	key := k.Key
	name := key[len(prefix):]
	return fmt.Sprintf("https://s3.amazonaws.com/%s/%s%s", bucketName, prefix, url.QueryEscape(name))
}

func makeParentDirs(filename string) error {
	dir, _ := filepath.Split(filename)
	exists, err := fileExists(dir)
	if err != nil {
		return err
	}

	if !exists {
		err = os.MkdirAll(dir, 0755)
		if err != nil {
			return err
		}
	}
	return nil
}

func fileExists(path string) (bool, error) {
	_, err := os.Stat(path)
	if err == nil {
		return true, nil
	}
	if os.IsNotExist(err) {
		return false, nil
	}
	return false, err
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"errors"
	"flag"
	"fmt"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/ec2"
	"github.com/olekukonko/tablewriter"
	"log"
	"net"
	"os"
	"os/exec"
	"regexp"
	"strconv"
	"strings"
)

var config Config

func main() {

	config.transparant = *flag.Bool("transparant", true, "Forward lookup failures")
	flag.Parse()
	if flag.NArg() < 1 {
		flag.PrintDefaults()
		os.Exit(1)
	}
	config.hostname = flag.Args()[0]

	var ip string
	if hostIp, err := parseInput(config.hostname); err == nil {
		ip = hostIp
	} else {
		ip = queryEC2(config.hostname)
	}

	fmt.Println("Connecting to", ip, "...\n")

	cmd := exec.Command("ssh", ip)
	cmd.Stdout = os.Stdout
	cmd.Stdin = os.Stdin
	cmd.Stderr = os.Stderr
	cmd.Run()
}

func parseInput(hostname string) (string, error) {
	if net.ParseIP(hostname) != nil {
		return hostname, nil
	}
	var err error
	r, _ := regexp.Compile("ip-([0-9]{1,3})-([0-9]{1,3})-([0-9]{1,3})-([0-9]{1,3})")
	m := r.FindAllStringSubmatch(hostname, 4)

	if len(m) != 1 {
		return "", errors.New("")
	}

	return fmt.Sprintf("%s.%s.%s.%s", m[0][1], m[0][2], m[0][3], m[0][4]), err
}

func queryEC2(filter string) string {
	awsRegion := "us-east-1"

	instances := findInstances(filter, awsRegion)
	if len(instances) == 1 {
		return instances[0].ip
	} else if len(instances) == 0 {
		if config.transparant {
			return filter
		}
		fmt.Println("Could not match any instances")
		os.Exit(1)
	}
	fmt.Printf("listing instances with tag %v in: %v\n", filter, awsRegion)
	printTable(instances)

	reader := bufio.NewReader(os.Stdin)
	fmt.Print("Pick a number: ")
	choice, _ := reader.ReadString('\n')
	choiceInt, err := strconv.Atoi(choice[:len(choice)-1])
	if err != nil {
		choiceInt = 0
	}

	return instances[choiceInt].ip
}

func printTable(instances []Instance) {
	table := tablewriter.NewWriter(os.Stdout)
	table.SetHeader([]string{"", "Name", "Ip"})

	for index, instance := range instances {
		table.Append([]string{strconv.Itoa(index), instance.name, instance.ip})
	}

	table.Render()
}

func findInstances(filter string, awsRegion string) []Instance {
	sess := session.Must(session.NewSession())

	svc := ec2.New(sess, &aws.Config{Region: aws.String(awsRegion)})
	params := &ec2.DescribeInstancesInput{
		Filters: []*ec2.Filter{
			{
				Name: aws.String("tag:Name"),
				Values: []*string{
					aws.String(strings.Join([]string{"*", filter, "*"}, "")),
				},
			},
			{
				Name:   aws.String("instance-state-name"),
				Values: []*string{aws.String("running"), aws.String("pending")},
			},
		},
	}
	resp, err := svc.DescribeInstances(params)
	if err != nil {
		fmt.Println("there was an error listing instances in", awsRegion, err.Error())
		log.Fatal(err.Error())
	}

	var instances []Instance
	for _, instance := range resp.Reservations {
		i := instance.Instances[0]
		instances = append(instances, Instance{name: findTag("Name", i.Tags), ip: *i.PrivateIpAddress})
	}
	return instances
}

func findTag(key string, tags []*ec2.Tag) string {

	for _, item := range tags {
		if *item.Key == key {
			return *item.Value
		}
	}

	return ""
}

type Instance struct {
	name string
	ip   string
}

type Target struct {
	ip   string
	user string
}

type Config struct {
	transparant bool
	hostname    string
}
<commit_msg>Add sorting<commit_after>package main

import (
	"bufio"
	"errors"
	"flag"
	"fmt"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/ec2"
	"github.com/olekukonko/tablewriter"
	"log"
	"net"
	"os"
	"os/exec"
	"regexp"
	"sort"
	"strconv"
	"strings"
)

var config Config

func main() {

	config.transparant = *flag.Bool("transparant", true, "Forward lookup failures")
	flag.Parse()
	if flag.NArg() < 1 {
		flag.PrintDefaults()
		os.Exit(1)
	}
	config.hostname = flag.Args()[0]

	var ip string
	if hostIp, err := parseInput(config.hostname); err == nil {
		ip = hostIp
	} else {
		ip = queryEC2(config.hostname)
	}

	fmt.Println("Connecting to", ip, "...\n")

	cmd := exec.Command("ssh", ip)
	cmd.Stdout = os.Stdout
	cmd.Stdin = os.Stdin
	cmd.Stderr = os.Stderr
	cmd.Run()
}

func parseInput(hostname string) (string, error) {
	if net.ParseIP(hostname) != nil {
		return hostname, nil
	}
	var err error
	r, _ := regexp.Compile("ip-([0-9]{1,3})-([0-9]{1,3})-([0-9]{1,3})-([0-9]{1,3})")
	m := r.FindAllStringSubmatch(hostname, 4)

	if len(m) != 1 {
		return "", errors.New("")
	}

	return fmt.Sprintf("%s.%s.%s.%s", m[0][1], m[0][2], m[0][3], m[0][4]), err
}

func queryEC2(filter string) string {
	awsRegion := "us-east-1"

	instances := findInstances(filter, awsRegion)
	if len(instances) == 1 {
		return instances[0].ip
	} else if len(instances) == 0 {
		if config.transparant {
			return filter
		}
		fmt.Println("Could not match any instances")
		os.Exit(1)
	}
	fmt.Printf("listing instances with tag %v in: %v\n", filter, awsRegion)
	printTable(instances)

	reader := bufio.NewReader(os.Stdin)
	fmt.Print("Pick a number: ")
	choice, _ := reader.ReadString('\n')
	choiceInt, err := strconv.Atoi(choice[:len(choice)-1])
	if err != nil {
		choiceInt = 0
	}

	return instances[choiceInt].ip
}

func printTable(instances []Instance) {
	table := tablewriter.NewWriter(os.Stdout)
	table.SetHeader([]string{"", "Name", "Ip"})

	for index, instance := range instances {
		table.Append([]string{strconv.Itoa(index), instance.name, instance.ip})
	}

	table.Render()
}

func findInstances(filter string, awsRegion string) []Instance {
	sess := session.Must(session.NewSession())

	svc := ec2.New(sess, &aws.Config{Region: aws.String(awsRegion)})
	params := &ec2.DescribeInstancesInput{
		Filters: []*ec2.Filter{
			{
				Name: aws.String("tag:Name"),
				Values: []*string{
					aws.String(strings.Join([]string{"*", filter, "*"}, "")),
				},
			},
			{
				Name:   aws.String("instance-state-name"),
				Values: []*string{aws.String("running"), aws.String("pending")},
			},
		},
	}
	resp, err := svc.DescribeInstances(params)
	if err != nil {
		fmt.Println("there was an error listing instances in", awsRegion, err.Error())
		log.Fatal(err.Error())
	}

	var instances []Instance
	for _, instance := range resp.Reservations {
		i := instance.Instances[0]
		instances = append(instances, Instance{name: findTag("Name", i.Tags), ip: *i.PrivateIpAddress})
	}

	// Sort by name
	sort.Slice(instances, func(i, j int) bool {
		if instances[i].name == instances[j].name {
			return instances[i].ip < instances[j].ip
		}
		return instances[i].name < instances[j].name
	})

	return instances
}

func findTag(key string, tags []*ec2.Tag) string {

	for _, item := range tags {
		if *item.Key == key {
			return *item.Value
		}
	}

	return ""
}

type Instance struct {
	name string
	ip   string
}

type Target struct {
	ip   string
	user string
}

type Config struct {
	transparant bool
	hostname    string
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"net/url"
	"sync/atomic"
	"time"

	"github.com/hagen1778/chproxy/config"
	"github.com/hagen1778/chproxy/log"
)

func (s *scope) String() string {
	return fmt.Sprintf("[ Id: %d; User %q(%d) proxying as %q(%d) to %q(%d) ]",
		s.id,
		s.user.name, s.user.runningQueries(),
		s.clusterUser.name, s.clusterUser.runningQueries(),
		s.host.addr.Host, s.host.runningQueries())
}

type scope struct {
	id          int64
	host        *host
	cluster     *cluster
	user        *user
	clusterUser *clusterUser
}

var scopeId = time.Now().UnixNano()

func newScope(u *user, cu *clusterUser, c *cluster) *scope {
	return &scope{
		id:          atomic.AddInt64(&scopeId, 1),
		host:        c.getHost(),
		cluster:     c,
		user:        u,
		clusterUser: cu,
	}
}

func (s *scope) inc() error {
	s.user.inc()
	s.clusterUser.inc()
	s.host.inc()

	var err error
	if s.user.maxConcurrentQueries > 0 && s.user.runningQueries() > s.user.maxConcurrentQueries {
		err = fmt.Errorf("limits for user %q are exceeded: maxConcurrentQueries limit: %d", s.user.name, s.user.maxConcurrentQueries)
	}

	if s.clusterUser.maxConcurrentQueries > 0 && s.clusterUser.runningQueries() > s.clusterUser.maxConcurrentQueries {
		err = fmt.Errorf("limits for cluster user %q are exceeded: maxConcurrentQueries limit: %d", s.clusterUser.name, s.clusterUser.maxConcurrentQueries)
	}

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

	return nil
}

func (s *scope) dec() {
	s.host.dec()
	s.user.dec()
	s.clusterUser.dec()
}

type user struct {
	toUser          string
	toCluster       string
	allowedNetworks config.Networks

	name, password       string
	maxExecutionTime     time.Duration
	maxConcurrentQueries uint32

	queryCounter
}

type clusterUser struct {
	name, password       string
	maxExecutionTime     time.Duration
	maxConcurrentQueries uint32

	queryCounter
}

type host struct {
	addr *url.URL

	queryCounter
}

type cluster struct {
	nextIdx uint32
	hosts   []*host
	users   map[string]*clusterUser
}

func newCluster(h []*host, cu map[string]*clusterUser) *cluster {
	return &cluster{
		hosts:   h,
		users:   cu,
		nextIdx: uint32(time.Now().UnixNano()),
	}
}

// We don't use query_id because of distributed processing, the query ID is not passed to remote servers
func (c *cluster) killQueries(ua string, elapsed float64) {
	q := fmt.Sprintf("KILL QUERY WHERE http_user_agent = '%s' AND elapsed >= %d", ua, int(elapsed))
	log.Debugf("ExecutionTime exceeded. Going to call query %q for hosts %v", q, c.hosts)
	for _, host := range c.hosts {
		if err := doQuery(q, host.addr.String()); err != nil {
			log.Errorf("error while killing queries older than %.2fs: %s", elapsed, err)
		}
	}
}

// get least loaded + round-robin host from cluster
func (c *cluster) getHost() *host {
	idx := atomic.AddUint32(&c.nextIdx, 1)

	l := uint32(len(c.hosts))
	idx = idx % l
	idle := c.hosts[idx]
	idleN := idle.runningQueries()

	if idleN == 0 {
		return idle
	}

	// round hosts checking
	// until the least loaded is found
	for i := (idx + 1) % l; i != idx; i = (i + 1) % l {
		h := c.hosts[i]
		n := h.runningQueries()
		if n == 0 {
			return h
		}
		if n < idleN {
			idle, idleN = h, n
		}
	}

	return idle
}

type queryCounter struct {
	value uint32
}

func (qc *queryCounter) runningQueries() uint32 {
	return atomic.LoadUint32(&qc.value)
}

func (qc *queryCounter) inc() {
	atomic.AddUint32(&qc.value, 1)
}

func (qc *queryCounter) dec() {
	atomic.AddUint32(&qc.value, ^uint32(0))
}
<commit_msg>cache results of inc() function to avoid possible race condition between increasing counter and checking it's state<commit_after>package main

import (
	"fmt"
	"net/url"
	"sync/atomic"
	"time"

	"github.com/hagen1778/chproxy/config"
	"github.com/hagen1778/chproxy/log"
)

func (s *scope) String() string {
	return fmt.Sprintf("[ Id: %d; User %q(%d) proxying as %q(%d) to %q(%d) ]",
		s.id,
		s.user.name, s.user.runningQueries(),
		s.clusterUser.name, s.clusterUser.runningQueries(),
		s.host.addr.Host, s.host.runningQueries())
}

type scope struct {
	id          int64
	host        *host
	cluster     *cluster
	user        *user
	clusterUser *clusterUser
}

var scopeId = time.Now().UnixNano()

func newScope(u *user, cu *clusterUser, c *cluster) *scope {
	return &scope{
		id:          atomic.AddInt64(&scopeId, 1),
		host:        c.getHost(),
		cluster:     c,
		user:        u,
		clusterUser: cu,
	}
}

func (s *scope) inc() error {
	uq := s.user.inc()
	cq := s.clusterUser.inc()
	s.host.inc()

	var err error
	if s.user.maxConcurrentQueries > 0 && uq > s.user.maxConcurrentQueries {
		err = fmt.Errorf("limits for user %q are exceeded: maxConcurrentQueries limit: %d", s.user.name, s.user.maxConcurrentQueries)
	}

	if s.clusterUser.maxConcurrentQueries > 0 && cq > s.clusterUser.maxConcurrentQueries {
		err = fmt.Errorf("limits for cluster user %q are exceeded: maxConcurrentQueries limit: %d", s.clusterUser.name, s.clusterUser.maxConcurrentQueries)
	}

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

	return nil
}

func (s *scope) dec() {
	s.host.dec()
	s.user.dec()
	s.clusterUser.dec()
}

type user struct {
	toUser          string
	toCluster       string
	allowedNetworks config.Networks

	name, password       string
	maxExecutionTime     time.Duration
	maxConcurrentQueries uint32

	queryCounter
}

type clusterUser struct {
	name, password       string
	maxExecutionTime     time.Duration
	maxConcurrentQueries uint32

	queryCounter
}

type host struct {
	addr *url.URL

	queryCounter
}

type cluster struct {
	nextIdx uint32
	hosts   []*host
	users   map[string]*clusterUser
}

func newCluster(h []*host, cu map[string]*clusterUser) *cluster {
	return &cluster{
		hosts:   h,
		users:   cu,
		nextIdx: uint32(time.Now().UnixNano()),
	}
}

// We don't use query_id because of distributed processing, the query ID is not passed to remote servers
func (c *cluster) killQueries(ua string, elapsed float64) {
	q := fmt.Sprintf("KILL QUERY WHERE http_user_agent = '%s' AND elapsed >= %d", ua, int(elapsed))
	log.Debugf("ExecutionTime exceeded. Going to call query %q for hosts %v", q, c.hosts)
	for _, host := range c.hosts {
		if err := doQuery(q, host.addr.String()); err != nil {
			log.Errorf("error while killing queries older than %.2fs: %s", elapsed, err)
		}
	}
}

// get least loaded + round-robin host from cluster
func (c *cluster) getHost() *host {
	idx := atomic.AddUint32(&c.nextIdx, 1)

	l := uint32(len(c.hosts))
	idx = idx % l
	idle := c.hosts[idx]
	idleN := idle.runningQueries()

	if idleN == 0 {
		return idle
	}

	// round hosts checking
	// until the least loaded is found
	for i := (idx + 1) % l; i != idx; i = (i + 1) % l {
		h := c.hosts[i]
		n := h.runningQueries()
		if n == 0 {
			return h
		}
		if n < idleN {
			idle, idleN = h, n
		}
	}

	return idle
}

type queryCounter struct {
	value uint32
}

func (qc *queryCounter) runningQueries() uint32 {
	return atomic.LoadUint32(&qc.value)
}

func (qc *queryCounter) inc() uint32 {
	return atomic.AddUint32(&qc.value, 1)
}

func (qc *queryCounter) dec() {
	atomic.AddUint32(&qc.value, ^uint32(0))
}
<|endoftext|>


<commit_before>package gorm

import (
	"fmt"
	"strings"
	"time"
)

func BeforeCreate(scope *Scope) {
	scope.CallMethod("BeforeSave")
	scope.CallMethod("BeforeCreate")
}

func UpdateTimeStampWhenCreate(scope *Scope) {
	if !scope.HasError() {
		scope.SetColumn("CreatedAt", time.Now())
		scope.SetColumn("UpdatedAt", time.Now())
	}
}

func Create(scope *Scope) {
	defer scope.Trace(time.Now())

	if !scope.HasError() {
		// set create sql
		var sqls, columns []string

		for _, field := range scope.Fields() {
			if field.DBName != scope.PrimaryKey() && len(field.SqlTag) > 0 && !field.IsIgnored {
				columns = append(columns, scope.Quote(field.DBName))
				sqls = append(sqls, scope.AddToVars(field.Value))
			}
		}

		scope.Raw(fmt.Sprintf(
			"INSERT INTO %v (%v) VALUES (%v) %v",
			scope.TableName(),
			strings.Join(columns, ","),
			strings.Join(sqls, ","),
			scope.Dialect().ReturningStr(scope.PrimaryKey()),
		))

		// execute create sql
		var id interface{}
		if scope.Dialect().SupportLastInsertId() {
			if sql_result, err := scope.DB().Exec(scope.Sql, scope.SqlVars...); scope.Err(err) == nil {
				id, err = sql_result.LastInsertId()
				scope.Err(err)
			}
		} else {
			scope.Err(scope.DB().QueryRow(scope.Sql, scope.SqlVars...).Scan(&id))
		}

		if !scope.HasError() {
			scope.SetColumn(scope.PrimaryKey(), id)
		}
	}
}

func AfterCreate(scope *Scope) {
	scope.CallMethod("AfterCreate")
	scope.CallMethod("AfterSave")
}

func init() {
	DefaultCallback.Create().Register("gorm:begin_transaction", BeginTransaction)
	DefaultCallback.Create().Register("gorm:before_create", BeforeCreate)
	DefaultCallback.Create().Register("gorm:save_before_associations", SaveBeforeAssociations)
	DefaultCallback.Create().Register("gorm:update_time_stamp_when_create", UpdateTimeStampWhenCreate)
	DefaultCallback.Create().Register("gorm:create", Create)
	DefaultCallback.Create().Register("gorm:save_after_associations", SaveAfterAssociations)
	DefaultCallback.Create().Register("gorm:after_create", AfterCreate)
	DefaultCallback.Create().Register("gorm:commit_or_rollback_transaction", CommitOrRollbackTransaction)
}
<commit_msg>Fix insert empty record into database<commit_after>package gorm

import (
	"fmt"
	"strings"
	"time"
)

func BeforeCreate(scope *Scope) {
	scope.CallMethod("BeforeSave")
	scope.CallMethod("BeforeCreate")
}

func UpdateTimeStampWhenCreate(scope *Scope) {
	if !scope.HasError() {
		scope.SetColumn("CreatedAt", time.Now())
		scope.SetColumn("UpdatedAt", time.Now())
	}
}

func Create(scope *Scope) {
	defer scope.Trace(time.Now())

	if !scope.HasError() {
		// set create sql
		var sqls, columns []string

		for _, field := range scope.Fields() {
			if field.DBName != scope.PrimaryKey() && len(field.SqlTag) > 0 && !field.IsIgnored {
				columns = append(columns, scope.Quote(field.DBName))
				sqls = append(sqls, scope.AddToVars(field.Value))
			}
		}

		if len(columns) == 0 {
			scope.Raw(fmt.Sprintf("INSERT INTO %v DEFAULT VALUES %v",
				scope.TableName(),
				scope.Dialect().ReturningStr(scope.PrimaryKey()),
			))
		} else {
			scope.Raw(fmt.Sprintf(
				"INSERT INTO %v (%v) VALUES (%v) %v",
				scope.TableName(),
				strings.Join(columns, ","),
				strings.Join(sqls, ","),
				scope.Dialect().ReturningStr(scope.PrimaryKey()),
			))
		}

		// execute create sql
		var id interface{}
		if scope.Dialect().SupportLastInsertId() {
			if sql_result, err := scope.DB().Exec(scope.Sql, scope.SqlVars...); scope.Err(err) == nil {
				id, err = sql_result.LastInsertId()
				scope.Err(err)
			}
		} else {
			scope.Err(scope.DB().QueryRow(scope.Sql, scope.SqlVars...).Scan(&id))
		}

		if !scope.HasError() {
			scope.SetColumn(scope.PrimaryKey(), id)
		}
	}
}

func AfterCreate(scope *Scope) {
	scope.CallMethod("AfterCreate")
	scope.CallMethod("AfterSave")
}

func init() {
	DefaultCallback.Create().Register("gorm:begin_transaction", BeginTransaction)
	DefaultCallback.Create().Register("gorm:before_create", BeforeCreate)
	DefaultCallback.Create().Register("gorm:save_before_associations", SaveBeforeAssociations)
	DefaultCallback.Create().Register("gorm:update_time_stamp_when_create", UpdateTimeStampWhenCreate)
	DefaultCallback.Create().Register("gorm:create", Create)
	DefaultCallback.Create().Register("gorm:save_after_associations", SaveAfterAssociations)
	DefaultCallback.Create().Register("gorm:after_create", AfterCreate)
	DefaultCallback.Create().Register("gorm:commit_or_rollback_transaction", CommitOrRollbackTransaction)
}
<|endoftext|>


<commit_before>package db

import (
	"container/list"
	"sync"
)

// A FIFO queue for table cache.
type TableQueue struct {
	list     *list.List
	capacity int
}

// Add a new table Id into queue. Return the oldest one in the queue, or
// -1 if there are less than @capacity elements in the queue
func (tq *TableQueue) Add(id int64) (int64, *list.Element) {
	element := &list.Element{Value: id}
	tq.list.PushFront(element)
	if tq.list.Len() < tq.capacity {
		return -1, element
	} else {
		tmp := tq.list.Back()
		ret := tmp.Value.(int64)
		tq.list.Remove(tmp)
		return ret, element
	}
}

// Move the accessed element to the front of the queue
func (tq *TableQueue) Access(element *list.Element) {
	tq.list.MoveToFront(element)
}

type tblInfo struct {
	table   *Table
	element *list.Element
	done    chan bool
}

// Manage cache of tables.
type TableCache struct {
	mutex    sync.Mutex
	tableMap map[int64]tblInfo
	queue    TableQueue
}

func MakeTableCache(capacity int) *TableCache {
	return &TableCache{
		tableMap: map[int64]tblInfo{},
		queue:    TableQueue{capacity: capacity},
	}
}

// Add a table into cache.
func (tc *TableCache) Add(table *Table, id int64) {
	tc.mutex.Lock()
	defer tc.mutex.Unlock()

	_, found := tc.tableMap[id]
	if found {
		panic("New table should not be in the cache!")
	}

	oldId, element := tc.queue.Add(id)
	tc.tableMap[id] = tblInfo{table: table, element: element, done: make(chan bool)}
	if oldId >= 0 {
		delete(tc.tableMap, oldId)
	}
}

func (tc *TableCache) Get(id int64) *Table {
	tc.mutex.Lock()

	// Find if the table is already in cache.
	info, found := tc.tableMap[id]
	if found && info.table != nil {
		tc.queue.Access(info.element)
		tc.mutex.Unlock()
		return info.table
	}

	// Someone else is trying to load the table, so let's wait.
	if found && info.table == nil {
		tc.mutex.Unlock()
		<-info.done
		return tc.Get(id)
	}

	// Cache item cannot be found, put a place holder so that
	// other go routine that uses it can wait.
	tc.tableMap[id] = tblInfo{done: make(chan bool)}
	tc.mutex.Unlock()

	return nil
}

type DbImpl struct {
	path     string
	env      Env
	comp     Comparator
	skipList *Skiplist
	tmpList  *Skiplist
	mutex    sync.RWMutex
}

// Freeze current skiplist, push it down to tmpList, create a new skiplist
func (db *DbImpl) RotateSkiplist() (*Skiplist, *Skiplist) {
	db.mutex.Lock()
	defer db.mutex.Unlock()

	if db.tmpList != nil {
		panic("tmpList is not empty during rotation!")
	}

	db.tmpList = db.skipList
	db.skipList = MakeSkiplist()

	return db.skipList, db.tmpList
}

// Get parent dir of this db.
func (db *DbImpl) GetPath() string {
	return db.path
}

func (db *DbImpl) GetEnv() Env {
	return db.env
}

func (db *DbImpl) GetComparator() Comparator {
	return db.comp
}
<commit_msg>Implement table cache<commit_after>package db

import (
	"container/list"
	"sync"
)

// A FIFO queue for table cache.
type TableQueue struct {
	list     *list.List
	capacity int
}

// Add a new table Id into queue. Return the oldest one in the queue, or
// -1 if there are less than @capacity elements in the queue
func (tq *TableQueue) Add(id int64) (int64, *list.Element) {
	element := &list.Element{Value: id}
	tq.list.PushFront(element)
	if tq.list.Len() < tq.capacity {
		return -1, element
	} else {
		tmp := tq.list.Back()
		ret := tmp.Value.(int64)
		tq.list.Remove(tmp)
		return ret, element
	}
}

// Move the accessed element to the front of the queue
func (tq *TableQueue) Access(element *list.Element) {
	tq.list.MoveToFront(element)
}

type tblInfo struct {
	table   *Table
	element *list.Element
	done    chan bool
}

// Manage cache of tables.
type TableCache struct {
	mutex    sync.Mutex
	tableMap map[int64]tblInfo
	queue    TableQueue
	impl     *DbImpl
}

func MakeTableCache(impl *DbImpl, capacity int) *TableCache {
	return &TableCache{
		tableMap: map[int64]tblInfo{},
		queue:    TableQueue{capacity: capacity},
		impl:     impl,
	}
}

// Add a table into cache.
func (tc *TableCache) Add(table *Table, id int64) {
	tc.mutex.Lock()
	defer tc.mutex.Unlock()

	_, found := tc.tableMap[id]
	if found {
		panic("New table should not be in the cache!")
	}

	oldId, element := tc.queue.Add(id)
	tc.tableMap[id] = tblInfo{table: table, element: element, done: make(chan bool)}
	if oldId >= 0 {
		delete(tc.tableMap, oldId)
	}
}

func (tc *TableCache) Get(id int64) *Table {
	tc.mutex.Lock()

	// Find if the table is already in cache.
	info, found := tc.tableMap[id]
	if found && info.table != nil {
		tc.queue.Access(info.element)
		tc.mutex.Unlock()
		return info.table
	}

	// Someone else is trying to load the table, so let's wait.
	if found && info.table == nil {
		tc.mutex.Unlock()
		<-info.done
		return tc.Get(id)
	}

	// Cache item cannot be found, put a place holder so that
	// other go routine that uses it can wait.
	tc.tableMap[id] = tblInfo{done: make(chan bool)}
	tc.mutex.Unlock()

	finfo, found := tc.impl.GetManifest().GetFileInfo(id)
	if !found {
		panic("Expected id does not found")
	}

	file, status := tc.impl.GetEnv().NewSequentialFile(finfo.Location)
	if !status.Ok() {
		panic("File does not exist!")
	}

	var fsize uint64
	fsize, status = tc.impl.GetEnv().GetFileSize(finfo.Location)
	if !status.Ok() {
		panic("Cannot stat a file!")
	}

	buf := make([]byte, fsize)
	tbl := RecoverTable(file, buf, tc.impl.GetComparator())

	if tbl == nil {
		panic("Fails to recover a table!")
	}

	// Add the table into cache.
	tc.mutex.Lock()
	defer tc.mutex.Unlock()

	info, found = tc.tableMap[id]
	if !found {
		panic("Fails to find previously reserved entry!")
	}

	oldId, element := tc.queue.Add(id)

	info.table = tbl
	info.element = element

	tc.tableMap[id] = info

	// Notify all pending go routines that this table is loaded.
	close(info.done)

	if oldId >= 0 {
		delete(tc.tableMap, oldId)
	}

	return tbl
}

type DbImpl struct {
	path     string
	env      Env
	comp     Comparator
	skipList *Skiplist
	tmpList  *Skiplist
	manifest *Manifest
	mutex    sync.RWMutex
}

// Freeze current skiplist, push it down to tmpList, create a new skiplist
func (db *DbImpl) RotateSkiplist() (*Skiplist, *Skiplist) {
	db.mutex.Lock()
	defer db.mutex.Unlock()

	if db.tmpList != nil {
		panic("tmpList is not empty during rotation!")
	}

	db.tmpList = db.skipList
	db.skipList = MakeSkiplist()

	return db.skipList, db.tmpList
}

// Get parent dir of this db.
func (db *DbImpl) GetPath() string {
	return db.path
}

func (db *DbImpl) GetEnv() Env {
	return db.env
}

func (db *DbImpl) GetComparator() Comparator {
	return db.comp
}

func (db *DbImpl) GetManifest() *Manifest {
	return db.manifest
}
<|endoftext|>


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

package fakestorage

import (
	"io/ioutil"
	"reflect"
	"testing"

	"cloud.google.com/go/storage"
	"golang.org/x/net/context"
	"google.golang.org/api/iterator"
)

func TestServerClientObjectAttrs(t *testing.T) {
	const bucketName = "some-bucket"
	const objectName = "img/hi-res/party-01.jpg"
	const content = "some nice content"
	server := NewServer([]Object{
		{BucketName: bucketName, Name: objectName, Content: []byte(content)},
	})
	defer server.Stop()
	client := server.Client()
	objHandle := client.Bucket(bucketName).Object(objectName)
	attrs, err := objHandle.Attrs(context.Background())
	if err != nil {
		t.Fatal(err)
	}
	if attrs.Bucket != bucketName {
		t.Errorf("wrong bucket name\nwant %q\ngot  %q", bucketName, attrs.Bucket)
	}
	if attrs.Name != objectName {
		t.Errorf("wrong object name\nwant %q\ngot  %q", objectName, attrs.Name)
	}
	if attrs.Size != int64(len(content)) {
		t.Errorf("wrong size returned\nwant %d\ngot  %d", len(content), attrs.Size)
	}
}

func TestServerClientObjectAttrsAfterCreateObject(t *testing.T) {
	const bucketName = "prod-bucket"
	const objectName = "video/hi-res/best_video_1080p.mp4"
	server := NewServer(nil)
	defer server.Stop()
	server.CreateObject(Object{BucketName: bucketName, Name: objectName})
	client := server.Client()
	objHandle := client.Bucket(bucketName).Object(objectName)
	attrs, err := objHandle.Attrs(context.Background())
	if err != nil {
		t.Fatal(err)
	}
	if attrs.Bucket != bucketName {
		t.Errorf("wrong bucket name\nwant %q\ngot  %q", bucketName, attrs.Bucket)
	}
	if attrs.Name != objectName {
		t.Errorf("wrong object name\n want %q\ngot  %q", objectName, attrs.Name)
	}
}

func TestServerClientObjectAttrsErrors(t *testing.T) {
	server := NewServer([]Object{
		{BucketName: "some-bucket", Name: "img/hi-res/party-01.jpg"},
	})
	defer server.Stop()
	var tests = []struct {
		testCase   string
		bucketName string
		objectName string
	}{
		{
			"bucket not found",
			"other-bucket",
			"whatever-object",
		},
		{
			"object not found",
			"some-bucket",
			"img/low-res/party-01.jpg",
		},
	}
	for _, test := range tests {
		t.Run(test.testCase, func(t *testing.T) {
			objHandle := server.Client().Bucket(test.bucketName).Object(test.objectName)
			attrs, err := objHandle.Attrs(context.Background())
			if err == nil {
				t.Error("unexpected <nil> error")
			}
			if attrs != nil {
				t.Errorf("unexpected non-nil attrs: %#v", attrs)
			}
		})
	}
}

func TestServerClientObjectReader(t *testing.T) {
	const bucketName = "some-bucket"
	const objectName = "items/data.txt"
	const content = "some nice content"
	server := NewServer([]Object{
		{
			BucketName: bucketName,
			Name:       objectName,
			Content:    []byte(content),
		},
	})
	defer server.Stop()
	client := server.Client()
	objHandle := client.Bucket(bucketName).Object(objectName)
	reader, err := objHandle.NewReader(context.Background())
	if err != nil {
		t.Fatal(err)
	}
	defer reader.Close()
	data, err := ioutil.ReadAll(reader)
	if err != nil {
		t.Fatal(err)
	}
	if string(data) != content {
		t.Errorf("wrong data returned\nwant %q\ngot  %q", content, string(data))
	}
}

func TestServerClientObjectRangeReader(t *testing.T) {
	const bucketName = "some-bucket"
	const objectName = "items/data.txt"
	const content = "some really nice but long content stored in my object"
	server := NewServer([]Object{
		{
			BucketName: bucketName,
			Name:       objectName,
			Content:    []byte(content),
		},
	})
	defer server.Stop()
	var tests = []struct {
		testCase string
		offset   int64
		length   int64
	}{
		{
			"no length, just offset",
			2,
			-1,
		},
		{
			"zero offset, length",
			0,
			10,
		},
		{
			"offset and length",
			4,
			10,
		},
	}
	for _, test := range tests {
		t.Run(test.testCase, func(t *testing.T) {
			length := test.length
			if length == -1 {
				length = int64(len(content)) - test.offset + 1
			}
			expectedData := content[test.offset : test.offset+length-1]
			client := server.Client()
			objHandle := client.Bucket(bucketName).Object(objectName)
			reader, err := objHandle.NewRangeReader(context.Background(), test.offset, test.length)
			if err != nil {
				t.Fatal(err)
			}
			defer reader.Close()
			data, err := ioutil.ReadAll(reader)
			if err != nil {
				t.Fatal(err)
			}
			if string(data) != expectedData {
				t.Errorf("wrong data returned\nwant %q\ngot  %q", expectedData, string(data))
			}
		})
	}
}

func TestServerClientObjectReaderAfterCreateObject(t *testing.T) {
	const bucketName = "staging-bucket"
	const objectName = "items/data-overwritten.txt"
	const content = "data inside the object"
	server := NewServer([]Object{
		{BucketName: bucketName, Name: objectName},
	})
	defer server.Stop()
	server.CreateObject(Object{BucketName: bucketName, Name: objectName, Content: []byte(content)})
	client := server.Client()
	objHandle := client.Bucket(bucketName).Object(objectName)
	reader, err := objHandle.NewReader(context.Background())
	if err != nil {
		t.Fatal(err)
	}
	defer reader.Close()
	data, err := ioutil.ReadAll(reader)
	if err != nil {
		t.Fatal(err)
	}
	if string(data) != content {
		t.Errorf("wrong data returned\nwant %q\ngot  %q", content, string(data))
	}
}

func TestServerClientObjectReaderError(t *testing.T) {
	server := NewServer([]Object{
		{BucketName: "some-bucket", Name: "img/hi-res/party-01.jpg"},
	})
	defer server.Stop()
	var tests = []struct {
		testCase   string
		bucketName string
		objectName string
	}{
		{
			"bucket not found",
			"other-bucket",
			"whatever-object",
		},
		{
			"object not found",
			"some-bucket",
			"img/low-res/party-01.jpg",
		},
	}
	for _, test := range tests {
		t.Run(test.testCase, func(t *testing.T) {
			objHandle := server.Client().Bucket(test.bucketName).Object(test.objectName)
			reader, err := objHandle.NewReader(context.Background())
			if err == nil {
				t.Error("unexpected <nil> error")
			}
			if reader != nil {
				t.Errorf("unexpected non-nil attrs: %#v", reader)
			}
		})
	}
}

func TestServiceClientListObjects(t *testing.T) {
	server := NewServer([]Object{
		{BucketName: "some-bucket", Name: "img/low-res/party-01.jpg"},
		{BucketName: "some-bucket", Name: "img/hi-res/party-01.jpg"},
		{BucketName: "some-bucket", Name: "img/low-res/party-02.jpg"},
		{BucketName: "some-bucket", Name: "img/hi-res/party-02.jpg"},
		{BucketName: "some-bucket", Name: "img/low-res/party-03.jpg"},
		{BucketName: "some-bucket", Name: "img/hi-res/party-03.jpg"},
		{BucketName: "some-bucket", Name: "img/brand.jpg"},
		{BucketName: "some-bucket", Name: "video/hi-res/some_video_1080p.mp4"},
		{BucketName: "other-bucket", Name: "static/css/style.css"},
	})
	defer server.Stop()
	server.CreateBucket("empty-bucket")
	var tests = []struct {
		testCase      string
		bucketName    string
		query         *storage.Query
		expectedNames []string
	}{
		{
			"no prefix, no delimiter, multiple objects",
			"some-bucket",
			nil,
			[]string{
				"img/brand.jpg",
				"img/hi-res/party-01.jpg",
				"img/hi-res/party-02.jpg",
				"img/hi-res/party-03.jpg",
				"img/low-res/party-01.jpg",
				"img/low-res/party-02.jpg",
				"img/low-res/party-03.jpg",
				"video/hi-res/some_video_1080p.mp4",
			},
		},
		{
			"no prefix, no delimiter, single object",
			"other-bucket",
			nil,
			[]string{"static/css/style.css"},
		},
		{
			"no prefix, no delimiter, no objects",
			"empty-bucket",
			nil,
			[]string{},
		},
		{
			"filtering prefix only",
			"some-bucket",
			&storage.Query{Prefix: "img/"},
			[]string{
				"img/brand.jpg",
				"img/hi-res/party-01.jpg",
				"img/hi-res/party-02.jpg",
				"img/hi-res/party-03.jpg",
				"img/low-res/party-01.jpg",
				"img/low-res/party-02.jpg",
				"img/low-res/party-03.jpg",
			},
		},
		{
			"filtering prefix and delimiter",
			"some-bucket",
			&storage.Query{Prefix: "img/", Delimiter: "/"},
			[]string{"img/brand.jpg"},
		},
		{
			"filtering prefix, no objects",
			"some-bucket",
			&storage.Query{Prefix: "static/"},
			[]string{},
		},
	}
	client := server.Client()
	for _, test := range tests {
		t.Run(test.testCase, func(t *testing.T) {
			iter := client.Bucket(test.bucketName).Objects(context.Background(), test.query)
			names := []string{}
			obj, err := iter.Next()
			for ; err == nil; obj, err = iter.Next() {
				names = append(names, obj.Name)
			}
			if err != iterator.Done {
				t.Fatal(err)
			}
			if !reflect.DeepEqual(names, test.expectedNames) {
				t.Errorf("wrong names returned\nwant %#v\ngot  %#v", test.expectedNames, names)
			}
		})
	}
}

func TestServiceClientListObjectsBucketNotFound(t *testing.T) {
	server := NewServer(nil)
	defer server.Stop()
	iter := server.Client().Bucket("some-bucket").Objects(context.Background(), nil)
	obj, err := iter.Next()
	if err == nil {
		t.Error("got unexpected <nil> error")
	}
	if obj != nil {
		t.Errorf("got unexpected non-nil obj: %#v", obj)
	}
}
<commit_msg>fakestorage: add a test where prefix matches object name<commit_after>// Copyright 2017 Francisco Souza. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package fakestorage

import (
	"io/ioutil"
	"reflect"
	"testing"

	"cloud.google.com/go/storage"
	"golang.org/x/net/context"
	"google.golang.org/api/iterator"
)

func TestServerClientObjectAttrs(t *testing.T) {
	const bucketName = "some-bucket"
	const objectName = "img/hi-res/party-01.jpg"
	const content = "some nice content"
	server := NewServer([]Object{
		{BucketName: bucketName, Name: objectName, Content: []byte(content)},
	})
	defer server.Stop()
	client := server.Client()
	objHandle := client.Bucket(bucketName).Object(objectName)
	attrs, err := objHandle.Attrs(context.Background())
	if err != nil {
		t.Fatal(err)
	}
	if attrs.Bucket != bucketName {
		t.Errorf("wrong bucket name\nwant %q\ngot  %q", bucketName, attrs.Bucket)
	}
	if attrs.Name != objectName {
		t.Errorf("wrong object name\nwant %q\ngot  %q", objectName, attrs.Name)
	}
	if attrs.Size != int64(len(content)) {
		t.Errorf("wrong size returned\nwant %d\ngot  %d", len(content), attrs.Size)
	}
}

func TestServerClientObjectAttrsAfterCreateObject(t *testing.T) {
	const bucketName = "prod-bucket"
	const objectName = "video/hi-res/best_video_1080p.mp4"
	server := NewServer(nil)
	defer server.Stop()
	server.CreateObject(Object{BucketName: bucketName, Name: objectName})
	client := server.Client()
	objHandle := client.Bucket(bucketName).Object(objectName)
	attrs, err := objHandle.Attrs(context.Background())
	if err != nil {
		t.Fatal(err)
	}
	if attrs.Bucket != bucketName {
		t.Errorf("wrong bucket name\nwant %q\ngot  %q", bucketName, attrs.Bucket)
	}
	if attrs.Name != objectName {
		t.Errorf("wrong object name\n want %q\ngot  %q", objectName, attrs.Name)
	}
}

func TestServerClientObjectAttrsErrors(t *testing.T) {
	server := NewServer([]Object{
		{BucketName: "some-bucket", Name: "img/hi-res/party-01.jpg"},
	})
	defer server.Stop()
	var tests = []struct {
		testCase   string
		bucketName string
		objectName string
	}{
		{
			"bucket not found",
			"other-bucket",
			"whatever-object",
		},
		{
			"object not found",
			"some-bucket",
			"img/low-res/party-01.jpg",
		},
	}
	for _, test := range tests {
		t.Run(test.testCase, func(t *testing.T) {
			objHandle := server.Client().Bucket(test.bucketName).Object(test.objectName)
			attrs, err := objHandle.Attrs(context.Background())
			if err == nil {
				t.Error("unexpected <nil> error")
			}
			if attrs != nil {
				t.Errorf("unexpected non-nil attrs: %#v", attrs)
			}
		})
	}
}

func TestServerClientObjectReader(t *testing.T) {
	const bucketName = "some-bucket"
	const objectName = "items/data.txt"
	const content = "some nice content"
	server := NewServer([]Object{
		{
			BucketName: bucketName,
			Name:       objectName,
			Content:    []byte(content),
		},
	})
	defer server.Stop()
	client := server.Client()
	objHandle := client.Bucket(bucketName).Object(objectName)
	reader, err := objHandle.NewReader(context.Background())
	if err != nil {
		t.Fatal(err)
	}
	defer reader.Close()
	data, err := ioutil.ReadAll(reader)
	if err != nil {
		t.Fatal(err)
	}
	if string(data) != content {
		t.Errorf("wrong data returned\nwant %q\ngot  %q", content, string(data))
	}
}

func TestServerClientObjectRangeReader(t *testing.T) {
	const bucketName = "some-bucket"
	const objectName = "items/data.txt"
	const content = "some really nice but long content stored in my object"
	server := NewServer([]Object{
		{
			BucketName: bucketName,
			Name:       objectName,
			Content:    []byte(content),
		},
	})
	defer server.Stop()
	var tests = []struct {
		testCase string
		offset   int64
		length   int64
	}{
		{
			"no length, just offset",
			2,
			-1,
		},
		{
			"zero offset, length",
			0,
			10,
		},
		{
			"offset and length",
			4,
			10,
		},
	}
	for _, test := range tests {
		t.Run(test.testCase, func(t *testing.T) {
			length := test.length
			if length == -1 {
				length = int64(len(content)) - test.offset + 1
			}
			expectedData := content[test.offset : test.offset+length-1]
			client := server.Client()
			objHandle := client.Bucket(bucketName).Object(objectName)
			reader, err := objHandle.NewRangeReader(context.Background(), test.offset, test.length)
			if err != nil {
				t.Fatal(err)
			}
			defer reader.Close()
			data, err := ioutil.ReadAll(reader)
			if err != nil {
				t.Fatal(err)
			}
			if string(data) != expectedData {
				t.Errorf("wrong data returned\nwant %q\ngot  %q", expectedData, string(data))
			}
		})
	}
}

func TestServerClientObjectReaderAfterCreateObject(t *testing.T) {
	const bucketName = "staging-bucket"
	const objectName = "items/data-overwritten.txt"
	const content = "data inside the object"
	server := NewServer([]Object{
		{BucketName: bucketName, Name: objectName},
	})
	defer server.Stop()
	server.CreateObject(Object{BucketName: bucketName, Name: objectName, Content: []byte(content)})
	client := server.Client()
	objHandle := client.Bucket(bucketName).Object(objectName)
	reader, err := objHandle.NewReader(context.Background())
	if err != nil {
		t.Fatal(err)
	}
	defer reader.Close()
	data, err := ioutil.ReadAll(reader)
	if err != nil {
		t.Fatal(err)
	}
	if string(data) != content {
		t.Errorf("wrong data returned\nwant %q\ngot  %q", content, string(data))
	}
}

func TestServerClientObjectReaderError(t *testing.T) {
	server := NewServer([]Object{
		{BucketName: "some-bucket", Name: "img/hi-res/party-01.jpg"},
	})
	defer server.Stop()
	var tests = []struct {
		testCase   string
		bucketName string
		objectName string
	}{
		{
			"bucket not found",
			"other-bucket",
			"whatever-object",
		},
		{
			"object not found",
			"some-bucket",
			"img/low-res/party-01.jpg",
		},
	}
	for _, test := range tests {
		t.Run(test.testCase, func(t *testing.T) {
			objHandle := server.Client().Bucket(test.bucketName).Object(test.objectName)
			reader, err := objHandle.NewReader(context.Background())
			if err == nil {
				t.Error("unexpected <nil> error")
			}
			if reader != nil {
				t.Errorf("unexpected non-nil attrs: %#v", reader)
			}
		})
	}
}

func TestServiceClientListObjects(t *testing.T) {
	server := NewServer([]Object{
		{BucketName: "some-bucket", Name: "img/low-res/party-01.jpg"},
		{BucketName: "some-bucket", Name: "img/hi-res/party-01.jpg"},
		{BucketName: "some-bucket", Name: "img/low-res/party-02.jpg"},
		{BucketName: "some-bucket", Name: "img/hi-res/party-02.jpg"},
		{BucketName: "some-bucket", Name: "img/low-res/party-03.jpg"},
		{BucketName: "some-bucket", Name: "img/hi-res/party-03.jpg"},
		{BucketName: "some-bucket", Name: "img/brand.jpg"},
		{BucketName: "some-bucket", Name: "video/hi-res/some_video_1080p.mp4"},
		{BucketName: "other-bucket", Name: "static/css/style.css"},
	})
	defer server.Stop()
	server.CreateBucket("empty-bucket")
	var tests = []struct {
		testCase      string
		bucketName    string
		query         *storage.Query
		expectedNames []string
	}{
		{
			"no prefix, no delimiter, multiple objects",
			"some-bucket",
			nil,
			[]string{
				"img/brand.jpg",
				"img/hi-res/party-01.jpg",
				"img/hi-res/party-02.jpg",
				"img/hi-res/party-03.jpg",
				"img/low-res/party-01.jpg",
				"img/low-res/party-02.jpg",
				"img/low-res/party-03.jpg",
				"video/hi-res/some_video_1080p.mp4",
			},
		},
		{
			"no prefix, no delimiter, single object",
			"other-bucket",
			nil,
			[]string{"static/css/style.css"},
		},
		{
			"no prefix, no delimiter, no objects",
			"empty-bucket",
			nil,
			[]string{},
		},
		{
			"filtering prefix only",
			"some-bucket",
			&storage.Query{Prefix: "img/"},
			[]string{
				"img/brand.jpg",
				"img/hi-res/party-01.jpg",
				"img/hi-res/party-02.jpg",
				"img/hi-res/party-03.jpg",
				"img/low-res/party-01.jpg",
				"img/low-res/party-02.jpg",
				"img/low-res/party-03.jpg",
			},
		},
		{
			"full prefix",
			"some-bucket",
			&storage.Query{Prefix: "img/brand.jpg"},
			[]string{"img/brand.jpg"},
		},
		{
			"filtering prefix and delimiter",
			"some-bucket",
			&storage.Query{Prefix: "img/", Delimiter: "/"},
			[]string{"img/brand.jpg"},
		},
		{
			"filtering prefix, no objects",
			"some-bucket",
			&storage.Query{Prefix: "static/"},
			[]string{},
		},
	}
	client := server.Client()
	for _, test := range tests {
		t.Run(test.testCase, func(t *testing.T) {
			iter := client.Bucket(test.bucketName).Objects(context.Background(), test.query)
			names := []string{}
			obj, err := iter.Next()
			for ; err == nil; obj, err = iter.Next() {
				names = append(names, obj.Name)
			}
			if err != iterator.Done {
				t.Fatal(err)
			}
			if !reflect.DeepEqual(names, test.expectedNames) {
				t.Errorf("wrong names returned\nwant %#v\ngot  %#v", test.expectedNames, names)
			}
		})
	}
}

func TestServiceClientListObjectsBucketNotFound(t *testing.T) {
	server := NewServer(nil)
	defer server.Stop()
	iter := server.Client().Bucket("some-bucket").Objects(context.Background(), nil)
	obj, err := iter.Next()
	if err == nil {
		t.Error("got unexpected <nil> error")
	}
	if obj != nil {
		t.Errorf("got unexpected non-nil obj: %#v", obj)
	}
}
<|endoftext|>


<commit_before>package camo

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

	"github.com/cactus/go-camo/camo/encoding"
	"github.com/cactus/go-camo/router"
	"github.com/stretchr/testify/assert"
)

var camoConfig = Config{
	HMACKey:        []byte("0x24FEEDFACEDEADBEEFCAFE"),
	MaxSize:        5120 * 1024,
	RequestTimeout: time.Duration(10) * time.Second,
	MaxRedirects:   3,
	ServerName:     "go-camo",
}

func makeReq(testURL string) (*http.Request, error) {
	k := []byte(camoConfig.HMACKey)
	hexURL := encoding.B64EncodeURL(k, testURL)
	out := "http://example.com" + hexURL
	req, err := http.NewRequest("GET", out, nil)
	if err != nil {
		return nil, fmt.Errorf("Error building req url '%s': %s", testURL, err.Error())
	}
	return req, nil
}

func processRequest(req *http.Request, status int, camoConfig Config) (*httptest.ResponseRecorder, error) {
	camoServer, err := New(camoConfig)
	if err != nil {
		return nil, fmt.Errorf("Error building Camo: %s", err.Error())
	}

	router := &router.DumbRouter{
		AddHeaders:  map[string]string{"X-Go-Camo": "test"},
		ServerName:  camoConfig.ServerName,
		CamoHandler: camoServer,
	}

	record := httptest.NewRecorder()
	router.ServeHTTP(record, req)
	if got, want := record.Code, status; got != want {
		return record, fmt.Errorf("response code = %d, wanted %d", got, want)
	}
	return record, nil
}

func makeTestReq(testURL string, status int) (*httptest.ResponseRecorder, error) {
	req, err := makeReq(testURL)
	if err != nil {
		return nil, err
	}
	record, err := processRequest(req, status, camoConfig)
	if err != nil {
		return nil, err
	}
	return record, nil
}

func TestNotFound(t *testing.T) {
	t.Parallel()
	req, err := http.NewRequest("GET", "http://example.com/favicon.ico", nil)
	assert.Nil(t, err)

	record, err := processRequest(req, 404, camoConfig)
	assert.Nil(t, err)
	assert.Equal(t, record.Code, 404, "Expected 404 but got '%d' instead", record.Code)
	assert.Equal(t, record.Body.String(), "404 Not Found\n", "Expected 404 response body but got '%s' instead", record.Body.String())
	// validate headers
	assert.Equal(t, record.HeaderMap.Get("X-Go-Camo"), "test", "Expected custom response header not found")
	assert.Equal(t, record.HeaderMap.Get("Server"), "go-camo", "Expected 'Server' response header not found")
}

func TestSimpleValidImageURL(t *testing.T) {
	t.Parallel()
	testURL := "http://www.google.com/images/srpr/logo11w.png"
	record, err := makeTestReq(testURL, 200)
	assert.Nil(t, err)
	// validate headers
	assert.Equal(t, record.HeaderMap.Get("X-Go-Camo"), "test", "Expected custom response header not found")
	assert.Equal(t, record.HeaderMap.Get("Server"), "go-camo", "Expected 'Server' response header not found")
}

func TestGoogleChartURL(t *testing.T) {
	t.Parallel()
	testURL := "http://chart.apis.google.com/chart?chs=920x200&chxl=0:%7C2010-08-13%7C2010-09-12%7C2010-10-12%7C2010-11-11%7C1:%7C0%7C0%7C0%7C0%7C0%7C0&chm=B,EBF5FB,0,0,0&chco=008Cd6&chls=3,1,0&chg=8.3,20,1,4&chd=s:AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA&chxt=x,y&cht=lc"
	_, err := makeTestReq(testURL, 200)
	assert.Nil(t, err)
}

func TestChunkedImageFile(t *testing.T) {
	t.Parallel()
	testURL := "http://www.igvita.com/posts/12/spdyproxy-diagram.png"
	_, err := makeTestReq(testURL, 200)
	assert.Nil(t, err)
}

func TestFollowRedirects(t *testing.T) {
	t.Parallel()
	testURL := "http://cl.ly/1K0X2Y2F1P0o3z140p0d/boom-headshot.gif"
	_, err := makeTestReq(testURL, 200)
	assert.Nil(t, err)
}

func TestStrangeFormatRedirects(t *testing.T) {
	t.Parallel()
	testURL := "http://cl.ly/DPcp/Screen%20Shot%202012-01-17%20at%203.42.32%20PM.png"
	_, err := makeTestReq(testURL, 200)
	assert.Nil(t, err)
}

func TestRedirectsWithPathOnly(t *testing.T) {
	t.Parallel()
	testURL := "http://httpbin.org/redirect-to?url=%2Fredirect-to%3Furl%3Dhttp%3A%2F%2Fwww.google.com%2Fimages%2Fsrpr%2Flogo11w.png"
	_, err := makeTestReq(testURL, 200)
	assert.Nil(t, err)
}

func TestFollowTempRedirects(t *testing.T) {
	t.Parallel()
	testURL := "http://httpbin.org/redirect-to?url=http://www.google.com/images/srpr/logo11w.png"
	_, err := makeTestReq(testURL, 200)
	assert.Nil(t, err)
}

func TestBadContentType(t *testing.T) {
	t.Parallel()
	testURL := "http://httpbin.org/response-headers?Content-Type=what"
	_, err := makeTestReq(testURL, 400)
	assert.Nil(t, err)
}

func Test404InfiniRedirect(t *testing.T) {
	t.Parallel()
	testURL := "http://httpbin.org/redirect/4"
	_, err := makeTestReq(testURL, 404)
	assert.Nil(t, err)
}

func Test404URLWithoutHTTPHost(t *testing.T) {
	t.Parallel()
	testURL := "/picture/Mincemeat/Pimp.jpg"
	_, err := makeTestReq(testURL, 404)
	assert.Nil(t, err)
}

func Test404ImageLargerThan5MB(t *testing.T) {
	t.Parallel()
	testURL := "http://apod.nasa.gov/apod/image/0505/larryslookout_spirit_big.jpg"
	_, err := makeTestReq(testURL, 404)
	assert.Nil(t, err)
}

func Test404HostNotFound(t *testing.T) {
	t.Parallel()
	testURL := "http://flabergasted.cx"
	_, err := makeTestReq(testURL, 404)
	assert.Nil(t, err)
}

func Test404OnExcludes(t *testing.T) {
	t.Parallel()
	testURL := "http://iphone.internal.example.org/foo.cgi"
	_, err := makeTestReq(testURL, 404)
	assert.Nil(t, err)
}

func Test404OnNonImageContent(t *testing.T) {
	t.Parallel()
	testURL := "https://github.com/atmos/cinderella/raw/master/bootstrap.sh"
	_, err := makeTestReq(testURL, 404)
	assert.Nil(t, err)
}

func Test404On10xIpRange(t *testing.T) {
	t.Parallel()
	testURL := "http://10.0.0.1/foo.cgi"
	_, err := makeTestReq(testURL, 404)
	assert.Nil(t, err)
}

func Test404On169Dot254Net(t *testing.T) {
	t.Parallel()
	testURL := "http://169.254.0.1/foo.cgi"
	_, err := makeTestReq(testURL, 404)
	assert.Nil(t, err)
}

func Test404On172Dot16Net(t *testing.T) {
	t.Parallel()
	for i := 16; i < 32; i++ {
		testURL := "http://172.%d.0.1/foo.cgi"
		_, err := makeTestReq(fmt.Sprintf(testURL, i), 404)
		assert.Nil(t, err)
	}
}

func Test404On192Dot168Net(t *testing.T) {
	t.Parallel()
	testURL := "http://192.168.0.1/foo.cgi"
	_, err := makeTestReq(testURL, 404)
	assert.Nil(t, err)
}

func TestSupplyAcceptIfNoneGiven(t *testing.T) {
	t.Parallel()
	testURL := "http://images.anandtech.com/doci/6673/OpenMoboAMD30_575px.png"
	req, err := makeReq(testURL)
	req.Header.Del("Accept")
	assert.Nil(t, err)
	_, err = processRequest(req, 200, camoConfig)
	assert.Nil(t, err)
}

func TestTimeoutBeforeHeader(t *testing.T) {
	t.Parallel()
	c := Config{
		HMACKey:        []byte("0x24FEEDFACEDEADBEEFCAFE"),
		MaxSize:        5120 * 1024,
		RequestTimeout: time.Duration(500) * time.Millisecond,
		MaxRedirects:   3,
		ServerName:     "go-camo",
	}
	cc := make(chan bool, 1)
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		<-cc
		return
	}))
	defer ts.Close()

	req, err := makeReq(ts.URL)
	assert.Nil(t, err)

	errc := make(chan error, 1)
	go func() {
		_, err := processRequest(req, 502, c)
		errc <- err
	}()

	select {
	case e := <-errc:
		assert.Nil(t, e)
		cc <- true
	case <-time.After(1 * time.Second):
		cc <- true
		fmt.Errorf("timeout didn't fire in time")
	}
}
<commit_msg>deal with race in httptest<commit_after>package camo

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

	"github.com/cactus/go-camo/camo/encoding"
	"github.com/cactus/go-camo/router"
	"github.com/stretchr/testify/assert"
)

var camoConfig = Config{
	HMACKey:        []byte("0x24FEEDFACEDEADBEEFCAFE"),
	MaxSize:        5120 * 1024,
	RequestTimeout: time.Duration(10) * time.Second,
	MaxRedirects:   3,
	ServerName:     "go-camo",
}

func makeReq(testURL string) (*http.Request, error) {
	k := []byte(camoConfig.HMACKey)
	hexURL := encoding.B64EncodeURL(k, testURL)
	out := "http://example.com" + hexURL
	req, err := http.NewRequest("GET", out, nil)
	if err != nil {
		return nil, fmt.Errorf("Error building req url '%s': %s", testURL, err.Error())
	}
	return req, nil
}

func processRequest(req *http.Request, status int, camoConfig Config) (*httptest.ResponseRecorder, error) {
	camoServer, err := New(camoConfig)
	if err != nil {
		return nil, fmt.Errorf("Error building Camo: %s", err.Error())
	}

	router := &router.DumbRouter{
		AddHeaders:  map[string]string{"X-Go-Camo": "test"},
		ServerName:  camoConfig.ServerName,
		CamoHandler: camoServer,
	}

	record := httptest.NewRecorder()
	router.ServeHTTP(record, req)
	if got, want := record.Code, status; got != want {
		return record, fmt.Errorf("response code = %d, wanted %d", got, want)
	}
	return record, nil
}

func makeTestReq(testURL string, status int) (*httptest.ResponseRecorder, error) {
	req, err := makeReq(testURL)
	if err != nil {
		return nil, err
	}
	record, err := processRequest(req, status, camoConfig)
	if err != nil {
		return nil, err
	}
	return record, nil
}

func TestNotFound(t *testing.T) {
	t.Parallel()
	req, err := http.NewRequest("GET", "http://example.com/favicon.ico", nil)
	assert.Nil(t, err)

	record, err := processRequest(req, 404, camoConfig)
	assert.Nil(t, err)
	assert.Equal(t, record.Code, 404, "Expected 404 but got '%d' instead", record.Code)
	assert.Equal(t, record.Body.String(), "404 Not Found\n", "Expected 404 response body but got '%s' instead", record.Body.String())
	// validate headers
	assert.Equal(t, record.HeaderMap.Get("X-Go-Camo"), "test", "Expected custom response header not found")
	assert.Equal(t, record.HeaderMap.Get("Server"), "go-camo", "Expected 'Server' response header not found")
}

func TestSimpleValidImageURL(t *testing.T) {
	t.Parallel()
	testURL := "http://www.google.com/images/srpr/logo11w.png"
	record, err := makeTestReq(testURL, 200)
	assert.Nil(t, err)
	// validate headers
	assert.Equal(t, record.HeaderMap.Get("X-Go-Camo"), "test", "Expected custom response header not found")
	assert.Equal(t, record.HeaderMap.Get("Server"), "go-camo", "Expected 'Server' response header not found")
}

func TestGoogleChartURL(t *testing.T) {
	t.Parallel()
	testURL := "http://chart.apis.google.com/chart?chs=920x200&chxl=0:%7C2010-08-13%7C2010-09-12%7C2010-10-12%7C2010-11-11%7C1:%7C0%7C0%7C0%7C0%7C0%7C0&chm=B,EBF5FB,0,0,0&chco=008Cd6&chls=3,1,0&chg=8.3,20,1,4&chd=s:AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA&chxt=x,y&cht=lc"
	_, err := makeTestReq(testURL, 200)
	assert.Nil(t, err)
}

func TestChunkedImageFile(t *testing.T) {
	t.Parallel()
	testURL := "http://www.igvita.com/posts/12/spdyproxy-diagram.png"
	_, err := makeTestReq(testURL, 200)
	assert.Nil(t, err)
}

func TestFollowRedirects(t *testing.T) {
	t.Parallel()
	testURL := "http://cl.ly/1K0X2Y2F1P0o3z140p0d/boom-headshot.gif"
	_, err := makeTestReq(testURL, 200)
	assert.Nil(t, err)
}

func TestStrangeFormatRedirects(t *testing.T) {
	t.Parallel()
	testURL := "http://cl.ly/DPcp/Screen%20Shot%202012-01-17%20at%203.42.32%20PM.png"
	_, err := makeTestReq(testURL, 200)
	assert.Nil(t, err)
}

func TestRedirectsWithPathOnly(t *testing.T) {
	t.Parallel()
	testURL := "http://httpbin.org/redirect-to?url=%2Fredirect-to%3Furl%3Dhttp%3A%2F%2Fwww.google.com%2Fimages%2Fsrpr%2Flogo11w.png"
	_, err := makeTestReq(testURL, 200)
	assert.Nil(t, err)
}

func TestFollowTempRedirects(t *testing.T) {
	t.Parallel()
	testURL := "http://httpbin.org/redirect-to?url=http://www.google.com/images/srpr/logo11w.png"
	_, err := makeTestReq(testURL, 200)
	assert.Nil(t, err)
}

func TestBadContentType(t *testing.T) {
	t.Parallel()
	testURL := "http://httpbin.org/response-headers?Content-Type=what"
	_, err := makeTestReq(testURL, 400)
	assert.Nil(t, err)
}

func Test404InfiniRedirect(t *testing.T) {
	t.Parallel()
	testURL := "http://httpbin.org/redirect/4"
	_, err := makeTestReq(testURL, 404)
	assert.Nil(t, err)
}

func Test404URLWithoutHTTPHost(t *testing.T) {
	t.Parallel()
	testURL := "/picture/Mincemeat/Pimp.jpg"
	_, err := makeTestReq(testURL, 404)
	assert.Nil(t, err)
}

func Test404ImageLargerThan5MB(t *testing.T) {
	t.Parallel()
	testURL := "http://apod.nasa.gov/apod/image/0505/larryslookout_spirit_big.jpg"
	_, err := makeTestReq(testURL, 404)
	assert.Nil(t, err)
}

func Test404HostNotFound(t *testing.T) {
	t.Parallel()
	testURL := "http://flabergasted.cx"
	_, err := makeTestReq(testURL, 404)
	assert.Nil(t, err)
}

func Test404OnExcludes(t *testing.T) {
	t.Parallel()
	testURL := "http://iphone.internal.example.org/foo.cgi"
	_, err := makeTestReq(testURL, 404)
	assert.Nil(t, err)
}

func Test404OnNonImageContent(t *testing.T) {
	t.Parallel()
	testURL := "https://github.com/atmos/cinderella/raw/master/bootstrap.sh"
	_, err := makeTestReq(testURL, 404)
	assert.Nil(t, err)
}

func Test404On10xIpRange(t *testing.T) {
	t.Parallel()
	testURL := "http://10.0.0.1/foo.cgi"
	_, err := makeTestReq(testURL, 404)
	assert.Nil(t, err)
}

func Test404On169Dot254Net(t *testing.T) {
	t.Parallel()
	testURL := "http://169.254.0.1/foo.cgi"
	_, err := makeTestReq(testURL, 404)
	assert.Nil(t, err)
}

func Test404On172Dot16Net(t *testing.T) {
	t.Parallel()
	for i := 16; i < 32; i++ {
		testURL := "http://172.%d.0.1/foo.cgi"
		_, err := makeTestReq(fmt.Sprintf(testURL, i), 404)
		assert.Nil(t, err)
	}
}

func Test404On192Dot168Net(t *testing.T) {
	t.Parallel()
	testURL := "http://192.168.0.1/foo.cgi"
	_, err := makeTestReq(testURL, 404)
	assert.Nil(t, err)
}

func TestSupplyAcceptIfNoneGiven(t *testing.T) {
	t.Parallel()
	testURL := "http://images.anandtech.com/doci/6673/OpenMoboAMD30_575px.png"
	req, err := makeReq(testURL)
	req.Header.Del("Accept")
	assert.Nil(t, err)
	_, err = processRequest(req, 200, camoConfig)
	assert.Nil(t, err)
}

func TestTimeout(t *testing.T) {
	t.Parallel()
	c := Config{
		HMACKey:        []byte("0x24FEEDFACEDEADBEEFCAFE"),
		MaxSize:        5120 * 1024,
		RequestTimeout: time.Duration(500) * time.Millisecond,
		MaxRedirects:   3,
		ServerName:     "go-camo",
	}
	cc := make(chan bool, 1)
	received := make(chan bool)
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		received <- true
		<-cc
		r.Close = true
		w.Write([]byte("ok"))
		return
	}))
	defer ts.Close()

	req, err := makeReq(ts.URL)
	assert.Nil(t, err)

	errc := make(chan error, 1)
	go func() {
		_, err := processRequest(req, 502, c)
		errc <- err
	}()

	<-received
	select {
	case e := <-errc:
		assert.Nil(t, e)
		cc <- true
	case <-time.After(1 * time.Second):
		cc <- true
		fmt.Errorf("timeout didn't fire in time")
	}
	close(cc)
}
<|endoftext|>


<commit_before>// Copyright 2011, 2012, 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package charm

import (
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"regexp"
	"strings"

	"github.com/binary132/gojsonschema"
	"launchpad.net/goyaml"
)

var prohibitedSchemaKeys = map[string]bool{"$ref": true, "$schema": true}

var actionNameRule = regexp.MustCompile("^[a-z](?:[a-z-]*[a-z])?$")

// Actions defines the available actions for the charm.  Additional params
// may be added as metadata at a future time (e.g. version.)
type Actions struct {
	ActionSpecs map[string]ActionSpec `yaml:"actions,omitempty" bson:",omitempty"`
}

// ActionSpec is a definition of the parameters and traits of an Action.
// The Params map is expected to conform to JSON-Schema Draft 4 as defined at
// http://json-schema.org/draft-04/schema# (see http://json-schema.org/latest/json-schema-core.html)
type ActionSpec struct {
	Description string
	Params      map[string]interface{}
}

func NewActions() *Actions {
	return &Actions{}
}

// ValidateParams tells us whether an unmarshaled JSON object conforms to the
// Params for the specific ActionSpec.
// Usage: ok, err := ch.Actions()["snapshot"].Validate(jsonParams)
func (spec *ActionSpec) ValidateParams(params interface{}) (bool, error) {

	specSchemaDoc, err := gojsonschema.NewJsonSchemaDocument(spec.Params)
	if err != nil {
		return false, err
	}

	results := specSchemaDoc.Validate(params)
	if results.Valid() {
		return true, nil
	}

	var errorStrings []string
	for _, validationError := range results.Errors() {
		errorStrings = append(errorStrings, validationError.String())
	}
	return false, fmt.Errorf("JSON validation failed: %s", strings.Join(errorStrings, "; "))
}

// ReadActions builds an Actions spec from a charm's actions.yaml.
func ReadActionsYaml(r io.Reader) (*Actions, error) {
	data, err := ioutil.ReadAll(r)
	if err != nil {
		return nil, err
	}
	var unmarshaledActions Actions
	if err := goyaml.Unmarshal(data, &unmarshaledActions); err != nil {
		return nil, err
	}

	for name, actionSpec := range unmarshaledActions.ActionSpecs {
		if valid := actionNameRule.MatchString(name); !valid {
			return nil, fmt.Errorf("bad action name %s", name)
		}

		// Clean any map[interface{}]interface{}s out so they don't
		// cause problems with BSON serialization later.
		cleansedParams, err := cleanse(actionSpec.Params)
		if err != nil {
			return nil, err
		}

		// JSON-Schema must be a map
		cleansedParamsMap, ok := cleansedParams.(map[string]interface{})
		if !ok {
			return nil, fmt.Errorf("the params failed to parse as a map")
		}

		// Now substitute the cleansed map into the original.
		var tempSpec = unmarshaledActions.ActionSpecs[name]
		tempSpec.Params = cleansedParamsMap
		unmarshaledActions.ActionSpecs[name] = tempSpec

		// Make sure the new Params doc conforms to JSON-Schema
		// Draft 4 (http://json-schema.org/latest/json-schema-core.html)
		_, err = gojsonschema.NewJsonSchemaDocument(unmarshaledActions.ActionSpecs[name].Params)
		if err != nil {
			return nil, fmt.Errorf("invalid params schema for action schema %s: %v", name, err)
		}

	}
	return &unmarshaledActions, nil
}

// cleanse rejects schemas containing references or maps keyed with non-
// strings, and coerces acceptable maps to contain only maps with string keys.
func cleanse(input interface{}) (interface{}, error) {
	switch typedInput := input.(type) {

	// In this case, recurse in.
	case map[string]interface{}:
		newMap := make(map[string]interface{})
		for key, value := range typedInput {

			if prohibitedSchemaKeys[key] {
				return nil, fmt.Errorf("schema key %q not compatible with this version of juju", key)
			}

			newValue, err := cleanse(value)
			if err != nil {
				return nil, err
			}
			newMap[key] = newValue
		}
		return newMap, nil

	// Coerce keys to strings and error out if there's a problem; then recurse.
	case map[interface{}]interface{}:
		newMap := make(map[string]interface{})
		for key, value := range typedInput {
			typedKey, ok := key.(string)
			if !ok {
				return nil, errors.New("map keyed with non-string value")
			}
			newMap[typedKey] = value
		}
		return cleanse(newMap)

	// Recurse
	case []interface{}:
		newSlice := make([]interface{}, 0)
		for _, sliceValue := range typedInput {
			newSliceValue, err := cleanse(sliceValue)
			if err != nil {
				return nil, errors.New("map keyed with non-string value")
			}
			newSlice = append(newSlice, newSliceValue)
		}
		return newSlice, nil

	// Other kinds of values are OK.
	default:
		return input, nil
	}
}
<commit_msg>Update dependencies to pull in new repos with licensing fixed<commit_after>// Copyright 2011, 2012, 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package charm

import (
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"regexp"
	"strings"

	"github.com/juju/gojsonschema"
	"launchpad.net/goyaml"
)

var prohibitedSchemaKeys = map[string]bool{"$ref": true, "$schema": true}

var actionNameRule = regexp.MustCompile("^[a-z](?:[a-z-]*[a-z])?$")

// Actions defines the available actions for the charm.  Additional params
// may be added as metadata at a future time (e.g. version.)
type Actions struct {
	ActionSpecs map[string]ActionSpec `yaml:"actions,omitempty" bson:",omitempty"`
}

// ActionSpec is a definition of the parameters and traits of an Action.
// The Params map is expected to conform to JSON-Schema Draft 4 as defined at
// http://json-schema.org/draft-04/schema# (see http://json-schema.org/latest/json-schema-core.html)
type ActionSpec struct {
	Description string
	Params      map[string]interface{}
}

func NewActions() *Actions {
	return &Actions{}
}

// ValidateParams tells us whether an unmarshaled JSON object conforms to the
// Params for the specific ActionSpec.
// Usage: ok, err := ch.Actions()["snapshot"].Validate(jsonParams)
func (spec *ActionSpec) ValidateParams(params interface{}) (bool, error) {

	specSchemaDoc, err := gojsonschema.NewJsonSchemaDocument(spec.Params)
	if err != nil {
		return false, err
	}

	results := specSchemaDoc.Validate(params)
	if results.Valid() {
		return true, nil
	}

	var errorStrings []string
	for _, validationError := range results.Errors() {
		errorStrings = append(errorStrings, validationError.String())
	}
	return false, fmt.Errorf("JSON validation failed: %s", strings.Join(errorStrings, "; "))
}

// ReadActions builds an Actions spec from a charm's actions.yaml.
func ReadActionsYaml(r io.Reader) (*Actions, error) {
	data, err := ioutil.ReadAll(r)
	if err != nil {
		return nil, err
	}
	var unmarshaledActions Actions
	if err := goyaml.Unmarshal(data, &unmarshaledActions); err != nil {
		return nil, err
	}

	for name, actionSpec := range unmarshaledActions.ActionSpecs {
		if valid := actionNameRule.MatchString(name); !valid {
			return nil, fmt.Errorf("bad action name %s", name)
		}

		// Clean any map[interface{}]interface{}s out so they don't
		// cause problems with BSON serialization later.
		cleansedParams, err := cleanse(actionSpec.Params)
		if err != nil {
			return nil, err
		}

		// JSON-Schema must be a map
		cleansedParamsMap, ok := cleansedParams.(map[string]interface{})
		if !ok {
			return nil, fmt.Errorf("the params failed to parse as a map")
		}

		// Now substitute the cleansed map into the original.
		var tempSpec = unmarshaledActions.ActionSpecs[name]
		tempSpec.Params = cleansedParamsMap
		unmarshaledActions.ActionSpecs[name] = tempSpec

		// Make sure the new Params doc conforms to JSON-Schema
		// Draft 4 (http://json-schema.org/latest/json-schema-core.html)
		_, err = gojsonschema.NewJsonSchemaDocument(unmarshaledActions.ActionSpecs[name].Params)
		if err != nil {
			return nil, fmt.Errorf("invalid params schema for action schema %s: %v", name, err)
		}

	}
	return &unmarshaledActions, nil
}

// cleanse rejects schemas containing references or maps keyed with non-
// strings, and coerces acceptable maps to contain only maps with string keys.
func cleanse(input interface{}) (interface{}, error) {
	switch typedInput := input.(type) {

	// In this case, recurse in.
	case map[string]interface{}:
		newMap := make(map[string]interface{})
		for key, value := range typedInput {

			if prohibitedSchemaKeys[key] {
				return nil, fmt.Errorf("schema key %q not compatible with this version of juju", key)
			}

			newValue, err := cleanse(value)
			if err != nil {
				return nil, err
			}
			newMap[key] = newValue
		}
		return newMap, nil

	// Coerce keys to strings and error out if there's a problem; then recurse.
	case map[interface{}]interface{}:
		newMap := make(map[string]interface{})
		for key, value := range typedInput {
			typedKey, ok := key.(string)
			if !ok {
				return nil, errors.New("map keyed with non-string value")
			}
			newMap[typedKey] = value
		}
		return cleanse(newMap)

	// Recurse
	case []interface{}:
		newSlice := make([]interface{}, 0)
		for _, sliceValue := range typedInput {
			newSliceValue, err := cleanse(sliceValue)
			if err != nil {
				return nil, errors.New("map keyed with non-string value")
			}
			newSlice = append(newSlice, newSliceValue)
		}
		return newSlice, nil

	// Other kinds of values are OK.
	default:
		return input, nil
	}
}
<|endoftext|>


<commit_before>package main

import (
	"database/sql"
	"errors"
	"fmt"
)

type Coordinate struct {
	Latitude  float64
	Longitude float64
}

type Location struct {
	Address    Address
	Coordinate Coordinate
}

type Address struct {
	Street      string
	House       string
	HouseNumber string
	City        string
	Country     string
	Similarity  float64
}

func GetFuzzyAddress(config Config, address string, count int) ([]Location, error) {
	newconf := Config(config)
	newconf.DbName = "trgm_test"

	db, err := sql.Open("postgres", newconf.String())
	if err != nil {
		return nil, err
	}

	q := fmt.Sprintf("SELECT name, city, coord[0], coord[1], sml from get_appr('%s') LIMIT %d", address, count)
	rows, err := db.Query(q)

	if err != nil {
		return []Location{}, err
	}

	var locations []Location

	for rows.Next() {
		var (
			street_name, city string
			lat, long, sml    float64
		)

		if err := rows.Scan(&street_name, &city, &lat, &long, &sml); err != nil {
			return []Location{}, err
		}
		locations = append(locations,
			Location{Address{Street: street_name, City: city, Similarity: sml},
				Coordinate{Latitude: lat, Longitude: long}})
	}

	if err := rows.Err(); err != nil {
		return []Location{}, err
	}

	return locations, nil
}

func ResolveLocation(config Config, location Location) (Location, error) {
	if IsMissingCoordinate(location) {
		street := location.Address.Street
		if street != "" {
			locs, err := GetFuzzyAddress(config, street, 1)
			if err != nil {
				return Location{}, err
			}

			if len(locs) == 0 {
				e := errors.New("Couldn't find any address for street " + street)
				return Location{}, e
			}

			location.Coordinate = locs[0].Coordinate
			return location, nil
		} else {
			return Location{}, errors.New("Street empty, cannot search.")
		}
	} else {
		if location.Address.Country == "" {
			return Location{}, errors.New("Must provide country in Location.Address!")
		}

		coord := location.Coordinate
		correctCoord, err := CorrectPoint(config,
			Coord{coord.Latitude, coord.Longitude}, location.Address.Country)

		if err != nil {
			return Location{}, err
		}

		location.Coordinate.Latitude = correctCoord.Coord[0]
		location.Coordinate.Longitude = correctCoord.Coord[1]

		return location, nil
	}
}
<commit_msg>Fix the address data DTO.<commit_after>package main

import (
	"database/sql"
	"errors"
	"fmt"
)

type Coordinate struct {
	Latitude  float64
	Longitude float64
}

type Location struct {
	Address    Address
	Coordinate Coordinate
}

type Address struct {
	Street     string
	City       string
	Country    string
	Region     string
	PostalCode string

	ApartmentLetter string
	ApartmentNumber int
	Similarity      float64
}

func GetFuzzyAddress(config Config, address string, count int) ([]Location, error) {
	newconf := Config(config)
	newconf.DbName = "trgm_test"

	db, err := sql.Open("postgres", newconf.String())
	if err != nil {
		return nil, err
	}

	q := fmt.Sprintf("SELECT name, city, coord[0], coord[1], sml from get_appr('%s') LIMIT %d", address, count)
	rows, err := db.Query(q)

	if err != nil {
		return []Location{}, err
	}

	var locations []Location

	for rows.Next() {
		var (
			street_name, city string
			lat, long, sml    float64
		)

		if err := rows.Scan(&street_name, &city, &lat, &long, &sml); err != nil {
			return []Location{}, err
		}
		locations = append(locations,
			Location{Address{Street: street_name, City: city, Similarity: sml},
				Coordinate{Latitude: lat, Longitude: long}})
	}

	if err := rows.Err(); err != nil {
		return []Location{}, err
	}

	return locations, nil
}

func ResolveLocation(config Config, location Location) (Location, error) {
	if IsMissingCoordinate(location) {
		street := location.Address.Street
		if street != "" {
			locs, err := GetFuzzyAddress(config, street, 1)
			if err != nil {
				return Location{}, err
			}

			if len(locs) == 0 {
				e := errors.New("Couldn't find any address for street " + street)
				return Location{}, e
			}

			location.Coordinate = locs[0].Coordinate
			return location, nil
		} else {
			return Location{}, errors.New("Street empty, cannot search.")
		}
	} else {
		if location.Address.Country == "" {
			return Location{}, errors.New("Must provide country in Location.Address!")
		}

		coord := location.Coordinate
		correctCoord, err := CorrectPoint(config,
			Coord{coord.Latitude, coord.Longitude}, location.Address.Country)

		if err != nil {
			return Location{}, err
		}

		location.Coordinate.Latitude = correctCoord.Coord[0]
		location.Coordinate.Longitude = correctCoord.Coord[1]

		return location, nil
	}
}
<|endoftext|>