<commit_before>package cors

import (
	"path/filepath"
	"strings"

	"goa.design/goa/codegen"
	"goa.design/goa/codegen/service"
	"goa.design/goa/eval"
	httpcodegen "goa.design/goa/http/codegen"
	httpdesign "goa.design/goa/http/design"
	"goa.design/plugins/cors/design"
)

// ServicesData holds the all the ServiceData indexed by service name.
var ServicesData = make(map[string]*ServiceData)

type (
	// ServiceData contains the data necessary to generate origin handlers
	ServiceData struct {
		// Name is the name of the service.
		Name string
		// Origins is a list of origin expressions defined in API and service levels.
		Origins []*design.OriginExpr
		// OriginHandler is the name of the handler function that sets CORS headers.
		OriginHandler string
		// PreflightPaths is the list of paths that should handle OPTIONS requests.
		PreflightPaths []string
		// Endpoint is the CORS endpoint data.
		Endpoint *httpcodegen.EndpointData
	}
)

const pluginName = "cors"

// Register the plugin Generator functions.
func init() {
	codegen.RegisterPlugin(pluginName, "gen", Generate)
	codegen.RegisterPlugin(pluginName, "example", Example)
}

// Generate produces server code that handle preflight requests and updates
// the HTTP responses with the appropriate CORS headers.
func Generate(genpkg string, roots []eval.Root, files []*codegen.File) ([]*codegen.File, error) {
	for _, root := range roots {
		switch r := root.(type) {
		case *httpdesign.RootExpr:
			for _, s := range r.HTTPServices {
				name := s.Name()
				ServicesData[name] = BuildServiceData(name)
			}
			for _, f := range files {
				ServerCORS(f)
			}
		}
	}
	return files, nil
}

// Example modifies the generated main function so that the services are
// created to handle CORS.
func Example(genpkg string, roots []eval.Root, files []*codegen.File) ([]*codegen.File, error) {
	for _, root := range roots {
		switch r := root.(type) {
		case *httpdesign.RootExpr:
			for _, s := range r.HTTPServices {
				name := s.Name()
				ServicesData[name] = BuildServiceData(name)
			}
		}
	}
	for _, f := range files {
		for _, s := range f.Section("service-main") {
			data := s.Data.(map[string]interface{})
			svcs := data["Services"].([]*httpcodegen.ServiceData)
			for _, sdata := range svcs {
				sdata.Endpoints = append(sdata.Endpoints, ServicesData[sdata.Service.Name].Endpoint)
			}
		}
	}
	return files, nil
}

// BuildServiceData builds the data needed to render the CORS handlers.
func BuildServiceData(name string) *ServiceData {
	preflights := design.PreflightPaths(name)
	data := ServiceData{
		Name:           name,
		Origins:        design.Origins(name),
		PreflightPaths: design.PreflightPaths(name),
		OriginHandler:  "handle" + codegen.Goify(name, true) + "Origin",
		Endpoint: &httpcodegen.EndpointData{
			Method: &service.MethodData{
				VarName: "CORS",
			},
			MountHandler: "MountCORSHandler",
			HandlerInit:  "NewCORSHandler",
		},
	}
	for _, p := range preflights {
		data.Endpoint.Routes = append(data.Endpoint.Routes, &httpcodegen.RouteData{Verb: "OPTIONS", Path: p})
	}
	return &data
}

// ServerCORS updates the HTTP server file to handle preflight paths and
// adds the required CORS headers to the response.
func ServerCORS(f *codegen.File) {
	if filepath.Base(f.Path) != "server.go" {
		return
	}

	var svcData *ServiceData
	for _, s := range f.Section("server-struct") {
		codegen.AddImport(f.SectionTemplates[0],
			&codegen.ImportSpec{Path: "goa.design/plugins/cors"})

		data := s.Data.(*httpcodegen.ServiceData)
		svcData = ServicesData[data.Service.Name]
		for _, o := range svcData.Origins {
			if o.Regexp {
				codegen.AddImport(f.SectionTemplates[0],
					&codegen.ImportSpec{Path: "regexp"})
				break
			}
		}
		data.Endpoints = append(data.Endpoints, svcData.Endpoint)
		fm := codegen.TemplateFuncs()
		f.SectionTemplates = append(f.SectionTemplates, &codegen.SectionTemplate{
			Name:    "mount-cors",
			Source:  mountCORST,
			Data:    svcData,
			FuncMap: fm,
		})
		f.SectionTemplates = append(f.SectionTemplates, &codegen.SectionTemplate{
			Name:    "cors-handler-init",
			Source:  corsHandlerInitT,
			Data:    svcData,
			FuncMap: fm,
		})
		fm["join"] = strings.Join
		f.SectionTemplates = append(f.SectionTemplates, &codegen.SectionTemplate{
			Name:    "handle-cors",
			Source:  handleCORST,
			Data:    svcData,
			FuncMap: fm,
		})
	}
	for _, s := range f.Section("server-init") {
		s.Source = strings.Replace(s.Source,
			"e.{{ .Method.VarName }}, mux, {{ if .MultipartRequestDecoder }}{{ .MultipartRequestDecoder.InitName }}(mux, {{ .MultipartRequestDecoder.VarName }}){{ else }}dec{{ end }}, enc, eh",
			`{{ if ne .Method.VarName "CORS" }}e.{{ .Method.VarName }}, mux, {{ if .MultipartRequestDecoder }}{{ .MultipartRequestDecoder.InitName }}({{ .MultipartRequestDecoder.VarName }}){{ else }}dec{{ end }}, enc, eh{{ end }}`,
			-1)
	}
	for _, s := range f.Section("server-handler") {
		s.Source = strings.Replace(s.Source, "h.(http.HandlerFunc)", svcData.OriginHandler+"(h).(http.HandlerFunc)", -1)
	}
	for _, s := range f.Section("server-files") {
		s.Source = strings.Replace(s.Source, "h.ServeHTTP", svcData.OriginHandler+"(h).ServeHTTP", -1)
	}
}

// Data: ServiceData
var corsHandlerInitT = `{{ printf "%s creates a HTTP handler which returns a simple 200 response." .Endpoint.HandlerInit | comment }}
func {{ .Endpoint.HandlerInit }}() http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(200)
	})
}
`

// Data: ServiceData
var mountCORST = `{{ printf "%s configures the mux to serve the CORS endpoints for the service %s." .Endpoint.MountHandler .Name | comment }}
func {{ .Endpoint.MountHandler }}(mux goahttp.Muxer, h http.Handler) {
	h = {{ .OriginHandler }}(h)
	f, ok := h.(http.HandlerFunc)
	if !ok {
		f = func(w http.ResponseWriter, r *http.Request) {
			h.ServeHTTP(w, r)
		}
	}
	{{- range $p := .PreflightPaths }}
	mux.Handle("OPTIONS", "{{ $p }}", f)
	{{- end }}
}
`

// Data: ServiceData
var handleCORST = `{{ printf "%s applies the CORS response headers corresponding to the origin for the service %s." .OriginHandler .Name | comment }}
func {{ .OriginHandler }}(h http.Handler) http.Handler {
{{- range $i, $policy := .Origins }}
	{{- if $policy.Regexp }}
	spec{{$i}} := regexp.MustCompile({{ printf "%q" $policy.Origin }})
	{{- end }}
{{- end }}
	origHndlr := h.(http.HandlerFunc)
  return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    origin := r.Header.Get("Origin")
    if origin == "" {
      // Not a CORS request
			origHndlr(w, r)
			return
    }
	{{- range $i, $policy := .Origins }}
		{{- if $policy.Regexp }}
		if cors.MatchOriginRegexp(origin, spec{{$i}}) {
		{{- else }}
		if cors.MatchOrigin(origin, {{ printf "%q" $policy.Origin }}) {
		{{- end }}
      w.Header().Set("Access-Control-Allow-Origin", origin)
			{{- if not (eq $policy.Origin "*") }}
			w.Header().Set("Vary", "Origin")
			{{- end }}
			{{- if $policy.Exposed }}
			w.Header().Set("Access-Control-Expose-Headers", "{{ join $policy.Exposed ", " }}")
			{{- end }}
			{{- if gt $policy.MaxAge 0 }}
			w.Header().Set("Access-Control-Max-Age", "{{ $policy.MaxAge }}")
			{{- end }}
			w.Header().Set("Access-Control-Allow-Credentials", "{{ $policy.Credentials }}")
      if acrm := r.Header.Get("Access-Control-Request-Method"); acrm != "" {
        // We are handling a preflight request
				{{- if $policy.Methods }}
				w.Header().Set("Access-Control-Allow-Methods", "{{ join $policy.Methods ", " }}")
				{{- end }}
				{{- if $policy.Headers }}
				w.Header().Set("Access-Control-Allow-Headers", "{{ join $policy.Headers ", " }}")
				{{- end }}
			}
			origHndlr(w, r)
			return
    }
	{{- end }}
		origHndlr(w, r)
		return
  })
}
`
<commit_msg>Fix arguments to multipart decoder<commit_after>package cors

import (
	"path/filepath"
	"strings"

	"goa.design/goa/codegen"
	"goa.design/goa/codegen/service"
	"goa.design/goa/eval"
	httpcodegen "goa.design/goa/http/codegen"
	httpdesign "goa.design/goa/http/design"
	"goa.design/plugins/cors/design"
)

// ServicesData holds the all the ServiceData indexed by service name.
var ServicesData = make(map[string]*ServiceData)

type (
	// ServiceData contains the data necessary to generate origin handlers
	ServiceData struct {
		// Name is the name of the service.
		Name string
		// Origins is a list of origin expressions defined in API and service levels.
		Origins []*design.OriginExpr
		// OriginHandler is the name of the handler function that sets CORS headers.
		OriginHandler string
		// PreflightPaths is the list of paths that should handle OPTIONS requests.
		PreflightPaths []string
		// Endpoint is the CORS endpoint data.
		Endpoint *httpcodegen.EndpointData
	}
)

const pluginName = "cors"

// Register the plugin Generator functions.
func init() {
	codegen.RegisterPlugin(pluginName, "gen", Generate)
	codegen.RegisterPlugin(pluginName, "example", Example)
}

// Generate produces server code that handle preflight requests and updates
// the HTTP responses with the appropriate CORS headers.
func Generate(genpkg string, roots []eval.Root, files []*codegen.File) ([]*codegen.File, error) {
	for _, root := range roots {
		switch r := root.(type) {
		case *httpdesign.RootExpr:
			for _, s := range r.HTTPServices {
				name := s.Name()
				ServicesData[name] = BuildServiceData(name)
			}
			for _, f := range files {
				ServerCORS(f)
			}
		}
	}
	return files, nil
}

// Example modifies the generated main function so that the services are
// created to handle CORS.
func Example(genpkg string, roots []eval.Root, files []*codegen.File) ([]*codegen.File, error) {
	for _, root := range roots {
		switch r := root.(type) {
		case *httpdesign.RootExpr:
			for _, s := range r.HTTPServices {
				name := s.Name()
				ServicesData[name] = BuildServiceData(name)
			}
		}
	}
	for _, f := range files {
		for _, s := range f.Section("service-main") {
			data := s.Data.(map[string]interface{})
			svcs := data["Services"].([]*httpcodegen.ServiceData)
			for _, sdata := range svcs {
				sdata.Endpoints = append(sdata.Endpoints, ServicesData[sdata.Service.Name].Endpoint)
			}
		}
	}
	return files, nil
}

// BuildServiceData builds the data needed to render the CORS handlers.
func BuildServiceData(name string) *ServiceData {
	preflights := design.PreflightPaths(name)
	data := ServiceData{
		Name:           name,
		Origins:        design.Origins(name),
		PreflightPaths: design.PreflightPaths(name),
		OriginHandler:  "handle" + codegen.Goify(name, true) + "Origin",
		Endpoint: &httpcodegen.EndpointData{
			Method: &service.MethodData{
				VarName: "CORS",
			},
			MountHandler: "MountCORSHandler",
			HandlerInit:  "NewCORSHandler",
		},
	}
	for _, p := range preflights {
		data.Endpoint.Routes = append(data.Endpoint.Routes, &httpcodegen.RouteData{Verb: "OPTIONS", Path: p})
	}
	return &data
}

// ServerCORS updates the HTTP server file to handle preflight paths and
// adds the required CORS headers to the response.
func ServerCORS(f *codegen.File) {
	if filepath.Base(f.Path) != "server.go" {
		return
	}

	var svcData *ServiceData
	for _, s := range f.Section("server-struct") {
		codegen.AddImport(f.SectionTemplates[0],
			&codegen.ImportSpec{Path: "goa.design/plugins/cors"})

		data := s.Data.(*httpcodegen.ServiceData)
		svcData = ServicesData[data.Service.Name]
		for _, o := range svcData.Origins {
			if o.Regexp {
				codegen.AddImport(f.SectionTemplates[0],
					&codegen.ImportSpec{Path: "regexp"})
				break
			}
		}
		data.Endpoints = append(data.Endpoints, svcData.Endpoint)
		fm := codegen.TemplateFuncs()
		f.SectionTemplates = append(f.SectionTemplates, &codegen.SectionTemplate{
			Name:    "mount-cors",
			Source:  mountCORST,
			Data:    svcData,
			FuncMap: fm,
		})
		f.SectionTemplates = append(f.SectionTemplates, &codegen.SectionTemplate{
			Name:    "cors-handler-init",
			Source:  corsHandlerInitT,
			Data:    svcData,
			FuncMap: fm,
		})
		fm["join"] = strings.Join
		f.SectionTemplates = append(f.SectionTemplates, &codegen.SectionTemplate{
			Name:    "handle-cors",
			Source:  handleCORST,
			Data:    svcData,
			FuncMap: fm,
		})
	}
	for _, s := range f.Section("server-init") {
		s.Source = strings.Replace(s.Source,
			"e.{{ .Method.VarName }}, mux, {{ if .MultipartRequestDecoder }}{{ .MultipartRequestDecoder.InitName }}(mux, {{ .MultipartRequestDecoder.VarName }}){{ else }}dec{{ end }}, enc, eh",
			`{{ if ne .Method.VarName "CORS" }}e.{{ .Method.VarName }}, mux, {{ if .MultipartRequestDecoder }}{{ .MultipartRequestDecoder.InitName }}(mux, {{ .MultipartRequestDecoder.VarName }}){{ else }}dec{{ end }}, enc, eh{{ end }}`,
			-1)
	}
	for _, s := range f.Section("server-handler") {
		s.Source = strings.Replace(s.Source, "h.(http.HandlerFunc)", svcData.OriginHandler+"(h).(http.HandlerFunc)", -1)
	}
	for _, s := range f.Section("server-files") {
		s.Source = strings.Replace(s.Source, "h.ServeHTTP", svcData.OriginHandler+"(h).ServeHTTP", -1)
	}
}

