<commit_before>// Copyright 2020 The LUCI 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 config

import (
	"context"

	"google.golang.org/protobuf/proto"

	"go.chromium.org/luci/common/clock"
	"go.chromium.org/luci/common/data/stringset"
	"go.chromium.org/luci/common/errors"
	"go.chromium.org/luci/common/logging"
	"go.chromium.org/luci/common/retry/transient"
	"go.chromium.org/luci/common/sync/parallel"
	"go.chromium.org/luci/config"
	"go.chromium.org/luci/config/cfgclient"
	pb "go.chromium.org/luci/cv/api/config/v2"
	"go.chromium.org/luci/gae/service/datastore"
	"go.chromium.org/luci/server/tq"
)

const configFileName = "commit-queue.cfg"

func init() {
	tq.RegisterTaskClass(tq.TaskClass{
		ID:        "refresh-project-config",
		Prototype: &RefreshProjectConfigTask{},
		Queue:     "refresh-project-config",
		Handler: func(ctx context.Context, payload proto.Message) error {
			task := payload.(*RefreshProjectConfigTask)
			action, actionFn := "update", updateProject
			if task.GetDisable() {
				action, actionFn = "disable", disableProject
			}
			project := task.GetProject()
			if err := actionFn(ctx, project); err != nil {
				errors.Log(ctx, err)
				// Explicitly not forwarding transient tag so that tq won't retry
				// because the refresh task is submitted every minute by AppEngine
				// Cron.
				return errors.Reason("failed to %s project %q: %s", action, project, err).Err()
			}
			return nil
		},
	})
}

// SubmitRefreshTasks submits tasks that update config for LUCI projects
// or disable projects that do not have CV config in LUCI Config.
func SubmitRefreshTasks(ctx context.Context) error {
	projects, err := cfgclient.ProjectsWithConfig(ctx, configFileName)
	if err != nil {
		return errors.Annotate(err, "failed to get projects that have %q from LUCI Config", configFileName).Tag(transient.Tag).Err()
	}
	tasks := make([]*tq.Task, len(projects))
	for i, p := range projects {
		tasks[i] = &tq.Task{
			Payload: &RefreshProjectConfigTask{
				Project: p,
			},
		}
	}

	curEnabledProjects, err := getAllProjectIDs(ctx, true)
	if err != nil {
		return err
	}
	projectsInLUCIConfig := stringset.NewFromSlice(projects...)
	for _, p := range curEnabledProjects {
		if !projectsInLUCIConfig.Has(p) {
			tasks = append(tasks, &tq.Task{
				Payload: &RefreshProjectConfigTask{
					Project: p,
					Disable: true,
				},
			})
		}
	}

	err = parallel.WorkPool(32, func(workCh chan<- func() error) {
		for _, task := range tasks {
			workCh <- func() error {
				if err := tq.AddTask(ctx, task); err != nil {
					logging.WithError(err).Errorf(ctx, "Failed to submit refresh task for project %q", task.Payload.(*RefreshProjectConfigTask).GetProject())
					return err
				}
				return nil
			}
		}
	})

	if err != nil {
		return err.(errors.MultiError).First()
	}
	return nil
}

// updateProject imports the latest CV Config for a given LUCI Project
// from LUCI Config if the config in CV is outdated.
func updateProject(ctx context.Context, project string) error {
	externalHash, err := getExternalContentHash(ctx, project)
	if err != nil {
		return err
	}
	existingPC := ProjectConfig{Project: project}
	switch err := datastore.Get(ctx, &existingPC); {
	case err != nil && err != datastore.ErrNoSuchEntity:
		return errors.Annotate(err, "failed to get ProjectConfig(project=%q)", project).Tag(transient.Tag).Err()
	case existingPC.ExternalHash == externalHash:
		return nil // up-to-date
	}

	cfg, err := fetchCfg(ctx, project, externalHash)
	if err != nil {
		return err
	}

	// Write out ConfigHashInfo if missing and all ConfigGroups.
	localHash := computeHash(cfg)
	cgNames := make([]string, len(cfg.GetConfigGroups()))
	for i, cg := range cfg.GetConfigGroups() {
		cgNames[i] = makeConfigGroupName(cg.GetName(), i)
	}
	targetEVersion := existingPC.EVersion + 1

	err = datastore.RunInTransaction(ctx, func(ctx context.Context) error {
		hashInfo := ConfigHashInfo{
			Hash:    localHash,
			Project: datastore.MakeKey(ctx, projectConfigKind, project),
		}
		switch err := datastore.Get(ctx, &hashInfo); {
		case err != nil && err != datastore.ErrNoSuchEntity:
			return errors.Annotate(err, "failed to get ConfigHashInfo(Hash=%q)", localHash).Tag(transient.Tag).Err()
		case err == nil && hashInfo.ProjectEVersion >= targetEVersion:
			return nil // Do not go backwards.
		default:
			hashInfo.ProjectEVersion = targetEVersion
			hashInfo.UpdateTime = datastore.RoundTime(clock.Now(ctx))
			hashInfo.ConfigGroupNames = cgNames
			return errors.Annotate(datastore.Put(ctx, &hashInfo), "failed to put ConfigHashInfo(Hash=%q)", localHash).Tag(transient.Tag).Err()
		}
	}, nil)
	if err != nil {
		return errors.Annotate(err, "failed to run transaction to update ConfigHashInfo").Tag(transient.Tag).Err()
	}

	if err := putConfigGroups(ctx, cfg, project, localHash); err != nil {
		return err
	}

	err = datastore.RunInTransaction(ctx, func(ctx context.Context) error {
		pc := ProjectConfig{Project: project}
		switch err := datastore.Get(ctx, &pc); {
		case err != nil && err != datastore.ErrNoSuchEntity:
			return errors.Annotate(err, "failed to get ProjectConfig(project=%q)", project).Tag(transient.Tag).Err()
		case pc.EVersion != existingPC.EVersion:
			return nil // Already updated by concurrent updateProject.
		default:
			now := datastore.RoundTime(clock.Now(ctx))
			pc = ProjectConfig{
				Project:          project,
				Enabled:          true,
				UpdateTime:       now,
				EVersion:         targetEVersion,
				Hash:             localHash,
				ExternalHash:     externalHash,
				ConfigGroupNames: cgNames,
			}
			return errors.Annotate(datastore.Put(ctx, &pc), "failed to put ProjectConfig(project=%q)", project).Tag(transient.Tag).Err()
			// TODO(yiwzhang): Notify ProjectManager
		}
	}, nil)
	return errors.Annotate(err, "failed to run transaction to update ProjectConfig").Tag(transient.Tag).Err()
}

func getExternalContentHash(ctx context.Context, project string) (string, error) {
	var meta config.Meta
	switch err := cfgclient.Get(ctx, config.ProjectSet(project), configFileName, nil, &meta); {
	case err != nil:
		return "", errors.Annotate(err, "failed to fetch meta from LUCI Config").Tag(transient.Tag).Err()
	case meta.ContentHash == "":
		return "", errors.Reason("LUCI Config returns empty content hash for project %q", project).Err()
	default:
		return meta.ContentHash, nil
	}
}

func fetchCfg(ctx context.Context, project string, contentHash string) (*pb.Config, error) {
	content, err := cfgclient.Client(ctx).GetConfigByHash(ctx, contentHash)
	if err != nil {
		return nil, errors.Annotate(err, "failed to get config by content hash").Tag(transient.Tag).Err()
	}
	ret := &pb.Config{}
	if err := cfgclient.ProtoText(ret)(content); err != nil {
		return nil, errors.Annotate(err, "failed to deserialize config content").Err()
	}
	// TODO(yiwzhang): validate the config here again to prevent ingesting a bad
	// version of config that accidentally slips into LUCI Config.
	// See: go.chromium.org/luci/cq/appengine/config
	return ret, nil
}

// disableProject disables the given LUCI Project if it is currently enabled.
func disableProject(ctx context.Context, project string) error {
	err := datastore.RunInTransaction(ctx, func(ctx context.Context) error {
		pc := ProjectConfig{Project: project}
		switch err := datastore.Get(ctx, &pc); {
		case datastore.IsErrNoSuchEntity(err):
			return nil // No-op when disabling non-existent Project
		case err != nil:
			return errors.Annotate(err, "failed to get existing ProjectConfig").Tag(transient.Tag).Err()
		case !pc.Enabled:
			return nil // Already disabled
		}
		pc.Enabled = false
		pc.UpdateTime = datastore.RoundTime(clock.Now(ctx))
		pc.EVersion++
		if err := datastore.Put(ctx, &pc); err != nil {
			return errors.Annotate(err, "failed to put ProjectConfig").Tag(transient.Tag).Err()
		}
		// TODO(yiwzhang): Notify ProjectManager
		return nil
	}, nil)
	return errors.Annotate(err, "failed to run transaction to disable project %q", project).Tag(transient.Tag).Err()
}
<commit_msg>cv: log when project config changes.<commit_after>// Copyright 2020 The LUCI 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 config

import (
	"context"

	"google.golang.org/protobuf/proto"

	"go.chromium.org/luci/common/clock"
	"go.chromium.org/luci/common/data/stringset"
	"go.chromium.org/luci/common/errors"
	"go.chromium.org/luci/common/logging"
	"go.chromium.org/luci/common/retry/transient"
	"go.chromium.org/luci/common/sync/parallel"
	"go.chromium.org/luci/config"
	"go.chromium.org/luci/config/cfgclient"
	pb "go.chromium.org/luci/cv/api/config/v2"
	"go.chromium.org/luci/gae/service/datastore"
	"go.chromium.org/luci/server/tq"
)

const configFileName = "commit-queue.cfg"

func init() {
	tq.RegisterTaskClass(tq.TaskClass{
		ID:        "refresh-project-config",
		Prototype: &RefreshProjectConfigTask{},
		Queue:     "refresh-project-config",
		Handler: func(ctx context.Context, payload proto.Message) error {
			task := payload.(*RefreshProjectConfigTask)
			action, actionFn := "update", updateProject
			if task.GetDisable() {
				action, actionFn = "disable", disableProject
			}
			project := task.GetProject()
			if err := actionFn(ctx, project); err != nil {
				errors.Log(ctx, err)
				// Explicitly not forwarding transient tag so that tq won't retry
				// because the refresh task is submitted every minute by AppEngine
				// Cron.
				return errors.Reason("failed to %s project %q: %s", action, project, err).Err()
			}
			return nil
		},
	})
}

// SubmitRefreshTasks submits tasks that update config for LUCI projects
// or disable projects that do not have CV config in LUCI Config.
func SubmitRefreshTasks(ctx context.Context) error {
	projects, err := cfgclient.ProjectsWithConfig(ctx, configFileName)
	if err != nil {
		return errors.Annotate(err, "failed to get projects that have %q from LUCI Config", configFileName).Tag(transient.Tag).Err()
	}
	tasks := make([]*tq.Task, len(projects))
	for i, p := range projects {
		tasks[i] = &tq.Task{
			Payload: &RefreshProjectConfigTask{
				Project: p,
			},
		}
	}

	curEnabledProjects, err := getAllProjectIDs(ctx, true)
	if err != nil {
		return err
	}
	projectsInLUCIConfig := stringset.NewFromSlice(projects...)
	for _, p := range curEnabledProjects {
		if !projectsInLUCIConfig.Has(p) {
			tasks = append(tasks, &tq.Task{
				Payload: &RefreshProjectConfigTask{
					Project: p,
					Disable: true,
				},
			})
		}
	}

	err = parallel.WorkPool(32, func(workCh chan<- func() error) {
		for _, task := range tasks {
			workCh <- func() error {
				if err := tq.AddTask(ctx, task); err != nil {
					logging.WithError(err).Errorf(ctx, "Failed to submit refresh task for project %q", task.Payload.(*RefreshProjectConfigTask).GetProject())
					return err
				}
				return nil
			}
		}
	})

	if err != nil {
		return err.(errors.MultiError).First()
	}
	return nil
}

