<commit_before>package muniverse

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"net"
	"os/exec"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/unixpickle/essentials"
	"github.com/unixpickle/muniverse/chrome"
)

const (
	portRange    = "9000-9999"
	defaultImage = "unixpickle/muniverse:0.56.0"
)

const (
	callTimeout           = time.Minute * 2
	chromeConnectAttempts = 10
)

// This error message occurs very infrequently when doing
// `docker run` on my machine running Ubuntu 16.04.1.
const occasionalDockerErr = "Error response from daemon: device or resource busy."

// An Env controls and observes an environment.
//
// It is not safe to run an methods on an Env from more
// than one Goroutine at a time.
//
// The lifecycle of an environment is as follows:
// First, Reset is called to start an episode.
// Then, Step and Observe may be called repeatedly in any
// order until Step returns done=true to signal that the
// episode has ended.
// Once the episode has ended, Observe may be called but
// Step may not be.
// Call Reset to start a new episode and begin the process
// over again.
//
// When you are done with an Env, you must close it to
// clean up resources associated with it.
type Env interface {
	// Spec returns details about the environment.
	Spec() *EnvSpec

	// Reset resets the environment to a start state.
	Reset() error

	// Step sends the given events and advances the
	// episode by the given amount of time.
	//
	// If done is true, then the episode has ended.
	// After an episode ends, Reset must be called once
	// before Step may be called again.
	// However, observations may be made even after the
	// episode has ended.
	//
	// Typical event types are *chrome.MouseEvent and
	// *chrome.KeyEvent.
	Step(t time.Duration, events ...interface{}) (reward float64,
		done bool, err error)

	// Observe produces an observation for the current
	// state of the environment.
	Observe() (Obs, error)

	// Close cleans up resources used by the environment.
	//
	// After Close is called, the Env should not be used
	// anymore by any Goroutine.
	Close() error

	// Log returns internal log messages.
	// For example, it might return information about 404
	// errors.
	//
	// The returned list is a copy and may be modified by
	// the caller.
	Log() []string
}

type rawEnv struct {
	spec     EnvSpec
	gameHost string

	containerID string
	devConn     *chrome.Conn
	lastScore   float64
	needsReset  bool

	// Used to garbage collect the container if we
	// exit ungracefully.
	killSocket net.Conn
}

// NewEnv creates a new environment inside the default
// Docker image.
// This may take a few minutes to run the first time,
// since it has to download a large Docker image.
func NewEnv(spec *EnvSpec) (Env, error) {
	return NewEnvContainer(defaultImage, spec)
}

// NewEnvContainer creates a new environment inside a
// new Docker container of the given Docker image.
func NewEnvContainer(image string, spec *EnvSpec) (env Env, err error) {
	defer essentials.AddCtxTo("create environment", &err)

	ctx, cancel := callCtx()
	defer cancel()

	var id string

	// Retry as a workaround for an occasional error given
	// by `docker run`.
	for i := 0; i < 3; i++ {
		id, err = dockerRun(ctx, image, spec)
		if err == nil || !strings.Contains(err.Error(), occasionalDockerErr) {
			break
		}
	}

	if err != nil {
		return
	}

	ports, err := dockerBoundPorts(ctx, id)
	if err != nil {
		return
	}

	conn, err := connectDevTools(ctx, "localhost:"+ports["9222/tcp"])
	if err != nil {
		return
	}

	killSock, err := (&net.Dialer{}).DialContext(ctx, "tcp",
		"localhost:"+ports["1337/tcp"])
	if err != nil {
		conn.Close()
		return
	}

	return &rawEnv{
		spec:        *spec,
		gameHost:    "localhost",
		containerID: id,
		devConn:     conn,
		killSocket:  killSock,
	}, nil
}

// NewEnvChrome connects to an existing Chrome DevTools
// server and runs an environment in there.
//
// The gameHost argument specifies where to load games.
// For example, gameHost might be "localhost:8080" if the
// game "Foobar" should be loaded from
// "http://localhost/Foobar".
//
// The Chrome instance must have at least one page open,
// since an open page is selected and used to run the
// environment.
func NewEnvChrome(host, gameHost string, spec *EnvSpec) (env Env, err error) {
	defer essentials.AddCtxTo("create environment", &err)

	ctx, cancel := callCtx()
	defer cancel()

	conn, err := connectDevTools(ctx, host)
	if err != nil {
		return
	}

	return &rawEnv{
		spec:       *spec,
		gameHost:   gameHost,
		devConn:    conn,
		needsReset: true,
	}, nil
}

func (r *rawEnv) Spec() *EnvSpec {
	res := r.spec
	return &res
}

func (r *rawEnv) Reset() (err error) {
	defer essentials.AddCtxTo("reset environment", &err)

	ctx, cancel := callCtx()
	defer cancel()

	err = r.devConn.NavigateSafe(ctx, r.envURL())
	if err != nil {
		return
	}

	initCode := "window.muniverse.init(" + r.spec.Options + ");"
	err = r.devConn.EvalPromise(ctx, initCode, nil)
	if err != nil {
		return
	}

	err = r.devConn.EvalPromise(ctx, "window.muniverse.score();", &r.lastScore)
	err = essentials.AddCtx("get score", err)

	if err == nil {
		r.needsReset = false
	}

	return
}

func (r *rawEnv) Step(t time.Duration, events ...interface{}) (reward float64,
	done bool, err error) {
	defer essentials.AddCtxTo("step environment", &err)

	if r.needsReset {
		err = errors.New("environment needs reset")
		return
	}

	ctx, cancel := callCtx()
	defer cancel()

	for _, event := range events {
		switch event := event.(type) {
		case *chrome.MouseEvent:
			err = r.devConn.DispatchMouseEvent(ctx, event)
		case *chrome.KeyEvent:
			err = r.devConn.DispatchKeyEvent(ctx, event)
		default:
			err = fmt.Errorf("unsupported event type: %T", event)
		}
		if err != nil {
			return
		}
	}

	millis := int(t / time.Millisecond)
	timeStr := strconv.Itoa(millis)
	err = r.devConn.EvalPromise(ctx, "window.muniverse.step("+timeStr+");", &done)
	if err != nil {
		return
	}

	if done {
		r.needsReset = true
	}

	lastScore := r.lastScore
	err = r.devConn.EvalPromise(ctx, "window.muniverse.score();", &r.lastScore)
	if err != nil {
		err = essentials.AddCtx("get score", err)
		return
	}
	reward = r.lastScore - lastScore

	return
}

func (r *rawEnv) Observe() (obs Obs, err error) {
	ctx, cancel := callCtx()
	defer cancel()

	pngData, err := r.devConn.ScreenshotPNG(ctx)
	if err != nil {
		return
	}
	return pngObs(pngData), nil
}

func (r *rawEnv) Close() (err error) {
	defer essentials.AddCtxTo("close environment", &err)

	ctx, cancel := callCtx()
	defer cancel()

	errs := []error{
		r.devConn.Close(),
	}
	if r.containerID != "" {
		_, e := dockerCommand(ctx, "kill", r.containerID)
		errs = append(errs, e)
	}

	if r.killSocket != nil {
		// TODO: look into if this can ever produce an error,
		// since the container might already have closed the
		// socket by now.
		//
		// We don't close this *before* stopping the container
		// since `docker kill` might fail if the container
		// already died and was cleaned up.
		r.killSocket.Close()
	}

	// Any calls after Close() should trigger simple errors.
	r.devConn = nil
	r.killSocket = nil

	for _, err := range errs {
		if err != nil {
			return err
		}
	}
	return nil
}

func (r *rawEnv) Log() []string {
	return r.devConn.ConsoleLog()
}

func (r *rawEnv) envURL() string {
	baseName := r.spec.Name
	if r.spec.VariantOf != "" {
		baseName = r.spec.VariantOf
	}
	return "http://" + r.gameHost + "/" + baseName
}

func callCtx() (context.Context, context.CancelFunc) {
	return context.WithTimeout(context.Background(), callTimeout)
}

func dockerRun(ctx context.Context, container string, spec *EnvSpec) (id string,
	err error) {
	args := []string{
		"run",
		"-p",
		portRange + ":9222",
		"-p",
		portRange + ":1337",
		"--shm-size=200m",
		"-d",   // Run in detached mode.
		"--rm", // Automatically delete the container.
		"-i",   // Give netcat a stdin to read from.
		container,
		fmt.Sprintf("--window-size=%d,%d", spec.Width, spec.Height),
	}

	output, err := dockerCommand(ctx, args...)
	if err != nil {
		return "", essentials.AddCtx("docker run",
			fmt.Errorf("%s (make sure docker is up-to-date)", err))
	}

	return strings.TrimSpace(string(output)), nil
}

func dockerBoundPorts(ctx context.Context,
	containerID string) (mapping map[string]string, err error) {
	defer essentials.AddCtxTo("docker inspect", &err)
	rawJSON, err := dockerCommand(ctx, "inspect", containerID)
	if err != nil {
		return nil, err
	}
	var info []struct {
		NetworkSettings struct {
			Ports map[string][]struct {
				HostPort string
			}
		}
	}
	if err := json.Unmarshal(rawJSON, &info); err != nil {
		return nil, err
	}
	if len(info) != 1 {
		return nil, errors.New("unexpected number of results")
	}
	rawMapping := info[0].NetworkSettings.Ports
	mapping = map[string]string{}
	for containerPort, hostPorts := range rawMapping {
		if len(hostPorts) != 1 {
			return nil, errors.New("unexpected number of host ports")
		}
		mapping[containerPort] = hostPorts[0].HostPort
	}
	return
}

var dockerLock sync.Mutex

func dockerCommand(ctx context.Context, args ...string) (output []byte, err error) {
	dockerLock.Lock()
	defer dockerLock.Unlock()
	output, err = exec.CommandContext(ctx, "docker", args...).Output()
	if err != nil {
		if eo, ok := err.(*exec.ExitError); ok && len(eo.Stderr) > 0 {
			stderrMsg := strings.TrimSpace(string(eo.Stderr))
			err = fmt.Errorf("%s: %s", eo.String(), stderrMsg)
		}
	}
	return
}

func connectDevTools(ctx context.Context, host string) (conn *chrome.Conn,
	err error) {
	for i := 0; i < chromeConnectAttempts; i++ {
		conn, err = attemptDevTools(ctx, host)
		if err == nil {
			return
		}
		select {
		case <-time.After(time.Second):
		case <-ctx.Done():
			return nil, ctx.Err()
		}
	}
	return
}

func attemptDevTools(ctx context.Context, host string) (conn *chrome.Conn,
	err error) {
	endpoints, err := chrome.Endpoints(ctx, host)
	if err != nil {
		return
	}

	for _, ep := range endpoints {
		if ep.Type == "page" && ep.WebSocketURL != "" {
			return chrome.NewConn(ctx, ep.WebSocketURL)
		}
	}

	return nil, errors.New("no Chrome page endpoint")
}
<commit_msg>use custom downloaded_games directory without new container<commit_after>package muniverse

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"net"
	"os/exec"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/unixpickle/essentials"
	"github.com/unixpickle/muniverse/chrome"
)

const (
	portRange    = "9000-9999"
	defaultImage = "unixpickle/muniverse:0.56.0"
)

const (
	callTimeout           = time.Minute * 2
	chromeConnectAttempts = 10
)

// This error message occurs very infrequently when doing
// `docker run` on my machine running Ubuntu 16.04.1.
const occasionalDockerErr = "Error response from daemon: device or resource busy."

// An Env controls and observes an environment.
//
// It is not safe to run an methods on an Env from more
// than one Goroutine at a time.
//
// The lifecycle of an environment is as follows:
// First, Reset is called to start an episode.
// Then, Step and Observe may be called repeatedly in any
// order until Step returns done=true to signal that the
// episode has ended.
// Once the episode has ended, Observe may be called but
// Step may not be.
// Call Reset to start a new episode and begin the process
// over again.
//
// When you are done with an Env, you must close it to
// clean up resources associated with it.
type Env interface {
	// Spec returns details about the environment.
	Spec() *EnvSpec

	// Reset resets the environment to a start state.
	Reset() error

	// Step sends the given events and advances the
	// episode by the given amount of time.
	//
	// If done is true, then the episode has ended.
	// After an episode ends, Reset must be called once
	// before Step may be called again.
	// However, observations may be made even after the
	// episode has ended.
	//
	// Typical event types are *chrome.MouseEvent and
	// *chrome.KeyEvent.
	Step(t time.Duration, events ...interface{}) (reward float64,
		done bool, err error)

	// Observe produces an observation for the current
	// state of the environment.
	Observe() (Obs, error)

	// Close cleans up resources used by the environment.
	//
	// After Close is called, the Env should not be used
	// anymore by any Goroutine.
	Close() error

	// Log returns internal log messages.
	// For example, it might return information about 404
	// errors.
	//
	// The returned list is a copy and may be modified by
	// the caller.
	Log() []string
}

type rawEnv struct {
	spec     EnvSpec
	gameHost string

	containerID string
	devConn     *chrome.Conn
	lastScore   float64
	needsReset  bool

	// Used to garbage collect the container if we
	// exit ungracefully.
	killSocket net.Conn
}

// NewEnv creates a new environment inside the default
// Docker image.
// This may take a few minutes to run the first time,
// since it has to download a large Docker image.
func NewEnv(spec *EnvSpec) (Env, error) {
	return NewEnvContainer(defaultImage, spec)
}

// NewEnvContainer creates a new environment inside a
// new Docker container of the given Docker image.
func NewEnvContainer(image string, spec *EnvSpec) (Env, error) {
	return newEnvDocker(image, "", spec)
}