// Data: ServiceData
var corsHandlerInitT = `{{ printf "%s creates a HTTP handler which returns a simple 200 response." .Endpoint.HandlerInit | comment }}
func {{ .Endpoint.HandlerInit }}() http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(200)
	})
}
`

// Data: ServiceData
var mountCORST = `{{ printf "%s configures the mux to serve the CORS endpoints for the service %s." .Endpoint.MountHandler .Name | comment }}
func {{ .Endpoint.MountHandler }}(mux goahttp.Muxer, h http.Handler) {
	h = {{ .OriginHandler }}(h)
	f, ok := h.(http.HandlerFunc)
	if !ok {
		f = func(w http.ResponseWriter, r *http.Request) {
			h.ServeHTTP(w, r)
		}
	}
	{{- range $p := .PreflightPaths }}
	mux.Handle("OPTIONS", "{{ $p }}", f)
	{{- end }}
}
`

// Data: ServiceData
var handleCORST = `{{ printf "%s applies the CORS response headers corresponding to the origin for the service %s." .OriginHandler .Name | comment }}
func {{ .OriginHandler }}(h http.Handler) http.Handler {
{{- range $i, $policy := .Origins }}
	{{- if $policy.Regexp }}
	spec{{$i}} := regexp.MustCompile({{ printf "%q" $policy.Origin }})
	{{- end }}
{{- end }}
	origHndlr := h.(http.HandlerFunc)
  return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    origin := r.Header.Get("Origin")
    if origin == "" {
      // Not a CORS request
			origHndlr(w, r)
			return
    }
	{{- range $i, $policy := .Origins }}
		{{- if $policy.Regexp }}
		if cors.MatchOriginRegexp(origin, spec{{$i}}) {
		{{- else }}
		if cors.MatchOrigin(origin, {{ printf "%q" $policy.Origin }}) {
		{{- end }}
      w.Header().Set("Access-Control-Allow-Origin", origin)
			{{- if not (eq $policy.Origin "*") }}
			w.Header().Set("Vary", "Origin")
			{{- end }}
			{{- if $policy.Exposed }}
			w.Header().Set("Access-Control-Expose-Headers", "{{ join $policy.Exposed ", " }}")
			{{- end }}
			{{- if gt $policy.MaxAge 0 }}
			w.Header().Set("Access-Control-Max-Age", "{{ $policy.MaxAge }}")
			{{- end }}
			w.Header().Set("Access-Control-Allow-Credentials", "{{ $policy.Credentials }}")
      if acrm := r.Header.Get("Access-Control-Request-Method"); acrm != "" {
        // We are handling a preflight request
				{{- if $policy.Methods }}
				w.Header().Set("Access-Control-Allow-Methods", "{{ join $policy.Methods ", " }}")
				{{- end }}
				{{- if $policy.Headers }}
				w.Header().Set("Access-Control-Allow-Headers", "{{ join $policy.Headers ", " }}")
				{{- end }}
			}
			origHndlr(w, r)
			return
    }
	{{- end }}
		origHndlr(w, r)
		return
  })
}
`
<|endoftext|>


<commit_before>package wrap

import (
	"bytes"
	"crypto/md5"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"strings"
	"time"

	"github.com/Songmu/wrapcommander"
	"github.com/mackerelio/mackerel-agent/config"
	mackerel "github.com/mackerelio/mackerel-client-go"
	"github.com/mackerelio/mkr/logger"
	"golang.org/x/sync/errgroup"
	cli "gopkg.in/urfave/cli.v1"
)

// CommandPlugin is definition of mkr wrap
var Command = cli.Command{
	Name:      "wrap",
	Usage:     "wrap command status",
	ArgsUsage: "[--verbose | -v] [--name | -n <name>] [--memo | -m <memo>] -- /path/to/batch",
	Description: `
    wrap command line
`,
	Action: doWrap,
	Flags: []cli.Flag{
		cli.StringFlag{Name: "name, n", Value: "", Usage: "monitored `check-name` which must be unique on a host"},
		cli.BoolFlag{Name: "verbose, v", Usage: "verbose output"},
		cli.StringFlag{Name: "memo, m", Value: "", Usage: "`memo` of the job"},
		cli.StringFlag{Name: "H, host", Value: "", Usage: "`hostID`"},
		cli.BoolFlag{Name: "warning, w", Usage: "alerts as warning"},
	},
}

func doWrap(c *cli.Context) error {
	confFile := c.GlobalString("conf")
	conf, err := config.LoadConfig(confFile)
	if err != nil {
		return err
	}
	apibase := c.GlobalString("apibase")
	if apibase == "" {
		apibase = conf.Apibase
	}

	apikey := conf.Apikey
	if apikey == "" {
		apikey = os.Getenv("MACKEREL_APIKEY")
	}
	if apikey == "" {
		logger.Log("error", "[mkr wrap] failed to detect Mackerel APIKey. Try to specify in mackerel-agent.conf or export MACKEREL_APIKEY='<Your apikey>'")
	}
	hostID, _ := conf.LoadHostID()
	if c.String("host") != "" {
		hostID = c.String("host")
	}
	if hostID == "" {
		logger.Log("error", "[mkr wrap] failed to load hostID. Try to specify -host option explicitly")
	}
	// Since command execution has the highest priority, even when apikey or
	// hostID is empty, we don't return errors and only output the log here.

	cmd := c.Args()
	if len(cmd) > 0 && cmd[0] == "--" {
		cmd = cmd[1:]
	}
	if len(cmd) < 1 {
		return fmt.Errorf("no commands specified")
	}

	return (&app{
		apibase: apibase,
		name:    c.String("name"),
		verbose: c.Bool("verbose"),
		memo:    c.String("memo"),
		warning: c.Bool("warning"),
		hostID:  hostID,
		apikey:  apikey,
		cmd:     cmd,
	}).run()
}

type app struct {
	apibase string
	name    string
	verbose bool
	memo    string
	warning bool
	hostID  string
	apikey  string
	cmd     []string
}

type result struct {
	Cmd        []string
	Name, Memo string

	Output, Stdout, Stderr string `json:"-"`
	Pid                    int
	ExitCode               int
	Signaled               bool
	StartAt, EndAt         time.Time

	Msg     string
	Success bool
}

var reg = regexp.MustCompile(`[^-a-zA-Z0-9_]`)

func normalizeName(devName string) string {
	return reg.ReplaceAllString(strings.TrimSpace(devName), "_")
}

func (re *result) checkName() string {
	if re.Name != "" {
		return re.Name
	}
	sum := md5.Sum([]byte(strings.Join(re.Cmd, " ")))
	return fmt.Sprintf("mkrwrap-%s-%x",
		normalizeName(filepath.Base(re.Cmd[0])),
		sum[0:3])
}

func (re *result) resultFile() string {
	return filepath.Join(os.TempDir(), fmt.Sprintf("mkrwrap-%s.json", re.checkName()))
}

func (re *result) loadLastResult() (*result, error) {
	prevRe := &result{}
	fname := re.resultFile()

	f, err := os.Open(fname)
	if err != nil {
		if os.IsNotExist(err) {
			return nil, nil
		}
		return nil, err
	}
	defer f.Close()

	err = json.NewDecoder(f).Decode(prevRe)
	return prevRe, err
}

func (re *result) saveResult() error {
	fname := re.resultFile()
	tmpf, err := ioutil.TempFile(filepath.Dir(fname), "tmp-mkrwrap")
	if err != nil {
		return err
	}
	defer func(tmpfname string) {
		tmpf.Close()
		os.Remove(tmpfname)
	}(tmpf.Name())

	if err := json.NewEncoder(tmpf).Encode(re); err != nil {
		return err
	}
	if err := tmpf.Close(); err != nil {
		return err
	}
	return os.Rename(tmpf.Name(), fname)
}

func (ap *app) run() error {
	re := ap.runCmd()
	// TODO keep original exit code
	return ap.report(re)
}

func (ap *app) runCmd() *result {
	cmd := exec.Command(ap.cmd[0], ap.cmd[1:]...)
	re := &result{
		Cmd:  ap.cmd,
		Name: ap.name,
		Memo: ap.memo,
	}

	stdoutPipe, err := cmd.StdoutPipe()
	if err != nil {
		re.Msg = fmt.Sprintf("command invocation failed with follwing error: %s", err)
		re.ExitCode = wrapcommander.ResolveExitCode(err)
		return re
	}
	defer stdoutPipe.Close()

	stderrPipe, err := cmd.StderrPipe()
	if err != nil {
		re.Msg = fmt.Sprintf("command invocation failed with follwing error: %s", err)
		re.ExitCode = wrapcommander.ResolveExitCode(err)
		return re
	}
	defer stderrPipe.Close()

	var (
		bufStdout = &bytes.Buffer{}
		bufStderr = &bytes.Buffer{}
		bufMerged = &bytes.Buffer{}
	)
	stdoutPipe2 := io.TeeReader(stdoutPipe, io.MultiWriter(bufStdout, bufMerged))
	stderrPipe2 := io.TeeReader(stderrPipe, io.MultiWriter(bufStderr, bufMerged))

	re.StartAt = time.Now()
	err = cmd.Start()
	if err != nil {
		re.Msg = fmt.Sprintf("command invocation failed with follwing error: %s", err)
		re.ExitCode = wrapcommander.ResolveExitCode(err)
		return re
	}
	re.Pid = cmd.Process.Pid
	eg := &errgroup.Group{}

	eg.Go(func() error {
		defer stdoutPipe.Close()
		_, err := io.Copy(os.Stdout, stdoutPipe2)
		return err
	})
	eg.Go(func() error {
		defer stderrPipe.Close()
		_, err := io.Copy(os.Stderr, stderrPipe2)
		return err
	})
	eg.Wait()

	cmdErr := cmd.Wait()
	re.EndAt = time.Now()
	re.ExitCode = wrapcommander.ResolveExitCode(cmdErr)
	if re.ExitCode > 128 {
		w, ok := wrapcommander.ErrorToWaitStatus(cmdErr)
		if ok {
			re.Signaled = w.Signaled()
		}
	}
	if !re.Signaled {
		re.Msg = fmt.Sprintf("command exited with code: %d", re.ExitCode)
	} else {
		re.Msg = fmt.Sprintf("command died with signal: %d", re.ExitCode&127)
	}
	re.Stdout = bufStdout.String()
	re.Stderr = bufStderr.String()
	re.Output = bufMerged.String()

	re.Success = re.ExitCode == 0
	return re
}

func (ap *app) report(re *result) error {
	if ap.apikey == "" || ap.hostID == "" {
		return fmt.Errorf("Both of apikey and hostID are needed to report result to Mackerel")
	}
	lastRe, err := re.loadLastResult()
	if err != nil {
		// resultFile something went wrong.
		// It may be no permission, broken json, not a normal file, and so on.
		// Though it rough, try to delete as workaround
		err := os.RemoveAll(re.resultFile())
		if err != nil {
			// XXX report result here?
			return err
		}
	}
	if lastRe == nil || !lastRe.Success || !re.Success {
		ap.doReport(re)
	}
	return re.saveResult()
}

func (ap *app) doReport(re *result) error {
	checkSt := mackerel.CheckStatusOK
	if !re.Success {
		if ap.warning {
			checkSt = mackerel.CheckStatusWarning
		} else {
			checkSt = mackerel.CheckStatusCritical
		}
	}
	msg := re.Msg
	if re.Memo != "" {
		msg += "\nMemo: " + re.Memo
	}
	msg += "\nCommand% " + strings.Join(re.Cmd, " ")
	if ap.verbose {
		msg += "\n" + re.Output
	}
	const messageLengthLimit = 1024
	runes := []rune(msg)
	if len(runes) > messageLengthLimit {
		msg = string(runes[0:messageLengthLimit])
	}
	crs := &mackerel.CheckReports{
		Reports: []*mackerel.CheckReport{
			{
				Source:     mackerel.NewCheckSourceHost(ap.hostID),
				Name:       re.checkName(),
				Status:     checkSt,
				OccurredAt: time.Now().Unix(),
				Message:    msg,
			},
		},
	}
	cli, err := mackerel.NewClientWithOptions(ap.apikey, ap.apibase, false)
	if err != nil {
		return err
	}
	return cli.PostCheckReports(crs)
}
<commit_msg>refactoring<commit_after>package wrap

import (
	"bytes"
	"crypto/md5"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"strings"
	"time"

	"github.com/Songmu/wrapcommander"
	"github.com/mackerelio/mackerel-agent/config"
	mackerel "github.com/mackerelio/mackerel-client-go"
	"github.com/mackerelio/mkr/logger"
	"golang.org/x/sync/errgroup"
	cli "gopkg.in/urfave/cli.v1"
)

// Command is definition of mkr wrap
var Command = cli.Command{
	Name:      "wrap",
	Usage:     "wrap command status",
	ArgsUsage: "[--verbose | -v] [--name | -n <name>] [--memo | -m <memo>] -- /path/to/batch",
	Description: `
    wrap command line
`,
	Action: doWrap,
	Flags: []cli.Flag{
		cli.StringFlag{Name: "name, n", Value: "", Usage: "monitored `check-name` which must be unique on a host"},
		cli.BoolFlag{Name: "verbose, v", Usage: "verbose output"},
		cli.StringFlag{Name: "memo, m", Value: "", Usage: "`memo` of the job"},
		cli.StringFlag{Name: "H, host", Value: "", Usage: "`hostID`"},
		cli.BoolFlag{Name: "warning, w", Usage: "alerts as warning"},
	},
}

