<commit_before>package bittorrent

import (
	"errors"
	"net/url"
	"strconv"
	"strings"
)

// Params is used to fetch (optional) request parameters from an Announce.
// For HTTP Announces this includes the request path and parsed query, for UDP
// Announces this is the extracted path and parsed query from optional URLData
// as specified in BEP41.
//
// See ParseURLData for specifics on parsing and limitations.
type Params interface {
	// String returns a string parsed from a query. Every key can be
	// returned as a string because they are encoded in the URL as strings.
	String(key string) (string, bool)

	// RawPath returns the raw path from the request URL.
	// The path returned can contain URL encoded data.
	// For a request of the form "/announce?port=1234" this would return
	// "/announce".
	RawPath() string

	// RawQuery returns the raw query from the request URL, excluding the
	// delimiter '?'.
	// For a request of the form "/announce?port=1234" this would return
	// "port=1234"
	RawQuery() string
}

// ErrKeyNotFound is returned when a provided key has no value associated with
// it.
var ErrKeyNotFound = errors.New("query: value for the provided key does not exist")

// ErrInvalidInfohash is returned when parsing a query encounters an infohash
// with invalid length.
var ErrInvalidInfohash = ClientError("provided invalid infohash")

// QueryParams parses a URL Query and implements the Params interface with some
// additional helpers.
type QueryParams struct {
	path       string
	query      string
	params     map[string]string
	infoHashes []InfoHash
}

// ParseURLData parses a request URL or UDP URLData as defined in BEP41.
// It expects a concatenated string of the request's path and query parts as
// defined in RFC 3986. As both the udp: and http: scheme used by BitTorrent
// include an authority part the path part must always begin with a slash.
// An example of the expected URLData would be "/announce?port=1234&uploaded=0"
// or "/?auth=0x1337".
// HTTP servers should pass (*http.Request).RequestURI, UDP servers should
// pass the concatenated, unchanged URLData as defined in BEP41.
//
// Note that, in the case of a key occurring multiple times in the query, only
// the last value for that key is kept.
// The only exception to this rule is the key "info_hash" which will attempt to
// parse each value as an InfoHash and return an error if parsing fails. All
// InfoHashes are collected and can later be retrieved by calling the InfoHashes
// method.
//
// Also note that any error that is encountered during parsing is returned as a
// ClientError, as this method is expected to be used to parse client-provided
// data.
func ParseURLData(urlData string) (*QueryParams, error) {
	var path, query string

	queryDelim := strings.IndexAny(urlData, "?")
	if queryDelim == -1 {
		path = urlData
	} else {
		path = urlData[:queryDelim]
		query = urlData[queryDelim+1:]
	}

	q, err := parseQuery(query)
	if err != nil {
		return nil, ClientError(err.Error())
	}
	q.path = path
	return q, nil
}

// parseQuery parses a URL query into QueryParams.
// The query is expected to exclude the delimiting '?'.
func parseQuery(query string) (q *QueryParams, err error) {
	// This is basically url.parseQuery, but with a map[string]string
	// instead of map[string][]string for the values.
	q = &QueryParams{
		query:      query,
		infoHashes: nil,
		params:     make(map[string]string),
	}

	for query != "" {
		key := query
		if i := strings.IndexAny(key, "&;"); i >= 0 {
			key, query = key[:i], key[i+1:]
		} else {
			query = ""
		}
		if key == "" {
			continue
		}
		value := ""
		if i := strings.Index(key, "="); i >= 0 {
			key, value = key[:i], key[i+1:]
		}
		key, err = url.QueryUnescape(key)
		if err != nil {
			return nil, err
		}
		value, err = url.QueryUnescape(value)
		if err != nil {
			return nil, err
		}

		if key == "info_hash" {
			if len(value) != 20 {
				return nil, ErrInvalidInfohash
			}
			q.infoHashes = append(q.infoHashes, InfoHashFromString(value))
		} else {
			q.params[strings.ToLower(key)] = value
		}
	}

	return q, nil
}

// String returns a string parsed from a query. Every key can be returned as a
// string because they are encoded in the URL as strings.
func (qp *QueryParams) String(key string) (string, bool) {
	value, ok := qp.params[key]
	return value, ok
}

// Uint64 returns a uint parsed from a query. After being called, it is safe to
// cast the uint64 to your desired length.
func (qp *QueryParams) Uint64(key string) (uint64, error) {
	str, exists := qp.params[key]
	if !exists {
		return 0, ErrKeyNotFound
	}

	val, err := strconv.ParseUint(str, 10, 64)
	if err != nil {
		return 0, err
	}

	return val, nil
}

// InfoHashes returns a list of requested infohashes.
func (qp *QueryParams) InfoHashes() []InfoHash {
	return qp.infoHashes
}

// RawPath returns the raw path from the parsed URL.
func (qp *QueryParams) RawPath() string {
	return qp.path
}

// RawQuery returns the raw query from the parsed URL.
func (qp *QueryParams) RawQuery() string {
	return qp.query
}
<commit_msg>bittorrent: make invalid query escape errors static<commit_after>package bittorrent

import (
	"errors"
	"net/url"
	"strconv"
	"strings"

	log "github.com/Sirupsen/logrus"
)

// Params is used to fetch (optional) request parameters from an Announce.
// For HTTP Announces this includes the request path and parsed query, for UDP
// Announces this is the extracted path and parsed query from optional URLData
// as specified in BEP41.
//
// See ParseURLData for specifics on parsing and limitations.
type Params interface {
	// String returns a string parsed from a query. Every key can be
	// returned as a string because they are encoded in the URL as strings.
	String(key string) (string, bool)

	// RawPath returns the raw path from the request URL.
	// The path returned can contain URL encoded data.
	// For a request of the form "/announce?port=1234" this would return
	// "/announce".
	RawPath() string

	// RawQuery returns the raw query from the request URL, excluding the
	// delimiter '?'.
	// For a request of the form "/announce?port=1234" this would return
	// "port=1234"
	RawQuery() string
}

// ErrKeyNotFound is returned when a provided key has no value associated with
// it.
var ErrKeyNotFound = errors.New("query: value for the provided key does not exist")

// ErrInvalidInfohash is returned when parsing a query encounters an infohash
// with invalid length.
var ErrInvalidInfohash = ClientError("provided invalid infohash")

// ErrInvalidQueryEscape is returned when a query string contains invalid
// escapes.
var ErrInvalidQueryEscape = ClientError("invalid query escape")

// QueryParams parses a URL Query and implements the Params interface with some
// additional helpers.
type QueryParams struct {
	path       string
	query      string
	params     map[string]string
	infoHashes []InfoHash
}

// ParseURLData parses a request URL or UDP URLData as defined in BEP41.
// It expects a concatenated string of the request's path and query parts as
// defined in RFC 3986. As both the udp: and http: scheme used by BitTorrent
// include an authority part the path part must always begin with a slash.
// An example of the expected URLData would be "/announce?port=1234&uploaded=0"
// or "/?auth=0x1337".
// HTTP servers should pass (*http.Request).RequestURI, UDP servers should
// pass the concatenated, unchanged URLData as defined in BEP41.
//
// Note that, in the case of a key occurring multiple times in the query, only
// the last value for that key is kept.
// The only exception to this rule is the key "info_hash" which will attempt to
// parse each value as an InfoHash and return an error if parsing fails. All
// InfoHashes are collected and can later be retrieved by calling the InfoHashes
// method.
//
// Also note that any error that is encountered during parsing is returned as a
// ClientError, as this method is expected to be used to parse client-provided
// data.
func ParseURLData(urlData string) (*QueryParams, error) {
	var path, query string

	queryDelim := strings.IndexAny(urlData, "?")
	if queryDelim == -1 {
		path = urlData
	} else {
		path = urlData[:queryDelim]
		query = urlData[queryDelim+1:]
	}

	q, err := parseQuery(query)
	if err != nil {
		return nil, ClientError(err.Error())
	}
	q.path = path
	return q, nil
}

// parseQuery parses a URL query into QueryParams.
// The query is expected to exclude the delimiting '?'.
func parseQuery(query string) (q *QueryParams, err error) {
	// This is basically url.parseQuery, but with a map[string]string
	// instead of map[string][]string for the values.
	q = &QueryParams{
		query:      query,
		infoHashes: nil,
		params:     make(map[string]string),
	}

	for query != "" {
		key := query
		if i := strings.IndexAny(key, "&;"); i >= 0 {
			key, query = key[:i], key[i+1:]
		} else {
			query = ""
		}
		if key == "" {
			continue
		}
		value := ""
		if i := strings.Index(key, "="); i >= 0 {
			key, value = key[:i], key[i+1:]
		}
		key, err = url.QueryUnescape(key)
		if err != nil {
			// QueryUnescape returns an error like "invalid escape: '%x'".
			// But frontends record these errors to prometheus, which generates
			// a lot of time series.
			// We log it here for debugging instead.
			log.WithFields(log.Fields{"error": err}).Debug("failed to unescape query param key")
			return nil, ErrInvalidQueryEscape
		}
		value, err = url.QueryUnescape(value)
		if err != nil {
			// QueryUnescape returns an error like "invalid escape: '%x'".
			// But frontends record these errors to prometheus, which generates
			// a lot of time series.
			// We log it here for debugging instead.
			log.WithFields(log.Fields{"error": err}).Debug("failed to unescape query param value")
			return nil, ErrInvalidQueryEscape
		}

		if key == "info_hash" {
			if len(value) != 20 {
				return nil, ErrInvalidInfohash
			}
			q.infoHashes = append(q.infoHashes, InfoHashFromString(value))
		} else {
			q.params[strings.ToLower(key)] = value
		}
	}

	return q, nil
}

// String returns a string parsed from a query. Every key can be returned as a
// string because they are encoded in the URL as strings.
func (qp *QueryParams) String(key string) (string, bool) {
	value, ok := qp.params[key]
	return value, ok
}

// Uint64 returns a uint parsed from a query. After being called, it is safe to
// cast the uint64 to your desired length.
func (qp *QueryParams) Uint64(key string) (uint64, error) {
	str, exists := qp.params[key]
	if !exists {
		return 0, ErrKeyNotFound
	}

	val, err := strconv.ParseUint(str, 10, 64)
	if err != nil {
		return 0, err
	}

	return val, nil
}

// InfoHashes returns a list of requested infohashes.
func (qp *QueryParams) InfoHashes() []InfoHash {
	return qp.infoHashes
}

// RawPath returns the raw path from the parsed URL.
func (qp *QueryParams) RawPath() string {
	return qp.path
}

// RawQuery returns the raw query from the parsed URL.
func (qp *QueryParams) RawQuery() string {
	return qp.query
}
<|endoftext|>


<commit_before>package main

import (
	"github.com/buildbox/agent/command"
	"github.com/codegangsta/cli"
)

var Commands []cli.Command

var AgentDescription = `Usage:

   buildbox agent start [arguments...]

Description:

   When a job is ready to run it will call the "bootstrap-script"
   and pass it all the environment variables required for the job to run.
   This script is responsible for checking out the code, and running the
   actual build script defined in the project.

   The agent will run any jobs within a PTY (pseudo terminal) if available.

Example:

   $ buildbox agent start --token xxx`

var DownloadHelpDescription = `Usage:

   buildbox artifact download [arguments...]

Description:

   Downloads artifacts from Buildbox to the local machine.

   You need to ensure that your search query is surrounded by quotes if
   using a wild card as the built-in shell path globbing will provide files,
   which will break the download.

Example:

   $ buildbox artifact download "pkg/*.tar.gz" . --build xxx

   This will search across all the artifacts for the build with files that match that part.
   The first argument is the search query, and the second argument is the download destination.

   If you're trying to download a specific file, and there are multiple artifacts from different
   jobs, you can target the paticular job you want to download the artifact from:

   $ buildbox artifact download "pkg/*.tar.gz" . --job "tests" --build xxx

   You can also use the job's id (provided by the environment variable $BUILDBOX_JOB_ID)`

var UploadHelpDescription = `Usage:

   buildbox artifact upload <pattern> <destination> [arguments...]

Description:

   Uploads files to a job as artifacts.

   You need to ensure that the paths are surrounded by quotes otherwise the
   built-in shell path globbing will provide the files, which is currently not
   supported.

Example:

   $ buildbox artifact upload "log/**/*.log"

   You can also upload directy to Amazon S3 if you'd like to host your own artifacts:

   $ export AWS_SECRET_ACCESS_KEY=yyy
   $ export AWS_ACCESS_KEY_ID=xxx
   $ buildbox artifact upload "log/**/*.log" s3://name-of-your-s3-bucket/$BUILDBOX_JOB_ID`

var SetHelpDescription = `Usage:

   buildbox data set <key> <value> [arguments...]

Description:

   Set arbitrary data on a build using a basic key/value store.

Example:

   $ buildbox data set "foo" "bar"`

var GetHelpDescription = `Usage:

   buildbox data get <key> [arguments...]

Description:

   Get data from a builds key/value store.

Example:

   $ buildbox data get "foo"`