// NewEnvGamesDir creates a new environment with a custom
// games directory.
// The directory should contain subdirectories for each
// base game, similar to the downloaded_games directory in
// the default image.
//
// The directory path is slightly restricted.
// In particular, it cannot contain a ':' (colon).
// See https://github.com/moby/moby/issues/8604 for more.
func NewEnvGamesDir(dir string, spec *EnvSpec) (Env, error) {
	return newEnvDocker(defaultImage, dir, spec)
}

func newEnvDocker(image, volume string, spec *EnvSpec) (env Env, err error) {
	defer essentials.AddCtxTo("create environment", &err)

	ctx, cancel := callCtx()
	defer cancel()

	var id string

	// Retry as a workaround for an occasional error given
	// by `docker run`.
	for i := 0; i < 3; i++ {
		id, err = dockerRun(ctx, image, volume, spec)
		if err == nil || !strings.Contains(err.Error(), occasionalDockerErr) {
			break
		}
	}

	if err != nil {
		return
	}

	ports, err := dockerBoundPorts(ctx, id)
	if err != nil {
		return
	}

	conn, err := connectDevTools(ctx, "localhost:"+ports["9222/tcp"])
	if err != nil {
		return
	}

	killSock, err := (&net.Dialer{}).DialContext(ctx, "tcp",
		"localhost:"+ports["1337/tcp"])
	if err != nil {
		conn.Close()
		return
	}

	return &rawEnv{
		spec:        *spec,
		gameHost:    "localhost",
		containerID: id,
		devConn:     conn,
		killSocket:  killSock,
	}, nil
}

// NewEnvChrome connects to an existing Chrome DevTools
// server and runs an environment in there.
//
// The gameHost argument specifies where to load games.
// For example, gameHost might be "localhost:8080" if the
// game "Foobar" should be loaded from
// "http://localhost/Foobar".
//
// The Chrome instance must have at least one page open,
// since an open page is selected and used to run the
// environment.
func NewEnvChrome(host, gameHost string, spec *EnvSpec) (env Env, err error) {
	defer essentials.AddCtxTo("create environment", &err)

	ctx, cancel := callCtx()
	defer cancel()

	conn, err := connectDevTools(ctx, host)
	if err != nil {
		return
	}

	return &rawEnv{
		spec:       *spec,
		gameHost:   gameHost,
		devConn:    conn,
		needsReset: true,
	}, nil
}

func (r *rawEnv) Spec() *EnvSpec {
	res := r.spec
	return &res
}

func (r *rawEnv) Reset() (err error) {
	defer essentials.AddCtxTo("reset environment", &err)

	ctx, cancel := callCtx()
	defer cancel()

	err = r.devConn.NavigateSafe(ctx, r.envURL())
	if err != nil {
		return
	}

	initCode := "window.muniverse.init(" + r.spec.Options + ");"
	err = r.devConn.EvalPromise(ctx, initCode, nil)
	if err != nil {
		return
	}

	err = r.devConn.EvalPromise(ctx, "window.muniverse.score();", &r.lastScore)
	err = essentials.AddCtx("get score", err)

	if err == nil {
		r.needsReset = false
	}

	return
}

func (r *rawEnv) Step(t time.Duration, events ...interface{}) (reward float64,
	done bool, err error) {
	defer essentials.AddCtxTo("step environment", &err)

	if r.needsReset {
		err = errors.New("environment needs reset")
		return
	}

	ctx, cancel := callCtx()
	defer cancel()

	for _, event := range events {
		switch event := event.(type) {
		case *chrome.MouseEvent:
			err = r.devConn.DispatchMouseEvent(ctx, event)
		case *chrome.KeyEvent:
			err = r.devConn.DispatchKeyEvent(ctx, event)
		default:
			err = fmt.Errorf("unsupported event type: %T", event)
		}
		if err != nil {
			return
		}
	}

	millis := int(t / time.Millisecond)
	timeStr := strconv.Itoa(millis)
	err = r.devConn.EvalPromise(ctx, "window.muniverse.step("+timeStr+");", &done)
	if err != nil {
		return
	}

	if done {
		r.needsReset = true
	}

	lastScore := r.lastScore
	err = r.devConn.EvalPromise(ctx, "window.muniverse.score();", &r.lastScore)
	if err != nil {
		err = essentials.AddCtx("get score", err)
		return
	}
	reward = r.lastScore - lastScore

	return
}

func (r *rawEnv) Observe() (obs Obs, err error) {
	ctx, cancel := callCtx()
	defer cancel()

	pngData, err := r.devConn.ScreenshotPNG(ctx)
	if err != nil {
		return
	}
	return pngObs(pngData), nil
}

func (r *rawEnv) Close() (err error) {
	defer essentials.AddCtxTo("close environment", &err)

	ctx, cancel := callCtx()
	defer cancel()

	errs := []error{
		r.devConn.Close(),
	}
	if r.containerID != "" {
		_, e := dockerCommand(ctx, "kill", r.containerID)
		errs = append(errs, e)
	}

	if r.killSocket != nil {
		// TODO: look into if this can ever produce an error,
		// since the container might already have closed the
		// socket by now.
		//
		// We don't close this *before* stopping the container
		// since `docker kill` might fail if the container
		// already died and was cleaned up.
		r.killSocket.Close()
	}

	// Any calls after Close() should trigger simple errors.
	r.devConn = nil
	r.killSocket = nil

	for _, err := range errs {
		if err != nil {
			return err
		}
	}
	return nil
}

func (r *rawEnv) Log() []string {
	return r.devConn.ConsoleLog()
}

func (r *rawEnv) envURL() string {
	baseName := r.spec.Name
	if r.spec.VariantOf != "" {
		baseName = r.spec.VariantOf
	}
	return "http://" + r.gameHost + "/" + baseName
}

func callCtx() (context.Context, context.CancelFunc) {
	return context.WithTimeout(context.Background(), callTimeout)
}

func dockerRun(ctx context.Context, container, volume string,
	spec *EnvSpec) (id string, err error) {
	args := []string{
		"run",
		"-p",
		portRange + ":9222",
		"-p",
		portRange + ":1337",
		"--shm-size=200m",
		"-d",   // Run in detached mode.
		"--rm", // Automatically delete the container.
		"-i",   // Give netcat a stdin to read from.
	}
	if volume != "" {
		if strings.Contains(volume, ":") {
			return "", errors.New("path contains colons: " + volume)
		}
		args = append(args, "-v", volume+":/downloaded_games")
	}
	args = append(args, container,
		fmt.Sprintf("--window-size=%d,%d", spec.Width, spec.Height))

	output, err := dockerCommand(ctx, args...)
	if err != nil {
		return "", essentials.AddCtx("docker run",
			fmt.Errorf("%s (make sure docker is up-to-date)", err))
	}

	return strings.TrimSpace(string(output)), nil
}

func dockerBoundPorts(ctx context.Context,
	containerID string) (mapping map[string]string, err error) {
	defer essentials.AddCtxTo("docker inspect", &err)
	rawJSON, err := dockerCommand(ctx, "inspect", containerID)
	if err != nil {
		return nil, err
	}
	var info []struct {
		NetworkSettings struct {
			Ports map[string][]struct {
				HostPort string
			}
		}
	}
	if err := json.Unmarshal(rawJSON, &info); err != nil {
		return nil, err
	}
	if len(info) != 1 {
		return nil, errors.New("unexpected number of results")
	}
	rawMapping := info[0].NetworkSettings.Ports
	mapping = map[string]string{}
	for containerPort, hostPorts := range rawMapping {
		if len(hostPorts) != 1 {
			return nil, errors.New("unexpected number of host ports")
		}
		mapping[containerPort] = hostPorts[0].HostPort
	}
	return
}

var dockerLock sync.Mutex

func dockerCommand(ctx context.Context, args ...string) (output []byte, err error) {
	dockerLock.Lock()
	defer dockerLock.Unlock()
	output, err = exec.CommandContext(ctx, "docker", args...).Output()
	if err != nil {
		if eo, ok := err.(*exec.ExitError); ok && len(eo.Stderr) > 0 {
			stderrMsg := strings.TrimSpace(string(eo.Stderr))
			err = fmt.Errorf("%s: %s", eo.String(), stderrMsg)
		}
	}
	return
}

func connectDevTools(ctx context.Context, host string) (conn *chrome.Conn,
	err error) {
	for i := 0; i < chromeConnectAttempts; i++ {
		conn, err = attemptDevTools(ctx, host)
		if err == nil {
			return
		}
		select {
		case <-time.After(time.Second):
		case <-ctx.Done():
			return nil, ctx.Err()
		}
	}
	return
}

func attemptDevTools(ctx context.Context, host string) (conn *chrome.Conn,
	err error) {
	endpoints, err := chrome.Endpoints(ctx, host)
	if err != nil {
		return
	}

	for _, ep := range endpoints {
		if ep.Type == "page" && ep.WebSocketURL != "" {
			return chrome.NewConn(ctx, ep.WebSocketURL)
		}
	}

	return nil, errors.New("no Chrome page endpoint")
}
<|endoftext|>


<commit_before>package gnmsys

import (
	"time"
	"net/http/cookiejar"
	"net/http"
	"log"
	"net/url"
	"strings"
	"fmt"
	"io/ioutil"
	"encoding/json"
	"os"
	"path"
	"path/filepath"
	"io"
	"github.com/gonum/plot"
)

type System interface {
	// function that will signal to the system to clean up and shutdown
	SignalTerm()
	// function that will write the reports to disk
	SignalFlush()
	// List all the Reports available
	GetReports() []Report
	// Get the output file of the report
	GetReportFile(report Report) string
}

type SystemSignal int

const (
	term = 1 + iota
	flush
	tick
)
type SysConfig struct {
	UrlStem,
	Username,
	Password,
	OutputDir string
	SampleConfigs []SampleConfig
}
type defaultSystem struct {
	config SysConfig
	signals chan SystemSignal
	client  http.Client
	reports []Report
}

func CreateSystem(config SysConfig, reportFactories ...ReportFactory) defaultSystem {
	if config.OutputDir == "" || config.OutputDir == "." {
		config.OutputDir = "gnm_reports"
	} else {
		os.MkdirAll(config.OutputDir, os.FileMode(0755))
	}
	options := cookiejar.Options{}
	jar, err := cookiejar.New(&options)
	if err != nil {
		log.Fatal(err)
	}
	reports := make([]Report, len(reportFactories) * len(config.SampleConfigs))
	if (len(reports) == 0) {
		log.Fatalf("No reports are configured\nReport Factories: %d\nSampleConfigs: %v\n", len(reportFactories), config.SampleConfigs)
	}
	for i, fac := range reportFactories {
		for j, sConf := range config.SampleConfigs {
			reports[(i * len(config.SampleConfigs)) + j] = fac(sConf)
		}
	}
	system := defaultSystem{
		config: config,
		reports: reports,
		signals: make(chan SystemSignal),
		client: http.Client{Jar: jar}}

	system.validate()

	go loop(system.signals)

	return system
}

func (sys defaultSystem) validate() {
	for _, conf := range sys.config.SampleConfigs {
		conf.Validate()
	}

	probeFile := filepath.Join(sys.config.OutputDir, "probe")
	if err := ioutil.WriteFile(probeFile, []byte("t"), os.FileMode(0664)); err != nil {
		log.Fatalf("Do not have write permissions to %s\n", sys.config.OutputDir)
	}
	os.Remove(probeFile)

	if _, err := plot.New(); err != nil {
		fmt.Printf("Error creating a test graph %q\n", err.Error())
		msg := "A likely problem is that the directory with font files cannot be found." +
		"  Copy https://github.com/gonum/plot/tree/master/vg/fonts to same directory as the executable"
		fmt.Printf(msg)

		log.Printf("Error creating a test graph %q\n", err.Error())
		log.Fatalf(msg)
	}
}

func loop(signals chan <- SystemSignal) {
	for {
		signals <- tick
		time.Sleep(time.Second)
	}
}

func (sys defaultSystem) GetReportFile(report Report) string {
	var catDirName string
	for _, sampConf := range sys.config.SampleConfigs {
		if sampConf.Name == report.GetCategory() {
			catDirName = sampConf.DirName
			break;
		}
	}
	return filepath.Join(sys.config.OutputDir, catDirName, report.GetFileName())
}
func (sys defaultSystem) GetReports() []Report {
	return sys.reports
}
func (sys defaultSystem) SignalTerm() {
	fmt.Printf("Finalizing reports and shutting down...\n")
	sys.signals <- term
}
func (sys defaultSystem) SignalFlush() {
	fmt.Printf("Saving/Flushing Reports to disk\n")
	sys.signals <- flush
}

type systemState struct {
	startTime time.Time
	initializationComplete, mustLogin bool
	urlStem string
	loginCredentials url.Values
}
const timeFmt = "Start Time: 2006 Jan _2 15:04:05"
func (state *systemState) initialize() {
	if !state.initializationComplete {
		state.initializationComplete = true
		state.startTime = time.Now()
		fmt.Println("System has started")
	}
}
func (sys defaultSystem) Run() {
	state := &systemState{
		urlStem: sys.config.UrlStem,
		loginCredentials: url.Values{"username":[]string{sys.config.Username}, "password":[]string{sys.config.Password}},
		initializationComplete: false,
		mustLogin: true}

	for sig := range sys.signals {
		switch sig {
		case term:
			goto shutdown
		case flush:
			sys.save(state.startTime.Format(timeFmt))
		case tick:
			sys.pollMetrics(state)
		}
	}

	shutdown:
	sys.save(state.startTime.Format(timeFmt))

	fmt.Printf("\nSystem has Cleanly shutdown\n\n[DONE]\n")
}