// updateProject imports the latest CV Config for a given LUCI Project
// from LUCI Config if the config in CV is outdated.
func updateProject(ctx context.Context, project string) error {
	externalHash, err := getExternalContentHash(ctx, project)
	if err != nil {
		return err
	}
	existingPC := ProjectConfig{Project: project}
	switch err := datastore.Get(ctx, &existingPC); {
	case err != nil && err != datastore.ErrNoSuchEntity:
		return errors.Annotate(err, "failed to get ProjectConfig(project=%q)", project).Tag(transient.Tag).Err()
	case existingPC.ExternalHash == externalHash:
		return nil // up-to-date
	}

	cfg, err := fetchCfg(ctx, project, externalHash)
	if err != nil {
		return err
	}

	// Write out ConfigHashInfo if missing and all ConfigGroups.
	localHash := computeHash(cfg)
	cgNames := make([]string, len(cfg.GetConfigGroups()))
	for i, cg := range cfg.GetConfigGroups() {
		cgNames[i] = makeConfigGroupName(cg.GetName(), i)
	}
	targetEVersion := existingPC.EVersion + 1

	err = datastore.RunInTransaction(ctx, func(ctx context.Context) error {
		hashInfo := ConfigHashInfo{
			Hash:    localHash,
			Project: datastore.MakeKey(ctx, projectConfigKind, project),
		}
		switch err := datastore.Get(ctx, &hashInfo); {
		case err != nil && err != datastore.ErrNoSuchEntity:
			return errors.Annotate(err, "failed to get ConfigHashInfo(Hash=%q)", localHash).Tag(transient.Tag).Err()
		case err == nil && hashInfo.ProjectEVersion >= targetEVersion:
			return nil // Do not go backwards.
		default:
			hashInfo.ProjectEVersion = targetEVersion
			hashInfo.UpdateTime = datastore.RoundTime(clock.Now(ctx))
			hashInfo.ConfigGroupNames = cgNames
			return errors.Annotate(datastore.Put(ctx, &hashInfo), "failed to put ConfigHashInfo(Hash=%q)", localHash).Tag(transient.Tag).Err()
		}
	}, nil)
	if err != nil {
		return errors.Annotate(err, "failed to run transaction to update ConfigHashInfo").Tag(transient.Tag).Err()
	}

	if err := putConfigGroups(ctx, cfg, project, localHash); err != nil {
		return err
	}

	updated := false
	err = datastore.RunInTransaction(ctx, func(ctx context.Context) error {
		updated = false
		pc := ProjectConfig{Project: project}
		switch err := datastore.Get(ctx, &pc); {
		case err != nil && err != datastore.ErrNoSuchEntity:
			return errors.Annotate(err, "failed to get ProjectConfig(project=%q)", project).Tag(transient.Tag).Err()
		case pc.EVersion != existingPC.EVersion:
			return nil // Already updated by concurrent updateProject.
		default:
			now := datastore.RoundTime(clock.Now(ctx))
			pc = ProjectConfig{
				Project:          project,
				Enabled:          true,
				UpdateTime:       now,
				EVersion:         targetEVersion,
				Hash:             localHash,
				ExternalHash:     externalHash,
				ConfigGroupNames: cgNames,
			}
			updated = true
			return errors.Annotate(datastore.Put(ctx, &pc), "failed to put ProjectConfig(project=%q)", project).Tag(transient.Tag).Err()
			// TODO(yiwzhang): Notify ProjectManager
		}
	}, nil)

	switch {
	case err != nil:
		return errors.Annotate(err, "failed to run transaction to update ProjectConfig").Tag(transient.Tag).Err()
	case updated:
		logging.Infof(ctx, "updated project %q to rev %s hash %s ", project, "TODO", localHash)
	}
	return nil
}

func getExternalContentHash(ctx context.Context, project string) (string, error) {
	var meta config.Meta
	switch err := cfgclient.Get(ctx, config.ProjectSet(project), configFileName, nil, &meta); {
	case err != nil:
		return "", errors.Annotate(err, "failed to fetch meta from LUCI Config").Tag(transient.Tag).Err()
	case meta.ContentHash == "":
		return "", errors.Reason("LUCI Config returns empty content hash for project %q", project).Err()
	default:
		return meta.ContentHash, nil
	}
}

func fetchCfg(ctx context.Context, project string, contentHash string) (*pb.Config, error) {
	content, err := cfgclient.Client(ctx).GetConfigByHash(ctx, contentHash)
	if err != nil {
		return nil, errors.Annotate(err, "failed to get config by content hash").Tag(transient.Tag).Err()
	}
	ret := &pb.Config{}
	if err := cfgclient.ProtoText(ret)(content); err != nil {
		return nil, errors.Annotate(err, "failed to deserialize config content").Err()
	}
	// TODO(yiwzhang): validate the config here again to prevent ingesting a bad
	// version of config that accidentally slips into LUCI Config.
	// See: go.chromium.org/luci/cq/appengine/config
	return ret, nil
}

// disableProject disables the given LUCI Project if it is currently enabled.
func disableProject(ctx context.Context, project string) error {
	disabled := false

	err := datastore.RunInTransaction(ctx, func(ctx context.Context) error {
		disabled = false
		pc := ProjectConfig{Project: project}
		switch err := datastore.Get(ctx, &pc); {
		case datastore.IsErrNoSuchEntity(err):
			return nil // No-op when disabling non-existent Project
		case err != nil:
			return errors.Annotate(err, "failed to get existing ProjectConfig").Tag(transient.Tag).Err()
		case !pc.Enabled:
			return nil // Already disabled
		}
		pc.Enabled = false
		pc.UpdateTime = datastore.RoundTime(clock.Now(ctx))
		pc.EVersion++
		if err := datastore.Put(ctx, &pc); err != nil {
			return errors.Annotate(err, "failed to put ProjectConfig").Tag(transient.Tag).Err()
		}
		disabled = true
		// TODO(yiwzhang): Notify ProjectManager
		return nil
	}, nil)

	switch {
	case err != nil:
		return errors.Annotate(err, "failed to run transaction to disable project %q", project).Tag(transient.Tag).Err()
	case disabled:
		logging.Infof(ctx, "disabled project %q", project)
	}
	return nil
}
<|endoftext|>


<commit_before>package java

import (
	"bytes"
	"fmt"
	"log"
	"os"
	"path/filepath"
	"strings"
	"text/template"

	"github.com/ezbuy/tgen/global"
	"github.com/ezbuy/tgen/langs"
	"github.com/ezbuy/tgen/tmpl"
	"github.com/samuel/go-thrift/parser"
)

const (
	javaLang = "java"
	javaExt  = ".java"
)

const (
	JavaTypeshort  = "short"
	JavaTypeint    = "int"
	JavaTypelong   = "long"
	JavaTypebool   = "boolean"
	JavaTypebyte   = "byte"
	JavaTypedouble = "double"

	JavaTypeString = "String"

	JavaTypeShort  = "Short"
	JavaTypeInt    = "Integer"
	JavaTypeLong   = "Long"
	JavaTypeBool   = "Boolean"
	JavaTypeByte   = "Byte"
	JavaTypeDouble = "Double"

	// other types (such as array, map, etc.) are implemented in the method 'Typecast'
)

const (
	TPL_ENUM    = "tgen/java/enum"
	TPL_STRUCT  = "tgen/java/struct"
	TPL_SERVICE = "tgen/java/service"
)

var plaintypemapping = map[string]string{
	langs.ThriftTypeI16:    JavaTypeshort,
	langs.ThriftTypeI32:    JavaTypeint,
	langs.ThriftTypeI64:    JavaTypelong,
	langs.ThriftTypeString: JavaTypeString,
	langs.ThriftTypeByte:   JavaTypebyte,
	langs.ThriftTypeBool:   JavaTypebool,
	langs.ThriftTypeDouble: JavaTypedouble,
}

var objecttypemapping = map[string]string{
	langs.ThriftTypeI16:    JavaTypeShort,
	langs.ThriftTypeI32:    JavaTypeInt,
	langs.ThriftTypeI64:    JavaTypeLong,
	langs.ThriftTypeString: JavaTypeString,
	langs.ThriftTypeByte:   JavaTypeByte,
	langs.ThriftTypeBool:   JavaTypeBool,
	langs.ThriftTypeDouble: JavaTypeDouble,
}

type JavaGen struct {
	langs.BaseGen
}

type BaseJava struct {
	Namespace string
	t         *parser.Thrift
	ts        map[string]*parser.Thrift
}

func (b *BaseJava) FilterVariableName(n string) string {
	if b.IsKeyword(n) {
		return fmt.Sprintf("t%s%s", strings.ToUpper(n[:1]), n[1:])
	}
	return n
}

func (b *BaseJava) IsKeyword(n string) bool {
	switch n {
	case "package", "int", "short", "long", "byte", "boolean", "case", "switch", "if ", "for", "else",
		"goto", "Integer", "Short", "Long", "Byte", "Boolean", "class", "break", "try", "catch",
		"double", "Double", "do", "while", "final", "finally", "continue", "interface", "private",
		"public", "protected", "return", "this", "throw", "static", "super", "throws",
		"true", "false", "float", "volatile", "synchronized", "abstract", "default", "extends",
		"native", "new":
		return true
	}
	return false
}

func (b *BaseJava) PlainTypecast(t *parser.Type) string {
	return b.typecast(t, true)
}

func (b *BaseJava) ObjectTypecast(t *parser.Type) string {
	return b.typecast(t, false)
}

func (b *BaseJava) typecast(t *parser.Type, isplain bool) string {
	if t == nil {
		if isplain {
			return "void"
		} else {
			return "Void"
		}
	}

	var typemapping map[string]string

	if isplain {
		typemapping = plaintypemapping
	} else {
		typemapping = objecttypemapping
	}

	if t, ok := typemapping[t.Name]; ok {
		return t
	}

	switch t.Name {
	case langs.ThriftTypeList, langs.ThriftTypeSet:
		return fmt.Sprintf("java.util.ArrayList<%s>", b.ObjectTypecast(t.ValueType))
	case langs.ThriftTypeMap:
		return fmt.Sprintf("java.util.Map<%s, %s>", b.ObjectTypecast(t.KeyType), b.ObjectTypecast(t.ValueType))
	default:
		s := strings.Split(t.Name, ".")
		if len(s) == 1 {
			return s[0]
		} else if len(s) == 2 {
			pkg := ""
			for k, v := range b.t.Includes {
				if k == s[0] {

					for p, t := range b.ts {
						if v == p {
							pkg = t.Namespaces[javaLang]
							break
						}
					}

					break
				}
			}

			if pkg == "" {
				return s[1]
			}
			return fmt.Sprintf("%s.%s", pkg, s[1])
		} else {
			return t.Name
		}
	}
}

func (b *BaseJava) AssembleParams(method *parser.Method) string {
	var buf bytes.Buffer

	for i, arg := range method.Arguments {
		if i != 0 {
			buf.WriteString(", ")
		}

		buf.WriteString(fmt.Sprintf("final %s %s", b.PlainTypecast(arg.Type), b.FilterVariableName(arg.Name)))
	}

	if len(method.Arguments) == 0 {
		buf.WriteString("")
	} else {
		buf.WriteString(", ")
	}

	buf.WriteString(fmt.Sprintf("final Listener<%s> listener", b.ObjectTypecast(method.ReturnType)))

	return buf.String()
}

func (b *BaseJava) GetInnerType(t *parser.Type) string {
	if t == nil {
		return "Void"
	}

	// map is ignored
	if t.Name == langs.ThriftTypeList || t.Name == langs.ThriftTypeSet {
		return b.GetInnerType(t.ValueType)
	}

	return b.ObjectTypecast(t)
}

// IsEnum checks whether a type is enum.
// it first checks in its own definition, than check from included files
func (b *BaseJava) IsEnum(t *parser.Type) bool {
	if t == nil {
		return false
	}

	names := strings.Split(t.Name, ".")

	if len(names) == 1 {
		for n := range b.t.Enums {
			if n == t.Name {
				return true
			}
		}

		return false
	}

	for path, thrift := range b.ts {
		if thrift == b.t {
			continue
		}

		filename := strings.TrimSuffix(filepath.Base(path), filepath.Ext(path))
		if filename != names[0] {
			continue
		}

		for n := range thrift.Enums {
			if n == names[1] {
				return true
			}
		}
	}

	return false
}

type javaEnum struct {
	*BaseJava
	*parser.Enum
}

func (e *javaEnum) GenerateProperties() string {
	size := len(e.Enum.Values)
	i := 0

	var buf bytes.Buffer
	for _, v := range e.Enum.Values {
		buf.WriteString(fmt.Sprintf("\t@SerializedName(\"%d\")\n", v.Value))
		buf.WriteString(fmt.Sprintf("\t%s(%d)", strings.ToUpper(v.Name), v.Value))
		if i == size-1 {
			buf.WriteString(";\n")
		} else {
			buf.WriteString(",\n")
		}
		i++
	}
	return buf.String()
}

type javaStruct struct {
	*BaseJava
	*parser.Struct
}

func (this *javaStruct) HasKeyword() bool {
	for _, f := range this.Struct.Fields {
		if this.BaseJava.IsKeyword(f.Name) {
			return true
		}
	}
	return false
}

type javaService struct {
	*BaseJava
	*parser.Service
}

func generateAll(gen *JavaGen, output string, parsedThrift map[string]*parser.Thrift) {
	generateWithModel(gen, global.MODE_REST, filepath.Join(output, global.MODE_REST), parsedThrift)
	generateWithModel(gen, global.MODE_JSONRPC, filepath.Join(output, global.MODE_JSONRPC), parsedThrift)
}

func (g *JavaGen) Generate(output string, parsedThrift map[string]*parser.Thrift) {
	if global.Mode != "" {
		generateWithModel(g, global.Mode, output, parsedThrift)
	} else {
		generateAll(g, output, parsedThrift)
	}
}

