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

package main

import (
	"bufio"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"os"
	"strings"

	"github.com/coreos/rkt/Godeps/_workspace/src/github.com/appc/spec/discovery"
	"github.com/coreos/rkt/Godeps/_workspace/src/golang.org/x/crypto/openpgp"
)

// getPubKeyLocations discovers one location at prefix
func getPubKeyLocations(prefix string, allowHTTP bool, debug bool) ([]string, error) {
	if prefix == "" {
		return nil, fmt.Errorf("empty prefix")
	}

	kls, err := metaDiscoverPubKeyLocations(prefix, allowHTTP, debug)
	if err != nil {
		return nil, fmt.Errorf("prefix meta discovery error: %v", err)
	}

	if len(kls) == 0 {
		return nil, fmt.Errorf("meta discovery on %s resulted in no keys", prefix)
	}

	return kls, nil
}

// addKeys adds the keys listed in pkls at prefix
func addKeys(pkls []string, prefix string, allowHTTP, forceAccept, allowOverride bool) error {
	ks := getKeystore()
	if ks == nil {
		panic("could not get the key store")
	}

	for _, pkl := range pkls {
		u, err := url.Parse(pkl)
		if err != nil {
			return err
		}
		pk, err := getPubKey(u.Scheme, pkl, allowHTTP)
		if err != nil {
			return fmt.Errorf("error accessing the key %s: %v", pkl, err)
		}
		defer pk.Close()

		exists, err := ks.TrustedKeyPrefixExists(prefix, pk)
		if err != nil {
			return fmt.Errorf("error reading the key %s: %v", pkl, err)
		}
		err = displayKey(prefix, pkl, pk)
		if err != nil {
			return fmt.Errorf("error displaying the key %s: %v", pkl, err)
		}
		if exists && !allowOverride {
			stderr("Key %q already in the keystore", pkl)
			continue
		}

		if globalFlags.TrustKeysFromHttps && u.Scheme == "https" {
			forceAccept = true
		}

		if !forceAccept {
			accepted, err := reviewKey()
			if err != nil {
				return fmt.Errorf("error reviewing key: %v", err)
			}
			if !accepted {
				stderr("Not trusting %q", pkl)
				continue
			}
		}

		if forceAccept {
			stderr("Trusting %q for prefix %q without fingerprint review.", pkl, prefix)
		} else {
			stderr("Trusting %q for prefix %q after fingerprint review.", pkl, prefix)
		}

		if prefix == "" {
			path, err := ks.StoreTrustedKeyRoot(pk)
			if err != nil {
				return fmt.Errorf("Error adding root key: %v", err)
			}
			stderr("Added root key at %q", path)
		} else {
			path, err := ks.StoreTrustedKeyPrefix(prefix, pk)
			if err != nil {
				return fmt.Errorf("Error adding key for prefix %q: %v", prefix, err)
			}
			stderr("Added key for prefix %q at %q", prefix, path)
		}
	}
	return nil
}

// metaDiscoverPubKeyLocations discovers the public key through ACDiscovery by applying prefix as an ACApp
func metaDiscoverPubKeyLocations(prefix string, allowHTTP bool, debug bool) ([]string, error) {
	app, err := discovery.NewAppFromString(prefix)
	if err != nil {
		return nil, err
	}

	ep, attempts, err := discovery.DiscoverPublicKeys(*app, allowHTTP)
	if err != nil {
		return nil, err
	}

	if debug {
		for _, a := range attempts {
			stderr("meta tag 'ac-discovery-pubkeys' not found on %s: %v", a.Prefix, a.Error)
		}
	}

	return ep.Keys, nil
}

// getPubKey retrieves a public key (if remote), and verifies it's a gpg key
func getPubKey(scheme, location string, allowHTTP bool) (*os.File, error) {
	switch scheme {
	case "":
		return os.Open(location)
	case "http":
		if !allowHTTP {
			return nil, fmt.Errorf("--insecure-allow-http required for http URLs")
		}
		fallthrough
	case "https":
		return downloadKey(location)
	}

	return nil, fmt.Errorf("only http and https urls supported")
}

// downloadKey retrieves the file, storing it in a deleted tempfile
func downloadKey(url string) (*os.File, error) {
	tf, err := ioutil.TempFile("", "")
	if err != nil {
		return nil, fmt.Errorf("error creating tempfile: %v", err)
	}
	os.Remove(tf.Name()) // no need to keep the tempfile around

	defer func() {
		if err != nil {
			tf.Close()
		}
	}()

	res, err := http.Get(url)
	if err != nil {
		return nil, fmt.Errorf("error getting key: %v", err)
	}
	defer res.Body.Close()

	if res.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("bad HTTP status code: %d", res.StatusCode)
	}

	if _, err := io.Copy(tf, res.Body); err != nil {
		return nil, fmt.Errorf("error copying key: %v", err)
	}

	if _, err = tf.Seek(0, os.SEEK_SET); err != nil {
		return nil, fmt.Errorf("error seeking: %v", err)
	}

	return tf, nil
}

// displayKey shows the key summary
func displayKey(prefix string, location string, key *os.File) error {
	defer key.Seek(0, os.SEEK_SET)

	kr, err := openpgp.ReadArmoredKeyRing(key)
	if err != nil {
		return fmt.Errorf("error reading key: %v", err)
	}

	stderr("prefix: %q\nkey: %q", prefix, location)
	for _, k := range kr {
		stderr("gpg key fingerprint is: %s", fingerToString(k.PrimaryKey.Fingerprint))
		for _, sk := range k.Subkeys {
			stderr("    subkey fingerprint: %s", fingerToString(sk.PublicKey.Fingerprint))
		}
		for n, _ := range k.Identities {
			stderr("\t%s", n)
		}
	}
	return nil
}

// reviewKey asks the user to accept the key
func reviewKey() (bool, error) {
	in := bufio.NewReader(os.Stdin)
	for {
		stderr("Are you sure you want to trust this key (yes/no)?")
		input, err := in.ReadString('\n')
		if err != nil {
			return false, fmt.Errorf("error reading input: %v", err)
		}
		switch input {
		case "yes\n":
			return true, nil
		case "no\n":
			return false, nil
		default:
			stderr("Please enter 'yes' or 'no'")
		}
	}
}

func fingerToString(fpr [20]byte) string {
	str := ""
	for i, b := range fpr {
		if i > 0 && i%2 == 0 {
			str += " "
			if i == 10 {
				str += " "
			}
		}
		str += strings.ToUpper(fmt.Sprintf("%.2x", b))
	}
	return str
}
<commit_msg>pubkey: Refactor the code<commit_after>// Copyright 2015 The rkt Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package pubkey

import (
	"bufio"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"os"
	"strings"

	"github.com/coreos/rkt/pkg/keystore"

	"github.com/coreos/rkt/Godeps/_workspace/src/github.com/appc/spec/discovery"
	"github.com/coreos/rkt/Godeps/_workspace/src/golang.org/x/crypto/openpgp"
)

type Manager struct {
	InsecureAllowHttp  bool
	TrustKeysFromHttps bool
	Ks                 *keystore.Keystore
	Debug              bool
}

type AcceptOption int
type OverrideOption int

const (
	AcceptForce AcceptOption = iota
	AcceptAsk
)

const (
	OverrideAllow OverrideOption = iota
	OverrideDeny
)

// GetPubKeyLocations discovers one location at prefix
func (m *Manager) GetPubKeyLocations(prefix string) ([]string, error) {
	if prefix == "" {
		return nil, fmt.Errorf("empty prefix")
	}

	kls, err := m.metaDiscoverPubKeyLocations(prefix)
	if err != nil {
		return nil, fmt.Errorf("prefix meta discovery error: %v", err)
	}

	if len(kls) == 0 {
		return nil, fmt.Errorf("meta discovery on %s resulted in no keys", prefix)
	}

	return kls, nil
}

// AddKeys adds the keys listed in pkls at prefix
func (m *Manager) AddKeys(pkls []string, prefix string, accept AcceptOption, override OverrideOption) error {
	if m.Ks == nil {
		return fmt.Errorf("no keystore available to add keys to")
	}

	for _, pkl := range pkls {
		u, err := url.Parse(pkl)
		if err != nil {
			return err
		}
		pk, err := m.getPubKey(u)
		if err != nil {
			return fmt.Errorf("error accessing the key %s: %v", pkl, err)
		}
		defer pk.Close()

		exists, err := m.Ks.TrustedKeyPrefixExists(prefix, pk)
		if err != nil {
			return fmt.Errorf("error reading the key %s: %v", pkl, err)
		}
		err = displayKey(prefix, pkl, pk)
		if err != nil {
			return fmt.Errorf("error displaying the key %s: %v", pkl, err)
		}
		if exists && override == OverrideDeny {
			stderr("Key %q already in the keystore", pkl)
			continue
		}

		if m.TrustKeysFromHttps && u.Scheme == "https" {
			accept = AcceptForce
		}

		if accept == AcceptAsk {
			accepted, err := reviewKey()
			if err != nil {
				return fmt.Errorf("error reviewing key: %v", err)
			}
			if !accepted {
				stderr("Not trusting %q", pkl)
				continue
			}
		}

		if accept == AcceptForce {
			stderr("Trusting %q for prefix %q without fingerprint review.", pkl, prefix)
		} else {
			stderr("Trusting %q for prefix %q after fingerprint review.", pkl, prefix)
		}

		if prefix == "" {
			path, err := m.Ks.StoreTrustedKeyRoot(pk)
			if err != nil {
				return fmt.Errorf("Error adding root key: %v", err)
			}
			stderr("Added root key at %q", path)
		} else {
			path, err := m.Ks.StoreTrustedKeyPrefix(prefix, pk)
			if err != nil {
				return fmt.Errorf("Error adding key for prefix %q: %v", prefix, err)
			}
			stderr("Added key for prefix %q at %q", prefix, path)
		}
	}
	return nil
}

// metaDiscoverPubKeyLocations discovers the public key through ACDiscovery by applying prefix as an ACApp
func (m *Manager) metaDiscoverPubKeyLocations(prefix string) ([]string, error) {
	app, err := discovery.NewAppFromString(prefix)
	if err != nil {
		return nil, err
	}

	// TODO(krnowak): we should probably apply credential headers
	// from config here
	ep, attempts, err := discovery.DiscoverPublicKeys(*app, m.InsecureAllowHttp)
	if err != nil {
		return nil, err
	}

	if m.Debug {
		for _, a := range attempts {
			stderr("meta tag 'ac-discovery-pubkeys' not found on %s: %v", a.Prefix, a.Error)
		}
	}

	return ep.Keys, nil
}

// getPubKey retrieves a public key (if remote), and verifies it's a gpg key
func (m *Manager) getPubKey(u *url.URL) (*os.File, error) {
	switch u.Scheme {
	case "":
		return os.Open(u.Path)
	case "http":
		if !m.InsecureAllowHttp {
			return nil, fmt.Errorf("--insecure-allow-http required for http URLs")
		}
		fallthrough
	case "https":
		return downloadKey(u)
	}

	return nil, fmt.Errorf("only local files and http or https URLs supported")
}

// downloadKey retrieves the file, storing it in a deleted tempfile
func downloadKey(u *url.URL) (*os.File, error) {
	tf, err := ioutil.TempFile("", "")
	if err != nil {
		return nil, fmt.Errorf("error creating tempfile: %v", err)
	}
	os.Remove(tf.Name()) // no need to keep the tempfile around

	defer func() {
		if tf != nil {
			tf.Close()
		}
	}()

	// TODO(krnowak): we should probably apply credential headers
	// from config here
	res, err := http.Get(u.String())
	if err != nil {
		return nil, fmt.Errorf("error getting key: %v", err)
	}
	defer res.Body.Close()

	if res.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("bad HTTP status code: %d", res.StatusCode)
	}

	if _, err := io.Copy(tf, res.Body); err != nil {
		return nil, fmt.Errorf("error copying key: %v", err)
	}

	if _, err = tf.Seek(0, os.SEEK_SET); err != nil {
		return nil, fmt.Errorf("error seeking: %v", err)
	}

	retTf := tf
	tf = nil
	return retTf, nil
}

// displayKey shows the key summary
func displayKey(prefix, location string, key *os.File) error {
	defer key.Seek(0, os.SEEK_SET)

	kr, err := openpgp.ReadArmoredKeyRing(key)
	if err != nil {
		return fmt.Errorf("error reading key: %v", err)
	}

	stderr("prefix: %q\nkey: %q", prefix, location)
	for _, k := range kr {
		stderr("gpg key fingerprint is: %s", fingerToString(k.PrimaryKey.Fingerprint))
		for _, sk := range k.Subkeys {
			stderr("    subkey fingerprint: %s", fingerToString(sk.PublicKey.Fingerprint))
		}
		for n, _ := range k.Identities {
			stderr("\t%s", n)
		}
	}
	return nil
}

// reviewKey asks the user to accept the key
func reviewKey() (bool, error) {
	in := bufio.NewReader(os.Stdin)
	for {
		stderr("Are you sure you want to trust this key (yes/no)?")
		input, err := in.ReadString('\n')
		if err != nil {
			return false, fmt.Errorf("error reading input: %v", err)
		}
		switch input {
		case "yes\n":
			return true, nil
		case "no\n":
			return false, nil
		default:
			stderr("Please enter 'yes' or 'no'")
		}
	}
}

func fingerToString(fpr [20]byte) string {
	str := ""
	for i, b := range fpr {
		if i > 0 && i%2 == 0 {
			str += " "
			if i == 10 {
				str += " "
			}
		}
		str += strings.ToUpper(fmt.Sprintf("%.2x", b))
	}
	return str
}

func stderr(format string, a ...interface{}) {
	out := fmt.Sprintf(format, a...)
	fmt.Fprintln(os.Stderr, strings.TrimSuffix(out, "\n"))
}
<|endoftext|>


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

package fileserver

import (
	"fmt"
	"net/http"
	"os"
	"path"
	"strings"
	"time"

	"github.com/caddyserver/caddy/v2"
	"github.com/caddyserver/caddy/v2/caddyconfig/caddyfile"
	"github.com/caddyserver/caddy/v2/modules/caddyhttp"
)