func (sys defaultSystem) pollMetrics(state *systemState) {
	defer func() {
		if r := recover(); r != nil {
			msg := "Recovering from Panic: %v\n"
			fmt.Printf(msg, r)
			log.Printf(msg, r)
			state.mustLogin = false
		}
	}()

	if (state.mustLogin) {
		loginUrl := state.urlStem+"/j_spring_security_check"
		log.Printf("Start Login: %s \n", loginUrl)
		resp, _ := sys.client.PostForm(loginUrl, state.loginCredentials)

		log.Printf("Login response: %q '%v': \n\n", resp.Status, resp.StatusCode)
		state.mustLogin = false
		if resp.StatusCode > 300 {
			loc, _ := resp.Location()
			if loc == nil || !strings.Contains(loc.Path, "home") {
				log.Panicf("Error %v", loc.Path)
			}
		}
	}
	state.initialize()

	requestTime := time.Now().Unix() - state.startTime.Unix()
	metricsUrl := sys.config.UrlStem+"/monitor/metrics"
	log.Printf("Making Metrics request %s", metricsUrl)
	resp, _ := sys.client.Get(metricsUrl)
	log.Printf("Metrics response: %q '%v'\n", resp.Status, resp.StatusCode)
	if resp.StatusCode > 300 {
		log.Panicf("Error obtaining metrics in: %q: '%v'\n", resp.Status, resp.StatusCode)
	}

	data, _ := ioutil.ReadAll(resp.Body)
	var jsonData map[string]interface{}

	err := json.Unmarshal(data, &jsonData)

	if err != nil {
		msg := "Metrics response was not valid json %v\n\n"
		log.Panicf(msg, err.Error())
	}
	metrics := Json{jsonData}

	for i, report := range sys.reports {
		log.Printf("Check if report should be updated: %q\n", report.GetName())
		log.Printf("report %d of %d\n", i, len(sys.reports))
		if timeToUpdate(int64(requestTime), report) {
			report.Update(int64(requestTime), metrics)
		}
	}

	if timeToWriteGraphs(requestTime, state.startTime) {
		sys.save(state.startTime.Format(timeFmt))
	}
}

func timeToWriteGraphs(requestTime int64, startTime time.Time) bool {
	timeDiff := (time.Now().Second() - startTime.Second())
	return requestTime > 60 && timeDiff == 0
}
func timeToUpdate(timeSeconds int64, report Report) bool {
	interval := int64(report.GetUpdateInterval())
	timeNano := timeSeconds * int64(time.Second)
	return timeNano % interval == 0
}

func (sys defaultSystem) save(titleModifier string) {
	tmpDir := path.Join(sys.config.OutputDir, "tmp")
	for _, report := range sys.reports {
		report.Save(titleModifier, tmpDir)
	}

	filepath.Walk(tmpDir, func(file string, info os.FileInfo, err error) error {
		if err == nil && !info.IsDir() {
			rel, err := filepath.Rel(tmpDir, file)
			if err == nil {
				dest := path.Join(sys.config.OutputDir, rel)
				os.Remove(dest)
				if _, err := os.Stat(dest); os.IsNotExist(err) {
					os.MkdirAll(filepath.Dir(dest), os.FileMode(0755))
					err = os.Rename(file, dest)
					if err == nil {
						log.Printf("Moved %s to %s\n", file, dest)
					} else {
						log.Printf("Error occurred when attempting to move %s to %s:\n%q\n", file, dest, err.Error())
						copyMove(file, dest)
					}
				} else {
					copyMove(file, dest)
				}
			}
		}
		return nil
	})

	os.RemoveAll(tmpDir)
}

func copyMove(source, dest string){
	sFile, err := os.Create(source)

	if err != nil {
		log.Printf("Failed to open/create source file %s in copy@sytem.go:\n%q\n", source, err.Error())
		return
	}
	defer sFile.Close()

	dFile, err := os.Create(dest)
	if err != nil {
		log.Printf("Failed to open/create dest file %s in copy@sytem.go:\n%q\n", dest, err.Error())
		return
	}
	defer dFile.Close()

	if _, err = io.Copy(dFile, sFile); err != nil {
		log.Printf("Error occurred trying to copy %s to %s:\n%q\n", source, dest, err.Error())
		return
	}
	os.Remove(source)
}<commit_msg>Add debug logs and attempt to improve robustness of updating graphs.  Trying to protect against panics<commit_after>package gnmsys

import (
	"time"
	"net/http/cookiejar"
	"net/http"
	"log"
	"net/url"
	"strings"
	"fmt"
	"io/ioutil"
	"encoding/json"
	"os"
	"path"
	"path/filepath"
	"io"
	"github.com/gonum/plot"
	"runtime/debug"
)

type System interface {
	// function that will signal to the system to clean up and shutdown
	SignalTerm()
	// function that will write the reports to disk
	SignalFlush()
	// List all the Reports available
	GetReports() []Report
	// Get the output file of the report
	GetReportFile(report Report) string
}

type SystemSignal int

const (
	term = 1 + iota
	flush
	tick
)
type SysConfig struct {
	UrlStem,
	Username,
	Password,
	OutputDir string
	SampleConfigs []SampleConfig
}
type defaultSystem struct {
	config SysConfig
	signals chan SystemSignal
	client  http.Client
	reports []Report
}

func CreateSystem(config SysConfig, reportFactories ...ReportFactory) defaultSystem {
	if config.OutputDir == "" || config.OutputDir == "." {
		config.OutputDir = "gnm_reports"
	} else {
		os.MkdirAll(config.OutputDir, os.FileMode(0755))
	}
	options := cookiejar.Options{}
	jar, err := cookiejar.New(&options)
	if err != nil {
		log.Fatal(err)
	}
	reports := make([]Report, len(reportFactories) * len(config.SampleConfigs))
	if (len(reports) == 0) {
		log.Fatalf("No reports are configured\nReport Factories: %d\nSampleConfigs: %v\n", len(reportFactories), config.SampleConfigs)
	}
	for i, fac := range reportFactories {
		for j, sConf := range config.SampleConfigs {
			reports[(i * len(config.SampleConfigs)) + j] = fac(sConf)
		}
	}
	system := defaultSystem{
		config: config,
		reports: reports,
		signals: make(chan SystemSignal),
		client: http.Client{Jar: jar}}

	system.validate()

	go loop(system.signals)

	return system
}

func (sys defaultSystem) validate() {
	for _, conf := range sys.config.SampleConfigs {
		conf.Validate()
	}

	probeFile := filepath.Join(sys.config.OutputDir, "probe")
	if err := ioutil.WriteFile(probeFile, []byte("t"), os.FileMode(0664)); err != nil {
		log.Fatalf("Do not have write permissions to %s\n", sys.config.OutputDir)
	}
	os.Remove(probeFile)

	if _, err := plot.New(); err != nil {
		fmt.Printf("Error creating a test graph %q\n", err.Error())
		msg := "A likely problem is that the directory with font files cannot be found." +
		"  Copy https://github.com/gonum/plot/tree/master/vg/fonts to same directory as the executable"
		fmt.Printf(msg)

		log.Printf("Error creating a test graph %q\n", err.Error())
		log.Fatalf(msg)
	}
}

func loop(signals chan <- SystemSignal) {
	for {
		signals <- tick
		time.Sleep(time.Second)
	}
}

func (sys defaultSystem) GetReportFile(report Report) string {
	var catDirName string
	for _, sampConf := range sys.config.SampleConfigs {
		if sampConf.Name == report.GetCategory() {
			catDirName = sampConf.DirName
			break;
		}
	}
	return filepath.Join(sys.config.OutputDir, catDirName, report.GetFileName())
}
func (sys defaultSystem) GetReports() []Report {
	return sys.reports
}
func (sys defaultSystem) SignalTerm() {
	fmt.Printf("Finalizing reports and shutting down...\n")
	sys.signals <- term
}
func (sys defaultSystem) SignalFlush() {
	fmt.Printf("Saving/Flushing Reports to disk\n")
	sys.signals <- flush
}

type systemState struct {
	startTime time.Time
	initializationComplete, mustLogin bool
	urlStem string
	loginCredentials url.Values
}
const timeFmt = "Start Time: 2006 Jan _2 15:04:05"
func (state *systemState) initialize() {
	if !state.initializationComplete {
		state.initializationComplete = true
		state.startTime = time.Now()
		fmt.Println("System has started")
	}
}
func (sys defaultSystem) Run() {
	state := &systemState{
		urlStem: sys.config.UrlStem,
		loginCredentials: url.Values{"username":[]string{sys.config.Username}, "password":[]string{sys.config.Password}},
		initializationComplete: false,
		mustLogin: true}

	for sig := range sys.signals {
		switch sig {
		case term:
			goto shutdown
		case flush:
			sys.save(state.startTime.Format(timeFmt))
		case tick:
			sys.pollMetrics(state)
		}
	}

	shutdown:
	sys.save(state.startTime.Format(timeFmt))

	fmt.Printf("\nSystem has Cleanly shutdown\n\n[DONE]\n")
}

func (sys defaultSystem) pollMetrics(state *systemState) {
	defer func() {
		if r := recover(); r != nil {
			msg := "Recovering from Panic: %v\n"
			fmt.Printf(msg, r)
			log.Printf(msg, r)
			debug.PrintStack()
			state.mustLogin = false
		}
	}()

	if (state.mustLogin) {
		loginUrl := state.urlStem+"/j_spring_security_check"
		log.Printf("Start Login: %s \n", loginUrl)
		resp, _ := sys.client.PostForm(loginUrl, state.loginCredentials)

		log.Printf("Login response: %q '%v': \n\n", resp.Status, resp.StatusCode)
		state.mustLogin = false
		if resp.StatusCode > 300 {
			loc, _ := resp.Location()
			if loc == nil || !strings.Contains(loc.Path, "home") {
				log.Panicf("Error %v", loc.Path)
			}
		}
	}
	state.initialize()

	requestTime := time.Now().Unix() - state.startTime.Unix()
	metricsUrl := sys.config.UrlStem+"/monitor/metrics"
	log.Printf("Making Metrics request %s", metricsUrl)
	resp, _ := sys.client.Get(metricsUrl)
	log.Printf("Metrics response: %q '%v'\n", resp.Status, resp.StatusCode)
	if resp.StatusCode > 300 {
		log.Panicf("Error obtaining metrics in: %q: '%v'\n", resp.Status, resp.StatusCode)
	}

	data, _ := ioutil.ReadAll(resp.Body)
	var jsonData map[string]interface{}

	err := json.Unmarshal(data, &jsonData)

	if err != nil {
		msg := "Metrics response was not valid json %v\n\n"
		log.Panicf(msg, err.Error())
	}
	metrics := Json{jsonData}

	for i, report := range sys.reports {
		log.Printf("Check if report should be updated: %q\n", report.GetName())
		log.Printf("report %d of %d\n", i, len(sys.reports))
		if timeToUpdate(int64(requestTime), report) {
			safeUpdateReport(report, metrics, requestTime)
		}
	}

	if timeToWriteGraphs(requestTime, state.startTime) {
		sys.save(state.startTime.Format(timeFmt))
	}
}

func safeUpdateReport(report Report, metrics Json, requestTime int64) {
	defer func() {
		if r := recover(); r != nil {
			msg := "Recovering from error in safeUpdateReport\n    Report %q\n    Metrics: %v\n    Error %v\n"
			fmt.Printf(msg, report.GetName(), metrics.Data, r)
			log.Printf(msg, r)
		}
	}()
	report.Update(requestTime, metrics)
}

func timeToWriteGraphs(requestTime int64, startTime time.Time) bool {
	timeDiff := (time.Now().Second() - startTime.Second())
	return requestTime > 60 && timeDiff == 0
}
func timeToUpdate(timeSeconds int64, report Report) bool {
	interval := int64(report.GetUpdateInterval())
	timeNano := timeSeconds * int64(time.Second)
	return timeNano % interval == 0
}

func (sys defaultSystem) save(titleModifier string) {
	tmpDir := path.Join(sys.config.OutputDir, "tmp")
	for _, report := range sys.reports {
		report.Save(titleModifier, tmpDir)
	}

	filepath.Walk(tmpDir, func(file string, info os.FileInfo, err error) error {
		if err == nil && !info.IsDir() {
			rel, err := filepath.Rel(tmpDir, file)
			if err == nil {
				dest := path.Join(sys.config.OutputDir, rel)
				os.Remove(dest)
				if _, err := os.Stat(dest); os.IsNotExist(err) {
					os.MkdirAll(filepath.Dir(dest), os.FileMode(0755))
					err = os.Rename(file, dest)
					if err == nil {
						log.Printf("Moved %s to %s\n", file, dest)
					} else {
						log.Printf("Error occurred when attempting to move %s to %s:\n%q\n", file, dest, err.Error())
						copyMove(file, dest)
					}
				} else {
					copyMove(file, dest)
				}
			}
		}
		return nil
	})

	os.RemoveAll(tmpDir)
}

func copyMove(source, dest string){
	sFile, err := os.Create(source)

	if err != nil {
		log.Printf("Failed to open/create source file %s in copy@sytem.go:\n%q\n", source, err.Error())
		return
	}
	defer sFile.Close()

	dFile, err := os.Create(dest)
	if err != nil {
		log.Printf("Failed to open/create dest file %s in copy@sytem.go:\n%q\n", dest, err.Error())
		return
	}
	defer dFile.Close()

	if _, err = io.Copy(dFile, sFile); err != nil {
		log.Printf("Error occurred trying to copy %s to %s:\n%q\n", source, dest, err.Error())
		return
	}
	os.Remove(source)
}<|endoftext|>


