<commit_before>package start

import (
	"fmt"
	"os"
	"testing"

	flag "github.com/ogier/pflag"
	. "github.com/smartystreets/goconvey/convey"
)

func TestAdd(t *testing.T) {
	Convey("When adding a command to the command map, then...", t, func() {
		Commands.Add(&Command{
			Name: "cmd1",
		})

		Convey("the command map should contain it", func() {
			So(Commands["cmd1"], ShouldNotBeNil)
			So(Commands["cmd1"].Name, ShouldEqual, "cmd1")

		})

		Convey("adding a command with the same name should fail", func() {
			err := Commands.Add(&Command{
				Name: "cmd1",
			})
			So(err, ShouldNotBeNil)
		})

		Convey("adding a subcommand to the command should work", func() {
			err := Commands.Add(&Command{
				Parent: "cmd1",
				Name:   "subcmd1",
			})
			So(err, ShouldBeNil)
			So(len(Commands["cmd1"].children), ShouldEqual, 1)
			So(Commands["cmd1"].children["subcmd1"], ShouldNotBeNil)
			So(Commands["cmd1"].children["subcmd1"].Name, ShouldEqual, "subcmd1")

			Convey("adding the same subcommand twice should fail", func() {
				err := Commands.Add(&Command{
					Parent: "cmd1",
					Name:   "subcmd1",
				})
				So(err, ShouldNotBeNil)
			})

		})

		Reset(func() {
			Commands = make(CommandMap)
		})
	})

}

func TestCommands(t *testing.T) {

	var yes bool
	var size int

	// ContinueOnError is required when running goconvey as server; otherwise, unrecognized
	// flags that are passed to the test executable will cause an error:
	// "unknown shorthand flag: 't' in -test.v=true"
	flag.CommandLine = flag.NewFlagSet(os.Args[0], flag.ContinueOnError)

	flag.BoolVarP(&yes, "yes", "y", false, "A boolean flag")
	flag.IntVarP(&size, "size", "s", 23, "An int flag")

	Commands = make(CommandMap)

	Convey("The flags should exist", t, func() {
		Parse()
		So(flag.Lookup("yes").Name, ShouldEqual, "yes")
		So(flag.Lookup("size").Name, ShouldEqual, "size")
	})

	Convey("When setting up some commands, then...", t, func() {
		Description = "This is the test application for the start package."

		Add(&Command{
			Name:  "test",
			Short: "A test command",
			Long:  "Command test helps testing the start package. It accepts all flags.",
			Cmd: func(cmd *Command) error {
				fmt.Println("This is the test command.")
				return nil
			},
		})

		Add(&Command{
			Name:  "flags",
			Flags: []string{"yes", "size"},
			Short: "A test command with flags",
			Long: `Command flags helps testing flags.
It accepts the flags --yes and --size.`,
			Cmd: func(cmd *Command) error {
				fmt.Println("This is the testflags command.")
				fmt.Println("--yes is %v", yes)
				fmt.Println("--size is %v", size)
				return nil
			},
		})

		Add(&Command{
			Name:  "do",
			Short: "A command with subcommands: something, nothing.",
			Long: `Command do helps testing subcommands.
Usage:
	do something
	do nothing`,
		})

		Add(&Command{
			Parent: "do",
			Name:   "something",
			Short:  "A subcommand that does something.",
			Long:   "do something does something",
			Cmd: func(cmd *Command) error {
				fmt.Println("This is the do something command.")
				return nil
			},
		})

		Add(&Command{
			Parent: "do",
			Name:   "nothing",
			Short:  "A subcommand that does nothing.",
			Long:   "do nothing does nothing",
			Cmd: func(cmd *Command) error {
				fmt.Println("This is the do nothing command.")
				return nil
			},
		})

		Parse()

		Convey("readCommand should identify all of them correctly", func() {
			cmd, err := readCommand([]string{"test", "arg1", "arg2"})
			So(cmd.Name, ShouldEqual, "test")
			So(cmd.Args, ShouldResemble, []string{"arg1", "arg2"})
			So(err, ShouldBeNil)

			cmd, err = readCommand([]string{"do", "something", "arg1"})
			So(cmd.Name, ShouldEqual, "something")
			So(cmd.Args, ShouldResemble, []string{"arg1"})
			So(err, ShouldBeNil)

			cmd, err = readCommand([]string{"do", "nothing"})
			So(cmd.Name, ShouldEqual, "nothing")
			So(cmd.Args, ShouldResemble, []string{})
			So(err, ShouldBeNil)
		})

		Convey("readCommand should return the Usage command if no valid command was passed in", func() {
			cmd, err := readCommand([]string{"invalid", "arg1"})
			So(cmd.Cmd, ShouldEqual, Usage)
			So(err, ShouldBeNil)
		})

		Convey("Usage() should print the usage", func() {
			Usage(nil)
			Usage(Commands["test"])
			Usage(Commands["flags"])
			Usage(Commands["do"])
			Usage(Commands["do"].children["something"])
			Usage(Commands["do"].children["nothing"])

		})
		Reset(func() {
			Commands = make(CommandMap)
		})
	})

}
<commit_msg>Minor changes.<commit_after>package start

import (
	"fmt"
	"os"
	"testing"

	flag "github.com/ogier/pflag"
	. "github.com/smartystreets/goconvey/convey"
)

func TestAdd(t *testing.T) {
	Convey("When adding a command to the command map, then...", t, func() {
		Commands.Add(&Command{
			Name: "cmd1",
		})

		Convey("the command map should contain it", func() {
			So(Commands["cmd1"], ShouldNotBeNil)
			So(Commands["cmd1"].Name, ShouldEqual, "cmd1")

		})

		Convey("adding a command with the same name should fail", func() {
			err := Commands.Add(&Command{
				Name: "cmd1",
			})
			So(err, ShouldNotBeNil)
		})

		Convey("adding a subcommand to the command should work", func() {
			err := Commands.Add(&Command{
				Parent: "cmd1",
				Name:   "subcmd1",
			})
			So(err, ShouldBeNil)
			So(len(Commands["cmd1"].children), ShouldEqual, 1)
			So(Commands["cmd1"].children["subcmd1"], ShouldNotBeNil)
			So(Commands["cmd1"].children["subcmd1"].Name, ShouldEqual, "subcmd1")

			Convey("adding the same subcommand twice should fail", func() {
				err := Commands.Add(&Command{
					Parent: "cmd1",
					Name:   "subcmd1",
				})
				So(err, ShouldNotBeNil)
			})

		})

		Reset(func() {
			Commands = make(CommandMap)
		})
	})

}

func TestCommands(t *testing.T) {

	var yes bool
	var size int

	// ContinueOnError is required when running goconvey as server; otherwise, unrecognized
	// flags that are passed to the test executable will cause an error:
	// "unknown shorthand flag: 't' in -test.v=true"
	flag.CommandLine = flag.NewFlagSet(os.Args[0], flag.ContinueOnError)

	flag.BoolVarP(&yes, "yes", "y", false, "A boolean flag")
	flag.IntVarP(&size, "size", "s", 23, "An int flag")

	Commands = make(CommandMap)

	Convey("The flags should exist", t, func() {
		Parse()
		So(flag.Lookup("yes").Name, ShouldEqual, "yes")
		So(flag.Lookup("size").Name, ShouldEqual, "size")
	})

	Convey("When setting up some commands, then...", t, func() {
		Description = "This is the test application for the start package."

		Add(&Command{
			Name:  "test",
			Short: "A test command",
			Long:  "Command test helps testing the start package. It accepts all flags.",
			Cmd: func(cmd *Command) error {
				fmt.Println("This is the test command.")
				return nil
			},
		})

		Add(&Command{
			Name:  "flags",
			Flags: []string{"yes", "size"},
			Short: "A test command with flags",
			Long:  "Command flags helps testing flags.",
			Cmd: func(cmd *Command) error {
				fmt.Println("This is the testflags command.")
				fmt.Println("--yes is %v", yes)
				fmt.Println("--size is %v", size)
				return nil
			},
		})

		Add(&Command{
			Name:  "do",
			Short: "A command with subcommands",
			Long: `Command do helps testing subcommands.
Usage:
	do something
	do nothing`,
		})

		Add(&Command{
			Parent: "do",
			Name:   "something",
			Short:  "A subcommand that does something.",
			Long:   "do something does something",
			Cmd: func(cmd *Command) error {
				fmt.Println("This is the do something command.")
				return nil
			},
		})

		Add(&Command{
			Parent: "do",
			Name:   "nothing",
			Short:  "A subcommand that does nothing.",
			Long:   "do nothing does nothing",
			Cmd: func(cmd *Command) error {
				fmt.Println("This is the do nothing command.")
				return nil
			},
		})

		Parse()

		Convey("readCommand should identify all of them correctly", func() {
			cmd, err := readCommand([]string{"test", "arg1", "arg2"})
			So(cmd.Name, ShouldEqual, "test")
			So(cmd.Args, ShouldResemble, []string{"arg1", "arg2"})
			So(err, ShouldBeNil)

			cmd, err = readCommand([]string{"do", "something", "arg1"})
			So(cmd.Name, ShouldEqual, "something")
			So(cmd.Args, ShouldResemble, []string{"arg1"})
			So(err, ShouldBeNil)

			cmd, err = readCommand([]string{"do", "nothing"})
			So(cmd.Name, ShouldEqual, "nothing")
			So(cmd.Args, ShouldResemble, []string{})
			So(err, ShouldBeNil)
		})

		Convey("readCommand should return the Usage command if no valid command was passed in", func() {
			cmd, err := readCommand([]string{"invalid", "arg1"})
			So(cmd.Cmd, ShouldEqual, Usage)
			So(err, ShouldBeNil)
		})

		Convey("Usage() should print the usage", func() {
			Usage(nil)
			Usage(Commands["test"])
			Usage(Commands["flags"])
			Usage(Commands["do"])
			Usage(Commands["do"].children["something"])
			Usage(Commands["do"].children["nothing"])

		})
		Reset(func() {
			Commands = make(CommandMap)
		})
	})

}
<|endoftext|>


<commit_before>package sirius

import (
	"github.com/kayex/sirius/slack"
	"testing"
)

func TestMessage_Command(t *testing.T) {
	cases := []struct {
		msg    Message
		search string
		match  bool
		cmd    *Command
	}{
		{
			msg:    NewMessage(slack.UserID{UserID: "123", TeamID: "abc"}, "!foo bar", "#channel", "0"),
			search: "foo",
			match:  true,
			cmd: &Command{
				Name: "foo",
				Args: []string{
					"bar",
				},
			},
		},
		{
			msg:    NewMessage(slack.UserID{UserID: "123", TeamID: "abc"}, "!foo", "#channel", "0"),
			search: "foo",
			match:  true,
			cmd: &Command{
				Name: "foo",
				Args: []string{},
			},
		},
		{
			msg:    NewMessage(slack.UserID{UserID: "123", TeamID: "abc"}, "!foo bar baz", "#channel", "0"),
			search: "foo",
			match:  true,
			cmd: &Command{
				Name: "foo",
				Args: []string{
					"bar",
					"baz",
				},
			},
		},
		{
			msg:    NewMessage(slack.UserID{UserID: "123", TeamID: "abc"}, "!foo", "#channel", "0"),
			search: "bar",
			match:  false,
			cmd:    nil,
		},
	}

	for _, c := range cases {
		cmd, match := c.msg.Command(c.search)

		if match != c.match {
			t.Fatalf("Expected Command(%q) to return (%#v, %v), got (%#v, %v)", c.msg.Text, c.cmd, c.match, cmd, match)
		}

		if c.cmd == nil {
			return
		}

		if len(cmd.Args) != len(c.cmd.Args) {
			t.Fatalf("Expected len(Command(%q).Args) to be %d, got %d", c.msg.Text, len(c.cmd.Args), len(cmd.Args))
		}

		for i := 0; i < len(c.cmd.Args); i++ {
			expArg := c.cmd.Args[i]
			actArg := cmd.Args[i]

			if actArg != expArg {
				t.Fatalf("Expected \"%s\" as argument number %v for message \"%v\", got \"%s\"", expArg, i, c.msg.Text, actArg)
			}
		}
	}
}
<commit_msg>Errorf instead of Fatalf when message command argument parsing fails<commit_after>package sirius

import (
	"github.com/kayex/sirius/slack"
	"testing"
)

func TestMessage_Command(t *testing.T) {
	cases := []struct {
		msg    Message
		search string
		match  bool
		cmd    *Command
	}{
		{
			msg:    NewMessage(slack.UserID{UserID: "123", TeamID: "abc"}, "!foo bar", "#channel", "0"),
			search: "foo",
			match:  true,
			cmd: &Command{
				Name: "foo",
				Args: []string{
					"bar",
				},
			},
		},
		{
			msg:    NewMessage(slack.UserID{UserID: "123", TeamID: "abc"}, "!foo", "#channel", "0"),
			search: "foo",
			match:  true,
			cmd: &Command{
				Name: "foo",
				Args: []string{},
			},
		},
		{
			msg:    NewMessage(slack.UserID{UserID: "123", TeamID: "abc"}, "!foo bar baz", "#channel", "0"),
			search: "foo",
			match:  true,
			cmd: &Command{
				Name: "foo",
				Args: []string{
					"bar",
					"baz",
				},
			},
		},
		{
			msg:    NewMessage(slack.UserID{UserID: "123", TeamID: "abc"}, "!foo", "#channel", "0"),
			search: "bar",
			match:  false,
			cmd:    nil,
		},
	}

	for _, c := range cases {
		cmd, match := c.msg.Command(c.search)

		if match != c.match {
			t.Fatalf("Expected Command(%q) to return (%#v, %v), got (%#v, %v)", c.msg.Text, c.cmd, c.match, cmd, match)
		}

		if c.cmd == nil {
			return
		}

		if len(cmd.Args) != len(c.cmd.Args) {
			t.Fatalf("Expected len(Command(%q).Args) to be %d, got %d", c.msg.Text, len(c.cmd.Args), len(cmd.Args))
		}

		for i := 0; i < len(c.cmd.Args); i++ {
			expArg := c.cmd.Args[i]
			actArg := cmd.Args[i]

			if actArg != expArg {
				t.Errorf("Expected Command(%q).Args[%d] to be %q, got %q", c.msg.Text, i, expArg, actArg)
			}
		}
	}
}
<|endoftext|>


