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

package vtgate

import (
	"fmt"
	"reflect"
	"testing"
	"time"

	mproto "github.com/youtube/vitess/go/mysql/proto"
	tproto "github.com/youtube/vitess/go/vt/tabletserver/proto"
	"github.com/youtube/vitess/go/vt/vtgate/proto"
)

// This file uses the sandbox_test framework.

func TestScatterConnExecute(t *testing.T) {
	testScatterConnGeneric(t, func(shards []string) (*mproto.QueryResult, error) {
		stc := NewScatterConn(new(sandboxTopo), "aa", 1*time.Millisecond, 3)
		return stc.Execute("query", nil, "", shards, "", nil)
	})
}

func TestScatterConnExecuteBatch(t *testing.T) {
	testScatterConnGeneric(t, func(shards []string) (*mproto.QueryResult, error) {
		stc := NewScatterConn(new(sandboxTopo), "aa", 1*time.Millisecond, 3)
		queries := []tproto.BoundQuery{{"query", nil}}
		qrs, err := stc.ExecuteBatch(queries, "", shards, "", nil)
		if err != nil {
			return nil, err
		}
		return &qrs.List[0], err
	})
}

func TestScatterConnStreamExecute(t *testing.T) {
	testScatterConnGeneric(t, func(shards []string) (*mproto.QueryResult, error) {
		stc := NewScatterConn(new(sandboxTopo), "aa", 1*time.Millisecond, 3)
		qr := new(mproto.QueryResult)
		err := stc.StreamExecute("query", nil, "", shards, "", nil, func(r interface{}) error {
			appendResult(qr, r.(*mproto.QueryResult))
			return nil
		})
		return qr, err
	})
}

func testScatterConnGeneric(t *testing.T, f func(shards []string) (*mproto.QueryResult, error)) {
	// no shard
	resetSandbox()
	qr, err := f(nil)
	if qr.RowsAffected != 0 {
		t.Errorf("want 0, got %v", qr.RowsAffected)
	}
	if err != nil {
		t.Errorf("want nil, got %v", err)
	}

	// single shard
	resetSandbox()
	sbc := &sandboxConn{mustFailServer: 1}
	testConns[0] = sbc
	qr, err = f([]string{"0"})
	want := "error: err, shard: (.0.), host: 0"
	// Verify server error string.
	if err == nil || err.Error() != want {
		t.Errorf("want %s, got %v", want, err)
	}
	// Ensure that we tried only once.
	if sbc.ExecCount != 1 {
		t.Errorf("want 1, got %v", sbc.ExecCount)
	}

	// two shards
	resetSandbox()
	sbc0 := &sandboxConn{mustFailServer: 1}
	testConns[0] = sbc0
	sbc1 := &sandboxConn{mustFailServer: 1}
	testConns[1] = sbc1
	_, err = f([]string{"0", "1"})
	// Verify server errors are consolidated.
	want = "error: err, shard: (.0.), host: 0\nerror: err, shard: (.1.), host: 1"
	if err == nil || err.Error() != want {
		t.Errorf("\nwant\n%s\ngot\n%v", want, err)
	}
	// Ensure that we tried only once.
	if sbc0.ExecCount != 1 {
		t.Errorf("want 1, got %v", sbc0.ExecCount)
	}
	if sbc1.ExecCount != 1 {
		t.Errorf("want 1, got %v", sbc1.ExecCount)
	}

	// duplicate shards
	resetSandbox()
	sbc = &sandboxConn{}
	testConns[0] = sbc
	qr, err = f([]string{"0", "0"})
	// Ensure that we executed only once.
	if sbc.ExecCount != 1 {
		t.Errorf("want 1, got %v", sbc.ExecCount)
	}

	// no errors
	resetSandbox()
	sbc0 = &sandboxConn{}
	testConns[0] = sbc0
	sbc1 = &sandboxConn{}
	testConns[1] = sbc1
	qr, err = f([]string{"0", "1"})
	if err != nil {
		t.Errorf("want nil, got %v", err)
	}
	if sbc0.ExecCount != 1 {
		t.Errorf("want 1, got %v", sbc0.ExecCount)
	}
	if sbc1.ExecCount != 1 {
		t.Errorf("want 1, got %v", sbc1.ExecCount)
	}
	if qr.RowsAffected != 2 {
		t.Errorf("want 2, got %v", qr.RowsAffected)
	}
	if len(qr.Rows) != 2 {
		t.Errorf("want 2, got %v", len(qr.Rows))
	}
}

func TestScatterConnStreamExecuteSendError(t *testing.T) {
	resetSandbox()
	sbc := &sandboxConn{}
	testConns[0] = sbc
	stc := NewScatterConn(new(sandboxTopo), "aa", 1*time.Millisecond, 3)
	err := stc.StreamExecute("query", nil, "", []string{"0"}, "", nil, func(interface{}) error {
		return fmt.Errorf("send error")
	})
	want := "send error"
	// Ensure that we handle send errors.
	if err == nil || err.Error() != want {
		t.Errorf("want %s, got %v", want, err)
	}
}

func TestScatterConnCommitSuccess(t *testing.T) {
	resetSandbox()
	sbc0 := &sandboxConn{}
	testConns[0] = sbc0
	sbc1 := &sandboxConn{mustFailTxPool: 1}
	testConns[1] = sbc1
	stc := NewScatterConn(new(sandboxTopo), "aa", 1*time.Millisecond, 3)

	// Sequence the executes to ensure commit order
	session := NewSafeSession(&proto.Session{InTransaction: true})
	stc.Execute("query1", nil, "", []string{"0"}, "", session)
	wantSession := proto.Session{
		InTransaction: true,
		ShardSessions: []*proto.ShardSession{{
			Keyspace:      "",
			Shard:         "0",
			TabletType:    "",
			TransactionId: 1,
		}},
	}
	if !reflect.DeepEqual(wantSession, *session.Session) {
		t.Errorf("want\n%#v, got\n%#v", wantSession, *session.Session)
	}
	stc.Execute("query1", nil, "", []string{"0", "1"}, "", session)
	wantSession = proto.Session{
		InTransaction: true,
		ShardSessions: []*proto.ShardSession{{
			Keyspace:      "",
			Shard:         "0",
			TabletType:    "",
			TransactionId: 1,
		}, {
			Keyspace:      "",
			Shard:         "1",
			TabletType:    "",
			TransactionId: 2,
		}},
	}
	if !reflect.DeepEqual(wantSession, *session.Session) {
		t.Errorf("want\n%#v, got\n%#v", wantSession, *session.Session)
	}
	sbc0.mustFailServer = 1
	err := stc.Commit(session)
	if err == nil {
		t.Errorf("want error, got nil")
	}
	wantSession = proto.Session{}
	if !reflect.DeepEqual(wantSession, *session.Session) {
		t.Errorf("want\n%#v, got\n%#v", wantSession, *session.Session)
	}
	if sbc0.CommitCount != 1 {
		t.Errorf("want 1, got %d", sbc0.CommitCount)
	}
	/*
		// Flaky: This test should be run manually.
		runtime.Gosched()
		if sbc1.RollbackCount != 1 {
			t.Errorf("want 1, got %d", sbc1.RollbackCount)
		}
	*/
}

func TestScatterConnRollback(t *testing.T) {
	resetSandbox()
	sbc0 := &sandboxConn{}
	testConns[0] = sbc0
	sbc1 := &sandboxConn{mustFailTxPool: 1}
	testConns[1] = sbc1
	stc := NewScatterConn(new(sandboxTopo), "aa", 1*time.Millisecond, 3)

	// Sequence the executes to ensure commit order
	session := NewSafeSession(&proto.Session{InTransaction: true})
	stc.Execute("query1", nil, "", []string{"0"}, "", session)
	stc.Execute("query1", nil, "", []string{"0", "1"}, "", session)
	err := stc.Rollback(session)
	if err != nil {
		t.Errorf("want nil, got %v", err)
	}
	wantSession := proto.Session{}
	if !reflect.DeepEqual(wantSession, *session.Session) {
		t.Errorf("want\n%#v, got\n%#v", wantSession, *session.Session)
	}
	/*
		// Flaky: This test should be run manually.
		runtime.Gosched()
		if sbc0.RollbackCount != 1 {
			t.Errorf("want 1, got %d", sbc0.RollbackCount)
		}
		if sbc1.RollbackCount != 1 {
			t.Errorf("want 1, got %d", sbc1.RollbackCount)
		}
	*/
}

func TestScatterConnClose(t *testing.T) {
	resetSandbox()
	sbc := &sandboxConn{}
	testConns[0] = sbc
	stc := NewScatterConn(new(sandboxTopo), "aa", 1*time.Millisecond, 3)
	stc.Execute("query1", nil, "", []string{"0"}, "", nil)
	stc.Close()
	/*
		// Flaky: This test should be run manually.
		runtime.Gosched()
		if sbc.CloseCount != 1 {
			t.Errorf("want 1, got %d", sbc.CommitCount)
		}
	*/
}
<commit_msg>Fix typo<commit_after>// Copyright 2012, Google Inc. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package vtgate

import (
	"fmt"
	"reflect"
	"testing"
	"time"

	mproto "github.com/youtube/vitess/go/mysql/proto"
	tproto "github.com/youtube/vitess/go/vt/tabletserver/proto"
	"github.com/youtube/vitess/go/vt/vtgate/proto"
)

// This file uses the sandbox_test framework.

func TestScatterConnExecute(t *testing.T) {
	testScatterConnGeneric(t, func(shards []string) (*mproto.QueryResult, error) {
		stc := NewScatterConn(new(sandboxTopo), "aa", 1*time.Millisecond, 3)
		return stc.Execute("query", nil, "", shards, "", nil)
	})
}

func TestScatterConnExecuteBatch(t *testing.T) {
	testScatterConnGeneric(t, func(shards []string) (*mproto.QueryResult, error) {
		stc := NewScatterConn(new(sandboxTopo), "aa", 1*time.Millisecond, 3)
		queries := []tproto.BoundQuery{{"query", nil}}
		qrs, err := stc.ExecuteBatch(queries, "", shards, "", nil)
		if err != nil {
			return nil, err
		}
		return &qrs.List[0], err
	})
}

func TestScatterConnStreamExecute(t *testing.T) {
	testScatterConnGeneric(t, func(shards []string) (*mproto.QueryResult, error) {
		stc := NewScatterConn(new(sandboxTopo), "aa", 1*time.Millisecond, 3)
		qr := new(mproto.QueryResult)
		err := stc.StreamExecute("query", nil, "", shards, "", nil, func(r interface{}) error {
			appendResult(qr, r.(*mproto.QueryResult))
			return nil
		})
		return qr, err
	})
}

func testScatterConnGeneric(t *testing.T, f func(shards []string) (*mproto.QueryResult, error)) {
	// no shard
	resetSandbox()
	qr, err := f(nil)
	if qr.RowsAffected != 0 {
		t.Errorf("want 0, got %v", qr.RowsAffected)
	}
	if err != nil {
		t.Errorf("want nil, got %v", err)
	}

	// single shard
	resetSandbox()
	sbc := &sandboxConn{mustFailServer: 1}
	testConns[0] = sbc
	qr, err = f([]string{"0"})
	want := "error: err, shard: (.0.), host: 0"
	// Verify server error string.
	if err == nil || err.Error() != want {
		t.Errorf("want %s, got %v", want, err)
	}
	// Ensure that we tried only once.
	if sbc.ExecCount != 1 {
		t.Errorf("want 1, got %v", sbc.ExecCount)
	}

	// two shards
	resetSandbox()
	sbc0 := &sandboxConn{mustFailServer: 1}
	testConns[0] = sbc0
	sbc1 := &sandboxConn{mustFailServer: 1}
	testConns[1] = sbc1
	_, err = f([]string{"0", "1"})
	// Verify server errors are consolidated.
	want = "error: err, shard: (.0.), host: 0\nerror: err, shard: (.1.), host: 1"
	if err == nil || err.Error() != want {
		t.Errorf("\nwant\n%s\ngot\n%v", want, err)
	}
	// Ensure that we tried only once.
	if sbc0.ExecCount != 1 {
		t.Errorf("want 1, got %v", sbc0.ExecCount)
	}
	if sbc1.ExecCount != 1 {
		t.Errorf("want 1, got %v", sbc1.ExecCount)
	}

	// duplicate shards
	resetSandbox()
	sbc = &sandboxConn{}
	testConns[0] = sbc
	qr, err = f([]string{"0", "0"})
	// Ensure that we executed only once.
	if sbc.ExecCount != 1 {
		t.Errorf("want 1, got %v", sbc.ExecCount)
	}

	// no errors
	resetSandbox()
	sbc0 = &sandboxConn{}
	testConns[0] = sbc0
	sbc1 = &sandboxConn{}
	testConns[1] = sbc1
	qr, err = f([]string{"0", "1"})
	if err != nil {
		t.Errorf("want nil, got %v", err)
	}
	if sbc0.ExecCount != 1 {
		t.Errorf("want 1, got %v", sbc0.ExecCount)
	}
	if sbc1.ExecCount != 1 {
		t.Errorf("want 1, got %v", sbc1.ExecCount)
	}
	if qr.RowsAffected != 2 {
		t.Errorf("want 2, got %v", qr.RowsAffected)
	}
	if len(qr.Rows) != 2 {
		t.Errorf("want 2, got %v", len(qr.Rows))
	}
}

func TestScatterConnStreamExecuteSendError(t *testing.T) {
	resetSandbox()
	sbc := &sandboxConn{}
	testConns[0] = sbc
	stc := NewScatterConn(new(sandboxTopo), "aa", 1*time.Millisecond, 3)
	err := stc.StreamExecute("query", nil, "", []string{"0"}, "", nil, func(interface{}) error {
		return fmt.Errorf("send error")
	})
	want := "send error"
	// Ensure that we handle send errors.
	if err == nil || err.Error() != want {
		t.Errorf("want %s, got %v", want, err)
	}
}