<commit_before>package gs

import (
	"archive/tar"
	"compress/gzip"
	"fmt"
	"io"
	"net/http"
	"os"
	"path/filepath"
	"strings"

	assert "github.com/stretchr/testify/require"
	"go.skia.org/infra/go/httputils"
	"go.skia.org/infra/go/testutils"
)

const (
	// GS bucket where we store test data. Add a folder to this bucket
	// with the tests for a particular component.
	GS_TEST_DATA_ROOT_URI = "http://storage.googleapis.com/skia-infra-testdata/"
)

func openUri(uriPath string) (*http.Response, error) {
	uri := GS_TEST_DATA_ROOT_URI + uriPath

	client := httputils.NewTimeoutClient()
	request, err := RequestForStorageURL(uri)
	if err != nil {
		return nil, err
	}

	resp, err := client.Do(request)
	if err != nil {
		return nil, err
	}
	if resp.StatusCode != 200 {
		return nil, fmt.Errorf("Downloading %s failed. Got response status: %d", uri, resp.StatusCode)
	}

	return resp, nil
}

// DownloadTestDataFile downloads a file with test data from Google Storage.
// The uriPath identifies what to download from the test bucket in GS.
// The content must be publicly accessible.
// The file will be downloaded and stored at provided target
// path (regardless of what the original name is).
// If the the uri ends with '.gz' it will be transparently unzipped.
func DownloadTestDataFile(t assert.TestingT, uriPath, targetPath string) error {
	dir, _ := filepath.Split(targetPath)
	if err := os.MkdirAll(dir, 0755); err != nil {
		return err
	}

	resp, err := openUri(uriPath)
	if err != nil {
		return err
	}
	defer testutils.CloseInTest(t, resp.Body)

	// Open the output
	var r io.ReadCloser = resp.Body
	if strings.HasSuffix(uriPath, ".gz") {
		r, err = gzip.NewReader(r)
		if err != nil {
			return err
		}
	}

	f, err := os.Create(targetPath)
	if err != nil {
		return err
	}
	defer testutils.CloseInTest(t, f)
	_, err = io.Copy(f, r)
	return err
}

// DownloadTestDataArchive downloads testfiles that are stored in
// a gz compressed tar archive and decompresses them into the provided
// target directory.
func DownloadTestDataArchive(t assert.TestingT, uriPath, targetDir string) error {
	if !strings.HasSuffix(uriPath, ".tar.gz") {
		return fmt.Errorf("Expected .tar.gz file. But got:%s", uriPath)
	}

	if err := os.MkdirAll(targetDir, 0755); err != nil {
		return err
	}

	resp, err := openUri(uriPath)
	if err != nil {
		return err
	}
	defer testutils.CloseInTest(t, resp.Body)

	// Open the output
	r, err := gzip.NewReader(resp.Body)
	if err != nil {
		return err
	}
	tarReader := tar.NewReader(r)

	for {
		hdr, err := tarReader.Next()
		if err == io.EOF {
			break
		}

		if err != nil {
			return err
		}

		targetPath := filepath.Join(targetDir, hdr.Name)

		if hdr.Typeflag == tar.TypeDir {
			if err := os.MkdirAll(targetPath, 0755); err != nil {
				return err
			}
		} else {
			f, err := os.Create(targetPath)
			if err != nil {
				return err
			}
			_, err = io.Copy(f, tarReader)
			if err != nil {
				return err
			}
			testutils.CloseInTest(t, f)
		}
	}

	return nil
}
<commit_msg>Remove testing as dependencie in skia_correctness<commit_after>package gs

import (
	"archive/tar"
	"compress/gzip"
	"fmt"
	"io"
	"net/http"
	"os"
	"path/filepath"
	"strings"

	assert "github.com/stretchr/testify/require"
	"go.skia.org/infra/go/httputils"
)

const (
	// GS bucket where we store test data. Add a folder to this bucket
	// with the tests for a particular component.
	GS_TEST_DATA_ROOT_URI = "http://storage.googleapis.com/skia-infra-testdata/"
)

func openUri(uriPath string) (*http.Response, error) {
	uri := GS_TEST_DATA_ROOT_URI + uriPath

	client := httputils.NewTimeoutClient()
	request, err := RequestForStorageURL(uri)
	if err != nil {
		return nil, err
	}

	resp, err := client.Do(request)
	if err != nil {
		return nil, err
	}
	if resp.StatusCode != 200 {
		return nil, fmt.Errorf("Downloading %s failed. Got response status: %d", uri, resp.StatusCode)
	}

	return resp, nil
}

// DownloadTestDataFile downloads a file with test data from Google Storage.
// The uriPath identifies what to download from the test bucket in GS.
// The content must be publicly accessible.
// The file will be downloaded and stored at provided target
// path (regardless of what the original name is).
// If the the uri ends with '.gz' it will be transparently unzipped.
func DownloadTestDataFile(t assert.TestingT, uriPath, targetPath string) error {
	dir, _ := filepath.Split(targetPath)
	if err := os.MkdirAll(dir, 0755); err != nil {
		return err
	}

	resp, err := openUri(uriPath)
	if err != nil {
		return err
	}
	defer func() { assert.Nil(t, resp.Body.Close()) }()

	// Open the output
	var r io.ReadCloser = resp.Body
	if strings.HasSuffix(uriPath, ".gz") {
		r, err = gzip.NewReader(r)
		if err != nil {
			return err
		}
	}

	f, err := os.Create(targetPath)
	if err != nil {
		return err
	}
	defer func() { assert.Nil(t, f.Close()) }()
	_, err = io.Copy(f, r)
	return err
}

// DownloadTestDataArchive downloads testfiles that are stored in
// a gz compressed tar archive and decompresses them into the provided
// target directory.
func DownloadTestDataArchive(t assert.TestingT, uriPath, targetDir string) error {
	if !strings.HasSuffix(uriPath, ".tar.gz") {
		return fmt.Errorf("Expected .tar.gz file. But got:%s", uriPath)
	}

	if err := os.MkdirAll(targetDir, 0755); err != nil {
		return err
	}

	resp, err := openUri(uriPath)
	if err != nil {
		return err
	}
	defer func() { assert.Nil(t, resp.Body.Close()) }()

	// Open the output
	r, err := gzip.NewReader(resp.Body)
	if err != nil {
		return err
	}
	tarReader := tar.NewReader(r)

	for {
		hdr, err := tarReader.Next()
		if err == io.EOF {
			break
		}

		if err != nil {
			return err
		}

		targetPath := filepath.Join(targetDir, hdr.Name)

		if hdr.Typeflag == tar.TypeDir {
			if err := os.MkdirAll(targetPath, 0755); err != nil {
				return err
			}
		} else {
			f, err := os.Create(targetPath)
			if err != nil {
				return err
			}
			_, err = io.Copy(f, tarReader)
			if err != nil {
				return err
			}
			defer func() { assert.Nil(t, f.Close()) }()
		}
	}

	return nil
}
<|endoftext|>


<commit_before>package darkfeed

import (
	"testing"

	"./schemas"
	"github.com/stretchr/testify/assert"
)

func TestPrice_SetPrecision(t *testing.T) {
	p := Price{
		Price:     1385216,
		Precision: -4,
		TickSize:  7,
		Currency:  schemas.CurrencyGBP,
	}
	p.SetPrecision(-6, 1)
	assert.Equal(t, 138521600, int(p.Price))
	assert.Equal(t, -6, int(p.Precision))
	assert.Equal(t, 1, int(p.TickSize))
	p.SetPrecision(-2, 5)
	assert.Equal(t, 13850, int(p.Price))
	assert.Equal(t, -2, int(p.Precision))
	assert.Equal(t, 5, int(p.TickSize))
}

func TestPriceFromFloat64(t *testing.T) {
	p := PriceFromFloat64(132.15, -2, 1, schemas.CurrencyUSD)
	assert.Equal(t, 13215, int(p.Price))
	assert.Equal(t, -2, int(p.Precision))
	assert.Equal(t, 1, int(p.TickSize))
	assert.Equal(t, uint8(schemas.CurrencyUSD), p.Currency)
	p2 := PriceFromFloat64(132.155, -3, 5, schemas.CurrencyCAD)
	assert.Equal(t, 132155, int(p2.Price))
	assert.Equal(t, -3, int(p2.Precision))
	assert.Equal(t, 5, int(p2.TickSize))
	assert.Equal(t, uint8(schemas.CurrencyCAD), p2.Currency)
}

func TestPriceFromUInt32(t *testing.T) {
	p := PriceFromUInt32(13852, -2, 1, schemas.CurrencyUSD)
	assert.Equal(t, 13852, int(p.Price))
	assert.Equal(t, -2, int(p.Precision))
	assert.Equal(t, 1, int(p.TickSize))
	assert.Equal(t, uint8(schemas.CurrencyUSD), p.Currency)

	p2 := PriceFromUInt32(1385216, -4, 7, schemas.CurrencyGBP)
	assert.Equal(t, 1385216, int(p2.Price))
	assert.Equal(t, -4, int(p2.Precision))
	assert.Equal(t, 7, int(p2.TickSize))
	assert.Equal(t, uint8(schemas.CurrencyGBP), p2.Currency)

}

func TestPrice_AsInt(t *testing.T) {
	p := PriceFromUInt32(1385221, -4, 7, schemas.CurrencyGBP)
	assert.Equal(t, 1385223, int(p.Price))
}

func TestPrice_AsFloat32(t *testing.T) {
	p := PriceFromUInt32(1385221, -4, 7, schemas.CurrencyGBP)
	fp := p.AsFloat32()
	assert.Equal(t, float32(138.5223), fp)
}

func TestPrice_AsFloat64(t *testing.T) {
	p := PriceFromUInt32(1385221, -4, 7, schemas.CurrencyGBP)
	fp := p.AsFloat64()
	assert.Equal(t, float64(138.5223), fp)
}

func TestPrice_LessThan(t *testing.T) {
	p := PriceFromFloat64(138.50, -2, 1, schemas.CurrencyUSD)
	gbp := PriceFromFloat64(138.50, -2, 1, schemas.CurrencyGBP)
	assert.False(t, p.LessThan(gbp))

	//higher price, equal precision
	p2 := PriceFromFloat64(138.52, -2, 1, schemas.CurrencyUSD)
	assert.True(t, p.LessThan(p2))
	//lower price, equal precision
	p3 := PriceFromFloat64(138.49, -2, 1, schemas.CurrencyUSD)
	assert.False(t, p.LessThan(p3))
	//higher price, greater precision
	p4 := PriceFromFloat64(138.52, -4, 1, schemas.CurrencyUSD)
	assert.True(t, p.LessThan(p4))
	//lower price, greater precision
	p5 := PriceFromFloat64(138.49, -4, 1, schemas.CurrencyUSD)
	assert.False(t, p.LessThan(p5))
	//higher price, lower precision
	p6 := PriceFromFloat64(138.6, -1, 1, schemas.CurrencyUSD)
	assert.True(t, p.LessThan(p6))
	//lower price, lower precision
	p7 := PriceFromFloat64(138.49, -1, 1, schemas.CurrencyUSD)
	assert.False(t, p.LessThan(p7))
}

func TestPrice_GreaterThan(t *testing.T) {
	p := PriceFromFloat64(138.50, -2, 1, schemas.CurrencyUSD)
	gbp := PriceFromFloat64(138.50, -2, 1, schemas.CurrencyGBP)
	assert.False(t, p.LessThan(gbp))

	//higher price, equal precision
	p2 := PriceFromFloat64(138.52, -2, 1, schemas.CurrencyUSD)
	assert.False(t, p.GreaterThan(p2))
	//lower price, equal precision
	p3 := PriceFromFloat64(138.49, -2, 1, schemas.CurrencyUSD)
	assert.True(t, p.GreaterThan(p3))
	//higher price, greater precision
	p4 := PriceFromFloat64(138.52, -4, 1, schemas.CurrencyUSD)
	assert.False(t, p.GreaterThan(p4))
	//lower price, greater precision
	p5 := PriceFromFloat64(138.49, -4, 1, schemas.CurrencyUSD)
	assert.True(t, p.GreaterThan(p5))
	//higher price, lower precision
	p6 := PriceFromFloat64(138.6, -1, 1, schemas.CurrencyUSD)
	assert.False(t, p.GreaterThan(p6))
	//lower price, lower precision
	p7 := PriceFromFloat64(138.49, -1, 1, schemas.CurrencyUSD)
	assert.True(t, p.GreaterThan(p7))
}