<commit_before>// Copyright 2016 The Neugram Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

import (
	"bufio"
	"bytes"
	"flag"
	"fmt"
	"io"
	"os"
	"os/exec"
	"os/signal"
	"path/filepath"
	"reflect"
	"strings"
	"time"

	"neugram.io/ng/eval"
	"neugram.io/ng/eval/environ"
	"neugram.io/ng/eval/shell"
	"neugram.io/ng/parser"
	"neugram.io/ng/tipe"

	"github.com/kr/pretty"
	"github.com/peterh/liner"
)

var (
	origMode liner.ModeApplier

	lineNg        *liner.State // ng-mode line reader
	historyNgFile = ""
	historyNg     = make(chan string, 1)
	historyShFile = ""
	historySh     = make(chan string, 1)
	sigint        = make(chan os.Signal, 1)

	p   *parser.Parser
	prg *eval.Program
)

func exit(code int) {
	lineNg.Close()
	os.Exit(code)
}

func exitf(format string, args ...interface{}) {
	fmt.Fprintf(os.Stderr, "ng: "+format, args...)
	exit(1)
}

func mode() liner.ModeApplier {
	m, err := liner.TerminalMode()
	if err != nil {
		exitf("terminal mode: %v", err)
	}
	return m
}

const usageLine = "ng [programfile | -e cmd] [arguments]"

func usage() {
	fmt.Fprintf(os.Stderr, `ng - neugram scripting language and shell

Usage:
	%s

Options:
`, usageLine)
	flag.PrintDefaults()
}

func main() {
	shell.Init()

	help := flag.Bool("h", false, "display help message and exit")
	e := flag.String("e", "", "program passed as a string")
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "usage: %s\n", usageLine)
		os.Exit(1)
	}
	flag.Parse()

	if *help {
		usage()
		os.Exit(0)
	}
	if *e != "" {
		initProgram(filepath.Join(cwd, "ng-arg"))
		res := p.ParseLine([]byte(*e))
		handleResult(res)
		return
	}

	origMode = mode()
	lineNg = liner.NewLiner()
	loop()
}

func setWindowSize(env map[interface{}]interface{}) {
	// TODO windowsize
	// TODO
	// TODO
	// TODO
	/*
		rows, cols, err := job.WindowSize(os.Stderr.Fd())
		if err != nil {
			fmt.Printf("ng: could not get window size: %v\n", err)
		} else {
			// TODO: these are meant to be shell variables, not
			// environment variables. But then, how do programs
			// like `ls` read them?
			env["LINES"] = strconv.Itoa(rows)
			env["COLUMNS"] = strconv.Itoa(cols)
		}
	*/
}

func ps1(env *environ.Environ) string {
	v := env.Get("PS1")
	if v == "" {
		return "ng$ "
	}
	if strings.IndexByte(v, '\\') == -1 {
		return v
	}
	var buf []byte
	for {
		i := strings.IndexByte(v, '\\')
		if i == -1 || i == len(v)-1 {
			break
		}
		buf = append(buf, v[:i]...)
		b := v[i+1]
		v = v[i+2:]
		switch b {
		case 'h', 'H':
			out, err := exec.Command("hostname").CombinedOutput()
			if err != nil {
				fmt.Fprintf(os.Stderr, "ng: %v\n", err)
				continue
			}
			if b == 'h' {
				if i := bytes.IndexByte(out, '.'); i >= 0 {
					out = out[:i]
				}
			}
			if len(out) > 0 && out[len(out)-1] == '\n' {
				out = out[:len(out)-1]
			}
			buf = append(buf, out...)
		case 'n':
			buf = append(buf, '\n')
		case 'w', 'W':
			cwd := env.Get("PWD")
			if home := env.Get("HOME"); home != "" {
				cwd = strings.Replace(cwd, home, "~", 1)
			}
			if b == 'W' {
				cwd = filepath.Base(cwd)
			}
			buf = append(buf, cwd...)
		}
		// TODO: '!', '#', '$', 'nnn', 's', 'j', and more.
	}
	buf = append(buf, v...)
	return string(buf)
}

var cwd string

func init() {
	var err error
	cwd, err = os.Getwd()
	if err != nil {
		panic(err)
	}
}

func initProgram(path string) {
	p = parser.New()
	prg = eval.New(path)
	shell.Env = prg.Environ()
	shell.Alias = prg.Alias()

	// TODO this env setup could be done in neugram code
	env := prg.Environ()
	for _, s := range os.Environ() {
		i := strings.Index(s, "=")
		env.Set(s[:i], s[i+1:])
	}
	wd, err := os.Getwd()
	if err == nil {
		env.Set("PWD", wd)
	}
	//setWindowSize(env)

	signal.Notify(sigint, os.Interrupt)
}

func loop() {
	// TODO: support starting via a shebang: #!/bin/ng.
	// When doing so, path = os.Args[0]
	path := filepath.Join(cwd, "ng-interactive")
	initProgram(path)

	state := parser.StateStmt
	if os.Args[0] == "ngsh" || os.Args[0] == "-ngsh" {
		initFile := filepath.Join(os.Getenv("HOME"), ".ngshinit")
		if f, err := os.Open(initFile); err == nil {
			scanner := bufio.NewScanner(f)
			for scanner.Scan() {
				res := p.ParseLine(scanner.Bytes())
				handleResult(res)
				state = res.State
			}
			if err := scanner.Err(); err != nil {
				exitf(".ngshinit: %v", err)
			}
			f.Close()
		}
		switch state {
		case parser.StateStmtPartial, parser.StateCmdPartial:
			exitf(".ngshinit: ends in a partial statement")
		case parser.StateStmt:
			res := p.ParseLine([]byte("$$"))
			handleResult(res)
			state = res.State
		}
	}

	lineNg.SetTabCompletionStyle(liner.TabPrints)
	lineNg.SetWordCompleter(completer)
	lineNg.SetCtrlCAborts(true)

	if f, err := os.Open(historyShFile); err == nil {
		lineNg.SetMode("sh")
		lineNg.ReadHistory(f)
		f.Close()
	}
	go historyWriter(historyShFile, historySh)

	if f, err := os.Open(historyNgFile); err == nil {
		lineNg.SetMode("ng")
		lineNg.ReadHistory(f)
		f.Close()
	}
	go historyWriter(historyNgFile, historyNg)

	for {
		var (
			mode    string
			prompt  string
			history chan string
		)
		switch state {
		case parser.StateUnknown:
			mode, prompt, history = "ng", "??> ", historyNg
		case parser.StateStmt:
			mode, prompt, history = "ng", "ng> ", historyNg
		case parser.StateStmtPartial:
			mode, prompt, history = "ng", "..> ", historyNg
		case parser.StateCmd:
			mode, prompt, history = "sh", ps1(prg.Environ()), historySh
		case parser.StateCmdPartial:
			mode, prompt, history = "sh", "..$ ", historySh
		default:
			exitf("unkown parser state: %v", state)
		}
		lineNg.SetMode(mode)
		data, err := lineNg.Prompt(prompt)
		if err == liner.ErrPromptAborted {
			switch state {
			case parser.StateStmtPartial:
				fmt.Printf("TODO interrupt partial statement\n")
			case parser.StateCmdPartial:
				fmt.Printf("TODO interrupt partial command\n")
			}
		} else if err != nil {
			if err == io.EOF {
				exit(0)
			}
			exitf("error reading input: %v", err)
		}
		if data == "" {
			continue
		}
		lineNg.AppendHistory(mode, data)
		history <- data
		select { // drain sigint
		case <-sigint:
		default:
		}
		res := p.ParseLine([]byte(data))
		handleResult(res)
		state = res.State
	}
}

func handleResult(res parser.Result) {
	for _, s := range res.Stmts {
		v, err := prg.Eval(s, sigint)
		if err != nil {
			fmt.Printf("ng: %v\n", err)
			continue
		}
		if len(v) > 1 {
			fmt.Print("(")
		}
		for i, val := range v {
			if i > 0 {
				fmt.Print(", ")
			}
			if val == (reflect.Value{}) {
				fmt.Print("<nil>")
			} else {
				pretty.Print(val.Interface())
			}
		}
		if len(v) > 1 {
			fmt.Println(")")
		} else if len(v) == 1 {
			fmt.Println("")
		}
	}
	for _, err := range res.Errs {
		fmt.Println(err.Error())
	}
	//editMode := mode()
	//origMode.ApplyMode()
	for _, cmd := range res.Cmds {
		j := &shell.Job{
			Cmd:    cmd,
			Params: prg,
			Stdin:  os.Stdin,
			Stdout: os.Stdout,
			Stderr: os.Stderr,
		}
		if err := j.Start(); err != nil {
			fmt.Println(err)
			continue
		}
		done, err := j.Wait()
		if err != nil {
			fmt.Println(err)
			continue
		}
		if !done {
			break // TODO not right, instead we should just have one cmd, not Cmds here.
		}
	}
	//editMode.ApplyMode()
}

func printValue(t tipe.Type, v interface{}) {
	// This is, effectively, a primitive type-aware printf implementation
	// that understands the neugram evaluator data layout. A far better
	// version of this would be an "ngfmt" package, that implemented the
	// printing command in neugram, using a "ngreflect" package. But it
	// will be a while until I build a reflect package, so this will have
	// to do.
	//
	// Still: avoid putting too much machinary in this. At some point soon
	// it's not worth the effort.
	/*switch t := tipe.Underlying(t).(type) {
	case *tipe.Struct:
	fmt.Print("{")
	for i, name := range t.FieldNames {
		fmt.Printf("%s: ", name)
		printValue(t.Fields[i], v.(*eval.StructVal).Fields[i].Value)
		if i < len(t.FieldNames)-1 {
			fmt.Print(", ")
		}
	}
	fmt.Print("}")
	default:
	}*/
	fmt.Print(v)
}

func init() {
	if home := os.Getenv("HOME"); home != "" {
		historyNgFile = filepath.Join(home, ".ng_history")
		historyShFile = filepath.Join(home, ".ngsh_history")
	}
}

func historyWriter(dst string, src <-chan string) {
	var batch []string
	ticker := time.Tick(250 * time.Millisecond)
	for {
		select {
		case line := <-src:
			batch = append(batch, line)
		case <-ticker:
			if len(batch) > 0 && dst != "" {
				// TODO: FcntlFlock
				f, err := os.OpenFile(dst, os.O_WRONLY|os.O_APPEND|os.O_CREATE, 0664)
				if err == nil {
					for _, line := range batch {
						fmt.Fprintf(f, "%s\n", line)
					}
					f.Close()
				}
			}
			batch = nil
		}
	}
}
<commit_msg>ng: shebang support<commit_after>// Copyright 2016 The Neugram Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

import (
	"bufio"
	"bytes"
	"flag"
	"fmt"
	"io"
	"os"
	"os/exec"
	"os/signal"
	"path/filepath"
	"reflect"
	"strings"
	"time"

	"neugram.io/ng/eval"
	"neugram.io/ng/eval/environ"
	"neugram.io/ng/eval/shell"
	"neugram.io/ng/parser"
	"neugram.io/ng/tipe"

	"github.com/kr/pretty"
	"github.com/peterh/liner"
)

var (
	origMode liner.ModeApplier

	lineNg        *liner.State // ng-mode line reader
	historyNgFile = ""
	historyNg     = make(chan string, 1)
	historyShFile = ""
	historySh     = make(chan string, 1)
	sigint        = make(chan os.Signal, 1)

	p   *parser.Parser
	prg *eval.Program
)

func exit(code int) {
	if lineNg != nil {
		lineNg.Close()
	}
	os.Exit(code)
}

func exitf(format string, args ...interface{}) {
	fmt.Fprintf(os.Stderr, "ng: "+format+"\n", args...)
	exit(1)
}

func mode() liner.ModeApplier {
	m, err := liner.TerminalMode()
	if err != nil {
		exitf("terminal mode: %v", err)
	}
	return m
}

const usageLine = "ng [programfile | -e cmd] [arguments]"

func usage() {
	fmt.Fprintf(os.Stderr, `ng - neugram scripting language and shell

Usage:
	%s

Options:
`, usageLine)
	flag.PrintDefaults()
}

func main() {
	shell.Init()

	help := flag.Bool("h", false, "display help message and exit")
	e := flag.String("e", "", "program passed as a string")
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "usage: %s\n", usageLine)
		os.Exit(1)
	}
	flag.Parse()

	if *help {
		usage()
		os.Exit(0)
	}
	if *e != "" {
		initProgram(filepath.Join(cwd, "ng-arg"))
		res := p.ParseLine([]byte(*e))
		handleResult(res)
		return
	}
	if args := flag.Args(); len(args) > 0 {
		// TODO: plumb through the rest of the args
		path := args[0]
		initProgram(path)
		f, err := os.Open(path)
		if err != nil {
			exitf("%v", err)
		}
		state, err := runFile(f)
		if err != nil {
			exitf("%v", err)
		}
		if state == parser.StateCmd {
			exitf("%s: ends in an unclosed shell statement", args[0])
		}
		return
	}

	origMode = mode()
	lineNg = liner.NewLiner()
	loop()
}

func setWindowSize(env map[interface{}]interface{}) {
	// TODO windowsize
	// TODO
	// TODO
	// TODO
	/*
		rows, cols, err := job.WindowSize(os.Stderr.Fd())
		if err != nil {
			fmt.Printf("ng: could not get window size: %v\n", err)
		} else {
			// TODO: these are meant to be shell variables, not
			// environment variables. But then, how do programs
			// like `ls` read them?
			env["LINES"] = strconv.Itoa(rows)
			env["COLUMNS"] = strconv.Itoa(cols)
		}
	*/
}