func TestScatterConnCommitSuccess(t *testing.T) {
	resetSandbox()
	sbc0 := &sandboxConn{}
	testConns[0] = sbc0
	sbc1 := &sandboxConn{mustFailTxPool: 1}
	testConns[1] = sbc1
	stc := NewScatterConn(new(sandboxTopo), "aa", 1*time.Millisecond, 3)

	// Sequence the executes to ensure commit order
	session := NewSafeSession(&proto.Session{InTransaction: true})
	stc.Execute("query1", nil, "", []string{"0"}, "", session)
	wantSession := proto.Session{
		InTransaction: true,
		ShardSessions: []*proto.ShardSession{{
			Keyspace:      "",
			Shard:         "0",
			TabletType:    "",
			TransactionId: 1,
		}},
	}
	if !reflect.DeepEqual(wantSession, *session.Session) {
		t.Errorf("want\n%#v, got\n%#v", wantSession, *session.Session)
	}
	stc.Execute("query1", nil, "", []string{"0", "1"}, "", session)
	wantSession = proto.Session{
		InTransaction: true,
		ShardSessions: []*proto.ShardSession{{
			Keyspace:      "",
			Shard:         "0",
			TabletType:    "",
			TransactionId: 1,
		}, {
			Keyspace:      "",
			Shard:         "1",
			TabletType:    "",
			TransactionId: 2,
		}},
	}
	if !reflect.DeepEqual(wantSession, *session.Session) {
		t.Errorf("want\n%#v, got\n%#v", wantSession, *session.Session)
	}
	sbc0.mustFailServer = 1
	err := stc.Commit(session)
	if err == nil {
		t.Errorf("want error, got nil")
	}
	wantSession = proto.Session{}
	if !reflect.DeepEqual(wantSession, *session.Session) {
		t.Errorf("want\n%#v, got\n%#v", wantSession, *session.Session)
	}
	if sbc0.CommitCount != 1 {
		t.Errorf("want 1, got %d", sbc0.CommitCount)
	}
	/*
		// Flaky: This test should be run manually.
		runtime.Gosched()
		if sbc1.RollbackCount != 1 {
			t.Errorf("want 1, got %d", sbc1.RollbackCount)
		}
	*/
}

func TestScatterConnRollback(t *testing.T) {
	resetSandbox()
	sbc0 := &sandboxConn{}
	testConns[0] = sbc0
	sbc1 := &sandboxConn{mustFailTxPool: 1}
	testConns[1] = sbc1
	stc := NewScatterConn(new(sandboxTopo), "aa", 1*time.Millisecond, 3)

	// Sequence the executes to ensure commit order
	session := NewSafeSession(&proto.Session{InTransaction: true})
	stc.Execute("query1", nil, "", []string{"0"}, "", session)
	stc.Execute("query1", nil, "", []string{"0", "1"}, "", session)
	err := stc.Rollback(session)
	if err != nil {
		t.Errorf("want nil, got %v", err)
	}
	wantSession := proto.Session{}
	if !reflect.DeepEqual(wantSession, *session.Session) {
		t.Errorf("want\n%#v, got\n%#v", wantSession, *session.Session)
	}
	/*
		// Flaky: This test should be run manually.
		runtime.Gosched()
		if sbc0.RollbackCount != 1 {
			t.Errorf("want 1, got %d", sbc0.RollbackCount)
		}
		if sbc1.RollbackCount != 1 {
			t.Errorf("want 1, got %d", sbc1.RollbackCount)
		}
	*/
}

func TestScatterConnClose(t *testing.T) {
	resetSandbox()
	sbc := &sandboxConn{}
	testConns[0] = sbc
	stc := NewScatterConn(new(sandboxTopo), "aa", 1*time.Millisecond, 3)
	stc.Execute("query1", nil, "", []string{"0"}, "", nil)
	stc.Close()
	/*
		// Flaky: This test should be run manually.
		runtime.Gosched()
		if sbc.CloseCount != 1 {
			t.Errorf("want 1, got %d", sbc.CloseCount)
		}
	*/
}
<|endoftext|>


<commit_before><commit_msg>player hand calculate value<commit_after><|endoftext|>


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

package main

import (
	"bytes"
	"encoding/hex"
	"fmt"
	"github.com/jacobsa/comeback/blob"
	"github.com/jacobsa/comeback/disk"
	"github.com/jacobsa/comeback/fs"
	"github.com/jacobsa/comeback/repr"
	"github.com/jacobsa/comeback/sys"
	"io"
	"log"
	"os"
	"path"
	"syscall"
	"time"
)

const (
	gTarget = "/tmp/restore_target"
)

var blobStore blob.Store

type score struct {
	hash []byte
}

func (s *score) Sha1Hash() []byte {
	return s.hash
}

func fromHexHash(h string) (blob.Score, error) {
	b, err := hex.DecodeString(h)
	if err != nil {
		return nil, fmt.Errorf("Invalid hex string: %s", h)
	}

	return &score{b}, nil
}

func chooseUserId(uid sys.UserId, username *string) (sys.UserId, error) {
	// If there is no symbolic username, just return the UID.
	if username == nil {
		return uid, nil
	}

	// Create a user registry.
	registry, err := sys.NewUserRegistry()
	if err != nil {
		return 0, fmt.Errorf("Creating user registry: %v", err)
	}

	// Attempt to look up the username. If it's not found, return the UID.
	betterUid, err := registry.FindByName(*username)

	if _, ok := err.(sys.NotFoundError); ok {
		return uid, nil
	} else if err != nil {
		return 0, fmt.Errorf("Looking up user: %v", err)
	}

	return betterUid, nil
}

func chooseGroupId(gid sys.GroupId, groupname *string) (sys.GroupId, error) {
	// If there is no symbolic groupname, just return the GID.
	if groupname == nil {
		return gid, nil
	}

	// Create a group registry.
	registry, err := sys.NewGroupRegistry()
	if err != nil {
		return 0, fmt.Errorf("Creating group registry: %v", err)
	}

	// Attempt to look up the groupname. If it's not found, return the GID.
	betterGid, err := registry.FindByName(*groupname)

	if _, ok := err.(sys.NotFoundError); ok {
		return gid, nil
	} else if err != nil {
		return 0, fmt.Errorf("Looking up group: %v", err)
	}

	return betterGid, nil
}

// Set the modification time for the supplied path without following symlinks
// (as syscall.Chtimes and therefore os.Chtimes do).
//
// c.f. http://stackoverflow.com/questions/10608724/set-modification-date-on-symbolic-link-in-cocoa
func setModTime(path string, mtime time.Time) error {
	// Open the file without following symlinks. Use O_NONBLOCK to allow opening
	// of named pipes without a writer.
	fd, err := syscall.Open(path, syscall.O_NONBLOCK|syscall.O_SYMLINK, 0)
	if err != nil {
		return err
	}

	defer syscall.Close(fd)

	// Call futimes.
	var utimes [2]syscall.Timeval
	atime := time.Now()
	atime_ns := atime.Unix()*1e9 + int64(atime.Nanosecond())
	mtime_ns := mtime.Unix()*1e9 + int64(mtime.Nanosecond())
	utimes[0] = syscall.NsecToTimeval(atime_ns)
	utimes[1] = syscall.NsecToTimeval(mtime_ns)

	err = syscall.Futimes(fd, utimes[0:])
	if err != nil {
		return err
	}

	return nil
}

// Restore the file whose contents are described by the referenced blobs to the
// supplied target, whose parent must already exist.
func restoreFile(target string, scores []blob.Score) error {
	// Open the file.
	f, err := os.Create(target)
	defer f.Close()

	if err != nil {
		return fmt.Errorf("Create: %v", err)
	}

	// Process each blob.
	for _, score := range scores {
		// Load the blob.
		blob, err := blobStore.Load(score)
		if err != nil {
			return fmt.Errorf("Loading blob: %v", err)
		}

		// Write out its contents.
		_, err = io.Copy(f, bytes.NewReader(blob))
		if err != nil {
			return fmt.Errorf("Copy: %v", err)
		}
	}

	return nil
}

// Like os.Chmod, but don't follow symlinks.
func setPermissions(path string, permissions os.FileMode) error {
	mode := syscallPermissions(permissions)

	// Open the file without following symlinks.
	fd, err := syscall.Open(path, syscall.O_RDONLY|syscall.O_SYMLINK, 0)
	if err != nil {
		return err
	}

	defer syscall.Close(fd)

	// Call fchmod.
	err = syscall.Fchmod(fd, mode)
	if err != nil {
		return err
	}

	return nil
}

// Restore the directory whose contents are described by the referenced blob to
// the supplied target, which must already exist.
func restoreDir(target string, score blob.Score) error {
	// Load the appropriate blob.
	blob, err := blobStore.Load(score)
	if err != nil {
		return fmt.Errorf("Loading blob: %v", err)
	}

	// Parse its contents.
	entries, err := repr.Unmarshal(blob)
	if err != nil {
		return fmt.Errorf("Parsing blob: %v", err)
	}

	// Deal with each entry.
	for _, entry := range entries {
		entryPath := path.Join(target, entry.Name)

		// Switch on type.
		switch entry.Type {
		case fs.TypeFile:
			if err := restoreFile(entryPath, entry.Scores); err != nil {
				return err
			}

			if err := setPermissions(entryPath, entry.Permissions); err != nil {
				return err
			}

		case fs.TypeDirectory:
			if len(entry.Scores) != 1 {
				return fmt.Errorf("Wrong number of scores: %v", entry)
			}

			if err = os.Mkdir(entryPath, 0700); err != nil {
				return err
			}

			if err = restoreDir(entryPath, entry.Scores[0]); err != nil {
				return err
			}

		case fs.TypeSymlink:
			err = os.Symlink(entry.Target, entryPath)
			if err != nil {
				return err
			}

		case fs.TypeNamedPipe:
			err = makeNamedPipe(entryPath, entry.Permissions)
			if err != nil {
				return err
			}

		case fs.TypeBlockDevice:
			err = makeBlockDevice(entryPath, entry.Permissions, entry.Device)
			if err != nil {
				return err
			}

		case fs.TypeCharDevice:
			err = makeCharDevice(entryPath, entry.Permissions, entry.Device)
			if err != nil {
				return err
			}

		default:
			return fmt.Errorf("Don't know how to deal with entry: %v", entry)
		}

		// Fix ownership.
		uid, err := chooseUserId(entry.Uid, entry.Username)
		if err != nil {
			return fmt.Errorf("chooseUserId: %v", err)
		}

		gid, err := chooseGroupId(entry.Gid, entry.Groupname)
		if err != nil {
			return fmt.Errorf("chooseGroupId: %v", err)
		}

		if err = os.Lchown(entryPath, int(uid), int(gid)); err != nil {
			return fmt.Errorf("Chown: %v", err)
		}

		// Fix modification time.
		if err = setModTime(entryPath, entry.MTime); err != nil {
			return fmt.Errorf("setModTime(%s): %v", entryPath, err)
		}
	}

	return nil
}

func syscallPermissions(permissions os.FileMode) (o uint32) {
	// Include r/w/x permission bits.
	o = uint32(permissions & os.ModePerm)

	// Also include setuid/setgid/sticky bits.
	if permissions&os.ModeSetuid != 0 {
		o |= syscall.S_ISUID
	}

	if permissions&os.ModeSetgid != 0 {
		o |= syscall.S_ISGID
	}

	if permissions&os.ModeSticky != 0 {
		o |= syscall.S_ISVTX
	}

	return
}

// Create a named pipe at the supplied path.
func makeNamedPipe(path string, permissions os.FileMode) error {
	return syscall.Mkfifo(path, syscallPermissions(permissions))
}

// Create a block device at the supplied path.
func makeBlockDevice(path string, permissions os.FileMode, dev int32) error {
	mode := syscallPermissions(permissions) | syscall.S_IFBLK
	if err := syscall.Mknod(path, mode, int(dev)); err != nil {
		return fmt.Errorf("syscall.Mknod: %v", err)
	}

	return nil
}

// Create a character device at the supplied path.
func makeCharDevice(path string, permissions os.FileMode, dev int32) error {
	mode := syscallPermissions(permissions) | syscall.S_IFCHR
	if err := syscall.Mknod(path, mode, int(dev)); err != nil {
		return fmt.Errorf("syscall.Mknod: %v", err)
	}

	return nil
}

func main() {
	var err error

	// Create the blob store.
	blobStore, err = disk.NewBlobStore("/tmp/blobs")
	if err != nil {
		log.Fatalf("Creating store: %v", err)
	}

	// Parse the score.
	score, err := fromHexHash("228a6254c7585525744192b51f099833fca8c654")
	if err != nil {
		log.Fatalf("Parsing score: %v", err)
	}

	// Make sure the target doesn't exist.
	err = os.RemoveAll(gTarget)
	if err != nil {
		log.Fatalf("RemoveAll: %v", err)
	}

	// Create the target.
	err = os.Mkdir("/tmp/restore_target", 0755)
	if err != nil {
		log.Fatalf("Mkdir: %v", err)
	}

	// Attempt a restore.
	err = restoreDir(gTarget, score)
	if err != nil {
		log.Fatalf("Restoring: %v", err)
	}
}
<commit_msg>Fixed mtimes on devices.<commit_after>// Copyright 2012 Aaron Jacobs. All Rights Reserved.
// Author: aaronjjacobs@gmail.com (Aaron Jacobs)
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package main

import (
	"bytes"
	"encoding/hex"
	"fmt"
	"github.com/jacobsa/comeback/blob"
	"github.com/jacobsa/comeback/disk"
	"github.com/jacobsa/comeback/fs"
	"github.com/jacobsa/comeback/repr"
	"github.com/jacobsa/comeback/sys"
	"io"
	"log"
	"os"
	"path"
	"syscall"
	"time"
)

const (
	gTarget = "/tmp/restore_target"
)

var blobStore blob.Store

type score struct {
	hash []byte
}

func (s *score) Sha1Hash() []byte {
	return s.hash
}

func fromHexHash(h string) (blob.Score, error) {
	b, err := hex.DecodeString(h)
	if err != nil {
		return nil, fmt.Errorf("Invalid hex string: %s", h)
	}

	return &score{b}, nil
}

func chooseUserId(uid sys.UserId, username *string) (sys.UserId, error) {
	// If there is no symbolic username, just return the UID.
	if username == nil {
		return uid, nil
	}

	// Create a user registry.
	registry, err := sys.NewUserRegistry()
	if err != nil {
		return 0, fmt.Errorf("Creating user registry: %v", err)
	}

	// Attempt to look up the username. If it's not found, return the UID.
	betterUid, err := registry.FindByName(*username)

	if _, ok := err.(sys.NotFoundError); ok {
		return uid, nil
	} else if err != nil {
		return 0, fmt.Errorf("Looking up user: %v", err)
	}

	return betterUid, nil
}

