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

package monitor

import (
	"atlantis/supervisor/containers/serialize"
	"atlantis/supervisor/rpc/types"
	"bytes"
	"fmt"
	"github.com/BurntSushi/toml"
	"github.com/jigish/go-flags"
	"os"
	"os/exec"
	"path"
	"path/filepath"
	"strings"
	"time"
)

const (
	OK = iota
	Warning
	Critical
	Uknown
)

type Config struct {
	ContainerFile   string `toml:"container_file"`
	ContainersDir   string `toml:"container_dir"`
	InventoryDir    string `toml:"inventory_dir"`
	SSHIdentity     string `toml:"ssh_identity"`
	SSHUser         string `toml:"ssh_user"`
	CheckName       string `toml:"check_name"`
	CheckDir        string `toml:"check_dir"`
	DefaultGroup    string `toml:"default_group"`
	TimeoutDuration uint   `toml:"timeout_duration"`
	Verbose         bool   `toml:"verbose"`
}

type Opts struct {
	ContainerFile   string `short:"f" long:"container-file" description:"file to get contianers information from"`
	ContainersDir   string `short:"s" long:"containers-dir" description:"directory where configs for each container live"`
	SSHIdentity     string `short:"i" long:"ssh-identity" description:"file containing the SSH key for all containers"`
	SSHUser         string `short:"u" long:"ssh-user" description:"user account to ssh into containers"`
	CheckName       string `short:"n" long:"check-name" description:"service name that will appear in Nagios for the monitor"`
	CheckDir        string `short:"d" long:"check-dir" description:"directory containing all the scripts for the monitoring checks"`
	DefaultGroup    string `short:"g" long:"default-group" description:"default contact group to use if there is no valid group provided"`
	Config          string `short:"c" long:"config-file" default:"/etc/atlantis/supervisor/monitor.toml" description:"the config file to use"`
	TimeoutDuration uint   `short:"t" long:"timeout-duration" description:"max number of seconds to wait for a monitoring check to finish"`
	Verbose         bool   `short:"v" long:"verbose" default:false description:"print verbose debug information"`
}

type ServiceCheck struct {
	Service  string
	User     string
	Identity string
	Host     string
	Port     uint16
	Script   string
}

//TODO(mchandra):Need defaults defined by constants
var config = &Config{
	ContainerFile:   "/etc/atlantis/supervisor/save/containers",
	ContainersDir:   "/etc/atlantis/containers",
	InventoryDir:    "/etc/atlantis/supervisor/inventory",
	SSHIdentity:     "/opt/atlantis/supervisor/master_id_rsa",
	SSHUser:         "root",
	CheckName:       "ContainerMonitor",
	CheckDir:        "/check_mk_checks",
	DefaultGroup:    "atlantis_orphan_apps",
	TimeoutDuration: 110,
	Verbose:         false,
}

func (s *ServiceCheck) cmd() *exec.Cmd {
	return silentSshCmd(s.User, s.Identity, s.Host, s.Script, s.Port)
}

func (s *ServiceCheck) timeOutMsg() string {
	return fmt.Sprintf("%d %s - Timeout occured during check\n", Critical, s.Service)
}

func (s *ServiceCheck) errMsg(err error) string {
	if err != nil {
		return fmt.Sprintf("%d %s - %s\n", Critical, s.Service, err.Error())
	} else {
		return fmt.Sprintf("%d %s - Error encountered while monitoring the service\n", Critical, s.Service)
	}
}

func (s *ServiceCheck) validate(msg string) string {
	m := strings.SplitN(msg, " ", 4)
	if len(m) > 1 && m[1] == s.Service {
		return msg
	}
	return s.errMsg(nil)
}

func (s *ServiceCheck) runCheck(done chan bool) {
	out, err := s.cmd().Output()
	if err != nil {
		fmt.Print(s.errMsg(err))
	} else {
		fmt.Print(s.validate(string(out)))
	}
	done <- true
}

func (s *ServiceCheck) checkWithTimeout(results chan bool, d time.Duration) {
	done := make(chan bool, 1)
	go s.runCheck(done)
	select {
	case <-done:
		results <- true
	case <-time.After(d):
		fmt.Print(s.timeOutMsg())
		results <- true
	}
}

type ContainerCheck struct {
	Name         string
	User         string
	Identity     string
	Directory    string
	Inventory    string
	ContactGroup string
	container    *types.Container
}

type ContainerConfig struct {
	Dependencies map[string]interface{}
}

func (c *ContainerCheck) getContactGroup() {
	c.ContactGroup = config.DefaultGroup
	config_file := filepath.Join(config.ContainersDir, c.container.ID, "config.json")
	var cont_config ContainerConfig
	if err := serialize.RetrieveObject(config_file, &cont_config); err != nil {
		fmt.Printf("%d %s - Could not retrieve container config %s: %s\n", Critical, config.CheckName, config_file, err)
	} else {
		dep, ok := cont_config.Dependencies["cmk"]
		if !ok {
			fmt.Printf("%d %s - cmk dep not present, defaulting to %s contact group!\n", OK, config.CheckName, config.DefaultGroup)
			return
		}
		cmk_dep, ok := dep.(map[string]interface{})
		if !ok {
			fmt.Printf("%d %s - cmk dep present, but value is not map[string]string!\n", Critical, config.CheckName)
			return
		}
		val, ok := cmk_dep["contact_group"]
		if !ok {
			fmt.Printf("%d %s - cmk dep present, but no contact_group key!\n", Critical, config.CheckName)
			return
		}
		group, ok := val.(string)
		if ok {
			c.ContactGroup = group
			c.verifyContactGroup()
		} else {
			fmt.Printf("%d %s - Value for contact_group key of cmk dep is not a string!\n", Critical, config.CheckName)
		}
	}
}

func (c *ContainerCheck) setContactGroup(name string) {
	if len(c.ContactGroup) == 0 {
		c.getContactGroup()
	}
	inventoryPath := path.Join(c.Inventory, name)
	if _, err := os.Stat(inventoryPath); os.IsNotExist(err) {
		output, err := exec.Command("/usr/bin/cmk_admin", "-s", name, "-a", c.ContactGroup).CombinedOutput()
		if err != nil {
			fmt.Printf("%d %s - Failure to update contact group for service %s. Error: %s\n", OK, config.CheckName, name, err.Error())
		} else {
			os.Create(inventoryPath)
		}
		if config.Verbose {
			fmt.Printf("\n/usr/bin/cmk_admin -s %s -a %s\n%s\n\n", name, c.ContactGroup, output)
		}
	}
}

func (c *ContainerCheck) verifyContactGroup() {
	var output bytes.Buffer
	var stderr bytes.Buffer
	cmk := exec.Command("/usr/bin/cmk_admin", "-l")
	egrep := exec.Command("egrep", "-x", fmt.Sprintf("[*]?\\s+(%s)", c.ContactGroup))
	var err error
	if egrep.Stdin, err = cmk.StdoutPipe(); err != nil {
		fmt.Printf("%d %s - Error piping output while verifying %s. Error: %s\n", Critical, config.CheckName, c.ContactGroup, err.Error())
		c.ContactGroup = config.DefaultGroup
		return
	}
	cmk.Stderr = &stderr
	egrep.Stdout, egrep.Stderr = &output, &stderr
	cmds := [2]*exec.Cmd{cmk, egrep}
	for _, cmd := range cmds {
		if err := cmd.Start(); err != nil {
			fmt.Printf("%d %s - Error starting %s while verifying %s. Error: %s\n", Critical, config.CheckName, cmd.Path, c.ContactGroup, err.Error())
			c.ContactGroup = config.DefaultGroup
			return
		}
	}
	for _, cmd := range cmds {
		if err := cmd.Wait(); err != nil {
			fmt.Printf("%d %s - Error running %s while verfying %s. Error: %s\n", Critical, config.CheckName, cmd.Path, c.ContactGroup, err.Error())
			c.ContactGroup = config.DefaultGroup
			return
		}
	}
}

func (c *ContainerCheck) Run(t time.Duration, done chan bool) {
	c.setContactGroup(c.Name)
	defer func() { done <- true }()
	o, err := silentSshCmd(c.User, c.Identity, c.container.Host, "ls "+c.Directory, c.container.SSHPort).Output()
	if err != nil {
		fmt.Printf("%d %s - Error getting checks for container: %s\n", Critical, c.Name, err.Error())
		return
	}
	fmt.Printf("%d %s - Got checks for container\n", OK, c.Name)
	scripts := strings.Split(strings.TrimSpace(string(o)), "\n")
	if len(scripts) == 0 || len(scripts[0]) == 0 {
		// nothing to check on this container, exit
		return
	}
	c.checkAll(scripts, t)
}

func (c *ContainerCheck) checkAll(scripts []string, t time.Duration) {
	results := make(chan bool, len(scripts))
	for _, s := range scripts {
		serviceName := fmt.Sprintf("%s_%s", strings.Split(s, ".")[0], c.container.ID)
		c.setContactGroup(serviceName)
		go c.serviceCheck(s).checkWithTimeout(results, t)
	}
	for _ = range scripts {
		<-results
	}
}

func (c *ContainerCheck) serviceCheck(script string) *ServiceCheck {
	// The full path to the script is required
	command := fmt.Sprintf("%s/%s %d %s", c.Directory, script, c.container.PrimaryPort, c.container.ID)
	// The service name is obtained be removing the file extension from the script and appending the container
	// id
	serviceName := fmt.Sprintf("%s_%s", strings.Split(script, ".")[0], c.container.ID)
	return &ServiceCheck{serviceName, c.User, c.Identity, c.container.Host, c.container.SSHPort, command}
}

func silentSshCmd(user, identity, host, cmd string, port uint16) *exec.Cmd {
	args := []string{"-q", user + "@" + host, "-i", identity, "-p", fmt.Sprintf("%d", port), "-o", "StrictHostKeyChecking=no", cmd}
	return exec.Command("ssh", args...)
}

func overlayConfig() {
	opts := &Opts{}
	flags.Parse(opts)
	if opts.Config != "" {
		_, err := toml.DecodeFile(opts.Config, config)
		if err != nil {
			// no need to panic here. we have reasonable defaults.
		}
	}
	if opts.ContainerFile != "" {
		config.ContainerFile = opts.ContainerFile
	}
	if opts.ContainersDir != "" {
		config.ContainersDir = opts.ContainersDir
	}
	if opts.SSHIdentity != "" {
		config.SSHIdentity = opts.SSHIdentity
	}
	if opts.SSHUser != "" {
		config.SSHUser = opts.SSHUser
	}
	if opts.CheckDir != "" {
		config.CheckDir = opts.CheckDir
	}
	if opts.CheckName != "" {
		config.CheckName = opts.CheckName
	}
	if opts.DefaultGroup != "" {
		config.DefaultGroup = opts.DefaultGroup
	}
	if opts.TimeoutDuration != 0 {
		config.TimeoutDuration = opts.TimeoutDuration
	}
	if opts.Verbose {
		config.Verbose = true
	}
}

//file containing containers and service name to show in Nagios for the monitor itself
func Run() {
	overlayConfig()
	var contMap map[string]*types.Container
	//Check if folder exists
	_, err := os.Stat(config.ContainerFile)
	if os.IsNotExist(err) {
		fmt.Printf("%d %s - Container file does not exists %s. Likely no live containers present.\n", OK, config.CheckName, config.ContainerFile)
		return
	}
	if err := serialize.RetrieveObject(config.ContainerFile, &contMap); err != nil {
		fmt.Printf("%d %s - Error retrieving %s: %s\n", Critical, config.CheckName, config.ContainerFile, err)
		return
	}
	done := make(chan bool, len(contMap))
	config.SSHIdentity = strings.Replace(config.SSHIdentity, "~", os.Getenv("HOME"), 1)
	for _, c := range contMap {
		if c.Host == "" {
			c.Host = "localhost"
		}
		check := &ContainerCheck{config.CheckName + "_" + c.ID, config.SSHUser, config.SSHIdentity, config.CheckDir, config.InventoryDir, "", c}
		go check.Run(time.Duration(config.TimeoutDuration)*time.Second, done)
	}
	for _ = range contMap {
		<-done
	}
	// Clean up inventories from containers that no longer exist
	err = filepath.Walk(config.InventoryDir, func(path string, _ os.FileInfo, _ error) error {
		if path == config.InventoryDir {
			return nil
		}
		var err error
		split := strings.Split(path, "_")
		cont := split[len(split)-1]
		if _, ok := contMap[cont]; !ok {
			err = os.Remove(path)
		}
		return err
	})
	if err != nil {
		fmt.Printf("%d %s - Error iterating over inventory to delete obsolete markers. Error: %s\n", OK, config.CheckName, err.Error())
	}
}
<commit_msg>simplify verify as per CR<commit_after>/* Copyright 2014 Ooyala, Inc. All rights reserved.
 *
 * This file is licensed under the Apache License, Version 2.0 (the "License"); you may not use this file
 * except in compliance with the License. You may obtain a copy of the License at
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software distributed under the License is
 * distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and limitations under the License.
 */

package monitor