func doWrap(c *cli.Context) error {
	confFile := c.GlobalString("conf")
	conf, err := config.LoadConfig(confFile)
	if err != nil {
		return err
	}
	apibase := c.GlobalString("apibase")
	if apibase == "" {
		apibase = conf.Apibase
	}

	apikey := conf.Apikey
	if apikey == "" {
		apikey = os.Getenv("MACKEREL_APIKEY")
	}
	if apikey == "" {
		logger.Log("error", "[mkr wrap] failed to detect Mackerel APIKey. Try to specify in mackerel-agent.conf or export MACKEREL_APIKEY='<Your apikey>'")
	}
	hostID, _ := conf.LoadHostID()
	if c.String("host") != "" {
		hostID = c.String("host")
	}
	if hostID == "" {
		logger.Log("error", "[mkr wrap] failed to load hostID. Try to specify -host option explicitly")
	}
	// Since command execution has the highest priority, even when apikey or
	// hostID is empty, we don't return errors and only output the log here.

	cmd := c.Args()
	if len(cmd) > 0 && cmd[0] == "--" {
		cmd = cmd[1:]
	}
	if len(cmd) < 1 {
		return fmt.Errorf("no commands specified")
	}

	return (&app{
		apibase: apibase,
		name:    c.String("name"),
		verbose: c.Bool("verbose"),
		memo:    c.String("memo"),
		warning: c.Bool("warning"),
		hostID:  hostID,
		apikey:  apikey,
		cmd:     cmd,
	}).run()
}

type app struct {
	apibase string
	name    string
	verbose bool
	memo    string
	warning bool
	hostID  string
	apikey  string
	cmd     []string
}

type result struct {
	Cmd        []string
	Name, Memo string

	Output, Stdout, Stderr string `json:"-"`
	Pid                    int
	ExitCode               int
	Signaled               bool
	StartAt, EndAt         time.Time

	Msg     string
	Success bool
}

var reg = regexp.MustCompile(`[^-a-zA-Z0-9_]`)

func normalizeName(devName string) string {
	return reg.ReplaceAllString(strings.TrimSpace(devName), "_")
}

func (re *result) checkName() string {
	if re.Name != "" {
		return re.Name
	}
	sum := md5.Sum([]byte(strings.Join(re.Cmd, " ")))
	return fmt.Sprintf("mkrwrap-%s-%x",
		normalizeName(filepath.Base(re.Cmd[0])),
		sum[0:3])
}

func (re *result) resultFile() string {
	return filepath.Join(os.TempDir(), fmt.Sprintf("mkrwrap-%s.json", re.checkName()))
}

func (re *result) loadLastResult() (*result, error) {
	prevRe := &result{}
	fname := re.resultFile()

	f, err := os.Open(fname)
	if err != nil {
		if os.IsNotExist(err) {
			return nil, nil
		}
		return nil, err
	}
	defer f.Close()

	err = json.NewDecoder(f).Decode(prevRe)
	return prevRe, err
}

func (re *result) saveResult() error {
	fname := re.resultFile()
	tmpf, err := ioutil.TempFile(filepath.Dir(fname), "tmp-mkrwrap")
	if err != nil {
		return err
	}
	defer func(tmpfname string) {
		tmpf.Close()
		os.Remove(tmpfname)
	}(tmpf.Name())

	if err := json.NewEncoder(tmpf).Encode(re); err != nil {
		return err
	}
	if err := tmpf.Close(); err != nil {
		return err
	}
	return os.Rename(tmpf.Name(), fname)
}

func (re *result) errorEnd(format string, err error) *result {
	re.Msg = fmt.Sprintf(format, err)
	re.ExitCode = wrapcommander.ResolveExitCode(err)
	return re
}

func (ap *app) run() error {
	re := ap.runCmd()
	// TODO keep original exit code
	return ap.report(re)
}

func (ap *app) runCmd() *result {
	cmd := exec.Command(ap.cmd[0], ap.cmd[1:]...)
	re := &result{
		Cmd:  ap.cmd,
		Name: ap.name,
		Memo: ap.memo,
	}

	stdoutPipe, err := cmd.StdoutPipe()
	if err != nil {
		return re.errorEnd("command invocation failed with follwing error: %s", err)
	}
	defer stdoutPipe.Close()

	stderrPipe, err := cmd.StderrPipe()
	if err != nil {
		return re.errorEnd("command invocation failed with follwing error: %s", err)
	}
	defer stderrPipe.Close()

	var (
		bufStdout = &bytes.Buffer{}
		bufStderr = &bytes.Buffer{}
		bufMerged = &bytes.Buffer{}
	)
	stdoutPipe2 := io.TeeReader(stdoutPipe, io.MultiWriter(bufStdout, bufMerged))
	stderrPipe2 := io.TeeReader(stderrPipe, io.MultiWriter(bufStderr, bufMerged))

	re.StartAt = time.Now()
	err = cmd.Start()
	if err != nil {
		return re.errorEnd("command invocation failed with follwing error: %s", err)
	}
	re.Pid = cmd.Process.Pid
	eg := &errgroup.Group{}

	eg.Go(func() error {
		defer stdoutPipe.Close()
		_, err := io.Copy(os.Stdout, stdoutPipe2)
		return err
	})
	eg.Go(func() error {
		defer stderrPipe.Close()
		_, err := io.Copy(os.Stderr, stderrPipe2)
		return err
	})
	eg.Wait()

	cmdErr := cmd.Wait()
	re.EndAt = time.Now()
	re.ExitCode = wrapcommander.ResolveExitCode(cmdErr)
	if re.ExitCode > 128 {
		w, ok := wrapcommander.ErrorToWaitStatus(cmdErr)
		if ok {
			re.Signaled = w.Signaled()
		}
	}
	if !re.Signaled {
		re.Msg = fmt.Sprintf("command exited with code: %d", re.ExitCode)
	} else {
		re.Msg = fmt.Sprintf("command died with signal: %d", re.ExitCode&127)
	}
	re.Stdout = bufStdout.String()
	re.Stderr = bufStderr.String()
	re.Output = bufMerged.String()

	re.Success = re.ExitCode == 0
	return re
}

func (ap *app) report(re *result) error {
	if ap.apikey == "" || ap.hostID == "" {
		return fmt.Errorf("Both of apikey and hostID are needed to report result to Mackerel")
	}
	lastRe, err := re.loadLastResult()
	if err != nil {
		// resultFile something went wrong.
		// It may be no permission, broken json, not a normal file, and so on.
		// Though it rough, try to delete as workaround
		err := os.RemoveAll(re.resultFile())
		if err != nil {
			// XXX report result here?
			return err
		}
	}
	if lastRe == nil || !lastRe.Success || !re.Success {
		ap.doReport(re)
	}
	return re.saveResult()
}

func (ap *app) doReport(re *result) error {
	checkSt := mackerel.CheckStatusOK
	if !re.Success {
		if ap.warning {
			checkSt = mackerel.CheckStatusWarning
		} else {
			checkSt = mackerel.CheckStatusCritical
		}
	}
	msg := re.Msg
	if re.Memo != "" {
		msg += "\nMemo: " + re.Memo
	}
	msg += "\nCommand% " + strings.Join(re.Cmd, " ")
	if ap.verbose {
		msg += "\n" + re.Output
	}
	const messageLengthLimit = 1024
	runes := []rune(msg)
	if len(runes) > messageLengthLimit {
		msg = string(runes[0:messageLengthLimit])
	}
	crs := &mackerel.CheckReports{
		Reports: []*mackerel.CheckReport{
			{
				Source:     mackerel.NewCheckSourceHost(ap.hostID),
				Name:       re.checkName(),
				Status:     checkSt,
				OccurredAt: time.Now().Unix(),
				Message:    msg,
			},
		},
	}
	cli, err := mackerel.NewClientWithOptions(ap.apikey, ap.apibase, false)
	if err != nil {
		return err
	}
	return cli.PostCheckReports(crs)
}
<|endoftext|>


<commit_before>package extkeys

import (
	"encoding/json"
	"fmt"
	"os"
	"testing"
)

type VectorsFile struct {
	Data    map[string][][6]string
	vectors []*Vector
}

type Vector struct {
	language, salt, password, input, mnemonic, seed, xprv string
}

// TestMnemonicPhrase
func TestMnemonicPhrase(t *testing.T) {

	mnemonic := NewMnemonic(Salt)

	// test strength validation
	strengths := []int{127, 129, 257}
	for _, s := range strengths {
		_, err := mnemonic.MnemonicPhrase(s, EnglishLanguage)
		if err != ErrInvalidEntropyStrength {
			t.Errorf("Entropy strength `%d` should be invalid", s)
		}
	}

	// test mnemonic generation
	t.Log("Test mnemonic generation:")
	for _, language := range mnemonic.AvailableLanguages() {
		phrase, err := mnemonic.MnemonicPhrase(128, language)
		t.Logf("Mnemonic (%s): %s", Languages[language], phrase)

		if err != nil {
			t.Errorf("Test failed: could not create seed: %s", err)
		}

		if !mnemonic.ValidMnemonic(phrase, language) {
			t.Error("Seed is not valid Mnenomic")
		}
	}

	// run against test vectors
	vectorsFile, err := LoadVectorsFile("mnemonic_vectors.json")
	if err != nil {
		t.Error(err)
	}

	t.Log("Test against pre-computed seed vectors:")
	stats := map[string]int{}
	for _, vector := range vectorsFile.vectors {
		stats[vector.language] += 1
		mnemonic := NewMnemonic(vector.salt)
		seed := mnemonic.MnemonicSeed(vector.mnemonic, vector.password)
		if fmt.Sprintf("%x", seed) != vector.seed {
			t.Errorf("Test failed (%s): incorrect seed (%x) generated (expected: %s)", vector.language, seed, vector.seed)
			return
		}
		//t.Logf("Test passed: correct seed (%x) generated (expected: %s)", seed, vector.seed)
	}
	for language, count := range stats {
		t.Logf("[%s]: %d tests completed", language, count)
	}
}

func LoadVectorsFile(path string) (*VectorsFile, error) {
	fp, err := os.Open(path)
	if err != nil {
		return nil, fmt.Errorf("Test failed: cannot open vectors file: %s", err)
	}

	var vectorsFile VectorsFile
	if err := json.NewDecoder(fp).Decode(&vectorsFile); err != nil {
		return nil, fmt.Errorf("Test failed: cannot parse vectors file: %s", err)
	}

	// load data into Vector structs
	for language, data := range vectorsFile.Data {
		for _, item := range data {
			vectorsFile.vectors = append(vectorsFile.vectors, &Vector{language, item[0], item[1], item[2], item[3], item[4], item[5]})
		}
	}

	return &vectorsFile, nil
}

func (v *Vector) String() string {
	return fmt.Sprintf("{salt: %s, password: %s, input: %s, mnemonic: %s, seed: %s, xprv: %s}",
		v.salt, v.password, v.input, v.mnemonic, v.seed, v.xprv)
}
<commit_msg>fix test output<commit_after>package extkeys

import (
	"encoding/json"
	"fmt"
	"os"
	"testing"
)

type VectorsFile struct {
	Data    map[string][][6]string
	vectors []*Vector
}

type Vector struct {
	language, salt, password, input, mnemonic, seed, xprv string
}

// TestMnemonicPhrase
func TestMnemonicPhrase(t *testing.T) {

	mnemonic := NewMnemonic(Salt)

	// test strength validation
	strengths := []int{127, 129, 257}
	for _, s := range strengths {
		_, err := mnemonic.MnemonicPhrase(s, EnglishLanguage)
		if err != ErrInvalidEntropyStrength {
			t.Errorf("Entropy strength '%d' should be invalid", s)
		}
	}

	// test mnemonic generation
	t.Log("Test mnemonic generation:")
	for _, language := range mnemonic.AvailableLanguages() {
		phrase, err := mnemonic.MnemonicPhrase(128, language)
		t.Logf("Mnemonic (%s): %s", Languages[language], phrase)

		if err != nil {
			t.Errorf("Test failed: could not create seed: %s", err)
		}

		if !mnemonic.ValidMnemonic(phrase, language) {
			t.Error("Seed is not valid Mnenomic")
		}
	}

	// run against test vectors
	vectorsFile, err := LoadVectorsFile("mnemonic_vectors.json")
	if err != nil {
		t.Error(err)
	}

	t.Log("Test against pre-computed seed vectors:")
	stats := map[string]int{}
	for _, vector := range vectorsFile.vectors {
		stats[vector.language] += 1
		mnemonic := NewMnemonic(vector.salt)
		seed := mnemonic.MnemonicSeed(vector.mnemonic, vector.password)
		if fmt.Sprintf("%x", seed) != vector.seed {
			t.Errorf("Test failed (%s): incorrect seed (%x) generated (expected: %s)", vector.language, seed, vector.seed)
			return
		}
		//t.Logf("Test passed: correct seed (%x) generated (expected: %s)", seed, vector.seed)
	}
	for language, count := range stats {
		t.Logf("[%s]: %d tests completed", language, count)
	}
}

func LoadVectorsFile(path string) (*VectorsFile, error) {
	fp, err := os.Open(path)
	if err != nil {
		return nil, fmt.Errorf("Test failed: cannot open vectors file: %s", err)
	}

	var vectorsFile VectorsFile
	if err := json.NewDecoder(fp).Decode(&vectorsFile); err != nil {
		return nil, fmt.Errorf("Test failed: cannot parse vectors file: %s", err)
	}

	// load data into Vector structs
	for language, data := range vectorsFile.Data {
		for _, item := range data {
			vectorsFile.vectors = append(vectorsFile.vectors, &Vector{language, item[0], item[1], item[2], item[3], item[4], item[5]})
		}
	}

	return &vectorsFile, nil
}