func chooseGroupId(gid sys.GroupId, groupname *string) (sys.GroupId, error) {
	// If there is no symbolic groupname, just return the GID.
	if groupname == nil {
		return gid, nil
	}

	// Create a group registry.
	registry, err := sys.NewGroupRegistry()
	if err != nil {
		return 0, fmt.Errorf("Creating group registry: %v", err)
	}

	// Attempt to look up the groupname. If it's not found, return the GID.
	betterGid, err := registry.FindByName(*groupname)

	if _, ok := err.(sys.NotFoundError); ok {
		return gid, nil
	} else if err != nil {
		return 0, fmt.Errorf("Looking up group: %v", err)
	}

	return betterGid, nil
}

// Set the modification time for the supplied path without following symlinks
// (as syscall.Chtimes and therefore os.Chtimes do).
//
// c.f. http://stackoverflow.com/questions/10608724/set-modification-date-on-symbolic-link-in-cocoa
func setModTime(path string, mtime time.Time) error {
	// Open the file without following symlinks. Use O_NONBLOCK to allow opening
	// of named pipes without a writer.
	fd, err := syscall.Open(path, syscall.O_NONBLOCK|syscall.O_SYMLINK, 0)
	if err != nil {
		return err
	}

	defer syscall.Close(fd)

	// Call futimes.
	var utimes [2]syscall.Timeval
	atime := time.Now()
	atime_ns := atime.Unix()*1e9 + int64(atime.Nanosecond())
	mtime_ns := mtime.Unix()*1e9 + int64(mtime.Nanosecond())
	utimes[0] = syscall.NsecToTimeval(atime_ns)
	utimes[1] = syscall.NsecToTimeval(mtime_ns)

	err = syscall.Futimes(fd, utimes[0:])
	if err != nil {
		return err
	}

	return nil
}

// Restore the file whose contents are described by the referenced blobs to the
// supplied target, whose parent must already exist.
func restoreFile(target string, scores []blob.Score) error {
	// Open the file.
	f, err := os.Create(target)
	defer f.Close()

	if err != nil {
		return fmt.Errorf("Create: %v", err)
	}

	// Process each blob.
	for _, score := range scores {
		// Load the blob.
		blob, err := blobStore.Load(score)
		if err != nil {
			return fmt.Errorf("Loading blob: %v", err)
		}

		// Write out its contents.
		_, err = io.Copy(f, bytes.NewReader(blob))
		if err != nil {
			return fmt.Errorf("Copy: %v", err)
		}
	}

	return nil
}

// Like os.Chmod, but don't follow symlinks.
func setPermissions(path string, permissions os.FileMode) error {
	mode := syscallPermissions(permissions)

	// Open the file without following symlinks.
	fd, err := syscall.Open(path, syscall.O_RDONLY|syscall.O_SYMLINK, 0)
	if err != nil {
		return err
	}

	defer syscall.Close(fd)

	// Call fchmod.
	err = syscall.Fchmod(fd, mode)
	if err != nil {
		return err
	}

	return nil
}

// Restore the directory whose contents are described by the referenced blob to
// the supplied target, which must already exist.
func restoreDir(target string, score blob.Score) error {
	// Load the appropriate blob.
	blob, err := blobStore.Load(score)
	if err != nil {
		return fmt.Errorf("Loading blob: %v", err)
	}

	// Parse its contents.
	entries, err := repr.Unmarshal(blob)
	if err != nil {
		return fmt.Errorf("Parsing blob: %v", err)
	}

	// Deal with each entry.
	for _, entry := range entries {
		entryPath := path.Join(target, entry.Name)

		// Switch on type.
		switch entry.Type {
		case fs.TypeFile:
			if err := restoreFile(entryPath, entry.Scores); err != nil {
				return err
			}

			if err := setPermissions(entryPath, entry.Permissions); err != nil {
				return err
			}

		case fs.TypeDirectory:
			if len(entry.Scores) != 1 {
				return fmt.Errorf("Wrong number of scores: %v", entry)
			}

			if err = os.Mkdir(entryPath, 0700); err != nil {
				return err
			}

			if err = restoreDir(entryPath, entry.Scores[0]); err != nil {
				return err
			}

		case fs.TypeSymlink:
			err = os.Symlink(entry.Target, entryPath)
			if err != nil {
				return err
			}

		case fs.TypeNamedPipe:
			err = makeNamedPipe(entryPath, entry.Permissions)
			if err != nil {
				return err
			}

		case fs.TypeBlockDevice:
			err = makeBlockDevice(entryPath, entry.Permissions, entry.Device)
			if err != nil {
				return err
			}

		case fs.TypeCharDevice:
			err = makeCharDevice(entryPath, entry.Permissions, entry.Device)
			if err != nil {
				return err
			}

		default:
			return fmt.Errorf("Don't know how to deal with entry: %v", entry)
		}

		// Fix ownership.
		uid, err := chooseUserId(entry.Uid, entry.Username)
		if err != nil {
			return fmt.Errorf("chooseUserId: %v", err)
		}

		gid, err := chooseGroupId(entry.Gid, entry.Groupname)
		if err != nil {
			return fmt.Errorf("chooseGroupId: %v", err)
		}

		if err = os.Lchown(entryPath, int(uid), int(gid)); err != nil {
			return fmt.Errorf("Chown: %v", err)
		}

		// Fix modification time, but not on devices (otherwise we get resource
		// busy errors).
		if entry.Type != fs.TypeBlockDevice && entry.Type != fs.TypeCharDevice {
			if err = setModTime(entryPath, entry.MTime); err != nil {
				return fmt.Errorf("setModTime(%s): %v", entryPath, err)
			}
		}
	}

	return nil
}

func syscallPermissions(permissions os.FileMode) (o uint32) {
	// Include r/w/x permission bits.
	o = uint32(permissions & os.ModePerm)

	// Also include setuid/setgid/sticky bits.
	if permissions&os.ModeSetuid != 0 {
		o |= syscall.S_ISUID
	}

	if permissions&os.ModeSetgid != 0 {
		o |= syscall.S_ISGID
	}

	if permissions&os.ModeSticky != 0 {
		o |= syscall.S_ISVTX
	}

	return
}

// Create a named pipe at the supplied path.
func makeNamedPipe(path string, permissions os.FileMode) error {
	return syscall.Mkfifo(path, syscallPermissions(permissions))
}

// Create a block device at the supplied path.
func makeBlockDevice(path string, permissions os.FileMode, dev int32) error {
	mode := syscallPermissions(permissions) | syscall.S_IFBLK
	if err := syscall.Mknod(path, mode, int(dev)); err != nil {
		return fmt.Errorf("syscall.Mknod: %v", err)
	}

	return nil
}

// Create a character device at the supplied path.
func makeCharDevice(path string, permissions os.FileMode, dev int32) error {
	mode := syscallPermissions(permissions) | syscall.S_IFCHR
	if err := syscall.Mknod(path, mode, int(dev)); err != nil {
		return fmt.Errorf("syscall.Mknod: %v", err)
	}

	return nil
}

func main() {
	var err error

	// Create the blob store.
	blobStore, err = disk.NewBlobStore("/tmp/blobs")
	if err != nil {
		log.Fatalf("Creating store: %v", err)
	}

	// Parse the score.
	score, err := fromHexHash("228a6254c7585525744192b51f099833fca8c654")
	if err != nil {
		log.Fatalf("Parsing score: %v", err)
	}

	// Make sure the target doesn't exist.
	err = os.RemoveAll(gTarget)
	if err != nil {
		log.Fatalf("RemoveAll: %v", err)
	}

	// Create the target.
	err = os.Mkdir("/tmp/restore_target", 0755)
	if err != nil {
		log.Fatalf("Mkdir: %v", err)
	}

	// Attempt a restore.
	err = restoreDir(gTarget, score)
	if err != nil {
		log.Fatalf("Restoring: %v", err)
	}
}
<|endoftext|>


<commit_before>package eventmaster

import (
	"database/sql"
	"encoding/json"
	"fmt"

	sq "github.com/Masterminds/squirrel"
	"github.com/lib/pq"
	_ "github.com/lib/pq"
	"github.com/pkg/errors"
	log "github.com/sirupsen/logrus"
	eventmaster "github.com/wish/eventmaster/proto"
)

// implements datastore
type PostgresStore struct {
	db *sql.DB
}

type PostgresConfig struct {
	Addr        string `json:"addr"`
	Port        int    `json:"port"`
	Database    string `json:"database"`
	ServiceName string `json:"service_name"`
	Username    string `json:"username"`
	Password    string `json:"password"`
}

func NewPostgresStore(c PostgresConfig) (*PostgresStore, error) {
	var host string
	if c.ServiceName != "" {
		host = c.ServiceName + ".service.consul"
	} else {
		host = c.Addr
	}
	log.Infof("Connecting to postgres: %v", host)
	psqlInfo := fmt.Sprintf("host=%s port=%d user=%s "+
		"password=%s dbname=%s sslmode=disable",
		host, c.Port, c.Username, c.Password, c.Database)
	db, err := sql.Open("postgres", psqlInfo)
	if err != nil {
		return nil, errors.Wrap(err, "Error creating postgres session")
	}
	err = db.Ping()
	if err != nil {
		return nil, errors.Wrap(err, "Error creating postgres session")
	}
	log.Infof("Successfully connected to postgres %s", host)

	return &PostgresStore{
		db: db,
	}, nil
}

func (p *PostgresStore) AddEvent(e *Event) error {
	data := "{}"
	if e.Data != nil {
		dataBytes, err := json.Marshal(e.Data)
		if err != nil {
			return err
		}
		data = string(dataBytes)
	}

	tx, err := p.db.Begin()
	if err != nil {
		return err
	}

	_, err = tx.Exec("INSERT INTO event "+
		"(event_id, parent_event_id, dc_id, topic_id, host, target_host_set, \"user\", event_time, tag_set, received_time)"+
		" VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10)",
		e.EventID, e.ParentEventID, e.DCID, e.TopicID, e.Host, pq.Array(e.TargetHosts), e.User, e.EventTime, pq.Array(e.Tags), e.ReceivedTime)

	if err != nil {
		tx.Rollback()
		return err
	}

	_, err = tx.Exec("INSERT INTO event_metadata (event_id, data_json) VALUES ($1, $2)", e.EventID, data)
	if err != nil {
		tx.Rollback()
		return err
	}

	return tx.Commit()
}

func getPlaceHolderArray(slice []string) string {
	expr := "ARRAY["
	first := true
	for i := 0; i < len(slice); i++ {
		if !first {
			expr += ", "
		}
		expr += "?"
		first = false
	}
	expr += "]"
	return expr
}

func stringsToIfaces(in []string) []interface{} {
	new := make([]interface{}, len(in))
	for i, v := range in {
		new[i] = v
	}
	return new
}

func (p *PostgresStore) Find(q *eventmaster.Query, topicIDs []string, dcIDs []string) (Events, error) {
	// use sql builder for better readability
	psql := sq.StatementBuilder.PlaceholderFormat(sq.Dollar)
	columns := []string{"event_id", "parent_event_id", "dc_id", "topic_id", "host", "target_host_set", "\"user\"", "event_time", "tag_set", "received_time"}
	table := "event"
	base_query := psql.Select(columns...).From(table)

	wheres := sq.And{}

	// add constarints, *1000 due to different time precision (sec to ms)
	wheres = append(wheres, sq.GtOrEq{"event_time": q.StartEventTime * 1000})
	wheres = append(wheres, sq.LtOrEq{"event_time": q.EndEventTime * 1000})

	if len(q.User) > 0 {
		wheres = append(wheres, sq.Eq{"user": q.User})
	}

	if len(q.ParentEventID) > 0 {
		wheres = append(wheres, sq.Eq{"parent_event_id": q.ParentEventID})
	}

	if len(q.Host) > 0 {
		wheres = append(wheres, sq.Eq{"host": q.Host})
	}

	if len(topicIDs) > 0 {
		wheres = append(wheres, sq.Eq{"topic_id": topicIDs})
	}

	if len(dcIDs) > 0 {
		wheres = append(wheres, sq.Eq{"dc_id": dcIDs})
	}

	if len(q.TargetHostSet) > 0 {
		// intersection of two arrays
		// squirrel doesn't support Postgres' array and repective methods. Building new SQL template manually
		// make identical slice of []interface{} type to pass to varadic functions
		clause := "target_host_set && " + getPlaceHolderArray((q.TargetHostSet))
		wheres = append(wheres, sq.Expr(clause, stringsToIfaces(q.TargetHostSet)...))
	}

	if len(q.TagSet) > 0 {
		var op string
		if q.TagAndOperator {
			// LHS array is the superset of RHS array
			op = ">&"
		} else {
			// LHS array has interestion with RHS array
			op = "&&"
		}
		clause := fmt.Sprintf("tag_set %s %s", op, getPlaceHolderArray(q.TagSet))
		wheres = append(wheres, sq.Expr(clause, stringsToIfaces(q.TagSet)...))
	}

	if len(q.ExcludeTagSet) > 0 {
		clause := fmt.Sprintf("NOT (tag_set && %s)", getPlaceHolderArray((q.ExcludeTagSet)))
		wheres = append(wheres, sq.Expr(clause, stringsToIfaces(q.ExcludeTagSet)...))
	}

	query := base_query.Where(wheres)

	var events []*Event

	var event_id string
	var parent_event_id string
	var dc_id string
	var topic_id string
	var host string
	var target_host_set []string
	var user string
	var event_time int64
	var tag_set []string
	var received_time int64

	rows, err := query.RunWith(p.db).Query()
	for rows.Next() {
		err = rows.Scan(&event_id, &parent_event_id, &dc_id, &topic_id, &host, pq.Array(&target_host_set), &user, &event_time, pq.Array(&tag_set), &received_time)
		if err != nil {
			return nil, err
		}
		events = append(events, &Event{
			EventID:       event_id,
			ParentEventID: parent_event_id,
			DCID:          dc_id,
			TopicID:       topic_id,
			Host:          host,
			TargetHosts:   target_host_set,
			User:          user,
			EventTime:     event_time,
			Tags:          tag_set,
			ReceivedTime:  received_time,
		})
	}

	return events, nil
}