func TestPrice_LessThanEq(t *testing.T) {
	p := PriceFromFloat64(138.50, -2, 1, schemas.CurrencyUSD)
	gbp := PriceFromFloat64(138.50, -2, 1, schemas.CurrencyGBP)
	assert.False(t, p.LessThanEq(gbp))

	//higher price, equal precision
	p2 := PriceFromFloat64(138.52, -2, 1, schemas.CurrencyUSD)
	assert.True(t, p.LessThanEq(p2))
	//lower price, equal precision
	p3 := PriceFromFloat64(138.49, -2, 1, schemas.CurrencyUSD)
	assert.False(t, p.LessThanEq(p3))
	//equal price, equal precision
	p4 := PriceFromFloat64(138.50, -2, 1, schemas.CurrencyUSD)
	assert.True(t, p.LessThanEq(p4))

	//higher price, greater precision
	p5 := PriceFromFloat64(138.52, -4, 1, schemas.CurrencyUSD)
	assert.True(t, p.LessThanEq(p5))
	//lower price, greater precision
	p6 := PriceFromFloat64(138.49, -4, 1, schemas.CurrencyUSD)
	assert.False(t, p.LessThanEq(p6))
	//equal price, greater precision
	p7 := PriceFromFloat64(138.50, -4, 1, schemas.CurrencyUSD)
	assert.True(t, p.LessThanEq(p7))

	//higher price, lower precision
	p8 := PriceFromFloat64(138.6, -1, 1, schemas.CurrencyUSD)
	assert.True(t, p.LessThanEq(p8))
	//lower price, lower precision
	p9 := PriceFromFloat64(138.49, -1, 1, schemas.CurrencyUSD)
	assert.False(t, p.LessThan(p9))
	//equal price, lower precision
	p10 := PriceFromFloat64(138.50, -1, 1, schemas.CurrencyUSD)
	assert.True(t, p.LessThanEq(p10))
}

func TestPrice_GreaterThanEq(t *testing.T) {
	p := PriceFromFloat64(138.50, -2, 1, schemas.CurrencyUSD)
	gbp := PriceFromFloat64(138.50, -2, 1, schemas.CurrencyGBP)
	assert.False(t, p.GreaterThanEq(gbp))

	//higher price, equal precision
	p2 := PriceFromFloat64(138.52, -2, 1, schemas.CurrencyUSD)
	assert.False(t, p.GreaterThanEq(p2))
	//lower price, equal precision
	p3 := PriceFromFloat64(138.49, -2, 1, schemas.CurrencyUSD)
	assert.True(t, p.GreaterThanEq(p3))
	//equal price, equal precision
	p4 := PriceFromFloat64(138.50, -2, 1, schemas.CurrencyUSD)
	assert.True(t, p.GreaterThanEq(p4))

	//higher price, greater precision
	p5 := PriceFromFloat64(138.52, -4, 1, schemas.CurrencyUSD)
	assert.False(t, p.GreaterThanEq(p5))
	//lower price, greater precision
	p6 := PriceFromFloat64(138.49, -4, 1, schemas.CurrencyUSD)
	assert.True(t, p.GreaterThanEq(p6))
	//equal price, greater precision
	p7 := PriceFromFloat64(138.50, -4, 1, schemas.CurrencyUSD)
	assert.True(t, p.GreaterThanEq(p7))

	//higher price, lower precision
	p8 := PriceFromFloat64(138.6, -1, 1, schemas.CurrencyUSD)
	assert.False(t, p.GreaterThanEq(p8))
	//lower price, lower precision
	p9 := PriceFromFloat64(138.49, -1, 1, schemas.CurrencyUSD)
	assert.True(t, p.GreaterThanEq(p9))
	//equal price, lower precision
	p10 := PriceFromFloat64(138.50, -1, 1, schemas.CurrencyUSD)
	assert.True(t, p.GreaterThanEq(p10))
}

func TestPrice_Equals(t *testing.T) {
	p := PriceFromFloat64(138.50, -2, 1, schemas.CurrencyUSD)
	gbp := PriceFromFloat64(138.50, -2, 1, schemas.CurrencyGBP)
	assert.False(t, p.Equals(gbp))

	//higher price, equal precision
	p2 := PriceFromFloat64(138.52, -2, 1, schemas.CurrencyUSD)
	assert.False(t, p.Equals(p2))
	//lower price, equal precision
	p3 := PriceFromFloat64(138.49, -2, 1, schemas.CurrencyUSD)
	assert.False(t, p.Equals(p3))
	//equal price, equal precision
	p4 := PriceFromFloat64(138.50, -2, 1, schemas.CurrencyUSD)
	assert.True(t, p.LessThanEq(p4))

	//higher price, greater precision
	p5 := PriceFromFloat64(138.52, -4, 1, schemas.CurrencyUSD)
	assert.False(t, p.Equals(p5))
	//lower price, greater precision
	p6 := PriceFromFloat64(138.49, -4, 1, schemas.CurrencyUSD)
	assert.False(t, p.Equals(p6))
	//equal price, greater precision
	p7 := PriceFromFloat64(138.50, -4, 1, schemas.CurrencyUSD)
	assert.True(t, p.LessThanEq(p7))

	//higher price, lower precision
	p8 := PriceFromFloat64(138.6, -1, 1, schemas.CurrencyUSD)
	assert.False(t, p.Equals(p8))
	//lower price, lower precision
	p9 := PriceFromFloat64(138.49, -1, 1, schemas.CurrencyUSD)
	assert.False(t, p.Equals(p9))
	//equal price, lower precision
	p10 := PriceFromFloat64(138.50, -1, 1, schemas.CurrencyUSD)
	assert.True(t, p.Equals(p10))
}
<commit_msg>[go] refactor tests to use "fb" prefix instead of "schemas" for flatbuffers code<commit_after>package darkfeed

import (
	"./schemas/fb"
	"github.com/stretchr/testify/assert"
	"testing"
)

func TestPrice_SetPrecision(t *testing.T) {
	p := Price{
		Price:     1385216,
		Precision: -4,
		TickSize:  7,
		Currency:  fb.CurrencyGBP,
	}
	p.SetPrecision(-6, 1)
	assert.Equal(t, 138521600, int(p.Price))
	assert.Equal(t, -6, int(p.Precision))
	assert.Equal(t, 1, int(p.TickSize))
	p.SetPrecision(-2, 5)
	assert.Equal(t, 13850, int(p.Price))
	assert.Equal(t, -2, int(p.Precision))
	assert.Equal(t, 5, int(p.TickSize))
}

func TestPriceFromFloat64(t *testing.T) {
	p := PriceFromFloat64(132.15, -2, 1, fb.CurrencyUSD)
	assert.Equal(t, 13215, int(p.Price))
	assert.Equal(t, -2, int(p.Precision))
	assert.Equal(t, 1, int(p.TickSize))
	assert.Equal(t, uint8(fb.CurrencyUSD), p.Currency)
	p2 := PriceFromFloat64(132.155, -3, 5, fb.CurrencyCAD)
	assert.Equal(t, 132155, int(p2.Price))
	assert.Equal(t, -3, int(p2.Precision))
	assert.Equal(t, 5, int(p2.TickSize))
	assert.Equal(t, uint8(fb.CurrencyCAD), p2.Currency)
}

func TestPriceFromUInt32(t *testing.T) {
	p := PriceFromUInt32(13852, -2, 1, fb.CurrencyUSD)
	assert.Equal(t, 13852, int(p.Price))
	assert.Equal(t, -2, int(p.Precision))
	assert.Equal(t, 1, int(p.TickSize))
	assert.Equal(t, uint8(fb.CurrencyUSD), p.Currency)

	p2 := PriceFromUInt32(1385216, -4, 7, fb.CurrencyGBP)
	assert.Equal(t, 1385216, int(p2.Price))
	assert.Equal(t, -4, int(p2.Precision))
	assert.Equal(t, 7, int(p2.TickSize))
	assert.Equal(t, uint8(fb.CurrencyGBP), p2.Currency)

}

func TestPrice_AsInt(t *testing.T) {
	p := PriceFromUInt32(1385221, -4, 7, fb.CurrencyGBP)
	assert.Equal(t, 1385223, int(p.Price))
}

func TestPrice_AsFloat32(t *testing.T) {
	p := PriceFromUInt32(1385221, -4, 7, fb.CurrencyGBP)
	fp := p.AsFloat32()
	assert.Equal(t, float32(138.5223), fp)
}

func TestPrice_AsFloat64(t *testing.T) {
	p := PriceFromUInt32(1385221, -4, 7, fb.CurrencyGBP)
	fp := p.AsFloat64()
	assert.Equal(t, float64(138.5223), fp)
}

func TestPrice_LessThan(t *testing.T) {
	p := PriceFromFloat64(138.50, -2, 1, fb.CurrencyUSD)
	gbp := PriceFromFloat64(138.50, -2, 1, fb.CurrencyGBP)
	assert.False(t, p.LessThan(gbp))

	//higher price, equal precision
	p2 := PriceFromFloat64(138.52, -2, 1, fb.CurrencyUSD)
	assert.True(t, p.LessThan(p2))
	//lower price, equal precision
	p3 := PriceFromFloat64(138.49, -2, 1, fb.CurrencyUSD)
	assert.False(t, p.LessThan(p3))
	//higher price, greater precision
	p4 := PriceFromFloat64(138.52, -4, 1, fb.CurrencyUSD)
	assert.True(t, p.LessThan(p4))
	//lower price, greater precision
	p5 := PriceFromFloat64(138.49, -4, 1, fb.CurrencyUSD)
	assert.False(t, p.LessThan(p5))
	//higher price, lower precision
	p6 := PriceFromFloat64(138.6, -1, 1, fb.CurrencyUSD)
	assert.True(t, p.LessThan(p6))
	//lower price, lower precision
	p7 := PriceFromFloat64(138.49, -1, 1, fb.CurrencyUSD)
	assert.False(t, p.LessThan(p7))
}

func TestPrice_GreaterThan(t *testing.T) {
	p := PriceFromFloat64(138.50, -2, 1, fb.CurrencyUSD)
	gbp := PriceFromFloat64(138.50, -2, 1, fb.CurrencyGBP)
	assert.False(t, p.LessThan(gbp))

	//higher price, equal precision
	p2 := PriceFromFloat64(138.52, -2, 1, fb.CurrencyUSD)
	assert.False(t, p.GreaterThan(p2))
	//lower price, equal precision
	p3 := PriceFromFloat64(138.49, -2, 1, fb.CurrencyUSD)
	assert.True(t, p.GreaterThan(p3))
	//higher price, greater precision
	p4 := PriceFromFloat64(138.52, -4, 1, fb.CurrencyUSD)
	assert.False(t, p.GreaterThan(p4))
	//lower price, greater precision
	p5 := PriceFromFloat64(138.49, -4, 1, fb.CurrencyUSD)
	assert.True(t, p.GreaterThan(p5))
	//higher price, lower precision
	p6 := PriceFromFloat64(138.6, -1, 1, fb.CurrencyUSD)
	assert.False(t, p.GreaterThan(p6))
	//lower price, lower precision
	p7 := PriceFromFloat64(138.49, -1, 1, fb.CurrencyUSD)
	assert.True(t, p.GreaterThan(p7))
}

func TestPrice_LessThanEq(t *testing.T) {
	p := PriceFromFloat64(138.50, -2, 1, fb.CurrencyUSD)
	gbp := PriceFromFloat64(138.50, -2, 1, fb.CurrencyGBP)
	assert.False(t, p.LessThanEq(gbp))

	//higher price, equal precision
	p2 := PriceFromFloat64(138.52, -2, 1, fb.CurrencyUSD)
	assert.True(t, p.LessThanEq(p2))
	//lower price, equal precision
	p3 := PriceFromFloat64(138.49, -2, 1, fb.CurrencyUSD)
	assert.False(t, p.LessThanEq(p3))
	//equal price, equal precision
	p4 := PriceFromFloat64(138.50, -2, 1, fb.CurrencyUSD)
	assert.True(t, p.LessThanEq(p4))

	//higher price, greater precision
	p5 := PriceFromFloat64(138.52, -4, 1, fb.CurrencyUSD)
	assert.True(t, p.LessThanEq(p5))
	//lower price, greater precision
	p6 := PriceFromFloat64(138.49, -4, 1, fb.CurrencyUSD)
	assert.False(t, p.LessThanEq(p6))
	//equal price, greater precision
	p7 := PriceFromFloat64(138.50, -4, 1, fb.CurrencyUSD)
	assert.True(t, p.LessThanEq(p7))

	//higher price, lower precision
	p8 := PriceFromFloat64(138.6, -1, 1, fb.CurrencyUSD)
	assert.True(t, p.LessThanEq(p8))
	//lower price, lower precision
	p9 := PriceFromFloat64(138.49, -1, 1, fb.CurrencyUSD)
	assert.False(t, p.LessThan(p9))
	//equal price, lower precision
	p10 := PriceFromFloat64(138.50, -1, 1, fb.CurrencyUSD)
	assert.True(t, p.LessThanEq(p10))
}

func TestPrice_GreaterThanEq(t *testing.T) {
	p := PriceFromFloat64(138.50, -2, 1, fb.CurrencyUSD)
	gbp := PriceFromFloat64(138.50, -2, 1, fb.CurrencyGBP)
	assert.False(t, p.GreaterThanEq(gbp))

	//higher price, equal precision
	p2 := PriceFromFloat64(138.52, -2, 1, fb.CurrencyUSD)
	assert.False(t, p.GreaterThanEq(p2))
	//lower price, equal precision
	p3 := PriceFromFloat64(138.49, -2, 1, fb.CurrencyUSD)
	assert.True(t, p.GreaterThanEq(p3))
	//equal price, equal precision
	p4 := PriceFromFloat64(138.50, -2, 1, fb.CurrencyUSD)
	assert.True(t, p.GreaterThanEq(p4))

	//higher price, greater precision
	p5 := PriceFromFloat64(138.52, -4, 1, fb.CurrencyUSD)
	assert.False(t, p.GreaterThanEq(p5))
	//lower price, greater precision
	p6 := PriceFromFloat64(138.49, -4, 1, fb.CurrencyUSD)
	assert.True(t, p.GreaterThanEq(p6))
	//equal price, greater precision
	p7 := PriceFromFloat64(138.50, -4, 1, fb.CurrencyUSD)
	assert.True(t, p.GreaterThanEq(p7))

	//higher price, lower precision
	p8 := PriceFromFloat64(138.6, -1, 1, fb.CurrencyUSD)
	assert.False(t, p.GreaterThanEq(p8))
	//lower price, lower precision
	p9 := PriceFromFloat64(138.49, -1, 1, fb.CurrencyUSD)
	assert.True(t, p.GreaterThanEq(p9))
	//equal price, lower precision
	p10 := PriceFromFloat64(138.50, -1, 1, fb.CurrencyUSD)
	assert.True(t, p.GreaterThanEq(p10))
}

