<commit_before>package frame

import (
	"context"
	"encoding/base64"
	"errors"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"runtime"
	"strings"

	"github.com/zetamatta/nyagos/commands"
	"github.com/zetamatta/nyagos/nodos"
	"github.com/zetamatta/nyagos/shell"
	"github.com/zetamatta/nyagos/texts"
)

// OptionNorc is true, then rcfiles are not executed.
var OptionNorc = false

// OptionGoColorable is true,
// then escape sequences are interpreted by go-colorable library.
var OptionGoColorable = true

// OptionEnableVirtualTerminalProcessing is true,
// then Windows10's ENABLE_VIRTUAL_TERMINAL_PROCESSING is enabled.
var OptionEnableVirtualTerminalProcessing = false

type ScriptEngineForOption interface {
	SetArg([]string)
	RunFile(context.Context, string) ([]byte, error)
	RunString(context.Context, string) error
}

type optionArg struct {
	args []string
	sh   *shell.Shell
	e    ScriptEngineForOption
	ctx  context.Context // ctx is the Context object at parsing
}

type optionT struct {
	F func()
	V func(*optionArg) (func(context.Context) error, error)
	U string
}

var optionMap = map[string]optionT{
	"--lua-first": {
		U: "\"LUACODE\"\nExecute \"LUACODE\" before processing any rcfiles and continue shell",
		V: func(p *optionArg) (func(context.Context) error, error) {
			if len(p.args) <= 0 {
				return nil, errors.New("--lua-first: requires parameters")
			}
			if err := p.e.RunString(p.ctx, p.args[0]); err != nil {
				fmt.Fprintln(os.Stderr, err.Error())
			}
			return nil, nil
		},
	},
	"--cmd-first": {
		U: "\"COMMAND\"\nExecute \"COMMAND\" before processing any rcfiles and continue shell",
		V: func(p *optionArg) (func(context.Context) error, error) {
			if len(p.args) <= 0 {
				return nil, errors.New("--cmd-first: requires parameters")
			}
			p.sh.Interpret(p.ctx, p.args[0])
			return nil, nil
		},
	},
	"-k": {
		U: "\"COMMAND\"\nExecute \"COMMAND\" and continue the command-line.",
		V: func(p *optionArg) (func(context.Context) error, error) {
			if len(p.args) <= 0 {
				return nil, errors.New("-k: requires parameters")
			}
			return func(ctx context.Context) error {
				p.sh.Interpret(ctx, p.args[0])
				return nil
			}, nil
		},
	},
	"-c": {
		U: "\"COMMAND\"\nExecute `COMMAND` and quit.",
		V: func(p *optionArg) (func(context.Context) error, error) {
			if len(p.args) <= 0 {
				return nil, errors.New("-c: requires parameters")
			}
			return func(ctx context.Context) error {
				p.sh.Interpret(ctx, p.args[0])
				return io.EOF
			}, nil
		},
	},
	"-b": {
		U: "\"BASE64edCOMMAND\"\nDecode and execute the command which is encoded with Base64.",
		V: func(p *optionArg) (func(context.Context) error, error) {
			if len(p.args) <= 0 {
				return nil, errors.New("-b: requires parameters")
			}
			data, err := base64.StdEncoding.DecodeString(p.args[0])
			if err != nil {
				return nil, err
			}
			text := string(data)
			return func(ctx context.Context) error {
				p.sh.Interpret(ctx, text)
				return io.EOF
			}, nil
		},
	},
	"-f": {
		U: "FILE ARG1 ARG2 ...\n" +
			"If FILE's suffix is .lua, execute Lua-code on it.\n" +
			"The script can refer arguments as `arg[]`.\n" +
			"Otherwise, read and execute commands on it.",
		V: func(p *optionArg) (func(context.Context) error, error) {
			if len(p.args) <= 0 {
				return nil, errors.New("-f: requires parameters")
			}
			if strings.HasSuffix(strings.ToLower(p.args[0]), ".lua") {
				// lua script
				return func(ctx context.Context) error {
					p.e.SetArg(p.args)
					_, err := p.e.RunFile(ctx, p.args[0])
					if err != nil {
						return err
					}
					return io.EOF
				}, nil
			}
			return func(ctx context.Context) error {
				// command script
				if err := p.sh.Source(ctx, p.args[0]); err != nil {
					return err
				}
				return io.EOF
			}, nil
		},
	},
	"-e": {
		U: "\"SCRIPTCODE\"\nExecute SCRIPTCODE with Lua interpreter and quit.",
		V: func(p *optionArg) (func(context.Context) error, error) {
			if len(p.args) <= 0 {
				return nil, errors.New("-e: requires parameters")
			}
			return func(ctx context.Context) error {
				p.e.SetArg(p.args)
				err := p.e.RunString(ctx, p.args[0])
				if err != nil {
					return err
				}
				return io.EOF
			}, nil
		},
	},
	"--lua-file": {
		U: "FILE ARG1 ARG2...\n" +
			"Execute FILE as Lua Script even if FILE's suffix is not .lua .\n" +
			"The script can refer arguments as `arg[]`.\n" +
			"Lines starting with `@` are ignored to embed into batchfile.",
		V: func(p *optionArg) (func(context.Context) error, error) {
			if len(p.args) <= 0 {
				return nil, errors.New("--lua-file: requires parameters")
			}
			return func(ctx context.Context) error {
				p.e.SetArg(p.args)
				_, err := p.e.RunFile(ctx, p.args[0])
				if err != nil {
					return err
				}
				return io.EOF
			}, nil
		},
	},
	"--show-version-only": {
		U: "\nshow version only",
		V: func(p *optionArg) (func(context.Context) error, error) {
			OptionNorc = true
			return func(context.Context) error {
				fmt.Printf("%s-%s-%s\n", Version, runtime.GOOS, runtime.GOARCH)
				return io.EOF
			}, nil
		},
	},
	"--disable-virtual-terminal-processing": {
		U: "\nDo not use Windows10's native ESCAPE SEQUENCE.",
		F: func() {
			OptionEnableVirtualTerminalProcessing = false
		},
	},
	"--enable-virtual-terminal-processing": {
		U: "\nEnable Windows10's native ESCAPE SEQUENCE.\nIt should be used with `--no-go-colorable`.",
		F: func() {
			OptionEnableVirtualTerminalProcessing = true
		},
	},
	"--no-go-colorable": {
		U: "\nDo not use the ESCAPE SEQUENCE emulation with go-colorable library.",
		F: func() {
			OptionGoColorable = false
		},
	},
	"--go-colorable": {
		U: "\nUse the ESCAPE SEQUENCE emulation with go-colorable library.",
		F: func() {
			OptionGoColorable = true
		},
	},
	"--norc": {
		U: "\nDo not load the startup-scripts: `~\\.nyagos` , `~\\_nyagos`\nand `(BINDIR)\\nyagos.d\\*`.",
		F: func() {
			OptionNorc = true
		},
	},
	"--look-curdir-first": {
		U: "\nSearch for the executable from the current directory before %PATH%.\n(compatible with CMD.EXE)",
		F: func() {
			shell.LookCurdirOrder = nodos.LookCurdirFirst
		},
	},
	"--look-curdir-last": {
		U: "\nSearch for the executable from the current directory after %PATH%.\n(compatible with PowerShell)",
		F: func() {
			shell.LookCurdirOrder = nodos.LookCurdirLast
		},
	},
	"--look-curdir-never": {
		U: "\nNever search for the executable from the current directory\nunless %PATH% contains.\n(compatible with UNIX Shells)",
		F: func() {
			shell.LookCurdirOrder = nodos.LookCurdirNever
		},
	},
}

func Title() {
	fmt.Printf("Nihongo Yet Another GOing Shell %s-%s by %s\n",
		VersionOrStamp(),
		runtime.GOARCH,
		runtime.Version())
	fmt.Println("(c) 2014-2019 NYAOS.ORG <http://www.nyaos.org>")
}

func help(p *optionArg) (func(context.Context) error, error) {
	OptionNorc = true
	return func(context.Context) error {
		Title()
		fmt.Println()
		for _, key := range texts.SortedKeys(optionMap) {
			val := optionMap[key]
			fmt.Printf("%s %s\n", key, strings.Replace(val.U, "\n", "\n\t", -1))
		}

		fmt.Println("\nThese script are called on startup")
		if me, err := os.Executable(); err == nil {
			binDir := filepath.Dir(me)
			nyagosD := filepath.Join(binDir, "nyagos.d")
			fmt.Printf("  %s\\*.lua\n", nyagosD)
			file1 := filepath.Join(binDir, ".nyagos")
			fmt.Printf("  %s (Lua)\n", file1)
			file1 = filepath.Join(binDir, "_nyagos")
			fmt.Printf("  %s (Command-lines)\n", file1)
		}

		home := strings.TrimSpace(os.Getenv("HOME"))
		if home == "" {
			home = os.Getenv("USERPROFILE")
		}
		file1 := filepath.Join(home, ".nyagos")
		fmt.Printf("  %s (Lua)\n", file1)
		file1 = filepath.Join(home, "_nyagos")
		fmt.Printf("  %s (Command-lines)\n", file1)

		return io.EOF
	}, nil
}

func isDefault(value bool) string {
	if value {
		return " [default]"
	}
	return ""
}

func OptionParse(_ctx context.Context, sh *shell.Shell, e ScriptEngineForOption) (func(context.Context) error, error) {
	args := os.Args[1:]
	optionMap["-h"] = optionT{V: help, U: "\nPrint this usage"}
	optionMap["--help"] = optionT{V: help, U: "\nPrint this usage"}

	for key, val := range commands.BoolOptions {
		_key := strings.Replace(key, "_", "-", -1)
		_val := val
		optionMap["--"+_key] = optionT{
			F: func() {
				*_val.V = true
			},
			U: fmt.Sprintf("(lua: `nyagos.option.%s=true`)%s\n%s",
				key,
				isDefault(*val.V),
				_val.Usage),
		}
		optionMap["--no-"+_key] = optionT{
			F: func() {
				*_val.V = false
			},
			U: fmt.Sprintf("(lua: `nyagos.option.%s=false`)%s\n%s",
				key,
				isDefault(!*val.V),
				_val.NoUsage),
		}
	}

	for i := 0; i < len(args); i++ {
		if f, ok := optionMap[args[i]]; ok {
			if f.F != nil {
				f.F()
			}
			if f.V != nil {
				return f.V(&optionArg{
					args: args[i+1:],
					sh:   sh,
					e:    e,
					ctx:  _ctx,
				})
			}
		} else {
			fmt.Fprintf(os.Stderr, "%s: unknown parameter\n", args[i])
		}
	}
	return nil, nil
}

var SilentMode = false
<commit_msg>Show GOOS on startup title<commit_after>package frame

import (
	"context"
	"encoding/base64"
	"errors"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"runtime"
	"strings"

	"github.com/zetamatta/nyagos/commands"
	"github.com/zetamatta/nyagos/nodos"
	"github.com/zetamatta/nyagos/shell"
	"github.com/zetamatta/nyagos/texts"
)

// OptionNorc is true, then rcfiles are not executed.
var OptionNorc = false

// OptionGoColorable is true,
// then escape sequences are interpreted by go-colorable library.
var OptionGoColorable = true

// OptionEnableVirtualTerminalProcessing is true,
// then Windows10's ENABLE_VIRTUAL_TERMINAL_PROCESSING is enabled.
var OptionEnableVirtualTerminalProcessing = false

type ScriptEngineForOption interface {
	SetArg([]string)
	RunFile(context.Context, string) ([]byte, error)
	RunString(context.Context, string) error
}

type optionArg struct {
	args []string
	sh   *shell.Shell
	e    ScriptEngineForOption
	ctx  context.Context // ctx is the Context object at parsing
}

type optionT struct {
	F func()
	V func(*optionArg) (func(context.Context) error, error)
	U string
}

