<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 types declares the data types and implements
// the algorithms for type-checking of Go packages.
// Use Check and Config.Check to invoke the type-checker.
//
// Type-checking consists of several interdependent phases:
//
// Name resolution maps each identifier (ast.Ident) in the program to the
// language object (Object) it denotes.
// Use Info.Objects, Info.Implicits for the results of name resolution.
//
// Constant folding computes the exact constant value (exact.Value) for
// every expression (ast.Expr) that is a compile-time constant.
// Use Info.Values for the results of constant folding.
//
// Type inference computes the type (Type) of every expression (ast.Expr)
// and checks for compliance with the language specification.
// Use Info.Types for the results of type evaluation.
//
package types

import (
	"bytes"
	"fmt"
	"go/ast"
	"go/token"

	"code.google.com/p/go.tools/go/exact"
)

// Check type-checks a package and returns the resulting complete package
// object, or a nil package and the first error. The package is specified
// by a list of *ast.Files and corresponding file set, and the import path
// the package is identified with. The clean path must not be empty or dot (".").
//
// For more control over type-checking and results, use Config.Check.
func Check(path string, fset *token.FileSet, files []*ast.File) (*Package, error) {
	var conf Config
	pkg, err := conf.Check(path, fset, files, nil)
	if err != nil {
		return nil, err
	}
	return pkg, nil
}

// An Error describes a type-checking error;
// it implements the error interface.
type Error struct {
	Fset *token.FileSet // file set for interpretation of Pos
	Pos  token.Pos      // error position
	Msg  string         // error message
}

// Error returns an error string formatted as follows:
// filename:line:column: message
func (err Error) Error() string {
	return fmt.Sprintf("%s: %s", err.Fset.Position(err.Pos), err.Msg)
}

// An importer resolves import paths to Packages.
// The imports map records packages already known,
// indexed by package path. The type-checker
// will invoke Import with Config.Packages.
// An importer must determine the canonical package path and
// check imports to see if it is already present in the map.
// If so, the Importer can return the map entry.  Otherwise,
// the importer must load the package data for the given path
// into a new *Package, record it in imports map, and return
// the package.
// TODO(gri) Need to be clearer about requirements of completeness.
type Importer func(map[string]*Package, string) (*Package, error)

// A Config specifies the configuration for type checking.
// The zero value for Config is a ready-to-use default configuration.
type Config struct {
	// If IgnoreFuncBodies is set, function bodies are not
	// type-checked.
	IgnoreFuncBodies bool

	// If FakeImportC is set, `import "C"` (for packages requiring Cgo)
	// declares an empty "C" package and errors are omitted for qualified
	// identifiers referring to package C (which won't find an object).
	// This feature is intended for the standard library cmd/api tool.
	//
	// Caution: Effects may be unpredictable due to follow-up errors.
	//          Do not use casually!
	FakeImportC bool

	// Packages is used to look up (and thus canonicalize) packages by
	// package path. If Packages is nil, it is set to a new empty map.
	// During type-checking, imported packages are added to the map.
	Packages map[string]*Package

	// If Error != nil, it is called with each error found
	// during type checking; err has dynamic type Error.
	Error func(err error)

	// If Import != nil, it is called for each imported package.
	// Otherwise, DefaultImport is called.
	Import Importer

	// If Sizes != nil, it provides the sizing functions for package unsafe.
	// Otherwise &StdSize{WordSize: 8, MaxAlign: 8} is used instead.
	Sizes Sizes
}

// DefaultImport is the default importer invoked if Config.Import == nil.
// The declaration:
//
//	import _ "code.google.com/p/go.tools/go/gcimporter"
//
// in a client of go/types will initialize DefaultImport to gcimporter.Import.
var DefaultImport Importer

// Info holds result type information for a type-checked package.
// Only the information for which a map is provided is collected.
// If the package has type errors, the collected information may
// be incomplete.
type Info struct {
	// Types maps expressions to their types. Identifiers on the
	// lhs of declarations are collected in Objects, not Types.
	//
	// For an expression denoting a predeclared built-in function
	// the recorded signature is call-site specific. If the call
	// result is not a constant, the recorded type is an argument-
	// specific signature. Otherwise, the recorded type is invalid.
	Types map[ast.Expr]Type

	// Values maps constant expressions to their values.
	Values map[ast.Expr]exact.Value

	// Objects maps identifiers to their corresponding objects (including
	// package names, dots "." of dot-imports, and blank "_" identifiers).
	// For identifiers that do not denote objects (e.g., the package name
	// in package clauses, blank identifiers on the lhs of assignments, or
	// symbolic variables t in t := x.(type) of type switch headers), the
	// corresponding objects are nil.
	Objects map[*ast.Ident]Object

	// Implicits maps nodes to their implicitly declared objects, if any.
	// The following node and object types may appear:
	//
	//	node               declared object
	//
	//	*ast.ImportSpec    *PkgName (imports w/o renames), or imported objects (dot-imports)
	//	*ast.CaseClause    type-specific *Var for each type switch case clause (incl. default)
	//      *ast.Field         anonymous struct field or parameter *Var
	//
	Implicits map[ast.Node]Object

	// Selections maps selector expressions to their corresponding selections.
	Selections map[*ast.SelectorExpr]*Selection

	// Scopes maps ast.Nodes to the scopes they define. Note that package scopes
	// are not associated with a specific node but with all files belonging to a
	// package. Thus, the package scope can be found in the type-checked package
	// object.
	//
	// The following node types may appear in Scopes:
	//
	//	*ast.File
	//	*ast.FuncType
	//	*ast.BlockStmt
	//	*ast.IfStmt
	//	*ast.SwitchStmt
	//	*ast.TypeSwitchStmt
	//	*ast.CaseClause
	//	*ast.CommClause
	//	*ast.ForStmt
	//	*ast.RangeStmt
	//
	Scopes map[ast.Node]*Scope

	// InitOrder is the list of package-level initializers in the order in which
	// they must be executed. Initializers referring to variables related by an
	// initialization dependency appear in topological order, the others appear
	// in source order. Variables without an initialization expression do not
	// appear in this list.
	InitOrder []*Initializer
}

// An Initializer describes a package-level variable, or a list of variables in case
// of a multi-valued initialization expression, and the corresponding initialization
// expression.
type Initializer struct {
	Lhs []*Var // var Lhs = Rhs
	Rhs ast.Expr
}

func (init *Initializer) String() string {
	var buf bytes.Buffer
	for i, lhs := range init.Lhs {
		if i > 0 {
			buf.WriteString(", ")
		}
		buf.WriteString(lhs.Name())
	}
	buf.WriteString(" = ")
	WriteExpr(&buf, init.Rhs)
	return buf.String()
}

// Check type-checks a package and returns the resulting package object,
// the first error if any, and if info != nil, additional type information.
// The package is marked as complete if no errors occurred, otherwise it is
// incomplete.
//
// The package is specified by a list of *ast.Files and corresponding
// file set, and the package path the package is identified with.
// The clean path must not be empty or dot (".").
func (conf *Config) Check(path string, fset *token.FileSet, files []*ast.File, info *Info) (*Package, error) {
	pkg, err := conf.check(path, fset, files, info)
	if err == nil {
		pkg.complete = true
	}
	return pkg, err
}

// IsAssignableTo reports whether a value of type V is assignable to a variable of type T.
func IsAssignableTo(V, T Type) bool {
	x := operand{mode: value, typ: V}
	return x.isAssignableTo(nil, T) // config not needed for non-constant x
}

// Implements reports whether a value of type V implements T, as follows:
//
// 1) For non-interface types V, or if static is set, V implements T if all
// methods of T are present in V. Informally, this reports whether V is a
// subtype of T.
//
// 2) For interface types V, and if static is not set, V implements T if all
// methods of T which are also present in V have matching types. Informally,
// this indicates whether a type assertion x.(T) where x is of type V would
// be legal (the concrete dynamic type of x may implement T even if V does
// not statically implement it).
//
func Implements(V Type, T *Interface, static bool) bool {
	f, _ := MissingMethod(V, T, static)
	return f == nil
}
<commit_msg>go.toools/go/types: clearer comment (per adonovan)<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 types declares the data types and implements
// the algorithms for type-checking of Go packages.
// Use Check and Config.Check to invoke the type-checker.
//
// Type-checking consists of several interdependent phases:
//
// Name resolution maps each identifier (ast.Ident) in the program to the
// language object (Object) it denotes.
// Use Info.Objects, Info.Implicits for the results of name resolution.
//
// Constant folding computes the exact constant value (exact.Value) for
// every expression (ast.Expr) that is a compile-time constant.
// Use Info.Values for the results of constant folding.
//
// Type inference computes the type (Type) of every expression (ast.Expr)
// and checks for compliance with the language specification.
// Use Info.Types for the results of type evaluation.
//
package types

import (
	"bytes"
	"fmt"
	"go/ast"
	"go/token"

	"code.google.com/p/go.tools/go/exact"
)

// Check type-checks a package and returns the resulting complete package
// object, or a nil package and the first error. The package is specified
// by a list of *ast.Files and corresponding file set, and the import path
// the package is identified with. The clean path must not be empty or dot (".").
//
// For more control over type-checking and results, use Config.Check.
func Check(path string, fset *token.FileSet, files []*ast.File) (*Package, error) {
	var conf Config
	pkg, err := conf.Check(path, fset, files, nil)
	if err != nil {
		return nil, err
	}
	return pkg, nil
}

// An Error describes a type-checking error;
// it implements the error interface.
type Error struct {
	Fset *token.FileSet // file set for interpretation of Pos
	Pos  token.Pos      // error position
	Msg  string         // error message
}

// Error returns an error string formatted as follows:
// filename:line:column: message
func (err Error) Error() string {
	return fmt.Sprintf("%s: %s", err.Fset.Position(err.Pos), err.Msg)
}

// An importer resolves import paths to Packages.
// The imports map records packages already known,
// indexed by package path. The type-checker
// will invoke Import with Config.Packages.
// An importer must determine the canonical package path and
// check imports to see if it is already present in the map.
// If so, the Importer can return the map entry.  Otherwise,
// the importer must load the package data for the given path
// into a new *Package, record it in imports map, and return
// the package.
// TODO(gri) Need to be clearer about requirements of completeness.
type Importer func(map[string]*Package, string) (*Package, error)

// A Config specifies the configuration for type checking.
// The zero value for Config is a ready-to-use default configuration.
type Config struct {
	// If IgnoreFuncBodies is set, function bodies are not
	// type-checked.
	IgnoreFuncBodies bool

	// If FakeImportC is set, `import "C"` (for packages requiring Cgo)
	// declares an empty "C" package and errors are omitted for qualified
	// identifiers referring to package C (which won't find an object).
	// This feature is intended for the standard library cmd/api tool.
	//
	// Caution: Effects may be unpredictable due to follow-up errors.
	//          Do not use casually!
	FakeImportC bool

	// Packages is used to look up (and thus canonicalize) packages by
	// package path. If Packages is nil, it is set to a new empty map.
	// During type-checking, imported packages are added to the map.
	Packages map[string]*Package

	// If Error != nil, it is called with each error found
	// during type checking; err has dynamic type Error.
	Error func(err error)

	// If Import != nil, it is called for each imported package.
	// Otherwise, DefaultImport is called.
	Import Importer

	// If Sizes != nil, it provides the sizing functions for package unsafe.
	// Otherwise &StdSize{WordSize: 8, MaxAlign: 8} is used instead.
	Sizes Sizes
}

// DefaultImport is the default importer invoked if Config.Import == nil.
// The declaration:
//
//	import _ "code.google.com/p/go.tools/go/gcimporter"
//
// in a client of go/types will initialize DefaultImport to gcimporter.Import.
var DefaultImport Importer

// Info holds result type information for a type-checked package.
// Only the information for which a map is provided is collected.
// If the package has type errors, the collected information may
// be incomplete.
type Info struct {
	// Types maps expressions to their types. Identifiers on the
	// lhs of declarations are collected in Objects, not Types.
	//
	// For an expression denoting a predeclared built-in function
	// the recorded signature is call-site specific. If the call
	// result is not a constant, the recorded type is an argument-
	// specific signature. Otherwise, the recorded type is invalid.
	Types map[ast.Expr]Type

	// Values maps constant expressions to their values.
	Values map[ast.Expr]exact.Value

	// Objects maps identifiers to their corresponding objects (including
	// package names, dots "." of dot-imports, and blank "_" identifiers).
	// For identifiers that do not denote objects (e.g., the package name
	// in package clauses, blank identifiers on the lhs of assignments, or
	// symbolic variables t in t := x.(type) of type switch headers), the
	// corresponding objects are nil.
	Objects map[*ast.Ident]Object

	// Implicits maps nodes to their implicitly declared objects, if any.
	// The following node and object types may appear:
	//
	//	node               declared object
	//
	//	*ast.ImportSpec    *PkgName for dot-imports and imports without renames
	//	*ast.CaseClause    type-specific *Var for each type switch case clause (incl. default)
	//      *ast.Field         anonymous struct field or parameter *Var
	//
	Implicits map[ast.Node]Object

	// Selections maps selector expressions to their corresponding selections.
	Selections map[*ast.SelectorExpr]*Selection

	// Scopes maps ast.Nodes to the scopes they define. Note that package scopes
	// are not associated with a specific node but with all files belonging to a
	// package. Thus, the package scope can be found in the type-checked package
	// object.
	//
	// The following node types may appear in Scopes:
	//
	//	*ast.File
	//	*ast.FuncType
	//	*ast.BlockStmt
	//	*ast.IfStmt
	//	*ast.SwitchStmt
	//	*ast.TypeSwitchStmt
	//	*ast.CaseClause
	//	*ast.CommClause
	//	*ast.ForStmt
	//	*ast.RangeStmt
	//
	Scopes map[ast.Node]*Scope

	// InitOrder is the list of package-level initializers in the order in which
	// they must be executed. Initializers referring to variables related by an
	// initialization dependency appear in topological order, the others appear
	// in source order. Variables without an initialization expression do not
	// appear in this list.
	InitOrder []*Initializer
}

