<commit_before>// Copyright 2020 PingCAP, Inc.
//
// 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,
// See the License for the specific language governing permissions and
// limitations under the License.

package tikv

import (
	"math/rand"
	"strings"
	"sync/atomic"
	"time"

	"github.com/pingcap/errors"
	"github.com/pingcap/failpoint"
	pb "github.com/pingcap/kvproto/pkg/kvrpcpb"
	"github.com/pingcap/tidb/kv"
	"github.com/pingcap/tidb/metrics"
	"github.com/pingcap/tidb/store/tikv/tikvrpc"
	"github.com/pingcap/tidb/util/logutil"
	"github.com/prometheus/client_golang/prometheus"
	"go.uber.org/zap"
)

type actionPessimisticLock struct {
	*kv.LockCtx
}
type actionPessimisticRollback struct{}

var (
	_ twoPhaseCommitAction = actionPessimisticLock{}
	_ twoPhaseCommitAction = actionPessimisticRollback{}

	tiKVTxnRegionsNumHistogramPessimisticLock     = metrics.TiKVTxnRegionsNumHistogram.WithLabelValues(metricsTag("pessimistic_lock"))
	tiKVTxnRegionsNumHistogramPessimisticRollback = metrics.TiKVTxnRegionsNumHistogram.WithLabelValues(metricsTag("pessimistic_rollback"))
)

func (actionPessimisticLock) String() string {
	return "pessimistic_lock"
}

func (actionPessimisticLock) tiKVTxnRegionsNumHistogram() prometheus.Observer {
	return tiKVTxnRegionsNumHistogramPessimisticLock
}

func (actionPessimisticRollback) String() string {
	return "pessimistic_rollback"
}

func (actionPessimisticRollback) tiKVTxnRegionsNumHistogram() prometheus.Observer {
	return tiKVTxnRegionsNumHistogramPessimisticRollback
}

func (action actionPessimisticLock) handleSingleBatch(c *twoPhaseCommitter, bo *Backoffer, batch batchMutations) error {
	m := batch.mutations
	mutations := make([]*pb.Mutation, m.Len())
	for i := 0; i < m.Len(); i++ {
		mut := &pb.Mutation{
			Op:  pb.Op_PessimisticLock,
			Key: m.GetKey(i),
		}
		if c.txn.us.HasPresumeKeyNotExists(m.GetKey(i)) || (c.doingAmend && m.GetOp(i) == pb.Op_Insert) {
			mut.Assertion = pb.Assertion_NotExist
		}
		mutations[i] = mut
	}
	elapsed := uint64(time.Since(c.txn.startTime) / time.Millisecond)
	req := tikvrpc.NewRequest(tikvrpc.CmdPessimisticLock, &pb.PessimisticLockRequest{
		Mutations:    mutations,
		PrimaryLock:  c.primary(),
		StartVersion: c.startTS,
		ForUpdateTs:  c.forUpdateTS,
		LockTtl:      elapsed + atomic.LoadUint64(&ManagedLockTTL),
		IsFirstLock:  c.isFirstLock,
		WaitTimeout:  action.LockWaitTime,
		ReturnValues: action.ReturnValues,
		MinCommitTs:  c.forUpdateTS + 1,
	}, pb.Context{Priority: c.priority, SyncLog: c.syncLog})
	lockWaitStartTime := action.WaitStartTime
	for {
		// if lockWaitTime set, refine the request `WaitTimeout` field based on timeout limit
		if action.LockWaitTime > 0 {
			timeLeft := action.LockWaitTime - (time.Since(lockWaitStartTime)).Milliseconds()
			if timeLeft <= 0 {
				req.PessimisticLock().WaitTimeout = kv.LockNoWait
			} else {
				req.PessimisticLock().WaitTimeout = timeLeft
			}
		}
		failpoint.Inject("PessimisticLockErrWriteConflict", func() error {
			time.Sleep(300 * time.Millisecond)
			return kv.ErrWriteConflict
		})
		startTime := time.Now()
		resp, err := c.store.SendReq(bo, req, batch.region, readTimeoutShort)
		if action.LockCtx.Stats != nil {
			atomic.AddInt64(&action.LockCtx.Stats.LockRPCTime, int64(time.Since(startTime)))
			atomic.AddInt64(&action.LockCtx.Stats.LockRPCCount, 1)
		}
		if err != nil {
			return errors.Trace(err)
		}
		regionErr, err := resp.GetRegionError()
		if err != nil {
			return errors.Trace(err)
		}
		if regionErr != nil {
			err = bo.Backoff(BoRegionMiss, errors.New(regionErr.String()))
			if err != nil {
				return errors.Trace(err)
			}
			err = c.pessimisticLockMutations(bo, action.LockCtx, batch.mutations)
			return errors.Trace(err)
		}
		if resp.Resp == nil {
			return errors.Trace(ErrBodyMissing)
		}
		lockResp := resp.Resp.(*pb.PessimisticLockResponse)
		keyErrs := lockResp.GetErrors()
		if len(keyErrs) == 0 {
			if action.ReturnValues {
				action.ValuesLock.Lock()
				for i, mutation := range mutations {
					action.Values[string(mutation.Key)] = kv.ReturnedValue{Value: lockResp.Values[i]}
				}
				action.ValuesLock.Unlock()
			}
			return nil
		}
		var locks []*Lock
		for _, keyErr := range keyErrs {
			// Check already exists error
			if alreadyExist := keyErr.GetAlreadyExist(); alreadyExist != nil {
				key := alreadyExist.GetKey()
				return c.extractKeyExistsErr(key)
			}
			if deadlock := keyErr.Deadlock; deadlock != nil {
				return &ErrDeadlock{Deadlock: deadlock}
			}

			// Extract lock from key error
			lock, err1 := extractLockFromKeyErr(keyErr)
			if err1 != nil {
				return errors.Trace(err1)
			}
			locks = append(locks, lock)
		}
		// Because we already waited on tikv, no need to Backoff here.
		// tikv default will wait 3s(also the maximum wait value) when lock error occurs
		startTime = time.Now()
		msBeforeTxnExpired, _, err := c.store.lockResolver.ResolveLocks(bo, 0, locks)
		if err != nil {
			return errors.Trace(err)
		}
		if action.LockCtx.Stats != nil {
			atomic.AddInt64(&action.LockCtx.Stats.ResolveLockTime, int64(time.Since(startTime)))
		}

		// If msBeforeTxnExpired is not zero, it means there are still locks blocking us acquiring
		// the pessimistic lock. We should return acquire fail with nowait set or timeout error if necessary.
		if msBeforeTxnExpired > 0 {
			if action.LockWaitTime == kv.LockNoWait {
				return ErrLockAcquireFailAndNoWaitSet
			} else if action.LockWaitTime == kv.LockAlwaysWait {
				// do nothing but keep wait
			} else {
				// the lockWaitTime is set, we should return wait timeout if we are still blocked by a lock
				if time.Since(lockWaitStartTime).Milliseconds() >= action.LockWaitTime {
					return errors.Trace(ErrLockWaitTimeout)
				}
			}
			if action.LockCtx.PessimisticLockWaited != nil {
				atomic.StoreInt32(action.LockCtx.PessimisticLockWaited, 1)
			}
		}

		// Handle the killed flag when waiting for the pessimistic lock.
		// When a txn runs into LockKeys() and backoff here, it has no chance to call
		// executor.Next() and check the killed flag.
		if action.Killed != nil {
			// Do not reset the killed flag here!
			// actionPessimisticLock runs on each region parallelly, we have to consider that
			// the error may be dropped.
			if atomic.LoadUint32(action.Killed) == 1 {
				return errors.Trace(ErrQueryInterrupted)
			}
		}
	}
}

func (actionPessimisticRollback) handleSingleBatch(c *twoPhaseCommitter, bo *Backoffer, batch batchMutations) error {
	req := tikvrpc.NewRequest(tikvrpc.CmdPessimisticRollback, &pb.PessimisticRollbackRequest{
		StartVersion: c.startTS,
		ForUpdateTs:  c.forUpdateTS,
		Keys:         batch.mutations.GetKeys(),
	})
	resp, err := c.store.SendReq(bo, req, batch.region, readTimeoutShort)
	if err != nil {
		return errors.Trace(err)
	}
	regionErr, err := resp.GetRegionError()
	if err != nil {
		return errors.Trace(err)
	}
	if regionErr != nil {
		err = bo.Backoff(BoRegionMiss, errors.New(regionErr.String()))
		if err != nil {
			return errors.Trace(err)
		}
		err = c.pessimisticRollbackMutations(bo, batch.mutations)
		return errors.Trace(err)
	}
	return nil
}

func (c *twoPhaseCommitter) pessimisticLockMutations(bo *Backoffer, lockCtx *kv.LockCtx, mutations CommitterMutations) error {
	if c.connID > 0 {
		failpoint.Inject("beforePessimisticLock", func(val failpoint.Value) {
			// Pass multiple instructions in one string, delimited by commas, to trigger multiple behaviors, like
			// `return("delay,fail")`. Then they will be executed sequentially at once.
			if v, ok := val.(string); ok {
				for _, action := range strings.Split(v, ",") {
					if action == "delay" {
						duration := time.Duration(rand.Int63n(int64(time.Second) * 5))
						logutil.Logger(bo.ctx).Info("[failpoint] injected delay at pessimistic lock",
							zap.Uint64("txnStartTS", c.startTS), zap.Duration("duration", duration))
						time.Sleep(duration)
					} else if action == "fail" {
						logutil.Logger(bo.ctx).Info("[failpoint] injected failure at pessimistic lock",
							zap.Uint64("txnStartTS", c.startTS))
						failpoint.Return(errors.New("injected failure at pessimistic lock"))
					}
				}
			}
		})
	}
	return c.doActionOnMutations(bo, actionPessimisticLock{lockCtx}, mutations)
}