func generateWithModel(gen *JavaGen, m string, output string, parsedThrift map[string]*parser.Thrift) {
	if m != global.MODE_REST && m != global.MODE_JSONRPC {
		log.Fatalf("mode '%s' is invalid", m)
	}

	gen.BaseGen.Init(javaLang, parsedThrift)

	if err := os.MkdirAll(output, 0755); err != nil {
		panic(fmt.Errorf("failed to create output directory %s", output))
	}

	// init templates
	enumTpl := initemplate(TPL_ENUM, "tmpl/java/enum.gojava")

	var structTpl *template.Template
	var serviceTpl *template.Template
	if m == global.MODE_REST {
		structTpl = initemplate(TPL_STRUCT, "tmpl/java/rest_struct.gojava")
		serviceTpl = initemplate(TPL_SERVICE, "tmpl/java/rest_service.gojava")
	} else if m == global.MODE_JSONRPC {
		structTpl = initemplate(TPL_STRUCT, "tmpl/java/jsonrpc_struct.gojava")
		serviceTpl = initemplate(TPL_SERVICE, "tmpl/java/jsonrpc_service.gojava")
	}

	// key is the absoule path of thrift file
	for f, t := range parsedThrift {
		if f != global.InputFile {
			continue // ignore
		}

		ns := t.Namespaces[javaLang]

		if len(t.Enums) > 0 {
			log.Printf("## enums")

			for _, e := range t.Enums {
				name := e.Name + javaExt

				// fix java file path
				p := filepath.Join(output, strings.Replace(ns, ".", "/", -1))
				if err := os.MkdirAll(p, 0755); err != nil {
					panic(fmt.Errorf("failed to create output directory %s", p))
				}

				path := filepath.Join(p, name)

				base := BaseJava{Namespace: ns, t: t, ts: parsedThrift}
				data := &javaEnum{BaseJava: &base, Enum: e}

				if err := outputfile(path, enumTpl, TPL_ENUM, data); err != nil {
					panic(fmt.Errorf("failed to write file %s. error: %v\n", path, err))
				}

				log.Printf("%s", path)
			}
		}

		log.Printf("## structs")

		for _, s := range t.Structs {
			// filename is the struct name
			name := s.Name + javaExt

			// fix java file path
			p := filepath.Join(output, strings.Replace(ns, ".", "/", -1))
			if err := os.MkdirAll(p, 0755); err != nil {
				panic(fmt.Errorf("failed to create output directory %s", p))
			}

			path := filepath.Join(p, name)

			base := BaseJava{Namespace: ns, t: t, ts: parsedThrift}
			data := &javaStruct{BaseJava: &base, Struct: s}

			if err := outputfile(path, structTpl, TPL_STRUCT, data); err != nil {
				panic(fmt.Errorf("failed to write file %s. error: %v\n", path, err))
			}

			log.Printf("%s", path)
		}

		log.Printf("## services")

		for _, s := range t.Services {
			// filename is the service name plus 'Service'
			name := s.Name + "Service" + javaExt

			// fix java file path
			p := filepath.Join(output, strings.Replace(ns, ".", "/", -1))
			if err := os.MkdirAll(p, 0755); err != nil {
				panic(fmt.Errorf("failed to create output directory %s", p))
			}

			path := filepath.Join(p, name)

			base := BaseJava{Namespace: ns, t: t, ts: parsedThrift}
			data := &javaService{BaseJava: &base, Service: s}

			if err := outputfile(path, serviceTpl, TPL_SERVICE, data); err != nil {
				panic(fmt.Errorf("failed to write file %s. error: %v\n", path, err))
			}

			log.Printf("%s", path)
		}
	}
}

func initemplate(n string, path string) *template.Template {
	data, err := tmpl.Asset(path)
	if err != nil {
		panic(err)
	}

	tpl := template.New(n)

	// register functions
	funcMap := template.FuncMap{
		"ToUpper": strings.ToUpper,
	}
	tpl.Funcs(funcMap)

	_, err = tpl.Parse(string(data))
	if err != nil {
		panic(err)
	}

	return tpl
}

func outputfile(fp string, t *template.Template, tplname string, data interface{}) error {
	file, err := os.OpenFile(fp, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0644)
	if err != nil {
		return err
	}

	defer file.Close()

	return t.ExecuteTemplate(file, tplname, data)
}

func init() {
	langs.Langs[javaLang] = &JavaGen{}
}
<commit_msg>remove unnecessary check<commit_after>package java

import (
	"bytes"
	"fmt"
	"log"
	"os"
	"path/filepath"
	"strings"
	"text/template"

	"github.com/ezbuy/tgen/global"
	"github.com/ezbuy/tgen/langs"
	"github.com/ezbuy/tgen/tmpl"
	"github.com/samuel/go-thrift/parser"
)

const (
	javaLang = "java"
	javaExt  = ".java"
)

const (
	JavaTypeshort  = "short"
	JavaTypeint    = "int"
	JavaTypelong   = "long"
	JavaTypebool   = "boolean"
	JavaTypebyte   = "byte"
	JavaTypedouble = "double"

	JavaTypeString = "String"

	JavaTypeShort  = "Short"
	JavaTypeInt    = "Integer"
	JavaTypeLong   = "Long"
	JavaTypeBool   = "Boolean"
	JavaTypeByte   = "Byte"
	JavaTypeDouble = "Double"

	// other types (such as array, map, etc.) are implemented in the method 'Typecast'
)

const (
	TPL_ENUM    = "tgen/java/enum"
	TPL_STRUCT  = "tgen/java/struct"
	TPL_SERVICE = "tgen/java/service"
)

var plaintypemapping = map[string]string{
	langs.ThriftTypeI16:    JavaTypeshort,
	langs.ThriftTypeI32:    JavaTypeint,
	langs.ThriftTypeI64:    JavaTypelong,
	langs.ThriftTypeString: JavaTypeString,
	langs.ThriftTypeByte:   JavaTypebyte,
	langs.ThriftTypeBool:   JavaTypebool,
	langs.ThriftTypeDouble: JavaTypedouble,
}

var objecttypemapping = map[string]string{
	langs.ThriftTypeI16:    JavaTypeShort,
	langs.ThriftTypeI32:    JavaTypeInt,
	langs.ThriftTypeI64:    JavaTypeLong,
	langs.ThriftTypeString: JavaTypeString,
	langs.ThriftTypeByte:   JavaTypeByte,
	langs.ThriftTypeBool:   JavaTypeBool,
	langs.ThriftTypeDouble: JavaTypeDouble,
}

type JavaGen struct {
	langs.BaseGen
}

type BaseJava struct {
	Namespace string
	t         *parser.Thrift
	ts        map[string]*parser.Thrift
}

func (b *BaseJava) FilterVariableName(n string) string {
	if b.IsKeyword(n) {
		return fmt.Sprintf("t%s%s", strings.ToUpper(n[:1]), n[1:])
	}
	return n
}

func (b *BaseJava) IsKeyword(n string) bool {
	switch n {
	case "package", "int", "short", "long", "byte", "boolean", "case", "switch", "if ", "for", "else",
		"goto", "Integer", "Short", "Long", "Byte", "Boolean", "class", "break", "try", "catch",
		"double", "Double", "do", "while", "final", "finally", "continue", "interface", "private",
		"public", "protected", "return", "this", "throw", "static", "super", "throws",
		"true", "false", "float", "volatile", "synchronized", "abstract", "default", "extends",
		"native", "new":
		return true
	}
	return false
}

func (b *BaseJava) PlainTypecast(t *parser.Type) string {
	return b.typecast(t, true)
}

func (b *BaseJava) ObjectTypecast(t *parser.Type) string {
	return b.typecast(t, false)
}

func (b *BaseJava) typecast(t *parser.Type, isplain bool) string {
	if t == nil {
		if isplain {
			return "void"
		} else {
			return "Void"
		}
	}

	var typemapping map[string]string

	if isplain {
		typemapping = plaintypemapping
	} else {
		typemapping = objecttypemapping
	}

	if t, ok := typemapping[t.Name]; ok {
		return t
	}

	switch t.Name {
	case langs.ThriftTypeList, langs.ThriftTypeSet:
		return fmt.Sprintf("java.util.ArrayList<%s>", b.ObjectTypecast(t.ValueType))
	case langs.ThriftTypeMap:
		return fmt.Sprintf("java.util.Map<%s, %s>", b.ObjectTypecast(t.KeyType), b.ObjectTypecast(t.ValueType))
	default:
		s := strings.Split(t.Name, ".")
		if len(s) == 1 {
			return s[0]
		} else if len(s) == 2 {
			pkg := ""
			for k, v := range b.t.Includes {
				if k == s[0] {

					for p, t := range b.ts {
						if v == p {
							pkg = t.Namespaces[javaLang]
							break
						}
					}

					break
				}
			}

			if pkg == "" {
				return s[1]
			}
			return fmt.Sprintf("%s.%s", pkg, s[1])
		} else {
			return t.Name
		}
	}
}

func (b *BaseJava) AssembleParams(method *parser.Method) string {
	var buf bytes.Buffer

	for i, arg := range method.Arguments {
		if i != 0 {
			buf.WriteString(", ")
		}

		buf.WriteString(fmt.Sprintf("final %s %s", b.PlainTypecast(arg.Type), b.FilterVariableName(arg.Name)))
	}

	if len(method.Arguments) == 0 {
		buf.WriteString("")
	} else {
		buf.WriteString(", ")
	}

	buf.WriteString(fmt.Sprintf("final Listener<%s> listener", b.ObjectTypecast(method.ReturnType)))

	return buf.String()
}

func (b *BaseJava) GetInnerType(t *parser.Type) string {
	if t == nil {
		return "Void"
	}

	// map is ignored
	if t.Name == langs.ThriftTypeList || t.Name == langs.ThriftTypeSet {
		return b.GetInnerType(t.ValueType)
	}

	return b.ObjectTypecast(t)
}

// IsEnum checks whether a type is enum.
// it first checks in its own definition, than check from included files
func (b *BaseJava) IsEnum(t *parser.Type) bool {
	if t == nil {
		return false
	}

	names := strings.Split(t.Name, ".")

	if len(names) == 1 {
		for n := range b.t.Enums {
			if n == t.Name {
				return true
			}
		}

		return false
	}

	for path, thrift := range b.ts {
		if thrift == b.t {
			continue
		}

		filename := strings.TrimSuffix(filepath.Base(path), filepath.Ext(path))
		if filename != names[0] {
			continue
		}

		for n := range thrift.Enums {
			if n == names[1] {
				return true
			}
		}
	}

	return false
}

type javaEnum struct {
	*BaseJava
	*parser.Enum
}

func (e *javaEnum) GenerateProperties() string {
	size := len(e.Enum.Values)
	i := 0

	var buf bytes.Buffer
	for _, v := range e.Enum.Values {
		buf.WriteString(fmt.Sprintf("\t@SerializedName(\"%d\")\n", v.Value))
		buf.WriteString(fmt.Sprintf("\t%s(%d)", strings.ToUpper(v.Name), v.Value))
		if i == size-1 {
			buf.WriteString(";\n")
		} else {
			buf.WriteString(",\n")
		}
		i++
	}
	return buf.String()
}

type javaStruct struct {
	*BaseJava
	*parser.Struct
}

func (this *javaStruct) HasKeyword() bool {
	for _, f := range this.Struct.Fields {
		if this.BaseJava.IsKeyword(f.Name) {
			return true
		}
	}
	return false
}

type javaService struct {
	*BaseJava
	*parser.Service
}

func generateAll(gen *JavaGen, output string, parsedThrift map[string]*parser.Thrift) {
	generateWithModel(gen, global.MODE_REST, filepath.Join(output, global.MODE_REST), parsedThrift)
	generateWithModel(gen, global.MODE_JSONRPC, filepath.Join(output, global.MODE_JSONRPC), parsedThrift)
}

func (g *JavaGen) Generate(output string, parsedThrift map[string]*parser.Thrift) {
	if global.Mode != "" {
		generateWithModel(g, global.Mode, output, parsedThrift)
	} else {
		generateAll(g, output, parsedThrift)
	}
}

func generateWithModel(gen *JavaGen, m string, output string, parsedThrift map[string]*parser.Thrift) {
	if m != global.MODE_REST && m != global.MODE_JSONRPC {
		log.Fatalf("mode '%s' is invalid", m)
	}

	gen.BaseGen.Init(javaLang, parsedThrift)

	if err := os.MkdirAll(output, 0755); err != nil {
		panic(fmt.Errorf("failed to create output directory %s", output))
	}

	// init templates
	enumTpl := initemplate(TPL_ENUM, "tmpl/java/enum.gojava")

	var structTpl *template.Template
	var serviceTpl *template.Template
	if m == global.MODE_REST {
		structTpl = initemplate(TPL_STRUCT, "tmpl/java/rest_struct.gojava")
		serviceTpl = initemplate(TPL_SERVICE, "tmpl/java/rest_service.gojava")
	} else if m == global.MODE_JSONRPC {
		structTpl = initemplate(TPL_STRUCT, "tmpl/java/jsonrpc_struct.gojava")
		serviceTpl = initemplate(TPL_SERVICE, "tmpl/java/jsonrpc_service.gojava")
	}

	// key is the absoule path of thrift file
	for _, t := range parsedThrift {
		// due to java's features,
		// we generate the struct and service in seperate template file

		ns := t.Namespaces["java"]

		log.Printf("## structs")

		for _, s := range t.Structs {
			// filename is the struct name
			name := s.Name + javaExt

			// fix java file path
			p := filepath.Join(output, strings.Replace(ns, ".", "/", -1))
			if err := os.MkdirAll(p, 0755); err != nil {
				panic(fmt.Errorf("failed to create output directory %s", p))
			}

			path := filepath.Join(p, name)

			base := BaseJava{Namespace: ns, t: t, ts: parsedThrift}
			data := &javaStruct{BaseJava: &base, Struct: s}

			if err := outputfile(path, structTpl, TPL_STRUCT, data); err != nil {
				panic(fmt.Errorf("failed to write file %s. error: %v\n", path, err))
			}

			log.Printf("%s", path)
		}

		log.Printf("## services")

		for _, s := range t.Services {
			// filename is the service name plus 'Service'
			name := s.Name + "Service" + javaExt

			// fix java file path
			p := filepath.Join(output, strings.Replace(ns, ".", "/", -1))
			if err := os.MkdirAll(p, 0755); err != nil {
				panic(fmt.Errorf("failed to create output directory %s", p))
			}

			path := filepath.Join(p, name)

			base := BaseJava{Namespace: ns, t: t, ts: parsedThrift}
			data := &javaService{BaseJava: &base, Service: s}

			if err := outputfile(path, serviceTpl, TPL_SERVICE, data); err != nil {
				panic(fmt.Errorf("failed to write file %s. error: %v\n", path, err))
			}

			log.Printf("%s", path)
		}
	}
}

func initemplate(n string, path string) *template.Template {
	data, err := tmpl.Asset(path)
	if err != nil {
		panic(err)
	}

	tpl := template.New(n)

	// register functions
	funcMap := template.FuncMap{
		"ToUpper": strings.ToUpper,
	}
	tpl.Funcs(funcMap)

	_, err = tpl.Parse(string(data))
	if err != nil {
		panic(err)
	}

	return tpl
}