// An Initializer describes a package-level variable, or a list of variables in case
// of a multi-valued initialization expression, and the corresponding initialization
// expression.
type Initializer struct {
	Lhs []*Var // var Lhs = Rhs
	Rhs ast.Expr
}

func (init *Initializer) String() string {
	var buf bytes.Buffer
	for i, lhs := range init.Lhs {
		if i > 0 {
			buf.WriteString(", ")
		}
		buf.WriteString(lhs.Name())
	}
	buf.WriteString(" = ")
	WriteExpr(&buf, init.Rhs)
	return buf.String()
}

// Check type-checks a package and returns the resulting package object,
// the first error if any, and if info != nil, additional type information.
// The package is marked as complete if no errors occurred, otherwise it is
// incomplete.
//
// The package is specified by a list of *ast.Files and corresponding
// file set, and the package path the package is identified with.
// The clean path must not be empty or dot (".").
func (conf *Config) Check(path string, fset *token.FileSet, files []*ast.File, info *Info) (*Package, error) {
	pkg, err := conf.check(path, fset, files, info)
	if err == nil {
		pkg.complete = true
	}
	return pkg, err
}

// IsAssignableTo reports whether a value of type V is assignable to a variable of type T.
func IsAssignableTo(V, T Type) bool {
	x := operand{mode: value, typ: V}
	return x.isAssignableTo(nil, T) // config not needed for non-constant x
}

// Implements reports whether a value of type V implements T, as follows:
//
// 1) For non-interface types V, or if static is set, V implements T if all
// methods of T are present in V. Informally, this reports whether V is a
// subtype of T.
//
// 2) For interface types V, and if static is not set, V implements T if all
// methods of T which are also present in V have matching types. Informally,
// this indicates whether a type assertion x.(T) where x is of type V would
// be legal (the concrete dynamic type of x may implement T even if V does
// not statically implement it).
//
func Implements(V Type, T *Interface, static bool) bool {
	f, _ := MissingMethod(V, T, static)
	return f == nil
}
<|endoftext|>


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

package main

import (
	"fmt"
	"os"
	"testing"
	"time"

	. "github.com/jacobsa/ogletest"
	"github.com/jgeewax/cli"
)

func TestFlags(t *testing.T) { RunTests(t) }

////////////////////////////////////////////////////////////////////////
// Boilerplate
////////////////////////////////////////////////////////////////////////

type FlagsTest struct {
}

func init() { RegisterTestSuite(&FlagsTest{}) }

func parseArgs(args []string) (flags *flagStorage) {
	// Create a CLI app, and abuse it to snoop on the flags.
	app := newApp()
	app.Action = func(appCtx *cli.Context) {
		flags = populateFlags(appCtx)
	}

	// Simulate argv.
	fullArgs := append([]string{"some_app"}, args...)

	err := app.Run(fullArgs)
	AssertEq(nil, err)

	return
}

////////////////////////////////////////////////////////////////////////
// Tests
////////////////////////////////////////////////////////////////////////

func (t *FlagsTest) Defaults() {
	f := parseArgs([]string{})

	// File system
	ExpectEq(0, len(f.MountOptions), "Options: %v", f.MountOptions)
	ExpectEq(os.FileMode(0755), f.DirMode)
	ExpectEq(os.FileMode(0644), f.FileMode)
	ExpectEq(-1, f.Uid)
	ExpectEq(-1, f.Gid)
	ExpectFalse(f.ImplicitDirs)

	// GCS
	ExpectEq("", f.KeyFile)
	ExpectEq(-1, f.EgressBandwidthLimitBytesPerSecond)
	ExpectEq(5, f.OpRateLimitHz)

	// Tuning
	ExpectEq(time.Minute, f.StatCacheTTL)
	ExpectEq(time.Minute, f.TypeCacheTTL)
	ExpectEq(1<<24, f.GCSChunkSize)
	ExpectEq("", f.TempDir)
	ExpectEq(1<<31, f.TempDirLimit)

	// Debugging
	ExpectFalse(f.DebugFuse)
	ExpectFalse(f.DebugGCS)
	ExpectFalse(f.DebugHTTP)
	ExpectFalse(f.DebugInvariants)
}

func (t *FlagsTest) Bools() {
	names := []string{
		"implicit-dirs",
		"debug_fuse",
		"debug_gcs",
		"debug_http",
		"debug_invariants",
	}

	var args []string
	var f *flagStorage

	// --foo form
	args = nil
	for _, n := range names {
		args = append(args, fmt.Sprintf("-%s", n))
	}

	f = parseArgs(args)
	ExpectTrue(f.ImplicitDirs)
	ExpectTrue(f.DebugFuse)
	ExpectTrue(f.DebugGCS)
	ExpectTrue(f.DebugHTTP)
	ExpectTrue(f.DebugInvariants)

	// --foo=false form
	args = nil
	for _, n := range names {
		args = append(args, fmt.Sprintf("-%s=false", n))
	}

	f = parseArgs(args)
	ExpectFalse(f.ImplicitDirs)
	ExpectFalse(f.DebugFuse)
	ExpectFalse(f.DebugGCS)
	ExpectFalse(f.DebugHTTP)
	ExpectFalse(f.DebugInvariants)

	// --foo=true form
	args = nil
	for _, n := range names {
		args = append(args, fmt.Sprintf("-%s=true", n))
	}

	f = parseArgs(args)
	ExpectTrue(f.ImplicitDirs)
	ExpectTrue(f.DebugFuse)
	ExpectTrue(f.DebugGCS)
	ExpectTrue(f.DebugHTTP)
	ExpectTrue(f.DebugInvariants)
}

func (t *FlagsTest) Numbers() {
	args := []string{
		"--dir-mode=0711",
		"--file-mode", "0611",
		"--uid=17",
		"--gid=19",
		"--limit-bytes-per-sec=123.4",
		"--limit-ops-per-sec=56.78",
		"--gcs-chunk-size=1000",
		"--temp-dir-bytes=2000",
	}

	f := parseArgs(args)
	ExpectEq(os.FileMode(0711), f.DirMode)
	ExpectEq(os.FileMode(0611), f.FileMode)
	ExpectEq(17, f.Uid)
	ExpectEq(19, f.Gid)
	ExpectEq(123.4, f.EgressBandwidthLimitBytesPerSecond)
	ExpectEq(56.78, f.OpRateLimitHz)
	ExpectEq(1000, f.GCSChunkSize)
	ExpectEq(2000, f.TempDirLimit)
}

func (t *FlagsTest) Strings() {
	args := []string{
		"--key-file", "-asdf",
		"--temp-dir=foobar",
	}

	f := parseArgs(args)
	ExpectEq("-asdf", f.KeyFile)
	ExpectEq("foobar", f.TempDir)
}

func (t *FlagsTest) Durations() {
	args := []string{
		"--stat-cache-ttl", "1m17s",
		"--type-cache-ttl", "19ns",
	}

	f := parseArgs(args)
	ExpectEq(77*time.Second, f.StatCacheTTL)
	ExpectEq(19*time.Nanosecond, f.TypeCacheTTL)
}

func (t *FlagsTest) Maps() {
	AssertTrue(false, "TODO")
}
<commit_msg>FlagsTest.Maps<commit_after>// Copyright 2015 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package main

import (
	"fmt"
	"os"
	"sort"
	"testing"
	"time"

	. "github.com/jacobsa/oglematchers"
	. "github.com/jacobsa/ogletest"
	"github.com/jgeewax/cli"
)

func TestFlags(t *testing.T) { RunTests(t) }

////////////////////////////////////////////////////////////////////////
// Boilerplate
////////////////////////////////////////////////////////////////////////

type FlagsTest struct {
}

func init() { RegisterTestSuite(&FlagsTest{}) }

func parseArgs(args []string) (flags *flagStorage) {
	// Create a CLI app, and abuse it to snoop on the flags.
	app := newApp()
	app.Action = func(appCtx *cli.Context) {
		flags = populateFlags(appCtx)
	}

	// Simulate argv.
	fullArgs := append([]string{"some_app"}, args...)

	err := app.Run(fullArgs)
	AssertEq(nil, err)

	return
}

////////////////////////////////////////////////////////////////////////
// Tests
////////////////////////////////////////////////////////////////////////

func (t *FlagsTest) Defaults() {
	f := parseArgs([]string{})

	// File system
	ExpectNe(nil, f.MountOptions)
	ExpectEq(0, len(f.MountOptions), "Options: %v", f.MountOptions)

	ExpectEq(os.FileMode(0755), f.DirMode)
	ExpectEq(os.FileMode(0644), f.FileMode)
	ExpectEq(-1, f.Uid)
	ExpectEq(-1, f.Gid)
	ExpectFalse(f.ImplicitDirs)

	// GCS
	ExpectEq("", f.KeyFile)
	ExpectEq(-1, f.EgressBandwidthLimitBytesPerSecond)
	ExpectEq(5, f.OpRateLimitHz)

	// Tuning
	ExpectEq(time.Minute, f.StatCacheTTL)
	ExpectEq(time.Minute, f.TypeCacheTTL)
	ExpectEq(1<<24, f.GCSChunkSize)
	ExpectEq("", f.TempDir)
	ExpectEq(1<<31, f.TempDirLimit)

	// Debugging
	ExpectFalse(f.DebugFuse)
	ExpectFalse(f.DebugGCS)
	ExpectFalse(f.DebugHTTP)
	ExpectFalse(f.DebugInvariants)
}

func (t *FlagsTest) Bools() {
	names := []string{
		"implicit-dirs",
		"debug_fuse",
		"debug_gcs",
		"debug_http",
		"debug_invariants",
	}

	var args []string
	var f *flagStorage

	// --foo form
	args = nil
	for _, n := range names {
		args = append(args, fmt.Sprintf("-%s", n))
	}

	f = parseArgs(args)
	ExpectTrue(f.ImplicitDirs)
	ExpectTrue(f.DebugFuse)
	ExpectTrue(f.DebugGCS)
	ExpectTrue(f.DebugHTTP)
	ExpectTrue(f.DebugInvariants)

	// --foo=false form
	args = nil
	for _, n := range names {
		args = append(args, fmt.Sprintf("-%s=false", n))
	}

	f = parseArgs(args)
	ExpectFalse(f.ImplicitDirs)
	ExpectFalse(f.DebugFuse)
	ExpectFalse(f.DebugGCS)
	ExpectFalse(f.DebugHTTP)
	ExpectFalse(f.DebugInvariants)

	// --foo=true form
	args = nil
	for _, n := range names {
		args = append(args, fmt.Sprintf("-%s=true", n))
	}

	f = parseArgs(args)
	ExpectTrue(f.ImplicitDirs)
	ExpectTrue(f.DebugFuse)
	ExpectTrue(f.DebugGCS)
	ExpectTrue(f.DebugHTTP)
	ExpectTrue(f.DebugInvariants)
}

func (t *FlagsTest) Numbers() {
	args := []string{
		"--dir-mode=0711",
		"--file-mode", "0611",
		"--uid=17",
		"--gid=19",
		"--limit-bytes-per-sec=123.4",
		"--limit-ops-per-sec=56.78",
		"--gcs-chunk-size=1000",
		"--temp-dir-bytes=2000",
	}

	f := parseArgs(args)
	ExpectEq(os.FileMode(0711), f.DirMode)
	ExpectEq(os.FileMode(0611), f.FileMode)
	ExpectEq(17, f.Uid)
	ExpectEq(19, f.Gid)
	ExpectEq(123.4, f.EgressBandwidthLimitBytesPerSecond)
	ExpectEq(56.78, f.OpRateLimitHz)
	ExpectEq(1000, f.GCSChunkSize)
	ExpectEq(2000, f.TempDirLimit)
}

func (t *FlagsTest) Strings() {
	args := []string{
		"--key-file", "-asdf",
		"--temp-dir=foobar",
	}

	f := parseArgs(args)
	ExpectEq("-asdf", f.KeyFile)
	ExpectEq("foobar", f.TempDir)
}

func (t *FlagsTest) Durations() {
	args := []string{
		"--stat-cache-ttl", "1m17s",
		"--type-cache-ttl", "19ns",
	}

	f := parseArgs(args)
	ExpectEq(77*time.Second, f.StatCacheTTL)
	ExpectEq(19*time.Nanosecond, f.TypeCacheTTL)
}