func TestPrice_Equals(t *testing.T) {
	p := PriceFromFloat64(138.50, -2, 1, fb.CurrencyUSD)
	gbp := PriceFromFloat64(138.50, -2, 1, fb.CurrencyGBP)
	assert.False(t, p.Equals(gbp))

	//higher price, equal precision
	p2 := PriceFromFloat64(138.52, -2, 1, fb.CurrencyUSD)
	assert.False(t, p.Equals(p2))
	//lower price, equal precision
	p3 := PriceFromFloat64(138.49, -2, 1, fb.CurrencyUSD)
	assert.False(t, p.Equals(p3))
	//equal price, equal precision
	p4 := PriceFromFloat64(138.50, -2, 1, fb.CurrencyUSD)
	assert.True(t, p.LessThanEq(p4))

	//higher price, greater precision
	p5 := PriceFromFloat64(138.52, -4, 1, fb.CurrencyUSD)
	assert.False(t, p.Equals(p5))
	//lower price, greater precision
	p6 := PriceFromFloat64(138.49, -4, 1, fb.CurrencyUSD)
	assert.False(t, p.Equals(p6))
	//equal price, greater precision
	p7 := PriceFromFloat64(138.50, -4, 1, fb.CurrencyUSD)
	assert.True(t, p.LessThanEq(p7))

	//higher price, lower precision
	p8 := PriceFromFloat64(138.6, -1, 1, fb.CurrencyUSD)
	assert.False(t, p.Equals(p8))
	//lower price, lower precision
	p9 := PriceFromFloat64(138.49, -1, 1, fb.CurrencyUSD)
	assert.False(t, p.Equals(p9))
	//equal price, lower precision
	p10 := PriceFromFloat64(138.50, -1, 1, fb.CurrencyUSD)
	assert.True(t, p.Equals(p10))
}
<|endoftext|>


<commit_before>package godiskcache

import (
	"crypto/sha256"
	"encoding/hex"
	"io/ioutil"
	"log"
	"os"
	"time"
) //import

type GoDiskCache struct {
	keys      map[string]cacheFile
	directory string
} //struct

type cacheFile struct {
	fileName string
	lifeTime int
} //struct

type Params struct {
	Directory string
} //struct

func New(p *Params) *GoDiskCache {
	var directory string = os.TempDir()

	if len(p.Directory) > 0 {
		directory = p.Directory
		err := os.Mkdir(directory, 0744)

		if err != nil {
			log.Fatalln(err)
		} //if
	} //if

	return &GoDiskCache{keys: make(map[string]cacheFile), directory: directory}
} //New

func NewParams() *Params {
	return &Params{}
} //NewParams

func (dc *GoDiskCache) Get(key string) (string, error) {
	var err error

	defer func() {
		if rec := recover(); rec != nil {
			log.Println(rec)
		} //if
	}() //func

	//open the cache file
	if file, err := os.Open(dc.directory + dc.keys[key].fileName); err == nil {
		//get stats about the file, need modified time
		if fi, err := file.Stat(); err == nil {
			//check that cache file is still valid
			if int(time.Since(fi.ModTime()).Seconds()) < dc.keys[key].lifeTime {
				//try reading entire file
				if data, err := ioutil.ReadAll(file); err == nil {
					return string(data), err
				} //if
			} //if
		} //if
	} //if

	return "", err
} //Get

func (dc *GoDiskCache) Set(key, data string, lifetime int) error {
	var err error

	defer func() {
		if rec := recover(); rec != nil {
			log.Println(rec)
		} //if
	}() //func

	//convert string to byte slice
	converted := []byte(key)

	//hash the byte slice and return the resulting string
	hasher := sha256.New()
	hasher.Write(converted)
	filename := "godiskcache_" + hex.EncodeToString(hasher.Sum(nil))

	//open the file
	if file, err := os.Create(dc.directory + filename); err == nil {
		_, err = file.Write([]byte(data))
		_ = file.Close()
	} //if

	if err == nil {
		dc.keys[key] = cacheFile{fileName: filename, lifeTime: lifetime}
	} //if

	return err
} //func
<commit_msg>Just printing error from os.Mkdir instead of dying<commit_after>package godiskcache

import (
	"crypto/sha256"
	"encoding/hex"
	"io/ioutil"
	"log"
	"os"
	"time"
) //import

type GoDiskCache struct {
	keys      map[string]cacheFile
	directory string
} //struct

type cacheFile struct {
	fileName string
	lifeTime int
} //struct

type Params struct {
	Directory string
} //struct

func New(p *Params) *GoDiskCache {
	var directory string = os.TempDir()

	if len(p.Directory) > 0 {
		directory = p.Directory
		err := os.Mkdir(directory, 0744)

		if err != nil {
			log.Println(err)
		} //if
	} //if

	return &GoDiskCache{keys: make(map[string]cacheFile), directory: directory}
} //New

func NewParams() *Params {
	return &Params{}
} //NewParams

func (dc *GoDiskCache) Get(key string) (string, error) {
	var err error

	defer func() {
		if rec := recover(); rec != nil {
			log.Println(rec)
		} //if
	}() //func

	//open the cache file
	if file, err := os.Open(dc.directory + dc.keys[key].fileName); err == nil {
		//get stats about the file, need modified time
		if fi, err := file.Stat(); err == nil {
			//check that cache file is still valid
			if int(time.Since(fi.ModTime()).Seconds()) < dc.keys[key].lifeTime {
				//try reading entire file
				if data, err := ioutil.ReadAll(file); err == nil {
					return string(data), err
				} //if
			} //if
		} //if
	} //if

	return "", err
} //Get

func (dc *GoDiskCache) Set(key, data string, lifetime int) error {
	var err error

	defer func() {
		if rec := recover(); rec != nil {
			log.Println(rec)
		} //if
	}() //func

	//convert string to byte slice
	converted := []byte(key)

	//hash the byte slice and return the resulting string
	hasher := sha256.New()
	hasher.Write(converted)
	filename := "godiskcache_" + hex.EncodeToString(hasher.Sum(nil))

	//open the file
	if file, err := os.Create(dc.directory + filename); err == nil {
		_, err = file.Write([]byte(data))
		_ = file.Close()
	} //if

	if err == nil {
		dc.keys[key] = cacheFile{fileName: filename, lifeTime: lifetime}
	} //if

	return err
} //func
<|endoftext|>


<commit_before>// Copyright 2011-2015 visualfc <visualfc@gmail.com>. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package gotest

import (
	"fmt"
	"go/build"
	"os"
	"os/exec"
	"strings"

	"github.com/visualfc/gotools/pkg/command"
)

var Command = &command.Command{
	Run:         runGotest,
	UsageLine:   "gotest -f filename [build/test flags]",
	Short:       "go test go filename",
	Long:        `go test go filename`,
	CustomFlags: true,
}

var testFileName string
var testFileArgs string

//func init() {
//	Command.Flag.StringVar(&testFileName, "f", "", "test go filename")
//}

func runGotest(cmd *command.Command, args []string) error {
	index := -1
	for n, arg := range args {
		if arg == "-f" {
			index = n
			break
		}
	}
	if index >= 0 && index < len(args) {
		testFileName = args[index+1]
		var r []string
		r = append(r, args[0:index]...)
		r = append(r, args[index+2:]...)
		args = r
	}

	if testFileName == "" {
		cmd.Usage()
		return os.ErrInvalid
	}
	if !strings.HasSuffix(testFileName, "_test.go") {
		fmt.Println("The test filename must xxx_test.go")
		return os.ErrInvalid
	}

	pkg, err := build.ImportDir(".", 0)
	if err != nil {
		fmt.Println("import dir error", err)
		return err
	}

	var testFiles []string

	for _, file := range pkg.XTestGoFiles {
		if file == testFileName {
			testFiles = append(testFiles, file)
			break
		}
	}
	for _, file := range pkg.TestGoFiles {
		if file == testFileName {
			testFiles = append(testFiles, pkg.GoFiles...)
			testFiles = append(testFiles, file)
			break
		}
	}

	gobin, err := exec.LookPath("go")
	if err != nil {
		fmt.Println("error look go", err)
		return err
	}

	var testArgs []string
	testArgs = append(testArgs, "test")
	if len(args) > 0 {
		testArgs = append(testArgs, args...)
	}
	testArgs = append(testArgs, testFiles...)

	command := exec.Command(gobin, testArgs...)
	command.Dir = pkg.Dir
	command.Stdin = os.Stdin
	command.Stdout = os.Stdout
	command.Stderr = os.Stderr

	return command.Run()
}
<commit_msg>apply build tags (-tags) to default go build context<commit_after>// Copyright 2011-2015 visualfc <visualfc@gmail.com>. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package gotest

import (
	"fmt"
	"go/build"
	"os"
	"os/exec"
	"strings"

	"github.com/visualfc/gotools/pkg/command"
)

var Command = &command.Command{
	Run:         runGotest,
	UsageLine:   "gotest -f filename [build/test flags]",
	Short:       "go test go filename",
	Long:        `go test go filename`,
	CustomFlags: true,
}

var testFileName string
var testFileArgs string

func init() {
	//Command.Flag.StringVar(&testFileName, "f", "", "test go filename")
	ApplyBuildTags()
}

func runGotest(cmd *command.Command, args []string) error {
	index := -1
	for n, arg := range args {
		if arg == "-f" {
			index = n
			break
		}
	}
	if index >= 0 && index < len(args) {
		testFileName = args[index+1]
		var r []string
		r = append(r, args[0:index]...)
		r = append(r, args[index+2:]...)
		args = r
	}

	if testFileName == "" {
		cmd.Usage()
		return os.ErrInvalid
	}
	if !strings.HasSuffix(testFileName, "_test.go") {
		fmt.Println("The test filename must xxx_test.go")
		return os.ErrInvalid
	}

	pkg, err := build.ImportDir(".", 0)
	if err != nil {
		fmt.Println("import dir error", err)
		return err
	}

	var testFiles []string

	for _, file := range pkg.XTestGoFiles {
		if file == testFileName {
			testFiles = append(testFiles, file)
			break
		}
	}
	for _, file := range pkg.TestGoFiles {
		if file == testFileName {
			testFiles = append(testFiles, pkg.GoFiles...)
			testFiles = append(testFiles, file)
			break
		}
	}

	gobin, err := exec.LookPath("go")
	if err != nil {
		fmt.Println("error look go", err)
		return err
	}

	var testArgs []string
	testArgs = append(testArgs, "test")
	if len(args) > 0 {
		testArgs = append(testArgs, args...)
	}
	testArgs = append(testArgs, testFiles...)

	command := exec.Command(gobin, testArgs...)
	command.Dir = pkg.Dir
	command.Stdin = os.Stdin
	command.Stdout = os.Stdout
	command.Stderr = os.Stderr

	return command.Run()
}

func ApplyBuildTags() {
	nexttag := false
	for _, arg := range os.Args[1:] {
		if nexttag {
			build.Default.BuildTags = strings.Split(arg, ",")
			nexttag = false
			continue
		}
		if arg == "-tags" {
			nexttag = true
		}
	}
}
<|endoftext|>


<commit_before>/*
Package gothic wraps common behaviour when using Goth. This makes it quick, and easy, to get up
and running with Goth. Of course, if you want complete control over how things flow, in regards
to the authentication process, feel free and use Goth directly.

See https://github.com/markbates/goth/examples/main.go to see this in action.
*/
package gothic

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

	"github.com/gorilla/sessions"
	"github.com/markbates/goth"
)

// SessionName is the key used to access the session store.
const SessionName = "_gothic_session"

// Store can/should be set by applications using gothic. The default is a cookie store.
var Store sessions.Store
var defaultStore sessions.Store

var keySet = false

func init() {
	key := []byte(os.Getenv("SESSION_SECRET"))
	keySet = len(key) != 0
	Store = sessions.NewCookieStore([]byte(key))
	defaultStore = Store
}

/*
BeginAuthHandler is a convienence handler for starting the authentication process.
It expects to be able to get the name of the provider from the query parameters
as either "provider" or ":provider".

BeginAuthHandler will redirect the user to the appropriate authentication end-point
for the requested provider.

See https://github.com/markbates/goth/examples/main.go to see this in action.
*/
func BeginAuthHandler(res http.ResponseWriter, req *http.Request) {
	url, err := GetAuthURL(res, req)
	if err != nil {
		res.WriteHeader(http.StatusBadRequest)
		fmt.Fprintln(res, err)
		return
	}

	http.Redirect(res, req, url, http.StatusTemporaryRedirect)
}

// SetState sets the state string associated with the given request.
// This state is sent to the provider and can be retrieved during the
// callback.
var SetState = func() string {
	return "state"
}

// GetState gets the state returned by the provider during the callback.
// This is used to prevent CSRF attacks, see
// http://tools.ietf.org/html/rfc6749#section-10.12
var GetState = func(req *http.Request) string {
	return req.URL.Query().Get("state")
}

