<commit_before>/*
   Copyright 2013 Juliano Martinez <juliano@martinez.io>

   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.

   @author: Juliano Martinez
*/

package main

import (
	"fmt"
	docopt "github.com/docopt/docopt.go"
	"github.com/fiorix/go-redis/redis"
	hpr_utils "github.com/ncode/hot-potato-router/utils"
	"strconv"
	"strings"
)

var (
	cfg = hpr_utils.NewConfig()
	rc  = redis.New(cfg.Options["redis"]["server_list"])
)

func main() {
	usage := `Hot Potato Router Control.

Usage:
  hprctl add <vhost> <backend_ip:port> [--weight=<n>]
  hprctl del <vhost> <backend_ip:port> [--weight=<n>]
  hprctl show <vhost>
  hprctl list

Args:
  add           add a new vhost and backend
  dell          del a vhost and a backend
  show          show all backends from a given vhost
  list          list all vhosts

Options:
  -h --help     Show this screen.
  --version     Show version.
  --weight=<n>  Weight in wrr [default: 1].
`

	arguments, _ := docopt.Parse(usage, nil, true, "Hot Potato Router 0.3.0", false)
	if arguments["add"] == true {
		_, err := rc.ZAdd(fmt.Sprintf("hpr-backends::%s", arguments["<vhost>"]), arguments["--weight"], arguments["<backend_ip:port>"])
		hpr_utils.CheckPanic(err, "Unable to write on hpr database")
		return
	}

	/* if arguments["del"] == true {
		_, err := rc.ZRem(fmt.Sprintf("hpr-backends::%s", arguments["<vhost>"]), arguments["--weight"], arguments["<backend_ip:port>"])
		hpr_utils.CheckPanic(err, "Unable to write on hpr database")
		return
	} */

	if arguments["show"] == true {
		bes, err := rc.ZRange(fmt.Sprintf("hpr-backends::%s", arguments["<vhost>"]), 0, -1, true)
		hpr_utils.CheckPanic(err, "Unable to write on hpr database")
		fmt.Printf(":: vhost [ %s ]\n", arguments["<vhost>"])
		var url string
		for _, be := range bes {
			count, err := strconv.Atoi(be)
			if err != nil {
				url = be
				continue
			}
			fmt.Printf("-- backend %s weight=%d\n", url, count)
		}
		return
	}

	if arguments["list"] == true {
		keys, err := rc.Keys("hpr-backends::*")
		hpr_utils.CheckPanic(err, "Unable to write on hpr database")
		for _, k := range keys {
			fmt.Printf(":: vhost [ %s ]\n", strings.TrimLeft(k, "hpr-backends::"))
		}
		return
	}
}
<commit_msg>del working :)<commit_after>/*
   Copyright 2013 Juliano Martinez <juliano@martinez.io>

   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.

   @author: Juliano Martinez
*/

package main

import (
	"fmt"
	docopt "github.com/docopt/docopt.go"
	"github.com/fiorix/go-redis/redis"
	hpr_utils "github.com/ncode/hot-potato-router/utils"
	"strconv"
	"strings"
)

var (
	cfg = hpr_utils.NewConfig()
	rc  = redis.New(cfg.Options["redis"]["server_list"])
)