func (t *FlagsTest) Maps() {
	args := []string{
		"-o", "rw,nodev",
		"-o", "user=jacobsa,noauto",
	}

	f := parseArgs(args)

	var keys sort.StringSlice
	for k := range f.MountOptions {
		keys = append(keys, k)
	}

	sort.Sort(keys)
	AssertThat(keys, ElementsAre("noauto", "nodev", "rw", "user"))

	ExpectEq("", f.MountOptions["noauto"])
	ExpectEq("", f.MountOptions["nodev"])
	ExpectEq("", f.MountOptions["rw"])
	ExpectEq("jacobsa", f.MountOptions["user"])
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"net/url"
	"os"

	functions "github.com/funcy/functions_go"
	"github.com/urfave/cli"
)

// Version of Functions CLI
var Version = "0.3.8"

func version() cli.Command {
	r := versionCmd{VersionApi: functions.NewVersionApi()}
	return cli.Command{
		Name:   "version",
		Usage:  "displays fn and functions daemon versions",
		Action: r.version,
	}
}

type versionCmd struct {
	*functions.VersionApi
}

func (r *versionCmd) version(c *cli.Context) error {
	apiURL := os.Getenv("API_URL")
	if apiURL == "" {
		apiURL = "http://localhost:8080"
	}

	u, err := url.Parse(apiURL)
	if err != nil {
		return err
	}
	r.Configuration.BasePath = u.String()

	fmt.Println("Client version:", Version)
	v, _, err := r.VersionGet()
	if err != nil {
		return err
	}
	fmt.Println("Server version", v.Version)
	return nil
}
<commit_msg>fn tool: 0.3.9 release [skip ci]<commit_after>package main

import (
	"fmt"
	"net/url"
	"os"

	functions "github.com/funcy/functions_go"
	"github.com/urfave/cli"
)

// Version of Functions CLI
var Version = "0.3.9"

func version() cli.Command {
	r := versionCmd{VersionApi: functions.NewVersionApi()}
	return cli.Command{
		Name:   "version",
		Usage:  "displays fn and functions daemon versions",
		Action: r.version,
	}
}

type versionCmd struct {
	*functions.VersionApi
}

func (r *versionCmd) version(c *cli.Context) error {
	apiURL := os.Getenv("API_URL")
	if apiURL == "" {
		apiURL = "http://localhost:8080"
	}

	u, err := url.Parse(apiURL)
	if err != nil {
		return err
	}
	r.Configuration.BasePath = u.String()

	fmt.Println("Client version:", Version)
	v, _, err := r.VersionGet()
	if err != nil {
		return err
	}
	fmt.Println("Server version", v.Version)
	return nil
}
<|endoftext|>


<commit_before>package models

const WEBSOCKET_BUFFER_SIZE = 1024 * 1024 * 32
const TCP_BUFFER_SIZE = 1024 * 128
<commit_msg>update<commit_after>package models

const WEBSOCKET_BUFFER_SIZE = 1024 * 1024 * 32
const TCP_BUFFER_SIZE = 1024 * 256
<|endoftext|>


<commit_before>package main

import (
	"crypto/sha256"
	"encoding/json"
	"flag"
	"fmt"
	"log"
	"net/http"
	"net/url"
	"os"
	"runtime"
	"strings"
	"time"

	"github.com/golang/groupcache"
	"github.com/pierrre/githubhook"
	"github.com/pierrre/imageserver"
	imageserver_cache "github.com/pierrre/imageserver/cache"
	imageserver_cache_groupcache "github.com/pierrre/imageserver/cache/groupcache"
	imageserver_cache_memory "github.com/pierrre/imageserver/cache/memory"
	imageserver_http "github.com/pierrre/imageserver/http"
	imageserver_http_nfntresize "github.com/pierrre/imageserver/http/nfntresize"
	imageserver_image "github.com/pierrre/imageserver/image"
	_ "github.com/pierrre/imageserver/image/bmp"
	imageserver_image_gamma "github.com/pierrre/imageserver/image/gamma"
	_ "github.com/pierrre/imageserver/image/gif"
	_ "github.com/pierrre/imageserver/image/jpeg"
	imageserver_image_nfntresize "github.com/pierrre/imageserver/image/nfntresize"
	_ "github.com/pierrre/imageserver/image/png"
	_ "github.com/pierrre/imageserver/image/tiff"
	imageserver_testdata "github.com/pierrre/imageserver/testdata"
)

const (
	groupcacheName = "imageserver"
)

var (
	flagHTTPAddr            = ":8080"
	flagGitHubWebhookSecret string
	flagHTTPExpires         = time.Duration(7 * 24 * time.Hour)
	flagGroupcache          = int64(512 * (1 << 20))
	flagGroupcacheSelf      string
	flagGroupcachePeers     string
	flagCacheMemory         = int64(64 * (1 << 20))
)

func main() {
	parseFlags()
	initLog()
	logEnv()
	startGroupcacheHTTPServer()
	startHTTPServer()
}

func parseFlags() {
	flag.StringVar(&flagHTTPAddr, "http", flagHTTPAddr, "HTTP addr")
	flag.DurationVar(&flagHTTPExpires, "http-expires", flagHTTPExpires, "HTTP expires")
	flag.StringVar(&flagGitHubWebhookSecret, "github-webhook-secret", flagGitHubWebhookSecret, "GitHub webhook secret")
	flag.Int64Var(&flagGroupcache, "groupcache", flagGroupcache, "Groupcache")
	flag.StringVar(&flagGroupcacheSelf, "groupcache-self", flagGroupcacheSelf, "Groupcache self")
	flag.StringVar(&flagGroupcachePeers, "groupcache-peers", flagGroupcachePeers, "Groupcache peers")
	flag.Int64Var(&flagCacheMemory, "cache-memory", flagCacheMemory, "Cache memory")
	flag.Parse()
}

func initLog() {
	log.SetFlags(log.Ldate | log.Ltime | log.Lmicroseconds)
	log.SetOutput(os.Stdout)
	log.Println("Start")
}

func logEnv() {
	log.Printf("Go version: %s", runtime.Version())
	log.Printf("Go max procs: %d", runtime.GOMAXPROCS(0))
}

func startHTTPServer() {
	log.Printf("Start HTTP server on %s", flagHTTPAddr)
	err := http.ListenAndServe(flagHTTPAddr, newHTTPHandler())
	if err != nil {
		log.Panic(err)
	}
}

func newHTTPHandler() http.Handler {
	mux := http.NewServeMux()
	mux.Handle("/", newImageHTTPHandler())
	mux.Handle("/favicon.ico", http.NotFoundHandler())
	if h := newGitHubWebhookHTTPHandler(); h != nil {
		mux.Handle("/github_webhook", h)
	}
	return mux
}

func newGitHubWebhookHTTPHandler() http.Handler {
	if flagGitHubWebhookSecret == "" {
		return nil
	}
	return &githubhook.Handler{
		Secret: flagGitHubWebhookSecret,
		Delivery: func(event string, deliveryID string, payload interface{}) {
			log.Printf("Received GitHub webhook: %s", event)
			if event == "push" {
				delay := time.Duration(5 * time.Second)
				log.Printf("Killing process in %s", delay)
				time.AfterFunc(delay, func() {
					log.Println("Killing process now")
					os.Exit(0)
				})
			}
		},
		Error: func(err error, req *http.Request) {
			log.Printf("GitHub webhook error: %s", err)
		},
	}
}

func newImageHTTPHandler() http.Handler {
	var handler http.Handler
	handler = &imageserver_http.Handler{
		Parser: &imageserver_http.ListParser{
			&imageserver_http.SourceParser{},
			&imageserver_http_nfntresize.Parser{},
			&imageserver_http.FormatParser{},
			&imageserver_http.QualityParser{},
			&imageserver_http.GammaCorrectionParser{},
		},
		Server:   newServer(),
		ETagFunc: imageserver_http.NewParamsHashETagFunc(sha256.New),
		ErrorFunc: func(err error, req *http.Request) {
			log.Printf("Internal error: %s", err)
		},
	}
	if flagHTTPExpires != 0 {
		handler = &imageserver_http.ExpiresHandler{
			Handler: handler,
			Expires: flagHTTPExpires,
		}
	}
	return handler
}

func newServer() imageserver.Server {
	srv := imageserver_testdata.Server
	srv = newServerImage(srv)
	srv = newServerLimit(srv)
	srv = newServerValidate(srv)
	srv = newServerGroupcache(srv)
	srv = newServerCacheMemory(srv)
	return srv
}

func newServerImage(srv imageserver.Server) imageserver.Server {
	return &imageserver_image.Server{
		Server: srv,
		Processor: imageserver_image_gamma.NewCorrectionProcessor(
			&imageserver_image_nfntresize.Processor{},
			true,
		),
	}
}

func newServerLimit(srv imageserver.Server) imageserver.Server {
	return imageserver.NewLimitServer(srv, runtime.GOMAXPROCS(0)*2)
}

func newServerValidate(srv imageserver.Server) imageserver.Server {
	return &imageserver_image_nfntresize.ValidateParamsServer{
		Server:    srv,
		WidthMax:  2048,
		HeightMax: 2048,
	}
}

func newServerGroupcache(srv imageserver.Server) imageserver.Server {
	if flagGroupcache <= 0 {
		return srv
	}
	return imageserver_cache_groupcache.NewServer(
		srv,
		imageserver_cache.NewParamsHashKeyGenerator(sha256.New),
		groupcacheName,
		flagGroupcache,
	)
}

func newServerCacheMemory(srv imageserver.Server) imageserver.Server {
	if flagCacheMemory <= 0 {
		return srv
	}
	return &imageserver_cache.Server{
		Server:       srv,
		Cache:        imageserver_cache_memory.New(flagCacheMemory),
		KeyGenerator: imageserver_cache.NewParamsHashKeyGenerator(sha256.New),
	}
}

func startGroupcacheHTTPServer() {
	if flagGroupcacheSelf == "" {
		return
	}
	go func() {
		log.Printf("Start groupcache HTTP server on %s", flagGroupcacheSelf)
		err := http.ListenAndServe(flagGroupcacheSelf, newGroupcacheHTTPHandler())
		if err != nil {
			log.Fatal(err)
		}
	}()
}

func newGroupcacheHTTPHandler() http.Handler {
	mux := http.NewServeMux()
	mux.Handle("/", newGroupcacheHTTPPool())
	mux.HandleFunc("/stats", groupcacheStatsHTTPHandler)
	mux.Handle("/favicon.ico", http.NotFoundHandler())
	return mux
}

func newGroupcacheHTTPPool() *groupcache.HTTPPool {
	self := (&url.URL{Scheme: "http", Host: flagGroupcacheSelf}).String()
	var peers []string
	peers = append(peers, self)
	for _, p := range strings.Split(flagGroupcachePeers, ",") {
		if p == "" {
			continue
		}
		peer := (&url.URL{Scheme: "http", Host: p}).String()
		peers = append(peers, peer)
	}
	opts := &groupcache.HTTPPoolOptions{BasePath: "/"}
	pool := groupcache.NewHTTPPoolOpts(self, opts)
	pool.Context = imageserver_cache_groupcache.HTTPPoolContext
	pool.Transport = imageserver_cache_groupcache.NewHTTPPoolTransport(nil)
	pool.Set(peers...)
	return pool
}

func groupcacheStatsHTTPHandler(w http.ResponseWriter, req *http.Request) {
	gp := groupcache.GetGroup(groupcacheName)
	if gp == nil {
		http.Error(w, fmt.Sprintf("group %s not found", groupcacheName), http.StatusServiceUnavailable)
		return
	}
	type cachesStats struct {
		Main groupcache.CacheStats
		Hot  groupcache.CacheStats
	}
	type stats struct {
		Group  groupcache.Stats
		Caches cachesStats
	}
	data, err := json.MarshalIndent(
		stats{
			Group: gp.Stats,
			Caches: cachesStats{
				Main: gp.CacheStats(groupcache.MainCache),
				Hot:  gp.CacheStats(groupcache.HotCache),
			},
		},
		"",
		"	",
	)
	if err != nil {
		log.Println(err)
		http.Error(w, http.StatusText(http.StatusInternalServerError), http.StatusInternalServerError)
		return
	}
	_, err = w.Write(data)
	if err != nil {
		log.Println(err)
		return
	}
}
<commit_msg>add HTTP log in advanced example<commit_after>package main