func init() {
	caddy.RegisterModule(MatchFile{})
}

// MatchFile is an HTTP request matcher that can match
// requests based upon file existence.
//
// Upon matching, three new placeholders will be made
// available:
//
// - `{http.matchers.file.relative}` The root-relative
// path of the file. This is often useful when rewriting
// requests.
// - `{http.matchers.file.absolute}` The absolute path
// of the matched file.
// - `{http.matchers.file.type}` Set to "directory" if
// the matched file is a directory, "file" otherwise.
// - `{http.matchers.file.remainder}` Set to the remainder
// of the path if the path was split by `split_path`.
type MatchFile struct {
	// The root directory, used for creating absolute
	// file paths, and required when working with
	// relative paths; if not specified, `{http.vars.root}`
	// will be used, if set; otherwise, the current
	// directory is assumed. Accepts placeholders.
	Root string `json:"root,omitempty"`

	// The list of files to try. Each path here is
	// considered related to Root. If nil, the request
	// URL's path will be assumed. Files and
	// directories are treated distinctly, so to match
	// a directory, the filepath MUST end in a forward
	// slash `/`. To match a regular file, there must
	// be no trailing slash. Accepts placeholders.
	TryFiles []string `json:"try_files,omitempty"`

	// How to choose a file in TryFiles. Can be:
	//
	// - first_exist
	// - smallest_size
	// - largest_size
	// - most_recently_modified
	//
	// Default is first_exist.
	TryPolicy string `json:"try_policy,omitempty"`

	// A list of delimiters to use to split the path in two
	// when trying files. If empty, no splitting will
	// occur, and the path will be tried as-is. For each
	// split value, the left-hand side of the split,
	// including the split value, will be the path tried.
	// For example, the path `/remote.php/dav/` using the
	// split value `.php` would try the file `/remote.php`.
	// Each delimiter must appear at the end of a URI path
	// component in order to be used as a split delimiter.
	SplitPath []string `json:"split_path,omitempty"`
}

// CaddyModule returns the Caddy module information.
func (MatchFile) CaddyModule() caddy.ModuleInfo {
	return caddy.ModuleInfo{
		ID:  "http.matchers.file",
		New: func() caddy.Module { return new(MatchFile) },
	}
}

// UnmarshalCaddyfile sets up the matcher from Caddyfile tokens. Syntax:
//
//     file <files...> {
//         root <path>
//         try_files <files...>
//         try_policy first_exist|smallest_size|largest_size|most_recently_modified
//     }
//
func (m *MatchFile) UnmarshalCaddyfile(d *caddyfile.Dispenser) error {
	for d.Next() {
		m.TryFiles = append(m.TryFiles, d.RemainingArgs()...)
		for d.NextBlock(0) {
			switch d.Val() {
			case "root":
				if !d.NextArg() {
					return d.ArgErr()
				}
				m.Root = d.Val()
			case "try_files":
				m.TryFiles = append(m.TryFiles, d.RemainingArgs()...)
				if len(m.TryFiles) == 0 {
					return d.ArgErr()
				}
			case "try_policy":
				if !d.NextArg() {
					return d.ArgErr()
				}
				m.TryPolicy = d.Val()
			case "split_path":
				m.SplitPath = d.RemainingArgs()
				if len(m.SplitPath) == 0 {
					return d.ArgErr()
				}
			default:
				return d.Errf("unrecognized subdirective: %s", d.Val())
			}
		}
	}
	return nil
}

// Provision sets up m's defaults.
func (m *MatchFile) Provision(_ caddy.Context) error {
	if m.Root == "" {
		m.Root = "{http.vars.root}"
	}
	return nil
}

// Validate ensures m has a valid configuration.
func (m MatchFile) Validate() error {
	switch m.TryPolicy {
	case "",
		tryPolicyFirstExist,
		tryPolicyLargestSize,
		tryPolicySmallestSize,
		tryPolicyMostRecentlyMod:
	default:
		return fmt.Errorf("unknown try policy %s", m.TryPolicy)
	}
	return nil
}

// Match returns true if r matches m. Returns true
// if a file was matched. If so, four placeholders
// will be available:
//    - http.matchers.file.relative
//    - http.matchers.file.absolute
//    - http.matchers.file.type
//    - http.matchers.file.remainder
func (m MatchFile) Match(r *http.Request) bool {
	return m.selectFile(r)
}

// selectFile chooses a file according to m.TryPolicy by appending
// the paths in m.TryFiles to m.Root, with placeholder replacements.
func (m MatchFile) selectFile(r *http.Request) (matched bool) {
	repl := r.Context().Value(caddy.ReplacerCtxKey).(*caddy.Replacer)

	root := repl.ReplaceAll(m.Root, ".")

	// if list of files to try was omitted entirely,
	// assume URL path
	if m.TryFiles == nil {
		// m is not a pointer, so this is safe
		m.TryFiles = []string{r.URL.Path}
	}

	// common preparation of the file into parts
	prepareFilePath := func(file string) (suffix, fullpath, remainder string) {
		suffix, remainder = m.firstSplit(path.Clean(repl.ReplaceAll(file, "")))
		if strings.HasSuffix(file, "/") {
			suffix += "/"
		}
		fullpath = sanitizedPathJoin(root, suffix)
		return
	}

	// sets up the placeholders for the matched file
	setPlaceholders := func(info os.FileInfo, rel string, abs string, remainder string) {
		repl.Set("http.matchers.file.relative", rel)
		repl.Set("http.matchers.file.absolute", abs)
		repl.Set("http.matchers.file.remainder", remainder)

		fileType := "file"
		if info.IsDir() {
			fileType = "directory"
		}
		repl.Set("http.matchers.file.type", fileType)
	}

	switch m.TryPolicy {
	case "", tryPolicyFirstExist:
		for _, f := range m.TryFiles {
			suffix, fullpath, remainder := prepareFilePath(f)
			if info, exists := strictFileExists(fullpath); exists {
				setPlaceholders(info, suffix, fullpath, remainder)
				return true
			}
		}

	case tryPolicyLargestSize:
		var largestSize int64
		var largestFilename string
		var largestSuffix string
		var remainder string
		var info os.FileInfo
		for _, f := range m.TryFiles {
			suffix, fullpath, splitRemainder := prepareFilePath(f)
			info, err := os.Stat(fullpath)
			if err == nil && info.Size() > largestSize {
				largestSize = info.Size()
				largestFilename = fullpath
				largestSuffix = suffix
				remainder = splitRemainder
			}
		}
		setPlaceholders(info, largestSuffix, largestFilename, remainder)
		return true

	case tryPolicySmallestSize:
		var smallestSize int64
		var smallestFilename string
		var smallestSuffix string
		var remainder string
		var info os.FileInfo
		for _, f := range m.TryFiles {
			suffix, fullpath, splitRemainder := prepareFilePath(f)
			info, err := os.Stat(fullpath)
			if err == nil && (smallestSize == 0 || info.Size() < smallestSize) {
				smallestSize = info.Size()
				smallestFilename = fullpath
				smallestSuffix = suffix
				remainder = splitRemainder
			}
		}
		setPlaceholders(info, smallestSuffix, smallestFilename, remainder)
		return true

	case tryPolicyMostRecentlyMod:
		var recentDate time.Time
		var recentFilename string
		var recentSuffix string
		var remainder string
		var info os.FileInfo
		for _, f := range m.TryFiles {
			suffix, fullpath, splitRemainder := prepareFilePath(f)
			info, err := os.Stat(fullpath)
			if err == nil &&
				(recentDate.IsZero() || info.ModTime().After(recentDate)) {
				recentDate = info.ModTime()
				recentFilename = fullpath
				recentSuffix = suffix
				remainder = splitRemainder
			}
		}
		setPlaceholders(info, recentSuffix, recentFilename, remainder)
		return true
	}

	return
}

// strictFileExists returns true if file exists
// and matches the convention of the given file
// path. If the path ends in a forward slash,
// the file must also be a directory; if it does
// NOT end in a forward slash, the file must NOT
// be a directory.
func strictFileExists(file string) (os.FileInfo, bool) {
	stat, err := os.Stat(file)
	if err != nil {
		// in reality, this can be any error
		// such as permission or even obscure
		// ones like "is not a directory" (when
		// trying to stat a file within a file);
		// in those cases we can't be sure if
		// the file exists, so we just treat any
		// error as if it does not exist; see
		// https://stackoverflow.com/a/12518877/1048862
		return nil, false
	}
	if strings.HasSuffix(file, separator) {
		// by convention, file paths ending
		// in a path separator must be a directory
		return stat, stat.IsDir()
	}
	// by convention, file paths NOT ending
	// in a path separator must NOT be a directory
	return stat, !stat.IsDir()
}

// firstSplit returns the first result where the path
// can be split in two by a value in m.SplitPath. The
// return values are the first piece of the path that
// ends with the split substring and the remainder.
// If the path cannot be split, the path is returned
// as-is (with no remainder).
func (m MatchFile) firstSplit(path string) (splitPart, remainder string) {
	for _, split := range m.SplitPath {
		if idx := indexFold(path, split); idx > -1 {
			pos := idx + len(split)
			// skip the split if it's not the final part of the filename
			if pos != len(path) && !strings.HasPrefix(path[pos:], "/") {
				continue
			}
			return path[:pos], path[pos:]
		}
	}
	return path, ""
}

// There is no strings.IndexFold() function like there is strings.EqualFold(),
// but we can use strings.EqualFold() to build our own case-insensitive
// substring search (as of Go 1.14).
func indexFold(haystack, needle string) int {
	nlen := len(needle)
	for i := 0; i+nlen < len(haystack); i++ {
		if strings.EqualFold(haystack[i:i+nlen], needle) {
			return i
		}
	}
	return -1
}

const (
	tryPolicyFirstExist      = "first_exist"
	tryPolicyLargestSize     = "largest_size"
	tryPolicySmallestSize    = "smallest_size"
	tryPolicyMostRecentlyMod = "most_recently_modified"
)