func main() {
	usage := `Hot Potato Router Control.

Usage:
  hprctl add <vhost> <backend_ip:port> [--weight=<n>]
  hprctl del <vhost> <backend_ip:port> [--weight=<n>]
  hprctl show <vhost>
  hprctl list

Args:
  add           add a new vhost and backend
  dell          del a vhost and a backend
  show          show all backends from a given vhost
  list          list all vhosts

Options:
  -h --help     Show this screen.
  --version     Show version.
  --weight=<n>  Weight in wrr [default: 1].
`

	arguments, _ := docopt.Parse(usage, nil, true, "Hot Potato Router 0.3.0", false)
	if arguments["add"] == true {
		_, err := rc.ZAdd(fmt.Sprintf("hpr-backends::%s", arguments["<vhost>"]), arguments["--weight"], arguments["<backend_ip:port>"])
		hpr_utils.CheckPanic(err, "Unable to write on hpr database")
		return
	}

	if arguments["del"] == true {
		_, err := rc.ZRem(fmt.Sprintf("hpr-backends::%s", arguments["<vhost>"]), arguments["--weight"], arguments["<backend_ip:port>"])
		hpr_utils.CheckPanic(err, "Unable to write on hpr database")
		return
	}

	if arguments["show"] == true {
		bes, err := rc.ZRange(fmt.Sprintf("hpr-backends::%s", arguments["<vhost>"]), 0, -1, true)
		hpr_utils.CheckPanic(err, "Unable to write on hpr database")
		fmt.Printf(":: vhost [ %s ]\n", arguments["<vhost>"])
		var url string
		for _, be := range bes {
			count, err := strconv.Atoi(be)
			if err != nil {
				url = be
				continue
			}
			fmt.Printf("-- backend %s weight=%d\n", url, count)
		}
		return
	}

	if arguments["list"] == true {
		keys, err := rc.Keys("hpr-backends::*")
		hpr_utils.CheckPanic(err, "Unable to write on hpr database")
		for _, k := range keys {
			fmt.Printf(":: vhost [ %s ]\n", strings.TrimLeft(k, "hpr-backends::"))
		}
		return
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"net/http"
)

func handler(w http.ResponseWriter, r *http.Request) {
	w.WriteHeader(http.StatusGone)
}

func main() {
	fmt.Printf("Starting http-sinkhole at 8080 port")
	http.HandleFunc("/", handler)
	http.ListenAndServe(":8080", nil)
}
<commit_msg>Added logging capabilities<commit_after>package main

import (
	"fmt"
	"gopkg.in/natefinch/lumberjack.v2"
	"log"
	"net/http"
	"os"
)

func handler(w http.ResponseWriter, r *http.Request) {
	w.WriteHeader(http.StatusGone)
	log.Printf("Referer: %s\n", r.Header.Get("Referer"))
}

func main() {
	log.SetOutput(&lumberjack.Logger{
		Filename:   os.Args[1],
		MaxSize:    250, // megabytes
		MaxBackups: 3,
		MaxAge:     1, //days
	})
	fmt.Printf("Starting http-sinkhole at 8080 port")
	http.HandleFunc("/", handler)
	http.ListenAndServe(":8080", nil)
}
<|endoftext|>


<commit_before>package image

import (
	gocontext "context"
	"fmt"
	"net/url"
	"runtime"

	lxd "github.com/lxc/lxd/client"
	lxdapi "github.com/lxc/lxd/shared/api"
	"github.com/sirupsen/logrus"
	"github.com/travis-ci/worker/context"
	"golang.org/x/sync/errgroup"
)

var (
	arch  = runtime.GOARCH
	infra = fmt.Sprintf("lxd-%s", arch)

	tags   = []string{"os:linux"}
	groups = []string{"group:stable", "group:edge", "group:dev"}
)

func NewManager(ctx gocontext.Context, selector *APISelector, imagesBaseURL *url.URL) (*Manager, error) {
	logger := context.LoggerFromContext(ctx).WithField("self", "image_manager")

	client, err := lxd.ConnectLXDUnix("", nil)
	if err != nil {
		return nil, err
	}

	return &Manager{
		client:        client,
		selector:      selector,
		logger:        logger,
		imagesBaseURL: imagesBaseURL,
	}, nil
}

type Manager struct {
	client        lxd.ContainerServer
	selector      *APISelector
	logger        *logrus.Entry
	imagesBaseURL *url.URL
}

func (m *Manager) Load(imageName string) error {
	ok, err := m.Exists(imageName)
	if err != nil {
		return err
	}

	if ok {
		m.logger.WithFields(logrus.Fields{
			"image_name": imageName,
		}).Info("image already present")

		return nil
	}

	imageURL := m.imageUrl(imageName)

	m.logger.WithFields(logrus.Fields{
		"image_name": imageName,
		"image_url":  imageURL,
	}).Info("importing image")

	return m.importImage(imageName, imageURL)
}

func (m *Manager) Exists(imageName string) (bool, error) {
	images, err := m.client.GetImages()
	if err != nil {
		return false, err
	}

	for _, img := range images {
		for _, alias := range img.Aliases {
			if alias.Name == imageName {
				return true, nil
			}
		}
	}

	return false, nil
}

func (m *Manager) Update(ctx gocontext.Context) error {
	m.logger.WithFields(logrus.Fields{
		"arch":  arch,
		"infra": infra,
	}).Info("updating lxc images")

	images := []*apiSelectorImageRef{}

	for _, group := range groups {
		imagesGroup, err := m.selector.SelectAll(ctx, infra, append(tags, group))
		if err != nil {
			return err
		}

		images = append(images, imagesGroup...)
	}

	g, _ := errgroup.WithContext(ctx)
	for _, img := range images {
		img := img

		g.Go(func() error {
			m.logger.WithFields(logrus.Fields{
				"image_name": img.Name,
				"infra":      infra,
				"group":      img.Group(),
			}).Info("updating image")

			return m.Load(img.Name)

		})

	}

	return g.Wait()
}

func (m *Manager) Cleanup() error {
	// TODO

	return nil
}

func (m *Manager) importImage(imageName, imgURL string) error {
	op, err := m.client.CreateImage(
		lxdapi.ImagesPost{
			Filename: imageName,
			Source: &lxdapi.ImagesPostSource{
				Type: "url",
				URL:  imgURL,
			},
			Aliases: []lxdapi.ImageAlias{
				lxdapi.ImageAlias{Name: imageName},
			},
		}, nil,
	)
	if err != nil {
		return err
	}

	err = op.Wait()
	if err != nil {
		return err
	}

	alias, _, err := m.client.GetImageAlias(imageName)
	if err != nil {
		return err
	}

	image, _, err := m.client.GetImage(alias.Target)
	if err != nil {
		return err
	}

	containerName := fmt.Sprintf("%s-warmup", imageName)

	rop, err := m.client.CreateContainerFromImage(
		m.client,
		*image,
		lxdapi.ContainersPost{
			Name: containerName,
		},
	)
	if err != nil {
		return err
	}

	err = rop.Wait()
	if err != nil {
		return err
	}

	defer func() {
		op, err := m.client.DeleteContainer(containerName)
		if err != nil {
			m.logger.WithFields(logrus.Fields{
				"error":          err,
				"container_name": containerName,
			}).Error("failed to delete container")
			return
		}

		err = op.Wait()
		if err != nil {
			m.logger.WithFields(logrus.Fields{
				"error":          err,
				"container_name": containerName,
			}).Error("failed to delete container")
			return
		}
	}()

	return err
}

func (m *Manager) imageUrl(name string) string {
	u := *m.imagesBaseURL
	u.Path = fmt.Sprintf("/images/%s", name)
	return u.String()
}
<commit_msg>Update LXD images server path<commit_after>package image

import (
	gocontext "context"
	"fmt"
	"net/url"
	"runtime"

	lxd "github.com/lxc/lxd/client"
	lxdapi "github.com/lxc/lxd/shared/api"
	"github.com/sirupsen/logrus"
	"github.com/travis-ci/worker/context"
	"golang.org/x/sync/errgroup"
)

var (
	arch  = runtime.GOARCH
	infra = fmt.Sprintf("lxd-%s", arch)

	tags   = []string{"os:linux"}
	groups = []string{"group:stable", "group:edge", "group:dev"}
)

func NewManager(ctx gocontext.Context, selector *APISelector, imagesServerURL *url.URL) (*Manager, error) {
	logger := context.LoggerFromContext(ctx).WithField("self", "image_manager")

	client, err := lxd.ConnectLXDUnix("", nil)
	if err != nil {
		return nil, err
	}

	return &Manager{
		client:          client,
		selector:        selector,
		logger:          logger,
		imagesServerURL: imagesServerURL,
	}, nil
}

type Manager struct {
	client          lxd.ContainerServer
	selector        *APISelector
	logger          *logrus.Entry
	imagesServerURL *url.URL
}

func (m *Manager) Load(imageName string) error {
	ok, err := m.Exists(imageName)
	if err != nil {
		return err
	}

	if ok {
		m.logger.WithFields(logrus.Fields{
			"image_name": imageName,
		}).Info("image already present")

		return nil
	}

	imageURL := m.imageUrl(imageName)

	m.logger.WithFields(logrus.Fields{
		"image_name": imageName,
		"image_url":  imageURL,
	}).Info("importing image")

	return m.importImage(imageName, imageURL)
}

func (m *Manager) Exists(imageName string) (bool, error) {
	images, err := m.client.GetImages()
	if err != nil {
		return false, err
	}

	for _, img := range images {
		for _, alias := range img.Aliases {
			if alias.Name == imageName {
				return true, nil
			}
		}
	}

	return false, nil
}

func (m *Manager) Update(ctx gocontext.Context) error {
	m.logger.WithFields(logrus.Fields{
		"arch":  arch,
		"infra": infra,
	}).Info("updating lxc images")

	images := []*apiSelectorImageRef{}

	for _, group := range groups {
		imagesGroup, err := m.selector.SelectAll(ctx, infra, append(tags, group))
		if err != nil {
			return err
		}

		images = append(images, imagesGroup...)
	}

	g, _ := errgroup.WithContext(ctx)
	for _, img := range images {
		img := img

		g.Go(func() error {
			m.logger.WithFields(logrus.Fields{
				"image_name": img.Name,
				"infra":      infra,
				"group":      img.Group(),
			}).Info("updating image")

			return m.Load(img.Name)

		})

	}

	return g.Wait()
}

func (m *Manager) Cleanup() error {
	// TODO

	return nil
}

func (m *Manager) importImage(imageName, imgURL string) error {
	op, err := m.client.CreateImage(
		lxdapi.ImagesPost{
			Filename: imageName,
			Source: &lxdapi.ImagesPostSource{
				Type: "url",
				URL:  imgURL,
			},
			Aliases: []lxdapi.ImageAlias{
				lxdapi.ImageAlias{Name: imageName},
			},
		}, nil,
	)
	if err != nil {
		return err
	}

	err = op.Wait()
	if err != nil {
		return err
	}

	alias, _, err := m.client.GetImageAlias(imageName)
	if err != nil {
		return err
	}

	image, _, err := m.client.GetImage(alias.Target)
	if err != nil {
		return err
	}

	containerName := fmt.Sprintf("%s-warmup", imageName)

	rop, err := m.client.CreateContainerFromImage(
		m.client,
		*image,
		lxdapi.ContainersPost{
			Name: containerName,
		},
	)
	if err != nil {
		return err
	}

	err = rop.Wait()
	if err != nil {
		return err
	}

	defer func() {
		op, err := m.client.DeleteContainer(containerName)
		if err != nil {
			m.logger.WithFields(logrus.Fields{
				"error":          err,
				"container_name": containerName,
			}).Error("failed to delete container")
			return
		}

		err = op.Wait()
		if err != nil {
			m.logger.WithFields(logrus.Fields{
				"error":          err,
				"container_name": containerName,
			}).Error("failed to delete container")
			return
		}
	}()

	return err
}

func (m *Manager) imageUrl(name string) string {
	u := *m.imagesServerURL
	u.Path = fmt.Sprintf("/images/travis/%s", name)
	return u.String()
}
<|endoftext|>


<commit_before>// Copyright 2015, David Howden
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package index defines functionality for creating and manipulating a Tchaik music index.
package index

import (
	"compress/gzip"
	"encoding/json"
	"fmt"
	"io"
	"strconv"
	"time"
)

// Library represents the tchaik music library.  Currently we don't have anything
// better than wrapping around the iTunes library (oh the shame!)
type Library interface {
	// Tracks returns a slice of all the tracks in the library.
	Tracks() []Track

	// Track returns the track from the given identifier, second return value true
	// if successful.
	Track(identifier string) (Track, bool)
}

// Track is an interface which defines a music file.
type Track interface {
	// GetString returns the string attribute with given name.
	GetString(string) string
	// GetInt returns the int attribute with given name.
	GetInt(string) int
	// GetTime returns the time attribute with given name.
	GetTime(string) time.Time
}

var trackStringFields = []string{
	"TrackID", // unique identifier for the track.
	"Name",
	"Album",
	"AlbumArtist",
	"Artist",
	"Composer",
	"Location", // location of the associated file
}

var trackIntFields = []string{
	"TotalTime",
	"Year",
	"DiscNumber",
	"TrackNumber",
	"TrackCount",
	"DiscCount",
}

var trackTimeFields = []string{
	"DateAdded",
	"DateModified",
}

// Convert reads all the data exported by the Library and writes into the standard
// tchaik Library implementation.
// NB: The identifier field is set to be the value of TrackID on every track, regardless
// of whether this value has already been set in the input Library.
func Convert(l Library, id string) *library {
	allTracks := l.Tracks()
	tracks := make(map[string]*track, len(allTracks))

	n := len(trackStringFields) + len(trackIntFields)
	for _, t := range allTracks {
		m := make(map[string]interface{}, n)
		for _, f := range trackStringFields {
			m[f] = t.GetString(f)
		}
		for _, f := range trackIntFields {
			m[f] = t.GetInt(f)
		}
		for _, f := range trackTimeFields {
			m[f] = t.GetTime(f)
		}

		identifier := t.GetString(id)
		m["TrackID"] = identifier
		tracks[identifier] = &track{
			flds: m,
		}
	}
	return &library{
		tracks,
	}
}

// library is the default internal implementation Library which acts as the data
// source for all media tracks.
type library struct {
	Trks map[string]*track // NB: Exported so that we can easily encode
}

// Tracks implements Library.
func (l *library) Tracks() []Track {
	tracks := make([]Track, 0, len(l.Trks))
	for _, t := range l.Trks {
		tracks = append(tracks, t)
	}
	return tracks
}

// Track implements Library.
func (l *library) Track(id string) (Track, bool) {
	t, ok := l.Trks[id]
	return t, ok
}

// WriteTo writes the Library data to the writer, currently using gzipped-JSON.
func WriteTo(l Library, w io.Writer) error {
	gzw, err := gzip.NewWriterLevel(w, gzip.BestSpeed)
	if err != nil {
		return err
	}
	defer gzw.Close()
	enc := json.NewEncoder(gzw)
	return enc.Encode(l)
}

// ReadFrom reads the gzipped-JSON representation of a Library.
func ReadFrom(r io.Reader) (Library, error) {
	gzr, err := gzip.NewReader(r)
	if err != nil {
		return nil, err
	}
	defer gzr.Close()
	dec := json.NewDecoder(gzr)
	l := &library{}
	err = dec.Decode(l)
	return l, err
}

// track is the default implementation of the Track interface.
type track struct {
	flds map[string]interface{} // NB: Exported so that we can easily encode.
}

func (t *track) MarshalJSON() ([]byte, error) {
	return json.Marshal(t.flds)
}

func (t *track) UnmarshalJSON(b []byte) error {
	t.flds = make(map[string]interface{})
	return json.Unmarshal(b, &t.flds)
}

// GetString implements Track.
func (t *track) GetString(name string) string {
	x, ok := t.flds[name]
	if !ok {
		panic(fmt.Sprintf("unknown string field '%v'", name))
	}
	if x == nil {
		panic(fmt.Sprintf("<nil> string field '%v'", name))
	}
	s, ok := x.(string)
	if !ok {
		panic(fmt.Sprintf("field '%v': expected string, got %#v (%T)", name, x, x))
	}
	return s
}

// GetInt implements Track.
func (t *track) GetInt(name string) int {
	x, ok := t.flds[name]
	if !ok {
		panic(fmt.Sprintf("unknown int field '%v'", name))
	}
	if x == nil {
		panic(fmt.Sprintf("<nil> int field '%v'", name))
	}
	switch x := x.(type) {
	case int:
		return x
	case float64:
		return int(x)
	case string:
		n, err := strconv.Atoi(x)
		if err != nil {
			panic(fmt.Sprintf("error converting string to int: %v", err))
		}
		return n
	}
	panic(fmt.Sprintf("unknown type '%T' for field '%v'", name, name))
}

// GetTime implements Track.
func (t *track) GetTime(name string) time.Time {
	x, ok := t.flds[name]
	if !ok {
		panic(fmt.Sprintf("unknown time field '%v'", name))
	}
	if x == nil {
		panic(fmt.Sprintf("<nil> time field '%v'", name))
	}
	if x, ok := x.(time.Time); ok {
		return x
	}
	panic(fmt.Sprintf("unknown type '%T' for field '%v'", x, name))
}
<commit_msg>Fix panic string for track GetInt<commit_after>// Copyright 2015, David Howden
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package index defines functionality for creating and manipulating a Tchaik music index.
package index

import (
	"compress/gzip"
	"encoding/json"
	"fmt"
	"io"
	"strconv"
	"time"
)

// Library represents the tchaik music library.  Currently we don't have anything
// better than wrapping around the iTunes library (oh the shame!)
type Library interface {
	// Tracks returns a slice of all the tracks in the library.
	Tracks() []Track

	// Track returns the track from the given identifier, second return value true
	// if successful.
	Track(identifier string) (Track, bool)
}

// Track is an interface which defines a music file.
type Track interface {
	// GetString returns the string attribute with given name.
	GetString(string) string
	// GetInt returns the int attribute with given name.
	GetInt(string) int
	// GetTime returns the time attribute with given name.
	GetTime(string) time.Time
}

var trackStringFields = []string{
	"TrackID", // unique identifier for the track.
	"Name",
	"Album",
	"AlbumArtist",
	"Artist",
	"Composer",
	"Location", // location of the associated file
}

var trackIntFields = []string{
	"TotalTime",
	"Year",
	"DiscNumber",
	"TrackNumber",
	"TrackCount",
	"DiscCount",
}

var trackTimeFields = []string{
	"DateAdded",
	"DateModified",
}

// Convert reads all the data exported by the Library and writes into the standard
// tchaik Library implementation.
// NB: The identifier field is set to be the value of TrackID on every track, regardless
// of whether this value has already been set in the input Library.
func Convert(l Library, id string) *library {
	allTracks := l.Tracks()
	tracks := make(map[string]*track, len(allTracks))

	n := len(trackStringFields) + len(trackIntFields)
	for _, t := range allTracks {
		m := make(map[string]interface{}, n)
		for _, f := range trackStringFields {
			m[f] = t.GetString(f)
		}
		for _, f := range trackIntFields {
			m[f] = t.GetInt(f)
		}
		for _, f := range trackTimeFields {
			m[f] = t.GetTime(f)
		}

		identifier := t.GetString(id)
		m["TrackID"] = identifier
		tracks[identifier] = &track{
			flds: m,
		}
	}
	return &library{
		tracks,
	}
}

// library is the default internal implementation Library which acts as the data
// source for all media tracks.
type library struct {
	Trks map[string]*track // NB: Exported so that we can easily encode
}

// Tracks implements Library.
func (l *library) Tracks() []Track {
	tracks := make([]Track, 0, len(l.Trks))
	for _, t := range l.Trks {
		tracks = append(tracks, t)
	}
	return tracks
}

// Track implements Library.
func (l *library) Track(id string) (Track, bool) {
	t, ok := l.Trks[id]
	return t, ok
}

// WriteTo writes the Library data to the writer, currently using gzipped-JSON.
func WriteTo(l Library, w io.Writer) error {
	gzw, err := gzip.NewWriterLevel(w, gzip.BestSpeed)
	if err != nil {
		return err
	}
	defer gzw.Close()
	enc := json.NewEncoder(gzw)
	return enc.Encode(l)
}

// ReadFrom reads the gzipped-JSON representation of a Library.
func ReadFrom(r io.Reader) (Library, error) {
	gzr, err := gzip.NewReader(r)
	if err != nil {
		return nil, err
	}
	defer gzr.Close()
	dec := json.NewDecoder(gzr)
	l := &library{}
	err = dec.Decode(l)
	return l, err
}

// track is the default implementation of the Track interface.
type track struct {
	flds map[string]interface{} // NB: Exported so that we can easily encode.
}

func (t *track) MarshalJSON() ([]byte, error) {
	return json.Marshal(t.flds)
}

func (t *track) UnmarshalJSON(b []byte) error {
	t.flds = make(map[string]interface{})
	return json.Unmarshal(b, &t.flds)
}

// GetString implements Track.
func (t *track) GetString(name string) string {
	x, ok := t.flds[name]
	if !ok {
		panic(fmt.Sprintf("unknown string field '%v'", name))
	}
	if x == nil {
		panic(fmt.Sprintf("<nil> string field '%v'", name))
	}
	s, ok := x.(string)
	if !ok {
		panic(fmt.Sprintf("field '%v': expected string, got %#v (%T)", name, x, x))
	}
	return s
}

// GetInt implements Track.
func (t *track) GetInt(name string) int {
	x, ok := t.flds[name]
	if !ok {
		panic(fmt.Sprintf("unknown int field '%v'", name))
	}
	if x == nil {
		panic(fmt.Sprintf("<nil> int field '%v'", name))
	}
	switch x := x.(type) {
	case int:
		return x
	case float64:
		return int(x)
	case string:
		n, err := strconv.Atoi(x)
		if err != nil {
			panic(fmt.Sprintf("error converting string to int: %v", err))
		}
		return n
	}
	panic(fmt.Sprintf("unknown type '%T' for field '%v'", x, name))
}

// GetTime implements Track.
func (t *track) GetTime(name string) time.Time {
	x, ok := t.flds[name]
	if !ok {
		panic(fmt.Sprintf("unknown time field '%v'", name))
	}
	if x == nil {
		panic(fmt.Sprintf("<nil> time field '%v'", name))
	}
	if x, ok := x.(time.Time); ok {
		return x
	}
	panic(fmt.Sprintf("unknown type '%T' for field '%v'", x, name))
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/binary"
	uc "github.com/unicorn-engine/unicorn/bindings/go/unicorn"
	"io/ioutil"
	"os"
	"unsafe"

	"github.com/lunixbochs/usercorn/go/cmd"
	"github.com/lunixbochs/usercorn/go/models"
)

/*
#include <stdlib.h>
#include <sys/shm.h>
#include <string.h>

void *afl_setup() {
	char *id = getenv("__AFL_SHM_ID");
	if (id == NULL) {
		return NULL;
	}
	void *afl_area = shmat(atoi(id), NULL, 0);
	if (afl_area == (void *)-1) {
		return NULL;
	}
	return afl_area;
}

*/
import "C"

var MAP_SIZE uint64 = 1 << 16
var FORKSRV_FD = 198
var aflHello = []byte{1, 2, 3, 4}

func main() {
	message := []byte("In fuzz main")
	ioutil.WriteFile("/tmp/outfile", message, 0444)

	forksrvCtrl := os.NewFile(uintptr(FORKSRV_FD), "afl_ctrl")
	forksrvStatus := os.NewFile(uintptr(FORKSRV_FD+1), "afl_status")

	c := cmd.NewUsercornCmd()
	var forkAddr *uint64
	var fuzzInterp *bool

	aflArea := C.afl_setup()
	if aflArea == nil {
		panic("could not set up AFL shared memory")
	}
	fuzzMap := []byte((*[1 << 30]byte)(unsafe.Pointer(aflArea))[:])

	var lastPos uint64
	blockTrace := func(_ uc.Unicorn, addr uint64, size uint32) {
		if lastPos == 0 {
			lastPos = addr >> 1
			return
		}
		loc := (addr >> 4) ^ (addr << 8)
		loc &= MAP_SIZE - 1
		fuzzMap[loc]++
		lastPos = addr >> 1
	}
	c.SetupFlags = func() error {
		forkAddr = c.Flags.Uint64("forkaddr", 0, "wait until this address to fork and begin fuzzing")
		fuzzInterp = c.Flags.Bool("fuzzinterp", false, "controls whether fuzzing is delayed until program's main entry point")
		return nil
	}
	c.SetupUsercorn = func() error {
		if _, err := c.Usercorn.HookAdd(uc.HOOK_BLOCK, blockTrace, 1, 0); err != nil {
			return err
		}
		return nil
	}
	c.RunUsercorn = func(args, env []string) error {
		u := c.Usercorn
		u.Println("Saving Usercorn state...")

		var savedRegEnums []int
		for _, enum := range u.Arch().Regs {
			savedRegEnums = append(savedRegEnums, enum)
		}
		savedRegs, _ := u.RegReadBatch(savedRegEnums)

		type serialMem struct {
			Addr, Size uint64
			Prot       int
			Data       []byte
		}
		var savedMem []serialMem
		for _, m := range u.Mappings() {
			mem, err := u.MemRead(m.Addr, m.Size)
			if err != nil {
				u.Printf("Warning: error saving memory at 0x%x-0x%x: %s\n", m.Addr, m.Addr+m.Size, err)
				continue
			}
			savedMem = append(savedMem, serialMem{
				Addr: m.Addr, Size: m.Size, Prot: m.Prot, Data: mem,
			})
		}

		u.Println("Starting Usercorn")
		if _, err := forksrvStatus.Write(aflHello); err != nil {
			u.Println("AFL hello failed.")
			return err
		}
		var aflMsg [4]byte
		// afl forkserver loop
		for {
			lastPos = 0
			if _, err := forksrvCtrl.Read(aflMsg[:]); err != nil {
				u.Printf("Failed to receive control signal from AFL: %s\n", err)
				return err
			}
			u.Println("AFL requested new child")

			// spawn a fake child so AFL has something other than us to kill
			// monitor it and if afl kills it, stop the current emulation
			args := []string{"/bin/cat"}
			var procAttr os.ProcAttr
			proc, err := os.StartProcess(args[0], args, &procAttr)
			if err != nil {
				u.Printf("Failed to spawn child: %s\n", err)
				return err
			}

			// restore register and memory state
			u.RegWriteBatch(savedRegEnums, savedRegs)
			for _, m := range u.Mappings() {
				u.MemUnmap(m.Addr, m.Size)
			}
			for _, m := range savedMem {
				u.MemMapProt(m.Addr, m.Size, m.Prot)
				u.MemWrite(m.Addr, m.Data)
			}

			binary.LittleEndian.PutUint32(aflMsg[:], uint32(proc.Pid))
			if _, err := forksrvStatus.Write(aflMsg[:]); err != nil {
				u.Printf("Failed to send pid to AFL: %s\n", err)
				return err
			}

			// Goroutine to stop usercorn if afl-fuzz kills our fake process
			go func() {
				proc.Wait()
				u.Stop()
				u.Println("Child+Usercorn stopped")
			}()

			status := 0
			err = u.Run(args, env)
			if _, ok := err.(models.ExitStatus); ok {
			} else if err != nil {
				u.Printf("Usercorn err: %s\n", err)
				status = 257
			}
			binary.LittleEndian.PutUint32(aflMsg[:], uint32(status))
			if _, err := forksrvStatus.Write(aflMsg[:]); err != nil {
				u.Printf("Failed to send status to AFL: %s\n", err)
				return err
			}

			proc.Kill()
			proc.Wait()
		}
	}

	c.Run(os.Args, os.Environ())
}
<commit_msg>don't print during normal fuzz runs<commit_after>package main

import (
	"encoding/binary"
	uc "github.com/unicorn-engine/unicorn/bindings/go/unicorn"
	"io/ioutil"
	"os"
	"unsafe"

	"github.com/lunixbochs/usercorn/go/cmd"
	"github.com/lunixbochs/usercorn/go/models"
)

/*
#include <stdlib.h>
#include <sys/shm.h>
#include <string.h>

void *afl_setup() {
	char *id = getenv("__AFL_SHM_ID");
	if (id == NULL) {
		return NULL;
	}
	void *afl_area = shmat(atoi(id), NULL, 0);
	if (afl_area == (void *)-1) {
		return NULL;
	}
	return afl_area;
}

*/
import "C"

var MAP_SIZE uint64 = 1 << 16
var FORKSRV_FD = 198
var aflHello = []byte{1, 2, 3, 4}

func main() {
	message := []byte("In fuzz main")
	ioutil.WriteFile("/tmp/outfile", message, 0444)

	forksrvCtrl := os.NewFile(uintptr(FORKSRV_FD), "afl_ctrl")
	forksrvStatus := os.NewFile(uintptr(FORKSRV_FD+1), "afl_status")

	c := cmd.NewUsercornCmd()
	var forkAddr *uint64
	var fuzzInterp *bool

	aflArea := C.afl_setup()
	if aflArea == nil {
		panic("could not set up AFL shared memory")
	}
	fuzzMap := []byte((*[1 << 30]byte)(unsafe.Pointer(aflArea))[:])

	var lastPos uint64
	blockTrace := func(_ uc.Unicorn, addr uint64, size uint32) {
		if lastPos == 0 {
			lastPos = addr >> 1
			return
		}
		loc := (addr >> 4) ^ (addr << 8)
		loc &= MAP_SIZE - 1
		fuzzMap[loc]++
		lastPos = addr >> 1
	}
	c.SetupFlags = func() error {
		forkAddr = c.Flags.Uint64("forkaddr", 0, "wait until this address to fork and begin fuzzing")
		fuzzInterp = c.Flags.Bool("fuzzinterp", false, "controls whether fuzzing is delayed until program's main entry point")
		return nil
	}
	c.SetupUsercorn = func() error {
		if _, err := c.Usercorn.HookAdd(uc.HOOK_BLOCK, blockTrace, 1, 0); err != nil {
			return err
		}
		return nil
	}
	c.RunUsercorn = func(args, env []string) error {
		u := c.Usercorn
		u.Println("Saving Usercorn state...")

		var savedRegEnums []int
		for _, enum := range u.Arch().Regs {
			savedRegEnums = append(savedRegEnums, enum)
		}
		savedRegs, _ := u.RegReadBatch(savedRegEnums)

		type serialMem struct {
			Addr, Size uint64
			Prot       int
			Data       []byte
		}
		var savedMem []serialMem
		for _, m := range u.Mappings() {
			mem, err := u.MemRead(m.Addr, m.Size)
			if err != nil {
				u.Printf("Warning: error saving memory at 0x%x-0x%x: %s\n", m.Addr, m.Addr+m.Size, err)
				continue
			}
			savedMem = append(savedMem, serialMem{
				Addr: m.Addr, Size: m.Size, Prot: m.Prot, Data: mem,
			})
		}

		u.Println("Starting Usercorn")
		if _, err := forksrvStatus.Write(aflHello); err != nil {
			u.Println("AFL hello failed.")
			return err
		}
		var aflMsg [4]byte
		// afl forkserver loop
		for {
			lastPos = 0
			if _, err := forksrvCtrl.Read(aflMsg[:]); err != nil {
				u.Printf("Failed to receive control signal from AFL: %s\n", err)
				return err
			}

			// spawn a fake child so AFL has something other than us to kill
			// monitor it and if afl kills it, stop the current emulation
			args := []string{"/bin/cat"}
			var procAttr os.ProcAttr
			proc, err := os.StartProcess(args[0], args, &procAttr)
			if err != nil {
				u.Printf("Failed to spawn child: %s\n", err)
				return err
			}

			// restore register and memory state
			u.RegWriteBatch(savedRegEnums, savedRegs)
			for _, m := range u.Mappings() {
				u.MemUnmap(m.Addr, m.Size)
			}
			for _, m := range savedMem {
				u.MemMapProt(m.Addr, m.Size, m.Prot)
				u.MemWrite(m.Addr, m.Data)
			}

			binary.LittleEndian.PutUint32(aflMsg[:], uint32(proc.Pid))
			if _, err := forksrvStatus.Write(aflMsg[:]); err != nil {
				u.Printf("Failed to send pid to AFL: %s\n", err)
				return err
			}

			// Goroutine to stop usercorn if afl-fuzz kills our fake process
			go func() {
				proc.Wait()
				u.Stop()
			}()

			status := 0
			err = u.Run(args, env)
			if _, ok := err.(models.ExitStatus); ok {
			} else if err != nil {
				u.Printf("Usercorn err: %s\n", err)
				status = 257
			}
			binary.LittleEndian.PutUint32(aflMsg[:], uint32(status))
			if _, err := forksrvStatus.Write(aflMsg[:]); err != nil {
				u.Printf("Failed to send status to AFL: %s\n", err)
				return err
			}

			proc.Kill()
			proc.Wait()
		}
	}

	c.Run(os.Args, os.Environ())
}
<|endoftext|>


<commit_before>package repl

var cmdRc = `
func _fallback(name, val)
    if type(val) == 'function' then
        val()
    else
        print val
    end
end

local _hook_types = {
	-- TODO: hook sys, map?
	code = cpu.HOOK_CODE,
	block = cpu.HOOK_BLOCK,
	intr = cpu.HOOK_INTR,
	read = cpu.HOOK_MEM_READ,
	write = cpu.HOOK_MEM_WRITE,
	fault = cpu.HOOK_MEM_ERR,
}

func _on_hook(name, fn, start, stop)
	local type = _hook_types[name]
	if type != nil then
		if start == nil then return u.hook_add(type, fn)
		else return u.hook_add(type, fn, start, stop) end
	end
	print 'unimplemented hook type %s' % name
    -- u.on_hook_add(name, fn, start, stop)
end

func off()
	if hh then
		u.hook_del(hh)
		hh = nil
	end
end

func read(addr, size)
	if size == nil then size = 16 end
	return u.mem_read(addr, size)
end

func write(addr, s)
	u.mem_write(addr, s)
end

func maps()
	print 'Memory map:'
	local m = us:Mappings()
	for i in m() do
	print m[i]:String()
	end
end

func map(addr, size, prot)
    if addr == nil then
		maps()
    else
        if prot == nil then prot = cpu.PROT_ALL end
        u.mem_map(addr, size, prot)
    end
end

func asm(src, addr)
	if addr == nil then addr = 0 end
	return u.asm(src, addr)
end

func patch(addr, src)
	local x = asm(src, addr)
	-- autonop
	local total = 0
	for _, ins in pairs(u.dis(addr, x:len() + 16)) do
		total = total + ins.bytes:len()
		if total == x:len() then break end
		if total > x:len() then
			x = x .. asm('nop', 0):rep(total - x:len())
			break
		end
	end
	write(addr, x)
end

func dis(addr, size, indent)
	if addr == nil then addr = pc end
	if size == nil then size = 16 end
	if indent == nil then indent = 0 end
	local pad = (' '):rep(indent)
	local d = u.dis(addr, size)
	local width = 0
	for _, ins in ipairs(d) do width = math.max(width, ins.bytes:len()) end
	for _, ins in ipairs(d) do
		local fmt = pad .. '0x%x: %' .. (width * 2) .. 'x %s %s'
		print fmt % {ins.addr, ins.bytes, ins.name, ins.op_str}
	end
end

func c() u.continue() end

func s(steps)
	if steps == nil then steps = 1 end
	u.step(steps)
end

func b(baddr)
	local hh = u.hook_add(cpu.HOOK_CODE, func()
		if bskip == baddr then
			bskip = nil
			return
		end
		bskip = baddr
		print 'Breakpoint N hit at 0x%x' % baddr
		u.stop()
	end, baddr, baddr)
	print 'Breakpoint N at 0x%x' % baddr
end

func runto(name)
    on code do
        if ins.name == name then
            print '[-] stopping at "%s"' % name
            dis
            u.stop()
            off
        end
    end
    c
end
`
<commit_msg>don't print bytes in dis<commit_after>package repl

var cmdRc = `
func _fallback(name, val)
    if type(val) == 'function' then
        val()
    else
        print val
    end
end

local _hook_types = {
	-- TODO: hook sys, map?
	code = cpu.HOOK_CODE,
	block = cpu.HOOK_BLOCK,
	intr = cpu.HOOK_INTR,
	read = cpu.HOOK_MEM_READ,
	write = cpu.HOOK_MEM_WRITE,
	fault = cpu.HOOK_MEM_ERR,
}

func _on_hook(name, fn, start, stop)
	local type = _hook_types[name]
	if type != nil then
		if start == nil then return u.hook_add(type, fn)
		else return u.hook_add(type, fn, start, stop) end
	end
	print 'unimplemented hook type %s' % name
    -- u.on_hook_add(name, fn, start, stop)
end

func off()
	if hh then
		u.hook_del(hh)
		hh = nil
	end
end

func read(addr, size)
	if size == nil then size = 16 end
	return u.mem_read(addr, size)
end

func write(addr, s)
	u.mem_write(addr, s)
end

func maps()
	print 'Memory map:'
	local m = us:Mappings()
	for i in m() do
	print m[i]:String()
	end
end

func map(addr, size, prot)
    if addr == nil then
		maps()
    else
        if prot == nil then prot = cpu.PROT_ALL end
        u.mem_map(addr, size, prot)
    end
end

func asm(src, addr)
	if addr == nil then addr = 0 end
	return u.asm(src, addr)
end

func patch(addr, src)
	local x = asm(src, addr)
	-- autonop
	local total = 0
	for _, ins in pairs(u.dis(addr, x:len() + 16)) do
		total = total + ins.bytes:len()
		if total == x:len() then break end
		if total > x:len() then
			x = x .. asm('nop', 0):rep(total - x:len())
			break
		end
	end
	write(addr, x)
end

func dis(addr, size, indent)
	if addr == nil then addr = pc end
	if size == nil then size = 16 end
	if indent == nil then indent = 0 end
	local pad = (' '):rep(indent)
	local d = u.dis(addr, size)
	local width = 0
	for _, ins in ipairs(d) do
		local fmt = pad .. '0x%x: %s %s'
		print fmt % {ins.addr, ins.name, ins.op_str}
	end
end

func c() u.continue() end

func s(steps)
	if steps == nil then steps = 1 end
	u.step(steps)
end

func b(baddr)
	local hh = u.hook_add(cpu.HOOK_CODE, func()
		if bskip == baddr then
			bskip = nil
			return
		end
		bskip = baddr
		print 'Breakpoint N hit at 0x%x' % baddr
		u.stop()
	end, baddr, baddr)
	print 'Breakpoint N at 0x%x' % baddr
end

func runto(name)
    on code do
        if ins.name == name then
            print '[-] stopping at "%s"' % name
            dis
            u.stop()
            off
        end
    end
    c
end
`
<|endoftext|>


<commit_before>package velox

import (
	"encoding/json"
	"errors"
	"net/http"
	"strconv"
	"time"

	"github.com/jpillora/eventsource"
)

type eventSourceTransport struct {
	writeTimeout time.Duration
	w            http.ResponseWriter
	isConnected  bool
	connected    chan struct{}
}

func (es *eventSourceTransport) connect(w http.ResponseWriter, r *http.Request) error {
	//connection controls
	es.isConnected = true
	es.connected = make(chan struct{})
	go func() {
		select {
		case <-es.connected:
		case <-r.Context().Done(): //client disconnected early
			es.close()
		}
	}()
	//eventsource headers
	w.Header().Set("Cache-Control", "no-cache")
	w.Header().Set("Vary", "Accept")
	w.Header().Set("Content-Type", "text/event-stream")
	//connection is now expecting a stream of events
	es.w = w
	return nil
}

func (es *eventSourceTransport) send(upd *update) error {
	b, err := json.Marshal(upd)
	if err != nil {
		return err
	}
	sent := make(chan error)
	go func() {
		if es.isConnected {
			err := eventsource.WriteEvent(es.w, eventsource.Event{
				ID:   strconv.FormatInt(upd.Version, 10),
				Data: b,
			})
			if f, ok := es.w.(http.Flusher); ok {
				f.Flush()
			}
			sent <- err
		}
	}()
	select {
	case <-time.After(es.writeTimeout):
		return errors.New("timeout")
	case err := <-sent:
		return err
	}
}

func (es *eventSourceTransport) wait() error {
	<-es.connected
	return nil
}

func (es *eventSourceTransport) close() error {
	if es.isConnected {
		//unblocking the wait, causes the http handler to return
		close(es.connected)
		es.isConnected = false
	}
	return nil
}
<commit_msg>fix: remove unnecessary flush<commit_after>package velox

import (
	"encoding/json"
	"errors"
	"net/http"
	"strconv"
	"time"

	"github.com/jpillora/eventsource"
)

type eventSourceTransport struct {
	writeTimeout time.Duration
	w            http.ResponseWriter
	isConnected  bool
	connected    chan struct{}
}

func (es *eventSourceTransport) connect(w http.ResponseWriter, r *http.Request) error {
	//connection controls
	es.isConnected = true
	es.connected = make(chan struct{})
	go func() {
		select {
		case <-es.connected:
		case <-r.Context().Done(): //client disconnected early
			es.close()
		}
	}()
	//eventsource headers
	w.Header().Set("Cache-Control", "no-cache")
	w.Header().Set("Vary", "Accept")
	w.Header().Set("Content-Type", "text/event-stream")
	//connection is now expecting a stream of events
	es.w = w
	return nil
}

func (es *eventSourceTransport) send(upd *update) error {
	b, err := json.Marshal(upd)
	if err != nil {
		return err
	}
	sent := make(chan error)
	go func() {
		if es.isConnected {
			err := eventsource.WriteEvent(es.w, eventsource.Event{
				ID:   strconv.FormatInt(upd.Version, 10),
				Data: b,
			})
			sent <- err
		}
	}()
	select {
	case <-time.After(es.writeTimeout):
		return errors.New("timeout")
	case err := <-sent:
		return err
	}
}

func (es *eventSourceTransport) wait() error {
	<-es.connected
	return nil
}

func (es *eventSourceTransport) close() error {
	if es.isConnected {
		//unblocking the wait, causes the http handler to return
		close(es.connected)
		es.isConnected = false
	}
	return nil
}
<|endoftext|>


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

package zk

import (
	"log"
	"math/rand"
	"sync"
	"time"

	"launchpad.net/gozk/zookeeper"
)

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

/* When you need to talk to multiple zk cells, you need a simple
abstraction so you aren't caching clients all over the place.

ConnCache guarantees that you have at most one zookeeper connection per cell.
*/

type cachedConn struct {
	mutex sync.Mutex // used to notify if multiple goroutine simultaneously want a connection
	zconn Conn
}

type ConnCache struct {
	mutex        sync.Mutex
	zconnCellMap map[string]*cachedConn // map cell name to connection
	useZkocc     bool
}

func (cc *ConnCache) ConnForPath(zkPath string) (cn Conn, err error) {
	zcell := ZkCellFromZkPath(zkPath)

	cc.mutex.Lock()
	if cc.zconnCellMap == nil {
		cc.mutex.Unlock()
		return nil, &zookeeper.Error{Op: "dial", Code: zookeeper.ZCLOSING}
	}

	conn, ok := cc.zconnCellMap[zcell]
	if !ok {
		conn = &cachedConn{}
		cc.zconnCellMap[zcell] = conn
	}
	cc.mutex.Unlock()

	// We only want one goroutine at a time trying to connect here, so keep the
	// lock during the zk dial process.
	conn.mutex.Lock()
	defer conn.mutex.Unlock()

	if conn.zconn != nil {
		return conn.zconn, nil
	}

	if cc.useZkocc {
		conn.zconn, err = DialZkocc(ZkPathToZkAddr(zkPath, true), *recvTimeout)
	} else {
		conn.zconn, err = cc.newZookeeperConn(zkPath, zcell)
	}
	return conn.zconn, err
}

func (cc *ConnCache) newZookeeperConn(zkPath, zcell string) (Conn, error) {
	conn, session, err := DialZk(ZkPathToZkAddr(zkPath, false), *recvTimeout)
	if err != nil {
		return nil, err
	}
	go cc.handleSessionEvents(zcell, conn, session)
	return conn, nil
}

func (cc *ConnCache) handleSessionEvents(cell string, conn Conn, session <-chan zookeeper.Event) {
	for event := range session {
		switch event.State {
		case zookeeper.STATE_EXPIRED_SESSION:
			conn.Close()
			fallthrough
		case zookeeper.STATE_CLOSED:
			cc.mutex.Lock()
			if cc.zconnCellMap != nil {
				delete(cc.zconnCellMap, cell)
			}
			cc.mutex.Unlock()
			log.Printf("zk conn cache: session for cell %v ended: %v", cell, event)
			return
		default:
			log.Printf("zk conn cache: session for cell %v event: %v", cell, event)
		}
	}
}

func (cc *ConnCache) Close() error {
	cc.mutex.Lock()
	defer cc.mutex.Unlock()
	for _, conn := range cc.zconnCellMap {
		conn.mutex.Lock()
		if conn.zconn != nil {
			conn.zconn.Close()
			conn.zconn = nil
		}
		conn.mutex.Unlock()
	}
	cc.zconnCellMap = nil
	return nil
}

func NewConnCache(useZkocc bool) *ConnCache {
	return &ConnCache{
		zconnCellMap: make(map[string]*cachedConn),
		useZkocc:     useZkocc}
}
<commit_msg>Fixing this file from last check-in, oops.<commit_after>// Copyright 2012, Google Inc. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package zk

import (
	"log"
	"math/rand"
	"sync"
	"time"

	"launchpad.net/gozk/zookeeper"
)

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

/* When you need to talk to multiple zk cells, you need a simple
abstraction so you aren't caching clients all over the place.

ConnCache guarantees that you have at most one zookeeper connection per cell.
*/

type cachedConn struct {
	mutex sync.Mutex // used to notify if multiple goroutine simultaneously want a connection
	zconn Conn
}

type ConnCache struct {
	mutex        sync.Mutex
	zconnCellMap map[string]*cachedConn // map cell name to connection
	useZkocc     bool
}

func (cc *ConnCache) ConnForPath(zkPath string) (cn Conn, err error) {
	zcell := ZkCellFromZkPath(zkPath)

	cc.mutex.Lock()
	if cc.zconnCellMap == nil {
		cc.mutex.Unlock()
		return nil, &zookeeper.Error{Op: "dial", Code: zookeeper.ZCLOSING}
	}

	conn, ok := cc.zconnCellMap[zcell]
	if !ok {
		conn = &cachedConn{}
		cc.zconnCellMap[zcell] = conn
	}
	cc.mutex.Unlock()

	// We only want one goroutine at a time trying to connect here, so keep the
	// lock during the zk dial process.
	conn.mutex.Lock()
	defer conn.mutex.Unlock()

	if conn.zconn != nil {
		return conn.zconn, nil
	}

	if cc.useZkocc {
		conn.zconn, err = DialZkocc(ZkPathToZkAddr(zkPath, true), *baseTimeout)
	} else {
		conn.zconn, err = cc.newZookeeperConn(zkPath, zcell)
	}
	return conn.zconn, err
}

func (cc *ConnCache) newZookeeperConn(zkPath, zcell string) (Conn, error) {
	conn, session, err := DialZk(ZkPathToZkAddr(zkPath, false), *baseTimeout)
	if err != nil {
		return nil, err
	}
	go cc.handleSessionEvents(zcell, conn, session)
	return conn, nil
}

func (cc *ConnCache) handleSessionEvents(cell string, conn Conn, session <-chan zookeeper.Event) {
	for event := range session {
		switch event.State {
		case zookeeper.STATE_EXPIRED_SESSION:
			conn.Close()
			fallthrough
		case zookeeper.STATE_CLOSED:
			cc.mutex.Lock()
			if cc.zconnCellMap != nil {
				delete(cc.zconnCellMap, cell)
			}
			cc.mutex.Unlock()
			log.Printf("zk conn cache: session for cell %v ended: %v", cell, event)
			return
		default:
			log.Printf("zk conn cache: session for cell %v event: %v", cell, event)
		}
	}
}

func (cc *ConnCache) Close() error {
	cc.mutex.Lock()
	defer cc.mutex.Unlock()
	for _, conn := range cc.zconnCellMap {
		conn.mutex.Lock()
		if conn.zconn != nil {
			conn.zconn.Close()
			conn.zconn = nil
		}
		conn.mutex.Unlock()
	}
	cc.zconnCellMap = nil
	return nil
}

func NewConnCache(useZkocc bool) *ConnCache {
	return &ConnCache{
		zconnCellMap: make(map[string]*cachedConn),
		useZkocc:     useZkocc}
}
<|endoftext|>


<commit_before>package chat

import (
	"bytes"
	"context"
	"crypto/md5"
	"encoding/hex"
	"fmt"
	"io"
	"time"

	"github.com/keybase/client/go/chat/s3"
	"github.com/keybase/client/go/libkb"
	"github.com/keybase/client/go/logger"
	"github.com/keybase/client/go/protocol/chat1"
	"golang.org/x/sync/errgroup"
)

const minMultiSize = 5 * 1024 * 1024 // can't use Multi API with parts less than 5MB
const blockSize = 5 * 1024 * 1024    // 5MB is the minimum Multi part size

// PutS3Result is the success result of calling PutS3.
type PutS3Result struct {
	Region   string
	Endpoint string
	Bucket   string
	Path     string
	Size     int64
}

// PutS3 uploads the data in Reader r to S3.  It chooses whether to use
// putSingle or putMultiPipeline based on the size of the object.
func PutS3(ctx context.Context, log logger.Logger, r io.Reader, size int64, params chat1.S3Params, signer s3.Signer, progress ProgressReporter, local chat1.LocalSource, previous *AttachmentInfo) (*PutS3Result, error) {
	region := s3.Region{
		Name:             params.RegionName,
		S3Endpoint:       params.RegionEndpoint,
		S3BucketEndpoint: params.RegionBucketEndpoint,
	}
	conn := s3.New(signer, region)
	conn.AccessKey = params.AccessKey

	b := conn.Bucket(params.Bucket)

	if size <= minMultiSize {
		if err := putSingle(ctx, log, r, size, params, b, progress); err != nil {
			return nil, err
		}
	} else {
		objectKey, err := putMultiPipeline(ctx, log, r, size, params, b, progress, local, previous)
		if err != nil {
			return nil, err
		}
		params.ObjectKey = objectKey
	}

	res := PutS3Result{
		Region:   params.RegionName,
		Endpoint: params.RegionEndpoint,
		Bucket:   params.Bucket,
		Path:     params.ObjectKey,
		Size:     size,
	}

	return &res, nil
}

// putSingle uploads data in r to S3 with the Put API.  It has to be
// used for anything less than 5MB.  It can be used for anything up
// to 5GB, but putMultiPipeline best for anything over 5MB.
func putSingle(ctx context.Context, log logger.Logger, r io.Reader, size int64, params chat1.S3Params, b *s3.Bucket, progress ProgressReporter) error {
	log.Debug("s3 putSingle (size = %d)", size)

	// In order to be able to retry the upload, need to read in the entire
	// attachment.  But putSingle is only called for attachments <= 5MB, so
	// this isn't horrible.
	buf := make([]byte, size)
	n, err := io.ReadFull(r, buf)
	if err != nil {
		return err
	}
	if int64(n) != size {
		return fmt.Errorf("invalid read attachment size: %d (expected %d)", n, size)
	}
	sr := bytes.NewReader(buf)

	progWriter := newProgressWriter(progress, int(size))
	tee := io.TeeReader(sr, progWriter)

	var lastErr error
	for i := 0; i < 10; i++ {
		select {
		case <-ctx.Done():
			return ctx.Err()
		case <-time.After(libkb.BackoffDefault.Duration(i)):
		}
		log.Debug("s3 putSingle attempt %d", i+1)
		err := b.PutReader(params.ObjectKey, tee, size, "application/octet-stream", s3.ACL(params.Acl), s3.Options{})
		if err == nil {
			log.Debug("putSingle attempt %d success", i+1)
			return nil
		}
		log.Debug("putSingle attempt %d error: %s", i+1, err)
		lastErr = err

		// move back to beginning of sr buffer for retry
		sr.Seek(0, io.SeekStart)
		progWriter = newProgressWriter(progress, int(size))
		tee = io.TeeReader(sr, progWriter)
	}
	return fmt.Errorf("failed putSingle (last error: %s)", lastErr)
}

// putMultiPipeline uploads data in r to S3 using the Multi API.  It uses a
// pipeline to upload 10 blocks of data concurrently.  Each block is 5MB.
// It returns the object key if no errors.  putMultiPipeline will return
// a different object key from params.ObjectKey if a previous Put is
// successfully resumed and completed.
func putMultiPipeline(ctx context.Context, log logger.Logger, r io.Reader, size int64, params chat1.S3Params, b *s3.Bucket, progress ProgressReporter, local chat1.LocalSource, previous *AttachmentInfo) (string, error) {
	log.Debug("s3 putMultiPipeline (size = %d)", size)

	if previous != nil {
		log.Warning("put multi, changing object key to %s", previous.ObjectKey)
		params.ObjectKey = previous.ObjectKey
	}

	log.Warning("put multi, object key = %s", params.ObjectKey)

	multi, err := b.Multi(params.ObjectKey, "application/octet-stream", s3.ACL(params.Acl))
	if err != nil {
		log.Debug("Multi error: %s", err)
		return "", err
	}

	log.Warning("init multi: %+v", multi)

	var previousParts []s3.Part
	if previous != nil {
		previousParts, err = multi.ListParts()
		if err != nil {
			log.Debug("ignoring multi.ListParts error: %s", err)
		}
		log.Warning("previousParts: %+v", previousParts)
	}

	type job struct {
		block []byte
		index int
	}
	eg, ctx := errgroup.WithContext(ctx)
	blockCh := make(chan job)
	retCh := make(chan s3.Part)
	eg.Go(func() error {
		defer close(blockCh)
		var partNumber int
		for {
			partNumber++
			block := make([]byte, blockSize)
			// must call io.ReadFull to ensure full block read
			n, err := io.ReadFull(r, block)
			// io.ErrUnexpectedEOF will be returned for last partial block,
			// which is ok.
			if err != nil && err != io.ErrUnexpectedEOF && err != io.EOF {
				return err
			}
			if n < blockSize {
				block = block[:n]
			}
			if n > 0 {
				select {
				case blockCh <- job{block: block, index: partNumber}:
				case <-ctx.Done():
					return ctx.Err()
				}
			}
			if err == io.EOF || err == io.ErrUnexpectedEOF {
				break
			}
		}
		return nil
	})
	for i := 0; i < 10; i++ {
		eg.Go(func() error {
			for b := range blockCh {
				log.Debug("start: upload part %d", b.index)
				if previous != nil && len(previousParts) > b.index {
					// check previousParts for this block
					p := previousParts[b.index-1] // part list starts at index 1
					md5sum := md5.Sum(b.block)
					md5hex := `"` + hex.EncodeToString(md5sum[:]) + `"`
					if int(p.Size) == len(b.block) && p.ETag == md5hex {
						log.Debug("part %d already uploaded", b.index)
						select {
						case retCh <- p:
						case <-ctx.Done():
							return ctx.Err()
						}
						continue
					} else {
						log.Debug("part %d mismatch:  size %d != expected %d or etag %s != expected %s", b.index, p.Size, len(b.block), p.ETag, md5hex)
					}
				}
				part, putErr := putRetry(ctx, log, multi, b.index, b.block)
				if putErr != nil {
					return putErr
				}
				select {
				case retCh <- part:
				case <-ctx.Done():
					return ctx.Err()
				}
				log.Debug("finish: upload part %d", b.index)
			}
			return nil
		})
	}

	go func() {
		eg.Wait()
		close(retCh)
	}()

	var complete int64
	var parts []s3.Part
	for p := range retCh {
		parts = append(parts, p)
		complete += p.Size
		if progress != nil {
			progress(int(complete), int(size))
		}
	}
	if err := eg.Wait(); err != nil {
		return "", err
	}

	log.Debug("s3 putMulti all parts uploaded, completing request")

	if err = multi.Complete(parts); err != nil {
		log.Debug("multi.Complete error: %s", err)
		return "", err
	}
	log.Debug("s3 putMulti success, %d parts", len(parts))

	return params.ObjectKey, nil
}

// putRetry sends a block to S3, retrying 10 times w/ backoff.
func putRetry(ctx context.Context, log logger.Logger, multi *s3.Multi, partNumber int, block []byte) (s3.Part, error) {
	var lastErr error
	for i := 0; i < 10; i++ {
		select {
		case <-ctx.Done():
			return s3.Part{}, ctx.Err()
		case <-time.After(libkb.BackoffDefault.Duration(i)):
		}
		log.Debug("attempt %d to upload part %d", i+1, partNumber)
		part, putErr := multi.PutPart(partNumber, bytes.NewReader(block))
		if putErr == nil {
			log.Debug("success in attempt %d to upload part %d", i+1, partNumber)
			log.Debug("part: %+v", part)
			return part, nil
		}
		log.Debug("error in attempt %d to upload part %d: %s", i+1, putErr)
		lastErr = putErr
	}
	return s3.Part{}, fmt.Errorf("failed to put part %d (last error: %s)", partNumber, lastErr)
}
<commit_msg>Clean out debug warnings<commit_after>package chat

import (
	"bytes"
	"context"
	"crypto/md5"
	"encoding/hex"
	"fmt"
	"io"
	"time"

	"github.com/keybase/client/go/chat/s3"
	"github.com/keybase/client/go/libkb"
	"github.com/keybase/client/go/logger"
	"github.com/keybase/client/go/protocol/chat1"
	"golang.org/x/sync/errgroup"
)

const minMultiSize = 5 * 1024 * 1024 // can't use Multi API with parts less than 5MB
const blockSize = 5 * 1024 * 1024    // 5MB is the minimum Multi part size

// PutS3Result is the success result of calling PutS3.
type PutS3Result struct {
	Region   string
	Endpoint string
	Bucket   string
	Path     string
	Size     int64
}

// PutS3 uploads the data in Reader r to S3.  It chooses whether to use
// putSingle or putMultiPipeline based on the size of the object.
func PutS3(ctx context.Context, log logger.Logger, r io.Reader, size int64, params chat1.S3Params, signer s3.Signer, progress ProgressReporter, local chat1.LocalSource, previous *AttachmentInfo) (*PutS3Result, error) {
	region := s3.Region{
		Name:             params.RegionName,
		S3Endpoint:       params.RegionEndpoint,
		S3BucketEndpoint: params.RegionBucketEndpoint,
	}
	conn := s3.New(signer, region)
	conn.AccessKey = params.AccessKey

	b := conn.Bucket(params.Bucket)

	if size <= minMultiSize {
		if err := putSingle(ctx, log, r, size, params, b, progress); err != nil {
			return nil, err
		}
	} else {
		objectKey, err := putMultiPipeline(ctx, log, r, size, params, b, progress, local, previous)
		if err != nil {
			return nil, err
		}
		params.ObjectKey = objectKey
	}

	res := PutS3Result{
		Region:   params.RegionName,
		Endpoint: params.RegionEndpoint,
		Bucket:   params.Bucket,
		Path:     params.ObjectKey,
		Size:     size,
	}

	return &res, nil
}

// putSingle uploads data in r to S3 with the Put API.  It has to be
// used for anything less than 5MB.  It can be used for anything up
// to 5GB, but putMultiPipeline best for anything over 5MB.
func putSingle(ctx context.Context, log logger.Logger, r io.Reader, size int64, params chat1.S3Params, b *s3.Bucket, progress ProgressReporter) error {
	log.Debug("s3 putSingle (size = %d)", size)

	// In order to be able to retry the upload, need to read in the entire
	// attachment.  But putSingle is only called for attachments <= 5MB, so
	// this isn't horrible.
	buf := make([]byte, size)
	n, err := io.ReadFull(r, buf)
	if err != nil {
		return err
	}
	if int64(n) != size {
		return fmt.Errorf("invalid read attachment size: %d (expected %d)", n, size)
	}
	sr := bytes.NewReader(buf)

	progWriter := newProgressWriter(progress, int(size))
	tee := io.TeeReader(sr, progWriter)

	var lastErr error
	for i := 0; i < 10; i++ {
		select {
		case <-ctx.Done():
			return ctx.Err()
		case <-time.After(libkb.BackoffDefault.Duration(i)):
		}
		log.Debug("s3 putSingle attempt %d", i+1)
		err := b.PutReader(params.ObjectKey, tee, size, "application/octet-stream", s3.ACL(params.Acl), s3.Options{})
		if err == nil {
			log.Debug("putSingle attempt %d success", i+1)
			return nil
		}
		log.Debug("putSingle attempt %d error: %s", i+1, err)
		lastErr = err

		// move back to beginning of sr buffer for retry
		sr.Seek(0, io.SeekStart)
		progWriter = newProgressWriter(progress, int(size))
		tee = io.TeeReader(sr, progWriter)
	}
	return fmt.Errorf("failed putSingle (last error: %s)", lastErr)
}

// putMultiPipeline uploads data in r to S3 using the Multi API.  It uses a
// pipeline to upload 10 blocks of data concurrently.  Each block is 5MB.
// It returns the object key if no errors.  putMultiPipeline will return
// a different object key from params.ObjectKey if a previous Put is
// successfully resumed and completed.
func putMultiPipeline(ctx context.Context, log logger.Logger, r io.Reader, size int64, params chat1.S3Params, b *s3.Bucket, progress ProgressReporter, local chat1.LocalSource, previous *AttachmentInfo) (string, error) {
	log.Debug("s3 putMultiPipeline (size = %d)", size)

	if previous != nil {
		log.Debug("put multi, changing object key to %s", previous.ObjectKey)
		params.ObjectKey = previous.ObjectKey
	}

	multi, err := b.Multi(params.ObjectKey, "application/octet-stream", s3.ACL(params.Acl))
	if err != nil {
		log.Debug("Multi error: %s", err)
		return "", err
	}

	var previousParts []s3.Part
	if previous != nil {
		previousParts, err = multi.ListParts()
		if err != nil {
			log.Debug("ignoring multi.ListParts error: %s", err)
		}
	}

	type job struct {
		block []byte
		index int
	}
	eg, ctx := errgroup.WithContext(ctx)
	blockCh := make(chan job)
	retCh := make(chan s3.Part)
	eg.Go(func() error {
		defer close(blockCh)
		var partNumber int
		for {
			partNumber++
			block := make([]byte, blockSize)
			// must call io.ReadFull to ensure full block read
			n, err := io.ReadFull(r, block)
			// io.ErrUnexpectedEOF will be returned for last partial block,
			// which is ok.
			if err != nil && err != io.ErrUnexpectedEOF && err != io.EOF {
				return err
			}
			if n < blockSize {
				block = block[:n]
			}
			if n > 0 {
				select {
				case blockCh <- job{block: block, index: partNumber}:
				case <-ctx.Done():
					return ctx.Err()
				}
			}
			if err == io.EOF || err == io.ErrUnexpectedEOF {
				break
			}
		}
		return nil
	})
	for i := 0; i < 10; i++ {
		eg.Go(func() error {
			for b := range blockCh {
				log.Debug("start: upload part %d", b.index)
				if previous != nil && len(previousParts) > b.index {
					// check previousParts for this block
					p := previousParts[b.index-1] // part list starts at index 1
					md5sum := md5.Sum(b.block)
					md5hex := `"` + hex.EncodeToString(md5sum[:]) + `"`
					if int(p.Size) == len(b.block) && p.ETag == md5hex {
						log.Debug("part %d already uploaded", b.index)
						select {
						case retCh <- p:
						case <-ctx.Done():
							return ctx.Err()
						}
						continue
					} else {
						log.Debug("part %d mismatch:  size %d != expected %d or etag %s != expected %s", b.index, p.Size, len(b.block), p.ETag, md5hex)
					}
				}
				part, putErr := putRetry(ctx, log, multi, b.index, b.block)
				if putErr != nil {
					return putErr
				}
				select {
				case retCh <- part:
				case <-ctx.Done():
					return ctx.Err()
				}
				log.Debug("finish: upload part %d", b.index)
			}
			return nil
		})
	}

	go func() {
		eg.Wait()
		close(retCh)
	}()

	var complete int64
	var parts []s3.Part
	for p := range retCh {
		parts = append(parts, p)
		complete += p.Size
		if progress != nil {
			progress(int(complete), int(size))
		}
	}
	if err := eg.Wait(); err != nil {
		return "", err
	}

	log.Debug("s3 putMulti all parts uploaded, completing request")

	if err = multi.Complete(parts); err != nil {
		log.Debug("multi.Complete error: %s", err)
		return "", err
	}
	log.Debug("s3 putMulti success, %d parts", len(parts))

	return params.ObjectKey, nil
}

// putRetry sends a block to S3, retrying 10 times w/ backoff.
func putRetry(ctx context.Context, log logger.Logger, multi *s3.Multi, partNumber int, block []byte) (s3.Part, error) {
	var lastErr error
	for i := 0; i < 10; i++ {
		select {
		case <-ctx.Done():
			return s3.Part{}, ctx.Err()
		case <-time.After(libkb.BackoffDefault.Duration(i)):
		}
		log.Debug("attempt %d to upload part %d", i+1, partNumber)
		part, putErr := multi.PutPart(partNumber, bytes.NewReader(block))
		if putErr == nil {
			log.Debug("success in attempt %d to upload part %d", i+1, partNumber)
			log.Debug("part: %+v", part)
			return part, nil
		}
		log.Debug("error in attempt %d to upload part %d: %s", i+1, putErr)
		lastErr = putErr
	}
	return s3.Part{}, fmt.Errorf("failed to put part %d (last error: %s)", partNumber, lastErr)
}
<|endoftext|>


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

package godoc

import (
	"bytes"
	"reflect"
	"sort"
	"strings"
	"testing"

	"code.google.com/p/go.tools/godoc/vfs/mapfs"
)

func newCorpus(t *testing.T) *Corpus {
	c := NewCorpus(mapfs.New(map[string]string{
		"src/pkg/foo/foo.go": `// Package foo is an example.
package foo

import "bar"

const Pi = 3.1415

var Foos []Foo

// Foo is stuff.
type Foo struct{}

func New() *Foo {
   return new(Foo)
}
`,
		"src/pkg/bar/bar.go": `// Package bar is another example to test races.
package bar
`,
		"src/pkg/other/bar/bar.go": `// Package bar is another bar package.
package bar
func X() {}
`,
		"src/pkg/skip/skip.go": `// Package skip should be skipped.
package skip
func Skip() {}
`,
		"src/pkg/bar/readme.txt": `Whitelisted text file.
`,
		"src/pkg/bar/baz.zzz": `Text file not whitelisted.
`,
	}))
	c.IndexEnabled = true
	c.IndexDirectory = func(dir string) bool {
		return !strings.Contains(dir, "skip")
	}

	if err := c.Init(); err != nil {
		t.Fatal(err)
	}
	return c
}

func TestIndex(t *testing.T) {
	for _, docs := range []bool{true, false} {
		for _, goCode := range []bool{true, false} {
			for _, fullText := range []bool{true, false} {
				c := newCorpus(t)
				c.IndexDocs = docs
				c.IndexGoCode = goCode
				c.IndexFullText = fullText
				c.UpdateIndex()
				ix, _ := c.CurrentIndex()
				if ix == nil {
					t.Fatal("no index")
				}
				t.Logf("docs, goCode, fullText = %v,%v,%v", docs, goCode, fullText)
				testIndex(t, c, ix)
			}
		}
	}
}

func TestIndexWriteRead(t *testing.T) {
	type key struct {
		docs, goCode, fullText bool
	}
	type val struct {
		buf *bytes.Buffer
		c   *Corpus
	}
	m := map[key]val{}

	for _, docs := range []bool{true, false} {
		for _, goCode := range []bool{true, false} {
			for _, fullText := range []bool{true, false} {
				k := key{docs, goCode, fullText}
				c := newCorpus(t)
				c.IndexDocs = docs
				c.IndexGoCode = goCode
				c.IndexFullText = fullText
				c.UpdateIndex()
				ix, _ := c.CurrentIndex()
				if ix == nil {
					t.Fatal("no index")
				}
				var buf bytes.Buffer
				nw, err := ix.WriteTo(&buf)
				if err != nil {
					t.Fatalf("Index.WriteTo: %v", err)
				}
				m[k] = val{bytes.NewBuffer(buf.Bytes()), c}
				ix2 := new(Index)
				nr, err := ix2.ReadFrom(&buf)
				if err != nil {
					t.Fatalf("Index.ReadFrom: %v", err)
				}
				if nr != nw {
					t.Errorf("Wrote %d bytes to index but read %d", nw, nr)
				}
				testIndex(t, c, ix)
			}
		}
	}
	// Test CompatibleWith
	for k1, v1 := range m {
		ix := new(Index)
		if _, err := ix.ReadFrom(v1.buf); err != nil {
			t.Fatalf("Index.ReadFrom: %v", err)
		}
		for k2, v2 := range m {
			if got, want := ix.CompatibleWith(v2.c), k1 == k2; got != want {
				t.Errorf("CompatibleWith = %v; want %v for %v, %v", got, want, k1, k2)
			}
		}
	}
}

func testIndex(t *testing.T, c *Corpus, ix *Index) {
	if _, ok := ix.words["Skip"]; ok {
		t.Errorf("the word Skip was found; expected it to be skipped")
	}
	checkStats(t, c, ix)
	checkImportCount(t, c, ix)
	checkPackagePath(t, c, ix)
	checkExports(t, c, ix)
	checkIdents(t, c, ix)
}

// checkStats checks the Index's statistics.
// Some statistics are only set when we're indexing Go code.
func checkStats(t *testing.T, c *Corpus, ix *Index) {
	want := Statistics{}
	if c.IndexFullText {
		want.Bytes = 314
		want.Files = 4
		want.Lines = 21
	} else if c.IndexDocs || c.IndexGoCode {
		want.Bytes = 291
		want.Files = 3
		want.Lines = 20
	}
	if c.IndexGoCode {
		want.Words = 8
		want.Spots = 12
	}
	if got := ix.Stats(); !reflect.DeepEqual(got, want) {
		t.Errorf("Stats = %#v; want %#v", got, want)
	}
}

// checkImportCount checks the Index's import count map.
// It is only set when we're indexing Go code.
func checkImportCount(t *testing.T, c *Corpus, ix *Index) {
	want := map[string]int{}
	if c.IndexGoCode {
		want = map[string]int{
			"bar": 1,
		}
	}
	if got := ix.ImportCount(); !reflect.DeepEqual(got, want) {
		t.Errorf("ImportCount = %v; want %v", got, want)
	}
}

// checkPackagePath checks the Index's package path map.
// It is set if at least one of the indexing options is enabled.
func checkPackagePath(t *testing.T, c *Corpus, ix *Index) {
	want := map[string]map[string]bool{}
	if c.IndexDocs || c.IndexGoCode || c.IndexFullText {
		want = map[string]map[string]bool{
			"foo": map[string]bool{
				"foo": true,
			},
			"bar": map[string]bool{
				"bar":       true,
				"other/bar": true,
			},
		}
	}
	if got := ix.PackagePath(); !reflect.DeepEqual(got, want) {
		t.Errorf("PackagePath = %v; want %v", got, want)
	}
}

// checkExports checks the Index's exports map.
// It is only set when we're indexing Go code.
func checkExports(t *testing.T, c *Corpus, ix *Index) {
	want := map[string]map[string]SpotKind{}
	if c.IndexGoCode {
		want = map[string]map[string]SpotKind{
			"foo": map[string]SpotKind{
				"Pi":   ConstDecl,
				"Foos": VarDecl,
				"Foo":  TypeDecl,
				"New":  FuncDecl,
			},
			"other/bar": map[string]SpotKind{
				"X": FuncDecl,
			},
		}
	}
	if got := ix.Exports(); !reflect.DeepEqual(got, want) {
		t.Errorf("Exports = %v; want %v", got, want)
	}
}

// checkIdents checks the Index's indents map.
// It is only set when we're indexing documentation.
func checkIdents(t *testing.T, c *Corpus, ix *Index) {
	want := map[SpotKind]map[string][]Ident{}
	if c.IndexDocs {
		want = map[SpotKind]map[string][]Ident{
			ConstDecl: map[string][]Ident{
				"Pi": []Ident{{"/src/pkg/foo", "foo", "Pi", ""}},
			},
			VarDecl: map[string][]Ident{
				"Foos": []Ident{{"/src/pkg/foo", "foo", "Foos", ""}},
			},
			TypeDecl: map[string][]Ident{
				"Foo": []Ident{{"/src/pkg/foo", "foo", "Foo", "Foo is stuff."}},
			},
			FuncDecl: map[string][]Ident{
				"New": []Ident{{"/src/pkg/foo", "foo", "New", ""}},
				"X":   []Ident{{"/src/pkg/other/bar", "bar", "X", ""}},
			},
		}
	}
	if got := ix.Idents(); !reflect.DeepEqual(got, want) {
		t.Errorf("Idents = %v; want %v", got, want)
	}
}

func TestIdentResultSort(t *testing.T) {
	for _, tc := range []struct {
		ir  []Ident
		exp []Ident
	}{
		{
			ir: []Ident{
				{"/a/b/pkg2", "pkg2", "MyFunc2", ""},
				{"/b/d/pkg3", "pkg3", "MyFunc3", ""},
				{"/a/b/pkg1", "pkg1", "MyFunc1", ""},
			},
			exp: []Ident{
				{"/a/b/pkg1", "pkg1", "MyFunc1", ""},
				{"/a/b/pkg2", "pkg2", "MyFunc2", ""},
				{"/b/d/pkg3", "pkg3", "MyFunc3", ""},
			},
		},
	} {
		if sort.Sort(byPackage(tc.ir)); !reflect.DeepEqual(tc.ir, tc.exp) {
			t.Errorf("got: %v, want %v", tc.ir, tc.exp)
		}
	}
}

func TestIdentPackageFilter(t *testing.T) {
	for _, tc := range []struct {
		ir  []Ident
		pak string
		exp []Ident
	}{
		{
			ir: []Ident{
				{"/a/b/pkg2", "pkg2", "MyFunc2", ""},
				{"/b/d/pkg3", "pkg3", "MyFunc3", ""},
				{"/a/b/pkg1", "pkg1", "MyFunc1", ""},
			},
			pak: "pkg2",
			exp: []Ident{
				{"/a/b/pkg2", "pkg2", "MyFunc2", ""},
			},
		},
	} {
		if res := byPackage(tc.ir).filter(tc.pak); !reflect.DeepEqual(res, tc.exp) {
			t.Errorf("got: %v, want %v", res, tc.exp)
		}
	}
}
<commit_msg>godoc: Update identifier indexing tests.<commit_after>// Copyright 2013 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package godoc

import (
	"bytes"
	"reflect"
	"sort"
	"strings"
	"testing"

	"code.google.com/p/go.tools/godoc/vfs/mapfs"
)

func newCorpus(t *testing.T) *Corpus {
	c := NewCorpus(mapfs.New(map[string]string{
		"src/pkg/foo/foo.go": `// Package foo is an example.
package foo

import "bar"

const Pi = 3.1415

var Foos []Foo

// Foo is stuff.
type Foo struct{}

func New() *Foo {
   return new(Foo)
}
`,
		"src/pkg/bar/bar.go": `// Package bar is another example to test races.
package bar
`,
		"src/pkg/other/bar/bar.go": `// Package bar is another bar package.
package bar
func X() {}
`,
		"src/pkg/skip/skip.go": `// Package skip should be skipped.
package skip
func Skip() {}
`,
		"src/pkg/bar/readme.txt": `Whitelisted text file.
`,
		"src/pkg/bar/baz.zzz": `Text file not whitelisted.
`,
	}))
	c.IndexEnabled = true
	c.IndexDirectory = func(dir string) bool {
		return !strings.Contains(dir, "skip")
	}

	if err := c.Init(); err != nil {
		t.Fatal(err)
	}
	return c
}

func TestIndex(t *testing.T) {
	for _, docs := range []bool{true, false} {
		for _, goCode := range []bool{true, false} {
			for _, fullText := range []bool{true, false} {
				c := newCorpus(t)
				c.IndexDocs = docs
				c.IndexGoCode = goCode
				c.IndexFullText = fullText
				c.UpdateIndex()
				ix, _ := c.CurrentIndex()
				if ix == nil {
					t.Fatal("no index")
				}
				t.Logf("docs, goCode, fullText = %v,%v,%v", docs, goCode, fullText)
				testIndex(t, c, ix)
			}
		}
	}
}

func TestIndexWriteRead(t *testing.T) {
	type key struct {
		docs, goCode, fullText bool
	}
	type val struct {
		buf *bytes.Buffer
		c   *Corpus
	}
	m := map[key]val{}

	for _, docs := range []bool{true, false} {
		for _, goCode := range []bool{true, false} {
			for _, fullText := range []bool{true, false} {
				k := key{docs, goCode, fullText}
				c := newCorpus(t)
				c.IndexDocs = docs
				c.IndexGoCode = goCode
				c.IndexFullText = fullText
				c.UpdateIndex()
				ix, _ := c.CurrentIndex()
				if ix == nil {
					t.Fatal("no index")
				}
				var buf bytes.Buffer
				nw, err := ix.WriteTo(&buf)
				if err != nil {
					t.Fatalf("Index.WriteTo: %v", err)
				}
				m[k] = val{bytes.NewBuffer(buf.Bytes()), c}
				ix2 := new(Index)
				nr, err := ix2.ReadFrom(&buf)
				if err != nil {
					t.Fatalf("Index.ReadFrom: %v", err)
				}
				if nr != nw {
					t.Errorf("Wrote %d bytes to index but read %d", nw, nr)
				}
				testIndex(t, c, ix)
			}
		}
	}
	// Test CompatibleWith
	for k1, v1 := range m {
		ix := new(Index)
		if _, err := ix.ReadFrom(v1.buf); err != nil {
			t.Fatalf("Index.ReadFrom: %v", err)
		}
		for k2, v2 := range m {
			if got, want := ix.CompatibleWith(v2.c), k1 == k2; got != want {
				t.Errorf("CompatibleWith = %v; want %v for %v, %v", got, want, k1, k2)
			}
		}
	}
}

func testIndex(t *testing.T, c *Corpus, ix *Index) {
	if _, ok := ix.words["Skip"]; ok {
		t.Errorf("the word Skip was found; expected it to be skipped")
	}
	checkStats(t, c, ix)
	checkImportCount(t, c, ix)
	checkPackagePath(t, c, ix)
	checkExports(t, c, ix)
	checkIdents(t, c, ix)
}

// checkStats checks the Index's statistics.
// Some statistics are only set when we're indexing Go code.
func checkStats(t *testing.T, c *Corpus, ix *Index) {
	want := Statistics{}
	if c.IndexFullText {
		want.Bytes = 314
		want.Files = 4
		want.Lines = 21
	} else if c.IndexDocs || c.IndexGoCode {
		want.Bytes = 291
		want.Files = 3
		want.Lines = 20
	}
	if c.IndexGoCode {
		want.Words = 8
		want.Spots = 12
	}
	if got := ix.Stats(); !reflect.DeepEqual(got, want) {
		t.Errorf("Stats = %#v; want %#v", got, want)
	}
}

// checkImportCount checks the Index's import count map.
// It is only set when we're indexing Go code.
func checkImportCount(t *testing.T, c *Corpus, ix *Index) {
	want := map[string]int{}
	if c.IndexGoCode {
		want = map[string]int{
			"bar": 1,
		}
	}
	if got := ix.ImportCount(); !reflect.DeepEqual(got, want) {
		t.Errorf("ImportCount = %v; want %v", got, want)
	}
}

// checkPackagePath checks the Index's package path map.
// It is set if at least one of the indexing options is enabled.
func checkPackagePath(t *testing.T, c *Corpus, ix *Index) {
	want := map[string]map[string]bool{}
	if c.IndexDocs || c.IndexGoCode || c.IndexFullText {
		want = map[string]map[string]bool{
			"foo": map[string]bool{
				"foo": true,
			},
			"bar": map[string]bool{
				"bar":       true,
				"other/bar": true,
			},
		}
	}
	if got := ix.PackagePath(); !reflect.DeepEqual(got, want) {
		t.Errorf("PackagePath = %v; want %v", got, want)
	}
}

// checkExports checks the Index's exports map.
// It is only set when we're indexing Go code.
func checkExports(t *testing.T, c *Corpus, ix *Index) {
	want := map[string]map[string]SpotKind{}
	if c.IndexGoCode {
		want = map[string]map[string]SpotKind{
			"foo": map[string]SpotKind{
				"Pi":   ConstDecl,
				"Foos": VarDecl,
				"Foo":  TypeDecl,
				"New":  FuncDecl,
			},
			"other/bar": map[string]SpotKind{
				"X": FuncDecl,
			},
		}
	}
	if got := ix.Exports(); !reflect.DeepEqual(got, want) {
		t.Errorf("Exports = %v; want %v", got, want)
	}
}

// checkIdents checks the Index's indents map.
// It is only set when we're indexing documentation.
func checkIdents(t *testing.T, c *Corpus, ix *Index) {
	want := map[SpotKind]map[string][]Ident{}
	if c.IndexDocs {
		want = map[SpotKind]map[string][]Ident{
			PackageClause: map[string][]Ident{
				"bar": []Ident{
					{"/src/pkg/bar", "bar", "bar", "Package bar is another example to test races."},
					{"/src/pkg/other/bar", "bar", "bar", "Package bar is another bar package."},
				},
				"foo": []Ident{{"/src/pkg/foo", "foo", "foo", "Package foo is an example."}},
			},
			ConstDecl: map[string][]Ident{
				"Pi": []Ident{{"/src/pkg/foo", "foo", "Pi", ""}},
			},
			VarDecl: map[string][]Ident{
				"Foos": []Ident{{"/src/pkg/foo", "foo", "Foos", ""}},
			},
			TypeDecl: map[string][]Ident{
				"Foo": []Ident{{"/src/pkg/foo", "foo", "Foo", "Foo is stuff."}},
			},
			FuncDecl: map[string][]Ident{
				"New": []Ident{{"/src/pkg/foo", "foo", "New", ""}},
				"X":   []Ident{{"/src/pkg/other/bar", "bar", "X", ""}},
			},
		}
	}
	if got := ix.Idents(); !reflect.DeepEqual(got, want) {
		t.Errorf("Idents = %v; want %v", got, want)
	}
}

func TestIdentResultSort(t *testing.T) {
	for _, tc := range []struct {
		ir  []Ident
		exp []Ident
	}{
		{
			ir: []Ident{
				{"/a/b/pkg2", "pkg2", "MyFunc2", ""},
				{"/b/d/pkg3", "pkg3", "MyFunc3", ""},
				{"/a/b/pkg1", "pkg1", "MyFunc1", ""},
			},
			exp: []Ident{
				{"/a/b/pkg1", "pkg1", "MyFunc1", ""},
				{"/a/b/pkg2", "pkg2", "MyFunc2", ""},
				{"/b/d/pkg3", "pkg3", "MyFunc3", ""},
			},
		},
	} {
		if sort.Sort(byPackage(tc.ir)); !reflect.DeepEqual(tc.ir, tc.exp) {
			t.Errorf("got: %v, want %v", tc.ir, tc.exp)
		}
	}
}

func TestIdentPackageFilter(t *testing.T) {
	for _, tc := range []struct {
		ir  []Ident
		pak string
		exp []Ident
	}{
		{
			ir: []Ident{
				{"/a/b/pkg2", "pkg2", "MyFunc2", ""},
				{"/b/d/pkg3", "pkg3", "MyFunc3", ""},
				{"/a/b/pkg1", "pkg1", "MyFunc1", ""},
			},
			pak: "pkg2",
			exp: []Ident{
				{"/a/b/pkg2", "pkg2", "MyFunc2", ""},
			},
		},
	} {
		if res := byPackage(tc.ir).filter(tc.pak); !reflect.DeepEqual(res, tc.exp) {
			t.Errorf("got: %v, want %v", res, tc.exp)
		}
	}
}
<|endoftext|>


<commit_before>// jekyll is the configuration of handlers and such specific to the org's requirements. This is what you should copy and customize.
package jekyll

import (
	"fmt"

	"github.com/parkr/auto-reply/affinity"
	"github.com/parkr/auto-reply/autopull"
	"github.com/parkr/auto-reply/chlog"
	"github.com/parkr/auto-reply/ctx"
	"github.com/parkr/auto-reply/hooks"
	"github.com/parkr/auto-reply/labeler"
	"github.com/parkr/auto-reply/lgtm"

	"github.com/google/go-github/github"
	"github.com/parkr/auto-reply/jekyll/deprecate"
	"github.com/parkr/auto-reply/jekyll/issuecomment"
)

var jekyllOrgEventHandlers = hooks.EventHandlerMap{
	hooks.CreateEvent: {chlog.CreateReleaseOnTagHandler},
	hooks.IssuesEvent: {deprecate.DeprecateOldRepos},
	hooks.IssueCommentEvent: {
		issuecomment.PendingFeedbackUnlabeler,
		issuecomment.StaleUnlabeler,
		chlog.MergeAndLabel,
	},
	hooks.PushEvent:        {autopull.AutomaticallyCreatePullRequest("jekyll/jekyll")},
	hooks.PullRequestEvent: {labeler.PendingRebaseNeedsWorkPRUnlabeler},
	hooks.StatusEvent:      {statStatus},
}

func statStatus(context *ctx.Context, payload interface{}) error {
	status, ok := payload.(*github.StatusEvent)
	if !ok {
		return context.NewError("statStatus: not an status event")
	}

	context.SetIssue(*status.Repo.Owner.Login, *status.Repo.Name, -1)

	if context.Statsd != nil {
		statName := fmt.Sprintf("status.%s", *status.State)
		context.Log("context.Statsd.Count(%s, 1, []string{context:%s, repo:%s}, 1)", statName, *status.Context, context.Issue)
		return context.Statsd.Incr(
			statName,
			[]string{
				"context:" + *status.Context,
				"repo:" + context.Issue.String(),
			},
			float64(1.0), // rate..?
		)
	}
	return nil
}

func jekyllAffinityHandler(context *ctx.Context) *affinity.Handler {
	handler := &affinity.Handler{}

	handler.AddRepo("jekyll", "jekyll")
	handler.AddRepo("jekyll", "minima")

	handler.AddTeam(context, 1961060) // @jekyll/build
	handler.AddTeam(context, 1961072) // @jekyll/documentation
	handler.AddTeam(context, 1961061) // @jekyll/ecosystem
	handler.AddTeam(context, 1961065) // @jekyll/performance
	handler.AddTeam(context, 1961059) // @jekyll/stability
	handler.AddTeam(context, 1116640) // @jekyll/windows

	context.Log("affinity teams: %q", handler.GetTeams())
	context.Log("affinity team repos: %q", handler.GetRepos())

	return handler
}

func newLgtmHandler() *lgtm.Handler {
	handler := &lgtm.Handler{}

	handler.AddRepo("jekyll", "jekyll", 2)
	handler.AddRepo("jekyll", "jekyll-coffeescript", 2)
	handler.AddRepo("jekyll", "jekyll-compose", 1)
	handler.AddRepo("jekyll", "jekyll-docs", 1)
	handler.AddRepo("jekyll", "jekyll-feed", 1)
	handler.AddRepo("jekyll", "jekyll-gist", 2)
	handler.AddRepo("jekyll", "jekyll-import", 1)
	handler.AddRepo("jekyll", "jekyll-mentions", 2)
	handler.AddRepo("jekyll", "jekyll-opal", 2)
	handler.AddRepo("jekyll", "jekyll-paginate", 2)
	handler.AddRepo("jekyll", "jekyll-redirect-from", 2)
	handler.AddRepo("jekyll", "jekyll-sass-converter", 2)
	handler.AddRepo("jekyll", "jekyll-sitemap", 2)
	handler.AddRepo("jekyll", "jekyll-textile-converter", 2)
	handler.AddRepo("jekyll", "jekyll-watch", 2)
	handler.AddRepo("jekyll", "github-metadata", 2)
	handler.AddRepo("jekyll", "jemoji", 1)
	handler.AddRepo("jekyll", "mercenary", 1)
	handler.AddRepo("jekyll", "minima", 1)

	return handler
}

func NewJekyllOrgHandler(context *ctx.Context) *hooks.GlobalHandler {
	affinityHandler := jekyllAffinityHandler(context)
	jekyllOrgEventHandlers.AddHandler(hooks.IssuesEvent, affinityHandler.AssignIssueToAffinityTeamCaptain)
	jekyllOrgEventHandlers.AddHandler(hooks.IssueCommentEvent, affinityHandler.AssignIssueToAffinityTeamCaptainFromComment)
	jekyllOrgEventHandlers.AddHandler(hooks.PullRequestEvent, affinityHandler.AssignPRToAffinityTeamCaptain)

	lgtmHandler := newLgtmHandler()
	jekyllOrgEventHandlers.AddHandler(hooks.PullRequestReviewEvent, lgtmHandler.PullRequestReviewHandler)

	return &hooks.GlobalHandler{
		Context:       context,
		EventHandlers: jekyllOrgEventHandlers,
	}
}
<commit_msg>Add jekyll-seo-tag<commit_after>// jekyll is the configuration of handlers and such specific to the org's requirements. This is what you should copy and customize.
package jekyll

import (
	"fmt"

	"github.com/parkr/auto-reply/affinity"
	"github.com/parkr/auto-reply/autopull"
	"github.com/parkr/auto-reply/chlog"
	"github.com/parkr/auto-reply/ctx"
	"github.com/parkr/auto-reply/hooks"
	"github.com/parkr/auto-reply/labeler"
	"github.com/parkr/auto-reply/lgtm"

	"github.com/google/go-github/github"
	"github.com/parkr/auto-reply/jekyll/deprecate"
	"github.com/parkr/auto-reply/jekyll/issuecomment"
)

var jekyllOrgEventHandlers = hooks.EventHandlerMap{
	hooks.CreateEvent: {chlog.CreateReleaseOnTagHandler},
	hooks.IssuesEvent: {deprecate.DeprecateOldRepos},
	hooks.IssueCommentEvent: {
		issuecomment.PendingFeedbackUnlabeler,
		issuecomment.StaleUnlabeler,
		chlog.MergeAndLabel,
	},
	hooks.PushEvent:        {autopull.AutomaticallyCreatePullRequest("jekyll/jekyll")},
	hooks.PullRequestEvent: {labeler.PendingRebaseNeedsWorkPRUnlabeler},
	hooks.StatusEvent:      {statStatus},
}

func statStatus(context *ctx.Context, payload interface{}) error {
	status, ok := payload.(*github.StatusEvent)
	if !ok {
		return context.NewError("statStatus: not an status event")
	}

	context.SetIssue(*status.Repo.Owner.Login, *status.Repo.Name, -1)

	if context.Statsd != nil {
		statName := fmt.Sprintf("status.%s", *status.State)
		context.Log("context.Statsd.Count(%s, 1, []string{context:%s, repo:%s}, 1)", statName, *status.Context, context.Issue)
		return context.Statsd.Incr(
			statName,
			[]string{
				"context:" + *status.Context,
				"repo:" + context.Issue.String(),
			},
			float64(1.0), // rate..?
		)
	}
	return nil
}

func jekyllAffinityHandler(context *ctx.Context) *affinity.Handler {
	handler := &affinity.Handler{}

	handler.AddRepo("jekyll", "jekyll")
	handler.AddRepo("jekyll", "minima")

	handler.AddTeam(context, 1961060) // @jekyll/build
	handler.AddTeam(context, 1961072) // @jekyll/documentation
	handler.AddTeam(context, 1961061) // @jekyll/ecosystem
	handler.AddTeam(context, 1961065) // @jekyll/performance
	handler.AddTeam(context, 1961059) // @jekyll/stability
	handler.AddTeam(context, 1116640) // @jekyll/windows

	context.Log("affinity teams: %q", handler.GetTeams())
	context.Log("affinity team repos: %q", handler.GetRepos())

	return handler
}

func newLgtmHandler() *lgtm.Handler {
	handler := &lgtm.Handler{}

	handler.AddRepo("jekyll", "jekyll", 2)
	handler.AddRepo("jekyll", "jekyll-coffeescript", 2)
	handler.AddRepo("jekyll", "jekyll-compose", 1)
	handler.AddRepo("jekyll", "jekyll-docs", 1)
	handler.AddRepo("jekyll", "jekyll-feed", 1)
	handler.AddRepo("jekyll", "jekyll-gist", 2)
	handler.AddRepo("jekyll", "jekyll-import", 1)
	handler.AddRepo("jekyll", "jekyll-mentions", 2)
	handler.AddRepo("jekyll", "jekyll-opal", 2)
	handler.AddRepo("jekyll", "jekyll-paginate", 2)
	handler.AddRepo("jekyll", "jekyll-redirect-from", 2)
	handler.AddRepo("jekyll", "jekyll-sass-converter", 2)
	handler.AddRepo("jekyll", "jekyll-seo-tag", 1)
	handler.AddRepo("jekyll", "jekyll-sitemap", 2)
	handler.AddRepo("jekyll", "jekyll-textile-converter", 2)
	handler.AddRepo("jekyll", "jekyll-watch", 2)
	handler.AddRepo("jekyll", "github-metadata", 2)
	handler.AddRepo("jekyll", "jemoji", 1)
	handler.AddRepo("jekyll", "mercenary", 1)
	handler.AddRepo("jekyll", "minima", 1)

	return handler
}

func NewJekyllOrgHandler(context *ctx.Context) *hooks.GlobalHandler {
	affinityHandler := jekyllAffinityHandler(context)
	jekyllOrgEventHandlers.AddHandler(hooks.IssuesEvent, affinityHandler.AssignIssueToAffinityTeamCaptain)
	jekyllOrgEventHandlers.AddHandler(hooks.IssueCommentEvent, affinityHandler.AssignIssueToAffinityTeamCaptainFromComment)
	jekyllOrgEventHandlers.AddHandler(hooks.PullRequestEvent, affinityHandler.AssignPRToAffinityTeamCaptain)

	lgtmHandler := newLgtmHandler()
	jekyllOrgEventHandlers.AddHandler(hooks.PullRequestReviewEvent, lgtmHandler.PullRequestReviewHandler)

	return &hooks.GlobalHandler{
		Context:       context,
		EventHandlers: jekyllOrgEventHandlers,
	}
}
<|endoftext|>


<commit_before>// Copyright 2013 The Pythia Authors.
// This file is part of Pythia.
//
// Pythia is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published by
// the Free Software Foundation, version 3 of the License.
//
// Pythia is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with Pythia.  If not, see <http://www.gnu.org/licenses/>.

package pythia

import (
	"encoding/json"
)

// Status of a task execution.
type Status string

const (
	Success  Status = "success"  // finished, output = stdout
	Timeout         = "timeout"  // timed out, output = stdout so far
	Overflow        = "overflow" // stdout too big, output = capped stdout
	Abort           = "abort"    // aborted by abort message, no output
	Crash           = "crash"    // sandbox crashed, output = stdout + exit code
	Error           = "error"    // unrecoverable error (e.g. misformatted task), output = error message
)

// Task is the description of a task to be run in a sandbox.
type Task struct {
	// Environment is the name of the root filesystem.
	Environment string `json:"environment"`

	// TaskFS is the relative path to the task filesystem.
	TaskFS string `json:"taskfs"`

	// Execution limits to be enforced in the sandbox.
	Limits struct {
		// Maximum execution time in seconds.
		Time int `json:"time"`

		// Total amount of main memory (in megabytes) allocated
		// to the sandbox VM.
		Memory int `json:"memory"`

		// Fraction (in percents) of main memory that can be used as disk space
		// in a tmpfs. Note that only used disk space is allocated.
		Disk int `json:"disk"`

		// Maximum size of the output (in bytes).
		Output int `json:"output"`
	} `json:"limits"`
}

func (task Task) String() string {
	s, err := json.Marshal(task)
	if err != nil {
		return "<" + err.Error() + ">"
	}
	return string(s)
}

// Message type.
type MsgType string

const (
	// Register a sandbox pool.
	// Pool->Queue
	RegisterPoolMsg MsgType = "register-pool"

	// Request execution of a task.
	// Frontend->Queue, Queue->Pool
	LaunchMsg = "launch"

	// Job done.
	// Pool->Queue, Queue->Frontend.
	DoneMsg = "done"

	// Abort job. The receiving end shall send a done message with status abort
	// (or another status if the job has ended meanwhile).
	// Frontend->Queue, Queue->Pool
	AbortMsg = "abort"
)

// A Message is the basic entity that is sent between components. Messages are
// serialized to JSON.
type Message struct {
	// The message.
	Message MsgType `json:"message"`

	// The capacity of the pool. Only for message register-pool.
	Capacity int `json:"capacity,omitempty"`

	// The task identifier. Only for messages launch, done and abort.
	Id string `json:"id,omitempty"`

	// The task to launch. Only for message launch.
	Task *Task `json:"task,omitempty"`

	// The input to feed to the task. Only for message launch.
	Input string `json:"input,omitempty"`

	// The result status of the execution. Only for message done.
	Status Status `json:"status,omitempty"`

	// The result output of the execution. Only for message done.
	Output string `json:"output,omitempty"`
}

func (msg Message) String() string {
	s, err := json.Marshal(msg)
	if err != nil {
		return "<" + err.Error() + ">"
	}
	return string(s)
}

// vim:set sw=4 ts=4 noet:
<commit_msg>Explicitely set type of constants<commit_after>// Copyright 2013 The Pythia Authors.
// This file is part of Pythia.
//
// Pythia is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published by
// the Free Software Foundation, version 3 of the License.
//
// Pythia is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with Pythia.  If not, see <http://www.gnu.org/licenses/>.

package pythia

import (
	"encoding/json"
)

// Status of a task execution.
type Status string

const (
	Success  Status = "success"  // finished, output = stdout
	Timeout  Status = "timeout"  // timed out, output = stdout so far
	Overflow Status = "overflow" // stdout too big, output = capped stdout
	Abort    Status = "abort"    // aborted by abort message, no output
	Crash    Status = "crash"    // sandbox crashed, output = stdout + exit code
	Error    Status = "error"    // unrecoverable error (e.g. misformatted task), output = error message
)

// Task is the description of a task to be run in a sandbox.
type Task struct {
	// Environment is the name of the root filesystem.
	Environment string `json:"environment"`

	// TaskFS is the relative path to the task filesystem.
	TaskFS string `json:"taskfs"`

	// Execution limits to be enforced in the sandbox.
	Limits struct {
		// Maximum execution time in seconds.
		Time int `json:"time"`

		// Total amount of main memory (in megabytes) allocated
		// to the sandbox VM.
		Memory int `json:"memory"`

		// Fraction (in percents) of main memory that can be used as disk space
		// in a tmpfs. Note that only used disk space is allocated.
		Disk int `json:"disk"`

		// Maximum size of the output (in bytes).
		Output int `json:"output"`
	} `json:"limits"`
}

func (task Task) String() string {
	s, err := json.Marshal(task)
	if err != nil {
		return "<" + err.Error() + ">"
	}
	return string(s)
}

// Message type.
type MsgType string

const (
	// Register a sandbox pool.
	// Pool->Queue
	RegisterPoolMsg MsgType = "register-pool"

	// Request execution of a task.
	// Frontend->Queue, Queue->Pool
	LaunchMsg MsgType = "launch"

	// Job done.
	// Pool->Queue, Queue->Frontend.
	DoneMsg MsgType = "done"

	// Abort job. The receiving end shall send a done message with status abort
	// (or another status if the job has ended meanwhile).
	// Frontend->Queue, Queue->Pool
	AbortMsg MsgType = "abort"
)

// A Message is the basic entity that is sent between components. Messages are
// serialized to JSON.
type Message struct {
	// The message.
	Message MsgType `json:"message"`

	// The capacity of the pool. Only for message register-pool.
	Capacity int `json:"capacity,omitempty"`

	// The task identifier. Only for messages launch, done and abort.
	Id string `json:"id,omitempty"`

	// The task to launch. Only for message launch.
	Task *Task `json:"task,omitempty"`

	// The input to feed to the task. Only for message launch.
	Input string `json:"input,omitempty"`

	// The result status of the execution. Only for message done.
	Status Status `json:"status,omitempty"`

	// The result output of the execution. Only for message done.
	Output string `json:"output,omitempty"`
}

func (msg Message) String() string {
	s, err := json.Marshal(msg)
	if err != nil {
		return "<" + err.Error() + ">"
	}
	return string(s)
}

// vim:set sw=4 ts=4 noet:
<|endoftext|>


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

// Package testing implements utilities to aid in testing.
package testing

import (
	"fmt"
	"strings"

	"github.com/google/chrome-ssh-agent/go/dom"
	"github.com/gopherjs/gopherjs/js"
)

// resultsAsString returns the test results as a single string suitable for
// display.
func resultsAsString(errs []error) string {
	var lines []string
	for _, err := range errs {
		lines = append(lines, err.Error())
	}
	return strings.Join(lines, "\n")
}

// getBody returns the body object, and undefined if none is present.
func getBody(d *dom.DOM) *js.Object {
	for _, e := range d.GetElementsByTag("body") {
		return e
	}
	return js.Undefined
}

// WriteResults adds elements to the supplied DOM summarizing the test results.
// The elements added are hidden, but are given identifiers such that the
// results can be queried by automation. The following elements are added:
// - failureCount: a div element, whose contained text is the number of tests
//     that failed.
// - failures: A pre element, whose contained text is a human-readable list
//     of the individual failures.
func WriteResults(d *dom.DOM, errs []error) {
	body := getBody(d)
	// Top-level container element into which we'll write results.
	d.AppendChild(body, d.NewElement("div"), func(results *js.Object) {
		// Results are hidden from view.  They may be queried
		// programmatically from the test framework.
		results.Set("style", "visiblity: hidden")

		// Indicate how many tests failed.
		d.AppendChild(results, d.NewElement("div"), func(failureCount *js.Object) {
			// Allow the element to be read by automation.
			failureCount.Set("id", "failureCount")
			d.AppendChild(failureCount, d.NewText(fmt.Sprintf("%d", len(errs))), nil)
		})

		// Enumerate the failures. This is a more readable list of the
		// individual tests that failed.
		d.AppendChild(results, d.NewElement("pre"), func(failures *js.Object) {
			// Allow element to be read by automation.
			failures.Set("id", "failures")
			d.AppendChild(failures, d.NewText(resultsAsString(errs)), nil)
		})
	})
}
<commit_msg>Don't hide results, since that cause automation to not be able to find them<commit_after>// Copyright 2018 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Package testing implements utilities to aid in testing.
package testing

import (
	"fmt"
	"strings"

	"github.com/google/chrome-ssh-agent/go/dom"
	"github.com/gopherjs/gopherjs/js"
)

// resultsAsString returns the test results as a single string suitable for
// display.
func resultsAsString(errs []error) string {
	var lines []string
	for _, err := range errs {
		lines = append(lines, err.Error())
	}
	return strings.Join(lines, "\n")
}

// getBody returns the body object, and undefined if none is present.
func getBody(d *dom.DOM) *js.Object {
	for _, e := range d.GetElementsByTag("body") {
		return e
	}
	return js.Undefined
}

// WriteResults adds elements to the supplied DOM summarizing the test results.
// The elements are given identifiers such that the results can be queried by
// automation. The following elements are added:
// - failureCount: a div element, whose contained text is the number of tests
//     that failed.
// - failures: A pre element, whose contained text is a human-readable list
//     of the individual failures.
func WriteResults(d *dom.DOM, errs []error) {
	body := getBody(d)
	// Top-level container element into which we'll write results.
	d.AppendChild(body, d.NewElement("div"), func(results *js.Object) {
		// Indicate how many tests failed.
		d.AppendChild(results, d.NewElement("div"), func(failureCount *js.Object) {
			// Allow the element to be read by automation.
			failureCount.Set("id", "failureCount")
			d.AppendChild(failureCount, d.NewText(fmt.Sprintf("%d", len(errs))), nil)
		})

		// Enumerate the failures. This is a more readable list of the
		// individual tests that failed.
		d.AppendChild(results, d.NewElement("pre"), func(failures *js.Object) {
			// Allow element to be read by automation.
			failures.Set("id", "failures")
			d.AppendChild(failures, d.NewText(resultsAsString(errs)), nil)
		})
	})
}
<|endoftext|>