var optionMap = map[string]optionT{
	"--lua-first": {
		U: "\"LUACODE\"\nExecute \"LUACODE\" before processing any rcfiles and continue shell",
		V: func(p *optionArg) (func(context.Context) error, error) {
			if len(p.args) <= 0 {
				return nil, errors.New("--lua-first: requires parameters")
			}
			if err := p.e.RunString(p.ctx, p.args[0]); err != nil {
				fmt.Fprintln(os.Stderr, err.Error())
			}
			return nil, nil
		},
	},
	"--cmd-first": {
		U: "\"COMMAND\"\nExecute \"COMMAND\" before processing any rcfiles and continue shell",
		V: func(p *optionArg) (func(context.Context) error, error) {
			if len(p.args) <= 0 {
				return nil, errors.New("--cmd-first: requires parameters")
			}
			p.sh.Interpret(p.ctx, p.args[0])
			return nil, nil
		},
	},
	"-k": {
		U: "\"COMMAND\"\nExecute \"COMMAND\" and continue the command-line.",
		V: func(p *optionArg) (func(context.Context) error, error) {
			if len(p.args) <= 0 {
				return nil, errors.New("-k: requires parameters")
			}
			return func(ctx context.Context) error {
				p.sh.Interpret(ctx, p.args[0])
				return nil
			}, nil
		},
	},
	"-c": {
		U: "\"COMMAND\"\nExecute `COMMAND` and quit.",
		V: func(p *optionArg) (func(context.Context) error, error) {
			if len(p.args) <= 0 {
				return nil, errors.New("-c: requires parameters")
			}
			return func(ctx context.Context) error {
				p.sh.Interpret(ctx, p.args[0])
				return io.EOF
			}, nil
		},
	},
	"-b": {
		U: "\"BASE64edCOMMAND\"\nDecode and execute the command which is encoded with Base64.",
		V: func(p *optionArg) (func(context.Context) error, error) {
			if len(p.args) <= 0 {
				return nil, errors.New("-b: requires parameters")
			}
			data, err := base64.StdEncoding.DecodeString(p.args[0])
			if err != nil {
				return nil, err
			}
			text := string(data)
			return func(ctx context.Context) error {
				p.sh.Interpret(ctx, text)
				return io.EOF
			}, nil
		},
	},
	"-f": {
		U: "FILE ARG1 ARG2 ...\n" +
			"If FILE's suffix is .lua, execute Lua-code on it.\n" +
			"The script can refer arguments as `arg[]`.\n" +
			"Otherwise, read and execute commands on it.",
		V: func(p *optionArg) (func(context.Context) error, error) {
			if len(p.args) <= 0 {
				return nil, errors.New("-f: requires parameters")
			}
			if strings.HasSuffix(strings.ToLower(p.args[0]), ".lua") {
				// lua script
				return func(ctx context.Context) error {
					p.e.SetArg(p.args)
					_, err := p.e.RunFile(ctx, p.args[0])
					if err != nil {
						return err
					}
					return io.EOF
				}, nil
			}
			return func(ctx context.Context) error {
				// command script
				if err := p.sh.Source(ctx, p.args[0]); err != nil {
					return err
				}
				return io.EOF
			}, nil
		},
	},
	"-e": {
		U: "\"SCRIPTCODE\"\nExecute SCRIPTCODE with Lua interpreter and quit.",
		V: func(p *optionArg) (func(context.Context) error, error) {
			if len(p.args) <= 0 {
				return nil, errors.New("-e: requires parameters")
			}
			return func(ctx context.Context) error {
				p.e.SetArg(p.args)
				err := p.e.RunString(ctx, p.args[0])
				if err != nil {
					return err
				}
				return io.EOF
			}, nil
		},
	},
	"--lua-file": {
		U: "FILE ARG1 ARG2...\n" +
			"Execute FILE as Lua Script even if FILE's suffix is not .lua .\n" +
			"The script can refer arguments as `arg[]`.\n" +
			"Lines starting with `@` are ignored to embed into batchfile.",
		V: func(p *optionArg) (func(context.Context) error, error) {
			if len(p.args) <= 0 {
				return nil, errors.New("--lua-file: requires parameters")
			}
			return func(ctx context.Context) error {
				p.e.SetArg(p.args)
				_, err := p.e.RunFile(ctx, p.args[0])
				if err != nil {
					return err
				}
				return io.EOF
			}, nil
		},
	},
	"--show-version-only": {
		U: "\nshow version only",
		V: func(p *optionArg) (func(context.Context) error, error) {
			OptionNorc = true
			return func(context.Context) error {
				fmt.Printf("%s-%s-%s\n", Version, runtime.GOOS, runtime.GOARCH)
				return io.EOF
			}, nil
		},
	},
	"--disable-virtual-terminal-processing": {
		U: "\nDo not use Windows10's native ESCAPE SEQUENCE.",
		F: func() {
			OptionEnableVirtualTerminalProcessing = false
		},
	},
	"--enable-virtual-terminal-processing": {
		U: "\nEnable Windows10's native ESCAPE SEQUENCE.\nIt should be used with `--no-go-colorable`.",
		F: func() {
			OptionEnableVirtualTerminalProcessing = true
		},
	},
	"--no-go-colorable": {
		U: "\nDo not use the ESCAPE SEQUENCE emulation with go-colorable library.",
		F: func() {
			OptionGoColorable = false
		},
	},
	"--go-colorable": {
		U: "\nUse the ESCAPE SEQUENCE emulation with go-colorable library.",
		F: func() {
			OptionGoColorable = true
		},
	},
	"--norc": {
		U: "\nDo not load the startup-scripts: `~\\.nyagos` , `~\\_nyagos`\nand `(BINDIR)\\nyagos.d\\*`.",
		F: func() {
			OptionNorc = true
		},
	},
	"--look-curdir-first": {
		U: "\nSearch for the executable from the current directory before %PATH%.\n(compatible with CMD.EXE)",
		F: func() {
			shell.LookCurdirOrder = nodos.LookCurdirFirst
		},
	},
	"--look-curdir-last": {
		U: "\nSearch for the executable from the current directory after %PATH%.\n(compatible with PowerShell)",
		F: func() {
			shell.LookCurdirOrder = nodos.LookCurdirLast
		},
	},
	"--look-curdir-never": {
		U: "\nNever search for the executable from the current directory\nunless %PATH% contains.\n(compatible with UNIX Shells)",
		F: func() {
			shell.LookCurdirOrder = nodos.LookCurdirNever
		},
	},
}

func Title() {
	fmt.Printf("Nihongo Yet Another GOing Shell %s-%s-%s by %s\n",
		VersionOrStamp(),
		runtime.GOOS,
		runtime.GOARCH,
		runtime.Version())
	fmt.Println("(c) 2014-2019 NYAOS.ORG <http://www.nyaos.org>")
}

func help(p *optionArg) (func(context.Context) error, error) {
	OptionNorc = true
	return func(context.Context) error {
		Title()
		fmt.Println()
		for _, key := range texts.SortedKeys(optionMap) {
			val := optionMap[key]
			fmt.Printf("%s %s\n", key, strings.Replace(val.U, "\n", "\n\t", -1))
		}

		fmt.Println("\nThese script are called on startup")
		if me, err := os.Executable(); err == nil {
			binDir := filepath.Dir(me)
			nyagosD := filepath.Join(binDir, "nyagos.d")
			fmt.Printf("  %s\\*.lua\n", nyagosD)
			file1 := filepath.Join(binDir, ".nyagos")
			fmt.Printf("  %s (Lua)\n", file1)
			file1 = filepath.Join(binDir, "_nyagos")
			fmt.Printf("  %s (Command-lines)\n", file1)
		}

		home := strings.TrimSpace(os.Getenv("HOME"))
		if home == "" {
			home = os.Getenv("USERPROFILE")
		}
		file1 := filepath.Join(home, ".nyagos")
		fmt.Printf("  %s (Lua)\n", file1)
		file1 = filepath.Join(home, "_nyagos")
		fmt.Printf("  %s (Command-lines)\n", file1)

		return io.EOF
	}, nil
}

func isDefault(value bool) string {
	if value {
		return " [default]"
	}
	return ""
}

func OptionParse(_ctx context.Context, sh *shell.Shell, e ScriptEngineForOption) (func(context.Context) error, error) {
	args := os.Args[1:]
	optionMap["-h"] = optionT{V: help, U: "\nPrint this usage"}
	optionMap["--help"] = optionT{V: help, U: "\nPrint this usage"}

	for key, val := range commands.BoolOptions {
		_key := strings.Replace(key, "_", "-", -1)
		_val := val
		optionMap["--"+_key] = optionT{
			F: func() {
				*_val.V = true
			},
			U: fmt.Sprintf("(lua: `nyagos.option.%s=true`)%s\n%s",
				key,
				isDefault(*val.V),
				_val.Usage),
		}
		optionMap["--no-"+_key] = optionT{
			F: func() {
				*_val.V = false
			},
			U: fmt.Sprintf("(lua: `nyagos.option.%s=false`)%s\n%s",
				key,
				isDefault(!*val.V),
				_val.NoUsage),
		}
	}

	for i := 0; i < len(args); i++ {
		if f, ok := optionMap[args[i]]; ok {
			if f.F != nil {
				f.F()
			}
			if f.V != nil {
				return f.V(&optionArg{
					args: args[i+1:],
					sh:   sh,
					e:    e,
					ctx:  _ctx,
				})
			}
		} else {
			fmt.Fprintf(os.Stderr, "%s: unknown parameter\n", args[i])
		}
	}
	return nil, nil
}

var SilentMode = false
<|endoftext|>


<commit_before>package lg

import (
	"context"
	"fmt"
	"io"
	"net/http"
	"time"

	"go.uber.org/zap"
	"go.uber.org/zap/zapcore"

	uuid "github.com/satori/go.uuid"
)

var (
	// log the zap logger
	log *zap.Logger
)

const requestIDKey = 0

// Rfc3339NanoEncoder to encode time field to RFC3339Nano format.
func Rfc3339NanoEncoder(t time.Time, enc zapcore.PrimitiveArrayEncoder) {
	enc.AppendString(t.Format(time.RFC3339Nano))
}

// InitLogger must be first to be called.
func InitLogger(debug bool) {
	var cfg zap.Config

	if debug {
		cfg = zap.NewDevelopmentConfig()
	} else {
		cfg = zap.NewProductionConfig()
		cfg.EncoderConfig.LevelKey = "lvl"
		cfg.EncoderConfig.MessageKey = "msg"
		cfg.EncoderConfig.TimeKey = "timestamp"
		cfg.EncoderConfig.EncodeTime = Rfc3339NanoEncoder
	}

	var err error
	if log, err = cfg.Build(); err != nil {
		panic("Error create logger.")
	}
}

// L to get the request logger.
func L(ctx context.Context) *zap.Logger {
	if ctx == nil {
		return log
	}
	if l, ok := ctx.Value(requestIDKey).(*zap.Logger); ok {
		return l
	}
	panic("Should have request logger.")
}

// LogRequest to log every request.
func LogRequest(next http.Handler) http.Handler {
	fn := func(w http.ResponseWriter, r *http.Request) {
		l := log.With(zap.String("requestID", uuid.NewV1().String()))
		ctx := r.Context()
		ctx = context.WithValue(ctx, requestIDKey, l)
		r = r.WithContext(ctx)

		// Start timer
		start := time.Now()

		// wrap the ResponseWriter
		lw := &basicWriter{ResponseWriter: w}

		// Process request
		next.ServeHTTP(lw, r)
		lw.maybeWriteHeader()

		// Stop timer
		end := time.Now()
		latency := end.Sub(start)
		statusCode := lw.Status()

		l.Info("request",
			zap.String("method", r.Method),
			zap.String("url", r.RequestURI),
			zap.Int("code", statusCode),
			zap.String("clientIP", r.RemoteAddr),
			zap.Int("bytes", lw.bytes),
			zap.Int64("duration", int64(latency)/int64(time.Microsecond)),
		)
	}

	return http.HandlerFunc(fn)
}

// Recoverer the recover middware.
func Recoverer(next http.Handler) http.Handler {
	fn := func(w http.ResponseWriter, r *http.Request) {
		defer func() {
			if err := recover(); err != nil {
				// stack := stack(3)
				L(r.Context()).Error(fmt.Sprint(err), zap.Stack("stack"))
				http.Error(w, http.StatusText(http.StatusInternalServerError), http.StatusInternalServerError)
			}
		}()

		next.ServeHTTP(w, r)
	}

	return http.HandlerFunc(fn)
}

// writerProxy is a proxy around an http.ResponseWriter that allows you to hook
// into various parts of the response process.
type writerProxy interface {
	http.ResponseWriter
	// Status returns the HTTP status of the request, or 0 if one has not
	// yet been sent.
	Status() int
	// BytesWritten returns the total number of bytes sent to the client.
	BytesWritten() int
	// Tee causes the response body to be written to the given io.Writer in
	// addition to proxying the writes through. Only one io.Writer can be
	// tee'd to at once: setting a second one will overwrite the first.
	// Writes will be sent to the proxy before being written to this
	// io.Writer. It is illegal for the tee'd writer to be modified
	// concurrently with writes.
	Tee(io.Writer)
	// Unwrap returns the original proxied target.
	Unwrap() http.ResponseWriter
}

// basicWriter wraps a http.ResponseWriter that implements the minimal
// http.ResponseWriter interface.
type basicWriter struct {
	http.ResponseWriter
	wroteHeader bool
	code        int
	bytes       int
	tee         io.Writer
}

func (b *basicWriter) WriteHeader(code int) {
	if !b.wroteHeader {
		b.code = code
		b.wroteHeader = true
		b.ResponseWriter.WriteHeader(code)
	}
}

func (b *basicWriter) Write(buf []byte) (int, error) {
	b.WriteHeader(http.StatusOK)
	n, err := b.ResponseWriter.Write(buf)
	if b.tee != nil {
		_, err2 := b.tee.Write(buf[:n])
		// Prefer errors generated by the proxied writer.
		if err == nil {
			err = err2
		}
	}
	b.bytes += n
	return n, err
}

func (b *basicWriter) maybeWriteHeader() {
	if !b.wroteHeader {
		b.WriteHeader(http.StatusOK)
	}
}
func (b *basicWriter) Status() int {
	return b.code
}
func (b *basicWriter) BytesWritten() int {
	return b.bytes
}
func (b *basicWriter) Tee(w io.Writer) {
	b.tee = w
}
func (b *basicWriter) Unwrap() http.ResponseWriter {
	return b.ResponseWriter
}
<commit_msg>Add a key type and change recover log field.<commit_after>package lg

import (
	"context"
	"io"
	"net/http"
	"time"

	"go.uber.org/zap"
	"go.uber.org/zap/zapcore"

	uuid "github.com/satori/go.uuid"
)

var (
	// log the zap logger
	log *zap.Logger
)

type key int

const requestIDKey key = 0

// Rfc3339NanoEncoder to encode time field to RFC3339Nano format.
func Rfc3339NanoEncoder(t time.Time, enc zapcore.PrimitiveArrayEncoder) {
	enc.AppendString(t.Format(time.RFC3339Nano))
}

// InitLogger must be first to be called.
func InitLogger(debug bool) {
	var cfg zap.Config

	if debug {
		cfg = zap.NewDevelopmentConfig()
	} else {
		cfg = zap.NewProductionConfig()
		cfg.EncoderConfig.LevelKey = "lvl"
		cfg.EncoderConfig.MessageKey = "msg"
		cfg.EncoderConfig.TimeKey = "timestamp"
		cfg.EncoderConfig.EncodeTime = Rfc3339NanoEncoder
	}

	var err error
	if log, err = cfg.Build(); err != nil {
		panic("Error create logger.")
	}
}

// L to get the request logger.
func L(ctx context.Context) *zap.Logger {
	if ctx == nil {
		return log
	}
	if l, ok := ctx.Value(requestIDKey).(*zap.Logger); ok {
		return l
	}
	panic("Should have request logger.")
}

