<commit_before>package render

import (
	"html/template"
	"io"
	"os"
	"path/filepath"
	"sort"
	"strings"

	"github.com/pkg/errors"
	"github.com/sirupsen/logrus"

	// this blank import is here because dep doesn't
	// handle transitive dependencies correctly
	_ "github.com/russross/blackfriday"
)

type templateRenderer struct {
	*Engine
	contentType string
	names       []string
	data        Data
}

func (s templateRenderer) ContentType() string {
	return s.contentType
}

func (s *templateRenderer) Render(w io.Writer, data Data) error {
	s.data = data
	var body template.HTML
	var err error
	for _, name := range s.names {
		body, err = s.exec(name, data)
		if err != nil {
			return err
		}
		data["yield"] = body
	}
	w.Write([]byte(body))
	return nil
}

func (s templateRenderer) partial(name string, dd Data) (template.HTML, error) {
	d, f := filepath.Split(name)
	name = filepath.Join(d, "_"+f)
	m := Data{}
	for k, v := range s.data {
		m[k] = v
	}
	for k, v := range dd {
		m[k] = v
	}
	return s.exec(name, m)
}

func (s templateRenderer) exec(name string, data Data) (template.HTML, error) {
	ct := strings.ToLower(s.contentType)
	data["contentType"] = ct

	if filepath.Ext(name) == "" {
		switch {
		case strings.Contains(ct, "html"):
			name += ".html"
		case strings.Contains(ct, "javascript"):
			name += ".js"
		case strings.Contains(ct, "markdown"):
			name += ".md"
		}
	}

	// Try to use localized version
	templateName := name
	if languages, ok := data["languages"].([]string); ok {
		ll := len(languages)
		if ll > 0 {
			// Default language is the last in the list
			defaultLanguage := languages[ll-1]
			ext := filepath.Ext(name)
			rawName := strings.TrimSuffix(name, ext)

			for _, l := range languages {
				var candidateName string
				if l == defaultLanguage {
					break
				}
				candidateName = rawName + "." + strings.ToLower(l) + ext
				if s.TemplatesBox.Has(candidateName) {
					// Replace name with the existing suffixed version
					templateName = candidateName
					break
				}
			}
		}
	}

	// Set current_template to context
	if _, ok := data["current_template"]; !ok {
		data["current_template"] = templateName
	}

	source, err := s.TemplatesBox.MustBytes(templateName)
	if err != nil {
		return "", err
	}

	helpers := map[string]interface{}{
		"partial": s.partial,
	}

	helpers = s.addAssetsHelpers(helpers)

	for k, v := range s.Helpers {
		helpers[k] = v
	}

	body := string(source)
	for _, ext := range s.exts(name) {
		te, ok := s.TemplateEngines[ext]
		if !ok {
			logrus.Errorf("could not find a template engine for %s\n", ext)
			continue
		}
		body, err = te(body, data, helpers)
		if err != nil {
			return "", errors.Wrap(err, name)
		}
	}

	return template.HTML(body), nil
}

func (s templateRenderer) exts(name string) []string {
	exts := []string{}
	for {
		ext := filepath.Ext(name)
		if ext == "" {
			break
		}
		name = strings.TrimSuffix(name, ext)
		exts = append(exts, strings.ToLower(ext[1:]))
	}
	if len(exts) == 0 {
		return []string{"html"}
	}
	sort.Sort(sort.Reverse(sort.StringSlice(exts)))
	return exts
}

func (s templateRenderer) assetPath(file string) (string, error) {

	if len(assetMap) == 0 || os.Getenv("GO_ENV") != "production" {
		manifest, err := s.AssetsBox.MustString("manifest.json")

		if err != nil {
			manifest, err = s.AssetsBox.MustString("assets/manifest.json")
			if err != nil {
				return assetPathFor(file), nil
			}
		}

		err = loadManifest(manifest)
		if err != nil {
			return assetPathFor(file), errors.Wrap(err, "your manifest.json is not correct")
		}
	}

	return assetPathFor(file), nil
}

// Template renders the named files using the specified
// content type and the github.com/gobuffalo/plush
// package for templating. If more than 1 file is provided
// the second file will be considered a "layout" file
// and the first file will be the "content" file which will
// be placed into the "layout" using "{{yield}}".
func Template(c string, names ...string) Renderer {
	e := New(Options{})
	return e.Template(c, names...)
}

// Template renders the named files using the specified
// content type and the github.com/gobuffalo/plush
// package for templating. If more than 1 file is provided
// the second file will be considered a "layout" file
// and the first file will be the "content" file which will
// be placed into the "layout" using "{{yield}}".
func (e *Engine) Template(c string, names ...string) Renderer {
	return &templateRenderer{
		Engine:      e,
		contentType: c,
		names:       names,
	}
}
<commit_msg>Remove logger need from the render package (#1186)<commit_after>package render

import (
	"fmt"
	"html/template"
	"io"
	"os"
	"path/filepath"
	"sort"
	"strings"

	"github.com/pkg/errors"

	// this blank import is here because dep doesn't
	// handle transitive dependencies correctly
	_ "github.com/russross/blackfriday"
)

type templateRenderer struct {
	*Engine
	contentType string
	names       []string
	data        Data
}

func (s templateRenderer) ContentType() string {
	return s.contentType
}

func (s *templateRenderer) Render(w io.Writer, data Data) error {
	s.data = data
	var body template.HTML
	var err error
	for _, name := range s.names {
		body, err = s.exec(name, data)
		if err != nil {
			return err
		}
		data["yield"] = body
	}
	w.Write([]byte(body))
	return nil
}

func (s templateRenderer) partial(name string, dd Data) (template.HTML, error) {
	d, f := filepath.Split(name)
	name = filepath.Join(d, "_"+f)
	m := Data{}
	for k, v := range s.data {
		m[k] = v
	}
	for k, v := range dd {
		m[k] = v
	}
	return s.exec(name, m)
}

func (s templateRenderer) exec(name string, data Data) (template.HTML, error) {
	ct := strings.ToLower(s.contentType)
	data["contentType"] = ct

	if filepath.Ext(name) == "" {
		switch {
		case strings.Contains(ct, "html"):
			name += ".html"
		case strings.Contains(ct, "javascript"):
			name += ".js"
		case strings.Contains(ct, "markdown"):
			name += ".md"
		}
	}

	// Try to use localized version
	templateName := name
	if languages, ok := data["languages"].([]string); ok {
		ll := len(languages)
		if ll > 0 {
			// Default language is the last in the list
			defaultLanguage := languages[ll-1]
			ext := filepath.Ext(name)
			rawName := strings.TrimSuffix(name, ext)

			for _, l := range languages {
				var candidateName string
				if l == defaultLanguage {
					break
				}
				candidateName = rawName + "." + strings.ToLower(l) + ext
				if s.TemplatesBox.Has(candidateName) {
					// Replace name with the existing suffixed version
					templateName = candidateName
					break
				}
			}
		}
	}

	// Set current_template to context
	if _, ok := data["current_template"]; !ok {
		data["current_template"] = templateName
	}

	source, err := s.TemplatesBox.MustBytes(templateName)
	if err != nil {
		return "", err
	}

	helpers := map[string]interface{}{
		"partial": s.partial,
	}

	helpers = s.addAssetsHelpers(helpers)

	for k, v := range s.Helpers {
		helpers[k] = v
	}

	body := string(source)
	for _, ext := range s.exts(name) {
		te, ok := s.TemplateEngines[ext]
		if !ok {
			return "", fmt.Errorf("could not find a template engine for %s", ext)
		}
		body, err = te(body, data, helpers)
		if err != nil {
			return "", errors.Wrap(err, name)
		}
	}

	return template.HTML(body), nil
}

func (s templateRenderer) exts(name string) []string {
	exts := []string{}
	for {
		ext := filepath.Ext(name)
		if ext == "" {
			break
		}
		name = strings.TrimSuffix(name, ext)
		exts = append(exts, strings.ToLower(ext[1:]))
	}
	if len(exts) == 0 {
		return []string{"html"}
	}
	sort.Sort(sort.Reverse(sort.StringSlice(exts)))
	return exts
}

func (s templateRenderer) assetPath(file string) (string, error) {

	if len(assetMap) == 0 || os.Getenv("GO_ENV") != "production" {
		manifest, err := s.AssetsBox.MustString("manifest.json")

		if err != nil {
			manifest, err = s.AssetsBox.MustString("assets/manifest.json")
			if err != nil {
				return assetPathFor(file), nil
			}
		}

		err = loadManifest(manifest)
		if err != nil {
			return assetPathFor(file), errors.Wrap(err, "your manifest.json is not correct")
		}
	}

	return assetPathFor(file), nil
}

// Template renders the named files using the specified
// content type and the github.com/gobuffalo/plush
// package for templating. If more than 1 file is provided
// the second file will be considered a "layout" file
// and the first file will be the "content" file which will
// be placed into the "layout" using "{{yield}}".
func Template(c string, names ...string) Renderer {
	e := New(Options{})
	return e.Template(c, names...)
}

// Template renders the named files using the specified
// content type and the github.com/gobuffalo/plush
// package for templating. If more than 1 file is provided
// the second file will be considered a "layout" file
// and the first file will be the "content" file which will
// be placed into the "layout" using "{{yield}}".
func (e *Engine) Template(c string, names ...string) Renderer {
	return &templateRenderer{
		Engine:      e,
		contentType: c,
		names:       names,
	}
}
<|endoftext|>


<commit_before>package cmd

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"
	"text/tabwriter"

	"github.com/mobingi/mobingi-cli/pkg/cli"
	"github.com/mobingilabs/mobingi-sdk-go/mobingi/alm"
	"github.com/mobingilabs/mobingi-sdk-go/pkg/cmdline"
	"github.com/mobingilabs/mobingi-sdk-go/pkg/pretty"
	"github.com/spf13/cobra"
)

func StackDescribeCmd() *cobra.Command {
	cmd := &cobra.Command{
		Use:   "describe",
		Short: "display stack details",
		Long: `Display stack details. If you specify the '--out=[filename]' option,
make sure you provide the full path of the file. If the path has
space(s) in it, make sure to surround it with double quotes.

Valid format values: min (default), json, raw

Examples:

  $ ` + cmdline.Args0() + ` stack describe --id=58c2297d25645-Y6NSE4VjP-tk
  $ ` + cmdline.Args0() + ` stack describe --id=58c2297d25645-Y6NSE4VjP-tk --fmt=json`,
		Run: describe,
	}

	cmd.Flags().String("id", "", "stack id")
	return cmd
}

func describe(cmd *cobra.Command, args []string) {
	sess, err := clisession()
	cli.ErrorExit(err, 1)

	svc := alm.New(sess)
	in := &alm.StackDescribeInput{
		StackId: cli.GetCliStringFlag(cmd, "id"),
	}

	resp, body, err := svc.Describe(in)
	cli.ErrorExit(err, 1)
	exitOn401(resp)

	// we process `--fmt=raw` option first
	out := cli.GetCliStringFlag(cmd, "out")
	pfmt := cli.GetCliStringFlag(cmd, "fmt")
	/*
		if sess.Config.ApiVersion == 3 {
			if pfmt == "min" || pfmt == "" {
				pfmt = "json"
			}
		}
	*/

	switch pfmt {
	case "raw":
		fmt.Println(string(body))
		if out != "" {
			err = ioutil.WriteFile(out, body, 0644)
			cli.ErrorExit(err, 1)
		}
	case "json":
		indent := cli.GetCliIntFlag(cmd, "indent")
		js := pretty.JSON(string(body), indent)
		fmt.Println(js)

		// write to file option
		if out != "" {
			err = ioutil.WriteFile(out, []byte(js), 0644)
			cli.ErrorExit(err, 1)
		}
	default:
		if pfmt == "min" || pfmt == "" {
			/*
				if sess.Config.ApiVersion >= 3 {
					err = v3DescribeStack(cmd, body)
					cli.ErrorExit(err, 1)
					return
				}
			*/

			type Configuration struct {
				// v3
				Description string          `json:"description,omitempty"`
				Label       string          `json:"label,omitempty"`
				Version     string          `json:"version,omitempty"`
				Vendor      json.RawMessage `json:"vendor,omitempty"`
				// v2
				AWS                 string      `json:"AWS,omitempty"`
				AWSAccountName      string      `json:"AWS_ACCOUNT_NAME,omitempty"`
				AssociatePublicIp   string      `json:"AssociatePublicIP,omitempty"`
				ELBOpen443Port      string      `json:"ELBOpen443Port,omitempty"`
				ELBOpen80Port       string      `json:"ELBOpen80Port,omitempty"`
				SpotInstanceMaxSize int         `json:"SpotInstanceMaxSize,omitempty"`
				SpotInstanceMinSize int         `json:"SpotInstanceMinSize,omitempty"`
				SpotPrice           string      `json:"SpotPrice,omitempty"`
				Architecture        string      `json:"architecture,omitempty"`
				Code                string      `json:"code,omitempty"`
				Image               string      `json:"image,omitempty"`
				Max                 interface{} `json:"max,omitempty"`
				MaxOrigin           interface{} `json:"maxOrigin,omitempty"`
				Min                 interface{} `json:"min,omitempty"`
				MinOrigin           interface{} `json:"minOrigin,omitempty"`
				Nickname            string      `json:"nickname,omitempty"`
				Region              string      `json:"region,omitempty"`
				Type                string      `json:"type,omitempty"`
			}

			type State struct {
				Code string `json:"Code,omitempty"`
				Name string `json:"Name,omitempty"`
			}

			type Instance struct {
				AmiLaunchIndex        string      `json:"AmiLaunchIndex,omitempty"`
				Architecture          string      `json:"Architecture,omitempty"`
				BlockDeviceMappings   interface{} `json:"BlockDeviceMappings,omitempty"`
				ClientToken           string      `json:"ClientToken,omitempty"`
				EbsOptimized          bool        `json:"EbsOptimized,omitempty"`
				Hypervisor            string      `json:"Hypervisor,omitempty"`
				ImageId               string      `json:"ImageId,omitempty"`
				InstanceId            string      `json:"InstanceId,omitempty"`
				InstanceType          string      `json:"InstanceType,omitempty"`
				InstanceLifecycle     string      `json:"InstanceLifecycle,omitempty"`
				SpotInstanceRequestId string      `json:"SpotInstanceRequestId,omitempty"`
				KeyName               string      `json:"KeyName,omitempty"`
				LaunchTime            string      `json:"LaunchTime,omitempty"`
				Monitoring            interface{} `json:"Monitoring,omitempty"`
				NetworkInterfaces     interface{} `json:"NetworkInterfaces,omitempty"`
				Placement             interface{} `json:"Placement,omitempty"`
				PrivateDnsName        string      `json:"PrivateDnsName,omitempty"`
				PrivateIpAddress      string      `json:"PrivateIpAddress,omitempty"`
				ProductCodes          []string    `json:"ProductCodes,omitempty"`
				PublicDnsName         string      `json:"PublicDnsName,omitempty"`
				PublicIpAddress       string      `json:"PublicIpAddress,omitempty"`
				Reservation           interface{} `json:"Reservation,omitempty"`
				RootDeviceName        string      `json:"RootDeviceName,omitempty"`
				RootDeviceType        string      `json:"RootDeviceType,omitempty"`
				SecurityGroups        interface{} `json:"SecurityGroups,omitempty"`
				SourceDestCheck       bool        `json:"SourceDestCheck,omitempty"`
				State                 State       `json:"State,omitempty"`
				StateTransitionReason string      `json:"StateTransitionReason,omitempty"`
				SubnetId              string      `json:"SubnetId,omitempty"`
				Tags                  interface{} `json:"Tags,omitempty"`
				VirtualizationType    string      `json:"VirtualizationType,omitempty"`
				VpcId                 string      `json:"VpcId,omitempty"`
				EnaSupport            string      `json:"enaSupport,omitempty"`
			}

			type DescribeStack struct {
				AuthToken     string        `json:"auth_token,omitempty"`
				Configuration Configuration `json:"configuration,omitempty"`
				CreateTime    string        `json:"create_time,omitempty"`
				Instances     []Instance    `json:"Instances,omitempty"`
				Nickname      string        `json:"nickname,omitempty"`
				StackId       string        `json:"stack_id,omitempty"`
				StackOutputs  interface{}   `json:"stack_outputs,omitempty"`
				StackStatus   string        `json:"stack_status,omitempty"`
				UserId        string        `json:"user_id,omitempty"`
			}

			var stacks []alm.DescribeStack
			var stack alm.DescribeStack

			switch sess.Config.ApiVersion {
			case 3:
				err = json.Unmarshal(body, &stack)
				cli.ErrorExit(err, 1)
			default:
				err = json.Unmarshal(body, &stacks)
				cli.ErrorExit(err, 1)
				stack = stacks[0]
			}

			w := tabwriter.NewWriter(os.Stdout, 0, 10, 5, ' ', 0)
			fmt.Fprintf(w, "INSTANCE ID\tINSTANCE TYPE\tINSTANCE MODEL\tPUBLIC IP\tPRIVATE IP\tSTATUS\n")
			for _, inst := range stack.Instances {
				instype := "on-demand"
				if inst.InstanceLifecycle == "spot" {
					instype = inst.InstanceLifecycle
				}

				fmt.Fprintf(w, "%s\t%s\t%s\t%s\t%s\t%s\n",
					inst.InstanceId,
					instype,
					inst.InstanceType,
					inst.PublicIpAddress,
					inst.PrivateIpAddress,
					inst.State.Name)
			}

			w.Flush()
		}
	}
}