func (p *PostgresStore) FindByID(id string, inclData bool) (*Event, error) {
	rows, err := p.db.Query("SELECT event_id, parent_event_id, dc_id, topic_id, host,"+
		" target_host_set, user, event_time, tag_set, received_time "+
		"FROM event WHERE event_id=$1 LIMIT 1", id)

	if err != nil {
		return nil, err
	}

	event := Event{}
	if rows.Next() {
		err = rows.Scan(&event.EventID, &event.ParentEventID, &event.DCID,
			&event.TopicID, &event.Host, pq.Array(&event.TargetHosts), &event.User,
			&event.EventTime, pq.Array(&event.Tags), &event.ReceivedTime)
		if err != nil {
			return nil, err
		}
	} else {
		return nil, errors.New("cannot find event entry with given id")
	}

	if err != nil {
		return nil, err
	}

	if inclData {
		rows, err = p.db.Query("SELECT data_json FROM event_metadata WHERE event_id=$1 LIMIT 1", id)
		if err != nil {
			return nil, err
		}
		var data []byte
		if rows.Next() {
			err = rows.Scan(&data)
			if err != nil {
				return nil, err
			}
			err = json.Unmarshal(data, &event.Data)
			if err != nil {
				return nil, err
			}
		} else {
			return nil, errors.New("cannot find data entry with given id")
		}
	}

	return &event, nil
}

func (p *PostgresStore) FindIDs(query *eventmaster.TimeQuery, handle HandleEvent) error {
	var order string
	if query.Ascending {
		order = "ASC"
	} else {
		order = "DESC"
	}

	// *1000 to convert time precision (sec to ms)
	rows, err := p.db.Query("SELECT event_id FROM event WHERE event_time >= $1 AND event_time <= $2 "+
		"ORDER BY event_time "+order+" LIMIT $3 ", query.StartEventTime*1000, query.EndEventTime*1000, query.Limit)
	if err != nil {
		return err
	}

	for rows.Next() {
		var id string
		err = rows.Scan(&id)
		if err != nil {
			return err
		}
		err = handle(id)
		if err != nil {
			rows.Close()
			return err
		}
	}
	return nil
}

func (p *PostgresStore) GetTopics() ([]Topic, error) {
	rows, err := p.db.Query("SELECT topic_id, topic_name, data_schema FROM event_topic")
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var topics []Topic
	var id string
	var name, schema string

	for rows.Next() {
		err = rows.Scan(&id, &name, &schema)
		if err != nil {
			return nil, err
		}
		var s map[string]interface{}
		err := json.Unmarshal([]byte(schema), &s)
		if err != nil {
			return nil, err
		}
		topics = append(topics, Topic{
			ID:     id,
			Name:   name,
			Schema: s,
		})
	}

	return topics, nil
}

func (p *PostgresStore) AddTopic(t RawTopic) error {
	_, err := p.db.Exec("INSERT INTO event_topic (topic_id, topic_name, data_schema) VALUES ($1, $2, $3)", t.ID, t.Name, t.Schema)
	return err
}

func (p *PostgresStore) UpdateTopic(t RawTopic) error {
	_, err := p.db.Exec("UPDATE event_topic SET topic_name=$1, data_schema=$2 WHERE topic_id=$3", t.Name, t.Schema, t.ID)
	return err
}

func (p *PostgresStore) DeleteTopic(id string) error {
	_, err := p.db.Exec("DELETE FROM event_topic WHERE topic_id=$1", id)
	return err
}

func (p *PostgresStore) GetDCs() ([]DC, error) {
	rows, err := p.db.Query("SELECT dc_id, dc from event_dc")
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var dcs []DC

	var dc_id, dc string
	for rows.Next() {
		err = rows.Scan(&dc_id, &dc)
		if err != nil {
			return nil, err
		}
		dcs = append(dcs, DC{
			ID:   dc_id,
			Name: dc,
		})
	}
	return dcs, nil
}

func (p *PostgresStore) AddDC(dc DC) error {
	stmt, err := p.db.Prepare("INSERT INTO event_dc (dc_id, dc) VALUES ($1, $2)")
	if err != nil {
		return errors.Wrap(err, "Failed preparing insert DC statement")
	}
	_, err = stmt.Exec(dc.ID, dc.Name)
	if err != nil {
		return err
	}

	return nil
}

func (p *PostgresStore) UpdateDC(id string, newName string) error {
	_, err := p.db.Exec("UPDATE event_dc SET dc=$1 WHERE dc_id=$2", newName, id)
	return err
}

func (p *PostgresStore) CloseSession() {
	p.db.Close()
}
<commit_msg>Add comments for public functions<commit_after>package eventmaster

import (
	"database/sql"
	"encoding/json"
	"fmt"

	sq "github.com/Masterminds/squirrel"
	"github.com/lib/pq"
	_ "github.com/lib/pq"
	"github.com/pkg/errors"
	log "github.com/sirupsen/logrus"
	eventmaster "github.com/wish/eventmaster/proto"
)

// implements datastore
type PostgresStore struct {
	db *sql.DB
}

type PostgresConfig struct {
	Addr        string `json:"addr"`
	Port        int    `json:"port"`
	Database    string `json:"database"`
	ServiceName string `json:"service_name"`
	Username    string `json:"username"`
	Password    string `json:"password"`
}

// Initializes new connection to Postgres DB
func NewPostgresStore(c PostgresConfig) (*PostgresStore, error) {
	var host string
	if c.ServiceName != "" {
		host = c.ServiceName + ".service.consul"
	} else {
		host = c.Addr
	}
	log.Infof("Connecting to postgres: %v", host)
	psqlInfo := fmt.Sprintf("host=%s port=%d user=%s "+
		"password=%s dbname=%s sslmode=disable",
		host, c.Port, c.Username, c.Password, c.Database)
	db, err := sql.Open("postgres", psqlInfo)
	if err != nil {
		return nil, errors.Wrap(err, "Error creating postgres session")
	}
	err = db.Ping()
	if err != nil {
		return nil, errors.Wrap(err, "Error creating postgres session")
	}
	log.Infof("Successfully connected to postgres %s", host)

	return &PostgresStore{
		db: db,
	}, nil
}

// Add a new event to the DB
func (p *PostgresStore) AddEvent(e *Event) error {
	data := "{}"
	if e.Data != nil {
		dataBytes, err := json.Marshal(e.Data)
		if err != nil {
			return err
		}
		data = string(dataBytes)
	}

	tx, err := p.db.Begin()
	if err != nil {
		return err
	}

	_, err = tx.Exec("INSERT INTO event "+
		"(event_id, parent_event_id, dc_id, topic_id, host, target_host_set, \"user\", event_time, tag_set, received_time)"+
		" VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10)",
		e.EventID, e.ParentEventID, e.DCID, e.TopicID, e.Host, pq.Array(e.TargetHosts), e.User, e.EventTime, pq.Array(e.Tags), e.ReceivedTime)

	if err != nil {
		rollBackErr := tx.Rollback()
		if rollBackErr != nil {
			return errors.Wrap(err, "Rollback failure")
		}
		return err
	}

	_, err = tx.Exec("INSERT INTO event_metadata (event_id, data_json) VALUES ($1, $2)", e.EventID, data)
	if err != nil {
		rollBackErr := tx.Rollback()
		if rollBackErr != nil {
			return errors.Wrap(err, "Rollback failure")
		}
		return err
	}

	return tx.Commit()
}

func getPlaceHolderArray(slice []string) string {
	expr := "ARRAY["
	first := true
	for i := 0; i < len(slice); i++ {
		if !first {
			expr += ", "
		}
		expr += "?"
		first = false
	}
	expr += "]"
	return expr
}

func stringsToIfaces(in []string) []interface{} {
	new := make([]interface{}, len(in))
	for i, v := range in {
		new[i] = v
	}
	return new
}

// Find respective events that match given criteria
func (p *PostgresStore) Find(q *eventmaster.Query, topicIDs []string, dcIDs []string) (Events, error) {
	// use sql builder for better readability
	psql := sq.StatementBuilder.PlaceholderFormat(sq.Dollar)
	columns := []string{"event_id", "parent_event_id", "dc_id", "topic_id", "host", "target_host_set", "\"user\"", "event_time", "tag_set", "received_time"}
	table := "event"
	base_query := psql.Select(columns...).From(table)

	wheres := sq.And{}

	// add constarints, *1000 due to different time precision (sec to ms)
	wheres = append(wheres, sq.GtOrEq{"event_time": q.StartEventTime * 1000})
	wheres = append(wheres, sq.LtOrEq{"event_time": q.EndEventTime * 1000})

	if len(q.User) > 0 {
		wheres = append(wheres, sq.Eq{"user": q.User})
	}

	if len(q.ParentEventID) > 0 {
		wheres = append(wheres, sq.Eq{"parent_event_id": q.ParentEventID})
	}

	if len(q.Host) > 0 {
		wheres = append(wheres, sq.Eq{"host": q.Host})
	}

	if len(topicIDs) > 0 {
		wheres = append(wheres, sq.Eq{"topic_id": topicIDs})
	}

	if len(dcIDs) > 0 {
		wheres = append(wheres, sq.Eq{"dc_id": dcIDs})
	}

	if len(q.TargetHostSet) > 0 {
		// intersection of two arrays
		// squirrel doesn't support Postgres' array and repective methods. Building new SQL template manually
		// make identical slice of []interface{} type to pass to varadic functions
		clause := "target_host_set && " + getPlaceHolderArray((q.TargetHostSet))
		wheres = append(wheres, sq.Expr(clause, stringsToIfaces(q.TargetHostSet)...))
	}

	if len(q.TagSet) > 0 {
		var op string
		if q.TagAndOperator {
			// LHS array is the superset of RHS array
			op = ">&"
		} else {
			// LHS array has interestion with RHS array
			op = "&&"
		}
		clause := fmt.Sprintf("tag_set %s %s", op, getPlaceHolderArray(q.TagSet))
		wheres = append(wheres, sq.Expr(clause, stringsToIfaces(q.TagSet)...))
	}

	if len(q.ExcludeTagSet) > 0 {
		clause := fmt.Sprintf("NOT (tag_set && %s)", getPlaceHolderArray((q.ExcludeTagSet)))
		wheres = append(wheres, sq.Expr(clause, stringsToIfaces(q.ExcludeTagSet)...))
	}

	query := base_query.Where(wheres)

	var events []*Event

	var event_id string
	var parent_event_id string
	var dc_id string
	var topic_id string
	var host string
	var target_host_set []string
	var user string
	var event_time int64
	var tag_set []string
	var received_time int64

	rows, err := query.RunWith(p.db).Query()
	for rows.Next() {
		err = rows.Scan(&event_id, &parent_event_id, &dc_id, &topic_id, &host, pq.Array(&target_host_set), &user, &event_time, pq.Array(&tag_set), &received_time)
		if err != nil {
			return nil, err
		}
		events = append(events, &Event{
			EventID:       event_id,
			ParentEventID: parent_event_id,
			DCID:          dc_id,
			TopicID:       topic_id,
			Host:          host,
			TargetHosts:   target_host_set,
			User:          user,
			EventTime:     event_time,
			Tags:          tag_set,
			ReceivedTime:  received_time,
		})
	}

	return events, nil
}

// Find the event with given ID, include detailed json data if needed
func (p *PostgresStore) FindByID(id string, inclData bool) (*Event, error) {
	rows, err := p.db.Query("SELECT event_id, parent_event_id, dc_id, topic_id, host,"+
		" target_host_set, user, event_time, tag_set, received_time "+
		"FROM event WHERE event_id=$1 LIMIT 1", id)

	if err != nil {
		return nil, err
	}

	event := Event{}
	if rows.Next() {
		err = rows.Scan(&event.EventID, &event.ParentEventID, &event.DCID,
			&event.TopicID, &event.Host, pq.Array(&event.TargetHosts), &event.User,
			&event.EventTime, pq.Array(&event.Tags), &event.ReceivedTime)
		if err != nil {
			return nil, err
		}
	} else {
		return nil, errors.New("cannot find event entry with given id")
	}

	if err != nil {
		return nil, err
	}

	if inclData {
		rows, err = p.db.Query("SELECT data_json FROM event_metadata WHERE event_id=$1 LIMIT 1", id)
		if err != nil {
			return nil, err
		}
		var data []byte
		if rows.Next() {
			err = rows.Scan(&data)
			if err != nil {
				return nil, err
			}
			err = json.Unmarshal(data, &event.Data)
			if err != nil {
				return nil, err
			}
		} else {
			return nil, errors.New("cannot find data entry with given id")
		}
	}

	return &event, nil
}

// Find eventIDs that match the given TimeQuery constraints
func (p *PostgresStore) FindIDs(query *eventmaster.TimeQuery, handle HandleEvent) error {
	var order string
	if query.Ascending {
		order = "ASC"
	} else {
		order = "DESC"
	}

	// *1000 to convert time precision (sec to ms)
	rows, err := p.db.Query("SELECT event_id FROM event WHERE event_time >= $1 AND event_time <= $2 "+
		"ORDER BY event_time "+order+" LIMIT $3 ", query.StartEventTime*1000, query.EndEventTime*1000, query.Limit)
	if err != nil {
		return err
	}

	for rows.Next() {
		var id string
		err = rows.Scan(&id)
		if err != nil {
			return err
		}
		err = handle(id)
		if err != nil {
			rows.Close()
			return err
		}
	}
	return nil
}

// Get all topics
func (p *PostgresStore) GetTopics() ([]Topic, error) {
	rows, err := p.db.Query("SELECT topic_id, topic_name, data_schema FROM event_topic")
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var topics []Topic
	var id string
	var name, schema string

	for rows.Next() {
		err = rows.Scan(&id, &name, &schema)
		if err != nil {
			return nil, err
		}
		var s map[string]interface{}
		err := json.Unmarshal([]byte(schema), &s)
		if err != nil {
			return nil, err
		}
		topics = append(topics, Topic{
			ID:     id,
			Name:   name,
			Schema: s,
		})
	}

	return topics, nil
}

// Add a topic to the DB
func (p *PostgresStore) AddTopic(t RawTopic) error {
	_, err := p.db.Exec("INSERT INTO event_topic (topic_id, topic_name, data_schema) VALUES ($1, $2, $3)", t.ID, t.Name, t.Schema)
	return err
}