func ps1(env *environ.Environ) string {
	v := env.Get("PS1")
	if v == "" {
		return "ng$ "
	}
	if strings.IndexByte(v, '\\') == -1 {
		return v
	}
	var buf []byte
	for {
		i := strings.IndexByte(v, '\\')
		if i == -1 || i == len(v)-1 {
			break
		}
		buf = append(buf, v[:i]...)
		b := v[i+1]
		v = v[i+2:]
		switch b {
		case 'h', 'H':
			out, err := exec.Command("hostname").CombinedOutput()
			if err != nil {
				fmt.Fprintf(os.Stderr, "ng: %v\n", err)
				continue
			}
			if b == 'h' {
				if i := bytes.IndexByte(out, '.'); i >= 0 {
					out = out[:i]
				}
			}
			if len(out) > 0 && out[len(out)-1] == '\n' {
				out = out[:len(out)-1]
			}
			buf = append(buf, out...)
		case 'n':
			buf = append(buf, '\n')
		case 'w', 'W':
			cwd := env.Get("PWD")
			if home := env.Get("HOME"); home != "" {
				cwd = strings.Replace(cwd, home, "~", 1)
			}
			if b == 'W' {
				cwd = filepath.Base(cwd)
			}
			buf = append(buf, cwd...)
		}
		// TODO: '!', '#', '$', 'nnn', 's', 'j', and more.
	}
	buf = append(buf, v...)
	return string(buf)
}

var cwd string

func init() {
	var err error
	cwd, err = os.Getwd()
	if err != nil {
		panic(err)
	}
}

func initProgram(path string) {
	p = parser.New()
	prg = eval.New(path)
	shell.Env = prg.Environ()
	shell.Alias = prg.Alias()

	// TODO this env setup could be done in neugram code
	env := prg.Environ()
	for _, s := range os.Environ() {
		i := strings.Index(s, "=")
		env.Set(s[:i], s[i+1:])
	}
	wd, err := os.Getwd()
	if err == nil {
		env.Set("PWD", wd)
	}
	//setWindowSize(env)

	signal.Notify(sigint, os.Interrupt)
}

func runFile(f *os.File) (parser.ParserState, error) {
	state := parser.StateStmt
	scanner := bufio.NewScanner(f)
	for i := 0; scanner.Scan(); i++ {
		b := scanner.Bytes()
		if i == 0 && len(b) > 2 && b[0] == '#' && b[1] == '!' { // shebang
			continue
		}
		res := p.ParseLine(b)
		handleResult(res)
		state = res.State
	}
	if err := scanner.Err(); err != nil {
		return state, fmt.Errorf("%s: %v", f.Name(), err)
	}
	switch state {
	case parser.StateStmtPartial, parser.StateCmdPartial:
		return state, fmt.Errorf("%s: ends in a partial statement", f.Name())
	default:
		return state, nil
	}
}

func loop() {
	path := filepath.Join(cwd, "ng-interactive")
	initProgram(path)

	state := parser.StateStmt
	if os.Args[0] == "ngsh" || os.Args[0] == "-ngsh" {
		initFile := filepath.Join(os.Getenv("HOME"), ".ngshinit")
		if f, err := os.Open(initFile); err == nil {
			var err error
			state, err = runFile(f)
			f.Close()
			if err != nil {
				exitf("%v", err)
			}
		}
		if state == parser.StateStmt {
			res := p.ParseLine([]byte("$$"))
			handleResult(res)
			state = res.State
		}
	}

	lineNg.SetTabCompletionStyle(liner.TabPrints)
	lineNg.SetWordCompleter(completer)
	lineNg.SetCtrlCAborts(true)

	if f, err := os.Open(historyShFile); err == nil {
		lineNg.SetMode("sh")
		lineNg.ReadHistory(f)
		f.Close()
	}
	go historyWriter(historyShFile, historySh)

	if f, err := os.Open(historyNgFile); err == nil {
		lineNg.SetMode("ng")
		lineNg.ReadHistory(f)
		f.Close()
	}
	go historyWriter(historyNgFile, historyNg)

	for {
		var (
			mode    string
			prompt  string
			history chan string
		)
		switch state {
		case parser.StateUnknown:
			mode, prompt, history = "ng", "??> ", historyNg
		case parser.StateStmt:
			mode, prompt, history = "ng", "ng> ", historyNg
		case parser.StateStmtPartial:
			mode, prompt, history = "ng", "..> ", historyNg
		case parser.StateCmd:
			mode, prompt, history = "sh", ps1(prg.Environ()), historySh
		case parser.StateCmdPartial:
			mode, prompt, history = "sh", "..$ ", historySh
		default:
			exitf("unkown parser state: %v", state)
		}
		lineNg.SetMode(mode)
		data, err := lineNg.Prompt(prompt)
		if err == liner.ErrPromptAborted {
			switch state {
			case parser.StateStmtPartial:
				fmt.Printf("TODO interrupt partial statement\n")
			case parser.StateCmdPartial:
				fmt.Printf("TODO interrupt partial command\n")
			}
		} else if err != nil {
			if err == io.EOF {
				exit(0)
			}
			exitf("error reading input: %v", err)
		}
		if data == "" {
			continue
		}
		lineNg.AppendHistory(mode, data)
		history <- data
		select { // drain sigint
		case <-sigint:
		default:
		}
		res := p.ParseLine([]byte(data))
		handleResult(res)
		state = res.State
	}
}

func handleResult(res parser.Result) {
	for _, s := range res.Stmts {
		v, err := prg.Eval(s, sigint)
		if err != nil {
			fmt.Printf("ng: %v\n", err)
			continue
		}
		if len(v) > 1 {
			fmt.Print("(")
		}
		for i, val := range v {
			if i > 0 {
				fmt.Print(", ")
			}
			if val == (reflect.Value{}) {
				fmt.Print("<nil>")
			} else {
				pretty.Print(val.Interface())
			}
		}
		if len(v) > 1 {
			fmt.Println(")")
		} else if len(v) == 1 {
			fmt.Println("")
		}
	}
	for _, err := range res.Errs {
		fmt.Println(err.Error())
	}
	//editMode := mode()
	//origMode.ApplyMode()
	for _, cmd := range res.Cmds {
		j := &shell.Job{
			Cmd:    cmd,
			Params: prg,
			Stdin:  os.Stdin,
			Stdout: os.Stdout,
			Stderr: os.Stderr,
		}
		if err := j.Start(); err != nil {
			fmt.Println(err)
			continue
		}
		done, err := j.Wait()
		if err != nil {
			fmt.Println(err)
			continue
		}
		if !done {
			break // TODO not right, instead we should just have one cmd, not Cmds here.
		}
	}
	//editMode.ApplyMode()
}

func printValue(t tipe.Type, v interface{}) {
	// This is, effectively, a primitive type-aware printf implementation
	// that understands the neugram evaluator data layout. A far better
	// version of this would be an "ngfmt" package, that implemented the
	// printing command in neugram, using a "ngreflect" package. But it
	// will be a while until I build a reflect package, so this will have
	// to do.
	//
	// Still: avoid putting too much machinary in this. At some point soon
	// it's not worth the effort.
	/*switch t := tipe.Underlying(t).(type) {
	case *tipe.Struct:
	fmt.Print("{")
	for i, name := range t.FieldNames {
		fmt.Printf("%s: ", name)
		printValue(t.Fields[i], v.(*eval.StructVal).Fields[i].Value)
		if i < len(t.FieldNames)-1 {
			fmt.Print(", ")
		}
	}
	fmt.Print("}")
	default:
	}*/
	fmt.Print(v)
}

func init() {
	if home := os.Getenv("HOME"); home != "" {
		historyNgFile = filepath.Join(home, ".ng_history")
		historyShFile = filepath.Join(home, ".ngsh_history")
	}
}

func historyWriter(dst string, src <-chan string) {
	var batch []string
	ticker := time.Tick(250 * time.Millisecond)
	for {
		select {
		case line := <-src:
			batch = append(batch, line)
		case <-ticker:
			if len(batch) > 0 && dst != "" {
				// TODO: FcntlFlock
				f, err := os.OpenFile(dst, os.O_WRONLY|os.O_APPEND|os.O_CREATE, 0664)
				if err == nil {
					for _, line := range batch {
						fmt.Fprintf(f, "%s\n", line)
					}
					f.Close()
				}
			}
			batch = nil
		}
	}
}
<|endoftext|>


<commit_before>// Copyright 2012 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package x509

// RFC 1423 describes the encryption of PEM blocks. The algorithm used to
// generate a key from the password was derived by looking at the OpenSSL
// implementation.

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/des"
	"crypto/md5"
	"encoding/hex"
	"encoding/pem"
	"errors"
	"io"
	"strings"
)

type PEMCipher int

// Possible values for the EncryptPEMBlock encryption algorithm.
const (
	_ PEMCipher = iota
	PEMCipherDES
	PEMCipher3DES
	PEMCipherAES128
	PEMCipherAES192
	PEMCipherAES256
)

// rfc1423Algo holds a method for enciphering a PEM block.
type rfc1423Algo struct {
	cipher     PEMCipher
	name       string
	cipherFunc func(key []byte) (cipher.Block, error)
	keySize    int
	blockSize  int
}

// rfc1423Algos holds a slice of the possible ways to encrypt a PEM
// block. The ivSize numbers were taken from the OpenSSL source.
var rfc1423Algos = []rfc1423Algo{{
	cipher:     PEMCipherDES,
	name:       "DES-CBC",
	cipherFunc: des.NewCipher,
	keySize:    8,
	blockSize:  des.BlockSize,
}, {
	cipher:     PEMCipher3DES,
	name:       "DES-EDE3-CBC",
	cipherFunc: des.NewTripleDESCipher,
	keySize:    24,
	blockSize:  des.BlockSize,
}, {
	cipher:     PEMCipherAES128,
	name:       "AES-128-CBC",
	cipherFunc: aes.NewCipher,
	keySize:    16,
	blockSize:  aes.BlockSize,
}, {
	cipher:     PEMCipherAES192,
	name:       "AES-192-CBC",
	cipherFunc: aes.NewCipher,
	keySize:    24,
	blockSize:  aes.BlockSize,
}, {
	cipher:     PEMCipherAES256,
	name:       "AES-256-CBC",
	cipherFunc: aes.NewCipher,
	keySize:    32,
	blockSize:  aes.BlockSize,
},
}

// deriveKey uses a key derivation function to stretch the password into a key
// with the number of bits our cipher requires. This algorithm was derived from
// the OpenSSL source.
func (c rfc1423Algo) deriveKey(password, salt []byte) []byte {
	hash := md5.New()
	out := make([]byte, c.keySize)
	var digest []byte

	for i := 0; i < len(out); i += len(digest) {
		hash.Reset()
		hash.Write(digest)
		hash.Write(password)
		hash.Write(salt)
		digest = hash.Sum(digest[:0])
		copy(out[i:], digest)
	}
	return out
}

// IsEncryptedPEMBlock returns if the PEM block is password encrypted.
func IsEncryptedPEMBlock(b *pem.Block) bool {
	_, ok := b.Headers["DEK-Info"]
	return ok
}

// IncorrectPasswordError is returned when an incorrect password is detected.
var IncorrectPasswordError = errors.New("x509: decryption password incorrect")

// DecryptPEMBlock takes a password encrypted PEM block and the password used to
// encrypt it and returns a slice of decrypted DER encoded bytes. It inspects
// the DEK-Info header to determine the algorithm used for decryption. If no
// DEK-Info header is present, an error is returned. If an incorrect password
// is detected an IncorrectPasswordError is returned. Because of deficiencies
// in the encrypted-PEM format, it's not always possible to detect an incorrect
// password. In these cases no error will be returned but the decrypted DER
// bytes will be random noise.
func DecryptPEMBlock(b *pem.Block, password []byte) ([]byte, error) {
	dek, ok := b.Headers["DEK-Info"]
	if !ok {
		return nil, errors.New("x509: no DEK-Info header in block")
	}

	idx := strings.Index(dek, ",")
	if idx == -1 {
		return nil, errors.New("x509: malformed DEK-Info header")
	}

	mode, hexIV := dek[:idx], dek[idx+1:]
	ciph := cipherByName(mode)
	if ciph == nil {
		return nil, errors.New("x509: unknown encryption mode")
	}
	iv, err := hex.DecodeString(hexIV)
	if err != nil {
		return nil, err
	}
	if len(iv) != ciph.blockSize {
		return nil, errors.New("x509: incorrect IV size")
	}

	// Based on the OpenSSL implementation. The salt is the first 8 bytes
	// of the initialization vector.
	key := ciph.deriveKey(password, iv[:8])
	block, err := ciph.cipherFunc(key)
	if err != nil {
		return nil, err
	}

	if len(b.Bytes)%block.BlockSize() != 0 {
		return nil, errors.New("x509: encrypted PEM data is not a multiple of the block size")
	}

	data := make([]byte, len(b.Bytes))
	dec := cipher.NewCBCDecrypter(block, iv)
	dec.CryptBlocks(data, b.Bytes)

	// Blocks are padded using a scheme where the last n bytes of padding are all
	// equal to n. It can pad from 1 to blocksize bytes inclusive. See RFC 1423.
	// For example:
	//	[x y z 2 2]
	//	[x y 7 7 7 7 7 7 7]
	// If we detect a bad padding, we assume it is an invalid password.
	dlen := len(data)
	if dlen == 0 || dlen%ciph.blockSize != 0 {
		return nil, errors.New("x509: invalid padding")
	}
	last := int(data[dlen-1])
	if dlen < last {
		return nil, IncorrectPasswordError
	}
	if last == 0 || last > ciph.blockSize {
		return nil, IncorrectPasswordError
	}
	for _, val := range data[dlen-last:] {
		if int(val) != last {
			return nil, IncorrectPasswordError
		}
	}
	return data[:dlen-last], nil
}