func v3DescribeStack(cmd *cobra.Command, body []byte) error {
	return nil
}
<commit_msg>Cleanup.<commit_after>package cmd

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"
	"text/tabwriter"

	"github.com/mobingi/mobingi-cli/pkg/cli"
	"github.com/mobingilabs/mobingi-sdk-go/mobingi/alm"
	"github.com/mobingilabs/mobingi-sdk-go/pkg/cmdline"
	"github.com/mobingilabs/mobingi-sdk-go/pkg/pretty"
	"github.com/spf13/cobra"
)

func StackDescribeCmd() *cobra.Command {
	cmd := &cobra.Command{
		Use:   "describe",
		Short: "display stack details",
		Long: `Display stack details. If you specify the '--out=[filename]' option,
make sure you provide the full path of the file. If the path has
space(s) in it, make sure to surround it with double quotes.

Valid format values: min (default), json, raw

Examples:

  $ ` + cmdline.Args0() + ` stack describe --id=58c2297d25645-Y6NSE4VjP-tk
  $ ` + cmdline.Args0() + ` stack describe --id=58c2297d25645-Y6NSE4VjP-tk --fmt=json`,
		Run: describe,
	}

	cmd.Flags().String("id", "", "stack id")
	return cmd
}

func describe(cmd *cobra.Command, args []string) {
	sess, err := clisession()
	cli.ErrorExit(err, 1)

	svc := alm.New(sess)
	in := &alm.StackDescribeInput{
		StackId: cli.GetCliStringFlag(cmd, "id"),
	}

	resp, body, err := svc.Describe(in)
	cli.ErrorExit(err, 1)
	exitOn401(resp)

	// we process `--fmt=raw` option first
	out := cli.GetCliStringFlag(cmd, "out")
	pfmt := cli.GetCliStringFlag(cmd, "fmt")
	/*
		if sess.Config.ApiVersion == 3 {
			if pfmt == "min" || pfmt == "" {
				pfmt = "json"
			}
		}
	*/

	switch pfmt {
	case "raw":
		fmt.Println(string(body))
		if out != "" {
			err = ioutil.WriteFile(out, body, 0644)
			cli.ErrorExit(err, 1)
		}
	case "json":
		indent := cli.GetCliIntFlag(cmd, "indent")
		js := pretty.JSON(string(body), indent)
		fmt.Println(js)

		// write to file option
		if out != "" {
			err = ioutil.WriteFile(out, []byte(js), 0644)
			cli.ErrorExit(err, 1)
		}
	default:
		if pfmt == "min" || pfmt == "" {
			var stacks []alm.DescribeStack
			var stack alm.DescribeStack

			switch sess.Config.ApiVersion {
			case 3:
				err = json.Unmarshal(body, &stack)
				cli.ErrorExit(err, 1)
			default:
				err = json.Unmarshal(body, &stacks)
				cli.ErrorExit(err, 1)
				stack = stacks[0]
			}

			w := tabwriter.NewWriter(os.Stdout, 0, 10, 5, ' ', 0)
			fmt.Fprintf(w, "INSTANCE ID\tINSTANCE TYPE\tINSTANCE MODEL\tPUBLIC IP\tPRIVATE IP\tSTATUS\n")
			for _, inst := range stack.Instances {
				instype := "on-demand"
				if inst.InstanceLifecycle == "spot" {
					instype = inst.InstanceLifecycle
				}

				fmt.Fprintf(w, "%s\t%s\t%s\t%s\t%s\t%s\n",
					inst.InstanceId,
					instype,
					inst.InstanceType,
					inst.PublicIpAddress,
					inst.PrivateIpAddress,
					inst.State.Name)
			}

			w.Flush()
		}
	}
}

func v3DescribeStack(cmd *cobra.Command, body []byte) error {
	return nil
}
<|endoftext|>


<commit_before>package player

import (
	"fmt"
	"log"
	"os"
	"os/exec"
	"strconv"
	"time"

	"github.com/lhz/considerate/hvsc"
	//"io/ioutil"
)

const (
	PLAY_COMMAND = iota
	STOP_COMMAND
	QUIT_COMMAND
)

type PlayerMsg struct {
	Command int
	Args    []string
}

var CurrentTune *hvsc.SidTune
var CurrentSong int
var StartTime time.Time
var MsgChan chan PlayerMsg
var Playing = false

func Run() {
	MsgChan = make(chan PlayerMsg)

	var playCmd *exec.Cmd

	for {
		select {
		case msg := <-MsgChan:
			switch msg.Command {
			case PLAY_COMMAND:
				playCmd = exec.Command("/usr/bin/sidplay2", "-o"+msg.Args[1], msg.Args[0])
				playCmd.Stdout = os.Stdout
				if err := playCmd.Start(); err != nil {
					log.Print("Failed to start player process: ", err)
				}
				StartTime = time.Now()
			case STOP_COMMAND:
				stopCommand(playCmd)
			case QUIT_COMMAND:
				stopCommand(playCmd)
				return
			}
		}
	}
}

func PlaySub(subTune int) {
	Stop()
	Playing = true
	CurrentSong = subTune
	MsgChan <- PlayerMsg{Command: PLAY_COMMAND, Args: []string{CurrentTune.FullPath(), strconv.Itoa(subTune)}}
}

func Play(index, subTune int) {
	Stop()
	tune := hvsc.FilteredTunes[index]
	CurrentTune = &tune
	CurrentSong = subTune
	Playing = true
	MsgChan <- PlayerMsg{Command: PLAY_COMMAND, Args: []string{tune.FullPath(), strconv.Itoa(subTune)}}
}

func PrevSong() {
	if Playing && CurrentSong > 1 {
		PlaySub(CurrentSong - 1)
	}
}

func NextSong() {
	if Playing && CurrentSong < CurrentTune.Header.Songs {
		PlaySub(CurrentSong + 1)
	}
}

func Stop() {
	Playing = false
	MsgChan <- PlayerMsg{Command: STOP_COMMAND, Args: []string{}}
}

func Quit() {
	Playing = false
	MsgChan <- PlayerMsg{Command: QUIT_COMMAND, Args: []string{}}
}

func Elapsed() string {
	return TimeFormat(time.Since(StartTime))
}

func SongLength() string {
	if CurrentTune == nil || CurrentSong < 1 {
		return ""
	}
	return TimeFormat(CurrentTune.SongLengths[CurrentSong-1])
}

func TimeFormat(duration time.Duration) string {
	seconds := int(duration.Seconds())
	return fmt.Sprintf("%02d:%02d", seconds/60, seconds%60)
}

func stopCommand(cmd *exec.Cmd) {
	if cmd != nil {
		cmd.Process.Signal(os.Interrupt)
		cmd.Wait()
	}
}
<commit_msg>Set StartTime when calling exec.<commit_after>package player

import (
	"fmt"
	"log"
	"os"
	"os/exec"
	"strconv"
	"time"

	"github.com/lhz/considerate/hvsc"
)

const (
	PLAY_COMMAND = iota
	STOP_COMMAND
	QUIT_COMMAND
)

type PlayerMsg struct {
	Command int
	Args    []string
}

var CurrentTune *hvsc.SidTune
var CurrentSong int
var StartTime time.Time
var MsgChan chan PlayerMsg
var Playing = false

func Run() {
	MsgChan = make(chan PlayerMsg)

	var playCmd *exec.Cmd

	for {
		select {
		case msg := <-MsgChan:
			switch msg.Command {
			case PLAY_COMMAND:
				StartTime = time.Now()
				playCmd = exec.Command("/usr/bin/sidplay2", "-o"+msg.Args[1], msg.Args[0])
				playCmd.Stdout = os.Stdout
				if err := playCmd.Start(); err != nil {
					log.Print("Failed to start player process: ", err)
				}
			case STOP_COMMAND:
				stopCommand(playCmd)
			case QUIT_COMMAND:
				stopCommand(playCmd)
				return
			}
		}
	}
}

func PlaySub(subTune int) {
	Stop()
	Playing = true
	CurrentSong = subTune
	MsgChan <- PlayerMsg{Command: PLAY_COMMAND, Args: []string{CurrentTune.FullPath(), strconv.Itoa(subTune)}}
}

func Play(index, subTune int) {
	Stop()
	tune := hvsc.FilteredTunes[index]
	CurrentTune = &tune
	CurrentSong = subTune
	Playing = true
	MsgChan <- PlayerMsg{Command: PLAY_COMMAND, Args: []string{tune.FullPath(), strconv.Itoa(subTune)}}
}

func PrevSong() {
	if Playing && CurrentSong > 1 {
		PlaySub(CurrentSong - 1)
	}
}

func NextSong() {
	if Playing && CurrentSong < CurrentTune.Header.Songs {
		PlaySub(CurrentSong + 1)
	}
}

func Stop() {
	Playing = false
	MsgChan <- PlayerMsg{Command: STOP_COMMAND, Args: []string{}}
}

func Quit() {
	Playing = false
	MsgChan <- PlayerMsg{Command: QUIT_COMMAND, Args: []string{}}
}

func Elapsed() string {
	return TimeFormat(time.Since(StartTime))
}

func SongLength() string {
	if CurrentTune == nil || CurrentSong < 1 {
		return ""
	}
	return TimeFormat(CurrentTune.SongLengths[CurrentSong-1])
}

func TimeFormat(duration time.Duration) string {
	seconds := int(duration.Seconds())
	return fmt.Sprintf("%02d:%02d", seconds/60, seconds%60)
}

func stopCommand(cmd *exec.Cmd) {
	if cmd != nil {
		cmd.Process.Signal(os.Interrupt)
		cmd.Wait()
	}
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"encoding/hex"
	"errors"
	"fmt"
	"log"
	"net"
	"net/http"
	"os"
	"os/signal"
	"sort"
	"sync"
	"time"

	"github.com/alexbakker/tox4go/bootstrap"
	"github.com/alexbakker/tox4go/crypto"
	"github.com/alexbakker/tox4go/dht"
	"github.com/alexbakker/tox4go/dht/ping"
	"github.com/alexbakker/tox4go/relay"
	"github.com/alexbakker/tox4go/transport"
	"github.com/didip/tollbooth"
)

const (
	enableIpv6  = true
	probeRate   = 1 * time.Minute
	refreshRate = 5 * time.Minute
)

var (
	lastScan    int64
	lastRefresh int64
	nodes       = []*toxNode{}
	nodesMutex  = sync.Mutex{}
	tcpPorts    = []int{443, 3389, 33445}
)