func (c *twoPhaseCommitter) pessimisticRollbackMutations(bo *Backoffer, mutations CommitterMutations) error {
	return c.doActionOnMutations(bo, actionPessimisticRollback{}, mutations)
}
<commit_msg>store/tikv: set short pessimistic lock TTL through failpoints (#22431)<commit_after>// Copyright 2020 PingCAP, Inc.
//
// 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,
// See the License for the specific language governing permissions and
// limitations under the License.

package tikv

import (
	"encoding/hex"
	"math/rand"
	"strings"
	"sync/atomic"
	"time"

	"github.com/pingcap/errors"
	"github.com/pingcap/failpoint"
	pb "github.com/pingcap/kvproto/pkg/kvrpcpb"
	"github.com/pingcap/tidb/kv"
	"github.com/pingcap/tidb/metrics"
	"github.com/pingcap/tidb/store/tikv/tikvrpc"
	"github.com/pingcap/tidb/util/logutil"
	"github.com/prometheus/client_golang/prometheus"
	"go.uber.org/zap"
)

type actionPessimisticLock struct {
	*kv.LockCtx
}
type actionPessimisticRollback struct{}

var (
	_ twoPhaseCommitAction = actionPessimisticLock{}
	_ twoPhaseCommitAction = actionPessimisticRollback{}

	tiKVTxnRegionsNumHistogramPessimisticLock     = metrics.TiKVTxnRegionsNumHistogram.WithLabelValues(metricsTag("pessimistic_lock"))
	tiKVTxnRegionsNumHistogramPessimisticRollback = metrics.TiKVTxnRegionsNumHistogram.WithLabelValues(metricsTag("pessimistic_rollback"))
)

func (actionPessimisticLock) String() string {
	return "pessimistic_lock"
}

func (actionPessimisticLock) tiKVTxnRegionsNumHistogram() prometheus.Observer {
	return tiKVTxnRegionsNumHistogramPessimisticLock
}

func (actionPessimisticRollback) String() string {
	return "pessimistic_rollback"
}

func (actionPessimisticRollback) tiKVTxnRegionsNumHistogram() prometheus.Observer {
	return tiKVTxnRegionsNumHistogramPessimisticRollback
}

func (action actionPessimisticLock) handleSingleBatch(c *twoPhaseCommitter, bo *Backoffer, batch batchMutations) error {
	m := batch.mutations
	mutations := make([]*pb.Mutation, m.Len())
	for i := 0; i < m.Len(); i++ {
		mut := &pb.Mutation{
			Op:  pb.Op_PessimisticLock,
			Key: m.GetKey(i),
		}
		if c.txn.us.HasPresumeKeyNotExists(m.GetKey(i)) || (c.doingAmend && m.GetOp(i) == pb.Op_Insert) {
			mut.Assertion = pb.Assertion_NotExist
		}
		mutations[i] = mut
	}
	elapsed := uint64(time.Since(c.txn.startTime) / time.Millisecond)
	ttl := elapsed + atomic.LoadUint64(&ManagedLockTTL)
	failpoint.Inject("shortPessimisticLockTTL", func() {
		ttl = 1
		keys := make([]string, 0, len(mutations))
		for _, m := range mutations {
			keys = append(keys, hex.EncodeToString(m.Key))
		}
		logutil.BgLogger().Info("[failpoint] injected lock ttl = 1 on pessimistic lock",
			zap.Uint64("txnStartTS", c.startTS), zap.Strings("keys", keys))
	})
	req := tikvrpc.NewRequest(tikvrpc.CmdPessimisticLock, &pb.PessimisticLockRequest{
		Mutations:    mutations,
		PrimaryLock:  c.primary(),
		StartVersion: c.startTS,
		ForUpdateTs:  c.forUpdateTS,
		LockTtl:      ttl,
		IsFirstLock:  c.isFirstLock,
		WaitTimeout:  action.LockWaitTime,
		ReturnValues: action.ReturnValues,
		MinCommitTs:  c.forUpdateTS + 1,
	}, pb.Context{Priority: c.priority, SyncLog: c.syncLog})
	lockWaitStartTime := action.WaitStartTime
	for {
		// if lockWaitTime set, refine the request `WaitTimeout` field based on timeout limit
		if action.LockWaitTime > 0 {
			timeLeft := action.LockWaitTime - (time.Since(lockWaitStartTime)).Milliseconds()
			if timeLeft <= 0 {
				req.PessimisticLock().WaitTimeout = kv.LockNoWait
			} else {
				req.PessimisticLock().WaitTimeout = timeLeft
			}
		}
		failpoint.Inject("PessimisticLockErrWriteConflict", func() error {
			time.Sleep(300 * time.Millisecond)
			return kv.ErrWriteConflict
		})
		startTime := time.Now()
		resp, err := c.store.SendReq(bo, req, batch.region, readTimeoutShort)
		if action.LockCtx.Stats != nil {
			atomic.AddInt64(&action.LockCtx.Stats.LockRPCTime, int64(time.Since(startTime)))
			atomic.AddInt64(&action.LockCtx.Stats.LockRPCCount, 1)
		}
		if err != nil {
			return errors.Trace(err)
		}
		regionErr, err := resp.GetRegionError()
		if err != nil {
			return errors.Trace(err)
		}
		if regionErr != nil {
			err = bo.Backoff(BoRegionMiss, errors.New(regionErr.String()))
			if err != nil {
				return errors.Trace(err)
			}
			err = c.pessimisticLockMutations(bo, action.LockCtx, batch.mutations)
			return errors.Trace(err)
		}
		if resp.Resp == nil {
			return errors.Trace(ErrBodyMissing)
		}
		lockResp := resp.Resp.(*pb.PessimisticLockResponse)
		keyErrs := lockResp.GetErrors()
		if len(keyErrs) == 0 {
			if action.ReturnValues {
				action.ValuesLock.Lock()
				for i, mutation := range mutations {
					action.Values[string(mutation.Key)] = kv.ReturnedValue{Value: lockResp.Values[i]}
				}
				action.ValuesLock.Unlock()
			}
			return nil
		}
		var locks []*Lock
		for _, keyErr := range keyErrs {
			// Check already exists error
			if alreadyExist := keyErr.GetAlreadyExist(); alreadyExist != nil {
				key := alreadyExist.GetKey()
				return c.extractKeyExistsErr(key)
			}
			if deadlock := keyErr.Deadlock; deadlock != nil {
				return &ErrDeadlock{Deadlock: deadlock}
			}

			// Extract lock from key error
			lock, err1 := extractLockFromKeyErr(keyErr)
			if err1 != nil {
				return errors.Trace(err1)
			}
			locks = append(locks, lock)
		}
		// Because we already waited on tikv, no need to Backoff here.
		// tikv default will wait 3s(also the maximum wait value) when lock error occurs
		startTime = time.Now()
		msBeforeTxnExpired, _, err := c.store.lockResolver.ResolveLocks(bo, 0, locks)
		if err != nil {
			return errors.Trace(err)
		}
		if action.LockCtx.Stats != nil {
			atomic.AddInt64(&action.LockCtx.Stats.ResolveLockTime, int64(time.Since(startTime)))
		}

		// If msBeforeTxnExpired is not zero, it means there are still locks blocking us acquiring
		// the pessimistic lock. We should return acquire fail with nowait set or timeout error if necessary.
		if msBeforeTxnExpired > 0 {
			if action.LockWaitTime == kv.LockNoWait {
				return ErrLockAcquireFailAndNoWaitSet
			} else if action.LockWaitTime == kv.LockAlwaysWait {
				// do nothing but keep wait
			} else {
				// the lockWaitTime is set, we should return wait timeout if we are still blocked by a lock
				if time.Since(lockWaitStartTime).Milliseconds() >= action.LockWaitTime {
					return errors.Trace(ErrLockWaitTimeout)
				}
			}
			if action.LockCtx.PessimisticLockWaited != nil {
				atomic.StoreInt32(action.LockCtx.PessimisticLockWaited, 1)
			}
		}

		// Handle the killed flag when waiting for the pessimistic lock.
		// When a txn runs into LockKeys() and backoff here, it has no chance to call
		// executor.Next() and check the killed flag.
		if action.Killed != nil {
			// Do not reset the killed flag here!
			// actionPessimisticLock runs on each region parallelly, we have to consider that
			// the error may be dropped.
			if atomic.LoadUint32(action.Killed) == 1 {
				return errors.Trace(ErrQueryInterrupted)
			}
		}
	}
}

func (actionPessimisticRollback) handleSingleBatch(c *twoPhaseCommitter, bo *Backoffer, batch batchMutations) error {
	req := tikvrpc.NewRequest(tikvrpc.CmdPessimisticRollback, &pb.PessimisticRollbackRequest{
		StartVersion: c.startTS,
		ForUpdateTs:  c.forUpdateTS,
		Keys:         batch.mutations.GetKeys(),
	})
	resp, err := c.store.SendReq(bo, req, batch.region, readTimeoutShort)
	if err != nil {
		return errors.Trace(err)
	}
	regionErr, err := resp.GetRegionError()
	if err != nil {
		return errors.Trace(err)
	}
	if regionErr != nil {
		err = bo.Backoff(BoRegionMiss, errors.New(regionErr.String()))
		if err != nil {
			return errors.Trace(err)
		}
		err = c.pessimisticRollbackMutations(bo, batch.mutations)
		return errors.Trace(err)
	}
	return nil
}

func (c *twoPhaseCommitter) pessimisticLockMutations(bo *Backoffer, lockCtx *kv.LockCtx, mutations CommitterMutations) error {
	if c.connID > 0 {
		failpoint.Inject("beforePessimisticLock", func(val failpoint.Value) {
			// Pass multiple instructions in one string, delimited by commas, to trigger multiple behaviors, like
			// `return("delay,fail")`. Then they will be executed sequentially at once.
			if v, ok := val.(string); ok {
				for _, action := range strings.Split(v, ",") {
					if action == "delay" {
						duration := time.Duration(rand.Int63n(int64(time.Second) * 5))
						logutil.Logger(bo.ctx).Info("[failpoint] injected delay at pessimistic lock",
							zap.Uint64("txnStartTS", c.startTS), zap.Duration("duration", duration))
						time.Sleep(duration)
					} else if action == "fail" {
						logutil.Logger(bo.ctx).Info("[failpoint] injected failure at pessimistic lock",
							zap.Uint64("txnStartTS", c.startTS))
						failpoint.Return(errors.New("injected failure at pessimistic lock"))
					}
				}
			}
		})
	}
	return c.doActionOnMutations(bo, actionPessimisticLock{lockCtx}, mutations)
}

func (c *twoPhaseCommitter) pessimisticRollbackMutations(bo *Backoffer, mutations CommitterMutations) error {
	return c.doActionOnMutations(bo, actionPessimisticRollback{}, mutations)
}
<|endoftext|>


<commit_before>package supervisor

import (
	"bufio"
	"io"
	"os"
	"strconv"
	"syscall"
	"testing"
)

// Compare two string/string maps.
func cmpMap(m1 map[string]string, m2 map[string]string) bool {
	if len(m1) != len(m2) {
		return false
	}
	for k, v := range m2 {
		if v != m1[k] {
			return false
		}
	}
	return true
}

// Compare two byte arrays.
func cmpBytes(p1 []byte, p2 []byte) bool {
	if p1 == nil {
		p1 = []byte{}
	}
	if p2 == nil {
		p2 = []byte{}
	}

	if len(p1) != len(p2) {
		return false
	}
	for i, v := range p2 {
		if v != p1[i] {
			return false
		}
	}
	return true
}

// Compare two events
func cmpEvents(e1 *Event, e2 *Event) bool {
	switch {
	case e1 == nil && e2 == nil:
		return true
	case e1 == nil || e2 == nil:
		return false
	default:
		return cmpMap(e1.Header, e2.Header) && cmpMap(e1.Meta, e2.Meta) && cmpBytes(e1.Payload, e2.Payload)
	}
}

// Redirect stdout.
func redirectStdout() (stdout *os.File, err error) {
	stdout, stdoutWriter, err := os.Pipe()
	if err == nil {
		err = syscall.Dup2(int(stdoutWriter.Fd()), int(os.Stdout.Fd()))
	}
	return
}

// Redirect stdin.
func redirectStdin() (stdin *os.File, err error) {
	stdinReader, stdin, err := os.Pipe()
	if err == nil {
		err = syscall.Dup2(int(stdinReader.Fd()), int(syscall.Stdin))
	}
	return
}

func unredirectStdout() (err error) {
	stdout := os.NewFile(uintptr(syscall.Stdout), "/dev/stdout")
	if err == nil {
		err = syscall.Dup2(int(stdout.Fd()), int(os.Stdout.Fd()))
	}
	return
}

// Construct an event.
func createEvent(serial int, eventname string, processname string, payload []byte) *Event {
	serialstr := strconv.Itoa(serial)
	return &Event{
		map[string]string{
			"ver":        "3.0",
			"server":     "supervisor",
			"eventname":  eventname,
			"serial":     serialstr,
			"pool":       "listener",
			"poolserial": serialstr,
		},
		map[string]string{
			"processname": processname,
			"groupname":   processname,
		},
		payload,
	}
}

// Test the ReadEvent function.
func TestRead(t *testing.T) {
	reader, writer := io.Pipe()
	bufReader := bufio.NewReader(reader)
	serial := 0

	sendAndVerify := func(eventname string, payload []byte) {
		sentEvent := createEvent(serial, eventname, "test", payload)
		serial++

		go func() {
			_, err := writer.Write(sentEvent.ToBytes())
			if err != nil {
				t.Error(err)
			}
		}()

		receiveEvent, err := ReadEvent(bufReader)
		if err != nil {
			t.Error(err)
		}

		if !cmpEvents(sentEvent, receiveEvent) {
			t.Error("invalid event received")
		}
	}

	sendAndVerify("EVENT_EMPTY_PAYLOAD", []byte{})
	sendAndVerify("EVENT_FULL_PAYLOAD", []byte("this is a payload test"))
}

// Test the WriteResult functions.
func TestWrite(t *testing.T) {
	reader, writer := io.Pipe()

	readAndVerify := func(expected string) {
		payload, err := ReadResult(reader)
		switch {
		case err != nil:
			t.Error(err)
		case string(payload) != expected:
			t.Errorf("Payload result invalid: %s != %s", payload, expected)
		}
	}

	payload := "some arbitrary data"
	go WriteResult(writer, []byte(payload))
	readAndVerify(payload)

	go WriteResultOK(writer)
	readAndVerify("OK")

	go WriteResultFail(writer)
	readAndVerify("FAIL")
}

// Test the Listen function.
func TestListen(t *testing.T) {
	stdin, stdinWriter := io.Pipe()
	stdoutReader, stdout := io.Pipe()

	ch := make(chan *Event, 1)
	reader := bufio.NewReader(stdoutReader)

	go func() {
		if err := Listen(stdin, stdout, ch); err != nil {
			t.Error(err)
		}
	}()

	serial := 0
	sendAndVerify := func(eventname string, payload []byte) {
		sentEvent := createEvent(serial, eventname, "test", payload)
		serial++

		bytes := sentEvent.ToBytes()
		_, err := stdinWriter.Write(bytes)
		if err != nil {
			t.Error(err)
		}

		result, err := ReadResult(reader)
		if err != nil {
			t.Error(err)
		}
		if string(result) != "OK" {
			t.Error("invalid result")
		}

		receiveEvent, ok := <-ch
		if !ok {
			t.Error("channel closed")
		} else if !cmpEvents(sentEvent, receiveEvent) {
			t.Error("invalid event received")
		}
	}

	sendAndVerify("PROCESS_STATE_RUNNING", []byte{})
	sendAndVerify("PROCESS_LOG_STDERR", []byte("some pretend log data"))
}
<commit_msg>Remove unused test functions.<commit_after>package supervisor

import (
	"bufio"
	"io"
	"strconv"
	"testing"
)

// Compare two string/string maps.
func cmpMap(m1 map[string]string, m2 map[string]string) bool {
	if len(m1) != len(m2) {
		return false
	}
	for k, v := range m2 {
		if v != m1[k] {
			return false
		}
	}
	return true
}

// Compare two byte arrays.
func cmpBytes(p1 []byte, p2 []byte) bool {
	if p1 == nil {
		p1 = []byte{}
	}
	if p2 == nil {
		p2 = []byte{}
	}

	if len(p1) != len(p2) {
		return false
	}
	for i, v := range p2 {
		if v != p1[i] {
			return false
		}
	}
	return true
}

// Compare two events
func cmpEvents(e1 *Event, e2 *Event) bool {
	switch {
	case e1 == nil && e2 == nil:
		return true
	case e1 == nil || e2 == nil:
		return false
	default:
		return cmpMap(e1.Header, e2.Header) && cmpMap(e1.Meta, e2.Meta) && cmpBytes(e1.Payload, e2.Payload)
	}
}

// Construct an event.
func createEvent(serial int, eventname string, processname string, payload []byte) *Event {
	serialstr := strconv.Itoa(serial)
	return &Event{
		map[string]string{
			"ver":        "3.0",
			"server":     "supervisor",
			"eventname":  eventname,
			"serial":     serialstr,
			"pool":       "listener",
			"poolserial": serialstr,
		},
		map[string]string{
			"processname": processname,
			"groupname":   processname,
		},
		payload,
	}
}

// Test the ReadEvent function.
func TestRead(t *testing.T) {
	reader, writer := io.Pipe()
	bufReader := bufio.NewReader(reader)
	serial := 0

	sendAndVerify := func(eventname string, payload []byte) {
		sentEvent := createEvent(serial, eventname, "test", payload)
		serial++

		go func() {
			_, err := writer.Write(sentEvent.ToBytes())
			if err != nil {
				t.Error(err)
			}
		}()

		receiveEvent, err := ReadEvent(bufReader)
		if err != nil {
			t.Error(err)
		}

		if !cmpEvents(sentEvent, receiveEvent) {
			t.Error("invalid event received")
		}
	}

	sendAndVerify("EVENT_EMPTY_PAYLOAD", []byte{})
	sendAndVerify("EVENT_FULL_PAYLOAD", []byte("this is a payload test"))
}

// Test the WriteResult functions.
func TestWrite(t *testing.T) {
	reader, writer := io.Pipe()

	readAndVerify := func(expected string) {
		payload, err := ReadResult(reader)
		switch {
		case err != nil:
			t.Error(err)
		case string(payload) != expected:
			t.Errorf("Payload result invalid: %s != %s", payload, expected)
		}
	}

	payload := "some arbitrary data"
	go WriteResult(writer, []byte(payload))
	readAndVerify(payload)

	go WriteResultOK(writer)
	readAndVerify("OK")

	go WriteResultFail(writer)
	readAndVerify("FAIL")
}

// Test the Listen function.
func TestListen(t *testing.T) {
	stdin, stdinWriter := io.Pipe()
	stdoutReader, stdout := io.Pipe()

	ch := make(chan *Event, 1)
	reader := bufio.NewReader(stdoutReader)

	go func() {
		if err := Listen(stdin, stdout, ch); err != nil {
			t.Error(err)
		}
	}()

	serial := 0
	sendAndVerify := func(eventname string, payload []byte) {
		sentEvent := createEvent(serial, eventname, "test", payload)
		serial++

		bytes := sentEvent.ToBytes()
		_, err := stdinWriter.Write(bytes)
		if err != nil {
			t.Error(err)
		}

		result, err := ReadResult(reader)
		if err != nil {
			t.Error(err)
		}
		if string(result) != "OK" {
			t.Error("invalid result")
		}

		receiveEvent, ok := <-ch
		if !ok {
			t.Error("channel closed")
		} else if !cmpEvents(sentEvent, receiveEvent) {
			t.Error("invalid event received")
		}
	}

	sendAndVerify("PROCESS_STATE_RUNNING", []byte{})
	sendAndVerify("PROCESS_LOG_STDERR", []byte("some pretend log data"))
}
<|endoftext|>


<commit_before>//	logplus.Logln("this is example")

package logplus

import (
	"fmt"
	"os"
	"path"
	"runtime"
	"strings"
	"time"
)

const (
	resetAll = "\033[0m"
)

var (
	timeFormat = time.RFC3339
	logLevel   = INFO
)

type Color int

const (
	ForegroundDefault Color = iota
	ForegroundBlack
	ForegroundRed
	ForegroundGreen
	ForegroundYellow
	ForegroundBlue
	ForegroundMagenta
	ForegroundCyan
	ForegroundWhite

	BackgroundDefault
	BackgroundBlack
	BackgroundRed
	BackgroundGreen
	BackgroundYellow
	BackgroundBlue
	BackgroundMagenta
	BackgroundCyan
	BackgroundWhite
)

func (color Color) String() string {
	switch color {
	case ForegroundDefault:
		return "\033[39m"
	case ForegroundBlack:
		return "\033[30m"
	case ForegroundRed:
		return "\033[31m"
	case ForegroundGreen:
		return "\033[32m"
	case ForegroundYellow:
		return "\033[33m"
	case ForegroundBlue:
		return "\033[34m"
	case ForegroundMagenta:
		return "\033[35m"
	case ForegroundCyan:
		return "\033[36m"
	case ForegroundWhite:
		return "\033[97m"
	case BackgroundDefault:
		return "\033[49m"
	case BackgroundBlack:
		return "\033[40m"
	case BackgroundRed:
		return "\033[41m"
	case BackgroundGreen:
		return "\033[42m"
	case BackgroundYellow:
		return "\033[43m"
	case BackgroundBlue:
		return "\033[44m"
	case BackgroundMagenta:
		return "\033[45m"
	case BackgroundCyan:
		return "\033[46m"
	case BackgroundWhite:
		return "\033[107m"
	}
	panic("Unknown value")
}

type LogLevel int

const (
	PANIC LogLevel = iota
	FATAL
	ERROR
	WARN
	INFO
	DEBUG
)

func (level LogLevel) String() string {
	switch level {
	case PANIC:
		return "PANIC"
	case FATAL:
		return "FATAL"
	case ERROR:
		return "ERROR"
	case WARN:
		return "WARN"
	case INFO:
		return "INFO"
	case DEBUG:
		return "DEBUG"
	}
	panic("Unknown value")
}

type CallInfo struct {
	PackageName string
	FileName    string
	FuncName    string
	Line        int
}

func (ci CallInfo) String() string {
	return fmt.Sprintf("%s#%s (%s:%d)", ci.PackageName, ci.FuncName, ci.FileName, ci.Line)
}

func SetTimeFormat(tf string) {
	timeFormat = tf
}

func SetLevel(level LogLevel) {
	logLevel = level
}

func isLogAvailable(level LogLevel) bool {
	if level <= logLevel {
		return true
	}
	return false
}

func Log(log ...interface{}) {
	info := getCallInfo()
	fmt.Printf(" %s [%s] <%d> ", time.Now().Format(timeFormat), info, os.Getpid)
	fmt.Print(log...)
}

func Logf(format string, log ...interface{}) {
	info := getCallInfo()
	fmt.Printf(" %s [%s] <%d> ", time.Now().Format(timeFormat), info, os.Getpid)
	fmt.Printf(format, log...)
}

func Logln(log ...interface{}) {
	info := getCallInfo()
	fmt.Printf(" %s [%s] <%d> ", time.Now().Format(timeFormat), info, os.Getpid)
	fmt.Println(log...)
}

func Colored(color Color, log ...interface{}) {
	info := getCallInfo()
	fmt.Printf("%s [%s] <%d> ", time.Now().Format(timeFormat), info, os.Getpid)
	fmt.Print(color)
	fmt.Print(log...)
	fmt.Print(resetAll)
}

func Coloredf(color Color, format string, log ...interface{}) {
	info := getCallInfo()
	fmt.Printf("%s [%s] <%d> ", time.Now().Format(timeFormat), info, os.Getpid)
	fmt.Print(color)
	fmt.Printf(format, log...)
	fmt.Print(resetAll)
}

func Coloredln(color Color, log ...interface{}) {
	info := getCallInfo()
	fmt.Printf("%s [%s] <%d> ", time.Now().Format(timeFormat), info, os.Getpid)
	fmt.Print(color)
	fmt.Print(log...)
	fmt.Println(resetAll)
}

func Panic(log ...interface{}) {
	if isLogAvailable(PANIC) {
		printLogInfo(BackgroundRed, PANIC)
		fmt.Print(log...)

		panic(fmt.Sprint(log...))
	}
}

func Panicf(format string, log ...interface{}) {
	if isLogAvailable(PANIC) {
		printLogInfo(BackgroundRed, PANIC)
		fmt.Printf(format, log...)

		panic(fmt.Sprintf(format, log...))
	}
}

func Panicln(log ...interface{}) {
	if isLogAvailable(PANIC) {
		printLogInfo(BackgroundRed, PANIC)
		fmt.Println(log...)

		panic(fmt.Sprintln(log...))
	}
}

func Fatal(log ...interface{}) {
	if isLogAvailable(FATAL) {
		printLogInfo(BackgroundRed, FATAL)
		fmt.Print(log...)

		os.Exit(1)
	}
}

func Fatalf(format string, log ...interface{}) {
	if isLogAvailable(FATAL) {
		printLogInfo(BackgroundRed, FATAL)
		fmt.Printf(format, log...)

		os.Exit(1)
	}
}

func Fatalln(log ...interface{}) {
	if isLogAvailable(FATAL) {
		printLogInfo(BackgroundRed, FATAL)
		fmt.Println(log...)

		os.Exit(1)
	}
}

func Error(log ...interface{}) {
	if isLogAvailable(ERROR) {
		printLogInfo(BackgroundMagenta, ERROR)
		fmt.Print(log...)
	}
}

func Errorf(format string, log ...interface{}) {
	if isLogAvailable(ERROR) {
		printLogInfo(BackgroundMagenta, ERROR)
		fmt.Printf(format, log...)
	}
}

func Errorln(log ...interface{}) {
	if isLogAvailable(ERROR) {
		printLogInfo(BackgroundMagenta, ERROR)
		fmt.Println(log...)
	}
}

func Warn(log ...interface{}) {
	if isLogAvailable(WARN) {
		printLogInfo(BackgroundYellow, WARN)
		fmt.Print(log...)
	}
}

func Warnf(format string, log ...interface{}) {
	if isLogAvailable(WARN) {
		printLogInfo(BackgroundYellow, WARN)
		fmt.Printf(format, log...)
	}
}

func Warnln(log ...interface{}) {
	if isLogAvailable(WARN) {
		printLogInfo(BackgroundYellow, WARN)
		fmt.Println(log...)
	}
}

func Info(log ...interface{}) {
	if isLogAvailable(INFO) {
		printLogInfo(BackgroundGreen, INFO)
		fmt.Print(log...)
	}
}

func Infof(format string, log ...interface{}) {
	if isLogAvailable(INFO) {
		printLogInfo(BackgroundGreen, INFO)
		fmt.Printf(format, log...)
	}
}

func Infoln(log ...interface{}) {
	if isLogAvailable(INFO) {
		printLogInfo(BackgroundGreen, INFO)
		fmt.Println(log...)
	}
}

func Debug(log ...interface{}) {
	if isLogAvailable(DEBUG) {
		printLogInfo(BackgroundCyan, DEBUG)
		fmt.Print(log...)
	}
}

func Debugf(format string, log ...interface{}) {
	if isLogAvailable(DEBUG) {
		printLogInfo(BackgroundCyan, DEBUG)
		fmt.Printf(format, log...)
	}
}

func Debugln(log ...interface{}) {
	if isLogAvailable(DEBUG) {
		printLogInfo(BackgroundCyan, DEBUG)
		fmt.Println(log...)
	}
}

func concat(sep string, strs ...string) string {
	var result = make([]byte, 0, 100)
	for i, _ := range strs {
		result = append(result, strs[i]...)
		if i < len(strs)-1 {
			result = append(result, sep...)
		}
	}
	return string(result)
}

func printLogInfo(color Color, level LogLevel) {
	fmt.Print(concat("", color.String(), level.String(), resetAll))
	fmt.Printf(" %s [%s] <%d> ", time.Now().Format(timeFormat), getCallInfo(), os.Getpid)
}

func getCallInfo() *CallInfo {
	pc, filePath, line, _ := runtime.Caller(2)
	_, fileName := path.Split(filePath)
	parts := strings.Split(runtime.FuncForPC(pc).Name(), ".")
	pl := len(parts)
	packageName := ""
	funcName := parts[pl-1]
	if parts[pl-2][0] == '(' {
		funcName = parts[pl-2] + "." + funcName
		packageName = strings.Join(parts[0:pl-2], ".")
	} else {
		packageName = strings.Join(parts[0:pl-1], ".")
	}

	return &CallInfo{
		PackageName: packageName,
		FileName:    fileName,
		FuncName:    funcName,
		Line:        line,
	}
}
<commit_msg>refactoring<commit_after>//	logplus.Logln("this is example")

package logplus

import (
	"fmt"
	"os"
	"path"
	"runtime"
	"strings"
	"time"
)

const (
	resetAll = "\033[0m"
)

var (
	timeFormat = time.RFC3339
	logLevel   = INFO
)

type Color int

const (
	ForegroundDefault Color = iota
	ForegroundBlack
	ForegroundRed
	ForegroundGreen
	ForegroundYellow
	ForegroundBlue
	ForegroundMagenta
	ForegroundCyan
	ForegroundWhite

	BackgroundDefault
	BackgroundBlack
	BackgroundRed
	BackgroundGreen
	BackgroundYellow
	BackgroundBlue
	BackgroundMagenta
	BackgroundCyan
	BackgroundWhite
)

func (color Color) String() string {
	switch color {
	case ForegroundDefault:
		return "\033[39m"
	case ForegroundBlack:
		return "\033[30m"
	case ForegroundRed:
		return "\033[31m"
	case ForegroundGreen:
		return "\033[32m"
	case ForegroundYellow:
		return "\033[33m"
	case ForegroundBlue:
		return "\033[34m"
	case ForegroundMagenta:
		return "\033[35m"
	case ForegroundCyan:
		return "\033[36m"
	case ForegroundWhite:
		return "\033[97m"
	case BackgroundDefault:
		return "\033[49m"
	case BackgroundBlack:
		return "\033[40m"
	case BackgroundRed:
		return "\033[41m"
	case BackgroundGreen:
		return "\033[42m"
	case BackgroundYellow:
		return "\033[43m"
	case BackgroundBlue:
		return "\033[44m"
	case BackgroundMagenta:
		return "\033[45m"
	case BackgroundCyan:
		return "\033[46m"
	case BackgroundWhite:
		return "\033[107m"
	}
	panic("Unknown value")
}

type LogLevel int

const (
	PANIC LogLevel = iota
	FATAL
	ERROR
	WARN
	INFO
	DEBUG
)

func (level LogLevel) String() string {
	switch level {
	case PANIC:
		return "PANIC"
	case FATAL:
		return "FATAL"
	case ERROR:
		return "ERROR"
	case WARN:
		return "WARN"
	case INFO:
		return "INFO"
	case DEBUG:
		return "DEBUG"
	}
	panic("Unknown value")
}

type CallInfo struct {
	PackageName string
	FileName    string
	FuncName    string
	Line        int
}

func (ci CallInfo) String() string {
	return fmt.Sprintf("%s#%s (%s:%d)", ci.PackageName, ci.FuncName, ci.FileName, ci.Line)
}

func SetTimeFormat(tf string) {
	timeFormat = tf
}

func SetLevel(level LogLevel) {
	logLevel = level
}

func isLogAvailable(level LogLevel) bool {
	if level <= logLevel {
		return true
	}
	return false
}

func Log(log ...interface{}) {
	printCaller(3)

	fmt.Print(log...)
}

func Logf(format string, log ...interface{}) {
	printCaller(3)

	fmt.Printf(format, log...)
}

func Logln(log ...interface{}) {
	printCaller(3)

	fmt.Println(log...)
}

func Colored(color Color, log ...interface{}) {
	printCaller(3)

	fmt.Print(color)
	fmt.Print(log...)
	fmt.Print(resetAll)
}

func Coloredf(color Color, format string, log ...interface{}) {
	printCaller(3)

	fmt.Print(color)
	fmt.Printf(format, log...)
	fmt.Print(resetAll)
}

func Coloredln(color Color, log ...interface{}) {
	printCaller(3)

	fmt.Print(color)
	fmt.Print(log...)
	fmt.Println(resetAll)
}

func Panic(log ...interface{}) {
	if isLogAvailable(PANIC) {
		printLogInfo(BackgroundRed, PANIC)

		fmt.Print(log...)

		panic(fmt.Sprint(log...))
	}
}

func Panicf(format string, log ...interface{}) {
	if isLogAvailable(PANIC) {
		printLogInfo(BackgroundRed, PANIC)

		fmt.Printf(format, log...)

		panic(fmt.Sprintf(format, log...))
	}
}

func Panicln(log ...interface{}) {
	if isLogAvailable(PANIC) {
		printLogInfo(BackgroundRed, PANIC)

		fmt.Println(log...)

		panic(fmt.Sprintln(log...))
	}
}

func Fatal(log ...interface{}) {
	if isLogAvailable(FATAL) {
		printLogInfo(BackgroundRed, FATAL)

		fmt.Print(log...)

		os.Exit(1)
	}
}

func Fatalf(format string, log ...interface{}) {
	if isLogAvailable(FATAL) {
		printLogInfo(BackgroundRed, FATAL)

		fmt.Printf(format, log...)

		os.Exit(1)
	}
}

func Fatalln(log ...interface{}) {
	if isLogAvailable(FATAL) {
		printLogInfo(BackgroundRed, FATAL)

		fmt.Println(log...)

		os.Exit(1)
	}
}

func Error(log ...interface{}) {
	if isLogAvailable(ERROR) {
		printLogInfo(BackgroundMagenta, ERROR)

		fmt.Print(log...)
	}
}

func Errorf(format string, log ...interface{}) {
	if isLogAvailable(ERROR) {
		printLogInfo(BackgroundMagenta, ERROR)

		fmt.Printf(format, log...)
	}
}

func Errorln(log ...interface{}) {
	if isLogAvailable(ERROR) {
		printLogInfo(BackgroundMagenta, ERROR)

		fmt.Println(log...)
	}
}

func Warn(log ...interface{}) {
	if isLogAvailable(WARN) {
		printLogInfo(BackgroundYellow, WARN)

		fmt.Print(log...)
	}
}

func Warnf(format string, log ...interface{}) {
	if isLogAvailable(WARN) {
		printLogInfo(BackgroundYellow, WARN)

		fmt.Printf(format, log...)
	}
}

func Warnln(log ...interface{}) {
	if isLogAvailable(WARN) {
		printLogInfo(BackgroundYellow, WARN)

		fmt.Println(log...)
	}
}

func Info(log ...interface{}) {
	if isLogAvailable(INFO) {
		printLogInfo(BackgroundGreen, INFO)

		fmt.Print(log...)
	}
}

func Infof(format string, log ...interface{}) {
	if isLogAvailable(INFO) {
		printLogInfo(BackgroundGreen, INFO)

		fmt.Printf(format, log...)
	}
}

func Infoln(log ...interface{}) {
	if isLogAvailable(INFO) {
		printLogInfo(BackgroundGreen, INFO)

		fmt.Println(log...)
	}
}

func Debug(log ...interface{}) {
	if isLogAvailable(DEBUG) {
		printLogInfo(BackgroundCyan, DEBUG)

		fmt.Print(log...)
	}
}

func Debugf(format string, log ...interface{}) {
	if isLogAvailable(DEBUG) {
		printLogInfo(BackgroundCyan, DEBUG)

		fmt.Printf(format, log...)
	}
}

func Debugln(log ...interface{}) {
	if isLogAvailable(DEBUG) {
		printLogInfo(BackgroundCyan, DEBUG)

		fmt.Println(log...)
	}
}

func concat(sep string, strs ...string) string {
	var result = make([]byte, 0, 100)
	for i, _ := range strs {
		result = append(result, strs[i]...)
		if i < len(strs)-1 {
			result = append(result, sep...)
		}
	}
	return string(result)
}

func printLogInfo(color Color, level LogLevel) {
	fmt.Print(concat("", color.String(), level.String(), resetAll))
	printCaller(4)
}

func printCaller(depth int) {
	t := time.Now().Format(timeFormat)
	info := getCallInfo(depth)
	fmt.Printf(" %s [%s] <%d> ", t, info, os.Getpid)
}

func getCallInfo(depth int) *CallInfo {
	pc, filePath, line, _ := runtime.Caller(depth)
	_, fileName := path.Split(filePath)
	parts := strings.Split(runtime.FuncForPC(pc).Name(), ".")
	pl := len(parts)
	packageName := ""
	funcName := parts[pl-1]
	if parts[pl-2][0] == '(' {
		funcName = parts[pl-2] + "." + funcName
		packageName = strings.Join(parts[0:pl-2], ".")
	} else {
		packageName = strings.Join(parts[0:pl-1], ".")
	}

	return &CallInfo{
		PackageName: packageName,
		FileName:    fileName,
		FuncName:    funcName,
		Line:        line,
	}
}
<|endoftext|>


<commit_before>// ls -- list files and directories
// Part of goutils (https://github.com/trevorparker/goutils)
//
// Copyright (c) 2014 Trevor Parker <trevor@trevorparker.com>
// All rights reserved
//
// Distributed under the terms of the Modified BSD License (see LICENSE)

package main

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"os"
	"strings"
)