// EncryptPEMBlock returns a PEM block of the specified type holding the
// given DER-encoded data encrypted with the specified algorithm and
// password.
func EncryptPEMBlock(rand io.Reader, blockType string, data, password []byte, alg PEMCipher) (*pem.Block, error) {
	ciph := cipherByKey(alg)
	if ciph == nil {
		return nil, errors.New("x509: unknown encryption mode")
	}
	iv := make([]byte, ciph.blockSize)
	if _, err := io.ReadFull(rand, iv); err != nil {
		return nil, errors.New("x509: cannot generate IV: " + err.Error())
	}
	// The salt is the first 8 bytes of the initialization vector,
	// matching the key derivation in DecryptPEMBlock.
	key := ciph.deriveKey(password, iv[:8])
	block, err := ciph.cipherFunc(key)
	if err != nil {
		return nil, err
	}
	enc := cipher.NewCBCEncrypter(block, iv)
	pad := ciph.blockSize - len(data)%ciph.blockSize
	encrypted := make([]byte, len(data), len(data)+pad)
	// We could save this copy by encrypting all the whole blocks in
	// the data separately, but it doesn't seem worth the additional
	// code.
	copy(encrypted, data)
	// See RFC 1423, Section 1.1.
	for i := 0; i < pad; i++ {
		encrypted = append(encrypted, byte(pad))
	}
	enc.CryptBlocks(encrypted, encrypted)

	return &pem.Block{
		Type: blockType,
		Headers: map[string]string{
			"Proc-Type": "4,ENCRYPTED",
			"DEK-Info":  ciph.name + "," + hex.EncodeToString(iv),
		},
		Bytes: encrypted,
	}, nil
}

func cipherByName(name string) *rfc1423Algo {
	for i := range rfc1423Algos {
		alg := &rfc1423Algos[i]
		if alg.name == name {
			return alg
		}
	}
	return nil
}

func cipherByKey(key PEMCipher) *rfc1423Algo {
	for i := range rfc1423Algos {
		alg := &rfc1423Algos[i]
		if alg.cipher == key {
			return alg
		}
	}
	return nil
}
<commit_msg>crypto/x509: deprecate legacy PEM encryption<commit_after>// Copyright 2012 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package x509

// RFC 1423 describes the encryption of PEM blocks. The algorithm used to
// generate a key from the password was derived by looking at the OpenSSL
// implementation.

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/des"
	"crypto/md5"
	"encoding/hex"
	"encoding/pem"
	"errors"
	"io"
	"strings"
)

type PEMCipher int

// Possible values for the EncryptPEMBlock encryption algorithm.
const (
	_ PEMCipher = iota
	PEMCipherDES
	PEMCipher3DES
	PEMCipherAES128
	PEMCipherAES192
	PEMCipherAES256
)

// rfc1423Algo holds a method for enciphering a PEM block.
type rfc1423Algo struct {
	cipher     PEMCipher
	name       string
	cipherFunc func(key []byte) (cipher.Block, error)
	keySize    int
	blockSize  int
}

// rfc1423Algos holds a slice of the possible ways to encrypt a PEM
// block. The ivSize numbers were taken from the OpenSSL source.
var rfc1423Algos = []rfc1423Algo{{
	cipher:     PEMCipherDES,
	name:       "DES-CBC",
	cipherFunc: des.NewCipher,
	keySize:    8,
	blockSize:  des.BlockSize,
}, {
	cipher:     PEMCipher3DES,
	name:       "DES-EDE3-CBC",
	cipherFunc: des.NewTripleDESCipher,
	keySize:    24,
	blockSize:  des.BlockSize,
}, {
	cipher:     PEMCipherAES128,
	name:       "AES-128-CBC",
	cipherFunc: aes.NewCipher,
	keySize:    16,
	blockSize:  aes.BlockSize,
}, {
	cipher:     PEMCipherAES192,
	name:       "AES-192-CBC",
	cipherFunc: aes.NewCipher,
	keySize:    24,
	blockSize:  aes.BlockSize,
}, {
	cipher:     PEMCipherAES256,
	name:       "AES-256-CBC",
	cipherFunc: aes.NewCipher,
	keySize:    32,
	blockSize:  aes.BlockSize,
},
}

// deriveKey uses a key derivation function to stretch the password into a key
// with the number of bits our cipher requires. This algorithm was derived from
// the OpenSSL source.
func (c rfc1423Algo) deriveKey(password, salt []byte) []byte {
	hash := md5.New()
	out := make([]byte, c.keySize)
	var digest []byte

	for i := 0; i < len(out); i += len(digest) {
		hash.Reset()
		hash.Write(digest)
		hash.Write(password)
		hash.Write(salt)
		digest = hash.Sum(digest[:0])
		copy(out[i:], digest)
	}
	return out
}

// IsEncryptedPEMBlock returns whether the PEM block is password encrypted
// according to RFC 1423.
//
// Deprecated: Legacy PEM encryption as specified in RFC 1423 is insecure by
// design. Since it does not authenticate the ciphertext, it is vulnerable to
// padding oracle attacks that can let an attacker recover the plaintext.
func IsEncryptedPEMBlock(b *pem.Block) bool {
	_, ok := b.Headers["DEK-Info"]
	return ok
}

// IncorrectPasswordError is returned when an incorrect password is detected.
var IncorrectPasswordError = errors.New("x509: decryption password incorrect")

// DecryptPEMBlock takes a PEM block encrypted according to RFC 1423 and the
// password used to encrypt it and returns a slice of decrypted DER encoded
// bytes. It inspects the DEK-Info header to determine the algorithm used for
// decryption. If no DEK-Info header is present, an error is returned. If an
// incorrect password is detected an IncorrectPasswordError is returned. Because
// of deficiencies in the format, it's not always possible to detect an
// incorrect password. In these cases no error will be returned but the
// decrypted DER bytes will be random noise.
//
// Deprecated: Legacy PEM encryption as specified in RFC 1423 is insecure by
// design. Since it does not authenticate the ciphertext, it is vulnerable to
// padding oracle attacks that can let an attacker recover the plaintext.
func DecryptPEMBlock(b *pem.Block, password []byte) ([]byte, error) {
	dek, ok := b.Headers["DEK-Info"]
	if !ok {
		return nil, errors.New("x509: no DEK-Info header in block")
	}

	idx := strings.Index(dek, ",")
	if idx == -1 {
		return nil, errors.New("x509: malformed DEK-Info header")
	}

	mode, hexIV := dek[:idx], dek[idx+1:]
	ciph := cipherByName(mode)
	if ciph == nil {
		return nil, errors.New("x509: unknown encryption mode")
	}
	iv, err := hex.DecodeString(hexIV)
	if err != nil {
		return nil, err
	}
	if len(iv) != ciph.blockSize {
		return nil, errors.New("x509: incorrect IV size")
	}

	// Based on the OpenSSL implementation. The salt is the first 8 bytes
	// of the initialization vector.
	key := ciph.deriveKey(password, iv[:8])
	block, err := ciph.cipherFunc(key)
	if err != nil {
		return nil, err
	}

	if len(b.Bytes)%block.BlockSize() != 0 {
		return nil, errors.New("x509: encrypted PEM data is not a multiple of the block size")
	}

	data := make([]byte, len(b.Bytes))
	dec := cipher.NewCBCDecrypter(block, iv)
	dec.CryptBlocks(data, b.Bytes)

	// Blocks are padded using a scheme where the last n bytes of padding are all
	// equal to n. It can pad from 1 to blocksize bytes inclusive. See RFC 1423.
	// For example:
	//	[x y z 2 2]
	//	[x y 7 7 7 7 7 7 7]
	// If we detect a bad padding, we assume it is an invalid password.
	dlen := len(data)
	if dlen == 0 || dlen%ciph.blockSize != 0 {
		return nil, errors.New("x509: invalid padding")
	}
	last := int(data[dlen-1])
	if dlen < last {
		return nil, IncorrectPasswordError
	}
	if last == 0 || last > ciph.blockSize {
		return nil, IncorrectPasswordError
	}
	for _, val := range data[dlen-last:] {
		if int(val) != last {
			return nil, IncorrectPasswordError
		}
	}
	return data[:dlen-last], nil
}

// EncryptPEMBlock returns a PEM block of the specified type holding the
// given DER encoded data encrypted with the specified algorithm and
// password according to RFC 1423.
//
// Deprecated: Legacy PEM encryption as specified in RFC 1423 is insecure by
// design. Since it does not authenticate the ciphertext, it is vulnerable to
// padding oracle attacks that can let an attacker recover the plaintext.
func EncryptPEMBlock(rand io.Reader, blockType string, data, password []byte, alg PEMCipher) (*pem.Block, error) {
	ciph := cipherByKey(alg)
	if ciph == nil {
		return nil, errors.New("x509: unknown encryption mode")
	}
	iv := make([]byte, ciph.blockSize)
	if _, err := io.ReadFull(rand, iv); err != nil {
		return nil, errors.New("x509: cannot generate IV: " + err.Error())
	}
	// The salt is the first 8 bytes of the initialization vector,
	// matching the key derivation in DecryptPEMBlock.
	key := ciph.deriveKey(password, iv[:8])
	block, err := ciph.cipherFunc(key)
	if err != nil {
		return nil, err
	}
	enc := cipher.NewCBCEncrypter(block, iv)
	pad := ciph.blockSize - len(data)%ciph.blockSize
	encrypted := make([]byte, len(data), len(data)+pad)
	// We could save this copy by encrypting all the whole blocks in
	// the data separately, but it doesn't seem worth the additional
	// code.
	copy(encrypted, data)
	// See RFC 1423, Section 1.1.
	for i := 0; i < pad; i++ {
		encrypted = append(encrypted, byte(pad))
	}
	enc.CryptBlocks(encrypted, encrypted)

	return &pem.Block{
		Type: blockType,
		Headers: map[string]string{
			"Proc-Type": "4,ENCRYPTED",
			"DEK-Info":  ciph.name + "," + hex.EncodeToString(iv),
		},
		Bytes: encrypted,
	}, nil
}

func cipherByName(name string) *rfc1423Algo {
	for i := range rfc1423Algos {
		alg := &rfc1423Algos[i]
		if alg.name == name {
			return alg
		}
	}
	return nil
}

func cipherByKey(key PEMCipher) *rfc1423Algo {
	for i := range rfc1423Algos {
		alg := &rfc1423Algos[i]
		if alg.cipher == key {
			return alg
		}
	}
	return nil
}
<|endoftext|>


<commit_before>package workers

import (
	"fmt"
	"github.com/APTrust/exchange/constants"
	"github.com/APTrust/exchange/context"
	"github.com/APTrust/exchange/models"
	"github.com/APTrust/exchange/network"
	"github.com/APTrust/exchange/stats"
	"github.com/APTrust/exchange/util"
    "github.com/aws/aws-sdk-go/service/s3"
	"net/url"
	"strconv"
	"time"
)

// If Config.BucketReaderCacheHours isn't set to something
// sensible, cache Ingest WorkItems up to this many hours old.
const DEFAULT_CACHE_HOURS = 24

// How many S3 keys should we fetch in each batch when
// we're getting the contents of a bucket?
const MAX_KEYS = 1000

// APTBucketReader scans APTrust receiving buckets for items
// that need to be ingested. It creates a WorkItem record and
// an NSQ entry for each qualifying bag, and is responsible for
// knowing whether items in the receiving buckets actually need
// to be queued.
type APTBucketReader struct {
	Context                *context.Context
	Institutions           map[string]*models.Institution
	RecentIngestItems      map[string]*models.WorkItem
	stats                  *stats.APTBucketReaderStats
	statsEnabled           bool
}

// Creates a new bucket reader with the given context.
// Param enableStats is generally used for integration tests.
// Enabling stats in production can cause high memory usage,
// so keep that off unless you're trying to diagnose specific problems.
func NewAPTBucketReader(context *context.Context, enableStats bool) (*APTBucketReader){
	reader := &APTBucketReader{
		Context: context,
		Institutions: make(map[string]*models.Institution),
		RecentIngestItems: make(map[string]*models.WorkItem),
		statsEnabled: enableStats,
	}
	if enableStats {
		reader.stats = stats.NewAPTBucketReaderStats()
	}
	return reader
}

func (reader *APTBucketReader) Run() (error) {
	err := reader.cacheInstitutions()
	if err != nil { return err }
	err = reader.cacheRecentIngestItems()
	if err != nil { return err }
	reader.readAllBuckets()
	return nil
}

// Cache a list of all institutions. There are < 20.
// Exit on failure.
func (reader *APTBucketReader) cacheInstitutions() (error) {
	params := url.Values{}
	params.Add("page", "1")
	params.Add("per_page", "100")
	resp := reader.Context.PharosClient.InstitutionList(params)
	if resp.Error != nil {
		if reader.stats != nil {
			reader.stats.AddError(resp.Error.Error())
		}
		return resp.Error
	}
	if resp.Response.StatusCode != 200 {
		return reader.processPharosError(resp)
	}
	for _, inst := range resp.Institutions() {
		reader.Institutions[inst.Identifier] = inst
		if reader.stats != nil {
			reader.stats.AddToInstitutionsCached(inst)
		}
	}
	reader.Context.MessageLog.Info("Loaded %d institutions", len(reader.Institutions))
	return nil
}