func main() {
	if parseFlags() {
		return
	}

	// load state if available
	state, err := loadState()
	if err != nil {
		log.Fatalf("error loading state: %s", err.Error())
	}
	lastScan = state.LastScan
	lastRefresh = state.LastRefresh
	nodes = state.Nodes

	if err := loadCountries(); err != nil {
		log.Fatalf("error loading countries.json: %s", err.Error())
	}

	inst, err := NewInstance(":33450")
	if err != nil {
		log.Fatalf("fatal: %s", err.Error())
	}
	inst.UDPTransport.Handle(dht.PacketIDSendNodes, inst.handleSendNodesPacket)
	inst.UDPTransport.Handle(bootstrap.PacketIDBootstrapInfo, handleBootstrapInfoPacket)

	//handle stop signal
	interruptChan := make(chan os.Signal)
	signal.Notify(interruptChan, os.Interrupt)

	//setup http server
	listener, err := net.Listen("tcp", fmt.Sprintf(":%d", httpListenPort))
	if err != nil {
		log.Fatalf("error in net.Listen: %s", err.Error())
	}
	limiter := tollbooth.NewLimiter(1, 2*time.Second)
	limiter.Methods = []string{"POST"}
	limiter.IPLookups = []string{"X-Forwarded-For", "RemoteAddr", "X-Real-IP"}
	serveMux := http.NewServeMux()
	serveMux.HandleFunc("/", handleHTTPRequest)
	serveMux.Handle("/test", tollbooth.LimitFuncHandler(limiter, handleHTTPRequest))
	serveMux.HandleFunc("/json", handleJSONRequest)
	go func() {
		err := http.Serve(listener, serveMux)
		if err != nil {
			log.Printf("http server error: %s\n", err.Error())
			interruptChan <- os.Interrupt
		}
	}()

	//listen for tox packets
	go func() {
		err := inst.UDPTransport.Listen()
		if err != nil {
			log.Printf("udp transport error: %s\n", err.Error())
			interruptChan <- os.Interrupt
		}
	}()
	//go tcpTransport.Listen()

	err = refreshNodes()
	if err != nil {
		log.Fatal(err.Error())
	}
	inst.probeNodes()

	probeTicker := time.NewTicker(probeRate)
	refreshTicker := time.NewTicker(refreshRate)
	updateTicker := time.NewTicker(30 * time.Second)
	run := true

	for run {
		select {
		case <-interruptChan:
			fmt.Printf("killing routines\n")
			probeTicker.Stop()
			refreshTicker.Stop()
			updateTicker.Stop()
			inst.UDPTransport.Stop()
			//tcpTransport.Stop()
			listener.Close()
			run = false
		case <-probeTicker.C:
			// we want an empty ping list at the start of every probe
			inst.PingsMutex.Lock()
			inst.Pings.Clear(false)
			inst.PingsMutex.Unlock()

			nodesMutex.Lock()
			err := inst.probeNodes()
			nodesMutex.Unlock()
			if err != nil {
				log.Printf("error while trying to probe nodes: %s", err.Error())
			}
		case <-refreshTicker.C:
			err := refreshNodes()
			if err != nil {
				log.Printf("error while trying to refresh nodes: %s", err.Error())
			}
		case <-updateTicker.C:
			inst.PingsMutex.Lock()
			inst.Pings.Clear(true)
			inst.PingsMutex.Unlock()

			nodesMutex.Lock()
			for _, node := range nodes {
				if time.Now().Sub(time.Unix(node.LastPing, 0)) > time.Minute*2 {
					node.UDPStatus = false
				}
			}
			sort.Stable(nodeSlice(nodes))

			state := getState()
			saveState(state)
			nodesMutex.Unlock()
		}
	}
}

func refreshNodes() error {
	parsedNodes, err := parseNodes()
	if err != nil {
		return err
	}

	nodesMutex.Lock()
	for _, freshNode := range parsedNodes {
		for _, node := range nodes {
			if freshNode.PublicKey == node.PublicKey {
				freshNode.LastPing = node.LastPing
				freshNode.UDPStatus = node.UDPStatus
				freshNode.TCPStatus = node.TCPStatus
				freshNode.TCPPorts = node.TCPPorts
				freshNode.MOTD = node.MOTD
				freshNode.Version = node.Version
				break
			}
		}
	}
	nodes = parsedNodes
	sort.Stable(nodeSlice(nodes))
	nodesMutex.Unlock()

	lastRefresh = time.Now().Unix()
	return nil
}

func (i *instance) probeNodes() error {
	for _, node := range nodes {
		err := i.getBootstrapInfo(node)
		if err != nil {
			fmt.Println(err.Error())
		}

		p, err := i.getNodes(node)
		if err != nil {
			fmt.Println(err.Error())
		} else {
			i.PingsMutex.Lock()
			err = i.Pings.Add(p)
			if err != nil {
				fmt.Println(err.Error())
			}
			i.PingsMutex.Unlock()
		}

		ports := tcpPorts
		exists := false
		for _, i := range ports {
			if i == node.Port {
				exists = true
			}
		}
		if !exists {
			ports = append(ports, node.Port)
		}

		go i.probeNodeTCPPorts(node, ports)
	}

	lastScan = time.Now().Unix()
	return nil
}

func (i *instance) probeNodeTCPPorts(node *toxNode, ports []int) {
	c := make(chan int)
	for _, port := range ports {
		go func(p int) {
			conn, err := connectTCP(node, p)
			if err != nil {
				fmt.Printf("%s\n", err.Error())
				c <- -1
				return
			}

			err = i.tcpHandshake(node, conn)
			if err != nil {
				fmt.Printf("%s\n", err.Error())
				c <- -1
			} else {
				c <- p
			}
			conn.Close()
		}(port)
	}

	newPorts := []int{}
	for i := 0; i < len(ports); i++ {
		port := <-c
		if port != -1 {
			fmt.Printf("tcp port for %s: %d\n", node.Maintainer, port)
			newPorts = append(newPorts, port)
		}
	}

	nodesMutex.Lock()
	node.TCPPorts = newPorts
	if len(node.TCPPorts) > 0 {
		node.LastPing = time.Now().Unix()
	}
	node.TCPStatus = len(node.TCPPorts) > 0
	sort.Stable(nodeSlice(nodes))
	nodesMutex.Unlock()
}

func (i *instance) tcpHandshake(node *toxNode, conn *net.TCPConn) error {
	nodePublicKey := new([crypto.PublicKeySize]byte)
	decPublicKey, err := hex.DecodeString(node.PublicKey)
	if err != nil {
		return err
	}
	copy(nodePublicKey[:], decPublicKey)

	relayConn, err := relay.NewConnection()
	if err != nil {
		return err
	}

	req, err := relayConn.StartHandshake()
	if err != nil {
		return err
	}

	reqBytes, err := req.MarshalBinary()
	if err != nil {
		return err
	}

	encryptedReqBytes, nonce, err := i.Ident.EncryptBlob(reqBytes, nodePublicKey)
	if err != nil {
		return err
	}

	reqPacket := &relay.HandshakeRequestPacket{
		PublicKey: i.Ident.PublicKey,
		Nonce:     nonce,
		Payload:   encryptedReqBytes,
	}

	reqPacketBytes, err := reqPacket.MarshalBinary()
	if err != nil {
		return err
	}

	conn.SetDeadline(time.Now().Add(2 * time.Second))
	_, err = conn.Write(reqPacketBytes)
	if err != nil {
		return err
	}

	buffer := make([]byte, 96)
	left := len(buffer)
	for left > 0 {
		read, readErr := conn.Read(buffer[len(buffer)-left:])
		if readErr != nil {
			return readErr
		}
		left -= read
	}

	res := relay.HandshakeResponsePacket{}
	err = res.UnmarshalBinary(buffer)
	if err != nil {
		return err
	}

	decryptedBytes, err := i.Ident.DecryptBlob(res.Payload, nodePublicKey, res.Nonce)
	if err != nil {
		return err
	}

	resPacket := &relay.HandshakePayload{}
	err = resPacket.UnmarshalBinary(decryptedBytes)
	if err != nil {
		return err
	}

	return relayConn.EndHandshake(resPacket)
}

func (i *instance) getNodes(node *toxNode) (*ping.Ping, error) {
	nodePublicKey := new([crypto.PublicKeySize]byte)
	decPublicKey, err := hex.DecodeString(node.PublicKey)
	if err != nil {
		return nil, err
	}
	copy(nodePublicKey[:], decPublicKey)

	p, err := ping.NewPing(nodePublicKey)
	if err != nil {
		return nil, err
	}

	packet := &dht.GetNodesPacket{
		PublicKey: i.Ident.PublicKey,
		PingID:    p.ID,
	}

	dhtPacket, err := i.Ident.EncryptPacket(transport.Packet(packet), nodePublicKey)
	if err != nil {
		return nil, err
	}

	payload, err := dhtPacket.MarshalBinary()
	if err != nil {
		return nil, err
	}

	if err := i.sendToUDP(payload, node); err != nil {
		return nil, err
	}

	return p, nil
}

func (i *instance) getBootstrapInfo(node *toxNode) error {
	packet, err := bootstrap.ConstructPacket(&bootstrap.InfoRequestPacket{})
	if err != nil {
		return err
	}

	payload, err := packet.MarshalBinary()
	if err != nil {
		return err
	}

	return i.sendToUDP(payload, node)
}

func (i *instance) sendToUDP(data []byte, node *toxNode) error {
	ip, err := getNodeIP(node)
	if err != nil {
		return err
	}

	return i.UDPTransport.Send(
		&transport.Message{
			Data: data,
			Addr: &net.UDPAddr{
				IP:   ip,
				Port: node.Port,
			},
		},
	)
}

func getNodeIP(node *toxNode) (net.IP, error) {
	if node.ip4 != nil {
		return node.ip4, nil
	} else if enableIpv6 && node.ip6 != nil {
		return node.ip6, nil
	}

	return nil, fmt.Errorf("no valid ip found for %s", node.Maintainer)
}

func connectTCP(node *toxNode, port int) (*net.TCPConn, error) {
	ip, err := getNodeIP(node)
	if err != nil {
		return nil, err
	}

	dialer := net.Dialer{}
	dialer.Deadline = time.Now().Add(2 * time.Second)

	tempConn, err := dialer.Dial("tcp", fmt.Sprintf("%s:%d", ip, port))
	if err != nil {
		return nil, err
	}

	conn, ok := tempConn.(*net.TCPConn)
	if !ok {
		return nil, errors.New("not a tcp conn")
	}

	return conn, nil
}

func (i *instance) handleSendNodesPacket(msg *transport.Message) error {
	dhtPacket := &dht.Packet{}
	err := dhtPacket.UnmarshalBinary(msg.Data)
	if err != nil {
		return err
	}

	decryptedPacket, err := i.Ident.DecryptPacket(dhtPacket)
	if err != nil {
		return err
	}

	packet, ok := decryptedPacket.(*dht.SendNodesPacket)
	if !ok {
		return nil
	}

	i.PingsMutex.Lock()
	ping := i.Pings.Find(dhtPacket.SenderPublicKey, packet.PingID, true)
	i.PingsMutex.Unlock()

	if ping == nil {
		err := fmt.Errorf("sendnodes packet from unknown node: %s:%d", msg.Addr.IP, msg.Addr.Port)
		fmt.Printf("error: %s", err.Error())
		return err
	}

	nodesMutex.Lock()
	for _, node := range nodes {
		publicKey, err := hex.DecodeString(node.PublicKey)
		if err != nil {
			continue
		}

		if bytes.Equal(publicKey, dhtPacket.SenderPublicKey[:]) {
			node.UDPStatus = true
			node.LastPing = time.Now().Unix()
			break
		}
	}
	sort.Stable(nodeSlice(nodes))
	nodesMutex.Unlock()

	return nil
}

func handleBootstrapInfoPacket(msg *transport.Message) error {
	bootstrapPacket := &bootstrap.Packet{}
	err := bootstrapPacket.UnmarshalBinary(msg.Data)
	if err != nil {
		return err
	}

	transPacket, err := bootstrap.DestructPacket(bootstrapPacket)
	if err != nil {
		return err
	}

	packet, ok := transPacket.(*bootstrap.InfoResponsePacket)
	if !ok {
		return errors.New("wtf")
	}

	nodesMutex.Lock()
	for _, node := range nodes {
		if (node.ip4 != nil && node.ip4.Equal(msg.Addr.IP)) ||
			(node.ip6 != nil && node.ip6.Equal(msg.Addr.IP)) {
			node.MOTD = packet.MOTD
			node.Version = fmt.Sprintf("%d", packet.Version)
			break
		}
	}
	nodesMutex.Unlock()

	return nil
}
<commit_msg>Only listen for IPv4 connections from localhost<commit_after>package main

import (
	"bytes"
	"encoding/hex"
	"errors"
	"fmt"
	"log"
	"net"
	"net/http"
	"os"
	"os/signal"
	"sort"
	"sync"
	"time"

	"github.com/alexbakker/tox4go/bootstrap"
	"github.com/alexbakker/tox4go/crypto"
	"github.com/alexbakker/tox4go/dht"
	"github.com/alexbakker/tox4go/dht/ping"
	"github.com/alexbakker/tox4go/relay"
	"github.com/alexbakker/tox4go/transport"
	"github.com/didip/tollbooth"
)

const (
	enableIpv6  = true
	probeRate   = 1 * time.Minute
	refreshRate = 5 * time.Minute
)

var (
	lastScan    int64
	lastRefresh int64
	nodes       = []*toxNode{}
	nodesMutex  = sync.Mutex{}
	tcpPorts    = []int{443, 3389, 33445}
)