func outputfile(fp string, t *template.Template, tplname string, data interface{}) error {
	file, err := os.OpenFile(fp, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0644)
	if err != nil {
		return err
	}

	defer file.Close()

	return t.ExecuteTemplate(file, tplname, data)
}

func init() {
	langs.Langs[javaLang] = &JavaGen{}
}
<|endoftext|>


<commit_before>package server

import (
	"github.com/LapisBlue/Lapitar/lapitar/mc"
	"github.com/LapisBlue/Lapitar/lapitar/server/cache"
	"github.com/LapisBlue/Lapitar/lapitar/util"
	"github.com/zenazn/goji/web"
	"image"
	"image/png"
	"log"
	"net/http"
	"strconv"
	"time"
)

func printError(err error, message ...interface{}) {
	if err == nil {
		return
	}

	log.Println(message...)
	log.Println(err)
}

func parseSize(c web.C, def int) (result int) {
	size := c.URLParams["size"]
	result, err := strconv.Atoi(size)
	if err != nil {
		printError(err, "Failed to parse size", size)
		return def
	}
	return
}

func loadSkinMeta(name string, watch *util.StopWatch) (skin cache.SkinMeta) {
	watch.Mark()

	if mc.IsName(name) {
		skin = cache.FetchByName(name)
	} else {
		name = mc.ParseUUID(name)
		if mc.IsUUID(name) {
			skin = cache.Fetch(name)
		} else {
			skin = cache.FallbackByName(name)
		}
	}

	log.Println("Loaded skin:", skin.Profile().Name(), watch)
	return
}

const (
	keepCache    = 24 * time.Hour
	cacheControl = "max-age=86400" // 24*60*60, one day in seconds
)

func serveCached(w http.ResponseWriter, r *http.Request, meta cache.SkinMeta) bool {
	if tag := r.Header.Get("If-None-Match"); tag == meta.ID() {
		prepareResponse(w, r, meta)
		w.WriteHeader(http.StatusNotModified)
		return true
	}

	return false
}

func prepareResponse(w http.ResponseWriter, r *http.Request, meta cache.SkinMeta) {
	w.Header().Add("Cache-Control", cacheControl)
	w.Header().Add("Expires", time.Now().Add(keepCache).UTC().Format(http.TimeFormat))
	w.Header().Add("ETag", meta.ID())
	//w.Header().Add("Last-Modified", meta.LastMod().UTC().Format(http.TimeFormat))
}

func sendResult(w http.ResponseWriter, player string, result image.Image, watch *util.StopWatch) (err error) {
	watch.Mark()
	w.Header().Add("Content-Type", "image/png")
	err = png.Encode(w, result)
	if err == nil {
		log.Println("Sent response:", player, watch)
	} else {
		printError(err, "Failed to send response:", player, watch)
	}

	return
}
<commit_msg>Add Last-Modified header again<commit_after>package server

import (
	"github.com/LapisBlue/Lapitar/lapitar/mc"
	"github.com/LapisBlue/Lapitar/lapitar/server/cache"
	"github.com/LapisBlue/Lapitar/lapitar/util"
	"github.com/zenazn/goji/web"
	"image"
	"image/png"
	"log"
	"net/http"
	"strconv"
	"time"
)

func printError(err error, message ...interface{}) {
	if err == nil {
		return
	}

	log.Println(message...)
	log.Println(err)
}

func parseSize(c web.C, def int) (result int) {
	size := c.URLParams["size"]
	result, err := strconv.Atoi(size)
	if err != nil {
		printError(err, "Failed to parse size", size)
		return def
	}
	return
}

func loadSkinMeta(name string, watch *util.StopWatch) (skin cache.SkinMeta) {
	watch.Mark()

	if mc.IsName(name) {
		skin = cache.FetchByName(name)
	} else {
		name = mc.ParseUUID(name)
		if mc.IsUUID(name) {
			skin = cache.Fetch(name)
		} else {
			skin = cache.FallbackByName(name)
		}
	}

	log.Println("Loaded skin:", skin.Profile().Name(), watch)
	return
}

const (
	keepCache    = 24 * time.Hour
	cacheControl = "max-age=86400" // 24*60*60, one day in seconds
)

func serveCached(w http.ResponseWriter, r *http.Request, meta cache.SkinMeta) bool {
	if tag := r.Header.Get("If-None-Match"); tag == meta.ID() {
		prepareResponse(w, r, meta)
		w.WriteHeader(http.StatusNotModified)
		return true
	}

	return false
}

func prepareResponse(w http.ResponseWriter, r *http.Request, meta cache.SkinMeta) {
	w.Header().Add("Cache-Control", cacheControl)
	w.Header().Add("Expires", time.Now().Add(keepCache).UTC().Format(http.TimeFormat))
	w.Header().Add("ETag", meta.ID())
	w.Header().Add("Last-Modified", meta.Timestamp().UTC().Format(http.TimeFormat))
}

func sendResult(w http.ResponseWriter, player string, result image.Image, watch *util.StopWatch) (err error) {
	watch.Mark()
	w.Header().Add("Content-Type", "image/png")
	err = png.Encode(w, result)
	if err == nil {
		log.Println("Sent response:", player, watch)
	} else {
		printError(err, "Failed to send response:", player, watch)
	}

	return
}
<|endoftext|>


<commit_before>package fnet

import (
	"errors"
	"fmt"
	"github.com/viphxin/xingo/iface"
	"github.com/viphxin/xingo/logger"
	"sync"
)

type ConnectionMgr struct {
	connections map[uint32]iface.Iconnection
	conMrgLock  sync.RWMutex
}

func (this *ConnectionMgr) Add(conn iface.Iconnection) {
	this.conMrgLock.Lock()
	defer this.conMrgLock.Unlock()
	this.connections[conn.GetSessionId()] = conn
	logger.Debug(fmt.Sprintf("Total connection: %d", len(this.connections)))
}

func (this *ConnectionMgr) Remove(conn iface.Iconnection) error {
	this.conMrgLock.Lock()
	defer this.conMrgLock.Unlock()
	_, ok := this.connections[conn.GetSessionId()]
	if ok {
		delete(this.connections, conn.GetSessionId())
		logger.Info(len(this.connections))
		return nil
	} else {
		return errors.New("not found!!")
	}

}

func (this *ConnectionMgr) Get(sid uint32) (iface.Iconnection, error) {
	v, ok := this.connections[sid]
	if ok {
		delete(this.connections, sid)
		return v, nil
	} else {
		return nil, errors.New("not found!!")
	}
}

func (this *ConnectionMgr) Len() int {
	return len(this.connections)
}

func NewConnectionMgr() *ConnectionMgr {
	return &ConnectionMgr{
		connections: make(map[uint32]iface.Iconnection),
	}
}
<commit_msg>fix multi goroutine error<commit_after>package fnet

import (
	"errors"
	"fmt"
	"github.com/viphxin/xingo/iface"
	"github.com/viphxin/xingo/logger"
	"sync"
)

type ConnectionMgr struct {
	connections map[uint32]iface.Iconnection
	conMrgLock  sync.RWMutex
}

func (this *ConnectionMgr) Add(conn iface.Iconnection) {
	this.conMrgLock.Lock()
	defer this.conMrgLock.Unlock()
	this.connections[conn.GetSessionId()] = conn
	logger.Debug(fmt.Sprintf("Total connection: %d", len(this.connections)))
}

func (this *ConnectionMgr) Remove(conn iface.Iconnection) error {
	this.conMrgLock.Lock()
	defer this.conMrgLock.Unlock()
	_, ok := this.connections[conn.GetSessionId()]
	if ok {
		delete(this.connections, conn.GetSessionId())
		logger.Info(len(this.connections))
		return nil
	} else {
		return errors.New("not found!!")
	}

}

func (this *ConnectionMgr) Get(sid uint32) (iface.Iconnection, error) {
	this.conMrgLock.Lock()
	defer this.conMrgLock.Unlock()
	v, ok := this.connections[sid]
	if ok {
		delete(this.connections, sid)
		return v, nil
	} else {
		return nil, errors.New("not found!!")
	}
}

func (this *ConnectionMgr) Len() int {
	this.conMrgLock.Lock()
	defer this.conMrgLock.Unlock()
	return len(this.connections)
}

func NewConnectionMgr() *ConnectionMgr {
	return &ConnectionMgr{
		connections: make(map[uint32]iface.Iconnection),
	}
}
<|endoftext|>


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

package inode_test

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

	"golang.org/x/net/context"

	"github.com/googlecloudplatform/gcsfuse/fs/inode"
	"github.com/googlecloudplatform/gcsfuse/gcsproxy"
	"github.com/googlecloudplatform/gcsfuse/lease"
	"github.com/googlecloudplatform/gcsfuse/timeutil"
	"github.com/jacobsa/fuse/fuseops"
	"github.com/jacobsa/gcloud/gcs"
	"github.com/jacobsa/gcloud/gcs/gcsfake"
	"github.com/jacobsa/gcloud/gcs/gcsutil"
	. "github.com/jacobsa/ogletest"
)

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

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

const uid = 123
const gid = 456

const fileInodeID = 17
const fileInodeName = "foo/bar"
const fileMode os.FileMode = 0641

type FileTest struct {
	ctx    context.Context
	bucket gcs.Bucket
	leaser lease.FileLeaser
	clock  timeutil.SimulatedClock

	initialContents string
	backingObj      *gcs.Object

	in *inode.FileInode
}

var _ SetUpInterface = &FileTest{}
var _ TearDownInterface = &FileTest{}

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

func (t *FileTest) SetUp(ti *TestInfo) {
	t.ctx = ti.Ctx
	t.clock.SetTime(time.Date(2012, 8, 15, 22, 56, 0, 0, time.Local))
	t.leaser = lease.NewFileLeaser("", math.MaxInt32, math.MaxInt64)
	t.bucket = gcsfake.NewFakeBucket(&t.clock, "some_bucket")

	// Set up the backing object.
	var err error

	t.initialContents = "taco"
	t.backingObj, err = gcsutil.CreateObject(
		t.ctx,
		t.bucket,
		fileInodeName,
		t.initialContents)

	AssertEq(nil, err)

	// Create the inode.
	t.in = inode.NewFileInode(
		fileInodeID,
		t.backingObj,
		fuseops.InodeAttributes{
			Uid:  uid,
			Gid:  gid,
			Mode: fileMode,
		},
		math.MaxUint64, // GCS chunk size
		t.bucket,
		t.leaser,
		gcsproxy.NewObjectSyncer(
			1, // Append threshold
			".gcsfuse_tmp/",
			t.bucket),
		&t.clock)

	t.in.Lock()
}

func (t *FileTest) TearDown() {
	t.in.Unlock()
}

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

func (t *FileTest) ID() {
	ExpectEq(fileInodeID, t.in.ID())
}

func (t *FileTest) Name() {
	ExpectEq(fileInodeName, t.in.Name())
}

func (t *FileTest) InitialSourceGeneration() {
	ExpectEq(t.backingObj.Generation, t.in.SourceGeneration())
}

func (t *FileTest) InitialAttributes() {
	attrs, err := t.in.Attributes(t.ctx)
	AssertEq(nil, err)

	ExpectEq(len(t.initialContents), attrs.Size)
	ExpectEq(1, attrs.Nlink)
	ExpectEq(uid, attrs.Uid)
	ExpectEq(gid, attrs.Gid)
	ExpectEq(fileMode, attrs.Mode)
	ExpectThat(attrs.Mtime, timeutil.TimeEq(t.backingObj.Updated))
}

func (t *FileTest) Read() {
	AssertEq("taco", t.initialContents)

	// Make several reads, checking the expected contents. We should never get an
	// EOF error, since fuseops.ReadFileOp is not supposed to see those.
	testCases := []struct {
		offset   int64
		size     int
		expected string
	}{
		{0, 1, "t"},
		{0, 2, "ta"},
		{0, 3, "tac"},
		{0, 4, "taco"},
		{0, 5, "taco"},

		{1, 1, "a"},
		{1, 2, "ac"},
		{1, 3, "aco"},
		{1, 4, "aco"},

		{3, 1, "o"},
		{3, 2, "o"},

		// Empty ranges
		{0, 0, ""},
		{3, 0, ""},
		{4, 0, ""},
		{4, 1, ""},
		{5, 0, ""},
		{5, 1, ""},
		{5, 2, ""},
	}

	for _, tc := range testCases {
		desc := fmt.Sprintf("offset: %d, size: %d", tc.offset, tc.size)

		data, err := t.in.Read(t.ctx, tc.offset, tc.size)
		AssertEq(nil, err, "%s", desc)
		ExpectEq(tc.expected, string(data), "%s", desc)
	}
}

func (t *FileTest) Write() {
	var data []byte
	var err error

	AssertEq("taco", t.initialContents)

	// Overwite a byte.
	err = t.in.Write(t.ctx, []byte("p"), 0)
	AssertEq(nil, err)

	// Add some data at the end.
	t.clock.AdvanceTime(time.Second)
	writeTime := t.clock.Now()

	err = t.in.Write(t.ctx, []byte("burrito"), 4)
	AssertEq(nil, err)

	t.clock.AdvanceTime(time.Second)

	// Read back the content.
	data, err = t.in.Read(t.ctx, 0, 1024)
	AssertEq(nil, err)
	ExpectEq("pacoburrito", string(data))

	// Check attributes.
	attrs, err := t.in.Attributes(t.ctx)
	AssertEq(nil, err)

	ExpectEq(len("pacoburrito"), attrs.Size)
	ExpectThat(attrs.Mtime, timeutil.TimeEq(writeTime))
}