// Cache a list of Ingest items that have been added to
// the list of WorkItems in the past X hours, so we won't
// have to do 1000 lookups. (X is usually around 24, but
// check Config.BucketReaderCacheHours for the exact value.
// This function exits on failure.
func (reader *APTBucketReader) cacheRecentIngestItems() (error) {
	hours := reader.Context.Config.BucketReaderCacheHours
	if hours < 1 {
		hours = DEFAULT_CACHE_HOURS
	}
	createdAfter := time.Now().Add(time.Duration(-1 * hours) * time.Hour).UTC()
	params := url.Values{}
	params.Add("page", "1")
	params.Add("per_page", "100")  // change to 4 to test hasMoreResults loop
	params.Add("item_action", constants.ActionIngest)
	params.Add("created_after", createdAfter.Format(time.RFC3339))
	hasMoreResults := true
	for hasMoreResults {
		resp := reader.Context.PharosClient.WorkItemList(params)
		if resp.Error != nil {
			if reader.stats != nil {
				reader.stats.AddError(resp.Error.Error())
			}
			return resp.Error
		}
		if resp.Response.StatusCode != 200 {
			return reader.processPharosError(resp)
		}
		reader.Context.MessageLog.Debug("%s", resp.Request.URL.String())
		for _, workItem := range resp.WorkItems() {
			hashKey := reader.makeHashKey(workItem.Name, workItem.ETag,
				workItem.BagDate.Format(time.RFC3339))
			reader.RecentIngestItems[hashKey] = workItem
			if reader.stats != nil {
				reader.stats.AddWorkItem("WorkItemsCached", workItem)
			}
		}
		if resp.Next != nil {
			pageNum, err := strconv.Atoi(params.Get("page"))
			if err != nil {
				msg := fmt.Sprintf("Page number '%s' doesn't look like a number!", params.Get("page"))
				if reader.stats != nil {
					reader.stats.AddError(msg)
				}
				return fmt.Errorf(msg)
			}
			params.Set("page", strconv.Itoa(pageNum + 1))
		} else {
			hasMoreResults = false
		}
	}
	reader.Context.MessageLog.Info("Loaded %d recent ingest WorkItems", len(reader.RecentIngestItems))
	return nil
}

func (reader *APTBucketReader) makeHashKey(key, etag, lastModified string) (string) {
	return fmt.Sprintf("%s|%s|%s", key, etag, lastModified)
}

func (reader *APTBucketReader) readAllBuckets() {
	for _, bucketName := range reader.Context.Config.ReceivingBuckets {
		reader.processBucket(bucketName)
	}
}

func (reader *APTBucketReader) processBucket(bucketName string) () {
	reader.Context.MessageLog.Debug("Checking bucket %s", bucketName)
	s3ObjList := network.NewS3ObjectList(reader.Context.Config.APTrustS3Region,
		bucketName, MAX_KEYS)
	keepFetching := true
	for keepFetching {
		s3ObjList.GetList()
		if s3ObjList.ErrorMessage != "" {
			if reader.stats != nil {
				reader.stats.AddError(s3ObjList.ErrorMessage)
			}
			reader.Context.MessageLog.Error(s3ObjList.ErrorMessage)
			break
		}
		for _, s3Object := range s3ObjList.Response.Contents {
			reader.stats.AddS3Item(fmt.Sprintf("%s/%s", bucketName, *s3Object.Key))
			reader.processS3Object(s3Object, bucketName)
		}
		keepFetching = *s3ObjList.Response.IsTruncated
	}
}

func (reader *APTBucketReader) processS3Object (s3Object *s3.Object, bucketName string) {
	if reader.Context.Config.MaxFileSize > int64(0) && *s3Object.Size > reader.Context.Config.MaxFileSize {
		msg := fmt.Sprintf("Skipping %s/%s because size %d is greater than " +
			"current max file size %d", bucketName, *s3Object.Key, *s3Object.Size,
			reader.Context.Config.MaxFileSize)
		if reader.stats != nil {
			reader.stats.AddWarning(msg)
		}
		reader.Context.MessageLog.Debug(msg)
		return
	}
	workItem, err := reader.findWorkItem(bucketName, *s3Object.Key, *s3Object.LastModified)
	if err != nil {
		// Don't create a work item, because one may already exist.
		// Error will be logged and added to stats at source.
		msg := fmt.Sprintf("Not creating WorkItem for %s/%s because " +
			"we can't tell if one already exists.", bucketName, *s3Object.Key)
		reader.Context.MessageLog.Warning(msg)
		if reader.stats != nil {
			reader.stats.AddWarning(msg)
		}
		return
	}
	if workItem == nil {
		workItem = reader.createWorkItem(bucketName, s3Object)
		if workItem == nil {
			// Error logged and statted at source.
			return
		}
	}
	if workItem.QueuedAt == nil || workItem.QueuedAt.IsZero() {
		reader.addToNSQ(workItem)
		reader.markAsQueued(workItem)
	}
}

func (reader *APTBucketReader) findWorkItem(key, etag string, lastModified time.Time) (*models.WorkItem, error) {
	hashKey := reader.makeHashKey(key, etag, lastModified.Format(time.RFC3339))
	if workItem, ok := reader.RecentIngestItems[hashKey]; ok {
		reader.Context.MessageLog.Debug("Found hash key '%s' in cache", hashKey)
		return workItem, nil
	}
	reader.Context.MessageLog.Debug("Looking up hash key '%s' in Pharos", hashKey)
	params := url.Values{}
	params.Add("page", "1")
	params.Add("per_page", "1")
	params.Add("item_action", constants.ActionIngest)
	params.Add("name", key)
	params.Add("etag", etag)
	params.Add("bag_date", lastModified.Format(time.RFC3339))
	resp := reader.Context.PharosClient.WorkItemList(params)
	if resp.Error != nil {
		errMsg := fmt.Sprintf("Error getting WorkItem for name '%s', etag '%s', time '%s': %v",
			params.Get("name"), params.Get("etag"), params.Get("bag_date"), resp.Error)
		reader.Context.MessageLog.Debug("%s", resp.Request.URL.String())
		reader.Context.MessageLog.Error(errMsg)
		if reader.stats != nil {
			reader.stats.AddError(errMsg)
		}
		return nil, fmt.Errorf(errMsg)
	}
	if resp.Response.StatusCode != 200 {
		err := reader.processPharosError(resp)
		return nil, err
	}
	workItem := resp.WorkItem()
	if workItem != nil {
		reader.Context.MessageLog.Debug("Found WorkItem for hash key '%s' in Pharos", hashKey)
		if reader.stats != nil {
			reader.stats.AddWorkItem("WorkItemsFetched", workItem)
		}
	} else {
		reader.Context.MessageLog.Debug("Did not find WorkItem for hash key '%s' in Pharos", hashKey)
	}
	return workItem, nil
}

func (reader *APTBucketReader) createWorkItem(bucket string, s3Object *s3.Object) (*models.WorkItem) {
	// Create a WorkItem in Pharos
	institution := reader.Institutions[util.OwnerOf(bucket)]
	if institution == nil {
		errMsg := fmt.Sprintf("Cannot find institution record for item %s/%s. " +
			"Owner computes to '%s'", bucket, *s3Object.Key, util.OwnerOf(bucket))
		reader.Context.MessageLog.Error(errMsg)
		if reader.stats != nil {
			reader.stats.AddError(errMsg)
		}
		return nil
	}
	workItem := &models.WorkItem{}
	workItem.Name = *s3Object.Key
	workItem.Bucket = bucket
	workItem.ETag = *s3Object.ETag
	workItem.Size = *s3Object.Size
	workItem.BagDate = *s3Object.LastModified
	workItem.InstitutionId = institution.Id
	workItem.User = constants.APTrustSystemUser
	workItem.Date = time.Now().UTC()
	workItem.Note = "Bag is in receiving bucket"
	workItem.Action = constants.ActionIngest
	workItem.Stage = constants.StageReceive
	workItem.Status = constants.StatusPending
	workItem.Outcome = "Item is pending ingest"
	workItem.Retry = true
	resp := reader.Context.PharosClient.WorkItemSave(workItem)

	if resp.Error != nil {
		errMsg := fmt.Sprintf("Error creating WorkItem for name '%s', etag '%s', time '%s': %v",
			workItem.Name, workItem.ETag, workItem.BagDate, resp.Error)
		reader.Context.MessageLog.Debug("%s", resp.Request.URL.String())
		reader.Context.MessageLog.Error(errMsg)
		if reader.stats != nil {
			reader.stats.AddError(errMsg)
		}
		return nil
	}
	if resp.Response.StatusCode != 201 {
		reader.processPharosError(resp)
		return nil
	}

	savedWorkItem := resp.WorkItem()
	reader.Context.MessageLog.Debug("Created WorkItem with id %d for %s/%s in Pharos",
		savedWorkItem.Id, bucket, *s3Object.Key)
	if reader.stats != nil {
		reader.stats.AddWorkItem("WorkItemsCreated", savedWorkItem)
	}
	return savedWorkItem
}

func (reader *APTBucketReader) addToNSQ(workItem *models.WorkItem) {
	client := network.NewNSQClient(reader.Context.Config.NsqdHttpAddress)
	err := client.Enqueue(reader.Context.Config.FetchWorker.NsqTopic, workItem.Id)
	if err != nil {
		msg := fmt.Sprintf("Error sending WorkItem %d to NSQ: %v", workItem.Id, err)
		if reader.stats != nil {
			reader.stats.AddError(msg)
		}
		reader.Context.MessageLog.Error(msg)
		return
	}
	reader.Context.MessageLog.Info("Added WorkItem id %d to NSQ (%s/%s)",
		workItem.Id, workItem.Bucket, workItem.Name)
	if reader.stats != nil {
		reader.stats.AddWorkItem("WorkItemsQueued", workItem)
	}
	return
}

func (reader *APTBucketReader) markAsQueued(workItem *models.WorkItem) (*models.WorkItem) {
	utcNow := time.Now().UTC()
	workItem.QueuedAt = &utcNow
	resp := reader.Context.PharosClient.WorkItemSave(workItem)
	if resp.Error != nil {
		errMsg := fmt.Sprintf("Error setting QueuedAt for WorkItem with id %d: %v",
			workItem.Id, resp.Error)
		reader.Context.MessageLog.Debug("%s", resp.Request.URL.String())
		reader.Context.MessageLog.Error(errMsg)
		if reader.stats != nil {
			reader.stats.AddError(errMsg)
		}
		return nil
	}
	if resp.Response.StatusCode != 200 {
		reader.processPharosError(resp)
		return nil
	}
	if reader.stats != nil {
		reader.stats.AddWorkItem("WorkItemsMarkedAsQueued", workItem)
	}
	return resp.WorkItem()
}

func (reader *APTBucketReader) processPharosError(resp *network.PharosResponse) (error) {
	respBody := ""
	bytesRead, readErr := resp.RawResponseData()
	if readErr == nil {
		respBody = string(bytesRead)
	} else {
		respBody = fmt.Sprintf("[Could not read response body: %v]", readErr)
	}
	msg := fmt.Sprintf("%s %s returned status code %d. Response body: %s",
		resp.Request.Method, resp.Request.URL, resp.Response.StatusCode, respBody)
	reader.Context.MessageLog.Error(msg)
	if reader.stats != nil {
		reader.stats.AddError(msg)
	}
	return fmt.Errorf(msg)
}

func (reader *APTBucketReader) GetStats() (*stats.APTBucketReaderStats) {
	return reader.stats
}
<commit_msg>Strip quotes from ETag<commit_after>package workers

import (
	"fmt"
	"github.com/APTrust/exchange/constants"
	"github.com/APTrust/exchange/context"
	"github.com/APTrust/exchange/models"
	"github.com/APTrust/exchange/network"
	"github.com/APTrust/exchange/stats"
	"github.com/APTrust/exchange/util"
    "github.com/aws/aws-sdk-go/service/s3"
	"net/url"
	"strconv"
	"strings"
	"time"
)

// If Config.BucketReaderCacheHours isn't set to something
// sensible, cache Ingest WorkItems up to this many hours old.
const DEFAULT_CACHE_HOURS = 24

// How many S3 keys should we fetch in each batch when
// we're getting the contents of a bucket?
const MAX_KEYS = 1000

// APTBucketReader scans APTrust receiving buckets for items
// that need to be ingested. It creates a WorkItem record and
// an NSQ entry for each qualifying bag, and is responsible for
// knowing whether items in the receiving buckets actually need
// to be queued.
type APTBucketReader struct {
	Context                *context.Context
	Institutions           map[string]*models.Institution
	RecentIngestItems      map[string]*models.WorkItem
	stats                  *stats.APTBucketReaderStats
	statsEnabled           bool
}

// Creates a new bucket reader with the given context.
// Param enableStats is generally used for integration tests.
// Enabling stats in production can cause high memory usage,
// so keep that off unless you're trying to diagnose specific problems.
func NewAPTBucketReader(context *context.Context, enableStats bool) (*APTBucketReader){
	reader := &APTBucketReader{
		Context: context,
		Institutions: make(map[string]*models.Institution),
		RecentIngestItems: make(map[string]*models.WorkItem),
		statsEnabled: enableStats,
	}
	if enableStats {
		reader.stats = stats.NewAPTBucketReaderStats()
	}
	return reader
}

func (reader *APTBucketReader) Run() (error) {
	err := reader.cacheInstitutions()
	if err != nil { return err }
	err = reader.cacheRecentIngestItems()
	if err != nil { return err }
	reader.readAllBuckets()
	return nil
}

// Cache a list of all institutions. There are < 20.
// Exit on failure.
func (reader *APTBucketReader) cacheInstitutions() (error) {
	params := url.Values{}
	params.Add("page", "1")
	params.Add("per_page", "100")
	resp := reader.Context.PharosClient.InstitutionList(params)
	if resp.Error != nil {
		if reader.stats != nil {
			reader.stats.AddError(resp.Error.Error())
		}
		return resp.Error
	}
	if resp.Response.StatusCode != 200 {
		return reader.processPharosError(resp)
	}
	for _, inst := range resp.Institutions() {
		reader.Institutions[inst.Identifier] = inst
		if reader.stats != nil {
			reader.stats.AddToInstitutionsCached(inst)
		}
	}
	reader.Context.MessageLog.Info("Loaded %d institutions", len(reader.Institutions))
	return nil
}