// Update topic with given topic ID
func (p *PostgresStore) UpdateTopic(t RawTopic) error {
	_, err := p.db.Exec("UPDATE event_topic SET topic_name=$1, data_schema=$2 WHERE topic_id=$3", t.Name, t.Schema, t.ID)
	return err
}

// Delete topic with given topic ID
func (p *PostgresStore) DeleteTopic(id string) error {
	_, err := p.db.Exec("DELETE FROM event_topic WHERE topic_id=$1", id)
	return err
}

// Get all Datacenters
func (p *PostgresStore) GetDCs() ([]DC, error) {
	rows, err := p.db.Query("SELECT dc_id, dc from event_dc")
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var dcs []DC

	var dc_id, dc string
	for rows.Next() {
		err = rows.Scan(&dc_id, &dc)
		if err != nil {
			return nil, err
		}
		dcs = append(dcs, DC{
			ID:   dc_id,
			Name: dc,
		})
	}
	return dcs, nil
}

// Add a datacenter to the table
func (p *PostgresStore) AddDC(dc DC) error {
	stmt, err := p.db.Prepare("INSERT INTO event_dc (dc_id, dc) VALUES ($1, $2)")
	if err != nil {
		return errors.Wrap(err, "Failed preparing insert DC statement")
	}
	_, err = stmt.Exec(dc.ID, dc.Name)
	if err != nil {
		return err
	}

	return nil
}

// Update datacenter with given ID
func (p *PostgresStore) UpdateDC(id string, newName string) error {
	_, err := p.db.Exec("UPDATE event_dc SET dc=$1 WHERE dc_id=$2", newName, id)
	return err
}

// Close the database session
func (p *PostgresStore) CloseSession() {
	p.db.Close()
}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"fmt"
)

type player struct {
	HP          int
	MP          int
	Simellas    int
	Armour      armour
	Weapon      weapon
	Shield      shield
	Consumables map[consumable]int
	Rods        map[rod]rodProps
	Aptitudes   map[aptitude]bool
	Statuses    map[status]int
	Pos         position
	Target      position
	LOS         map[position]bool
	Rays        rayMap
	Bored       int
}

func (p *player) HPMax() int {
	hpmax := 42
	if p.Aptitudes[AptHealthy] {
		hpmax += 10
	}
	hpmax -= 3 * p.Bored
	if hpmax < 21 {
		hpmax = 21
	}
	return hpmax
}

func (p *player) MPMax() int {
	mpmax := 3
	if p.Aptitudes[AptMagic] {
		mpmax += 2
	}
	if p.Armour == CelmistRobe {
		mpmax += 2
	}
	return mpmax
}

func (p *player) Accuracy() int {
	acc := 15
	if p.Aptitudes[AptAccurate] {
		acc += 2
	}
	return acc
}

func (p *player) RangedAccuracy() int {
	acc := 15
	if p.Aptitudes[AptAccurate] {
		acc += 10
	}
	return acc
}

func (p *player) Armor() int {
	ar := 0
	switch p.Armour {
	case LeatherArmour:
		ar += 3
	case ChainMail:
		ar += 6
	case PonderousnessPlates:
		ar += 10
	}
	if p.Aptitudes[AptScales] {
		ar += 2
	}
	if p.HasStatus(StatusLignification) {
		ar = 9 + ar/2
	}
	if p.HasStatus(StatusCorrosion) {
		ar -= 2 * p.Statuses[StatusCorrosion]
		if ar < 0 {
			ar = 0
		}
	}
	return ar
}

func (p *player) Attack() int {
	attack := p.Weapon.Attack()
	if p.Aptitudes[AptStrong] {
		attack += attack / 5
	}
	if p.HasStatus(StatusCorrosion) {
		penalty := p.Statuses[StatusCorrosion]
		if penalty > 5 {
			penalty = 5
		}
		attack -= penalty
	}
	return attack
}

func (p *player) Block() int {
	block := p.Shield.Block()
	if p.HasStatus(StatusDisabledShield) {
		block /= 3
	}
	return block
}

func (p *player) Evasion() int {
	ev := 15
	if p.Aptitudes[AptAgile] {
		ev += 3
	}
	switch p.Armour {
	case ChainMail:
		ev -= 1
	case PonderousnessPlates:
		ev -= 2
	}
	if p.HasStatus(StatusAgile) {
		ev += 7
	}
	return ev
}

func (p *player) HasStatus(st status) bool {
	return p.Statuses[st] > 0
}

func (p *player) AptitudeCount() int {
	count := 0
	for _, b := range p.Aptitudes {
		if b {
			count++
		}
	}
	return count
}

func (g *game) AutoToDir(ev event) bool {
	if g.MonsterInLOS() == nil {
		err := g.MovePlayer(g.Player.Pos.To(g.AutoDir), ev)
		if err != nil {
			g.Print(err.Error())
			g.AutoDir = NoDir
			return false
		}
		return true
	}
	g.AutoDir = NoDir
	return false
}

func (g *game) GoToDir(dir direction, ev event) error {
	if g.MonsterInLOS() != nil {
		g.AutoDir = NoDir
		return errors.New("You cannot travel while there are monsters in view.")
	}
	err := g.MovePlayer(g.Player.Pos.To(dir), ev)
	if err != nil {
		return err
	}
	g.AutoDir = dir
	return nil
}

func (g *game) MoveToTarget(ev event) bool {
	if !g.AutoTarget.valid() {
		return false
	}
	path := g.PlayerPath(g.Player.Pos, g.AutoTarget)
	if g.MonsterInLOS() != nil {
		g.AutoTarget = InvalidPos
	}
	if len(path) < 1 {
		g.AutoTarget = InvalidPos
		return false
	}
	var err error
	if len(path) > 1 {
		err = g.MovePlayer(path[len(path)-2], ev)
		if g.ExclusionsMap[path[len(path)-2]] {
			g.AutoTarget = InvalidPos
		}
	} else {
		g.WaitTurn(ev)
	}
	if err != nil {
		g.Print(err.Error())
		g.AutoTarget = InvalidPos
		return false
	}
	if g.AutoTarget.valid() && g.Player.Pos == g.AutoTarget {
		g.AutoTarget = InvalidPos
	}
	return true
}

func (g *game) WaitTurn(ev event) {
	// XXX Really wait for 10 ?
	grade := 3
	if len(g.Noise) > 0 || g.StatusRest() {
		grade = 1
	}
	g.BoredomAction(ev, grade)
	ev.Renew(g, 10)
}

func (g *game) MonsterCount() (count int) {
	for _, mons := range g.Monsters {
		if mons.Exists() {
			count++
		}
	}
	return count
}

func (g *game) BoredomAction(ev event, grade int) {
	obor := g.Boredom
	if g.MonsterInLOS() == nil {
		g.Boredom += grade
	} else {
		g.Boredom--
		if g.Boredom < 0 {
			g.Boredom = 0
			g.Player.Bored = 0
		}
		return
	}
	if g.Boredom >= 120 && obor < 120 {
		if g.MonsterCount() > 4 {
			g.PrintStyled("You feel a little bored, your health could decline.", logCritic)
			g.StopAuto()
		}
	}
	if g.Boredom >= 130 && (obor/10 != g.Boredom/10) {
		g.Player.Bored++
		g.PrintStyled("You feel unhealthy.", logCritic)
		g.StopAuto()
		if g.Player.HP > g.Player.HPMax() {
			g.Player.HP -= 3
		}
	}
}

func (g *game) FunAction() {
	g.Boredom -= 15
	if g.Boredom < 0 {
		g.Boredom = 0
		g.Player.Bored = 0
	}
}

func (g *game) Rest(ev event) error {
	if g.MonsterInLOS() != nil {
		return fmt.Errorf("You cannot sleep while monsters are in view.")
	}
	if cld, ok := g.Clouds[g.Player.Pos]; ok && cld == CloudFire {
		return errors.New("You cannot rest on flames.")
	}
	if !g.NeedsRegenRest() && !g.StatusRest() {
		return errors.New("You do not need to rest.")
	}
	g.WaitTurn(ev)
	g.Resting = true
	g.RestingTurns = 0
	if g.StatusRest() {
		g.RestingTurns = -1 // not true resting, just waiting for status end
	}
	g.FunAction()
	return nil
}

func (g *game) StatusRest() bool {
	for _, q := range g.Player.Statuses {
		if q > 0 {
			return true
		}
	}
	return false
}

func (g *game) NeedsRegenRest() bool {
	return g.Player.HP < g.Player.HPMax() || g.Player.MP < g.Player.MPMax()
}

func (g *game) Equip(ev event) error {
	if eq, ok := g.Equipables[g.Player.Pos]; ok {
		eq.Equip(g)
		ev.Renew(g, 10)
		return nil
	}
	return errors.New("Found nothing to equip here.")
}

func (g *game) Teleportation(ev event) {
	var pos position
	i := 0
	count := 0
	for {
		count++
		if count > 1000 {
			panic("Teleportation")
		}
		pos = g.FreeCell()
		if pos.Distance(g.Player.Pos) < 15 && i < 1000 {
			i++
			continue
		}
		break

	}
	if pos.valid() {
		// should always happen
		opos := g.Player.Pos
		g.Print("You feel yourself teleported away.")
		g.ui.TeleportAnimation(g, opos, pos, true)
		g.PlacePlayerAt(pos)
	} else {
		// should not happen
		g.Print("Something went wrong with the teleportation.")
	}
}

func (g *game) CollectGround() {
	pos := g.Player.Pos
	if g.Simellas[pos] > 0 {
		g.Player.Simellas += g.Simellas[pos]
		if g.Simellas[pos] == 1 {
			g.Print("You pick up a simella.")
		} else {
			g.Printf("You pick up %d simellas.", g.Simellas[pos])
		}
		g.DijkstraMapRebuild = true
		delete(g.Simellas, pos)
	}
	if c, ok := g.Collectables[pos]; ok {
		g.Player.Consumables[c.Consumable] += c.Quantity
		g.DijkstraMapRebuild = true
		delete(g.Collectables, pos)
		if c.Quantity > 1 {
			g.Printf("You take %d %s.", c.Quantity, c.Consumable.Plural())
		} else {
			g.Printf("You take %s.", Indefinite(c.Consumable.String(), false))
		}
	}
	if r, ok := g.Rods[pos]; ok {
		g.Player.Rods[r] = rodProps{Charge: r.MaxCharge() - 1}
		g.DijkstraMapRebuild = true
		delete(g.Rods, pos)
		g.Printf("You take a %s.", r)
		g.StoryPrintf("You found and took a %s.", r)
	}
	if eq, ok := g.Equipables[pos]; ok {
		g.Printf("You stand over %s.", Indefinite(eq.String(), false))
	} else if _, ok := g.Stairs[pos]; ok {
		g.Print("You stand over stairs.")
	} else if g.Doors[pos] {
		g.Print("You stand at the door.")
	}
}

func (g *game) MovePlayer(pos position, ev event) error {
	if !pos.valid() {
		return errors.New("You cannot move there.")
	}
	c := g.Dungeon.Cell(pos)
	if c.T == WallCell && !g.Player.HasStatus(StatusDig) {
		return errors.New("You cannot move into a wall.")
	}
	if g.Player.HasStatus(StatusConfusion) {
		switch pos.Dir(g.Player.Pos) {
		case E, N, W, S:
		default:
			return errors.New("You cannot use diagonal movements while confused.")
		}
	}
	delay := 10
	mons := g.MonsterAt(pos)
	if !mons.Exists() {
		if g.Player.HasStatus(StatusLignification) {
			return errors.New("You cannot move while lignified")
		}
		if c.T == WallCell {
			g.Dungeon.SetCell(pos, FreeCell)
			g.MakeNoise(WallNoise, pos)
			g.Print(g.CrackSound())
			g.Fog(pos, 1, ev)
			g.Stats.Digs++
		}
		g.PlacePlayerAt(pos)
		if !g.Autoexploring {
			g.BoredomAction(ev, 1)
		}
		if g.Player.Aptitudes[AptFast] {
			// only fast for movement
			delay -= 2
		}
		switch g.Player.Armour {
		case PonderousnessPlates:
			delay += 2
		case SpeedRobe:
			delay -= 3
		}
		if g.Player.HasStatus(StatusSwift) {
			// only fast for movement
			delay -= 3
		}
		g.Stats.Moves++
	} else {
		g.FunAction()
		g.AttackMonster(mons, ev)
	}
	if g.Player.HasStatus(StatusBerserk) {
		delay -= 3
	}
	if g.Player.HasStatus(StatusSlow) {
		delay += 3
	}
	if delay < 3 {
		delay = 3
	}
	ev.Renew(g, delay)
	return nil
}

func (g *game) HealPlayer(ev event) {
	if g.Player.HP < g.Player.HPMax() {
		g.Player.HP++
	}
	delay := 50
	ev.Renew(g, delay)
}

func (g *game) MPRegen(ev event) {
	if g.Player.MP < g.Player.MPMax() {
		g.Player.MP++
	}
	delay := 100
	ev.Renew(g, delay)
}

func (g *game) Smoke(ev event) {
	dij := &normalPath{game: g}
	nm := Dijkstra(dij, []position{g.Player.Pos}, 2)
	for pos := range nm {
		_, ok := g.Clouds[pos]
		if !ok {
			g.Clouds[pos] = CloudFog
			g.PushEvent(&cloudEvent{ERank: ev.Rank() + 100 + RandInt(100), EAction: CloudEnd, Pos: pos})
		}
	}
	g.Player.Statuses[StatusSwift]++
	g.PushEvent(&simpleEvent{ERank: ev.Rank() + 20 + RandInt(10), EAction: HasteEnd})
	g.ComputeLOS()
	g.Print("You feel an energy burst and smoking coming out from you.")
}

func (g *game) Corrosion(ev event) {
	g.Player.Statuses[StatusCorrosion]++
	g.PushEvent(&simpleEvent{ERank: ev.Rank() + 80 + RandInt(40), EAction: CorrosionEnd})
	g.Print("Your equipment is corroded.")
}

func (g *game) Confusion(ev event) {
	if !g.Player.HasStatus(StatusConfusion) {
		g.Player.Statuses[StatusConfusion]++
		g.PushEvent(&simpleEvent{ERank: ev.Rank() + 100 + RandInt(100), EAction: ConfusionEnd})
		g.Print("You feel confused.")
	}
}