func init() {
	Commands = []cli.Command{
		{
			Name:  "agent",
			Usage: "Starts a Buildbox agent",
			Subcommands: []cli.Command{
				{
					Name:        "start",
					Usage:       "Starts a Buildbox agent",
					Description: AgentDescription,
					Flags: []cli.Flag{
						cli.StringFlag{
							Name:   "token",
							Value:  "",
							Usage:  "Your account agent token",
							EnvVar: "BUILDBOX_AGENT_TOKEN",
						},
						cli.StringFlag{
							Name:   "access-token",
							Value:  "",
							Usage:  "DEPRECATED: The agents access token",
							EnvVar: "BUILDBOX_AGENT_ACCESS_TOKEN",
						},
						cli.StringFlag{
							Name:   "name",
							Value:  "",
							Usage:  "The name of the agent",
							EnvVar: "BUILDBOX_AGENT_NAME",
						},
						cli.StringSliceFlag{
							Name:   "meta-data",
							Value:  &cli.StringSlice{},
							Usage:  "Meta data for the agent",
							EnvVar: "BUILDBOX_AGENT_META_DATA",
						},
						cli.StringFlag{
							Name:   "bootstrap-script",
							Value:  "$HOME/.buildbox/bootstrap.sh",
							Usage:  "Path to the bootstrap script",
							EnvVar: "BUILDBOX_BOOTSTRAP_SCRIPT_PATH",
						},
						cli.StringFlag{
							Name:  "url",
							Value: "https://agent.buildbox.io/v2",
							Usage: "The agent API endpoint",
						},
						cli.BoolFlag{
							Name:  "meta-data-ec2-tags",
							Usage: "Populate the meta data from the current instances EC2 Tags",
						},
						cli.BoolFlag{
							Name:  "debug",
							Usage: "Enable debug mode.",
						},
					},
					Action: command.AgentStartCommandAction,
				},
			},
		},
		{
			Name:  "artifact",
			Usage: "Upload/download artifacts from Buildbox jobs",
			Subcommands: []cli.Command{
				{
					Name:        "download",
					Usage:       "Downloads artifacts from Buildbox to the local machine.",
					Description: DownloadHelpDescription,
					Flags: []cli.Flag{
						cli.StringFlag{
							Name:   "job",
							Value:  "",
							Usage:  "Which job should the artifacts be downloaded from",
							EnvVar: "BUILDBOX_JOB_ID",
						},
						cli.StringFlag{
							Name:   "build",
							Value:  "",
							Usage:  "Which build should the artifacts be downloaded from",
							EnvVar: "BUILDBOX_BUILD_ID",
						},
						cli.StringFlag{
							Name:   "agent-access-token",
							Value:  "",
							Usage:  "The access token used to identify the agent",
							EnvVar: "BUILDBOX_AGENT_ACCESS_TOKEN",
						},
						cli.StringFlag{
							Name:  "url",
							Value: "https://agent.buildbox.io/v2",
							Usage: "The agent API endpoint",
						},
						cli.BoolFlag{
							Name:  "debug",
							Usage: "Enable debug mode",
						},
					},
					Action: command.ArtifactDownloadCommandAction,
				},
				{
					Name:        "upload",
					Usage:       "Uploads files to a job as artifacts.",
					Description: UploadHelpDescription,
					Flags: []cli.Flag{
						cli.StringFlag{
							Name:   "job",
							Value:  "",
							Usage:  "Which job should the artifacts be downloaded from",
							EnvVar: "BUILDBOX_JOB_ID",
						},
						cli.StringFlag{
							Name:   "agent-access-token",
							Value:  "",
							Usage:  "The access token used to identify the agent",
							EnvVar: "BUILDBOX_AGENT_ACCESS_TOKEN",
						},
						cli.StringFlag{
							Name:  "url",
							Value: "https://agent.buildbox.io/v2",
							Usage: "The agent API endpoint",
						},
						cli.BoolFlag{
							Name:  "debug",
							Usage: "Enable debug mode",
						},
					},
					Action: command.ArtifactUploadCommandAction,
				},
			},
		},
		{
			Name:  "data",
			Usage: "Get/set data from Buildbox jobs",
			Subcommands: []cli.Command{
				{
					Name:        "set",
					Usage:       "Set data on a build",
					Description: SetHelpDescription,
					Flags: []cli.Flag{
						cli.StringFlag{
							Name:   "job",
							Value:  "",
							Usage:  "Which job should the artifacts be downloaded from",
							EnvVar: "BUILDBOX_JOB_ID",
						},
						cli.StringFlag{
							Name:   "agent-access-token",
							Value:  "",
							Usage:  "The access token used to identify the agent",
							EnvVar: "BUILDBOX_AGENT_ACCESS_TOKEN",
						},
						cli.StringFlag{
							Name:  "url",
							Value: "https://agent.buildbox.io/v2",
							Usage: "The agent API endpoint",
						},
						cli.BoolFlag{
							Name:  "debug",
							Usage: "Enable debug mode",
						},
					},
					Action: command.DataSetCommandAction,
				},
				{
					Name:        "get",
					Usage:       "Get data from a build",
					Description: GetHelpDescription,
					Flags: []cli.Flag{
						cli.StringFlag{
							Name:   "job",
							Value:  "",
							Usage:  "Which job should the artifacts be downloaded from",
							EnvVar: "BUILDBOX_JOB_ID",
						},
						cli.StringFlag{
							Name:   "agent-access-token",
							Value:  "",
							Usage:  "The access token used to identify the agent",
							EnvVar: "BUILDBOX_AGENT_ACCESS_TOKEN",
						},
						cli.StringFlag{
							Name:  "url",
							Value: "https://agent.buildbox.io/v2",
							Usage: "The agent API endpoint",
						},
						cli.BoolFlag{
							Name:  "debug",
							Usage: "Enable debug mode",
						},
					},
					Action: command.DataGetCommandAction,
				},
			},
		},
	}
}
<commit_msg>Allow debugging and the agent api URL to be set via ENV variables.<commit_after>package main

import (
	"github.com/buildbox/agent/command"
	"github.com/codegangsta/cli"
)

var Commands []cli.Command

var AgentDescription = `Usage:

   buildbox agent start [arguments...]

Description:

   When a job is ready to run it will call the "bootstrap-script"
   and pass it all the environment variables required for the job to run.
   This script is responsible for checking out the code, and running the
   actual build script defined in the project.

   The agent will run any jobs within a PTY (pseudo terminal) if available.

Example:

   $ buildbox agent start --token xxx`

var DownloadHelpDescription = `Usage:

   buildbox artifact download [arguments...]

Description:

   Downloads artifacts from Buildbox to the local machine.

   You need to ensure that your search query is surrounded by quotes if
   using a wild card as the built-in shell path globbing will provide files,
   which will break the download.

Example:

   $ buildbox artifact download "pkg/*.tar.gz" . --build xxx

   This will search across all the artifacts for the build with files that match that part.
   The first argument is the search query, and the second argument is the download destination.

   If you're trying to download a specific file, and there are multiple artifacts from different
   jobs, you can target the paticular job you want to download the artifact from:

   $ buildbox artifact download "pkg/*.tar.gz" . --job "tests" --build xxx

   You can also use the job's id (provided by the environment variable $BUILDBOX_JOB_ID)`

var UploadHelpDescription = `Usage:

   buildbox artifact upload <pattern> <destination> [arguments...]

Description:

   Uploads files to a job as artifacts.

   You need to ensure that the paths are surrounded by quotes otherwise the
   built-in shell path globbing will provide the files, which is currently not
   supported.

Example:

   $ buildbox artifact upload "log/**/*.log"

   You can also upload directy to Amazon S3 if you'd like to host your own artifacts:

   $ export AWS_SECRET_ACCESS_KEY=yyy
   $ export AWS_ACCESS_KEY_ID=xxx
   $ buildbox artifact upload "log/**/*.log" s3://name-of-your-s3-bucket/$BUILDBOX_JOB_ID`

var SetHelpDescription = `Usage:

   buildbox data set <key> <value> [arguments...]

Description:

   Set arbitrary data on a build using a basic key/value store.

Example:

   $ buildbox data set "foo" "bar"`

var GetHelpDescription = `Usage:

   buildbox data get <key> [arguments...]

Description:

   Get data from a builds key/value store.

Example:

   $ buildbox data get "foo"`