func (t *FileTest) Truncate() {
	var attrs fuseops.InodeAttributes
	var data []byte
	var err error

	AssertEq("taco", t.initialContents)

	// Truncate downward.
	t.clock.AdvanceTime(time.Second)
	truncateTime := t.clock.Now()

	err = t.in.Truncate(t.ctx, 2)
	AssertEq(nil, err)

	t.clock.AdvanceTime(time.Second)

	// Read the contents.
	data, err = t.in.Read(t.ctx, 0, 1024)
	AssertEq(nil, err)
	ExpectEq("ta", string(data))

	// Check attributes.
	attrs, err = t.in.Attributes(t.ctx)
	AssertEq(nil, err)

	ExpectEq(len("ta"), attrs.Size)
	ExpectThat(attrs.Mtime, timeutil.TimeEq(truncateTime))
}

func (t *FileTest) Sync_NotClobbered() {
	var attrs fuseops.InodeAttributes
	var err error

	// Truncate downward.
	t.clock.AdvanceTime(time.Second)

	err = t.in.Truncate(t.ctx, 2)
	AssertEq(nil, err)

	t.clock.AdvanceTime(time.Second)

	// Sync.
	err = t.in.Sync(t.ctx)
	AssertEq(nil, err)

	// The generation should have advanced.
	ExpectLt(t.backingObj.Generation, t.in.SourceGeneration())

	// Stat the current object in the bucket.
	statReq := &gcs.StatObjectRequest{Name: t.in.Name()}
	o, err := t.bucket.StatObject(t.ctx, statReq)

	AssertEq(nil, err)
	ExpectEq(t.in.SourceGeneration(), o.Generation)
	ExpectEq(2, o.Size)

	// Check attributes.
	attrs, err = t.in.Attributes(t.ctx)
	AssertEq(nil, err)

	ExpectEq(2, attrs.Size)
	ExpectThat(attrs.Mtime, timeutil.TimeEq(o.Updated))
}

func (t *FileTest) Sync_Clobbered() {
	var err error

	// Truncate downward.
	err = t.in.Truncate(t.ctx, 2)
	AssertEq(nil, err)

	// Clobber the backing object.
	newObj, err := gcsutil.CreateObject(t.ctx, t.bucket, t.in.Name(), "burrito")
	AssertEq(nil, err)

	// Sync. The call should succeed, but nothing should change.
	err = t.in.Sync(t.ctx)

	AssertEq(nil, err)
	ExpectEq(t.backingObj.Generation, t.in.SourceGeneration())

	// The object in the bucket should not have been changed.
	statReq := &gcs.StatObjectRequest{Name: t.in.Name()}
	o, err := t.bucket.StatObject(t.ctx, statReq)

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

package inode_test

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

	"golang.org/x/net/context"

	"github.com/googlecloudplatform/gcsfuse/fs/inode"
	"github.com/googlecloudplatform/gcsfuse/gcsproxy"
	"github.com/googlecloudplatform/gcsfuse/lease"
	"github.com/googlecloudplatform/gcsfuse/timeutil"
	"github.com/jacobsa/fuse/fuseops"
	"github.com/jacobsa/gcloud/gcs"
	"github.com/jacobsa/gcloud/gcs/gcsfake"
	"github.com/jacobsa/gcloud/gcs/gcsutil"
	. "github.com/jacobsa/ogletest"
)

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

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

const uid = 123
const gid = 456

const fileInodeID = 17
const fileInodeName = "foo/bar"
const fileMode os.FileMode = 0641

type FileTest struct {
	ctx    context.Context
	bucket gcs.Bucket
	leaser lease.FileLeaser
	clock  timeutil.SimulatedClock

	initialContents string
	backingObj      *gcs.Object

	in *inode.FileInode
}

var _ SetUpInterface = &FileTest{}
var _ TearDownInterface = &FileTest{}

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

func (t *FileTest) SetUp(ti *TestInfo) {
	t.ctx = ti.Ctx
	t.clock.SetTime(time.Date(2012, 8, 15, 22, 56, 0, 0, time.Local))
	t.leaser = lease.NewFileLeaser("", math.MaxInt32, math.MaxInt64)
	t.bucket = gcsfake.NewFakeBucket(&t.clock, "some_bucket")

	// Set up the backing object.
	var err error

	t.initialContents = "taco"
	t.backingObj, err = gcsutil.CreateObject(
		t.ctx,
		t.bucket,
		fileInodeName,
		t.initialContents)

	AssertEq(nil, err)

	// Create the inode.
	t.in = inode.NewFileInode(
		fileInodeID,
		t.backingObj,
		fuseops.InodeAttributes{
			Uid:  uid,
			Gid:  gid,
			Mode: fileMode,
		},
		math.MaxUint64, // GCS chunk size
		t.bucket,
		t.leaser,
		gcsproxy.NewObjectSyncer(
			1, // Append threshold
			".gcsfuse_tmp/",
			t.bucket),
		&t.clock)

	t.in.Lock()
}

func (t *FileTest) TearDown() {
	t.in.Unlock()
}

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

func (t *FileTest) ID() {
	ExpectEq(fileInodeID, t.in.ID())
}

func (t *FileTest) Name() {
	ExpectEq(fileInodeName, t.in.Name())
}

func (t *FileTest) InitialSourceGeneration() {
	ExpectEq(t.backingObj.Generation, t.in.SourceGeneration())
}

func (t *FileTest) InitialAttributes() {
	attrs, err := t.in.Attributes(t.ctx)
	AssertEq(nil, err)

	ExpectEq(len(t.initialContents), attrs.Size)
	ExpectEq(1, attrs.Nlink)
	ExpectEq(uid, attrs.Uid)
	ExpectEq(gid, attrs.Gid)
	ExpectEq(fileMode, attrs.Mode)
	ExpectThat(attrs.Mtime, timeutil.TimeEq(t.backingObj.Updated))
}

func (t *FileTest) Read() {
	AssertEq("taco", t.initialContents)

	// Make several reads, checking the expected contents. We should never get an
	// EOF error, since fuseops.ReadFileOp is not supposed to see those.
	testCases := []struct {
		offset   int64
		size     int
		expected string
	}{
		{0, 1, "t"},
		{0, 2, "ta"},
		{0, 3, "tac"},
		{0, 4, "taco"},
		{0, 5, "taco"},

		{1, 1, "a"},
		{1, 2, "ac"},
		{1, 3, "aco"},
		{1, 4, "aco"},

		{3, 1, "o"},
		{3, 2, "o"},

		// Empty ranges
		{0, 0, ""},
		{3, 0, ""},
		{4, 0, ""},
		{4, 1, ""},
		{5, 0, ""},
		{5, 1, ""},
		{5, 2, ""},
	}

	for _, tc := range testCases {
		desc := fmt.Sprintf("offset: %d, size: %d", tc.offset, tc.size)

		data, err := t.in.Read(t.ctx, tc.offset, tc.size)
		AssertEq(nil, err, "%s", desc)
		ExpectEq(tc.expected, string(data), "%s", desc)
	}
}

func (t *FileTest) Write() {
	var data []byte
	var err error

	AssertEq("taco", t.initialContents)

	// Overwite a byte.
	err = t.in.Write(t.ctx, []byte("p"), 0)
	AssertEq(nil, err)

	// Add some data at the end.
	t.clock.AdvanceTime(time.Second)
	writeTime := t.clock.Now()

	err = t.in.Write(t.ctx, []byte("burrito"), 4)
	AssertEq(nil, err)

	t.clock.AdvanceTime(time.Second)

	// Read back the content.
	data, err = t.in.Read(t.ctx, 0, 1024)
	AssertEq(nil, err)
	ExpectEq("pacoburrito", string(data))

	// Check attributes.
	attrs, err := t.in.Attributes(t.ctx)
	AssertEq(nil, err)

	ExpectEq(len("pacoburrito"), attrs.Size)
	ExpectThat(attrs.Mtime, timeutil.TimeEq(writeTime))
}

func (t *FileTest) Truncate() {
	var attrs fuseops.InodeAttributes
	var data []byte
	var err error

	AssertEq("taco", t.initialContents)

	// Truncate downward.
	t.clock.AdvanceTime(time.Second)
	truncateTime := t.clock.Now()

	err = t.in.Truncate(t.ctx, 2)
	AssertEq(nil, err)

	t.clock.AdvanceTime(time.Second)

	// Read the contents.
	data, err = t.in.Read(t.ctx, 0, 1024)
	AssertEq(nil, err)
	ExpectEq("ta", string(data))

	// Check attributes.
	attrs, err = t.in.Attributes(t.ctx)
	AssertEq(nil, err)

	ExpectEq(len("ta"), attrs.Size)
	ExpectThat(attrs.Mtime, timeutil.TimeEq(truncateTime))
}

func (t *FileTest) WriteThenSync() {
	AssertTrue(false, "TODO")
}

func (t *FileTest) AppendThenSync() {
	AssertTrue(false, "TODO")
}

func (t *FileTest) TruncateDownwardThenSync() {
	var attrs fuseops.InodeAttributes
	var err error

	// Truncate downward.
	t.clock.AdvanceTime(time.Second)

	err = t.in.Truncate(t.ctx, 2)
	AssertEq(nil, err)

	t.clock.AdvanceTime(time.Second)

	// Sync.
	err = t.in.Sync(t.ctx)
	AssertEq(nil, err)

	// The generation should have advanced.
	ExpectLt(t.backingObj.Generation, t.in.SourceGeneration())

	// Stat the current object in the bucket.
	statReq := &gcs.StatObjectRequest{Name: t.in.Name()}
	o, err := t.bucket.StatObject(t.ctx, statReq)

	AssertEq(nil, err)
	ExpectEq(t.in.SourceGeneration(), o.Generation)
	ExpectEq(2, o.Size)

	// Check attributes.
	attrs, err = t.in.Attributes(t.ctx)
	AssertEq(nil, err)

	ExpectEq(2, attrs.Size)
	ExpectThat(attrs.Mtime, timeutil.TimeEq(o.Updated))
}

func (t *FileTest) TruncateUpwardThenSync() {
	AssertTrue(false, "TODO")
}

func (t *FileTest) Sync_Clobbered() {
	var err error

	// Truncate downward.
	err = t.in.Truncate(t.ctx, 2)
	AssertEq(nil, err)

	// Clobber the backing object.
	newObj, err := gcsutil.CreateObject(t.ctx, t.bucket, t.in.Name(), "burrito")
	AssertEq(nil, err)

	// Sync. The call should succeed, but nothing should change.
	err = t.in.Sync(t.ctx)

	AssertEq(nil, err)
	ExpectEq(t.backingObj.Generation, t.in.SourceGeneration())

	// The object in the bucket should not have been changed.
	statReq := &gcs.StatObjectRequest{Name: t.in.Name()}
	o, err := t.bucket.StatObject(t.ctx, statReq)

	AssertEq(nil, err)
	ExpectEq(newObj.Generation, o.Generation)
	ExpectEq(newObj.Size, o.Size)
}
<|endoftext|>


<commit_before>package pluginutil

import (
	"crypto/sha256"
	"fmt"
	"os"
	"os/exec"
	"time"

	plugin "github.com/hashicorp/go-plugin"
	"github.com/hashicorp/vault/helper/mlock"
)

var (
	// PluginUnwrapTokenEnv is the ENV name used to pass unwrap tokens to the
	// plugin.
	PluginMlockEnabled = "VAULT_PLUGIN_MLOCK_ENABLED"
)

type Looker interface {
	LookupPlugin(string) (*PluginRunner, error)
}

type Wrapper interface {
	ResponseWrapData(data map[string]interface{}, ttl time.Duration, jwt bool) (string, error)
	MlockDisabled() bool
}

type LookWrapper interface {
	Looker
	Wrapper
}

type PluginRunner struct {
	Name    string   `json:"name"`
	Command string   `json:"command"`
	Args    []string `json:"args"`
	Sha256  []byte   `json:"sha256"`
	Builtin bool     `json:"builtin"`
}

func (r *PluginRunner) Run(wrapper Wrapper, pluginMap map[string]plugin.Plugin, hs plugin.HandshakeConfig, env []string) (*plugin.Client, error) {
	// Get a CA TLS Certificate
	CACertBytes, CACert, CAKey, err := GenerateCACert()
	if err != nil {
		return nil, err
	}

	// Use CA to sign a client cert and return a configured TLS config
	clientTLSConfig, err := CreateClientTLSConfig(CACert, CAKey)
	if err != nil {
		return nil, err
	}

	// Use CA to sign a server cert and wrap the values in a response wrapped
	// token.
	wrapToken, err := WrapServerConfig(wrapper, CACertBytes, CACert, CAKey)
	if err != nil {
		return nil, err
	}

	mlock := "true"
	if wrapper.MlockDisabled() {
		mlock = "false"
	}

	cmd := exec.Command(r.Command, r.Args...)
	cmd.Env = append(cmd.Env, env...)
	// Add the response wrap token to the ENV of the plugin
	cmd.Env = append(cmd.Env, fmt.Sprintf("%s=%s", PluginUnwrapTokenEnv, wrapToken))
	// Add the mlock setting to the ENV of the plugin
	cmd.Env = append(cmd.Env, fmt.Sprintf("%s=%s", PluginMlockEnabled, mlock))

	secureConfig := &plugin.SecureConfig{
		Checksum: r.Sha256,
		Hash:     sha256.New(),
	}

	client := plugin.NewClient(&plugin.ClientConfig{
		HandshakeConfig: hs,
		Plugins:         pluginMap,
		Cmd:             cmd,
		TLSConfig:       clientTLSConfig,
		SecureConfig:    secureConfig,
	})

	return client, nil
}