func main() {
	if parseFlags() {
		return
	}

	// load state if available
	state, err := loadState()
	if err != nil {
		log.Fatalf("error loading state: %s", err.Error())
	}
	lastScan = state.LastScan
	lastRefresh = state.LastRefresh
	nodes = state.Nodes

	if err := loadCountries(); err != nil {
		log.Fatalf("error loading countries.json: %s", err.Error())
	}

	inst, err := NewInstance(":33450")
	if err != nil {
		log.Fatalf("fatal: %s", err.Error())
	}
	inst.UDPTransport.Handle(dht.PacketIDSendNodes, inst.handleSendNodesPacket)
	inst.UDPTransport.Handle(bootstrap.PacketIDBootstrapInfo, handleBootstrapInfoPacket)

	//handle stop signal
	interruptChan := make(chan os.Signal)
	signal.Notify(interruptChan, os.Interrupt)

	//setup http server
	listener, err := net.Listen("tcp4", fmt.Sprintf("localhost:%d", httpListenPort))
	if err != nil {
		log.Fatalf("error in net.Listen: %s", err.Error())
	}
	limiter := tollbooth.NewLimiter(1, 2*time.Second)
	limiter.Methods = []string{"POST"}
	limiter.IPLookups = []string{"X-Forwarded-For", "RemoteAddr", "X-Real-IP"}
	serveMux := http.NewServeMux()
	serveMux.HandleFunc("/", handleHTTPRequest)
	serveMux.Handle("/test", tollbooth.LimitFuncHandler(limiter, handleHTTPRequest))
	serveMux.HandleFunc("/json", handleJSONRequest)
	go func() {
		err := http.Serve(listener, serveMux)
		if err != nil {
			log.Printf("http server error: %s\n", err.Error())
			interruptChan <- os.Interrupt
		}
	}()

	//listen for tox packets
	go func() {
		err := inst.UDPTransport.Listen()
		if err != nil {
			log.Printf("udp transport error: %s\n", err.Error())
			interruptChan <- os.Interrupt
		}
	}()
	//go tcpTransport.Listen()

	err = refreshNodes()
	if err != nil {
		log.Fatal(err.Error())
	}
	inst.probeNodes()

	probeTicker := time.NewTicker(probeRate)
	refreshTicker := time.NewTicker(refreshRate)
	updateTicker := time.NewTicker(30 * time.Second)
	run := true

	for run {
		select {
		case <-interruptChan:
			fmt.Printf("killing routines\n")
			probeTicker.Stop()
			refreshTicker.Stop()
			updateTicker.Stop()
			inst.UDPTransport.Stop()
			//tcpTransport.Stop()
			listener.Close()
			run = false
		case <-probeTicker.C:
			// we want an empty ping list at the start of every probe
			inst.PingsMutex.Lock()
			inst.Pings.Clear(false)
			inst.PingsMutex.Unlock()

			nodesMutex.Lock()
			err := inst.probeNodes()
			nodesMutex.Unlock()
			if err != nil {
				log.Printf("error while trying to probe nodes: %s", err.Error())
			}
		case <-refreshTicker.C:
			err := refreshNodes()
			if err != nil {
				log.Printf("error while trying to refresh nodes: %s", err.Error())
			}
		case <-updateTicker.C:
			inst.PingsMutex.Lock()
			inst.Pings.Clear(true)
			inst.PingsMutex.Unlock()

			nodesMutex.Lock()
			for _, node := range nodes {
				if time.Now().Sub(time.Unix(node.LastPing, 0)) > time.Minute*2 {
					node.UDPStatus = false
				}
			}
			sort.Stable(nodeSlice(nodes))

			state := getState()
			saveState(state)
			nodesMutex.Unlock()
		}
	}
}

func refreshNodes() error {
	parsedNodes, err := parseNodes()
	if err != nil {
		return err
	}

	nodesMutex.Lock()
	for _, freshNode := range parsedNodes {
		for _, node := range nodes {
			if freshNode.PublicKey == node.PublicKey {
				freshNode.LastPing = node.LastPing
				freshNode.UDPStatus = node.UDPStatus
				freshNode.TCPStatus = node.TCPStatus
				freshNode.TCPPorts = node.TCPPorts
				freshNode.MOTD = node.MOTD
				freshNode.Version = node.Version
				break
			}
		}
	}
	nodes = parsedNodes
	sort.Stable(nodeSlice(nodes))
	nodesMutex.Unlock()

	lastRefresh = time.Now().Unix()
	return nil
}

func (i *instance) probeNodes() error {
	for _, node := range nodes {
		err := i.getBootstrapInfo(node)
		if err != nil {
			fmt.Println(err.Error())
		}

		p, err := i.getNodes(node)
		if err != nil {
			fmt.Println(err.Error())
		} else {
			i.PingsMutex.Lock()
			err = i.Pings.Add(p)
			if err != nil {
				fmt.Println(err.Error())
			}
			i.PingsMutex.Unlock()
		}

		ports := tcpPorts
		exists := false
		for _, i := range ports {
			if i == node.Port {
				exists = true
			}
		}
		if !exists {
			ports = append(ports, node.Port)
		}

		go i.probeNodeTCPPorts(node, ports)
	}

	lastScan = time.Now().Unix()
	return nil
}

func (i *instance) probeNodeTCPPorts(node *toxNode, ports []int) {
	c := make(chan int)
	for _, port := range ports {
		go func(p int) {
			conn, err := connectTCP(node, p)
			if err != nil {
				fmt.Printf("%s\n", err.Error())
				c <- -1
				return
			}

			err = i.tcpHandshake(node, conn)
			if err != nil {
				fmt.Printf("%s\n", err.Error())
				c <- -1
			} else {
				c <- p
			}
			conn.Close()
		}(port)
	}

	newPorts := []int{}
	for i := 0; i < len(ports); i++ {
		port := <-c
		if port != -1 {
			fmt.Printf("tcp port for %s: %d\n", node.Maintainer, port)
			newPorts = append(newPorts, port)
		}
	}

	nodesMutex.Lock()
	node.TCPPorts = newPorts
	if len(node.TCPPorts) > 0 {
		node.LastPing = time.Now().Unix()
	}
	node.TCPStatus = len(node.TCPPorts) > 0
	sort.Stable(nodeSlice(nodes))
	nodesMutex.Unlock()
}

func (i *instance) tcpHandshake(node *toxNode, conn *net.TCPConn) error {
	nodePublicKey := new([crypto.PublicKeySize]byte)
	decPublicKey, err := hex.DecodeString(node.PublicKey)
	if err != nil {
		return err
	}
	copy(nodePublicKey[:], decPublicKey)

	relayConn, err := relay.NewConnection()
	if err != nil {
		return err
	}

	req, err := relayConn.StartHandshake()
	if err != nil {
		return err
	}

	reqBytes, err := req.MarshalBinary()
	if err != nil {
		return err
	}

	encryptedReqBytes, nonce, err := i.Ident.EncryptBlob(reqBytes, nodePublicKey)
	if err != nil {
		return err
	}

	reqPacket := &relay.HandshakeRequestPacket{
		PublicKey: i.Ident.PublicKey,
		Nonce:     nonce,
		Payload:   encryptedReqBytes,
	}

	reqPacketBytes, err := reqPacket.MarshalBinary()
	if err != nil {
		return err
	}

	conn.SetDeadline(time.Now().Add(2 * time.Second))
	_, err = conn.Write(reqPacketBytes)
	if err != nil {
		return err
	}

	buffer := make([]byte, 96)
	left := len(buffer)
	for left > 0 {
		read, readErr := conn.Read(buffer[len(buffer)-left:])
		if readErr != nil {
			return readErr
		}
		left -= read
	}

	res := relay.HandshakeResponsePacket{}
	err = res.UnmarshalBinary(buffer)
	if err != nil {
		return err
	}

	decryptedBytes, err := i.Ident.DecryptBlob(res.Payload, nodePublicKey, res.Nonce)
	if err != nil {
		return err
	}

	resPacket := &relay.HandshakePayload{}
	err = resPacket.UnmarshalBinary(decryptedBytes)
	if err != nil {
		return err
	}

	return relayConn.EndHandshake(resPacket)
}

func (i *instance) getNodes(node *toxNode) (*ping.Ping, error) {
	nodePublicKey := new([crypto.PublicKeySize]byte)
	decPublicKey, err := hex.DecodeString(node.PublicKey)
	if err != nil {
		return nil, err
	}
	copy(nodePublicKey[:], decPublicKey)

	p, err := ping.NewPing(nodePublicKey)
	if err != nil {
		return nil, err
	}

	packet := &dht.GetNodesPacket{
		PublicKey: i.Ident.PublicKey,
		PingID:    p.ID,
	}

	dhtPacket, err := i.Ident.EncryptPacket(transport.Packet(packet), nodePublicKey)
	if err != nil {
		return nil, err
	}

	payload, err := dhtPacket.MarshalBinary()
	if err != nil {
		return nil, err
	}

	if err := i.sendToUDP(payload, node); err != nil {
		return nil, err
	}

	return p, nil
}

func (i *instance) getBootstrapInfo(node *toxNode) error {
	packet, err := bootstrap.ConstructPacket(&bootstrap.InfoRequestPacket{})
	if err != nil {
		return err
	}

	payload, err := packet.MarshalBinary()
	if err != nil {
		return err
	}

	return i.sendToUDP(payload, node)
}

func (i *instance) sendToUDP(data []byte, node *toxNode) error {
	ip, err := getNodeIP(node)
	if err != nil {
		return err
	}

	return i.UDPTransport.Send(
		&transport.Message{
			Data: data,
			Addr: &net.UDPAddr{
				IP:   ip,
				Port: node.Port,
			},
		},
	)
}

func getNodeIP(node *toxNode) (net.IP, error) {
	if node.ip4 != nil {
		return node.ip4, nil
	} else if enableIpv6 && node.ip6 != nil {
		return node.ip6, nil
	}

	return nil, fmt.Errorf("no valid ip found for %s", node.Maintainer)
}

func connectTCP(node *toxNode, port int) (*net.TCPConn, error) {
	ip, err := getNodeIP(node)
	if err != nil {
		return nil, err
	}

	dialer := net.Dialer{}
	dialer.Deadline = time.Now().Add(2 * time.Second)

	tempConn, err := dialer.Dial("tcp", fmt.Sprintf("%s:%d", ip, port))
	if err != nil {
		return nil, err
	}

	conn, ok := tempConn.(*net.TCPConn)
	if !ok {
		return nil, errors.New("not a tcp conn")
	}

	return conn, nil
}

func (i *instance) handleSendNodesPacket(msg *transport.Message) error {
	dhtPacket := &dht.Packet{}
	err := dhtPacket.UnmarshalBinary(msg.Data)
	if err != nil {
		return err
	}

	decryptedPacket, err := i.Ident.DecryptPacket(dhtPacket)
	if err != nil {
		return err
	}

	packet, ok := decryptedPacket.(*dht.SendNodesPacket)
	if !ok {
		return nil
	}

	i.PingsMutex.Lock()
	ping := i.Pings.Find(dhtPacket.SenderPublicKey, packet.PingID, true)
	i.PingsMutex.Unlock()

	if ping == nil {
		err := fmt.Errorf("sendnodes packet from unknown node: %s:%d", msg.Addr.IP, msg.Addr.Port)
		fmt.Printf("error: %s", err.Error())
		return err
	}

	nodesMutex.Lock()
	for _, node := range nodes {
		publicKey, err := hex.DecodeString(node.PublicKey)
		if err != nil {
			continue
		}

		if bytes.Equal(publicKey, dhtPacket.SenderPublicKey[:]) {
			node.UDPStatus = true
			node.LastPing = time.Now().Unix()
			break
		}
	}
	sort.Stable(nodeSlice(nodes))
	nodesMutex.Unlock()

	return nil
}

func handleBootstrapInfoPacket(msg *transport.Message) error {
	bootstrapPacket := &bootstrap.Packet{}
	err := bootstrapPacket.UnmarshalBinary(msg.Data)
	if err != nil {
		return err
	}

	transPacket, err := bootstrap.DestructPacket(bootstrapPacket)
	if err != nil {
		return err
	}

	packet, ok := transPacket.(*bootstrap.InfoResponsePacket)
	if !ok {
		return errors.New("wtf")
	}

	nodesMutex.Lock()
	for _, node := range nodes {
		if (node.ip4 != nil && node.ip4.Equal(msg.Addr.IP)) ||
			(node.ip6 != nil && node.ip6.Equal(msg.Addr.IP)) {
			node.MOTD = packet.MOTD
			node.Version = fmt.Sprintf("%d", packet.Version)
			break
		}
	}
	nodesMutex.Unlock()

	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"flag"
	"github.com/robfig/cron"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"strings"
)

func main() {
	configPath := flag.String("c", "wl520cron.json", "配置文件路径")
	wlc := flag.String("wlc", "", "wl520配置文件路径")
	flag.Parse()
	wl520Crons := readConfig(*configPath)
	c := cron.New()
	for _, v := range wl520Crons {
		cmdCloned := v.Cmd
		if *wlc != "" {
			cmdCloned = cmdCloned + " -c=" + *wlc
		}
		log.Printf("%s:%s\n", v.Cron, cmdCloned)
		c.AddFunc(v.Cron, func() {
			cmd := exec.Command("wl520", strings.Split(cmdCloned, " ")...)
			out, err := cmd.CombinedOutput()
			if err != nil {
				panic(err)
			}
			log.Println(string(out))
		})
	}
	c.Start()
	select {}
}

//读取配置文件
func readConfig(configPath string) []Wl520Cron {
	bytes, err := ioutil.ReadFile(configPath)
	if err != nil {
		createCronFile(configPath)
		log.Fatalf("配置文件不存在，已创建默认配置文件%s\n", configPath)
	}
	wl520Crons := make([]Wl520Cron, 0)
	json.Unmarshal(bytes, &wl520Crons)
	return wl520Crons
}

type Wl520Cron struct {
	Cron string `json:"cron"`
	Cmd  string `json:"cmd"`
}

//创建默认配置文件
func createCronFile(path string) {
	var f, _ = os.OpenFile(path, os.O_CREATE|os.O_RDWR, os.ModeAppend)
	defer f.Close()
	wl520Crons := []Wl520Cron{{Cron: "* */30 * * * *", Cmd: "-a -p"},
		{Cron: "* 0 1,13 * * *", Cmd: "-t -v=20 -farm-sign"}}
	bytes, _ := json.MarshalIndent(wl520Crons, "", "  ")
	f.Write(bytes)
}
<commit_msg>Error Handle<commit_after>package main

import (
	"encoding/json"
	"flag"
	"github.com/robfig/cron"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"strings"
)

