<commit_before>package transloadit

import (
	"crypto/hmac"
	"crypto/sha1"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"strings"
	"time"
)

type Config struct {
	AuthKey    string
	AuthSecret string
	Region     string
	Endpoint   string
}

var DefaultConfig = Config{
	Region:   "us-east-1",
	Endpoint: "http://api2.transloadit.com",
}

type Client struct {
	config     Config
	httpClient *http.Client
}

// Options when retrieving a list.
// Look at the documentation which properties are accepted
// and to see their meaining, e.g. https://transloadit.com/docs/api-docs#retrieve-assembly-list
// for listing assemblies.
type ListOptions struct {
	Page       int        `json:"page,omitempty"`
	PageSize   int        `json:"pagesize,omitempty"`
	Sort       string     `json:"sort,omitempty"`
	Order      string     `json:"order,omitempty"`
	Fields     []string   `json:"fields,omitempty"`
	Type       string     `json:"type,omitempty"`
	Keywords   []string   `json:"keyword,omitempty"`
	AssemblyId string     `json:"assembly_id,omitempty"`
	FromDate   *time.Time `json:"fromdate,omitempty"`
	ToDate     *time.Time `json:"todate,omitempty"`
}

type authListOptions struct {
	*ListOptions

	// For internal use only!
	Auth struct {
		Key     string `json:"key"`
		Expires string `json:"expires"`
	} `json:"auth"`
}

type RequestError struct {
	Code    string `json:"error"`
	Message string `json:"message"`
}

func (err RequestError) Error() string {
	return fmt.Sprintf("request failed due to %s: %s", err.Code, err.Message)
}

// Create a new client using the provided configuration object.
// An error will be returned if no AuthKey or AuthSecret is found in config.
func NewClient(config Config) (*Client, error) {
	if config.AuthKey == "" {
		return nil, errors.New("failed to create client: missing AuthKey")
	}

	if config.AuthSecret == "" {
		return nil, errors.New("failed to create client: missing AuthSecret")
	}

	client := &Client{
		config:     config,
		httpClient: &http.Client{},
	}

	return client, nil
}

func (client *Client) sign(params map[string]interface{}) (string, string, error) {
	params["auth"] = map[string]string{
		"key":     client.config.AuthKey,
		"expires": getExpireString(),
	}

	b, err := json.Marshal(params)
	if err != nil {
		return "", "", fmt.Errorf("unable to create signature: %s", err)
	}

	hash := hmac.New(sha1.New, []byte(client.config.AuthSecret))
	hash.Write(b)
	return string(b), hex.EncodeToString(hash.Sum(nil)), nil
}

func (client *Client) doRequest(req *http.Request, result interface{}) error {
	req.Header.Set("User-Agent", "Transloadit Go SDK "+Version)

	res, err := client.httpClient.Do(req)
	if err != nil {
		return fmt.Errorf("failed execute http request: %s", err)
	}
	defer res.Body.Close()

	// Limit response to 128MB
	reader := io.LimitReader(res.Body, 128*1024*1024)
	body, err := ioutil.ReadAll(reader)
	if err != nil {
		return fmt.Errorf("failed execute http request: %s", err)
	}

	if !(res.StatusCode >= 200 && res.StatusCode < 300) {
		var reqErr RequestError
		if err := json.Unmarshal(body, &reqErr); err != nil {
			return fmt.Errorf("failed unmarshal http request: %s", err)
		}

		return reqErr
	}

	if result != nil {
		err = json.Unmarshal(body, result)
	}

	return err
}

func (client *Client) request(method string, path string, content map[string]interface{}, result interface{}) error {
	uri := path
	// Don't add host for absolute urls
	if u, err := url.Parse(path); err == nil && u.Scheme == "" {
		uri = client.config.Endpoint + "/" + path
	}

	// Ensure content is a map
	if content == nil {
		content = make(map[string]interface{})
	}

	// Create signature
	params, signature, err := client.sign(content)
	if err != nil {
		return fmt.Errorf("request: %s", err)
	}

	v := url.Values{}
	v.Set("params", params)
	v.Set("signature", signature)

	var body io.Reader
	if method == "GET" {
		uri += "?" + v.Encode()
	} else {
		body = strings.NewReader(v.Encode())
	}
	req, err := http.NewRequest(method, uri, body)
	if err != nil {
		return fmt.Errorf("request: %s", err)
	}

	if method != "GET" {
		// Add content type header
		req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	}

	return client.doRequest(req, result)
}

func (client *Client) listRequest(path string, listOptions *ListOptions, result interface{}) error {
	uri := client.config.Endpoint + "/" + path

	options := authListOptions{
		ListOptions: listOptions,
		Auth: struct {
			Key     string `json:"key"`
			Expires string `json:"expires"`
		}{
			Key:     client.config.AuthKey,
			Expires: getExpireString(),
		},
	}

	b, err := json.Marshal(options)
	if err != nil {
		return fmt.Errorf("unable to create signature: %s", err)
	}

	hash := hmac.New(sha1.New, []byte(client.config.AuthSecret))
	hash.Write(b)

	params := string(b)
	signature := hex.EncodeToString(hash.Sum(nil))

	v := url.Values{}
	v.Set("params", params)
	v.Set("signature", signature)

	uri += "?" + v.Encode()

	req, err := http.NewRequest("GET", uri, nil)
	if err != nil {
		return fmt.Errorf("request: %s", err)
	}

	return client.doRequest(req, result)
}

func getExpireString() string {
	// Expires in 1 hour
	expires := time.Now().UTC().Add(time.Hour)
	expiresStr := fmt.Sprintf("%04d/%02d/%02d %02d:%02d:%02d+00:00", expires.Year(), expires.Month(), expires.Day(), expires.Hour(), expires.Minute(), expires.Second())
	return string(expiresStr)
}
<commit_msg>Do not read entire response before decoding<commit_after>package transloadit

import (
	"crypto/hmac"
	"crypto/sha1"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"strings"
	"time"
)

type Config struct {
	AuthKey    string
	AuthSecret string
	Region     string
	Endpoint   string
}

var DefaultConfig = Config{
	Region:   "us-east-1",
	Endpoint: "http://api2.transloadit.com",
}

type Client struct {
	config     Config
	httpClient *http.Client
}

// Options when retrieving a list.
// Look at the documentation which properties are accepted
// and to see their meaining, e.g. https://transloadit.com/docs/api-docs#retrieve-assembly-list
// for listing assemblies.
type ListOptions struct {
	Page       int        `json:"page,omitempty"`
	PageSize   int        `json:"pagesize,omitempty"`
	Sort       string     `json:"sort,omitempty"`
	Order      string     `json:"order,omitempty"`
	Fields     []string   `json:"fields,omitempty"`
	Type       string     `json:"type,omitempty"`
	Keywords   []string   `json:"keyword,omitempty"`
	AssemblyId string     `json:"assembly_id,omitempty"`
	FromDate   *time.Time `json:"fromdate,omitempty"`
	ToDate     *time.Time `json:"todate,omitempty"`
}

type authListOptions struct {
	*ListOptions

	// For internal use only!
	Auth struct {
		Key     string `json:"key"`
		Expires string `json:"expires"`
	} `json:"auth"`
}

type RequestError struct {
	Code    string `json:"error"`
	Message string `json:"message"`
}

func (err RequestError) Error() string {
	return fmt.Sprintf("request failed due to %s: %s", err.Code, err.Message)
}

// Create a new client using the provided configuration object.
// An error will be returned if no AuthKey or AuthSecret is found in config.
func NewClient(config Config) (*Client, error) {
	if config.AuthKey == "" {
		return nil, errors.New("failed to create client: missing AuthKey")
	}

	if config.AuthSecret == "" {
		return nil, errors.New("failed to create client: missing AuthSecret")
	}

	client := &Client{
		config:     config,
		httpClient: &http.Client{},
	}

	return client, nil
}

func (client *Client) sign(params map[string]interface{}) (string, string, error) {
	params["auth"] = map[string]string{
		"key":     client.config.AuthKey,
		"expires": getExpireString(),
	}

	b, err := json.Marshal(params)
	if err != nil {
		return "", "", fmt.Errorf("unable to create signature: %s", err)
	}

	hash := hmac.New(sha1.New, []byte(client.config.AuthSecret))
	hash.Write(b)
	return string(b), hex.EncodeToString(hash.Sum(nil)), nil
}

func (client *Client) doRequest(req *http.Request, result interface{}) error {
	req.Header.Set("User-Agent", "Transloadit Go SDK "+Version)

	res, err := client.httpClient.Do(req)
	if err != nil {
		return fmt.Errorf("failed execute http request: %s", err)
	}
	defer res.Body.Close()

	// Limit response to 128MB
	reader := io.LimitReader(res.Body, 128*1024*1024)
	decoder := json.NewDecoder(reader)

	if !(res.StatusCode >= 200 && res.StatusCode < 300) {
		var reqErr RequestError
		if err := decoder.Decode(&reqErr); err != nil {
			return fmt.Errorf("failed unmarshal http request: %s", err)
		}

		return reqErr
	}

	if result != nil {
		if err := decoder.Decode(result); err != nil {
			return fmt.Errorf("failed unmarshal http request: %s", err)
		}
	}

	return nil
}

func (client *Client) request(method string, path string, content map[string]interface{}, result interface{}) error {
	uri := path
	// Don't add host for absolute urls
	if u, err := url.Parse(path); err == nil && u.Scheme == "" {
		uri = client.config.Endpoint + "/" + path
	}

	// Ensure content is a map
	if content == nil {
		content = make(map[string]interface{})
	}

	// Create signature
	params, signature, err := client.sign(content)
	if err != nil {
		return fmt.Errorf("request: %s", err)
	}

	v := url.Values{}
	v.Set("params", params)
	v.Set("signature", signature)

	var body io.Reader
	if method == "GET" {
		uri += "?" + v.Encode()
	} else {
		body = strings.NewReader(v.Encode())
	}
	req, err := http.NewRequest(method, uri, body)
	if err != nil {
		return fmt.Errorf("request: %s", err)
	}

	if method != "GET" {
		// Add content type header
		req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	}

	return client.doRequest(req, result)
}

func (client *Client) listRequest(path string, listOptions *ListOptions, result interface{}) error {
	uri := client.config.Endpoint + "/" + path

	options := authListOptions{
		ListOptions: listOptions,
		Auth: struct {
			Key     string `json:"key"`
			Expires string `json:"expires"`
		}{
			Key:     client.config.AuthKey,
			Expires: getExpireString(),
		},
	}

	b, err := json.Marshal(options)
	if err != nil {
		return fmt.Errorf("unable to create signature: %s", err)
	}

	hash := hmac.New(sha1.New, []byte(client.config.AuthSecret))
	hash.Write(b)

	params := string(b)
	signature := hex.EncodeToString(hash.Sum(nil))

	v := url.Values{}
	v.Set("params", params)
	v.Set("signature", signature)

	uri += "?" + v.Encode()

	req, err := http.NewRequest("GET", uri, nil)
	if err != nil {
		return fmt.Errorf("request: %s", err)
	}

	return client.doRequest(req, result)
}

func getExpireString() string {
	// Expires in 1 hour
	expires := time.Now().UTC().Add(time.Hour)
	expiresStr := fmt.Sprintf("%04d/%02d/%02d %02d:%02d:%02d+00:00", expires.Year(), expires.Month(), expires.Day(), expires.Hour(), expires.Minute(), expires.Second())
	return string(expiresStr)
}
<|endoftext|>


<commit_before>package mongo

import (
	"bytes"
	"crypto/rand"
	"encoding/base64"
	"fmt"
	"io/ioutil"
	"net"
	"os"
	"os/exec"
	"path"
	"path/filepath"

	"github.com/juju/loggo"
	"labix.org/v2/mgo"

	"launchpad.net/juju-core/instance"
	"launchpad.net/juju-core/replicaset"
	"launchpad.net/juju-core/state/api/params"
	"launchpad.net/juju-core/upstart"
	"launchpad.net/juju-core/utils"
	"launchpad.net/juju-core/version"
)

const (
	maxFiles = 65000
	maxProcs = 20000

	serviceName = "juju-db"

	// SharedSecretFile is the name of the Mongo shared secret file
	// located within the Juju data directory.
	SharedSecretFile = "shared-secret"

	ReplicaSetName = "juju"
)

var (
	logger          = loggo.GetLogger("juju.agent.mongo")
	mongoConfigPath = "/etc/default/mongodb"

	// JujuMongodPath holds the default path to the juju-specific mongod.
	JujuMongodPath = "/usr/lib/juju/bin/mongod"

	upstartConfInstall          = (*upstart.Conf).Install
	upstartServiceStopAndRemove = (*upstart.Service).StopAndRemove
)

// WithAddresses represents an entity that has a set of
// addresses. e.g. a state Machine object
type WithAddresses interface {
	Addresses() []instance.Address
}

// IsMaster returns a boolean that represents whether the given
// machine's peer address is the primary mongo host for the replicaset
func IsMaster(session *mgo.Session, obj WithAddresses) (bool, error) {
	addrs := obj.Addresses()

	masterHostPort, err := replicaset.MasterHostPort(session)

	// If the replica set has not been configured, then we
	// can have only one master and the caller must
	// be that master.
	if err == replicaset.ErrMasterNotConfigured {
		return true, nil
	}

	if err != nil {
		return false, err
	}

	masterAddr, _, err := net.SplitHostPort(masterHostPort)
	if err != nil {
		return false, err
	}

	machinePeerAddr := SelectPeerAddress(addrs)
	return machinePeerAddr == masterAddr, nil
}