func OptionallyEnableMlock() error {
	if os.Getenv(PluginMlockEnabled) == "true" {
		return mlock.LockMemory()
	}

	return nil
}
<commit_msg>Add comments to the plugin runner<commit_after>package pluginutil

import (
	"crypto/sha256"
	"fmt"
	"os"
	"os/exec"
	"time"

	plugin "github.com/hashicorp/go-plugin"
	"github.com/hashicorp/vault/helper/mlock"
)

var (
	// PluginUnwrapTokenEnv is the ENV name used to pass unwrap tokens to the
	// plugin.
	PluginMlockEnabled = "VAULT_PLUGIN_MLOCK_ENABLED"
)

// Looker defines the plugin Lookup function that looks into the plugin catalog
// for availible plugins and returns a PluginRunner
type Looker interface {
	LookupPlugin(string) (*PluginRunner, error)
}

// Wrapper interface defines the functions needed by the runner to wrap the
// metadata needed to run a plugin process. This includes looking up Mlock
// configuration and wrapping data in a respose wrapped token.
type Wrapper interface {
	ResponseWrapData(data map[string]interface{}, ttl time.Duration, jwt bool) (string, error)
	MlockDisabled() bool
}

// LookWrapper defines the functions for both Looker and Wrapper
type LookWrapper interface {
	Looker
	Wrapper
}

// PluginRunner defines the metadata needed to run a plugin securely with
// go-plugin.
type PluginRunner struct {
	Name    string   `json:"name"`
	Command string   `json:"command"`
	Args    []string `json:"args"`
	Sha256  []byte   `json:"sha256"`
	Builtin bool     `json:"builtin"`
}

// Run takes a wrapper instance, and the go-plugin paramaters and executes a
// plugin.
func (r *PluginRunner) Run(wrapper Wrapper, pluginMap map[string]plugin.Plugin, hs plugin.HandshakeConfig, env []string) (*plugin.Client, error) {
	// Get a CA TLS Certificate
	CACertBytes, CACert, CAKey, err := GenerateCACert()
	if err != nil {
		return nil, err
	}

	// Use CA to sign a client cert and return a configured TLS config
	clientTLSConfig, err := CreateClientTLSConfig(CACert, CAKey)
	if err != nil {
		return nil, err
	}

	// Use CA to sign a server cert and wrap the values in a response wrapped
	// token.
	wrapToken, err := WrapServerConfig(wrapper, CACertBytes, CACert, CAKey)
	if err != nil {
		return nil, err
	}

	mlock := "true"
	if wrapper.MlockDisabled() {
		mlock = "false"
	}

	cmd := exec.Command(r.Command, r.Args...)
	cmd.Env = append(cmd.Env, env...)
	// Add the response wrap token to the ENV of the plugin
	cmd.Env = append(cmd.Env, fmt.Sprintf("%s=%s", PluginUnwrapTokenEnv, wrapToken))
	// Add the mlock setting to the ENV of the plugin
	cmd.Env = append(cmd.Env, fmt.Sprintf("%s=%s", PluginMlockEnabled, mlock))

	secureConfig := &plugin.SecureConfig{
		Checksum: r.Sha256,
		Hash:     sha256.New(),
	}

	client := plugin.NewClient(&plugin.ClientConfig{
		HandshakeConfig: hs,
		Plugins:         pluginMap,
		Cmd:             cmd,
		TLSConfig:       clientTLSConfig,
		SecureConfig:    secureConfig,
	})

	return client, nil
}

// OptionallyEnableMlock determines if mlock should be called, and if so enables
// mlock.
func OptionallyEnableMlock() error {
	if os.Getenv(PluginMlockEnabled) == "true" {
		return mlock.LockMemory()
	}

	return nil
}
<|endoftext|>


<commit_before>// Copyright (c) 2013-2015 Conformal Systems LLC.
// Use of this source code is governed by an ISC
// license that can be found in the LICENSE file.

package wire

import (
	"bytes"
	"fmt"
	"io"
	"unicode/utf8"
)

// MessageHeaderSize is the number of bytes in a bitcoin message header.
// Bitcoin network (magic) 4 bytes + command 12 bytes + payload length 4 bytes +
// checksum 4 bytes.
const MessageHeaderSize = 24

// CommandSize is the fixed size of all commands in the common bitcoin message
// header.  Shorter commands must be zero padded.
const CommandSize = 12

// MaxMessagePayload is the maximum bytes a message can be regardless of other
// individual limits imposed by messages themselves.
const MaxMessagePayload = (1024 * 1024 * 32) // 32MB

// Commands used in bitcoin message headers which describe the type of message.
const (
	CmdVersion     = "version"
	CmdVerAck      = "verack"
	CmdGetAddr     = "getaddr"
	CmdAddr        = "addr"
	CmdGetBlocks   = "getblocks"
	CmdInv         = "inv"
	CmdGetData     = "getdata"
	CmdNotFound    = "notfound"
	CmdBlock       = "block" // Factoid Block
	CmdTx          = "tx"    // Factoid Tx
	CmdGetHeaders  = "getheaders"
	CmdHeaders     = "headers"
	CmdPing        = "ping"
	CmdPong        = "pong"
	CmdAlert       = "alert"
	CmdMemPool     = "mempool"
	CmdFilterAdd   = "filteradd"
	CmdFilterClear = "filterclear"
	CmdFilterLoad  = "filterload"
	CmdMerkleBlock = "merkleblock"
	CmdReject      = "reject"

	// Factom additions:

	// Factom downloads: block & entry
	CmdGetDirBlocks  = "getdirblocks"
	CmdDirInv        = "dirInv"
	CmdGetDirData    = "getdirdata"
	CmdDirBlock      = "dirblock"
	CmdGetNonDirData = "getnondirdata"
	CmdCBlock        = "ecblock"    // Entry Credit Block
	CmdEBlock        = "entryblock" // Entry Block
	CmdGetEntryData  = "getentrydata"
	CmdEntry         = "entry"

	CmdCommitChain = "commitchain"
	CmdRevealChain = "revealchain"
	CmdCommitEntry = "commitentry"
	CmdRevealEntry = "revealentry"

	CmdAcknowledgement = "confirmation"
	CmdMHashReveal     = "mhashreveal"
)

// MaxAppMsgPayload is the maximum bytes a factom app message can be in bytes.
const MaxAppMsgPayload = (10 * 1024) // 10Kib

// Message is an interface that describes a bitcoin message.  A type that
// implements Message has complete control over the representation of its data
// and may therefore contain additional or fewer fields than those which
// are used directly in the protocol encoded message.
type Message interface {
	BtcDecode(io.Reader, uint32) error
	BtcEncode(io.Writer, uint32) error
	Command() string
	MaxPayloadLength(uint32) uint32
}

// makeEmptyMessage creates a message of the appropriate concrete type based
// on the command.
func makeEmptyMessage(command string) (Message, error) {
	var msg Message
	switch command {
	case CmdVersion:
		msg = &MsgVersion{}

	case CmdVerAck:
		msg = &MsgVerAck{}

	case CmdGetAddr:
		msg = &MsgGetAddr{}

	case CmdAddr:
		msg = &MsgAddr{}

	case CmdGetBlocks:
		msg = &MsgGetBlocks{}

	case CmdBlock:
		msg = &MsgBlock{}

	case CmdInv:
		msg = &MsgInv{}

	case CmdGetData:
		msg = &MsgGetData{}

	case CmdNotFound:
		msg = &MsgNotFound{}

	case CmdTx:
		msg = &MsgTx{}

	case CmdPing:
		msg = &MsgPing{}

	case CmdPong:
		msg = &MsgPong{}

	case CmdGetHeaders:
		msg = &MsgGetHeaders{}

	case CmdHeaders:
		msg = &MsgHeaders{}

	case CmdAlert:
		msg = &MsgAlert{}

	case CmdMemPool:
		msg = &MsgMemPool{}

	case CmdFilterAdd:
		msg = &MsgFilterAdd{}

	case CmdFilterClear:
		msg = &MsgFilterClear{}

	case CmdFilterLoad:
		msg = &MsgFilterLoad{}

	case CmdMerkleBlock:
		msg = &MsgMerkleBlock{}

	case CmdReject:
		msg = &MsgReject{}

	case CmdRevealEntry:
		msg = &MsgRevealEntry{}

	case CmdCommitEntry:
		msg = &MsgCommitEntry{}

	case CmdRevealChain:
		msg = &MsgRevealChain{}

	case CmdCommitChain:
		msg = &MsgCommitChain{}

	case CmdAcknowledgement:
		msg = &MsgAcknowledgement{}

	case CmdGetDirBlocks:
		msg = &MsgGetDirBlocks{}

	case CmdDirInv:
		msg = &MsgDirInv{}

	case CmdGetDirData:
		msg = &MsgGetDirData{}

	case CmdCBlock:
		msg = &MsgCBlock{}

	case CmdEBlock:
		//msg = &MsgEBlock{}

	case CmdEntry:
		//msg = &MsgEntry{}

	default:
		return nil, fmt.Errorf("unhandled command [%s]", command)
	}
	return msg, nil
}

// messageHeader defines the header structure for all bitcoin protocol messages.
type messageHeader struct {
	magic    BitcoinNet // 4 bytes
	command  string     // 12 bytes
	length   uint32     // 4 bytes
	checksum [4]byte    // 4 bytes
}

// readMessageHeader reads a bitcoin message header from r.
func readMessageHeader(r io.Reader) (int, *messageHeader, error) {
	// Since readElements doesn't return the amount of bytes read, attempt
	// to read the entire header into a buffer first in case there is a
	// short read so the proper amount of read bytes are known.  This works
	// since the header is a fixed size.
	var headerBytes [MessageHeaderSize]byte
	n, err := io.ReadFull(r, headerBytes[:])
	if err != nil {
		return n, nil, err
	}
	hr := bytes.NewReader(headerBytes[:])

	// Create and populate a messageHeader struct from the raw header bytes.
	hdr := messageHeader{}
	var command [CommandSize]byte
	readElements(hr, &hdr.magic, &command, &hdr.length, &hdr.checksum)

	// Strip trailing zeros from command string.
	hdr.command = string(bytes.TrimRight(command[:], string(0)))

	return n, &hdr, nil
}

// discardInput reads n bytes from reader r in chunks and discards the read
// bytes.  This is used to skip payloads when various errors occur and helps
// prevent rogue nodes from causing massive memory allocation through forging
// header length.
func discardInput(r io.Reader, n uint32) {
	maxSize := uint32(10 * 1024) // 10k at a time
	numReads := n / maxSize
	bytesRemaining := n % maxSize
	if n > 0 {
		buf := make([]byte, maxSize)
		for i := uint32(0); i < numReads; i++ {
			io.ReadFull(r, buf)
		}
	}
	if bytesRemaining > 0 {
		buf := make([]byte, bytesRemaining)
		io.ReadFull(r, buf)
	}
}

// WriteMessageN writes a bitcoin Message to w including the necessary header
// information and returns the number of bytes written.    This function is the
// same as WriteMessage except it also returns the number of bytes written.
func WriteMessageN(w io.Writer, msg Message, pver uint32, btcnet BitcoinNet) (int, error) {
	totalBytes := 0

	// Enforce max command size.
	var command [CommandSize]byte
	cmd := msg.Command()
	if len(cmd) > CommandSize {
		str := fmt.Sprintf("command [%s] is too long [max %v]",
			cmd, CommandSize)
		return totalBytes, messageError("WriteMessage", str)
	}
	copy(command[:], []byte(cmd))

	// Encode the message payload.
	var bw bytes.Buffer
	err := msg.BtcEncode(&bw, pver)
	if err != nil {
		return totalBytes, err
	}
	payload := bw.Bytes()
	lenp := len(payload)

	// Enforce maximum overall message payload.
	if lenp > MaxMessagePayload {
		str := fmt.Sprintf("message payload is too large - encoded "+
			"%d bytes, but maximum message payload is %d bytes",
			lenp, MaxMessagePayload)
		return totalBytes, messageError("WriteMessage", str)
	}

	// Enforce maximum message payload based on the message type.
	mpl := msg.MaxPayloadLength(pver)
	if uint32(lenp) > mpl {
		str := fmt.Sprintf("message payload is too large - encoded "+
			"%d bytes, but maximum message payload size for "+
			"messages of type [%s] is %d.", lenp, cmd, mpl)
		return totalBytes, messageError("WriteMessage", str)
	}

	// Create header for the message.
	hdr := messageHeader{}
	hdr.magic = btcnet
	hdr.command = cmd
	hdr.length = uint32(lenp)
	copy(hdr.checksum[:], DoubleSha256(payload)[0:4])

	// Encode the header for the message.  This is done to a buffer
	// rather than directly to the writer since writeElements doesn't
	// return the number of bytes written.
	hw := bytes.NewBuffer(make([]byte, 0, MessageHeaderSize))
	writeElements(hw, hdr.magic, command, hdr.length, hdr.checksum)

	// Write header.
	n, err := w.Write(hw.Bytes())
	if err != nil {
		totalBytes += n
		return totalBytes, err
	}
	totalBytes += n

	// Write payload.
	n, err = w.Write(payload)
	if err != nil {
		totalBytes += n
		return totalBytes, err
	}
	totalBytes += n

	return totalBytes, nil
}

// WriteMessage writes a bitcoin Message to w including the necessary header
// information.  This function is the same as WriteMessageN except it doesn't
// doesn't return the number of bytes written.  This function is mainly provided
// for backwards compatibility with the original API, but it's also useful for
// callers that don't care about byte counts.
func WriteMessage(w io.Writer, msg Message, pver uint32, btcnet BitcoinNet) error {
	_, err := WriteMessageN(w, msg, pver, btcnet)
	return err
}