func (g *game) PlacePlayerAt(pos position) {
	g.Player.Pos = pos
	g.CollectGround()
	g.ComputeLOS()
	g.MakeMonstersAware()
}
<commit_msg>restore old value for wait turn boredom<commit_after>package main

import (
	"errors"
	"fmt"
)

type player struct {
	HP          int
	MP          int
	Simellas    int
	Armour      armour
	Weapon      weapon
	Shield      shield
	Consumables map[consumable]int
	Rods        map[rod]rodProps
	Aptitudes   map[aptitude]bool
	Statuses    map[status]int
	Pos         position
	Target      position
	LOS         map[position]bool
	Rays        rayMap
	Bored       int
}

func (p *player) HPMax() int {
	hpmax := 42
	if p.Aptitudes[AptHealthy] {
		hpmax += 10
	}
	hpmax -= 3 * p.Bored
	if hpmax < 21 {
		hpmax = 21
	}
	return hpmax
}

func (p *player) MPMax() int {
	mpmax := 3
	if p.Aptitudes[AptMagic] {
		mpmax += 2
	}
	if p.Armour == CelmistRobe {
		mpmax += 2
	}
	return mpmax
}

func (p *player) Accuracy() int {
	acc := 15
	if p.Aptitudes[AptAccurate] {
		acc += 2
	}
	return acc
}

func (p *player) RangedAccuracy() int {
	acc := 15
	if p.Aptitudes[AptAccurate] {
		acc += 10
	}
	return acc
}

func (p *player) Armor() int {
	ar := 0
	switch p.Armour {
	case LeatherArmour:
		ar += 3
	case ChainMail:
		ar += 6
	case PonderousnessPlates:
		ar += 10
	}
	if p.Aptitudes[AptScales] {
		ar += 2
	}
	if p.HasStatus(StatusLignification) {
		ar = 9 + ar/2
	}
	if p.HasStatus(StatusCorrosion) {
		ar -= 2 * p.Statuses[StatusCorrosion]
		if ar < 0 {
			ar = 0
		}
	}
	return ar
}

func (p *player) Attack() int {
	attack := p.Weapon.Attack()
	if p.Aptitudes[AptStrong] {
		attack += attack / 5
	}
	if p.HasStatus(StatusCorrosion) {
		penalty := p.Statuses[StatusCorrosion]
		if penalty > 5 {
			penalty = 5
		}
		attack -= penalty
	}
	return attack
}

func (p *player) Block() int {
	block := p.Shield.Block()
	if p.HasStatus(StatusDisabledShield) {
		block /= 3
	}
	return block
}

func (p *player) Evasion() int {
	ev := 15
	if p.Aptitudes[AptAgile] {
		ev += 3
	}
	switch p.Armour {
	case ChainMail:
		ev -= 1
	case PonderousnessPlates:
		ev -= 2
	}
	if p.HasStatus(StatusAgile) {
		ev += 7
	}
	return ev
}

func (p *player) HasStatus(st status) bool {
	return p.Statuses[st] > 0
}

func (p *player) AptitudeCount() int {
	count := 0
	for _, b := range p.Aptitudes {
		if b {
			count++
		}
	}
	return count
}

func (g *game) AutoToDir(ev event) bool {
	if g.MonsterInLOS() == nil {
		err := g.MovePlayer(g.Player.Pos.To(g.AutoDir), ev)
		if err != nil {
			g.Print(err.Error())
			g.AutoDir = NoDir
			return false
		}
		return true
	}
	g.AutoDir = NoDir
	return false
}

func (g *game) GoToDir(dir direction, ev event) error {
	if g.MonsterInLOS() != nil {
		g.AutoDir = NoDir
		return errors.New("You cannot travel while there are monsters in view.")
	}
	err := g.MovePlayer(g.Player.Pos.To(dir), ev)
	if err != nil {
		return err
	}
	g.AutoDir = dir
	return nil
}

func (g *game) MoveToTarget(ev event) bool {
	if !g.AutoTarget.valid() {
		return false
	}
	path := g.PlayerPath(g.Player.Pos, g.AutoTarget)
	if g.MonsterInLOS() != nil {
		g.AutoTarget = InvalidPos
	}
	if len(path) < 1 {
		g.AutoTarget = InvalidPos
		return false
	}
	var err error
	if len(path) > 1 {
		err = g.MovePlayer(path[len(path)-2], ev)
		if g.ExclusionsMap[path[len(path)-2]] {
			g.AutoTarget = InvalidPos
		}
	} else {
		g.WaitTurn(ev)
	}
	if err != nil {
		g.Print(err.Error())
		g.AutoTarget = InvalidPos
		return false
	}
	if g.AutoTarget.valid() && g.Player.Pos == g.AutoTarget {
		g.AutoTarget = InvalidPos
	}
	return true
}

func (g *game) WaitTurn(ev event) {
	// XXX Really wait for 10 ?
	grade := 1
	if len(g.Noise) > 0 || g.StatusRest() {
		grade = 1
	}
	g.BoredomAction(ev, grade)
	ev.Renew(g, 10)
}

func (g *game) MonsterCount() (count int) {
	for _, mons := range g.Monsters {
		if mons.Exists() {
			count++
		}
	}
	return count
}

func (g *game) BoredomAction(ev event, grade int) {
	obor := g.Boredom
	if g.MonsterInLOS() == nil {
		g.Boredom += grade
	} else {
		g.Boredom--
		if g.Boredom < 0 {
			g.Boredom = 0
			g.Player.Bored = 0
		}
		return
	}
	if g.Boredom >= 120 && obor < 120 {
		if g.MonsterCount() > 4 {
			g.PrintStyled("You feel a little bored, your health could decline.", logCritic)
			g.StopAuto()
		}
	}
	if g.Boredom >= 130 && (obor/10 != g.Boredom/10) {
		g.Player.Bored++
		g.PrintStyled("You feel unhealthy.", logCritic)
		g.StopAuto()
		if g.Player.HP > g.Player.HPMax() {
			g.Player.HP -= 3
		}
	}
}

func (g *game) FunAction() {
	g.Boredom -= 15
	if g.Boredom < 0 {
		g.Boredom = 0
		g.Player.Bored = 0
	}
}

func (g *game) Rest(ev event) error {
	if g.MonsterInLOS() != nil {
		return fmt.Errorf("You cannot sleep while monsters are in view.")
	}
	if cld, ok := g.Clouds[g.Player.Pos]; ok && cld == CloudFire {
		return errors.New("You cannot rest on flames.")
	}
	if !g.NeedsRegenRest() && !g.StatusRest() {
		return errors.New("You do not need to rest.")
	}
	g.WaitTurn(ev)
	g.Resting = true
	g.RestingTurns = 0
	if g.StatusRest() {
		g.RestingTurns = -1 // not true resting, just waiting for status end
	}
	g.FunAction()
	return nil
}

func (g *game) StatusRest() bool {
	for _, q := range g.Player.Statuses {
		if q > 0 {
			return true
		}
	}
	return false
}

func (g *game) NeedsRegenRest() bool {
	return g.Player.HP < g.Player.HPMax() || g.Player.MP < g.Player.MPMax()
}

func (g *game) Equip(ev event) error {
	if eq, ok := g.Equipables[g.Player.Pos]; ok {
		eq.Equip(g)
		ev.Renew(g, 10)
		return nil
	}
	return errors.New("Found nothing to equip here.")
}

func (g *game) Teleportation(ev event) {
	var pos position
	i := 0
	count := 0
	for {
		count++
		if count > 1000 {
			panic("Teleportation")
		}
		pos = g.FreeCell()
		if pos.Distance(g.Player.Pos) < 15 && i < 1000 {
			i++
			continue
		}
		break

	}
	if pos.valid() {
		// should always happen
		opos := g.Player.Pos
		g.Print("You feel yourself teleported away.")
		g.ui.TeleportAnimation(g, opos, pos, true)
		g.PlacePlayerAt(pos)
	} else {
		// should not happen
		g.Print("Something went wrong with the teleportation.")
	}
}

func (g *game) CollectGround() {
	pos := g.Player.Pos
	if g.Simellas[pos] > 0 {
		g.Player.Simellas += g.Simellas[pos]
		if g.Simellas[pos] == 1 {
			g.Print("You pick up a simella.")
		} else {
			g.Printf("You pick up %d simellas.", g.Simellas[pos])
		}
		g.DijkstraMapRebuild = true
		delete(g.Simellas, pos)
	}
	if c, ok := g.Collectables[pos]; ok {
		g.Player.Consumables[c.Consumable] += c.Quantity
		g.DijkstraMapRebuild = true
		delete(g.Collectables, pos)
		if c.Quantity > 1 {
			g.Printf("You take %d %s.", c.Quantity, c.Consumable.Plural())
		} else {
			g.Printf("You take %s.", Indefinite(c.Consumable.String(), false))
		}
	}
	if r, ok := g.Rods[pos]; ok {
		g.Player.Rods[r] = rodProps{Charge: r.MaxCharge() - 1}
		g.DijkstraMapRebuild = true
		delete(g.Rods, pos)
		g.Printf("You take a %s.", r)
		g.StoryPrintf("You found and took a %s.", r)
	}
	if eq, ok := g.Equipables[pos]; ok {
		g.Printf("You stand over %s.", Indefinite(eq.String(), false))
	} else if _, ok := g.Stairs[pos]; ok {
		g.Print("You stand over stairs.")
	} else if g.Doors[pos] {
		g.Print("You stand at the door.")
	}
}

func (g *game) MovePlayer(pos position, ev event) error {
	if !pos.valid() {
		return errors.New("You cannot move there.")
	}
	c := g.Dungeon.Cell(pos)
	if c.T == WallCell && !g.Player.HasStatus(StatusDig) {
		return errors.New("You cannot move into a wall.")
	}
	if g.Player.HasStatus(StatusConfusion) {
		switch pos.Dir(g.Player.Pos) {
		case E, N, W, S:
		default:
			return errors.New("You cannot use diagonal movements while confused.")
		}
	}
	delay := 10
	mons := g.MonsterAt(pos)
	if !mons.Exists() {
		if g.Player.HasStatus(StatusLignification) {
			return errors.New("You cannot move while lignified")
		}
		if c.T == WallCell {
			g.Dungeon.SetCell(pos, FreeCell)
			g.MakeNoise(WallNoise, pos)
			g.Print(g.CrackSound())
			g.Fog(pos, 1, ev)
			g.Stats.Digs++
		}
		g.PlacePlayerAt(pos)
		if !g.Autoexploring {
			g.BoredomAction(ev, 1)
		}
		if g.Player.Aptitudes[AptFast] {
			// only fast for movement
			delay -= 2
		}
		switch g.Player.Armour {
		case PonderousnessPlates:
			delay += 2
		case SpeedRobe:
			delay -= 3
		}
		if g.Player.HasStatus(StatusSwift) {
			// only fast for movement
			delay -= 3
		}
		g.Stats.Moves++
	} else {
		g.FunAction()
		g.AttackMonster(mons, ev)
	}
	if g.Player.HasStatus(StatusBerserk) {
		delay -= 3
	}
	if g.Player.HasStatus(StatusSlow) {
		delay += 3
	}
	if delay < 3 {
		delay = 3
	}
	ev.Renew(g, delay)
	return nil
}

func (g *game) HealPlayer(ev event) {
	if g.Player.HP < g.Player.HPMax() {
		g.Player.HP++
	}
	delay := 50
	ev.Renew(g, delay)
}

func (g *game) MPRegen(ev event) {
	if g.Player.MP < g.Player.MPMax() {
		g.Player.MP++
	}
	delay := 100
	ev.Renew(g, delay)
}

func (g *game) Smoke(ev event) {
	dij := &normalPath{game: g}
	nm := Dijkstra(dij, []position{g.Player.Pos}, 2)
	for pos := range nm {
		_, ok := g.Clouds[pos]
		if !ok {
			g.Clouds[pos] = CloudFog
			g.PushEvent(&cloudEvent{ERank: ev.Rank() + 100 + RandInt(100), EAction: CloudEnd, Pos: pos})
		}
	}
	g.Player.Statuses[StatusSwift]++
	g.PushEvent(&simpleEvent{ERank: ev.Rank() + 20 + RandInt(10), EAction: HasteEnd})
	g.ComputeLOS()
	g.Print("You feel an energy burst and smoking coming out from you.")
}

func (g *game) Corrosion(ev event) {
	g.Player.Statuses[StatusCorrosion]++
	g.PushEvent(&simpleEvent{ERank: ev.Rank() + 80 + RandInt(40), EAction: CorrosionEnd})
	g.Print("Your equipment is corroded.")
}

func (g *game) Confusion(ev event) {
	if !g.Player.HasStatus(StatusConfusion) {
		g.Player.Statuses[StatusConfusion]++
		g.PushEvent(&simpleEvent{ERank: ev.Rank() + 100 + RandInt(100), EAction: ConfusionEnd})
		g.Print("You feel confused.")
	}
}

func (g *game) PlacePlayerAt(pos position) {
	g.Player.Pos = pos
	g.CollectGround()
	g.ComputeLOS()
	g.MakeMonstersAware()
}
<|endoftext|>


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

package main

import (
	"bytes"
	"encoding/hex"
	"fmt"
	"github.com/jacobsa/comeback/blob"
	"github.com/jacobsa/comeback/disk"
	"github.com/jacobsa/comeback/fs"
	"github.com/jacobsa/comeback/repr"
	"io"
	"log"
	"os"
	"path"
	"syscall"
)

const (
	gTarget = "/tmp/restore_target"
)

var blobStore blob.Store

type score struct {
	hash []byte
}

func (s *score) Sha1Hash() []byte {
	return s.hash
}

func fromHexHash(h string) (blob.Score, error) {
	b, err := hex.DecodeString(h)
	if err != nil {
		return nil, fmt.Errorf("Invalid hex string: %s", h)
	}

	return &score{b}, nil
}

// Restore the file whose contents are described by the referenced blobs to the
// supplied target, whose parent must already exist.
func restoreFile(target string, scores []blob.Score) error {
	// Open the file.
	f, err := os.Create(target)
	defer f.Close()

	if err != nil {
		return fmt.Errorf("Create: %v", err)
	}

	// Process each blob.
	for _, score := range scores {
		// Load the blob.
		blob, err := blobStore.Load(score)
		if err != nil {
			return fmt.Errorf("Loading blob: %v", err)
		}

		// Write out its contents.
		_, err = io.Copy(f, bytes.NewReader(blob))
		if err != nil {
			return fmt.Errorf("Copy: %v", err)
		}
	}

	return nil
}