// SelectPeerAddress returns the address to use as the
// mongo replica set peer address by selecting it from the given addresses.
func SelectPeerAddress(addrs []instance.Address) string {
	return instance.SelectInternalAddress(addrs, false)
}

// SelectPeerHostPort returns the HostPort to use as the
// mongo replica set peer by selecting it from the given hostPorts.
func SelectPeerHostPort(hostPorts []instance.HostPort) string {
	return instance.SelectInternalHostPort(hostPorts, false)
}

// GenerateSharedSecret generates a pseudo-random shared secret (keyfile)
// for use with Mongo replica sets.
func GenerateSharedSecret() (string, error) {
	// "A key’s length must be between 6 and 1024 characters and may
	// only contain characters in the base64 set."
	//   -- http://docs.mongodb.org/manual/tutorial/generate-key-file/
	buf := make([]byte, base64.StdEncoding.DecodedLen(1024))
	if _, err := rand.Read(buf); err != nil {
		return "", fmt.Errorf("cannot read random secret: %v", err)
	}
	return base64.StdEncoding.EncodeToString(buf), nil
}

// MongoPath returns the executable path to be used to run mongod on this
// machine. If the juju-bundled version of mongo exists, it will return that
// path, otherwise it will return the command to run mongod from the path.
func MongodPath() (string, error) {
	if _, err := os.Stat(JujuMongodPath); err == nil {
		return JujuMongodPath, nil
	}

	path, err := exec.LookPath("mongod")
	if err != nil {
		return "", err
	}
	return path, nil
}

// RemoveService removes the mongoDB upstart service from this machine.
func RemoveService(namespace string) error {
	svc := upstart.NewService(ServiceName(namespace))
	return upstartServiceStopAndRemove(svc)
}

// EnsureMongoServer ensures that the correct mongo upstart script is installed
// and running.
//
// This method will remove old versions of the mongo upstart script as necessary
// before installing the new version.
//
// The namespace is a unique identifier to prevent multiple instances of mongo
// on this machine from colliding. This should be empty unless using
// the local provider.
func EnsureMongoServer(dataDir string, namespace string, info params.StateServingInfo) error {
	logger.Infof("Ensuring mongo server is running; dataDir %s; port %d", dataDir, info.StatePort)
	dbDir := filepath.Join(dataDir, "db")

	if err := os.MkdirAll(dbDir, 0700); err != nil {
		return fmt.Errorf("cannot create mongo dbdir: %v", err)
	}

	certKey := info.Cert + "\n" + info.PrivateKey
	err := utils.AtomicWriteFile(sslKeyPath(dataDir), []byte(certKey), 0600)
	if err != nil {
		return fmt.Errorf("cannot write SSL key: %v", err)
	}

	err = utils.AtomicWriteFile(sharedSecretPath(dataDir), []byte(info.SharedSecret), 0600)
	if err != nil {
		return fmt.Errorf("cannot write mongod shared secret: %v", err)
	}

	// Disable the default mongodb installed by the mongodb-server package.
	// Only do this if the file doesn't exist already, so users can run
	// their own mongodb server if they wish to.
	if _, err := os.Stat(mongoConfigPath); os.IsNotExist(err) {
		err = ioutil.WriteFile(
			mongoConfigPath,
			[]byte("ENABLE_MONGODB=no"),
			0644,
		)
		if err != nil {
			return err
		}
	}

	if err := aptGetInstallMongod(); err != nil {
		return fmt.Errorf("cannot install mongod: %v", err)
	}

	upstartConf, err := mongoUpstartService(namespace, dataDir, dbDir, info.StatePort)
	if err != nil {
		return err
	}

	// TODO(natefinch) 2014-04-12 https://launchpad.net/bugs/1306902
	// remove this once we support upgrading to HA
	if service.Installed() {
		return nil
	}

	if err := makeJournalDirs(dbDir); err != nil {
		return fmt.Errorf("Error creating journal directories: %v", err)
	}
	return upstartConfInstall(upstartConf)
}

// ServiceName returns the name of the upstart service config for mongo using
// the given namespace.
func ServiceName(namespace string) string {
	if namespace != "" {
		return fmt.Sprintf("%s-%s", serviceName, namespace)
	}
	return serviceName
}

func makeJournalDirs(dataDir string) error {
	journalDir := path.Join(dataDir, "journal")

	if err := os.MkdirAll(journalDir, 0700); err != nil {
		logger.Errorf("failed to make mongo journal dir %s: %v", journalDir, err)
		return err
	}

	// manually create the prealloc files, since otherwise they get created as 100M files.
	zeroes := make([]byte, 64*1024) // should be enough for anyone
	for x := 0; x < 3; x++ {
		name := fmt.Sprintf("prealloc.%d", x)
		filename := filepath.Join(journalDir, name)
		f, err := os.OpenFile(filename, os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0700)
		// TODO(jam) 2014-04-12 https://launchpad.net/bugs/1306902
		// When we support upgrading Mongo into Replica mode, we should
		// start rewriting the upstart config
		if os.IsExist(err) {
			// already exists, don't overwrite
			continue
		}
		if err != nil {
			return fmt.Errorf("failed to open mongo prealloc file %q: %v", filename, err)
		}
		defer f.Close()
		for total := 0; total < 1024*1024; {
			n, err := f.Write(zeroes)
			if err != nil {
				return fmt.Errorf("failed to write to mongo prealloc file %q: %v", filename, err)
			}
			total += n
		}
	}
	return nil
}

func sslKeyPath(dataDir string) string {
	return filepath.Join(dataDir, "server.pem")
}

func sharedSecretPath(dataDir string) string {
	return filepath.Join(dataDir, SharedSecretFile)
}

// mongoUpstartService returns the upstart config for the mongo state service.
//
func mongoUpstartService(namespace, dataDir, dbDir string, port int) (*upstart.Conf, error) {
	svc := upstart.NewService(ServiceName(namespace))

	mongoPath, err := MongodPath()
	if err != nil {
		return nil, err
	}

	conf := &upstart.Conf{
		Service: *svc,
		Desc:    "juju state database",
		Limit: map[string]string{
			"nofile": fmt.Sprintf("%d %d", maxFiles, maxFiles),
			"nproc":  fmt.Sprintf("%d %d", maxProcs, maxProcs),
		},
		Cmd: mongoPath + " --auth" +
			" --dbpath=" + dbDir +
			" --sslOnNormalPorts" +
			" --sslPEMKeyFile " + utils.ShQuote(sslKeyPath(dataDir)) +
			" --sslPEMKeyPassword ignored" +
			" --bind_ip 0.0.0.0" +
			" --port " + fmt.Sprint(port) +
			" --noprealloc" +
			" --syslog" +
			" --smallfiles" +
			" --replSet " + ReplicaSetName +
			" --keyFile " + utils.ShQuote(sharedSecretPath(dataDir)),
	}
	return conf, nil
}

func aptGetInstallMongod() error {
	// Only Quantal requires the PPA.
	if version.Current.Series == "quantal" {
		if err := addAptRepository("ppa:juju/stable"); err != nil {
			return err
		}
	}
	cmds := utils.AptGetPreparePackages([]string{"mongodb-server"}, version.Current.Series)
	logger.Infof("installing mongodb-server")
	for _, cmd := range cmds {
		if err := utils.AptGetInstall(cmd...); err != nil {
			return err
		}
	}
	return nil
}

func addAptRepository(name string) error {
	// add-apt-repository requires python-software-properties
	cmds := utils.AptGetPreparePackages(
		[]string{"python-software-properties"},
		version.Current.Series,
	)
	logger.Infof("installing python-software-properties")
	for _, cmd := range cmds {
		if err := utils.AptGetInstall(cmd...); err != nil {
			return err
		}
	}

	logger.Infof("adding apt repository %q", name)
	cmd := exec.Command("add-apt-repository", "-y", name)
	out, err := cmd.CombinedOutput()
	if err != nil {
		return fmt.Errorf("cannot add apt repository: %v (output %s)", err, bytes.TrimSpace(out))
	}
	return nil
}
<commit_msg>agent/mongo: minor fix<commit_after>package mongo

import (
	"bytes"
	"crypto/rand"
	"encoding/base64"
	"fmt"
	"io/ioutil"
	"net"
	"os"
	"os/exec"
	"path"
	"path/filepath"

	"github.com/juju/loggo"
	"labix.org/v2/mgo"

	"launchpad.net/juju-core/instance"
	"launchpad.net/juju-core/replicaset"
	"launchpad.net/juju-core/state/api/params"
	"launchpad.net/juju-core/upstart"
	"launchpad.net/juju-core/utils"
	"launchpad.net/juju-core/version"
)

const (
	maxFiles = 65000
	maxProcs = 20000

	serviceName = "juju-db"

	// SharedSecretFile is the name of the Mongo shared secret file
	// located within the Juju data directory.
	SharedSecretFile = "shared-secret"

	ReplicaSetName = "juju"
)

var (
	logger          = loggo.GetLogger("juju.agent.mongo")
	mongoConfigPath = "/etc/default/mongodb"

	// JujuMongodPath holds the default path to the juju-specific mongod.
	JujuMongodPath = "/usr/lib/juju/bin/mongod"

	upstartConfInstall          = (*upstart.Conf).Install
	upstartServiceStopAndRemove = (*upstart.Service).StopAndRemove
)

// WithAddresses represents an entity that has a set of
// addresses. e.g. a state Machine object
type WithAddresses interface {
	Addresses() []instance.Address
}

// IsMaster returns a boolean that represents whether the given
// machine's peer address is the primary mongo host for the replicaset
func IsMaster(session *mgo.Session, obj WithAddresses) (bool, error) {
	addrs := obj.Addresses()

	masterHostPort, err := replicaset.MasterHostPort(session)

	// If the replica set has not been configured, then we
	// can have only one master and the caller must
	// be that master.
	if err == replicaset.ErrMasterNotConfigured {
		return true, nil
	}

	if err != nil {
		return false, err
	}

	masterAddr, _, err := net.SplitHostPort(masterHostPort)
	if err != nil {
		return false, err
	}

	machinePeerAddr := SelectPeerAddress(addrs)
	return machinePeerAddr == masterAddr, nil
}

// SelectPeerAddress returns the address to use as the
// mongo replica set peer address by selecting it from the given addresses.
func SelectPeerAddress(addrs []instance.Address) string {
	return instance.SelectInternalAddress(addrs, false)
}

// SelectPeerHostPort returns the HostPort to use as the
// mongo replica set peer by selecting it from the given hostPorts.
func SelectPeerHostPort(hostPorts []instance.HostPort) string {
	return instance.SelectInternalHostPort(hostPorts, false)
}

// GenerateSharedSecret generates a pseudo-random shared secret (keyfile)
// for use with Mongo replica sets.
func GenerateSharedSecret() (string, error) {
	// "A key’s length must be between 6 and 1024 characters and may
	// only contain characters in the base64 set."
	//   -- http://docs.mongodb.org/manual/tutorial/generate-key-file/
	buf := make([]byte, base64.StdEncoding.DecodedLen(1024))
	if _, err := rand.Read(buf); err != nil {
		return "", fmt.Errorf("cannot read random secret: %v", err)
	}
	return base64.StdEncoding.EncodeToString(buf), nil
}

// MongoPath returns the executable path to be used to run mongod on this
// machine. If the juju-bundled version of mongo exists, it will return that
// path, otherwise it will return the command to run mongod from the path.
func MongodPath() (string, error) {
	if _, err := os.Stat(JujuMongodPath); err == nil {
		return JujuMongodPath, nil
	}

	path, err := exec.LookPath("mongod")
	if err != nil {
		return "", err
	}
	return path, nil
}

// RemoveService removes the mongoDB upstart service from this machine.
func RemoveService(namespace string) error {
	svc := upstart.NewService(ServiceName(namespace))
	return upstartServiceStopAndRemove(svc)
}

// EnsureMongoServer ensures that the correct mongo upstart script is installed
// and running.
//
// This method will remove old versions of the mongo upstart script as necessary
// before installing the new version.
//
// The namespace is a unique identifier to prevent multiple instances of mongo
// on this machine from colliding. This should be empty unless using
// the local provider.
func EnsureMongoServer(dataDir string, namespace string, info params.StateServingInfo) error {
	logger.Infof("Ensuring mongo server is running; dataDir %s; port %d", dataDir, info.StatePort)
	dbDir := filepath.Join(dataDir, "db")

	if err := os.MkdirAll(dbDir, 0700); err != nil {
		return fmt.Errorf("cannot create mongo dbdir: %v", err)
	}

	certKey := info.Cert + "\n" + info.PrivateKey
	err := utils.AtomicWriteFile(sslKeyPath(dataDir), []byte(certKey), 0600)
	if err != nil {
		return fmt.Errorf("cannot write SSL key: %v", err)
	}

	err = utils.AtomicWriteFile(sharedSecretPath(dataDir), []byte(info.SharedSecret), 0600)
	if err != nil {
		return fmt.Errorf("cannot write mongod shared secret: %v", err)
	}

	// Disable the default mongodb installed by the mongodb-server package.
	// Only do this if the file doesn't exist already, so users can run
	// their own mongodb server if they wish to.
	if _, err := os.Stat(mongoConfigPath); os.IsNotExist(err) {
		err = ioutil.WriteFile(
			mongoConfigPath,
			[]byte("ENABLE_MONGODB=no"),
			0644,
		)
		if err != nil {
			return err
		}
	}

	if err := aptGetInstallMongod(); err != nil {
		return fmt.Errorf("cannot install mongod: %v", err)
	}

	upstartConf, err := mongoUpstartService(namespace, dataDir, dbDir, info.StatePort)
	if err != nil {
		return err
	}

	// TODO(natefinch) 2014-04-12 https://launchpad.net/bugs/1306902
	// remove this once we support upgrading to HA
	if upstartConf.Installed() {
		return nil
	}

	if err := makeJournalDirs(dbDir); err != nil {
		return fmt.Errorf("Error creating journal directories: %v", err)
	}
	return upstartConfInstall(upstartConf)
}