// Cache a list of Ingest items that have been added to
// the list of WorkItems in the past X hours, so we won't
// have to do 1000 lookups. (X is usually around 24, but
// check Config.BucketReaderCacheHours for the exact value.
// This function exits on failure.
func (reader *APTBucketReader) cacheRecentIngestItems() (error) {
	hours := reader.Context.Config.BucketReaderCacheHours
	if hours < 1 {
		hours = DEFAULT_CACHE_HOURS
	}
	createdAfter := time.Now().Add(time.Duration(-1 * hours) * time.Hour).UTC()
	params := url.Values{}
	params.Add("page", "1")
	params.Add("per_page", "100")  // change to 4 to test hasMoreResults loop
	params.Add("item_action", constants.ActionIngest)
	params.Add("created_after", createdAfter.Format(time.RFC3339))
	hasMoreResults := true
	for hasMoreResults {
		resp := reader.Context.PharosClient.WorkItemList(params)
		if resp.Error != nil {
			if reader.stats != nil {
				reader.stats.AddError(resp.Error.Error())
			}
			return resp.Error
		}
		if resp.Response.StatusCode != 200 {
			return reader.processPharosError(resp)
		}
		reader.Context.MessageLog.Debug("%s", resp.Request.URL.String())
		for _, workItem := range resp.WorkItems() {
			hashKey := reader.makeHashKey(workItem.Name, workItem.ETag,
				workItem.BagDate.Format(time.RFC3339))
			reader.RecentIngestItems[hashKey] = workItem
			if reader.stats != nil {
				reader.stats.AddWorkItem("WorkItemsCached", workItem)
			}
		}
		if resp.Next != nil {
			pageNum, err := strconv.Atoi(params.Get("page"))
			if err != nil {
				msg := fmt.Sprintf("Page number '%s' doesn't look like a number!", params.Get("page"))
				if reader.stats != nil {
					reader.stats.AddError(msg)
				}
				return fmt.Errorf(msg)
			}
			params.Set("page", strconv.Itoa(pageNum + 1))
		} else {
			hasMoreResults = false
		}
	}
	reader.Context.MessageLog.Info("Loaded %d recent ingest WorkItems", len(reader.RecentIngestItems))
	return nil
}

func (reader *APTBucketReader) makeHashKey(key, etag, lastModified string) (string) {
	return fmt.Sprintf("%s|%s|%s", key, etag, lastModified)
}

func (reader *APTBucketReader) readAllBuckets() {
	for _, bucketName := range reader.Context.Config.ReceivingBuckets {
		reader.processBucket(bucketName)
	}
}

func (reader *APTBucketReader) processBucket(bucketName string) () {
	reader.Context.MessageLog.Debug("Checking bucket %s", bucketName)
	s3ObjList := network.NewS3ObjectList(reader.Context.Config.APTrustS3Region,
		bucketName, MAX_KEYS)
	keepFetching := true
	for keepFetching {
		s3ObjList.GetList()
		if s3ObjList.ErrorMessage != "" {
			if reader.stats != nil {
				reader.stats.AddError(s3ObjList.ErrorMessage)
			}
			reader.Context.MessageLog.Error(s3ObjList.ErrorMessage)
			break
		}
		for _, s3Object := range s3ObjList.Response.Contents {
			reader.stats.AddS3Item(fmt.Sprintf("%s/%s", bucketName, *s3Object.Key))
			reader.processS3Object(s3Object, bucketName)
		}
		keepFetching = *s3ObjList.Response.IsTruncated
	}
}

func (reader *APTBucketReader) processS3Object (s3Object *s3.Object, bucketName string) {
	if reader.Context.Config.MaxFileSize > int64(0) && *s3Object.Size > reader.Context.Config.MaxFileSize {
		msg := fmt.Sprintf("Skipping %s/%s because size %d is greater than " +
			"current max file size %d", bucketName, *s3Object.Key, *s3Object.Size,
			reader.Context.Config.MaxFileSize)
		if reader.stats != nil {
			reader.stats.AddWarning(msg)
		}
		reader.Context.MessageLog.Debug(msg)
		return
	}
	workItem, err := reader.findWorkItem(bucketName, *s3Object.Key, *s3Object.LastModified)
	if err != nil {
		// Don't create a work item, because one may already exist.
		// Error will be logged and added to stats at source.
		msg := fmt.Sprintf("Not creating WorkItem for %s/%s because " +
			"we can't tell if one already exists.", bucketName, *s3Object.Key)
		reader.Context.MessageLog.Warning(msg)
		if reader.stats != nil {
			reader.stats.AddWarning(msg)
		}
		return
	}
	if workItem == nil {
		workItem = reader.createWorkItem(bucketName, s3Object)
		if workItem == nil {
			// Error logged and statted at source.
			return
		}
	}
	if workItem.QueuedAt == nil || workItem.QueuedAt.IsZero() {
		reader.addToNSQ(workItem)
		reader.markAsQueued(workItem)
	}
}

func (reader *APTBucketReader) findWorkItem(key, etag string, lastModified time.Time) (*models.WorkItem, error) {
	hashKey := reader.makeHashKey(key, etag, lastModified.Format(time.RFC3339))
	if workItem, ok := reader.RecentIngestItems[hashKey]; ok {
		reader.Context.MessageLog.Debug("Found hash key '%s' in cache", hashKey)
		return workItem, nil
	}
	reader.Context.MessageLog.Debug("Looking up hash key '%s' in Pharos", hashKey)
	params := url.Values{}
	params.Add("page", "1")
	params.Add("per_page", "1")
	params.Add("item_action", constants.ActionIngest)
	params.Add("name", key)
	params.Add("etag", etag)
	params.Add("bag_date", lastModified.Format(time.RFC3339))
	resp := reader.Context.PharosClient.WorkItemList(params)
	if resp.Error != nil {
		errMsg := fmt.Sprintf("Error getting WorkItem for name '%s', etag '%s', time '%s': %v",
			params.Get("name"), params.Get("etag"), params.Get("bag_date"), resp.Error)
		reader.Context.MessageLog.Debug("%s", resp.Request.URL.String())
		reader.Context.MessageLog.Error(errMsg)
		if reader.stats != nil {
			reader.stats.AddError(errMsg)
		}
		return nil, fmt.Errorf(errMsg)
	}
	if resp.Response.StatusCode != 200 {
		err := reader.processPharosError(resp)
		return nil, err
	}
	workItem := resp.WorkItem()
	if workItem != nil {
		reader.Context.MessageLog.Debug("Found WorkItem for hash key '%s' in Pharos", hashKey)
		if reader.stats != nil {
			reader.stats.AddWorkItem("WorkItemsFetched", workItem)
		}
	} else {
		reader.Context.MessageLog.Debug("Did not find WorkItem for hash key '%s' in Pharos", hashKey)
	}
	return workItem, nil
}

func (reader *APTBucketReader) createWorkItem(bucket string, s3Object *s3.Object) (*models.WorkItem) {
	// Create a WorkItem in Pharos
	institution := reader.Institutions[util.OwnerOf(bucket)]
	if institution == nil {
		errMsg := fmt.Sprintf("Cannot find institution record for item %s/%s. " +
			"Owner computes to '%s'", bucket, *s3Object.Key, util.OwnerOf(bucket))
		reader.Context.MessageLog.Error(errMsg)
		if reader.stats != nil {
			reader.stats.AddError(errMsg)
		}
		return nil
	}
	workItem := &models.WorkItem{}
	workItem.Name = *s3Object.Key
	workItem.Bucket = bucket
	workItem.ETag = strings.Replace(*s3Object.ETag, "\"", "", -1)
	workItem.Size = *s3Object.Size
	workItem.BagDate = *s3Object.LastModified
	workItem.InstitutionId = institution.Id
	workItem.User = constants.APTrustSystemUser
	workItem.Date = time.Now().UTC()
	workItem.Note = "Bag is in receiving bucket"
	workItem.Action = constants.ActionIngest
	workItem.Stage = constants.StageReceive
	workItem.Status = constants.StatusPending
	workItem.Outcome = "Item is pending ingest"
	workItem.Retry = true
	resp := reader.Context.PharosClient.WorkItemSave(workItem)

	if resp.Error != nil {
		errMsg := fmt.Sprintf("Error creating WorkItem for name '%s', etag '%s', time '%s': %v",
			workItem.Name, workItem.ETag, workItem.BagDate, resp.Error)
		reader.Context.MessageLog.Debug("%s", resp.Request.URL.String())
		reader.Context.MessageLog.Error(errMsg)
		if reader.stats != nil {
			reader.stats.AddError(errMsg)
		}
		return nil
	}
	if resp.Response.StatusCode != 201 {
		reader.processPharosError(resp)
		return nil
	}

	savedWorkItem := resp.WorkItem()
	reader.Context.MessageLog.Debug("Created WorkItem with id %d for %s/%s in Pharos",
		savedWorkItem.Id, bucket, *s3Object.Key)
	if reader.stats != nil {
		reader.stats.AddWorkItem("WorkItemsCreated", savedWorkItem)
	}
	return savedWorkItem
}

func (reader *APTBucketReader) addToNSQ(workItem *models.WorkItem) {
	client := network.NewNSQClient(reader.Context.Config.NsqdHttpAddress)
	err := client.Enqueue(reader.Context.Config.FetchWorker.NsqTopic, workItem.Id)
	if err != nil {
		msg := fmt.Sprintf("Error sending WorkItem %d to NSQ: %v", workItem.Id, err)
		if reader.stats != nil {
			reader.stats.AddError(msg)
		}
		reader.Context.MessageLog.Error(msg)
		return
	}
	reader.Context.MessageLog.Info("Added WorkItem id %d to NSQ (%s/%s)",
		workItem.Id, workItem.Bucket, workItem.Name)
	if reader.stats != nil {
		reader.stats.AddWorkItem("WorkItemsQueued", workItem)
	}
	return
}

func (reader *APTBucketReader) markAsQueued(workItem *models.WorkItem) (*models.WorkItem) {
	utcNow := time.Now().UTC()
	workItem.QueuedAt = &utcNow
	resp := reader.Context.PharosClient.WorkItemSave(workItem)
	if resp.Error != nil {
		errMsg := fmt.Sprintf("Error setting QueuedAt for WorkItem with id %d: %v",
			workItem.Id, resp.Error)
		reader.Context.MessageLog.Debug("%s", resp.Request.URL.String())
		reader.Context.MessageLog.Error(errMsg)
		if reader.stats != nil {
			reader.stats.AddError(errMsg)
		}
		return nil
	}
	if resp.Response.StatusCode != 200 {
		reader.processPharosError(resp)
		return nil
	}
	if reader.stats != nil {
		reader.stats.AddWorkItem("WorkItemsMarkedAsQueued", workItem)
	}
	return resp.WorkItem()
}

func (reader *APTBucketReader) processPharosError(resp *network.PharosResponse) (error) {
	respBody := ""
	bytesRead, readErr := resp.RawResponseData()
	if readErr == nil {
		respBody = string(bytesRead)
	} else {
		respBody = fmt.Sprintf("[Could not read response body: %v]", readErr)
	}
	msg := fmt.Sprintf("%s %s returned status code %d. Response body: %s",
		resp.Request.Method, resp.Request.URL, resp.Response.StatusCode, respBody)
	reader.Context.MessageLog.Error(msg)
	if reader.stats != nil {
		reader.stats.AddError(msg)
	}
	return fmt.Errorf(msg)
}

func (reader *APTBucketReader) GetStats() (*stats.APTBucketReaderStats) {
	return reader.stats
}
<|endoftext|>


<commit_before>// Package camoproxy provides an HTTP proxy server with content type
// restrictions as well as regex host allow and deny list support.
package camoproxy

import (
	"code.google.com/p/gorilla/mux"
	"crypto/hmac"
	"crypto/sha1"
	"encoding/hex"
	"errors"
	"fmt"
	"github.com/cactus/gologit"
	"io"
	"log"
	"net"
	"net/http"
	"net/url"
	"regexp"
	"strings"
	"sync"
	"time"
)

// Headers that are acceptible to pass from the client to the remote
// server. Only those present and true, are forwarded.
var validReqHeaders = map[string]bool{
	"Accept":            true,
	"Accept-Charset":    true,
	"Accept-Encoding":   true,
	"Accept-Language":   true,
	"Cache-Control":     true,
	"If-None-Match":     true,
	"If-Modified-Since": true,
}

// Headers that are acceptible to pass from the remote server to the
// client. Only those present and true, are forwarded. Empty implies
// no filtering.
var validRespHeaders = map[string]bool{}

type proxyStatus struct {
	sync.Mutex
	clientsServed uint64
	bytesServed   uint64
}

func (ps *proxyStatus) AddServed() {
	ps.Lock()
	defer ps.Unlock()
	ps.clientsServed += 1
}

func (ps *proxyStatus) AddBytes(bc int64) {
	ps.Lock()
	defer ps.Unlock()
	if bc <= 0 {
		return
	}
	ps.bytesServed += uint64(bc)
}

func (ps *proxyStatus) GetStats() (b uint64, c uint64) {
	return ps.clientsServed, ps.bytesServed
}

// A ProxyHandler is a Camo like HTTP proxy, that provides content type
// restrictions as well as regex host allow and deny list support
type ProxyHandler struct {
	Client    *http.Client
	HMacKey   []byte
	Allowlist []*regexp.Regexp
	Denylist  []*regexp.Regexp
	MaxSize   int64
	log       *gologit.DebugLogger
	stats     *proxyStatus
}

// StatsHandler returns an http.Handler that returns running totals and stats
// about the server.
func (p *ProxyHandler) StatsHandler() http.Handler {
	return http.HandlerFunc(
		func(w http.ResponseWriter, r *http.Request) {
			w.Header().Set("Content-Type", "text/plain; charset=utf-8")
			w.WriteHeader(200)
			c, b := p.stats.GetStats()
			fmt.Fprintf(w, "ClientsServed, BytesServed\n%d, %d\n", c, b)
		})
}