func (v *Vector) String() string {
	return fmt.Sprintf("{salt: %s, password: %s, input: %s, mnemonic: %s, seed: %s, xprv: %s}",
		v.salt, v.password, v.input, v.mnemonic, v.seed, v.xprv)
}
<|endoftext|>


<commit_before>package client

import (
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"net"
	"strings"
	"time"

	"github.com/jcmturner/gokrb5/v8/iana/errorcode"
	"github.com/jcmturner/gokrb5/v8/messages"
)

// SendToKDC performs network actions to send data to the KDC.
func (cl *Client) sendToKDC(b []byte, realm string) ([]byte, error) {
	var rb []byte
	if cl.Config.LibDefaults.UDPPreferenceLimit == 1 {
		//1 means we should always use TCP
		rb, errtcp := cl.sendKDCTCP(realm, b)
		if errtcp != nil {
			if e, ok := errtcp.(messages.KRBError); ok {
				return rb, e
			}
			return rb, fmt.Errorf("communication error with KDC via TCP: %v", errtcp)
		}
		return rb, nil
	}
	if len(b) <= cl.Config.LibDefaults.UDPPreferenceLimit {
		//Try UDP first, TCP second
		rb, errudp := cl.sendKDCUDP(realm, b)
		if errudp != nil {
			if e, ok := errudp.(messages.KRBError); ok && e.ErrorCode != errorcode.KRB_ERR_RESPONSE_TOO_BIG {
				// Got a KRBError from KDC
				// If this is not a KRB_ERR_RESPONSE_TOO_BIG we will return immediately otherwise will try TCP.
				return rb, e
			}
			// Try TCP
			r, errtcp := cl.sendKDCTCP(realm, b)
			if errtcp != nil {
				if e, ok := errtcp.(messages.KRBError); ok {
					// Got a KRBError
					return r, e
				}
				return r, fmt.Errorf("failed to communicate with KDC. Attempts made with UDP (%v) and then TCP (%v)", errudp, errtcp)
			}
			rb = r
		}
		return rb, nil
	}
	//Try TCP first, UDP second
	rb, errtcp := cl.sendKDCTCP(realm, b)
	if errtcp != nil {
		if e, ok := errtcp.(messages.KRBError); ok {
			// Got a KRBError from KDC so returning and not trying UDP.
			return rb, e
		}
		rb, errudp := cl.sendKDCUDP(realm, b)
		if errudp != nil {
			if e, ok := errudp.(messages.KRBError); ok {
				// Got a KRBError
				return rb, e
			}
			return rb, fmt.Errorf("failed to communicate with KDC. Attempts made with TCP (%v) and then UDP (%v)", errtcp, errudp)
		}
	}
	return rb, nil
}

// sendKDCUDP sends bytes to the KDC via UDP.
func (cl *Client) sendKDCUDP(realm string, b []byte) ([]byte, error) {
	var r []byte
	_, kdcs, err := cl.Config.GetKDCs(realm, false)
	if err != nil {
		return r, err
	}
	r, err = dialSendUDP(kdcs, b)
	if err != nil {
		return r, err
	}
	return checkForKRBError(r)
}

// dialSendUDP establishes a UDP connection to a KDC.
func dialSendUDP(kdcs map[int]string, b []byte) ([]byte, error) {
	var errs []string
	for i := 1; i <= len(kdcs); i++ {
		udpAddr, err := net.ResolveUDPAddr("udp", kdcs[i])
		if err != nil {
			errs = append(errs, fmt.Sprintf("error resolving KDC address: %v", err))
			continue
		}

		conn, err := net.DialTimeout("udp", udpAddr.String(), 5*time.Second)
		if err != nil {
			errs = append(errs, fmt.Sprintf("error setting dial timeout on connection to %s: %v", kdcs[i], err))
			continue
		}
		if err := conn.SetDeadline(time.Now().Add(5 * time.Second)); err != nil {
			errs = append(errs, fmt.Sprintf("error setting deadline on connection to %s: %v", kdcs[i], err))
			continue
		}
		// conn is guaranteed to be a UDPConn
		rb, err := sendUDP(conn.(*net.UDPConn), b)
		if err != nil {
			errs = append(errs, fmt.Sprintf("error sneding to %s: %v", kdcs[i], err))
			continue
		}
		return rb, nil
	}
	return nil, fmt.Errorf("error sending to a KDC: %s", strings.Join(errs, "; "))
}

// sendUDP sends bytes to connection over UDP.
func sendUDP(conn *net.UDPConn, b []byte) ([]byte, error) {
	var r []byte
	defer conn.Close()
	_, err := conn.Write(b)
	if err != nil {
		return r, fmt.Errorf("error sending to (%s): %v", conn.RemoteAddr().String(), err)
	}
	udpbuf := make([]byte, 4096)
	n, _, err := conn.ReadFrom(udpbuf)
	r = udpbuf[:n]
	if err != nil {
		return r, fmt.Errorf("sending over UDP failed to %s: %v", conn.RemoteAddr().String(), err)
	}
	if len(r) < 1 {
		return r, fmt.Errorf("no response data from %s", conn.RemoteAddr().String())
	}
	return r, nil
}

// sendKDCTCP sends bytes to the KDC via TCP.
func (cl *Client) sendKDCTCP(realm string, b []byte) ([]byte, error) {
	var r []byte
	_, kdcs, err := cl.Config.GetKDCs(realm, true)
	if err != nil {
		return r, err
	}
	r, err = dialSendTCP(kdcs, b)
	if err != nil {
		return r, err
	}
	return checkForKRBError(r)
}

// dialKDCTCP establishes a TCP connection to a KDC.
func dialSendTCP(kdcs map[int]string, b []byte) ([]byte, error) {
	var errs []string
	for i := 1; i <= len(kdcs); i++ {
		tcpAddr, err := net.ResolveTCPAddr("tcp", kdcs[i])
		if err != nil {
			errs = append(errs, fmt.Sprintf("error resolving KDC address: %v", err))
			continue
		}

		conn, err := net.DialTimeout("tcp", tcpAddr.String(), 5*time.Second)
		if err != nil {
			errs = append(errs, fmt.Sprintf("error setting dial timeout on connection to %s: %v", kdcs[i], err))
			continue
		}
		if err := conn.SetDeadline(time.Now().Add(5 * time.Second)); err != nil {
			errs = append(errs, fmt.Sprintf("error setting deadline on connection to %s: %v", kdcs[i], err))
			continue
		}
		// conn is guaranteed to be a TCPConn
		rb, err := sendTCP(conn.(*net.TCPConn), b)
		if err != nil {
			errs = append(errs, fmt.Sprintf("error sneding to %s: %v", kdcs[i], err))
			continue
		}
		return rb, nil
	}
	if len(errs) > 0 {
		errorString := strings.Join(errs, ",")
		return nil, fmt.Errorf("error in getting a TCP connection to any of the KDCs: %v", errorString)
	}
	return nil, errors.New("error in getting a TCP connection to any of the KDCs")
}

// sendTCP sends bytes to connection over TCP.
func sendTCP(conn *net.TCPConn, b []byte) ([]byte, error) {
	defer conn.Close()
	var r []byte
	// RFC 4120 7.2.2 specifies the first 4 bytes indicate the length of the message in big endian order.
	hb := make([]byte, 4, 4)
	binary.BigEndian.PutUint32(hb, uint32(len(b)))
	b = append(hb, b...)

	_, err := conn.Write(b)
	if err != nil {
		return r, fmt.Errorf("error sending to KDC (%s): %v", conn.RemoteAddr().String(), err)
	}

	sh := make([]byte, 4, 4)
	_, err = conn.Read(sh)
	if err != nil {
		return r, fmt.Errorf("error reading response size header: %v", err)
	}
	s := binary.BigEndian.Uint32(sh)

	rb := make([]byte, s, s)
	_, err = io.ReadFull(conn, rb)
	if err != nil {
		return r, fmt.Errorf("error reading response: %v", err)
	}
	if len(rb) < 1 {
		return r, fmt.Errorf("no response data from KDC %s", conn.RemoteAddr().String())
	}
	return rb, nil
}

// checkForKRBError checks if the response bytes from the KDC are a KRBError.
func checkForKRBError(b []byte) ([]byte, error) {
	var KRBErr messages.KRBError
	if err := KRBErr.Unmarshal(b); err == nil {
		return b, KRBErr
	}
	return b, nil
}
<commit_msg>Update network.go<commit_after>package client

import (
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"net"
	"strings"
	"time"

	"github.com/jcmturner/gokrb5/v8/iana/errorcode"
	"github.com/jcmturner/gokrb5/v8/messages"
)

// SendToKDC performs network actions to send data to the KDC.
func (cl *Client) sendToKDC(b []byte, realm string) ([]byte, error) {
	var rb []byte
	if cl.Config.LibDefaults.UDPPreferenceLimit == 1 {
		//1 means we should always use TCP
		rb, errtcp := cl.sendKDCTCP(realm, b)
		if errtcp != nil {
			if e, ok := errtcp.(messages.KRBError); ok {
				return rb, e
			}
			return rb, fmt.Errorf("communication error with KDC via TCP: %v", errtcp)
		}
		return rb, nil
	}
	if len(b) <= cl.Config.LibDefaults.UDPPreferenceLimit {
		//Try UDP first, TCP second
		rb, errudp := cl.sendKDCUDP(realm, b)
		if errudp != nil {
			if e, ok := errudp.(messages.KRBError); ok && e.ErrorCode != errorcode.KRB_ERR_RESPONSE_TOO_BIG {
				// Got a KRBError from KDC
				// If this is not a KRB_ERR_RESPONSE_TOO_BIG we will return immediately otherwise will try TCP.
				return rb, e
			}
			// Try TCP
			r, errtcp := cl.sendKDCTCP(realm, b)
			if errtcp != nil {
				if e, ok := errtcp.(messages.KRBError); ok {
					// Got a KRBError
					return r, e
				}
				return r, fmt.Errorf("failed to communicate with KDC. Attempts made with UDP (%v) and then TCP (%v)", errudp, errtcp)
			}
			rb = r
		}
		return rb, nil
	}
	//Try TCP first, UDP second
	rb, errtcp := cl.sendKDCTCP(realm, b)
	if errtcp != nil {
		if e, ok := errtcp.(messages.KRBError); ok {
			// Got a KRBError from KDC so returning and not trying UDP.
			return rb, e
		}
		rb, errudp := cl.sendKDCUDP(realm, b)
		if errudp != nil {
			if e, ok := errudp.(messages.KRBError); ok {
				// Got a KRBError
				return rb, e
			}
			return rb, fmt.Errorf("failed to communicate with KDC. Attempts made with TCP (%v) and then UDP (%v)", errtcp, errudp)
		}
	}
	return rb, nil
}

// sendKDCUDP sends bytes to the KDC via UDP.
func (cl *Client) sendKDCUDP(realm string, b []byte) ([]byte, error) {
	var r []byte
	_, kdcs, err := cl.Config.GetKDCs(realm, false)
	if err != nil {
		return r, err
	}
	r, err = dialSendUDP(kdcs, b)
	if err != nil {
		return r, err
	}
	return checkForKRBError(r)
}

// dialSendUDP establishes a UDP connection to a KDC.
func dialSendUDP(kdcs map[int]string, b []byte) ([]byte, error) {
	var errs []string
	for i := 1; i <= len(kdcs); i++ {
		udpAddr, err := net.ResolveUDPAddr("udp", kdcs[i])
		if err != nil {
			errs = append(errs, fmt.Sprintf("error resolving KDC address: %v", err))
			continue
		}

		conn, err := net.DialTimeout("udp", udpAddr.String(), 5*time.Second)
		if err != nil {
			errs = append(errs, fmt.Sprintf("error setting dial timeout on connection to %s: %v", kdcs[i], err))
			continue
		}
		if err := conn.SetDeadline(time.Now().Add(5 * time.Second)); err != nil {
			errs = append(errs, fmt.Sprintf("error setting deadline on connection to %s: %v", kdcs[i], err))
			continue
		}
		// conn is guaranteed to be a UDPConn
		rb, err := sendUDP(conn.(*net.UDPConn), b)
		if err != nil {
			errs = append(errs, fmt.Sprintf("error sneding to %s: %v", kdcs[i], err))
			continue
		}
		return rb, nil
	}
	return nil, fmt.Errorf("error sending to a KDC: %s", strings.Join(errs, "; "))
}

// sendUDP sends bytes to connection over UDP.
func sendUDP(conn *net.UDPConn, b []byte) ([]byte, error) {
	var r []byte
	defer conn.Close()
	_, err := conn.Write(b)
	if err != nil {
		return r, fmt.Errorf("error sending to (%s): %v", conn.RemoteAddr().String(), err)
	}
	udpbuf := make([]byte, 4096)
	n, _, err := conn.ReadFrom(udpbuf)
	r = udpbuf[:n]
	if err != nil {
		return r, fmt.Errorf("sending over UDP failed to %s: %v", conn.RemoteAddr().String(), err)
	}
	if len(r) < 1 {
		return r, fmt.Errorf("no response data from %s", conn.RemoteAddr().String())
	}
	return r, nil
}

// sendKDCTCP sends bytes to the KDC via TCP.
func (cl *Client) sendKDCTCP(realm string, b []byte) ([]byte, error) {
	var r []byte
	_, kdcs, err := cl.Config.GetKDCs(realm, true)
	if err != nil {
		return r, err
	}
	r, err = dialSendTCP(kdcs, b)
	if err != nil {
		return r, err
	}
	return checkForKRBError(r)
}