// ServiceName returns the name of the upstart service config for mongo using
// the given namespace.
func ServiceName(namespace string) string {
	if namespace != "" {
		return fmt.Sprintf("%s-%s", serviceName, namespace)
	}
	return serviceName
}

func makeJournalDirs(dataDir string) error {
	journalDir := path.Join(dataDir, "journal")

	if err := os.MkdirAll(journalDir, 0700); err != nil {
		logger.Errorf("failed to make mongo journal dir %s: %v", journalDir, err)
		return err
	}

	// manually create the prealloc files, since otherwise they get created as 100M files.
	zeroes := make([]byte, 64*1024) // should be enough for anyone
	for x := 0; x < 3; x++ {
		name := fmt.Sprintf("prealloc.%d", x)
		filename := filepath.Join(journalDir, name)
		f, err := os.OpenFile(filename, os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0700)
		// TODO(jam) 2014-04-12 https://launchpad.net/bugs/1306902
		// When we support upgrading Mongo into Replica mode, we should
		// start rewriting the upstart config
		if os.IsExist(err) {
			// already exists, don't overwrite
			continue
		}
		if err != nil {
			return fmt.Errorf("failed to open mongo prealloc file %q: %v", filename, err)
		}
		defer f.Close()
		for total := 0; total < 1024*1024; {
			n, err := f.Write(zeroes)
			if err != nil {
				return fmt.Errorf("failed to write to mongo prealloc file %q: %v", filename, err)
			}
			total += n
		}
	}
	return nil
}

func sslKeyPath(dataDir string) string {
	return filepath.Join(dataDir, "server.pem")
}

func sharedSecretPath(dataDir string) string {
	return filepath.Join(dataDir, SharedSecretFile)
}

// mongoUpstartService returns the upstart config for the mongo state service.
//
func mongoUpstartService(namespace, dataDir, dbDir string, port int) (*upstart.Conf, error) {
	svc := upstart.NewService(ServiceName(namespace))

	mongoPath, err := MongodPath()
	if err != nil {
		return nil, err
	}

	conf := &upstart.Conf{
		Service: *svc,
		Desc:    "juju state database",
		Limit: map[string]string{
			"nofile": fmt.Sprintf("%d %d", maxFiles, maxFiles),
			"nproc":  fmt.Sprintf("%d %d", maxProcs, maxProcs),
		},
		Cmd: mongoPath + " --auth" +
			" --dbpath=" + dbDir +
			" --sslOnNormalPorts" +
			" --sslPEMKeyFile " + utils.ShQuote(sslKeyPath(dataDir)) +
			" --sslPEMKeyPassword ignored" +
			" --bind_ip 0.0.0.0" +
			" --port " + fmt.Sprint(port) +
			" --noprealloc" +
			" --syslog" +
			" --smallfiles" +
			" --replSet " + ReplicaSetName +
			" --keyFile " + utils.ShQuote(sharedSecretPath(dataDir)),
	}
	return conf, nil
}

func aptGetInstallMongod() error {
	// Only Quantal requires the PPA.
	if version.Current.Series == "quantal" {
		if err := addAptRepository("ppa:juju/stable"); err != nil {
			return err
		}
	}
	cmds := utils.AptGetPreparePackages([]string{"mongodb-server"}, version.Current.Series)
	logger.Infof("installing mongodb-server")
	for _, cmd := range cmds {
		if err := utils.AptGetInstall(cmd...); err != nil {
			return err
		}
	}
	return nil
}

func addAptRepository(name string) error {
	// add-apt-repository requires python-software-properties
	cmds := utils.AptGetPreparePackages(
		[]string{"python-software-properties"},
		version.Current.Series,
	)
	logger.Infof("installing python-software-properties")
	for _, cmd := range cmds {
		if err := utils.AptGetInstall(cmd...); err != nil {
			return err
		}
	}

	logger.Infof("adding apt repository %q", name)
	cmd := exec.Command("add-apt-repository", "-y", name)
	out, err := cmd.CombinedOutput()
	if err != nil {
		return fmt.Errorf("cannot add apt repository: %v (output %s)", err, bytes.TrimSpace(out))
	}
	return nil
}
<|endoftext|>


<commit_before>package analyses

import (
	"time"

	"github.com/ion-channel/ionic/rulesets"
	"github.com/ion-channel/ionic/scans"
)

// Analysis is a representation of an Ion Channel Analysis within the system
type Analysis struct {
	ID            string       `json:"id" xml:"id"`
	LatestID      string       `json:"latest_id" xml:"latest_id"`
	TeamID        string       `json:"team_id" xml:"team_id"`
	ProjectID     string       `json:"project_id" xml:"project_id"`
	Name          string       `json:"name" xml:"name"`
	Text          *string      `json:"text" xml:"text"`
	Type          string       `json:"type" xml:"type"`
	Source        string       `json:"source" xml:"source"`
	Branch        string       `json:"branch" xml:"branch"`
	Description   string       `json:"description" xml:"description"`
	Status        string       `json:"status" xml:"status"`
	RulesetID     string       `json:"ruleset_id" xml:"ruleset_id"`
	CreatedAt     time.Time    `json:"created_at" xml:"created_at"`
	UpdatedAt     time.Time    `json:"updated_at" xml:"updated_at"`
	Duration      float64      `json:"duration" xml:"duration"`
	TriggerHash   string       `json:"trigger_hash" xml:"trigger_hash"`
	TriggerText   string       `json:"trigger_text" xml:"trigger_text"`
	TriggerAuthor string       `json:"trigger_author" xml:"trigger_author"`
	ScanSummaries []scans.Scan `json:"scan_summaries" xml:"scan_summaries"`
	Public        bool         `json:"public" xml:"public"`
}

// Summary is a representation of a summarized Ion Channel Analysis
// within the system
type Summary struct {
	ID            string    `json:"id"`
	AnalysisID    string    `json:"analysis_id"`
	TeamID        string    `json:"team_id"`
	Branch        string    `json:"branch"`
	Description   string    `json:"description"`
	Risk          string    `json:"risk"`
	Summary       string    `json:"summary"`
	Passed        bool      `json:"passed"`
	RulesetID     string    `json:"ruleset_id"`
	RulesetName   string    `json:"ruleset_name"`
	Duration      float64   `json:"duration"`
	CreatedAt     time.Time `json:"created_at"`
	TriggerHash   string    `json:"trigger_hash"`
	TriggerText   string    `json:"trigger_text"`
	TriggerAuthor string    `json:"trigger_author"`
	Trigger       string    `json:"trigger"`
}

// NewSummary takes an Analysis and AppliedRulesetSummary to calculate and
// return a Summary of the Analysis
func NewSummary(a *Analysis, appliedRuleset *rulesets.AppliedRulesetSummary) *Summary {
	if a != nil {
		rulesetName := "N/A"
		risk := "high"
		passed := false

		if appliedRuleset != nil {
			risk, passed = appliedRuleset.SummarizeEvaluation()

			if appliedRuleset.RuleEvaluationSummary != nil && appliedRuleset.RuleEvaluationSummary.RulesetName != "" {
				rulesetName = appliedRuleset.RuleEvaluationSummary.RulesetName
			}
		}

		return &Summary{
			ID:            a.ID,
			AnalysisID:    a.ID,
			TeamID:        a.TeamID,
			Branch:        a.Branch,
			Description:   a.Description,
			Risk:          risk,
			Summary:       "",
			Passed:        passed,
			RulesetID:     a.RulesetID,
			RulesetName:   rulesetName,
			Duration:      a.Duration,
			CreatedAt:     a.CreatedAt,
			TriggerHash:   a.TriggerHash,
			TriggerText:   a.TriggerText,
			TriggerAuthor: a.TriggerAuthor,
			Trigger:       "source commit",
		}
	}

	return &Summary{}
}
<commit_msg>creating new analysis navigation<commit_after>package analyses

import (
	"time"

	"github.com/ion-channel/ionic/rulesets"
	"github.com/ion-channel/ionic/scans"
)

// Analysis is a representation of an Ion Channel Analysis within the system
type Analysis struct {
	ID            string       `json:"id" xml:"id"`
	TeamID        string       `json:"team_id" xml:"team_id"`
	ProjectID     string       `json:"project_id" xml:"project_id"`
	Name          string       `json:"name" xml:"name"`
	Text          *string      `json:"text" xml:"text"`
	Type          string       `json:"type" xml:"type"`
	Source        string       `json:"source" xml:"source"`
	Branch        string       `json:"branch" xml:"branch"`
	Description   string       `json:"description" xml:"description"`
	Status        string       `json:"status" xml:"status"`
	RulesetID     string       `json:"ruleset_id" xml:"ruleset_id"`
	CreatedAt     time.Time    `json:"created_at" xml:"created_at"`
	UpdatedAt     time.Time    `json:"updated_at" xml:"updated_at"`
	Duration      float64      `json:"duration" xml:"duration"`
	TriggerHash   string       `json:"trigger_hash" xml:"trigger_hash"`
	TriggerText   string       `json:"trigger_text" xml:"trigger_text"`
	TriggerAuthor string       `json:"trigger_author" xml:"trigger_author"`
	ScanSummaries []scans.Scan `json:"scan_summaries" xml:"scan_summaries"`
	Public        bool         `json:"public" xml:"public"`
}

// Summary is a representation of a summarized Ion Channel Analysis
// within the system
type Summary struct {
	ID            string    `json:"id"`
	AnalysisID    string    `json:"analysis_id"`
	TeamID        string    `json:"team_id"`
	Branch        string    `json:"branch"`
	Description   string    `json:"description"`
	Risk          string    `json:"risk"`
	Summary       string    `json:"summary"`
	Passed        bool      `json:"passed"`
	RulesetID     string    `json:"ruleset_id"`
	RulesetName   string    `json:"ruleset_name"`
	Duration      float64   `json:"duration"`
	CreatedAt     time.Time `json:"created_at"`
	TriggerHash   string    `json:"trigger_hash"`
	TriggerText   string    `json:"trigger_text"`
	TriggerAuthor string    `json:"trigger_author"`
	Trigger       string    `json:"trigger"`
}

// Navigation represents a navigational meta data reference to given analysis
type Navigation struct {
	Analysis       *Analysis `json:"analysis"`
	LatestAnalysis *Analysis `json:"latest_analysis"`
}

// NewSummary takes an Analysis and AppliedRulesetSummary to calculate and
// return a Summary of the Analysis
func NewSummary(a *Analysis, appliedRuleset *rulesets.AppliedRulesetSummary) *Summary {
	if a != nil {
		rulesetName := "N/A"
		risk := "high"
		passed := false

		if appliedRuleset != nil {
			risk, passed = appliedRuleset.SummarizeEvaluation()

			if appliedRuleset.RuleEvaluationSummary != nil && appliedRuleset.RuleEvaluationSummary.RulesetName != "" {
				rulesetName = appliedRuleset.RuleEvaluationSummary.RulesetName
			}
		}

		return &Summary{
			ID:            a.ID,
			AnalysisID:    a.ID,
			TeamID:        a.TeamID,
			Branch:        a.Branch,
			Description:   a.Description,
			Risk:          risk,
			Summary:       "",
			Passed:        passed,
			RulesetID:     a.RulesetID,
			RulesetName:   rulesetName,
			Duration:      a.Duration,
			CreatedAt:     a.CreatedAt,
			TriggerHash:   a.TriggerHash,
			TriggerText:   a.TriggerText,
			TriggerAuthor: a.TriggerAuthor,
			Trigger:       "source commit",
		}
	}

	return &Summary{}
}
<|endoftext|>


<commit_before>package search

import (
	"container/heap"
	"fmt"
	"github.com/balzaczyy/golucene/index"
	"math"
)

/** Holds one hit in {@link TopDocs}. */
type ScoreDoc struct {
	/** The score of this document for the query. */
	score float32
	/** A hit document's number.
	 * @see IndexSearcher#doc(int) */
	doc int
	/** Only set by {@link TopDocs#merge} */
	shardIndex int
}

func newScoreDoc(doc int, score float32) ScoreDoc {
	return ScoreDoc{score: score, doc: doc}
}

func newShardedScoreDoc(doc int, score float32, shardIndex int) ScoreDoc {
	return ScoreDoc{score, doc, shardIndex}
}

func (d ScoreDoc) String() string {
	return fmt.Sprintf("doc=%v score=%v shardIndex=%v", d.doc, d.score, d.shardIndex)
}

type PriorityQueue struct {
	items []interface{}
	less  func(i, j int) bool
}

func (pq PriorityQueue) Len() int            { return len(pq.items) }
func (pq PriorityQueue) Less(i, j int) bool  { return pq.less(i, j) }
func (pq PriorityQueue) Swap(i, j int)       { pq.items[i], pq.items[j] = pq.items[j], pq.items[i] }
func (pq *PriorityQueue) Push(x interface{}) { pq.items = append(pq.items, x) }
func (pq *PriorityQueue) Pop() interface{} {
	n := pq.Len()
	ans := pq.items[n-1]
	pq.items = pq.items[0 : n-1]
	return ans
}

type TopDocs struct {
	totalHits int
	scoreDocs []ScoreDoc
	maxScore  float64
}

type Collector interface {
	SetScorer(s *Scorer)
	Collect(doc int)
	SetNextReader(ctx index.AtomicReaderContext)
	AcceptsDocsOutOfOrder() bool
}