// LogRequest to log every request.
func LogRequest(next http.Handler) http.Handler {
	fn := func(w http.ResponseWriter, r *http.Request) {
		l := log.With(zap.String("requestID", uuid.NewV1().String()))
		ctx := r.Context()
		ctx = context.WithValue(ctx, requestIDKey, l)
		r = r.WithContext(ctx)

		// Start timer
		start := time.Now()

		// wrap the ResponseWriter
		lw := &basicWriter{ResponseWriter: w}

		// Process request
		next.ServeHTTP(lw, r)
		lw.maybeWriteHeader()

		// Stop timer
		end := time.Now()
		latency := end.Sub(start)
		statusCode := lw.Status()

		l.Info("request",
			zap.String("method", r.Method),
			zap.String("url", r.RequestURI),
			zap.Int("code", statusCode),
			zap.String("clientIP", r.RemoteAddr),
			zap.Int("bytes", lw.bytes),
			zap.Int64("duration", int64(latency)/int64(time.Microsecond)),
		)
	}

	return http.HandlerFunc(fn)
}

// Recoverer the recover middware.
func Recoverer(next http.Handler) http.Handler {
	fn := func(w http.ResponseWriter, r *http.Request) {
		defer func() {
			if err := recover(); err != nil {
				// stack := stack(3)
				L(r.Context()).Error("panic", zap.Error(err.(error)), zap.Stack("stack"))
				http.Error(w, http.StatusText(http.StatusInternalServerError), http.StatusInternalServerError)
			}
		}()

		next.ServeHTTP(w, r)
	}

	return http.HandlerFunc(fn)
}

// writerProxy is a proxy around an http.ResponseWriter that allows you to hook
// into various parts of the response process.
type writerProxy interface {
	http.ResponseWriter
	// Status returns the HTTP status of the request, or 0 if one has not
	// yet been sent.
	Status() int
	// BytesWritten returns the total number of bytes sent to the client.
	BytesWritten() int
	// Tee causes the response body to be written to the given io.Writer in
	// addition to proxying the writes through. Only one io.Writer can be
	// tee'd to at once: setting a second one will overwrite the first.
	// Writes will be sent to the proxy before being written to this
	// io.Writer. It is illegal for the tee'd writer to be modified
	// concurrently with writes.
	Tee(io.Writer)
	// Unwrap returns the original proxied target.
	Unwrap() http.ResponseWriter
}

// basicWriter wraps a http.ResponseWriter that implements the minimal
// http.ResponseWriter interface.
type basicWriter struct {
	http.ResponseWriter
	wroteHeader bool
	code        int
	bytes       int
	tee         io.Writer
}

func (b *basicWriter) WriteHeader(code int) {
	if !b.wroteHeader {
		b.code = code
		b.wroteHeader = true
		b.ResponseWriter.WriteHeader(code)
	}
}

func (b *basicWriter) Write(buf []byte) (int, error) {
	b.WriteHeader(http.StatusOK)
	n, err := b.ResponseWriter.Write(buf)
	if b.tee != nil {
		_, err2 := b.tee.Write(buf[:n])
		// Prefer errors generated by the proxied writer.
		if err == nil {
			err = err2
		}
	}
	b.bytes += n
	return n, err
}

func (b *basicWriter) maybeWriteHeader() {
	if !b.wroteHeader {
		b.WriteHeader(http.StatusOK)
	}
}
func (b *basicWriter) Status() int {
	return b.code
}
func (b *basicWriter) BytesWritten() int {
	return b.bytes
}
func (b *basicWriter) Tee(w io.Writer) {
	b.tee = w
}
func (b *basicWriter) Unwrap() http.ResponseWriter {
	return b.ResponseWriter
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"errors"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net"
	"strings"
	"sync"
	"sync/atomic"
	"time"
)

type backendStatus int

const (
	stateDown backendStatus = iota
	statePaused
	statePlaying
)

var (
	port         = flag.Int("port", 6600, "Port to run on")
	backends     = flag.String("backends", "", "Backends to support")
	backendsFile = flag.String("backends_file", "", "File with more backends to support")

	backendAddrs     []string
	backendPasswords []string
	backendStates    []backendStatus

	availableBackends int
	mtx               sync.Mutex
	unavailable       = sync.NewCond(&mtx)
	sock              net.Listener
	recentToggles     int32
	preferredBackend  int
	trafficStats      int64
)

func main() {
	flag.Parse()
	b := strings.Split(*backends, ",")
	if *backendsFile != "" {
		d, err := ioutil.ReadFile(*backendsFile)
		if err != nil {
			log.Fatal(err)
		}
		lines := strings.Split(string(d), "\n")
		b = append(b, lines...)
	}
	for _, addr := range b {
		if addr == "" {
			continue
		}
		ex := strings.Split(addr, "@")
		if len(ex) == 2 {
			backendAddrs = append(backendAddrs, ex[1])
			backendPasswords = append(backendPasswords, ex[0])
		} else {
			backendPasswords = append(backendPasswords, "")
			backendAddrs = append(backendAddrs, ex[0])
		}
		backendStates = append(backendStates, stateDown)
	}

	if len(backendAddrs) == 0 {
		log.Fatal("No backends specified")
	}

	for id := range backendAddrs {
		go handleBackend(id)
	}

	for {
		mtx.Lock()
		for availableBackends == 0 {
			unavailable.Wait()
		}

		var err error
		sock, err = net.Listen("tcp", fmt.Sprintf(":%d", *port))
		if err != nil {
			log.Fatalf("Listen failed: %v", err)
		}
		mtx.Unlock()

		for {
			conn, err := sock.Accept()
			if err != nil {
				if strings.Contains(err.Error(), "use of closed network connection") {
					break
				}
				log.Printf("Accept failed: %s\n", err)
			}
			log.Printf("[%s] connected", conn.RemoteAddr())
			go func(conn net.Conn) {
				defer conn.Close()
				err := handleConnection(conn)
				if err != nil {
					log.Printf("[%s] Closed: %v", conn.RemoteAddr(), err)
				}
			}(conn)
		}
		sock.Close()
	}
}

func setAvailability(id int, state backendStatus) {
	mtx.Lock()
	defer mtx.Unlock()
	fmt.Printf("State[%s #%d]: %d\n", backendAddrs[id], id, state)
	if state != backendStates[id] {
		if state == stateDown {
			availableBackends--
			if availableBackends == 0 {
				sock.Close()
			}
		} else {
			availableBackends++
			if availableBackends == 1 {
				unavailable.Broadcast()
			}
		}
	}
	backendStates[id] = state
}

func handleBackend(id int) {
	addr := backendAddrs[id]
	for {
		state, err := getBackendStatus(id)
		fmt.Printf("State[%s]: %d (%v)\n", addr, state, err)
		if err != nil {
			log.Printf("[%s] Failed: %s", addr, err)
		}
		setAvailability(id, state)
		time.Sleep(time.Minute)
	}
}

func getBackendStatus(id int) (backendStatus, error) {
	conn, err := net.Dial("tcp", backendAddrs[id])
	if err != nil {
		return stateDown, err
	}
	defer conn.Close()
	scanner := bufio.NewScanner(conn)
	if !scanner.Scan() {
		return stateDown, scanner.Err()
	}
	line := scanner.Text()
	if !strings.HasPrefix(line, "OK MPD ") {
		return stateDown, fmt.Errorf("Got a weird greeting: %s", line)
	}
	if backendPasswords[id] != "" {
		fmt.Fprintf(conn, "password %s\n", backendPasswords[id])
		if !scanner.Scan() {
			return stateDown, scanner.Err()
		}
	}
	io.WriteString(conn, "status\n")
	state := statePaused
	for scanner.Scan() {
		line := scanner.Text()
		if line == "state: play" {
			state = statePlaying
		}
		if line == "OK" {
			break
		}
		if strings.HasPrefix(line, "ACK") {
			return state, fmt.Errorf("unexpected error: %s", line)
		}
	}
	return state, scanner.Err()
}

func handleConnection(conn net.Conn) error {
	choice := preferredBackend
	var mpd net.Conn
	for {
		mtx.Lock()
	Loop:
		for id, state := range backendStates {
			switch state {
			case statePlaying:
				choice = id
				break Loop
			case statePaused:
				if backendStates[choice] == stateDown {
					choice = id
				}
			}
		}
		mtx.Unlock()
		if backendStates[choice] == stateDown {
			return errors.New("No backend available")
		}
		var err error
		mpd, err = net.Dial("tcp", backendAddrs[choice])
		if err != nil {
			setAvailability(choice, stateDown)
			return err
		}
		defer mpd.Close()
		break
	}
	done := make(chan error, 1)
	inStatus := false
	go func() {
		scanner := bufio.NewScanner(io.TeeReader(mpd, conn))
		for scanner.Scan() {
			line := scanner.Text()
			atomic.AddInt64(&trafficStats, int64(len(line)+1))
			// fmt.Println(line)
			if line == "OK" || strings.HasPrefix(line, "ACK") {
				inStatus = false
			}
			if inStatus && strings.HasPrefix(line, "state: ") {
				if line == "state: play" {
					setAvailability(choice, statePlaying)
				} else {
					setAvailability(choice, statePaused)
				}
			}
		}
		done <- scanner.Err()
	}()
	go func() {
		scanner := bufio.NewScanner(io.TeeReader(conn, mpd))
		for scanner.Scan() {
			line := scanner.Text()
			atomic.AddInt64(&trafficStats, int64(len(line)+1))
			fmt.Println(line)
			if line == "status" {
				inStatus = true
			}
			if line == "mpdswitcher_stats" {
				fmt.Fprintf(conn, "traffic_bytes: %d\n", atomic.LoadInt64(&trafficStats))
			}
			if strings.HasPrefix(line, "play") || strings.HasPrefix(line, "pause") {
				mtx.Lock()
				wantPlaying := backendStates[choice] != statePlaying
				mtx.Unlock()
				if strings.HasPrefix(line, "play") != strings.Contains(line, "1") {
					wantPlaying = false
					setAvailability(choice, statePaused)
				} else if strings.Contains(line, "0") {
					wantPlaying = true
					mtx.Lock()
					preferredBackend = choice
					mtx.Unlock()
					setAvailability(choice, statePlaying)
				}

				v := atomic.AddInt32(&recentToggles, 1)
				go func() {
					time.Sleep(time.Second)
					atomic.AddInt32(&recentToggles, -1)
				}()
				if v >= 4 && v <= 5 && !wantPlaying {
					mtx.Lock()
					preferredBackend = (choice + 1) % len(backendAddrs)
					fmt.Printf("New preferred backend: %d\n", preferredBackend)
					mtx.Unlock()
					break
				}
			}
		}
		done <- scanner.Err()
	}()
	return <-done
}
<commit_msg>Better distinguish play/pause commands<commit_after>package main

import (
	"bufio"
	"errors"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net"
	"strings"
	"sync"
	"sync/atomic"
	"time"
)

type backendStatus int

const (
	stateDown backendStatus = iota
	statePaused
	statePlaying
)

var (
	port         = flag.Int("port", 6600, "Port to run on")
	backends     = flag.String("backends", "", "Backends to support")
	backendsFile = flag.String("backends_file", "", "File with more backends to support")

	backendAddrs     []string
	backendPasswords []string
	backendStates    []backendStatus

	availableBackends int
	mtx               sync.Mutex
	unavailable       = sync.NewCond(&mtx)
	sock              net.Listener
	recentToggles     int32
	preferredBackend  int
	trafficStats      int64
)

func main() {
	flag.Parse()
	b := strings.Split(*backends, ",")
	if *backendsFile != "" {
		d, err := ioutil.ReadFile(*backendsFile)
		if err != nil {
			log.Fatal(err)
		}
		lines := strings.Split(string(d), "\n")
		b = append(b, lines...)
	}
	for _, addr := range b {
		if addr == "" {
			continue
		}
		ex := strings.Split(addr, "@")
		if len(ex) == 2 {
			backendAddrs = append(backendAddrs, ex[1])
			backendPasswords = append(backendPasswords, ex[0])
		} else {
			backendPasswords = append(backendPasswords, "")
			backendAddrs = append(backendAddrs, ex[0])
		}
		backendStates = append(backendStates, stateDown)
	}

	if len(backendAddrs) == 0 {
		log.Fatal("No backends specified")
	}

	for id := range backendAddrs {
		go handleBackend(id)
	}

	for {
		mtx.Lock()
		for availableBackends == 0 {
			unavailable.Wait()
		}

		var err error
		sock, err = net.Listen("tcp", fmt.Sprintf(":%d", *port))
		if err != nil {
			log.Fatalf("Listen failed: %v", err)
		}
		mtx.Unlock()

		for {
			conn, err := sock.Accept()
			if err != nil {
				if strings.Contains(err.Error(), "use of closed network connection") {
					break
				}
				log.Printf("Accept failed: %s\n", err)
			}
			log.Printf("[%s] connected", conn.RemoteAddr())
			go func(conn net.Conn) {
				defer conn.Close()
				err := handleConnection(conn)
				if err != nil {
					log.Printf("[%s] Closed: %v", conn.RemoteAddr(), err)
				}
			}(conn)
		}
		sock.Close()
	}
}

func setAvailability(id int, state backendStatus) {
	mtx.Lock()
	defer mtx.Unlock()
	fmt.Printf("State[%s #%d]: %d\n", backendAddrs[id], id, state)
	if state != backendStates[id] {
		if state == stateDown {
			availableBackends--
			if availableBackends == 0 {
				sock.Close()
			}
		} else {
			availableBackends++
			if availableBackends == 1 {
				unavailable.Broadcast()
			}
		}
	}
	backendStates[id] = state
}

func handleBackend(id int) {
	addr := backendAddrs[id]
	for {
		state, err := getBackendStatus(id)
		fmt.Printf("State[%s]: %d (%v)\n", addr, state, err)
		if err != nil {
			log.Printf("[%s] Failed: %s", addr, err)
		}
		setAvailability(id, state)
		time.Sleep(time.Minute)
	}
}