func init() {
	Commands = []cli.Command{
		{
			Name:  "agent",
			Usage: "Starts a Buildbox agent",
			Subcommands: []cli.Command{
				{
					Name:        "start",
					Usage:       "Starts a Buildbox agent",
					Description: AgentDescription,
					Flags: []cli.Flag{
						cli.StringFlag{
							Name:   "token",
							Value:  "",
							Usage:  "Your account agent token",
							EnvVar: "BUILDBOX_AGENT_TOKEN",
						},
						cli.StringFlag{
							Name:   "access-token",
							Value:  "",
							Usage:  "DEPRECATED: The agents access token",
							EnvVar: "BUILDBOX_AGENT_ACCESS_TOKEN",
						},
						cli.StringFlag{
							Name:   "name",
							Value:  "",
							Usage:  "The name of the agent",
							EnvVar: "BUILDBOX_AGENT_NAME",
						},
						cli.StringSliceFlag{
							Name:   "meta-data",
							Value:  &cli.StringSlice{},
							Usage:  "Meta data for the agent",
							EnvVar: "BUILDBOX_AGENT_META_DATA",
						},
						cli.StringFlag{
							Name:   "bootstrap-script",
							Value:  "$HOME/.buildbox/bootstrap.sh",
							Usage:  "Path to the bootstrap script",
							EnvVar: "BUILDBOX_BOOTSTRAP_SCRIPT_PATH",
						},
						cli.StringFlag{
							Name:   "url",
							Value:  "https://agent.buildbox.io/v2",
							Usage:  "The agent API endpoint",
							EnvVar: "BUILDBOX_AGENT_API_URL",
						},
						cli.BoolFlag{
							Name:  "meta-data-ec2-tags",
							Usage: "Populate the meta data from the current instances EC2 Tags",
						},
						cli.BoolFlag{
							Name:   "debug",
							Usage:  "Enable debug mode.",
							EnvVar: "BUILDBOX_AGENT_DEBUG",
						},
					},
					Action: command.AgentStartCommandAction,
				},
			},
		},
		{
			Name:  "artifact",
			Usage: "Upload/download artifacts from Buildbox jobs",
			Subcommands: []cli.Command{
				{
					Name:        "download",
					Usage:       "Downloads artifacts from Buildbox to the local machine.",
					Description: DownloadHelpDescription,
					Flags: []cli.Flag{
						cli.StringFlag{
							Name:   "job",
							Value:  "",
							Usage:  "Which job should the artifacts be downloaded from",
							EnvVar: "BUILDBOX_JOB_ID",
						},
						cli.StringFlag{
							Name:   "build",
							Value:  "",
							Usage:  "Which build should the artifacts be downloaded from",
							EnvVar: "BUILDBOX_BUILD_ID",
						},
						cli.StringFlag{
							Name:   "agent-access-token",
							Value:  "",
							Usage:  "The access token used to identify the agent",
							EnvVar: "BUILDBOX_AGENT_ACCESS_TOKEN",
						},
						cli.StringFlag{
							Name:   "url",
							Value:  "https://agent.buildbox.io/v2",
							Usage:  "The agent API endpoint",
							EnvVar: "BUILDBOX_AGENT_API_URL",
						},
						cli.BoolFlag{
							Name:   "debug",
							Usage:  "Enable debug mode",
							EnvVar: "BUILDBOX_AGENT_DEBUG",
						},
					},
					Action: command.ArtifactDownloadCommandAction,
				},
				{
					Name:        "upload",
					Usage:       "Uploads files to a job as artifacts.",
					Description: UploadHelpDescription,
					Flags: []cli.Flag{
						cli.StringFlag{
							Name:   "job",
							Value:  "",
							Usage:  "Which job should the artifacts be downloaded from",
							EnvVar: "BUILDBOX_JOB_ID",
						},
						cli.StringFlag{
							Name:   "agent-access-token",
							Value:  "",
							Usage:  "The access token used to identify the agent",
							EnvVar: "BUILDBOX_AGENT_ACCESS_TOKEN",
						},
						cli.StringFlag{
							Name:   "url",
							Value:  "https://agent.buildbox.io/v2",
							Usage:  "The agent API endpoint",
							EnvVar: "BUILDBOX_AGENT_API_URL",
						},
						cli.BoolFlag{
							Name:   "debug",
							Usage:  "Enable debug mode",
							EnvVar: "BUILDBOX_AGENT_DEBUG",
						},
					},
					Action: command.ArtifactUploadCommandAction,
				},
			},
		},
		{
			Name:  "data",
			Usage: "Get/set data from Buildbox jobs",
			Subcommands: []cli.Command{
				{
					Name:        "set",
					Usage:       "Set data on a build",
					Description: SetHelpDescription,
					Flags: []cli.Flag{
						cli.StringFlag{
							Name:   "job",
							Value:  "",
							Usage:  "Which job should the artifacts be downloaded from",
							EnvVar: "BUILDBOX_JOB_ID",
						},
						cli.StringFlag{
							Name:   "agent-access-token",
							Value:  "",
							Usage:  "The access token used to identify the agent",
							EnvVar: "BUILDBOX_AGENT_ACCESS_TOKEN",
						},
						cli.StringFlag{
							Name:   "url",
							Value:  "https://agent.buildbox.io/v2",
							Usage:  "The agent API endpoint",
							EnvVar: "BUILDBOX_AGENT_API_URL",
						},
						cli.BoolFlag{
							Name:   "debug",
							Usage:  "Enable debug mode",
							EnvVar: "BUILDBOX_AGENT_DEBUG",
						},
					},
					Action: command.DataSetCommandAction,
				},
				{
					Name:        "get",
					Usage:       "Get data from a build",
					Description: GetHelpDescription,
					Flags: []cli.Flag{
						cli.StringFlag{
							Name:   "job",
							Value:  "",
							Usage:  "Which job should the artifacts be downloaded from",
							EnvVar: "BUILDBOX_JOB_ID",
						},
						cli.StringFlag{
							Name:   "agent-access-token",
							Value:  "",
							Usage:  "The access token used to identify the agent",
							EnvVar: "BUILDBOX_AGENT_ACCESS_TOKEN",
						},
						cli.StringFlag{
							Name:   "url",
							Value:  "https://agent.buildbox.io/v2",
							Usage:  "The agent API endpoint",
							EnvVar: "BUILDBOX_AGENT_API_URL",
						},
						cli.BoolFlag{
							Name:   "debug",
							Usage:  "Enable debug mode",
							EnvVar: "BUILDBOX_AGENT_DEBUG",
						},
					},
					Action: command.DataGetCommandAction,
				},
			},
		},
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"strings"
)

type model interface {
	incr(key string) error //increment the given key, setting it to zero if it doesn't exist
	getInt(key string) int //get a key as an int, defaulting to 0 if it doesn't exist
}

// a dabopobo command, consisting of a regex to match against and a commandHandler to run if it matches
type cmd struct {
	regex   string
	handler commandHandler
}

/*
A commandHandler is a function that handles a dabopobo command and returns a slice of bytes to respond with and possibly an error.

m is the data model for this process.
submatches is the output of FindAllStringSubmatch() when running the message text on the corresponding cmd's regex.
username is the username of the user who sent the message.
response is the response to be sent back to slack.
err is non-nil if an error is produced. response should be empty in this case.
*/
type commandHandler func(m model, submatches [][]string, username string) (response string, err error)

var mutateKarmaCmd = cmd{"(\\(.+\\)|[^ ]+?)(\\+\\++|--+|\\+-|-\\+)", mutateKarma}

//handles identifier++
func mutateKarma(m model, mutations [][]string, username string) (s string, err error) {
	if username == "slackbot" { // ignore if message is from the bot
		return
	}
	for _, mutation := range mutations {
		identifier := maybeRemoveParens(mutation[1])
		op := mutation[2]
		if identifier != "" && identifier != username { // users may not mutate themselves
			suffix := canonicalizeSuffix(op)
			key := strings.ToLower(identifier) + suffix // canonicalize identifier as lowercase to prevent confusion
			err = m.incr(key)
			if err != nil {
				err = nil
				return
			}
			fmt.Println(key)
		}
	}
	return
}

var getKarmaCmd = cmd{"^!karma +([^ ].+)", getKarma}

//handles !karma identifier
func getKarma(m model, identifier [][]string, username string) (text string, err error) {
	name := identifier[0][1] //since the regex has a beginning of string hook, there should only be one match, so we only care about index 0.
	karmaset := getKarmaSet(m, name)
	text = fmt.Sprintf("%v's karma is %v %v", name, karmaset.value(), karmaset)
	fmt.Println(text)
	return
}

var helpCmd = cmd{"^!karma(|help)$", help}

func help(m model, s [][]string, u string) (string, error) {
	fmt.Println("help message")
	return strings.Join(
		[]string{
			"!karma thing\tdisplays things karma. It can be anything, including with spaces.",
			"thing++\t(with at least 2 pluses) gives positive karma",
			"thing--\t(with at least 2 minuses) gives negative karma",
			"thing+-\t(either order) gives neutral karma",
			"(thing with spaces)++\t(with any of the above) gives karma to a thing with spaces in it.",
			"!karma or !karmahelp\tdisplays this help",
		},
		"\n",
	), nil
}

var mentionedCmd = cmd{"dabopobo", mentioned}

func mentioned(m model, s [][]string, u string) (string, error) {
	return "don't touch me", nil
}

// getKarmaSet loads the karma for a given key
func getKarmaSet(m model, name string) (k karmaSet) {
	name = strings.ToLower(name)
	k.plusplus = m.getInt(name + "++")
	k.minusminus = m.getInt(name + "--")
	k.plusminus = m.getInt(name + "+-")
	return
}
<commit_msg>Removed redundant bot message filtering<commit_after>package main

import (
	"fmt"
	"strings"
)

type model interface {
	incr(key string) error //increment the given key, setting it to zero if it doesn't exist
	getInt(key string) int //get a key as an int, defaulting to 0 if it doesn't exist
}

// a dabopobo command, consisting of a regex to match against and a commandHandler to run if it matches
type cmd struct {
	regex   string
	handler commandHandler
}

/*
A commandHandler is a function that handles a dabopobo command and returns a slice of bytes to respond with and possibly an error.

m is the data model for this process.
submatches is the output of FindAllStringSubmatch() when running the message text on the corresponding cmd's regex.
username is the username of the user who sent the message.
response is the response to be sent back to slack.
err is non-nil if an error is produced. response should be empty in this case.
*/
type commandHandler func(m model, submatches [][]string, username string) (response string, err error)

var mutateKarmaCmd = cmd{"(\\(.+\\)|[^ ]+?)(\\+\\++|--+|\\+-|-\\+)", mutateKarma}

//handles identifier++
func mutateKarma(m model, mutations [][]string, username string) (s string, err error) {
	for _, mutation := range mutations {
		identifier := maybeRemoveParens(mutation[1])
		op := mutation[2]
		if identifier != "" && identifier != username { // users may not mutate themselves
			suffix := canonicalizeSuffix(op)
			key := strings.ToLower(identifier) + suffix // canonicalize identifier as lowercase to prevent confusion
			err = m.incr(key)
			if err != nil {
				err = nil
				return
			}
			fmt.Println(key)
		}
	}
	return
}

var getKarmaCmd = cmd{"^!karma +([^ ].+)", getKarma}

//handles !karma identifier
func getKarma(m model, identifier [][]string, username string) (text string, err error) {
	name := identifier[0][1] //since the regex has a beginning of string hook, there should only be one match, so we only care about index 0.
	karmaset := getKarmaSet(m, name)
	text = fmt.Sprintf("%v's karma is %v %v", name, karmaset.value(), karmaset)
	fmt.Println(text)
	return
}

var helpCmd = cmd{"^!karma(|help)$", help}

func help(m model, s [][]string, u string) (string, error) {
	fmt.Println("help message")
	return strings.Join(
		[]string{
			"!karma thing\tdisplays things karma. It can be anything, including with spaces.",
			"thing++\t(with at least 2 pluses) gives positive karma",
			"thing--\t(with at least 2 minuses) gives negative karma",
			"thing+-\t(either order) gives neutral karma",
			"(thing with spaces)++\t(with any of the above) gives karma to a thing with spaces in it.",
			"!karma or !karmahelp\tdisplays this help",
		},
		"\n",
	), nil
}

var mentionedCmd = cmd{"dabopobo", mentioned}

func mentioned(m model, s [][]string, u string) (string, error) {
	return "don't touch me", nil
}

// getKarmaSet loads the karma for a given key
func getKarmaSet(m model, name string) (k karmaSet) {
	name = strings.ToLower(name)
	k.plusplus = m.getInt(name + "++")
	k.minusminus = m.getInt(name + "--")
	k.plusminus = m.getInt(name + "+-")
	return
}
<|endoftext|>


<commit_before>package picol

import (
	"fmt"
	"strconv"
)

func ArityErr(i *Interp, name string, argv []string) error {
	return fmt.Errorf("Wrong number of args for %s %s", name, argv)
}

func CommandMath(i *Interp, argv []string, pd interface{}) (string, error) {
	if len(argv) != 3 {
		return "", ArityErr(i, argv[0], argv)
	}
	a, _ := strconv.Atoi(argv[1])
	b, _ := strconv.Atoi(argv[2])
	var c int
	switch {
	case argv[0] == "+":
		c = a + b
	case argv[0] == "-":
		c = a - b
	case argv[0] == "*":
		c = a * b
	case argv[0] == "/":
		c = a / b
	case argv[0] == ">":
		if a > b {
			c = 1
		}
	case argv[0] == ">=":
		if a >= b {
			c = 1
		}
	case argv[0] == "<":
		if a < b {
			c = 1
		}
	case argv[0] == "<=":
		if a <= b {
			c = 1
		}
	case argv[0] == "==":
		if a == b {
			c = 1
		}
	case argv[0] == "!=":
		if a != b {
			c = 1
		}
	default: // FIXME I hate warnings
		c = 0
	}
	return fmt.Sprintf("%d", c), nil
}

func CommandSet(i *Interp, argv []string, pd interface{}) (string, error) {
	if len(argv) != 3 {
		return "", ArityErr(i, argv[0], argv)
	}
	i.SetVar(argv[1], argv[2])
	return argv[2], nil
}

func CommandIf(i *Interp, argv []string, pd interface{}) (string, error) {
	if len(argv) != 3 && len(argv) != 5 {
		return "", ArityErr(i, argv[0], argv)
	}

	result, err := i.Eval(argv[1])
	if err != nil {
		return "", err
	}

	if r, _ := strconv.Atoi(result); r != 0 {
		return i.Eval(argv[2])
	} else if len(argv) == 5 {
		return i.Eval(argv[4])
	}

	return result, nil
}

func CommandWhile(i *Interp, argv []string, pd interface{}) (string, error) {
	if len(argv) != 3 {
		return "", ArityErr(i, argv[0], argv)
	}

	for {
		result, err := i.Eval(argv[1])
		if err != nil {
			return "", err
		}
		if r, _ := strconv.Atoi(result); r != 0 {
			result, err := i.Eval(argv[2])
			switch err {
			case PICOL_CONTINUE, nil:
				//pass
			case PICOL_BREAK:
				return result, nil
			default:
				return result, err
			}
		} else {
			return result, nil
		}
	}
}

func CommandRetCodes(i *Interp, argv []string, pd interface{}) (string, error) {
	if len(argv) != 1 {
		return "", ArityErr(i, argv[0], argv)
	}
	switch argv[0] {
	case "break":
		return "", PICOL_BREAK
	case "continue":
		return "", PICOL_CONTINUE
	}
	return "", nil
}

func DropCallFrame(i *Interp) {
	i.callframe = i.callframe.parent
}

func CommandCallProc(i *Interp, argv []string, pd interface{}) (string, error) {
	var x []string

	if pd, ok := pd.([]string); ok {
		x = pd
	} else {
		return "", nil
	}

	alist := x[0]
	body := x[1]
	p := alist[:]
	arity := 0

	done := false
	i.callframe = &CallFrame{vars: make(map[string]Var), parent: i.callframe}
	defer DropCallFrame(i) // remove the called proc callframe

	for {
		start := p
		for len(p) != 0 && p[0] != ' ' {
			p = p[1:]
		}
		if len(p) != 0 && p == start {
			p = p[1:]
			continue
		}

		if p == start {
			break
		}
		if len(p) == 0 {
			done = true
		} else {
			p = p[1:1]
		}
		arity++
		if arity > len(argv)-1 {
			return "", fmt.Errorf("Proc '%s' called with wrong arg num", argv[0])
		}
		i.SetVar(start, argv[arity])
		if len(p) != 0 {
			p = p[1:]
		}
		if done {
			break
		}
	}

	if arity != len(argv)-1 {
		return "", fmt.Errorf("Proc '%s' called with wrong arg num", argv[0])
	}

	result, err := i.Eval(body)
	if err == PICOL_RETURN {
		err = nil
	}
	return result, err
}

func CommandProc(i *Interp, argv []string, pd interface{}) (string, error) {
	if len(argv) != 4 {
		return "", ArityErr(i, argv[0], argv)
	}
	return "", i.RegisterCommand(argv[1], CommandCallProc, []string{argv[2], argv[3]})
}

func CommandReturn(i *Interp, argv []string, pd interface{}) (string, error) {
	if len(argv) != 1 && len(argv) != 2 {
		return "", ArityErr(i, argv[0], argv)
	}
	var r string
	if len(argv) == 2 {
		r = argv[1]
	}
	return r, PICOL_RETURN
}

func CommandError(i *Interp, argv []string, pd interface{}) (string, error) {
	if len(argv) != 1 && len(argv) != 2 {
		return "", ArityErr(i, argv[0], argv)
	}
	return "", fmt.Errorf(argv[1])
}

func (i *Interp) RegisterCoreCommands() {
	name := [...]string{"+", "-", "*", "/", ">", ">=", "<", "<=", "==", "!="}
	for _, n := range name {
		i.RegisterCommand(n, CommandMath, nil)
	}
	i.RegisterCommand("set", CommandSet, nil)
	i.RegisterCommand("if", CommandIf, nil)
	i.RegisterCommand("while", CommandWhile, nil)
	i.RegisterCommand("break", CommandRetCodes, nil)
	i.RegisterCommand("continue", CommandRetCodes, nil)
	i.RegisterCommand("proc", CommandProc, nil)
	i.RegisterCommand("return", CommandReturn, nil)
	i.RegisterCommand("error", CommandReturn, nil)
}
<commit_msg>fix commands<commit_after>package picol

import (
	"fmt"
	"strconv"
)

func ArityErr(i *Interp, name string, argv []string) error {
	return fmt.Errorf("Wrong number of args for %s %s", name, argv)
}

func CommandMath(i *Interp, argv []string, pd interface{}) (string, error) {
	if len(argv) != 3 {
		return "", ArityErr(i, argv[0], argv)
	}
	a, _ := strconv.Atoi(argv[1])
	b, _ := strconv.Atoi(argv[2])
	var c int
	switch {
	case argv[0] == "+":
		c = a + b
	case argv[0] == "-":
		c = a - b
	case argv[0] == "*":
		c = a * b
	case argv[0] == "/":
		c = a / b
	case argv[0] == ">":
		if a > b {
			c = 1
		}
	case argv[0] == ">=":
		if a >= b {
			c = 1
		}
	case argv[0] == "<":
		if a < b {
			c = 1
		}
	case argv[0] == "<=":
		if a <= b {
			c = 1
		}
	case argv[0] == "==":
		if a == b {
			c = 1
		}
	case argv[0] == "!=":
		if a != b {
			c = 1
		}
	default: // FIXME I hate warnings
		c = 0
	}
	return fmt.Sprintf("%d", c), nil
}

func CommandSet(i *Interp, argv []string, pd interface{}) (string, error) {
	if len(argv) != 3 {
		return "", ArityErr(i, argv[0], argv)
	}
	i.SetVar(argv[1], argv[2])
	return argv[2], nil
}

func CommandIf(i *Interp, argv []string, pd interface{}) (string, error) {
	if len(argv) != 3 && len(argv) != 5 {
		return "", ArityErr(i, argv[0], argv)
	}

	result, err := i.Eval(argv[1])
	if err != nil {
		return "", err
	}

	if r, _ := strconv.Atoi(result); r != 0 {
		return i.Eval(argv[2])
	} else if len(argv) == 5 {
		return i.Eval(argv[4])
	}

	return result, nil
}

func CommandWhile(i *Interp, argv []string, pd interface{}) (string, error) {
	if len(argv) != 3 {
		return "", ArityErr(i, argv[0], argv)
	}

	for {
		result, err := i.Eval(argv[1])
		if err != nil {
			return "", err
		}
		if r, _ := strconv.Atoi(result); r != 0 {
			result, err := i.Eval(argv[2])
			switch err {
			case PICOL_CONTINUE, nil:
				//pass
			case PICOL_BREAK:
				return result, nil
			default:
				return result, err
			}
		} else {
			return result, nil
		}
	}
}

func CommandRetCodes(i *Interp, argv []string, pd interface{}) (string, error) {
	if len(argv) != 1 {
		return "", ArityErr(i, argv[0], argv)
	}
	switch argv[0] {
	case "break":
		return "", PICOL_BREAK
	case "continue":
		return "", PICOL_CONTINUE
	}
	return "", nil
}

func CommandCallProc(i *Interp, argv []string, pd interface{}) (string, error) {
	var x []string

	if pd, ok := pd.([]string); ok {
		x = pd
	} else {
		return "", nil
	}

	alist := x[0]
	body := x[1]
	p := alist[:]
	arity := 0

	fmt.Println("CallProc alist", alist, "body", body)

	done := false
	i.callframe = &CallFrame{vars: make(map[string]Var), parent: i.callframe}
	defer func() { i.callframe = i.callframe.parent }() // remove the called proc callframe

	for {
		start := p
		for len(p) != 0 && p[0] != ' ' {
			p = p[1:]
		}
		if len(p) != 0 && p == start {
			p = p[1:]
			continue
		}

		if p == start {
			break
		}
		if len(p) == 0 {
			done = true
		} else {
			p = p[1:1]
		}
		arity++
		if arity > len(argv)-1 {
			return "", fmt.Errorf("Proc '%s' called with wrong arg num", argv[0])
		}
		i.SetVar(start, argv[arity])
		if len(p) != 0 {
			p = p[1:]
		}
		if done {
			break
		}
	}

	if arity != len(argv)-1 {
		return "", fmt.Errorf("Proc '%s' called with wrong arg num", argv[0])
	}

	result, err := i.Eval(body)
	if err == PICOL_RETURN {
		err = nil
	}
	return result, err
}

func CommandProc(i *Interp, argv []string, pd interface{}) (string, error) {
	if len(argv) != 4 {
		return "", ArityErr(i, argv[0], argv)
	}
	return "", i.RegisterCommand(argv[1], CommandCallProc, []string{argv[2], argv[3]})
}

func CommandReturn(i *Interp, argv []string, pd interface{}) (string, error) {
	if len(argv) != 1 && len(argv) != 2 {
		return "", ArityErr(i, argv[0], argv)
	}
	var r string
	if len(argv) == 2 {
		r = argv[1]
	}
	return r, PICOL_RETURN
}

func CommandError(i *Interp, argv []string, pd interface{}) (string, error) {
	if len(argv) != 1 && len(argv) != 2 {
		return "", ArityErr(i, argv[0], argv)
	}
	return "", fmt.Errorf(argv[1])
}

func (i *Interp) RegisterCoreCommands() {
	name := [...]string{"+", "-", "*", "/", ">", ">=", "<", "<=", "==", "!="}
	for _, n := range name {
		i.RegisterCommand(n, CommandMath, nil)
	}
	i.RegisterCommand("set", CommandSet, nil)
	i.RegisterCommand("if", CommandIf, nil)
	i.RegisterCommand("while", CommandWhile, nil)
	i.RegisterCommand("break", CommandRetCodes, nil)
	i.RegisterCommand("continue", CommandRetCodes, nil)
	i.RegisterCommand("proc", CommandProc, nil)
	i.RegisterCommand("return", CommandReturn, nil)
	i.RegisterCommand("error", CommandError, nil)
}
<|endoftext|>


<commit_before>package engine

import (
	"context"
	"errors"
	"fmt"
	"os"
	"runtime/debug"
	"sync"
	"time"

	"code.cloudfoundry.org/lager"
	"code.cloudfoundry.org/lager/lagerctx"
	"github.com/concourse/concourse/atc"
	"github.com/concourse/concourse/atc/creds"
	"github.com/concourse/concourse/atc/db"
	"github.com/concourse/concourse/atc/event"
	"github.com/concourse/concourse/atc/exec"
	"github.com/concourse/concourse/atc/metric"
	"github.com/concourse/concourse/tracing"
)

//go:generate counterfeiter . Engine

type Engine interface {
	NewBuild(db.Build) Runnable

	Drain(context.Context)
}

//go:generate counterfeiter . Runnable

type Runnable interface {
	Run(context.Context)
}

//go:generate counterfeiter . StepBuilder

type StepBuilder interface {
	BuildStepper(db.Build) (exec.Stepper, error)
}

func NewEngine(
	builder StepBuilder,
	secrets creds.Secrets,
	varSourcePool creds.VarSourcePool,
) Engine {
	return &engine{
		builder:       builder,
		release:       make(chan bool),
		trackedStates: new(sync.Map),
		waitGroup:     new(sync.WaitGroup),

		globalSecrets: secrets,
		varSourcePool: varSourcePool,
	}
}

type engine struct {
	builder       StepBuilder
	release       chan bool
	trackedStates *sync.Map
	waitGroup     *sync.WaitGroup

	globalSecrets creds.Secrets
	varSourcePool creds.VarSourcePool
}

func (engine *engine) Drain(ctx context.Context) {
	logger := lagerctx.FromContext(ctx)

	logger.Info("start")
	defer logger.Info("done")

	close(engine.release)

	logger.Info("waiting")

	engine.waitGroup.Wait()
}

func (engine *engine) NewBuild(build db.Build) Runnable {
	return NewBuild(
		build,
		engine.builder,
		engine.globalSecrets,
		engine.varSourcePool,
		engine.release,
		engine.trackedStates,
		engine.waitGroup,
	)
}

func NewBuild(
	build db.Build,
	builder StepBuilder,
	globalSecrets creds.Secrets,
	varSourcePool creds.VarSourcePool,
	release chan bool,
	trackedStates *sync.Map,
	waitGroup *sync.WaitGroup,
) Runnable {
	return &engineBuild{
		build:   build,
		builder: builder,

		globalSecrets: globalSecrets,
		varSourcePool: varSourcePool,

		release:       release,
		trackedStates: trackedStates,
		waitGroup:     waitGroup,
	}
}

type engineBuild struct {
	build   db.Build
	builder StepBuilder

	globalSecrets creds.Secrets
	varSourcePool creds.VarSourcePool

	release       chan bool
	trackedStates *sync.Map
	waitGroup     *sync.WaitGroup
}

func (b *engineBuild) Run(ctx context.Context) {
	b.waitGroup.Add(1)
	defer b.waitGroup.Done()

	logger := lagerctx.FromContext(ctx).WithData(b.build.LagerData())

	lock, acquired, err := b.build.AcquireTrackingLock(logger, time.Minute)
	if err != nil {
		logger.Error("failed-to-get-lock", err)
		return
	}

	if !acquired {
		logger.Debug("build-already-tracked")
		return
	}

	defer lock.Release()

	found, err := b.build.Reload()
	if err != nil {
		logger.Error("failed-to-load-build-from-db", err)
		return
	}

	if !found {
		logger.Info("build-not-found")
		return
	}

	if !b.build.IsRunning() {
		logger.Info("build-already-finished")
		return
	}

	notifier, err := b.build.AbortNotifier()
	if err != nil {
		logger.Error("failed-to-listen-for-aborts", err)
		return
	}

	defer notifier.Close()

	ctx, span := tracing.StartSpanFollowing(ctx, b.build, "build", b.build.TracingAttrs())
	defer span.End()

	stepper, err := b.builder.BuildStepper(b.build)
	if err != nil {
		logger.Error("failed-to-construct-build-stepper", err)

		// Fails the build if BuildStep returned an error because such unrecoverable
		// errors will cause a build to never start to run.
		b.buildStepErrored(logger, err.Error())
		b.finish(logger.Session("finish"), err, false)

		return
	}

	b.trackStarted(logger)
	defer b.trackFinished(logger)

	logger.Info("running")

	state, err := b.runState(logger, stepper)
	if err != nil {
		logger.Error("failed-to-create-run-state", err)

		// Fails the build if fetching the pipeline variables fails, as these errors
		// are unrecoverable - e.g. if pipeline var_sources is wrong
		b.buildStepErrored(logger, err.Error())
		b.finish(logger.Session("finish"), err, false)

		return
	}
	defer b.clearRunState()

	ctx, cancel := context.WithCancel(ctx)

	noleak := make(chan bool)
	defer close(noleak)

	go func() {
		select {
		case <-noleak:
		case <-notifier.Notify():
			logger.Info("aborting")
			cancel()
		}
	}()

	succeeded := false

	done := make(chan struct{})
	go func() {
		defer close(done)
		defer func() {
			if r := recover(); r != nil {
				err = fmt.Errorf("panic in engine build step run %d: %v", b.build.ID(), r)

				fmt.Fprintf(os.Stderr, "%s\n %s\n", err.Error(), string(debug.Stack()))
				logger.Error("panic-in-engine-build-step-run", err)
			}
		}()

		succeeded, err = state.Run(lagerctx.NewContext(ctx, logger), b.build.PrivatePlan())
	}()

	select {
	case <-b.release:
		logger.Info("releasing")

	case <-done:
		if err != nil {
			if ok := errors.As(err, &exec.Retriable{}); ok {
				return
			}
		}

		b.finish(logger.Session("finish"), err, succeeded)
	}
}

func (b *engineBuild) buildStepErrored(logger lager.Logger, message string) {
	err := b.build.SaveEvent(event.Error{
		Message: message,
		Origin: event.Origin{
			ID: event.OriginID(b.build.PrivatePlan().ID),
		},
		Time: time.Now().Unix(),
	})
	if err != nil {
		logger.Error("failed-to-save-error-event", err)
	}
}

func (b *engineBuild) finish(logger lager.Logger, err error, succeeded bool) {
	if errors.Is(err, context.Canceled) {
		b.saveStatus(logger, atc.StatusAborted)
		logger.Info("aborted")

	} else if err != nil {
		b.saveStatus(logger, atc.StatusErrored)
		logger.Info("errored", lager.Data{"error": err.Error()})

	} else if succeeded {
		b.saveStatus(logger, atc.StatusSucceeded)
		logger.Info("succeeded")

	} else {
		b.saveStatus(logger, atc.StatusFailed)
		logger.Info("failed")
	}
}

func (b *engineBuild) saveStatus(logger lager.Logger, status atc.BuildStatus) {
	if err := b.build.Finish(db.BuildStatus(status)); err != nil {
		logger.Error("failed-to-finish-build", err)
	}
}

func (b *engineBuild) trackStarted(logger lager.Logger) {
	metric.BuildStarted{
		Build: b.build,
	}.Emit(logger)
}

func (b *engineBuild) trackFinished(logger lager.Logger) {
	found, err := b.build.Reload()
	if err != nil {
		logger.Error("failed-to-load-build-from-db", err)
		return
	}

	if !found {
		logger.Info("build-removed")
		return
	}

	if !b.build.IsRunning() {
		metric.BuildFinished{
			Build: b.build,
		}.Emit(logger)
	}
}

func (b *engineBuild) runState(logger lager.Logger, stepper exec.Stepper) (exec.RunState, error) {
	id := fmt.Sprintf("build:%v", b.build.ID())
	existingState, ok := b.trackedStates.Load(id)
	if ok {
		return existingState.(exec.RunState), nil
	}
	credVars, err := b.build.Variables(logger, b.globalSecrets, b.varSourcePool)
	if err != nil {
		return nil, err
	}
	state, _ := b.trackedStates.LoadOrStore(id, exec.NewRunState(stepper, credVars, atc.EnableRedactSecrets))
	return state.(exec.RunState), nil
}

func (b *engineBuild) clearRunState() {
	id := fmt.Sprintf("build:%v", b.build.ID())
	b.trackedStates.Delete(id)
}
<commit_msg>atc: structure: slight var initialization tweak<commit_after>package engine

import (
	"context"
	"errors"
	"fmt"
	"os"
	"runtime/debug"
	"sync"
	"time"

	"code.cloudfoundry.org/lager"
	"code.cloudfoundry.org/lager/lagerctx"
	"github.com/concourse/concourse/atc"
	"github.com/concourse/concourse/atc/creds"
	"github.com/concourse/concourse/atc/db"
	"github.com/concourse/concourse/atc/event"
	"github.com/concourse/concourse/atc/exec"
	"github.com/concourse/concourse/atc/metric"
	"github.com/concourse/concourse/tracing"
)

//go:generate counterfeiter . Engine

type Engine interface {
	NewBuild(db.Build) Runnable

	Drain(context.Context)
}

//go:generate counterfeiter . Runnable

type Runnable interface {
	Run(context.Context)
}

//go:generate counterfeiter . StepBuilder

type StepBuilder interface {
	BuildStepper(db.Build) (exec.Stepper, error)
}

func NewEngine(
	builder StepBuilder,
	secrets creds.Secrets,
	varSourcePool creds.VarSourcePool,
) Engine {
	return &engine{
		builder:       builder,
		release:       make(chan bool),
		trackedStates: new(sync.Map),
		waitGroup:     new(sync.WaitGroup),

		globalSecrets: secrets,
		varSourcePool: varSourcePool,
	}
}

type engine struct {
	builder       StepBuilder
	release       chan bool
	trackedStates *sync.Map
	waitGroup     *sync.WaitGroup

	globalSecrets creds.Secrets
	varSourcePool creds.VarSourcePool
}

func (engine *engine) Drain(ctx context.Context) {
	logger := lagerctx.FromContext(ctx)

	logger.Info("start")
	defer logger.Info("done")

	close(engine.release)

	logger.Info("waiting")

	engine.waitGroup.Wait()
}

func (engine *engine) NewBuild(build db.Build) Runnable {
	return NewBuild(
		build,
		engine.builder,
		engine.globalSecrets,
		engine.varSourcePool,
		engine.release,
		engine.trackedStates,
		engine.waitGroup,
	)
}

func NewBuild(
	build db.Build,
	builder StepBuilder,
	globalSecrets creds.Secrets,
	varSourcePool creds.VarSourcePool,
	release chan bool,
	trackedStates *sync.Map,
	waitGroup *sync.WaitGroup,
) Runnable {
	return &engineBuild{
		build:   build,
		builder: builder,

		globalSecrets: globalSecrets,
		varSourcePool: varSourcePool,

		release:       release,
		trackedStates: trackedStates,
		waitGroup:     waitGroup,
	}
}

type engineBuild struct {
	build   db.Build
	builder StepBuilder

	globalSecrets creds.Secrets
	varSourcePool creds.VarSourcePool

	release       chan bool
	trackedStates *sync.Map
	waitGroup     *sync.WaitGroup
}

func (b *engineBuild) Run(ctx context.Context) {
	b.waitGroup.Add(1)
	defer b.waitGroup.Done()

	logger := lagerctx.FromContext(ctx).WithData(b.build.LagerData())

	lock, acquired, err := b.build.AcquireTrackingLock(logger, time.Minute)
	if err != nil {
		logger.Error("failed-to-get-lock", err)
		return
	}

	if !acquired {
		logger.Debug("build-already-tracked")
		return
	}

	defer lock.Release()

	found, err := b.build.Reload()
	if err != nil {
		logger.Error("failed-to-load-build-from-db", err)
		return
	}

	if !found {
		logger.Info("build-not-found")
		return
	}

	if !b.build.IsRunning() {
		logger.Info("build-already-finished")
		return
	}

	notifier, err := b.build.AbortNotifier()
	if err != nil {
		logger.Error("failed-to-listen-for-aborts", err)
		return
	}

	defer notifier.Close()

	ctx, span := tracing.StartSpanFollowing(ctx, b.build, "build", b.build.TracingAttrs())
	defer span.End()

	stepper, err := b.builder.BuildStepper(b.build)
	if err != nil {
		logger.Error("failed-to-construct-build-stepper", err)

		// Fails the build if BuildStep returned an error because such unrecoverable
		// errors will cause a build to never start to run.
		b.buildStepErrored(logger, err.Error())
		b.finish(logger.Session("finish"), err, false)

		return
	}

	b.trackStarted(logger)
	defer b.trackFinished(logger)

	logger.Info("running")

	state, err := b.runState(logger, stepper)
	if err != nil {
		logger.Error("failed-to-create-run-state", err)

		// Fails the build if fetching the pipeline variables fails, as these errors
		// are unrecoverable - e.g. if pipeline var_sources is wrong
		b.buildStepErrored(logger, err.Error())
		b.finish(logger.Session("finish"), err, false)

		return
	}
	defer b.clearRunState()

	ctx, cancel := context.WithCancel(ctx)

	noleak := make(chan bool)
	defer close(noleak)

	go func() {
		select {
		case <-noleak:
		case <-notifier.Notify():
			logger.Info("aborting")
			cancel()
		}
	}()

	var succeeded bool
	var runErr error

	done := make(chan struct{})
	go func() {
		defer close(done)
		defer func() {
			if r := recover(); r != nil {
				runErr = fmt.Errorf("panic in engine build step run %d: %v", b.build.ID(), r)

				fmt.Fprintf(os.Stderr, "%s\n %s\n", runErr.Error(), string(debug.Stack()))
				logger.Error("panic-in-engine-build-step-run", err)
			}
		}()

		succeeded, runErr = state.Run(lagerctx.NewContext(ctx, logger), b.build.PrivatePlan())
	}()

	select {
	case <-b.release:
		logger.Info("releasing")

	case <-done:
		if errors.As(runErr, &exec.Retriable{}) {
			return
		}

		b.finish(logger.Session("finish"), runErr, succeeded)
	}
}

func (b *engineBuild) buildStepErrored(logger lager.Logger, message string) {
	err := b.build.SaveEvent(event.Error{
		Message: message,
		Origin: event.Origin{
			ID: event.OriginID(b.build.PrivatePlan().ID),
		},
		Time: time.Now().Unix(),
	})
	if err != nil {
		logger.Error("failed-to-save-error-event", err)
	}
}

func (b *engineBuild) finish(logger lager.Logger, err error, succeeded bool) {
	if errors.Is(err, context.Canceled) {
		b.saveStatus(logger, atc.StatusAborted)
		logger.Info("aborted")

	} else if err != nil {
		b.saveStatus(logger, atc.StatusErrored)
		logger.Info("errored", lager.Data{"error": err.Error()})

	} else if succeeded {
		b.saveStatus(logger, atc.StatusSucceeded)
		logger.Info("succeeded")

	} else {
		b.saveStatus(logger, atc.StatusFailed)
		logger.Info("failed")
	}
}

func (b *engineBuild) saveStatus(logger lager.Logger, status atc.BuildStatus) {
	if err := b.build.Finish(db.BuildStatus(status)); err != nil {
		logger.Error("failed-to-finish-build", err)
	}
}

func (b *engineBuild) trackStarted(logger lager.Logger) {
	metric.BuildStarted{
		Build: b.build,
	}.Emit(logger)
}

func (b *engineBuild) trackFinished(logger lager.Logger) {
	found, err := b.build.Reload()
	if err != nil {
		logger.Error("failed-to-load-build-from-db", err)
		return
	}

	if !found {
		logger.Info("build-removed")
		return
	}

	if !b.build.IsRunning() {
		metric.BuildFinished{
			Build: b.build,
		}.Emit(logger)
	}
}

func (b *engineBuild) runState(logger lager.Logger, stepper exec.Stepper) (exec.RunState, error) {
	id := fmt.Sprintf("build:%v", b.build.ID())
	existingState, ok := b.trackedStates.Load(id)
	if ok {
		return existingState.(exec.RunState), nil
	}
	credVars, err := b.build.Variables(logger, b.globalSecrets, b.varSourcePool)
	if err != nil {
		return nil, err
	}
	state, _ := b.trackedStates.LoadOrStore(id, exec.NewRunState(stepper, credVars, atc.EnableRedactSecrets))
	return state.(exec.RunState), nil
}

func (b *engineBuild) clearRunState() {
	id := fmt.Sprintf("build:%v", b.build.ID())
	b.trackedStates.Delete(id)
}
<|endoftext|>


<commit_before>package hub

import (
	"code.google.com/p/go.net/websocket"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"log"
	"math/rand"
	"net/http"
	"os"
	"sync"
	"time"
)

const (
	defaultBufferSize     = 1024
	defaultHeartbeat      = time.Second * 5
	defaultSessionTimeout = time.Second * 30
	idLength              = 16
	bufLength             = 4096
)

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

type Authorizer func(uri, token string, write bool) (authorized bool, err error)

func prettify(obj interface{}) string {
	b, err := json.MarshalIndent(obj, "", "  ")
	if err != nil {
		panic(err)
	}
	return string(b)
}

type Server struct {
	heartbeat      time.Duration
	sessionTimeout time.Duration
	sessions       map[string]*Session
	subscriptions  map[string]map[string]*Session
	subscribers    map[string]map[string]bool
	lock           *sync.RWMutex
	bufferSize     int
	logger         *log.Logger
	authorizer     Authorizer
}

func NewServer() *Server {
	return &Server{
		heartbeat:      defaultHeartbeat,
		bufferSize:     defaultBufferSize,
		sessionTimeout: defaultSessionTimeout,
		sessions:       map[string]*Session{},
		subscriptions:  map[string]map[string]*Session{},
		subscribers:    map[string]map[string]bool{},
		lock:           &sync.RWMutex{},
		logger:         log.New(os.Stdout, "pusher: ", 0),
		authorizer: func(uri, token string, write bool) (bool, error) {
			return true, nil
		},
	}
}

func (self *Server) Authorizer(f Authorizer) *Server {
	self.authorizer = f
	return self
}

func (self *Server) addSubscription(sess *Session, uri string) {
	self.lock.Lock()
	defer self.lock.Unlock()

	if _, found := self.subscriptions[uri]; !found {
		self.subscriptions[uri] = map[string]*Session{}
	}
	self.subscriptions[uri][sess.id] = sess

	if _, found := self.subscribers[sess.id]; !found {
		self.subscribers[sess.id] = map[string]bool{}
	}
	self.subscribers[sess.id][uri] = true

	self.logger.Printf("%v\t[ws]\t%v\t%v\t%v\t[subscribe]", time.Now(), uri, sess.RemoteAddr, sess.id)
}

func (self *Server) Emit(message Message) {
	self.lock.Lock()
	defer self.lock.Unlock()

	for _, sess := range self.subscriptions[message.URI] {
		sess.send(message)
	}
}

func (self *Server) removeSubscription(id, uri string, withLocking bool) {
	if withLocking {
		self.lock.Lock()
		defer self.lock.Unlock()
	}

	delete(self.subscriptions[uri], id)
	if len(self.subscriptions[uri]) == 0 {
		delete(self.subscriptions, uri)
	}

	delete(self.subscribers[id], uri)
	if len(self.subscribers[id]) == 0 {
		delete(self.subscribers, id)
	}
	self.logger.Printf("%v\t[ws]\t%v\t%v\t-\t[unsubscribe]", time.Now(), uri, id)
}

func (self *Server) randomId() string {
	buf := make([]byte, idLength)
	for index, _ := range buf {
		buf[index] = byte(rand.Int31())
	}
	return (base64.URLEncoding.EncodeToString(buf))
}

func (self *Server) removeSession(id string) {
	self.lock.Lock()
	defer self.lock.Unlock()

	delete(self.sessions, id)

	for uri, _ := range self.subscribers[id] {
		self.removeSubscription(id, uri, false)
	}
	self.logger.Printf("%v\t[ws]\t[cleanup]\t%v", time.Now(), id)
}

func (self *Server) getSession(id string) (result *Session) {
	self.lock.Lock()
	defer self.lock.Unlock()

	// if the id is not found (because the server restarted, or the client gave the initial empty id) we just create a new id and insert a new session
	if result = self.sessions[id]; result == nil {
		result = &Session{
			output:         make(chan Message, self.bufferSize),
			id:             self.randomId(),
			server:         self,
			authorizations: map[string]bool{},
			lock:           &sync.RWMutex{},
		}
		result.cleanupTimer = &time.Timer{}
		self.sessions[result.id] = result
	}
	return
}

func (self *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	self.logger.Printf("%v\t%v\t%v\t%v", time.Now(), r.Method, r.URL, r.RemoteAddr)
	websocket.Handler(self.getSession(r.URL.Query().Get("session_id")).handle).ServeHTTP(w, r)
}

type MessageType string
type ErrorType string

const (
	TypeError       = "Error"
	TypeHeartbeat   = "Heartbeat"
	TypeWelcome     = "Welcome"
	TypeSubscribe   = "Subscribe"
	TypeUnsubscribe = "Unsubscribe"
	TypeMessage     = "Message"
	TypeAuthorize   = "Authorize"
	TypeAck         = "Ack"
)

const (
	TypeJSONError          = "JSONError"
	TypeAuthorizationError = "AuthorizationError"
	TypeSyntaxError        = "SyntaxError"
)

type Welcome struct {
	Heartbeat      time.Duration
	SessionTimeout time.Duration
	Id             string
}

type Error struct {
	Message string
	Type    ErrorType
}

type Message struct {
	Type    MessageType
	Id      string      `json:",omitempty"`
	Welcome *Welcome    `json:",omitempty"`
	Error   *Error      `json:",omitempty"`
	Data    interface{} `json:",omitempty"`
	URI     string      `json:",omitempty"`
	Token   string      `json:",omitempty"`
	Write   bool        `json:",omitempty"`
}

type Session struct {
	ws             *websocket.Conn
	id             string
	RemoteAddr     string
	input          chan Message
	output         chan Message
	closing        chan struct{}
	server         *Server
	cleanupTimer   *time.Timer
	authorizations map[string]bool
	lock           *sync.RWMutex
}

func (self *Session) parseMessage(b []byte) (result Message, err error) {
	err = json.Unmarshal(b, &result)
	return
}

func (self *Session) readLoop() {
	defer self.terminate()
	buf := make([]byte, bufLength)
	n, err := self.ws.Read(buf)
	for err == nil {
		if message, err := self.parseMessage(buf[:n]); err == nil {
			self.input <- message
		} else {
			self.send(Message{
				Type: TypeError,
				Error: &Error{
					Message: err.Error(),
					Type:    TypeJSONError,
				},
				Data: string(buf[:n])})
		}
		n, err = self.ws.Read(buf)
	}
}

func (self *Session) writeLoop() {
	defer self.terminate()
	var message Message
	var err error
	var n int
	var encoded []byte
	for {
		select {
		case message = <-self.output:
			if encoded, err = json.Marshal(message); err == nil {
				if n, err = self.ws.Write(encoded); err != nil {
					self.server.logger.Printf("Error sending %s on %+v: %v", encoded, self.ws, err)
					return
				} else if n != len(encoded) {
					self.server.logger.Printf("Unable to send all of %s on %+v: only sent %v bytes", encoded, self.ws, n)
					return
				}
			} else {
				self.server.logger.Printf("Unable to JSON marshal %+v: %v", message, err)
				return
			}
		case <-self.closing:
			return
		}
	}
}

func (self *Session) heartbeatLoop() {
	for {
		select {
		case <-self.closing:
			return
		case <-time.After(self.server.heartbeat / 2):
			self.send(Message{Type: TypeHeartbeat})
		}
	}
}

func (self *Session) authorized(uri string, wantWrite bool) bool {
	self.lock.RLock()
	defer self.lock.RUnlock()
	hasWrite, found := self.authorizations[uri]
	if !found {
		return false
	}
	return !wantWrite || hasWrite
}

func (self *Session) send(message Message) {
	self.pushOutput(message)
}

func (self *Session) pushOutput(message Message) {
	select {
	case self.output <- message:
	default:
		self.server.logger.Printf("Unable to send %+v to %+v, output buffer full", message, self)
	}
}

func (self *Session) handleMessage(message Message) {
	switch message.Type {
	case TypeHeartbeat:
	case TypeMessage:
		if !self.authorized(message.URI, true) {
			self.send(Message{
				Type: TypeError,
				Id:   message.Id,
				Error: &Error{
					Message: fmt.Sprintf("%v not authorized for writing to %v", self.id, message.URI),
					Type:    TypeAuthorizationError,
				},
				Data: message,
			})
			return
		}
		self.server.Emit(message)
	case TypeUnsubscribe:
		self.server.removeSubscription(self.id, message.URI, true)
	case TypeSubscribe:
		if !self.authorized(message.URI, false) {
			self.send(Message{
				Type: TypeError,
				Id:   message.Id,
				Error: &Error{
					Message: fmt.Sprintf("%v not authorized for subscribing to %v", self.id, message.URI),
					Type:    TypeAuthorizationError,
				},
				Data: message,
			})
			return
		}
		self.server.addSubscription(self, message.URI)
	case TypeAuthorize:
		ok, err := self.server.authorizer(message.URI, message.Token, message.Write)
		if err != nil {
			self.send(Message{
				Type: TypeError,
				Id:   message.Id,
				Error: &Error{
					Message: err.Error(),
					Type:    TypeAuthorizationError,
				}, Data: message,
			})
			return
		}
		if !ok {
			self.send(Message{
				Type: TypeError,
				Id:   message.Id,
				Error: &Error{
					Message: fmt.Sprintf("%v does not provide authorization for %v", message.Token, message.URI),
					Type:    TypeAuthorizationError,
				}, Data: message,
			})
			return
		}
		self.lock.Lock()
		defer self.lock.Unlock()
		self.authorizations[message.URI] = (self.authorizations[message.URI] || message.Write)
		self.server.logger.Printf("%v\t[ws]\t[authorize]\t%v\t%v\t%v", time.Now(), self.RemoteAddr, self.id, message.URI)
	default:
		self.send(Message{
			Type: TypeError,
			Id:   message.Id,
			Error: &Error{
				Message: fmt.Sprintf("Unknown message type %#v", message.Type),
				Type:    TypeSyntaxError,
			},
			Data: message,
		})
		return
	}
	if message.Id != "" {
		self.send(Message{
			Type: TypeAck,
			Id:   message.Id,
		})
	}

}

func (self *Session) remove() {
	self.server.removeSession(self.id)
}

func (self *Session) terminate() {
	self.ws.Close()
	select {
	case _ = <-self.closing:
	default:
		close(self.closing)
	}
	self.server.logger.Printf("%v\t[ws]\t[disconnect]\t%v\t%v", time.Now(), self.RemoteAddr, self.id)
	self.cleanupTimer = time.AfterFunc(self.server.sessionTimeout, self.remove)
}

func (self *Session) handle(ws *websocket.Conn) {
	self.server.logger.Printf("%v\t[ws]\t[connect]\t%v\t%v", time.Now(), ws.Request().RemoteAddr, self.id)

	defer self.terminate()

	self.ws = ws
	self.RemoteAddr = self.ws.Request().RemoteAddr
	self.cleanupTimer.Stop()
	self.input = make(chan Message)
	self.closing = make(chan struct{})

	go self.readLoop()
	go self.writeLoop()
	go self.heartbeatLoop()

	self.send(Message{
		Type: TypeWelcome,
		Welcome: &Welcome{
			Heartbeat:      self.server.heartbeat / time.Millisecond,
			SessionTimeout: self.server.sessionTimeout / time.Millisecond,
			Id:             self.id,
		},
	})
	var message Message
	for {
		select {
		case _ = <-self.closing:
			return
		case message = <-self.input:
			self.handleMessage(message)
		case <-time.After(self.server.heartbeat):
			return
		}
	}
}
<commit_msg>removed the uncancelled cleanup timer problem<commit_after>package hub

import (
	"code.google.com/p/go.net/websocket"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"log"
	"math/rand"
	"net/http"
	"os"
	"sync"
	"time"
)

const (
	defaultBufferSize     = 1024
	defaultHeartbeat      = time.Second * 5
	defaultSessionTimeout = time.Second * 30
	idLength              = 16
	bufLength             = 4096
)

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

type Authorizer func(uri, token string, write bool) (authorized bool, err error)

func prettify(obj interface{}) string {
	b, err := json.MarshalIndent(obj, "", "  ")
	if err != nil {
		panic(err)
	}
	return string(b)
}

type Server struct {
	heartbeat      time.Duration
	sessionTimeout time.Duration
	sessions       map[string]*Session
	subscriptions  map[string]map[string]*Session
	subscribers    map[string]map[string]bool
	lock           *sync.RWMutex
	bufferSize     int
	logger         *log.Logger
	authorizer     Authorizer
}

func NewServer() *Server {
	return &Server{
		heartbeat:      defaultHeartbeat,
		bufferSize:     defaultBufferSize,
		sessionTimeout: defaultSessionTimeout,
		sessions:       map[string]*Session{},
		subscriptions:  map[string]map[string]*Session{},
		subscribers:    map[string]map[string]bool{},
		lock:           &sync.RWMutex{},
		logger:         log.New(os.Stdout, "pusher: ", 0),
		authorizer: func(uri, token string, write bool) (bool, error) {
			return true, nil
		},
	}
}

func (self *Server) Authorizer(f Authorizer) *Server {
	self.authorizer = f
	return self
}

func (self *Server) addSubscription(sess *Session, uri string) {
	self.lock.Lock()
	defer self.lock.Unlock()

	if _, found := self.subscriptions[uri]; !found {
		self.subscriptions[uri] = map[string]*Session{}
	}
	self.subscriptions[uri][sess.id] = sess

	if _, found := self.subscribers[sess.id]; !found {
		self.subscribers[sess.id] = map[string]bool{}
	}
	self.subscribers[sess.id][uri] = true

	self.logger.Printf("%v\t[ws]\t%v\t%v\t%v\t[subscribe]", time.Now(), uri, sess.RemoteAddr, sess.id)
}

func (self *Server) Emit(message Message) {
	self.lock.Lock()
	defer self.lock.Unlock()

	for _, sess := range self.subscriptions[message.URI] {
		sess.send(message)
	}
}

func (self *Server) removeSubscription(id, uri string, withLocking bool) {
	if withLocking {
		self.lock.Lock()
		defer self.lock.Unlock()
	}

	delete(self.subscriptions[uri], id)
	if len(self.subscriptions[uri]) == 0 {
		delete(self.subscriptions, uri)
	}

	delete(self.subscribers[id], uri)
	if len(self.subscribers[id]) == 0 {
		delete(self.subscribers, id)
	}
	self.logger.Printf("%v\t[ws]\t%v\t%v\t-\t[unsubscribe]", time.Now(), uri, id)
}

func (self *Server) randomId() string {
	buf := make([]byte, idLength)
	for index, _ := range buf {
		buf[index] = byte(rand.Int31())
	}
	return (base64.URLEncoding.EncodeToString(buf))
}

func (self *Server) removeSession(id string) {
	self.lock.Lock()
	defer self.lock.Unlock()

	delete(self.sessions, id)

	for uri, _ := range self.subscribers[id] {
		self.removeSubscription(id, uri, false)
	}
	self.logger.Printf("%v\t[ws]\t[cleanup]\t%v", time.Now(), id)
}

func (self *Server) getSession(id string) (result *Session) {
	self.lock.Lock()
	defer self.lock.Unlock()

	// if the id is not found (because the server restarted, or the client gave the initial empty id) we just create a new id and insert a new session
	if result = self.sessions[id]; result == nil {
		result = &Session{
			output:         make(chan Message, self.bufferSize),
			id:             self.randomId(),
			server:         self,
			authorizations: map[string]bool{},
			lock:           &sync.RWMutex{},
		}
		result.cleanupTimer = &time.Timer{}
		self.sessions[result.id] = result
	}
	return
}

func (self *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	self.logger.Printf("%v\t%v\t%v\t%v", time.Now(), r.Method, r.URL, r.RemoteAddr)
	websocket.Handler(self.getSession(r.URL.Query().Get("session_id")).handle).ServeHTTP(w, r)
}

type MessageType string
type ErrorType string

const (
	TypeError       = "Error"
	TypeHeartbeat   = "Heartbeat"
	TypeWelcome     = "Welcome"
	TypeSubscribe   = "Subscribe"
	TypeUnsubscribe = "Unsubscribe"
	TypeMessage     = "Message"
	TypeAuthorize   = "Authorize"
	TypeAck         = "Ack"
)

const (
	TypeJSONError          = "JSONError"
	TypeAuthorizationError = "AuthorizationError"
	TypeSyntaxError        = "SyntaxError"
)

type Welcome struct {
	Heartbeat      time.Duration
	SessionTimeout time.Duration
	Id             string
}

type Error struct {
	Message string
	Type    ErrorType
}

type Message struct {
	Type    MessageType
	Id      string      `json:",omitempty"`
	Welcome *Welcome    `json:",omitempty"`
	Error   *Error      `json:",omitempty"`
	Data    interface{} `json:",omitempty"`
	URI     string      `json:",omitempty"`
	Token   string      `json:",omitempty"`
	Write   bool        `json:",omitempty"`
}

type Session struct {
	ws             *websocket.Conn
	id             string
	RemoteAddr     string
	input          chan Message
	output         chan Message
	closing        chan struct{}
	server         *Server
	cleanupTimer   *time.Timer
	authorizations map[string]bool
	lock           *sync.RWMutex
}

func (self *Session) parseMessage(b []byte) (result Message, err error) {
	err = json.Unmarshal(b, &result)
	return
}

func (self *Session) readLoop() {
	defer self.terminate()
	buf := make([]byte, bufLength)
	n, err := self.ws.Read(buf)
	for err == nil {
		if message, err := self.parseMessage(buf[:n]); err == nil {
			self.input <- message
		} else {
			self.send(Message{
				Type: TypeError,
				Error: &Error{
					Message: err.Error(),
					Type:    TypeJSONError,
				},
				Data: string(buf[:n])})
		}
		n, err = self.ws.Read(buf)
	}
}

func (self *Session) writeLoop() {
	defer self.terminate()
	var message Message
	var err error
	var n int
	var encoded []byte
	for {
		select {
		case message = <-self.output:
			if encoded, err = json.Marshal(message); err == nil {
				if n, err = self.ws.Write(encoded); err != nil {
					self.server.logger.Printf("Error sending %s on %+v: %v", encoded, self.ws, err)
					return
				} else if n != len(encoded) {
					self.server.logger.Printf("Unable to send all of %s on %+v: only sent %v bytes", encoded, self.ws, n)
					return
				}
			} else {
				self.server.logger.Printf("Unable to JSON marshal %+v: %v", message, err)
				return
			}
		case <-self.closing:
			return
		}
	}
}

func (self *Session) heartbeatLoop() {
	for {
		select {
		case <-self.closing:
			return
		case <-time.After(self.server.heartbeat / 2):
			self.send(Message{Type: TypeHeartbeat})
		}
	}
}

func (self *Session) authorized(uri string, wantWrite bool) bool {
	self.lock.RLock()
	defer self.lock.RUnlock()
	hasWrite, found := self.authorizations[uri]
	if !found {
		return false
	}
	return !wantWrite || hasWrite
}

func (self *Session) send(message Message) {
	self.pushOutput(message)
}

func (self *Session) pushOutput(message Message) {
	select {
	case self.output <- message:
	default:
		self.server.logger.Printf("Unable to send %+v to %+v, output buffer full", message, self)
	}
}

func (self *Session) handleMessage(message Message) {
	switch message.Type {
	case TypeHeartbeat:
	case TypeMessage:
		if !self.authorized(message.URI, true) {
			self.send(Message{
				Type: TypeError,
				Id:   message.Id,
				Error: &Error{
					Message: fmt.Sprintf("%v not authorized for writing to %v", self.id, message.URI),
					Type:    TypeAuthorizationError,
				},
				Data: message,
			})
			return
		}
		self.server.Emit(message)
	case TypeUnsubscribe:
		self.server.removeSubscription(self.id, message.URI, true)
	case TypeSubscribe:
		if !self.authorized(message.URI, false) {
			self.send(Message{
				Type: TypeError,
				Id:   message.Id,
				Error: &Error{
					Message: fmt.Sprintf("%v not authorized for subscribing to %v", self.id, message.URI),
					Type:    TypeAuthorizationError,
				},
				Data: message,
			})
			return
		}
		self.server.addSubscription(self, message.URI)
	case TypeAuthorize:
		ok, err := self.server.authorizer(message.URI, message.Token, message.Write)
		if err != nil {
			self.send(Message{
				Type: TypeError,
				Id:   message.Id,
				Error: &Error{
					Message: err.Error(),
					Type:    TypeAuthorizationError,
				}, Data: message,
			})
			return
		}
		if !ok {
			self.send(Message{
				Type: TypeError,
				Id:   message.Id,
				Error: &Error{
					Message: fmt.Sprintf("%v does not provide authorization for %v", message.Token, message.URI),
					Type:    TypeAuthorizationError,
				}, Data: message,
			})
			return
		}
		self.lock.Lock()
		defer self.lock.Unlock()
		self.authorizations[message.URI] = (self.authorizations[message.URI] || message.Write)
		self.server.logger.Printf("%v\t[ws]\t[authorize]\t%v\t%v\t%v", time.Now(), self.RemoteAddr, self.id, message.URI)
	default:
		self.send(Message{
			Type: TypeError,
			Id:   message.Id,
			Error: &Error{
				Message: fmt.Sprintf("Unknown message type %#v", message.Type),
				Type:    TypeSyntaxError,
			},
			Data: message,
		})
		return
	}
	if message.Id != "" {
		self.send(Message{
			Type: TypeAck,
			Id:   message.Id,
		})
	}

}

func (self *Session) remove() {
	self.server.removeSession(self.id)
}

func (self *Session) terminate() {
	self.ws.Close()
	select {
	case _ = <-self.closing:
	default:
		close(self.closing)
	}
	self.server.logger.Printf("%v\t[ws]\t[disconnect]\t%v\t%v", time.Now(), self.RemoteAddr, self.id)
	self.cleanupTimer.Stop()
	self.cleanupTimer = time.AfterFunc(self.server.sessionTimeout, self.remove)
}

func (self *Session) handle(ws *websocket.Conn) {
	self.server.logger.Printf("%v\t[ws]\t[connect]\t%v\t%v", time.Now(), ws.Request().RemoteAddr, self.id)

	defer self.terminate()

	self.ws = ws
	self.RemoteAddr = self.ws.Request().RemoteAddr
	self.cleanupTimer.Stop()
	self.input = make(chan Message)
	self.closing = make(chan struct{})

	go self.readLoop()
	go self.writeLoop()
	go self.heartbeatLoop()

	self.send(Message{
		Type: TypeWelcome,
		Welcome: &Welcome{
			Heartbeat:      self.server.heartbeat / time.Millisecond,
			SessionTimeout: self.server.sessionTimeout / time.Millisecond,
			Id:             self.id,
		},
	})
	var message Message
	for {
		select {
		case _ = <-self.closing:
			return
		case message = <-self.input:
			self.handleMessage(message)
		case <-time.After(self.server.heartbeat):
			return
		}
	}
}
<|endoftext|>


<commit_before>package revel

import (
	"bytes"
	"encoding/json"
	"encoding/xml"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"reflect"
	"strconv"
	"time"
)

type Result interface {
	Apply(req *Request, resp *Response)
}

// This result handles all kinds of error codes (500, 404, ..).
// It renders the relevant error page (errors/CODE.format, e.g. errors/500.json).
// If RunMode is "dev", this results in a friendly error page.
type ErrorResult struct {
	RenderArgs map[string]interface{}
	Error      error
}

func (r ErrorResult) Apply(req *Request, resp *Response) {
	format := req.Format
	status := resp.Status
	if status == 0 {
		status = http.StatusInternalServerError
	}

	contentType := ContentTypeByFilename("xxx." + format)
	if contentType == DefaultFileContentType {
		contentType = "text/plain"
	}

	// Get the error template.
	var err error
	templatePath := fmt.Sprintf("errors/%d.%s", status, format)
	tmpl, err := MainTemplateLoader.Template(templatePath)

	// This func shows a plaintext error message, in case the template rendering
	// doesn't work.
	showPlaintext := func(err error) {
		PlaintextErrorResult{fmt.Errorf("Server Error:\n%s\n\n"+
			"Additionally, an error occurred when rendering the error page:\n%s",
			r.Error, err)}.Apply(req, resp)
	}

	if tmpl == nil {
		if err == nil {
			err = fmt.Errorf("Couldn't find template %s", templatePath)
		}
		showPlaintext(err)
		return
	}

	// If it's not a revel error, wrap it in one.
	var revelError *Error
	switch e := r.Error.(type) {
	case *Error:
		revelError = e
	case error:
		revelError = &Error{
			Title:       "Server Error",
			Description: e.Error(),
		}
	}

	if revelError == nil {
		panic("no error provided")
	}

	if r.RenderArgs == nil {
		r.RenderArgs = make(map[string]interface{})
	}
	r.RenderArgs["RunMode"] = RunMode
	r.RenderArgs["Error"] = revelError
	r.RenderArgs["Router"] = MainRouter

	// Render it.
	var b bytes.Buffer
	err = tmpl.Render(&b, r.RenderArgs)

	// If there was an error, print it in plain text.
	if err != nil {
		showPlaintext(err)
		return
	}

	resp.WriteHeader(status, contentType)
	b.WriteTo(resp.Out)
}

type PlaintextErrorResult struct {
	Error error
}

// This method is used when the template loader or error template is not available.
func (r PlaintextErrorResult) Apply(req *Request, resp *Response) {
	resp.WriteHeader(http.StatusInternalServerError, "text/plain")
	resp.Out.Write([]byte(r.Error.Error()))
}

// Action methods return this result to request a template be rendered.
type RenderTemplateResult struct {
	Template   Template
	RenderArgs map[string]interface{}
}

func (r *RenderTemplateResult) Apply(req *Request, resp *Response) {
	// Handle panics when rendering templates.
	defer func() {
		if err := recover(); err != nil {
			ERROR.Println(err)
			PlaintextErrorResult{fmt.Errorf("Template Execution Panic in %s:\n%s",
				r.Template.Name(), err)}.Apply(req, resp)
		}
	}()

	chunked := Config.BoolDefault("results.chunked", false)

	// If it's a HEAD request, throw away the bytes.
	out := io.Writer(resp.Out)
	if req.Method == "HEAD" {
		out = ioutil.Discard
	}

	// In a prod mode, write the status, render, and hope for the best.
	// (In a dev mode, always render to a temporary buffer first to avoid having
	// error pages distorted by HTML already written)
	if chunked && !DevMode {
		resp.WriteHeader(http.StatusOK, "text/html")
		r.render(req, resp, out)
		return
	}

	// Render the template into a temporary buffer, to see if there was an error
	// rendering the template.  If not, then copy it into the response buffer.
	// Otherwise, template render errors may result in unpredictable HTML (and
	// would carry a 200 status code)
	var b bytes.Buffer
	r.render(req, resp, &b)
	if !chunked {
		resp.Out.Header().Set("Content-Length", strconv.Itoa(b.Len()))
	}
	resp.WriteHeader(http.StatusOK, "text/html")
	b.WriteTo(out)
}

func (r *RenderTemplateResult) render(req *Request, resp *Response, wr io.Writer) {
	err := r.Template.Render(wr, r.RenderArgs)
	if err == nil {
		return
	}

	var templateContent []string
	templateName, line, description := parseTemplateError(err)
	if templateName == "" {
		templateName = r.Template.Name()
		templateContent = r.Template.Content()
	} else {
		if tmpl, err := MainTemplateLoader.Template(templateName); err == nil {
			templateContent = tmpl.Content()
		}
	}
	compileError := &Error{
		Title:       "Template Execution Error",
		Path:        templateName,
		Description: description,
		Line:        line,
		SourceLines: templateContent,
	}
	resp.Status = 500
	ERROR.Printf("Template Execution Error (in %s): %s", templateName, description)
	ErrorResult{r.RenderArgs, compileError}.Apply(req, resp)
}

type RenderHtmlResult struct {
	html string
}

func (r RenderHtmlResult) Apply(req *Request, resp *Response) {
	resp.WriteHeader(http.StatusOK, "text/html")
	resp.Out.Write([]byte(r.html))
}

type RenderJsonResult struct {
	obj      interface{}
	callback string
}

func (r RenderJsonResult) Apply(req *Request, resp *Response) {
	var b []byte
	var err error
	if Config.BoolDefault("results.pretty", false) {
		b, err = json.MarshalIndent(r.obj, "", "  ")
	} else {
		b, err = json.Marshal(r.obj)
	}

	if err != nil {
		ErrorResult{Error: err}.Apply(req, resp)
		return
	}

	if r.callback == "" {
		resp.WriteHeader(http.StatusOK, "application/json")
		resp.Out.Write(b)
		return
	}

	resp.WriteHeader(http.StatusOK, "application/javascript")
	resp.Out.Write([]byte(r.callback + "("))
	resp.Out.Write(b)
	resp.Out.Write([]byte(");"))
}

type RenderXmlResult struct {
	obj interface{}
}

func (r RenderXmlResult) Apply(req *Request, resp *Response) {
	var b []byte
	var err error
	if Config.BoolDefault("results.pretty", false) {
		b, err = xml.MarshalIndent(r.obj, "", "  ")
	} else {
		b, err = xml.Marshal(r.obj)
	}

	if err != nil {
		ErrorResult{Error: err}.Apply(req, resp)
		return
	}

	resp.WriteHeader(http.StatusOK, "application/xml")
	resp.Out.Write(b)
}

type RenderTextResult struct {
	text string
}

func (r RenderTextResult) Apply(req *Request, resp *Response) {
	resp.WriteHeader(http.StatusOK, "text/plain")
	resp.Out.Write([]byte(r.text))
}

type ContentDisposition string

var (
	Attachment ContentDisposition = "attachment"
	Inline     ContentDisposition = "inline"
)

type BinaryResult struct {
	Reader   io.Reader
	Name     string
	Length   int64
	Delivery ContentDisposition
	ModTime  time.Time
}

func (r *BinaryResult) Apply(req *Request, resp *Response) {
	disposition := string(r.Delivery)
	if r.Name != "" {
		disposition += fmt.Sprintf("; filename=%s", r.Name)
	}
	resp.Out.Header().Set("Content-Disposition", disposition)

	// If we have a ReadSeeker, delegate to http.ServeContent
	if rs, ok := r.Reader.(io.ReadSeeker); ok {
		http.ServeContent(resp.Out, req.Request, r.Name, r.ModTime, rs)
	} else {
		// Else, do a simple io.Copy.
		if r.Length != -1 {
			resp.Out.Header().Set("Content-Length", strconv.FormatInt(r.Length, 10))
		}
		resp.WriteHeader(http.StatusOK, ContentTypeByFilename(r.Name))
		io.Copy(resp.Out, r.Reader)
	}

	// Close the Reader if we can
	if v, ok := r.Reader.(io.Closer); ok {
		v.Close()
	}
}

type RedirectToUrlResult struct {
	url string
}

func (r *RedirectToUrlResult) Apply(req *Request, resp *Response) {
	resp.Out.Header().Set("Location", r.url)
	resp.WriteHeader(http.StatusFound, "")
}

type RedirectToActionResult struct {
	val interface{}
}

func (r *RedirectToActionResult) Apply(req *Request, resp *Response) {
	url, err := getRedirectUrl(r.val)
	if err != nil {
		ERROR.Println("Couldn't resolve redirect:", err.Error())
		ErrorResult{Error: err}.Apply(req, resp)
		return
	}
	resp.Out.Header().Set("Location", url)
	resp.WriteHeader(http.StatusFound, "")
}

func getRedirectUrl(item interface{}) (string, error) {
	// Handle strings
	if url, ok := item.(string); ok {
		return url, nil
	}

	// Handle funcs
	val := reflect.ValueOf(item)
	typ := reflect.TypeOf(item)
	if typ.Kind() == reflect.Func && typ.NumIn() > 0 {
		// Get the Controller Method
		recvType := typ.In(0)
		method := FindMethod(recvType, val)
		if method == nil {
			return "", errors.New("couldn't find method")
		}

		// Construct the action string (e.g. "Controller.Method")
		if recvType.Kind() == reflect.Ptr {
			recvType = recvType.Elem()
		}
		action := recvType.Name() + "." + method.Name
		actionDef := MainRouter.Reverse(action, make(map[string]string))
		if actionDef == nil {
			return "", errors.New("no route for action " + action)
		}

		return actionDef.String(), nil
	}

	// Out of guesses
	return "", errors.New("didn't recognize type: " + typ.String())
}
<commit_msg>Add utf-8 charset enconding in Result - Fixes #363<commit_after>package revel

import (
	"bytes"
	"encoding/json"
	"encoding/xml"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"reflect"
	"strconv"
	"time"
)

type Result interface {
	Apply(req *Request, resp *Response)
}

// This result handles all kinds of error codes (500, 404, ..).
// It renders the relevant error page (errors/CODE.format, e.g. errors/500.json).
// If RunMode is "dev", this results in a friendly error page.
type ErrorResult struct {
	RenderArgs map[string]interface{}
	Error      error
}

func (r ErrorResult) Apply(req *Request, resp *Response) {
	format := req.Format
	status := resp.Status
	if status == 0 {
		status = http.StatusInternalServerError
	}

	contentType := ContentTypeByFilename("xxx." + format)
	if contentType == DefaultFileContentType {
		contentType = "text/plain"
	}

	// Get the error template.
	var err error
	templatePath := fmt.Sprintf("errors/%d.%s", status, format)
	tmpl, err := MainTemplateLoader.Template(templatePath)

	// This func shows a plaintext error message, in case the template rendering
	// doesn't work.
	showPlaintext := func(err error) {
		PlaintextErrorResult{fmt.Errorf("Server Error:\n%s\n\n"+
			"Additionally, an error occurred when rendering the error page:\n%s",
			r.Error, err)}.Apply(req, resp)
	}

	if tmpl == nil {
		if err == nil {
			err = fmt.Errorf("Couldn't find template %s", templatePath)
		}
		showPlaintext(err)
		return
	}

	// If it's not a revel error, wrap it in one.
	var revelError *Error
	switch e := r.Error.(type) {
	case *Error:
		revelError = e
	case error:
		revelError = &Error{
			Title:       "Server Error",
			Description: e.Error(),
		}
	}

	if revelError == nil {
		panic("no error provided")
	}

	if r.RenderArgs == nil {
		r.RenderArgs = make(map[string]interface{})
	}
	r.RenderArgs["RunMode"] = RunMode
	r.RenderArgs["Error"] = revelError
	r.RenderArgs["Router"] = MainRouter

	// Render it.
	var b bytes.Buffer
	err = tmpl.Render(&b, r.RenderArgs)

	// If there was an error, print it in plain text.
	if err != nil {
		showPlaintext(err)
		return
	}

	resp.WriteHeader(status, contentType)
	b.WriteTo(resp.Out)
}

type PlaintextErrorResult struct {
	Error error
}

// This method is used when the template loader or error template is not available.
func (r PlaintextErrorResult) Apply(req *Request, resp *Response) {
	resp.WriteHeader(http.StatusInternalServerError, "text/plain; charset=utf-8")
	resp.Out.Write([]byte(r.Error.Error()))
}

// Action methods return this result to request a template be rendered.
type RenderTemplateResult struct {
	Template   Template
	RenderArgs map[string]interface{}
}

func (r *RenderTemplateResult) Apply(req *Request, resp *Response) {
	// Handle panics when rendering templates.
	defer func() {
		if err := recover(); err != nil {
			ERROR.Println(err)
			PlaintextErrorResult{fmt.Errorf("Template Execution Panic in %s:\n%s",
				r.Template.Name(), err)}.Apply(req, resp)
		}
	}()

	chunked := Config.BoolDefault("results.chunked", false)

	// If it's a HEAD request, throw away the bytes.
	out := io.Writer(resp.Out)
	if req.Method == "HEAD" {
		out = ioutil.Discard
	}

	// In a prod mode, write the status, render, and hope for the best.
	// (In a dev mode, always render to a temporary buffer first to avoid having
	// error pages distorted by HTML already written)
	if chunked && !DevMode {
		resp.WriteHeader(http.StatusOK, "text/html; charset=utf-8")
		r.render(req, resp, out)
		return
	}

	// Render the template into a temporary buffer, to see if there was an error
	// rendering the template.  If not, then copy it into the response buffer.
	// Otherwise, template render errors may result in unpredictable HTML (and
	// would carry a 200 status code)
	var b bytes.Buffer
	r.render(req, resp, &b)
	if !chunked {
		resp.Out.Header().Set("Content-Length", strconv.Itoa(b.Len()))
	}
	resp.WriteHeader(http.StatusOK, "text/html; charset=utf-8")
	b.WriteTo(out)
}

func (r *RenderTemplateResult) render(req *Request, resp *Response, wr io.Writer) {
	err := r.Template.Render(wr, r.RenderArgs)
	if err == nil {
		return
	}

	var templateContent []string
	templateName, line, description := parseTemplateError(err)
	if templateName == "" {
		templateName = r.Template.Name()
		templateContent = r.Template.Content()
	} else {
		if tmpl, err := MainTemplateLoader.Template(templateName); err == nil {
			templateContent = tmpl.Content()
		}
	}
	compileError := &Error{
		Title:       "Template Execution Error",
		Path:        templateName,
		Description: description,
		Line:        line,
		SourceLines: templateContent,
	}
	resp.Status = 500
	ERROR.Printf("Template Execution Error (in %s): %s", templateName, description)
	ErrorResult{r.RenderArgs, compileError}.Apply(req, resp)
}

type RenderHtmlResult struct {
	html string
}

func (r RenderHtmlResult) Apply(req *Request, resp *Response) {
	resp.WriteHeader(http.StatusOK, "text/html; charset=utf-8")
	resp.Out.Write([]byte(r.html))
}

type RenderJsonResult struct {
	obj      interface{}
	callback string
}

func (r RenderJsonResult) Apply(req *Request, resp *Response) {
	var b []byte
	var err error
	if Config.BoolDefault("results.pretty", false) {
		b, err = json.MarshalIndent(r.obj, "", "  ")
	} else {
		b, err = json.Marshal(r.obj)
	}

	if err != nil {
		ErrorResult{Error: err}.Apply(req, resp)
		return
	}

	if r.callback == "" {
		resp.WriteHeader(http.StatusOK, "application/json; charset=utf-8")
		resp.Out.Write(b)
		return
	}

	resp.WriteHeader(http.StatusOK, "application/javascript; charset=utf-8")
	resp.Out.Write([]byte(r.callback + "("))
	resp.Out.Write(b)
	resp.Out.Write([]byte(");"))
}

type RenderXmlResult struct {
	obj interface{}
}

func (r RenderXmlResult) Apply(req *Request, resp *Response) {
	var b []byte
	var err error
	if Config.BoolDefault("results.pretty", false) {
		b, err = xml.MarshalIndent(r.obj, "", "  ")
	} else {
		b, err = xml.Marshal(r.obj)
	}

	if err != nil {
		ErrorResult{Error: err}.Apply(req, resp)
		return
	}

	resp.WriteHeader(http.StatusOK, "application/xml; charset=utf-8")
	resp.Out.Write(b)
}

type RenderTextResult struct {
	text string
}

func (r RenderTextResult) Apply(req *Request, resp *Response) {
	resp.WriteHeader(http.StatusOK, "text/plain; charset=utf-8")
	resp.Out.Write([]byte(r.text))
}

type ContentDisposition string

var (
	Attachment ContentDisposition = "attachment"
	Inline     ContentDisposition = "inline"
)

type BinaryResult struct {
	Reader   io.Reader
	Name     string
	Length   int64
	Delivery ContentDisposition
	ModTime  time.Time
}

func (r *BinaryResult) Apply(req *Request, resp *Response) {
	disposition := string(r.Delivery)
	if r.Name != "" {
		disposition += fmt.Sprintf("; filename=%s", r.Name)
	}
	resp.Out.Header().Set("Content-Disposition", disposition)

	// If we have a ReadSeeker, delegate to http.ServeContent
	if rs, ok := r.Reader.(io.ReadSeeker); ok {
		http.ServeContent(resp.Out, req.Request, r.Name, r.ModTime, rs)
	} else {
		// Else, do a simple io.Copy.
		if r.Length != -1 {
			resp.Out.Header().Set("Content-Length", strconv.FormatInt(r.Length, 10))
		}
		resp.WriteHeader(http.StatusOK, ContentTypeByFilename(r.Name))
		io.Copy(resp.Out, r.Reader)
	}

	// Close the Reader if we can
	if v, ok := r.Reader.(io.Closer); ok {
		v.Close()
	}
}

type RedirectToUrlResult struct {
	url string
}

func (r *RedirectToUrlResult) Apply(req *Request, resp *Response) {
	resp.Out.Header().Set("Location", r.url)
	resp.WriteHeader(http.StatusFound, "")
}

type RedirectToActionResult struct {
	val interface{}
}

func (r *RedirectToActionResult) Apply(req *Request, resp *Response) {
	url, err := getRedirectUrl(r.val)
	if err != nil {
		ERROR.Println("Couldn't resolve redirect:", err.Error())
		ErrorResult{Error: err}.Apply(req, resp)
		return
	}
	resp.Out.Header().Set("Location", url)
	resp.WriteHeader(http.StatusFound, "")
}

func getRedirectUrl(item interface{}) (string, error) {
	// Handle strings
	if url, ok := item.(string); ok {
		return url, nil
	}

	// Handle funcs
	val := reflect.ValueOf(item)
	typ := reflect.TypeOf(item)
	if typ.Kind() == reflect.Func && typ.NumIn() > 0 {
		// Get the Controller Method
		recvType := typ.In(0)
		method := FindMethod(recvType, val)
		if method == nil {
			return "", errors.New("couldn't find method")
		}

		// Construct the action string (e.g. "Controller.Method")
		if recvType.Kind() == reflect.Ptr {
			recvType = recvType.Elem()
		}
		action := recvType.Name() + "." + method.Name
		actionDef := MainRouter.Reverse(action, make(map[string]string))
		if actionDef == nil {
			return "", errors.New("no route for action " + action)
		}

		return actionDef.String(), nil
	}

	// Out of guesses
	return "", errors.New("didn't recognize type: " + typ.String())
}
<|endoftext|>