import (
	"crypto/sha256"
	"encoding/json"
	"flag"
	"fmt"
	"log"
	"net/http"
	"net/url"
	"os"
	"runtime"
	"strings"
	"time"

	"github.com/golang/groupcache"
	"github.com/pierrre/githubhook"
	"github.com/pierrre/imageserver"
	imageserver_cache "github.com/pierrre/imageserver/cache"
	imageserver_cache_groupcache "github.com/pierrre/imageserver/cache/groupcache"
	imageserver_cache_memory "github.com/pierrre/imageserver/cache/memory"
	imageserver_http "github.com/pierrre/imageserver/http"
	imageserver_http_nfntresize "github.com/pierrre/imageserver/http/nfntresize"
	imageserver_image "github.com/pierrre/imageserver/image"
	_ "github.com/pierrre/imageserver/image/bmp"
	imageserver_image_gamma "github.com/pierrre/imageserver/image/gamma"
	_ "github.com/pierrre/imageserver/image/gif"
	_ "github.com/pierrre/imageserver/image/jpeg"
	imageserver_image_nfntresize "github.com/pierrre/imageserver/image/nfntresize"
	_ "github.com/pierrre/imageserver/image/png"
	_ "github.com/pierrre/imageserver/image/tiff"
	imageserver_testdata "github.com/pierrre/imageserver/testdata"
)

const (
	groupcacheName = "imageserver"
)

var (
	flagHTTPAddr            = ":8080"
	flagGitHubWebhookSecret string
	flagHTTPExpires         = time.Duration(7 * 24 * time.Hour)
	flagGroupcache          = int64(512 * (1 << 20))
	flagGroupcacheSelf      string
	flagGroupcachePeers     string
	flagCacheMemory         = int64(64 * (1 << 20))
)

func main() {
	parseFlags()
	initLog()
	logEnv()
	startGroupcacheHTTPServer()
	startHTTPServer()
}

func parseFlags() {
	flag.StringVar(&flagHTTPAddr, "http", flagHTTPAddr, "HTTP addr")
	flag.DurationVar(&flagHTTPExpires, "http-expires", flagHTTPExpires, "HTTP expires")
	flag.StringVar(&flagGitHubWebhookSecret, "github-webhook-secret", flagGitHubWebhookSecret, "GitHub webhook secret")
	flag.Int64Var(&flagGroupcache, "groupcache", flagGroupcache, "Groupcache")
	flag.StringVar(&flagGroupcacheSelf, "groupcache-self", flagGroupcacheSelf, "Groupcache self")
	flag.StringVar(&flagGroupcachePeers, "groupcache-peers", flagGroupcachePeers, "Groupcache peers")
	flag.Int64Var(&flagCacheMemory, "cache-memory", flagCacheMemory, "Cache memory")
	flag.Parse()
}

func initLog() {
	log.SetFlags(log.Ldate | log.Ltime | log.Lmicroseconds)
	log.SetOutput(os.Stdout)
	log.Println("Start")
}

func logEnv() {
	log.Printf("Go version: %s", runtime.Version())
	log.Printf("Go max procs: %d", runtime.GOMAXPROCS(0))
}

func startHTTPServer() {
	log.Printf("Start HTTP server on %s", flagHTTPAddr)
	err := http.ListenAndServe(flagHTTPAddr, newHTTPHandler())
	if err != nil {
		log.Panic(err)
	}
}

func newHTTPHandler() http.Handler {
	mux := http.NewServeMux()
	mux.Handle("/", newImageHTTPHandler())
	mux.Handle("/favicon.ico", http.NotFoundHandler())
	if h := newGitHubWebhookHTTPHandler(); h != nil {
		mux.Handle("/github_webhook", h)
	}
	h := newLoggerHTTPHandler(mux)
	return h
}

func newLoggerHTTPHandler(h http.Handler) http.Handler {
	return http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
		lrw := &logResponseWriter{ResponseWriter: rw}
		start := time.Now()
		h.ServeHTTP(lrw, req)
		end := time.Now()
		dur := end.Sub(start)
		log.Printf(
			"HTTP %s %s (%s %s) => %d %d %s",
			req.Method, req.URL,
			req.RemoteAddr, req.UserAgent(),
			lrw.Code, lrw.Size, dur,
		)
	})
}

type logResponseWriter struct {
	http.ResponseWriter
	Code int
	Size int
}

func (rw *logResponseWriter) WriteHeader(code int) {
	rw.ResponseWriter.WriteHeader(code)
	rw.Code = code
}

func (rw *logResponseWriter) Write(b []byte) (int, error) {
	if rw.Code == 0 {
		rw.WriteHeader(http.StatusOK)
	}
	size, err := rw.ResponseWriter.Write(b)
	if err == nil {
		rw.Size += size
	}
	return size, err
}

func newGitHubWebhookHTTPHandler() http.Handler {
	if flagGitHubWebhookSecret == "" {
		return nil
	}
	return &githubhook.Handler{
		Secret: flagGitHubWebhookSecret,
		Delivery: func(event string, deliveryID string, payload interface{}) {
			log.Printf("Received GitHub webhook: %s", event)
			if event == "push" {
				delay := time.Duration(5 * time.Second)
				log.Printf("Killing process in %s", delay)
				time.AfterFunc(delay, func() {
					log.Println("Killing process now")
					os.Exit(0)
				})
			}
		},
		Error: func(err error, req *http.Request) {
			log.Printf("GitHub webhook error: %s", err)
		},
	}
}

func newImageHTTPHandler() http.Handler {
	var handler http.Handler
	handler = &imageserver_http.Handler{
		Parser: &imageserver_http.ListParser{
			&imageserver_http.SourceParser{},
			&imageserver_http_nfntresize.Parser{},
			&imageserver_http.FormatParser{},
			&imageserver_http.QualityParser{},
			&imageserver_http.GammaCorrectionParser{},
		},
		Server:   newServer(),
		ETagFunc: imageserver_http.NewParamsHashETagFunc(sha256.New),
		ErrorFunc: func(err error, req *http.Request) {
			log.Printf("Internal error: %s", err)
		},
	}
	if flagHTTPExpires != 0 {
		handler = &imageserver_http.ExpiresHandler{
			Handler: handler,
			Expires: flagHTTPExpires,
		}
	}
	return handler
}

func newServer() imageserver.Server {
	srv := imageserver_testdata.Server
	srv = newServerImage(srv)
	srv = newServerLimit(srv)
	srv = newServerValidate(srv)
	srv = newServerGroupcache(srv)
	srv = newServerCacheMemory(srv)
	return srv
}

func newServerImage(srv imageserver.Server) imageserver.Server {
	return &imageserver_image.Server{
		Server: srv,
		Processor: imageserver_image_gamma.NewCorrectionProcessor(
			&imageserver_image_nfntresize.Processor{},
			true,
		),
	}
}

func newServerLimit(srv imageserver.Server) imageserver.Server {
	return imageserver.NewLimitServer(srv, runtime.GOMAXPROCS(0)*2)
}

func newServerValidate(srv imageserver.Server) imageserver.Server {
	return &imageserver_image_nfntresize.ValidateParamsServer{
		Server:    srv,
		WidthMax:  2048,
		HeightMax: 2048,
	}
}

func newServerGroupcache(srv imageserver.Server) imageserver.Server {
	if flagGroupcache <= 0 {
		return srv
	}
	return imageserver_cache_groupcache.NewServer(
		srv,
		imageserver_cache.NewParamsHashKeyGenerator(sha256.New),
		groupcacheName,
		flagGroupcache,
	)
}

func newServerCacheMemory(srv imageserver.Server) imageserver.Server {
	if flagCacheMemory <= 0 {
		return srv
	}
	return &imageserver_cache.Server{
		Server:       srv,
		Cache:        imageserver_cache_memory.New(flagCacheMemory),
		KeyGenerator: imageserver_cache.NewParamsHashKeyGenerator(sha256.New),
	}
}

func startGroupcacheHTTPServer() {
	if flagGroupcacheSelf == "" {
		return
	}
	go func() {
		log.Printf("Start groupcache HTTP server on %s", flagGroupcacheSelf)
		err := http.ListenAndServe(flagGroupcacheSelf, newGroupcacheHTTPHandler())
		if err != nil {
			log.Fatal(err)
		}
	}()
}

func newGroupcacheHTTPHandler() http.Handler {
	mux := http.NewServeMux()
	mux.Handle("/", newGroupcacheHTTPPool())
	mux.HandleFunc("/stats", groupcacheStatsHTTPHandler)
	mux.Handle("/favicon.ico", http.NotFoundHandler())
	return mux
}

func newGroupcacheHTTPPool() *groupcache.HTTPPool {
	self := (&url.URL{Scheme: "http", Host: flagGroupcacheSelf}).String()
	var peers []string
	peers = append(peers, self)
	for _, p := range strings.Split(flagGroupcachePeers, ",") {
		if p == "" {
			continue
		}
		peer := (&url.URL{Scheme: "http", Host: p}).String()
		peers = append(peers, peer)
	}
	opts := &groupcache.HTTPPoolOptions{BasePath: "/"}
	pool := groupcache.NewHTTPPoolOpts(self, opts)
	pool.Context = imageserver_cache_groupcache.HTTPPoolContext
	pool.Transport = imageserver_cache_groupcache.NewHTTPPoolTransport(nil)
	pool.Set(peers...)
	return pool
}

func groupcacheStatsHTTPHandler(w http.ResponseWriter, req *http.Request) {
	gp := groupcache.GetGroup(groupcacheName)
	if gp == nil {
		http.Error(w, fmt.Sprintf("group %s not found", groupcacheName), http.StatusServiceUnavailable)
		return
	}
	type cachesStats struct {
		Main groupcache.CacheStats
		Hot  groupcache.CacheStats
	}
	type stats struct {
		Group  groupcache.Stats
		Caches cachesStats
	}
	data, err := json.MarshalIndent(
		stats{
			Group: gp.Stats,
			Caches: cachesStats{
				Main: gp.CacheStats(groupcache.MainCache),
				Hot:  gp.CacheStats(groupcache.HotCache),
			},
		},
		"",
		"	",
	)
	if err != nil {
		log.Println(err)
		http.Error(w, http.StatusText(http.StatusInternalServerError), http.StatusInternalServerError)
		return
	}
	_, err = w.Write(data)
	if err != nil {
		log.Println(err)
		return
	}
}
<|endoftext|>


<commit_before>package log

import (
	"testing"

	"github.com/stretchr/testify/assert"
)

func TestWaitingTaskDisplaysWaitingStatus(t *testing.T) {
	task := NewWaitingTask("example")

	assert.Equal(t, "example: ...", <-task.Updates())
}

func TestWaitingTaskCallsDoneWhenComplete(t *testing.T) {
	task := NewWaitingTask("example")

	select {
	case v, ok := <-task.Updates():
		if ok {
			assert.Equal(t, "example: ...", v)
		} else {
			t.Fatal("expected channel to be open")
		}
	default:
	}

	task.Complete()

	if _, ok := <-task.Updates(); ok {
		t.Fatalf("expected channel to be closed")
	}
}

func TestWaitingTaskPanicsWithMultipleDoneCalls(t *testing.T) {
	task := NewWaitingTask("example")

	task.Complete()

	defer func() {
		if err := recover(); err == nil {
			t.Fatal("githistory/log: expected panic()")
		} else {
			if s, ok := err.(error); ok {
				assert.Equal(t, "close of closed channel", s.Error())
			} else {
				t.Fatal("githistory/log: expected panic() to implement error")
			}
		}
	}()

	task.Complete()
}
<commit_msg>git/githistory/log: test WaitingTask Throttle()-d<commit_after>package log

import (
	"testing"

	"github.com/stretchr/testify/assert"
)

func TestWaitingTaskDisplaysWaitingStatus(t *testing.T) {
	task := NewWaitingTask("example")

	assert.Equal(t, "example: ...", <-task.Updates())
}

func TestWaitingTaskCallsDoneWhenComplete(t *testing.T) {
	task := NewWaitingTask("example")

	select {
	case v, ok := <-task.Updates():
		if ok {
			assert.Equal(t, "example: ...", v)
		} else {
			t.Fatal("expected channel to be open")
		}
	default:
	}

	task.Complete()

	if _, ok := <-task.Updates(); ok {
		t.Fatalf("expected channel to be closed")
	}
}

func TestWaitingTaskPanicsWithMultipleDoneCalls(t *testing.T) {
	task := NewWaitingTask("example")

	task.Complete()

	defer func() {
		if err := recover(); err == nil {
			t.Fatal("githistory/log: expected panic()")
		} else {
			if s, ok := err.(error); ok {
				assert.Equal(t, "close of closed channel", s.Error())
			} else {
				t.Fatal("githistory/log: expected panic() to implement error")
			}
		}
	}()

	task.Complete()
}

func TestWaitingTaskIsThrottled(t *testing.T) {
	task := NewWaitingTask("example")

	throttled := task.Throttled()

	assert.True(t, throttled,
		"git/githistory/log: expected *WaitingTask to be Throttle()-d")
}
<|endoftext|>


<commit_before>// Copyright 2011 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.

// +build !appengine,!secure

package main

import (
	"flag"
	"fmt"
	"go/build"
	"io"
	"log"
	"net"
	"net/http"
	"net/url"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"strings"
	"time"

	"github.com/Go-zh/tools/playground/socket"

	// Imports so that go build/install automatically installs them.
	_ "github.com/Go-zh/tour/pic"
	_ "github.com/Go-zh/tour/tree"
	_ "github.com/Go-zh/tour/wc"
)

const (
	basePkg    = "github.com/Go-zh/tour"
	socketPath = "/socket"
)