import (
	"atlantis/supervisor/containers/serialize"
	"atlantis/supervisor/rpc/types"
	"bytes"
	"fmt"
	"github.com/BurntSushi/toml"
	"github.com/jigish/go-flags"
	"os"
	"os/exec"
	"path"
	"path/filepath"
	"strings"
	"time"
)

const (
	OK = iota
	Warning
	Critical
	Uknown
)

type Config struct {
	ContainerFile   string `toml:"container_file"`
	ContainersDir   string `toml:"container_dir"`
	InventoryDir    string `toml:"inventory_dir"`
	SSHIdentity     string `toml:"ssh_identity"`
	SSHUser         string `toml:"ssh_user"`
	CheckName       string `toml:"check_name"`
	CheckDir        string `toml:"check_dir"`
	DefaultGroup    string `toml:"default_group"`
	TimeoutDuration uint   `toml:"timeout_duration"`
	Verbose         bool   `toml:"verbose"`
}

type Opts struct {
	ContainerFile   string `short:"f" long:"container-file" description:"file to get container information"`
	ContainersDir   string `short:"s" long:"containers-dir" description:"directory containing configs for each container"`
	SSHIdentity     string `short:"i" long:"ssh-identity" description:"file containing the SSH key for all containers"`
	SSHUser         string `short:"u" long:"ssh-user" description:"user account to ssh into containers"`
	CheckName       string `short:"n" long:"check-name" description:"service name that will appear in Nagios for the monitor"`
	CheckDir        string `short:"d" long:"check-dir" description:"directory containing all the scripts for the monitoring checks"`
	DefaultGroup    string `short:"g" long:"default-group" description:"default contact group to use if there is no valid group provided"`
	Config          string `short:"c" long:"config-file" default:"/etc/atlantis/supervisor/monitor.toml" description:"the config file to use"`
	TimeoutDuration uint   `short:"t" long:"timeout-duration" description:"max number of seconds to wait for a monitoring check to finish"`
	Verbose         bool   `short:"v" long:"verbose" default:false description:"print verbose debug information"`
}

type ServiceCheck struct {
	Service  string
	User     string
	Identity string
	Host     string
	Port     uint16
	Script   string
}

//TODO(mchandra):Need defaults defined by constants
var config = &Config{
	ContainerFile:   "/etc/atlantis/supervisor/save/containers",
	ContainersDir:   "/etc/atlantis/containers",
	InventoryDir:    "/etc/atlantis/supervisor/inventory",
	SSHIdentity:     "/opt/atlantis/supervisor/master_id_rsa",
	SSHUser:         "root",
	CheckName:       "ContainerMonitor",
	CheckDir:        "/check_mk_checks",
	DefaultGroup:    "atlantis_orphan_apps",
	TimeoutDuration: 110,
	Verbose:         false,
}

func (s *ServiceCheck) cmd() *exec.Cmd {
	return silentSshCmd(s.User, s.Identity, s.Host, s.Script, s.Port)
}

func (s *ServiceCheck) timeOutMsg() string {
	return fmt.Sprintf("%d %s - Timeout occured during check\n", Critical, s.Service)
}

func (s *ServiceCheck) errMsg(err error) string {
	if err != nil {
		return fmt.Sprintf("%d %s - %s\n", Critical, s.Service, err.Error())
	} else {
		return fmt.Sprintf("%d %s - Error encountered while monitoring the service\n", Critical, s.Service)
	}
}

func (s *ServiceCheck) validate(msg string) string {
	m := strings.SplitN(msg, " ", 4)
	if len(m) > 1 && m[1] == s.Service {
		return msg
	}
	return s.errMsg(nil)
}

func (s *ServiceCheck) runCheck(done chan bool) {
	out, err := s.cmd().Output()
	if err != nil {
		fmt.Print(s.errMsg(err))
	} else {
		fmt.Print(s.validate(string(out)))
	}
	done <- true
}

func (s *ServiceCheck) checkWithTimeout(results chan bool, d time.Duration) {
	done := make(chan bool, 1)
	go s.runCheck(done)
	select {
	case <-done:
		results <- true
	case <-time.After(d):
		fmt.Print(s.timeOutMsg())
		results <- true
	}
}

type ContainerCheck struct {
	Name         string
	User         string
	Identity     string
	Directory    string
	Inventory    string
	ContactGroup string
	container    *types.Container
}

type ContainerConfig struct {
	Dependencies map[string]interface{}
}

func (c *ContainerCheck) verifyContactGroup(group string) bool {
	output, err := exec.Command("/usr/bin/cmk_admin", "-l").Output()
	for _, l := range strings.Split(output, "\n") {
		cg := strings.TrimSpace(strings.TrimPrefix(l, "*"))
		if cg == c.ContactGroup {
			return true
		}
	}
	return false
}

func (c *ContainerCheck) getContactGroup() {
	c.ContactGroup = config.DefaultGroup
	config_file := filepath.Join(config.ContainersDir, c.container.ID, "config.json")
	var cont_config ContainerConfig
	if err := serialize.RetrieveObject(config_file, &cont_config); err != nil {
		fmt.Printf("%d %s - Could not retrieve container config %s: %s\n", Critical, config.CheckName, config_file, err)
	} else {
		dep, ok := cont_config.Dependencies["cmk"]
		if !ok {
			fmt.Printf("%d %s - cmk dep not present, defaulting to %s contact group!\n", OK, config.CheckName, config.DefaultGroup)
			return
		}
		cmk_dep, ok := dep.(map[string]interface{})
		if !ok {
			fmt.Printf("%d %s - cmk dep present, but value is not map[string]string!\n", Critical, config.CheckName)
			return
		}
		val, ok := cmk_dep["contact_group"]
		if !ok {
			fmt.Printf("%d %s - cmk dep present, but no contact_group key!\n", Critical, config.CheckName)
			return
		}
		group, ok := val.(string)
		if ok {
			group = strings.ToLower(group)
			if c.verifyContactGroup(group) {
				c.ContactGroup = group
			} else {
				fmt.Printf("%d %s - Specified contact_group does not exist in cmk! Falling back to default group %s.\n", Critical, config.CheckName, config.DefaultGroup)
			}
		} else {
			fmt.Printf("%d %s - Value for contact_group key of cmk dep is not a string!\n", Critical, config.CheckName)
		}
	}
}

func (c *ContainerCheck) setContactGroup(name string) {
	if len(c.ContactGroup) == 0 {
		c.getContactGroup()
	}
	inventoryPath := path.Join(c.Inventory, name)
	if _, err := os.Stat(inventoryPath); os.IsNotExist(err) {
		output, err := exec.Command("/usr/bin/cmk_admin", "-s", name, "-a", c.ContactGroup).CombinedOutput()
		if err != nil {
			fmt.Printf("%d %s - Failure to update contact group for service %s. Error: %s\n", OK, config.CheckName, name, err.Error())
		} else {
			os.Create(inventoryPath)
		}
		if config.Verbose {
			fmt.Printf("\n/usr/bin/cmk_admin -s %s -a %s\n%s\n\n", name, c.ContactGroup, output)
		}
	}
}

func (c *ContainerCheck) Run(t time.Duration, done chan bool) {
	c.setContactGroup(c.Name)
	defer func() { done <- true }()
	o, err := silentSshCmd(c.User, c.Identity, c.container.Host, "ls "+c.Directory, c.container.SSHPort).Output()
	if err != nil {
		fmt.Printf("%d %s - Error getting checks for container: %s\n", Critical, c.Name, err.Error())
		return
	}
	fmt.Printf("%d %s - Got checks for container\n", OK, c.Name)
	scripts := strings.Split(strings.TrimSpace(string(o)), "\n")
	if len(scripts) == 0 || len(scripts[0]) == 0 {
		// nothing to check on this container, exit
		return
	}
	c.checkAll(scripts, t)
}

func (c *ContainerCheck) checkAll(scripts []string, t time.Duration) {
	results := make(chan bool, len(scripts))
	for _, s := range scripts {
		serviceName := fmt.Sprintf("%s_%s", strings.Split(s, ".")[0], c.container.ID)
		c.setContactGroup(serviceName)
		go c.serviceCheck(s).checkWithTimeout(results, t)
	}
	for _ = range scripts {
		<-results
	}
}

func (c *ContainerCheck) serviceCheck(script string) *ServiceCheck {
	// The full path to the script is required
	command := fmt.Sprintf("%s/%s %d %s", c.Directory, script, c.container.PrimaryPort, c.container.ID)
	// The service name is obtained be removing the file extension from the script and appending the container
	// id
	serviceName := fmt.Sprintf("%s_%s", strings.Split(script, ".")[0], c.container.ID)
	return &ServiceCheck{serviceName, c.User, c.Identity, c.container.Host, c.container.SSHPort, command}
}

func silentSshCmd(user, identity, host, cmd string, port uint16) *exec.Cmd {
	args := []string{"-q", user + "@" + host, "-i", identity, "-p", fmt.Sprintf("%d", port), "-o", "StrictHostKeyChecking=no", cmd}
	return exec.Command("ssh", args...)
}

func overlayConfig() {
	opts := &Opts{}
	flags.Parse(opts)
	if opts.Config != "" {
		_, err := toml.DecodeFile(opts.Config, config)
		if err != nil {
			// no need to panic here. we have reasonable defaults.
		}
	}
	if opts.ContainerFile != "" {
		config.ContainerFile = opts.ContainerFile
	}
	if opts.ContainersDir != "" {
		config.ContainersDir = opts.ContainersDir
	}
	if opts.SSHIdentity != "" {
		config.SSHIdentity = opts.SSHIdentity
	}
	if opts.SSHUser != "" {
		config.SSHUser = opts.SSHUser
	}
	if opts.CheckDir != "" {
		config.CheckDir = opts.CheckDir
	}
	if opts.CheckName != "" {
		config.CheckName = opts.CheckName
	}
	if opts.DefaultGroup != "" {
		config.DefaultGroup = opts.DefaultGroup
	}
	if opts.TimeoutDuration != 0 {
		config.TimeoutDuration = opts.TimeoutDuration
	}
	if opts.Verbose {
		config.Verbose = true
	}
}