// dialKDCTCP establishes a TCP connection to a KDC.
func dialSendTCP(kdcs map[int]string, b []byte) ([]byte, error) {
	var errs []string
	for i := 1; i <= len(kdcs); i++ {
		tcpAddr, err := net.ResolveTCPAddr("tcp", kdcs[i])
		if err != nil {
			errs = append(errs, fmt.Sprintf("error resolving KDC address: %v", err))
			continue
		}

		conn, err := net.DialTimeout("tcp", tcpAddr.String(), 5*time.Second)
		if err != nil {
			errs = append(errs, fmt.Sprintf("error setting dial timeout on connection to %s: %v", kdcs[i], err))
			continue
		}
		if err := conn.SetDeadline(time.Now().Add(5 * time.Second)); err != nil {
			errs = append(errs, fmt.Sprintf("error setting deadline on connection to %s: %v", kdcs[i], err))
			continue
		}
		// conn is guaranteed to be a TCPConn
		rb, err := sendTCP(conn.(*net.TCPConn), b)
		if err != nil {
			errs = append(errs, fmt.Sprintf("error sneding to %s: %v", kdcs[i], err))
			continue
		}
		return rb, nil
	}
	return nil, fmt.Errorf("error sending to a KDC: %s", strings.Join(errs, "; "))
}

// sendTCP sends bytes to connection over TCP.
func sendTCP(conn *net.TCPConn, b []byte) ([]byte, error) {
	defer conn.Close()
	var r []byte
	// RFC 4120 7.2.2 specifies the first 4 bytes indicate the length of the message in big endian order.
	hb := make([]byte, 4, 4)
	binary.BigEndian.PutUint32(hb, uint32(len(b)))
	b = append(hb, b...)

	_, err := conn.Write(b)
	if err != nil {
		return r, fmt.Errorf("error sending to KDC (%s): %v", conn.RemoteAddr().String(), err)
	}

	sh := make([]byte, 4, 4)
	_, err = conn.Read(sh)
	if err != nil {
		return r, fmt.Errorf("error reading response size header: %v", err)
	}
	s := binary.BigEndian.Uint32(sh)

	rb := make([]byte, s, s)
	_, err = io.ReadFull(conn, rb)
	if err != nil {
		return r, fmt.Errorf("error reading response: %v", err)
	}
	if len(rb) < 1 {
		return r, fmt.Errorf("no response data from KDC %s", conn.RemoteAddr().String())
	}
	return rb, nil
}

// checkForKRBError checks if the response bytes from the KDC are a KRBError.
func checkForKRBError(b []byte) ([]byte, error) {
	var KRBErr messages.KRBError
	if err := KRBErr.Unmarshal(b); err == nil {
		return b, KRBErr
	}
	return b, nil
}
<|endoftext|>