// ServerHTTP handles the client request, validates the request is validly
// HMAC signed, filters based on the Allow/Deny list, and then proxies
// valid requests to the desired endpoint. Responses are filtered for
// proper image content types.
func (p *ProxyHandler) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	p.log.Debugln("Request:", req.URL)
	p.stats.AddServed()

	vars := mux.Vars(req)
	surl, ok := p.decodeUrl(vars["sigHash"], vars["encodedUrl"])
	if !ok {
		http.Error(w, "Bad Signature", http.StatusForbidden)
		return
	}
	p.log.Debugln("URL:", surl)

	u, err := url.Parse(surl)
	if err != nil {
		log.Println(err)
		http.Error(w, "Bad url", http.StatusBadRequest)
		return
	}

	if u.Host == "" {
		http.Error(w, "Bad url", http.StatusNotFound)
		return
	}

	// if Allowlist is set, require match
	matchFound := true
	if len(p.Allowlist) > 0 {
		matchFound = false
		for _, rgx := range p.Allowlist {
			if rgx.MatchString(u.Host) {
				matchFound = true
			}
		}
	}
	if !matchFound {
		http.Error(w, "Allowlist host failure", http.StatusNotFound)
		return
	}

	// filter out Denylist urls based on regexes. Do this second
	// as Denylist takes precedence
	for _, rgx := range p.Denylist {
		if rgx.MatchString(u.Host) {
			http.Error(w, "Denylist host failure", http.StatusNotFound)
			return
		}
	}

	nreq, err := http.NewRequest("GET", surl, nil)
	if err != nil {
		p.log.Debugln("Could not create NewRequest", err)
		http.Error(w, "Error Fetching Resource", http.StatusBadGateway)
		return
	}

	// filter headers
	p.copyHeader(&nreq.Header, &req.Header, &validReqHeaders)
	nreq.Header.Add("connection", "close")
	nreq.Header.Add("user-agent", "pew pew pew")

	resp, err := p.Client.Do(nreq)
	if err != nil {
		p.log.Debugln("Could not connect to endpoint", err)
		if strings.Contains(err.Error(), "timeout") {
			http.Error(w, "Error Fetching Resource", http.StatusBadGateway)
		} else {
			http.Error(w, "Error Fetching Resource", http.StatusNotFound)
		}
		return
	}
	defer resp.Body.Close()

	// check for too large a response
	if resp.ContentLength > p.MaxSize {
		p.log.Debugln("Content length exceeded", surl)
		http.Error(w, "Content length exceeded", http.StatusNotFound)
		return
	}

	switch resp.StatusCode {
	case 200:
		// check content type
		ct, ok := resp.Header[http.CanonicalHeaderKey("content-type")]
		if !ok || ct[0][:6] != "image/" {
			log.Println("Non-Image content-type returned", u)
			http.Error(w, "Non-Image content-type returned",
				http.StatusBadRequest)
			return
		}
	case 300:
		log.Println("Multiple choices not supported")
		http.Error(w, "Multiple choices not supported", http.StatusNotFound)
		return
	case 301, 302, 303:
		// if we get a redirect here, we either disabled following,
		// or followed until max depth and still got one (redirect loop)
		http.Error(w, "Not Found", http.StatusNotFound)
		return
	case 304:
		h := w.Header()
		h.Set("X-Content-Type-Options", "nosniff")
		w.WriteHeader(304)
		return
	case 404:
		http.Error(w, "Not Found", http.StatusNotFound)
		return
	case 500, 502, 503, 504:
		// upstream errors should probably just 502. client can try later.
		http.Error(w, "Error Fetching Resource", http.StatusBadGateway)
		return
	default:
		http.Error(w, "Not Found", http.StatusNotFound)
		return
	}

	h := w.Header()
	p.copyHeader(&h, &resp.Header, &validRespHeaders)
	h.Set("X-Content-Type-Options", "nosniff")
	w.WriteHeader(resp.StatusCode)
	bW, err := io.Copy(w, resp.Body)
	if err != nil {
		p.log.Println("Error writing response:", err)
		return
	}
	p.stats.AddBytes(bW)
	p.log.Debugln(req, resp.StatusCode)
}

// validateURL ensures the url is properly verified via HMAC, and then
// unencodes the url, returning the url (if valid) and whether the
// HMAC was verified.
func (p *ProxyHandler) decodeUrl(hexdig string, hexurl string) (surl string, valid bool) {
	urlBytes, err := hex.DecodeString(hexurl)
	if err != nil {
		p.log.Println("Bad Hex Decode", hexurl)
		return
	}
	surl = string(urlBytes)
	mac := hmac.New(sha1.New, p.HMacKey)
	mac.Write([]byte(surl))
	macSum := hex.EncodeToString(mac.Sum([]byte{}))
	if macSum != hexdig {
		p.log.Debugf("Bad signature: %s != %s\n", macSum, hexdig)
		return
	}
	valid = true
	return
}

// copy headers from src into dst
// empty filter map will result in no filtering being done
func (p *ProxyHandler) copyHeader(dst, src *http.Header, filter *map[string]bool) {
	f := *filter
	filtering := false
	if len(f) > 0 {
		filtering = true
	}

	for k, vv := range *src {
		if x, ok := f[k]; filtering && !ok && !x {
			continue
		}
		for _, v := range vv {
			dst.Add(k, v)
		}
	}
}

type ProxyConfig struct {
	HmacKey         string
	AllowList       []string
	DenyList        []string
	MaxSize         int64
	FollowRedirects bool
	RequestTimeout  uint
}

func New(pc ProxyConfig, logger *gologit.DebugLogger) *ProxyHandler {
	tr := &http.Transport{
		Dial: func(netw, addr string) (net.Conn, error) {
			// 2 second timeout on requests
			timeout := time.Second * time.Duration(pc.RequestTimeout)
			c, err := net.DialTimeout(netw, addr, timeout)
			if err != nil {
				return nil, err
			}
			// also set time limit on reading
			c.SetDeadline(time.Now().Add(timeout))
			return c, nil
		}}

	// spawn an idle conn trimmer
	go func() {
		time.Sleep(5 * time.Minute)
		tr.CloseIdleConnections()
	}()

	// build/compile regex
	client := &http.Client{Transport: tr}
	if !pc.FollowRedirects {
		client.CheckRedirect = func(req *http.Request, via []*http.Request) error {
			return errors.New("Not following redirect")
		}
	}

	allow := make([]*regexp.Regexp, 0)
	deny := make([]*regexp.Regexp, 0)

	var c *regexp.Regexp
	var err error
	for _, v := range pc.DenyList {
		c, err = regexp.Compile(v)
		if err != nil {
			log.Fatal(err)
		}
		deny = append(deny, c)
	}
	for _, v := range pc.AllowList {
		c, err = regexp.Compile(v)
		if err != nil {
			log.Fatal(err)
		}
		allow = append(allow, c)
	}

	return &ProxyHandler{
		Client:    client,
		HMacKey:   []byte(pc.HmacKey),
		Allowlist: allow,
		Denylist:  deny,
		MaxSize:   pc.MaxSize,
		log:       logger,
		stats:     &proxyStatus{}}
}
<commit_msg>only calculate stats if enabled<commit_after>// Package camoproxy provides an HTTP proxy server with content type
// restrictions as well as regex host allow and deny list support.
package camoproxy

import (
	"code.google.com/p/gorilla/mux"
	"crypto/hmac"
	"crypto/sha1"
	"encoding/hex"
	"errors"
	"fmt"
	"github.com/cactus/gologit"
	"io"
	"log"
	"net"
	"net/http"
	"net/url"
	"regexp"
	"strings"
	"sync"
	"time"
)

// Headers that are acceptible to pass from the client to the remote
// server. Only those present and true, are forwarded.
var validReqHeaders = map[string]bool{
	"Accept":            true,
	"Accept-Charset":    true,
	"Accept-Encoding":   true,
	"Accept-Language":   true,
	"Cache-Control":     true,
	"If-None-Match":     true,
	"If-Modified-Since": true,
}

// Headers that are acceptible to pass from the remote server to the
// client. Only those present and true, are forwarded. Empty implies
// no filtering.
var validRespHeaders = map[string]bool{}

type proxyStatus struct {
	sync.Mutex
	clientsServed uint64
	bytesServed   uint64
	Enable        bool
}

func (ps *proxyStatus) AddServed() {
	ps.Lock()
	defer ps.Unlock()
	ps.clientsServed += 1
}

func (ps *proxyStatus) AddBytes(bc int64) {
	ps.Lock()
	defer ps.Unlock()
	if bc <= 0 {
		return
	}
	ps.bytesServed += uint64(bc)
}

func (ps *proxyStatus) GetStats() (b uint64, c uint64) {
	return ps.clientsServed, ps.bytesServed
}

// A ProxyHandler is a Camo like HTTP proxy, that provides content type
// restrictions as well as regex host allow and deny list support
type ProxyHandler struct {
	Client    *http.Client
	HMacKey   []byte
	Allowlist []*regexp.Regexp
	Denylist  []*regexp.Regexp
	MaxSize   int64
	log       *gologit.DebugLogger
	stats     *proxyStatus
}

// StatsHandler returns an http.Handler that returns running totals and stats
// about the server.
func (p *ProxyHandler) StatsHandler() http.Handler {
	p.stats.Enable = true
	return http.HandlerFunc(
		func(w http.ResponseWriter, r *http.Request) {
			w.Header().Set("Content-Type", "text/plain; charset=utf-8")
			w.WriteHeader(200)
			c, b := p.stats.GetStats()
			fmt.Fprintf(w, "ClientsServed, BytesServed\n%d, %d\n", c, b)
		})
}

// ServerHTTP handles the client request, validates the request is validly
// HMAC signed, filters based on the Allow/Deny list, and then proxies
// valid requests to the desired endpoint. Responses are filtered for
// proper image content types.
func (p *ProxyHandler) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	p.log.Debugln("Request:", req.URL)
	if p.stats.Enable {
		p.stats.AddServed()
	}

	vars := mux.Vars(req)
	surl, ok := p.decodeUrl(vars["sigHash"], vars["encodedUrl"])
	if !ok {
		http.Error(w, "Bad Signature", http.StatusForbidden)
		return
	}
	p.log.Debugln("URL:", surl)

	u, err := url.Parse(surl)
	if err != nil {
		log.Println(err)
		http.Error(w, "Bad url", http.StatusBadRequest)
		return
	}

	if u.Host == "" {
		http.Error(w, "Bad url", http.StatusNotFound)
		return
	}

	// if Allowlist is set, require match
	matchFound := true
	if len(p.Allowlist) > 0 {
		matchFound = false
		for _, rgx := range p.Allowlist {
			if rgx.MatchString(u.Host) {
				matchFound = true
			}
		}
	}
	if !matchFound {
		http.Error(w, "Allowlist host failure", http.StatusNotFound)
		return
	}

	// filter out Denylist urls based on regexes. Do this second
	// as Denylist takes precedence
	for _, rgx := range p.Denylist {
		if rgx.MatchString(u.Host) {
			http.Error(w, "Denylist host failure", http.StatusNotFound)
			return
		}
	}

	nreq, err := http.NewRequest("GET", surl, nil)
	if err != nil {
		p.log.Debugln("Could not create NewRequest", err)
		http.Error(w, "Error Fetching Resource", http.StatusBadGateway)
		return
	}

	// filter headers
	p.copyHeader(&nreq.Header, &req.Header, &validReqHeaders)
	nreq.Header.Add("connection", "close")
	nreq.Header.Add("user-agent", "pew pew pew")

	resp, err := p.Client.Do(nreq)
	if err != nil {
		p.log.Debugln("Could not connect to endpoint", err)
		if strings.Contains(err.Error(), "timeout") {
			http.Error(w, "Error Fetching Resource", http.StatusBadGateway)
		} else {
			http.Error(w, "Error Fetching Resource", http.StatusNotFound)
		}
		return
	}
	defer resp.Body.Close()

	// check for too large a response
	if resp.ContentLength > p.MaxSize {
		p.log.Debugln("Content length exceeded", surl)
		http.Error(w, "Content length exceeded", http.StatusNotFound)
		return
	}

	switch resp.StatusCode {
	case 200:
		// check content type
		ct, ok := resp.Header[http.CanonicalHeaderKey("content-type")]
		if !ok || ct[0][:6] != "image/" {
			log.Println("Non-Image content-type returned", u)
			http.Error(w, "Non-Image content-type returned",
				http.StatusBadRequest)
			return
		}
	case 300:
		log.Println("Multiple choices not supported")
		http.Error(w, "Multiple choices not supported", http.StatusNotFound)
		return
	case 301, 302, 303:
		// if we get a redirect here, we either disabled following,
		// or followed until max depth and still got one (redirect loop)
		http.Error(w, "Not Found", http.StatusNotFound)
		return
	case 304:
		h := w.Header()
		h.Set("X-Content-Type-Options", "nosniff")
		w.WriteHeader(304)
		return
	case 404:
		http.Error(w, "Not Found", http.StatusNotFound)
		return
	case 500, 502, 503, 504:
		// upstream errors should probably just 502. client can try later.
		http.Error(w, "Error Fetching Resource", http.StatusBadGateway)
		return
	default:
		http.Error(w, "Not Found", http.StatusNotFound)
		return
	}

	h := w.Header()
	p.copyHeader(&h, &resp.Header, &validRespHeaders)
	h.Set("X-Content-Type-Options", "nosniff")
	w.WriteHeader(resp.StatusCode)
	bW, err := io.Copy(w, resp.Body)
	if err != nil {
		p.log.Println("Error writing response:", err)
		return
	}
	if p.stats.Enable {
		p.stats.AddBytes(bW)
	}
	p.log.Debugln(req, resp.StatusCode)
}