func getBackendStatus(id int) (backendStatus, error) {
	conn, err := net.Dial("tcp", backendAddrs[id])
	if err != nil {
		return stateDown, err
	}
	defer conn.Close()
	scanner := bufio.NewScanner(conn)
	if !scanner.Scan() {
		return stateDown, scanner.Err()
	}
	line := scanner.Text()
	if !strings.HasPrefix(line, "OK MPD ") {
		return stateDown, fmt.Errorf("Got a weird greeting: %s", line)
	}
	if backendPasswords[id] != "" {
		fmt.Fprintf(conn, "password %s\n", backendPasswords[id])
		if !scanner.Scan() {
			return stateDown, scanner.Err()
		}
	}
	io.WriteString(conn, "status\n")
	state := statePaused
	for scanner.Scan() {
		line := scanner.Text()
		if line == "state: play" {
			state = statePlaying
		}
		if line == "OK" {
			break
		}
		if strings.HasPrefix(line, "ACK") {
			return state, fmt.Errorf("unexpected error: %s", line)
		}
	}
	return state, scanner.Err()
}

func handleConnection(conn net.Conn) error {
	choice := preferredBackend
	var mpd net.Conn
	for {
		mtx.Lock()
	Loop:
		for id, state := range backendStates {
			switch state {
			case statePlaying:
				choice = id
				break Loop
			case statePaused:
				if backendStates[choice] == stateDown {
					choice = id
				}
			}
		}
		mtx.Unlock()
		if backendStates[choice] == stateDown {
			return errors.New("No backend available")
		}
		var err error
		mpd, err = net.Dial("tcp", backendAddrs[choice])
		if err != nil {
			setAvailability(choice, stateDown)
			return err
		}
		defer mpd.Close()
		break
	}
	done := make(chan error, 1)
	inStatus := false
	go func() {
		scanner := bufio.NewScanner(io.TeeReader(mpd, conn))
		for scanner.Scan() {
			line := scanner.Text()
			atomic.AddInt64(&trafficStats, int64(len(line)+1))
			// fmt.Println(line)
			if line == "OK" || strings.HasPrefix(line, "ACK") {
				inStatus = false
			}
			if inStatus && strings.HasPrefix(line, "state: ") {
				if line == "state: play" {
					setAvailability(choice, statePlaying)
				} else {
					setAvailability(choice, statePaused)
				}
			}
		}
		done <- scanner.Err()
	}()
	go func() {
		scanner := bufio.NewScanner(io.TeeReader(conn, mpd))
		for scanner.Scan() {
			line := scanner.Text()
			atomic.AddInt64(&trafficStats, int64(len(line)+1))
			fmt.Println(line)
			if line == "status" {
				inStatus = true
			}
			if line == "mpdswitcher_stats" {
				fmt.Fprintf(conn, "traffic_bytes: %d\n", atomic.LoadInt64(&trafficStats))
			}
			if line == "play" || strings.HasPrefix(line, "playid ") || strings.HasPrefix(line, "pause") {
				mtx.Lock()
				wantPlaying := backendStates[choice] != statePlaying
				mtx.Unlock()
				if (line == "play" || strings.HasPrefix(line, "playid")) != (strings.HasPrefix(line, "pause") && strings.Contains(line, "1")) {
					wantPlaying = false
					setAvailability(choice, statePaused)
				} else if strings.HasPrefix(line, "pause") && strings.Contains(line, "0") {
					wantPlaying = true
					mtx.Lock()
					preferredBackend = choice
					mtx.Unlock()
					setAvailability(choice, statePlaying)
				}

				v := atomic.AddInt32(&recentToggles, 1)
				go func() {
					time.Sleep(time.Second)
					atomic.AddInt32(&recentToggles, -1)
				}()
				if v >= 4 && v <= 5 && !wantPlaying {
					mtx.Lock()
					preferredBackend = (choice + 1) % len(backendAddrs)
					fmt.Printf("New preferred backend: %d\n", preferredBackend)
					mtx.Unlock()
					break
				}
			}
		}
		done <- scanner.Err()
	}()
	return <-done
}
<|endoftext|>


<commit_before>// Copyright (c) 2015 Mute Communications Ltd.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package msg

import (
	"bytes"
	"crypto/aes"
	"crypto/hmac"
	"crypto/rand"
	"crypto/sha512"
	"io"
	"io/ioutil"
	"math/big"

	"github.com/agl/ed25519"
	"github.com/mutecomm/mute/cipher"
	"github.com/mutecomm/mute/encode/base64"
	"github.com/mutecomm/mute/log"
	"github.com/mutecomm/mute/msg/padding"
	"github.com/mutecomm/mute/msg/session"
	"github.com/mutecomm/mute/uid"
)

func rootKeyAgreementSender(
	senderHeaderPub *[32]byte,
	senderIdentity, recipientIdentity string,
	senderSession, senderID, recipientKI, recipientID *uid.KeyEntry,
	previousRootKeyHash *[64]byte,
	numOfKeys uint64,
	keyStore session.Store,
) error {
	senderIdentityPub := senderID.PublicKey32()
	senderIdentityPriv := senderID.PrivateKey32()
	senderSessionPub := senderSession.PublicKey32()
	senderSessionPriv := senderSession.PrivateKey32()
	recipientIdentityPub := recipientID.PublicKey32()
	recipientKeyInitPub := recipientKI.PublicKey32()

	log.Debugf("senderIdentityPub:    %s", base64.Encode(senderIdentityPub[:]))
	log.Debugf("senderSessionPub:     %s", base64.Encode(senderSessionPub[:]))
	log.Debugf("recipientIdentityPub: %s", base64.Encode(recipientIdentityPub[:]))
	log.Debugf("recipientKeyInitPub:  %s", base64.Encode(recipientKeyInitPub[:]))

	// check keys to prevent reflection attacks and replays
	err := checkKeys(senderHeaderPub, senderIdentityPub, senderSessionPub,
		recipientIdentityPub, recipientKeyInitPub)
	if err != nil {
		return err
	}

	// compute t1
	t1, err := cipher.ECDH(senderIdentityPriv, recipientKeyInitPub, senderIdentityPub)
	if err != nil {
		return err
	}

	// compute t2
	t2, err := cipher.ECDH(senderSessionPriv, recipientKeyInitPub, senderSessionPub)
	if err != nil {
		return err
	}

	// compute t3
	t3, err := cipher.ECDH(senderSessionPriv, recipientIdentityPub, senderSessionPub)
	if err != nil {
		return err
	}

	// derive root key
	rootKey, err := deriveRootKey(t1, t2, t3, previousRootKeyHash)
	if err != nil {
		return err
	}

	// generate message keys
	err = generateMessageKeys(senderIdentity, recipientIdentity,
		senderID.HASH, recipientID.HASH, rootKey, false, senderSessionPub,
		recipientKeyInitPub, numOfKeys, keyStore)
	if err != nil {
		return err
	}
	return nil
}

// EncryptArgs contains all arguments for a message encryption.
type EncryptArgs struct {
	Writer                 io.Writer     // encrypted messagte is written here (base64 encoded)
	From                   *uid.Message  // sender UID
	To                     *uid.Message  // recipient UID
	SenderLastKeychainHash string        // last hash chain entry known to the sender
	PrivateSigKey          *[64]byte     // if this is s not nil the message is signed with the key
	Reader                 io.Reader     // data to encrypt is read here (only for StatusCode == StatusOK)
	NumOfKeys              uint64        // number of generated sessions keys (default: NumOfFutureKeys)
	AvgSessionSize         uint          // average session size (default: AverageSessionSize)
	Rand                   io.Reader     // random source
	KeyStore               session.Store // for managing session keys
	StatusCode             StatusCode    // status code of the encrypted message
}

// Encrypt encrypts a message with the argument given in args and returns the
// nymAddress the message should be delivered to.
func Encrypt(args *EncryptArgs) (nymAddress string, err error) {
	log.Debugf("msg.Encrypt(): %s -> %s", args.From.Identity(), args.To.Identity())

	// set defaults
	if args.NumOfKeys == 0 {
		args.NumOfKeys = NumOfFutureKeys
	}
	if args.AvgSessionSize == 0 {
		args.AvgSessionSize = AverageSessionSize
	}

	// create sender key
	senderHeaderKey, err := cipher.Curve25519Generate(cipher.RandReader)
	if err != nil {
		return "", log.Error(err)
	}

	// create pre-header
	ph := newPreHeader(senderHeaderKey.PublicKey()[:])

	// create base64 encoder
	var out bytes.Buffer
	wc := base64.NewEncoder(&out)

	// write pre-header
	var buf bytes.Buffer
	var count uint32
	if err := ph.write(&buf); err != nil {
		return "", err
	}
	oh := newOuterHeader(preHeaderPacket, count, buf.Bytes())
	if err := oh.write(wc, true); err != nil {
		return "", err
	}
	count++

	// get session state
	sender := args.From.Identity()
	recipient := args.To.Identity()
	sessionStateKey := session.CalcStateKey(args.From.PubKey().PublicKey32(),
		args.To.PubKey().PublicKey32())
	ss, err := args.KeyStore.GetSessionState(sessionStateKey)
	if err != nil {
		return "", err
	}
	if ss == nil {
		// no session found -> start first session
		log.Debug("no session found -> start first session")
		var recipientTemp *uid.KeyEntry
		recipientTemp, nymAddress, err = args.KeyStore.GetPublicKeyEntry(args.To)
		if err != nil {
			return "", err
		}
		// create session key
		var senderSession uid.KeyEntry
		if err := senderSession.InitDHKey(args.Rand); err != nil {
			return "", err
		}
		// store session key
		if err := addSessionKey(args.KeyStore, &senderSession); err != nil {
			return "", err
		}
		// root key agreement
		err = rootKeyAgreementSender(senderHeaderKey.PublicKey(),
			args.From.Identity(), args.To.Identity(), &senderSession,
			args.From.PubKey(), recipientTemp, args.To.PubKey(), nil,
			args.NumOfKeys, args.KeyStore)
		if err != nil {
			return "", err
		}
		// set session state
		ss = &session.State{
			SenderSessionCount:          0,
			SenderMessageCount:          0,
			MaxRecipientCount:           0,
			RecipientTemp:               *recipientTemp,
			SenderSessionPub:            senderSession,
			NextSenderSessionPub:        nil,
			NextRecipientSessionPubSeen: nil,
			NymAddress:                  nymAddress,
		}
		log.Debugf("set session: %s", ss.SenderSessionPub.HASH)
		err = args.KeyStore.SetSessionState(sessionStateKey, ss)
		if err != nil {
			return "", err
		}
	} else {
		log.Debug("session found")
		log.Debugf("got session: %s", ss.SenderSessionPub.HASH)
		nymAddress = ss.NymAddress
		// start new session in randomized fashion
		n, err := rand.Int(cipher.RandReader, big.NewInt(int64(args.AvgSessionSize)))
		if err != nil {
			return "", err
		}
		if n.Int64() == 0 {
			// create next session key
			var nextSenderSession uid.KeyEntry
			if err := nextSenderSession.InitDHKey(args.Rand); err != nil {
				return "", err
			}
			ss.NextSenderSessionPub = &nextSenderSession
			log.Debugf("set session: %s", ss.SenderSessionPub.HASH)
			err = args.KeyStore.SetSessionState(sessionStateKey, ss)
			if err != nil {
				return "", err
			}
		}
	}

	// create header
	log.Debugf("senderID:    %s", args.From.UIDContent.PUBKEYS[0].HASH)
	log.Debugf("recipientID: %s", args.To.UIDContent.PUBKEYS[0].HASH)
	log.Debugf("ss.SenderSessionCount: %d", ss.SenderSessionCount)
	log.Debugf("ss.SenderMessageCount: %d", ss.SenderMessageCount)
	log.Debugf("ss.RecipientTempHash:  %s", ss.RecipientTemp.HASH)
	h, err := newHeader(args.From, args.To, ss.RecipientTemp.HASH,
		&ss.SenderSessionPub, ss.NextSenderSessionPub,
		ss.NextRecipientSessionPubSeen, ss.SenderSessionCount,
		ss.SenderMessageCount, args.SenderLastKeychainHash, args.Rand,
		args.StatusCode)
	if err != nil {
		return "", err
	}
	log.Debugf("h.SenderSessionPub:             %s", h.SenderSessionPub.HASH)
	if h.NextSenderSessionPub != nil {
		log.Debugf("h.NextSenderSessionPub:         %s", h.NextSenderSessionPub.HASH)
	}
	if h.NextRecipientSessionPubSeen != nil {
		log.Debugf("h.NextRecipientSessionPubSeen:  %s",
			h.NextRecipientSessionPubSeen.HASH)
	}

	// create (encrypted) header packet
	recipientIdentityPub, err := args.To.PublicKey()
	if err != nil {
		return "", err
	}

	hp, err := newHeaderPacket(h, recipientIdentityPub,
		senderHeaderKey.PrivateKey(), args.Rand)
	if err != nil {
		return "", err
	}

	// write (encrypted) header packet
	buf.Reset()
	if err := hp.write(&buf); err != nil {
		return "", err
	}
	oh = newOuterHeader(encryptedHeader, count, buf.Bytes())
	if err := oh.write(wc, true); err != nil {
		return "", err
	}
	count++

	sessionKey := session.CalcKey(args.From.PubKey().HASH,
		args.To.PubKey().HASH, ss.SenderSessionPub.HASH, ss.RecipientTemp.HASH)

	// make sure we got enough message keys
	n, err := args.KeyStore.NumMessageKeys(sessionKey)
	if err != nil {
		return "", err
	}
	if ss.SenderMessageCount >= n {
		// generate more message keys
		log.Debugf("generate more message keys (ss.SenderMessageCount=%d)",
			ss.SenderMessageCount)
		chainKey, err := args.KeyStore.GetChainKey(sessionKey)
		if err != nil {
			return "", err
		}

		err = generateMessageKeys(sender, recipient, args.From.PubKey().HASH,
			args.To.PubKey().HASH, chainKey, false,
			ss.SenderSessionPub.PublicKey32(), ss.RecipientTemp.PublicKey32(),
			args.NumOfKeys, args.KeyStore)
		if err != nil {
			return "", err
		}
	}

	// get message key
	messageKey, err := args.KeyStore.GetMessageKey(sessionKey, true,
		ss.SenderMessageCount)
	if err != nil {
		return "", err
	}

	// derive symmetric keys
	cryptoKey, hmacKey, err := deriveSymmetricKeys(messageKey)
	if err != nil {
		return "", err
	}

	// write crypto setup packet
	iv := make([]byte, aes.BlockSize)
	if _, err := io.ReadFull(args.Rand, iv); err != nil {
		return "", log.Error(err)
	}
	oh = newOuterHeader(cryptoSetup, count, iv)

	if err := oh.write(wc, true); err != nil {
		return "", err
	}
	count++

	// start HMAC calculation
	mac := hmac.New(sha512.New, hmacKey)
	if err := oh.write(mac, true); err != nil {
		return "", err
	}

	// actual encryption
	content, err := ioutil.ReadAll(args.Reader)
	if err != nil {
		return "", log.Error(err)
	}
	// enforce maximum content length
	if len(content) > MaxContentLength {
		return "", log.Errorf("len(content) = %d > %d = MaxContentLength)",
			len(content), MaxContentLength)
	}

	// encrypted packet
	var contentHash []byte
	var innerType uint8
	if args.PrivateSigKey != nil {
		contentHash = cipher.SHA512(content)
		innerType = dataType | signType
	} else {
		innerType = dataType
	}
	ih := newInnerHeader(innerType, false, content)
	buf.Reset()
	if err := ih.write(&buf); err != nil {
		return "", err
	}
	stream := cipher.AES256CTRStream(cryptoKey, iv)
	stream.XORKeyStream(buf.Bytes(), buf.Bytes())
	oh = newOuterHeader(encryptedPacket, count, buf.Bytes())
	if err := oh.write(wc, true); err != nil {
		return "", err
	}
	count++

	// continue HMAC calculation
	if err := oh.write(mac, true); err != nil {
		return "", err
	}

	// signature header & padding
	buf.Reset()
	if args.PrivateSigKey != nil {
		sig := ed25519.Sign(args.PrivateSigKey, contentHash)
		// signature
		ih = newInnerHeader(signatureType, true, sig[:])
		if err := ih.write(&buf); err != nil {
			return "", err
		}
		// padding
		padLen := MaxContentLength - len(content)
		pad, err := padding.Generate(padLen, cipher.RandReader)
		if err != nil {
			return "", err
		}
		ih = newInnerHeader(paddingType, false, pad)
		if err := ih.write(&buf); err != nil {
			return "", err
		}
	} else {
		// just padding
		padLen := MaxContentLength + signatureSize - encryptedPacketSize +
			innerHeaderSize - len(content)
		pad, err := padding.Generate(padLen, cipher.RandReader)
		if err != nil {
			return "", err
		}
		ih = newInnerHeader(paddingType, false, pad)
		if err := ih.write(&buf); err != nil {
			return "", err
		}
	}
	// encrypt inner header
	stream.XORKeyStream(buf.Bytes(), buf.Bytes())
	oh = newOuterHeader(encryptedPacket, count, buf.Bytes())
	if err := oh.write(wc, true); err != nil {
		return "", err
	}
	count++

	// continue HMAC calculation
	if err := oh.write(mac, true); err != nil {
		return "", err
	}

	// create HMAC header
	oh = newOuterHeader(hmacPacket, count, nil)
	oh.PLen = sha512.Size
	if err := oh.write(mac, false); err != nil {
		return "", err
	}
	oh.inner = mac.Sum(oh.inner)
	log.Debugf("HMAC:       %s", base64.Encode(oh.inner))
	if err := oh.write(wc, true); err != nil {
		return "", err
	}
	count++

	// write output
	wc.Close()
	if out.Len() != EncodedMsgSize {
		return "", log.Errorf("out.Len() = %d != %d = EncodedMsgSize)",
			out.Len(), EncodedMsgSize)
	}
	if _, err := io.Copy(args.Writer, &out); err != nil {
		return "", log.Error(err)
	}

	// delete message key
	err = args.KeyStore.DelMessageKey(sessionKey, true, ss.SenderMessageCount)
	if err != nil {
		return "", err
	}
	// increase SenderMessageCount
	ss.SenderMessageCount++
	err = args.KeyStore.SetSessionState(sessionStateKey, ss)
	if err != nil {
		return "", err
	}

	return
}
<commit_msg>msg: store next session state<commit_after>// Copyright (c) 2015 Mute Communications Ltd.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package msg