type TopDocsCollector struct {
	Collector
	self      interface{}
	pq        *PriorityQueue // PriorityQueue
	TotalHits int
	// populateResults func(results []ScoreDoc, howMany int)
	// newTopDocs      func(results []ScoreDoc, start int) TopDocs
	// topDocsSize     func() int
	// TopDocs         func() TopDocs
}

func newTopDocsCollector(self interface{}, pq *PriorityQueue) *TopDocsCollector {
	return &TopDocsCollector{self: self, pq: pq}
}

func (c *TopDocsCollector) populateResults(results []ScoreDoc, howMany int) {
	for i := howMany - 1; i >= 0; i-- {
		results[i] = *(heap.Pop(c.pq).(*ScoreDoc))
	}
}

func (c *TopDocsCollector) newTopDocs(results []ScoreDoc, start int) TopDocs {
	panic("not implemented yet")
	// if results == nil {
	// 	return TopDocs{0, []ScoreDoc{}, math.NaN()}
	// }
	// return TopDocs{c.TotalHits, results, math.NaN()}
}

func (c *TopDocsCollector) topDocsSize() int {
	// In case pq was populated with sentinel values, there might be less
	// results than pq.size(). Therefore return all results until either
	// pq.size() or totalHits.
	if n := c.pq.Len(); c.TotalHits >= n {
		return n
	}
	return c.TotalHits
}

func (c *TopDocsCollector) TopDocs() TopDocs {
	// In case pq was populated with sentinel values, there might be less
	// results than pq.size(). Therefore return all results until either
	// pq.size() or totalHits.
	return c.TopDocsRange(0, c.topDocsSize())
}

func (c *TopDocsCollector) TopDocsRange(start, howMany int) TopDocs {
	// In case pq was populated with sentinel values, there might be less
	// results than pq.size(). Therefore return all results until either
	// pq.size() or totalHits.
	size := c.topDocsSize()

	// Don't bother to throw an exception, just return an empty TopDocs in case
	// the parameters are invalid or out of range.
	// TODO: shouldn't we throw IAE if apps give bad params here so they dont
	// have sneaky silent bugs?
	if start < 0 || start >= size || howMany <= 0 {
		return c.newTopDocs(nil, start)
	}

	// We know that start < pqsize, so just fix howMany.
	if size-start < howMany {
		howMany = size - start
	}
	results := make([]ScoreDoc, howMany)

	// pq's pop() returns the 'least' element in the queue, therefore need
	// to discard the first ones, until we reach the requested range.
	// Note that this loop will usually not be executed, since the common usage
	// should be that the caller asks for the last howMany results. However it's
	// needed here for completeness.
	for i := c.pq.Len() - start - howMany; i > 0; i-- {
		heap.Pop(c.pq)
	}

	// Get the requested results from pq.
	c.populateResults(results, howMany)

	return c.newTopDocs(results, start)
}

type TopScoreDocCollector struct {
	*TopDocsCollector
	pqTop   ScoreDoc
	docBase int
	scorer  Scorer
}

func newTocScoreDocCollector(numHits int) *TopScoreDocCollector {
	docs := make([]interface{}, numHits)
	for i, _ := range docs {
		docs[i] = newScoreDoc(math.MaxInt32, -math.MaxFloat32)
	}
	pq := &PriorityQueue{items: docs}
	pq.less = func(i, j int) bool {
		hitA := pq.items[i].(ScoreDoc)
		hitB := pq.items[j].(ScoreDoc)
		if hitA.score == hitB.score {
			return hitA.doc > hitB.doc
		}
		return hitA.score < hitB.score
	}
	heap.Init(pq)

	pqTop := heap.Pop(pq).(ScoreDoc)
	heap.Push(pq, pqTop)
	c := &TopScoreDocCollector{pqTop: pqTop}
	c.TopDocsCollector = newTopDocsCollector(c, pq)
	return c
}

func (c *TopScoreDocCollector) newTopDocs(results []ScoreDoc, start int) TopDocs {
	if results == nil {
		return TopDocs{0, []ScoreDoc{}, math.NaN()}
	}

	// We need to compute maxScore in order to set it in TopDocs. If start == 0,
	// it means the largest element is already in results, use its score as
	// maxScore. Otherwise pop everything else, until the largest element is
	// extracted and use its score as maxScore.
	maxScore := math.NaN()
	if start == 0 {
		maxScore = float64(results[0].score)
	} else {
		pq := c.TopDocsCollector.pq
		for i := pq.Len(); i > 1; i-- {
			heap.Pop(pq)
		}
		maxScore = float64(heap.Pop(pq).(ScoreDoc).score)
	}

	return TopDocs{c.TopDocsCollector.TotalHits, results, maxScore}
}

func (c *TopScoreDocCollector) SetNextReader(ctx index.AtomicReaderContext) {
	c.docBase = ctx.DocBase
}

func NewTopScoreDocCollector(numHits int, after ScoreDoc, docsScoredInOrder bool) *TopDocsCollector {
	if numHits < 0 {
		panic("numHits must be > 0; please use TotalHitCountCollector if you just need the total hit count")
	}

	if docsScoredInOrder {
		return NewInOrderTopScoreDocCollector(numHits).TopScoreDocCollector.TopDocsCollector.Collector.(*TopDocsCollector)
		// TODO support paging
	} else {
		panic("not supported yet")
	}
}

// Assumes docs are scored in order.
type InOrderTopScoreDocCollector struct {
	*TopScoreDocCollector
}

func NewInOrderTopScoreDocCollector(numHits int) *InOrderTopScoreDocCollector {
	return &InOrderTopScoreDocCollector{newTocScoreDocCollector(numHits)}
}

func (c *InOrderTopScoreDocCollector) Collect(doc int) {
	score := c.scorer.Score()

	// This collector cannot handle these scores:
	// assert score != -math.MaxFloat64
	// assert !math.IsNaN(score)

	c.TotalHits++
	if score <= c.pqTop.score {
		// Since docs are returned in-order (i.e., increasing doc Id), a document
		// with equal score to pqTop.score cannot compete since HitQueue favors
		// documents with lower doc Ids. Therefore reject those docs too.
		return
	}
	c.pqTop.doc = doc + c.docBase
	c.pqTop.score = score
	heap.Pop(c.pq)
	heap.Push(c.pq, c.pqTop)
}

func (c *InOrderTopScoreDocCollector) AcceptsDocsOutOfOrder() bool {
	return false
}
<commit_msg>convert TopDocsCollector to interface<commit_after>package search

import (
	"container/heap"
	"fmt"
	"github.com/balzaczyy/golucene/index"
	"math"
)

/** Holds one hit in {@link TopDocs}. */
type ScoreDoc struct {
	/** The score of this document for the query. */
	score float32
	/** A hit document's number.
	 * @see IndexSearcher#doc(int) */
	doc int
	/** Only set by {@link TopDocs#merge} */
	shardIndex int
}

func newScoreDoc(doc int, score float32) ScoreDoc {
	return ScoreDoc{score: score, doc: doc}
}

func newShardedScoreDoc(doc int, score float32, shardIndex int) ScoreDoc {
	return ScoreDoc{score, doc, shardIndex}
}

func (d ScoreDoc) String() string {
	return fmt.Sprintf("doc=%v score=%v shardIndex=%v", d.doc, d.score, d.shardIndex)
}

type PriorityQueue struct {
	items []interface{}
	less  func(i, j int) bool
}

func (pq PriorityQueue) Len() int            { return len(pq.items) }
func (pq PriorityQueue) Less(i, j int) bool  { return pq.less(i, j) }
func (pq PriorityQueue) Swap(i, j int)       { pq.items[i], pq.items[j] = pq.items[j], pq.items[i] }
func (pq *PriorityQueue) Push(x interface{}) { pq.items = append(pq.items, x) }
func (pq *PriorityQueue) Pop() interface{} {
	n := pq.Len()
	ans := pq.items[n-1]
	pq.items = pq.items[0 : n-1]
	return ans
}

type TopDocs struct {
	totalHits int
	scoreDocs []ScoreDoc
	maxScore  float64
}

type Collector interface {
	SetScorer(s *Scorer)
	Collect(doc int)
	SetNextReader(ctx index.AtomicReaderContext)
	AcceptsDocsOutOfOrder() bool
}

// search/TopDocsCollector.java
/**
 * A base class for all collectors that return a {@link TopDocs} output. This
 * collector allows easy extension by providing a single constructor which
 * accepts a {@link PriorityQueue} as well as protected members for that
 * priority queue and a counter of the number of total hits.

 * Extending classes can override any of the methods to provide their own
 * implementation, as well as avoid the use of the priority queue entirely by
 * passing null to {@link #TopDocsCollector(PriorityQueue)}. In that case
 * however, you might want to consider overriding all methods, in order to avoid
 * a NullPointerException.
 */
type TopDocsCollector interface {
	Collector
	/** Returns the top docs that were collected by this collector. */
	TopDocs() TopDocs
	/**
	 * Returns the documents in the rage [start .. start+howMany) that were
	 * collected by this collector. Note that if start >= pq.size(), an empty
	 * TopDocs is returned, and if pq.size() - start < howMany, then only the
	 * available documents in [start .. pq.size()) are returned.

	 * This method is useful to call in case pagination of search results is
	 * allowed by the search application, as well as it attempts to optimize the
	 * memory used by allocating only as much as requested by howMany.

	 * NOTE: you cannot call this method more than once for each search
	 * execution. If you need to call it more than once, passing each time a
	 * different range, you should call {@link #topDocs()} and work with the
	 * returned {@link TopDocs} object, which will contain all the results this
	 * search execution collected.
	 */
	TopDocsRange(start, howMany int) TopDocs
}

type TopDocsCreator interface {
	/**
	 * Populates the results array with the ScoreDoc instances. This can be
	 * overridden in case a different ScoreDoc type should be returned.
	 */
	populateResults(results []ScoreDoc, howMany int)
	/**
	 * Returns a {@link TopDocs} instance containing the given results. If
	 * results is null it means there are no results to return,
	 * either because there were 0 calls to collect() or because the arguments to
	 * topDocs were invalid.
	 */
	newTopDocs(results []ScoreDoc, start int) TopDocs
	/** The number of valid PQ entries */
	topDocsSize() int
}

type abstractTopDocsCollector struct {
	Collector
	TopDocsCreator
	pq        *PriorityQueue // PriorityQueue
	TotalHits int
}

func newTopDocsCollector(self interface{}, pq *PriorityQueue) *abstractTopDocsCollector {
	return &abstractTopDocsCollector{
		Collector:      self.(Collector),
		TopDocsCreator: self.(TopDocsCreator),
		pq:             pq,
	}
}

func (c *abstractTopDocsCollector) AcceptsDocsOutOfOrder() bool {
	return false
}

func (c *abstractTopDocsCollector) populateResults(results []ScoreDoc, howMany int) {
	for i := howMany - 1; i >= 0; i-- {
		results[i] = heap.Pop(c.pq).(ScoreDoc)
	}
}

func (c *abstractTopDocsCollector) topDocsSize() int {
	// In case pq was populated with sentinel values, there might be less
	// results than pq.size(). Therefore return all results until either
	// pq.size() or totalHits.
	if n := c.pq.Len(); c.TotalHits >= n {
		return n
	}
	return c.TotalHits
}

func (c *abstractTopDocsCollector) TopDocs() TopDocs {
	// In case pq was populated with sentinel values, there might be less
	// results than pq.size(). Therefore return all results until either
	// pq.size() or totalHits.
	return c.TopDocsRange(0, c.topDocsSize())
}

func (c *abstractTopDocsCollector) TopDocsRange(start, howMany int) TopDocs {
	// In case pq was populated with sentinel values, there might be less
	// results than pq.size(). Therefore return all results until either
	// pq.size() or totalHits.
	size := c.topDocsSize()

	// Don't bother to throw an exception, just return an empty TopDocs in case
	// the parameters are invalid or out of range.
	// TODO: shouldn't we throw IAE if apps give bad params here so they dont
	// have sneaky silent bugs?
	if start < 0 || start >= size || howMany <= 0 {
		return c.newTopDocs(nil, start)
	}

	// We know that start < pqsize, so just fix howMany.
	if size-start < howMany {
		howMany = size - start
	}
	results := make([]ScoreDoc, howMany)

	// pq's pop() returns the 'least' element in the queue, therefore need
	// to discard the first ones, until we reach the requested range.
	// Note that this loop will usually not be executed, since the common usage
	// should be that the caller asks for the last howMany results. However it's
	// needed here for completeness.
	for i := c.pq.Len() - start - howMany; i > 0; i-- {
		heap.Pop(c.pq)
	}

	// Get the requested results from pq.
	c.populateResults(results, howMany)

	return c.newTopDocs(results, start)
}

type TopScoreDocCollector struct {
	*abstractTopDocsCollector
	pqTop   ScoreDoc
	docBase int
	scorer  Scorer
}

func newTocScoreDocCollector(numHits int) *TopScoreDocCollector {
	docs := make([]interface{}, numHits)
	for i, _ := range docs {
		docs[i] = newScoreDoc(math.MaxInt32, -math.MaxFloat32)
	}
	pq := &PriorityQueue{items: docs}
	pq.less = func(i, j int) bool {
		hitA := pq.items[i].(ScoreDoc)
		hitB := pq.items[j].(ScoreDoc)
		if hitA.score == hitB.score {
			return hitA.doc > hitB.doc
		}
		return hitA.score < hitB.score
	}
	heap.Init(pq)

	pqTop := heap.Pop(pq).(ScoreDoc)
	heap.Push(pq, pqTop)
	c := &TopScoreDocCollector{pqTop: pqTop}
	c.abstractTopDocsCollector = newTopDocsCollector(c, pq)
	return c
}