/*
GetAuthURL starts the authentication process with the requested provided.
It will return a URL that should be used to send users to.

It expects to be able to get the name of the provider from the query parameters
as either "provider" or ":provider".

I would recommend using the BeginAuthHandler instead of doing all of these steps
yourself, but that's entirely up to you.
*/
func GetAuthURL(res http.ResponseWriter, req *http.Request) (string, error) {

	if !keySet && defaultStore == Store {
		fmt.Println("goth/gothic: no SESSION_SECRET environment variable is set. The default cookie store is not available and any calls will fail. Ignore this warning if you are using a different store.")
	}

	providerName, err := GetProviderName(req)
	if err != nil {
		return "", err
	}

	provider, err := goth.GetProvider(providerName)
	if err != nil {
		return "", err
	}
	sess, err := provider.BeginAuth(SetState())
	if err != nil {
		return "", err
	}

	url, err := sess.GetAuthURL()
	if err != nil {
		return "", err
	}

	session, _ := Store.Get(req, SessionName)
	session.Values[SessionName] = sess.Marshal()
	err = session.Save(req, res)
	if err != nil {
		return "", err
	}

	return url, err
}

/*
CompleteUserAuth does what it says on the tin. It completes the authentication
process and fetches all of the basic information about the user from the provider.

It expects to be able to get the name of the provider from the query parameters
as either "provider" or ":provider".

See https://github.com/markbates/goth/examples/main.go to see this in action.
*/
var CompleteUserAuth = func(res http.ResponseWriter, req *http.Request) (goth.User, error) {

	if !keySet && defaultStore == Store {
		fmt.Println("goth/gothic: no SESSION_SECRET environment variable is set. The default cookie store is not available and any calls will fail. Ignore this warning if you are using a different store.")
	}

	providerName, err := GetProviderName(req)
	if err != nil {
		return goth.User{}, err
	}

	provider, err := goth.GetProvider(providerName)
	if err != nil {
		return goth.User{}, err
	}

	session, _ := Store.Get(req, SessionName)

	if session.Values[SessionName] == nil {
		return goth.User{}, errors.New("could not find a matching session for this request")
	}

	sess, err := provider.UnmarshalSession(session.Values[SessionName].(string))
	if err != nil {
		return goth.User{}, err
	}

	_, err = sess.Authorize(provider, req.URL.Query())

	if err != nil {
		return goth.User{}, err
	}

	return provider.FetchUser(sess)
}

// GetProviderName is a function used to get the name of a provider
// for a given request. By default, this provider is fetched from
// the URL query string. If you provide it in a different way,
// assign your own function to this variable that returns the provider
// name for your request.
var GetProviderName = getProviderName

func getProviderName(req *http.Request) (string, error) {
	provider := req.URL.Query().Get("provider")
	if provider == "" {
		provider = req.URL.Query().Get(":provider")
	}
	if provider == "" {
		return provider, errors.New("you must select a provider")
	}
	return provider, nil
}
<commit_msg>Fix LinkedIn state passing<commit_after>/*
Package gothic wraps common behaviour when using Goth. This makes it quick, and easy, to get up
and running with Goth. Of course, if you want complete control over how things flow, in regards
to the authentication process, feel free and use Goth directly.

See https://github.com/markbates/goth/examples/main.go to see this in action.
*/
package gothic

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

	"github.com/gorilla/sessions"
	"github.com/markbates/goth"
)

// SessionName is the key used to access the session store.
const SessionName = "_gothic_session"

// Store can/should be set by applications using gothic. The default is a cookie store.
var Store sessions.Store
var defaultStore sessions.Store

var keySet = false

func init() {
	key := []byte(os.Getenv("SESSION_SECRET"))
	keySet = len(key) != 0
	Store = sessions.NewCookieStore([]byte(key))
	defaultStore = Store
}

/*
BeginAuthHandler is a convienence handler for starting the authentication process.
It expects to be able to get the name of the provider from the query parameters
as either "provider" or ":provider".

BeginAuthHandler will redirect the user to the appropriate authentication end-point
for the requested provider.

See https://github.com/markbates/goth/examples/main.go to see this in action.
*/
func BeginAuthHandler(res http.ResponseWriter, req *http.Request) {
	url, err := GetAuthURL(res, req)
	if err != nil {
		res.WriteHeader(http.StatusBadRequest)
		fmt.Fprintln(res, err)
		return
	}

	http.Redirect(res, req, url, http.StatusTemporaryRedirect)
}

// SetState sets the state string associated with the given request.
// This state is sent to the provider and can be retrieved during the
// callback.
var SetState = func(req *http.Request) string {
	state := req.URL.Query().Get("state")
	if len(state) > 0 {
		return state
	}
	return "state"
}

// GetState gets the state returned by the provider during the callback.
// This is used to prevent CSRF attacks, see
// http://tools.ietf.org/html/rfc6749#section-10.12
var GetState = func(req *http.Request) string {
	return req.URL.Query().Get("state")
}

/*
GetAuthURL starts the authentication process with the requested provided.
It will return a URL that should be used to send users to.

It expects to be able to get the name of the provider from the query parameters
as either "provider" or ":provider".

I would recommend using the BeginAuthHandler instead of doing all of these steps
yourself, but that's entirely up to you.
*/
func GetAuthURL(res http.ResponseWriter, req *http.Request) (string, error) {

	if !keySet && defaultStore == Store {
		fmt.Println("goth/gothic: no SESSION_SECRET environment variable is set. The default cookie store is not available and any calls will fail. Ignore this warning if you are using a different store.")
	}

	providerName, err := GetProviderName(req)
	if err != nil {
		return "", err
	}

	provider, err := goth.GetProvider(providerName)
	if err != nil {
		return "", err
	}
	sess, err := provider.BeginAuth(SetState(req))
	if err != nil {
		return "", err
	}

	url, err := sess.GetAuthURL()
	if err != nil {
		return "", err
	}

	session, _ := Store.Get(req, SessionName)
	session.Values[SessionName] = sess.Marshal()
	err = session.Save(req, res)
	if err != nil {
		return "", err
	}

	return url, err
}

/*
CompleteUserAuth does what it says on the tin. It completes the authentication
process and fetches all of the basic information about the user from the provider.

It expects to be able to get the name of the provider from the query parameters
as either "provider" or ":provider".

See https://github.com/markbates/goth/examples/main.go to see this in action.
*/
var CompleteUserAuth = func(res http.ResponseWriter, req *http.Request) (goth.User, error) {

	if !keySet && defaultStore == Store {
		fmt.Println("goth/gothic: no SESSION_SECRET environment variable is set. The default cookie store is not available and any calls will fail. Ignore this warning if you are using a different store.")
	}

	providerName, err := GetProviderName(req)
	if err != nil {
		return goth.User{}, err
	}

	provider, err := goth.GetProvider(providerName)
	if err != nil {
		return goth.User{}, err
	}

	session, _ := Store.Get(req, SessionName)

	if session.Values[SessionName] == nil {
		return goth.User{}, errors.New("could not find a matching session for this request")
	}

	sess, err := provider.UnmarshalSession(session.Values[SessionName].(string))
	if err != nil {
		return goth.User{}, err
	}

	_, err = sess.Authorize(provider, req.URL.Query())

	if err != nil {
		return goth.User{}, err
	}

	return provider.FetchUser(sess)
}

// GetProviderName is a function used to get the name of a provider
// for a given request. By default, this provider is fetched from
// the URL query string. If you provide it in a different way,
// assign your own function to this variable that returns the provider
// name for your request.
var GetProviderName = getProviderName

func getProviderName(req *http.Request) (string, error) {
	provider := req.URL.Query().Get("provider")
	if provider == "" {
		provider = req.URL.Query().Get(":provider")
	}
	if provider == "" {
		return provider, errors.New("you must select a provider")
	}
	return provider, nil
}
<|endoftext|>


<commit_before>/*
Package gothic wraps common behaviour when using Goth. This makes it quick, and easy, to get up
and running with Goth. Of course, if you want complete control over how things flow, in regards
to the authentication process, feel free and use Goth directly.

See https://github.com/markbates/goth/blob/master/examples/main.go to see this in action.
*/
package gothic

import (
	"bytes"
	"compress/gzip"
	"context"
	"crypto/rand"
	"encoding/base64"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"os"
	"strings"

	"github.com/gorilla/mux"
	"github.com/gorilla/sessions"
	"github.com/markbates/goth"
)

// SessionName is the key used to access the session store.
const SessionName = "_gothic_session"

// Store can/should be set by applications using gothic. The default is a cookie store.
var Store sessions.Store
var defaultStore sessions.Store

var keySet = false

type key int

// ProviderParamKey can be used as a key in context when passing in a provider
const ProviderParamKey key = iota

func init() {
	key := []byte(os.Getenv("SESSION_SECRET"))
	keySet = len(key) != 0

	cookieStore := sessions.NewCookieStore([]byte(key))
	cookieStore.Options.HttpOnly = true
	Store = cookieStore
	defaultStore = Store
}

/*
BeginAuthHandler is a convenience handler for starting the authentication process.
It expects to be able to get the name of the provider from the query parameters
as either "provider" or ":provider".

BeginAuthHandler will redirect the user to the appropriate authentication end-point
for the requested provider.

See https://github.com/markbates/goth/examples/main.go to see this in action.
*/
func BeginAuthHandler(res http.ResponseWriter, req *http.Request) {
	url, err := GetAuthURL(res, req)
	if err != nil {
		res.WriteHeader(http.StatusBadRequest)
		fmt.Fprintln(res, err)
		return
	}

	http.Redirect(res, req, url, http.StatusTemporaryRedirect)
}

// SetState sets the state string associated with the given request.
// If no state string is associated with the request, one will be generated.
// This state is sent to the provider and can be retrieved during the
// callback.
var SetState = func(req *http.Request) string {
	state := req.URL.Query().Get("state")
	if len(state) > 0 {
		return state
	}

	// If a state query param is not passed in, generate a random
	// base64-encoded nonce so that the state on the auth URL
	// is unguessable, preventing CSRF attacks, as described in
	//
	// https://auth0.com/docs/protocols/oauth2/oauth-state#keep-reading
	nonceBytes := make([]byte, 64)
	_, err := io.ReadFull(rand.Reader, nonceBytes)
	if err != nil {
		panic("gothic: source of randomness unavailable: " + err.Error())
	}
	return base64.URLEncoding.EncodeToString(nonceBytes)
}

// GetState gets the state returned by the provider during the callback.
// This is used to prevent CSRF attacks, see
// http://tools.ietf.org/html/rfc6749#section-10.12
var GetState = func(req *http.Request) string {
	return req.URL.Query().Get("state")
}

/*
GetAuthURL starts the authentication process with the requested provided.
It will return a URL that should be used to send users to.

It expects to be able to get the name of the provider from the query parameters
as either "provider" or ":provider".

I would recommend using the BeginAuthHandler instead of doing all of these steps
yourself, but that's entirely up to you.
*/
func GetAuthURL(res http.ResponseWriter, req *http.Request) (string, error) {
	if !keySet && defaultStore == Store {
		fmt.Println("goth/gothic: no SESSION_SECRET environment variable is set. The default cookie store is not available and any calls will fail. Ignore this warning if you are using a different store.")
	}

	providerName, err := GetProviderName(req)
	if err != nil {
		return "", err
	}

	provider, err := goth.GetProvider(providerName)
	if err != nil {
		return "", err
	}
	sess, err := provider.BeginAuth(SetState(req))
	if err != nil {
		return "", err
	}

	url, err := sess.GetAuthURL()
	if err != nil {
		return "", err
	}

	err = StoreInSession(providerName, sess.Marshal(), req, res)

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

	return url, err
}

/*
CompleteUserAuth does what it says on the tin. It completes the authentication
process and fetches all of the basic information about the user from the provider.

It expects to be able to get the name of the provider from the query parameters
as either "provider" or ":provider".

See https://github.com/markbates/goth/examples/main.go to see this in action.
*/
var CompleteUserAuth = func(res http.ResponseWriter, req *http.Request) (goth.User, error) {
	defer Logout(res, req)
	if !keySet && defaultStore == Store {
		fmt.Println("goth/gothic: no SESSION_SECRET environment variable is set. The default cookie store is not available and any calls will fail. Ignore this warning if you are using a different store.")
	}

	providerName, err := GetProviderName(req)
	if err != nil {
		return goth.User{}, err
	}

	provider, err := goth.GetProvider(providerName)
	if err != nil {
		return goth.User{}, err
	}

	value, err := GetFromSession(providerName, req)
	if err != nil {
		return goth.User{}, err
	}

	sess, err := provider.UnmarshalSession(value)
	if err != nil {
		return goth.User{}, err
	}

	err = validateState(req, sess)
	if err != nil {
		return goth.User{}, err
	}

	user, err := provider.FetchUser(sess)
	if err == nil {
		// user can be found with existing session data
		return user, err
	}

	params := req.URL.Query()
	if params.Encode() == "" && req.Method == "POST" {
		req.ParseForm()
		params = req.Form
	}

	// get new token and retry fetch
	_, err = sess.Authorize(provider, params)
	if err != nil {
		return goth.User{}, err
	}

	err = StoreInSession(providerName, sess.Marshal(), req, res)

	if err != nil {
		return goth.User{}, err
	}

	gu, err := provider.FetchUser(sess)
	return gu, err
}

// validateState ensures that the state token param from the original
// AuthURL matches the one included in the current (callback) request.
func validateState(req *http.Request, sess goth.Session) error {
	rawAuthURL, err := sess.GetAuthURL()
	if err != nil {
		return err
	}

	authURL, err := url.Parse(rawAuthURL)
	if err != nil {
		return err
	}

	reqState := GetState(req)

	originalState := authURL.Query().Get("state")
	if originalState != "" && (originalState != reqState) {
		return errors.New("state token mismatch")
	}
	return nil
}