//file containing containers and service name to show in Nagios for the monitor itself
func Run() {
	overlayConfig()
	var contMap map[string]*types.Container
	//Check if folder exists
	_, err := os.Stat(config.ContainerFile)
	if os.IsNotExist(err) {
		fmt.Printf("%d %s - Container file does not exists %s. Likely no live containers present.\n", OK, config.CheckName, config.ContainerFile)
		return
	}
	if err := serialize.RetrieveObject(config.ContainerFile, &contMap); err != nil {
		fmt.Printf("%d %s - Error retrieving %s: %s\n", Critical, config.CheckName, config.ContainerFile, err)
		return
	}
	done := make(chan bool, len(contMap))
	config.SSHIdentity = strings.Replace(config.SSHIdentity, "~", os.Getenv("HOME"), 1)
	for _, c := range contMap {
		if c.Host == "" {
			c.Host = "localhost"
		}
		check := &ContainerCheck{config.CheckName + "_" + c.ID, config.SSHUser, config.SSHIdentity, config.CheckDir, config.InventoryDir, "", c}
		go check.Run(time.Duration(config.TimeoutDuration)*time.Second, done)
	}
	for _ = range contMap {
		<-done
	}
	// Clean up inventories from containers that no longer exist
	err = filepath.Walk(config.InventoryDir, func(path string, _ os.FileInfo, _ error) error {
		if path == config.InventoryDir {
			return nil
		}
		var err error
		split := strings.Split(path, "_")
		cont := split[len(split)-1]
		if _, ok := contMap[cont]; !ok {
			err = os.Remove(path)
		}
		return err
	})
	if err != nil {
		fmt.Printf("%d %s - Error iterating over inventory to delete obsolete markers. Error: %s\n", OK, config.CheckName, err.Error())
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"net/http"
	"os"
	"path/filepath"
	"strconv"
	"strings"
)

const (
	fileTypeGallery = 0 //only image extension
	fileTypeGlobal  = 1 //only image extension
	fileTypeVideo   = 2 //only video extension
)

//serveFile serves image or video files
func serveFile(w http.ResponseWriter, r *http.Request) {

	path := filepath.Join(filepath.Dir(os.Args[0]), galleryPath, r.FormValue("gallery"))
	file := strings.Trim(r.FormValue("name"), `./\`)
	filetype, _ := strconv.Atoi(r.FormValue("filetype"))
	thumb, _ := strconv.ParseBool(r.FormValue("thumb"))

	switch filetype {
	case fileTypeGlobal:
		//do nothing
	case fileTypeGallery:
		path = filepath.Join(path, galleryImgPath)
	case fileTypeVideo:
		path = filepath.Join(path, galleryVidPath)
	}
	if thumb {
		path = filepath.Join(path, thumbPath)
	}
	path = filepath.Join(path, file)
	//extra safety to serve only allowed filetypes
	ext := filepath.Ext(path)
	if (thumb || filetype == fileTypeGallery || filetype == fileTypeGlobal) && !isPictureExt(ext) {
		fmt.Fprintln(w, "Not an image file")
		return
	}
	if (!thumb && filetype == fileTypeVideo) && !isVideoExt(ext) {
		fmt.Fprintln(w, "Not a video file")
		return
	}

	//Or ioutil.ReadFile...
	fh, err := os.Open(path)
	if err != nil {
		fmt.Fprintf(w, "Cannot open file %s", r.FormValue("file"))
		return
	}
	defer fh.Close()

	//serve file
	stat, err := fh.Stat()
	if err != nil {
		fmt.Fprintf(w, "Cannot open file %s", r.FormValue("file"))
		return
	}
	http.ServeContent(w, r, r.FormValue("file"), stat.ModTime(), fh)
	return
}

func isPictureExt(ext string) bool {
	switch strings.ToLower(ext) {
	case ".jpg", ".jpeg", ".gif", ".png":
		return true
	}
	return false
}

func isVideoExt(ext string) bool {
	switch strings.ToLower(ext) {
	case ".mp4", ".webm", ".avi", ".mkv", ".mpeg", ".mpg":
		return true
	}
	return false
}
<commit_msg>use filename var<commit_after>package main

import (
	"fmt"
	"net/http"
	"os"
	"path/filepath"
	"strconv"
	"strings"
)

const (
	fileTypeGallery = 0 //only image extension
	fileTypeGlobal  = 1 //only image extension
	fileTypeVideo   = 2 //only video extension
)

//serveFile serves image or video files
func serveFile(w http.ResponseWriter, r *http.Request) {

	path := filepath.Join(filepath.Dir(os.Args[0]), galleryPath, r.FormValue("gallery"))
	file := strings.Trim(r.FormValue("name"), `./\`)
	filetype, _ := strconv.Atoi(r.FormValue("filetype"))
	thumb, _ := strconv.ParseBool(r.FormValue("thumb"))

	switch filetype {
	case fileTypeGlobal:
		//do nothing
	case fileTypeGallery:
		path = filepath.Join(path, galleryImgPath)
	case fileTypeVideo:
		path = filepath.Join(path, galleryVidPath)
	}
	if thumb {
		path = filepath.Join(path, thumbPath)
	}
	path = filepath.Join(path, file)
	//extra safety to serve only allowed filetypes
	ext := filepath.Ext(path)
	if (thumb || filetype == fileTypeGallery || filetype == fileTypeGlobal) && !isPictureExt(ext) {
		fmt.Fprintln(w, "Not an image file")
		return
	}
	if (!thumb && filetype == fileTypeVideo) && !isVideoExt(ext) {
		fmt.Fprintln(w, "Not a video file")
		return
	}

	//Open file
	fh, err := os.Open(path)
	if err != nil {
		fmt.Fprintf(w, "Cannot open file %s", r.FormValue("file"))
		return
	}
	defer fh.Close()

	//Serve file
	stat, err := fh.Stat()
	if err != nil {
		fmt.Fprintf(w, "Cannot open file %s", r.FormValue("file"))
		return
	}
	http.ServeContent(w, r, file, stat.ModTime(), fh)
	return
}

func isPictureExt(ext string) bool {
	switch strings.ToLower(ext) {
	case ".jpg", ".jpeg", ".gif", ".png":
		return true
	}
	return false
}

func isVideoExt(ext string) bool {
	switch strings.ToLower(ext) {
	case ".mp4", ".webm", ".avi", ".mkv", ".mpeg", ".mpg":
		return true
	}
	return false
}
<|endoftext|>


<commit_before>package asset

import (
	"encoding/json"
	"time"

	"golang.org/x/net/context"

	"chain/api/appdb"
	chainjson "chain/encoding/json"
	"chain/errors"
	"chain/fedchain/bc"
)

type nodeTx struct {
	ID       bc.Hash            `json:"id"`
	Time     time.Time          `json:"transaction_time"`
	Inputs   []nodeTxInput      `json:"inputs"`
	Outputs  []nodeTxOutput     `json:"outputs"`
	Metadata chainjson.HexBytes `json:"metadata"`
}

type nodeTxInput struct {
	Type         string             `json:"type"`
	TxID         bc.Hash            `json:"transaction_id,omitempty"`
	TxOut        uint32             `json:"transaction_output,omitempty"`
	AssetID      bc.AssetID         `json:"asset_id"`
	AssetLabel   string             `json:"asset_label,omitempty"`
	Amount       uint64             `json:"amount"`
	Address      chainjson.HexBytes `json:"address,omitempty"`
	AccountID    string             `json:"account_id,omitempty"`
	AccountLabel string             `json:"account_label"`
	Metadata     chainjson.HexBytes `json:"metadata"`

	mNodeID string
}

type nodeTxOutput struct {
	AssetID      bc.AssetID         `json:"asset_id"`
	AssetLabel   string             `json:"asset_label,omitempty"`
	Amount       uint64             `json:"amount"`
	Address      chainjson.HexBytes `json:"address,omitempty"`
	AccountID    string             `json:"account_id,omitempty"`
	AccountLabel string             `json:"account_label,omitempty"`
	Metadata     chainjson.HexBytes `json:"metadata"`

	mNodeID string
}

// WriteNodeTxs persists a transaction along with its metadata
// for every node (issuer, manager) associated with the transaction.
func WriteNodeTxs(ctx context.Context, tx *bc.Tx, ts time.Time) error {
	ins, outs, err := appdb.GetActUTXOs(ctx, tx)
	if err != nil {
		return errors.Wrap(err, "fetching utxo account data")
	}

	assetMap := make(map[string]*appdb.ActAsset)
	var assetIDs []string
	for _, out := range tx.Outputs {
		assetIDs = append(assetIDs, out.AssetID.String())
	}
	assets, err := appdb.GetActAssets(ctx, assetIDs)
	if err != nil {
		return errors.Wrap(err, "getting assets")
	}
	for _, asset := range assets {
		assetMap[asset.ID] = asset
	}

	accountMap := make(map[string]*appdb.ActAccount)
	nodeAccounts := make(map[string][]string)
	var accountIDs []string
	for _, utxo := range append(ins, outs...) {
		if utxo.AccountID != "" {
			accountIDs = append(accountIDs, utxo.AccountID)
		}
	}
	accounts, err := appdb.GetActAccounts(ctx, accountIDs)
	if err != nil {
		return errors.Wrap(err, "getting accounts")
	}
	for _, acc := range accounts {
		accountMap[acc.ID] = acc
		nodeAccounts[acc.ManagerNodeID] = append(nodeAccounts[acc.ManagerNodeID], acc.ID)
	}

	nodeTx, err := generateNodeTx(tx, ins, outs, assetMap, accountMap, ts)
	if err != nil {
		return errors.Wrap(err, "generating master")
	}

	if tx.IsIssuance() {
		filteredTx, err := json.Marshal(filterAccounts(nodeTx, ""))
		if err != nil {
			return errors.Wrap(err, "filtering tx")
		}
		err = appdb.WriteIssuerTx(ctx, tx.Hash.String(), filteredTx, assets[0].IssuerNodeID, assets[0].ID)
		if err != nil {
			return errors.Wrap(err, "writing issuer tx")
		}
	}

	for nodeID, accountIDs := range nodeAccounts {
		filteredTx, err := json.Marshal(filterAccounts(nodeTx, nodeID))
		if err != nil {
			return errors.Wrap(err, "filtering tx")
		}
		err = appdb.WriteManagerTx(ctx, tx.Hash.String(), filteredTx, nodeID, accountIDs)
		if err != nil {
			return errors.Wrap(err, "writing manager tx")
		}
	}

	return nil
}

func generateNodeTx(
	tx *bc.Tx,
	ins []*appdb.ActUTXO,
	outs []*appdb.ActUTXO,
	assetMap map[string]*appdb.ActAsset,
	accountMap map[string]*appdb.ActAccount,
	ts time.Time,
) (*nodeTx, error) {
	actTx := &nodeTx{
		ID:       tx.Hash,
		Metadata: tx.Metadata,
		Time:     ts,
	}
	for i, in := range tx.Inputs {
		if in.IsIssuance() {
			var (
				amt     uint64
				assetID bc.AssetID
			)
			for _, o := range tx.Outputs {
				amt += o.Value
				assetID = o.AssetID
			}
			asset := assetMap[assetID.String()]
			var label string
			if asset != nil {
				label = asset.Label
			}
			actTx.Inputs = append(actTx.Inputs, nodeTxInput{
				Type:       "issuance",
				AssetID:    assetID,
				AssetLabel: label,
				Amount:     amt,
				Metadata:   in.Metadata,
			})
			continue
		}

		assetID, err := bc.ParseHash(ins[i].AssetID)
		if err != nil {
			return nil, errors.Wrap(err, "parsing utxo asset id")
		}
		asset := assetMap[ins[i].AssetID]
		var assetLabel string
		if asset != nil {
			assetLabel = asset.Label
		}

		account := accountMap[ins[i].AccountID]
		var accountLabel, mNodeID string
		if account != nil {
			accountLabel = account.Label
			mNodeID = account.ManagerNodeID
		}

		actTx.Inputs = append(actTx.Inputs, nodeTxInput{
			Type:         "transfer",
			TxID:         in.Previous.Hash,
			TxOut:        in.Previous.Index,
			AssetID:      bc.AssetID(assetID),
			AssetLabel:   assetLabel,
			Amount:       ins[i].Amount,
			AccountID:    ins[i].AccountID,
			AccountLabel: accountLabel,
			Address:      ins[i].Script,
			Metadata:     in.Metadata,
			mNodeID:      mNodeID,
		})
	}

	for i, out := range tx.Outputs {
		asset := assetMap[out.AssetID.String()]
		var assetLabel string
		if asset != nil {
			assetLabel = asset.Label
		}

		account := accountMap[outs[i].AccountID]
		var accountLabel, mNodeID string
		if account != nil {
			accountLabel = account.Label
			mNodeID = account.ManagerNodeID
		}

		actTx.Outputs = append(actTx.Outputs, nodeTxOutput{
			AssetID:      out.AssetID,
			AssetLabel:   assetLabel,
			Amount:       out.Value,
			Address:      out.Script,
			AccountID:    outs[i].AccountID,
			AccountLabel: accountLabel,
			Metadata:     out.Metadata,
			mNodeID:      mNodeID,
		})
	}
	return actTx, nil
}

func filterAccounts(tx *nodeTx, keepID string) *nodeTx {
	filteredTx := new(nodeTx)
	*filteredTx = *tx
	filteredTx.Inputs = make([]nodeTxInput, len(tx.Inputs))
	copy(filteredTx.Inputs, tx.Inputs)
	filteredTx.Outputs = make([]nodeTxOutput, len(tx.Outputs))
	copy(filteredTx.Outputs, tx.Outputs)
	for i := range filteredTx.Inputs {
		in := &filteredTx.Inputs[i]
		if in.mNodeID != keepID {
			in.AccountLabel = ""
			in.AccountID = ""
		}
	}
	for i := range filteredTx.Outputs {
		out := &filteredTx.Outputs[i]
		if out.mNodeID != keepID {
			out.AccountLabel = ""
			out.AccountID = ""
		}
	}
	return filteredTx
}
<commit_msg>api/asset: add omitempty json tag to account_label<commit_after>package asset

import (
	"encoding/json"
	"time"

	"golang.org/x/net/context"

	"chain/api/appdb"
	chainjson "chain/encoding/json"
	"chain/errors"
	"chain/fedchain/bc"
)

type nodeTx struct {
	ID       bc.Hash            `json:"id"`
	Time     time.Time          `json:"transaction_time"`
	Inputs   []nodeTxInput      `json:"inputs"`
	Outputs  []nodeTxOutput     `json:"outputs"`
	Metadata chainjson.HexBytes `json:"metadata"`
}

type nodeTxInput struct {
	Type         string             `json:"type"`
	TxID         bc.Hash            `json:"transaction_id,omitempty"`
	TxOut        uint32             `json:"transaction_output,omitempty"`
	AssetID      bc.AssetID         `json:"asset_id"`
	AssetLabel   string             `json:"asset_label,omitempty"`
	Amount       uint64             `json:"amount"`
	Address      chainjson.HexBytes `json:"address,omitempty"`
	AccountID    string             `json:"account_id,omitempty"`
	AccountLabel string             `json:"account_label,omitempty"`
	Metadata     chainjson.HexBytes `json:"metadata"`

	mNodeID string
}

type nodeTxOutput struct {
	AssetID      bc.AssetID         `json:"asset_id"`
	AssetLabel   string             `json:"asset_label,omitempty"`
	Amount       uint64             `json:"amount"`
	Address      chainjson.HexBytes `json:"address,omitempty"`
	AccountID    string             `json:"account_id,omitempty"`
	AccountLabel string             `json:"account_label,omitempty"`
	Metadata     chainjson.HexBytes `json:"metadata"`

	mNodeID string
}

// WriteNodeTxs persists a transaction along with its metadata
// for every node (issuer, manager) associated with the transaction.
func WriteNodeTxs(ctx context.Context, tx *bc.Tx, ts time.Time) error {
	ins, outs, err := appdb.GetActUTXOs(ctx, tx)
	if err != nil {
		return errors.Wrap(err, "fetching utxo account data")
	}

	assetMap := make(map[string]*appdb.ActAsset)
	var assetIDs []string
	for _, out := range tx.Outputs {
		assetIDs = append(assetIDs, out.AssetID.String())
	}
	assets, err := appdb.GetActAssets(ctx, assetIDs)
	if err != nil {
		return errors.Wrap(err, "getting assets")
	}
	for _, asset := range assets {
		assetMap[asset.ID] = asset
	}

	accountMap := make(map[string]*appdb.ActAccount)
	nodeAccounts := make(map[string][]string)
	var accountIDs []string
	for _, utxo := range append(ins, outs...) {
		if utxo.AccountID != "" {
			accountIDs = append(accountIDs, utxo.AccountID)
		}
	}
	accounts, err := appdb.GetActAccounts(ctx, accountIDs)
	if err != nil {
		return errors.Wrap(err, "getting accounts")
	}
	for _, acc := range accounts {
		accountMap[acc.ID] = acc
		nodeAccounts[acc.ManagerNodeID] = append(nodeAccounts[acc.ManagerNodeID], acc.ID)
	}

	nodeTx, err := generateNodeTx(tx, ins, outs, assetMap, accountMap, ts)
	if err != nil {
		return errors.Wrap(err, "generating master")
	}

	if tx.IsIssuance() {
		filteredTx, err := json.Marshal(filterAccounts(nodeTx, ""))
		if err != nil {
			return errors.Wrap(err, "filtering tx")
		}
		err = appdb.WriteIssuerTx(ctx, tx.Hash.String(), filteredTx, assets[0].IssuerNodeID, assets[0].ID)
		if err != nil {
			return errors.Wrap(err, "writing issuer tx")
		}
	}

	for nodeID, accountIDs := range nodeAccounts {
		filteredTx, err := json.Marshal(filterAccounts(nodeTx, nodeID))
		if err != nil {
			return errors.Wrap(err, "filtering tx")
		}
		err = appdb.WriteManagerTx(ctx, tx.Hash.String(), filteredTx, nodeID, accountIDs)
		if err != nil {
			return errors.Wrap(err, "writing manager tx")
		}
	}

	return nil
}

func generateNodeTx(
	tx *bc.Tx,
	ins []*appdb.ActUTXO,
	outs []*appdb.ActUTXO,
	assetMap map[string]*appdb.ActAsset,
	accountMap map[string]*appdb.ActAccount,
	ts time.Time,
) (*nodeTx, error) {
	actTx := &nodeTx{
		ID:       tx.Hash,
		Metadata: tx.Metadata,
		Time:     ts,
	}
	for i, in := range tx.Inputs {
		if in.IsIssuance() {
			var (
				amt     uint64
				assetID bc.AssetID
			)
			for _, o := range tx.Outputs {
				amt += o.Value
				assetID = o.AssetID
			}
			asset := assetMap[assetID.String()]
			var label string
			if asset != nil {
				label = asset.Label
			}
			actTx.Inputs = append(actTx.Inputs, nodeTxInput{
				Type:       "issuance",
				AssetID:    assetID,
				AssetLabel: label,
				Amount:     amt,
				Metadata:   in.Metadata,
			})
			continue
		}

		assetID, err := bc.ParseHash(ins[i].AssetID)
		if err != nil {
			return nil, errors.Wrap(err, "parsing utxo asset id")
		}
		asset := assetMap[ins[i].AssetID]
		var assetLabel string
		if asset != nil {
			assetLabel = asset.Label
		}

		account := accountMap[ins[i].AccountID]
		var accountLabel, mNodeID string
		if account != nil {
			accountLabel = account.Label
			mNodeID = account.ManagerNodeID
		}

		actTx.Inputs = append(actTx.Inputs, nodeTxInput{
			Type:         "transfer",
			TxID:         in.Previous.Hash,
			TxOut:        in.Previous.Index,
			AssetID:      bc.AssetID(assetID),
			AssetLabel:   assetLabel,
			Amount:       ins[i].Amount,
			AccountID:    ins[i].AccountID,
			AccountLabel: accountLabel,
			Address:      ins[i].Script,
			Metadata:     in.Metadata,
			mNodeID:      mNodeID,
		})
	}

	for i, out := range tx.Outputs {
		asset := assetMap[out.AssetID.String()]
		var assetLabel string
		if asset != nil {
			assetLabel = asset.Label
		}

		account := accountMap[outs[i].AccountID]
		var accountLabel, mNodeID string
		if account != nil {
			accountLabel = account.Label
			mNodeID = account.ManagerNodeID
		}

		actTx.Outputs = append(actTx.Outputs, nodeTxOutput{
			AssetID:      out.AssetID,
			AssetLabel:   assetLabel,
			Amount:       out.Value,
			Address:      out.Script,
			AccountID:    outs[i].AccountID,
			AccountLabel: accountLabel,
			Metadata:     out.Metadata,
			mNodeID:      mNodeID,
		})
	}
	return actTx, nil
}

func filterAccounts(tx *nodeTx, keepID string) *nodeTx {
	filteredTx := new(nodeTx)
	*filteredTx = *tx
	filteredTx.Inputs = make([]nodeTxInput, len(tx.Inputs))
	copy(filteredTx.Inputs, tx.Inputs)
	filteredTx.Outputs = make([]nodeTxOutput, len(tx.Outputs))
	copy(filteredTx.Outputs, tx.Outputs)
	for i := range filteredTx.Inputs {
		in := &filteredTx.Inputs[i]
		if in.mNodeID != keepID {
			in.AccountLabel = ""
			in.AccountID = ""
		}
	}
	for i := range filteredTx.Outputs {
		out := &filteredTx.Outputs[i]
		if out.mNodeID != keepID {
			out.AccountLabel = ""
			out.AccountID = ""
		}
	}
	return filteredTx
}
<|endoftext|>


<commit_before>package database

import (
	"encoding/json"
	"fmt"
	"io"
	"os"
	"path"
	"sync"
	"time"

	"github.com/asdine/storm"
	"github.com/asdine/storm/index"
	"github.com/boltdb/bolt"
	"github.com/hashicorp/raft"
)

type (
	// Database is the lowest level of the gansoi database, it represent the
	// on-disk database. Database implements raft.FSM.
	Database struct {
		dbMutex       sync.RWMutex
		db            *storm.DB
		listenersLock sync.RWMutex
		listeners     []Listener
	}
)

var (
	// ErrNotFound is returned when the specified record is not saved.
	ErrNotFound = storm.ErrNotFound
)

const (
	// stableBucket is the bucket name used for raft stable storage.
	stableBucket = "raft.StableStore"

	// logBucket is used for raft log storage.
	logBucket = "raft.LogStore"
)

// NewDatabase will instantiate a new database placed in filepath.
func NewDatabase(filepath string) (*Database, error) {
	d := &Database{}

	err := d.open(path.Join(filepath, "gansoi.db"))
	if err != nil {
		return nil, err
	}

	return d, nil
}

// Close will close the database. Accessing the database after this will
// result in a deadlock.
func (d *Database) Close() error {
	d.dbMutex.RLock()
	return d.db.Close()
}

// open will open the underlying file storage.
func (d *Database) open(filepath string) error {
	db, err := storm.Open(
		filepath,
		storm.BoltOptions(0600, &bolt.Options{Timeout: 1 * time.Second}),
		storm.AutoIncrement(),
	)
	if err != nil {
		return err
	}

	d.db = db

	d.Storm().Bolt.Update(func(tx *bolt.Tx) error {
		tx.CreateBucketIfNotExists([]byte(logBucket))

		return nil
	})

	return nil
}

// Storm will return the underlying Storm database.
func (d *Database) Storm() *storm.DB {
	d.dbMutex.RLock()
	defer d.dbMutex.RUnlock()

	return d.db
}

// ProcessLogEntry will process the log entry and apply whatever needs doing.
func (d *Database) ProcessLogEntry(entry *LogEntry) error {
	var err error

	var v interface{}

	switch entry.Command {
	case CommandSave:
		v, _ = entry.Payload()
		err = d.Save(v)
	case CommandDelete:
		v, _ = entry.Payload()
		err = d.db.DeleteStruct(v)
	default:
		err = fmt.Errorf("not implemented")
	}

	go func(command Command, data interface{}, err error) {
		d.listenersLock.RLock()

		for _, listener := range d.listeners {
			listener.PostLocalApply(command, data, err)
		}

		d.listenersLock.RUnlock()
	}(entry.Command, v, err)

	return err
}

// Apply implements raft.FSM.
func (d *Database) Apply(l *raft.Log) interface{} {
	entry := &LogEntry{}
	err := json.Unmarshal(l.Data, entry)
	if err != nil {
		// This should not happen..?
		fmt.Printf("%s: '%s'\n", err.Error(), string(l.Data))
		return nil
	}

	return d.ProcessLogEntry(entry)
}

// Snapshot implements raft.FSM.
func (d *Database) Snapshot() (raft.FSMSnapshot, error) {
	return &Snapshot{db: d}, nil
}

// Restore implements raft.FSM.
func (d *Database) Restore(source io.ReadCloser) error {
	db := d.Storm().Bolt
	d.dbMutex.Lock()
	defer d.dbMutex.Unlock()
	defer source.Close()

	path := db.Path()
	restorePath := path + ".restoretmp"

	file, err := os.Create(restorePath)
	if err != nil {
		return err
	}
	defer file.Close()

	_, err = io.Copy(file, source)
	if err != nil {
		return err
	}

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

	err = os.Remove(path)
	if err != nil {
		return err
	}

	err = os.Rename(restorePath, path)
	if err != nil {
		return err
	}

	err = d.open(path)
	if err != nil {
		return err
	}

	return nil
}

// Save will save an object to the database.
func (d *Database) Save(data interface{}) error {
	d.dbMutex.RLock()
	defer d.dbMutex.RUnlock()

	return d.db.Save(data)
}

// One will retrieve one (or zero) record from the database.
func (d *Database) One(fieldName string, value interface{}, to interface{}) error {
	d.dbMutex.RLock()
	defer d.dbMutex.RUnlock()

	return d.db.One(fieldName, value, to)
}

// All lists all kinds of a type.
func (d *Database) All(to interface{}, limit int, skip int, reverse bool) error {
	d.dbMutex.RLock()
	defer d.dbMutex.RUnlock()

	return d.db.All(to, func(opts *index.Options) {
		opts.Limit = limit
		opts.Skip = skip
		opts.Reverse = reverse
	})
}

// RegisterListener will register a listener for new changes to the database.
func (d *Database) RegisterListener(listener Listener) {
	d.listenersLock.Lock()
	defer d.listenersLock.Unlock()

	d.listeners = append(d.listeners, listener)
}

// Set implements raft.StableStore.
func (d *Database) Set(k, v []byte) error {
	return d.Storm().Set(stableBucket, k, v)
}

// Get implements raft.StableStore.
func (d *Database) Get(k []byte) ([]byte, error) {
	var value []byte
	err := d.Storm().Get(stableBucket, k, &value)
	if err != nil {
		return nil, err
	}

	return value, nil
}

// SetUint64 is like Set, but handles uint64 values
func (d *Database) SetUint64(key []byte, val uint64) error {
	return d.Set(key, uint64ToBytes(val))
}

// GetUint64 is like Get, but handles uint64 values
func (d *Database) GetUint64(key []byte) (uint64, error) {
	val, err := d.Get(key)
	if err != nil {
		return 0, err
	}

	return bytesToUint64(val), nil
}

// FirstIndex implements raft.LogStore.
func (d *Database) FirstIndex() (uint64, error) {
	tx, err := d.Storm().Bolt.Begin(false)
	if err != nil {
		return 0, err
	}
	defer tx.Rollback()

	curs := tx.Bucket([]byte(logBucket)).Cursor()

	first, _ := curs.First()
	if first == nil {
		return 0, nil
	}

	return bytesToUint64(first), nil
}

// LastIndex implements raft.LogStore.
func (d *Database) LastIndex() (uint64, error) {
	tx, err := d.Storm().Bolt.Begin(false)
	if err != nil {
		return 0, err
	}
	defer tx.Rollback()

	curs := tx.Bucket([]byte(logBucket)).Cursor()

	last, _ := curs.Last()
	if last == nil {
		return 0, nil
	}

	return bytesToUint64(last), nil
}

// GetLog implements raft.LogStore.
func (d *Database) GetLog(idx uint64, log *raft.Log) error {
	tx, err := d.Storm().Bolt.Begin(false)
	if err != nil {
		return err
	}
	defer tx.Rollback()

	value := tx.Bucket([]byte(logBucket)).Get(uint64ToBytes(idx))

	return json.Unmarshal(value, log)
}

// StoreLog implements raft.LogStore.
func (d *Database) StoreLog(log *raft.Log) error {
	return d.StoreLogs([]*raft.Log{log})
}

// StoreLogs implements raft.LogStore.
func (d *Database) StoreLogs(logs []*raft.Log) error {
	tx, err := d.Storm().Bolt.Begin(true)
	if err != nil {
		return err
	}
	defer tx.Rollback()

	for _, log := range logs {
		key := uint64ToBytes(log.Index)
		value, err := json.Marshal(log)
		if err != nil {
			return err
		}

		bucket := tx.Bucket([]byte(logBucket))
		err = bucket.Put(key, value)
		if err != nil {
			return err
		}
	}

	return tx.Commit()
}

// DeleteRange implements raft.LogStore.
func (d *Database) DeleteRange(min, max uint64) error {
	minKey := uint64ToBytes(min)

	// We can safely use standar Bolt functions, we're not interested in the
	// encoded values, we're only using delete, so Storm encoding/decoding
	// doesn't matter.
	tx, err := d.Storm().Bolt.Begin(true)
	if err != nil {
		return err
	}
	defer tx.Rollback()

	curs := tx.Bucket([]byte(logBucket)).Cursor()
	for k, _ := curs.Seek(minKey); k != nil; k, _ = curs.Next() {
		if bytesToUint64(k) > max {
			break
		}

		if err := curs.Delete(); err != nil {
			return err
		}
	}

	return tx.Commit()
}
<commit_msg>Avoid deadlock in open() when restoring from snapshot.<commit_after>package database

import (
	"encoding/json"
	"fmt"
	"io"
	"os"
	"path"
	"sync"
	"time"

	"github.com/asdine/storm"
	"github.com/asdine/storm/index"
	"github.com/boltdb/bolt"
	"github.com/hashicorp/raft"
)

type (
	// Database is the lowest level of the gansoi database, it represent the
	// on-disk database. Database implements raft.FSM.
	Database struct {
		dbMutex       sync.RWMutex
		db            *storm.DB
		listenersLock sync.RWMutex
		listeners     []Listener
	}
)

var (
	// ErrNotFound is returned when the specified record is not saved.
	ErrNotFound = storm.ErrNotFound
)

const (
	// stableBucket is the bucket name used for raft stable storage.
	stableBucket = "raft.StableStore"

	// logBucket is used for raft log storage.
	logBucket = "raft.LogStore"
)

// NewDatabase will instantiate a new database placed in filepath.
func NewDatabase(filepath string) (*Database, error) {
	d := &Database{}

	err := d.open(path.Join(filepath, "gansoi.db"))
	if err != nil {
		return nil, err
	}

	return d, nil
}

// Close will close the database. Accessing the database after this will
// result in a deadlock.
func (d *Database) Close() error {
	d.dbMutex.RLock()
	return d.db.Close()
}

// open will open the underlying file storage.
func (d *Database) open(filepath string) error {
	db, err := storm.Open(
		filepath,
		storm.BoltOptions(0600, &bolt.Options{Timeout: 1 * time.Second}),
		storm.AutoIncrement(),
	)
	if err != nil {
		return err
	}

	d.db = db

	d.db.Bolt.Update(func(tx *bolt.Tx) error {
		tx.CreateBucketIfNotExists([]byte(logBucket))

		return nil
	})

	return nil
}

// Storm will return the underlying Storm database.
func (d *Database) Storm() *storm.DB {
	d.dbMutex.RLock()
	defer d.dbMutex.RUnlock()

	return d.db
}

// ProcessLogEntry will process the log entry and apply whatever needs doing.
func (d *Database) ProcessLogEntry(entry *LogEntry) error {
	var err error

	var v interface{}

	switch entry.Command {
	case CommandSave:
		v, _ = entry.Payload()
		err = d.Save(v)
	case CommandDelete:
		v, _ = entry.Payload()
		err = d.db.DeleteStruct(v)
	default:
		err = fmt.Errorf("not implemented")
	}

	go func(command Command, data interface{}, err error) {
		d.listenersLock.RLock()

		for _, listener := range d.listeners {
			listener.PostLocalApply(command, data, err)
		}

		d.listenersLock.RUnlock()
	}(entry.Command, v, err)

	return err
}

// Apply implements raft.FSM.
func (d *Database) Apply(l *raft.Log) interface{} {
	entry := &LogEntry{}
	err := json.Unmarshal(l.Data, entry)
	if err != nil {
		// This should not happen..?
		fmt.Printf("%s: '%s'\n", err.Error(), string(l.Data))
		return nil
	}

	return d.ProcessLogEntry(entry)
}

// Snapshot implements raft.FSM.
func (d *Database) Snapshot() (raft.FSMSnapshot, error) {
	return &Snapshot{db: d}, nil
}

// Restore implements raft.FSM.
func (d *Database) Restore(source io.ReadCloser) error {
	db := d.Storm().Bolt
	d.dbMutex.Lock()
	defer d.dbMutex.Unlock()
	defer source.Close()

	path := db.Path()
	restorePath := path + ".restoretmp"

	file, err := os.Create(restorePath)
	if err != nil {
		return err
	}
	defer file.Close()

	_, err = io.Copy(file, source)
	if err != nil {
		return err
	}

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

	err = os.Remove(path)
	if err != nil {
		return err
	}

	err = os.Rename(restorePath, path)
	if err != nil {
		return err
	}

	err = d.open(path)
	if err != nil {
		return err
	}

	return nil
}

// Save will save an object to the database.
func (d *Database) Save(data interface{}) error {
	d.dbMutex.RLock()
	defer d.dbMutex.RUnlock()

	return d.db.Save(data)
}

// One will retrieve one (or zero) record from the database.
func (d *Database) One(fieldName string, value interface{}, to interface{}) error {
	d.dbMutex.RLock()
	defer d.dbMutex.RUnlock()

	return d.db.One(fieldName, value, to)
}

// All lists all kinds of a type.
func (d *Database) All(to interface{}, limit int, skip int, reverse bool) error {
	d.dbMutex.RLock()
	defer d.dbMutex.RUnlock()

	return d.db.All(to, func(opts *index.Options) {
		opts.Limit = limit
		opts.Skip = skip
		opts.Reverse = reverse
	})
}

// RegisterListener will register a listener for new changes to the database.
func (d *Database) RegisterListener(listener Listener) {
	d.listenersLock.Lock()
	defer d.listenersLock.Unlock()

	d.listeners = append(d.listeners, listener)
}

// Set implements raft.StableStore.
func (d *Database) Set(k, v []byte) error {
	return d.Storm().Set(stableBucket, k, v)
}

// Get implements raft.StableStore.
func (d *Database) Get(k []byte) ([]byte, error) {
	var value []byte
	err := d.Storm().Get(stableBucket, k, &value)
	if err != nil {
		return nil, err
	}

	return value, nil
}

// SetUint64 is like Set, but handles uint64 values
func (d *Database) SetUint64(key []byte, val uint64) error {
	return d.Set(key, uint64ToBytes(val))
}

// GetUint64 is like Get, but handles uint64 values
func (d *Database) GetUint64(key []byte) (uint64, error) {
	val, err := d.Get(key)
	if err != nil {
		return 0, err
	}

	return bytesToUint64(val), nil
}

// FirstIndex implements raft.LogStore.
func (d *Database) FirstIndex() (uint64, error) {
	tx, err := d.Storm().Bolt.Begin(false)
	if err != nil {
		return 0, err
	}
	defer tx.Rollback()

	curs := tx.Bucket([]byte(logBucket)).Cursor()

	first, _ := curs.First()
	if first == nil {
		return 0, nil
	}

	return bytesToUint64(first), nil
}

// LastIndex implements raft.LogStore.
func (d *Database) LastIndex() (uint64, error) {
	tx, err := d.Storm().Bolt.Begin(false)
	if err != nil {
		return 0, err
	}
	defer tx.Rollback()

	curs := tx.Bucket([]byte(logBucket)).Cursor()

	last, _ := curs.Last()
	if last == nil {
		return 0, nil
	}

	return bytesToUint64(last), nil
}

// GetLog implements raft.LogStore.
func (d *Database) GetLog(idx uint64, log *raft.Log) error {
	tx, err := d.Storm().Bolt.Begin(false)
	if err != nil {
		return err
	}
	defer tx.Rollback()

	value := tx.Bucket([]byte(logBucket)).Get(uint64ToBytes(idx))

	return json.Unmarshal(value, log)
}

// StoreLog implements raft.LogStore.
func (d *Database) StoreLog(log *raft.Log) error {
	return d.StoreLogs([]*raft.Log{log})
}

// StoreLogs implements raft.LogStore.
func (d *Database) StoreLogs(logs []*raft.Log) error {
	tx, err := d.Storm().Bolt.Begin(true)
	if err != nil {
		return err
	}
	defer tx.Rollback()

	for _, log := range logs {
		key := uint64ToBytes(log.Index)
		value, err := json.Marshal(log)
		if err != nil {
			return err
		}

		bucket := tx.Bucket([]byte(logBucket))
		err = bucket.Put(key, value)
		if err != nil {
			return err
		}
	}

	return tx.Commit()
}

// DeleteRange implements raft.LogStore.
func (d *Database) DeleteRange(min, max uint64) error {
	minKey := uint64ToBytes(min)

	// We can safely use standar Bolt functions, we're not interested in the
	// encoded values, we're only using delete, so Storm encoding/decoding
	// doesn't matter.
	tx, err := d.Storm().Bolt.Begin(true)
	if err != nil {
		return err
	}
	defer tx.Rollback()

	curs := tx.Bucket([]byte(logBucket)).Cursor()
	for k, _ := curs.Seek(minKey); k != nil; k, _ = curs.Next() {
		if bytesToUint64(k) > max {
			break
		}

		if err := curs.Delete(); err != nil {
			return err
		}
	}

	return tx.Commit()
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"net/http"
	"net/http/httputil"
	"net/url"
	"io/ioutil"
	"log"
	"github.com/BurntSushi/toml"
)

type Config struct {
	Upstream string
	ReqFanFactor int
	TimeoutMS int
}

func snd (c chan string, s string) {
	go func () { c <- s }()
}

// start starts Charm up and returns a done channel for the done message
func start(confPath string) (chan string) {
	done := make(chan string)

	// print a welcome message
	log.Print("Charm is starting up.")
	// read the config file
	log.Printf(".   . Reading %s", confPath)
	tomlData, err := ioutil.ReadFile(confPath)
	if err != nil {
		snd(done, fmt.Sprintf("Could not read file at %s", confPath))
		return done
	}
	// populate the config struct
	log.Print(".   . Loading config")
	var conf Config
	_, err = toml.Decode(string(tomlData), &conf)
	if err != nil {
		snd(done, "Could not decode config")
		return done
	}
	// report on the configuration
	log.Print("Charm is configured!")
	log.Printf(".   . Stabilizing %v", conf.Upstream)
	log.Printf(".   . with %v duplicate requests", conf.ReqFanFactor)
	log.Printf(".   . and a %v milisecond timeout.", conf.TimeoutMS)
	go run(conf, done)
	return done
}

func First() {}

type stableTransport struct {
	wrappedTransport http.RoundTripper
	reqFanFactor int
}

// stableTransport.RoundTrip makes many round trips and returns the first
// response
func (t *stableTransport) RoundTrip(r *http.Request) (*http.Response, error) {
	c := make(chan *http.Response)
	if t.wrappedTransport == nil {
		t.wrappedTransport = http.DefaultTransport
	}
	// fan out requests, send responses to the channel, log errors
	for i := 0; i < t.reqFanFactor; i++ {
		go func () {
			resp, err := t.wrappedTransport.RoundTrip(r)
			if err != nil {
				fmt.Printf("stableTransport.RoundTrip error (%)\n", err)
			} else {
				c <- resp
			}
		}()
	}

	// return the first good response
	return <-c, nil
}

type Stabilizer struct {
	upstreamURL *url.URL
	reqFanFactor int
}

// Stabilizer.ServeHTTP proxies with a stable transport
func (st *Stabilizer) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	proxy := httputil.NewSingleHostReverseProxy(st.upstreamURL)
	proxy.Transport = &stableTransport{proxy.Transport, st.reqFanFactor}
	proxy.ServeHTTP(w, r)
}

// run
func run(conf Config, done chan string) {
	// make a stabilizer
	upstreamURL := url.Parse(conf.Upstream)
	stabilizer := &Stabilizer{upstreamURL, conf.ReqFanFactor}

	// serve that stabilizer under a timeout
	timeout := conf.TimeoutMS * time.Milliseconds
	log.Fatal(http.ListenAndServe(
		":8000",
		http.TimeoutHandler(stabilizer, timeout, "upstream timeout"),
	))
}

func main() {
	// start Charm,
	done := start("/secret/charm.conf")
	// when Charm is done, log the message and quit.
        log.Print(<-done)
}
<commit_msg>initial working version<commit_after>package main

import (
	"fmt"
	"net/http"
	"net/http/httputil"
	"net/url"
	"io/ioutil"
	"log"
	"time"
	"github.com/BurntSushi/toml"
)

type Config struct {
	Upstream string
	ReqFanFactor int
	TimeoutMS int
}

func snd (c chan string, s string) {
	go func () { c <- s }()
}

// start starts Charm up and returns a done channel for the done message
func start(confPath string) (chan string) {
	done := make(chan string)

	// print a welcome message
	log.Print("Charm is starting up.")
	// read the config file
	log.Printf(".   . Reading %s", confPath)
	tomlData, err := ioutil.ReadFile(confPath)
	if err != nil {
		snd(done, fmt.Sprintf("Could not read file at %s", confPath))
		return done
	}
	// populate the config struct
	log.Print(".   . Loading config")
	var conf Config
	_, err = toml.Decode(string(tomlData), &conf)
	if err != nil {
		snd(done, "Could not decode config")
		return done
	}
	// report on the configuration
	log.Print("Charm is configured!")
	log.Printf(".   . Stabilizing %v", conf.Upstream)
	log.Printf(".   . with %v duplicate requests", conf.ReqFanFactor)
	log.Printf(".   . and a %v milisecond timeout.", conf.TimeoutMS)
	go run(conf, done)
	return done
}

func First() {}

type stableTransport struct {
	wrappedTransport http.RoundTripper
	reqFanFactor int
}

// stableTransport.RoundTrip makes many round trips and returns the first
// response
func (t *stableTransport) RoundTrip(r *http.Request) (*http.Response, error) {
	c := make(chan *http.Response)
	if t.wrappedTransport == nil {
		t.wrappedTransport = http.DefaultTransport
	}
	// fan out requests, send responses to the channel, log errors
	for i := 0; i < t.reqFanFactor; i++ {
		time.Sleep(5)
		go func () {
			resp, err := t.wrappedTransport.RoundTrip(r)
			if err != nil {
				fmt.Printf("stableTransport.RoundTrip error (%)\n", err)
			} else {
				c <- resp
			}
		}()
	}

	// return the first good response
	return <-c, nil
}

type Stabilizer struct {
	upstreamURL *url.URL
	reqFanFactor int
}

// Stabilizer.ServeHTTP proxies with a stable transport
func (st *Stabilizer) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	proxy := httputil.NewSingleHostReverseProxy(st.upstreamURL)
	proxy.Transport = &stableTransport{proxy.Transport, st.reqFanFactor}
	proxy.ServeHTTP(w, r)
}

// run
func run(conf Config, done chan string) {
	// make a stabilizer
	upstreamURL, err := url.Parse(conf.Upstream)
	if err != nil {
		log.Fatal("error parsing Upstream URL", conf.Upstream)
	}
	stabilizer := &Stabilizer{upstreamURL, conf.ReqFanFactor}

	// serve that stabilizer under a timeout
	timeout := time.Duration(conf.TimeoutMS) * time.Millisecond
	log.Fatal(http.ListenAndServe(
		":8000",
		http.TimeoutHandler(stabilizer, timeout, "upstream timeout"),
	))
}

func main() {
	// start Charm,
	done := start("/secret/charm.conf")
	// when Charm is done, log the message and quit.
        log.Print(<-done)
}
<|endoftext|>


<commit_before>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.
//
// Contributor: Julien Vehent jvehent@mozilla.com [:ulfr]

package main

import (
	"crypto/sha256"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"path"
	"time"

	log "github.com/sirupsen/logrus"
	"go.mozilla.org/autograph/signer"
)

// a signaturerequest is sent by an autograph client to request
// a signature on input data
type signaturerequest struct {
	Input   string `json:"input"`
	KeyID   string `json:"keyid,omitempty"`
	Options interface{}
}

// a signatureresponse is returned by autograph to a client with
// a signature computed on input data
type signatureresponse struct {
	Ref        string `json:"ref"`
	Type       string `json:"type"`
	Mode       string `json:"mode"`
	SignerID   string `json:"signer_id"`
	PublicKey  string `json:"public_key"`
	Signature  string `json:"signature,omitempty"`
	SignedFile string `json:"signed_file,omitempty"`
	X5U        string `json:"x5u,omitempty"`
}

// handleSignature endpoint accepts a list of signature requests in a HAWK authenticated POST request
// and calls the signers to generate signature responses.
func (a *autographer) handleSignature(w http.ResponseWriter, r *http.Request) {
	rid := getRequestID(r)
	starttime := time.Now()
	body, err := ioutil.ReadAll(r.Body)
	if err != nil {
		httpError(w, r, http.StatusBadRequest, "failed to read request body: %s", err)
		return
	}
	if r.Header.Get("Content-Type") != "application/json" {
		httpError(w, r, http.StatusBadRequest, "invalid content type, expected application/json")
		return
	}
	if len(body) < 10 {
		// it's impossible to have a valid request body smaller than 10 bytes
		httpError(w, r, http.StatusBadRequest, "empty or invalid request request body")
		return
	}
	if len(body) > 1048576000 {
		// the max body size is hardcoded to 1GB. Seriously, what are you trying to sign?
		httpError(w, r, http.StatusBadRequest, "request exceeds max size of 1GB")
		return
	}
	userid, authorized, err := a.authorize(r, body)
	if err != nil || !authorized {
		httpError(w, r, http.StatusUnauthorized, "authorization verification failed: %v", err)
		return
	}
	var sigreqs []signaturerequest
	err = json.Unmarshal(body, &sigreqs)
	if err != nil {
		httpError(w, r, http.StatusBadRequest, "failed to parse request body: %v", err)
		return
	}
	for i, sigreq := range sigreqs {
		if sigreq.Input == "" {
			httpError(w, r, http.StatusBadRequest, fmt.Sprintf("missing input in signature request %d", i))
		}
	}
	if a.debug {
		fmt.Printf("signature request\n-----------------\n%s\n", body)
	}
	sigresps := make([]signatureresponse, len(sigreqs))
	// Each signature requested in the http request body is processed individually.
	// For each, a signer is looked up, and used to compute a raw signature
	// the signature is then encoded appropriately, and added to the response slice
	for i, sigreq := range sigreqs {
		var (
			input      []byte
			sig        signer.Signature
			signedfile []byte
			hashlog    string
		)

		// Decode the base64 input data
		input, err = base64.StdEncoding.DecodeString(sigreq.Input)
		if err != nil {
			httpError(w, r, http.StatusBadRequest, "%v", err)
			return
		}

		// Find the ID of the requested signer
		// Return an error if the signer is not found or if the user is not allowed
		// to use this signer
		signerID, err := a.getSignerID(userid, sigreq.KeyID)
		if err != nil || signerID < 0 {
			httpError(w, r, http.StatusUnauthorized, "%v", err)
			return
		}
		sigresps[i] = signatureresponse{
			Ref:        id(),
			Type:       a.signers[signerID].Config().Type,
			Mode:       a.signers[signerID].Config().Mode,
			SignerID:   a.signers[signerID].Config().ID,
			PublicKey:  a.signers[signerID].Config().PublicKey,
			SignedFile: base64.StdEncoding.EncodeToString(signedfile),
			X5U:        a.signers[signerID].Config().X5U,
		}
		// Make sure the signer implements the right interface, then sign the data
		switch r.URL.RequestURI() {
		case "/sign/hash":
			hashSigner, ok := a.signers[signerID].(signer.HashSigner)
			if !ok {
				httpError(w, r, http.StatusBadRequest, "requested signer does not implement hash signing")
				return
			}
			sig, err = hashSigner.SignHash(input, sigreq.Options)
			if err != nil {
				httpError(w, r, http.StatusInternalServerError, "signing failed with error: %v", err)
				return
			}
			sigresps[i].Signature, err = sig.(signer.Signature).Marshal()
			if err != nil {
				httpError(w, r, http.StatusInternalServerError, "encoding failed with error: %v", err)
				return
			}
			// convert the input hash to hexadecimal for logging
			hashlog = fmt.Sprintf("%X", input)

		case "/sign/data":
			dataSigner, ok := a.signers[signerID].(signer.DataSigner)
			if !ok {
				httpError(w, r, http.StatusBadRequest, "requested signer does not implement data signing")
				return
			}
			sig, err = dataSigner.SignData(input, sigreq.Options)
			if err != nil {
				httpError(w, r, http.StatusInternalServerError, "signing failed with error: %v", err)
				return
			}
			sigresps[i].Signature, err = sig.(signer.Signature).Marshal()
			if err != nil {
				httpError(w, r, http.StatusInternalServerError, "encoding failed with error: %v", err)
				return
			}
			// calculate a hash of the input to store in the signing logs
			md := sha256.New()
			md.Write(input)
			hashlog = fmt.Sprintf("%X", md.Sum(nil))

		case "/sign/file":
			fileSigner, ok := a.signers[signerID].(signer.FileSigner)
			if !ok {
				httpError(w, r, http.StatusBadRequest, "requested signer does not implement file signing")
				return
			}
			signedfile, err = fileSigner.SignFile(input, sigreq.Options)
			if err != nil {
				httpError(w, r, http.StatusInternalServerError, "signing failed with error: %v", err)
				return
			}
			sigresps[i].SignedFile = base64.StdEncoding.EncodeToString(signedfile)
			// calculate a hash of the input to store in the signing logs
			md := sha256.New()
			md.Write(input)
			hashlog = fmt.Sprintf("%X", md.Sum(nil))
		}
		log.WithFields(log.Fields{
			"rid":        rid,
			"ref":        sigresps[i].Ref,
			"type":       sigresps[i].Type,
			"signer_id":  sigresps[i].SignerID,
			"input_hash": hashlog,
			"user_id":    userid,
			"t":          int32(time.Since(starttime) / time.Millisecond), //  request processing time in ms
		}).Info("signing operation succeeded")
	}
	respdata, err := json.Marshal(sigresps)
	if err != nil {
		httpError(w, r, http.StatusInternalServerError, "signing failed with error: %v", err)
		return
	}
	if a.debug {
		fmt.Printf("signature response\n------------------\n%s\n", respdata)
	}
	w.Header().Add("Content-Type", "application/json")
	w.WriteHeader(http.StatusCreated)
	w.Write(respdata)
	log.WithFields(log.Fields{"rid": rid}).Info("signing request completed successfully")
}

// handleHeartbeat returns a simple message indicating that the API is alive and well
func handleHeartbeat(w http.ResponseWriter, r *http.Request) {
	if r.Method != "GET" {
		httpError(w, r, http.StatusMethodNotAllowed, "%s method not allowed; endpoint accepts GET only", r.Method)
		return
	}
	w.Write([]byte("ohai"))
}

func handleVersion(w http.ResponseWriter, r *http.Request) {
	if r.Method != "GET" {
		httpError(w, r, http.StatusMethodNotAllowed, "%s method not allowed; endpoint accepts GET only", r.Method)
		return
	}
	dir, err := os.Getwd()
	if err != nil {
		httpError(w, r, http.StatusInternalServerError, "Could not get CWD")
		return
	}
	filename := path.Clean(dir + string(os.PathSeparator) + "version.json")
	f, err := os.Open(filename)
	if err != nil {
		httpError(w, r, http.StatusNotFound, "version.json file not found")
		return
	}
	stat, err := f.Stat()
	if err != nil {
		httpError(w, r, http.StatusInternalServerError, "stat failed on version.json")
		return
	}
	w.Header().Set("Content-Type", "application/json")
	http.ServeContent(w, r, "version.json", stat.ModTime(), f)
}
<commit_msg>handlers: time signatures relative to req start time<commit_after>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.
//
// Contributor: Julien Vehent jvehent@mozilla.com [:ulfr]

package main

import (
	"crypto/sha256"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"path"
	"time"

	log "github.com/sirupsen/logrus"
	"go.mozilla.org/autograph/signer"
)

// a signaturerequest is sent by an autograph client to request
// a signature on input data
type signaturerequest struct {
	Input   string `json:"input"`
	KeyID   string `json:"keyid,omitempty"`
	Options interface{}
}

// a signatureresponse is returned by autograph to a client with
// a signature computed on input data
type signatureresponse struct {
	Ref        string `json:"ref"`
	Type       string `json:"type"`
	Mode       string `json:"mode"`
	SignerID   string `json:"signer_id"`
	PublicKey  string `json:"public_key"`
	Signature  string `json:"signature,omitempty"`
	SignedFile string `json:"signed_file,omitempty"`
	X5U        string `json:"x5u,omitempty"`
}

// handleSignature endpoint accepts a list of signature requests in a HAWK authenticated POST request
// and calls the signers to generate signature responses.
func (a *autographer) handleSignature(w http.ResponseWriter, r *http.Request) {
	rid := getRequestID(r)
	starttime := getRequestStartTime(r)
	body, err := ioutil.ReadAll(r.Body)
	if err != nil {
		httpError(w, r, http.StatusBadRequest, "failed to read request body: %s", err)
		return
	}
	if r.Header.Get("Content-Type") != "application/json" {
		httpError(w, r, http.StatusBadRequest, "invalid content type, expected application/json")
		return
	}
	if len(body) < 10 {
		// it's impossible to have a valid request body smaller than 10 bytes
		httpError(w, r, http.StatusBadRequest, "empty or invalid request request body")
		return
	}
	if len(body) > 1048576000 {
		// the max body size is hardcoded to 1GB. Seriously, what are you trying to sign?
		httpError(w, r, http.StatusBadRequest, "request exceeds max size of 1GB")
		return
	}
	userid, authorized, err := a.authorize(r, body)
	if err != nil || !authorized {
		httpError(w, r, http.StatusUnauthorized, "authorization verification failed: %v", err)
		return
	}
	var sigreqs []signaturerequest
	err = json.Unmarshal(body, &sigreqs)
	if err != nil {
		httpError(w, r, http.StatusBadRequest, "failed to parse request body: %v", err)
		return
	}
	for i, sigreq := range sigreqs {
		if sigreq.Input == "" {
			httpError(w, r, http.StatusBadRequest, fmt.Sprintf("missing input in signature request %d", i))
		}
	}
	if a.debug {
		fmt.Printf("signature request\n-----------------\n%s\n", body)
	}
	sigresps := make([]signatureresponse, len(sigreqs))
	// Each signature requested in the http request body is processed individually.
	// For each, a signer is looked up, and used to compute a raw signature
	// the signature is then encoded appropriately, and added to the response slice
	for i, sigreq := range sigreqs {
		var (
			input      []byte
			sig        signer.Signature
			signedfile []byte
			hashlog    string
		)

		// Decode the base64 input data
		input, err = base64.StdEncoding.DecodeString(sigreq.Input)
		if err != nil {
			httpError(w, r, http.StatusBadRequest, "%v", err)
			return
		}

		// Find the ID of the requested signer
		// Return an error if the signer is not found or if the user is not allowed
		// to use this signer
		signerID, err := a.getSignerID(userid, sigreq.KeyID)
		if err != nil || signerID < 0 {
			httpError(w, r, http.StatusUnauthorized, "%v", err)
			return
		}
		sigresps[i] = signatureresponse{
			Ref:        id(),
			Type:       a.signers[signerID].Config().Type,
			Mode:       a.signers[signerID].Config().Mode,
			SignerID:   a.signers[signerID].Config().ID,
			PublicKey:  a.signers[signerID].Config().PublicKey,
			SignedFile: base64.StdEncoding.EncodeToString(signedfile),
			X5U:        a.signers[signerID].Config().X5U,
		}
		// Make sure the signer implements the right interface, then sign the data
		switch r.URL.RequestURI() {
		case "/sign/hash":
			hashSigner, ok := a.signers[signerID].(signer.HashSigner)
			if !ok {
				httpError(w, r, http.StatusBadRequest, "requested signer does not implement hash signing")
				return
			}
			sig, err = hashSigner.SignHash(input, sigreq.Options)
			if err != nil {
				httpError(w, r, http.StatusInternalServerError, "signing failed with error: %v", err)
				return
			}
			sigresps[i].Signature, err = sig.(signer.Signature).Marshal()
			if err != nil {
				httpError(w, r, http.StatusInternalServerError, "encoding failed with error: %v", err)
				return
			}
			// convert the input hash to hexadecimal for logging
			hashlog = fmt.Sprintf("%X", input)

		case "/sign/data":
			dataSigner, ok := a.signers[signerID].(signer.DataSigner)
			if !ok {
				httpError(w, r, http.StatusBadRequest, "requested signer does not implement data signing")
				return
			}
			sig, err = dataSigner.SignData(input, sigreq.Options)
			if err != nil {
				httpError(w, r, http.StatusInternalServerError, "signing failed with error: %v", err)
				return
			}
			sigresps[i].Signature, err = sig.(signer.Signature).Marshal()
			if err != nil {
				httpError(w, r, http.StatusInternalServerError, "encoding failed with error: %v", err)
				return
			}
			// calculate a hash of the input to store in the signing logs
			md := sha256.New()
			md.Write(input)
			hashlog = fmt.Sprintf("%X", md.Sum(nil))

		case "/sign/file":
			fileSigner, ok := a.signers[signerID].(signer.FileSigner)
			if !ok {
				httpError(w, r, http.StatusBadRequest, "requested signer does not implement file signing")
				return
			}
			signedfile, err = fileSigner.SignFile(input, sigreq.Options)
			if err != nil {
				httpError(w, r, http.StatusInternalServerError, "signing failed with error: %v", err)
				return
			}
			sigresps[i].SignedFile = base64.StdEncoding.EncodeToString(signedfile)
			// calculate a hash of the input to store in the signing logs
			md := sha256.New()
			md.Write(input)
			hashlog = fmt.Sprintf("%X", md.Sum(nil))
		}
		log.WithFields(log.Fields{
			"rid":        rid,
			"ref":        sigresps[i].Ref,
			"type":       sigresps[i].Type,
			"signer_id":  sigresps[i].SignerID,
			"input_hash": hashlog,
			"user_id":    userid,
			"t":          int32(time.Since(starttime) / time.Millisecond), //  request processing time in ms
		}).Info("signing operation succeeded")
	}
	respdata, err := json.Marshal(sigresps)
	if err != nil {
		httpError(w, r, http.StatusInternalServerError, "signing failed with error: %v", err)
		return
	}
	if a.debug {
		fmt.Printf("signature response\n------------------\n%s\n", respdata)
	}
	w.Header().Add("Content-Type", "application/json")
	w.WriteHeader(http.StatusCreated)
	w.Write(respdata)
	log.WithFields(log.Fields{"rid": rid}).Info("signing request completed successfully")
}

// handleHeartbeat returns a simple message indicating that the API is alive and well
func handleHeartbeat(w http.ResponseWriter, r *http.Request) {
	if r.Method != "GET" {
		httpError(w, r, http.StatusMethodNotAllowed, "%s method not allowed; endpoint accepts GET only", r.Method)
		return
	}
	w.Write([]byte("ohai"))
}

func handleVersion(w http.ResponseWriter, r *http.Request) {
	if r.Method != "GET" {
		httpError(w, r, http.StatusMethodNotAllowed, "%s method not allowed; endpoint accepts GET only", r.Method)
		return
	}
	dir, err := os.Getwd()
	if err != nil {
		httpError(w, r, http.StatusInternalServerError, "Could not get CWD")
		return
	}
	filename := path.Clean(dir + string(os.PathSeparator) + "version.json")
	f, err := os.Open(filename)
	if err != nil {
		httpError(w, r, http.StatusNotFound, "version.json file not found")
		return
	}
	stat, err := f.Stat()
	if err != nil {
		httpError(w, r, http.StatusInternalServerError, "stat failed on version.json")
		return
	}
	w.Header().Set("Content-Type", "application/json")
	http.ServeContent(w, r, "version.json", stat.ModTime(), f)
}
<|endoftext|>


<commit_before>package exporter

import (
	"io"
	"encoding/json"
	"strconv"

	"github.com/google/cayley/graph"
)

type Exporter struct {
	wr io.Writer
	qstore graph.QuadStore
	qi graph.Iterator
	err error
	count int32
}

func NewExporter(writer io.Writer, qstore graph.QuadStore) *Exporter {
	return NewExporterForIterator(writer, qstore, qstore.QuadsAllIterator())
}

func NewExporterForIterator(writer io.Writer, qstore graph.QuadStore, qi graph.Iterator) *Exporter {
	return &Exporter{wr: writer, qstore: qstore, qi: qi}
}

// number of records
func (exp *Exporter) Count() int32 {
	return exp.count
}

func (exp *Exporter) ExportQuad() {
	exp.qi.Reset()
	for it := exp.qi; graph.Next(it); {
		exp.count++
		quad := exp.qstore.Quad(it.Result())

		exp.WriteEscString(quad.Subject)
		exp.Write(" ")
		exp.WriteEscString(quad.Predicate)
		exp.Write(" ")
		exp.WriteEscString(quad.Object)
		if quad.Label != "" {
			exp.Write(" ")
			exp.WriteEscString(quad.Label)
		}
		exp.Write(" .\n")
	}
}

func (exp *Exporter) ExportJson() {
	var jstr []byte
	exp.Write("[")
	exp.qi.Reset()
	for it := exp.qi; graph.Next(it); {
		exp.count++
		if exp.count > 1 {
			exp.Write(",")
		}

		jstr, exp.err = json.Marshal(exp.qstore.Quad(it.Result()))
		if exp.err != nil {
			return
		}
		exp.Write(string(jstr[:]))
	}
	exp.Write("]\n")
}

//experimental
func (exp *Exporter) ExportGml() {
	var seen map[string]int32 // todo eliminate this for large dbs
	var id int32

	exp.Write("Creator Cayley\ngraph\n[\n")

	seen = make(map[string]int32)
	exp.qi.Reset()
	for it := exp.qi; graph.Next(it); {
		cur := exp.qstore.Quad(it.Result())
		if _, ok := seen[cur.Subject]; !ok {
			exp.Write("  node\n  [\n    id ")
			seen[cur.Subject] = id
			exp.Write(strconv.FormatInt(int64(id), 10))
			exp.Write("\n    label ")
			exp.WriteEscString(cur.Subject)
			exp.Write("\n  ]\n")
			id++
		}
		if _, ok := seen[cur.Object]; !ok {
			exp.Write("  node\n  [\n    id ")
			seen[cur.Object] = id
			exp.Write(strconv.FormatInt(int64(id), 10))
			exp.Write("\n    label ")
			exp.WriteEscString(cur.Object)
			exp.Write("\n  ]\n")
			id++
		}
		exp.count++
	}

	exp.qi.Reset()
	for it := exp.qi; graph.Next(it); {
		cur := exp.qstore.Quad(it.Result())
		exp.Write("  edge\n  [\n    source ")
		exp.Write(strconv.FormatInt(int64(seen[cur.Subject]), 10))
		exp.Write("\n    target ")
		exp.Write(strconv.FormatInt(int64(seen[cur.Object]), 10))
		exp.Write("\n    label ")
		exp.WriteEscString(cur.Predicate)
		exp.Write("\n  ]\n")
		exp.count++
	}
	exp.Write("]\n")
}

//experimental
func (exp *Exporter) ExportGraphml() {
	var seen map[string]bool // eliminate this for large databases

	exp.Write("<?xml version=\"1.0\" encoding=\"UTF-8\"?>\n")
	exp.Write("<graphml xmlns=\"http://graphml.graphdrawing.org/xmlns\"\n")
	exp.Write("   xmlns:xsi=\"http://www.w3.org/2001/XMLSchema-instance\"\n")
	exp.Write("   xsi:schemaLocation=\"http://graphml.graphdrawing.org/xmlns/1.0/graphml.xsd\">\n")
	exp.Write("  <graph id=\"Caylay\" edgedefault=\"directed\">\n")

	seen = make(map[string]bool)
	exp.qi.Reset()
	for it := exp.qi; graph.Next(it); {
		cur := exp.qstore.Quad(it.Result())
		if found := seen[cur.Subject]; !found {
			seen[cur.Subject] = true
			exp.Write("    <node id=")
			exp.WriteEscString(cur.Subject)
			exp.Write(" />\n")
		}
		if found := seen[cur.Object]; !found {
			seen[cur.Object] = true
			exp.Write("    <node id=")
			exp.WriteEscString(cur.Object)
			exp.Write(" />\n")
		}
		exp.count++
	}

	exp.qi.Reset()
	for it := exp.qi; graph.Next(it); {
		cur := exp.qstore.Quad(it.Result())
		exp.Write("    <edge source=")
		exp.WriteEscString(cur.Subject)
		exp.Write(" target=")
                exp.WriteEscString(cur.Object)
		exp.Write(">\n")
		exp.Write("      <data key=\"predicate\">")
		exp.Write(cur.Predicate)
		exp.Write("</data>\n    </edge>\n")
		exp.count++
	}
	exp.Write("  </graph>\n</graphml>\n");
}

//print out the string quoted, escaped
func (exp *Exporter) WriteEscString(str string) {
	var esc []byte

	if exp.err != nil {
		return
	}
	esc, exp.err = json.Marshal(str)
	if exp.err != nil {
		return
	}
	_, exp.err = exp.wr.Write(esc)
}

func (exp *Exporter) Write(str string) {
	if exp.err != nil {
		return
	}
	_, exp.err = exp.wr.Write([]byte(str))
}

func (exp *Exporter) Err() error {
	return exp.err
}
<commit_msg>changed count to just int<commit_after>package exporter

import (
	"io"
	"encoding/json"
	"strconv"

	"github.com/google/cayley/graph"
)

type Exporter struct {
	wr io.Writer
	qstore graph.QuadStore
	qi graph.Iterator
	err error
	count int
}

func NewExporter(writer io.Writer, qstore graph.QuadStore) *Exporter {
	return NewExporterForIterator(writer, qstore, qstore.QuadsAllIterator())
}

func NewExporterForIterator(writer io.Writer, qstore graph.QuadStore, qi graph.Iterator) *Exporter {
	return &Exporter{wr: writer, qstore: qstore, qi: qi}
}

// number of records
func (exp *Exporter) Count() int {
	return exp.count
}

func (exp *Exporter) ExportQuad() {
	exp.qi.Reset()
	for it := exp.qi; graph.Next(it); {
		exp.count++
		quad := exp.qstore.Quad(it.Result())

		exp.WriteEscString(quad.Subject)
		exp.Write(" ")
		exp.WriteEscString(quad.Predicate)
		exp.Write(" ")
		exp.WriteEscString(quad.Object)
		if quad.Label != "" {
			exp.Write(" ")
			exp.WriteEscString(quad.Label)
		}
		exp.Write(" .\n")
	}
}

func (exp *Exporter) ExportJson() {
	var jstr []byte
	exp.Write("[")
	exp.qi.Reset()
	for it := exp.qi; graph.Next(it); {
		exp.count++
		if exp.count > 1 {
			exp.Write(",")
		}

		jstr, exp.err = json.Marshal(exp.qstore.Quad(it.Result()))
		if exp.err != nil {
			return
		}
		exp.Write(string(jstr[:]))
	}
	exp.Write("]\n")
}

//experimental
func (exp *Exporter) ExportGml() {
	var seen map[string]int32 // todo eliminate this for large dbs
	var id int32

	exp.Write("Creator Cayley\ngraph\n[\n")

	seen = make(map[string]int32)
	exp.qi.Reset()
	for it := exp.qi; graph.Next(it); {
		cur := exp.qstore.Quad(it.Result())
		if _, ok := seen[cur.Subject]; !ok {
			exp.Write("  node\n  [\n    id ")
			seen[cur.Subject] = id
			exp.Write(strconv.FormatInt(int64(id), 10))
			exp.Write("\n    label ")
			exp.WriteEscString(cur.Subject)
			exp.Write("\n  ]\n")
			id++
		}
		if _, ok := seen[cur.Object]; !ok {
			exp.Write("  node\n  [\n    id ")
			seen[cur.Object] = id
			exp.Write(strconv.FormatInt(int64(id), 10))
			exp.Write("\n    label ")
			exp.WriteEscString(cur.Object)
			exp.Write("\n  ]\n")
			id++
		}
		exp.count++
	}

	exp.qi.Reset()
	for it := exp.qi; graph.Next(it); {
		cur := exp.qstore.Quad(it.Result())
		exp.Write("  edge\n  [\n    source ")
		exp.Write(strconv.FormatInt(int64(seen[cur.Subject]), 10))
		exp.Write("\n    target ")
		exp.Write(strconv.FormatInt(int64(seen[cur.Object]), 10))
		exp.Write("\n    label ")
		exp.WriteEscString(cur.Predicate)
		exp.Write("\n  ]\n")
		exp.count++
	}
	exp.Write("]\n")
}

//experimental
func (exp *Exporter) ExportGraphml() {
	var seen map[string]bool // eliminate this for large databases

	exp.Write("<?xml version=\"1.0\" encoding=\"UTF-8\"?>\n")
	exp.Write("<graphml xmlns=\"http://graphml.graphdrawing.org/xmlns\"\n")
	exp.Write("   xmlns:xsi=\"http://www.w3.org/2001/XMLSchema-instance\"\n")
	exp.Write("   xsi:schemaLocation=\"http://graphml.graphdrawing.org/xmlns/1.0/graphml.xsd\">\n")
	exp.Write("  <graph id=\"Caylay\" edgedefault=\"directed\">\n")

	seen = make(map[string]bool)
	exp.qi.Reset()
	for it := exp.qi; graph.Next(it); {
		cur := exp.qstore.Quad(it.Result())
		if found := seen[cur.Subject]; !found {
			seen[cur.Subject] = true
			exp.Write("    <node id=")
			exp.WriteEscString(cur.Subject)
			exp.Write(" />\n")
		}
		if found := seen[cur.Object]; !found {
			seen[cur.Object] = true
			exp.Write("    <node id=")
			exp.WriteEscString(cur.Object)
			exp.Write(" />\n")
		}
		exp.count++
	}

	exp.qi.Reset()
	for it := exp.qi; graph.Next(it); {
		cur := exp.qstore.Quad(it.Result())
		exp.Write("    <edge source=")
		exp.WriteEscString(cur.Subject)
		exp.Write(" target=")
                exp.WriteEscString(cur.Object)
		exp.Write(">\n")
		exp.Write("      <data key=\"predicate\">")
		exp.Write(cur.Predicate)
		exp.Write("</data>\n    </edge>\n")
		exp.count++
	}
	exp.Write("  </graph>\n</graphml>\n");
}

//print out the string quoted, escaped
func (exp *Exporter) WriteEscString(str string) {
	var esc []byte

	if exp.err != nil {
		return
	}
	esc, exp.err = json.Marshal(str)
	if exp.err != nil {
		return
	}
	_, exp.err = exp.wr.Write(esc)
}

func (exp *Exporter) Write(str string) {
	if exp.err != nil {
		return
	}
	_, exp.err = exp.wr.Write([]byte(str))
}

func (exp *Exporter) Err() error {
	return exp.err
}
<|endoftext|>


<commit_before>package executor

import (
	log "github.com/mesosphere/kubernetes-mesos/3rdparty/github.com/golang/glog"
	"github.com/mesosphere/kubernetes-mesos/3rdparty/github.com/mesosphere/mesos-go/mesos"
)

// KuberneteExecutor is an mesos executor that runs pods
// in a minion machine.
type KuberneteExecutor struct {
}

// New creates a new kubernete executor.
func New() *KuberneteExecutor {
	return &KuberneteExecutor{}
}

// Registered is called when the executor is successfully registered with the slave.
func (k *KuberneteExecutor) Registered(driver mesos.ExecutorDriver,
	executorInfo *mesos.ExecutorInfo, frameworkInfo *mesos.FrameworkInfo, slaveInfo *mesos.SlaveInfo) {
	log.Infof("Executor %v of framework %v registered with slave %v\n",
		executorInfo, frameworkInfo, slaveInfo)

	// TODO(yifan): Update the executor's status.
}

// Reregistered is called when the executor is successfully re-registered with the slave.
// This can happen when the slave fails over.
func (k *KuberneteExecutor) Reregistered(driver mesos.ExecutorDriver,
	executorInfo *mesos.ExecutorInfo, frameworkInfo *mesos.FrameworkInfo, slaveInfo *mesos.SlaveInfo) {
	log.Infof("Executor %v of framework %v reregistered with slave %v\n",
		executorInfo, frameworkInfo, slaveInfo)

	// TODO(yifan): Update the executor's status.
}

// Disconnected is called when the executor is disconnected with the slave.
func (k *KuberneteExecutor) Disconnected(driver mesos.ExecutorDriver) {
	log.Infof("Slave is disconnected\n")

	// TODO(yifan): Update the executor's status.
}

// LaunchTask is called when the executor receives a request to launch a task.
func (k *KuberneteExecutor) LaunchTask(driver mesos.ExecutorDriver, taskInfo *mesos.TaskInfo) {
	log.Infof("Launch task %v\n", taskInfo)

	// TODO(yifan): Launch a pod by update the pod manifest, and do a sync.
	// Then send back status update.
}

// KillTask is called when the executor receives a request to kill a task.
func (k *KuberneteExecutor) KillTask(driver mesos.ExecutorDriver, taskId *mesos.TaskID) {
	log.Infof("Kill task %v\n", taskId)

	// TODO(yifan): kill the container by updating the pod manifest, and do a sync.
	// Then send back status update.
}

// FrameworkMessage is called when the framework sends some message to the executor
func (k *KuberneteExecutor) FrameworkMessage(driver mesos.ExecutorDriver, message string) {
	log.Infof("Receives message from framework %v\n", message)
}

// Shutdown is called when the executor receives a shutdown request.
func (k *KuberneteExecutor) Shutdown(driver mesos.ExecutorDriver) {
	log.Infof("Shutdown the executor\n")

	// TODO(yifan): kill all running containers by updating the pod manifest, and do a sync.
	// Then send back status update.
}

// Error is called when some error happens.
func (k *KuberneteExecutor) Error(driver mesos.ExecutorDriver, message string) {
	log.Errorf("Executor error: %v\n", message)
}
<commit_msg>Enable kuberlete executor to interact with kuberlete.<commit_after>package executor

import (
	"github.com/GoogleCloudPlatform/kubernetes/pkg/api"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/kubelet"
	log "github.com/mesosphere/kubernetes-mesos/3rdparty/github.com/golang/glog"
	"github.com/mesosphere/kubernetes-mesos/3rdparty/github.com/mesosphere/mesos-go/mesos"
)

type kuberTask struct {
	mesosTaskInfo     *mesos.TaskInfo
	containerManifest *api.ContainerManifest
}

// KuberneteExecutor is an mesos executor that runs pods
// in a minion machine.
type KuberneteExecutor struct {
	*kubelet.Kubelet // the kubelet instance.
	registered       bool
	tasks            map[string]*kuberTask
}

// New creates a new kubernete executor.
func New() *KuberneteExecutor {
	return &KuberneteExecutor{
		kubelet.New(),
		false,
		make(map[string]*kuberTask),
	}
}

func (k *KuberneteExecutor) gatherContainerManifests() []api.ContainerManifest {
	var manifests []api.ContainerManifest
	for _, task := range k.tasks {
		manifests = append(manifests, *task.containerManifest)
	}
	return manifests
}

// Registered is called when the executor is successfully registered with the slave.
func (k *KuberneteExecutor) Registered(driver mesos.ExecutorDriver,
	executorInfo *mesos.ExecutorInfo, frameworkInfo *mesos.FrameworkInfo, slaveInfo *mesos.SlaveInfo) {
	log.Infof("Executor %v of framework %v registered with slave %v\n",
		executorInfo, frameworkInfo, slaveInfo)

	k.registered = true
}

// Reregistered is called when the executor is successfully re-registered with the slave.
// This can happen when the slave fails over.
func (k *KuberneteExecutor) Reregistered(driver mesos.ExecutorDriver,
	executorInfo *mesos.ExecutorInfo, frameworkInfo *mesos.FrameworkInfo, slaveInfo *mesos.SlaveInfo) {
	log.Infof("Executor %v of framework %v reregistered with slave %v\n",
		executorInfo, frameworkInfo, slaveInfo)

	k.registered = true
}

// Disconnected is called when the executor is disconnected with the slave.
func (k *KuberneteExecutor) Disconnected(driver mesos.ExecutorDriver) {
	log.Infof("Slave is disconnected\n")

	k.registered = false
}

// LaunchTask is called when the executor receives a request to launch a task.
func (k *KuberneteExecutor) LaunchTask(driver mesos.ExecutorDriver, taskInfo *mesos.TaskInfo) {
	log.Infof("Launch task %v\n", taskInfo)

	if !k.registered {
		// TODO(yifan): Send back status update task lost.
		return
	}

	taskId := taskInfo.GetTaskId().GetValue()
	if _, found := k.tasks[taskId]; found {
		log.Warningf("Task already launched\n")
		// Not to send back TASK_RUNNING here, because
		// may be duplicated messages or duplicated task id.
		return
	}

	// Get the container manifest from the taskInfo.
	var manifest api.ContainerManifest
	if err := k.ExtractYAMLData(taskInfo.GetData(), &manifest); err != nil {
		log.Warningf("Failed to extract yaml data from the taskInfo.data %v\n", err)

		// TODO(yifan): Send back status update task lost.
		return
	}

	// Add the task.
	k.tasks[taskId] = &kuberTask{
		mesosTaskInfo:     taskInfo,
		containerManifest: &manifest,
	}

	// Call SyncManifests to force kuberlete to run docker containers.
	if err := k.SyncManifests(k.gatherContainerManifests()); err != nil {
		log.Warningf("Failed to extract yaml data from the taskInfo.data %v\n", err)

		// TODO(yifan): Send back status update task lost.
		return
	}

	// TODO(yifan): Send back TASK_RUNNING.
}

// KillTask is called when the executor receives a request to kill a task.
func (k *KuberneteExecutor) KillTask(driver mesos.ExecutorDriver, taskId *mesos.TaskID) {
	log.Infof("Kill task %v\n", taskId)

	tid := taskId.GetValue()
	if _, ok := k.tasks[tid]; !ok {
		log.Infof("Failed to kill task, unknown task %v\n", tid)
		return
	}
	delete(k.tasks, tid)

	// Call SyncManifests to force kuberlete to kill deleted containers.
	if err := k.SyncManifests(k.gatherContainerManifests()); err != nil {
		log.Warningf("Failed to extract yaml data from the taskInfo.data %v\n", err)

		// TODO(yifan): Check the error type, (if already killed?)
		return
	}

	// TODO(yifan): Send back TASK_LOST.
}

// FrameworkMessage is called when the framework sends some message to the executor
func (k *KuberneteExecutor) FrameworkMessage(driver mesos.ExecutorDriver, message string) {
	log.Infof("Receives message from framework %v\n", message)
}

// Shutdown is called when the executor receives a shutdown request.
func (k *KuberneteExecutor) Shutdown(driver mesos.ExecutorDriver) {
	log.Infof("Shutdown the executor\n")

	for tid := range k.tasks {
		delete(k.tasks, tid)
		// Call SyncManifests to force kuberlete to kill deleted containers.
		if err := k.SyncManifests(k.gatherContainerManifests()); err != nil {
			log.Warningf("Failed to extract yaml data from the taskInfo.data %v\n", err)

			// TODO(yifan): Check the error type, (if already killed?)
		}
	}
	// Then send back TASK_LOST.
}

// Error is called when some error happens.
func (k *KuberneteExecutor) Error(driver mesos.ExecutorDriver, message string) {
	log.Errorf("Executor error: %v\n", message)
}
<|endoftext|>