func (c *TopScoreDocCollector) newTopDocs(results []ScoreDoc, start int) TopDocs {
	if results == nil {
		return TopDocs{0, []ScoreDoc{}, math.NaN()}
	}

	// We need to compute maxScore in order to set it in TopDocs. If start == 0,
	// it means the largest element is already in results, use its score as
	// maxScore. Otherwise pop everything else, until the largest element is
	// extracted and use its score as maxScore.
	maxScore := math.NaN()
	if start == 0 {
		maxScore = float64(results[0].score)
	} else {
		pq := c.pq
		for i := pq.Len(); i > 1; i-- {
			heap.Pop(pq)
		}
		maxScore = float64(heap.Pop(pq).(ScoreDoc).score)
	}

	return TopDocs{c.TotalHits, results, maxScore}
}

func (c *TopScoreDocCollector) SetNextReader(ctx index.AtomicReaderContext) {
	c.docBase = ctx.DocBase
}

func NewTopScoreDocCollector(numHits int, after ScoreDoc, docsScoredInOrder bool) TopDocsCollector {
	if numHits < 0 {
		panic("numHits must be > 0; please use TotalHitCountCollector if you just need the total hit count")
	}

	if docsScoredInOrder {
		return NewInOrderTopScoreDocCollector(numHits)
		// TODO support paging
	} else {
		panic("not supported yet")
	}
}

// Assumes docs are scored in order.
type InOrderTopScoreDocCollector struct {
	*TopScoreDocCollector
}

func NewInOrderTopScoreDocCollector(numHits int) *InOrderTopScoreDocCollector {
	return &InOrderTopScoreDocCollector{newTocScoreDocCollector(numHits)}
}

func (c *InOrderTopScoreDocCollector) Collect(doc int) {
	score := c.scorer.Score()

	// This collector cannot handle these scores:
	// assert score != -math.MaxFloat64
	// assert !math.IsNaN(score)

	c.TotalHits++
	if score <= c.pqTop.score {
		// Since docs are returned in-order (i.e., increasing doc Id), a document
		// with equal score to pqTop.score cannot compete since HitQueue favors
		// documents with lower doc Ids. Therefore reject those docs too.
		return
	}
	c.pqTop.doc = doc + c.docBase
	c.pqTop.score = score
	heap.Pop(c.pq)
	heap.Push(c.pq, c.pqTop)
}

func (c *InOrderTopScoreDocCollector) AcceptsDocsOutOfOrder() bool {
	return false
}
<|endoftext|>


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

package server

import (
	"fmt"
	"math"
	"path"
	"strings"
	"sync"
	"time"

	"github.com/coreos/etcd/clientv3"
	"github.com/golang/protobuf/proto"
	"github.com/juju/errors"
	"github.com/ngaut/log"
	"github.com/pingcap/kvproto/pkg/metapb"
	"github.com/pingcap/kvproto/pkg/pdpb"
	"golang.org/x/net/context"
)

var (
	errClusterNotBootstrapped = errors.New("cluster is not bootstrapped")
)

const (
	maxBatchRegionCount = 10000
)

// RaftCluster is used for cluster config management.
// Raft cluster key format:
// cluster 1 -> /1/raft, value is metapb.Cluster
// cluster 2 -> /2/raft
// For cluster 1
// store 1 -> /1/raft/s/1, value is metapb.Store
// region 1 -> /1/raft/r/1, value is metapb.Region
type RaftCluster struct {
	sync.RWMutex

	s *Server

	running bool

	clusterID   uint64
	clusterRoot string

	// cached cluster info
	cachedCluster *clusterInfo

	// balancer worker
	balancerWorker *balancerWorker
}

func (c *RaftCluster) start(meta metapb.Cluster) error {
	c.Lock()
	defer c.Unlock()

	if c.running {
		log.Warn("raft cluster has already been started")
		return nil
	}

	c.running = true

	c.cachedCluster = newClusterInfo(c.clusterRoot)
	c.cachedCluster.idAlloc = c.s.idAlloc

	c.cachedCluster.setMeta(&meta)

	// Cache all stores when start the cluster. We don't have
	// many stores, so it is OK to cache them all.
	// And we should use these cache for later ChangePeer too.
	if err := c.cacheAllStores(); err != nil {
		return errors.Trace(err)
	}

	if err := c.cacheAllRegions(); err != nil {
		return errors.Trace(err)
	}

	c.balancerWorker = newBalancerWorker(c.cachedCluster, c.s.cfg.BalanceCfg)
	c.balancerWorker.run()

	return nil
}

func (c *RaftCluster) stop() {
	c.Lock()
	defer c.Unlock()

	if !c.running {
		return
	}

	c.balancerWorker.stop()

	c.running = false
}

func (c *RaftCluster) isRunning() bool {
	c.RLock()
	defer c.RUnlock()

	return c.running
}

func (s *Server) getClusterRootPath() string {
	return path.Join(s.rootPath, "raft")
}

// GetRaftCluster gets raft cluster.
func (s *Server) GetRaftCluster() (*RaftCluster, error) {
	if s.cluster.isRunning() {
		return s.cluster, nil
	}

	return nil, nil
}

func (s *Server) createRaftCluster() error {
	if s.cluster.isRunning() {
		return nil
	}

	value, err := getValue(s.client, s.getClusterRootPath())
	if err != nil {
		return errors.Trace(err)
	}
	if value == nil {
		return nil
	}

	clusterMeta := metapb.Cluster{}
	if err = proto.Unmarshal(value, &clusterMeta); err != nil {
		return errors.Trace(err)
	}

	if err = s.cluster.start(clusterMeta); err != nil {
		return errors.Trace(err)
	}

	return nil
}

func makeStoreKey(clusterRootPath string, storeID uint64) string {
	return strings.Join([]string{clusterRootPath, "s", fmt.Sprintf("%020d", storeID)}, "/")
}

func makeRegionKey(clusterRootPath string, regionID uint64) string {
	return strings.Join([]string{clusterRootPath, "r", fmt.Sprintf("%020d", regionID)}, "/")
}

func makeStoreKeyPrefix(clusterRootPath string) string {
	return strings.Join([]string{clusterRootPath, "s", ""}, "/")
}

func checkBootstrapRequest(clusterID uint64, req *pdpb.BootstrapRequest) error {
	// TODO: do more check for request fields validation.

	storeMeta := req.GetStore()
	if storeMeta == nil {
		return errors.Errorf("missing store meta for bootstrap %d", clusterID)
	} else if storeMeta.GetId() == 0 {
		return errors.New("invalid zero store id")
	}

	regionMeta := req.GetRegion()
	if regionMeta == nil {
		return errors.Errorf("missing region meta for bootstrap %d", clusterID)
	} else if len(regionMeta.GetStartKey()) > 0 || len(regionMeta.GetEndKey()) > 0 {
		// first region start/end key must be empty
		return errors.Errorf("invalid first region key range, must all be empty for bootstrap %d", clusterID)
	} else if regionMeta.GetId() == 0 {
		return errors.New("invalid zero region id")
	}

	peers := regionMeta.GetPeers()
	if len(peers) != 1 {
		return errors.Errorf("invalid first region peer count %d, must be 1 for bootstrap %d", len(peers), clusterID)
	}

	peer := peers[0]
	if peer.GetStoreId() != storeMeta.GetId() {
		return errors.Errorf("invalid peer store id %d != %d for bootstrap %d", peer.GetStoreId(), storeMeta.GetId(), clusterID)
	}
	if peer.GetId() == 0 {
		return errors.New("invalid zero peer id")
	}

	return nil
}

func (s *Server) bootstrapCluster(req *pdpb.BootstrapRequest) (*pdpb.Response, error) {
	clusterID := s.cfg.ClusterID

	log.Infof("try to bootstrap raft cluster %d with %v", clusterID, req)

	if err := checkBootstrapRequest(clusterID, req); err != nil {
		return nil, errors.Trace(err)
	}

	clusterMeta := metapb.Cluster{
		Id:           proto.Uint64(clusterID),
		MaxPeerCount: proto.Uint32(uint32(s.cfg.MaxPeerCount)),
	}

	// Set cluster meta
	clusterValue, err := proto.Marshal(&clusterMeta)
	if err != nil {
		return nil, errors.Trace(err)
	}
	clusterRootPath := s.getClusterRootPath()

	var ops []clientv3.Op
	ops = append(ops, clientv3.OpPut(clusterRootPath, string(clusterValue)))

	// Set store meta
	storeMeta := req.GetStore()
	storePath := makeStoreKey(clusterRootPath, storeMeta.GetId())
	storeValue, err := proto.Marshal(storeMeta)
	if err != nil {
		return nil, errors.Trace(err)
	}
	ops = append(ops, clientv3.OpPut(storePath, string(storeValue)))

	regionValue, err := proto.Marshal(req.GetRegion())
	if err != nil {
		return nil, errors.Trace(err)
	}

	// Set region meta with region id.
	regionPath := makeRegionKey(clusterRootPath, req.GetRegion().GetId())
	ops = append(ops, clientv3.OpPut(regionPath, string(regionValue)))

	// TODO: we must figure out a better way to handle bootstrap failed, maybe intervene manually.
	bootstrapCmp := clientv3.Compare(clientv3.CreateRevision(clusterRootPath), "=", 0)
	ctx, cancel := context.WithTimeout(context.Background(), requestTimeout)
	resp, err := s.slowLogTxn(ctx).
		If(bootstrapCmp).
		Then(ops...).
		Commit()
	cancel()
	if err != nil {
		return nil, errors.Trace(err)
	}
	if !resp.Succeeded {
		log.Warnf("cluster %d already bootstrapped", clusterID)
		return newBootstrappedError(), nil
	}

	log.Infof("bootstrap cluster %d ok", clusterID)

	if err = s.cluster.start(clusterMeta); err != nil {
		return nil, errors.Trace(err)
	}

	return &pdpb.Response{
		Bootstrap: &pdpb.BootstrapResponse{},
	}, nil
}

func (c *RaftCluster) cacheAllStores() error {
	start := time.Now()

	key := makeStoreKeyPrefix(c.clusterRoot)
	resp, err := kvGet(c.s.client, key, clientv3.WithPrefix())
	if err != nil {
		return errors.Trace(err)
	}

	for _, kv := range resp.Kvs {
		store := &metapb.Store{}
		if err = proto.Unmarshal(kv.Value, store); err != nil {
			return errors.Trace(err)
		}

		c.cachedCluster.addStore(store)
	}
	log.Infof("cache all %d stores cost %s", len(resp.Kvs), time.Now().Sub(start))
	return nil
}

func (c *RaftCluster) cacheAllRegions() error {
	start := time.Now()

	nextID := uint64(0)
	endRegionKey := makeRegionKey(c.clusterRoot, math.MaxUint64)

	c.cachedCluster.regions.Lock()
	defer c.cachedCluster.regions.Unlock()

	for {
		key := makeRegionKey(c.clusterRoot, nextID)
		resp, err := kvGet(c.s.client, key, clientv3.WithRange(endRegionKey))
		if err != nil {
			return errors.Trace(err)
		}

		if len(resp.Kvs) == 0 {
			// No more data
			break
		}

		for _, kv := range resp.Kvs {
			region := &metapb.Region{}
			if err = proto.Unmarshal(kv.Value, region); err != nil {
				return errors.Trace(err)
			}

			nextID = region.GetId() + 1
			c.cachedCluster.regions.addRegion(region)
		}
	}

	log.Infof("cache all %d regions cost %s", len(c.cachedCluster.regions.regions), time.Now().Sub(start))
	return nil
}

func (c *RaftCluster) getRegion(regionKey []byte) (*metapb.Region, error) {
	return c.cachedCluster.regions.getRegion(regionKey), nil
}

// GetRegionAndLeader gets region and leader peer by regionID from cluster.
func (c *RaftCluster) GetRegionAndLeader(regionID uint64) (*metapb.Region, *metapb.Peer) {
	return c.cachedCluster.regions.getRegionAndLeader(regionID)
}

// GetRegions gets regions from cluster.
func (c *RaftCluster) GetRegions() []*metapb.Region {
	return c.cachedCluster.regions.getRegions()
}

// GetStores gets stores from cluster.
func (c *RaftCluster) GetStores() []*metapb.Store {
	return c.cachedCluster.getMetaStores()
}

// GetStore gets store from cluster.
func (c *RaftCluster) GetStore(storeID uint64) (*metapb.Store, *StoreStatus, error) {
	if storeID == 0 {
		return nil, nil, errors.New("invalid zero store id")
	}

	store := c.cachedCluster.getStore(storeID)
	if store == nil {
		return nil, nil, errors.Errorf("invalid store ID %d, not found", storeID)
	}

	return store.store, store.stats, nil
}

func (c *RaftCluster) putStore(store *metapb.Store) error {
	if store.GetId() == 0 {
		return errors.Errorf("invalid put store %v", store)
	}

	storeValue, err := proto.Marshal(store)
	if err != nil {
		return errors.Trace(err)
	}

	storePath := makeStoreKey(c.clusterRoot, store.GetId())
	ctx, cancel := context.WithTimeout(context.Background(), requestTimeout)
	resp, err := c.s.slowLogTxn(ctx).
		If(c.s.leaderCmp()).
		Then(clientv3.OpPut(storePath, string(storeValue))).
		Commit()
	cancel()
	if err != nil {
		return errors.Trace(err)
	}
	if !resp.Succeeded {
		return errors.Errorf("put store %v fail", store)
	}

	c.cachedCluster.addStore(store)

	return nil
}

// GetConfig gets config from cluster.
func (c *RaftCluster) GetConfig() *metapb.Cluster {
	return c.cachedCluster.getMeta()
}