// Restore the directory whose contents are described by the referenced blob to
// the supplied target, which must already exist.
func restoreDir(target string, score blob.Score) error {
	// Load the appropriate blob.
	blob, err := blobStore.Load(score)
	if err != nil {
		return fmt.Errorf("Loading blob: %v", err)
	}

	// Parse its contents.
	entries, err := repr.Unmarshal(blob)
	if err != nil {
		return fmt.Errorf("Parsing blob: %v", err)
	}

	// Deal with each entry.
	for _, entry := range entries {
		entryPath := path.Join(target, entry.Name)

		// Switch on type.
		switch entry.Type {
		case fs.TypeFile:
			err = restoreFile(entryPath, entry.Scores)
			if err != nil {
				return err
			}

		case fs.TypeDirectory:
			if len(entry.Scores) != 1 {
				return fmt.Errorf("Wrong number of scores: %v", entry)
			}

			if err = os.Mkdir(entryPath, 0700); err != nil {
				return err
			}

			if err = restoreDir(entryPath, entry.Scores[0]); err != nil {
				return err
			}

		case fs.TypeSymlink:
			err = os.Symlink(entry.Target, entryPath)
			if err != nil {
				return err
			}

		case fs.TypeNamedPipe:
			err = makeNamedPipe(entryPath, entry.Permissions)
			if err != nil {
				return err
			}

		case fs.TypeBlockDevice:
			err = makeBlockDevice(entryPath, entry.Permissions, entry.Device)
			if err != nil {
				return err
			}

		default:
			return fmt.Errorf("Don't know how to deal with entry: %v", entry)
		}
	}

	return nil
}

func syscallPermissions(permissions os.FileMode) (o uint32) {
	// Include r/w/x permission bits.
	o = uint32(permissions & os.ModePerm)

	// Also include setuid/setgid/sticky bits.
	if permissions & os.ModeSetuid != 0 {
		o |= syscall.S_ISUID
	}

	if permissions & os.ModeSetgid != 0 {
		o |= syscall.S_ISGID
	}

	if permissions & os.ModeSticky != 0 {
		o |= syscall.S_ISVTX
	}

	return
}

// Create a named pipe at the supplied path.
func makeNamedPipe(path string, permissions os.FileMode) error {
	return syscall.Mkfifo(path, syscallPermissions(permissions))
}

// Create a block device at the supplied path.
func makeBlockDevice(path string, permissions os.FileMode, dev int32) error {
	if err := syscall.Mknod(path, syscallPermissions(permissions), int(dev)); err != nil {
		return fmt.Errorf("syscall.Mknod: %v", err)
	}

	return nil
}

func main() {
	var err error

	// Create the blob store.
	blobStore, err = disk.NewBlobStore("/tmp/blobs")
	if err != nil {
		log.Fatalf("Creating store: %v", err)
	}

	// Parse the score.
	score, err := fromHexHash("feb26f438b5706fe93281dd6d2c453eca9416c42")
	if err != nil {
		log.Fatalf("Parsing score: %v", err)
	}

	// Make sure the target doesn't exist.
	err = os.RemoveAll(gTarget)
	if err != nil {
		log.Fatalf("RemoveAll: %v", err)
	}

	// Create the target.
	err = os.Mkdir("/tmp/restore_target", 0700)
	if err != nil {
		log.Fatalf("Mkdir: %v", err)
	}

	// Attempt a restore.
	err = restoreDir(gTarget, score)
	if err != nil {
		log.Fatalf("Restoring: %v", err)
	}
}
<commit_msg>Updated hash.<commit_after>// Copyright 2012 Aaron Jacobs. All Rights Reserved.
// Author: aaronjjacobs@gmail.com (Aaron Jacobs)
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package main

import (
	"bytes"
	"encoding/hex"
	"fmt"
	"github.com/jacobsa/comeback/blob"
	"github.com/jacobsa/comeback/disk"
	"github.com/jacobsa/comeback/fs"
	"github.com/jacobsa/comeback/repr"
	"io"
	"log"
	"os"
	"path"
	"syscall"
)

const (
	gTarget = "/tmp/restore_target"
)

var blobStore blob.Store

type score struct {
	hash []byte
}

func (s *score) Sha1Hash() []byte {
	return s.hash
}

func fromHexHash(h string) (blob.Score, error) {
	b, err := hex.DecodeString(h)
	if err != nil {
		return nil, fmt.Errorf("Invalid hex string: %s", h)
	}

	return &score{b}, nil
}

// Restore the file whose contents are described by the referenced blobs to the
// supplied target, whose parent must already exist.
func restoreFile(target string, scores []blob.Score) error {
	// Open the file.
	f, err := os.Create(target)
	defer f.Close()

	if err != nil {
		return fmt.Errorf("Create: %v", err)
	}

	// Process each blob.
	for _, score := range scores {
		// Load the blob.
		blob, err := blobStore.Load(score)
		if err != nil {
			return fmt.Errorf("Loading blob: %v", err)
		}

		// Write out its contents.
		_, err = io.Copy(f, bytes.NewReader(blob))
		if err != nil {
			return fmt.Errorf("Copy: %v", err)
		}
	}

	return nil
}

// Restore the directory whose contents are described by the referenced blob to
// the supplied target, which must already exist.
func restoreDir(target string, score blob.Score) error {
	// Load the appropriate blob.
	blob, err := blobStore.Load(score)
	if err != nil {
		return fmt.Errorf("Loading blob: %v", err)
	}

	// Parse its contents.
	entries, err := repr.Unmarshal(blob)
	if err != nil {
		return fmt.Errorf("Parsing blob: %v", err)
	}

	// Deal with each entry.
	for _, entry := range entries {
		entryPath := path.Join(target, entry.Name)

		// Switch on type.
		switch entry.Type {
		case fs.TypeFile:
			err = restoreFile(entryPath, entry.Scores)
			if err != nil {
				return err
			}

		case fs.TypeDirectory:
			if len(entry.Scores) != 1 {
				return fmt.Errorf("Wrong number of scores: %v", entry)
			}

			if err = os.Mkdir(entryPath, 0700); err != nil {
				return err
			}

			if err = restoreDir(entryPath, entry.Scores[0]); err != nil {
				return err
			}

		case fs.TypeSymlink:
			err = os.Symlink(entry.Target, entryPath)
			if err != nil {
				return err
			}

		case fs.TypeNamedPipe:
			err = makeNamedPipe(entryPath, entry.Permissions)
			if err != nil {
				return err
			}

		case fs.TypeBlockDevice:
			err = makeBlockDevice(entryPath, entry.Permissions, entry.Device)
			if err != nil {
				return err
			}

		default:
			return fmt.Errorf("Don't know how to deal with entry: %v", entry)
		}
	}

	return nil
}

func syscallPermissions(permissions os.FileMode) (o uint32) {
	// Include r/w/x permission bits.
	o = uint32(permissions & os.ModePerm)

	// Also include setuid/setgid/sticky bits.
	if permissions & os.ModeSetuid != 0 {
		o |= syscall.S_ISUID
	}

	if permissions & os.ModeSetgid != 0 {
		o |= syscall.S_ISGID
	}

	if permissions & os.ModeSticky != 0 {
		o |= syscall.S_ISVTX
	}

	return
}

// Create a named pipe at the supplied path.
func makeNamedPipe(path string, permissions os.FileMode) error {
	return syscall.Mkfifo(path, syscallPermissions(permissions))
}

// Create a block device at the supplied path.
func makeBlockDevice(path string, permissions os.FileMode, dev int32) error {
	if err := syscall.Mknod(path, syscallPermissions(permissions), int(dev)); err != nil {
		return fmt.Errorf("syscall.Mknod: %v", err)
	}

	return nil
}

func main() {
	var err error

	// Create the blob store.
	blobStore, err = disk.NewBlobStore("/tmp/blobs")
	if err != nil {
		log.Fatalf("Creating store: %v", err)
	}

	// Parse the score.
	score, err := fromHexHash("697e002f699399f859514f621953c64e14a0e110")
	if err != nil {
		log.Fatalf("Parsing score: %v", err)
	}

	// Make sure the target doesn't exist.
	err = os.RemoveAll(gTarget)
	if err != nil {
		log.Fatalf("RemoveAll: %v", err)
	}

	// Create the target.
	err = os.Mkdir("/tmp/restore_target", 0700)
	if err != nil {
		log.Fatalf("Mkdir: %v", err)
	}

	// Attempt a restore.
	err = restoreDir(gTarget, score)
	if err != nil {
		log.Fatalf("Restoring: %v", err)
	}
}
<|endoftext|>


<commit_before><commit_msg>Update to proper slice element type<commit_after><|endoftext|>


<commit_before>// Package jsutil provides various convenience functions for Javascript environment
package jsutil

import (
	"io/ioutil"

	"github.com/robertkrimen/otto"
)

func ConfigureRequire(vm *otto.Otto) {
	vm.Set("require", func(call otto.FunctionCall) otto.Value {

		const requireHeader = `(function() {
    var module = { exports: {} };
    var exports = module.exports;`

		const requireFooter = `   return module.exports;
}());`
		path, err := call.Argument(0).ToString()
		if err != nil {
			return otto.UndefinedValue()
		}

		body, err := ioutil.ReadFile(path)
		if err != nil {
			return otto.UndefinedValue()
		}

		val, err := vm.Run(requireHeader + string(body) + requireFooter)
		if err != nil {
			return otto.UndefinedValue()
		}

		return val
	})
}
<commit_msg>Able to load JS modules from $NODE_PATH<commit_after>// Package jsutil provides various convenience functions for Javascript environment
package jsutil

import (
	"io/ioutil"
	"os"
	"path/filepath"

	"github.com/robertkrimen/otto"
)

func ConfigureRequire(vm *otto.Otto) {
	vm.Set("require", func(call otto.FunctionCall) otto.Value {

		const requireHeader = `(function() {
    var module = { exports: {} };
    var exports = module.exports;`

		const requireFooter = `   return module.exports;
}());`
		modulePath, err := call.Argument(0).ToString()
		if err != nil {
			return otto.UndefinedValue()
		}

		ableToReadModule := false

		// Check current directory to load module
		body, err := ioutil.ReadFile(modulePath)
		if err == nil {
			ableToReadModule = true
		}

		if !ableToReadModule {
			// Check $NODE_PATH to load module
			nodePaths := os.Getenv("NODE_PATH")
			if nodePaths != "" {
				for _, nodePath := range filepath.SplitList(nodePaths) {
					body, err = ioutil.ReadFile(filepath.Join(nodePath, modulePath))
					if err == nil {
						ableToReadModule = true
						break
					}
				}
			}
		}

		if !ableToReadModule {
			return otto.UndefinedValue()
		}

		val, err := vm.Run(requireHeader + string(body) + requireFooter)
		if err != nil {
			return otto.UndefinedValue()
		}

		return val
	})
}
<|endoftext|>


<commit_before>package chained

import (
	"fmt"
	"io"
	"net"
	"net/http"
)

// Server provides the upstream side of a chained proxy setup. It can be run as
// a standalone HTTP server using Serve() or plugged into an existing HTTP
// server as an http.Handler.
type Server struct {
	// Dial: function for dialing destination
	Dial func(network, address string) (net.Conn, error)
}

// Serve provides a convenience function for starting an HTTP server using this
// Server as the Handler.
func (s *Server) Serve(l net.Listener) error {
	server := http.Server{
		Handler: s,
	}
	return server.Serve(l)
}

// ServeHTTP implements the method from http.Handler.
func (s *Server) ServeHTTP(resp http.ResponseWriter, req *http.Request) {
	hj, ok := resp.(http.Hijacker)
	if !ok {
		panic("Response doesn't allow hijacking!")
	}

	fl, ok := resp.(http.Flusher)
	if !ok {
		panic("Response doesn't allow flushing!")
	}

	if req.Method != CONNECT {
		resp.WriteHeader(405)
		fmt.Fprintf(resp, "Method %s not allowed", req.Method)
		return
	}

	address := req.Host
	connOut, err := s.Dial("tcp", address)
	if err != nil {
		resp.WriteHeader(502)
		fmt.Fprintf(resp, "Unable to dial %s : %s", address, err)
		return
	}
	resp.WriteHeader(200)
	fmt.Fprint(resp, "CONNECT OK")
	fl.Flush()

	connIn, _, err := hj.Hijack()
	if err != nil {
		log.Errorf("Unable to hijack connection: %s", err)
		connOut.Close()
		return
	}

	go io.Copy(connOut, connIn)
	go io.Copy(connIn, connOut)
}
<commit_msg>Making sure to close connIn and connOut in chained Server<commit_after>package chained

import (
	"fmt"
	"io"
	"net"
	"net/http"
	"sync"
)

// Server provides the upstream side of a chained proxy setup. It can be run as
// a standalone HTTP server using Serve() or plugged into an existing HTTP
// server as an http.Handler.
type Server struct {
	// Dial: function for dialing destination
	Dial func(network, address string) (net.Conn, error)
}

// Serve provides a convenience function for starting an HTTP server using this
// Server as the Handler.
func (s *Server) Serve(l net.Listener) error {
	server := http.Server{
		Handler: s,
	}
	return server.Serve(l)
}

// ServeHTTP implements the method from http.Handler.
func (s *Server) ServeHTTP(resp http.ResponseWriter, req *http.Request) {
	hj, ok := resp.(http.Hijacker)
	if !ok {
		panic("Response doesn't allow hijacking!")
	}

	fl, ok := resp.(http.Flusher)
	if !ok {
		panic("Response doesn't allow flushing!")
	}

	if req.Method != CONNECT {
		resp.WriteHeader(405)
		fmt.Fprintf(resp, "Method %s not allowed", req.Method)
		return
	}

	address := req.Host
	connOut, err := s.Dial("tcp", address)
	if err != nil {
		resp.WriteHeader(502)
		fmt.Fprintf(resp, "Unable to dial %s : %s", address, err)
		return
	}
	defer connOut.Close()
	resp.WriteHeader(200)
	fmt.Fprint(resp, "CONNECT OK")
	fl.Flush()

	connIn, _, err := hj.Hijack()
	if err != nil {
		log.Errorf("Unable to hijack connection: %s", err)
		return
	}
	defer connIn.Close()

	var wg sync.WaitGroup
	wg.Add(2)
	go func() {
		io.Copy(connOut, connIn)
		wg.Done()
	}()
	go func() {
		go io.Copy(connIn, connOut)
		wg.Done()
	}()
	wg.Wait()
}
<|endoftext|>