// Logout invalidates a user session.
func Logout(res http.ResponseWriter, req *http.Request) error {
	session, err := Store.Get(req, SessionName)
	if err != nil {
		return err
	}
	session.Options.MaxAge = -1
	session.Values = make(map[interface{}]interface{})
	err = session.Save(req, res)
	if err != nil {
		return errors.New("Could not delete user session ")
	}
	return nil
}

// GetProviderName is a function used to get the name of a provider
// for a given request. By default, this provider is fetched from
// the URL query string. If you provide it in a different way,
// assign your own function to this variable that returns the provider
// name for your request.
var GetProviderName = getProviderName

func getProviderName(req *http.Request) (string, error) {

	// try to get it from the url param "provider"
	if p := req.URL.Query().Get("provider"); p != "" {
		return p, nil
	}

	// try to get it from the url param ":provider"
	if p := req.URL.Query().Get(":provider"); p != "" {
		return p, nil
	}

	// try to get it from the context's value of "provider" key
	if p, ok := mux.Vars(req)["provider"]; ok {
		return p, nil
	}

	//  try to get it from the go-context's value of "provider" key
	if p, ok := req.Context().Value("provider").(string); ok {
		return p, nil
	}

	// try to get it from the go-context's value of providerContextKey key
	if p, ok := req.Context().Value(ProviderParamKey).(string); ok {
		return p, nil
	}

	// As a fallback, loop over the used providers, if we already have a valid session for any provider (ie. user has already begun authentication with a provider), then return that provider name
	providers := goth.GetProviders()
	session, _ := Store.Get(req, SessionName)
	for _, provider := range providers {
		p := provider.Name()
		value := session.Values[p]
		if _, ok := value.(string); ok {
			return p, nil
		}
	}

	// if not found then return an empty string with the corresponding error
	return "", errors.New("you must select a provider")
}

// GetContextWithProvider returns a new request context containing the provider
func GetContextWithProvider(req *http.Request, provider string) *http.Request {
	return req.WithContext(context.WithValue(req.Context(), ProviderParamKey, provider))
}

// StoreInSession stores a specified key/value pair in the session.
func StoreInSession(key string, value string, req *http.Request, res http.ResponseWriter) error {
	session, _ := Store.New(req, SessionName)

	if err := updateSessionValue(session, key, value); err != nil {
		return err
	}

	return session.Save(req, res)
}

// GetFromSession retrieves a previously-stored value from the session.
// If no value has previously been stored at the specified key, it will return an error.
func GetFromSession(key string, req *http.Request) (string, error) {
	session, _ := Store.Get(req, SessionName)
	value, err := getSessionValue(session, key)
	if err != nil {
		return "", errors.New("could not find a matching session for this request")
	}

	return value, nil
}

func getSessionValue(session *sessions.Session, key string) (string, error) {
	value := session.Values[key]
	if value == nil {
		return "", fmt.Errorf("could not find a matching session for this request")
	}

	rdata := strings.NewReader(value.(string))
	r, err := gzip.NewReader(rdata)
	if err != nil {
		return "", err
	}
	s, err := ioutil.ReadAll(r)
	if err != nil {
		return "", err
	}

	return string(s), nil
}

func updateSessionValue(session *sessions.Session, key, value string) error {
	var b bytes.Buffer
	gz := gzip.NewWriter(&b)
	if _, err := gz.Write([]byte(value)); err != nil {
		return err
	}
	if err := gz.Flush(); err != nil {
		return err
	}
	if err := gz.Close(); err != nil {
		return err
	}

	session.Values[key] = b.String()
	return nil
}
<commit_msg>fetch state from post param instead of url query param<commit_after>/*
Package gothic wraps common behaviour when using Goth. This makes it quick, and easy, to get up
and running with Goth. Of course, if you want complete control over how things flow, in regards
to the authentication process, feel free and use Goth directly.

See https://github.com/markbates/goth/blob/master/examples/main.go to see this in action.
*/
package gothic

import (
	"bytes"
	"compress/gzip"
	"context"
	"crypto/rand"
	"encoding/base64"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"os"
	"strings"

	"github.com/gorilla/mux"
	"github.com/gorilla/sessions"
	"github.com/markbates/goth"
)

// SessionName is the key used to access the session store.
const SessionName = "_gothic_session"

// Store can/should be set by applications using gothic. The default is a cookie store.
var Store sessions.Store
var defaultStore sessions.Store

var keySet = false

type key int

// ProviderParamKey can be used as a key in context when passing in a provider
const ProviderParamKey key = iota

func init() {
	key := []byte(os.Getenv("SESSION_SECRET"))
	keySet = len(key) != 0

	cookieStore := sessions.NewCookieStore([]byte(key))
	cookieStore.Options.HttpOnly = true
	Store = cookieStore
	defaultStore = Store
}

/*
BeginAuthHandler is a convenience handler for starting the authentication process.
It expects to be able to get the name of the provider from the query parameters
as either "provider" or ":provider".

BeginAuthHandler will redirect the user to the appropriate authentication end-point
for the requested provider.

See https://github.com/markbates/goth/examples/main.go to see this in action.
*/
func BeginAuthHandler(res http.ResponseWriter, req *http.Request) {
	url, err := GetAuthURL(res, req)
	if err != nil {
		res.WriteHeader(http.StatusBadRequest)
		fmt.Fprintln(res, err)
		return
	}

	http.Redirect(res, req, url, http.StatusTemporaryRedirect)
}

// SetState sets the state string associated with the given request.
// If no state string is associated with the request, one will be generated.
// This state is sent to the provider and can be retrieved during the
// callback.
var SetState = func(req *http.Request) string {
	state := req.URL.Query().Get("state")
	if len(state) > 0 {
		return state
	}

	// If a state query param is not passed in, generate a random
	// base64-encoded nonce so that the state on the auth URL
	// is unguessable, preventing CSRF attacks, as described in
	//
	// https://auth0.com/docs/protocols/oauth2/oauth-state#keep-reading
	nonceBytes := make([]byte, 64)
	_, err := io.ReadFull(rand.Reader, nonceBytes)
	if err != nil {
		panic("gothic: source of randomness unavailable: " + err.Error())
	}
	return base64.URLEncoding.EncodeToString(nonceBytes)
}

// GetState gets the state returned by the provider during the callback.
// This is used to prevent CSRF attacks, see
// http://tools.ietf.org/html/rfc6749#section-10.12
var GetState = func(req *http.Request) string {

	params := req.URL.Query()
	if params.Encode() == "" && req.Method == "POST" {
		req.ParseForm()
		params = req.Form
		return params.Get("state")
	}

	return req.URL.Query().Get("state")
}

/*
GetAuthURL starts the authentication process with the requested provided.
It will return a URL that should be used to send users to.

It expects to be able to get the name of the provider from the query parameters
as either "provider" or ":provider".

I would recommend using the BeginAuthHandler instead of doing all of these steps
yourself, but that's entirely up to you.
*/
func GetAuthURL(res http.ResponseWriter, req *http.Request) (string, error) {
	if !keySet && defaultStore == Store {
		fmt.Println("goth/gothic: no SESSION_SECRET environment variable is set. The default cookie store is not available and any calls will fail. Ignore this warning if you are using a different store.")
	}

	providerName, err := GetProviderName(req)
	if err != nil {
		return "", err
	}

	provider, err := goth.GetProvider(providerName)
	if err != nil {
		return "", err
	}
	sess, err := provider.BeginAuth(SetState(req))
	if err != nil {
		return "", err
	}

	url, err := sess.GetAuthURL()
	if err != nil {
		return "", err
	}

	err = StoreInSession(providerName, sess.Marshal(), req, res)

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

	return url, err
}

/*
CompleteUserAuth does what it says on the tin. It completes the authentication
process and fetches all of the basic information about the user from the provider.

It expects to be able to get the name of the provider from the query parameters
as either "provider" or ":provider".

See https://github.com/markbates/goth/examples/main.go to see this in action.
*/
var CompleteUserAuth = func(res http.ResponseWriter, req *http.Request) (goth.User, error) {
	defer Logout(res, req)
	if !keySet && defaultStore == Store {
		fmt.Println("goth/gothic: no SESSION_SECRET environment variable is set. The default cookie store is not available and any calls will fail. Ignore this warning if you are using a different store.")
	}

	providerName, err := GetProviderName(req)
	if err != nil {
		return goth.User{}, err
	}

	provider, err := goth.GetProvider(providerName)
	if err != nil {
		return goth.User{}, err
	}

	value, err := GetFromSession(providerName, req)
	if err != nil {
		return goth.User{}, err
	}

	sess, err := provider.UnmarshalSession(value)
	if err != nil {
		return goth.User{}, err
	}

	err = validateState(req, sess)
	if err != nil {
		return goth.User{}, err
	}

	user, err := provider.FetchUser(sess)
	if err == nil {
		// user can be found with existing session data
		return user, err
	}

	params := req.URL.Query()
	if params.Encode() == "" && req.Method == "POST" {
		req.ParseForm()
		params = req.Form
	}

	// get new token and retry fetch
	_, err = sess.Authorize(provider, params)
	if err != nil {
		return goth.User{}, err
	}

	err = StoreInSession(providerName, sess.Marshal(), req, res)

	if err != nil {
		return goth.User{}, err
	}

	gu, err := provider.FetchUser(sess)
	return gu, err
}

// validateState ensures that the state token param from the original
// AuthURL matches the one included in the current (callback) request.
func validateState(req *http.Request, sess goth.Session) error {
	rawAuthURL, err := sess.GetAuthURL()
	if err != nil {
		return err
	}

	authURL, err := url.Parse(rawAuthURL)
	if err != nil {
		return err
	}

	reqState := GetState(req)

	originalState := authURL.Query().Get("state")
	if originalState != "" && (originalState != reqState) {
		return errors.New("state token mismatch")
	}
	return nil
}

// Logout invalidates a user session.
func Logout(res http.ResponseWriter, req *http.Request) error {
	session, err := Store.Get(req, SessionName)
	if err != nil {
		return err
	}
	session.Options.MaxAge = -1
	session.Values = make(map[interface{}]interface{})
	err = session.Save(req, res)
	if err != nil {
		return errors.New("Could not delete user session ")
	}
	return nil
}

// GetProviderName is a function used to get the name of a provider
// for a given request. By default, this provider is fetched from
// the URL query string. If you provide it in a different way,
// assign your own function to this variable that returns the provider
// name for your request.
var GetProviderName = getProviderName

func getProviderName(req *http.Request) (string, error) {

	// try to get it from the url param "provider"
	if p := req.URL.Query().Get("provider"); p != "" {
		return p, nil
	}

	// try to get it from the url param ":provider"
	if p := req.URL.Query().Get(":provider"); p != "" {
		return p, nil
	}

	// try to get it from the context's value of "provider" key
	if p, ok := mux.Vars(req)["provider"]; ok {
		return p, nil
	}

	//  try to get it from the go-context's value of "provider" key
	if p, ok := req.Context().Value("provider").(string); ok {
		return p, nil
	}

	// try to get it from the go-context's value of providerContextKey key
	if p, ok := req.Context().Value(ProviderParamKey).(string); ok {
		return p, nil
	}

	// As a fallback, loop over the used providers, if we already have a valid session for any provider (ie. user has already begun authentication with a provider), then return that provider name
	providers := goth.GetProviders()
	session, _ := Store.Get(req, SessionName)
	for _, provider := range providers {
		p := provider.Name()
		value := session.Values[p]
		if _, ok := value.(string); ok {
			return p, nil
		}
	}

	// if not found then return an empty string with the corresponding error
	return "", errors.New("you must select a provider")
}

// GetContextWithProvider returns a new request context containing the provider
func GetContextWithProvider(req *http.Request, provider string) *http.Request {
	return req.WithContext(context.WithValue(req.Context(), ProviderParamKey, provider))
}

// StoreInSession stores a specified key/value pair in the session.
func StoreInSession(key string, value string, req *http.Request, res http.ResponseWriter) error {
	session, _ := Store.New(req, SessionName)

	if err := updateSessionValue(session, key, value); err != nil {
		return err
	}

	return session.Save(req, res)
}

// GetFromSession retrieves a previously-stored value from the session.
// If no value has previously been stored at the specified key, it will return an error.
func GetFromSession(key string, req *http.Request) (string, error) {
	session, _ := Store.Get(req, SessionName)
	value, err := getSessionValue(session, key)
	if err != nil {
		return "", errors.New("could not find a matching session for this request")
	}

	return value, nil
}

func getSessionValue(session *sessions.Session, key string) (string, error) {
	value := session.Values[key]
	if value == nil {
		return "", fmt.Errorf("could not find a matching session for this request")
	}

	rdata := strings.NewReader(value.(string))
	r, err := gzip.NewReader(rdata)
	if err != nil {
		return "", err
	}
	s, err := ioutil.ReadAll(r)
	if err != nil {
		return "", err
	}

	return string(s), nil
}

func updateSessionValue(session *sessions.Session, key, value string) error {
	var b bytes.Buffer
	gz := gzip.NewWriter(&b)
	if _, err := gz.Write([]byte(value)); err != nil {
		return err
	}
	if err := gz.Flush(); err != nil {
		return err
	}
	if err := gz.Close(); err != nil {
		return err
	}

	session.Values[key] = b.String()
	return nil
}
<|endoftext|>