func main() {
	configPath := flag.String("c", "wl520cron.json", "配置文件路径")
	wlc := flag.String("wlc", "", "wl520配置文件路径")
	flag.Parse()
	wl520Crons := readConfig(*configPath)
	c := cron.New()
	for _, v := range wl520Crons {
		cmdCloned := v.Cmd
		if *wlc != "" {
			cmdCloned = cmdCloned + " -c=" + *wlc
		}
		log.Printf("%s:%s\n", v.Cron, cmdCloned)
		c.AddFunc(v.Cron, func() {
			cmd := exec.Command("wl520", strings.Split(cmdCloned, " ")...)
			out, err := cmd.CombinedOutput()
			if err != nil {
				panic(err)
			}
			log.Println(string(out))
		})
	}
	c.Start()
	select {}
}

//读取配置文件
func readConfig(configPath string) []Wl520Cron {
	bytes, err := ioutil.ReadFile(configPath)
	if err != nil {
		createCronFile(configPath)
		log.Fatalf("配置文件不存在，已创建默认配置文件%s\n", configPath)
	}
	wl520Crons := make([]Wl520Cron, 0)
	json.Unmarshal(bytes, &wl520Crons)
	return wl520Crons
}

type Wl520Cron struct {
	Cron string `json:"cron"`
	Cmd  string `json:"cmd"`
}

//创建默认配置文件
func createCronFile(path string) {
	var f, _ = os.OpenFile(path, os.O_CREATE|os.O_RDWR, os.ModeAppend)
	defer f.Close()
	wl520Crons := []Wl520Cron{{Cron: "* */30 * * * *", Cmd: "-a -p"},
		{Cron: "* 0 1,13 * * *", Cmd: "-t -v=20 -farm-sign"}}
	bytes, _ := json.MarshalIndent(wl520Crons, "", "  ")
	_, err := f.Write(bytes)
	if err != nil {
		panic(err)
	}
}
<|endoftext|>


<commit_before>/* Upgrade a MIG agent

Version: MPL 1.1/GPL 2.0/LGPL 2.1

The contents of this file are subject to the Mozilla Public License Version
1.1 (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.mozilla.org/MPL/

Software distributed under the License is distributed on an "AS IS" basis,
WITHOUT WARRANTY OF ANY KIND, either express or implied. See the License
for the specific language governing rights and limitations under the
License.

The Initial Developer of the Original Code is
Mozilla Corporation
Portions created by the Initial Developer are Copyright (C) 2014
the Initial Developer. All Rights Reserved.

Contributor(s):
Julien Vehent jvehent@mozilla.com [:ulfr]

Alternatively, the contents of this file may be used under the terms of
either the GNU General Public License Version 2 or later (the "GPL"), or
the GNU Lesser General Public License Version 2.1 or later (the "LGPL"),
in which case the provisions of the GPL or the LGPL are applicable instead
of those above. If you wish to allow use of your version of this file only
under the terms of either the GPL or the LGPL, and not to allow others to
use your version of this file under the terms of the MPL, indicate your
decision by deleting the provisions above and replace them with the notice
and other provisions required by the GPL or the LGPL. If you do not delete
the provisions above, a recipient may use your version of this file under
the terms of any one of the MPL, the GPL or the LGPL.
*/

package upgrade

import (
	"bitbucket.org/kardianos/osext"
	"crypto/sha256"
	"encoding/json"
	"fmt"
	"hash"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"os/exec"
	"regexp"
	"runtime"
	"time"
)

type Parameters struct {
	Elements map[string]map[string]string `json:"elements"`
}

func NewParameters() (p Parameters) {
	return
}

type Results struct {
	Success    bool       `json:"success"`
	OldPID     int        `json:"oldpid"`
	Error      string     `json:"error,omitempty"`
	Statistics Statistics `json:"statistics,omitempty"`
}

func (p Parameters) Validate() (err error) {
	versionre := regexp.MustCompile(`^[a-z0-9]{7}-[0-9]{12}$`)
	locre := regexp.MustCompile(`^https?://`)
	checksumre := regexp.MustCompile(`^[a-zA-Z0-9]{64}$`)
	for k, el := range p.Elements {
		if !versionre.MatchString(el["to_version"]) {
			return fmt.Errorf("In %s, parameter 'to_version' with value '%s' is invalid. Expecting version.", k, el["to_version"])
		}
		if !locre.MatchString(el["location"]) {
			return fmt.Errorf("In %s, parameter 'location' with value '%s' is invalid. Expecting URL.", k, el["location"])
		}
		if !checksumre.MatchString(el["checksum"]) {
			return fmt.Errorf("In %s, parameter 'checksum' with value '%s' is invalid. Expecting SHA256 checksum.", k, el["checksum"])
		}
	}
	return
}

var stats Statistics

type Statistics struct {
	DownloadTime string `json:"downloadtime"`
	DownloadSize int64  `json:"downloadsize"`
}

func Run(Args []byte) string {
	p := NewParameters()

	err := json.Unmarshal(Args, &p.Elements)
	if err != nil {
		panic(err)
	}

	err = p.Validate()
	if err != nil {
		return buildResults(p, fmt.Sprintf("%v", err))
	}

	// Extract the parameters that apply to this OS and Arch
	key := fmt.Sprintf("%s/%s", runtime.GOOS, runtime.GOARCH)
	el, ok := p.Elements[key]
	if !ok {
		return buildResults(p, fmt.Sprintf("No parameter found for %s", key))
	}

	// Verify that the version we're told to upgrade to isn't the current one
	cversion, err := getCurrentVersion()
	if err != nil {
		return buildResults(p, fmt.Sprintf("%v", err))
	}
	if cversion == el["to_version"] {
		return buildResults(p, fmt.Sprintf("Agent is already running version '%s'", cversion))
	}

	// Download new agent binary from provided location
	binfd, err := downloadBinary(el["location"])
	if err != nil {
		return buildResults(p, fmt.Sprintf("%v", err))
	}

	// Verify checksum of the binary
	err = verifyChecksum(binfd, el["checksum"])
	if err != nil {
		return buildResults(p, fmt.Sprintf("%v", err))
	}

	// grab the path before closing the file descriptor
	binPath := binfd.Name()

	err = binfd.Close()
	if err != nil {
		return buildResults(p, fmt.Sprintf("%v", err))
	}

	// Dry run of the binary to verify that the version is correct
	// but also that it can run without error
	err = verifyVersion(binPath, el["to_version"])
	if err != nil {
		return buildResults(p, fmt.Sprintf("%v", err))
	}

	// Move the binary of the new agent from tmp, to the correct destination
	agentBinPath, err := moveBinary(binPath, el["to_version"])
	if err != nil {
		return buildResults(p, fmt.Sprintf("%v", err))
	}

	// Launch the new agent and exit the module
	_, err = exec.Command(agentBinPath).Output()
	if err != nil {
		return buildResults(p, fmt.Sprintf("%v", err))
	}

	return buildResults(p, "")
}

// Run the agent binary to obtain the current version
func getCurrentVersion() (cversion string, err error) {
	defer func() {
		if e := recover(); e != nil {
			err = fmt.Errorf("getCurrentVersion() -> %v", e)
		}
	}()
	bin, err := osext.Executable()
	if err != nil {
		panic(err)
	}
	out, err := exec.Command(bin, "-V").Output()
	if err != nil {
		panic(err)
	}
	if len(out) < 2 {
		panic("Failed to retrieve agent version.")
	}
	cversion = string(out[:len(out)-1])
	return
}

// downloadBinary retrieves the data from a location and saves it to a temp file
func downloadBinary(loc string) (tmpfd *os.File, err error) {
	defer func() {
		if e := recover(); e != nil {
			err = fmt.Errorf("downloadBinary() -> %v", e)
		}
	}()
	tmpfd, err = ioutil.TempFile("", "")
	if err != nil {
		panic(err)
	}
	start := time.Now()
	resp, err := http.Get(loc)
	if err != nil {
		panic(err)
	}
	stats.DownloadSize, err = io.Copy(tmpfd, resp.Body)
	stats.DownloadTime = time.Since(start).String()
	resp.Body.Close()
	return
}

// verifyChecksum computes the hash of a file and compares it
// to a checksum. If comparison fails, it returns an error.
func verifyChecksum(fd *os.File, checksum string) (err error) {
	defer func() {
		if e := recover(); e != nil {
			err = fmt.Errorf("verifyChecksum() -> %v", e)
		}
	}()
	var h hash.Hash
	h = sha256.New()
	buf := make([]byte, 4096)
	var offset int64 = 0
	for {
		block, err := fd.ReadAt(buf, offset)
		if err != nil && err != io.EOF {
			panic(err)
		}
		if block == 0 {
			break
		}
		h.Write(buf[:block])
		offset += int64(block)
	}
	hexhash := fmt.Sprintf("%x", h.Sum(nil))
	if hexhash != checksum {
		return fmt.Errorf("Checksum validation failed. Got '%s', Expected '%s'.",
			hexhash, checksum)
	}
	return
}

// verifyVersion runs a binary and compares the returned version
func verifyVersion(binPath, expectedVersion string) (err error) {
	defer func() {
		if e := recover(); e != nil {
			err = fmt.Errorf("verifyVersion() -> %v", e)
		}
	}()
	os.Chmod(binPath, 0750)
	out, err := exec.Command(binPath, "-V").Output()
	if err != nil {
		panic(err)
	}
	binVersion := string(out[:len(out)-1])
	if binVersion != expectedVersion {
		return fmt.Errorf("Version mismatch. Got '%s', Expected '%s'.",
			binVersion, expectedVersion)
	}
	return
}

func moveBinary(binPath, version string) (linkloc string, err error) {
	defer func() {
		if e := recover(); e != nil {
			err = fmt.Errorf("moveBinary() -> %v", e)
		}
	}()
	var target string
	switch runtime.GOOS {
	case "linux", "darwin", "freebsd", "openbsd", "netbsd":
		target = fmt.Sprintf("/sbin/mig-agent-%s", version)
		linkloc = "/sbin/mig-agent"
	case "windows":
		target = fmt.Sprintf("C:/Windows/mig-agent-%s.exe", version)
		linkloc = "C:/Windows/mig-agent"
	default:
		err = fmt.Errorf("'%s' isn't a supported OS", runtime.GOOS)
		return
	}
	// copy the file (rename may not work if we're crossing partitions)
	srcfd, err := os.Open(binPath)
	if err != nil {
		panic(err)
	}
	dstfd, err := os.Create(target)
	if err != nil {
		panic(err)
	}
	_, err = io.Copy(dstfd, srcfd)
	if err != nil {
		panic(err)
	}
	srcfd.Close()
	dstfd.Close()
	err = os.Remove(binPath)
	if err != nil {
		panic(err)
	}
	err = os.Chmod(target, 0750)
	if err != nil {
		panic(err)
	}
	// don't fail on removal of existing link, it may not exist
	os.Remove(linkloc)
	// create a symlink
	err = os.Symlink(target, linkloc)
	if err != nil {
		panic(err)
	}
	return
}

// buildResults transforms the ConnectedIPs map into a Results
// map that is serialized in JSON and returned as a string
func buildResults(params Parameters, errors string) string {
	var results Results
	results.OldPID = os.Getppid()
	if errors != "" {
		results.Error = errors
	} else {
		results.Success = true
	}
	results.Statistics = stats
	jsonOutput, err := json.Marshal(results)
	if err != nil {
		panic(err)
	}
	return string(jsonOutput[:])
}
<commit_msg>[minor] set permissions of mig-agent to 500 during upgrades<commit_after>/* Upgrade a MIG agent

Version: MPL 1.1/GPL 2.0/LGPL 2.1

The contents of this file are subject to the Mozilla Public License Version
1.1 (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.mozilla.org/MPL/

Software distributed under the License is distributed on an "AS IS" basis,
WITHOUT WARRANTY OF ANY KIND, either express or implied. See the License
for the specific language governing rights and limitations under the
License.

The Initial Developer of the Original Code is
Mozilla Corporation
Portions created by the Initial Developer are Copyright (C) 2014
the Initial Developer. All Rights Reserved.

Contributor(s):
Julien Vehent jvehent@mozilla.com [:ulfr]

Alternatively, the contents of this file may be used under the terms of
either the GNU General Public License Version 2 or later (the "GPL"), or
the GNU Lesser General Public License Version 2.1 or later (the "LGPL"),
in which case the provisions of the GPL or the LGPL are applicable instead
of those above. If you wish to allow use of your version of this file only
under the terms of either the GPL or the LGPL, and not to allow others to
use your version of this file under the terms of the MPL, indicate your
decision by deleting the provisions above and replace them with the notice
and other provisions required by the GPL or the LGPL. If you do not delete
the provisions above, a recipient may use your version of this file under
the terms of any one of the MPL, the GPL or the LGPL.
*/

package upgrade

import (
	"crypto/sha256"
	"encoding/json"
	"fmt"
	"hash"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"os/exec"
	"regexp"
	"runtime"
	"time"

	"bitbucket.org/kardianos/osext"
)

type Parameters struct {
	Elements map[string]map[string]string `json:"elements"`
}

func NewParameters() (p Parameters) {
	return
}

type Results struct {
	Success    bool       `json:"success"`
	OldPID     int        `json:"oldpid"`
	Error      string     `json:"error,omitempty"`
	Statistics Statistics `json:"statistics,omitempty"`
}

func (p Parameters) Validate() (err error) {
	versionre := regexp.MustCompile(`^[a-z0-9]{7}-[0-9]{12}$`)
	locre := regexp.MustCompile(`^https?://`)
	checksumre := regexp.MustCompile(`^[a-zA-Z0-9]{64}$`)
	for k, el := range p.Elements {
		if !versionre.MatchString(el["to_version"]) {
			return fmt.Errorf("In %s, parameter 'to_version' with value '%s' is invalid. Expecting version.", k, el["to_version"])
		}
		if !locre.MatchString(el["location"]) {
			return fmt.Errorf("In %s, parameter 'location' with value '%s' is invalid. Expecting URL.", k, el["location"])
		}
		if !checksumre.MatchString(el["checksum"]) {
			return fmt.Errorf("In %s, parameter 'checksum' with value '%s' is invalid. Expecting SHA256 checksum.", k, el["checksum"])
		}
	}
	return
}

var stats Statistics

type Statistics struct {
	DownloadTime string `json:"downloadtime"`
	DownloadSize int64  `json:"downloadsize"`
}