var (
	httpListen  = flag.String("http", "127.0.0.1:3999", "host:port to listen on")
	openBrowser = flag.Bool("openbrowser", true, "open browser automatically")
)

var (
	// GOPATH containing the tour packages
	gopath = os.Getenv("GOPATH")

	httpAddr string
)

// isRoot reports whether path is the root directory of the tour tree.
// To be the root, it must have content and template subdirectories.
func isRoot(path string) bool {
	_, err := os.Stat(filepath.Join(path, "content", "welcome.article"))
	if err == nil {
		_, err = os.Stat(filepath.Join(path, "template", "index.tmpl"))
	}
	return err == nil
}

func findRoot() (string, error) {
	ctx := build.Default
	p, err := ctx.Import(basePkg, "", build.FindOnly)
	if err == nil && isRoot(p.Dir) {
		return p.Dir, nil
	}
	tourRoot := filepath.Join(runtime.GOROOT(), "misc", "tour")
	ctx.GOPATH = tourRoot
	p, err = ctx.Import(basePkg, "", build.FindOnly)
	if err == nil && isRoot(tourRoot) {
		gopath = tourRoot
		return tourRoot, nil
	}
	return "", fmt.Errorf("could not find go-tour content; check $GOROOT and $GOPATH")
}

func main() {
	flag.Parse()

	// find and serve the go tour files
	root, err := findRoot()
	if err != nil {
		log.Fatalf("Couldn't find tour files: %v", err)
	}

	log.Println("Serving content from", root)

	host, port, err := net.SplitHostPort(*httpListen)
	if err != nil {
		log.Fatal(err)
	}
	if host == "" {
		host = "localhost"
	}
	if host != "127.0.0.1" && host != "localhost" {
		log.Print(localhostWarning)
	}
	httpAddr = host + ":" + port

	if err := initTour(root, "SocketTransport"); err != nil {
		log.Fatal(err)
	}

	http.HandleFunc("/", rootHandler)
	http.HandleFunc("/lesson/", lessonHandler)

	origin := &url.URL{Scheme: "http", Host: host + ":" + port}
	http.Handle(socketPath, socket.NewHandler(origin))

	// Keep these static file handlers in sync with ../app.yaml.
	static := http.FileServer(http.Dir(root))
	http.Handle("/content/img/", static)
	http.Handle("/static/", static)
	imgDir := filepath.Join(root, "static", "img")
	http.Handle("/favicon.ico", http.FileServer(http.Dir(imgDir)))

	go func() {
		url := "http://" + httpAddr
		if waitServer(url) && *openBrowser && startBrowser(url) {
			log.Printf("A browser window should open. If not, please visit %s", url)
		} else {
			log.Printf("Please open your web browser and visit %s", url)
		}
	}()
	log.Fatal(http.ListenAndServe(httpAddr, nil))
}

// rootHandler returns a handler for all the requests except the ones for lessons.
func rootHandler(w http.ResponseWriter, r *http.Request) {
	if err := renderUI(w); err != nil {
		log.Println(err)
	}
}

// lessonHandler handler the HTTP requests for lessons.
func lessonHandler(w http.ResponseWriter, r *http.Request) {
	lesson := strings.TrimPrefix(r.URL.Path, "/lesson/")
	if err := writeLesson(lesson, w); err != nil {
		if err == lessonNotFound {
			http.NotFound(w, r)
		} else {
			log.Println(err)
		}
	}
}

const localhostWarning = `
WARNING!  WARNING!  WARNING!

I appear to be listening on an address that is not localhost.
Anyone with access to this address and port will have access
to this machine as the user running gotour.

If you don't understand this message, hit Control-C to terminate this process.

WARNING!  WARNING!  WARNING!
`

type response struct {
	Output string `json:"output"`
	Errors string `json:"compile_errors"`
}

func init() {
	socket.Environ = environ
}

// environ returns the original execution environment with GOPATH
// replaced (or added) with the value of the global var gopath.
func environ() (env []string) {
	for _, v := range os.Environ() {
		if !strings.HasPrefix(v, "GOPATH=") {
			env = append(env, v)
		}
	}
	env = append(env, "GOPATH="+gopath)
	return
}

// waitServer waits some time for the http Server to start
// serving url. The return value reports whether it starts.
func waitServer(url string) bool {
	tries := 20
	for tries > 0 {
		resp, err := http.Get(url)
		if err == nil {
			resp.Body.Close()
			return true
		}
		time.Sleep(100 * time.Millisecond)
		tries--
	}
	return false
}

// startBrowser tries to open the URL in a browser, and returns
// whether it succeed.
func startBrowser(url string) bool {
	// try to start the browser
	var args []string
	switch runtime.GOOS {
	case "darwin":
		args = []string{"open"}
	case "windows":
		args = []string{"cmd", "/c", "start"}
	default:
		args = []string{"xdg-open"}
	}
	cmd := exec.Command(args[0], append(args[1:], url)...)
	return cmd.Start() == nil
}

// prepContent for the local tour simply returns the content as-is.
func prepContent(r io.Reader) io.Reader { return r }

// socketAddr returns the WebSocket handler address.
func socketAddr() string { return "ws://" + httpAddr + socketPath }
<commit_msg>gotour: use vendored copy of golang.org/x/tour packages<commit_after>// Copyright 2011 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.

// +build !appengine,!secure

package main

import (
	"flag"
	"fmt"
	"go/build"
	"io"
	"log"
	"net"
	"net/http"
	"net/url"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"strings"
	"time"

	"github.com/Go-zh/tools/playground/socket"

	// Imports so that go build/install automatically installs them.
	// For Go 1.5 and above, will use our vendored copy.
	_ "golang.org/x/tour/pic"
	_ "golang.org/x/tour/tree"
	_ "golang.org/x/tour/wc"
)

const (
	basePkg    = "github.com/Go-zh/tour"
	socketPath = "/socket"
)

var (
	httpListen  = flag.String("http", "127.0.0.1:3999", "host:port to listen on")
	openBrowser = flag.Bool("openbrowser", true, "open browser automatically")
)

var (
	// GOPATH containing the tour packages
	gopath = os.Getenv("GOPATH")

	httpAddr string
)

// isRoot reports whether path is the root directory of the tour tree.
// To be the root, it must have content and template subdirectories.
func isRoot(path string) bool {
	_, err := os.Stat(filepath.Join(path, "content", "welcome.article"))
	if err == nil {
		_, err = os.Stat(filepath.Join(path, "template", "index.tmpl"))
	}
	return err == nil
}

func findRoot() (string, error) {
	ctx := build.Default
	p, err := ctx.Import(basePkg, "", build.FindOnly)
	if err == nil && isRoot(p.Dir) {
		return p.Dir, nil
	}
	tourRoot := filepath.Join(runtime.GOROOT(), "misc", "tour")
	ctx.GOPATH = tourRoot
	p, err = ctx.Import(basePkg, "", build.FindOnly)
	if err == nil && isRoot(tourRoot) {
		gopath = tourRoot
		return tourRoot, nil
	}
	return "", fmt.Errorf("could not find go-tour content; check $GOROOT and $GOPATH")
}

func main() {
	flag.Parse()

	// find and serve the go tour files
	root, err := findRoot()
	if err != nil {
		log.Fatalf("Couldn't find tour files: %v", err)
	}

	log.Println("Serving content from", root)

	host, port, err := net.SplitHostPort(*httpListen)
	if err != nil {
		log.Fatal(err)
	}
	if host == "" {
		host = "localhost"
	}
	if host != "127.0.0.1" && host != "localhost" {
		log.Print(localhostWarning)
	}
	httpAddr = host + ":" + port

	if err := initTour(root, "SocketTransport"); err != nil {
		log.Fatal(err)
	}

	http.HandleFunc("/", rootHandler)
	http.HandleFunc("/lesson/", lessonHandler)

	origin := &url.URL{Scheme: "http", Host: host + ":" + port}
	http.Handle(socketPath, socket.NewHandler(origin))

	// Keep these static file handlers in sync with ../app.yaml.
	static := http.FileServer(http.Dir(root))
	http.Handle("/content/img/", static)
	http.Handle("/static/", static)
	imgDir := filepath.Join(root, "static", "img")
	http.Handle("/favicon.ico", http.FileServer(http.Dir(imgDir)))

	go func() {
		url := "http://" + httpAddr
		if waitServer(url) && *openBrowser && startBrowser(url) {
			log.Printf("A browser window should open. If not, please visit %s", url)
		} else {
			log.Printf("Please open your web browser and visit %s", url)
		}
	}()
	log.Fatal(http.ListenAndServe(httpAddr, nil))
}

// rootHandler returns a handler for all the requests except the ones for lessons.
func rootHandler(w http.ResponseWriter, r *http.Request) {
	if err := renderUI(w); err != nil {
		log.Println(err)
	}
}

// lessonHandler handler the HTTP requests for lessons.
func lessonHandler(w http.ResponseWriter, r *http.Request) {
	lesson := strings.TrimPrefix(r.URL.Path, "/lesson/")
	if err := writeLesson(lesson, w); err != nil {
		if err == lessonNotFound {
			http.NotFound(w, r)
		} else {
			log.Println(err)
		}
	}
}

const localhostWarning = `
WARNING!  WARNING!  WARNING!

I appear to be listening on an address that is not localhost.
Anyone with access to this address and port will have access
to this machine as the user running gotour.

If you don't understand this message, hit Control-C to terminate this process.

WARNING!  WARNING!  WARNING!
`

type response struct {
	Output string `json:"output"`
	Errors string `json:"compile_errors"`
}

func init() {
	socket.Environ = environ
}

// environ returns the original execution environment with GOPATH
// replaced (or added) with the value of the global var gopath.
func environ() (env []string) {
	for _, v := range os.Environ() {
		if !strings.HasPrefix(v, "GOPATH=") {
			env = append(env, v)
		}
	}
	env = append(env, "GOPATH="+gopath)
	return
}

// waitServer waits some time for the http Server to start
// serving url. The return value reports whether it starts.
func waitServer(url string) bool {
	tries := 20
	for tries > 0 {
		resp, err := http.Get(url)
		if err == nil {
			resp.Body.Close()
			return true
		}
		time.Sleep(100 * time.Millisecond)
		tries--
	}
	return false
}

// startBrowser tries to open the URL in a browser, and returns
// whether it succeed.
func startBrowser(url string) bool {
	// try to start the browser
	var args []string
	switch runtime.GOOS {
	case "darwin":
		args = []string{"open"}
	case "windows":
		args = []string{"cmd", "/c", "start"}
	default:
		args = []string{"xdg-open"}
	}
	cmd := exec.Command(args[0], append(args[1:], url)...)
	return cmd.Start() == nil
}

// prepContent for the local tour simply returns the content as-is.
func prepContent(r io.Reader) io.Reader { return r }

// socketAddr returns the WebSocket handler address.
func socketAddr() string { return "ws://" + httpAddr + socketPath }
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"log"

	eio "engine.io"
)

func main() {
	flag.Parse()
	server := eio.NewEngineBuilder().Build()
	server.OnConnect(func(socket eio.Socket) {
		socket.Send("Hello world!")
		socket.OnMessage(func(data []byte) {
			log.Println("recieve:", string(data))
		})
		socket.OnClose(func(reason string) {
			log.Println("socket closed: ", socket.Id())
		})
	})
	log.Fatalln(server.Listen(":3000"))
}
<commit_msg>这个人很懒, 向你丢了一只随机皮卡丘(战斗力: 6768)!<commit_after><|endoftext|>


<commit_before>package bongo

import (
	"encoding/json"
	"errors"
	"fmt"
	"strconv"
	"strings"

	"github.com/jinzhu/gorm"
)

// Fetch fetches the data from db by given parameters(fields of the struct)
func (b *Bongo) Fetch(i Modellable) error {
	if i.GetId() == 0 {
		return errors.New(fmt.Sprintf("Id is not set for %s", i.TableName()))
	}

	if err := b.DB.Table(i.TableName()).
		Find(i).
		Error; err != nil {
		return err
	}

	return nil
}

// ById Fetches data from db by it's id
func (b *Bongo) ById(i Modellable, id int64) error {
	if err := b.DB.
		Table(i.TableName()).
		Where("id = ?", id).
		Find(i).
		Error; err != nil {
		return err
	}

	return nil
}

// Creates a new record with the given struct and its fields
func (b *Bongo) Create(i Modellable) error {
	if err := b.DB.Save(i).Error; err != nil {
		return err
	}

	return nil
}

// Update updates all fields of a struct with assigned data
func (b *Bongo) Update(i Modellable) error {
	if i.GetId() == 0 {
		return errors.New(fmt.Sprintf("Id is not set for %s", i.TableName()))
	}

	// Update and Create is using the Save method, so they are
	// same functions but GORM handles, AfterCreate and AfterUpdate
	// in correct manner
	return b.Create(i)
}

func (b *Bongo) Delete(i Modellable) error {
	if i.GetId() == 0 {
		return errors.New(fmt.Sprintf("Id is not set for %s", i.TableName()))
	}

	if err := b.DB.Delete(i).Error; err != nil {
		return err
	}

	return nil
}