<commit_before>/*
Copyright 2021 The cert-manager Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package validation

import (
	"context"
	"fmt"
	"strings"

	admissionv1 "k8s.io/api/admission/v1"
	authzv1 "k8s.io/api/authorization/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/util/validation/field"
	authzclient "k8s.io/client-go/kubernetes/typed/authorization/v1"

	"github.com/jetstack/cert-manager/pkg/apis/certmanager"
	cmapi "github.com/jetstack/cert-manager/pkg/internal/apis/certmanager"
	"github.com/jetstack/cert-manager/pkg/internal/apis/certmanager/validation/util"
)

// ReviewApproval will review whether the client is able to approve or deny the
// given request, if indeed they are attempting to. A SubjectAccessReview will
// be performed if the client is attempting to approve/deny the request. An
// error will be returned if the SubjectAccessReview fails, or if they do not
// have permissions to perform the approval/denial.
func ReviewApproval(client authzclient.SubjectAccessReviewInterface,
	req *admissionv1.AdmissionRequest, oldObj, newObj runtime.Object) field.ErrorList {
	oldCR := oldObj.(*cmapi.CertificateRequest)
	newCR := newObj.(*cmapi.CertificateRequest)

	if !isApprovalRequest(oldCR, newCR) {
		return nil
	}

	ok, err := reviewRequest(client, req, newCR)
	if err != nil {
		return field.ErrorList{
			field.InternalError(field.NewPath("status.conditions"), err),
		}
	}

	if !ok {
		return field.ErrorList{
			field.Forbidden(field.NewPath("status.conditions"),
				fmt.Sprintf("user %q does not have permissions to set approved/denied conditions for issuer %v", req.UserInfo.Username, newCR.Spec.IssuerRef),
			),
		}
	}

	return nil
}

// reviewRequest will perform a SubjectAccessReview with the UserInfo fields of
// the client against the issuer of the CertificateRequest. A client must have
// the "approve" verb, for the resource "signer", in the CertificateRequest
// namespace, for the name "<issuer-kind>.<issuer-group>/<issuer-name>".
func reviewRequest(client authzclient.SubjectAccessReviewInterface, req *admissionv1.AdmissionRequest, cr *cmapi.CertificateRequest) (bool, error) {
	extra := make(map[string]authzv1.ExtraValue)
	for k, v := range req.UserInfo.Extra {
		extra[k] = authzv1.ExtraValue(v)
	}

	kind := cr.Spec.IssuerRef.Kind
	if len(kind) == 0 {
		kind = cmapi.IssuerKind
	}

	group := cr.Spec.IssuerRef.Group
	if len(group) == 0 {
		group = certmanager.GroupName
	}

	for _, name := range []string{
		fmt.Sprintf("%s.%s/*", strings.ToLower(kind), group),
		fmt.Sprintf("%s.%s/%s", strings.ToLower(kind), group, cr.Spec.IssuerRef.Name),
	} {
		resp, err := client.Create(context.TODO(), &authzv1.SubjectAccessReview{
			Spec: authzv1.SubjectAccessReviewSpec{
				User:   req.UserInfo.Username,
				Groups: req.UserInfo.Groups,
				Extra:  extra,
				UID:    req.UserInfo.UID,

				ResourceAttributes: &authzv1.ResourceAttributes{
					Group:    certmanager.GroupName,
					Resource: "signers",
					Name:     name,
					Verb:     "approve",
					Version:  "*",
				},
			},
		}, metav1.CreateOptions{})
		if err != nil {
			return false, err
		}

		if resp.Status.Allowed {
			return true, nil
		}
	}

	return false, nil
}

// isApprovalRequest will return true if the request is given a new approved or
// denied condition. This check is strictly concerned with these conditions
// being _added_. We do this to reduce the number of SAR calls made, since
// removal or changing of these conditions will be rejected elsewhere in the
// validation chain locally.
func isApprovalRequest(oldCR, newCR *cmapi.CertificateRequest) bool {
	oldCRApproving := util.GetCertificateRequestCondition(oldCR.Status.Conditions, cmapi.CertificateRequestConditionApproved)
	newCRApproving := util.GetCertificateRequestCondition(newCR.Status.Conditions, cmapi.CertificateRequestConditionApproved)

	if oldCRApproving == nil && newCRApproving != nil {
		return true
	}

	oldCRDenying := util.GetCertificateRequestCondition(oldCR.Status.Conditions, cmapi.CertificateRequestConditionDenied)
	newCRDenying := util.GetCertificateRequestCondition(newCR.Status.Conditions, cmapi.CertificateRequestConditionDenied)

	if oldCRDenying == nil && newCRDenying != nil {
		return true
	}

	return false
}
<commit_msg>Updates approval review comment to correctly state cluster scope and issuer name<commit_after>/*
Copyright 2021 The cert-manager Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package validation

import (
	"context"
	"fmt"
	"strings"

	admissionv1 "k8s.io/api/admission/v1"
	authzv1 "k8s.io/api/authorization/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/util/validation/field"
	authzclient "k8s.io/client-go/kubernetes/typed/authorization/v1"

	"github.com/jetstack/cert-manager/pkg/apis/certmanager"
	cmapi "github.com/jetstack/cert-manager/pkg/internal/apis/certmanager"
	"github.com/jetstack/cert-manager/pkg/internal/apis/certmanager/validation/util"
)

// ReviewApproval will review whether the client is able to approve or deny the
// given request, if indeed they are attempting to. A SubjectAccessReview will
// be performed if the client is attempting to approve/deny the request. An
// error will be returned if the SubjectAccessReview fails, or if they do not
// have permissions to perform the approval/denial.
func ReviewApproval(client authzclient.SubjectAccessReviewInterface,
	req *admissionv1.AdmissionRequest, oldObj, newObj runtime.Object) field.ErrorList {
	oldCR := oldObj.(*cmapi.CertificateRequest)
	newCR := newObj.(*cmapi.CertificateRequest)

	if !isApprovalRequest(oldCR, newCR) {
		return nil
	}

	ok, err := reviewRequest(client, req, newCR)
	if err != nil {
		return field.ErrorList{
			field.InternalError(field.NewPath("status.conditions"), err),
		}
	}

	if !ok {
		return field.ErrorList{
			field.Forbidden(field.NewPath("status.conditions"),
				fmt.Sprintf("user %q does not have permissions to set approved/denied conditions for issuer %v", req.UserInfo.Username, newCR.Spec.IssuerRef),
			),
		}
	}

	return nil
}

// reviewRequest will perform a SubjectAccessReview with the UserInfo fields of
// the client against the issuer of the CertificateRequest. A client must have
// the "approve" verb, for the resource "signer", at the Cluster scope, for the
// name "<issuer-kind>.<issuer-group>/<issuer-name>", or
// "<issuer-kind>.<issuer-group>/*".
func reviewRequest(client authzclient.SubjectAccessReviewInterface, req *admissionv1.AdmissionRequest, cr *cmapi.CertificateRequest) (bool, error) {
	extra := make(map[string]authzv1.ExtraValue)
	for k, v := range req.UserInfo.Extra {
		extra[k] = authzv1.ExtraValue(v)
	}

	kind := cr.Spec.IssuerRef.Kind
	if len(kind) == 0 {
		kind = cmapi.IssuerKind
	}

	group := cr.Spec.IssuerRef.Group
	if len(group) == 0 {
		group = certmanager.GroupName
	}

	for _, name := range []string{
		fmt.Sprintf("%s.%s/*", strings.ToLower(kind), group),
		fmt.Sprintf("%s.%s/%s", strings.ToLower(kind), group, cr.Spec.IssuerRef.Name),
	} {
		resp, err := client.Create(context.TODO(), &authzv1.SubjectAccessReview{
			Spec: authzv1.SubjectAccessReviewSpec{
				User:   req.UserInfo.Username,
				Groups: req.UserInfo.Groups,
				Extra:  extra,
				UID:    req.UserInfo.UID,

				ResourceAttributes: &authzv1.ResourceAttributes{
					Group:    certmanager.GroupName,
					Resource: "signers",
					Name:     name,
					Verb:     "approve",
					Version:  "*",
				},
			},
		}, metav1.CreateOptions{})
		if err != nil {
			return false, err
		}

		if resp.Status.Allowed {
			return true, nil
		}
	}

	return false, nil
}

// isApprovalRequest will return true if the request is given a new approved or
// denied condition. This check is strictly concerned with these conditions
// being _added_. We do this to reduce the number of SAR calls made, since
// removal or changing of these conditions will be rejected elsewhere in the
// validation chain locally.
func isApprovalRequest(oldCR, newCR *cmapi.CertificateRequest) bool {
	oldCRApproving := util.GetCertificateRequestCondition(oldCR.Status.Conditions, cmapi.CertificateRequestConditionApproved)
	newCRApproving := util.GetCertificateRequestCondition(newCR.Status.Conditions, cmapi.CertificateRequestConditionApproved)

	if oldCRApproving == nil && newCRApproving != nil {
		return true
	}

	oldCRDenying := util.GetCertificateRequestCondition(oldCR.Status.Conditions, cmapi.CertificateRequestConditionDenied)
	newCRDenying := util.GetCertificateRequestCondition(newCR.Status.Conditions, cmapi.CertificateRequestConditionDenied)

	if oldCRDenying == nil && newCRDenying != nil {
		return true
	}

	return false
}
<|endoftext|>


<commit_before>// Package mux provides a simple Discord message route multiplexer that
// parses messages and then executes a matching registered handler, if found.
// mux can be used with both Disgord and the DiscordGo library.
package mux

import (
	"fmt"
	"log"
	"regexp"
	"strings"

	"github.com/bwmarrin/discordgo"
)

// Route holds information about a specific message route handler
type Route struct {
	Pattern     string      // match pattern that should trigger this route handler
	Description string      // short description of this route
	Help        string      // detailed help string for this route
	Run         HandlerFunc // route handler function to call
}

// Context holds a bit of extra data we pass along to route handlers
// This way processing some of this only needs to happen once.
type Context struct {
	Fields          []string
	Content         string
	IsDirected      bool
	IsPrivate       bool
	HasPrefix       bool
	HasMention      bool
	HasMentionFirst bool
}

// HandlerFunc is the function signature required for a message route handler.
type HandlerFunc func(*discordgo.Session, *discordgo.Message, *Context)

// Mux is the main struct for all mux methods.
type Mux struct {
	Routes  []*Route
	Default *Route
	Prefix  string
}

// New returns a new Discord message route mux
func New() *Mux {
	m := &Mux{}
	m.Prefix = "-dg "
	return m
}

// Route allows you to register a route
func (m *Mux) Route(pattern, desc string, cb HandlerFunc) (*Route, error) {

	r := Route{}
	r.Pattern = pattern
	r.Description = desc
	r.Run = cb
	m.Routes = append(m.Routes, &r)

	return &r, nil
}

// FuzzyMatch attepts to find the best route match for a givin message.
func (m *Mux) FuzzyMatch(msg string) (*Route, []string) {

	// Tokenize the msg string into a slice of words
	fields := strings.Fields(msg)

	// no point to continue if there's no fields
	if len(fields) == 0 {
		return nil, nil
	}

	// Search though the command list for a match
	var r *Route
	var rank int

	var fk int
	for fk, fv := range fields {

		for _, rv := range m.Routes {

			// If we find an exact match, return that immediately.
			if rv.Pattern == fv {
				return rv, fields[fk:]
			}

			// Some "Fuzzy" searching...
			if strings.HasPrefix(rv.Pattern, fv) {
				if len(fv) > rank {
					r = rv
					rank = len(fv)
				}
			}
		}
	}
	return r, fields[fk:]
}

// OnMessageCreate is a DiscordGo Event Handler function.  This must be
// registered using the DiscordGo.Session.AddHandler function.  This function
// will receive all Discord messages and parse them for matches to registered
// routes.
func (m *Mux) OnMessageCreate(ds *discordgo.Session, mc *discordgo.MessageCreate) {

	var err error

	// Ignore all messages created by the Bot account itself
	if mc.Author.ID == ds.State.User.ID {
		return
	}

	// Create Context struct that we can put various infos into
	ctx := &Context{
		Content: strings.TrimSpace(mc.Content),
	}

	// Fetch the channel for this Message
	var c *discordgo.Channel
	c, err = ds.State.Channel(mc.ChannelID)
	if err != nil {
		// Try fetching via REST API
		c, err = ds.Channel(mc.ChannelID)
		if err != nil {
			log.Printf("unable to fetch Channel for Message,", err)
		} else {
			// Attempt to add this channel into our State
			err = ds.State.ChannelAdd(c)
			if err != nil {
				log.Printf("error updating State with Channel,", err)
			}
			// Add Channel info into Context
			if c.Type == discordgo.ChannelTypeDM {
				ctx.IsPrivate = true
				ctx.IsDirected = true
			}
		}
	}

	// Detect @name or @nick mentions
	if !ctx.IsDirected {

		// Detect if Bot was @mentioned
		for _, v := range mc.Mentions {

			if v.ID == ds.State.User.ID {

				ctx.IsDirected, ctx.HasMention = true, true

				reg := regexp.MustCompile(fmt.Sprintf("<@!?(%s)>", ds.State.User.ID))

				// Was the @mention the first part of the string?
				if reg.FindStringIndex(ctx.Content)[0] == 0 {
					ctx.HasMentionFirst = true
				}

				// strip bot mention tags from content string
				ctx.Content = reg.ReplaceAllString(ctx.Content, "")

				break
			}
		}
	}

	// Detect prefix mention
	if !ctx.IsDirected && len(m.Prefix) > 0 {

		// TODO : Must be changed to support a per-guild user defined prefix
		if strings.HasPrefix(ctx.Content, m.Prefix) {
			ctx.IsDirected, ctx.HasPrefix, ctx.HasMentionFirst = true, true, true
			ctx.Content = strings.TrimPrefix(ctx.Content, m.Prefix)
		}
	}

	// For now, if we're not specifically mentioned we do nothing.
	// later I might add an option for global non-mentioned command words
	if !ctx.IsDirected {
		return
	}

	// Try to find the "best match" command out of the message.
	r, fl := m.FuzzyMatch(ctx.Content)
	if r != nil {
		ctx.Fields = fl
		r.Run(ds, mc.Message, ctx)
		return
	}

	// If no command match was found, call the default.
	// Ignore if only @mentioned in the middle of a message
	if m.Default != nil && (ctx.HasMentionFirst) {
		// TODO: This could use a ratelimit
		// or should the ratelimit be inside the cmd handler?..
		// In the case of "talking" to another bot, this can create an endless
		// loop.  Probably most common in private messages.
		m.Default.Run(ds, mc.Message, ctx)
	}

}
<commit_msg>Fix errors and properly assign Channel data to ctx<commit_after>// Package mux provides a simple Discord message route multiplexer that
// parses messages and then executes a matching registered handler, if found.
// mux can be used with both Disgord and the DiscordGo library.
package mux

import (
	"fmt"
	"log"
	"regexp"
	"strings"

	"github.com/bwmarrin/discordgo"
)

// Route holds information about a specific message route handler
type Route struct {
	Pattern     string      // match pattern that should trigger this route handler
	Description string      // short description of this route
	Help        string      // detailed help string for this route
	Run         HandlerFunc // route handler function to call
}

// Context holds a bit of extra data we pass along to route handlers
// This way processing some of this only needs to happen once.
type Context struct {
	Fields          []string
	Content         string
	IsDirected      bool
	IsPrivate       bool
	HasPrefix       bool
	HasMention      bool
	HasMentionFirst bool
}

// HandlerFunc is the function signature required for a message route handler.
type HandlerFunc func(*discordgo.Session, *discordgo.Message, *Context)

// Mux is the main struct for all mux methods.
type Mux struct {
	Routes  []*Route
	Default *Route
	Prefix  string
}

// New returns a new Discord message route mux
func New() *Mux {
	m := &Mux{}
	m.Prefix = "-dg "
	return m
}

// Route allows you to register a route
func (m *Mux) Route(pattern, desc string, cb HandlerFunc) (*Route, error) {

	r := Route{}
	r.Pattern = pattern
	r.Description = desc
	r.Run = cb
	m.Routes = append(m.Routes, &r)

	return &r, nil
}

// FuzzyMatch attepts to find the best route match for a givin message.
func (m *Mux) FuzzyMatch(msg string) (*Route, []string) {

	// Tokenize the msg string into a slice of words
	fields := strings.Fields(msg)

	// no point to continue if there's no fields
	if len(fields) == 0 {
		return nil, nil
	}

	// Search though the command list for a match
	var r *Route
	var rank int

	var fk int
	for fk, fv := range fields {

		for _, rv := range m.Routes {

			// If we find an exact match, return that immediately.
			if rv.Pattern == fv {
				return rv, fields[fk:]
			}

			// Some "Fuzzy" searching...
			if strings.HasPrefix(rv.Pattern, fv) {
				if len(fv) > rank {
					r = rv
					rank = len(fv)
				}
			}
		}
	}
	return r, fields[fk:]
}

// OnMessageCreate is a DiscordGo Event Handler function.  This must be
// registered using the DiscordGo.Session.AddHandler function.  This function
// will receive all Discord messages and parse them for matches to registered
// routes.
func (m *Mux) OnMessageCreate(ds *discordgo.Session, mc *discordgo.MessageCreate) {

	var err error

	// Ignore all messages created by the Bot account itself
	if mc.Author.ID == ds.State.User.ID {
		return
	}

	// Create Context struct that we can put various infos into
	ctx := &Context{
		Content: strings.TrimSpace(mc.Content),
	}

	// Fetch the channel for this Message
	var c *discordgo.Channel
	c, err = ds.State.Channel(mc.ChannelID)
	if err != nil {
		// Try fetching via REST API
		c, err = ds.Channel(mc.ChannelID)
		if err != nil {
			log.Printf("unable to fetch Channel for Message, %s", err)
		} else {
			// Attempt to add this channel into our State
			err = ds.State.ChannelAdd(c)
			if err != nil {
				log.Printf("error updating State with Channel, %s", err)
			}
		}
	}
	// Add Channel info into Context (if we successfully got the channel)
	if c != nil {
		if c.Type == discordgo.ChannelTypeDM {
			ctx.IsPrivate = true
			ctx.IsDirected = true
		}
	}

	// Detect @name or @nick mentions
	if !ctx.IsDirected {

		// Detect if Bot was @mentioned
		for _, v := range mc.Mentions {

			if v.ID == ds.State.User.ID {

				ctx.IsDirected, ctx.HasMention = true, true

				reg := regexp.MustCompile(fmt.Sprintf("<@!?(%s)>", ds.State.User.ID))

				// Was the @mention the first part of the string?
				if reg.FindStringIndex(ctx.Content)[0] == 0 {
					ctx.HasMentionFirst = true
				}

				// strip bot mention tags from content string
				ctx.Content = reg.ReplaceAllString(ctx.Content, "")

				break
			}
		}
	}

	// Detect prefix mention
	if !ctx.IsDirected && len(m.Prefix) > 0 {

		// TODO : Must be changed to support a per-guild user defined prefix
		if strings.HasPrefix(ctx.Content, m.Prefix) {
			ctx.IsDirected, ctx.HasPrefix, ctx.HasMentionFirst = true, true, true
			ctx.Content = strings.TrimPrefix(ctx.Content, m.Prefix)
		}
	}

	// For now, if we're not specifically mentioned we do nothing.
	// later I might add an option for global non-mentioned command words
	if !ctx.IsDirected {
		return
	}

	// Try to find the "best match" command out of the message.
	r, fl := m.FuzzyMatch(ctx.Content)
	if r != nil {
		ctx.Fields = fl
		r.Run(ds, mc.Message, ctx)
		return
	}

	// If no command match was found, call the default.
	// Ignore if only @mentioned in the middle of a message
	if m.Default != nil && (ctx.HasMentionFirst) {
		// TODO: This could use a ratelimit
		// or should the ratelimit be inside the cmd handler?..
		// In the case of "talking" to another bot, this can create an endless
		// loop.  Probably most common in private messages.
		m.Default.Run(ds, mc.Message, ctx)
	}

}
<|endoftext|>


<commit_before>package arn

import (
	"github.com/aerogo/aerospike"
	"github.com/aerogo/api"
)

// DB is the main database client.
var DB = aerospike.NewDatabase(
	"arn-db",
	3000,
	"arn",
	[]interface{}{
		(*Analytics)(nil),
		(*Anime)(nil),
		(*AnimeCharacters)(nil),
		(*AnimeEpisodes)(nil),
		(*AnimeRelations)(nil),
		(*AnimeList)(nil),
		(*AniListToAnime)(nil),
		(*Character)(nil),
		(*DraftIndex)(nil),
		(*MyAnimeListToAnime)(nil),
		(*EmailToUser)(nil),
		(*FacebookToUser)(nil),
		(*GoogleToUser)(nil),
		(*Group)(nil),
		(*Item)(nil),
		(*Inventory)(nil),
		(*NickToUser)(nil),
		(*PayPalPayment)(nil),
		(*Post)(nil),
		(*Purchase)(nil),
		(*PushSubscriptions)(nil),
		(*SearchIndex)(nil),
		(*Settings)(nil),
		(*SoundTrack)(nil),
		(*Thread)(nil),
		(*TwitterToUser)(nil),
		(*User)(nil),
		(*UserFollows)(nil),
	},
)

// API ...
var API = api.New("/api/", DB)
<commit_msg>Added DBTypes<commit_after>package arn

import (
	"github.com/aerogo/aerospike"
	"github.com/aerogo/api"
)

// DB is the main database client.
var DB = aerospike.NewDatabase(
	"arn-db",
	3000,
	"arn",
	DBTypes,
)

// DBTypes ...
var DBTypes = []interface{}{
	(*Analytics)(nil),
	(*Anime)(nil),
	(*AnimeCharacters)(nil),
	(*AnimeEpisodes)(nil),
	(*AnimeRelations)(nil),
	(*AnimeList)(nil),
	(*AniListToAnime)(nil),
	(*Character)(nil),
	(*DraftIndex)(nil),
	(*MyAnimeListToAnime)(nil),
	(*EmailToUser)(nil),
	(*FacebookToUser)(nil),
	(*GoogleToUser)(nil),
	(*Group)(nil),
	(*Item)(nil),
	(*Inventory)(nil),
	(*NickToUser)(nil),
	(*PayPalPayment)(nil),
	(*Post)(nil),
	(*Purchase)(nil),
	(*PushSubscriptions)(nil),
	(*SearchIndex)(nil),
	(*Settings)(nil),
	(*SoundTrack)(nil),
	(*Thread)(nil),
	(*TwitterToUser)(nil),
	(*User)(nil),
	(*UserFollows)(nil),
}

// API ...
var API = api.New("/api/", DB)
<|endoftext|>


<commit_before>// Copyright (c) Ilia Kravets, 2015. All rights reserved. PROVIDED "AS IS"
// WITHOUT ANY WARRANTY, EXPRESS OR IMPLIED. See LICENSE file for details.

package rec

import (
	"log"
	"time"

	"code.google.com/p/gopacket"

	"my/itto/verify/packet"
	"my/itto/verify/sim"
)

type Stream struct {
	message *sim.SimMessage
	seconds map[gopacket.Flow]int
	seqNum  map[gopacket.Flow]uint64
}

func NewStream() *Stream {
	l := &Stream{
		seconds: make(map[gopacket.Flow]int),
		seqNum:  make(map[gopacket.Flow]uint64),
	}
	return l
}
func (l *Stream) MessageArrived(idm *sim.SimMessage) {
	l.message = idm

	flow := l.message.Pam.Flow()
	seq := l.message.Pam.SequenceNumber()
	if prevSeq, ok := l.seqNum[flow]; ok && prevSeq+1 != seq {
		log.Printf("seqNum gap; expected %d actual %d\n", prevSeq+1, seq)
	}
	l.seqNum[flow] = seq

	if m, ok := l.message.Pam.Layer().(packet.SecondsMessage); ok {
		l.seconds[flow] = m.Seconds()
	}
}
func (l *Stream) getSeqNum() uint64 {
	flow := l.message.Pam.Flow()
	return l.seqNum[flow]
}
func (l *Stream) getTimestamp() uint64 {
	flow := l.message.Pam.Flow()
	return uint64(l.seconds[flow])*1e9 + uint64(l.message.Pam.Layer().(packet.ExchangeMessage).Nanoseconds())
}
func (l *Stream) getExchangeMessage() packet.ExchangeMessage {
	return l.message.Pam.Layer().(packet.ExchangeMessage)
}
func (l *Stream) getPacketTimestamp() time.Time {
	return l.message.Pam.Timestamp()
}
<commit_msg>rec:Stream: ignore zero sequence number<commit_after>// Copyright (c) Ilia Kravets, 2015. All rights reserved. PROVIDED "AS IS"
// WITHOUT ANY WARRANTY, EXPRESS OR IMPLIED. See LICENSE file for details.

package rec

import (
	"log"
	"time"

	"code.google.com/p/gopacket"

	"my/itto/verify/packet"
	"my/itto/verify/sim"
)

type Stream struct {
	message *sim.SimMessage
	seconds map[gopacket.Flow]int
	seqNum  map[gopacket.Flow]uint64
}

func NewStream() *Stream {
	l := &Stream{
		seconds: make(map[gopacket.Flow]int),
		seqNum:  make(map[gopacket.Flow]uint64),
	}
	return l
}
func (l *Stream) MessageArrived(idm *sim.SimMessage) {
	l.message = idm

	flow := l.message.Pam.Flow()
	seq := l.message.Pam.SequenceNumber()
	if seq != 0 {
		if prevSeq, ok := l.seqNum[flow]; ok && prevSeq+1 != seq {
			log.Printf("seqNum gap; expected %d actual %d\n", prevSeq+1, seq)
		}
		l.seqNum[flow] = seq
	}

	if m, ok := l.message.Pam.Layer().(packet.SecondsMessage); ok {
		l.seconds[flow] = m.Seconds()
	}
}
func (l *Stream) getSeqNum() uint64 {
	flow := l.message.Pam.Flow()
	return l.seqNum[flow]
}
func (l *Stream) getTimestamp() uint64 {
	flow := l.message.Pam.Flow()
	return uint64(l.seconds[flow])*1e9 + uint64(l.message.Pam.Layer().(packet.ExchangeMessage).Nanoseconds())
}
func (l *Stream) getExchangeMessage() packet.ExchangeMessage {
	return l.message.Pam.Layer().(packet.ExchangeMessage)
}
func (l *Stream) getPacketTimestamp() time.Time {
	return l.message.Pam.Timestamp()
}
<|endoftext|>


<commit_before>package fat

import (
	"bytes"
	"encoding/binary"
	"errors"
	"fmt"
	"github.com/mitchellh/go-fs"
	"math"
	"time"
	"unicode/utf16"
)

type DirectoryAttr uint8

const (
	AttrReadOnly  DirectoryAttr = 0x01
	AttrHidden                  = 0x02
	AttrSystem                  = 0x04
	AttrVolumeId                = 0x08
	AttrDirectory               = 0x10
	AttrArchive                 = 0x20
	AttrLongName                = AttrReadOnly | AttrHidden | AttrSystem | AttrVolumeId
)

// The size in bytes of a single directory entry.
const DirectoryEntrySize = 32

// Mask applied to the ord of the last long entry.
const LastLongEntryMask = 0x40

// DirectoryCluster represents a cluster on the disk that contains
// entries/contents.
type DirectoryCluster struct {
	entries      []*DirectoryClusterEntry
	fat16Root    bool
	startCluster uint32
}

// DirectoryClusterEntry is a single 32-byte entry that is part of the
// chain of entries in a directory cluster.
type DirectoryClusterEntry struct {
	name       string
	ext        string
	attr       DirectoryAttr
	createTime time.Time
	accessTime time.Time
	writeTime  time.Time
	cluster    uint32
	fileSize   uint32
	deleted    bool

	longOrd      uint8
	longName     string
	longChecksum uint8
}

func DecodeDirectoryCluster(startCluster uint32, device fs.BlockDevice, fat *FAT) (*DirectoryCluster, error) {
	bs := fat.bs
	chain := fat.Chain(startCluster)
	data := make([]byte, uint32(len(chain))*bs.BytesPerCluster())
	for i, clusterNumber := range chain {
		dataOffset := uint32(i) * bs.BytesPerCluster()
		devOffset := int64(bs.ClusterOffset(int(clusterNumber)))
		chainData := data[dataOffset : dataOffset+bs.BytesPerCluster()]

		if _, err := device.ReadAt(chainData, devOffset); err != nil {
			return nil, err
		}
	}

	result, err := decodeDirectoryCluster(data, bs)
	if err != nil {
		return nil, err
	}

	result.startCluster = startCluster
	return result, nil
}

// DecodeFAT16RootDirectory decodes the FAT16 root directory structure
// from the device.
func DecodeFAT16RootDirectoryCluster(device fs.BlockDevice, bs *BootSectorCommon) (*DirectoryCluster, error) {
	data := make([]byte, DirectoryEntrySize*bs.RootEntryCount)
	if _, err := device.ReadAt(data, int64(bs.RootDirOffset())); err != nil {
		return nil, err
	}

	result, err := decodeDirectoryCluster(data, bs)
	if err != nil {
		return nil, err
	}

	result.fat16Root = true
	return result, nil
}

func decodeDirectoryCluster(data []byte, bs *BootSectorCommon) (*DirectoryCluster, error) {
	entries := make([]*DirectoryClusterEntry, 0, bs.RootEntryCount)
	for i := uint16(0); i < uint16(len(data)/DirectoryEntrySize); i++ {
		offset := i * DirectoryEntrySize
		entryData := data[offset : offset+DirectoryEntrySize]
		if entryData[0] == 0 {
			break
		}

		entry, err := DecodeDirectoryClusterEntry(entryData)
		if err != nil {
			return nil, err
		}

		entries = append(entries, entry)
	}

	result := &DirectoryCluster{
		entries: entries,
	}

	return result, nil
}

func NewDirectoryCluster(start uint32, parent uint32, t time.Time) *DirectoryCluster {
	cluster := new(DirectoryCluster)
	cluster.startCluster = start

	// Create the "." and ".." entries
	cluster.entries = []*DirectoryClusterEntry{
		&DirectoryClusterEntry{
			accessTime: t,
			attr:       AttrDirectory,
			cluster:    start,
			createTime: t,
			name:       ".",
			writeTime:  t,
		},
		&DirectoryClusterEntry{
			accessTime: t,
			attr:       AttrDirectory,
			cluster:    parent,
			createTime: t,
			name:       "..",
			writeTime:  t,
		},
	}

	return cluster
}

// NewFat16RootDirectory creates a new DirectoryCluster that is meant only
// to be the root directory of a FAT12/FAT16 filesystem.
func NewFat16RootDirectoryCluster(bs *BootSectorCommon, label string) (*DirectoryCluster, error) {
	if bs.RootEntryCount == 0 {
		return nil, errors.New("root entry count is 0 in boot sector")
	}

	result := &DirectoryCluster{
		entries: make([]*DirectoryClusterEntry, 1, bs.RootEntryCount),
	}

	// Create the volume ID entry
	result.entries[0] = &DirectoryClusterEntry{
		attr: AttrVolumeId,
		name: label,
		cluster: 0,
	}

	return result, nil
}

// Bytes returns the on-disk byte data for this directory structure.
func (d *DirectoryCluster) Bytes() []byte {
	result := make([]byte, cap(d.entries)*DirectoryEntrySize)

	for i, entry := range d.entries {
		offset := i * DirectoryEntrySize
		entryBytes := entry.Bytes()
		copy(result[offset:offset+DirectoryEntrySize], entryBytes)
	}

	return result
}

// WriteToDevice writes the cluster to the device.
func (d *DirectoryCluster) WriteToDevice(device fs.BlockDevice, fat *FAT) error {
	if d.fat16Root {
		// Write the cluster to the FAT16 root directory location
		offset := int64(fat.bs.RootDirOffset())
		if _, err := device.WriteAt(d.Bytes(), offset); err != nil {
			return err
		}
	} else {
		chain := &ClusterChain{
			device:       device,
			fat:          fat,
			startCluster: d.startCluster,
		}

		if _, err := chain.Write(d.Bytes()); err != nil {
			return err
		}
	}

	return nil
}

// Bytes returns the on-disk byte data for this directory entry.
func (d *DirectoryClusterEntry) Bytes() []byte {
	var result [DirectoryEntrySize]byte

	if d.longName != "" {
		runes := bytes.Runes([]byte(d.longName))

		// LDIR_Ord
		result[0] = d.longOrd

		// LDIR_Name1
		for i := 0; i < int(math.Min(float64(len(runes)), 5)); i++ {
			offset := 1 + (i * 2)
			data := result[offset : offset+2]
			binary.LittleEndian.PutUint16(data, uint16(runes[i]))
		}

		// LDIR_Attr
		result[11] = byte(AttrLongName)

		// LDIR_Type
		result[12] = 0

		// LDIR_Chksum
		result[13] = d.longChecksum

		// LDIR_Name2
		if len(runes) > 5 {
			limit := int(math.Min(float64(len(runes)), 11)) - 5
			for i := 0; i < limit; i++ {
				offset := 14 + (i * 2)
				data := result[offset : offset+2]
				binary.LittleEndian.PutUint16(data, uint16(runes[i+5]))
			}
		}

		// LDIR_FstClusLO
		result[26] = 0
		result[27] = 0

		// LDIR_Name3
		if len(runes) > 11 {
			limit := int(math.Min(float64(len(runes)), 13)) - 11
			for i := 0; i < limit; i++ {
				offset := 28 + (i * 2)
				data := result[offset : offset+2]
				binary.LittleEndian.PutUint16(data, uint16(runes[i+11]))
			}
		}
	} else {
		// DIR_Name
		var simpleName string
		if d.name == "." || d.name == ".." {
			simpleName = d.name
		} else {
			simpleName = fmt.Sprintf("%s.%s", d.name, d.ext)
		}
		copy(result[0:11], shortNameEntryValue(simpleName))

		// DIR_Attr
		result[11] = byte(d.attr)

		// DIR_CrtTime
		crtDate, crtTime, crtTenths := encodeDOSTime(d.createTime)
		result[13] = crtTenths
		binary.LittleEndian.PutUint16(result[14:16], crtTime)
		binary.LittleEndian.PutUint16(result[16:18], crtDate)

		// DIR_LstAccDate
		accDate, _, _ := encodeDOSTime(d.accessTime)
		binary.LittleEndian.PutUint16(result[18:20], accDate)

		// DIR_FstClusHI
		binary.LittleEndian.PutUint16(result[20:22], uint16(d.cluster>>16))

		// DIR_WrtTime and DIR_WrtDate
		wrtDate, wrtTime, _ := encodeDOSTime(d.writeTime)
		binary.LittleEndian.PutUint16(result[22:24], wrtTime)
		binary.LittleEndian.PutUint16(result[24:26], wrtDate)

		// DIR_FstClusLO
		binary.LittleEndian.PutUint16(result[26:28], uint16(d.cluster&0xFFFF))

		// DIR_FileSize
		binary.LittleEndian.PutUint32(result[28:32], d.fileSize)
	}

	return result[:]
}

// IsLong returns true if this is a long entry.
func (d *DirectoryClusterEntry) IsLong() bool {
	return (d.attr & AttrLongName) == AttrLongName
}

func (d *DirectoryClusterEntry) IsVolumeId() bool {
	return (d.attr & AttrVolumeId) == AttrVolumeId
}

// DecodeDirectoryClusterEntry decodes a single directory entry in the
// Directory structure.
func DecodeDirectoryClusterEntry(data []byte) (*DirectoryClusterEntry, error) {
	var result DirectoryClusterEntry

	// Do the attributes so we can determine if we're dealing with long names
	result.attr = DirectoryAttr(data[11])
	if (result.attr & AttrLongName) == AttrLongName {
		result.longOrd = data[0]

		chars := make([]uint16, 13)
		for i := 0; i < 5; i++ {
			offset := 1 + (i * 2)
			chars[i] = binary.LittleEndian.Uint16(data[offset : offset+2])
		}

		for i := 0; i < 6; i++ {
			offset := 14 + (i * 2)
			chars[i+5] = binary.LittleEndian.Uint16(data[offset : offset+2])
		}

		for i := 0; i < 2; i++ {
			offset := 28 + (i * 2)
			chars[i+11] = binary.LittleEndian.Uint16(data[offset : offset+2])
		}

		result.longName = string(utf16.Decode(chars))
		result.longChecksum = data[13]
	} else {
		result.deleted = data[0] == 0xE5

		// Basic attributes
		if data[0] == 0x05 {
			data[0] = 0xE5
		}

		result.name = string(data[0:8])
		result.ext = string(data[8:11])

		// Creation time
		createTimeTenths := data[13]
		createTimeWord := binary.LittleEndian.Uint16(data[14:16])
		createDateWord := binary.LittleEndian.Uint16(data[16:18])
		result.createTime = decodeDOSTime(createDateWord, createTimeWord, createTimeTenths)

		// Access time
		accessDateWord := binary.LittleEndian.Uint16(data[18:20])
		result.accessTime = decodeDOSTime(accessDateWord, 0, 0)

		// Write time
		writeTimeWord := binary.LittleEndian.Uint16(data[22:24])
		writeDateWord := binary.LittleEndian.Uint16(data[24:26])
		result.writeTime = decodeDOSTime(writeDateWord, writeTimeWord, 0)

		// Cluster
		result.cluster = uint32(binary.LittleEndian.Uint16(data[20:22]))
		result.cluster <<= 4
		result.cluster |= uint32(binary.LittleEndian.Uint16(data[26:28]))

		// File size
		result.fileSize = binary.LittleEndian.Uint32(data[28:32])
	}

	return &result, nil
}

// NewLongDirectoryClusterEntry returns the series of directory cluster
// entries that need to be written for a long directory entry. This list
// of entries does NOT contain the short name entry.
func NewLongDirectoryClusterEntry(name string, shortName string) ([]*DirectoryClusterEntry, error) {
	// Split up the shortName properly
	checksum := checksumShortName(shortNameEntryValue(shortName))

	// Calcualte the number of entries we'll actually need to store
	// the long name.
	numLongEntries := len(name) / 13
	if len(name)%13 != 0 {
		numLongEntries++
	}

	entries := make([]*DirectoryClusterEntry, numLongEntries)
	for i := 0; i < numLongEntries; i++ {
		entries[i] = new(DirectoryClusterEntry)
		entry := entries[i]
		entry.attr = AttrLongName
		entry.longOrd = uint8(numLongEntries - i)

		if i == 0 {
			entry.longOrd |= LastLongEntryMask
		}

		// Calculate the offsets of the string for this entry
		i := (len(name) % 13) + (i * 13)
		j := len(name) - i
		k := int(math.Min(float64(j+13), float64(len(name))))

		entry.longChecksum = checksum
		entry.longName = name[j:k]
	}

	return entries, nil
}

func decodeDOSTime(date, dosTime uint16, tenths uint8) time.Time {
	return time.Date(
		1980+int(date>>9),
		time.Month((date>>5)&0x0F),
		int(date&0x1F),
		int(dosTime>>11),
		int((dosTime>>5)&0x3F),
		int((dosTime&0x1F)*2),
		int(tenths)*10*int(time.Millisecond),
		time.Local)
}

func encodeDOSTime(t time.Time) (uint16, uint16, uint8) {
	var date uint16 = uint16((t.Year() - 1980) << 9)
	date |= uint16(t.Month()) << 5
	date += uint16(t.Day() & 0xFF)

	var time uint16 = uint16(t.Hour() << 11)
	time |= uint16(t.Minute() << 5)
	time += uint16(t.Second() / 2)

	var tenths uint8
	// TODO(mitchellh): Do tenths

	return date, time, tenths
}
<commit_msg>fat: properly pad long names with 0xFF<commit_after>package fat

import (
	"bytes"
	"encoding/binary"
	"errors"
	"fmt"
	"github.com/mitchellh/go-fs"
	"math"
	"time"
	"unicode/utf16"
)

type DirectoryAttr uint8

const (
	AttrReadOnly  DirectoryAttr = 0x01
	AttrHidden                  = 0x02
	AttrSystem                  = 0x04
	AttrVolumeId                = 0x08
	AttrDirectory               = 0x10
	AttrArchive                 = 0x20
	AttrLongName                = AttrReadOnly | AttrHidden | AttrSystem | AttrVolumeId
)

// The size in bytes of a single directory entry.
const DirectoryEntrySize = 32

// Mask applied to the ord of the last long entry.
const LastLongEntryMask = 0x40

// DirectoryCluster represents a cluster on the disk that contains
// entries/contents.
type DirectoryCluster struct {
	entries      []*DirectoryClusterEntry
	fat16Root    bool
	startCluster uint32
}

// DirectoryClusterEntry is a single 32-byte entry that is part of the
// chain of entries in a directory cluster.
type DirectoryClusterEntry struct {
	name       string
	ext        string
	attr       DirectoryAttr
	createTime time.Time
	accessTime time.Time
	writeTime  time.Time
	cluster    uint32
	fileSize   uint32
	deleted    bool

	longOrd      uint8
	longName     string
	longChecksum uint8
}

func DecodeDirectoryCluster(startCluster uint32, device fs.BlockDevice, fat *FAT) (*DirectoryCluster, error) {
	bs := fat.bs
	chain := fat.Chain(startCluster)
	data := make([]byte, uint32(len(chain))*bs.BytesPerCluster())
	for i, clusterNumber := range chain {
		dataOffset := uint32(i) * bs.BytesPerCluster()
		devOffset := int64(bs.ClusterOffset(int(clusterNumber)))
		chainData := data[dataOffset : dataOffset+bs.BytesPerCluster()]

		if _, err := device.ReadAt(chainData, devOffset); err != nil {
			return nil, err
		}
	}

	result, err := decodeDirectoryCluster(data, bs)
	if err != nil {
		return nil, err
	}

	result.startCluster = startCluster
	return result, nil
}

// DecodeFAT16RootDirectory decodes the FAT16 root directory structure
// from the device.
func DecodeFAT16RootDirectoryCluster(device fs.BlockDevice, bs *BootSectorCommon) (*DirectoryCluster, error) {
	data := make([]byte, DirectoryEntrySize*bs.RootEntryCount)
	if _, err := device.ReadAt(data, int64(bs.RootDirOffset())); err != nil {
		return nil, err
	}

	result, err := decodeDirectoryCluster(data, bs)
	if err != nil {
		return nil, err
	}

	result.fat16Root = true
	return result, nil
}

func decodeDirectoryCluster(data []byte, bs *BootSectorCommon) (*DirectoryCluster, error) {
	entries := make([]*DirectoryClusterEntry, 0, bs.RootEntryCount)
	for i := uint16(0); i < uint16(len(data)/DirectoryEntrySize); i++ {
		offset := i * DirectoryEntrySize
		entryData := data[offset : offset+DirectoryEntrySize]
		if entryData[0] == 0 {
			break
		}

		entry, err := DecodeDirectoryClusterEntry(entryData)
		if err != nil {
			return nil, err
		}

		entries = append(entries, entry)
	}

	result := &DirectoryCluster{
		entries: entries,
	}

	return result, nil
}

func NewDirectoryCluster(start uint32, parent uint32, t time.Time) *DirectoryCluster {
	cluster := new(DirectoryCluster)
	cluster.startCluster = start

	// Create the "." and ".." entries
	cluster.entries = []*DirectoryClusterEntry{
		&DirectoryClusterEntry{
			accessTime: t,
			attr:       AttrDirectory,
			cluster:    start,
			createTime: t,
			name:       ".",
			writeTime:  t,
		},
		&DirectoryClusterEntry{
			accessTime: t,
			attr:       AttrDirectory,
			cluster:    parent,
			createTime: t,
			name:       "..",
			writeTime:  t,
		},
	}

	return cluster
}

// NewFat16RootDirectory creates a new DirectoryCluster that is meant only
// to be the root directory of a FAT12/FAT16 filesystem.
func NewFat16RootDirectoryCluster(bs *BootSectorCommon, label string) (*DirectoryCluster, error) {
	if bs.RootEntryCount == 0 {
		return nil, errors.New("root entry count is 0 in boot sector")
	}

	result := &DirectoryCluster{
		entries: make([]*DirectoryClusterEntry, 1, bs.RootEntryCount),
	}

	// Create the volume ID entry
	result.entries[0] = &DirectoryClusterEntry{
		attr: AttrVolumeId,
		name: label,
		cluster: 0,
	}

	return result, nil
}

// Bytes returns the on-disk byte data for this directory structure.
func (d *DirectoryCluster) Bytes() []byte {
	result := make([]byte, cap(d.entries)*DirectoryEntrySize)

	for i, entry := range d.entries {
		offset := i * DirectoryEntrySize
		entryBytes := entry.Bytes()
		copy(result[offset:offset+DirectoryEntrySize], entryBytes)
	}

	return result
}

// WriteToDevice writes the cluster to the device.
func (d *DirectoryCluster) WriteToDevice(device fs.BlockDevice, fat *FAT) error {
	if d.fat16Root {
		// Write the cluster to the FAT16 root directory location
		offset := int64(fat.bs.RootDirOffset())
		if _, err := device.WriteAt(d.Bytes(), offset); err != nil {
			return err
		}
	} else {
		chain := &ClusterChain{
			device:       device,
			fat:          fat,
			startCluster: d.startCluster,
		}

		if _, err := chain.Write(d.Bytes()); err != nil {
			return err
		}
	}

	return nil
}

// Bytes returns the on-disk byte data for this directory entry.
func (d *DirectoryClusterEntry) Bytes() []byte {
	var result [DirectoryEntrySize]byte

	if d.longName != "" {
		runes := bytes.Runes([]byte(d.longName))

		// The name must be zero-terminated then padded with 0xFF
		// up to 13 characters
		if len(runes) < 13 {
			runes = append(runes, 0)
			for len(runes) < 13 {
				runes = append(runes, 0xFFFF)
			}
		}

		// LDIR_Ord
		result[0] = d.longOrd

		// LDIR_Name1
		for i := 0; i < int(math.Min(float64(len(runes)), 5)); i++ {
			offset := 1 + (i * 2)
			data := result[offset : offset+2]
			binary.LittleEndian.PutUint16(data, uint16(runes[i]))
		}

		// LDIR_Attr
		result[11] = byte(AttrLongName)

		// LDIR_Type
		result[12] = 0

		// LDIR_Chksum
		result[13] = d.longChecksum

		// LDIR_Name2
		for i := 0; i < 6; i++ {
			offset := 14 + (i * 2)
			data := result[offset : offset+2]
			binary.LittleEndian.PutUint16(data, uint16(runes[i+5]))
		}

		// LDIR_FstClusLO
		result[26] = 0
		result[27] = 0

		// LDIR_Name3
		for i := 0; i < 2; i++ {
			offset := 28 + (i * 2)
			data := result[offset : offset+2]
			binary.LittleEndian.PutUint16(data, uint16(runes[i+11]))
		}
	} else {
		// DIR_Name
		var simpleName string
		if d.name == "." || d.name == ".." {
			simpleName = d.name
		} else {
			simpleName = fmt.Sprintf("%s.%s", d.name, d.ext)
		}
		copy(result[0:11], shortNameEntryValue(simpleName))

		// DIR_Attr
		result[11] = byte(d.attr)

		// DIR_CrtTime
		crtDate, crtTime, crtTenths := encodeDOSTime(d.createTime)
		result[13] = crtTenths
		binary.LittleEndian.PutUint16(result[14:16], crtTime)
		binary.LittleEndian.PutUint16(result[16:18], crtDate)

		// DIR_LstAccDate
		accDate, _, _ := encodeDOSTime(d.accessTime)
		binary.LittleEndian.PutUint16(result[18:20], accDate)

		// DIR_FstClusHI
		binary.LittleEndian.PutUint16(result[20:22], uint16(d.cluster>>16))

		// DIR_WrtTime and DIR_WrtDate
		wrtDate, wrtTime, _ := encodeDOSTime(d.writeTime)
		binary.LittleEndian.PutUint16(result[22:24], wrtTime)
		binary.LittleEndian.PutUint16(result[24:26], wrtDate)

		// DIR_FstClusLO
		binary.LittleEndian.PutUint16(result[26:28], uint16(d.cluster&0xFFFF))

		// DIR_FileSize
		binary.LittleEndian.PutUint32(result[28:32], d.fileSize)
	}

	return result[:]
}

// IsLong returns true if this is a long entry.
func (d *DirectoryClusterEntry) IsLong() bool {
	return (d.attr & AttrLongName) == AttrLongName
}

func (d *DirectoryClusterEntry) IsVolumeId() bool {
	return (d.attr & AttrVolumeId) == AttrVolumeId
}

// DecodeDirectoryClusterEntry decodes a single directory entry in the
// Directory structure.
func DecodeDirectoryClusterEntry(data []byte) (*DirectoryClusterEntry, error) {
	var result DirectoryClusterEntry

	// Do the attributes so we can determine if we're dealing with long names
	result.attr = DirectoryAttr(data[11])
	if (result.attr & AttrLongName) == AttrLongName {
		result.longOrd = data[0]

		chars := make([]uint16, 13)
		for i := 0; i < 5; i++ {
			offset := 1 + (i * 2)
			chars[i] = binary.LittleEndian.Uint16(data[offset : offset+2])
		}

		for i := 0; i < 6; i++ {
			offset := 14 + (i * 2)
			chars[i+5] = binary.LittleEndian.Uint16(data[offset : offset+2])
		}

		for i := 0; i < 2; i++ {
			offset := 28 + (i * 2)
			chars[i+11] = binary.LittleEndian.Uint16(data[offset : offset+2])
		}

		result.longName = string(utf16.Decode(chars))
		result.longChecksum = data[13]
	} else {
		result.deleted = data[0] == 0xE5

		// Basic attributes
		if data[0] == 0x05 {
			data[0] = 0xE5
		}

		result.name = string(data[0:8])
		result.ext = string(data[8:11])

		// Creation time
		createTimeTenths := data[13]
		createTimeWord := binary.LittleEndian.Uint16(data[14:16])
		createDateWord := binary.LittleEndian.Uint16(data[16:18])
		result.createTime = decodeDOSTime(createDateWord, createTimeWord, createTimeTenths)

		// Access time
		accessDateWord := binary.LittleEndian.Uint16(data[18:20])
		result.accessTime = decodeDOSTime(accessDateWord, 0, 0)

		// Write time
		writeTimeWord := binary.LittleEndian.Uint16(data[22:24])
		writeDateWord := binary.LittleEndian.Uint16(data[24:26])
		result.writeTime = decodeDOSTime(writeDateWord, writeTimeWord, 0)

		// Cluster
		result.cluster = uint32(binary.LittleEndian.Uint16(data[20:22]))
		result.cluster <<= 4
		result.cluster |= uint32(binary.LittleEndian.Uint16(data[26:28]))

		// File size
		result.fileSize = binary.LittleEndian.Uint32(data[28:32])
	}

	return &result, nil
}

// NewLongDirectoryClusterEntry returns the series of directory cluster
// entries that need to be written for a long directory entry. This list
// of entries does NOT contain the short name entry.
func NewLongDirectoryClusterEntry(name string, shortName string) ([]*DirectoryClusterEntry, error) {
	// Split up the shortName properly
	checksum := checksumShortName(shortNameEntryValue(shortName))

	// Calcualte the number of entries we'll actually need to store
	// the long name.
	numLongEntries := len(name) / 13
	if len(name)%13 != 0 {
		numLongEntries++
	}

	entries := make([]*DirectoryClusterEntry, numLongEntries)
	for i := 0; i < numLongEntries; i++ {
		entries[i] = new(DirectoryClusterEntry)
		entry := entries[i]
		entry.attr = AttrLongName
		entry.longOrd = uint8(numLongEntries - i)

		if i == 0 {
			entry.longOrd |= LastLongEntryMask
		}

		// Calculate the offsets of the string for this entry
		i := (len(name) % 13) + (i * 13)
		j := len(name) - i
		k := int(math.Min(float64(j+13), float64(len(name))))

		entry.longChecksum = checksum
		entry.longName = name[j:k]
	}

	return entries, nil
}

func decodeDOSTime(date, dosTime uint16, tenths uint8) time.Time {
	return time.Date(
		1980+int(date>>9),
		time.Month((date>>5)&0x0F),
		int(date&0x1F),
		int(dosTime>>11),
		int((dosTime>>5)&0x3F),
		int((dosTime&0x1F)*2),
		int(tenths)*10*int(time.Millisecond),
		time.Local)
}

func encodeDOSTime(t time.Time) (uint16, uint16, uint8) {
	var date uint16 = uint16((t.Year() - 1980) << 9)
	date |= uint16(t.Month()) << 5
	date += uint16(t.Day() & 0xFF)

	var time uint16 = uint16(t.Hour() << 11)
	time |= uint16(t.Minute() << 5)
	time += uint16(t.Second() / 2)

	var tenths uint8
	// TODO(mitchellh): Do tenths

	return date, time, tenths
}
<|endoftext|>