// validateURL ensures the url is properly verified via HMAC, and then
// unencodes the url, returning the url (if valid) and whether the
// HMAC was verified.
func (p *ProxyHandler) decodeUrl(hexdig string, hexurl string) (surl string, valid bool) {
	urlBytes, err := hex.DecodeString(hexurl)
	if err != nil {
		p.log.Println("Bad Hex Decode", hexurl)
		return
	}
	surl = string(urlBytes)
	mac := hmac.New(sha1.New, p.HMacKey)
	mac.Write([]byte(surl))
	macSum := hex.EncodeToString(mac.Sum([]byte{}))
	if macSum != hexdig {
		p.log.Debugf("Bad signature: %s != %s\n", macSum, hexdig)
		return
	}
	valid = true
	return
}

// copy headers from src into dst
// empty filter map will result in no filtering being done
func (p *ProxyHandler) copyHeader(dst, src *http.Header, filter *map[string]bool) {
	f := *filter
	filtering := false
	if len(f) > 0 {
		filtering = true
	}

	for k, vv := range *src {
		if x, ok := f[k]; filtering && !ok && !x {
			continue
		}
		for _, v := range vv {
			dst.Add(k, v)
		}
	}
}

type ProxyConfig struct {
	HmacKey         string
	AllowList       []string
	DenyList        []string
	MaxSize         int64
	FollowRedirects bool
	RequestTimeout  uint
}

func New(pc ProxyConfig, logger *gologit.DebugLogger) *ProxyHandler {
	tr := &http.Transport{
		Dial: func(netw, addr string) (net.Conn, error) {
			// 2 second timeout on requests
			timeout := time.Second * time.Duration(pc.RequestTimeout)
			c, err := net.DialTimeout(netw, addr, timeout)
			if err != nil {
				return nil, err
			}
			// also set time limit on reading
			c.SetDeadline(time.Now().Add(timeout))
			return c, nil
		}}

	// spawn an idle conn trimmer
	go func() {
		time.Sleep(5 * time.Minute)
		tr.CloseIdleConnections()
	}()

	// build/compile regex
	client := &http.Client{Transport: tr}
	if !pc.FollowRedirects {
		client.CheckRedirect = func(req *http.Request, via []*http.Request) error {
			return errors.New("Not following redirect")
		}
	}

	allow := make([]*regexp.Regexp, 0)
	deny := make([]*regexp.Regexp, 0)

	var c *regexp.Regexp
	var err error
	for _, v := range pc.DenyList {
		c, err = regexp.Compile(v)
		if err != nil {
			log.Fatal(err)
		}
		deny = append(deny, c)
	}
	for _, v := range pc.AllowList {
		c, err = regexp.Compile(v)
		if err != nil {
			log.Fatal(err)
		}
		allow = append(allow, c)
	}

	return &ProxyHandler{
		Client:    client,
		HMacKey:   []byte(pc.HmacKey),
		Allowlist: allow,
		Denylist:  deny,
		MaxSize:   pc.MaxSize,
		log:       logger,
		stats:     &proxyStatus{}}
}
<|endoftext|>


<commit_before>// Package p9 exists so GoDoc will be generated :)
package p9
<commit_msg>GoDoc generated. Can be removed.<commit_after><|endoftext|>


<commit_before>package pb

import (
	"fmt"
	"io"
	"strings"
	"sync/atomic"
	"time"
)

var (
	// Default refresh rate - 200ms
	DefaultRefreshRate = time.Millisecond * 200

	BarStart = "["
	BarEnd   = "]"
	Empty    = "_"
	Current  = "="
	CurrentN = ">"
)

// Create new progress bar object
func New(total int) *ProgressBar {
	return &ProgressBar{
		Total:        int64(total),
		RefreshRate:  DefaultRefreshRate,
		ShowPercent:  true,
		ShowCounters: true,
		ShowBar:      true,
		ShowTimeLeft: true,
	}
}

// Create new object and start
func StartNew(total int) (pb *ProgressBar) {
	pb = New(total)
	pb.Start()
	return
}

// Callback for custom output
// For example:
// bar.Callback = func(s string) {
//     mySuperPrint(s)
// }
//
type Callback func(out string)

type ProgressBar struct {
	current int64 // current must be first member of struct (https://code.google.com/p/go/issues/detail?id=5278)

	Total                            int64
	RefreshRate                      time.Duration
	ShowPercent, ShowCounters        bool
	ShowSpeed, ShowTimeLeft, ShowBar bool
	Output                           io.Writer
	Callback                         Callback
	NotPrint                         bool
	Units                            int

	isFinish  bool
	startTime time.Time
}

// Start print
func (pb *ProgressBar) Start() {
	pb.startTime = time.Now()
	if pb.Total == 0 {
		pb.ShowBar = false
		pb.ShowTimeLeft = false
		pb.ShowPercent = false
	}
	go pb.writer()
}

// Increment current value
func (pb *ProgressBar) Increment() int {
	return pb.Add(1)
}

// Set current value
func (pb *ProgressBar) Set(current int) {
	atomic.StoreInt64(&pb.current, int64(current))
}

// Add to current value
func (pb *ProgressBar) Add(add int) int {
	return int(atomic.AddInt64(&pb.current, int64(add)))
}

// End print
func (pb *ProgressBar) Finish() {
	pb.isFinish = true
	pb.write(atomic.LoadInt64(&pb.current))
	if !pb.NotPrint {
		fmt.Println()
	}
}

// End print and write string 'str'
func (pb *ProgressBar) FinishPrint(str string) {
	pb.Finish()
	fmt.Println(str)
}

// implement io.Writer
func (pb *ProgressBar) Write(p []byte) (n int, err error) {
	n = len(p)
	pb.Add(n)
	return
}

// implement io.Reader
func (pb *ProgressBar) Read(p []byte) (n int, err error) {
	n = len(p)
	pb.Add(n)
	return
}

func (pb *ProgressBar) write(current int64) {
	width, _ := terminalWidth()
	var percentBox, countersBox, timeLeftBox, speedBox, barBox, end, out string

	// percents
	if pb.ShowPercent {
		percent := float64(current) / (float64(pb.Total) / float64(100))
		percentBox = fmt.Sprintf(" %#.02f %% ", percent)
	}

	// counters
	if pb.ShowCounters {
		if pb.Total > 0 {
			countersBox = fmt.Sprintf("%s / %s ", Format(current, pb.Units), Format(pb.Total, pb.Units))
		} else {
			countersBox = Format(current, pb.Units) + " "
		}
	}

	// time left
	if pb.ShowTimeLeft && current > 0 {
		fromStart := time.Now().Sub(pb.startTime)
		perEntry := fromStart / time.Duration(current)
		left := time.Duration(pb.Total-current) * perEntry
		left = (left / time.Second) * time.Second
		if left > 0 {
			timeLeftBox = left.String()
		}
	}

	// speed
	if pb.ShowSpeed && current > 0 {
		fromStart := time.Now().Sub(pb.startTime)
		speed := float64(current) / (float64(fromStart) / float64(time.Second))
		speedBox = Format(int64(speed), pb.Units) + "/s "
	}

	// bar
	if pb.ShowBar {
		size := width - len(countersBox+BarStart+BarEnd+percentBox+timeLeftBox+speedBox)
		if size > 0 {
			curCount := int(float64(current) / (float64(pb.Total) / float64(size)))
			emptCount := size - curCount
			barBox = BarStart
			if emptCount < 0 {
				emptCount = 0
			}
			if curCount > size {
				curCount = size
			}
			if emptCount <= 0 {
				barBox += strings.Repeat(Current, curCount)
			} else if curCount > 0 {
				barBox += strings.Repeat(Current, curCount-1) + CurrentN
			}

			barBox += strings.Repeat(Empty, emptCount) + BarEnd
		}
	}

	// check len
	out = countersBox + barBox + percentBox + speedBox + timeLeftBox
	if len(out) < width {
		end = strings.Repeat(" ", width-len(out))
	}

	out = countersBox + barBox + percentBox + speedBox + timeLeftBox

	// and print!
	switch {
	case pb.Output != nil:
		fmt.Fprint(pb.Output, out+end)
	case pb.Callback != nil:
		pb.Callback(out + end)
	case !pb.NotPrint:
		fmt.Print("\r" + out + end)
	}
}

func (pb *ProgressBar) writer() {
	var c, oc int64
	for {
		if pb.isFinish {
			break
		}
		c = atomic.LoadInt64(&pb.current)
		if c != oc {
			pb.write(c)
			oc = c
		}
		time.Sleep(pb.RefreshRate)
	}
}

type window struct {
	Row    uint16
	Col    uint16
	Xpixel uint16
	Ypixel uint16
}
<commit_msg>Fixing progress bar's rendering for 100% state.<commit_after>package pb

import (
	"fmt"
	"io"
	"math"
	"strings"
	"sync/atomic"
	"time"
)

var (
	// Default refresh rate - 200ms
	DefaultRefreshRate = time.Millisecond * 200

	BarStart = "["
	BarEnd   = "]"
	Empty    = "_"
	Current  = "="
	CurrentN = ">"
)

// Create new progress bar object
func New(total int) *ProgressBar {
	return &ProgressBar{
		Total:        int64(total),
		RefreshRate:  DefaultRefreshRate,
		ShowPercent:  true,
		ShowCounters: true,
		ShowBar:      true,
		ShowTimeLeft: true,
	}
}

// Create new object and start
func StartNew(total int) (pb *ProgressBar) {
	pb = New(total)
	pb.Start()
	return
}

// Callback for custom output
// For example:
// bar.Callback = func(s string) {
//     mySuperPrint(s)
// }
//
type Callback func(out string)

type ProgressBar struct {
	current int64 // current must be first member of struct (https://code.google.com/p/go/issues/detail?id=5278)

	Total                            int64
	RefreshRate                      time.Duration
	ShowPercent, ShowCounters        bool
	ShowSpeed, ShowTimeLeft, ShowBar bool
	Output                           io.Writer
	Callback                         Callback
	NotPrint                         bool
	Units                            int

	isFinish  bool
	startTime time.Time
}

// Start print
func (pb *ProgressBar) Start() {
	pb.startTime = time.Now()
	if pb.Total == 0 {
		pb.ShowBar = false
		pb.ShowTimeLeft = false
		pb.ShowPercent = false
	}
	go pb.writer()
}

// Increment current value
func (pb *ProgressBar) Increment() int {
	return pb.Add(1)
}

// Set current value
func (pb *ProgressBar) Set(current int) {
	atomic.StoreInt64(&pb.current, int64(current))
}

// Add to current value
func (pb *ProgressBar) Add(add int) int {
	return int(atomic.AddInt64(&pb.current, int64(add)))
}

// End print
func (pb *ProgressBar) Finish() {
	pb.isFinish = true
	pb.write(atomic.LoadInt64(&pb.current))
	if !pb.NotPrint {
		fmt.Println()
	}
}

// End print and write string 'str'
func (pb *ProgressBar) FinishPrint(str string) {
	pb.Finish()
	fmt.Println(str)
}

// implement io.Writer
func (pb *ProgressBar) Write(p []byte) (n int, err error) {
	n = len(p)
	pb.Add(n)
	return
}

// implement io.Reader
func (pb *ProgressBar) Read(p []byte) (n int, err error) {
	n = len(p)
	pb.Add(n)
	return
}

func (pb *ProgressBar) write(current int64) {
	width, _ := terminalWidth()
	var percentBox, countersBox, timeLeftBox, speedBox, barBox, end, out string

	// percents
	if pb.ShowPercent {
		percent := float64(current) / (float64(pb.Total) / float64(100))
		percentBox = fmt.Sprintf(" %#.02f %% ", percent)
	}

	// counters
	if pb.ShowCounters {
		if pb.Total > 0 {
			countersBox = fmt.Sprintf("%s / %s ", Format(current, pb.Units), Format(pb.Total, pb.Units))
		} else {
			countersBox = Format(current, pb.Units) + " "
		}
	}

	// time left
	if pb.ShowTimeLeft && current > 0 {
		fromStart := time.Now().Sub(pb.startTime)
		perEntry := fromStart / time.Duration(current)
		left := time.Duration(pb.Total-current) * perEntry
		left = (left / time.Second) * time.Second
		if left > 0 {
			timeLeftBox = left.String()
		}
	}

	// speed
	if pb.ShowSpeed && current > 0 {
		fromStart := time.Now().Sub(pb.startTime)
		speed := float64(current) / (float64(fromStart) / float64(time.Second))
		speedBox = Format(int64(speed), pb.Units) + "/s "
	}

	// bar
	if pb.ShowBar {
		size := width - len(countersBox+BarStart+BarEnd+percentBox+timeLeftBox+speedBox)
		if size > 0 {
			curCount := int(math.Ceil((float64(current) / float64(pb.Total)) * float64(size)))
			emptCount := size - curCount
			barBox = BarStart
			if emptCount < 0 {
				emptCount = 0
			}
			if curCount > size {
				curCount = size
			}
			if emptCount <= 0 {
				barBox += strings.Repeat(Current, curCount)
			} else if curCount > 0 {
				barBox += strings.Repeat(Current, curCount-1) + CurrentN
			}

			barBox += strings.Repeat(Empty, emptCount) + BarEnd
		}
	}

	// check len
	out = countersBox + barBox + percentBox + speedBox + timeLeftBox
	if len(out) < width {
		end = strings.Repeat(" ", width-len(out))
	}

	out = countersBox + barBox + percentBox + speedBox + timeLeftBox

	// and print!
	switch {
	case pb.Output != nil:
		fmt.Fprint(pb.Output, out+end)
	case pb.Callback != nil:
		pb.Callback(out + end)
	case !pb.NotPrint:
		fmt.Print("\r" + out + end)
	}
}

func (pb *ProgressBar) writer() {
	var c, oc int64
	for {
		if pb.isFinish {
			break
		}
		c = atomic.LoadInt64(&pb.current)
		if c != oc {
			pb.write(c)
			oc = c
		}
		time.Sleep(pb.RefreshRate)
	}
}

type window struct {
	Row    uint16
	Col    uint16
	Xpixel uint16
	Ypixel uint16
}
<|endoftext|>