func (c *RaftCluster) putConfig(meta *metapb.Cluster) error {
	if meta.GetId() != c.clusterID {
		return errors.Errorf("invalid cluster %v, mismatch cluster id %d", meta, c.clusterID)
	}

	metaValue, err := proto.Marshal(meta)
	if err != nil {
		return errors.Trace(err)
	}

	ctx, cancel := context.WithTimeout(context.Background(), requestTimeout)
	resp, err := c.s.slowLogTxn(ctx).
		If(c.s.leaderCmp()).
		Then(clientv3.OpPut(c.clusterRoot, string(metaValue))).
		Commit()
	cancel()
	if err != nil {
		return errors.Trace(err)
	}
	if !resp.Succeeded {
		return errors.Errorf("put cluster meta %v error", meta)
	}

	c.cachedCluster.setMeta(meta)

	return nil
}

// GetBalanceOperators gets the balance operators from cluster.
func (c *RaftCluster) GetBalanceOperators() map[uint64]Operator {
	return c.balancerWorker.getBalanceOperators()
}

// GetHistoryOperators gets the history operators from cluster.
func (c *RaftCluster) GetHistoryOperators() []Operator {
	return c.balancerWorker.getHistoryOperators()
}

// GetScore gets store score from balancer.
func (c *RaftCluster) GetScore(store *metapb.Store, status *StoreStatus) int {
	storeInfo := &storeInfo{
		store: store,
		stats: status,
	}

	regionCount := c.cachedCluster.regions.regionCount()
	return c.balancerWorker.storeScore(storeInfo, regionCount)
}