<commit_before>package cli

import (
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"runtime"
	"strings"
	"time"

	"github.com/schollz/croc/v6/src/croc"
	"github.com/schollz/croc/v6/src/tcp"
	"github.com/schollz/croc/v6/src/utils"
	"github.com/urfave/cli"
)

var Version string

func Run() (err error) {
	// use all of the processors
	runtime.GOMAXPROCS(runtime.NumCPU())

	app := cli.NewApp()
	app.Name = "croc"
	if Version == "" {
		Version = "v6.0.0-cb69e49"
	}
	app.Version = Version
	app.Compiled = time.Now()
	app.Usage = "easily and securely transfer stuff from one computer to another"
	app.UsageText = "croc allows any two computers to directly and securely transfer files"
	// app.ArgsUsage = "[args and such]"
	app.Commands = []cli.Command{
		{
			Name:        "send",
			Usage:       "send a file",
			Description: "send a file over the relay",
			ArgsUsage:   "[filename]",
			Flags: []cli.Flag{
				cli.StringFlag{Name: "code, c", Usage: "codephrase used to connect to relay"},
				cli.BoolFlag{Name: "no-local", Usage: "disable local relay when sending"},
				cli.StringFlag{Name: "ports", Value: "9009,9010,9011,9012,9013", Usage: "ports of the local relay (optional)"},
			},
			HelpName: "croc send",
			Action: func(c *cli.Context) error {
				return send(c)
			},
		},
		{
			Name:        "relay",
			Description: "start relay",
			HelpName:    "croc relay",
			Action: func(c *cli.Context) error {
				return relay(c)
			},
			Flags: []cli.Flag{
				cli.StringFlag{Name: "ports", Value: "9009,9010,9011,9012,9013", Usage: "ports of the relay"},
			},
		},
	}
	app.Flags = []cli.Flag{
		cli.BoolFlag{Name: "debug", Usage: "increase verbosity (a lot)"},
		cli.BoolFlag{Name: "yes", Usage: "automatically agree to all prompts"},
		cli.BoolFlag{Name: "stdout", Usage: "redirect file to stdout"},
		cli.StringFlag{Name: "relay", Value: "198.199.67.130:9009", Usage: "address of the relay"},
		cli.StringFlag{Name: "out", Value: ".", Usage: "specify an output folder to receive the file"},
	}
	app.EnableBashCompletion = true
	app.HideHelp = false
	app.HideVersion = false
	app.BashComplete = func(c *cli.Context) {
		fmt.Fprintf(c.App.Writer, "send\nreceive\relay")
	}
	app.Action = func(c *cli.Context) error {
		// if trying to send but forgot send, let the user know
		if c.Args().First() != "" && utils.Exists(c.Args().First()) {
			_, fname := filepath.Split(c.Args().First())
			yn := utils.GetInput(fmt.Sprintf("Did you mean to send '%s'? (y/n) ", fname))
			if strings.ToLower(yn) == "y" {
				return send(c)
			}
		}
		return receive(c)
	}

	return app.Run(os.Args)
}

// func saveDefaultConfig(c *cli.Context) error {
// 	return croc.SaveDefaultConfig()
// }

func send(c *cli.Context) (err error) {

	var fnames []string
	stat, _ := os.Stdin.Stat()
	if (stat.Mode() & os.ModeCharDevice) == 0 {
		f, err := ioutil.TempFile(".", "croc-stdin-")
		if err != nil {
			return err
		}
		_, err = io.Copy(f, os.Stdin)
		if err != nil {
			return err
		}
		err = f.Close()
		if err != nil {
			return err
		}
		fnames = []string{f.Name()}
		defer func() {
			err = os.Remove(fnames[0])
			if err != nil {
				log.Println(err)
			}
		}()
	} else {
		fnames = append([]string{c.Args().First()}, c.Args().Tail()...)
	}
	if len(fnames) == 0 {
		return errors.New("must specify file: croc send [filename]")
	}

	var sharedSecret string
	if c.String("code") != "" {
		sharedSecret = c.String("code")
	}
	// cr.LoadConfig()
	if len(sharedSecret) == 0 {
		// generate code phrase
		sharedSecret = utils.GetRandomName()
	}

	haveFolder := false
	paths := []string{}
	for _, fname := range fnames {
		stat, err := os.Stat(fname)
		if err != nil {
			return err
		}
		if stat.IsDir() {
			haveFolder = true
			err = filepath.Walk(fname,
				func(pathName string, info os.FileInfo, err error) error {
					if err != nil {
						return err
					}
					if !info.IsDir() {
						paths = append(paths, filepath.ToSlash(pathName))
					}
					return nil
				})
			if err != nil {
				return err
			}
		} else {
			paths = append(paths, filepath.ToSlash(fname))
		}
	}
	cr, err := croc.New(croc.Options{
		SharedSecret: sharedSecret,
		IsSender:     true,
		Debug:        c.GlobalBool("debug"),
		NoPrompt:     c.GlobalBool("yes"),
		RelayAddress: c.GlobalString("relay"),
		Stdout:       c.GlobalBool("stdout"),
		DisableLocal: c.Bool("no-local"),
		RelayPorts:   strings.Split(c.String("ports"), ","),
	})
	if err != nil {
		return
	}

	err = cr.Send(croc.TransferOptions{
		PathToFiles:      paths,
		KeepPathInRemote: haveFolder,
	})

	return
}

func receive(c *cli.Context) (err error) {
	var sharedSecret string
	if c.GlobalString("code") != "" {
		sharedSecret = c.GlobalString("code")
	}
	if c.Args().First() != "" {
		sharedSecret = c.Args().First()
	}
	if sharedSecret == "" {
		sharedSecret = utils.GetInput("Enter receive code: ")
	}
	if c.GlobalString("out") != "" {
		os.Chdir(c.GlobalString("out"))
	}

	cr, err := croc.New(croc.Options{
		SharedSecret: sharedSecret,
		IsSender:     false,
		Debug:        c.GlobalBool("debug"),
		NoPrompt:     c.GlobalBool("yes"),
		RelayAddress: c.GlobalString("relay"),
		Stdout:       c.GlobalBool("stdout"),
	})
	if err != nil {
		return
	}
	err = cr.Receive()
	return
}

func relay(c *cli.Context) (err error) {
	debugString := "info"
	if c.GlobalBool("debug") {
		debugString = "debug"
	}
	ports := strings.Split(c.String("ports"), ",")
	tcpPorts := strings.Join(ports[1:], ",")
	for i, port := range ports {
		if i == 0 {
			continue
		}
		go func(portStr string) {
			err = tcp.Run(debugString, portStr)
			if err != nil {
				panic(err)
			}
		}(port)
	}
	return tcp.Run(debugString, ports[0], tcpPorts)
}

// func dirSize(path string) (int64, error) {
// 	var size int64
// 	err := filepath.Walk(path, func(_ string, info os.FileInfo, err error) error {
// 		if !info.IsDir() {
// 			size += info.Size()
// 		}
// 		return err
// 	})
// 	return size, err
// }
<commit_msg>update hash<commit_after>package cli

import (
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"runtime"
	"strings"
	"time"

	"github.com/schollz/croc/v6/src/croc"
	"github.com/schollz/croc/v6/src/tcp"
	"github.com/schollz/croc/v6/src/utils"
	"github.com/urfave/cli"
)

var Version string

func Run() (err error) {
	// use all of the processors
	runtime.GOMAXPROCS(runtime.NumCPU())

	app := cli.NewApp()
	app.Name = "croc"
	if Version == "" {
		Version = "v6.0.0-bd9c7b5"
	}
	app.Version = Version
	app.Compiled = time.Now()
	app.Usage = "easily and securely transfer stuff from one computer to another"
	app.UsageText = "croc allows any two computers to directly and securely transfer files"
	// app.ArgsUsage = "[args and such]"
	app.Commands = []cli.Command{
		{
			Name:        "send",
			Usage:       "send a file",
			Description: "send a file over the relay",
			ArgsUsage:   "[filename]",
			Flags: []cli.Flag{
				cli.StringFlag{Name: "code, c", Usage: "codephrase used to connect to relay"},
				cli.BoolFlag{Name: "no-local", Usage: "disable local relay when sending"},
				cli.StringFlag{Name: "ports", Value: "9009,9010,9011,9012,9013", Usage: "ports of the local relay (optional)"},
			},
			HelpName: "croc send",
			Action: func(c *cli.Context) error {
				return send(c)
			},
		},
		{
			Name:        "relay",
			Description: "start relay",
			HelpName:    "croc relay",
			Action: func(c *cli.Context) error {
				return relay(c)
			},
			Flags: []cli.Flag{
				cli.StringFlag{Name: "ports", Value: "9009,9010,9011,9012,9013", Usage: "ports of the relay"},
			},
		},
	}
	app.Flags = []cli.Flag{
		cli.BoolFlag{Name: "debug", Usage: "increase verbosity (a lot)"},
		cli.BoolFlag{Name: "yes", Usage: "automatically agree to all prompts"},
		cli.BoolFlag{Name: "stdout", Usage: "redirect file to stdout"},
		cli.StringFlag{Name: "relay", Value: "198.199.67.130:9009", Usage: "address of the relay"},
		cli.StringFlag{Name: "out", Value: ".", Usage: "specify an output folder to receive the file"},
	}
	app.EnableBashCompletion = true
	app.HideHelp = false
	app.HideVersion = false
	app.BashComplete = func(c *cli.Context) {
		fmt.Fprintf(c.App.Writer, "send\nreceive\relay")
	}
	app.Action = func(c *cli.Context) error {
		// if trying to send but forgot send, let the user know
		if c.Args().First() != "" && utils.Exists(c.Args().First()) {
			_, fname := filepath.Split(c.Args().First())
			yn := utils.GetInput(fmt.Sprintf("Did you mean to send '%s'? (y/n) ", fname))
			if strings.ToLower(yn) == "y" {
				return send(c)
			}
		}
		return receive(c)
	}

	return app.Run(os.Args)
}

// func saveDefaultConfig(c *cli.Context) error {
// 	return croc.SaveDefaultConfig()
// }

func send(c *cli.Context) (err error) {

	var fnames []string
	stat, _ := os.Stdin.Stat()
	if (stat.Mode() & os.ModeCharDevice) == 0 {
		f, err := ioutil.TempFile(".", "croc-stdin-")
		if err != nil {
			return err
		}
		_, err = io.Copy(f, os.Stdin)
		if err != nil {
			return err
		}
		err = f.Close()
		if err != nil {
			return err
		}
		fnames = []string{f.Name()}
		defer func() {
			err = os.Remove(fnames[0])
			if err != nil {
				log.Println(err)
			}
		}()
	} else {
		fnames = append([]string{c.Args().First()}, c.Args().Tail()...)
	}
	if len(fnames) == 0 {
		return errors.New("must specify file: croc send [filename]")
	}

	var sharedSecret string
	if c.String("code") != "" {
		sharedSecret = c.String("code")
	}
	// cr.LoadConfig()
	if len(sharedSecret) == 0 {
		// generate code phrase
		sharedSecret = utils.GetRandomName()
	}

	haveFolder := false
	paths := []string{}
	for _, fname := range fnames {
		stat, err := os.Stat(fname)
		if err != nil {
			return err
		}
		if stat.IsDir() {
			haveFolder = true
			err = filepath.Walk(fname,
				func(pathName string, info os.FileInfo, err error) error {
					if err != nil {
						return err
					}
					if !info.IsDir() {
						paths = append(paths, filepath.ToSlash(pathName))
					}
					return nil
				})
			if err != nil {
				return err
			}
		} else {
			paths = append(paths, filepath.ToSlash(fname))
		}
	}
	cr, err := croc.New(croc.Options{
		SharedSecret: sharedSecret,
		IsSender:     true,
		Debug:        c.GlobalBool("debug"),
		NoPrompt:     c.GlobalBool("yes"),
		RelayAddress: c.GlobalString("relay"),
		Stdout:       c.GlobalBool("stdout"),
		DisableLocal: c.Bool("no-local"),
		RelayPorts:   strings.Split(c.String("ports"), ","),
	})
	if err != nil {
		return
	}

	err = cr.Send(croc.TransferOptions{
		PathToFiles:      paths,
		KeepPathInRemote: haveFolder,
	})

	return
}

func receive(c *cli.Context) (err error) {
	var sharedSecret string
	if c.GlobalString("code") != "" {
		sharedSecret = c.GlobalString("code")
	}
	if c.Args().First() != "" {
		sharedSecret = c.Args().First()
	}
	if sharedSecret == "" {
		sharedSecret = utils.GetInput("Enter receive code: ")
	}
	if c.GlobalString("out") != "" {
		os.Chdir(c.GlobalString("out"))
	}

	cr, err := croc.New(croc.Options{
		SharedSecret: sharedSecret,
		IsSender:     false,
		Debug:        c.GlobalBool("debug"),
		NoPrompt:     c.GlobalBool("yes"),
		RelayAddress: c.GlobalString("relay"),
		Stdout:       c.GlobalBool("stdout"),
	})
	if err != nil {
		return
	}
	err = cr.Receive()
	return
}

func relay(c *cli.Context) (err error) {
	debugString := "info"
	if c.GlobalBool("debug") {
		debugString = "debug"
	}
	ports := strings.Split(c.String("ports"), ",")
	tcpPorts := strings.Join(ports[1:], ",")
	for i, port := range ports {
		if i == 0 {
			continue
		}
		go func(portStr string) {
			err = tcp.Run(debugString, portStr)
			if err != nil {
				panic(err)
			}
		}(port)
	}
	return tcp.Run(debugString, ports[0], tcpPorts)
}

// func dirSize(path string) (int64, error) {
// 	var size int64
// 	err := filepath.Walk(path, func(_ string, info os.FileInfo, err error) error {
// 		if !info.IsDir() {
// 			size += info.Size()
// 		}
// 		return err
// 	})
// 	return size, err
// }
<|endoftext|>