type arg struct {
	file            []string
	almost_all      bool
	comma_separated bool
	one_per_line    bool
}

const (
	usage_message string = "usage: ls [OPTION ...] [FILE ...]"
	help_message  string = `List files and directories, and information about them.

  -A, --almost-all     include entries beginning with a dot, except
                       implied . and ..
  -m                   print a comma-separated list of entries
  -1                   print one entry per line
  -h, --help           print this help message and exit
`
)

func usage(error string) {
	fmt.Fprintf(os.Stderr, "ls: %s\n%s\n", error, usage_message)
	os.Exit(1)
}

func help() {
	fmt.Printf("%s\n%s", usage_message, help_message)
	os.Exit(0)
}

func ls(file string, args arg) {
	entries := make([]os.FileInfo, 0)

	// Determine if this is a file or directory, then call out
	// to ReadDir if it's a directory. Otherwise, we're can just
	// pass the file info on.
	fi, err := os.Stat(file)
	if err != nil {
		panic(err)
	} else if fi.IsDir() {
		e, err := ioutil.ReadDir(file)
		if err != nil {
			panic(err)
		}
		entries = e
	} else {
		entries = append(entries, fi)
	}
	printEntries(&entries, &args)
}

func printEntries(entries *[]os.FileInfo, args *arg) {
	var out bytes.Buffer

	filtered_entries := filterEntries(entries, args)

	if args.one_per_line {
		for _, e := range filtered_entries {
			out.WriteString(fmt.Sprintf("%s\n", e.Name()))
		}
		fmt.Print(out.String())
	} else if args.comma_separated {
		for i, e := range filtered_entries {
			out.WriteString(e.Name())
			if i < len(filtered_entries)-1 {
				out.WriteString(", ")
			}
		}
		fmt.Println(out.String())
	} else {
		longest_entry := 1
		for _, e := range filtered_entries {
			length := len(e.Name())
			if length > longest_entry {
				longest_entry = length + 1
			}
		}

		columns := int(78 / longest_entry)
		formatted_string := fmt.Sprintf("%%-%ds", longest_entry)
		for i, e := range filtered_entries {
			out.WriteString(fmt.Sprintf(formatted_string, e.Name()))
			if i%columns == columns-1 {
				out.WriteString("\n")
			}
		}
		fmt.Println(out.String())
	}
}