func Run(Args []byte) string {
	p := NewParameters()

	err := json.Unmarshal(Args, &p.Elements)
	if err != nil {
		panic(err)
	}

	err = p.Validate()
	if err != nil {
		return buildResults(p, fmt.Sprintf("%v", err))
	}

	// Extract the parameters that apply to this OS and Arch
	key := fmt.Sprintf("%s/%s", runtime.GOOS, runtime.GOARCH)
	el, ok := p.Elements[key]
	if !ok {
		return buildResults(p, fmt.Sprintf("No parameter found for %s", key))
	}

	// Verify that the version we're told to upgrade to isn't the current one
	cversion, err := getCurrentVersion()
	if err != nil {
		return buildResults(p, fmt.Sprintf("%v", err))
	}
	if cversion == el["to_version"] {
		return buildResults(p, fmt.Sprintf("Agent is already running version '%s'", cversion))
	}

	// Download new agent binary from provided location
	binfd, err := downloadBinary(el["location"])
	if err != nil {
		return buildResults(p, fmt.Sprintf("%v", err))
	}

	// Verify checksum of the binary
	err = verifyChecksum(binfd, el["checksum"])
	if err != nil {
		return buildResults(p, fmt.Sprintf("%v", err))
	}

	// grab the path before closing the file descriptor
	binPath := binfd.Name()

	err = binfd.Close()
	if err != nil {
		return buildResults(p, fmt.Sprintf("%v", err))
	}

	// Dry run of the binary to verify that the version is correct
	// but also that it can run without error
	err = verifyVersion(binPath, el["to_version"])
	if err != nil {
		return buildResults(p, fmt.Sprintf("%v", err))
	}

	// Move the binary of the new agent from tmp, to the correct destination
	agentBinPath, err := moveBinary(binPath, el["to_version"])
	if err != nil {
		return buildResults(p, fmt.Sprintf("%v", err))
	}

	// Launch the new agent and exit the module
	_, err = exec.Command(agentBinPath).Output()
	if err != nil {
		return buildResults(p, fmt.Sprintf("%v", err))
	}

	return buildResults(p, "")
}

// Run the agent binary to obtain the current version
func getCurrentVersion() (cversion string, err error) {
	defer func() {
		if e := recover(); e != nil {
			err = fmt.Errorf("getCurrentVersion() -> %v", e)
		}
	}()
	bin, err := osext.Executable()
	if err != nil {
		panic(err)
	}
	out, err := exec.Command(bin, "-V").Output()
	if err != nil {
		panic(err)
	}
	if len(out) < 2 {
		panic("Failed to retrieve agent version.")
	}
	cversion = string(out[:len(out)-1])
	return
}

// downloadBinary retrieves the data from a location and saves it to a temp file
func downloadBinary(loc string) (tmpfd *os.File, err error) {
	defer func() {
		if e := recover(); e != nil {
			err = fmt.Errorf("downloadBinary() -> %v", e)
		}
	}()
	tmpfd, err = ioutil.TempFile("", "")
	if err != nil {
		panic(err)
	}
	start := time.Now()
	resp, err := http.Get(loc)
	if err != nil {
		panic(err)
	}
	stats.DownloadSize, err = io.Copy(tmpfd, resp.Body)
	stats.DownloadTime = time.Since(start).String()
	resp.Body.Close()
	return
}

// verifyChecksum computes the hash of a file and compares it
// to a checksum. If comparison fails, it returns an error.
func verifyChecksum(fd *os.File, checksum string) (err error) {
	defer func() {
		if e := recover(); e != nil {
			err = fmt.Errorf("verifyChecksum() -> %v", e)
		}
	}()
	var h hash.Hash
	h = sha256.New()
	buf := make([]byte, 4096)
	var offset int64 = 0
	for {
		block, err := fd.ReadAt(buf, offset)
		if err != nil && err != io.EOF {
			panic(err)
		}
		if block == 0 {
			break
		}
		h.Write(buf[:block])
		offset += int64(block)
	}
	hexhash := fmt.Sprintf("%x", h.Sum(nil))
	if hexhash != checksum {
		return fmt.Errorf("Checksum validation failed. Got '%s', Expected '%s'.",
			hexhash, checksum)
	}
	return
}

// verifyVersion runs a binary and compares the returned version
func verifyVersion(binPath, expectedVersion string) (err error) {
	defer func() {
		if e := recover(); e != nil {
			err = fmt.Errorf("verifyVersion() -> %v", e)
		}
	}()
	os.Chmod(binPath, 0500)
	out, err := exec.Command(binPath, "-V").Output()
	if err != nil {
		panic(err)
	}
	binVersion := string(out[:len(out)-1])
	if binVersion != expectedVersion {
		return fmt.Errorf("Version mismatch. Got '%s', Expected '%s'.",
			binVersion, expectedVersion)
	}
	return
}

func moveBinary(binPath, version string) (linkloc string, err error) {
	defer func() {
		if e := recover(); e != nil {
			err = fmt.Errorf("moveBinary() -> %v", e)
		}
	}()
	var target string
	switch runtime.GOOS {
	case "linux", "darwin", "freebsd", "openbsd", "netbsd":
		target = fmt.Sprintf("/sbin/mig-agent-%s", version)
		linkloc = "/sbin/mig-agent"
	case "windows":
		target = fmt.Sprintf("C:/Windows/mig-agent-%s.exe", version)
		linkloc = "C:/Windows/mig-agent"
	default:
		err = fmt.Errorf("'%s' isn't a supported OS", runtime.GOOS)
		return
	}
	// copy the file (rename may not work if we're crossing partitions)
	srcfd, err := os.Open(binPath)
	if err != nil {
		panic(err)
	}
	dstfd, err := os.Create(target)
	if err != nil {
		panic(err)
	}
	_, err = io.Copy(dstfd, srcfd)
	if err != nil {
		panic(err)
	}
	srcfd.Close()
	dstfd.Close()
	err = os.Remove(binPath)
	if err != nil {
		panic(err)
	}
	err = os.Chmod(target, 0750)
	if err != nil {
		panic(err)
	}
	// don't fail on removal of existing link, it may not exist
	os.Remove(linkloc)
	// create a symlink
	err = os.Symlink(target, linkloc)
	if err != nil {
		panic(err)
	}
	return
}

// buildResults transforms the ConnectedIPs map into a Results
// map that is serialized in JSON and returned as a string
func buildResults(params Parameters, errors string) string {
	var results Results
	results.OldPID = os.Getppid()
	if errors != "" {
		results.Error = errors
	} else {
		results.Success = true
	}
	results.Statistics = stats
	jsonOutput, err := json.Marshal(results)
	if err != nil {
		panic(err)
	}
	return string(jsonOutput[:])
}
<|endoftext|>


<commit_before>package reporters

import (
	"log"

	"github.com/garyburd/redigo/redis"
	"github.com/ryanlower/drain/parser"
)

// Redis is a reporter that publishes hit statues to redis
// The address of the redis server to be published can be
// controlled using the REDIS_ADDRESS env (defaults to "localhost:6379")
type Redis struct {
	Reporter
}

// Report publishes the parsed log lines status code to redis
func (r *Redis) Report(hit *parser.ParsedLogLine) {
	conn := connect()

	conn.Do("PUBLISH", "drain.statuses", hit.Status)
}

func connect() redis.Conn {
	address := envOrDefault("REDIS_ADDRESS", "localhost:6379")
	// TODO, auth handling
	conn, err := redis.Dial("tcp", address)
	// TODO, how to handle errors in reporters?
	if err != nil {
		log.Panic(err)
	}

	return conn
}
<commit_msg>Publish hit as JSON to drain.hits channel from redis reporter<commit_after>package reporters

import (
	"encoding/json"
	"log"

	"github.com/garyburd/redigo/redis"
	"github.com/ryanlower/drain/parser"
)

// Redis is a reporter that publishes hit statues to redis
// The address of the redis server to be published can be
// controlled using the REDIS_ADDRESS env (defaults to "localhost:6379")
type Redis struct {
	Reporter
}

// Report publishes the parsed log lines status code to redis
func (r *Redis) Report(hit *parser.ParsedLogLine) {
	conn := connect()

	hitJSON, err := json.Marshal(hit)
	if err != nil {
		log.Panic(err)
	} else {
		conn.Do("PUBLISH", "drain.hits", hitJSON)
	}
}

func connect() redis.Conn {
	address := envOrDefault("REDIS_ADDRESS", "localhost:6379")
	// TODO, auth handling
	conn, err := redis.Dial("tcp", address)
	// TODO, how to handle errors in reporters?
	if err != nil {
		log.Panic(err)
	}

	return conn
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"encoding/json"
	"flag"
	"log"
	"net/http"
	"os"
	"strings"
	"sync"

	"github.com/miekg/dns"
	"github.com/nxadm/tail"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promhttp"

	"github.com/abh/geodns/countries"
	"github.com/abh/geodns/querylog"
)

// TODO:
// Add vendor yes/no
// add server region tag (identifier)?

var version string = "2.1"

func main() {

	log.Printf("Starting geodns-logs/%q", version)

	identifierFlag := flag.String("identifier", "", "identifier (hostname, pop name or similar)")
	// verboseFlag := flag.Bool("verbose", false, "verbose output")
	flag.Parse()

	var serverID string
	// var serverGroups []string

	if len(*identifierFlag) > 0 {
		ids := strings.Split(*identifierFlag, ",")
		serverID = ids[0]
		// if len(ids) > 1 {
		// serverGroups = ids[1:]
		// }
	}

	if len(serverID) == 0 {
		var err error
		serverID, err = os.Hostname()
		if err != nil {
			log.Printf("Could not get hostname: %s", err)
			os.Exit(2)
		}
	}

	queries = prometheus.NewCounterVec(
		prometheus.CounterOpts{
			Name: "dns_logs_total",
			Help: "Number of served queries",
		},
		[]string{"zone", "vendor", "usercc", "poolcc", "qtype"},
	)
	prometheus.MustRegister(queries)

	buildInfo := prometheus.NewGaugeVec(
		prometheus.GaugeOpts{
			Name: "geodns_logs_build_info",
			Help: "GeoDNS logs build information (in labels)",
		},
		[]string{"Version"},
	)
	prometheus.MustRegister(buildInfo)
	buildInfo.WithLabelValues("geodns-logs/" + version).Set(1)

	http.Handle("/metrics", promhttp.Handler())
	go func() {
		err := http.ListenAndServe(":8054", nil)
		if err != nil {
			log.Printf("could not start http server: %s", err)
		}
	}()

	if len(flag.Args()) < 1 {
		log.Printf("filename to process required")
		os.Exit(2)
	}

	filename := flag.Arg(0)

	logf, err := tail.TailFile(filename, tail.Config{
		// Location:  &tail.SeekInfo{-1, 0},
		Poll:      true, // inotify is flaky on EL6, so try this ...
		ReOpen:    true,
		MustExist: false,
		Follow:    true,
	})
	if err != nil {
		log.Printf("Could not tail '%s': %s", filename, err)
	}

	in := make(chan string)
	go processChan(in, nil)

	for line := range logf.Lines {
		if line.Err != nil {
			log.Printf("Error tailing file: %s", line.Err)
		}
		in <- line.Text
	}

}

var extraValidLabels = map[string]struct{}{
	"uk":       struct{}{},
	"_status":  struct{}{},
	"_country": struct{}{},
	"www":      struct{}{},
	"nag-test": struct{}{},
}

func validCC(label string) bool {
	if _, ok := countries.CountryContinent[label]; ok {
		return true
	}
	if _, ok := countries.ContinentCountries[label]; ok {
		return true
	}
	if _, ok := countries.RegionGroupRegions[label]; ok {
		return true
	}
	if _, ok := countries.RegionGroups[label]; ok {
		return true
	}
	if _, ok := extraValidLabels[label]; ok {
		return true
	}
	return false
}

func getPoolCC(label string) (string, bool) {
	l := dns.SplitDomainName(label)
	// log.Printf("LABEL: %+v", l)
	if len(l) == 0 {
		return "", true
	}

	for _, cc := range l {
		if validCC(cc) {
			return cc, true
		}
	}

	if len(l[0]) == 1 && strings.ContainsAny(l[0], "01234") {
		if len(l) == 1 {
			return "", true
		}
	}

	// log.Printf("LABEL '%s' unhandled cc...", label)
	return "", false
}

func processChan(in chan string, wg *sync.WaitGroup) error {
	e := querylog.Entry{}

	stats := NewStats()

	for line := range in {
		err := json.Unmarshal([]byte(line), &e)
		if err != nil {
			log.Printf("Can't unmarshal '%s': %s", line, err)
			return err
		}
		e.Name = strings.ToLower(e.Name)

		// fmt.Printf("%s %s\n", e.Origin, e.Name)

		err = stats.Add(&e)
		if err != nil {
			return err
		}
	}

	if wg != nil {
		wg.Done()
	}
	return nil
}

func processFile(file string, out chan<- *Stats) error {
	fh, err := os.Open(file)
	if err != nil {
		return err
	}

	in := make(chan string)

	wg := sync.WaitGroup{}
	wg.Add(1)
	go processChan(in, &wg)

	scanner := bufio.NewScanner(fh)

	for scanner.Scan() {
		in <- scanner.Text()
	}
	if err := scanner.Err(); err != nil {
		log.Println("reading standard input:", err)
	}

	close(in)

	wg.Wait()

	return nil
}
<commit_msg>geodns-logs: don't stop processing on errors<commit_after>package main

import (
	"bufio"
	"encoding/json"
	"flag"
	"log"
	"net/http"
	"os"
	"strings"
	"sync"

	"github.com/miekg/dns"
	"github.com/nxadm/tail"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promhttp"

	"github.com/abh/geodns/countries"
	"github.com/abh/geodns/querylog"
)

// TODO:
// Add vendor yes/no
// add server region tag (identifier)?

var version string = "2.1"