import (
	"bytes"
	"crypto/aes"
	"crypto/hmac"
	"crypto/rand"
	"crypto/sha512"
	"io"
	"io/ioutil"
	"math/big"

	"github.com/agl/ed25519"
	"github.com/mutecomm/mute/cipher"
	"github.com/mutecomm/mute/encode/base64"
	"github.com/mutecomm/mute/log"
	"github.com/mutecomm/mute/msg/padding"
	"github.com/mutecomm/mute/msg/session"
	"github.com/mutecomm/mute/uid"
)

func rootKeyAgreementSender(
	senderHeaderPub *[32]byte,
	senderIdentity, recipientIdentity string,
	senderSession, senderID, recipientKI, recipientID *uid.KeyEntry,
	previousRootKeyHash *[64]byte,
	numOfKeys uint64,
	keyStore session.Store,
) error {
	senderIdentityPub := senderID.PublicKey32()
	senderIdentityPriv := senderID.PrivateKey32()
	senderSessionPub := senderSession.PublicKey32()
	senderSessionPriv := senderSession.PrivateKey32()
	recipientIdentityPub := recipientID.PublicKey32()
	recipientKeyInitPub := recipientKI.PublicKey32()

	log.Debugf("senderIdentityPub:    %s", base64.Encode(senderIdentityPub[:]))
	log.Debugf("senderSessionPub:     %s", base64.Encode(senderSessionPub[:]))
	log.Debugf("recipientIdentityPub: %s", base64.Encode(recipientIdentityPub[:]))
	log.Debugf("recipientKeyInitPub:  %s", base64.Encode(recipientKeyInitPub[:]))

	// check keys to prevent reflection attacks and replays
	err := checkKeys(senderHeaderPub, senderIdentityPub, senderSessionPub,
		recipientIdentityPub, recipientKeyInitPub)
	if err != nil {
		return err
	}

	// compute t1
	t1, err := cipher.ECDH(senderIdentityPriv, recipientKeyInitPub, senderIdentityPub)
	if err != nil {
		return err
	}

	// compute t2
	t2, err := cipher.ECDH(senderSessionPriv, recipientKeyInitPub, senderSessionPub)
	if err != nil {
		return err
	}

	// compute t3
	t3, err := cipher.ECDH(senderSessionPriv, recipientIdentityPub, senderSessionPub)
	if err != nil {
		return err
	}

	// derive root key
	rootKey, err := deriveRootKey(t1, t2, t3, previousRootKeyHash)
	if err != nil {
		return err
	}

	// generate message keys
	err = generateMessageKeys(senderIdentity, recipientIdentity,
		senderID.HASH, recipientID.HASH, rootKey, false, senderSessionPub,
		recipientKeyInitPub, numOfKeys, keyStore)
	if err != nil {
		return err
	}
	return nil
}

// EncryptArgs contains all arguments for a message encryption.
type EncryptArgs struct {
	Writer                 io.Writer     // encrypted messagte is written here (base64 encoded)
	From                   *uid.Message  // sender UID
	To                     *uid.Message  // recipient UID
	SenderLastKeychainHash string        // last hash chain entry known to the sender
	PrivateSigKey          *[64]byte     // if this is s not nil the message is signed with the key
	Reader                 io.Reader     // data to encrypt is read here (only for StatusCode == StatusOK)
	NumOfKeys              uint64        // number of generated sessions keys (default: NumOfFutureKeys)
	AvgSessionSize         uint          // average session size (default: AverageSessionSize)
	Rand                   io.Reader     // random source
	KeyStore               session.Store // for managing session keys
	StatusCode             StatusCode    // status code of the encrypted message
}

// Encrypt encrypts a message with the argument given in args and returns the
// nymAddress the message should be delivered to.
func Encrypt(args *EncryptArgs) (nymAddress string, err error) {
	log.Debugf("msg.Encrypt(): %s -> %s", args.From.Identity(), args.To.Identity())

	// set defaults
	if args.NumOfKeys == 0 {
		args.NumOfKeys = NumOfFutureKeys
	}
	if args.AvgSessionSize == 0 {
		args.AvgSessionSize = AverageSessionSize
	}

	// create sender key
	senderHeaderKey, err := cipher.Curve25519Generate(cipher.RandReader)
	if err != nil {
		return "", log.Error(err)
	}

	// create pre-header
	ph := newPreHeader(senderHeaderKey.PublicKey()[:])

	// create base64 encoder
	var out bytes.Buffer
	wc := base64.NewEncoder(&out)

	// write pre-header
	var buf bytes.Buffer
	var count uint32
	if err := ph.write(&buf); err != nil {
		return "", err
	}
	oh := newOuterHeader(preHeaderPacket, count, buf.Bytes())
	if err := oh.write(wc, true); err != nil {
		return "", err
	}
	count++

	// get session state
	sender := args.From.Identity()
	recipient := args.To.Identity()
	sessionStateKey := session.CalcStateKey(args.From.PubKey().PublicKey32(),
		args.To.PubKey().PublicKey32())
	ss, err := args.KeyStore.GetSessionState(sessionStateKey)
	if err != nil {
		return "", err
	}
	if ss == nil {
		// no session found -> start first session
		log.Debug("no session found -> start first session")
		var recipientTemp *uid.KeyEntry
		recipientTemp, nymAddress, err = args.KeyStore.GetPublicKeyEntry(args.To)
		if err != nil {
			return "", err
		}
		// create session key
		var senderSession uid.KeyEntry
		if err := senderSession.InitDHKey(args.Rand); err != nil {
			return "", err
		}
		// store session key
		if err := addSessionKey(args.KeyStore, &senderSession); err != nil {
			return "", err
		}
		// root key agreement
		err = rootKeyAgreementSender(senderHeaderKey.PublicKey(),
			args.From.Identity(), args.To.Identity(), &senderSession,
			args.From.PubKey(), recipientTemp, args.To.PubKey(), nil,
			args.NumOfKeys, args.KeyStore)
		if err != nil {
			return "", err
		}
		// set session state
		ss = &session.State{
			SenderSessionCount:          0,
			SenderMessageCount:          0,
			MaxRecipientCount:           0,
			RecipientTemp:               *recipientTemp,
			SenderSessionPub:            senderSession,
			NextSenderSessionPub:        nil,
			NextRecipientSessionPubSeen: nil,
			NymAddress:                  nymAddress,
		}
		log.Debugf("set session: %s", ss.SenderSessionPub.HASH)
		err = args.KeyStore.SetSessionState(sessionStateKey, ss)
		if err != nil {
			return "", err
		}
	} else {
		log.Debug("session found")
		log.Debugf("got session: %s", ss.SenderSessionPub.HASH)
		nymAddress = ss.NymAddress
		// start new session in randomized fashion
		n, err := rand.Int(cipher.RandReader, big.NewInt(int64(args.AvgSessionSize)))
		if err != nil {
			return "", err
		}
		if n.Int64() == 0 {
			// create next session key
			var nextSenderSession uid.KeyEntry
			if err := nextSenderSession.InitDHKey(args.Rand); err != nil {
				return "", err
			}
			// store next session key
			if err := addSessionKey(args.KeyStore, &nextSenderSession); err != nil {
				return "", err
			}
			// update session state
			ss.NextSenderSessionPub = &nextSenderSession
			log.Debugf("set session: %s", ss.SenderSessionPub.HASH)
			err = args.KeyStore.SetSessionState(sessionStateKey, ss)
			if err != nil {
				return "", err
			}
		}
	}

	// create header
	log.Debugf("senderID:    %s", args.From.UIDContent.PUBKEYS[0].HASH)
	log.Debugf("recipientID: %s", args.To.UIDContent.PUBKEYS[0].HASH)
	log.Debugf("ss.SenderSessionCount: %d", ss.SenderSessionCount)
	log.Debugf("ss.SenderMessageCount: %d", ss.SenderMessageCount)
	log.Debugf("ss.RecipientTempHash:  %s", ss.RecipientTemp.HASH)
	h, err := newHeader(args.From, args.To, ss.RecipientTemp.HASH,
		&ss.SenderSessionPub, ss.NextSenderSessionPub,
		ss.NextRecipientSessionPubSeen, ss.SenderSessionCount,
		ss.SenderMessageCount, args.SenderLastKeychainHash, args.Rand,
		args.StatusCode)
	if err != nil {
		return "", err
	}
	log.Debugf("h.SenderSessionPub:             %s", h.SenderSessionPub.HASH)
	if h.NextSenderSessionPub != nil {
		log.Debugf("h.NextSenderSessionPub:         %s", h.NextSenderSessionPub.HASH)
	}
	if h.NextRecipientSessionPubSeen != nil {
		log.Debugf("h.NextRecipientSessionPubSeen:  %s",
			h.NextRecipientSessionPubSeen.HASH)
	}

	// create (encrypted) header packet
	recipientIdentityPub, err := args.To.PublicKey()
	if err != nil {
		return "", err
	}

	hp, err := newHeaderPacket(h, recipientIdentityPub,
		senderHeaderKey.PrivateKey(), args.Rand)
	if err != nil {
		return "", err
	}

	// write (encrypted) header packet
	buf.Reset()
	if err := hp.write(&buf); err != nil {
		return "", err
	}
	oh = newOuterHeader(encryptedHeader, count, buf.Bytes())
	if err := oh.write(wc, true); err != nil {
		return "", err
	}
	count++

	sessionKey := session.CalcKey(args.From.PubKey().HASH,
		args.To.PubKey().HASH, ss.SenderSessionPub.HASH, ss.RecipientTemp.HASH)

	// make sure we got enough message keys
	n, err := args.KeyStore.NumMessageKeys(sessionKey)
	if err != nil {
		return "", err
	}
	if ss.SenderMessageCount >= n {
		// generate more message keys
		log.Debugf("generate more message keys (ss.SenderMessageCount=%d)",
			ss.SenderMessageCount)
		chainKey, err := args.KeyStore.GetChainKey(sessionKey)
		if err != nil {
			return "", err
		}

		err = generateMessageKeys(sender, recipient, args.From.PubKey().HASH,
			args.To.PubKey().HASH, chainKey, false,
			ss.SenderSessionPub.PublicKey32(), ss.RecipientTemp.PublicKey32(),
			args.NumOfKeys, args.KeyStore)
		if err != nil {
			return "", err
		}
	}

	// get message key
	messageKey, err := args.KeyStore.GetMessageKey(sessionKey, true,
		ss.SenderMessageCount)
	if err != nil {
		return "", err
	}

	// derive symmetric keys
	cryptoKey, hmacKey, err := deriveSymmetricKeys(messageKey)
	if err != nil {
		return "", err
	}

	// write crypto setup packet
	iv := make([]byte, aes.BlockSize)
	if _, err := io.ReadFull(args.Rand, iv); err != nil {
		return "", log.Error(err)
	}
	oh = newOuterHeader(cryptoSetup, count, iv)

	if err := oh.write(wc, true); err != nil {
		return "", err
	}
	count++

	// start HMAC calculation
	mac := hmac.New(sha512.New, hmacKey)
	if err := oh.write(mac, true); err != nil {
		return "", err
	}

	// actual encryption
	content, err := ioutil.ReadAll(args.Reader)
	if err != nil {
		return "", log.Error(err)
	}
	// enforce maximum content length
	if len(content) > MaxContentLength {
		return "", log.Errorf("len(content) = %d > %d = MaxContentLength)",
			len(content), MaxContentLength)
	}

	// encrypted packet
	var contentHash []byte
	var innerType uint8
	if args.PrivateSigKey != nil {
		contentHash = cipher.SHA512(content)
		innerType = dataType | signType
	} else {
		innerType = dataType
	}
	ih := newInnerHeader(innerType, false, content)
	buf.Reset()
	if err := ih.write(&buf); err != nil {
		return "", err
	}
	stream := cipher.AES256CTRStream(cryptoKey, iv)
	stream.XORKeyStream(buf.Bytes(), buf.Bytes())
	oh = newOuterHeader(encryptedPacket, count, buf.Bytes())
	if err := oh.write(wc, true); err != nil {
		return "", err
	}
	count++

	// continue HMAC calculation
	if err := oh.write(mac, true); err != nil {
		return "", err
	}

	// signature header & padding
	buf.Reset()
	if args.PrivateSigKey != nil {
		sig := ed25519.Sign(args.PrivateSigKey, contentHash)
		// signature
		ih = newInnerHeader(signatureType, true, sig[:])
		if err := ih.write(&buf); err != nil {
			return "", err
		}
		// padding
		padLen := MaxContentLength - len(content)
		pad, err := padding.Generate(padLen, cipher.RandReader)
		if err != nil {
			return "", err
		}
		ih = newInnerHeader(paddingType, false, pad)
		if err := ih.write(&buf); err != nil {
			return "", err
		}
	} else {
		// just padding
		padLen := MaxContentLength + signatureSize - encryptedPacketSize +
			innerHeaderSize - len(content)
		pad, err := padding.Generate(padLen, cipher.RandReader)
		if err != nil {
			return "", err
		}
		ih = newInnerHeader(paddingType, false, pad)
		if err := ih.write(&buf); err != nil {
			return "", err
		}
	}
	// encrypt inner header
	stream.XORKeyStream(buf.Bytes(), buf.Bytes())
	oh = newOuterHeader(encryptedPacket, count, buf.Bytes())
	if err := oh.write(wc, true); err != nil {
		return "", err
	}
	count++

	// continue HMAC calculation
	if err := oh.write(mac, true); err != nil {
		return "", err
	}

	// create HMAC header
	oh = newOuterHeader(hmacPacket, count, nil)
	oh.PLen = sha512.Size
	if err := oh.write(mac, false); err != nil {
		return "", err
	}
	oh.inner = mac.Sum(oh.inner)
	log.Debugf("HMAC:       %s", base64.Encode(oh.inner))
	if err := oh.write(wc, true); err != nil {
		return "", err
	}
	count++

	// write output
	wc.Close()
	if out.Len() != EncodedMsgSize {
		return "", log.Errorf("out.Len() = %d != %d = EncodedMsgSize)",
			out.Len(), EncodedMsgSize)
	}
	if _, err := io.Copy(args.Writer, &out); err != nil {
		return "", log.Error(err)
	}

	// delete message key
	err = args.KeyStore.DelMessageKey(sessionKey, true, ss.SenderMessageCount)
	if err != nil {
		return "", err
	}
	// increase SenderMessageCount
	ss.SenderMessageCount++
	err = args.KeyStore.SetSessionState(sessionStateKey, ss)
	if err != nil {
		return "", err
	}

	return
}
<|endoftext|>