func (b *Bongo) FetchByIds(i Modellable, data interface{}, ids []int64) error {
	if len(ids) == 0 {
		return nil
	}

	orderByQuery := ""
	comma := ""
	for _, id := range ids {
		orderByQuery = orderByQuery + comma + " id = " + strconv.FormatInt(id, 10) + " desc"
		comma = ","
	}
	return b.DB.
		Table(i.TableName()).
		Order(orderByQuery).
		Where(ids).
		Find(data).
		Error

}

func (b *Bongo) UpdatePartial(i Modellable, set map[string]interface{}) error {
	if i.GetId() == 0 {
		return errors.New(fmt.Sprintf("Id is not set for %s", i.TableName()))
	}

	query := b.DB.Table(i.TableName())

	query = query.Where(i.GetId())

	if err := query.Update(set).Error; err != nil {
		return err
	}

	if err := b.Fetch(i); err != nil {
		return err
	}

	b.AfterUpdate(i)
	return nil
}

// selector, set
func (b *Bongo) UpdateMulti(i Modellable, rest ...map[string]interface{}) error {
	var set, selector map[string]interface{}

	switch len(rest) {
	case 1:
		set = rest[0]
		selector = nil
	case 2:
		selector = rest[0]
		set = rest[1]
	default:
		return errors.New("Update partial parameter list is wrong")
	}

	query := b.DB.Table(i.TableName())

	//add selector
	query = addWhere(query, selector)

	if err := query.Update(set).Error; err != nil {
		return err
	}

	return nil
}

func (b *Bongo) Count(i Modellable, where ...interface{}) (int, error) {
	var count int

	// init query
	query := b.DB

	// add table name
	query = query.Table(i.TableName())

	// add query
	query = query.Where(where[0], where[1:len(where)]...)

	return count, query.Count(&count).Error
}

func (b *Bongo) CountWithQuery(i Modellable, q *Query) (int, error) {
	query := b.BuildQuery(i, q)
	var count int
	return count, query.Count(&count).Error
}

type Query struct {
	Selector   map[string]interface{}
	Sort       map[string]string
	Pluck      string
	Pagination Pagination
}

type Pagination struct {
	Limit int
	Skip  int
}

func NewPagination(limit int, skip int) *Pagination {
	return &Pagination{
		Limit: limit,
		Skip:  skip,
	}
}

func NewQS(selector map[string]interface{}) *Query {
	return &Query{
		Selector: selector,
	}
}

// selector, sort, limit, pluck,
func (b *Bongo) Some(i Modellable, data interface{}, q *Query) error {
	err := b.executeQuery(i, data, q)
	if err == gorm.RecordNotFound {
		return nil
	}
	return err
}

func (b *Bongo) One(i Modellable, data interface{}, q *Query) error {
	q.Pagination.Limit = 1
	return b.executeQuery(i, data, q)
}

func (b *Bongo) BuildQuery(i Modellable, q *Query) *gorm.DB {
	// init query
	query := b.DB

	// add table name
	query = query.Table(i.TableName())

	// add sort options
	query = addSort(query, q.Sort)

	query = addSkip(query, q.Pagination.Skip)

	query = addLimit(query, q.Pagination.Limit)

	// add selector
	query = addWhere(query, q.Selector)

	return query
}

func (b *Bongo) executeQuery(i Modellable, data interface{}, q *Query) error {
	// init query
	query := b.BuildQuery(i, q)

	var err error
	// TODO refactor this part
	if q.Pluck != "" {
		if strings.Contains(q.Pluck, ",") {
			// add pluck data
			query = addPluck(query, q.Pluck)

			err = query.Find(data).Error
		} else {
			err = query.Pluck(q.Pluck, data).Error
		}
	} else {
		err = query.Find(data).Error
	}

	return err
}

func (b *Bongo) PublishEvent(eventName string, i Modellable) error {
	data, err := json.Marshal(i)
	if err != nil {
		b.log.Error("Error while marshalling for publish %s", err)
		return err
	}

	err = b.Broker.Publish(i.TableName()+"_"+eventName, data)
	if err != nil {
		b.log.Error("Error while publishing %s", err)
		return err
	}

	return nil
}

func (b *Bongo) AfterCreate(i Modellable) {
	b.PublishEvent("created", i)
}

func (b *Bongo) AfterUpdate(i Modellable) {
	b.PublishEvent("updated", i)
}

func (b *Bongo) AfterDelete(i Modellable) {
	b.PublishEvent("deleted", i)
}

func addSort(query *gorm.DB, options map[string]string) *gorm.DB {

	if options == nil {
		return query
	}

	if len(options) == 0 {
		return query
	}

	var opts []string
	for key, val := range options {
		opts = append(opts, fmt.Sprintf("%s %v", key, val))
	}
	return query.Order(strings.Join(opts, ","))
}

func addPluck(query *gorm.DB, plucked string) *gorm.DB {
	if plucked == "" {
		return query
	}

	return query.Select(plucked)
}

func addWhere(query *gorm.DB, selector map[string]interface{}) *gorm.DB {
	if selector == nil {
		return query
	}

	// instead sending one selector, do chaining here
	return query.Where(selector)
}

func addSkip(query *gorm.DB, skip int) *gorm.DB {
	if skip > 0 {
		return query.Offset(skip)
	}

	return query
}

func addLimit(query *gorm.DB, limit int) *gorm.DB {
	// if limit is minus or 0 ignore
	if limit > 0 {
		return query.Limit(limit)
	}

	return query
}
<commit_msg>Social: add a new function for checking errors, which will suppress the recordNotFound error<commit_after>package bongo

import (
	"encoding/json"
	"errors"
	"fmt"
	"strconv"
	"strings"

	"github.com/jinzhu/gorm"
)

// Fetch fetches the data from db by given parameters(fields of the struct)
func (b *Bongo) Fetch(i Modellable) error {
	if i.GetId() == 0 {
		return errors.New(fmt.Sprintf("Id is not set for %s", i.TableName()))
	}

	if err := b.DB.Table(i.TableName()).
		Find(i).
		Error; err != nil {
		return err
	}

	return nil
}

// ById Fetches data from db by it's id
func (b *Bongo) ById(i Modellable, id int64) error {
	if err := b.DB.
		Table(i.TableName()).
		Where("id = ?", id).
		Find(i).
		Error; err != nil {
		return err
	}

	return nil
}

// Creates a new record with the given struct and its fields
func (b *Bongo) Create(i Modellable) error {
	if err := b.DB.Save(i).Error; err != nil {
		return err
	}

	return nil
}

// Update updates all fields of a struct with assigned data
func (b *Bongo) Update(i Modellable) error {
	if i.GetId() == 0 {
		return errors.New(fmt.Sprintf("Id is not set for %s", i.TableName()))
	}

	// Update and Create is using the Save method, so they are
	// same functions but GORM handles, AfterCreate and AfterUpdate
	// in correct manner
	return b.Create(i)
}

func (b *Bongo) Delete(i Modellable) error {
	if i.GetId() == 0 {
		return errors.New(fmt.Sprintf("Id is not set for %s", i.TableName()))
	}

	if err := b.DB.Delete(i).Error; err != nil {
		return err
	}

	return nil
}

func (b *Bongo) FetchByIds(i Modellable, data interface{}, ids []int64) error {
	if len(ids) == 0 {
		return nil
	}

	orderByQuery := ""
	comma := ""
	for _, id := range ids {
		orderByQuery = orderByQuery + comma + " id = " + strconv.FormatInt(id, 10) + " desc"
		comma = ","
	}
	return b.DB.
		Table(i.TableName()).
		Order(orderByQuery).
		Where(ids).
		Find(data).
		Error

}

func (b *Bongo) UpdatePartial(i Modellable, set map[string]interface{}) error {
	if i.GetId() == 0 {
		return errors.New(fmt.Sprintf("Id is not set for %s", i.TableName()))
	}

	query := b.DB.Table(i.TableName())

	query = query.Where(i.GetId())

	if err := query.Update(set).Error; err != nil {
		return err
	}

	if err := b.Fetch(i); err != nil {
		return err
	}

	b.AfterUpdate(i)
	return nil
}

// selector, set
func (b *Bongo) UpdateMulti(i Modellable, rest ...map[string]interface{}) error {
	var set, selector map[string]interface{}

	switch len(rest) {
	case 1:
		set = rest[0]
		selector = nil
	case 2:
		selector = rest[0]
		set = rest[1]
	default:
		return errors.New("Update partial parameter list is wrong")
	}

	query := b.DB.Table(i.TableName())

	//add selector
	query = addWhere(query, selector)

	if err := query.Update(set).Error; err != nil {
		return err
	}

	return nil
}

func (b *Bongo) Count(i Modellable, where ...interface{}) (int, error) {
	var count int

	// init query
	query := b.DB

	// add table name
	query = query.Table(i.TableName())

	// add query
	query = query.Where(where[0], where[1:len(where)]...)

	return count, query.Count(&count).Error
}

func (b *Bongo) CountWithQuery(i Modellable, q *Query) (int, error) {
	query := b.BuildQuery(i, q)
	var count int
	return count, query.Count(&count).Error
}

type Query struct {
	Selector   map[string]interface{}
	Sort       map[string]string
	Pluck      string
	Pagination Pagination
}

type Pagination struct {
	Limit int
	Skip  int
}

func NewPagination(limit int, skip int) *Pagination {
	return &Pagination{
		Limit: limit,
		Skip:  skip,
	}
}

func NewQS(selector map[string]interface{}) *Query {
	return &Query{
		Selector: selector,
	}
}

// selector, sort, limit, pluck,
func (b *Bongo) Some(i Modellable, data interface{}, q *Query) error {
	err := b.executeQuery(i, data, q)
	if err == gorm.RecordNotFound {
		return nil
	}
	return err
}

func (b *Bongo) One(i Modellable, data interface{}, q *Query) error {
	q.Pagination.Limit = 1
	return b.executeQuery(i, data, q)
}

func (b *Bongo) BuildQuery(i Modellable, q *Query) *gorm.DB {
	// init query
	query := b.DB

	// add table name
	query = query.Table(i.TableName())

	// add sort options
	query = addSort(query, q.Sort)

	query = addSkip(query, q.Pagination.Skip)

	query = addLimit(query, q.Pagination.Limit)

	// add selector
	query = addWhere(query, q.Selector)

	return query
}

func (b *Bongo) executeQuery(i Modellable, data interface{}, q *Query) error {
	// init query
	query := b.BuildQuery(i, q)

	var err error
	// TODO refactor this part
	if q.Pluck != "" {
		if strings.Contains(q.Pluck, ",") {
			// add pluck data
			query = addPluck(query, q.Pluck)

			err = query.Find(data).Error
		} else {
			err = query.Pluck(q.Pluck, data).Error
		}
	} else {
		err = query.Find(data).Error
	}

	return err
}

func (b *Bongo) PublishEvent(eventName string, i Modellable) error {
	data, err := json.Marshal(i)
	if err != nil {
		b.log.Error("Error while marshalling for publish %s", err)
		return err
	}

	err = b.Broker.Publish(i.TableName()+"_"+eventName, data)
	if err != nil {
		b.log.Error("Error while publishing %s", err)
		return err
	}

	return nil
}

func (b *Bongo) AfterCreate(i Modellable) {
	b.PublishEvent("created", i)
}

func (b *Bongo) AfterUpdate(i Modellable) {
	b.PublishEvent("updated", i)
}

func (b *Bongo) AfterDelete(i Modellable) {
	b.PublishEvent("deleted", i)
}

func addSort(query *gorm.DB, options map[string]string) *gorm.DB {

	if options == nil {
		return query
	}

	if len(options) == 0 {
		return query
	}

	var opts []string
	for key, val := range options {
		opts = append(opts, fmt.Sprintf("%s %v", key, val))
	}
	return query.Order(strings.Join(opts, ","))
}

func addPluck(query *gorm.DB, plucked string) *gorm.DB {
	if plucked == "" {
		return query
	}

	return query.Select(plucked)
}

func addWhere(query *gorm.DB, selector map[string]interface{}) *gorm.DB {
	if selector == nil {
		return query
	}

	// instead sending one selector, do chaining here
	return query.Where(selector)
}

func addSkip(query *gorm.DB, skip int) *gorm.DB {
	if skip > 0 {
		return query.Offset(skip)
	}

	return query
}

func addLimit(query *gorm.DB, limit int) *gorm.DB {
	// if limit is minus or 0 ignore
	if limit > 0 {
		return query.Limit(limit)
	}

	return query
}

func CheckErr(res *gorm.DB) error {
	if res == nil {
		return nil
	}

	if res.Error == nil {
		return nil
	}

	if res.Error == gorm.RecordNotFound {
		return nil
	}

	return res.Error
}
<|endoftext|>


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

package socket

import (
	"encoding/binary"
	"flag"
	"fmt"
	"log"
	"net"
	"os"
	"os/user"
	"strconv"
	"syscall"

  "github.com/google/testimony/go/testimonyd/internal/vlog"
)

var (
	confFilename = flag.String("config", "/etc/testimony.conf", "Testimony config")
	logToSyslog  = flag.Bool("syslog", true, "log messages to syslog")
)