func filterEntries(entries *[]os.FileInfo, args *arg) []os.FileInfo {
	filtered_entries := make([]os.FileInfo, 0)
	for _, e := range *entries {
		if !args.almost_all && strings.HasPrefix(e.Name(), ".") {
			continue
		}
		filtered_entries = append(filtered_entries, e)
	}

	return filtered_entries
}

func main() {
	args := arg{}
	reached_files := false

	for i := 1; i < len(os.Args); i++ {
		if reached_files == false {
			if os.Args[i] == "-h" || os.Args[i] == "--help" {
				help()
			}
			if os.Args[i] == "-A" {
				args.almost_all = true
				continue
			}
			if os.Args[i] == "-m" {
				args.comma_separated = true
				continue
			}
			if os.Args[i] == "-1" {
				args.one_per_line = true
				continue
			}
			if os.Args[i] == "--" {
				reached_files = true
				continue
			}
			if os.Args[i] == "-" {
				reached_files = true
			}
		}
		reached_files = true
		arg_v := os.Args[i]
		args.file = append(args.file, arg_v)
	}

	if len(args.file) == 0 {
		ls("./", args)
	} else {
		for i := range args.file {
			ls(args.file[i], args)
		}
	}
}
<commit_msg>Add support for `-Q, --quote-name`<commit_after>// ls -- list files and directories
// Part of goutils (https://github.com/trevorparker/goutils)
//
// Copyright (c) 2014 Trevor Parker <trevor@trevorparker.com>
// All rights reserved
//
// Distributed under the terms of the Modified BSD License (see LICENSE)

package main

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"os"
	"strings"
)

type arg struct {
	file            []string
	almost_all      bool
	comma_separated bool
	quote_name      bool
	one_per_line    bool
}

const (
	usage_message string = "usage: ls [OPTION ...] [FILE ...]"
	help_message  string = `List files and directories, and information about them.

  -A, --almost-all     include entries beginning with a dot, except
                       implied . and ..
  -m                   print a comma-separated list of entries
  -Q, --quote-name     print each entry surrounded by double quotes
  -1                   print one entry per line
  -h, --help           print this help message and exit
`
)

func usage(error string) {
	fmt.Fprintf(os.Stderr, "ls: %s\n%s\n", error, usage_message)
	os.Exit(1)
}

func help() {
	fmt.Printf("%s\n%s", usage_message, help_message)
	os.Exit(0)
}

func ls(file string, args arg) {
	entries := make([]os.FileInfo, 0)

	// Determine if this is a file or directory, then call out
	// to ReadDir if it's a directory. Otherwise, we're can just
	// pass the file info on.
	fi, err := os.Stat(file)
	if err != nil {
		panic(err)
	} else if fi.IsDir() {
		e, err := ioutil.ReadDir(file)
		if err != nil {
			panic(err)
		}
		entries = e
	} else {
		entries = append(entries, fi)
	}
	printEntries(&entries, &args)
}