<commit_before>// Copyright (c) 2014 The SkyDNS Authors. All rights reserved.
// Use of this source code is governed by The MIT License (MIT) that can be
// found in the LICENSE file.

package msg

import (
	"net"
	"path"
	"strings"

	"github.com/miekg/dns"
)

// This *is* the rdata from a SRV record, but with a twist.
// Host (Target in SRV) must be a domain name, but if it looks like an IP
// address (4/6), we will treat it like an IP address.
type Service struct {
	Host     string `json:"host,omitempty"`
	Port     int    `json:"port,omitempty"`
	Priority int    `json:"priority,omitempty"`
	Weight   int    `json:"weight,omitempty"`
	Text     string `json:"text,omitempty"`
	Mail     bool   `json:"mail,omitempty"` // Be an MX record. Priority becomes Preference.
	Ttl      uint32 `json:"ttl,omitempty"`

	// When a SRV record with a "Host: IP-address" is added, we synthesize
	// a srv.Target domain name.  Normally we convert the full Key where
	// the record lives to a DNS name and use this as the srv.Target.  When
	// TargetStrip > 0 we strip the left most TargetStrip labels from the
	// DNS name.
	TargetStrip int `json:"targetstrip",omitempty"`

	// Etcd key where we found this service and ignored from json un-/marshalling
	Key string `json:"-"`
}

// NewSRV returns a new SRV record based on the Service.
func (s *Service) NewSRV(name string, weight uint16) *dns.SRV {
	host := dns.Fqdn(s.Host)

	offset, end := 0, false
	for i := 0; i < s.TargetStrip; i++ {
		offset, end = dns.NextLabel(host, offset)
	}
	if end {
		// We overshot the name, use the orignal one.
		offset = 0
	}
	host = host[offset:]

	return &dns.SRV{Hdr: dns.RR_Header{Name: name, Rrtype: dns.TypeSRV, Class: dns.ClassINET, Ttl: s.Ttl},
		Priority: uint16(s.Priority), Weight: weight, Port: uint16(s.Port), Target: host}
}

// NewMX returns a new MX record based on the Service.
func (s *Service) NewMX(name string, weight uint16) *dns.MX {
	host := dns.Fqdn(s.Host)

	offset, end := 0, false
	for i := 0; i < s.TargetStrip; i++ {
		offset, end = dns.NextLabel(host, offset)
	}
	if end {
		// We overshot the name, use the orignal one.
		offset = 0
	}
	host = host[offset:]

	return &dns.MX{Hdr: dns.RR_Header{Name: name, Rrtype: dns.TypeMX, Class: dns.ClassINET, Ttl: s.Ttl},
		Preference: uint16(s.Priority), Mx: host}
}

// NewA returns a new A record based on the Service.
func (s *Service) NewA(name string, ip net.IP) *dns.A {
	return &dns.A{Hdr: dns.RR_Header{Name: name, Rrtype: dns.TypeA, Class: dns.ClassINET, Ttl: s.Ttl}, A: ip}
}

// NewAAAA returns a new AAAA record based on the Service.
func (s *Service) NewAAAA(name string, ip net.IP) *dns.AAAA {
	return &dns.AAAA{Hdr: dns.RR_Header{Name: name, Rrtype: dns.TypeAAAA, Class: dns.ClassINET, Ttl: s.Ttl}, AAAA: ip}
}

// NewCNAME returns a new CNAME record based on the Service.
func (s *Service) NewCNAME(name string, target string) *dns.CNAME {
	return &dns.CNAME{Hdr: dns.RR_Header{Name: name, Rrtype: dns.TypeCNAME, Class: dns.ClassINET, Ttl: s.Ttl}, Target: target}
}

// NewNS returns a new NS record based on the Service.
func (s *Service) NewNS(name string, target string) *dns.NS {
	return &dns.NS{Hdr: dns.RR_Header{Name: name, Rrtype: dns.TypeNS, Class: dns.ClassINET, Ttl: s.Ttl}, Ns: target}
}

// NewTXT returns a new TXT record based on the Service.
func (s *Service) NewTXT(name string) *dns.TXT {
	return &dns.TXT{Hdr: dns.RR_Header{Name: name, Rrtype: dns.TypeTXT, Class: dns.ClassINET, Ttl: s.Ttl}, Txt: split255(s.Text)}
}

// NewPTR returns a new PTR record based on the Service.
func (s *Service) NewPTR(name string, ttl uint32) *dns.PTR {
	return &dns.PTR{Hdr: dns.RR_Header{Name: name, Rrtype: dns.TypePTR, Class: dns.ClassINET, Ttl: ttl}, Ptr: dns.Fqdn(s.Host)}
}

// As Path, but
// if a name contains wildcards (*), the name will be chopped of before the (first) wildcard, and
// we do a highler evel search and later find the matching names.
// So service.*.skydns.local, will look for all services under skydns.local and will later check
// for names that match service.*.skydns.local.  If a wildcard is found the returned bool is true.
func PathWithWildcard(s string) (string, bool) {
	l := dns.SplitDomainName(s)
	for i, j := 0, len(l)-1; i < j; i, j = i+1, j-1 {
		l[i], l[j] = l[j], l[i]
	}
	for i, k := range l {
		if k == "*" {
			return path.Join(append([]string{"/skydns/"}, l[:i]...)...), true
		}
	}
	return path.Join(append([]string{"/skydns/"}, l...)...), false
}

// Path converts a domainname to an etcd path. If s looks like service.staging.skydns.local.,
// the resulting key will be /skydns/local/skydns/staging/service .
func Path(s string) string {
	l := dns.SplitDomainName(s)
	for i, j := 0, len(l)-1; i < j; i, j = i+1, j-1 {
		l[i], l[j] = l[j], l[i]
	}
	return path.Join(append([]string{"/skydns/"}, l...)...)
}

// Domain is the opposite of Path.
func Domain(s string) string {
	l := strings.Split(s, "/")
	// start with 1, to strip /skydns
	for i, j := 1, len(l)-1; i < j; i, j = i+1, j-1 {
		l[i], l[j] = l[j], l[i]
	}
	return dns.Fqdn(strings.Join(l[1:len(l)-1], "."))
}

// Split255 splits a string into 255 byte chunks.
func split255(s string) []string {
	if len(s) < 255 {
		return []string{s}
	}
	sx := []string{}
	p, i := 0, 255
	for {
		if i <= len(s) {
			sx = append(sx, s[p:i])
		} else {
			sx = append(sx, s[p:])
			break

		}
		p, i = p+255, i+255
	}

	return sx
}
<commit_msg>Use the correct function from the merge conflict<commit_after>// Copyright (c) 2014 The SkyDNS Authors. All rights reserved.
// Use of this source code is governed by The MIT License (MIT) that can be
// found in the LICENSE file.

package msg

import (
	"net"
	"path"
	"strings"

	"github.com/miekg/dns"
)

// This *is* the rdata from a SRV record, but with a twist.
// Host (Target in SRV) must be a domain name, but if it looks like an IP
// address (4/6), we will treat it like an IP address.
type Service struct {
	Host     string `json:"host,omitempty"`
	Port     int    `json:"port,omitempty"`
	Priority int    `json:"priority,omitempty"`
	Weight   int    `json:"weight,omitempty"`
	Text     string `json:"text,omitempty"`
	Mail     bool   `json:"mail,omitempty"` // Be an MX record. Priority becomes Preference.
	Ttl      uint32 `json:"ttl,omitempty"`

	// When a SRV record with a "Host: IP-address" is added, we synthesize
	// a srv.Target domain name.  Normally we convert the full Key where
	// the record lives to a DNS name and use this as the srv.Target.  When
	// TargetStrip > 0 we strip the left most TargetStrip labels from the
	// DNS name.
	TargetStrip int `json:"targetstrip",omitempty"`

	// Etcd key where we found this service and ignored from json un-/marshalling
	Key string `json:"-"`
}

// NewSRV returns a new SRV record based on the Service.
func (s *Service) NewSRV(name string, weight uint16) *dns.SRV {
	host := dns.Fqdn(s.Host)

	offset, end := 0, false
	for i := 0; i < s.TargetStrip; i++ {
		offset, end = dns.NextLabel(host, offset)
	}
	if end {
		// We overshot the name, use the orignal one.
		offset = 0
	}
	host = host[offset:]

	return &dns.SRV{Hdr: dns.RR_Header{Name: name, Rrtype: dns.TypeSRV, Class: dns.ClassINET, Ttl: s.Ttl},
		Priority: uint16(s.Priority), Weight: weight, Port: uint16(s.Port), Target: host}
}