// FetchEvents fetches the operator events.
func (c *RaftCluster) FetchEvents(key uint64, all bool) []LogEvent {
	return c.balancerWorker.fetchEvents(key, all)
}
<commit_msg>server: fix panic bug when start cluster failed. (#184)<commit_after>// Copyright 2016 PingCAP, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// See the License for the specific language governing permissions and
// limitations under the License.

package server

import (
	"fmt"
	"math"
	"path"
	"strings"
	"sync"
	"time"

	"github.com/coreos/etcd/clientv3"
	"github.com/golang/protobuf/proto"
	"github.com/juju/errors"
	"github.com/ngaut/log"
	"github.com/pingcap/kvproto/pkg/metapb"
	"github.com/pingcap/kvproto/pkg/pdpb"
	"golang.org/x/net/context"
)

var (
	errClusterNotBootstrapped = errors.New("cluster is not bootstrapped")
)

const (
	maxBatchRegionCount = 10000
)

// RaftCluster is used for cluster config management.
// Raft cluster key format:
// cluster 1 -> /1/raft, value is metapb.Cluster
// cluster 2 -> /2/raft
// For cluster 1
// store 1 -> /1/raft/s/1, value is metapb.Store
// region 1 -> /1/raft/r/1, value is metapb.Region
type RaftCluster struct {
	sync.RWMutex

	s *Server

	running bool

	clusterID   uint64
	clusterRoot string

	// cached cluster info
	cachedCluster *clusterInfo

	// balancer worker
	balancerWorker *balancerWorker
}

func (c *RaftCluster) start(meta metapb.Cluster) error {
	c.Lock()
	defer c.Unlock()

	if c.running {
		log.Warn("raft cluster has already been started")
		return nil
	}

	c.cachedCluster = newClusterInfo(c.clusterRoot)
	c.cachedCluster.idAlloc = c.s.idAlloc

	c.cachedCluster.setMeta(&meta)

	// Cache all stores when start the cluster. We don't have
	// many stores, so it is OK to cache them all.
	// And we should use these cache for later ChangePeer too.
	if err := c.cacheAllStores(); err != nil {
		return errors.Trace(err)
	}

	if err := c.cacheAllRegions(); err != nil {
		return errors.Trace(err)
	}

	c.balancerWorker = newBalancerWorker(c.cachedCluster, c.s.cfg.BalanceCfg)
	c.balancerWorker.run()

	c.running = true

	return nil
}

func (c *RaftCluster) stop() {
	c.Lock()
	defer c.Unlock()

	if !c.running {
		return
	}

	c.balancerWorker.stop()

	c.running = false
}

func (c *RaftCluster) isRunning() bool {
	c.RLock()
	defer c.RUnlock()

	return c.running
}

func (s *Server) getClusterRootPath() string {
	return path.Join(s.rootPath, "raft")
}

// GetRaftCluster gets raft cluster.
func (s *Server) GetRaftCluster() (*RaftCluster, error) {
	if s.cluster.isRunning() {
		return s.cluster, nil
	}

	return nil, nil
}

func (s *Server) createRaftCluster() error {
	if s.cluster.isRunning() {
		return nil
	}

	value, err := getValue(s.client, s.getClusterRootPath())
	if err != nil {
		return errors.Trace(err)
	}
	if value == nil {
		return nil
	}

	clusterMeta := metapb.Cluster{}
	if err = proto.Unmarshal(value, &clusterMeta); err != nil {
		return errors.Trace(err)
	}

	if err = s.cluster.start(clusterMeta); err != nil {
		return errors.Trace(err)
	}

	return nil
}

func makeStoreKey(clusterRootPath string, storeID uint64) string {
	return strings.Join([]string{clusterRootPath, "s", fmt.Sprintf("%020d", storeID)}, "/")
}

func makeRegionKey(clusterRootPath string, regionID uint64) string {
	return strings.Join([]string{clusterRootPath, "r", fmt.Sprintf("%020d", regionID)}, "/")
}

func makeStoreKeyPrefix(clusterRootPath string) string {
	return strings.Join([]string{clusterRootPath, "s", ""}, "/")
}

func checkBootstrapRequest(clusterID uint64, req *pdpb.BootstrapRequest) error {
	// TODO: do more check for request fields validation.

	storeMeta := req.GetStore()
	if storeMeta == nil {
		return errors.Errorf("missing store meta for bootstrap %d", clusterID)
	} else if storeMeta.GetId() == 0 {
		return errors.New("invalid zero store id")
	}

	regionMeta := req.GetRegion()
	if regionMeta == nil {
		return errors.Errorf("missing region meta for bootstrap %d", clusterID)
	} else if len(regionMeta.GetStartKey()) > 0 || len(regionMeta.GetEndKey()) > 0 {
		// first region start/end key must be empty
		return errors.Errorf("invalid first region key range, must all be empty for bootstrap %d", clusterID)
	} else if regionMeta.GetId() == 0 {
		return errors.New("invalid zero region id")
	}

	peers := regionMeta.GetPeers()
	if len(peers) != 1 {
		return errors.Errorf("invalid first region peer count %d, must be 1 for bootstrap %d", len(peers), clusterID)
	}

	peer := peers[0]
	if peer.GetStoreId() != storeMeta.GetId() {
		return errors.Errorf("invalid peer store id %d != %d for bootstrap %d", peer.GetStoreId(), storeMeta.GetId(), clusterID)
	}
	if peer.GetId() == 0 {
		return errors.New("invalid zero peer id")
	}

	return nil
}

func (s *Server) bootstrapCluster(req *pdpb.BootstrapRequest) (*pdpb.Response, error) {
	clusterID := s.cfg.ClusterID

	log.Infof("try to bootstrap raft cluster %d with %v", clusterID, req)

	if err := checkBootstrapRequest(clusterID, req); err != nil {
		return nil, errors.Trace(err)
	}

	clusterMeta := metapb.Cluster{
		Id:           proto.Uint64(clusterID),
		MaxPeerCount: proto.Uint32(uint32(s.cfg.MaxPeerCount)),
	}

	// Set cluster meta
	clusterValue, err := proto.Marshal(&clusterMeta)
	if err != nil {
		return nil, errors.Trace(err)
	}
	clusterRootPath := s.getClusterRootPath()

	var ops []clientv3.Op
	ops = append(ops, clientv3.OpPut(clusterRootPath, string(clusterValue)))

	// Set store meta
	storeMeta := req.GetStore()
	storePath := makeStoreKey(clusterRootPath, storeMeta.GetId())
	storeValue, err := proto.Marshal(storeMeta)
	if err != nil {
		return nil, errors.Trace(err)
	}
	ops = append(ops, clientv3.OpPut(storePath, string(storeValue)))

	regionValue, err := proto.Marshal(req.GetRegion())
	if err != nil {
		return nil, errors.Trace(err)
	}

	// Set region meta with region id.
	regionPath := makeRegionKey(clusterRootPath, req.GetRegion().GetId())
	ops = append(ops, clientv3.OpPut(regionPath, string(regionValue)))

	// TODO: we must figure out a better way to handle bootstrap failed, maybe intervene manually.
	bootstrapCmp := clientv3.Compare(clientv3.CreateRevision(clusterRootPath), "=", 0)
	ctx, cancel := context.WithTimeout(context.Background(), requestTimeout)
	resp, err := s.slowLogTxn(ctx).
		If(bootstrapCmp).
		Then(ops...).
		Commit()
	cancel()
	if err != nil {
		return nil, errors.Trace(err)
	}
	if !resp.Succeeded {
		log.Warnf("cluster %d already bootstrapped", clusterID)
		return newBootstrappedError(), nil
	}

	log.Infof("bootstrap cluster %d ok", clusterID)

	if err = s.cluster.start(clusterMeta); err != nil {
		return nil, errors.Trace(err)
	}

	return &pdpb.Response{
		Bootstrap: &pdpb.BootstrapResponse{},
	}, nil
}

func (c *RaftCluster) cacheAllStores() error {
	start := time.Now()

	key := makeStoreKeyPrefix(c.clusterRoot)
	resp, err := kvGet(c.s.client, key, clientv3.WithPrefix())
	if err != nil {
		return errors.Trace(err)
	}

	for _, kv := range resp.Kvs {
		store := &metapb.Store{}
		if err = proto.Unmarshal(kv.Value, store); err != nil {
			return errors.Trace(err)
		}

		c.cachedCluster.addStore(store)
	}
	log.Infof("cache all %d stores cost %s", len(resp.Kvs), time.Now().Sub(start))
	return nil
}

func (c *RaftCluster) cacheAllRegions() error {
	start := time.Now()

	nextID := uint64(0)
	endRegionKey := makeRegionKey(c.clusterRoot, math.MaxUint64)

	c.cachedCluster.regions.Lock()
	defer c.cachedCluster.regions.Unlock()

	for {
		key := makeRegionKey(c.clusterRoot, nextID)
		resp, err := kvGet(c.s.client, key, clientv3.WithRange(endRegionKey))
		if err != nil {
			return errors.Trace(err)
		}

		if len(resp.Kvs) == 0 {
			// No more data
			break
		}

		for _, kv := range resp.Kvs {
			region := &metapb.Region{}
			if err = proto.Unmarshal(kv.Value, region); err != nil {
				return errors.Trace(err)
			}

			nextID = region.GetId() + 1
			c.cachedCluster.regions.addRegion(region)
		}
	}

	log.Infof("cache all %d regions cost %s", len(c.cachedCluster.regions.regions), time.Now().Sub(start))
	return nil
}

func (c *RaftCluster) getRegion(regionKey []byte) (*metapb.Region, error) {
	return c.cachedCluster.regions.getRegion(regionKey), nil
}

// GetRegionAndLeader gets region and leader peer by regionID from cluster.
func (c *RaftCluster) GetRegionAndLeader(regionID uint64) (*metapb.Region, *metapb.Peer) {
	return c.cachedCluster.regions.getRegionAndLeader(regionID)
}

// GetRegions gets regions from cluster.
func (c *RaftCluster) GetRegions() []*metapb.Region {
	return c.cachedCluster.regions.getRegions()
}

// GetStores gets stores from cluster.
func (c *RaftCluster) GetStores() []*metapb.Store {
	return c.cachedCluster.getMetaStores()
}

// GetStore gets store from cluster.
func (c *RaftCluster) GetStore(storeID uint64) (*metapb.Store, *StoreStatus, error) {
	if storeID == 0 {
		return nil, nil, errors.New("invalid zero store id")
	}

	store := c.cachedCluster.getStore(storeID)
	if store == nil {
		return nil, nil, errors.Errorf("invalid store ID %d, not found", storeID)
	}

	return store.store, store.stats, nil
}

func (c *RaftCluster) putStore(store *metapb.Store) error {
	if store.GetId() == 0 {
		return errors.Errorf("invalid put store %v", store)
	}

	storeValue, err := proto.Marshal(store)
	if err != nil {
		return errors.Trace(err)
	}

	storePath := makeStoreKey(c.clusterRoot, store.GetId())
	ctx, cancel := context.WithTimeout(context.Background(), requestTimeout)
	resp, err := c.s.slowLogTxn(ctx).
		If(c.s.leaderCmp()).
		Then(clientv3.OpPut(storePath, string(storeValue))).
		Commit()
	cancel()
	if err != nil {
		return errors.Trace(err)
	}
	if !resp.Succeeded {
		return errors.Errorf("put store %v fail", store)
	}

	c.cachedCluster.addStore(store)

	return nil
}

// GetConfig gets config from cluster.
func (c *RaftCluster) GetConfig() *metapb.Cluster {
	return c.cachedCluster.getMeta()
}

func (c *RaftCluster) putConfig(meta *metapb.Cluster) error {
	if meta.GetId() != c.clusterID {
		return errors.Errorf("invalid cluster %v, mismatch cluster id %d", meta, c.clusterID)
	}

	metaValue, err := proto.Marshal(meta)
	if err != nil {
		return errors.Trace(err)
	}

	ctx, cancel := context.WithTimeout(context.Background(), requestTimeout)
	resp, err := c.s.slowLogTxn(ctx).
		If(c.s.leaderCmp()).
		Then(clientv3.OpPut(c.clusterRoot, string(metaValue))).
		Commit()
	cancel()
	if err != nil {
		return errors.Trace(err)
	}
	if !resp.Succeeded {
		return errors.Errorf("put cluster meta %v error", meta)
	}

	c.cachedCluster.setMeta(meta)

	return nil
}

// GetBalanceOperators gets the balance operators from cluster.
func (c *RaftCluster) GetBalanceOperators() map[uint64]Operator {
	return c.balancerWorker.getBalanceOperators()
}

// GetHistoryOperators gets the history operators from cluster.
func (c *RaftCluster) GetHistoryOperators() []Operator {
	return c.balancerWorker.getHistoryOperators()
}

// GetScore gets store score from balancer.
func (c *RaftCluster) GetScore(store *metapb.Store, status *StoreStatus) int {
	storeInfo := &storeInfo{
		store: store,
		stats: status,
	}

	regionCount := c.cachedCluster.regions.regionCount()
	return c.balancerWorker.storeScore(storeInfo, regionCount)
}

// FetchEvents fetches the operator events.
func (c *RaftCluster) FetchEvents(key uint64, all bool) []LogEvent {
	return c.balancerWorker.fetchEvents(key, all)
}
<|endoftext|>


<commit_before>package server

import (
	"bytes"
	"fmt"

	. "github.com/siddontang/go-mysql/mysql"
	"github.com/siddontang/go/hack"
)

type Handler interface {
	//handle COM_INIT_DB command, you can check whether the dbName is valid, or other.
	UseDB(dbName string) error
	//handle COM_QUERY comamnd, like SELECT, INSERT, UPDATE, etc...
	//If Result has a Resultset (SELECT, SHOW, etc...), we will send this as the repsonse, otherwise, we will send Result
	HandleQuery(query string) (*Result, error)
	//handle COM_FILED_LIST command
	HandleFieldList(table string, fieldWildcard string) ([]*Field, error)
	//handle COM_STMT_PREPARE, params is the param number for this statement, columns is the column number
	//context will be used later for statement execute
	HandleStmtPrepare(query string) (params int, columns int, context interface{}, err error)
	//handle COM_STMT_EXECUTE, context is the previous one set in prepare
	//query is the statement prepare query, and args is the params for this statement
	HandleStmtExecute(context interface{}, query string, args []interface{}) (*Result, error)
	//handle COM_STMT_CLOSE, context is the previous one set in prepare
	//this handler has no response
	HandleStmtClose(context interface{}) (error)
}

func (c *Conn) HandleCommand() error {
	if c.Conn == nil {
		return fmt.Errorf("connection closed")
	}

	data, err := c.ReadPacket()
	if err != nil {
		c.Close()
		return err
	}

	v := c.dispatch(data)

	err = c.writeValue(v)

	if c.Conn != nil {
		c.ResetSequence()
	}

	if err != nil {
		c.Close()
	}
	return err
}

func (c *Conn) dispatch(data []byte) interface{} {
	cmd := data[0]
	data = data[1:]

	switch cmd {
	case COM_QUIT:
		c.Close()
		return noResponse{}
	case COM_QUERY:
		if r, err := c.h.HandleQuery(hack.String(data)); err != nil {
			return err
		} else {
			return r
		}
	case COM_PING:
		return nil
	case COM_INIT_DB:
		if err := c.h.UseDB(hack.String(data)); err != nil {
			return err
		} else {
			return nil
		}
	case COM_FIELD_LIST:
		index := bytes.IndexByte(data, 0x00)
		table := hack.String(data[0:index])
		wildcard := hack.String(data[index+1:])

		if fs, err := c.h.HandleFieldList(table, wildcard); err != nil {
			return err
		} else {
			return fs
		}
	case COM_STMT_PREPARE:
		c.stmtID++
		st := new(Stmt)
		st.ID = c.stmtID
		st.Query = hack.String(data)
		var err error
		if st.Params, st.Columns, st.Context, err = c.h.HandleStmtPrepare(st.Query); err != nil {
			return err
		} else {
			st.ResetParams()
			c.stmts[c.stmtID] = st
			return st
		}
	case COM_STMT_EXECUTE:
		if r, err := c.handleStmtExecute(data); err != nil {
			return err
		} else {
			return r
		}
	case COM_STMT_CLOSE:
		c.handleStmtClose(data)
		return noResponse{}
	case COM_STMT_SEND_LONG_DATA:
		c.handleStmtSendLongData(data)
		return noResponse{}
	case COM_STMT_RESET:
		if r, err := c.handleStmtReset(data); err != nil {
			return err
		} else {
			return r
		}
	default:
		msg := fmt.Sprintf("command %d is not supported now", cmd)
		return NewError(ER_UNKNOWN_ERROR, msg)
	}

	return fmt.Errorf("command %d is not handled correctly", cmd)
}

type EmptyHandler struct {
}

func (h EmptyHandler) UseDB(dbName string) error {
	return nil
}
func (h EmptyHandler) HandleQuery(query string) (*Result, error) {
	return nil, fmt.Errorf("not supported now")
}

func (h EmptyHandler) HandleFieldList(table string, fieldWildcard string) ([]*Field, error) {
	return nil, fmt.Errorf("not supported now")
}
func (h EmptyHandler) HandleStmtPrepare(query string) (int, int, interface{}, error) {
	return 0, 0, nil, fmt.Errorf("not supported now")
}
func (h EmptyHandler) HandleStmtExecute(context interface{}, query string, args []interface{}) (*Result, error) {
	return nil, fmt.Errorf("not supported now")
}
<commit_msg>修改EmptyHandler的接口未实现完整的bug (#76)<commit_after>package server

import (
	"bytes"
	"fmt"

	. "github.com/siddontang/go-mysql/mysql"
	"github.com/siddontang/go/hack"
)

type Handler interface {
	//handle COM_INIT_DB command, you can check whether the dbName is valid, or other.
	UseDB(dbName string) error
	//handle COM_QUERY comamnd, like SELECT, INSERT, UPDATE, etc...
	//If Result has a Resultset (SELECT, SHOW, etc...), we will send this as the repsonse, otherwise, we will send Result
	HandleQuery(query string) (*Result, error)
	//handle COM_FILED_LIST command
	HandleFieldList(table string, fieldWildcard string) ([]*Field, error)
	//handle COM_STMT_PREPARE, params is the param number for this statement, columns is the column number
	//context will be used later for statement execute
	HandleStmtPrepare(query string) (params int, columns int, context interface{}, err error)
	//handle COM_STMT_EXECUTE, context is the previous one set in prepare
	//query is the statement prepare query, and args is the params for this statement
	HandleStmtExecute(context interface{}, query string, args []interface{}) (*Result, error)
	//handle COM_STMT_CLOSE, context is the previous one set in prepare
	//this handler has no response
	HandleStmtClose(context interface{}) error
}

func (c *Conn) HandleCommand() error {
	if c.Conn == nil {
		return fmt.Errorf("connection closed")
	}

	data, err := c.ReadPacket()
	if err != nil {
		c.Close()
		return err
	}

	v := c.dispatch(data)

	err = c.writeValue(v)

	if c.Conn != nil {
		c.ResetSequence()
	}

	if err != nil {
		c.Close()
	}
	return err
}

func (c *Conn) dispatch(data []byte) interface{} {
	cmd := data[0]
	data = data[1:]

	switch cmd {
	case COM_QUIT:
		c.Close()
		return noResponse{}
	case COM_QUERY:
		if r, err := c.h.HandleQuery(hack.String(data)); err != nil {
			return err
		} else {
			return r
		}
	case COM_PING:
		return nil
	case COM_INIT_DB:
		if err := c.h.UseDB(hack.String(data)); err != nil {
			return err
		} else {
			return nil
		}
	case COM_FIELD_LIST:
		index := bytes.IndexByte(data, 0x00)
		table := hack.String(data[0:index])
		wildcard := hack.String(data[index+1:])

		if fs, err := c.h.HandleFieldList(table, wildcard); err != nil {
			return err
		} else {
			return fs
		}
	case COM_STMT_PREPARE:
		c.stmtID++
		st := new(Stmt)
		st.ID = c.stmtID
		st.Query = hack.String(data)
		var err error
		if st.Params, st.Columns, st.Context, err = c.h.HandleStmtPrepare(st.Query); err != nil {
			return err
		} else {
			st.ResetParams()
			c.stmts[c.stmtID] = st
			return st
		}
	case COM_STMT_EXECUTE:
		if r, err := c.handleStmtExecute(data); err != nil {
			return err
		} else {
			return r
		}
	case COM_STMT_CLOSE:
		c.handleStmtClose(data)
		return noResponse{}
	case COM_STMT_SEND_LONG_DATA:
		c.handleStmtSendLongData(data)
		return noResponse{}
	case COM_STMT_RESET:
		if r, err := c.handleStmtReset(data); err != nil {
			return err
		} else {
			return r
		}
	default:
		msg := fmt.Sprintf("command %d is not supported now", cmd)
		return NewError(ER_UNKNOWN_ERROR, msg)
	}

	return fmt.Errorf("command %d is not handled correctly", cmd)
}

type EmptyHandler struct {
}

func (h EmptyHandler) UseDB(dbName string) error {
	return nil
}
func (h EmptyHandler) HandleQuery(query string) (*Result, error) {
	return nil, fmt.Errorf("not supported now")
}

func (h EmptyHandler) HandleFieldList(table string, fieldWildcard string) ([]*Field, error) {
	return nil, fmt.Errorf("not supported now")
}
func (h EmptyHandler) HandleStmtPrepare(query string) (int, int, interface{}, error) {
	return 0, 0, nil, fmt.Errorf("not supported now")
}
func (h EmptyHandler) HandleStmtExecute(context interface{}, query string, args []interface{}) (*Result, error) {
	return nil, fmt.Errorf("not supported now")
}

func (h EmptyHandler) HandleStmtClose(context interface{}) error {
	return nil
}
<|endoftext|>


<commit_before>package server

import (
	"encoding/json"
	"errors"
	log "github.com/Sirupsen/logrus"
	"github.com/gorilla/mux"
	"github.com/gorilla/schema"
	"github.com/oxfeeefeee/appgo"
	"github.com/oxfeeefeee/appgo/auth"
	"github.com/unrolled/render"
	"net/http"
	"reflect"
)

const (
	UserIdFieldName      = "UserId__"
	AdminUserIdFieldName = "AdminUserId__"
	ResIdFieldName       = "ResourceId__"
	ContentFieldName     = "Content__"
)

const (
	_ HandlerType = iota
	HandlerTypeJson
	HandlerTypeHtml
)

var decoder = schema.NewDecoder()

type HandlerType int

type httpFunc struct {
	requireAuth    bool
	requireAdmin   bool
	hasResId       bool
	hasContent     bool
	dummyInput     bool
	allowAnonymous bool
	inputType      reflect.Type
	contentType    reflect.Type
	funcValue      reflect.Value
}

type handler struct {
	htype    HandlerType
	path     string
	template string
	funcs    map[string]*httpFunc
	supports []string
	ts       TokenStore
	renderer *render.Render
}

func init() {
	decoder.IgnoreUnknownKeys(true)
}

func (h *handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	f := h.funcs[r.Method]
	var input reflect.Value
	if f.dummyInput {
		input = reflect.ValueOf((*appgo.DummyInput)(nil))
	} else {
		input = reflect.New(f.inputType)
		if err := decoder.Decode(input.Interface(), r.URL.Query()); err != nil {
			h.renderError(w, appgo.NewApiErr(appgo.ECodeBadRequest, err.Error()))
			return
		}
	}
	if f.requireAuth {
		user, _ := h.authByHeader(r)
		s := input.Elem()
		field := s.FieldByName(UserIdFieldName)
		if user == 0 {
			if f.allowAnonymous {
				field.SetInt(appgo.AnonymousId)
			} else {
				h.renderError(w, appgo.NewApiErr(
					appgo.ECodeUnauthorized,
					"either remove UserId__ in your input define, or add allowAnonymous tag",
				))
				return
			}
		} else {
			field.SetInt(int64(user))
		}
	} else if f.requireAdmin {
		user, role := h.authByHeader(r)
		s := input.Elem()
		f := s.FieldByName(AdminUserIdFieldName)
		if user == 0 || role != appgo.RoleWebAdmin {
			h.renderError(w, appgo.NewApiErr(
				appgo.ECodeUnauthorized,
				"admin role required, you could remove AdminUserId__ in your input define"))
			return
		}
		f.SetInt(int64(user))
	}
	if f.hasResId {
		vars := mux.Vars(r)
		id := appgo.IdFromStr(vars["id"])
		if id == 0 {
			h.renderError(w, appgo.NewApiErr(
				appgo.ECodeNotFound,
				"ResourceId ('{id}' in url) required, you could remove ResourceId__ in your input define"))
			return
		}
		s := input.Elem()
		f := s.FieldByName(ResIdFieldName)
		f.SetInt(int64(id))
	}
	if f.hasContent {
		content := reflect.New(f.contentType.Elem())
		if err := json.NewDecoder(r.Body).Decode(content.Interface()); err != nil {
			h.renderError(w, appgo.NewApiErr(appgo.ECodeBadRequest, err.Error()))
			return
		}
		s := input.Elem()
		f := s.FieldByName(ContentFieldName)
		f.Set(content)
	}
	argsIn := []reflect.Value{input}
	returns := f.funcValue.Call(argsIn)
	if len(returns) == 0 || len(returns) > 2 {
		h.renderError(w, appgo.NewApiErr(appgo.ECodeInternal, "Bad api-func format"))
		return
	}
	// Either returns (reply, error) or returns (error)
	var retErr reflect.Value
	if len(returns) == 1 {
		retErr = returns[0]
	} else {
		retErr = returns[1]
	}
	// First check is err is nil
	if retErr.IsNil() {
		if len(returns) == 2 {
			h.renderData(w, returns[0].Interface())
		} else { // Empty return
			h.renderData(w, map[string]string{})
		}
	} else {
		if aerr, ok := retErr.Interface().(*appgo.ApiError); !ok {
			aerr = appgo.NewApiErr(appgo.ECodeInternal, "Bad api-func format")
		} else {
			h.renderError(w, aerr)
		}
	}
}

func (h *handler) authByHeader(r *http.Request) (appgo.Id, appgo.Role) {
	token := auth.Token(r.Header.Get(appgo.CustomTokenHeaderName))
	user, role := token.Validate()
	if user == 0 {
		return 0, 0
	}
	if !h.ts.Validate(token) {
		return 0, 0
	}
	return user, role
}

func newHandler(funcSet interface{}, htype HandlerType,
	ts TokenStore, renderer *render.Render) *handler {
	funcs := make(map[string]*httpFunc)
	// Let if panic if funSet's type is not right
	path := ""
	template := ""
	t := reflect.TypeOf(funcSet).Elem()
	if field, ok := t.FieldByName("META"); !ok {
		log.Panicln("Bad META setting (path, template)")
	} else {
		if p := field.Tag.Get("path"); p == "" {
			log.Panicln("Empty API path")
		} else {
			path = p
		}
		if htype == HandlerTypeHtml {
			if t := field.Tag.Get("template"); t == "" {
				log.Panicln("Empty HTML template")
			} else {
				template = t
			}
		}
	}
	structVal := reflect.Indirect(reflect.ValueOf(funcSet))
	supports := make([]string, 0, 4)
	if htype == HandlerTypeJson {
		methods := []string{"GET", "POST", "PUT", "DELETE"}
		for _, m := range methods {
			if fun, err := newHttpFunc(structVal, m); err != nil {
				log.Panicln(err)
			} else if fun != nil {
				funcs[m] = fun
				supports = append(supports, m)
			}
		}
		if len(supports) == 0 {
			log.Panicln("API supports no HTTP method")
		}
	} else if htype == HandlerTypeHtml {
		if fun, err := newHttpFunc(structVal, "HTML"); err != nil {
			log.Panicln(err)
		} else if fun == nil {
			log.Panicln("No HTML function for html")
		} else {
			funcs["GET"] = fun
		}
	} else {
		log.Panicln("Bad handler type")
	}
	return &handler{htype, path, template, funcs, supports, ts, renderer}
}

func newHttpFunc(structVal reflect.Value, fieldName string) (*httpFunc, error) {
	fieldVal := structVal.MethodByName(fieldName)
	if !fieldVal.IsValid() {
		return nil, nil
	}
	ftype := fieldVal.Type()
	inNum := ftype.NumIn()
	if inNum != 1 {
		return nil, errors.New("API func needs to have exact 1 parameter")
	}
	inputType := ftype.In(0)
	dummyInput := false
	if inputType.Kind() != reflect.Ptr {
		return nil, errors.New("API func's parameter needs to be a pointer")
	}
	if inputType == reflect.TypeOf((*appgo.DummyInput)(nil)) {
		dummyInput = true
	}
	inputType = inputType.Elem()
	requireAuth := false
	allowAnonymous := false
	if fromIdField, ok := inputType.FieldByName(UserIdFieldName); ok {
		requireAuth = true
		if fromIdField.Type.Kind() != reflect.Int64 {
			return nil, errors.New("API func's 2nd parameter needs to be Int64")
		}
		aa := fromIdField.Tag.Get("allowAnonymous")
		allowAnonymous = (aa == "true")
	}
	requireAdmin := false
	if fromIdType, ok := inputType.FieldByName(AdminUserIdFieldName); ok {
		requireAdmin = true
		if fromIdType.Type.Kind() != reflect.Int64 {
			return nil, errors.New("API func's 2nd parameter needs to be Int64")
		}
	}
	hasResId := false
	if resIdType, ok := inputType.FieldByName(ResIdFieldName); ok {
		hasResId = true
		if resIdType.Type.Kind() != reflect.Int64 {
			return nil, errors.New("ResId needs to be Int64")
		}
	}
	hasContent := false
	var contentType reflect.Type
	if ctype, ok := inputType.FieldByName(ContentFieldName); ok {
		hasContent = true
		contentType = ctype.Type
		if ctype.Type.Kind() != reflect.Ptr {
			return nil, errors.New("Content needs to be a pointer")
		}
	}
	return &httpFunc{requireAuth, requireAdmin, hasResId, hasContent,
		dummyInput, allowAnonymous, inputType, contentType, fieldVal}, nil
}
<commit_msg>add request__<commit_after>package server

import (
	"encoding/json"
	"errors"
	log "github.com/Sirupsen/logrus"
	"github.com/gorilla/mux"
	"github.com/gorilla/schema"
	"github.com/oxfeeefeee/appgo"
	"github.com/oxfeeefeee/appgo/auth"
	"github.com/unrolled/render"
	"net/http"
	"reflect"
)

const (
	UserIdFieldName      = "UserId__"
	AdminUserIdFieldName = "AdminUserId__"
	ResIdFieldName       = "ResourceId__"
	ContentFieldName     = "Content__"
	RequestFieldName     = "Request__"
)

const (
	_ HandlerType = iota
	HandlerTypeJson
	HandlerTypeHtml
)

var decoder = schema.NewDecoder()

type HandlerType int

type httpFunc struct {
	requireAuth    bool
	requireAdmin   bool
	hasResId       bool
	hasContent     bool
	hasRequest     bool
	dummyInput     bool
	allowAnonymous bool
	inputType      reflect.Type
	contentType    reflect.Type
	funcValue      reflect.Value
}

type handler struct {
	htype    HandlerType
	path     string
	template string
	funcs    map[string]*httpFunc
	supports []string
	ts       TokenStore
	renderer *render.Render
}

func init() {
	decoder.IgnoreUnknownKeys(true)
}

func (h *handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	f := h.funcs[r.Method]
	var input reflect.Value
	if f.dummyInput {
		input = reflect.ValueOf((*appgo.DummyInput)(nil))
	} else {
		input = reflect.New(f.inputType)
		if err := decoder.Decode(input.Interface(), r.URL.Query()); err != nil {
			h.renderError(w, appgo.NewApiErr(appgo.ECodeBadRequest, err.Error()))
			return
		}
	}
	if f.requireAuth {
		user, _ := h.authByHeader(r)
		s := input.Elem()
		field := s.FieldByName(UserIdFieldName)
		if user == 0 {
			if f.allowAnonymous {
				field.SetInt(appgo.AnonymousId)
			} else {
				h.renderError(w, appgo.NewApiErr(
					appgo.ECodeUnauthorized,
					"either remove UserId__ in your input define, or add allowAnonymous tag",
				))
				return
			}
		} else {
			field.SetInt(int64(user))
		}
	} else if f.requireAdmin {
		user, role := h.authByHeader(r)
		s := input.Elem()
		f := s.FieldByName(AdminUserIdFieldName)
		if user == 0 || role != appgo.RoleWebAdmin {
			h.renderError(w, appgo.NewApiErr(
				appgo.ECodeUnauthorized,
				"admin role required, you could remove AdminUserId__ in your input define"))
			return
		}
		f.SetInt(int64(user))
	}
	if f.hasResId {
		vars := mux.Vars(r)
		id := appgo.IdFromStr(vars["id"])
		if id == 0 {
			h.renderError(w, appgo.NewApiErr(
				appgo.ECodeNotFound,
				"ResourceId ('{id}' in url) required, you could remove ResourceId__ in your input define"))
			return
		}
		s := input.Elem()
		f := s.FieldByName(ResIdFieldName)
		f.SetInt(int64(id))
	}
	if f.hasContent {
		content := reflect.New(f.contentType.Elem())
		if err := json.NewDecoder(r.Body).Decode(content.Interface()); err != nil {
			h.renderError(w, appgo.NewApiErr(appgo.ECodeBadRequest, err.Error()))
			return
		}
		s := input.Elem()
		f := s.FieldByName(ContentFieldName)
		f.Set(content)
	}
	if f.hasRequest {
		s := input.Elem()
		f := s.FieldByName(RequestFieldName)
		f.Set(reflect.ValueOf(r))
	}
	argsIn := []reflect.Value{input}
	returns := f.funcValue.Call(argsIn)
	if len(returns) == 0 || len(returns) > 2 {
		h.renderError(w, appgo.NewApiErr(appgo.ECodeInternal, "Bad api-func format"))
		return
	}
	// Either returns (reply, error) or returns (error)
	var retErr reflect.Value
	if len(returns) == 1 {
		retErr = returns[0]
	} else {
		retErr = returns[1]
	}
	// First check is err is nil
	if retErr.IsNil() {
		if len(returns) == 2 {
			h.renderData(w, returns[0].Interface())
		} else { // Empty return
			h.renderData(w, map[string]string{})
		}
	} else {
		if aerr, ok := retErr.Interface().(*appgo.ApiError); !ok {
			aerr = appgo.NewApiErr(appgo.ECodeInternal, "Bad api-func format")
		} else {
			h.renderError(w, aerr)
		}
	}
}

func (h *handler) authByHeader(r *http.Request) (appgo.Id, appgo.Role) {
	token := auth.Token(r.Header.Get(appgo.CustomTokenHeaderName))
	user, role := token.Validate()
	if user == 0 {
		return 0, 0
	}
	if !h.ts.Validate(token) {
		return 0, 0
	}
	return user, role
}

func newHandler(funcSet interface{}, htype HandlerType,
	ts TokenStore, renderer *render.Render) *handler {
	funcs := make(map[string]*httpFunc)
	// Let if panic if funSet's type is not right
	path := ""
	template := ""
	t := reflect.TypeOf(funcSet).Elem()
	if field, ok := t.FieldByName("META"); !ok {
		log.Panicln("Bad META setting (path, template)")
	} else {
		if p := field.Tag.Get("path"); p == "" {
			log.Panicln("Empty API path")
		} else {
			path = p
		}
		if htype == HandlerTypeHtml {
			if t := field.Tag.Get("template"); t == "" {
				log.Panicln("Empty HTML template")
			} else {
				template = t
			}
		}
	}
	structVal := reflect.Indirect(reflect.ValueOf(funcSet))
	supports := make([]string, 0, 4)
	if htype == HandlerTypeJson {
		methods := []string{"GET", "POST", "PUT", "DELETE"}
		for _, m := range methods {
			if fun, err := newHttpFunc(structVal, m); err != nil {
				log.Panicln(err)
			} else if fun != nil {
				funcs[m] = fun
				supports = append(supports, m)
			}
		}
		if len(supports) == 0 {
			log.Panicln("API supports no HTTP method")
		}
	} else if htype == HandlerTypeHtml {
		if fun, err := newHttpFunc(structVal, "HTML"); err != nil {
			log.Panicln(err)
		} else if fun == nil {
			log.Panicln("No HTML function for html")
		} else {
			funcs["GET"] = fun
		}
	} else {
		log.Panicln("Bad handler type")
	}
	return &handler{htype, path, template, funcs, supports, ts, renderer}
}

func newHttpFunc(structVal reflect.Value, fieldName string) (*httpFunc, error) {
	fieldVal := structVal.MethodByName(fieldName)
	if !fieldVal.IsValid() {
		return nil, nil
	}
	ftype := fieldVal.Type()
	inNum := ftype.NumIn()
	if inNum != 1 {
		return nil, errors.New("API func needs to have exact 1 parameter")
	}
	inputType := ftype.In(0)
	dummyInput := false
	if inputType.Kind() != reflect.Ptr {
		return nil, errors.New("API func's parameter needs to be a pointer")
	}
	if inputType == reflect.TypeOf((*appgo.DummyInput)(nil)) {
		dummyInput = true
	}
	inputType = inputType.Elem()
	requireAuth := false
	allowAnonymous := false
	if fromIdField, ok := inputType.FieldByName(UserIdFieldName); ok {
		requireAuth = true
		if fromIdField.Type.Kind() != reflect.Int64 {
			return nil, errors.New("API func's 2nd parameter needs to be Int64")
		}
		aa := fromIdField.Tag.Get("allowAnonymous")
		allowAnonymous = (aa == "true")
	}
	requireAdmin := false
	if fromIdType, ok := inputType.FieldByName(AdminUserIdFieldName); ok {
		requireAdmin = true
		if fromIdType.Type.Kind() != reflect.Int64 {
			return nil, errors.New("API func's 2nd parameter needs to be Int64")
		}
	}
	hasResId := false
	if resIdType, ok := inputType.FieldByName(ResIdFieldName); ok {
		hasResId = true
		if resIdType.Type.Kind() != reflect.Int64 {
			return nil, errors.New("ResId needs to be Int64")
		}
	}
	hasContent := false
	var contentType reflect.Type
	if ctype, ok := inputType.FieldByName(ContentFieldName); ok {
		hasContent = true
		contentType = ctype.Type
		if ctype.Type.Kind() != reflect.Ptr {
			return nil, errors.New("Content needs to be a pointer")
		}
	}
	hasRequest := false
	if ctype, ok := inputType.FieldByName(RequestFieldName); ok {
		hasRequest = true
		if ctype.Type.Kind() != reflect.Ptr {
			return nil, errors.New("Request needs to be a pointer to http.Request")
		}
		if ctype.Type.Elem() != reflect.TypeOf((*http.Request)(nil)).Elem() {
			return nil, errors.New("Request needs to be a pointer to http.Request")
		}
	}
	return &httpFunc{requireAuth, requireAdmin, hasResId, hasContent, hasRequest,
		dummyInput, allowAnonymous, inputType, contentType, fieldVal}, nil
}
<|endoftext|>