func printEntries(entries *[]os.FileInfo, args *arg) {
	var out bytes.Buffer

	filtered_entries := filterEntries(entries, args)

	if args.one_per_line {
		for _, e := range filtered_entries {
			name := e.Name()
			if args.quote_name {
				name = fmt.Sprintf("\"%s\"", e.Name())
			}
			out.WriteString(fmt.Sprintf("%s\n", name))
		}
		fmt.Print(out.String())
	} else if args.comma_separated {
		for i, e := range filtered_entries {
			name := e.Name()
			if args.quote_name {
				name = fmt.Sprintf("\"%s\"", e.Name())
			}
			out.WriteString(name)
			if i < len(filtered_entries)-1 {
				out.WriteString(", ")
			}
		}
		fmt.Println(out.String())
	} else {
		longest_entry := 1
		for _, e := range filtered_entries {
			name := e.Name()
			if args.quote_name {
				name = fmt.Sprintf("\"%s\"", e.Name())
			}
			length := len(name)
			if length > longest_entry {
				longest_entry = length + 1
			}
		}

		columns := int(78 / longest_entry)
		formatted_string := fmt.Sprintf("%%-%ds", longest_entry)
		for i, e := range filtered_entries {
			name := e.Name()
			if args.quote_name {
				name = fmt.Sprintf("\"%s\"", e.Name())
			}
			out.WriteString(fmt.Sprintf(formatted_string, name))
			if i%columns == columns-1 {
				out.WriteString("\n")
			}
		}
		fmt.Println(out.String())
	}
}

func filterEntries(entries *[]os.FileInfo, args *arg) []os.FileInfo {
	filtered_entries := make([]os.FileInfo, 0)
	for _, e := range *entries {
		if !args.almost_all && strings.HasPrefix(e.Name(), ".") {
			continue
		}
		filtered_entries = append(filtered_entries, e)
	}

	return filtered_entries
}

func main() {
	args := arg{}
	reached_files := false

	for i := 1; i < len(os.Args); i++ {
		if reached_files == false {
			if os.Args[i] == "-h" || os.Args[i] == "--help" {
				help()
			}
			if os.Args[i] == "-A" {
				args.almost_all = true
				continue
			}
			if os.Args[i] == "-m" {
				args.comma_separated = true
				continue
			}
			if os.Args[i] == "-Q" || os.Args[i] == "--quote-name" {
				args.quote_name = true
				continue
			}
			if os.Args[i] == "-1" {
				args.one_per_line = true
				continue
			}
			if os.Args[i] == "--" {
				reached_files = true
				continue
			}
			if os.Args[i] == "-" {
				reached_files = true
			}
		}
		reached_files = true
		arg_v := os.Args[i]
		args.file = append(args.file, arg_v)
	}

	if len(args.file) == 0 {
		ls("./", args)
	} else {
		for i := range args.file {
			ls(args.file[i], args)
		}
	}
}
<|endoftext|>


<commit_before>package goopenzwave

// #include "gzw_manager.h"
// #include "gzw_notification.h"
// #include <stdlib.h>
import "C"
import (
	"fmt"
	"unsafe"
)

var (
	cmanager C.manager_t
)

// createManager creates the Manager singleton object. The Manager provides the
// public interface to OpenZWave, exposing all the functionality required to add
// Z-Wave support to an application. There can be only one Manager in an
// OpenZWave application. An Options object must be created and Locked first,
// otherwise the call to Manager::Create will fail. Once the Manager has been
// created, call AddWatcher to install a notification callback handler, and then
// call the AddDriver method for each attached PC Z-Wave controller in turn.
func createManager() error {
	cmanager = C.manager_create()
	if cmanager == nil {
		return fmt.Errorf("libopenzwave returned NULL pointer")
	}
	return nil
}

// getManager gets a pointer to the Manager object.
func getManager() C.manager_t {
	return cmanager
}

// destroyManager deletes the Manager and cleans up any associated objects.
func destroyManager() {
	C.manager_destroy()
}

// GetManagerVersionAsString Get the Version Number of OZW as a string.
func GetManagerVersionAsString() string {
	cstr := C.manager_getVersionAsString()
	defer C.free(unsafe.Pointer(cstr))
	return C.GoString(cstr)
}

// GetManagerVersionLongAsString Get the Version Number including Git commit of OZW as a string.
func GetManagerVersionLongAsString() string {
	cstr := C.manager_getVersionLongAsString()
	defer C.free(unsafe.Pointer(cstr))
	return C.GoString(cstr)
}

// ManagerVersion represents the OpenZWave library version as major and minor
// integers.
type ManagerVersion struct {
	Major int
	Minor int
}

// GetManagerVersion Get the Version Number as the Version Struct (Only Major/Minor returned).
func GetManagerVersion() ManagerVersion {
	var cMajor C.uint16_t
	var cMinor C.uint16_t
	C.manager_getVersion(&cMajor, &cMinor)
	return ManagerVersion{
		Major: int(cMajor),
		Minor: int(cMinor),
	}
}
<commit_msg>updated manager method names<commit_after>package goopenzwave

// #include "gzw_manager.h"
// #include "gzw_notification.h"
// #include <stdlib.h>
import "C"
import (
	"fmt"
	"unsafe"
)

var (
	cmanager C.manager_t
)

// createManager creates the Manager singleton object. The Manager provides the
// public interface to OpenZWave, exposing all the functionality required to add
// Z-Wave support to an application. There can be only one Manager in an
// OpenZWave application. An Options object must be created and Locked first,
// otherwise the call to Manager::Create will fail. Once the Manager has been
// created, call AddWatcher to install a notification callback handler, and then
// call the AddDriver method for each attached PC Z-Wave controller in turn.
func createManager() error {
	cmanager = C.manager_create()
	if cmanager == nil {
		return fmt.Errorf("libopenzwave returned NULL pointer")
	}
	return nil
}

// getManager gets a pointer to the Manager object.
func getManager() C.manager_t {
	return cmanager
}

// destroyManager deletes the Manager and cleans up any associated objects.
func destroyManager() {
	C.manager_destroy()
}

// GetVersionAsString returns the Version Number of OZW as a string.
func GetVersionAsString() string {
	cstr := C.manager_getVersionAsString()
	defer C.free(unsafe.Pointer(cstr))
	return C.GoString(cstr)
}

// GetVersionLongAsString returns the Version Number including Git commit of OZW
// as a string.
func GetVersionLongAsString() string {
	cstr := C.manager_getVersionLongAsString()
	defer C.free(unsafe.Pointer(cstr))
	return C.GoString(cstr)
}

// Version represents the OpenZWave library version as major and minor integers.
type Version struct {
	Major int
	Minor int
}

// GetVersion returns the Version Number as the Version Struct (Only Major/Minor
// returned).
func GetVersion() Version {
	var cMajor C.uint16_t
	var cMinor C.uint16_t
	C.manager_getVersion(&cMajor, &cMinor)
	return Version{
		Major: int(cMajor),
		Minor: int(cMinor),
	}
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"math"
	"net"
	"os"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/gosuri/uiprogress"
	log "github.com/sirupsen/logrus"
)

type Connection struct {
	Server              string
	File                FileMetaData
	NumberOfConnections int
	Code                string
	HashedCode          string
	IsSender            bool
	Debug               bool
	DontEncrypt         bool
	bars                []*uiprogress.Bar
}

type FileMetaData struct {
	Name  string
	Size  int
	Hash  string
	IV    string
	Salt  string
	bytes []byte
}

func NewConnection(flags *Flags) *Connection {
	c := new(Connection)
	c.Debug = flags.Debug
	c.DontEncrypt = flags.DontEncrypt
	c.Server = flags.Server
	c.Code = flags.Code
	c.NumberOfConnections = flags.NumberOfConnections
	if len(flags.File) > 0 {
		c.File.Name = flags.File
		c.IsSender = true
	} else {
		c.IsSender = false
	}
	return c
}

func (c *Connection) Run() {
	if len(c.Code) == 0 {
		if !c.IsSender {
			c.Code = getInput("Enter receive code: ")
		}
		if len(c.Code) < 5 {
			c.Code = GetRandomName()
		}
	}

	log.SetFormatter(&log.TextFormatter{})
	if c.Debug {
		log.SetLevel(log.DebugLevel)
	} else {
		log.SetLevel(log.WarnLevel)
	}

	if c.IsSender {
		// encrypt the file
		log.Debug("encrypting...")
		fdata, err := ioutil.ReadFile(c.File.Name)
		if err != nil {
			log.Fatal(err)
			return
		}
		c.File.bytes, c.File.Salt, c.File.IV = Encrypt(fdata, c.Code, c.DontEncrypt)
		log.Debug("...finished encryption")
		c.File.Hash = HashBytes(fdata)
		c.File.Size = len(c.File.bytes)
		if c.Debug {
			ioutil.WriteFile(c.File.Name+".encrypted", c.File.bytes, 0644)
		}
		fmt.Printf("Sending %d byte file named '%s'\n", c.File.Size, c.File.Name)
		fmt.Printf("Code is: %s\n", c.Code)
	}

	c.runClient()
}

// runClient spawns threads for parallel uplink/downlink via TCP
func (c *Connection) runClient() {
	logger := log.WithFields(log.Fields{
		"code":    c.Code,
		"sender?": c.IsSender,
	})

	c.HashedCode = Hash(c.Code)

	var wg sync.WaitGroup
	wg.Add(c.NumberOfConnections)

	uiprogress.Start()
	if !c.Debug {
		c.bars = make([]*uiprogress.Bar, c.NumberOfConnections)
	}
	gotOK := false
	for id := 0; id < c.NumberOfConnections; id++ {
		go func(id int) {
			defer wg.Done()
			port := strconv.Itoa(27001 + id)
			connection, err := net.Dial("tcp", c.Server+":"+port)
			if err != nil {
				panic(err)
			}
			defer connection.Close()

			message := receiveMessage(connection)
			logger.Debugf("relay says: %s", message)
			if c.IsSender {
				logger.Debugf("telling relay: %s", "s."+c.Code)
				metaData, err := json.Marshal(c.File)
				if err != nil {
					log.Error(err)
				}
				encryptedMetaData, salt, iv := Encrypt(metaData, c.Code)
				sendMessage("s."+c.HashedCode+"."+hex.EncodeToString(encryptedMetaData)+"-"+salt+"-"+iv, connection)
			} else {
				logger.Debugf("telling relay: %s", "r."+c.Code)
				sendMessage("r."+c.HashedCode+".0.0.0", connection)
			}
			if c.IsSender { // this is a sender
				if id == 0 {
					fmt.Printf("\nSending (<-%s)..\n", connection.RemoteAddr().String())
				}
				logger.Debug("waiting for ok from relay")
				message = receiveMessage(connection)
				logger.Debug("got ok from relay")
				// wait for pipe to be made
				time.Sleep(100 * time.Millisecond)
				// Write data from file
				logger.Debug("send file")
				c.sendFile(id, connection)
			} else { // this is a receiver
				logger.Debug("waiting for meta data from sender")
				message = receiveMessage(connection)
				m := strings.Split(message, "-")
				encryptedData, salt, iv := m[0], m[1], m[2]
				encryptedBytes, err := hex.DecodeString(encryptedData)
				if err != nil {
					log.Error(err)
					return
				}
				decryptedBytes, _ := Decrypt(encryptedBytes, c.Code, salt, iv, c.DontEncrypt)
				err = json.Unmarshal(decryptedBytes, &c.File)
				if err != nil {
					log.Error(err)
					return
				}
				log.Debugf("meta data received: %v", c.File)
				// have the main thread ask for the okay
				if id == 0 {
					fmt.Printf("Receiving file (%d bytes) into: %s\n", c.File.Size, c.File.Name)
					getOk := getInput("ok? (y/n): ")
					if getOk == "y" {
						gotOK = true
					} else {
						return
					}
				}
				// wait for the main thread to get the okay
				for limit := 0; limit < 1000; limit++ {
					if gotOK {
						break
					}
					time.Sleep(10 * time.Millisecond)
				}
				if !gotOK {
					return
				}
				sendMessage("ok", connection)
				logger.Debug("receive file")
				c.receiveFile(id, connection)
			}
		}(id)
	}
	wg.Wait()

	if !c.IsSender {
		c.catFile(c.File.Name)
		encrypted, err := ioutil.ReadFile(c.File.Name + ".encrypted")
		if err != nil {
			log.Error(err)
			return
		}
		fmt.Println("\n\ndecrypting...")
		log.Debugf("Code: [%s]", c.Code)
		log.Debugf("Salt: [%s]", c.File.Salt)
		log.Debugf("IV: [%s]", c.File.IV)
		decrypted, err := Decrypt(encrypted, c.Code, c.File.Salt, c.File.IV, c.DontEncrypt)
		if err != nil {
			log.Error(err)
			return
		}
		log.Debugf("writing %d bytes to %s", len(decrypted), c.File.Name)
		err = ioutil.WriteFile(c.File.Name, decrypted, 0644)
		if err != nil {
			log.Error(err)
		}
		if !c.Debug {
			os.Remove(c.File.Name + ".encrypted")
		}
		log.Debugf("\n\n\ndownloaded hash: [%s]", HashBytes(decrypted))
		log.Debugf("\n\n\nrelayed hash: [%s]", c.File.Hash)

		if c.File.Hash != HashBytes(decrypted) {
			fmt.Printf("\nUh oh! %s is corrupted! Sorry, try again.\n", c.File.Name)
		} else {
			fmt.Printf("\nDownloaded %s!", c.File.Name)
		}
	}
}