// ReadMessageN reads, validates, and parses the next bitcoin Message from r for
// the provided protocol version and bitcoin network.  It returns the number of
// bytes read in addition to the parsed Message and raw bytes which comprise the
// message.  This function is the same as ReadMessage except it also returns the
// number of bytes read.
func ReadMessageN(r io.Reader, pver uint32, btcnet BitcoinNet) (int, Message, []byte, error) {
	totalBytes := 0
	n, hdr, err := readMessageHeader(r)
	if err != nil {
		totalBytes += n
		return totalBytes, nil, nil, err
	}
	totalBytes += n

	// Enforce maximum message payload.
	if hdr.length > MaxMessagePayload {
		str := fmt.Sprintf("message payload is too large - header "+
			"indicates %d bytes, but max message payload is %d "+
			"bytes.", hdr.length, MaxMessagePayload)
		return totalBytes, nil, nil, messageError("ReadMessage", str)

	}

	// Check for messages from the wrong bitcoin network.
	if hdr.magic != btcnet {
		discardInput(r, hdr.length)
		str := fmt.Sprintf("message from other network [%v]", hdr.magic)
		return totalBytes, nil, nil, messageError("ReadMessage", str)
	}

	// Check for malformed commands.
	command := hdr.command
	if !utf8.ValidString(command) {
		discardInput(r, hdr.length)
		str := fmt.Sprintf("invalid command %v", []byte(command))
		return totalBytes, nil, nil, messageError("ReadMessage", str)
	}

	// Create struct of appropriate message type based on the command.
	msg, err := makeEmptyMessage(command)
	if err != nil {
		discardInput(r, hdr.length)
		return totalBytes, nil, nil, messageError("ReadMessage",
			err.Error())
	}

	// Check for maximum length based on the message type as a malicious client
	// could otherwise create a well-formed header and set the length to max
	// numbers in order to exhaust the machine's memory.
	mpl := msg.MaxPayloadLength(pver)
	if hdr.length > mpl {
		discardInput(r, hdr.length)
		str := fmt.Sprintf("payload exceeds max length - header "+
			"indicates %v bytes, but max payload size for "+
			"messages of type [%v] is %v.", hdr.length, command, mpl)
		return totalBytes, nil, nil, messageError("ReadMessage", str)
	}

	// Read payload.
	payload := make([]byte, hdr.length)
	n, err = io.ReadFull(r, payload)
	if err != nil {
		totalBytes += n
		return totalBytes, nil, nil, err
	}
	totalBytes += n

	// Test checksum.
	checksum := DoubleSha256(payload)[0:4]
	if !bytes.Equal(checksum[:], hdr.checksum[:]) {
		str := fmt.Sprintf("payload checksum failed - header "+
			"indicates %v, but actual checksum is %v.",
			hdr.checksum, checksum)
		return totalBytes, nil, nil, messageError("ReadMessage", str)
	}

	// Unmarshal message.  NOTE: This must be a *bytes.Buffer since the
	// MsgVersion BtcDecode function requires it.
	pr := bytes.NewBuffer(payload)
	err = msg.BtcDecode(pr, pver)
	if err != nil {
		return totalBytes, nil, nil, err
	}

	return totalBytes, msg, payload, nil
}

// ReadMessage reads, validates, and parses the next bitcoin Message from r for
// the provided protocol version and bitcoin network.  It returns the parsed
// Message and raw bytes which comprise the message.  This function only differs
// from ReadMessageN in that it doesn't return the number of bytes read.  This
// function is mainly provided for backwards compatibility with the original
// API, but it's also useful for callers that don't care about byte counts.
func ReadMessage(r io.Reader, pver uint32, btcnet BitcoinNet) (Message, []byte, error) {
	_, msg, buf, err := ReadMessageN(r, pver, btcnet)
	return msg, buf, err
}
<commit_msg>fix bug<commit_after>// Copyright (c) 2013-2015 Conformal Systems LLC.
// Use of this source code is governed by an ISC
// license that can be found in the LICENSE file.

package wire

import (
	"bytes"
	"fmt"
	"io"
	"unicode/utf8"
)

// MessageHeaderSize is the number of bytes in a bitcoin message header.
// Bitcoin network (magic) 4 bytes + command 12 bytes + payload length 4 bytes +
// checksum 4 bytes.
const MessageHeaderSize = 24

// CommandSize is the fixed size of all commands in the common bitcoin message
// header.  Shorter commands must be zero padded.
const CommandSize = 12

// MaxMessagePayload is the maximum bytes a message can be regardless of other
// individual limits imposed by messages themselves.
const MaxMessagePayload = (1024 * 1024 * 32) // 32MB

// Commands used in bitcoin message headers which describe the type of message.
const (
	CmdVersion     = "version"
	CmdVerAck      = "verack"
	CmdGetAddr     = "getaddr"
	CmdAddr        = "addr"
	CmdGetBlocks   = "getblocks"
	CmdInv         = "inv"
	CmdGetData     = "getdata"
	CmdNotFound    = "notfound"
	CmdBlock       = "block" // Factoid Block
	CmdTx          = "tx"    // Factoid Tx
	CmdGetHeaders  = "getheaders"
	CmdHeaders     = "headers"
	CmdPing        = "ping"
	CmdPong        = "pong"
	CmdAlert       = "alert"
	CmdMemPool     = "mempool"
	CmdFilterAdd   = "filteradd"
	CmdFilterClear = "filterclear"
	CmdFilterLoad  = "filterload"
	CmdMerkleBlock = "merkleblock"
	CmdReject      = "reject"

	// Factom additions:

	// Factom downloads: block & entry
	CmdGetDirBlocks  = "getdirblocks"
	CmdDirInv        = "dirInv"
	CmdGetDirData    = "getdirdata"
	CmdDirBlock      = "dirblock"
	CmdGetNonDirData = "getnondirdata"
	CmdCBlock        = "ecblock"    // Entry Credit Block
	CmdEBlock        = "entryblock" // Entry Block
	CmdGetEntryData  = "getentrydata"
	CmdEntry         = "entry"

	CmdCommitChain = "commitchain"
	CmdRevealChain = "revealchain"
	CmdCommitEntry = "commitentry"
	CmdRevealEntry = "revealentry"

	CmdAcknowledgement = "confirmation"
	CmdMHashReveal     = "mhashreveal"
)

// MaxAppMsgPayload is the maximum bytes a factom app message can be in bytes.
const MaxAppMsgPayload = (10 * 1024) // 10Kib

// Message is an interface that describes a bitcoin message.  A type that
// implements Message has complete control over the representation of its data
// and may therefore contain additional or fewer fields than those which
// are used directly in the protocol encoded message.
type Message interface {
	BtcDecode(io.Reader, uint32) error
	BtcEncode(io.Writer, uint32) error
	Command() string
	MaxPayloadLength(uint32) uint32
}

// makeEmptyMessage creates a message of the appropriate concrete type based
// on the command.
func makeEmptyMessage(command string) (Message, error) {
	var msg Message
	switch command {
	case CmdVersion:
		msg = &MsgVersion{}

	case CmdVerAck:
		msg = &MsgVerAck{}

	case CmdGetAddr:
		msg = &MsgGetAddr{}

	case CmdAddr:
		msg = &MsgAddr{}

	case CmdGetBlocks:
		msg = &MsgGetBlocks{}

	case CmdBlock:
		msg = &MsgBlock{}

	case CmdInv:
		msg = &MsgInv{}

	case CmdGetData:
		msg = &MsgGetData{}

	case CmdNotFound:
		msg = &MsgNotFound{}

	case CmdTx:
		msg = &MsgTx{}

	case CmdPing:
		msg = &MsgPing{}

	case CmdPong:
		msg = &MsgPong{}

	case CmdGetHeaders:
		msg = &MsgGetHeaders{}

	case CmdHeaders:
		msg = &MsgHeaders{}

	case CmdAlert:
		msg = &MsgAlert{}

	case CmdMemPool:
		msg = &MsgMemPool{}

	case CmdFilterAdd:
		msg = &MsgFilterAdd{}

	case CmdFilterClear:
		msg = &MsgFilterClear{}

	case CmdFilterLoad:
		msg = &MsgFilterLoad{}

	case CmdMerkleBlock:
		msg = &MsgMerkleBlock{}

	case CmdReject:
		msg = &MsgReject{}

	case CmdRevealEntry:
		msg = &MsgRevealEntry{}

	case CmdCommitEntry:
		msg = &MsgCommitEntry{}

	case CmdRevealChain:
		msg = &MsgRevealChain{}

	case CmdCommitChain:
		msg = &MsgCommitChain{}

	case CmdAcknowledgement:
		msg = &MsgAcknowledgement{}

	case CmdGetDirBlocks:
		msg = &MsgGetDirBlocks{}

	case CmdDirInv:
		msg = &MsgDirInv{}

	case CmdGetDirData:
		msg = &MsgGetDirData{}

	case CmdDirBlock:
		msg = &MsgDirBlock{}

	case CmdCBlock:
		msg = &MsgCBlock{}

	case CmdEBlock:
		msg = &MsgEBlock{}

	case CmdEntry:
		msg = &MsgEntry{}

	default:
		return nil, fmt.Errorf("unhandled command [%s]", command)
	}
	return msg, nil
}

// messageHeader defines the header structure for all bitcoin protocol messages.
type messageHeader struct {
	magic    BitcoinNet // 4 bytes
	command  string     // 12 bytes
	length   uint32     // 4 bytes
	checksum [4]byte    // 4 bytes
}

// readMessageHeader reads a bitcoin message header from r.
func readMessageHeader(r io.Reader) (int, *messageHeader, error) {
	// Since readElements doesn't return the amount of bytes read, attempt
	// to read the entire header into a buffer first in case there is a
	// short read so the proper amount of read bytes are known.  This works
	// since the header is a fixed size.
	var headerBytes [MessageHeaderSize]byte
	n, err := io.ReadFull(r, headerBytes[:])
	if err != nil {
		return n, nil, err
	}
	hr := bytes.NewReader(headerBytes[:])

	// Create and populate a messageHeader struct from the raw header bytes.
	hdr := messageHeader{}
	var command [CommandSize]byte
	readElements(hr, &hdr.magic, &command, &hdr.length, &hdr.checksum)

	// Strip trailing zeros from command string.
	hdr.command = string(bytes.TrimRight(command[:], string(0)))

	return n, &hdr, nil
}

// discardInput reads n bytes from reader r in chunks and discards the read
// bytes.  This is used to skip payloads when various errors occur and helps
// prevent rogue nodes from causing massive memory allocation through forging
// header length.
func discardInput(r io.Reader, n uint32) {
	maxSize := uint32(10 * 1024) // 10k at a time
	numReads := n / maxSize
	bytesRemaining := n % maxSize
	if n > 0 {
		buf := make([]byte, maxSize)
		for i := uint32(0); i < numReads; i++ {
			io.ReadFull(r, buf)
		}
	}
	if bytesRemaining > 0 {
		buf := make([]byte, bytesRemaining)
		io.ReadFull(r, buf)
	}
}

// WriteMessageN writes a bitcoin Message to w including the necessary header
// information and returns the number of bytes written.    This function is the
// same as WriteMessage except it also returns the number of bytes written.
func WriteMessageN(w io.Writer, msg Message, pver uint32, btcnet BitcoinNet) (int, error) {
	totalBytes := 0

	// Enforce max command size.
	var command [CommandSize]byte
	cmd := msg.Command()
	if len(cmd) > CommandSize {
		str := fmt.Sprintf("command [%s] is too long [max %v]",
			cmd, CommandSize)
		return totalBytes, messageError("WriteMessage", str)
	}
	copy(command[:], []byte(cmd))

	// Encode the message payload.
	var bw bytes.Buffer
	err := msg.BtcEncode(&bw, pver)
	if err != nil {
		return totalBytes, err
	}
	payload := bw.Bytes()
	lenp := len(payload)

	// Enforce maximum overall message payload.
	if lenp > MaxMessagePayload {
		str := fmt.Sprintf("message payload is too large - encoded "+
			"%d bytes, but maximum message payload is %d bytes",
			lenp, MaxMessagePayload)
		return totalBytes, messageError("WriteMessage", str)
	}

	// Enforce maximum message payload based on the message type.
	mpl := msg.MaxPayloadLength(pver)
	if uint32(lenp) > mpl {
		str := fmt.Sprintf("message payload is too large - encoded "+
			"%d bytes, but maximum message payload size for "+
			"messages of type [%s] is %d.", lenp, cmd, mpl)
		return totalBytes, messageError("WriteMessage", str)
	}

	// Create header for the message.
	hdr := messageHeader{}
	hdr.magic = btcnet
	hdr.command = cmd
	hdr.length = uint32(lenp)
	copy(hdr.checksum[:], DoubleSha256(payload)[0:4])

	// Encode the header for the message.  This is done to a buffer
	// rather than directly to the writer since writeElements doesn't
	// return the number of bytes written.
	hw := bytes.NewBuffer(make([]byte, 0, MessageHeaderSize))
	writeElements(hw, hdr.magic, command, hdr.length, hdr.checksum)

	// Write header.
	n, err := w.Write(hw.Bytes())
	if err != nil {
		totalBytes += n
		return totalBytes, err
	}
	totalBytes += n

	// Write payload.
	n, err = w.Write(payload)
	if err != nil {
		totalBytes += n
		return totalBytes, err
	}
	totalBytes += n

	return totalBytes, nil
}