func main() {

	log.Printf("Starting geodns-logs/%q", version)

	identifierFlag := flag.String("identifier", "", "identifier (hostname, pop name or similar)")
	// verboseFlag := flag.Bool("verbose", false, "verbose output")
	flag.Parse()

	var serverID string
	// var serverGroups []string

	if len(*identifierFlag) > 0 {
		ids := strings.Split(*identifierFlag, ",")
		serverID = ids[0]
		// if len(ids) > 1 {
		// serverGroups = ids[1:]
		// }
	}

	if len(serverID) == 0 {
		var err error
		serverID, err = os.Hostname()
		if err != nil {
			log.Printf("Could not get hostname: %s", err)
			os.Exit(2)
		}
	}

	queries = prometheus.NewCounterVec(
		prometheus.CounterOpts{
			Name: "dns_logs_total",
			Help: "Number of served queries",
		},
		[]string{"zone", "vendor", "usercc", "poolcc", "qtype"},
	)
	prometheus.MustRegister(queries)

	buildInfo := prometheus.NewGaugeVec(
		prometheus.GaugeOpts{
			Name: "geodns_logs_build_info",
			Help: "GeoDNS logs build information (in labels)",
		},
		[]string{"Version"},
	)
	prometheus.MustRegister(buildInfo)
	buildInfo.WithLabelValues("geodns-logs/" + version).Set(1)

	http.Handle("/metrics", promhttp.Handler())
	go func() {
		err := http.ListenAndServe(":8054", nil)
		if err != nil {
			log.Printf("could not start http server: %s", err)
		}
	}()

	if len(flag.Args()) < 1 {
		log.Printf("filename to process required")
		os.Exit(2)
	}

	filename := flag.Arg(0)

	logf, err := tail.TailFile(filename, tail.Config{
		// Location:  &tail.SeekInfo{-1, 0},
		Poll:      true, // inotify is flaky on EL6, so try this ...
		ReOpen:    true,
		MustExist: false,
		Follow:    true,
	})
	if err != nil {
		log.Printf("Could not tail '%s': %s", filename, err)
	}

	in := make(chan string)
	go processChan(in, nil)

	for line := range logf.Lines {
		if line.Err != nil {
			log.Printf("Error tailing file: %s", line.Err)
		}
		in <- line.Text
	}

}

var extraValidLabels = map[string]struct{}{
	"uk":       struct{}{},
	"_status":  struct{}{},
	"_country": struct{}{},
	"www":      struct{}{},
	"nag-test": struct{}{},
}

func validCC(label string) bool {
	if _, ok := countries.CountryContinent[label]; ok {
		return true
	}
	if _, ok := countries.ContinentCountries[label]; ok {
		return true
	}
	if _, ok := countries.RegionGroupRegions[label]; ok {
		return true
	}
	if _, ok := countries.RegionGroups[label]; ok {
		return true
	}
	if _, ok := extraValidLabels[label]; ok {
		return true
	}
	return false
}

func getPoolCC(label string) (string, bool) {
	l := dns.SplitDomainName(label)
	// log.Printf("LABEL: %+v", l)
	if len(l) == 0 {
		return "", true
	}

	for _, cc := range l {
		if validCC(cc) {
			return cc, true
		}
	}

	if len(l[0]) == 1 && strings.ContainsAny(l[0], "01234") {
		if len(l) == 1 {
			return "", true
		}
	}

	// log.Printf("LABEL '%s' unhandled cc...", label)
	return "", false
}

func processChan(in chan string, wg *sync.WaitGroup) error {
	e := querylog.Entry{}

	stats := NewStats()

	for line := range in {
		err := json.Unmarshal([]byte(line), &e)
		if err != nil {
			log.Printf("unmarshal error '%s': %s", line, err)
			continue
		}
		e.Name = strings.ToLower(e.Name)

		// fmt.Printf("%s %s\n", e.Origin, e.Name)

		err = stats.Add(&e)
		if err != nil {
			log.Printf("stats error: %s", err)
			continue
		}
	}

	if wg != nil {
		wg.Done()
	}
	return nil
}

func processFile(file string, out chan<- *Stats) error {
	fh, err := os.Open(file)
	if err != nil {
		return err
	}

	in := make(chan string)

	wg := sync.WaitGroup{}
	wg.Add(1)
	go processChan(in, &wg)

	scanner := bufio.NewScanner(fh)

	for scanner.Scan() {
		in <- scanner.Text()
	}
	if err := scanner.Err(); err != nil {
		log.Println("reading standard input:", err)
	}

	close(in)

	wg.Wait()

	return nil
}
<|endoftext|>


<commit_before>package main

import "testing"

func TestSqrt(t *testing.T) {

  t.Error("an error")
}
<commit_msg>test ok<commit_after>package main

import "testing"

func TestSqrt(t *testing.T) {
  //t.Error("an error")
}
<|endoftext|>


<commit_before>// Copyright (c) 2016 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package cpustat

// #include <linux/taskstats.h>
// #include <linux/genetlink.h>
import "C"

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"os"
	"syscall"
	"time"
)

// convert a byte slice of a null terminated C string into a Go string
func stringFromBytes(c []byte) string {
	nullPos := 0
	i := 0
	for ; i < len(c); i++ {
		if c[i] == 0 {
			nullPos = i
			break
		}
	}
	return string(c[:nullPos])
}

func readGetTaskstatsMessage(conn *NLConn) (*TaskStats, string, error) {
	inBytes, err := conn.Read()
	if err != nil {
		return nil, "", err
	}
	if len(inBytes) <= 0 {
		return nil, "", fmt.Errorf("short read requesting taskstats info: %d bytes", len(inBytes))
	}
	nlmsgs, err := syscall.ParseNetlinkMessage(inBytes)
	if err != nil {
		return nil, "", err
	}

	if len(nlmsgs) != 1 {
		panic(fmt.Sprint("got unexpected response size from get genl taskstats request: ", len(nlmsgs)))
	}

	if nlmsgs[0].Header.Type == syscall.NLMSG_ERROR {
		var errno int32
		buf := bytes.NewBuffer(nlmsgs[0].Data)
		_ = binary.Read(buf, binary.LittleEndian, &errno)
		if errno == -1 {
			panic("no permission")
		}
		return nil, "", fmt.Errorf("Netlink error code %d getting taskstats for %d", errno, nlmsgs[0].Header.Pid)
	}

	var offset int
	payload := nlmsgs[0].Data
	endian := binary.LittleEndian

	var stats TaskStats
	stats.Capturetime = time.Now()

	// these offsets and padding will break if struct taskstats ever changes
	// gen header 0-3
	// attr 4-7
	// attr 8-11
	tgid := endian.Uint32(payload[12:16])
	// attr 16-19

	offset = 20

	offset += 2 // version
	offset += 2 // 2 byte padding
	offset += 4 // exit code
	offset++    // flag
	offset++    // nice
	offset += 6 // 6 byte padding
	stats.Cpudelaycount = endian.Uint64(payload[offset : offset+8])
	offset += 8
	stats.Cpudelaytotal = endian.Uint64(payload[offset : offset+8])
	offset += 8
	stats.Blkiodelaycount = endian.Uint64(payload[offset : offset+8])
	offset += 8
	stats.Blkiodelaytotal = endian.Uint64(payload[offset : offset+8])
	offset += 8
	stats.Swapindelaycount = endian.Uint64(payload[offset : offset+8])
	offset += 8
	stats.Swapindelaytotal = endian.Uint64(payload[offset : offset+8])
	offset += 8
	offset += 8 // cpu run real total
	offset += 8 // cpu run virtual total
	comm := stringFromBytes(payload[offset : offset+32])
	offset += 32 // comm
	offset++     // sched
	offset += 7  // 7 byte padding
	offset += 4  // uid
	offset += 4  // gid
	pid := endian.Uint32(payload[offset : offset+4])
	offset += 4
	if pid != tgid {
		fmt.Printf("read value for unexpected pid %d != %d %+v\n", pid, tgid, stats)
	}
	offset += 4 // etime
	offset += 4 // btime
	offset += 4 // 4 byte padding
	offset += 8 // etime
	offset += 8 // utime
	offset += 8 // stime
	stats.Minflt = endian.Uint64(payload[offset : offset+8])
	offset += 8
	stats.Majflt = endian.Uint64(payload[offset : offset+8])
	offset += 8
	offset += 8 // coremem
	offset += 8 // virtmem
	offset += 8 // hiwater rss
	offset += 8 // hiwater vsz
	stats.Readchar = endian.Uint64(payload[offset : offset+8])
	offset += 8
	stats.Writechar = endian.Uint64(payload[offset : offset+8])
	offset += 8
	stats.Readsyscalls = endian.Uint64(payload[offset : offset+8])
	offset += 8
	stats.Writesyscalls = endian.Uint64(payload[offset : offset+8])
	offset += 8
	stats.Readbytes = endian.Uint64(payload[offset : offset+8])
	offset += 8
	stats.Writebytes = endian.Uint64(payload[offset : offset+8])
	offset += 8
	stats.Cancelledwritebytes = endian.Uint64(payload[offset : offset+8])
	offset += 8
	stats.Nvcsw = endian.Uint64(payload[offset : offset+8])
	offset += 8
	stats.Nivcsw = endian.Uint64(payload[offset : offset+8])
	offset += 8
	offset += 8 // utimescaled
	offset += 8 // stimescaled
	offset += 8 // cputimescaled
	stats.Freepagesdelaycount = endian.Uint64(payload[offset : offset+8])
	offset += 8
	stats.Freepagesdelaytotal = endian.Uint64(payload[offset : offset+8])
	offset += 8

	return &stats, comm, nil
}

var (
	systemEndianness = binary.LittleEndian
	globalSeq        = uint32(0)
)

// Send a genl taskstats message and hope that Linux doesn't change this layout in the future
func sendGetTaskstatsMessage(conn *NLConn, pid int) error {
	globalSeq++

	// this packet: is nl header(16) + genl header(4) + attribute(8) = 28
	outBytes := make([]byte, 28)

	// NL header
	binary.LittleEndian.PutUint32(outBytes, uint32(syscall.NLMSG_HDRLEN+4+8)) // len: 4 for genl, 8 for attr
	binary.LittleEndian.PutUint16(outBytes[4:], conn.genlFamily)              // type
	binary.LittleEndian.PutUint16(outBytes[6:], syscall.NLM_F_REQUEST)        // flags
	binary.LittleEndian.PutUint32(outBytes[8:], globalSeq)                    // seq
	binary.LittleEndian.PutUint32(outBytes[12:], uint32(conn.pid))            // pid

	// genl header
	outBytes[16] = C.TASKSTATS_CMD_GET      // command
	outBytes[17] = C.TASKSTATS_GENL_VERSION // version
	// 18 and 19 are reserved

	// attribute can be many things, but this one is 8 bytes of pure joy:
	//    len uint16 (always 8)
	//    cmd uint16 (always C.TASKSTATS_CMD_ATTR_PID)
	//    pid uint32 actual pid we want
	binary.LittleEndian.PutUint16(outBytes[20:], 8)
	binary.LittleEndian.PutUint16(outBytes[22:], C.TASKSTATS_CMD_ATTR_PID)
	binary.LittleEndian.PutUint32(outBytes[24:], uint32(pid))

	_, err := conn.Write(outBytes)
	return err
}

func TaskStatsLookupPid(conn *NLConn, pid int) (*TaskStats, string, error) {
	sendGetTaskstatsMessage(conn, pid)
	return readGetTaskstatsMessage(conn)
}

func readGetFamilyMessage(conn *NLConn) (uint16, error) {
	inBytes, err := conn.Read()
	if err != nil {
		return 0, err
	}
	if len(inBytes) <= 0 {
		return 0, fmt.Errorf("short read requesting genl family name: %d bytes", len(inBytes))
	}
	nlmsgs, err := syscall.ParseNetlinkMessage(inBytes)
	if err != nil {
		return 0, err
	}

	if len(nlmsgs) != 1 {
		panic(fmt.Sprint("got unexpected response size from get genl family request: ", len(nlmsgs)))
	}

	if nlmsgs[0].Header.Type == syscall.NLMSG_ERROR {
		var errno int32
		buf := bytes.NewBuffer(nlmsgs[0].Data)
		_ = binary.Read(buf, binary.LittleEndian, &errno)
		return 0, fmt.Errorf("Netlink error code %d getting TASKSTATS family id", errno)
	}
	skipLen := binary.LittleEndian.Uint16(nlmsgs[0].Data[4:])
	payloadType := binary.LittleEndian.Uint16(nlmsgs[0].Data[skipLen+8:])
	if payloadType != C.CTRL_ATTR_FAMILY_ID {
		return 0, fmt.Errorf("Netlink error: got unexpected genl attribute: %d", payloadType)
	}
	genlFamily := binary.LittleEndian.Uint16(nlmsgs[0].Data[skipLen+8+2:])
	return genlFamily, nil
}

// Send a genl taskstats message to get all genl families
func sendGetFamilyCmdMessage(conn *NLConn) error {
	globalSeq++
	genlName := []byte("TASKSTATS")
	genlName = append(genlName, 0, 0, 0)

	// this packet: is nl header(16) + genl header(4) + attribute(16) = 36
	outBytes := make([]byte, 36)

	// NL header
	binary.LittleEndian.PutUint32(outBytes, uint32(syscall.NLMSG_HDRLEN+4+16)) // len: 4 for genl, 16 for attr
	binary.LittleEndian.PutUint16(outBytes[4:], conn.family)                   // type
	binary.LittleEndian.PutUint16(outBytes[6:], syscall.NLM_F_REQUEST)         // flags
	binary.LittleEndian.PutUint32(outBytes[8:], globalSeq)                     // seq
	binary.LittleEndian.PutUint32(outBytes[12:], uint32(conn.pid))             // pid

	// genl header
	outBytes[16] = C.CTRL_CMD_GETFAMILY     // command
	outBytes[17] = C.TASKSTATS_GENL_VERSION // version
	// 18 and 19 are reserved

	// attribute can be many things, but this one is 8 bytes of pure joy:
	//    len uint16 (always 8)
	//    cmd uint16 (always genl.TASKSTATS_CMD_ATTR_PID)
	//    pid uint32 actual pid we want
	binary.LittleEndian.PutUint16(outBytes[20:], 11+syscall.NLA_HDRLEN)
	binary.LittleEndian.PutUint16(outBytes[22:], C.CTRL_ATTR_FAMILY_NAME)
	copy(outBytes[24:], genlName)
	_, err := conn.Write(outBytes)
	return err
}