func (c *Connection) catFile(fname string) {
	// cat the file
	os.Remove(fname)
	finished, err := os.Create(fname + ".encrypted")
	defer finished.Close()
	if err != nil {
		log.Fatal(err)
	}
	for id := 0; id < c.NumberOfConnections; id++ {
		fh, err := os.Open(fname + "." + strconv.Itoa(id))
		if err != nil {
			log.Fatal(err)
		}

		_, err = io.Copy(finished, fh)
		if err != nil {
			log.Fatal(err)
		}
		fh.Close()
		os.Remove(fname + "." + strconv.Itoa(id))
	}

}

func (c *Connection) receiveFile(id int, connection net.Conn) error {
	logger := log.WithFields(log.Fields{
		"function": "receiveFile #" + strconv.Itoa(id),
	})

	logger.Debug("waiting for chunk size from sender")
	fileSizeBuffer := make([]byte, 10)
	connection.Read(fileSizeBuffer)
	fileDataString := strings.Trim(string(fileSizeBuffer), ":")
	fileSizeInt, _ := strconv.Atoi(fileDataString)
	chunkSize := int64(fileSizeInt)
	logger.Debugf("chunk size: %d", chunkSize)

	os.Remove(c.File.Name + "." + strconv.Itoa(id))
	newFile, err := os.Create(c.File.Name + "." + strconv.Itoa(id))
	if err != nil {
		panic(err)
	}
	defer newFile.Close()

	if !c.Debug {
		c.bars[id] = uiprogress.AddBar(int(chunkSize)/1024 + 1).AppendCompleted().PrependElapsed()
	}

	logger.Debug("waiting for file")
	var receivedBytes int64
	for {
		if !c.Debug {
			c.bars[id].Incr()
		}
		if (chunkSize - receivedBytes) < BUFFERSIZE {
			logger.Debug("at the end")
			io.CopyN(newFile, connection, (chunkSize - receivedBytes))
			// Empty the remaining bytes that we don't need from the network buffer
			if (receivedBytes+BUFFERSIZE)-chunkSize < BUFFERSIZE {
				logger.Debug("empty remaining bytes from network buffer")
				connection.Read(make([]byte, (receivedBytes+BUFFERSIZE)-chunkSize))
			}
			break
		}
		io.CopyN(newFile, connection, BUFFERSIZE)
		receivedBytes += BUFFERSIZE
	}
	logger.Debug("received file")
	return nil
}

func (c *Connection) sendFile(id int, connection net.Conn) {
	logger := log.WithFields(log.Fields{
		"function": "sendFile #" + strconv.Itoa(id),
	})
	defer connection.Close()

	var err error

	numChunks := math.Ceil(float64(c.File.Size) / float64(BUFFERSIZE))
	chunksPerWorker := int(math.Ceil(numChunks / float64(c.NumberOfConnections)))

	chunkSize := int64(chunksPerWorker * BUFFERSIZE)
	if id+1 == c.NumberOfConnections {
		chunkSize = int64(c.File.Size) - int64(c.NumberOfConnections-1)*chunkSize
	}

	if id == 0 || id == c.NumberOfConnections-1 {
		logger.Debugf("numChunks: %v", numChunks)
		logger.Debugf("chunksPerWorker: %v", chunksPerWorker)
		logger.Debugf("bytesPerchunkSizeConnection: %v", chunkSize)
	}

	logger.Debugf("sending chunk size: %d", chunkSize)
	connection.Write([]byte(fillString(strconv.FormatInt(int64(chunkSize), 10), 10)))

	sendBuffer := make([]byte, BUFFERSIZE)
	file := bytes.NewBuffer(c.File.bytes)
	chunkI := 0
	if !c.Debug {
		c.bars[id] = uiprogress.AddBar(chunksPerWorker).AppendCompleted().PrependElapsed()
	}
	for {
		_, err = file.Read(sendBuffer)
		if err == io.EOF {
			//End of file reached, break out of for loop
			logger.Debug("EOF")
			break
		}
		if (chunkI >= chunksPerWorker*id && chunkI < chunksPerWorker*id+chunksPerWorker) || (id == c.NumberOfConnections-1 && chunkI >= chunksPerWorker*id) {
			connection.Write(sendBuffer)
			if !c.Debug {
				c.bars[id].Incr()
			}
		}
		chunkI++
	}
	logger.Debug("file is sent")
	return
}
<commit_msg>Better ui<commit_after>package main

import (
	"bytes"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"math"
	"net"
	"os"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/gosuri/uiprogress"
	log "github.com/sirupsen/logrus"
)

type Connection struct {
	Server              string
	File                FileMetaData
	NumberOfConnections int
	Code                string
	HashedCode          string
	IsSender            bool
	Debug               bool
	DontEncrypt         bool
	bars                []*uiprogress.Bar
}

type FileMetaData struct {
	Name  string
	Size  int
	Hash  string
	IV    string
	Salt  string
	bytes []byte
}

func NewConnection(flags *Flags) *Connection {
	c := new(Connection)
	c.Debug = flags.Debug
	c.DontEncrypt = flags.DontEncrypt
	c.Server = flags.Server
	c.Code = flags.Code
	c.NumberOfConnections = flags.NumberOfConnections
	if len(flags.File) > 0 {
		c.File.Name = flags.File
		c.IsSender = true
	} else {
		c.IsSender = false
	}
	return c
}

func (c *Connection) Run() {
	if len(c.Code) == 0 {
		if !c.IsSender {
			c.Code = getInput("Enter receive code: ")
		}
		if len(c.Code) < 5 {
			c.Code = GetRandomName()
		}
	}

	log.SetFormatter(&log.TextFormatter{})
	if c.Debug {
		log.SetLevel(log.DebugLevel)
	} else {
		log.SetLevel(log.WarnLevel)
	}

	if c.IsSender {
		// encrypt the file
		log.Debug("encrypting...")
		fdata, err := ioutil.ReadFile(c.File.Name)
		if err != nil {
			log.Fatal(err)
			return
		}
		c.File.bytes, c.File.Salt, c.File.IV = Encrypt(fdata, c.Code, c.DontEncrypt)
		log.Debug("...finished encryption")
		c.File.Hash = HashBytes(fdata)
		c.File.Size = len(c.File.bytes)
		if c.Debug {
			ioutil.WriteFile(c.File.Name+".encrypted", c.File.bytes, 0644)
		}
		fmt.Printf("Sending %d byte file named '%s'\n", c.File.Size, c.File.Name)
		fmt.Printf("Code is: %s\n", c.Code)
	}

	c.runClient()
}

// runClient spawns threads for parallel uplink/downlink via TCP
func (c *Connection) runClient() {
	logger := log.WithFields(log.Fields{
		"code":    c.Code,
		"sender?": c.IsSender,
	})

	c.HashedCode = Hash(c.Code)

	var wg sync.WaitGroup
	wg.Add(c.NumberOfConnections)

	uiprogress.Start()
	if !c.Debug {
		c.bars = make([]*uiprogress.Bar, c.NumberOfConnections)
	}
	gotOK := false
	for id := 0; id < c.NumberOfConnections; id++ {
		go func(id int) {
			defer wg.Done()
			port := strconv.Itoa(27001 + id)
			connection, err := net.Dial("tcp", c.Server+":"+port)
			if err != nil {
				panic(err)
			}
			defer connection.Close()

			message := receiveMessage(connection)
			logger.Debugf("relay says: %s", message)
			if c.IsSender {
				logger.Debugf("telling relay: %s", "s."+c.Code)
				metaData, err := json.Marshal(c.File)
				if err != nil {
					log.Error(err)
				}
				encryptedMetaData, salt, iv := Encrypt(metaData, c.Code)
				sendMessage("s."+c.HashedCode+"."+hex.EncodeToString(encryptedMetaData)+"-"+salt+"-"+iv, connection)
			} else {
				logger.Debugf("telling relay: %s", "r."+c.Code)
				sendMessage("r."+c.HashedCode+".0.0.0", connection)
			}
			if c.IsSender { // this is a sender
				if id == 0 {
					fmt.Printf("\nSending (<-%s)..\n", connection.RemoteAddr().String())
				}
				logger.Debug("waiting for ok from relay")
				message = receiveMessage(connection)
				logger.Debug("got ok from relay")
				// wait for pipe to be made
				time.Sleep(100 * time.Millisecond)
				// Write data from file
				logger.Debug("send file")
				c.sendFile(id, connection)
			} else { // this is a receiver
				logger.Debug("waiting for meta data from sender")
				message = receiveMessage(connection)
				m := strings.Split(message, "-")
				encryptedData, salt, iv := m[0], m[1], m[2]
				encryptedBytes, err := hex.DecodeString(encryptedData)
				if err != nil {
					log.Error(err)
					return
				}
				decryptedBytes, _ := Decrypt(encryptedBytes, c.Code, salt, iv, c.DontEncrypt)
				err = json.Unmarshal(decryptedBytes, &c.File)
				if err != nil {
					log.Error(err)
					return
				}
				log.Debugf("meta data received: %v", c.File)
				// have the main thread ask for the okay
				if id == 0 {
					fmt.Printf("Receiving file (%d bytes) into: %s\n", c.File.Size, c.File.Name)
					getOk := getInput("ok? (y/n): ")
					if getOk == "y" {
						gotOK = true
					} else {
						return
					}
				}
				// wait for the main thread to get the okay
				for limit := 0; limit < 1000; limit++ {
					if gotOK {
						break
					}
					time.Sleep(10 * time.Millisecond)
				}
				if !gotOK {
					return
				}
				sendMessage("ok", connection)
				logger.Debug("receive file")
				c.receiveFile(id, connection)
			}
		}(id)
	}
	wg.Wait()

	if !c.IsSender {
		c.catFile(c.File.Name)
		encrypted, err := ioutil.ReadFile(c.File.Name + ".encrypted")
		if err != nil {
			log.Error(err)
			return
		}
		fmt.Println("\n\ndecrypting...")
		log.Debugf("Code: [%s]", c.Code)
		log.Debugf("Salt: [%s]", c.File.Salt)
		log.Debugf("IV: [%s]", c.File.IV)
		decrypted, err := Decrypt(encrypted, c.Code, c.File.Salt, c.File.IV, c.DontEncrypt)
		if err != nil {
			log.Error(err)
			return
		}
		log.Debugf("writing %d bytes to %s", len(decrypted), c.File.Name)
		err = ioutil.WriteFile(c.File.Name, decrypted, 0644)
		if err != nil {
			log.Error(err)
		}
		if !c.Debug {
			os.Remove(c.File.Name + ".encrypted")
		}
		log.Debugf("\n\n\ndownloaded hash: [%s]", HashBytes(decrypted))
		log.Debugf("\n\n\nrelayed hash: [%s]", c.File.Hash)

		if c.File.Hash != HashBytes(decrypted) {
			fmt.Printf("\nUh oh! %s is corrupted! Sorry, try again.\n", c.File.Name)
		} else {
			fmt.Printf("\nReceived file written to %s", c.File.Name)
		}
	} else {
		fmt.Println("File sent.")
		// TODO: Add confirmation
	}
}

func (c *Connection) catFile(fname string) {
	// cat the file
	os.Remove(fname)
	finished, err := os.Create(fname + ".encrypted")
	defer finished.Close()
	if err != nil {
		log.Fatal(err)
	}
	for id := 0; id < c.NumberOfConnections; id++ {
		fh, err := os.Open(fname + "." + strconv.Itoa(id))
		if err != nil {
			log.Fatal(err)
		}

		_, err = io.Copy(finished, fh)
		if err != nil {
			log.Fatal(err)
		}
		fh.Close()
		os.Remove(fname + "." + strconv.Itoa(id))
	}

}