// NewMX returns a new MX record based on the Service.
func (s *Service) NewMX(name string) *dns.MX {
	host := dns.Fqdn(s.Host)

	offset, end := 0, false
	for i := 0; i < s.TargetStrip; i++ {
		offset, end = dns.NextLabel(host, offset)
	}
	if end {
		// We overshot the name, use the orignal one.
		offset = 0
	}
	host = host[offset:]

	return &dns.MX{Hdr: dns.RR_Header{Name: name, Rrtype: dns.TypeMX, Class: dns.ClassINET, Ttl: s.Ttl},
		Preference: uint16(s.Priority), Mx: host}
}

// NewA returns a new A record based on the Service.
func (s *Service) NewA(name string, ip net.IP) *dns.A {
	return &dns.A{Hdr: dns.RR_Header{Name: name, Rrtype: dns.TypeA, Class: dns.ClassINET, Ttl: s.Ttl}, A: ip}
}

// NewAAAA returns a new AAAA record based on the Service.
func (s *Service) NewAAAA(name string, ip net.IP) *dns.AAAA {
	return &dns.AAAA{Hdr: dns.RR_Header{Name: name, Rrtype: dns.TypeAAAA, Class: dns.ClassINET, Ttl: s.Ttl}, AAAA: ip}
}

// NewCNAME returns a new CNAME record based on the Service.
func (s *Service) NewCNAME(name string, target string) *dns.CNAME {
	return &dns.CNAME{Hdr: dns.RR_Header{Name: name, Rrtype: dns.TypeCNAME, Class: dns.ClassINET, Ttl: s.Ttl}, Target: target}
}

// NewNS returns a new NS record based on the Service.
func (s *Service) NewNS(name string, target string) *dns.NS {
	return &dns.NS{Hdr: dns.RR_Header{Name: name, Rrtype: dns.TypeNS, Class: dns.ClassINET, Ttl: s.Ttl}, Ns: target}
}

// NewTXT returns a new TXT record based on the Service.
func (s *Service) NewTXT(name string) *dns.TXT {
	return &dns.TXT{Hdr: dns.RR_Header{Name: name, Rrtype: dns.TypeTXT, Class: dns.ClassINET, Ttl: s.Ttl}, Txt: split255(s.Text)}
}

// NewPTR returns a new PTR record based on the Service.
func (s *Service) NewPTR(name string, ttl uint32) *dns.PTR {
	return &dns.PTR{Hdr: dns.RR_Header{Name: name, Rrtype: dns.TypePTR, Class: dns.ClassINET, Ttl: ttl}, Ptr: dns.Fqdn(s.Host)}
}

// As Path, but
// if a name contains wildcards (*), the name will be chopped of before the (first) wildcard, and
// we do a highler evel search and later find the matching names.
// So service.*.skydns.local, will look for all services under skydns.local and will later check
// for names that match service.*.skydns.local.  If a wildcard is found the returned bool is true.
func PathWithWildcard(s string) (string, bool) {
	l := dns.SplitDomainName(s)
	for i, j := 0, len(l)-1; i < j; i, j = i+1, j-1 {
		l[i], l[j] = l[j], l[i]
	}
	for i, k := range l {
		if k == "*" {
			return path.Join(append([]string{"/skydns/"}, l[:i]...)...), true
		}
	}
	return path.Join(append([]string{"/skydns/"}, l...)...), false
}

// Path converts a domainname to an etcd path. If s looks like service.staging.skydns.local.,
// the resulting key will be /skydns/local/skydns/staging/service .
func Path(s string) string {
	l := dns.SplitDomainName(s)
	for i, j := 0, len(l)-1; i < j; i, j = i+1, j-1 {
		l[i], l[j] = l[j], l[i]
	}
	return path.Join(append([]string{"/skydns/"}, l...)...)
}

// Domain is the opposite of Path.
func Domain(s string) string {
	l := strings.Split(s, "/")
	// start with 1, to strip /skydns
	for i, j := 1, len(l)-1; i < j; i, j = i+1, j-1 {
		l[i], l[j] = l[j], l[i]
	}
	return dns.Fqdn(strings.Join(l[1:len(l)-1], "."))
}

// Split255 splits a string into 255 byte chunks.
func split255(s string) []string {
	if len(s) < 255 {
		return []string{s}
	}
	sx := []string{}
	p, i := 0, 255
	for {
		if i <= len(s) {
			sx = append(sx, s[p:i])
		} else {
			sx = append(sx, s[p:])
			break

		}
		p, i = p+255, i+255
	}

	return sx
}
<|endoftext|>


<commit_before>// MurmurHash3 was written by Austin Appleby, and is placed in the public
package murmurhash3

import (
	"encoding/binary"
	"hash"
)

type (
	murmurhash3A uint32
	murmurhash3C uint32
	murmurhash3F uint64
)

func New3A() hash.Hash32 {
	var m murmurhash3A = 0
	return &m
}
func (m *murmurhash3A) Reset() { *m = 0 }
func (m *murmurhash3A) Size() int {
	return 4
}
func (m *murmurhash3A) BlockSize() int {
	return 4
}
func (m *murmurhash3A) Write(p []byte) (n int, err error) {
	*m = murmurhash3A(Murmur3A(p, uint32(*m)))
	return len(p), nil
}
func (m *murmurhash3A) Sum32() uint32 {
	return uint32(*m)
}
func (m *murmurhash3A) Sum(in []byte) []byte {
	v := uint32(*m)
	return append(in, byte(v>>24), byte(v>>16), byte(v>>8), byte(v))
}

func New3C() hash.Hash32 {
	var m murmurhash3C = 0
	return &m
}
func (m *murmurhash3C) Reset() { *m = 0 }
func (m *murmurhash3C) Size() int {
	return 4
}
func (m *murmurhash3C) BlockSize() int {
	return 16
}
func (m *murmurhash3C) Write(p []byte) (n int, err error) {
	*m = murmurhash3C(Murmur3C(p, uint32(*m))[0])
	return len(p), nil
}
func (m *murmurhash3C) Sum32() uint32 {
	return uint32(*m)
}
func (m *murmurhash3C) Sum(in []byte) []byte {
	v := uint32(*m)
	return append(in, byte(v>>24), byte(v>>16), byte(v>>8), byte(v))
}

func New3F() hash.Hash64 {
	var m murmurhash3F = 0
	return &m
}
func (m *murmurhash3F) Reset() { *m = 0 }
func (m *murmurhash3F) Size() int {
	return 8
}
func (m *murmurhash3F) BlockSize() int {
	return 16
}
func (m *murmurhash3F) Write(p []byte) (n int, err error) {
	*m = murmurhash3F(Murmur3F(p, uint64(*m))[0])
	return len(p), nil
}
func (m *murmurhash3F) Sum64() uint64 {
	return uint64(*m)
}
func (m *murmurhash3F) Sum(in []byte) []byte {
	v := uint64(*m)
	return append(in, byte(v>>56), byte(v>>48), byte(v>>40), byte(v>>32), byte(v>>24), byte(v>>16), byte(v>>8), byte(v))
}

func rotl32(x uint32, r uint8) uint32 {
	return (x << r) | (x >> (32 - r))
}

func rotl64(x uint64, r uint8) uint64 {
	return (x << r) | (x >> (64 - r))
}

func fmix32(h uint32) uint32 {
	h ^= h >> 16
	h *= 0x85ebca6b
	h ^= h >> 13
	h *= 0xc2b2ae35
	h ^= h >> 16

	return h
}

func fmix64(h uint64) uint64 {
	h ^= h >> 33
	h *= 0xff51afd7ed558ccd
	h ^= h >> 33
	h *= 0xc4ceb9fe1a85ec53
	h ^= h >> 33

	return h
}

// MurmurHash3 for x86, 32-bit (MurmurHash3_x86_32)
func Murmur3A(key []byte, seed uint32) uint32 {
	nblocks := len(key) / 4
	var h1 uint32 = seed

	var c1 uint32 = 0xcc9e2d51
	var c2 uint32 = 0x1b873593

	// body
	for i := 0; i < nblocks; i++ {
		k1 := binary.LittleEndian.Uint32(key[i*4:]) // TODO Validate

		k1 *= c1
		k1 = rotl32(k1, 15)
		k1 *= c2

		h1 ^= k1
		h1 = rotl32(h1, 13)
		h1 = h1*5 + 0xe6546b64
	}

	// tail
	var tail []byte = key[nblocks*4:] // TODO Validate
	var k1 uint32 = 0
	switch len(key) & 3 {
	case 3:
		k1 ^= uint32(tail[2]) << 16
		fallthrough
	case 2:
		k1 ^= uint32(tail[1]) << 8
		fallthrough
	case 1:
		k1 ^= uint32(tail[0])
		k1 *= c1
		k1 = rotl32(k1, 15)
		k1 *= c2
		h1 ^= k1
	}

	//finalization
	h1 ^= uint32(len(key))

	h1 = fmix32(h1)

	return h1
}

// MurmurHash3 for x86, 128-bit (MurmurHash3_x86_128)
func Murmur3C(key []byte, seed uint32) [4]uint32 {
	nblocks := len(key) / 16
	var h1 uint32 = seed
	var h2 uint32 = seed
	var h3 uint32 = seed
	var h4 uint32 = seed

	var c1 uint32 = 0x239b961b
	var c2 uint32 = 0xab0e9789
	var c3 uint32 = 0x38b34ae5
	var c4 uint32 = 0xa1e38b93

	// body
	for i := 0; i < nblocks; i++ {
		k1 := binary.LittleEndian.Uint32(key[(i*4+0)*4:]) // TODO Validate
		k2 := binary.LittleEndian.Uint32(key[(i*4+1)*4:])
		k3 := binary.LittleEndian.Uint32(key[(i*4+2)*4:])
		k4 := binary.LittleEndian.Uint32(key[(i*4+3)*4:])

		k1 *= c1
		k1 = rotl32(k1, 15)
		k1 *= c2
		h1 ^= k1

		h1 = rotl32(h1, 19)
		h1 += h2
		h1 = h1*5 + 0x561ccd1b

		k2 *= c2
		k2 = rotl32(k2, 16)
		k2 *= c3
		h2 ^= k2

		h2 = rotl32(h2, 17)
		h2 += h3
		h2 = h2*5 + 0x0bcaa747

		k3 *= c3
		k3 = rotl32(k3, 17)
		k3 *= c4
		h3 ^= k3

		h3 = rotl32(h3, 15)
		h3 += h4
		h3 = h3*5 + 0x96cd1c35

		k4 *= c4
		k4 = rotl32(k4, 18)
		k4 *= c1
		h4 ^= k4

		h4 = rotl32(h4, 13)
		h4 += h1
		h4 = h4*5 + 0x32ac3b17
	}

	// tail
	var tail []byte = key[nblocks*16:] // TODO Validate
	var k1 uint32 = 0
	var k2 uint32 = 0
	var k3 uint32 = 0
	var k4 uint32 = 0
	switch len(key) & 15 {
	case 15:
		k4 ^= uint32(tail[14]) << 16
		fallthrough
	case 14:
		k4 ^= uint32(tail[13]) << 8
		fallthrough
	case 13:
		k4 ^= uint32(tail[12]) << 0
		k4 *= c4
		k4 = rotl32(k4, 18)
		k4 *= c1
		h4 ^= k4
		fallthrough
	case 12:
		k3 ^= uint32(tail[11]) << 24
		fallthrough
	case 11:
		k3 ^= uint32(tail[10]) << 16
		fallthrough
	case 10:
		k3 ^= uint32(tail[9]) << 8
		fallthrough
	case 9:
		k3 ^= uint32(tail[8]) << 0
		k3 *= c3
		k3 = rotl32(k3, 17)
		k3 *= c4
		h3 ^= k3
		fallthrough
	case 8:
		k2 ^= uint32(tail[7]) << 24
		fallthrough
	case 7:
		k2 ^= uint32(tail[6]) << 16
		fallthrough
	case 6:
		k2 ^= uint32(tail[5]) << 8
		fallthrough
	case 5:
		k2 ^= uint32(tail[4]) << 0
		k2 *= c2
		k2 = rotl32(k2, 16)
		k2 *= c3
		h2 ^= k2
		fallthrough
	case 4:
		k1 ^= uint32(tail[3]) << 24
		fallthrough
	case 3:
		k1 ^= uint32(tail[2]) << 16
		fallthrough
	case 2:
		k1 ^= uint32(tail[1]) << 8
		fallthrough
	case 1:
		k1 ^= uint32(tail[0]) << 0
		k1 *= c1
		k1 = rotl32(k1, 15)
		k1 *= c2
		h1 ^= k1
	}

	//finalization
	h1 ^= uint32(len(key))
	h2 ^= uint32(len(key))
	h3 ^= uint32(len(key))
	h4 ^= uint32(len(key))

	h1 += h2
	h1 += h3
	h1 += h4
	h2 += h1
	h3 += h1
	h4 += h1

	h1 = fmix32(h1)
	h2 = fmix32(h2)
	h3 = fmix32(h3)
	h4 = fmix32(h4)

	h1 += h2
	h1 += h3
	h1 += h4
	h2 += h1
	h3 += h1
	h4 += h1

	return [4]uint32{h1, h2, h3, h4}
}