// Interface guards
var (
	_ caddy.Validator          = (*MatchFile)(nil)
	_ caddyhttp.RequestMatcher = (*MatchFile)(nil)
)
<commit_msg>fileserver: Fix `file` matcher with empty `try_files` (#4147)<commit_after>// Copyright 2015 Matthew Holt and The Caddy Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package fileserver

import (
	"fmt"
	"net/http"
	"os"
	"path"
	"strings"
	"time"

	"github.com/caddyserver/caddy/v2"
	"github.com/caddyserver/caddy/v2/caddyconfig/caddyfile"
	"github.com/caddyserver/caddy/v2/modules/caddyhttp"
)

func init() {
	caddy.RegisterModule(MatchFile{})
}

// MatchFile is an HTTP request matcher that can match
// requests based upon file existence.
//
// Upon matching, three new placeholders will be made
// available:
//
// - `{http.matchers.file.relative}` The root-relative
// path of the file. This is often useful when rewriting
// requests.
// - `{http.matchers.file.absolute}` The absolute path
// of the matched file.
// - `{http.matchers.file.type}` Set to "directory" if
// the matched file is a directory, "file" otherwise.
// - `{http.matchers.file.remainder}` Set to the remainder
// of the path if the path was split by `split_path`.
type MatchFile struct {
	// The root directory, used for creating absolute
	// file paths, and required when working with
	// relative paths; if not specified, `{http.vars.root}`
	// will be used, if set; otherwise, the current
	// directory is assumed. Accepts placeholders.
	Root string `json:"root,omitempty"`

	// The list of files to try. Each path here is
	// considered related to Root. If nil, the request
	// URL's path will be assumed. Files and
	// directories are treated distinctly, so to match
	// a directory, the filepath MUST end in a forward
	// slash `/`. To match a regular file, there must
	// be no trailing slash. Accepts placeholders.
	TryFiles []string `json:"try_files,omitempty"`

	// How to choose a file in TryFiles. Can be:
	//
	// - first_exist
	// - smallest_size
	// - largest_size
	// - most_recently_modified
	//
	// Default is first_exist.
	TryPolicy string `json:"try_policy,omitempty"`

	// A list of delimiters to use to split the path in two
	// when trying files. If empty, no splitting will
	// occur, and the path will be tried as-is. For each
	// split value, the left-hand side of the split,
	// including the split value, will be the path tried.
	// For example, the path `/remote.php/dav/` using the
	// split value `.php` would try the file `/remote.php`.
	// Each delimiter must appear at the end of a URI path
	// component in order to be used as a split delimiter.
	SplitPath []string `json:"split_path,omitempty"`
}

// CaddyModule returns the Caddy module information.
func (MatchFile) CaddyModule() caddy.ModuleInfo {
	return caddy.ModuleInfo{
		ID:  "http.matchers.file",
		New: func() caddy.Module { return new(MatchFile) },
	}
}

// UnmarshalCaddyfile sets up the matcher from Caddyfile tokens. Syntax:
//
//     file <files...> {
//         root <path>
//         try_files <files...>
//         try_policy first_exist|smallest_size|largest_size|most_recently_modified
//     }
//
func (m *MatchFile) UnmarshalCaddyfile(d *caddyfile.Dispenser) error {
	for d.Next() {
		m.TryFiles = append(m.TryFiles, d.RemainingArgs()...)
		for d.NextBlock(0) {
			switch d.Val() {
			case "root":
				if !d.NextArg() {
					return d.ArgErr()
				}
				m.Root = d.Val()
			case "try_files":
				m.TryFiles = append(m.TryFiles, d.RemainingArgs()...)
				if len(m.TryFiles) == 0 {
					return d.ArgErr()
				}
			case "try_policy":
				if !d.NextArg() {
					return d.ArgErr()
				}
				m.TryPolicy = d.Val()
			case "split_path":
				m.SplitPath = d.RemainingArgs()
				if len(m.SplitPath) == 0 {
					return d.ArgErr()
				}
			default:
				return d.Errf("unrecognized subdirective: %s", d.Val())
			}
		}
	}
	return nil
}

// Provision sets up m's defaults.
func (m *MatchFile) Provision(_ caddy.Context) error {
	if m.Root == "" {
		m.Root = "{http.vars.root}"
	}
	// if list of files to try was omitted entirely, assume URL path
	// (use placeholder instead of r.URL.Path; see issue #4146)
	if m.TryFiles == nil {
		m.TryFiles = []string{"{http.request.uri.path}"}
	}
	return nil
}

// Validate ensures m has a valid configuration.
func (m MatchFile) Validate() error {
	switch m.TryPolicy {
	case "",
		tryPolicyFirstExist,
		tryPolicyLargestSize,
		tryPolicySmallestSize,
		tryPolicyMostRecentlyMod:
	default:
		return fmt.Errorf("unknown try policy %s", m.TryPolicy)
	}
	return nil
}

// Match returns true if r matches m. Returns true
// if a file was matched. If so, four placeholders
// will be available:
//    - http.matchers.file.relative
//    - http.matchers.file.absolute
//    - http.matchers.file.type
//    - http.matchers.file.remainder
func (m MatchFile) Match(r *http.Request) bool {
	return m.selectFile(r)
}

// selectFile chooses a file according to m.TryPolicy by appending
// the paths in m.TryFiles to m.Root, with placeholder replacements.
func (m MatchFile) selectFile(r *http.Request) (matched bool) {
	repl := r.Context().Value(caddy.ReplacerCtxKey).(*caddy.Replacer)

	root := repl.ReplaceAll(m.Root, ".")

	// common preparation of the file into parts
	prepareFilePath := func(file string) (suffix, fullpath, remainder string) {
		suffix, remainder = m.firstSplit(path.Clean(repl.ReplaceAll(file, "")))
		if strings.HasSuffix(file, "/") {
			suffix += "/"
		}
		fullpath = sanitizedPathJoin(root, suffix)
		return
	}

	// sets up the placeholders for the matched file
	setPlaceholders := func(info os.FileInfo, rel string, abs string, remainder string) {
		repl.Set("http.matchers.file.relative", rel)
		repl.Set("http.matchers.file.absolute", abs)
		repl.Set("http.matchers.file.remainder", remainder)

		fileType := "file"
		if info.IsDir() {
			fileType = "directory"
		}
		repl.Set("http.matchers.file.type", fileType)
	}

	switch m.TryPolicy {
	case "", tryPolicyFirstExist:
		for _, f := range m.TryFiles {
			suffix, fullpath, remainder := prepareFilePath(f)
			if info, exists := strictFileExists(fullpath); exists {
				setPlaceholders(info, suffix, fullpath, remainder)
				return true
			}
		}

	case tryPolicyLargestSize:
		var largestSize int64
		var largestFilename string
		var largestSuffix string
		var remainder string
		var info os.FileInfo
		for _, f := range m.TryFiles {
			suffix, fullpath, splitRemainder := prepareFilePath(f)
			info, err := os.Stat(fullpath)
			if err == nil && info.Size() > largestSize {
				largestSize = info.Size()
				largestFilename = fullpath
				largestSuffix = suffix
				remainder = splitRemainder
			}
		}
		setPlaceholders(info, largestSuffix, largestFilename, remainder)
		return true

	case tryPolicySmallestSize:
		var smallestSize int64
		var smallestFilename string
		var smallestSuffix string
		var remainder string
		var info os.FileInfo
		for _, f := range m.TryFiles {
			suffix, fullpath, splitRemainder := prepareFilePath(f)
			info, err := os.Stat(fullpath)
			if err == nil && (smallestSize == 0 || info.Size() < smallestSize) {
				smallestSize = info.Size()
				smallestFilename = fullpath
				smallestSuffix = suffix
				remainder = splitRemainder
			}
		}
		setPlaceholders(info, smallestSuffix, smallestFilename, remainder)
		return true

	case tryPolicyMostRecentlyMod:
		var recentDate time.Time
		var recentFilename string
		var recentSuffix string
		var remainder string
		var info os.FileInfo
		for _, f := range m.TryFiles {
			suffix, fullpath, splitRemainder := prepareFilePath(f)
			info, err := os.Stat(fullpath)
			if err == nil &&
				(recentDate.IsZero() || info.ModTime().After(recentDate)) {
				recentDate = info.ModTime()
				recentFilename = fullpath
				recentSuffix = suffix
				remainder = splitRemainder
			}
		}
		setPlaceholders(info, recentSuffix, recentFilename, remainder)
		return true
	}

	return
}

// strictFileExists returns true if file exists
// and matches the convention of the given file
// path. If the path ends in a forward slash,
// the file must also be a directory; if it does
// NOT end in a forward slash, the file must NOT
// be a directory.
func strictFileExists(file string) (os.FileInfo, bool) {
	stat, err := os.Stat(file)
	if err != nil {
		// in reality, this can be any error
		// such as permission or even obscure
		// ones like "is not a directory" (when
		// trying to stat a file within a file);
		// in those cases we can't be sure if
		// the file exists, so we just treat any
		// error as if it does not exist; see
		// https://stackoverflow.com/a/12518877/1048862
		return nil, false
	}
	if strings.HasSuffix(file, separator) {
		// by convention, file paths ending
		// in a path separator must be a directory
		return stat, stat.IsDir()
	}
	// by convention, file paths NOT ending
	// in a path separator must NOT be a directory
	return stat, !stat.IsDir()
}

// firstSplit returns the first result where the path
// can be split in two by a value in m.SplitPath. The
// return values are the first piece of the path that
// ends with the split substring and the remainder.
// If the path cannot be split, the path is returned
// as-is (with no remainder).
func (m MatchFile) firstSplit(path string) (splitPart, remainder string) {
	for _, split := range m.SplitPath {
		if idx := indexFold(path, split); idx > -1 {
			pos := idx + len(split)
			// skip the split if it's not the final part of the filename
			if pos != len(path) && !strings.HasPrefix(path[pos:], "/") {
				continue
			}
			return path[:pos], path[pos:]
		}
	}
	return path, ""
}

// There is no strings.IndexFold() function like there is strings.EqualFold(),
// but we can use strings.EqualFold() to build our own case-insensitive
// substring search (as of Go 1.14).
func indexFold(haystack, needle string) int {
	nlen := len(needle)
	for i := 0; i+nlen < len(haystack); i++ {
		if strings.EqualFold(haystack[i:i+nlen], needle) {
			return i
		}
	}
	return -1
}

const (
	tryPolicyFirstExist      = "first_exist"
	tryPolicyLargestSize     = "largest_size"
	tryPolicySmallestSize    = "smallest_size"
	tryPolicyMostRecentlyMod = "most_recently_modified"
)

// Interface guards
var (
	_ caddy.Validator          = (*MatchFile)(nil)
	_ caddyhttp.RequestMatcher = (*MatchFile)(nil)
)
<|endoftext|>


<commit_before>package joblogger

import (
	"fmt"
	"net/http"
	"os"
	"time"

	"github.com/garyburd/redigo/redis"
	"github.com/octoblu/vulcand-job-logger/wrapper"
)

var redisConnections map[string]redis.Conn

// Handler implements http.Handler
type Handler struct {
	redisURI       string
	redisQueueName string
	backendID      string
	next           http.Handler
}

// NewHandler constructs a new handler
func NewHandler(redisURI, redisQueueName, backendID string, next http.Handler) *Handler {
	return &Handler{redisURI, redisQueueName, backendID, next}
}

// ServeHTTP will be called each time the request
// hits the location with this middleware activated
func (handler *Handler) ServeHTTP(rw http.ResponseWriter, r *http.Request) {
	redisChannel := make(chan []byte)

	go handler.logRequest(redisChannel)
	wrapped := wrapper.New(rw, redisChannel, time.Now(), handler.backendID)
	handler.next.ServeHTTP(wrapped, r)
}

func (handler *Handler) logRequest(logChannel chan []byte) {
	redisConn, err := handler.redisConn()
	if err != nil {
		logError("handler.redisConn() failed: %v\n", err)
		return
	}
	logEntryBytes := <-logChannel
	_, err = redisConn.Do("LPUSH", handler.redisQueueName, logEntryBytes)
	logError("Redis LPUSH failed: %v\n", err)
}

func (handler *Handler) redisConn() (redis.Conn, error) {
	key := fmt.Sprintf("%v/%v", handler.redisURI, handler.redisQueueName)
	conn, ok := redisConnections[key]
	if ok {
		return conn, nil
	}

	conn, err := redis.DialURL(handler.redisURI)
	if err != nil {
		return nil, err
	}

	redisConnections[key] = conn
	return conn, nil
}

func logError(fmtMessage string, err error) {
	if err == nil {
		return
	}
	fmt.Fprintf(os.Stderr, fmtMessage, err.Error())
}
<commit_msg>Do the logging<commit_after>package joblogger

import (
	"fmt"
	"net/http"
	"os"
	"time"

	"github.com/garyburd/redigo/redis"
	"github.com/octoblu/vulcand-job-logger/wrapper"
)

var redisConnections map[string]redis.Conn

// Handler implements http.Handler
type Handler struct {
	redisURI       string
	redisQueueName string
	backendID      string
	next           http.Handler
}

// NewHandler constructs a new handler
func NewHandler(redisURI, redisQueueName, backendID string, next http.Handler) *Handler {
	return &Handler{redisURI, redisQueueName, backendID, next}
}

// ServeHTTP will be called each time the request
// hits the location with this middleware activated
func (handler *Handler) ServeHTTP(rw http.ResponseWriter, r *http.Request) {
	redisChannel := make(chan []byte)

	go handler.logRequest(redisChannel)
	wrapped := wrapper.New(rw, redisChannel, time.Now(), handler.backendID)
	handler.next.ServeHTTP(wrapped, r)
}

func (handler *Handler) logRequest(logChannel chan []byte) {
	redisConn, err := handler.redisConn()
	if err != nil {
		logError("handler.redisConn() failed: %v\n", err)
		return
	}
	logEntryBytes := <-logChannel
	_, err = redisConn.Do("LPUSH", handler.redisQueueName, logEntryBytes)
	logError("Redis LPUSH failed: %v\n", err)
}

func (handler *Handler) redisConn() (redis.Conn, error) {
	if redisConnections == nil {
		redisConnections = make(map[string]redis.Conn)
	}

	key := fmt.Sprintf("%v/%v", handler.redisURI, handler.redisQueueName)
	conn, ok := redisConnections[key]
	if ok {
		return conn, nil
	}

	conn, err := redis.DialURL(handler.redisURI)
	if err != nil {
		return nil, err
	}

	redisConnections[key] = conn
	return conn, nil
}

func logError(fmtMessage string, err error) {
	if err == nil {
		return
	}
	fmt.Fprintf(os.Stderr, fmtMessage, err.Error())
}
<|endoftext|>


<commit_before>// Copyright 2014 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 trace

import (
	"sort"
	"strings"
)

// GDesc contains statistics and execution details of a single goroutine.
type GDesc struct {
	ID           uint64
	Name         string
	PC           uint64
	CreationTime int64
	StartTime    int64
	EndTime      int64

	// List of regions in the goroutine, sorted based on the start time.
	Regions []*UserRegionDesc

	// Statistics of execution time during the goroutine execution.
	GExecutionStat

	*gdesc // private part.
}

// UserRegionDesc represents a region and goroutine execution stats
// while the region was active.
type UserRegionDesc struct {
	TaskID uint64
	Name   string

	// Region start event. Normally EvUserRegion start event or nil,
	// but can be EvGoCreate event if the region is a synthetic
	// region representing task inheritance from the parent goroutine.
	Start *Event

	// Region end event. Normally EvUserRegion end event or nil,
	// but can be EvGoStop or EvGoEnd event if the goroutine
	// terminated without explicitly ending the region.
	End *Event

	GExecutionStat
}

// GExecutionStat contains statistics about a goroutine's execution
// during a period of time.
type GExecutionStat struct {
	ExecTime      int64
	SchedWaitTime int64
	IOTime        int64
	BlockTime     int64
	SyscallTime   int64
	GCTime        int64
	SweepTime     int64
	TotalTime     int64
}

// sub returns the stats v-s.
func (s GExecutionStat) sub(v GExecutionStat) (r GExecutionStat) {
	r = s
	r.ExecTime -= v.ExecTime
	r.SchedWaitTime -= v.SchedWaitTime
	r.IOTime -= v.IOTime
	r.BlockTime -= v.BlockTime
	r.SyscallTime -= v.SyscallTime
	r.GCTime -= v.GCTime
	r.SweepTime -= v.SweepTime
	r.TotalTime -= v.TotalTime
	return r
}

// snapshotStat returns the snapshot of the goroutine execution statistics.
// This is called as we process the ordered trace event stream. lastTs and
// activeGCStartTime are used to process pending statistics if this is called
// before any goroutine end event.
func (g *GDesc) snapshotStat(lastTs, activeGCStartTime int64) (ret GExecutionStat) {
	ret = g.GExecutionStat

	if g.gdesc == nil {
		return ret // finalized GDesc. No pending state.
	}

	if activeGCStartTime != 0 { // terminating while GC is active
		if g.CreationTime < activeGCStartTime {
			ret.GCTime += lastTs - activeGCStartTime
		} else {
			// The goroutine's lifetime completely overlaps
			// with a GC.
			ret.GCTime += lastTs - g.CreationTime
		}
	}

	if g.TotalTime == 0 {
		ret.TotalTime = lastTs - g.CreationTime
	}

	if g.lastStartTime != 0 {
		ret.ExecTime += lastTs - g.lastStartTime
	}
	if g.blockNetTime != 0 {
		ret.IOTime += lastTs - g.blockNetTime
	}
	if g.blockSyncTime != 0 {
		ret.BlockTime += lastTs - g.blockSyncTime
	}
	if g.blockSyscallTime != 0 {
		ret.SyscallTime += lastTs - g.blockSyscallTime
	}
	if g.blockSchedTime != 0 {
		ret.SchedWaitTime += lastTs - g.blockSchedTime
	}
	if g.blockSweepTime != 0 {
		ret.SweepTime += lastTs - g.blockSweepTime
	}
	return ret
}

// finalize is called when processing a goroutine end event or at
// the end of trace processing. This finalizes the execution stat
// and any active regions in the goroutine, in which case trigger is nil.
func (g *GDesc) finalize(lastTs, activeGCStartTime int64, trigger *Event) {
	if trigger != nil {
		g.EndTime = trigger.Ts
	}
	finalStat := g.snapshotStat(lastTs, activeGCStartTime)

	g.GExecutionStat = finalStat

	// System goroutines are never part of regions, even though they
	// "inherit" a task due to creation (EvGoCreate) from within a region.
	// This may happen e.g. if the first GC is triggered within a region,
	// starting the GC worker goroutines.
	if !IsSystemGoroutine(g.Name) {
		for _, s := range g.activeRegions {
			s.End = trigger
			s.GExecutionStat = finalStat.sub(s.GExecutionStat)
			g.Regions = append(g.Regions, s)
		}
	}
	*(g.gdesc) = gdesc{}
}

// gdesc is a private part of GDesc that is required only during analysis.
type gdesc struct {
	lastStartTime    int64
	blockNetTime     int64
	blockSyncTime    int64
	blockSyscallTime int64
	blockSweepTime   int64
	blockGCTime      int64
	blockSchedTime   int64

	activeRegions []*UserRegionDesc // stack of active regions
}

// GoroutineStats generates statistics for all goroutines in the trace.
func GoroutineStats(events []*Event) map[uint64]*GDesc {
	gs := make(map[uint64]*GDesc)
	var lastTs int64
	var gcStartTime int64 // gcStartTime == 0 indicates gc is inactive.
	for _, ev := range events {
		lastTs = ev.Ts
		switch ev.Type {
		case EvGoCreate:
			g := &GDesc{ID: ev.Args[0], CreationTime: ev.Ts, gdesc: new(gdesc)}
			g.blockSchedTime = ev.Ts
			// When a goroutine is newly created, inherit the task
			// of the active region. For ease handling of this
			// case, we create a fake region description with the
			// task id. This isn't strictly necessary as this
			// goroutine may not be assosciated with the task, but
			// it can be convenient to see all children created
			// during a region.
			if creatorG := gs[ev.G]; creatorG != nil && len(creatorG.gdesc.activeRegions) > 0 {
				regions := creatorG.gdesc.activeRegions
				s := regions[len(regions)-1]
				if s.TaskID != 0 {
					g.gdesc.activeRegions = []*UserRegionDesc{
						{TaskID: s.TaskID, Start: ev},
					}
				}
			}
			gs[g.ID] = g
		case EvGoStart, EvGoStartLabel:
			g := gs[ev.G]
			if g.PC == 0 {
				g.PC = ev.Stk[0].PC
				g.Name = ev.Stk[0].Fn
			}
			g.lastStartTime = ev.Ts
			if g.StartTime == 0 {
				g.StartTime = ev.Ts
			}
			if g.blockSchedTime != 0 {
				g.SchedWaitTime += ev.Ts - g.blockSchedTime
				g.blockSchedTime = 0
			}
		case EvGoEnd, EvGoStop:
			g := gs[ev.G]
			g.finalize(ev.Ts, gcStartTime, ev)
		case EvGoBlockSend, EvGoBlockRecv, EvGoBlockSelect,
			EvGoBlockSync, EvGoBlockCond:
			g := gs[ev.G]
			g.ExecTime += ev.Ts - g.lastStartTime
			g.lastStartTime = 0
			g.blockSyncTime = ev.Ts
		case EvGoSched, EvGoPreempt:
			g := gs[ev.G]
			g.ExecTime += ev.Ts - g.lastStartTime
			g.lastStartTime = 0
			g.blockSchedTime = ev.Ts
		case EvGoSleep, EvGoBlock:
			g := gs[ev.G]
			g.ExecTime += ev.Ts - g.lastStartTime
			g.lastStartTime = 0
		case EvGoBlockNet:
			g := gs[ev.G]
			g.ExecTime += ev.Ts - g.lastStartTime
			g.lastStartTime = 0
			g.blockNetTime = ev.Ts
		case EvGoBlockGC:
			g := gs[ev.G]
			g.ExecTime += ev.Ts - g.lastStartTime
			g.lastStartTime = 0
			g.blockGCTime = ev.Ts
		case EvGoUnblock:
			g := gs[ev.Args[0]]
			if g.blockNetTime != 0 {
				g.IOTime += ev.Ts - g.blockNetTime
				g.blockNetTime = 0
			}
			if g.blockSyncTime != 0 {
				g.BlockTime += ev.Ts - g.blockSyncTime
				g.blockSyncTime = 0
			}
			g.blockSchedTime = ev.Ts
		case EvGoSysBlock:
			g := gs[ev.G]
			g.ExecTime += ev.Ts - g.lastStartTime
			g.lastStartTime = 0
			g.blockSyscallTime = ev.Ts
		case EvGoSysExit:
			g := gs[ev.G]
			if g.blockSyscallTime != 0 {
				g.SyscallTime += ev.Ts - g.blockSyscallTime
				g.blockSyscallTime = 0
			}
			g.blockSchedTime = ev.Ts
		case EvGCSweepStart:
			g := gs[ev.G]
			if g != nil {
				// Sweep can happen during GC on system goroutine.
				g.blockSweepTime = ev.Ts
			}
		case EvGCSweepDone:
			g := gs[ev.G]
			if g != nil && g.blockSweepTime != 0 {
				g.SweepTime += ev.Ts - g.blockSweepTime
				g.blockSweepTime = 0
			}
		case EvGCStart:
			gcStartTime = ev.Ts
		case EvGCDone:
			for _, g := range gs {
				if g.EndTime != 0 {
					continue
				}
				if gcStartTime < g.CreationTime {
					g.GCTime += ev.Ts - g.CreationTime
				} else {
					g.GCTime += ev.Ts - gcStartTime
				}
			}
			gcStartTime = 0 // indicates gc is inactive.
		case EvUserRegion:
			g := gs[ev.G]
			switch mode := ev.Args[1]; mode {
			case 0: // region start
				g.activeRegions = append(g.activeRegions, &UserRegionDesc{
					Name:           ev.SArgs[0],
					TaskID:         ev.Args[0],
					Start:          ev,
					GExecutionStat: g.snapshotStat(lastTs, gcStartTime),
				})
			case 1: // region end
				var sd *UserRegionDesc
				if regionStk := g.activeRegions; len(regionStk) > 0 {
					n := len(regionStk)
					sd = regionStk[n-1]
					regionStk = regionStk[:n-1] // pop
					g.activeRegions = regionStk
				} else {
					sd = &UserRegionDesc{
						Name:   ev.SArgs[0],
						TaskID: ev.Args[0],
					}
				}
				sd.GExecutionStat = g.snapshotStat(lastTs, gcStartTime).sub(sd.GExecutionStat)
				sd.End = ev
				g.Regions = append(g.Regions, sd)
			}
		}
	}

	for _, g := range gs {
		g.finalize(lastTs, gcStartTime, nil)

		// sort based on region start time
		sort.Slice(g.Regions, func(i, j int) bool {
			x := g.Regions[i].Start
			y := g.Regions[j].Start
			if x == nil {
				return true
			}
			if y == nil {
				return false
			}
			return x.Ts < y.Ts
		})

		g.gdesc = nil
	}

	return gs
}

// RelatedGoroutines finds a set of goroutines related to goroutine goid.
func RelatedGoroutines(events []*Event, goid uint64) map[uint64]bool {
	// BFS of depth 2 over "unblock" edges
	// (what goroutines unblock goroutine goid?).
	gmap := make(map[uint64]bool)
	gmap[goid] = true
	for i := 0; i < 2; i++ {
		gmap1 := make(map[uint64]bool)
		for g := range gmap {
			gmap1[g] = true
		}
		for _, ev := range events {
			if ev.Type == EvGoUnblock && gmap[ev.Args[0]] {
				gmap1[ev.G] = true
			}
		}
		gmap = gmap1
	}
	gmap[0] = true // for GC events
	return gmap
}

func IsSystemGoroutine(entryFn string) bool {
	// This mimics runtime.isSystemGoroutine as closely as
	// possible.
	return entryFn != "runtime.main" && strings.HasPrefix(entryFn, "runtime.")
}
<commit_msg>internal/trace: fix typo in goroutines.go<commit_after>// Copyright 2014 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 trace

import (
	"sort"
	"strings"
)

// GDesc contains statistics and execution details of a single goroutine.
type GDesc struct {
	ID           uint64
	Name         string
	PC           uint64
	CreationTime int64
	StartTime    int64
	EndTime      int64

	// List of regions in the goroutine, sorted based on the start time.
	Regions []*UserRegionDesc

	// Statistics of execution time during the goroutine execution.
	GExecutionStat

	*gdesc // private part.
}

// UserRegionDesc represents a region and goroutine execution stats
// while the region was active.
type UserRegionDesc struct {
	TaskID uint64
	Name   string

	// Region start event. Normally EvUserRegion start event or nil,
	// but can be EvGoCreate event if the region is a synthetic
	// region representing task inheritance from the parent goroutine.
	Start *Event

	// Region end event. Normally EvUserRegion end event or nil,
	// but can be EvGoStop or EvGoEnd event if the goroutine
	// terminated without explicitly ending the region.
	End *Event

	GExecutionStat
}

// GExecutionStat contains statistics about a goroutine's execution
// during a period of time.
type GExecutionStat struct {
	ExecTime      int64
	SchedWaitTime int64
	IOTime        int64
	BlockTime     int64
	SyscallTime   int64
	GCTime        int64
	SweepTime     int64
	TotalTime     int64
}

// sub returns the stats v-s.
func (s GExecutionStat) sub(v GExecutionStat) (r GExecutionStat) {
	r = s
	r.ExecTime -= v.ExecTime
	r.SchedWaitTime -= v.SchedWaitTime
	r.IOTime -= v.IOTime
	r.BlockTime -= v.BlockTime
	r.SyscallTime -= v.SyscallTime
	r.GCTime -= v.GCTime
	r.SweepTime -= v.SweepTime
	r.TotalTime -= v.TotalTime
	return r
}

// snapshotStat returns the snapshot of the goroutine execution statistics.
// This is called as we process the ordered trace event stream. lastTs and
// activeGCStartTime are used to process pending statistics if this is called
// before any goroutine end event.
func (g *GDesc) snapshotStat(lastTs, activeGCStartTime int64) (ret GExecutionStat) {
	ret = g.GExecutionStat

	if g.gdesc == nil {
		return ret // finalized GDesc. No pending state.
	}

	if activeGCStartTime != 0 { // terminating while GC is active
		if g.CreationTime < activeGCStartTime {
			ret.GCTime += lastTs - activeGCStartTime
		} else {
			// The goroutine's lifetime completely overlaps
			// with a GC.
			ret.GCTime += lastTs - g.CreationTime
		}
	}

	if g.TotalTime == 0 {
		ret.TotalTime = lastTs - g.CreationTime
	}

	if g.lastStartTime != 0 {
		ret.ExecTime += lastTs - g.lastStartTime
	}
	if g.blockNetTime != 0 {
		ret.IOTime += lastTs - g.blockNetTime
	}
	if g.blockSyncTime != 0 {
		ret.BlockTime += lastTs - g.blockSyncTime
	}
	if g.blockSyscallTime != 0 {
		ret.SyscallTime += lastTs - g.blockSyscallTime
	}
	if g.blockSchedTime != 0 {
		ret.SchedWaitTime += lastTs - g.blockSchedTime
	}
	if g.blockSweepTime != 0 {
		ret.SweepTime += lastTs - g.blockSweepTime
	}
	return ret
}

// finalize is called when processing a goroutine end event or at
// the end of trace processing. This finalizes the execution stat
// and any active regions in the goroutine, in which case trigger is nil.
func (g *GDesc) finalize(lastTs, activeGCStartTime int64, trigger *Event) {
	if trigger != nil {
		g.EndTime = trigger.Ts
	}
	finalStat := g.snapshotStat(lastTs, activeGCStartTime)

	g.GExecutionStat = finalStat

	// System goroutines are never part of regions, even though they
	// "inherit" a task due to creation (EvGoCreate) from within a region.
	// This may happen e.g. if the first GC is triggered within a region,
	// starting the GC worker goroutines.
	if !IsSystemGoroutine(g.Name) {
		for _, s := range g.activeRegions {
			s.End = trigger
			s.GExecutionStat = finalStat.sub(s.GExecutionStat)
			g.Regions = append(g.Regions, s)
		}
	}
	*(g.gdesc) = gdesc{}
}

// gdesc is a private part of GDesc that is required only during analysis.
type gdesc struct {
	lastStartTime    int64
	blockNetTime     int64
	blockSyncTime    int64
	blockSyscallTime int64
	blockSweepTime   int64
	blockGCTime      int64
	blockSchedTime   int64

	activeRegions []*UserRegionDesc // stack of active regions
}

// GoroutineStats generates statistics for all goroutines in the trace.
func GoroutineStats(events []*Event) map[uint64]*GDesc {
	gs := make(map[uint64]*GDesc)
	var lastTs int64
	var gcStartTime int64 // gcStartTime == 0 indicates gc is inactive.
	for _, ev := range events {
		lastTs = ev.Ts
		switch ev.Type {
		case EvGoCreate:
			g := &GDesc{ID: ev.Args[0], CreationTime: ev.Ts, gdesc: new(gdesc)}
			g.blockSchedTime = ev.Ts
			// When a goroutine is newly created, inherit the task
			// of the active region. For ease handling of this
			// case, we create a fake region description with the
			// task id. This isn't strictly necessary as this
			// goroutine may not be associated with the task, but
			// it can be convenient to see all children created
			// during a region.
			if creatorG := gs[ev.G]; creatorG != nil && len(creatorG.gdesc.activeRegions) > 0 {
				regions := creatorG.gdesc.activeRegions
				s := regions[len(regions)-1]
				if s.TaskID != 0 {
					g.gdesc.activeRegions = []*UserRegionDesc{
						{TaskID: s.TaskID, Start: ev},
					}
				}
			}
			gs[g.ID] = g
		case EvGoStart, EvGoStartLabel:
			g := gs[ev.G]
			if g.PC == 0 {
				g.PC = ev.Stk[0].PC
				g.Name = ev.Stk[0].Fn
			}
			g.lastStartTime = ev.Ts
			if g.StartTime == 0 {
				g.StartTime = ev.Ts
			}
			if g.blockSchedTime != 0 {
				g.SchedWaitTime += ev.Ts - g.blockSchedTime
				g.blockSchedTime = 0
			}
		case EvGoEnd, EvGoStop:
			g := gs[ev.G]
			g.finalize(ev.Ts, gcStartTime, ev)
		case EvGoBlockSend, EvGoBlockRecv, EvGoBlockSelect,
			EvGoBlockSync, EvGoBlockCond:
			g := gs[ev.G]
			g.ExecTime += ev.Ts - g.lastStartTime
			g.lastStartTime = 0
			g.blockSyncTime = ev.Ts
		case EvGoSched, EvGoPreempt:
			g := gs[ev.G]
			g.ExecTime += ev.Ts - g.lastStartTime
			g.lastStartTime = 0
			g.blockSchedTime = ev.Ts
		case EvGoSleep, EvGoBlock:
			g := gs[ev.G]
			g.ExecTime += ev.Ts - g.lastStartTime
			g.lastStartTime = 0
		case EvGoBlockNet:
			g := gs[ev.G]
			g.ExecTime += ev.Ts - g.lastStartTime
			g.lastStartTime = 0
			g.blockNetTime = ev.Ts
		case EvGoBlockGC:
			g := gs[ev.G]
			g.ExecTime += ev.Ts - g.lastStartTime
			g.lastStartTime = 0
			g.blockGCTime = ev.Ts
		case EvGoUnblock:
			g := gs[ev.Args[0]]
			if g.blockNetTime != 0 {
				g.IOTime += ev.Ts - g.blockNetTime
				g.blockNetTime = 0
			}
			if g.blockSyncTime != 0 {
				g.BlockTime += ev.Ts - g.blockSyncTime
				g.blockSyncTime = 0
			}
			g.blockSchedTime = ev.Ts
		case EvGoSysBlock:
			g := gs[ev.G]
			g.ExecTime += ev.Ts - g.lastStartTime
			g.lastStartTime = 0
			g.blockSyscallTime = ev.Ts
		case EvGoSysExit:
			g := gs[ev.G]
			if g.blockSyscallTime != 0 {
				g.SyscallTime += ev.Ts - g.blockSyscallTime
				g.blockSyscallTime = 0
			}
			g.blockSchedTime = ev.Ts
		case EvGCSweepStart:
			g := gs[ev.G]
			if g != nil {
				// Sweep can happen during GC on system goroutine.
				g.blockSweepTime = ev.Ts
			}
		case EvGCSweepDone:
			g := gs[ev.G]
			if g != nil && g.blockSweepTime != 0 {
				g.SweepTime += ev.Ts - g.blockSweepTime
				g.blockSweepTime = 0
			}
		case EvGCStart:
			gcStartTime = ev.Ts
		case EvGCDone:
			for _, g := range gs {
				if g.EndTime != 0 {
					continue
				}
				if gcStartTime < g.CreationTime {
					g.GCTime += ev.Ts - g.CreationTime
				} else {
					g.GCTime += ev.Ts - gcStartTime
				}
			}
			gcStartTime = 0 // indicates gc is inactive.
		case EvUserRegion:
			g := gs[ev.G]
			switch mode := ev.Args[1]; mode {
			case 0: // region start
				g.activeRegions = append(g.activeRegions, &UserRegionDesc{
					Name:           ev.SArgs[0],
					TaskID:         ev.Args[0],
					Start:          ev,
					GExecutionStat: g.snapshotStat(lastTs, gcStartTime),
				})
			case 1: // region end
				var sd *UserRegionDesc
				if regionStk := g.activeRegions; len(regionStk) > 0 {
					n := len(regionStk)
					sd = regionStk[n-1]
					regionStk = regionStk[:n-1] // pop
					g.activeRegions = regionStk
				} else {
					sd = &UserRegionDesc{
						Name:   ev.SArgs[0],
						TaskID: ev.Args[0],
					}
				}
				sd.GExecutionStat = g.snapshotStat(lastTs, gcStartTime).sub(sd.GExecutionStat)
				sd.End = ev
				g.Regions = append(g.Regions, sd)
			}
		}
	}

	for _, g := range gs {
		g.finalize(lastTs, gcStartTime, nil)

		// sort based on region start time
		sort.Slice(g.Regions, func(i, j int) bool {
			x := g.Regions[i].Start
			y := g.Regions[j].Start
			if x == nil {
				return true
			}
			if y == nil {
				return false
			}
			return x.Ts < y.Ts
		})

		g.gdesc = nil
	}

	return gs
}

// RelatedGoroutines finds a set of goroutines related to goroutine goid.
func RelatedGoroutines(events []*Event, goid uint64) map[uint64]bool {
	// BFS of depth 2 over "unblock" edges
	// (what goroutines unblock goroutine goid?).
	gmap := make(map[uint64]bool)
	gmap[goid] = true
	for i := 0; i < 2; i++ {
		gmap1 := make(map[uint64]bool)
		for g := range gmap {
			gmap1[g] = true
		}
		for _, ev := range events {
			if ev.Type == EvGoUnblock && gmap[ev.Args[0]] {
				gmap1[ev.G] = true
			}
		}
		gmap = gmap1
	}
	gmap[0] = true // for GC events
	return gmap
}

func IsSystemGoroutine(entryFn string) bool {
	// This mimics runtime.isSystemGoroutine as closely as
	// possible.
	return entryFn != "runtime.main" && strings.HasPrefix(entryFn, "runtime.")
}
<|endoftext|>


<commit_before>package netsec

import (
	"errors"
	"sync"
)

type ContainerSecurity struct {
	veth           string
	mark           int
	ID             string
	Pid            int
	SecurityGroups map[string]uint16 // ipgroup name -> port
}

type NetworkSecurity struct {
	sync.Mutex
	deniedIPs  map[string]bool              // list of denied IPs. map for easy existence check
	ipGroups   map[string][]string          // group name -> list of infrastructure IPs to blanket deny
	containers map[string]*ContainerSecurty // container id -> ContainerSecurity
}

func (n *NetworkSecurity) UpdateIPGroup(name string, ips []string) error {
	n.Lock()
	defer n.Unlock()
	current, exists := n.ipGroups[ips]
	toRemove := []string{}
	if exists {
		// figure out what is being removed
		incomingMap := map[string]bool{}
		for _, ip := range ips {
			incomingMap[ip] = true
		}
		for _, ip := range current {
			if !incomingMap[ip] {
				toRemove = append(toRemove, ip)
			}
		}
	}
	// add blanket deny rule for new IPs
	newIPs := []string{}
	for _, ip := range ips {
		if deniedIPs[ip] {
			// already exists
			continue
		}
		if err := n.blanketDeny(true, ip); err != nil {
			return err
		}
		deniedIPs[ip] = true
		newIPs = append(newIPs, ip)
	}
	// remove blanket deny rule for removed IPs
	for _, ip := range toRemove {
		if !deniedIPs[ip] {
			// already removed
			continue
		}
		n.blanketDeny(false, ip)
		delete(deniedIPs, ip)
	}
	// add/remove forward rules for new IPs for everything that uses the name
	for _, contSec := range n.containers {
		port, exists := contSec.SecurityGroups[name]
		if !exists {
			// this container does not use this ip group
			continue
		}
		// add new ips
		for _, ip := range newIPs {
			n.openPort(true, contSec.mark, ip, port)
		}
		// remove old ips
		for _, ip := range toRemove {
			n.openPort(false, contSec.mark, ip, port)
		}
	}
	// update ipGroups
	n.ipGroups[name] = ips
	return nil
}

func (n *NetworkSecurity) DeleteIPGroup(name string) error {
	if err := n.UpdateIPGroup(name, []string{}); err != nil {
		return err
	}
	n.Lock()
	defer n.Unlock()
	delete(n.ipGroups, name)
	return nil
}

func (n *NetworkSecurity) AddContainerSecurity(id string, pid int, sgs map[string]uint16) error {
	n.Lock()
	defer n.Unlock()
	if _, exists := n.containers[id]; exists {
		// we already have security set up for this id. don't do it and return an error.
		return errors.New("Container "+id+" already has Network Security set up.")
	}
	// make sure all groups exist
	for group, _ := range sgs {
		_, exists := n.ipGroups[group]
		if !exists {
			return errors.New("IP Group "+group+" does not exist")
		}
	}

	// fetch network info
	contSec := &ContainerSecurity{
		ID: id,
		Pid: pid,
		SecurityGroups: sgs
	}
	if err := contSec.GetNetworkInfo(); err != nil {
		return err
	}
	// create prerouting mark
	if err := n.markVeth(true, contSec.veth, contSec.mark); err != nil {
		return err
	}
	// add forward rules
	for group, port := range sgs {
		ips := n.ipGroups[group]
		for _, ip := range ips {
			if err := n.openPort(true, contSec.mark, ip, port); err != nil {
				defer n.RemoveContainerSecurity(id) // cleanup created references when we error out
				return err
			}
		}
	}
}

func (n *NetworkSecurity) RemoveContainerSecurity(id string) error {
	n.Lock()
	defer n.Unlock()
	var (
		contSec *ContainerSecurity
		exists bool
	)
	if contSec, exists = n.containers[id]; !exists {
		// no container security here, nothing to remove
		return nil
	}
	// remove prerouting mark
	n.markVeth(false, contSec.veth, contSec.mark)
	// remove forward rules
	for group, port := range sgs {
		ips := n.ipGroups[group]
		for _, ip := range ips {
			n.openPort(false, contSec.mark, ip, port)
		}
	}
	return nil
}

func (n *NetworkSecurity) blanketDeny(add bool, ip string) error {
	// TODO[jigish] implement
	return errors.New("unimplemented")
}

func (n *NetworkSecurity) openPort(add bool, mark int, ip string, port uint16) error {
	// TODO[jigish] implement
	return errors.New("unimplemented")
}

func (n *NetworkSecurity) markVeth(add bool, veth string, mark int) error {
	// TODO[jigish] implement
	return errors.New("unimplemented")
}
<commit_msg>add skeleton GetNetworkInfo<commit_after>package netsec

import (
	"errors"
	"sync"
)

type ContainerSecurity struct {
	veth           string
	mark           int
	ID             string
	Pid            int
	SecurityGroups map[string]uint16 // ipgroup name -> port
}

func (c *ContainerSecurity) GetNetworkInfo() error {
	// this will call out to guano to figure out what the veth and mark should be
	// TODO[jigish] implement
	return errors.New("unimplemented")
}

type NetworkSecurity struct {
	sync.Mutex
	deniedIPs  map[string]bool              // list of denied IPs. map for easy existence check
	ipGroups   map[string][]string          // group name -> list of infrastructure IPs to blanket deny
	containers map[string]*ContainerSecurty // container id -> ContainerSecurity
}

func (n *NetworkSecurity) UpdateIPGroup(name string, ips []string) error {
	n.Lock()
	defer n.Unlock()
	current, exists := n.ipGroups[ips]
	toRemove := []string{}
	if exists {
		// figure out what is being removed
		incomingMap := map[string]bool{}
		for _, ip := range ips {
			incomingMap[ip] = true
		}
		for _, ip := range current {
			if !incomingMap[ip] {
				toRemove = append(toRemove, ip)
			}
		}
	}
	// add blanket deny rule for new IPs
	newIPs := []string{}
	for _, ip := range ips {
		if deniedIPs[ip] {
			// already exists
			continue
		}
		if err := n.blanketDeny(true, ip); err != nil {
			return err
		}
		deniedIPs[ip] = true
		newIPs = append(newIPs, ip)
	}
	// remove blanket deny rule for removed IPs
	for _, ip := range toRemove {
		if !deniedIPs[ip] {
			// already removed
			continue
		}
		n.blanketDeny(false, ip)
		delete(deniedIPs, ip)
	}
	// add/remove forward rules for new IPs for everything that uses the name
	for _, contSec := range n.containers {
		port, exists := contSec.SecurityGroups[name]
		if !exists {
			// this container does not use this ip group
			continue
		}
		// add new ips
		for _, ip := range newIPs {
			n.openPort(true, contSec.mark, ip, port)
		}
		// remove old ips
		for _, ip := range toRemove {
			n.openPort(false, contSec.mark, ip, port)
		}
	}
	// update ipGroups
	n.ipGroups[name] = ips
	return nil
}

func (n *NetworkSecurity) DeleteIPGroup(name string) error {
	if err := n.UpdateIPGroup(name, []string{}); err != nil {
		return err
	}
	n.Lock()
	defer n.Unlock()
	delete(n.ipGroups, name)
	return nil
}

func (n *NetworkSecurity) AddContainerSecurity(id string, pid int, sgs map[string]uint16) error {
	n.Lock()
	defer n.Unlock()
	if _, exists := n.containers[id]; exists {
		// we already have security set up for this id. don't do it and return an error.
		return errors.New("Container "+id+" already has Network Security set up.")
	}
	// make sure all groups exist
	for group, _ := range sgs {
		_, exists := n.ipGroups[group]
		if !exists {
			return errors.New("IP Group "+group+" does not exist")
		}
	}

	// fetch network info
	contSec := &ContainerSecurity{
		ID: id,
		Pid: pid,
		SecurityGroups: sgs
	}
	if err := contSec.GetNetworkInfo(); err != nil {
		return err
	}
	// create prerouting mark
	if err := n.markVeth(true, contSec.veth, contSec.mark); err != nil {
		return err
	}
	// add forward rules
	for group, port := range sgs {
		ips := n.ipGroups[group]
		for _, ip := range ips {
			if err := n.openPort(true, contSec.mark, ip, port); err != nil {
				defer n.RemoveContainerSecurity(id) // cleanup created references when we error out
				return err
			}
		}
	}
}

func (n *NetworkSecurity) RemoveContainerSecurity(id string) error {
	n.Lock()
	defer n.Unlock()
	var (
		contSec *ContainerSecurity
		exists bool
	)
	if contSec, exists = n.containers[id]; !exists {
		// no container security here, nothing to remove
		return nil
	}
	// remove prerouting mark
	n.markVeth(false, contSec.veth, contSec.mark)
	// remove forward rules
	for group, port := range sgs {
		ips := n.ipGroups[group]
		for _, ip := range ips {
			n.openPort(false, contSec.mark, ip, port)
		}
	}
	return nil
}

func (n *NetworkSecurity) blanketDeny(add bool, ip string) error {
	// TODO[jigish] implement
	return errors.New("unimplemented")
}

func (n *NetworkSecurity) openPort(add bool, mark int, ip string, port uint16) error {
	// TODO[jigish] implement
	return errors.New("unimplemented")
}

func (n *NetworkSecurity) markVeth(add bool, veth string, mark int) error {
	// TODO[jigish] implement
	return errors.New("unimplemented")
}
<|endoftext|>


<commit_before>package server

import (
	"../bitset"
	"../hashtree"
	"../network"
	"bufio"
	"bytes"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"os"
)

var ERROR_NOT_LOCAL = errors.New("file is not locally available")
var ERROR_LEVEL_LOW = errors.New("the inner hash level is lower than cached")
var ERROR_INDEX_OFF = errors.New("the index of inner hashes is out of range for file of this size")
var ERROR_ALREADY_EXIST = errors.New("the thing was already added or started")

type FileState int

const (
	FILE_UNKNOW FileState = iota
	FILE_PART
	FILE_COMPLETE
)

type Database interface {
	LowestInnerHashes() hashtree.Level
	ImportFromReader(r io.Reader) network.StaticId
	GetState(id network.StaticId) FileState
	GetAt(b []byte, id network.StaticId, off hashtree.Bytes) (int, error)
	GetInnerHashes(id network.StaticId, req network.InnerHashes) (network.InnerHashes, error)
	StartPart(id network.StaticId) error
	PutAt(b []byte, id network.StaticId, off hashtree.Bytes) error
	PutInnerHashes(id network.StaticId, set network.InnerHashes) error
	Remove(id network.StaticId)
}

func ImportLocalFile(d Database, location string) (id network.StaticId) {
	f, _ := os.Open(location)
	r := bufio.NewReader(f)
	id = d.ImportFromReader(r)
	f.Close()
	return
}

type simpleDatabase struct {
	datafolder        *os.File
	dirname           string
	lowestInnerHashes hashtree.Level
}

func OpenSimpleDatabase(dirname string, lowestInnerHashes hashtree.Level) Database {
	os.MkdirAll(dirname, 0777)
	dir, err := os.Open(dirname)
	if err != nil {
		panic(err)
	}
	d := &simpleDatabase{
		datafolder:        dir,
		dirname:           dirname,
		lowestInnerHashes: lowestInnerHashes,
	}

	return d
}

func (d *simpleDatabase) LowestInnerHashes() hashtree.Level {
	return d.lowestInnerHashes
}

var refHash = hashtree.NewFile()

func (d *simpleDatabase) hashNumber(leafs hashtree.Nodes, l hashtree.Level, n hashtree.Nodes) int64 {
	sum := hashtree.Nodes(0)
	for i := hashtree.Level(0); i < l; i++ {
		sum += refHash.LevelWidth(leafs, i)
	}
	return int64(sum + n)
}
func (d *simpleDatabase) hashTopNumber(leafs hashtree.Nodes) int64 {
	sum := hashtree.Nodes(0)
	l := hashtree.Levels(leafs)
	for i := hashtree.Level(0); i < l; i++ {
		sum += refHash.LevelWidth(leafs, i)
	}
	return int64(sum)
}

func (d *simpleDatabase) hashPosition(leafs hashtree.Nodes, l hashtree.Level, n hashtree.Nodes) int64 {
	return d.hashNumber(leafs, l, n) * int64(refHash.Size())
}

func (d *simpleDatabase) innerHashListenerFile(hasher hashtree.HashTree, len hashtree.Bytes) *os.File {
	leafs := hasher.Nodes(len)
	top := hasher.Levels(leafs) - 1
	file, err := ioutil.TempFile(d.dirname, "listener-")
	if err != nil {
		panic(err)
	}
	listener := func(l hashtree.Level, i hashtree.Nodes, h *hashtree.H256) {
		//TODO: don't save levels lower than needed
		if l == top {
			return //don't need the root here
		}
		b := h.ToBytes()
		off := d.hashPosition(leafs, l, i)
		file.WriteAt(b, off)
	}
	hasher.SetInnerHashListener(listener)
	return file
}

func (d *simpleDatabase) ImportFromReader(r io.Reader) network.StaticId {
	f, err := ioutil.TempFile(d.dirname, "import-")
	if err != nil {
		panic(err)
	}
	len, err2 := io.Copy(f, r)
	if err2 != nil {
		panic(err2)
	}
	hasher := hashtree.NewFile()
	hashFile := d.innerHashListenerFile(hasher, hashtree.Bytes(len))

	f.Seek(0, os.SEEK_SET)
	io.Copy(hasher, f)
	id := network.StaticId{
		Hash:   hasher.Sum(nil),
		Length: &len,
	}
	f.Close()
	hashFile.Close()

	err = os.Rename(f.Name(), d.fileNameForId(id))
	if err != nil {
		if os.IsExist(err) {
			os.Remove(f.Name())
		} else {
			panic(err)
		}
	}
	err = os.Rename(hashFile.Name(), d.hashFileNameForId(id))
	if err != nil {
		if os.IsExist(err) {
			os.Remove(hashFile.Name())
		} else {
			panic(err)
		}
	}
	return id
}

func (d *simpleDatabase) GetState(id network.StaticId) FileState {
	_, err := os.Stat(d.fileNameForId(id))
	if os.IsNotExist(err) {
		_, err = os.Stat(d.hashFileNameForId(id))
		if os.IsNotExist(err) {
			return FILE_UNKNOW
		} else {
			return FILE_PART
		}
	} else {
		return FILE_COMPLETE
	}
}

func (d *simpleDatabase) GetAt(b []byte, id network.StaticId, off hashtree.Bytes) (int, error) {
	f, err := os.Open(d.fileNameForId(id))
	if err != nil {
		return 0, ERROR_NOT_LOCAL
	}
	defer f.Close()
	return f.ReadAt(b, int64(off))
}

func (d *simpleDatabase) GetInnerHashes(id network.StaticId, req network.InnerHashes) (network.InnerHashes, error) {
	leaf := refHash.Nodes(hashtree.Bytes(id.GetLength()))
	level := hashtree.Level(req.GetHeight())
	from := hashtree.Nodes(req.GetFrom())
	nodes := hashtree.Nodes(req.GetLength())
	if level < d.lowestInnerHashes {
		return req, ERROR_LEVEL_LOW
	} else if from+nodes > refHash.LevelWidth(leaf, level) {
		return req, ERROR_INDEX_OFF
	}
	f, err := os.Open(d.hashFileNameForId(id))
	if err != nil {
		return req, ERROR_NOT_LOCAL
	}
	defer f.Close()
	off := d.hashPosition(leaf, level, from)
	b := make([]byte, refHash.Size()*int(nodes))
	f.ReadAt(b, off)
	req.Hashes = b
	return req, nil
}

func (d *simpleDatabase) StartPart(id network.StaticId) error {
	_, err := os.Stat(d.hashFileNameForId(id))
	if os.IsNotExist(err) {
		pf, err := os.Create(d.partFileNameForId(id))
		if err != nil {
			return err
		}
		defer pf.Close()
		hf, err2 := os.Create(d.hashFileNameForId(id))
		if err2 != nil {
			return err2
		}
		defer hf.Close()
		return nil
	} else {
		return ERROR_ALREADY_EXIST
	}

}
func (d *simpleDatabase) PutAt(b []byte, id network.StaticId, off hashtree.Bytes) error {
	return nil
}
func (d *simpleDatabase) PutInnerHashes(id network.StaticId, set network.InnerHashes) error {
	leafs := id.Blocks()
	bits := bitset.OpenCountingFileBacked(d.haveHashNameForId(id), int(d.hashTopNumber(leafs)-1))
	defer bits.Close()
	f, err := os.Open(d.hashFileNameForId(id))
	if err != nil {
		return ERROR_NOT_LOCAL
	}
	defer f.Close()
	hashBuffer := make([]byte, refHash.Size())
	splited := set.SplitLocalSummable(&id)
	for _, hashes := range splited {
		rootL, rootN := hashes.LocalRoot(leafs)
		key := int(d.hashNumber(leafs, rootL, rootN))
		if key == bits.Capacity() {
			//this is root
		} else if !bits.Get(key) {
			continue // this part of hashes can not be verified, skiped
		}
		sum := hashes.LocalSum()
		if key == bits.Capacity() {
			copy(hashBuffer, id.GetHash())
		} else {
			off := d.hashPosition(leafs, rootL, rootN)
			if _, err := f.ReadAt(hashBuffer, off); err != nil {
				panic(err)
			}
		}
		if bytes.Equal(sum, hashBuffer) {
			//verified, now save
			listener := func(l hashtree.Level, i hashtree.Nodes, h *hashtree.H256) {
				realL := set.GetHeightL() + l
				realN := i >> uint32(l)
				if realL == rootL {
					return //don't need the root here, already verified
				}
				b := h.ToBytes()
				off := d.hashPosition(leafs, realL, realN)
				f.WriteAt(b, off)
				bits.Set(int(d.hashNumber(leafs, realL, realN)))

				for realL+1 != rootL && realN != 0 && realN%2 == 0 && realN+1 == hashtree.LevelWidth(leafs, realL) {
					//the node and it's parent have the same hash, so, also write to parent
					realL += 1
					realN /= 2

					off = d.hashPosition(leafs, realL, realN)
					f.WriteAt(b, off)
					bits.Set(int(d.hashNumber(leafs, realL, realN)))
				}
			}
			hasher := hashtree.NewFile()
			hasher.SetInnerHashListener(listener)
			hasher.Write(hashes.GetHashes())
			hasher.Sum(nil)
		}
	}
	return nil
}

func remove(filename string) {
	err := os.Remove(filename)
	if err != nil && !os.IsNotExist(err) {
		panic(err)
	}
}

func (d *simpleDatabase) Remove(id network.StaticId) {
	remove(d.havePartNameForId(id))
	remove(d.haveHashNameForId(id))
	remove(d.partFileNameForId(id))
	remove(d.fileNameForId(id))
	remove(d.hashFileNameForId(id))
}

func (d *simpleDatabase) fileNameForId(id network.StaticId) string {
	return fmt.Sprintf("%s/F-%s", d.datafolder.Name(), id.CompactId())
}
func (d *simpleDatabase) hashFileNameForId(id network.StaticId) string {
	return fmt.Sprintf("%s/H-%s", d.datafolder.Name(), id.CompactId())
}
func (d *simpleDatabase) partFileNameForId(id network.StaticId) string {
	return fmt.Sprintf("%s/P-%s", d.datafolder.Name(), id.CompactId())
}
func (d *simpleDatabase) haveHashNameForId(id network.StaticId) string {
	return fmt.Sprintf("%s/hH-%s", d.datafolder.Name(), id.CompactId())
}
func (d *simpleDatabase) havePartNameForId(id network.StaticId) string {
	return fmt.Sprintf("%s/hP-%s", d.datafolder.Name(), id.CompactId())
}
<commit_msg>return if have all inner hashes<commit_after>package server

import (
	"../bitset"
	"../hashtree"
	"../network"
	"bufio"
	"bytes"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"os"
)

var ERROR_NOT_LOCAL = errors.New("file is not locally available")
var ERROR_LEVEL_LOW = errors.New("the inner hash level is lower than cached")
var ERROR_INDEX_OFF = errors.New("the index of inner hashes is out of range for file of this size")
var ERROR_ALREADY_EXIST = errors.New("the thing was already added or started")

type FileState int

const (
	FILE_UNKNOW FileState = iota
	FILE_PART
	FILE_COMPLETE
)

type Database interface {
	LowestInnerHashes() hashtree.Level
	ImportFromReader(r io.Reader) network.StaticId
	GetState(id network.StaticId) FileState
	GetAt(b []byte, id network.StaticId, off hashtree.Bytes) (int, error)
	GetInnerHashes(id network.StaticId, req network.InnerHashes) (network.InnerHashes, error)
	StartPart(id network.StaticId) error
	PutAt(b []byte, id network.StaticId, off hashtree.Bytes) error
	PutInnerHashes(id network.StaticId, set network.InnerHashes) (bool, error)
	Remove(id network.StaticId)
}

func ImportLocalFile(d Database, location string) (id network.StaticId) {
	f, _ := os.Open(location)
	r := bufio.NewReader(f)
	id = d.ImportFromReader(r)
	f.Close()
	return
}

type simpleDatabase struct {
	datafolder        *os.File
	dirname           string
	lowestInnerHashes hashtree.Level
}

func OpenSimpleDatabase(dirname string, lowestInnerHashes hashtree.Level) Database {
	os.MkdirAll(dirname, 0777)
	dir, err := os.Open(dirname)
	if err != nil {
		panic(err)
	}
	d := &simpleDatabase{
		datafolder:        dir,
		dirname:           dirname,
		lowestInnerHashes: lowestInnerHashes,
	}

	return d
}

func (d *simpleDatabase) LowestInnerHashes() hashtree.Level {
	return d.lowestInnerHashes
}

var refHash = hashtree.NewFile()

func (d *simpleDatabase) hashNumber(leafs hashtree.Nodes, l hashtree.Level, n hashtree.Nodes) int64 {
	sum := hashtree.Nodes(0)
	for i := hashtree.Level(0); i < l; i++ {
		sum += refHash.LevelWidth(leafs, i)
	}
	return int64(sum + n)
}
func (d *simpleDatabase) hashTopNumber(leafs hashtree.Nodes) int64 {
	sum := hashtree.Nodes(0)
	l := hashtree.Levels(leafs)
	for i := hashtree.Level(0); i < l; i++ {
		sum += refHash.LevelWidth(leafs, i)
	}
	return int64(sum)
}

func (d *simpleDatabase) hashPosition(leafs hashtree.Nodes, l hashtree.Level, n hashtree.Nodes) int64 {
	return d.hashNumber(leafs, l, n) * int64(refHash.Size())
}

func (d *simpleDatabase) innerHashListenerFile(hasher hashtree.HashTree, len hashtree.Bytes) *os.File {
	leafs := hasher.Nodes(len)
	top := hasher.Levels(leafs) - 1
	file, err := ioutil.TempFile(d.dirname, "listener-")
	if err != nil {
		panic(err)
	}
	listener := func(l hashtree.Level, i hashtree.Nodes, h *hashtree.H256) {
		//TODO: don't save levels lower than needed
		if l == top {
			return //don't need the root here
		}
		b := h.ToBytes()
		off := d.hashPosition(leafs, l, i)
		file.WriteAt(b, off)
	}
	hasher.SetInnerHashListener(listener)
	return file
}

func (d *simpleDatabase) ImportFromReader(r io.Reader) network.StaticId {
	f, err := ioutil.TempFile(d.dirname, "import-")
	if err != nil {
		panic(err)
	}
	len, err2 := io.Copy(f, r)
	if err2 != nil {
		panic(err2)
	}
	hasher := hashtree.NewFile()
	hashFile := d.innerHashListenerFile(hasher, hashtree.Bytes(len))

	f.Seek(0, os.SEEK_SET)
	io.Copy(hasher, f)
	id := network.StaticId{
		Hash:   hasher.Sum(nil),
		Length: &len,
	}
	f.Close()
	hashFile.Close()

	err = os.Rename(f.Name(), d.fileNameForId(id))
	if err != nil {
		if os.IsExist(err) {
			os.Remove(f.Name())
		} else {
			panic(err)
		}
	}
	err = os.Rename(hashFile.Name(), d.hashFileNameForId(id))
	if err != nil {
		if os.IsExist(err) {
			os.Remove(hashFile.Name())
		} else {
			panic(err)
		}
	}
	return id
}

func (d *simpleDatabase) GetState(id network.StaticId) FileState {
	_, err := os.Stat(d.fileNameForId(id))
	if os.IsNotExist(err) {
		_, err = os.Stat(d.hashFileNameForId(id))
		if os.IsNotExist(err) {
			return FILE_UNKNOW
		} else {
			return FILE_PART
		}
	} else {
		return FILE_COMPLETE
	}
}

func (d *simpleDatabase) GetAt(b []byte, id network.StaticId, off hashtree.Bytes) (int, error) {
	f, err := os.Open(d.fileNameForId(id))
	if err != nil {
		return 0, ERROR_NOT_LOCAL
	}
	defer f.Close()
	return f.ReadAt(b, int64(off))
}

func (d *simpleDatabase) GetInnerHashes(id network.StaticId, req network.InnerHashes) (network.InnerHashes, error) {
	leaf := refHash.Nodes(hashtree.Bytes(id.GetLength()))
	level := hashtree.Level(req.GetHeight())
	from := hashtree.Nodes(req.GetFrom())
	nodes := hashtree.Nodes(req.GetLength())
	if level < d.lowestInnerHashes {
		return req, ERROR_LEVEL_LOW
	} else if from+nodes > refHash.LevelWidth(leaf, level) {
		return req, ERROR_INDEX_OFF
	}
	f, err := os.Open(d.hashFileNameForId(id))
	if err != nil {
		return req, ERROR_NOT_LOCAL
	}
	defer f.Close()
	off := d.hashPosition(leaf, level, from)
	b := make([]byte, refHash.Size()*int(nodes))
	f.ReadAt(b, off)
	req.Hashes = b
	return req, nil
}

func (d *simpleDatabase) StartPart(id network.StaticId) error {
	_, err := os.Stat(d.hashFileNameForId(id))
	if os.IsNotExist(err) {
		pf, err := os.Create(d.partFileNameForId(id))
		if err != nil {
			return err
		}
		defer pf.Close()
		hf, err2 := os.Create(d.hashFileNameForId(id))
		if err2 != nil {
			return err2
		}
		defer hf.Close()
		return nil
	} else {
		return ERROR_ALREADY_EXIST
	}

}
func (d *simpleDatabase) PutAt(b []byte, id network.StaticId, off hashtree.Bytes) error {
	return nil
}
func (d *simpleDatabase) PutInnerHashes(id network.StaticId, set network.InnerHashes) (complete bool, err error) {
	leafs := id.Blocks()
	bits := bitset.OpenCountingFileBacked(d.haveHashNameForId(id), int(d.hashTopNumber(leafs)-1))
	defer bits.Close()
	f, err := os.Open(d.hashFileNameForId(id))
	if err != nil {
		return false, ERROR_NOT_LOCAL
	}
	defer f.Close()
	hashBuffer := make([]byte, refHash.Size())
	splited := set.SplitLocalSummable(&id)
	for _, hashes := range splited {
		rootL, rootN := hashes.LocalRoot(leafs)
		key := int(d.hashNumber(leafs, rootL, rootN))
		if key == bits.Capacity() {
			//this is root
		} else if !bits.Get(key) {
			continue // this part of hashes can not be verified, skiped
		}
		sum := hashes.LocalSum()
		if key == bits.Capacity() {
			copy(hashBuffer, id.GetHash())
		} else {
			off := d.hashPosition(leafs, rootL, rootN)
			if _, err := f.ReadAt(hashBuffer, off); err != nil {
				panic(err)
			}
		}
		if bytes.Equal(sum, hashBuffer) {
			//verified, now save
			listener := func(l hashtree.Level, i hashtree.Nodes, h *hashtree.H256) {
				realL := set.GetHeightL() + l
				realN := i >> uint32(l)
				if realL == rootL {
					return //don't need the root here, already verified
				}
				b := h.ToBytes()
				off := d.hashPosition(leafs, realL, realN)
				f.WriteAt(b, off)
				bits.Set(int(d.hashNumber(leafs, realL, realN)))

				for realL+1 != rootL && realN != 0 && realN%2 == 0 && realN+1 == hashtree.LevelWidth(leafs, realL) {
					//the node and it's parent have the same hash, so, also write to parent
					realL += 1
					realN /= 2

					off = d.hashPosition(leafs, realL, realN)
					f.WriteAt(b, off)
					bits.Set(int(d.hashNumber(leafs, realL, realN)))
				}
			}
			hasher := hashtree.NewFile()
			hasher.SetInnerHashListener(listener)
			hasher.Write(hashes.GetHashes())
			hasher.Sum(nil)
		}
	}
	return bits.Full(), nil
}

func remove(filename string) {
	err := os.Remove(filename)
	if err != nil && !os.IsNotExist(err) {
		panic(err)
	}
}

func (d *simpleDatabase) Remove(id network.StaticId) {
	remove(d.havePartNameForId(id))
	remove(d.haveHashNameForId(id))
	remove(d.partFileNameForId(id))
	remove(d.fileNameForId(id))
	remove(d.hashFileNameForId(id))
}

func (d *simpleDatabase) fileNameForId(id network.StaticId) string {
	return fmt.Sprintf("%s/F-%s", d.datafolder.Name(), id.CompactId())
}
func (d *simpleDatabase) hashFileNameForId(id network.StaticId) string {
	return fmt.Sprintf("%s/H-%s", d.datafolder.Name(), id.CompactId())
}
func (d *simpleDatabase) partFileNameForId(id network.StaticId) string {
	return fmt.Sprintf("%s/P-%s", d.datafolder.Name(), id.CompactId())
}
func (d *simpleDatabase) haveHashNameForId(id network.StaticId) string {
	return fmt.Sprintf("%s/hH-%s", d.datafolder.Name(), id.CompactId())
}
func (d *simpleDatabase) havePartNameForId(id network.StaticId) string {
	return fmt.Sprintf("%s/hP-%s", d.datafolder.Name(), id.CompactId())
}
<|endoftext|>


<commit_before>package main

import (
	"eaciit/gdrj/model"
	"eaciit/gdrj/modules"
	"os"

	"strings"
	"sync"
	"time"

	"github.com/eaciit/dbox"
	"github.com/eaciit/orm/v1"
	"github.com/eaciit/toolkit"

	"flag"
)

var mutex = new(sync.Mutex)
var conn dbox.IConnection

func setinitialconnection() {
	var err error
	conn, err = modules.GetDboxIConnection("db_godrej")

	if err != nil {
		toolkit.Println("Initial connection found : ", err)
		os.Exit(1)
	}

	err = gdrj.SetDb(conn)
	if err != nil {
		toolkit.Println("Initial connection found : ", err)
		os.Exit(1)
	}
}

var (
	masters    = toolkit.M{}
	t0         time.Time
	fiscalyear int
)

func buildmap(holder interface{},
	fnModel func() orm.IModel,
	filter *dbox.Filter,
	fnIter func(holder interface{}, obj interface{})) interface{} {
	crx, _ := gdrj.Find(fnModel(), filter, nil)
	defer crx.Close()
	for {
		s := fnModel()
		e := crx.Fetch(s, 1, false)
		if e != nil {
			break
		}
		fnIter(holder, s)
	}
	return holder
}

func getCursor(obj orm.IModel) dbox.ICursor {
	c, e := gdrj.Find(obj,
		nil, nil)
	//toolkit.M{}.Set("take", 10))
	if e != nil {
		return nil
	}
	return c
}

func prepmaster() {
	toolkit.Println("--> PL Model")
	masters.Set("plmodel", buildmap(map[string]*gdrj.PLModel{},
		func() orm.IModel {
			return new(gdrj.PLModel)
		},
		nil,
		func(holder, obj interface{}) {
			h := holder.(map[string]*gdrj.PLModel)
			o := obj.(*gdrj.PLModel)
			h[o.ID] = o
		}).(map[string]*gdrj.PLModel))

	toolkit.Println("--> Sub Channel")
	subchannels := toolkit.M{}
	csr, _ := conn.NewQuery().From("subchannels").Cursor(nil)
	defer csr.Close()
	for {
		m := toolkit.M{}
		e := csr.Fetch(&m, 1, false)
		if e != nil {
			break
		}
		subchannels.Set(m.GetString("_id"), m.GetString("title"))
	}
	masters.Set("subchannels", subchannels)

	toolkit.Println("--> Customer")
	customers := toolkit.M{}
	ccb := getCursor(new(gdrj.Customer))
	defer ccb.Close()
	for {
		cust := new(gdrj.Customer)
		e := ccb.Fetch(cust, 1, false)
		if e != nil {
			break
		}

		customers.Set(cust.ID, cust)
	}
	masters.Set("customers", customers)

	branchs := toolkit.M{}
	cmb := getCursor(new(gdrj.MasterBranch))
	defer cmb.Close()
	for {
		stx := toolkit.M{}
		e := cmb.Fetch(&stx, 1, false)
		if e != nil {
			break
		}

		branchs.Set(stx.Get("_id", "").(string), stx)
	}
	masters.Set("branchs", branchs)

	toolkit.Println("--> Trx Gross Proc")
	globalgross, globalgrossvdist := float64(0), float64(0)
	grossbybranch, grossbybrand, grossbysku, grossbychannel := toolkit.M{}, toolkit.M{}, toolkit.M{}, toolkit.M{}
	grossbymonthvdist, grossbymonth, grossbymonthsku, grossbymonthchannel, grossbymonthbrandchannel := toolkit.M{}, toolkit.M{}, toolkit.M{}, toolkit.M{}, toolkit.M{}
	csr01, _ := conn.NewQuery().From("salestrxs-grossproc").Cursor(nil)
	defer csr01.Close()
	for {
		m := toolkit.M{}
		e := csr01.Fetch(&m, 1, false)
		if e != nil {
			break
		}

		agross := m.GetFloat64("gross")
		globalgross += agross
		if strings.ToUpper(m.GetString("src")) == "VDIST" {
			globalgrossvdist += agross
		}

		key := toolkit.Sprintf("%s", m.GetString("brand"))
		tempval := grossbybrand.GetFloat64(key) + agross
		grossbybrand.Set(key, tempval)

		key = toolkit.Sprintf("%s", m.GetString("branchid"))
		tempval = grossbybranch.GetFloat64(key) + agross
		grossbybranch.Set(key, tempval)

		key = toolkit.Sprintf("%s", m.GetString("skuid"))
		tempval = grossbysku.GetFloat64(key) + agross
		grossbysku.Set(key, tempval)

		key = toolkit.Sprintf("%s", m.GetString("channelcheck"))
		tempval = grossbychannel.GetFloat64(key) + agross
		grossbychannel.Set(key, tempval)

		key = toolkit.Sprintf("%d_%d", m.GetInt("year"), m.GetInt("month"))
		tempval = grossbymonth.GetFloat64(key) + agross
		grossbymonth.Set(key, tempval)

		if strings.ToUpper(m.GetString("src")) == "VDIST" {
			key = toolkit.Sprintf("%d_%d", m.GetInt("year"), m.GetInt("month"))
			tempval = grossbymonthvdist.GetFloat64(key) + agross
			grossbymonthvdist.Set(key, tempval)
		}

		key = toolkit.Sprintf("%d_%d_%s", m.GetInt("year"), m.GetInt("month"), m.GetString("skuid"))
		tempval = grossbymonthsku.GetFloat64(key) + agross
		grossbymonthsku.Set(key, tempval)

		key = toolkit.Sprintf("%d_%d_%s", m.GetInt("year"), m.GetInt("month"), m.GetString("channelcheck"))
		tempval = grossbymonthchannel.GetFloat64(key) + agross
		grossbymonthchannel.Set(key, tempval)

		key = toolkit.Sprintf("%d_%d_%s_%s", m.GetInt("year"), m.GetInt("month"), m.GetString("brand"), m.GetString("channelcheck"))
		tempval = grossbymonthbrandchannel.GetFloat64(key) + agross
		grossbymonthbrandchannel.Set(key, tempval)

	}

	masters.Set("globalgross", globalgross).Set("globalgrossvdist", globalgrossvdist)
	// tkm
	masters.Set("grossbybranch", grossbybranch).Set("grossbybrand", grossbybrand).Set("grossbysku", grossbysku).Set("grossbychannel", grossbychannel)
	masters.Set("grossbymonth", grossbymonth).Set("grossbymonthvdist", grossbymonthvdist).Set("grossbymonthsku", grossbymonthsku).
		Set("grossbymonthchannel", grossbymonthchannel).Set("grossbymonthbrandchannel", grossbymonthbrandchannel)
}

func main() {

	t0 = time.Now()
	flag.IntVar(&fiscalyear, "year", 2015, "fiscal year to process. default is 2015")
	flag.Parse()

	eperiode := time.Date(fiscalyear, 4, 1, 0, 0, 0, 0, time.UTC)
	speriode := eperiode.AddDate(-1, 0, 0)

	setinitialconnection()
	defer gdrj.CloseDb()

	var f *dbox.Filter
	f = dbox.And(dbox.Gte("date", speriode), dbox.Lt("date", eperiode))
	// f = dbox.Or(dbox.Eq("_id", "RD_2014_4_40007354_24260"),dbox.Eq("_id", "EXPORT_2014_4_40007767_13"),dbox.Eq("_id", "VDIST/2015-2014/FK/IPR/14003129_CD04_2"))

	toolkit.Printfn("Prepare : %v", t0)
	prepmaster()

	toolkit.Printfn("Run : %v", t0)

	c, _ := gdrj.Find(new(gdrj.SalesTrx), f, nil)
	defer c.Close()

	count := c.Count()
	jobs := make(chan *gdrj.SalesTrx, count)
	result := make(chan string, count)
	for wi := 0; wi < 10; wi++ {
		go workerproc(wi, jobs, result)
	}

	step := count / 100
	if step == 0 {
		step = 10
	}

	i := 0
	toolkit.Printfn("START ... %d records ", count)
	for {
		stx := new(gdrj.SalesTrx)
		e := c.Fetch(stx, 1, false)
		if e != nil {
			break
		}

		i++
		jobs <- stx
		if i%step == 0 {
			toolkit.Printfn("Sending %d of %d (%d) in %s",
				i, count, i/step,
				time.Since(t0).String())
		}
	}

	close(jobs)

	for ri := 0; ri < i; ri++ {
		<-result

		if ri%step == 0 {
			toolkit.Printfn("Saving %d of %d (%d pct) in %s",
				ri, count, ri/step, time.Since(t0).String())
		}
	}
}

func workerproc(wi int, jobs <-chan *gdrj.SalesTrx, result chan<- string) {
	workerconn, _ := modules.GetDboxIConnection("db_godrej")
	defer workerconn.Close()

	trx := new(gdrj.SalesTrx)
	for trx = range jobs {
		pl := new(gdrj.SalesPL)

		pl.ID = trx.ID
		pl.SKUID = trx.SKUID
		pl.SKUID_VDIST = trx.SKUID_VDIST
		pl.OutletID = trx.OutletID

		pl.Date = gdrj.SetDate(trx.Date)

		pl.SalesQty = trx.SalesQty
		pl.GrossAmount = trx.GrossAmount
		pl.DiscountAmount = trx.DiscountAmount
		pl.TaxAmount = trx.TaxAmount

		pl.Customer = trx.Customer
		pl.Product = trx.Product

		pl.CleanAndClasify(masters)
		pl.RatioCalc(masters)

		pl.CalcSales(masters)
		pl.CalcSum(masters)

		tablename := toolkit.Sprintf("%v-1", pl.TableName())
		workerconn.NewQuery().From(tablename).
			Save().Exec(toolkit.M{}.Set("data", pl))

		result <- pl.ID
	}
}
<commit_msg>make global<commit_after>package main

import (
	"eaciit/gdrj/model"
	"eaciit/gdrj/modules"
	"os"

	"strings"
	"sync"
	"time"

	"github.com/eaciit/dbox"
	"github.com/eaciit/orm/v1"
	"github.com/eaciit/toolkit"

	"flag"
)

var mutex = new(sync.Mutex)
var conn dbox.IConnection

func setinitialconnection() {
	var err error
	conn, err = modules.GetDboxIConnection("db_godrej")

	if err != nil {
		toolkit.Println("Initial connection found : ", err)
		os.Exit(1)
	}

	err = gdrj.SetDb(conn)
	if err != nil {
		toolkit.Println("Initial connection found : ", err)
		os.Exit(1)
	}
}

var (
	masters                                                 = toolkit.M{}
	t0                                                      time.Time
	fiscalyear                                              int
	subchannels                                             = toolkit.M{}
	customers                                               = toolkit.M{}
	branchs                                                 = toolkit.M{}
	globalgross, globalgrossvdist                           = float64(0), float64(0)
	grossbybranch, grossbybrand, grossbysku, grossbychannel = toolkit.M{}, toolkit.M{}, toolkit.M{}, toolkit.M{}
	grossbymonthvdist, grossbymonth, grossbymonthsku        = toolkit.M{}, toolkit.M{}, toolkit.M{}
	grossbymonthchannel, grossbymonthbrandchannel           = toolkit.M{}, toolkit.M{}
)

func buildmap(holder interface{},
	fnModel func() orm.IModel,
	filter *dbox.Filter,
	fnIter func(holder interface{}, obj interface{})) interface{} {
	crx, _ := gdrj.Find(fnModel(), filter, nil)
	defer crx.Close()
	for {
		s := fnModel()
		e := crx.Fetch(s, 1, false)
		if e != nil {
			break
		}
		fnIter(holder, s)
	}
	return holder
}

func getCursor(obj orm.IModel) dbox.ICursor {
	c, e := gdrj.Find(obj,
		nil, nil)
	//toolkit.M{}.Set("take", 10))
	if e != nil {
		return nil
	}
	return c
}

func prepmaster() {
	toolkit.Println("--> PL Model")
	masters.Set("plmodel", buildmap(map[string]*gdrj.PLModel{},
		func() orm.IModel {
			return new(gdrj.PLModel)
		},
		nil,
		func(holder, obj interface{}) {
			h := holder.(map[string]*gdrj.PLModel)
			o := obj.(*gdrj.PLModel)
			h[o.ID] = o
		}).(map[string]*gdrj.PLModel))

	toolkit.Println("--> Sub Channel")
	csr, _ := conn.NewQuery().From("subchannels").Cursor(nil)
	defer csr.Close()
	for {
		m := toolkit.M{}
		e := csr.Fetch(&m, 1, false)
		if e != nil {
			break
		}
		subchannels.Set(m.GetString("_id"), m.GetString("title"))
	}
	masters.Set("subchannels", subchannels)

	toolkit.Println("--> Customer")
	ccb := getCursor(new(gdrj.Customer))
	defer ccb.Close()
	for {
		cust := new(gdrj.Customer)
		e := ccb.Fetch(cust, 1, false)
		if e != nil {
			break
		}

		customers.Set(cust.ID, cust)
	}
	masters.Set("customers", customers)

	cmb := getCursor(new(gdrj.MasterBranch))
	defer cmb.Close()
	for {
		stx := toolkit.M{}
		e := cmb.Fetch(&stx, 1, false)
		if e != nil {
			break
		}

		branchs.Set(stx.Get("_id", "").(string), stx)
	}
	masters.Set("branchs", branchs)

	toolkit.Println("--> Trx Gross Proc")
	csr01, _ := conn.NewQuery().From("salestrxs-grossproc").Cursor(nil)
	defer csr01.Close()
	for {
		m := toolkit.M{}
		e := csr01.Fetch(&m, 1, false)
		if e != nil {
			break
		}

		agross := m.GetFloat64("gross")
		globalgross += agross
		if strings.ToUpper(m.GetString("src")) == "VDIST" {
			globalgrossvdist += agross
		}

		key := toolkit.Sprintf("%s", m.GetString("brand"))
		tempval := grossbybrand.GetFloat64(key) + agross
		grossbybrand.Set(key, tempval)

		key = toolkit.Sprintf("%s", m.GetString("branchid"))
		tempval = grossbybranch.GetFloat64(key) + agross
		grossbybranch.Set(key, tempval)

		key = toolkit.Sprintf("%s", m.GetString("skuid"))
		tempval = grossbysku.GetFloat64(key) + agross
		grossbysku.Set(key, tempval)

		key = toolkit.Sprintf("%s", m.GetString("channelcheck"))
		tempval = grossbychannel.GetFloat64(key) + agross
		grossbychannel.Set(key, tempval)

		key = toolkit.Sprintf("%d_%d", m.GetInt("year"), m.GetInt("month"))
		tempval = grossbymonth.GetFloat64(key) + agross
		grossbymonth.Set(key, tempval)

		if strings.ToUpper(m.GetString("src")) == "VDIST" {
			key = toolkit.Sprintf("%d_%d", m.GetInt("year"), m.GetInt("month"))
			tempval = grossbymonthvdist.GetFloat64(key) + agross
			grossbymonthvdist.Set(key, tempval)
		}

		key = toolkit.Sprintf("%d_%d_%s", m.GetInt("year"), m.GetInt("month"), m.GetString("skuid"))
		tempval = grossbymonthsku.GetFloat64(key) + agross
		grossbymonthsku.Set(key, tempval)

		key = toolkit.Sprintf("%d_%d_%s", m.GetInt("year"), m.GetInt("month"), m.GetString("channelcheck"))
		tempval = grossbymonthchannel.GetFloat64(key) + agross
		grossbymonthchannel.Set(key, tempval)

		key = toolkit.Sprintf("%d_%d_%s_%s", m.GetInt("year"), m.GetInt("month"), m.GetString("brand"), m.GetString("channelcheck"))
		tempval = grossbymonthbrandchannel.GetFloat64(key) + agross
		grossbymonthbrandchannel.Set(key, tempval)

	}

	masters.Set("globalgross", globalgross).Set("globalgrossvdist", globalgrossvdist)
	// tkm
	masters.Set("grossbybranch", grossbybranch).Set("grossbybrand", grossbybrand).Set("grossbysku", grossbysku).Set("grossbychannel", grossbychannel)
	masters.Set("grossbymonth", grossbymonth).Set("grossbymonthvdist", grossbymonthvdist).Set("grossbymonthsku", grossbymonthsku).
		Set("grossbymonthchannel", grossbymonthchannel).Set("grossbymonthbrandchannel", grossbymonthbrandchannel)
}

func main() {

	t0 = time.Now()
	flag.IntVar(&fiscalyear, "year", 2015, "fiscal year to process. default is 2015")
	flag.Parse()

	eperiode := time.Date(fiscalyear, 4, 1, 0, 0, 0, 0, time.UTC)
	speriode := eperiode.AddDate(-1, 0, 0)

	setinitialconnection()
	defer gdrj.CloseDb()

	var f *dbox.Filter
	f = dbox.And(dbox.Gte("date", speriode), dbox.Lt("date", eperiode))
	// f = dbox.Or(dbox.Eq("_id", "RD_2014_4_40007354_24260"),dbox.Eq("_id", "EXPORT_2014_4_40007767_13"),dbox.Eq("_id", "VDIST/2015-2014/FK/IPR/14003129_CD04_2"))

	toolkit.Printfn("Prepare : %v", t0)
	prepmaster()

	toolkit.Printfn("Run : %v", t0)

	c, _ := gdrj.Find(new(gdrj.SalesTrx), f, nil)
	defer c.Close()

	count := c.Count()
	jobs := make(chan *gdrj.SalesTrx, count)
	result := make(chan string, count)
	for wi := 0; wi < 10; wi++ {
		go workerproc(wi, jobs, result)
	}

	step := count / 100
	if step == 0 {
		step = 10
	}

	i := 0
	toolkit.Printfn("START ... %d records ", count)
	for {
		stx := new(gdrj.SalesTrx)
		e := c.Fetch(stx, 1, false)
		if e != nil {
			break
		}

		i++
		jobs <- stx
		if i%step == 0 {
			toolkit.Printfn("Sending %d of %d (%d) in %s",
				i, count, i/step,
				time.Since(t0).String())
		}
	}

	close(jobs)

	for ri := 0; ri < i; ri++ {
		<-result

		if ri%step == 0 {
			toolkit.Printfn("Saving %d of %d (%d pct) in %s",
				ri, count, ri/step, time.Since(t0).String())
		}
	}
}

func workerproc(wi int, jobs <-chan *gdrj.SalesTrx, result chan<- string) {
	workerconn, _ := modules.GetDboxIConnection("db_godrej")
	defer workerconn.Close()

	trx := new(gdrj.SalesTrx)
	for trx = range jobs {
		pl := new(gdrj.SalesPL)

		pl.ID = trx.ID
		pl.SKUID = trx.SKUID
		pl.SKUID_VDIST = trx.SKUID_VDIST
		pl.OutletID = trx.OutletID

		pl.Date = gdrj.SetDate(trx.Date)

		pl.SalesQty = trx.SalesQty
		pl.GrossAmount = trx.GrossAmount
		pl.DiscountAmount = trx.DiscountAmount
		pl.TaxAmount = trx.TaxAmount

		pl.Customer = trx.Customer
		pl.Product = trx.Product

		pl.CleanAndClasify(masters)
		pl.RatioCalc(masters)

		pl.CalcSales(masters)
		pl.CalcSum(masters)

		tablename := toolkit.Sprintf("%v-1", pl.TableName())
		workerconn.NewQuery().From(tablename).
			Save().Exec(toolkit.M{}.Set("data", pl))

		result <- pl.ID
	}
}
<|endoftext|>