func (c *Connection) receiveFile(id int, connection net.Conn) error {
	logger := log.WithFields(log.Fields{
		"function": "receiveFile #" + strconv.Itoa(id),
	})

	logger.Debug("waiting for chunk size from sender")
	fileSizeBuffer := make([]byte, 10)
	connection.Read(fileSizeBuffer)
	fileDataString := strings.Trim(string(fileSizeBuffer), ":")
	fileSizeInt, _ := strconv.Atoi(fileDataString)
	chunkSize := int64(fileSizeInt)
	logger.Debugf("chunk size: %d", chunkSize)

	os.Remove(c.File.Name + "." + strconv.Itoa(id))
	newFile, err := os.Create(c.File.Name + "." + strconv.Itoa(id))
	if err != nil {
		panic(err)
	}
	defer newFile.Close()

	if !c.Debug {
		c.bars[id] = uiprogress.AddBar(int(chunkSize)/1024 + 1).AppendCompleted().PrependElapsed()
	}

	logger.Debug("waiting for file")
	var receivedBytes int64
	for {
		if !c.Debug {
			c.bars[id].Incr()
		}
		if (chunkSize - receivedBytes) < BUFFERSIZE {
			logger.Debug("at the end")
			io.CopyN(newFile, connection, (chunkSize - receivedBytes))
			// Empty the remaining bytes that we don't need from the network buffer
			if (receivedBytes+BUFFERSIZE)-chunkSize < BUFFERSIZE {
				logger.Debug("empty remaining bytes from network buffer")
				connection.Read(make([]byte, (receivedBytes+BUFFERSIZE)-chunkSize))
			}
			break
		}
		io.CopyN(newFile, connection, BUFFERSIZE)
		receivedBytes += BUFFERSIZE
	}
	logger.Debug("received file")
	return nil
}

func (c *Connection) sendFile(id int, connection net.Conn) {
	logger := log.WithFields(log.Fields{
		"function": "sendFile #" + strconv.Itoa(id),
	})
	defer connection.Close()

	var err error

	numChunks := math.Ceil(float64(c.File.Size) / float64(BUFFERSIZE))
	chunksPerWorker := int(math.Ceil(numChunks / float64(c.NumberOfConnections)))

	chunkSize := int64(chunksPerWorker * BUFFERSIZE)
	if id+1 == c.NumberOfConnections {
		chunkSize = int64(c.File.Size) - int64(c.NumberOfConnections-1)*chunkSize
	}

	if id == 0 || id == c.NumberOfConnections-1 {
		logger.Debugf("numChunks: %v", numChunks)
		logger.Debugf("chunksPerWorker: %v", chunksPerWorker)
		logger.Debugf("bytesPerchunkSizeConnection: %v", chunkSize)
	}

	logger.Debugf("sending chunk size: %d", chunkSize)
	connection.Write([]byte(fillString(strconv.FormatInt(int64(chunkSize), 10), 10)))

	sendBuffer := make([]byte, BUFFERSIZE)
	file := bytes.NewBuffer(c.File.bytes)
	chunkI := 0
	if !c.Debug {
		c.bars[id] = uiprogress.AddBar(chunksPerWorker).AppendCompleted().PrependElapsed()
	}
	for {
		_, err = file.Read(sendBuffer)
		if err == io.EOF {
			//End of file reached, break out of for loop
			logger.Debug("EOF")
			break
		}
		if (chunkI >= chunksPerWorker*id && chunkI < chunksPerWorker*id+chunksPerWorker) || (id == c.NumberOfConnections-1 && chunkI >= chunksPerWorker*id) {
			connection.Write(sendBuffer)
			if !c.Debug {
				c.bars[id].Incr()
			}
		}
		chunkI++
	}
	logger.Debug("file is sent")
	return
}
<|endoftext|>


<commit_before>//Package gokeepasslib is a library written in go which provides functionality to decrypt and parse keepass 2 files (kdbx)
package gokeepasslib

import (
	"encoding/xml"
	"time"
)
//Container for all elements of a keepass database
type DBContent struct {
	XMLName xml.Name  `xml:"KeePassFile"`
	Meta    *MetaData `xml:"Meta"`
	Root    *RootData `xml:"Root"`
}

//The metadata headers at the top of kdbx files, continas things like the name of the database
type MetaData struct {
	Generator                  string        `xml:"Generator"`
	HeaderHash                 string        `xml:"HeaderHash"`
	DatabaseName               string        `xml:"DatabaseName"`
	DatabaseNameChanged        *time.Time    `xml:"DatabaseNameChanged"`
	DatabaseDescription        string        `xml:"DatabaseDescription"`
	DatabaseDescriptionChanged *time.Time    `xml:"DatabaseDescriptionChanged"`
	DefaultUserName            string        `xml:"DefaultUserName"`
	DefaultUserNameChanged     *time.Time    `xml:"DefaultUserNameChanged"`
	MaintenanceHistoryDays     string        `xml:"MaintenanceHistoryDays"`
	Color                      string        `xml:"Color"`
	MasterKeyChanged           *time.Time    `xml:"MasterKeyChanged"`
	MasterKeyChangeRec         int64         `xml:"MasterKeyChangeRec"`
	MasterKeyChangeForce       int64         `xml:"MasterKeyChangeForce"`
	MemoryProtection           MemProtection `xml:"MemoryProtection"`
	RecycleBinEnabled          boolWrapper   `xml:"RecycleBinEnabled"`
	RecycleBinUUID             string        `xml:"RecycleBinUUID"`
	RecycleBinChanged          *time.Time    `xml:"RecycleBinChanged"`
	EntryTemplatesGroup        string        `xml:"EntryTemplatesGroup"`
	EntryTemplatesGroupChanged *time.Time    `xml:"EntryTemplatesGroupChanged"`
	HistoryMaxItems            int64         `xml:"HistoryMaxItems"`
	HistoryMaxSize             int64         `xml:"HistoryMaxSize"`
	LastSelectedGroup          string        `xml:"LastSelectedGroup"`
	LastTopVisibleGroup        string        `xml:"LastTopVisibleGroup"`
	Binaries                   string        `xml:"Binaries"`
	CustomData                 string        `xml:"CustomData"`
}

type MemProtection struct {
	ProtectTitle    boolWrapper `xml:"ProtectTitle"`
	ProtectUserName boolWrapper `xml:"ProtectUserName"`
	ProtectPassword boolWrapper `xml:"ProtectPassword"`
	ProtectURL      boolWrapper `xml:"ProtectURL"`
	ProtectNotes    boolWrapper `xml:"ProtectNotes"`
}

//Stores the actual content of a database (all enteries sorted into groups and the recycle bin)
type RootData struct {
	Groups         []Group             `xml:"Group"`
	DeletedObjects []DeletedObjectData `xml:"DeletedObjects>DeletedObject"`
}

//Structure to store entries in their named groups for organization
type Group struct {
	UUID                    string      `xml:"UUID"`
	Name                    string      `xml:"Name"`
	Notes                   string      `xml:"Notes"`
	IconID                  int64       `xml:"IconID"`
	Times                   TimeData    `xml:"Times"`
	IsExpanded              boolWrapper `xml:"IsExpanded"`
	DefaultAutoTypeSequence string      `xml:"DefaultAutoTypeSequence"`
	EnableAutoType          string      `xml:"EnableAutoType"`
	EnableSearching         string      `xml:"EnableSearching"`
	LastTopVisibleEntry     string      `xml:"LastTopVisibleEntry"`
	Groups                  []Group     `xml:"Group,omitempty"`
	Entries                 []Entry     `xml:"Entry,omitempty"`
}

//All metadata relating to times for groups and entries, such as last modification time
type TimeData struct {
	CreationTime         *time.Time  `xml:"CreationTime"`
	LastModificationTime *time.Time  `xml:"LastModificationTime"`
	LastAcessTime        *time.Time  `xml:"LastAcessTime"`
	ExpiryTime           *time.Time  `xml:"ExpiryTime"`
	Expires              boolWrapper `xml:"Expires"`
	UsageCount           int64       `xml:"UsageCount"`
	LocationChanged      *time.Time  `xml:"LocationChanged"`
}
//structure for each parsed entry in a keepass database
type Entry struct {
	UUID            string       `xml:"UUID"`
	IconID          int64        `xml:"IconID"`
	ForegroundColor string       `xml:"ForegroundColor"`
	BackgroundColor string       `xml:"BackgroundColor"`
	OverrideURL     string       `xml:"OverrideURL"`
	Tags            string       `xml:"Tags"`
	Times           TimeData     `xml:"Times"`
	Values          []ValueData  `xml:"String,omitempty"`
	AutoType        AutoTypeData `xml:"AutoType"`
	Histories       []History    `xml:"History"`
	Password        []byte       `xml:"-"`
}

func (e *Entry) protected() bool {
	for _, v := range e.Values {
		if v.Key == "Password" && bool(v.Value.Protected) {
			return true
		}
	}
	return false
}

//Gets the value in e corresponding with key k, or an empty string otherwise
func (e *Entry) get(k string) string {
	var val string
	for _, v := range e.Values {
		if v.Key == k {
			val = v.Value.Content
		}
	}
	return val
}
func (e *Entry) getPassword() string {
	return e.get("password")
}
func (e *Entry) getPasswordIndex() int {
	for i, v := range e.Values {
		if v.Key == "Password" {
			return i
		}
	}
	return 0
}

func (e *Entry) GetTitle() string {
	return e.get("Title")
}

type History struct {
	Entries []Entry `xml:"Entry"`
}

type ValueData struct {
	Key   string `xml:"Key"`
	Value V      `xml:"Value"`
}

type V struct {
	Content   string      `xml:",innerxml"`
	Protected boolWrapper `xml:"Protected,attr,omitempty"`
}

type AutoTypeData struct {
	Enabled                 boolWrapper         `xml:"Enabled"`
	DataTransferObfuscation int64               `xml:"DataTransferObfuscation"`
	Association             AutoTypeAssociation `xml:"Association"`
}

type AutoTypeAssociation struct {
	Window            string `xml:"Window"`
	KeystrokeSequence string `xml:"KeystrokeSequence"`
}

type DeletedObjectData struct {
	XMLName      xml.Name   `xml:"DeletedObject"`
	UUID         string     `xml:"UUID"`
	DeletionTime *time.Time `xml:"DeletionTime"`
}
<commit_msg>fixed typo<commit_after>//Package gokeepasslib is a library written in go which provides functionality to decrypt and parse keepass 2 files (kdbx)
package gokeepasslib

import (
	"encoding/xml"
	"time"
)
//Container for all elements of a keepass database
type DBContent struct {
	XMLName xml.Name  `xml:"KeePassFile"`
	Meta    *MetaData `xml:"Meta"`
	Root    *RootData `xml:"Root"`
}

//The metadata headers at the top of kdbx files, continas things like the name of the database
type MetaData struct {
	Generator                  string        `xml:"Generator"`
	HeaderHash                 string        `xml:"HeaderHash"`
	DatabaseName               string        `xml:"DatabaseName"`
	DatabaseNameChanged        *time.Time    `xml:"DatabaseNameChanged"`
	DatabaseDescription        string        `xml:"DatabaseDescription"`
	DatabaseDescriptionChanged *time.Time    `xml:"DatabaseDescriptionChanged"`
	DefaultUserName            string        `xml:"DefaultUserName"`
	DefaultUserNameChanged     *time.Time    `xml:"DefaultUserNameChanged"`
	MaintenanceHistoryDays     string        `xml:"MaintenanceHistoryDays"`
	Color                      string        `xml:"Color"`
	MasterKeyChanged           *time.Time    `xml:"MasterKeyChanged"`
	MasterKeyChangeRec         int64         `xml:"MasterKeyChangeRec"`
	MasterKeyChangeForce       int64         `xml:"MasterKeyChangeForce"`
	MemoryProtection           MemProtection `xml:"MemoryProtection"`
	RecycleBinEnabled          boolWrapper   `xml:"RecycleBinEnabled"`
	RecycleBinUUID             string        `xml:"RecycleBinUUID"`
	RecycleBinChanged          *time.Time    `xml:"RecycleBinChanged"`
	EntryTemplatesGroup        string        `xml:"EntryTemplatesGroup"`
	EntryTemplatesGroupChanged *time.Time    `xml:"EntryTemplatesGroupChanged"`
	HistoryMaxItems            int64         `xml:"HistoryMaxItems"`
	HistoryMaxSize             int64         `xml:"HistoryMaxSize"`
	LastSelectedGroup          string        `xml:"LastSelectedGroup"`
	LastTopVisibleGroup        string        `xml:"LastTopVisibleGroup"`
	Binaries                   string        `xml:"Binaries"`
	CustomData                 string        `xml:"CustomData"`
}

type MemProtection struct {
	ProtectTitle    boolWrapper `xml:"ProtectTitle"`
	ProtectUserName boolWrapper `xml:"ProtectUserName"`
	ProtectPassword boolWrapper `xml:"ProtectPassword"`
	ProtectURL      boolWrapper `xml:"ProtectURL"`
	ProtectNotes    boolWrapper `xml:"ProtectNotes"`
}

//Stores the actual content of a database (all enteries sorted into groups and the recycle bin)
type RootData struct {
	Groups         []Group             `xml:"Group"`
	DeletedObjects []DeletedObjectData `xml:"DeletedObjects>DeletedObject"`
}

//Structure to store entries in their named groups for organization
type Group struct {
	UUID                    string      `xml:"UUID"`
	Name                    string      `xml:"Name"`
	Notes                   string      `xml:"Notes"`
	IconID                  int64       `xml:"IconID"`
	Times                   TimeData    `xml:"Times"`
	IsExpanded              boolWrapper `xml:"IsExpanded"`
	DefaultAutoTypeSequence string      `xml:"DefaultAutoTypeSequence"`
	EnableAutoType          string      `xml:"EnableAutoType"`
	EnableSearching         string      `xml:"EnableSearching"`
	LastTopVisibleEntry     string      `xml:"LastTopVisibleEntry"`
	Groups                  []Group     `xml:"Group,omitempty"`
	Entries                 []Entry     `xml:"Entry,omitempty"`
}