// Testimony is the configuration parsed from the config file.
type Testimony []SocketConfig

const protocolVersion = 1

// SocketConfig defines how an individual socket should be set up.
type SocketConfig struct {
	SocketName         string // filename for the socket
	Interface          string // interface to sniff packets on
	BlockSize          int    // block size (in bytes) of a single packet block
	NumBlocks          int    // number of packet blocks in the memory region
	BlockTimeoutMillis int    // timeout for filling up a single block
	FanoutType         int    // which type of fanout to use (see linux/if_packet.h)
	FanoutSize         int    // number of threads to fan out to
	User               string // user to provide the socket to (will chown it)
	Filter             string // BPF filter to apply to this socket
}

func (s SocketConfig) uid() (int, error) {
	var u *user.User
	var err error
	if s.User == "" {
		u, err = user.Current()
	} else {
		u, err = user.Lookup(s.User)
	}
	if err != nil {
		return 0, fmt.Errorf("could not get user: %v", err)
	}
	return strconv.Atoi(u.Uid)
}

// RunTestimony runs the testimonyd server given the passed in configuration.
func RunTestimony(t Testimony) {
	fanoutID := 0
	names := map[string]bool{}
	for _, sc := range t {
		// Check for duplicate socket names
		if names[sc.SocketName] {
			log.Fatalf("invalid config: duplicate socket name %q", sc.SocketName)
		}
		names[sc.SocketName] = true

		// Set up FanoutSize sockets and start goroutines to manage each.
		var socks []*socket
		fanoutID++
		for i := 0; i < sc.FanoutSize; i++ {
			sock, err := newSocket(sc, fanoutID, i)
			if err != nil {
				log.Fatalf("invalid config %+v: %v", sc, err)
			}
			socks = append(socks, sock)
			go sock.run()
		}

		// Set up UNIX socket to serve these AF_PACKET sockets on, and start
		// goroutine to manage its connections.
		_ignore_error_ := os.Remove(sc.SocketName)
		_ = _ignore_error_
		list, err := net.ListenUnix("unix", &net.UnixAddr{Net: "unix", Name: sc.SocketName})
		if err != nil {
			log.Fatalf("failed to listen on socket: %v", err)
		} else if err := setPermissions(sc); err != nil {
			log.Fatalf("failed to set socket permissions: %v", err)
		}
		go t.run(list, sc, socks)
	}
	// We'd love to drop privs here, but thanks to
	// https://github.com/golang/go/issues/1435 we can't :(
	select {} // Block (serving) forever.
}

func setPermissions(sc SocketConfig) error {
	uid, err := sc.uid()
	if err != nil {
		return fmt.Errorf("could not get uid to change to: %v", err)
	}
	vlog.V(1, "chowning %q to %d", sc.SocketName, uid)
	if err := syscall.Chown(sc.SocketName, uid, 0); err != nil {
		return fmt.Errorf("unable to chown to (%d, 0): %v", uid, err)
	}
	return nil
}

func (t Testimony) run(list *net.UnixListener, sc SocketConfig, socks []*socket) {
	for {
		c, err := list.AcceptUnix()
		if err != nil {
			log.Fatalf("failed to accept connection: %v", err)
		}
		go t.handle(socks, c)
	}
}

func (t Testimony) handle(socks []*socket, c *net.UnixConn) {
	defer func() {
		if c != nil {
			c.Close()
		}
	}()
	log.Printf("Received new connection %v", c.RemoteAddr())
	var buf [13]byte
	buf[0] = protocolVersion
	binary.BigEndian.PutUint32(buf[1:], uint32(len(socks)))
	binary.BigEndian.PutUint32(buf[5:], uint32(socks[0].conf.BlockSize))
	binary.BigEndian.PutUint32(buf[9:], uint32(socks[0].conf.NumBlocks))
	if _, err := c.Write(buf[:]); err != nil {
		log.Printf("new conn failed to write version: %v", err)
		return
	}
	var fanoutMsg [4]byte
	if n, err := c.Read(fanoutMsg[:]); n != len(fanoutMsg) || err != nil {
		log.Printf("new conn failed to read conf: %v", err)
		return
	}
	idx := int(binary.BigEndian.Uint32(fanoutMsg[:]))
	if idx < 0 || idx >= len(socks) {
		log.Printf("new conn invalid index %v", idx)
		return
	}
	sock := socks[idx]
	fdMsg := syscall.UnixRights(sock.fd)
	var msg [1]byte // dummy byte
	n, n2, err := c.WriteMsgUnix(
		msg[:], fdMsg, nil)
	if err != nil || n != len(msg) || n2 != len(fdMsg) {
		log.Printf("new conn failed to send file descriptor: %v", err)
		return
	}
	vlog.V(2, "new conn spun up, passing off to socket")
	sock.newConns <- c
	c = nil // so it doesn't get closed by deferred func.
}
<commit_msg>Fix panic.<commit_after>// Copyright 2015 Google Inc. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package socket

import (
	"encoding/binary"
	"fmt"
	"log"
	"net"
	"os"
	"os/user"
	"strconv"
	"syscall"

  "github.com/google/testimony/go/testimonyd/internal/vlog"
)

// Testimony is the configuration parsed from the config file.
type Testimony []SocketConfig

const protocolVersion = 1

// SocketConfig defines how an individual socket should be set up.
type SocketConfig struct {
	SocketName         string // filename for the socket
	Interface          string // interface to sniff packets on
	BlockSize          int    // block size (in bytes) of a single packet block
	NumBlocks          int    // number of packet blocks in the memory region
	BlockTimeoutMillis int    // timeout for filling up a single block
	FanoutType         int    // which type of fanout to use (see linux/if_packet.h)
	FanoutSize         int    // number of threads to fan out to
	User               string // user to provide the socket to (will chown it)
	Filter             string // BPF filter to apply to this socket
}

func (s SocketConfig) uid() (int, error) {
	var u *user.User
	var err error
	if s.User == "" {
		u, err = user.Current()
	} else {
		u, err = user.Lookup(s.User)
	}
	if err != nil {
		return 0, fmt.Errorf("could not get user: %v", err)
	}
	return strconv.Atoi(u.Uid)
}

// RunTestimony runs the testimonyd server given the passed in configuration.
func RunTestimony(t Testimony) {
	fanoutID := 0
	names := map[string]bool{}
	for _, sc := range t {
		// Check for duplicate socket names
		if names[sc.SocketName] {
			log.Fatalf("invalid config: duplicate socket name %q", sc.SocketName)
		}
		names[sc.SocketName] = true

		// Set up FanoutSize sockets and start goroutines to manage each.
		var socks []*socket
		fanoutID++
		for i := 0; i < sc.FanoutSize; i++ {
			sock, err := newSocket(sc, fanoutID, i)
			if err != nil {
				log.Fatalf("invalid config %+v: %v", sc, err)
			}
			socks = append(socks, sock)
			go sock.run()
		}

		// Set up UNIX socket to serve these AF_PACKET sockets on, and start
		// goroutine to manage its connections.
		_ignore_error_ := os.Remove(sc.SocketName)
		_ = _ignore_error_
		list, err := net.ListenUnix("unix", &net.UnixAddr{Net: "unix", Name: sc.SocketName})
		if err != nil {
			log.Fatalf("failed to listen on socket: %v", err)
		} else if err := setPermissions(sc); err != nil {
			log.Fatalf("failed to set socket permissions: %v", err)
		}
		go t.run(list, sc, socks)
	}
	// We'd love to drop privs here, but thanks to
	// https://github.com/golang/go/issues/1435 we can't :(
	select {} // Block (serving) forever.
}

func setPermissions(sc SocketConfig) error {
	uid, err := sc.uid()
	if err != nil {
		return fmt.Errorf("could not get uid to change to: %v", err)
	}
	vlog.V(1, "chowning %q to %d", sc.SocketName, uid)
	if err := syscall.Chown(sc.SocketName, uid, 0); err != nil {
		return fmt.Errorf("unable to chown to (%d, 0): %v", uid, err)
	}
	return nil
}

func (t Testimony) run(list *net.UnixListener, sc SocketConfig, socks []*socket) {
	for {
		c, err := list.AcceptUnix()
		if err != nil {
			log.Fatalf("failed to accept connection: %v", err)
		}
		go t.handle(socks, c)
	}
}

func (t Testimony) handle(socks []*socket, c *net.UnixConn) {
	defer func() {
		if c != nil {
			c.Close()
		}
	}()
	log.Printf("Received new connection %v", c.RemoteAddr())
	var buf [13]byte
	buf[0] = protocolVersion
	binary.BigEndian.PutUint32(buf[1:], uint32(len(socks)))
	binary.BigEndian.PutUint32(buf[5:], uint32(socks[0].conf.BlockSize))
	binary.BigEndian.PutUint32(buf[9:], uint32(socks[0].conf.NumBlocks))
	if _, err := c.Write(buf[:]); err != nil {
		log.Printf("new conn failed to write version: %v", err)
		return
	}
	var fanoutMsg [4]byte
	if n, err := c.Read(fanoutMsg[:]); n != len(fanoutMsg) || err != nil {
		log.Printf("new conn failed to read conf: %v", err)
		return
	}
	idx := int(binary.BigEndian.Uint32(fanoutMsg[:]))
	if idx < 0 || idx >= len(socks) {
		log.Printf("new conn invalid index %v", idx)
		return
	}
	sock := socks[idx]
	fdMsg := syscall.UnixRights(sock.fd)
	var msg [1]byte // dummy byte
	n, n2, err := c.WriteMsgUnix(
		msg[:], fdMsg, nil)
	if err != nil || n != len(msg) || n2 != len(fdMsg) {
		log.Printf("new conn failed to send file descriptor: %v", err)
		return
	}
	vlog.V(2, "new conn spun up, passing off to socket")
	sock.newConns <- c
	c = nil // so it doesn't get closed by deferred func.
}
<|endoftext|>


<commit_before>// Copyright 2015 Keybase, Inc. All rights reserved. Use of
// this source code is governed by the included BSD license.

package saltpack

import (
	"bytes"
	"fmt"
	basex "github.com/keybase/client/go/encoding/basex"
	"io"
	"io/ioutil"
	"strings"
)

// ArmorParams specify armor formatting, encoding and punctuation.
type ArmorParams struct {
	// BytesPerWord is the number of characters in each "word" of output.
	// We'll put spaces between words.
	BytesPerWord int
	// WordsPerLine is the number of words for each line of output. We'll
	// put newlines between two subsequent lines of output.
	WordsPerLine int
	// Punctuation is the byte inserted after the three "sentences" of
	// our encoding -- the header, the body and the footer.
	Punctuation byte
	// Encoding is the basex encoding to use, including strictness parameters
	Encoding *basex.Encoding
}

type armorEncoderStream struct {
	buf     *bytes.Buffer
	footer  string
	encoded io.Writer
	encoder io.WriteCloser
	nWords  int
	params  ArmorParams
}

func (s *armorEncoderStream) Write(b []byte) (n int, err error) {
	n, err = s.encoder.Write(b)
	if err != nil {
		return n, err
	}
	s.spaceAndOutputBuffer()
	return n, err
}

func (s *armorEncoderStream) spaceAndOutputBuffer() error {
	for s.buf.Len() > s.params.BytesPerWord {
		buf := s.buf.Next(s.params.BytesPerWord)
		s.nWords++
		sep := byte(' ')
		if s.nWords%s.params.WordsPerLine == 0 {
			sep = byte('\n')
		}
		if _, err := s.encoded.Write(buf); err != nil {
			return err
		}
		if _, err := s.encoded.Write([]byte{sep}); err != nil {
			return err
		}
	}
	return nil
}

func (s *armorEncoderStream) Close() (err error) {
	if err = s.encoder.Close(); err != nil {
		return err
	}
	if err := s.spaceAndOutputBuffer(); err != nil {
		return err
	}
	lst := s.buf.Bytes()
	if _, err := s.encoded.Write([]byte(lst)); err != nil {
		return err
	}
	s.nWords++
	pad := ""
	if len(lst) == s.params.BytesPerWord {
		if s.nWords%s.params.WordsPerLine == 0 {
			pad = "\n"
		} else {
			pad = " "
		}
	}
	if _, err := fmt.Fprintf(s.encoded, "%s%c\n\n%s%c\n", pad, s.params.Punctuation, s.footer, s.params.Punctuation); err != nil {
		return err
	}
	return nil
}

// NewArmorEncoderStream makes a new Armor encoding stream, using the given encoding
// Pass it an `encoded` stream writer to write the
// encoded stream to.  Also pass a header, and a footer string.  It will
// return an io.WriteCloser on success, that you can write raw (unencoded) data to.
// An error will be returned if there is trouble writing the header to encoded.
//
// To make the output look pretty, a space is inserted every 15 characters of output,
// and a newline is inserted every 200 words.
func NewArmorEncoderStream(encoded io.Writer, header string, footer string, params ArmorParams) (io.WriteCloser, error) {
	ret := &armorEncoderStream{
		buf:     new(bytes.Buffer),
		encoded: encoded,
		footer:  footer,
		params:  params,
	}
	ret.encoder = basex.NewEncoder(params.Encoding, ret.buf)
	if _, err := fmt.Fprintf(encoded, "%s%c\n\n", header, params.Punctuation); err != nil {
		return nil, err
	}
	return ret, nil
}