// WriteMessage writes a bitcoin Message to w including the necessary header
// information.  This function is the same as WriteMessageN except it doesn't
// doesn't return the number of bytes written.  This function is mainly provided
// for backwards compatibility with the original API, but it's also useful for
// callers that don't care about byte counts.
func WriteMessage(w io.Writer, msg Message, pver uint32, btcnet BitcoinNet) error {
	_, err := WriteMessageN(w, msg, pver, btcnet)
	return err
}

// ReadMessageN reads, validates, and parses the next bitcoin Message from r for
// the provided protocol version and bitcoin network.  It returns the number of
// bytes read in addition to the parsed Message and raw bytes which comprise the
// message.  This function is the same as ReadMessage except it also returns the
// number of bytes read.
func ReadMessageN(r io.Reader, pver uint32, btcnet BitcoinNet) (int, Message, []byte, error) {
	totalBytes := 0
	n, hdr, err := readMessageHeader(r)
	if err != nil {
		totalBytes += n
		return totalBytes, nil, nil, err
	}
	totalBytes += n

	// Enforce maximum message payload.
	if hdr.length > MaxMessagePayload {
		str := fmt.Sprintf("message payload is too large - header "+
			"indicates %d bytes, but max message payload is %d "+
			"bytes.", hdr.length, MaxMessagePayload)
		return totalBytes, nil, nil, messageError("ReadMessage", str)

	}

	// Check for messages from the wrong bitcoin network.
	if hdr.magic != btcnet {
		discardInput(r, hdr.length)
		str := fmt.Sprintf("message from other network [%v]", hdr.magic)
		return totalBytes, nil, nil, messageError("ReadMessage", str)
	}

	// Check for malformed commands.
	command := hdr.command
	if !utf8.ValidString(command) {
		discardInput(r, hdr.length)
		str := fmt.Sprintf("invalid command %v", []byte(command))
		return totalBytes, nil, nil, messageError("ReadMessage", str)
	}

	// Create struct of appropriate message type based on the command.
	msg, err := makeEmptyMessage(command)
	if err != nil {
		discardInput(r, hdr.length)
		return totalBytes, nil, nil, messageError("ReadMessage",
			err.Error())
	}

	// Check for maximum length based on the message type as a malicious client
	// could otherwise create a well-formed header and set the length to max
	// numbers in order to exhaust the machine's memory.
	mpl := msg.MaxPayloadLength(pver)
	if hdr.length > mpl {
		discardInput(r, hdr.length)
		str := fmt.Sprintf("payload exceeds max length - header "+
			"indicates %v bytes, but max payload size for "+
			"messages of type [%v] is %v.", hdr.length, command, mpl)
		return totalBytes, nil, nil, messageError("ReadMessage", str)
	}

	// Read payload.
	payload := make([]byte, hdr.length)
	n, err = io.ReadFull(r, payload)
	if err != nil {
		totalBytes += n
		return totalBytes, nil, nil, err
	}
	totalBytes += n

	// Test checksum.
	checksum := DoubleSha256(payload)[0:4]
	if !bytes.Equal(checksum[:], hdr.checksum[:]) {
		str := fmt.Sprintf("payload checksum failed - header "+
			"indicates %v, but actual checksum is %v.",
			hdr.checksum, checksum)
		return totalBytes, nil, nil, messageError("ReadMessage", str)
	}

	// Unmarshal message.  NOTE: This must be a *bytes.Buffer since the
	// MsgVersion BtcDecode function requires it.
	pr := bytes.NewBuffer(payload)
	err = msg.BtcDecode(pr, pver)
	if err != nil {
		return totalBytes, nil, nil, err
	}

	return totalBytes, msg, payload, nil
}

// ReadMessage reads, validates, and parses the next bitcoin Message from r for
// the provided protocol version and bitcoin network.  It returns the parsed
// Message and raw bytes which comprise the message.  This function only differs
// from ReadMessageN in that it doesn't return the number of bytes read.  This
// function is mainly provided for backwards compatibility with the original
// API, but it's also useful for callers that don't care about byte counts.
func ReadMessage(r io.Reader, pver uint32, btcnet BitcoinNet) (Message, []byte, error) {
	_, msg, buf, err := ReadMessageN(r, pver, btcnet)
	return msg, buf, err
}
<|endoftext|>


<commit_before>package update

import (
	"archive/tar"
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path"
	"sync"

	"github.com/fsouza/go-dockerclient"
)

type jsonstreamMessage struct {
	Status string `json:"status"`
	ID     string `json:"id"`
	Error  string `json:"error"`
}

// jsonstreamErrorDetector is a io.Writer wrapper whose purpose is to err out
// if the stream written to it will contain a Docker image
type jsonstreamErrorDetector struct {
	buffer *bytes.Buffer
}

func (jsed *jsonstreamErrorDetector) Write(p []byte) (n int, err error) {
	if jsed.buffer == nil {
		jsed.buffer = bytes.NewBuffer([]byte{})
	}

	n, err = jsed.buffer.Write(p)
	if err != nil {
		return 0, err
	}

	var msg jsonstreamMessage
	dec := json.NewDecoder(jsed.buffer)

	for {
		err = dec.Decode(&msg)
		if err != nil {
			if err == io.EOF {
				break
			}
			return 0, err
		}
	}

	if msg.Error != "" {
		return 0, fmt.Errorf("Docker error: %s", msg.Error)
	}

	return n, nil
}

func pullImage(repository, tag string, authCfg io.Reader) error {
	var jsed jsonstreamErrorDetector

	opts := docker.PullImageOptions{
		Repository:    repository,
		Tag:           tag,
		OutputStream:  &jsed,
		RawJSONStream: true,
	}

	client, err := docker.NewClientFromEnv()
	if err != nil {
		return err
	}

	var auth *docker.AuthConfigurations
	if authCfg == nil {
		auth, err = docker.NewAuthConfigurationsFromDockerCfg()
	} else {
		auth, err = docker.NewAuthConfigurations(authCfg)
	}
	if err != nil {
		return err
	}

	err = client.PullImage(opts, auth.Configs["quay.io"])
	if err != nil {
		return err
	}

	return nil
}

// exportDockerImage writes a TAR archive of a given Docker image's rootfs
// into the output writer.
//
// If returned error is not nil then the final state of the output writer
// is not defined, i.e. the data might have been written partially.
func exportDockerImage(repository, tag string, output io.Writer) error {
	c, err := docker.NewClientFromEnv()
	if err != nil {
		return err
	}

	createContainerOpts := docker.CreateContainerOptions{
		Config: &docker.Config{
			Image: fmt.Sprintf("%s:%s", repository, tag),
		},
	}

	containter, err := c.CreateContainer(createContainerOpts)
	if err != nil {
		return err
	}
	defer c.RemoveContainer(docker.RemoveContainerOptions{ID: containter.ID, RemoveVolumes: true, Force: true})

	exportContainerOptions := docker.ExportContainerOptions{
		ID:           containter.ID,
		OutputStream: output,
	}

	return c.ExportContainer(exportContainerOptions)
}

// extractDockerImage writes a given image's rootfs to the target folder,
// ignoring everything that isn't a regular file or directory
func extractDockerImage(repository, tag, extractDir string) error {
	// this pipe connects exportDockerImage to the tar reader
	pipeReader, pipeWriter := io.Pipe()

	// dump the TAR-ed image to one end of the pipe
	var wg sync.WaitGroup
	wg.Add(1)
	var extractErr error
	go func() {
		extractErr = exportDockerImage(repository, tag, pipeWriter)
		// need to close this end or the other end will block
		pipeWriter.Close()
		wg.Done()
	}()

	// wrap a TAR reader around the other and extract
	tarReader := tar.NewReader(pipeReader)
	for {
		header, nextErr := tarReader.Next()
		if nextErr != nil {
			if nextErr == io.EOF {
				break
			}
			return nextErr
		}

		switch header.Typeflag {
		case tar.TypeDir:
			mkdirErr := os.MkdirAll(path.Join(extractDir, header.Name), 0755)
			if mkdirErr != nil {
				return mkdirErr
			}
			io.CopyN(ioutil.Discard, tarReader, header.Size)
		case tar.TypeReg:
			f, err := os.OpenFile(path.Join(extractDir, header.Name), os.O_CREATE|os.O_WRONLY, 0644)
			if err != nil {
				return err
			}
			_, err = io.CopyN(f, tarReader, header.Size)
			if err != nil {
				return err
			}
			f.Close()
		case tar.TypeSymlink:
			linkErr := os.Symlink(header.Linkname, path.Join(extractDir, header.Name))
			if linkErr != nil {
				return linkErr
			}
		default:
			// ignore
			io.CopyN(ioutil.Discard, tarReader, header.Size)
		}
	}

	wg.Wait()

	if extractErr != nil {
		return extractErr
	}

	return nil
}
<commit_msg>hardcode the path to the Docker config file<commit_after>package update

import (
	"archive/tar"
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path"
	"sync"

	"github.com/fsouza/go-dockerclient"
)

type jsonstreamMessage struct {
	Status string `json:"status"`
	ID     string `json:"id"`
	Error  string `json:"error"`
}

// jsonstreamErrorDetector is a io.Writer wrapper whose purpose is to err out
// if the stream written to it will contain a Docker image
type jsonstreamErrorDetector struct {
	buffer *bytes.Buffer
}

func (jsed *jsonstreamErrorDetector) Write(p []byte) (n int, err error) {
	if jsed.buffer == nil {
		jsed.buffer = bytes.NewBuffer([]byte{})
	}

	n, err = jsed.buffer.Write(p)
	if err != nil {
		return 0, err
	}

	var msg jsonstreamMessage
	dec := json.NewDecoder(jsed.buffer)

	for {
		err = dec.Decode(&msg)
		if err != nil {
			if err == io.EOF {
				break
			}
			return 0, err
		}
	}

	if msg.Error != "" {
		return 0, fmt.Errorf("Docker error: %s", msg.Error)
	}

	return n, nil
}

func pullImage(repository, tag string, authCfg io.Reader) error {
	var jsed jsonstreamErrorDetector

	opts := docker.PullImageOptions{
		Repository:    repository,
		Tag:           tag,
		OutputStream:  &jsed,
		RawJSONStream: true,
	}

	client, err := docker.NewClientFromEnv()
	if err != nil {
		return err
	}

	var auth *docker.AuthConfigurations
	if authCfg == nil {
		// the full path to the config file is provided,
		// since running from a systemd unit doesn't provide $HOME
		auth, err = docker.NewAuthConfigurationsFromFile("/root/.docker/config.json")
	} else {
		auth, err = docker.NewAuthConfigurations(authCfg)
	}
	if err != nil {
		return err
	}

	err = client.PullImage(opts, auth.Configs["quay.io"])
	if err != nil {
		return err
	}

	return nil
}

// exportDockerImage writes a TAR archive of a given Docker image's rootfs
// into the output writer.
//
// If returned error is not nil then the final state of the output writer
// is not defined, i.e. the data might have been written partially.
func exportDockerImage(repository, tag string, output io.Writer) error {
	c, err := docker.NewClientFromEnv()
	if err != nil {
		return err
	}

	createContainerOpts := docker.CreateContainerOptions{
		Config: &docker.Config{
			Image: fmt.Sprintf("%s:%s", repository, tag),
		},
	}

	containter, err := c.CreateContainer(createContainerOpts)
	if err != nil {
		return err
	}
	defer c.RemoveContainer(docker.RemoveContainerOptions{ID: containter.ID, RemoveVolumes: true, Force: true})

	exportContainerOptions := docker.ExportContainerOptions{
		ID:           containter.ID,
		OutputStream: output,
	}

	return c.ExportContainer(exportContainerOptions)
}

// extractDockerImage writes a given image's rootfs to the target folder,
// ignoring everything that isn't a regular file or directory
func extractDockerImage(repository, tag, extractDir string) error {
	// this pipe connects exportDockerImage to the tar reader
	pipeReader, pipeWriter := io.Pipe()

	// dump the TAR-ed image to one end of the pipe
	var wg sync.WaitGroup
	wg.Add(1)
	var extractErr error
	go func() {
		extractErr = exportDockerImage(repository, tag, pipeWriter)
		// need to close this end or the other end will block
		pipeWriter.Close()
		wg.Done()
	}()

	// wrap a TAR reader around the other and extract
	tarReader := tar.NewReader(pipeReader)
	for {
		header, nextErr := tarReader.Next()
		if nextErr != nil {
			if nextErr == io.EOF {
				break
			}
			return nextErr
		}

		switch header.Typeflag {
		case tar.TypeDir:
			mkdirErr := os.MkdirAll(path.Join(extractDir, header.Name), 0755)
			if mkdirErr != nil {
				return mkdirErr
			}
			io.CopyN(ioutil.Discard, tarReader, header.Size)
		case tar.TypeReg:
			f, err := os.OpenFile(path.Join(extractDir, header.Name), os.O_CREATE|os.O_WRONLY, 0644)
			if err != nil {
				return err
			}
			_, err = io.CopyN(f, tarReader, header.Size)
			if err != nil {
				return err
			}
			f.Close()
		case tar.TypeSymlink:
			linkErr := os.Symlink(header.Linkname, path.Join(extractDir, header.Name))
			if linkErr != nil {
				return linkErr
			}
		default:
			// ignore
			io.CopyN(ioutil.Discard, tarReader, header.Size)
		}
	}

	wg.Wait()

	if extractErr != nil {
		return extractErr
	}

	return nil
}
<|endoftext|>


<commit_before>// +build !arm,!ppc64,!ppc64le

package netlink

func ifrDataByte(b byte) int8 {
	return int8(b)
}
<commit_msg>Update netlink_linux_notarm.go<commit_after>// +build !arm,!ppc64,!ppc64le,!s390x

package netlink

func ifrDataByte(b byte) int8 {
	return int8(b)
}
<|endoftext|>