// MurmurHash3 for x64, 128-bit (MurmurHash3_x64_128)
func Murmur3F(key []byte, seed uint64) [2]uint64 {
	nblocks := len(key) / 16
	var h1 uint64 = seed
	var h2 uint64 = seed

	var c1 uint64 = 0x87c37b91114253d5
	var c2 uint64 = 0x4cf5ad432745937f

	// body
	for i := 0; i < nblocks; i++ {
		k1 := binary.LittleEndian.Uint64(key[(i*2+0)*8:]) // TODO Validate
		k2 := binary.LittleEndian.Uint64(key[(i*2+1)*8:])

		k1 *= c1
		k1 = rotl64(k1, 31)
		k1 *= c2
		h1 ^= k1

		h1 = rotl64(h1, 27)
		h1 += h2
		h1 = h1*5 + 0x52dce729

		k2 *= c2
		k2 = rotl64(k2, 33)
		k2 *= c1
		h2 ^= k2

		h2 = rotl64(h2, 31)
		h2 += h1
		h2 = h2*5 + 0x38495ab5
	}

	// tail
	var tail []byte = key[nblocks*16:] // TODO Validate
	var k1 uint64 = 0
	var k2 uint64 = 0
	switch len(key) & 15 {
	case 15:
		k2 ^= uint64(tail[14]) << 48
		fallthrough
	case 14:
		k2 ^= uint64(tail[13]) << 40
		fallthrough
	case 13:
		k2 ^= uint64(tail[12]) << 32
		fallthrough
	case 12:
		k2 ^= uint64(tail[11]) << 24
		fallthrough
	case 11:
		k2 ^= uint64(tail[10]) << 16
		fallthrough
	case 10:
		k2 ^= uint64(tail[9]) << 8
		fallthrough
	case 9:
		k2 ^= uint64(tail[8]) << 0
		k2 *= c2
		k2 = rotl64(k2, 33)
		k2 *= c1
		h2 ^= k2
		fallthrough
	case 8:
		k1 ^= uint64(tail[7]) << 56
		fallthrough
	case 7:
		k1 ^= uint64(tail[6]) << 48
		fallthrough
	case 6:
		k1 ^= uint64(tail[5]) << 40
		fallthrough
	case 5:
		k1 ^= uint64(tail[4]) << 32
		fallthrough
	case 4:
		k1 ^= uint64(tail[3]) << 24
		fallthrough
	case 3:
		k1 ^= uint64(tail[2]) << 16
		fallthrough
	case 2:
		k1 ^= uint64(tail[1]) << 8
		fallthrough
	case 1:
		k1 ^= uint64(tail[0]) << 0
		k1 *= c1
		k1 = rotl64(k1, 31)
		k1 *= c2
		h1 ^= k1
	}

	//finalization
	h1 ^= uint64(len(key))
	h2 ^= uint64(len(key))

	h1 += h2
	h2 += h1

	h1 = fmix64(h1)
	h2 = fmix64(h2)

	h1 += h2
	h2 += h1

	return [2]uint64{h1, h2}
}
<commit_msg>Golint applied.<commit_after>// MurmurHash3 was written by Austin Appleby, and is placed in the public
package murmurhash3

import (
	"encoding/binary"
	"hash"
)

type (
	murmurhash3A uint32
	murmurhash3C uint32
	murmurhash3F uint64
)

func New3A() hash.Hash32 {
	var m murmurhash3A
	return &m
}
func (m *murmurhash3A) Reset() { *m = 0 }
func (m *murmurhash3A) Size() int {
	return 4
}
func (m *murmurhash3A) BlockSize() int {
	return 4
}
func (m *murmurhash3A) Write(p []byte) (n int, err error) {
	*m = murmurhash3A(Murmur3A(p, uint32(*m)))
	return len(p), nil
}
func (m *murmurhash3A) Sum32() uint32 {
	return uint32(*m)
}
func (m *murmurhash3A) Sum(in []byte) []byte {
	v := uint32(*m)
	return append(in, byte(v>>24), byte(v>>16), byte(v>>8), byte(v))
}

func New3C() hash.Hash32 {
	var m murmurhash3C
	return &m
}
func (m *murmurhash3C) Reset() { *m = 0 }
func (m *murmurhash3C) Size() int {
	return 4
}
func (m *murmurhash3C) BlockSize() int {
	return 16
}
func (m *murmurhash3C) Write(p []byte) (n int, err error) {
	*m = murmurhash3C(Murmur3C(p, uint32(*m))[0])
	return len(p), nil
}
func (m *murmurhash3C) Sum32() uint32 {
	return uint32(*m)
}
func (m *murmurhash3C) Sum(in []byte) []byte {
	v := uint32(*m)
	return append(in, byte(v>>24), byte(v>>16), byte(v>>8), byte(v))
}

func New3F() hash.Hash64 {
	var m murmurhash3F
	return &m
}
func (m *murmurhash3F) Reset() { *m = 0 }
func (m *murmurhash3F) Size() int {
	return 8
}
func (m *murmurhash3F) BlockSize() int {
	return 16
}
func (m *murmurhash3F) Write(p []byte) (n int, err error) {
	*m = murmurhash3F(Murmur3F(p, uint64(*m))[0])
	return len(p), nil
}
func (m *murmurhash3F) Sum64() uint64 {
	return uint64(*m)
}
func (m *murmurhash3F) Sum(in []byte) []byte {
	v := uint64(*m)
	return append(in, byte(v>>56), byte(v>>48), byte(v>>40), byte(v>>32), byte(v>>24), byte(v>>16), byte(v>>8), byte(v))
}

func rotl32(x uint32, r uint8) uint32 {
	return (x << r) | (x >> (32 - r))
}

func rotl64(x uint64, r uint8) uint64 {
	return (x << r) | (x >> (64 - r))
}

func fmix32(h uint32) uint32 {
	h ^= h >> 16
	h *= 0x85ebca6b
	h ^= h >> 13
	h *= 0xc2b2ae35
	h ^= h >> 16

	return h
}

func fmix64(h uint64) uint64 {
	h ^= h >> 33
	h *= 0xff51afd7ed558ccd
	h ^= h >> 33
	h *= 0xc4ceb9fe1a85ec53
	h ^= h >> 33

	return h
}

// MurmurHash3 for x86, 32-bit (MurmurHash3_x86_32)
func Murmur3A(key []byte, seed uint32) uint32 {
	nblocks := len(key) / 4
	var h1 = seed

	var c1 uint32 = 0xcc9e2d51
	var c2 uint32 = 0x1b873593

	// body
	for i := 0; i < nblocks; i++ {
		k1 := binary.LittleEndian.Uint32(key[i*4:]) // TODO Validate

		k1 *= c1
		k1 = rotl32(k1, 15)
		k1 *= c2

		h1 ^= k1
		h1 = rotl32(h1, 13)
		h1 = h1*5 + 0xe6546b64
	}

	// tail
	var tail = key[nblocks*4:] // TODO Validate
	var k1 uint32
	switch len(key) & 3 {
	case 3:
		k1 ^= uint32(tail[2]) << 16
		fallthrough
	case 2:
		k1 ^= uint32(tail[1]) << 8
		fallthrough
	case 1:
		k1 ^= uint32(tail[0])
		k1 *= c1
		k1 = rotl32(k1, 15)
		k1 *= c2
		h1 ^= k1
	}

	//finalization
	h1 ^= uint32(len(key))

	h1 = fmix32(h1)

	return h1
}

// MurmurHash3 for x86, 128-bit (MurmurHash3_x86_128)
func Murmur3C(key []byte, seed uint32) [4]uint32 {
	nblocks := len(key) / 16
	var h1 = seed
	var h2 = seed
	var h3 = seed
	var h4 = seed

	var c1 uint32 = 0x239b961b
	var c2 uint32 = 0xab0e9789
	var c3 uint32 = 0x38b34ae5
	var c4 uint32 = 0xa1e38b93

	// body
	for i := 0; i < nblocks; i++ {
		k1 := binary.LittleEndian.Uint32(key[(i*4+0)*4:]) // TODO Validate
		k2 := binary.LittleEndian.Uint32(key[(i*4+1)*4:])
		k3 := binary.LittleEndian.Uint32(key[(i*4+2)*4:])
		k4 := binary.LittleEndian.Uint32(key[(i*4+3)*4:])

		k1 *= c1
		k1 = rotl32(k1, 15)
		k1 *= c2
		h1 ^= k1

		h1 = rotl32(h1, 19)
		h1 += h2
		h1 = h1*5 + 0x561ccd1b

		k2 *= c2
		k2 = rotl32(k2, 16)
		k2 *= c3
		h2 ^= k2

		h2 = rotl32(h2, 17)
		h2 += h3
		h2 = h2*5 + 0x0bcaa747

		k3 *= c3
		k3 = rotl32(k3, 17)
		k3 *= c4
		h3 ^= k3

		h3 = rotl32(h3, 15)
		h3 += h4
		h3 = h3*5 + 0x96cd1c35

		k4 *= c4
		k4 = rotl32(k4, 18)
		k4 *= c1
		h4 ^= k4

		h4 = rotl32(h4, 13)
		h4 += h1
		h4 = h4*5 + 0x32ac3b17
	}

	// tail
	var tail = key[nblocks*16:] // TODO Validate
	var k1 uint32
	var k2 uint32
	var k3 uint32
	var k4 uint32
	switch len(key) & 15 {
	case 15:
		k4 ^= uint32(tail[14]) << 16
		fallthrough
	case 14:
		k4 ^= uint32(tail[13]) << 8
		fallthrough
	case 13:
		k4 ^= uint32(tail[12]) << 0
		k4 *= c4
		k4 = rotl32(k4, 18)
		k4 *= c1
		h4 ^= k4
		fallthrough
	case 12:
		k3 ^= uint32(tail[11]) << 24
		fallthrough
	case 11:
		k3 ^= uint32(tail[10]) << 16
		fallthrough
	case 10:
		k3 ^= uint32(tail[9]) << 8
		fallthrough
	case 9:
		k3 ^= uint32(tail[8]) << 0
		k3 *= c3
		k3 = rotl32(k3, 17)
		k3 *= c4
		h3 ^= k3
		fallthrough
	case 8:
		k2 ^= uint32(tail[7]) << 24
		fallthrough
	case 7:
		k2 ^= uint32(tail[6]) << 16
		fallthrough
	case 6:
		k2 ^= uint32(tail[5]) << 8
		fallthrough
	case 5:
		k2 ^= uint32(tail[4]) << 0
		k2 *= c2
		k2 = rotl32(k2, 16)
		k2 *= c3
		h2 ^= k2
		fallthrough
	case 4:
		k1 ^= uint32(tail[3]) << 24
		fallthrough
	case 3:
		k1 ^= uint32(tail[2]) << 16
		fallthrough
	case 2:
		k1 ^= uint32(tail[1]) << 8
		fallthrough
	case 1:
		k1 ^= uint32(tail[0]) << 0
		k1 *= c1
		k1 = rotl32(k1, 15)
		k1 *= c2
		h1 ^= k1
	}

	//finalization
	h1 ^= uint32(len(key))
	h2 ^= uint32(len(key))
	h3 ^= uint32(len(key))
	h4 ^= uint32(len(key))

	h1 += h2
	h1 += h3
	h1 += h4
	h2 += h1
	h3 += h1
	h4 += h1

	h1 = fmix32(h1)
	h2 = fmix32(h2)
	h3 = fmix32(h3)
	h4 = fmix32(h4)

	h1 += h2
	h1 += h3
	h1 += h4
	h2 += h1
	h3 += h1
	h4 += h1

	return [4]uint32{h1, h2, h3, h4}
}

// MurmurHash3 for x64, 128-bit (MurmurHash3_x64_128)
func Murmur3F(key []byte, seed uint64) [2]uint64 {
	nblocks := len(key) / 16
	var h1 = seed
	var h2 = seed

	var c1 uint64 = 0x87c37b91114253d5
	var c2 uint64 = 0x4cf5ad432745937f

	// body
	for i := 0; i < nblocks; i++ {
		k1 := binary.LittleEndian.Uint64(key[(i*2+0)*8:]) // TODO Validate
		k2 := binary.LittleEndian.Uint64(key[(i*2+1)*8:])

		k1 *= c1
		k1 = rotl64(k1, 31)
		k1 *= c2
		h1 ^= k1

		h1 = rotl64(h1, 27)
		h1 += h2
		h1 = h1*5 + 0x52dce729

		k2 *= c2
		k2 = rotl64(k2, 33)
		k2 *= c1
		h2 ^= k2

		h2 = rotl64(h2, 31)
		h2 += h1
		h2 = h2*5 + 0x38495ab5
	}

	// tail
	var tail = key[nblocks*16:] // TODO Validate
	var k1 uint64
	var k2 uint64
	switch len(key) & 15 {
	case 15:
		k2 ^= uint64(tail[14]) << 48
		fallthrough
	case 14:
		k2 ^= uint64(tail[13]) << 40
		fallthrough
	case 13:
		k2 ^= uint64(tail[12]) << 32
		fallthrough
	case 12:
		k2 ^= uint64(tail[11]) << 24
		fallthrough
	case 11:
		k2 ^= uint64(tail[10]) << 16
		fallthrough
	case 10:
		k2 ^= uint64(tail[9]) << 8
		fallthrough
	case 9:
		k2 ^= uint64(tail[8]) << 0
		k2 *= c2
		k2 = rotl64(k2, 33)
		k2 *= c1
		h2 ^= k2
		fallthrough
	case 8:
		k1 ^= uint64(tail[7]) << 56
		fallthrough
	case 7:
		k1 ^= uint64(tail[6]) << 48
		fallthrough
	case 6:
		k1 ^= uint64(tail[5]) << 40
		fallthrough
	case 5:
		k1 ^= uint64(tail[4]) << 32
		fallthrough
	case 4:
		k1 ^= uint64(tail[3]) << 24
		fallthrough
	case 3:
		k1 ^= uint64(tail[2]) << 16
		fallthrough
	case 2:
		k1 ^= uint64(tail[1]) << 8
		fallthrough
	case 1:
		k1 ^= uint64(tail[0]) << 0
		k1 *= c1
		k1 = rotl64(k1, 31)
		k1 *= c2
		h1 ^= k1
	}

	//finalization
	h1 ^= uint64(len(key))
	h2 ^= uint64(len(key))

	h1 += h2
	h2 += h1

	h1 = fmix64(h1)
	h2 = fmix64(h2)

	h1 += h2
	h2 += h1

	return [2]uint64{h1, h2}
}
<|endoftext|>