// ArmorSeal takes an input plaintext and returns and output armor encoding
// as a string, or an error if a problem was encountered. Also provide a header
// and a footer to frame the message.
func ArmorSeal(plaintext []byte, header string, footer string, params ArmorParams) (string, error) {
	var buf bytes.Buffer
	enc, err := NewArmorEncoderStream(&buf, header, footer, params)
	if err != nil {
		return "", err
	}
	if _, err := enc.Write(plaintext); err != nil {
		return "", err
	}
	if err := enc.Close(); err != nil {
		return "", err
	}
	return buf.String(), nil
}

// Frame is a way to read the frame out of a Decoder stream.
type Frame interface {
	// GetHeader() returns the frame associated with this stream, or an error
	GetHeader() (string, error)
	// GetFooter() returns the frame associated with this stream, or an error
	GetFooter() (string, error)
}

type fdsState int

const (
	fdsHeader fdsState = iota
	fdsBody
	fdsFooter
	fdsEndOfStream
)

type framedDecoderStream struct {
	header []byte
	footer []byte
	state  fdsState
	params ArmorParams
	r      *PunctuatedReader
}

// Read from a framedDeecoderStream. The frame is the "BEGIN FOO." block
// at the footer, and the "END FOO." block at the end.
func (s *framedDecoderStream) Read(p []byte) (n int, err error) {

	// The largest frame we'll accept before we show an overflow.
	frameLim := 8192

	if s.state == fdsHeader {
		s.header, err = s.r.ReadUntilPunctuation(frameLim)
		if err != nil {
			return 0, err
		}
		s.state = fdsBody
	}

	if s.state == fdsBody {
		n, err = s.r.Read(p)
		if err == ErrPunctuated {
			err = nil
			s.state = fdsFooter
		}
		if err == io.EOF {
			err = io.ErrUnexpectedEOF
		}
		if err != nil {
			return 0, err
		}
	}

	if s.state == fdsFooter {
		s.footer, err = s.r.ReadUntilPunctuation(frameLim)
		if err != nil {
			return 0, err
		}
		s.state = fdsEndOfStream
	}

	if s.state == fdsEndOfStream {
		err = s.consumeUntilEOF()
		if err == io.EOF && n > 0 {
			err = nil
		}
	}

	return n, err
}

// consume the stream until we hit an EOF. For all data we consume, make
// sure that it's a valid byte as far as our underlying decoder is concerned.
// We might considering clamping down here on the number of characters we're willing
// to accept after the message is over. But for now, we're quite liberal.
func (s *framedDecoderStream) consumeUntilEOF() error {
	var buf [4096]byte
	for {
		n, err := s.r.Read(buf[:])
		if err != nil {
			return err
		}
		if n == 0 {
			return io.EOF
		}
		if !s.isValidByteSequence(buf[0:n]) {
			return ErrTrailingGarbage
		}
	}
}

// isValidByteSequence checks if the byte sequence is valid as far as our
// underlying encoder is concerned. This is usually quite liberal, and excludes
// non-ASCII and ASCII control characters, but leaves everything else in.
func (s *framedDecoderStream) isValidByteSequence(p []byte) bool {
	for _, b := range p {
		if !s.params.Encoding.IsValidByte(b) {
			return false
		}
	}
	return true
}

func (s *framedDecoderStream) toASCII(buf []byte) (string, error) {
	if !s.isValidByteSequence(buf) {
		return "", makeErrBadFrame("invalid ASCII sequence")
	}
	return strings.TrimSpace(string(buf)), nil
}

func (s *framedDecoderStream) GetFooter() (string, error) { return s.toASCII(s.footer) }
func (s *framedDecoderStream) GetHeader() (string, error) { return s.toASCII(s.header) }

// NewArmorDecoderStream is used to decode armored encoding. It returns a stream you
// can read from, and also a Frame you can query to see what the open/close
// frame markers were.
func NewArmorDecoderStream(r io.Reader, params ArmorParams) (io.Reader, Frame, error) {
	fds := &framedDecoderStream{r: NewPunctuatedReader(r, params.Punctuation), params: params}
	ret := basex.NewDecoder(params.Encoding, fds)
	return ret, fds, nil
}

// ArmorOpen runs armor stream decoding, but on a string, and it outputs a string.
func ArmorOpen(msg string, params ArmorParams) (body []byte, header string, footer string, err error) {
	var dec io.Reader
	var frame Frame
	buf := bytes.NewBufferString(msg)
	dec, frame, err = NewArmorDecoderStream(buf, params)
	if err != nil {
		return
	}
	body, err = ioutil.ReadAll(dec)
	if err != nil {
		return
	}
	if header, err = frame.GetHeader(); err != nil {
		return
	}
	if footer, err = frame.GetFooter(); err != nil {
		return
	}
	return body, header, footer, nil
}
<commit_msg>make the armor header and footer inline<commit_after>// Copyright 2015 Keybase, Inc. All rights reserved. Use of
// this source code is governed by the included BSD license.

package saltpack

import (
	"bytes"
	"fmt"
	basex "github.com/keybase/client/go/encoding/basex"
	"io"
	"io/ioutil"
	"strings"
)

// ArmorParams specify armor formatting, encoding and punctuation.
type ArmorParams struct {
	// BytesPerWord is the number of characters in each "word" of output.
	// We'll put spaces between words.
	BytesPerWord int
	// WordsPerLine is the number of words for each line of output. We'll
	// put newlines between two subsequent lines of output.
	WordsPerLine int
	// Punctuation is the byte inserted after the three "sentences" of
	// our encoding -- the header, the body and the footer.
	Punctuation byte
	// Encoding is the basex encoding to use, including strictness parameters
	Encoding *basex.Encoding
}

type armorEncoderStream struct {
	buf     *bytes.Buffer
	footer  string
	encoded io.Writer
	encoder io.WriteCloser
	nWords  int
	params  ArmorParams
}

func (s *armorEncoderStream) Write(b []byte) (n int, err error) {
	n, err = s.encoder.Write(b)
	if err != nil {
		return n, err
	}
	s.spaceAndOutputBuffer()
	return n, err
}

func (s *armorEncoderStream) spaceAndOutputBuffer() error {
	for s.buf.Len() > s.params.BytesPerWord {
		buf := s.buf.Next(s.params.BytesPerWord)
		s.nWords++
		sep := byte(' ')
		if s.nWords%s.params.WordsPerLine == 0 {
			sep = byte('\n')
		}
		if _, err := s.encoded.Write(buf); err != nil {
			return err
		}
		if _, err := s.encoded.Write([]byte{sep}); err != nil {
			return err
		}
	}
	return nil
}

func (s *armorEncoderStream) Close() (err error) {
	if err = s.encoder.Close(); err != nil {
		return err
	}
	if err := s.spaceAndOutputBuffer(); err != nil {
		return err
	}
	lst := s.buf.Bytes()
	if _, err := s.encoded.Write([]byte(lst)); err != nil {
		return err
	}
	s.nWords++
	pad := ""
	if len(lst) == s.params.BytesPerWord {
		if s.nWords%s.params.WordsPerLine == 0 {
			pad = "\n"
		} else {
			pad = " "
		}
	}
	if _, err := fmt.Fprintf(s.encoded, "%s%c %s%c\n", pad, s.params.Punctuation, s.footer, s.params.Punctuation); err != nil {
		return err
	}
	return nil
}

// NewArmorEncoderStream makes a new Armor encoding stream, using the given encoding
// Pass it an `encoded` stream writer to write the
// encoded stream to.  Also pass a header, and a footer string.  It will
// return an io.WriteCloser on success, that you can write raw (unencoded) data to.
// An error will be returned if there is trouble writing the header to encoded.
//
// To make the output look pretty, a space is inserted every 15 characters of output,
// and a newline is inserted every 200 words.
func NewArmorEncoderStream(encoded io.Writer, header string, footer string, params ArmorParams) (io.WriteCloser, error) {
	ret := &armorEncoderStream{
		buf:     new(bytes.Buffer),
		encoded: encoded,
		footer:  footer,
		params:  params,
	}
	ret.encoder = basex.NewEncoder(params.Encoding, ret.buf)
	if _, err := fmt.Fprintf(encoded, "%s%c ", header, params.Punctuation); err != nil {
		return nil, err
	}
	return ret, nil
}

// ArmorSeal takes an input plaintext and returns and output armor encoding
// as a string, or an error if a problem was encountered. Also provide a header
// and a footer to frame the message.
func ArmorSeal(plaintext []byte, header string, footer string, params ArmorParams) (string, error) {
	var buf bytes.Buffer
	enc, err := NewArmorEncoderStream(&buf, header, footer, params)
	if err != nil {
		return "", err
	}
	if _, err := enc.Write(plaintext); err != nil {
		return "", err
	}
	if err := enc.Close(); err != nil {
		return "", err
	}
	return buf.String(), nil
}

// Frame is a way to read the frame out of a Decoder stream.
type Frame interface {
	// GetHeader() returns the frame associated with this stream, or an error
	GetHeader() (string, error)
	// GetFooter() returns the frame associated with this stream, or an error
	GetFooter() (string, error)
}

type fdsState int

const (
	fdsHeader fdsState = iota
	fdsBody
	fdsFooter
	fdsEndOfStream
)

type framedDecoderStream struct {
	header []byte
	footer []byte
	state  fdsState
	params ArmorParams
	r      *PunctuatedReader
}

// Read from a framedDeecoderStream. The frame is the "BEGIN FOO." block
// at the footer, and the "END FOO." block at the end.
func (s *framedDecoderStream) Read(p []byte) (n int, err error) {

	// The largest frame we'll accept before we show an overflow.
	frameLim := 8192

	if s.state == fdsHeader {
		s.header, err = s.r.ReadUntilPunctuation(frameLim)
		if err != nil {
			return 0, err
		}
		s.state = fdsBody
	}

	if s.state == fdsBody {
		n, err = s.r.Read(p)
		if err == ErrPunctuated {
			err = nil
			s.state = fdsFooter
		}
		if err == io.EOF {
			err = io.ErrUnexpectedEOF
		}
		if err != nil {
			return 0, err
		}
	}

	if s.state == fdsFooter {
		s.footer, err = s.r.ReadUntilPunctuation(frameLim)
		if err != nil {
			return 0, err
		}
		s.state = fdsEndOfStream
	}

	if s.state == fdsEndOfStream {
		err = s.consumeUntilEOF()
		if err == io.EOF && n > 0 {
			err = nil
		}
	}

	return n, err
}

// consume the stream until we hit an EOF. For all data we consume, make
// sure that it's a valid byte as far as our underlying decoder is concerned.
// We might considering clamping down here on the number of characters we're willing
// to accept after the message is over. But for now, we're quite liberal.
func (s *framedDecoderStream) consumeUntilEOF() error {
	var buf [4096]byte
	for {
		n, err := s.r.Read(buf[:])
		if err != nil {
			return err
		}
		if n == 0 {
			return io.EOF
		}
		if !s.isValidByteSequence(buf[0:n]) {
			return ErrTrailingGarbage
		}
	}
}

// isValidByteSequence checks if the byte sequence is valid as far as our
// underlying encoder is concerned. This is usually quite liberal, and excludes
// non-ASCII and ASCII control characters, but leaves everything else in.
func (s *framedDecoderStream) isValidByteSequence(p []byte) bool {
	for _, b := range p {
		if !s.params.Encoding.IsValidByte(b) {
			return false
		}
	}
	return true
}

func (s *framedDecoderStream) toASCII(buf []byte) (string, error) {
	if !s.isValidByteSequence(buf) {
		return "", makeErrBadFrame("invalid ASCII sequence")
	}
	return strings.TrimSpace(string(buf)), nil
}

func (s *framedDecoderStream) GetFooter() (string, error) { return s.toASCII(s.footer) }
func (s *framedDecoderStream) GetHeader() (string, error) { return s.toASCII(s.header) }

// NewArmorDecoderStream is used to decode armored encoding. It returns a stream you
// can read from, and also a Frame you can query to see what the open/close
// frame markers were.
func NewArmorDecoderStream(r io.Reader, params ArmorParams) (io.Reader, Frame, error) {
	fds := &framedDecoderStream{r: NewPunctuatedReader(r, params.Punctuation), params: params}
	ret := basex.NewDecoder(params.Encoding, fds)
	return ret, fds, nil
}

// ArmorOpen runs armor stream decoding, but on a string, and it outputs a string.
func ArmorOpen(msg string, params ArmorParams) (body []byte, header string, footer string, err error) {
	var dec io.Reader
	var frame Frame
	buf := bytes.NewBufferString(msg)
	dec, frame, err = NewArmorDecoderStream(buf, params)
	if err != nil {
		return
	}
	body, err = ioutil.ReadAll(dec)
	if err != nil {
		return
	}
	if header, err = frame.GetHeader(); err != nil {
		return
	}
	if footer, err = frame.GetFooter(); err != nil {
		return
	}
	return body, header, footer, nil
}
<|endoftext|>