//All metadata relating to times for groups and entries, such as last modification time
type TimeData struct {
	CreationTime         *time.Time  `xml:"CreationTime"`
	LastModificationTime *time.Time  `xml:"LastModificationTime"`
	LastAcessTime        *time.Time  `xml:"LastAcessTime"`
	ExpiryTime           *time.Time  `xml:"ExpiryTime"`
	Expires              boolWrapper `xml:"Expires"`
	UsageCount           int64       `xml:"UsageCount"`
	LocationChanged      *time.Time  `xml:"LocationChanged"`
}
//structure for each parsed entry in a keepass database
type Entry struct {
	UUID            string       `xml:"UUID"`
	IconID          int64        `xml:"IconID"`
	ForegroundColor string       `xml:"ForegroundColor"`
	BackgroundColor string       `xml:"BackgroundColor"`
	OverrideURL     string       `xml:"OverrideURL"`
	Tags            string       `xml:"Tags"`
	Times           TimeData     `xml:"Times"`
	Values          []ValueData  `xml:"String,omitempty"`
	AutoType        AutoTypeData `xml:"AutoType"`
	Histories       []History    `xml:"History"`
	Password        []byte       `xml:"-"`
}

func (e *Entry) protected() bool {
	for _, v := range e.Values {
		if v.Key == "Password" && bool(v.Value.Protected) {
			return true
		}
	}
	return false
}

//Gets the value in e corresponding with key k, or an empty string otherwise
func (e *Entry) get(k string) string {
	var val string
	for _, v := range e.Values {
		if v.Key == k {
			val = v.Value.Content
		}
	}
	return val
}
func (e *Entry) getPassword() string {
	return e.get("Password")
}

func (e *Entry) getPasswordIndex() int {
	for i, v := range e.Values {
		if v.Key == "Password" {
			return i
		}
	}
	return 0
}

func (e *Entry) GetTitle() string {
	return e.get("Title")
}

type History struct {
	Entries []Entry `xml:"Entry"`
}

type ValueData struct {
	Key   string `xml:"Key"`
	Value V      `xml:"Value"`
}

type V struct {
	Content   string      `xml:",innerxml"`
	Protected boolWrapper `xml:"Protected,attr,omitempty"`
}

type AutoTypeData struct {
	Enabled                 boolWrapper         `xml:"Enabled"`
	DataTransferObfuscation int64               `xml:"DataTransferObfuscation"`
	Association             AutoTypeAssociation `xml:"Association"`
}

type AutoTypeAssociation struct {
	Window            string `xml:"Window"`
	KeystrokeSequence string `xml:"KeystrokeSequence"`
}

type DeletedObjectData struct {
	XMLName      xml.Name   `xml:"DeletedObject"`
	UUID         string     `xml:"UUID"`
	DeletionTime *time.Time `xml:"DeletionTime"`
}
<|endoftext|>


<commit_before>package gokeepasslib

import (
	"bytes"
	"encoding/binary"
	"encoding/xml"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
)

// Inner header bytes
const (
	InnerHeaderTerminator byte = 0x00 // Inner header terminator byte
	InnerHeaderIRSID      byte = 0x01 // Inner header InnerRandomStreamID byte
	InnerHeaderIRSKey     byte = 0x02 // Inner header InnerRandomStreamKey byte
	InnerHeaderBinary     byte = 0x03 // Inner header binary byte
)

// DBContent is a container for all elements of a keepass database
type DBContent struct {
	RawData     []byte       `xml:"-"` // Encrypted data
	InnerHeader *InnerHeader `xml:"-"`
	XMLName     xml.Name     `xml:"KeePassFile"`
	Meta        *MetaData    `xml:"Meta"`
	Root        *RootData    `xml:"Root"`
}

// InnerHeader is the container of crypt options and binaries, only for Kdbx v4
type InnerHeader struct {
	InnerRandomStreamID  uint32
	InnerRandomStreamKey []byte
	Binaries             Binaries
}

// NewContent creates a new database content with some good defaults
func NewContent() *DBContent {
	// Not necessary create InnerHeader because this will be a KDBX v3.1
	return &DBContent{
		Meta: NewMetaData(),
		Root: NewRootData(),
	}
}

// readFrom reads the InnerHeader from an io.Reader
func (ih *InnerHeader) readFrom(r io.Reader) error {
	binaryCount := 0 // Var used to count and index every binary
	for {
		var typ byte
		var length int32
		var data []byte

		if err := binary.Read(r, binary.LittleEndian, &typ); err != nil {
			return err
		}
		if err := binary.Read(r, binary.LittleEndian, &length); err != nil {
			return err
		}
		data = make([]byte, length)
		if err := binary.Read(r, binary.LittleEndian, &data); err != nil {
			return err
		}

		if typ == InnerHeaderTerminator {
			// End of inner header
			break
		} else if typ == InnerHeaderIRSID {
			// Found InnerRandomStream ID
			ih.InnerRandomStreamID = binary.LittleEndian.Uint32(data)
		} else if typ == InnerHeaderIRSKey {
			// Found InnerRandomStream Key
			ih.InnerRandomStreamKey = data
		} else if typ == InnerHeaderBinary {
			// Found a binary
			var protection byte
			reader := bytes.NewReader(data)

			binary.Read(reader, binary.LittleEndian, &protection) // Read memory protection flag
			content, _ := ioutil.ReadAll(reader)                  // Read content

			ih.Binaries = append(ih.Binaries, Binary{
				ID:               binaryCount,
				MemoryProtection: protection,
				Content:          content,
			})

			binaryCount = binaryCount + 1
		} else {
			return ErrUnknownInnerHeaderID(typ)
		}
	}
	return nil
}

// writeTo the InnerHeader to the given io.Writer
func (ih *InnerHeader) writeTo(w io.Writer) error {
	irsID := make([]byte, 4)
	binary.LittleEndian.PutUint32(irsID, ih.InnerRandomStreamID)

	if err := writeToInnerHeader(w, InnerHeaderIRSID, irsID); err != nil {
		return err
	}
	if err := writeToInnerHeader(w, InnerHeaderIRSKey, ih.InnerRandomStreamKey); err != nil {
		return err
	}

	for _, item := range ih.Binaries {
		buf := []byte{item.MemoryProtection}
		buf = append(buf, item.Content...)
		if err := writeToInnerHeader(w, InnerHeaderBinary, buf); err != nil {
			return err
		}
	}
	// End inner header
	if err := binary.Write(w, binary.LittleEndian, uint8(InnerHeaderTerminator)); err != nil {
		return err
	}
	if err := binary.Write(w, binary.LittleEndian, uint32(0)); err != nil {
		return err
	}
	return nil
}

// writeToInnerHeader is an helper to write an inner header item to the given io.Writer
func writeToInnerHeader(w io.Writer, id uint8, data []byte) error {
	if len(data) > 0 {
		if err := binary.Write(w, binary.LittleEndian, id); err != nil {
			return err
		}
		if err := binary.Write(w, binary.LittleEndian, uint32(len(data))); err != nil {
			return err
		}
		if err := binary.Write(w, binary.LittleEndian, data); err != nil {
			return err
		}
	}
	return nil
}

func (ih InnerHeader) String() string {
	return fmt.Sprintf(
		"1) InnerRandomStreamID: %d\n"+
			"2) InnerRandomStreamKey: %x\n"+
			"3) Binaries: %s\n",
		ih.InnerRandomStreamID,
		ih.InnerRandomStreamKey,
		ih.Binaries,
	)
}

// ErrEndOfInnerHeaders is the error returned when the end of inner header is read
var ErrEndOfInnerHeaders = errors.New("gokeepasslib: inner header id was 0, end of inner headers")

// ErrUnknownInnerHeaderID is the error returned if an unknown inner header is read
type ErrUnknownInnerHeaderID byte

func (i ErrUnknownInnerHeaderID) Error() string {
	return fmt.Sprintf("gokeepasslib: unknown inner header ID of %x", i)
}
<commit_msg>Fixed return redundancy<commit_after>package gokeepasslib

import (
	"bytes"
	"encoding/binary"
	"encoding/xml"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
)

// Inner header bytes
const (
	InnerHeaderTerminator byte = 0x00 // Inner header terminator byte
	InnerHeaderIRSID      byte = 0x01 // Inner header InnerRandomStreamID byte
	InnerHeaderIRSKey     byte = 0x02 // Inner header InnerRandomStreamKey byte
	InnerHeaderBinary     byte = 0x03 // Inner header binary byte
)

// DBContent is a container for all elements of a keepass database
type DBContent struct {
	RawData     []byte       `xml:"-"` // Encrypted data
	InnerHeader *InnerHeader `xml:"-"`
	XMLName     xml.Name     `xml:"KeePassFile"`
	Meta        *MetaData    `xml:"Meta"`
	Root        *RootData    `xml:"Root"`
}

// InnerHeader is the container of crypt options and binaries, only for Kdbx v4
type InnerHeader struct {
	InnerRandomStreamID  uint32
	InnerRandomStreamKey []byte
	Binaries             Binaries
}

// NewContent creates a new database content with some good defaults
func NewContent() *DBContent {
	// Not necessary create InnerHeader because this will be a KDBX v3.1
	return &DBContent{
		Meta: NewMetaData(),
		Root: NewRootData(),
	}
}

// readFrom reads the InnerHeader from an io.Reader
func (ih *InnerHeader) readFrom(r io.Reader) error {
	binaryCount := 0 // Var used to count and index every binary
	for {
		var typ byte
		var length int32
		var data []byte

		if err := binary.Read(r, binary.LittleEndian, &typ); err != nil {
			return err
		}
		if err := binary.Read(r, binary.LittleEndian, &length); err != nil {
			return err
		}
		data = make([]byte, length)
		if err := binary.Read(r, binary.LittleEndian, &data); err != nil {
			return err
		}

		if typ == InnerHeaderTerminator {
			// End of inner header
			break
		} else if typ == InnerHeaderIRSID {
			// Found InnerRandomStream ID
			ih.InnerRandomStreamID = binary.LittleEndian.Uint32(data)
		} else if typ == InnerHeaderIRSKey {
			// Found InnerRandomStream Key
			ih.InnerRandomStreamKey = data
		} else if typ == InnerHeaderBinary {
			// Found a binary
			var protection byte
			reader := bytes.NewReader(data)

			binary.Read(reader, binary.LittleEndian, &protection) // Read memory protection flag
			content, _ := ioutil.ReadAll(reader)                  // Read content

			ih.Binaries = append(ih.Binaries, Binary{
				ID:               binaryCount,
				MemoryProtection: protection,
				Content:          content,
			})

			binaryCount = binaryCount + 1
		} else {
			return ErrUnknownInnerHeaderID(typ)
		}
	}
	return nil
}

// writeTo the InnerHeader to the given io.Writer
func (ih *InnerHeader) writeTo(w io.Writer) error {
	irsID := make([]byte, 4)
	binary.LittleEndian.PutUint32(irsID, ih.InnerRandomStreamID)

	if err := writeToInnerHeader(w, InnerHeaderIRSID, irsID); err != nil {
		return err
	}
	if err := writeToInnerHeader(w, InnerHeaderIRSKey, ih.InnerRandomStreamKey); err != nil {
		return err
	}

	for _, item := range ih.Binaries {
		buf := []byte{item.MemoryProtection}
		buf = append(buf, item.Content...)
		if err := writeToInnerHeader(w, InnerHeaderBinary, buf); err != nil {
			return err
		}
	}
	// End inner header
	if err := binary.Write(w, binary.LittleEndian, uint8(InnerHeaderTerminator)); err != nil {
		return err
	}
	return binary.Write(w, binary.LittleEndian, uint32(0))
}

// writeToInnerHeader is an helper to write an inner header item to the given io.Writer
func writeToInnerHeader(w io.Writer, id uint8, data []byte) error {
	if len(data) > 0 {
		if err := binary.Write(w, binary.LittleEndian, id); err != nil {
			return err
		}
		if err := binary.Write(w, binary.LittleEndian, uint32(len(data))); err != nil {
			return err
		}
		if err := binary.Write(w, binary.LittleEndian, data); err != nil {
			return err
		}
	}
	return nil
}

func (ih InnerHeader) String() string {
	return fmt.Sprintf(
		"1) InnerRandomStreamID: %d\n"+
			"2) InnerRandomStreamKey: %x\n"+
			"3) Binaries: %s\n",
		ih.InnerRandomStreamID,
		ih.InnerRandomStreamKey,
		ih.Binaries,
	)
}

// ErrEndOfInnerHeaders is the error returned when the end of inner header is read
var ErrEndOfInnerHeaders = errors.New("gokeepasslib: inner header id was 0, end of inner headers")

// ErrUnknownInnerHeaderID is the error returned if an unknown inner header is read
type ErrUnknownInnerHeaderID byte

func (i ErrUnknownInnerHeaderID) Error() string {
	return fmt.Sprintf("gokeepasslib: unknown inner header ID of %x", i)
}
<|endoftext|>