func getGenlFamily(conn *NLConn) uint16 {
	err := sendGetFamilyCmdMessage(conn)
	if err != nil {
		panic(err)
	}
	gfamily, err := readGetFamilyMessage(conn)
	if err != nil {
		panic(err)
	}
	return gfamily
}

// NLConn holds the context necessary to pass around to external callers
type NLConn struct {
	fd         int
	family     uint16
	genlFamily uint16
	addr       syscall.SockaddrNetlink
	pid        int
	readBuf    []byte
}

func (s NLConn) Read() ([]byte, error) {
	n, _, err := syscall.Recvfrom(s.fd, s.readBuf, 0)
	return s.readBuf[:n], os.NewSyscallError("recvfrom", err)
}

func (s NLConn) Write(b []byte) (n int, err error) {
	e := syscall.Sendto(s.fd, b, 0, &s.addr)
	return len(b), os.NewSyscallError("sendto", e)
}

func (s NLConn) Close() error {
	e := syscall.Close(s.fd)
	return os.NewSyscallError("close", e)
}

func (s NLConn) String() string {
	return fmt.Sprintf("fd=%d family=%d genlFamily=%d pid=%d", s.fd, s.family, s.genlFamily, s.pid)
}

// NLInit sets up a new taskstats netlink socket
// All errors are fatal.
func NLInit() *NLConn {
	fd, err := syscall.Socket(syscall.AF_NETLINK, syscall.SOCK_DGRAM, syscall.NETLINK_GENERIC)
	if err != nil {
		panic(os.NewSyscallError("socket", err))
	}
	conn := NLConn{}
	conn.fd = fd
	conn.family = syscall.NETLINK_GENERIC
	conn.addr.Family = syscall.AF_NETLINK
	conn.addr.Pid = 0
	conn.addr.Groups = 0
	conn.pid = os.Getpid()
	conn.readBuf = make([]byte, 4096)
	err = syscall.Bind(fd, &conn.addr)
	if err != nil {
		panic(os.NewSyscallError("bind", err))
	}

	conn.genlFamily = getGenlFamily(&conn)

	return &conn
}
<commit_msg>Fix constants for older Linux<commit_after>// Copyright (c) 2016 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package cpustat

// #include <linux/taskstats.h>
import "C"

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"os"
	"syscall"
	"time"
)

// On older Linux systems, including linux/genetlink.h and taskstats.h doesn't compile.
// To fix this, we define these three symbols here. Note that they just happen to be
// sequential, but they are from 3 different enums.
const (
	CTRL_ATTR_FAMILY_ID   = 1
	CTRL_ATTR_FAMILY_NAME = 2
	CTRL_CMD_GETFAMILY    = 3
)

// convert a byte slice of a null terminated C string into a Go string
func stringFromBytes(c []byte) string {
	nullPos := 0
	i := 0
	for ; i < len(c); i++ {
		if c[i] == 0 {
			nullPos = i
			break
		}
	}
	return string(c[:nullPos])
}

func readGetTaskstatsMessage(conn *NLConn) (*TaskStats, string, error) {
	inBytes, err := conn.Read()
	if err != nil {
		return nil, "", err
	}
	if len(inBytes) <= 0 {
		return nil, "", fmt.Errorf("short read requesting taskstats info: %d bytes", len(inBytes))
	}
	nlmsgs, err := syscall.ParseNetlinkMessage(inBytes)
	if err != nil {
		return nil, "", err
	}

	if len(nlmsgs) != 1 {
		panic(fmt.Sprint("got unexpected response size from get genl taskstats request: ", len(nlmsgs)))
	}

	if nlmsgs[0].Header.Type == syscall.NLMSG_ERROR {
		var errno int32
		buf := bytes.NewBuffer(nlmsgs[0].Data)
		_ = binary.Read(buf, binary.LittleEndian, &errno)
		if errno == -1 {
			panic("no permission")
		}
		return nil, "", fmt.Errorf("Netlink error code %d getting taskstats for %d", errno, nlmsgs[0].Header.Pid)
	}

	var offset int
	payload := nlmsgs[0].Data
	endian := binary.LittleEndian

	var stats TaskStats
	stats.Capturetime = time.Now()

	// these offsets and padding will break if struct taskstats ever changes
	// gen header 0-3
	// attr 4-7
	// attr 8-11
	tgid := endian.Uint32(payload[12:16])
	// attr 16-19

	offset = 20

	offset += 2 // version
	offset += 2 // 2 byte padding
	offset += 4 // exit code
	offset++    // flag
	offset++    // nice
	offset += 6 // 6 byte padding
	stats.Cpudelaycount = endian.Uint64(payload[offset : offset+8])
	offset += 8
	stats.Cpudelaytotal = endian.Uint64(payload[offset : offset+8])
	offset += 8
	stats.Blkiodelaycount = endian.Uint64(payload[offset : offset+8])
	offset += 8
	stats.Blkiodelaytotal = endian.Uint64(payload[offset : offset+8])
	offset += 8
	stats.Swapindelaycount = endian.Uint64(payload[offset : offset+8])
	offset += 8
	stats.Swapindelaytotal = endian.Uint64(payload[offset : offset+8])
	offset += 8
	offset += 8 // cpu run real total
	offset += 8 // cpu run virtual total
	comm := stringFromBytes(payload[offset : offset+32])
	offset += 32 // comm
	offset++     // sched
	offset += 7  // 7 byte padding
	offset += 4  // uid
	offset += 4  // gid
	pid := endian.Uint32(payload[offset : offset+4])
	offset += 4
	if pid != tgid {
		fmt.Printf("read value for unexpected pid %d != %d %+v\n", pid, tgid, stats)
	}
	offset += 4 // etime
	offset += 4 // btime
	offset += 4 // 4 byte padding
	offset += 8 // etime
	offset += 8 // utime
	offset += 8 // stime
	stats.Minflt = endian.Uint64(payload[offset : offset+8])
	offset += 8
	stats.Majflt = endian.Uint64(payload[offset : offset+8])
	offset += 8
	offset += 8 // coremem
	offset += 8 // virtmem
	offset += 8 // hiwater rss
	offset += 8 // hiwater vsz
	stats.Readchar = endian.Uint64(payload[offset : offset+8])
	offset += 8
	stats.Writechar = endian.Uint64(payload[offset : offset+8])
	offset += 8
	stats.Readsyscalls = endian.Uint64(payload[offset : offset+8])
	offset += 8
	stats.Writesyscalls = endian.Uint64(payload[offset : offset+8])
	offset += 8
	stats.Readbytes = endian.Uint64(payload[offset : offset+8])
	offset += 8
	stats.Writebytes = endian.Uint64(payload[offset : offset+8])
	offset += 8
	stats.Cancelledwritebytes = endian.Uint64(payload[offset : offset+8])
	offset += 8
	stats.Nvcsw = endian.Uint64(payload[offset : offset+8])
	offset += 8
	stats.Nivcsw = endian.Uint64(payload[offset : offset+8])
	offset += 8
	offset += 8 // utimescaled
	offset += 8 // stimescaled
	offset += 8 // cputimescaled
	stats.Freepagesdelaycount = endian.Uint64(payload[offset : offset+8])
	offset += 8
	stats.Freepagesdelaytotal = endian.Uint64(payload[offset : offset+8])
	offset += 8

	return &stats, comm, nil
}

var (
	systemEndianness = binary.LittleEndian
	globalSeq        = uint32(0)
)

// Send a genl taskstats message and hope that Linux doesn't change this layout in the future
func sendGetTaskstatsMessage(conn *NLConn, pid int) error {
	globalSeq++

	// this packet: is nl header(16) + genl header(4) + attribute(8) = 28
	outBytes := make([]byte, 28)

	// NL header
	binary.LittleEndian.PutUint32(outBytes, uint32(syscall.NLMSG_HDRLEN+4+8)) // len: 4 for genl, 8 for attr
	binary.LittleEndian.PutUint16(outBytes[4:], conn.genlFamily)              // type
	binary.LittleEndian.PutUint16(outBytes[6:], syscall.NLM_F_REQUEST)        // flags
	binary.LittleEndian.PutUint32(outBytes[8:], globalSeq)                    // seq
	binary.LittleEndian.PutUint32(outBytes[12:], uint32(conn.pid))            // pid

	// genl header
	outBytes[16] = C.TASKSTATS_CMD_GET      // command
	outBytes[17] = C.TASKSTATS_GENL_VERSION // version
	// 18 and 19 are reserved

	// attribute can be many things, but this one is 8 bytes of pure joy:
	//    len uint16 (always 8)
	//    cmd uint16 (always C.TASKSTATS_CMD_ATTR_PID)
	//    pid uint32 actual pid we want
	binary.LittleEndian.PutUint16(outBytes[20:], 8)
	binary.LittleEndian.PutUint16(outBytes[22:], C.TASKSTATS_CMD_ATTR_PID)
	binary.LittleEndian.PutUint32(outBytes[24:], uint32(pid))

	_, err := conn.Write(outBytes)
	return err
}

func TaskStatsLookupPid(conn *NLConn, pid int) (*TaskStats, string, error) {
	sendGetTaskstatsMessage(conn, pid)
	return readGetTaskstatsMessage(conn)
}

func readGetFamilyMessage(conn *NLConn) (uint16, error) {
	inBytes, err := conn.Read()
	if err != nil {
		return 0, err
	}
	if len(inBytes) <= 0 {
		return 0, fmt.Errorf("short read requesting genl family name: %d bytes", len(inBytes))
	}
	nlmsgs, err := syscall.ParseNetlinkMessage(inBytes)
	if err != nil {
		return 0, err
	}

	if len(nlmsgs) != 1 {
		panic(fmt.Sprint("got unexpected response size from get genl family request: ", len(nlmsgs)))
	}

	if nlmsgs[0].Header.Type == syscall.NLMSG_ERROR {
		var errno int32
		buf := bytes.NewBuffer(nlmsgs[0].Data)
		_ = binary.Read(buf, binary.LittleEndian, &errno)
		return 0, fmt.Errorf("Netlink error code %d getting TASKSTATS family id", errno)
	}
	skipLen := binary.LittleEndian.Uint16(nlmsgs[0].Data[4:])
	payloadType := binary.LittleEndian.Uint16(nlmsgs[0].Data[skipLen+8:])
	if payloadType != CTRL_ATTR_FAMILY_ID {
		return 0, fmt.Errorf("Netlink error: got unexpected genl attribute: %d", payloadType)
	}
	genlFamily := binary.LittleEndian.Uint16(nlmsgs[0].Data[skipLen+8+2:])
	return genlFamily, nil
}

// Send a genl taskstats message to get all genl families
func sendGetFamilyCmdMessage(conn *NLConn) error {
	globalSeq++
	genlName := []byte("TASKSTATS")
	genlName = append(genlName, 0, 0, 0)

	// this packet: is nl header(16) + genl header(4) + attribute(16) = 36
	outBytes := make([]byte, 36)

	// NL header
	binary.LittleEndian.PutUint32(outBytes, uint32(syscall.NLMSG_HDRLEN+4+16)) // len: 4 for genl, 16 for attr
	binary.LittleEndian.PutUint16(outBytes[4:], conn.family)                   // type
	binary.LittleEndian.PutUint16(outBytes[6:], syscall.NLM_F_REQUEST)         // flags
	binary.LittleEndian.PutUint32(outBytes[8:], globalSeq)                     // seq
	binary.LittleEndian.PutUint32(outBytes[12:], uint32(conn.pid))             // pid

	// genl header
	outBytes[16] = CTRL_CMD_GETFAMILY       // command
	outBytes[17] = C.TASKSTATS_GENL_VERSION // version
	// 18 and 19 are reserved

	// attribute can be many things, but this one is 8 bytes of pure joy:
	//    len uint16 (always 8)
	//    cmd uint16 (always genl.TASKSTATS_CMD_ATTR_PID)
	//    pid uint32 actual pid we want
	binary.LittleEndian.PutUint16(outBytes[20:], 11+syscall.NLA_HDRLEN)
	binary.LittleEndian.PutUint16(outBytes[22:], CTRL_ATTR_FAMILY_NAME)
	copy(outBytes[24:], genlName)
	_, err := conn.Write(outBytes)
	return err
}

func getGenlFamily(conn *NLConn) uint16 {
	err := sendGetFamilyCmdMessage(conn)
	if err != nil {
		panic(err)
	}
	gfamily, err := readGetFamilyMessage(conn)
	if err != nil {
		panic(err)
	}
	return gfamily
}

// NLConn holds the context necessary to pass around to external callers
type NLConn struct {
	fd         int
	family     uint16
	genlFamily uint16
	addr       syscall.SockaddrNetlink
	pid        int
	readBuf    []byte
}

func (s NLConn) Read() ([]byte, error) {
	n, _, err := syscall.Recvfrom(s.fd, s.readBuf, 0)
	return s.readBuf[:n], os.NewSyscallError("recvfrom", err)
}

func (s NLConn) Write(b []byte) (n int, err error) {
	e := syscall.Sendto(s.fd, b, 0, &s.addr)
	return len(b), os.NewSyscallError("sendto", e)
}

func (s NLConn) Close() error {
	e := syscall.Close(s.fd)
	return os.NewSyscallError("close", e)
}

func (s NLConn) String() string {
	return fmt.Sprintf("fd=%d family=%d genlFamily=%d pid=%d", s.fd, s.family, s.genlFamily, s.pid)
}

// NLInit sets up a new taskstats netlink socket
// All errors are fatal.
func NLInit() *NLConn {
	fd, err := syscall.Socket(syscall.AF_NETLINK, syscall.SOCK_DGRAM, syscall.NETLINK_GENERIC)
	if err != nil {
		panic(os.NewSyscallError("socket", err))
	}
	conn := NLConn{}
	conn.fd = fd
	conn.family = syscall.NETLINK_GENERIC
	conn.addr.Family = syscall.AF_NETLINK
	conn.addr.Pid = 0
	conn.addr.Groups = 0
	conn.pid = os.Getpid()
	conn.readBuf = make([]byte, 4096)
	err = syscall.Bind(fd, &conn.addr)
	if err != nil {
		panic(os.NewSyscallError("bind", err))
	}

	conn.genlFamily = getGenlFamily(&conn)

	return &conn
}
<|endoftext|>