<commit_before>package agent

import (
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"math/rand"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"runtime"
	"strconv"
	"strings"
	"time"

	"github.com/hashicorp/consul/api"
	"github.com/hashicorp/consul/consul"
	"github.com/hashicorp/consul/consul/structs"
	"github.com/hashicorp/consul/logger"
	"github.com/hashicorp/consul/testutil/retry"
	"github.com/hashicorp/consul/types"
	"github.com/hashicorp/consul/version"
	uuid "github.com/hashicorp/go-uuid"
)

func init() {
	rand.Seed(time.Now().UnixNano()) // seed random number generator
}

// TempDir defines the base dir for temporary directories.
var TempDir = os.TempDir()

// TestAgent encapsulates an Agent with a default configuration and
// startup procedure suitable for testing. It panics if there are errors
// during creation or startup instead of returning errors. It manages a
// temporary data directory which is removed after shutdown.
//
// todo(fs): do we need the temp data dir if we run in dev mode?
type TestAgent struct {
	// Name is an optional name of the agent.
	Name string

	// Config is the agent configuration. If Config is nil then
	// TestConfig() is used. If Config.DataDir is set then it is
	// the callers responsibility to clean up the data directory.
	// Otherwise, a temporary data directory is created and removed
	// when Shutdown() is called.
	Config *Config

	// LogOutput is the sink for the logs. If nil, logs are written
	// to os.Stderr.
	LogOutput io.Writer

	// LogWriter is used for streaming logs.
	LogWriter *logger.LogWriter

	// DataDir is the data directory which is used when Config.DataDir
	// is not set. It is created automatically and removed when
	// Shutdown() is called.
	DataDir string

	// Key is the optional encryption key for the LAN and WAN keyring.
	Key string

	// NoInitialSync determines whether an anti-entropy run
	// will be scheduled after the agent started.
	NoInitialSync bool

	// dns is a reference to the first started DNS endpoint.
	// It is valid after Start().
	dns *DNSServer

	// srv is a reference to the first started HTTP endpoint.
	// It is valid after Start().
	srv *HTTPServer

	// Agent is the embedded consul agent.
	// It is valid after Start().
	*Agent
}

// NewTestAgent returns a started agent with the given name and
// configuration. It panics if the agent could not be started. The
// caller should call Shutdown() to stop the agent and remove temporary
// directories.
func NewTestAgent(name string, c *Config) *TestAgent {
	a := &TestAgent{Name: name, Config: c}
	a.Start()
	return a
}

type panicFailer struct{}

func (f *panicFailer) Log(args ...interface{}) { fmt.Println(args...) }
func (f *panicFailer) FailNow()                { panic("failed") }

// Start starts a test agent. It panics if the agent could not be started.
func (a *TestAgent) Start() *TestAgent {
	if a.Agent != nil {
		panic("TestAgent already started")
	}
	if a.Config == nil {
		a.Config = TestConfig()
	}
	if a.Config.DNSRecursor != "" {
		a.Config.DNSRecursors = append(a.Config.DNSRecursors, a.Config.DNSRecursor)
	}
	if a.Config.DataDir == "" {
		name := "agent"
		if a.Name != "" {
			name = a.Name + "-agent"
		}
		name = strings.Replace(name, "/", "_", -1)
		d, err := ioutil.TempDir(TempDir, name)
		if err != nil {
			panic(fmt.Sprintf("Error creating data dir %s: %s", filepath.Join(TempDir, name), err))
		}
		a.DataDir = d
		a.Config.DataDir = d
	}
	id := UniqueID()

	// since the ports are written into the data files we pick random ports
	// only once and try a number of times to start with them.
	pickRandomPorts(a.Config)

	for i := 10; i >= 0; i-- {
		// write the keyring
		if a.Key != "" {
			writeKey := func(key, filename string) {
				path := filepath.Join(a.Config.DataDir, filename)
				if err := initKeyring(path, key); err != nil {
					panic(fmt.Sprintf("Error creating keyring %s: %s", path, err))
				}
			}
			writeKey(a.Key, serfLANKeyring)
			writeKey(a.Key, serfWANKeyring)
		}

		agent, err := NewAgent(a.Config)
		if err != nil {
			panic(fmt.Sprintf("Error creating agent: %s", err))
		}

		logOutput := a.LogOutput
		if logOutput == nil {
			logOutput = os.Stderr
		}
		agent.LogOutput = logOutput
		agent.LogWriter = a.LogWriter
		agent.logger = log.New(logOutput, a.Name+" - ", log.LstdFlags)

		// we need the err var in the next exit condition
		if err := agent.Start(); err == nil {
			a.Agent = agent
			break
		} else if i == 0 {
			fmt.Println(id, a.Name, "Error starting agent:", err)
			runtime.Goexit()
		} else {
			agent.Shutdown()
			wait := time.Duration(rand.Int31n(2000)) * time.Millisecond
			fmt.Println(id, a.Name, "retrying in", wait)
			time.Sleep(wait)
		}
	}
	if !a.NoInitialSync {
		a.Agent.StartSync()
	}

	var out structs.IndexedNodes
	retry.Run(&panicFailer{}, func(r *retry.R) {
		if len(a.httpServers) == 0 {
			r.Fatal(a.Name, "waiting for server")
		}
		if a.Config.Bootstrap && a.Config.Server {
			// Ensure we have a leader and a node registration.
			args := &structs.DCSpecificRequest{
				Datacenter: a.Config.Datacenter,
				QueryOptions: structs.QueryOptions{
					MinQueryIndex: out.Index,
					MaxQueryTime:  25 * time.Millisecond,
				},
			}
			if err := a.RPC("Catalog.ListNodes", args, &out); err != nil {
				r.Fatal(a.Name, "Catalog.ListNodes failed:", err)
			}
			if !out.QueryMeta.KnownLeader {
				r.Fatal(a.Name, "No leader")
			}
			if out.Index == 0 {
				r.Fatal(a.Name, "Consul index is 0")
			}
		} else {
			req, _ := http.NewRequest("GET", "/v1/agent/self", nil)
			resp := httptest.NewRecorder()
			_, err := a.httpServers[0].AgentSelf(resp, req)
			if err != nil || resp.Code != 200 {
				r.Fatal(a.Name, "failed OK respose", err)
			}
		}
	})
	a.dns = a.dnsServers[0]
	a.srv = a.httpServers[0]
	return a
}

// Shutdown stops the agent and removes the data directory if it is
// managed by the test agent.
func (a *TestAgent) Shutdown() error {
	defer func() {
		if a.DataDir != "" {
			os.RemoveAll(a.DataDir)
		}
	}()
	return a.Agent.Shutdown()
}

func (a *TestAgent) HTTPAddr() string {
	if a.srv == nil {
		return ""
	}
	return a.srv.Addr
}

func (a *TestAgent) Client() *api.Client {
	conf := api.DefaultConfig()
	conf.Address = a.HTTPAddr()
	c, err := api.NewClient(conf)
	if err != nil {
		panic(fmt.Sprintf("Error creating consul API client: %s", err))
	}
	return c
}

func (a *TestAgent) consulConfig() *consul.Config {
	c, err := a.Agent.consulConfig()
	if err != nil {
		panic(err)
	}
	return c
}

func UniqueID() string {
	id := strconv.FormatUint(rand.Uint64(), 36)
	for len(id) < 16 {
		id += " "
	}
	return id
}

// TenPorts returns the first port number of a block of
// ten random ports.
func TenPorts() int {
	return 1030 + int(rand.Int31n(6440))*10
}

// pickRandomPorts selects random ports from fixed size random blocks of
// ports. This does not eliminate the chance for port conflict but
// reduces it significanltly with little overhead. Furthermore, asking
// the kernel for a random port by binding to port 0 prolongs the test
// execution (in our case +20sec) while also not fully eliminating the
// chance of port conflicts for concurrently executed test binaries.
// Instead of relying on one set of ports to be sufficient we retry
// starting the agent with different ports on port conflict.
func pickRandomPorts(c *Config) {
	port := TenPorts()
	c.Ports.DNS = port + 1
	c.Ports.HTTP = port + 2
	// when we enable HTTPS then we need to fix finding the
	// "first" HTTP server since that might be HTTPS server
	// c.Ports.HTTPS = port + 3
	c.Ports.SerfLan = port + 4
	c.Ports.SerfWan = port + 5
	c.Ports.Server = port + 6
}

// TestConfig returns a unique default configuration for testing an
// agent.
func TestConfig() *Config {
	nodeID, err := uuid.GenerateUUID()
	if err != nil {
		panic(err)
	}

	cfg := DefaultConfig()

	cfg.Version = version.Version
	cfg.VersionPrerelease = "c.d"

	cfg.NodeID = types.NodeID(nodeID)
	cfg.NodeName = "Node " + nodeID
	cfg.BindAddr = "127.0.0.1"
	cfg.AdvertiseAddr = "127.0.0.1"
	cfg.Datacenter = "dc1"
	cfg.Bootstrap = true
	cfg.Server = true

	ccfg := consul.DefaultConfig()
	cfg.ConsulConfig = ccfg

	ccfg.SerfLANConfig.MemberlistConfig.SuspicionMult = 3
	ccfg.SerfLANConfig.MemberlistConfig.ProbeTimeout = 100 * time.Millisecond
	ccfg.SerfLANConfig.MemberlistConfig.ProbeInterval = 100 * time.Millisecond
	ccfg.SerfLANConfig.MemberlistConfig.GossipInterval = 100 * time.Millisecond

	ccfg.SerfWANConfig.MemberlistConfig.SuspicionMult = 3
	ccfg.SerfWANConfig.MemberlistConfig.ProbeTimeout = 100 * time.Millisecond
	ccfg.SerfWANConfig.MemberlistConfig.ProbeInterval = 100 * time.Millisecond
	ccfg.SerfWANConfig.MemberlistConfig.GossipInterval = 100 * time.Millisecond

	ccfg.RaftConfig.LeaderLeaseTimeout = 20 * time.Millisecond
	ccfg.RaftConfig.HeartbeatTimeout = 40 * time.Millisecond
	ccfg.RaftConfig.ElectionTimeout = 40 * time.Millisecond

	ccfg.CoordinateUpdatePeriod = 100 * time.Millisecond
	ccfg.ServerHealthInterval = 10 * time.Millisecond
	return cfg
}

// TestACLConfig returns a default configuration for testing an agent
// with ACLs.
func TestACLConfig() *Config {
	cfg := TestConfig()
	cfg.ACLDatacenter = cfg.Datacenter
	cfg.ACLDefaultPolicy = "deny"
	cfg.ACLMasterToken = "root"
	cfg.ACLAgentToken = "root"
	cfg.ACLAgentMasterToken = "towel"
	cfg.ACLEnforceVersion8 = Bool(true)
	return cfg
}
<commit_msg>test: we will not use dev mode for testing for now<commit_after>package agent

import (
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"math/rand"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"runtime"
	"strconv"
	"strings"
	"time"

	"github.com/hashicorp/consul/api"
	"github.com/hashicorp/consul/consul"
	"github.com/hashicorp/consul/consul/structs"
	"github.com/hashicorp/consul/logger"
	"github.com/hashicorp/consul/testutil/retry"
	"github.com/hashicorp/consul/types"
	"github.com/hashicorp/consul/version"
	uuid "github.com/hashicorp/go-uuid"
)

func init() {
	rand.Seed(time.Now().UnixNano()) // seed random number generator
}

// TempDir defines the base dir for temporary directories.
var TempDir = os.TempDir()

// TestAgent encapsulates an Agent with a default configuration and
// startup procedure suitable for testing. It panics if there are errors
// during creation or startup instead of returning errors. It manages a
// temporary data directory which is removed after shutdown.
type TestAgent struct {
	// Name is an optional name of the agent.
	Name string

	// Config is the agent configuration. If Config is nil then
	// TestConfig() is used. If Config.DataDir is set then it is
	// the callers responsibility to clean up the data directory.
	// Otherwise, a temporary data directory is created and removed
	// when Shutdown() is called.
	Config *Config

	// LogOutput is the sink for the logs. If nil, logs are written
	// to os.Stderr.
	LogOutput io.Writer

	// LogWriter is used for streaming logs.
	LogWriter *logger.LogWriter

	// DataDir is the data directory which is used when Config.DataDir
	// is not set. It is created automatically and removed when
	// Shutdown() is called.
	DataDir string

	// Key is the optional encryption key for the LAN and WAN keyring.
	Key string

	// NoInitialSync determines whether an anti-entropy run
	// will be scheduled after the agent started.
	NoInitialSync bool

	// dns is a reference to the first started DNS endpoint.
	// It is valid after Start().
	dns *DNSServer

	// srv is a reference to the first started HTTP endpoint.
	// It is valid after Start().
	srv *HTTPServer

	// Agent is the embedded consul agent.
	// It is valid after Start().
	*Agent
}

// NewTestAgent returns a started agent with the given name and
// configuration. It panics if the agent could not be started. The
// caller should call Shutdown() to stop the agent and remove temporary
// directories.
func NewTestAgent(name string, c *Config) *TestAgent {
	a := &TestAgent{Name: name, Config: c}
	a.Start()
	return a
}

type panicFailer struct{}

func (f *panicFailer) Log(args ...interface{}) { fmt.Println(args...) }
func (f *panicFailer) FailNow()                { panic("failed") }

// Start starts a test agent. It panics if the agent could not be started.
func (a *TestAgent) Start() *TestAgent {
	if a.Agent != nil {
		panic("TestAgent already started")
	}
	if a.Config == nil {
		a.Config = TestConfig()
	}
	if a.Config.DNSRecursor != "" {
		a.Config.DNSRecursors = append(a.Config.DNSRecursors, a.Config.DNSRecursor)
	}
	if a.Config.DataDir == "" {
		name := "agent"
		if a.Name != "" {
			name = a.Name + "-agent"
		}
		name = strings.Replace(name, "/", "_", -1)
		d, err := ioutil.TempDir(TempDir, name)
		if err != nil {
			panic(fmt.Sprintf("Error creating data dir %s: %s", filepath.Join(TempDir, name), err))
		}
		a.DataDir = d
		a.Config.DataDir = d
	}
	id := UniqueID()

	// since the ports are written into the data files we pick random ports
	// only once and try a number of times to start with them.
	pickRandomPorts(a.Config)

	for i := 10; i >= 0; i-- {
		// write the keyring
		if a.Key != "" {
			writeKey := func(key, filename string) {
				path := filepath.Join(a.Config.DataDir, filename)
				if err := initKeyring(path, key); err != nil {
					panic(fmt.Sprintf("Error creating keyring %s: %s", path, err))
				}
			}
			writeKey(a.Key, serfLANKeyring)
			writeKey(a.Key, serfWANKeyring)
		}

		agent, err := NewAgent(a.Config)
		if err != nil {
			panic(fmt.Sprintf("Error creating agent: %s", err))
		}

		logOutput := a.LogOutput
		if logOutput == nil {
			logOutput = os.Stderr
		}
		agent.LogOutput = logOutput
		agent.LogWriter = a.LogWriter
		agent.logger = log.New(logOutput, a.Name+" - ", log.LstdFlags)

		// we need the err var in the next exit condition
		if err := agent.Start(); err == nil {
			a.Agent = agent
			break
		} else if i == 0 {
			fmt.Println(id, a.Name, "Error starting agent:", err)
			runtime.Goexit()
		} else {
			agent.Shutdown()
			wait := time.Duration(rand.Int31n(2000)) * time.Millisecond
			fmt.Println(id, a.Name, "retrying in", wait)
			time.Sleep(wait)
		}
	}
	if !a.NoInitialSync {
		a.Agent.StartSync()
	}

	var out structs.IndexedNodes
	retry.Run(&panicFailer{}, func(r *retry.R) {
		if len(a.httpServers) == 0 {
			r.Fatal(a.Name, "waiting for server")
		}
		if a.Config.Bootstrap && a.Config.Server {
			// Ensure we have a leader and a node registration.
			args := &structs.DCSpecificRequest{
				Datacenter: a.Config.Datacenter,
				QueryOptions: structs.QueryOptions{
					MinQueryIndex: out.Index,
					MaxQueryTime:  25 * time.Millisecond,
				},
			}
			if err := a.RPC("Catalog.ListNodes", args, &out); err != nil {
				r.Fatal(a.Name, "Catalog.ListNodes failed:", err)
			}
			if !out.QueryMeta.KnownLeader {
				r.Fatal(a.Name, "No leader")
			}
			if out.Index == 0 {
				r.Fatal(a.Name, "Consul index is 0")
			}
		} else {
			req, _ := http.NewRequest("GET", "/v1/agent/self", nil)
			resp := httptest.NewRecorder()
			_, err := a.httpServers[0].AgentSelf(resp, req)
			if err != nil || resp.Code != 200 {
				r.Fatal(a.Name, "failed OK respose", err)
			}
		}
	})
	a.dns = a.dnsServers[0]
	a.srv = a.httpServers[0]
	return a
}

// Shutdown stops the agent and removes the data directory if it is
// managed by the test agent.
func (a *TestAgent) Shutdown() error {
	defer func() {
		if a.DataDir != "" {
			os.RemoveAll(a.DataDir)
		}
	}()
	return a.Agent.Shutdown()
}

func (a *TestAgent) HTTPAddr() string {
	if a.srv == nil {
		return ""
	}
	return a.srv.Addr
}

func (a *TestAgent) Client() *api.Client {
	conf := api.DefaultConfig()
	conf.Address = a.HTTPAddr()
	c, err := api.NewClient(conf)
	if err != nil {
		panic(fmt.Sprintf("Error creating consul API client: %s", err))
	}
	return c
}

func (a *TestAgent) consulConfig() *consul.Config {
	c, err := a.Agent.consulConfig()
	if err != nil {
		panic(err)
	}
	return c
}

func UniqueID() string {
	id := strconv.FormatUint(rand.Uint64(), 36)
	for len(id) < 16 {
		id += " "
	}
	return id
}

// TenPorts returns the first port number of a block of
// ten random ports.
func TenPorts() int {
	return 1030 + int(rand.Int31n(6440))*10
}

// pickRandomPorts selects random ports from fixed size random blocks of
// ports. This does not eliminate the chance for port conflict but
// reduces it significanltly with little overhead. Furthermore, asking
// the kernel for a random port by binding to port 0 prolongs the test
// execution (in our case +20sec) while also not fully eliminating the
// chance of port conflicts for concurrently executed test binaries.
// Instead of relying on one set of ports to be sufficient we retry
// starting the agent with different ports on port conflict.
func pickRandomPorts(c *Config) {
	port := TenPorts()
	c.Ports.DNS = port + 1
	c.Ports.HTTP = port + 2
	// when we enable HTTPS then we need to fix finding the
	// "first" HTTP server since that might be HTTPS server
	// c.Ports.HTTPS = port + 3
	c.Ports.SerfLan = port + 4
	c.Ports.SerfWan = port + 5
	c.Ports.Server = port + 6
}

// TestConfig returns a unique default configuration for testing an
// agent.
func TestConfig() *Config {
	nodeID, err := uuid.GenerateUUID()
	if err != nil {
		panic(err)
	}

	cfg := DefaultConfig()

	cfg.Version = version.Version
	cfg.VersionPrerelease = "c.d"

	cfg.NodeID = types.NodeID(nodeID)
	cfg.NodeName = "Node " + nodeID
	cfg.BindAddr = "127.0.0.1"
	cfg.AdvertiseAddr = "127.0.0.1"
	cfg.Datacenter = "dc1"
	cfg.Bootstrap = true
	cfg.Server = true

	ccfg := consul.DefaultConfig()
	cfg.ConsulConfig = ccfg

	ccfg.SerfLANConfig.MemberlistConfig.SuspicionMult = 3
	ccfg.SerfLANConfig.MemberlistConfig.ProbeTimeout = 100 * time.Millisecond
	ccfg.SerfLANConfig.MemberlistConfig.ProbeInterval = 100 * time.Millisecond
	ccfg.SerfLANConfig.MemberlistConfig.GossipInterval = 100 * time.Millisecond

	ccfg.SerfWANConfig.MemberlistConfig.SuspicionMult = 3
	ccfg.SerfWANConfig.MemberlistConfig.ProbeTimeout = 100 * time.Millisecond
	ccfg.SerfWANConfig.MemberlistConfig.ProbeInterval = 100 * time.Millisecond
	ccfg.SerfWANConfig.MemberlistConfig.GossipInterval = 100 * time.Millisecond

	ccfg.RaftConfig.LeaderLeaseTimeout = 20 * time.Millisecond
	ccfg.RaftConfig.HeartbeatTimeout = 40 * time.Millisecond
	ccfg.RaftConfig.ElectionTimeout = 40 * time.Millisecond

	ccfg.CoordinateUpdatePeriod = 100 * time.Millisecond
	ccfg.ServerHealthInterval = 10 * time.Millisecond
	return cfg
}

// TestACLConfig returns a default configuration for testing an agent
// with ACLs.
func TestACLConfig() *Config {
	cfg := TestConfig()
	cfg.ACLDatacenter = cfg.Datacenter
	cfg.ACLDefaultPolicy = "deny"
	cfg.ACLMasterToken = "root"
	cfg.ACLAgentToken = "root"
	cfg.ACLAgentMasterToken = "towel"
	cfg.ACLEnforceVersion8 = Bool(true)
	return cfg
}
<|endoftext|>


<commit_before>package commands

import (
	"fmt"
	"io"
	"os"

	"github.com/git-lfs/git-lfs/errors"
	"github.com/git-lfs/git-lfs/filepathfilter"
	"github.com/git-lfs/git-lfs/git"
	"github.com/git-lfs/git-lfs/lfs"
	"github.com/git-lfs/git-lfs/localstorage"
	"github.com/git-lfs/git-lfs/tools"
	"github.com/git-lfs/git-lfs/tools/humanize"
	"github.com/git-lfs/git-lfs/tq"
	"github.com/spf13/cobra"
)

var (
	// smudgeSkip is a command-line flag belonging to the "git-lfs smudge"
	// command specifying whether to skip the smudge process.
	smudgeSkip = false
)

// delayedSmudge performs a 'delayed' smudge, adding the LFS pointer to the
// `*tq.TransferQueue` "q" if the file is not present locally, passes the given
// filepathfilter, and is not skipped. If the pointer is malformed, or already
// exists, it streams the contents to be written into the working copy to "to".
//
// delayedSmudge returns the number of bytes written, whether the checkout was
// delayed, the *lfs.Pointer that was smudged, and an error, if one occurred.
func delayedSmudge(s *git.FilterProcessScanner, to io.Writer, from io.Reader, q *tq.TransferQueue, filename string, skip bool, filter *filepathfilter.Filter) (int64, bool, *lfs.Pointer, error) {
	ptr, pbuf, perr := lfs.DecodeFrom(from)
	if perr != nil {
		// Write 'statusFromErr(nil)', even though 'perr != nil', since
		// we are about to write non-delayed smudged contents to "to".
		if err := s.WriteStatus(statusFromErr(nil)); err != nil {
			return 0, false, nil, err
		}

		n, err := tools.Spool(to, pbuf, localstorage.Objects().TempDir)
		if err != nil {
			return n, false, nil, errors.Wrap(err, perr.Error())
		}

		if n != 0 {
			return 0, false, nil, errors.NewNotAPointerError(errors.Errorf(
				"Unable to parse pointer at: %q", filename,
			))
		}
		return 0, false, nil, nil
	}

	lfs.LinkOrCopyFromReference(ptr.Oid, ptr.Size)

	path, err := lfs.LocalMediaPath(ptr.Oid)
	if err != nil {
		return 0, false, nil, err
	}

	if !skip && filter.Allows(filename) {
		if _, statErr := os.Stat(path); err == nil || !os.IsNotExist(statErr) {
			q.Add(filename, path, ptr.Oid, ptr.Size)
			return 0, true, ptr, nil
		}
	}

	if err := s.WriteStatus(statusFromErr(nil)); err != nil {
		return 0, false, nil, err
	}

	n, err := ptr.Encode(to)
	return int64(n), false, ptr, err
}

// smudge smudges the given `*lfs.Pointer`, "ptr", and writes its objects
// contents to the `io.Writer`, "to".
//
// If the encoded LFS pointer is not parse-able as a pointer, the contents of
// that file will instead be spooled to a temporary location on disk and then
// copied out back to Git. If the pointer file is empty, an empty file will be
// written with no error.
//
// If the smudged object did not "pass" the include and exclude filterset, it
// will not be downloaded, and the object will remain a pointer on disk, as if
// the smudge filter had not been applied at all.
//
// Any errors encountered along the way will be returned immediately if they
// were non-fatal, otherwise execution will halt and the process will be
// terminated by using the `commands.Panic()` func.
func smudge(to io.Writer, from io.Reader, filename string, skip bool, filter *filepathfilter.Filter) (int64, error) {
	ptr, pbuf, perr := lfs.DecodeFrom(from)
	if perr != nil {
		n, err := tools.Spool(to, pbuf, localstorage.Objects().TempDir)
		if err != nil {
			return 0, errors.Wrap(err, perr.Error())
		}

		if n != 0 {
			return 0, errors.NewNotAPointerError(errors.Errorf(
				"Unable to parse pointer at: %q", filename,
			))
		}
		return 0, nil
	}

	lfs.LinkOrCopyFromReference(ptr.Oid, ptr.Size)
	cb, file, err := lfs.CopyCallbackFile("download", filename, 1, 1)
	if err != nil {
		return 0, err
	}

	download := !skip
	if download {
		download = filter.Allows(filename)
	}

	n, err := ptr.Smudge(to, filename, download, getTransferManifest(), cb)
	if file != nil {
		file.Close()
	}

	if err != nil {
		ptr.Encode(to)
		// Download declined error is ok to skip if we weren't requesting download
		if !(errors.IsDownloadDeclinedError(err) && !download) {
			var oid string = ptr.Oid
			if len(oid) >= 7 {
				oid = oid[:7]
			}

			LoggedError(err, "Error downloading object: %s (%s): %s", filename, oid, err)
			if !cfg.SkipDownloadErrors() {
				os.Exit(2)
			}
		}
	}

	return n, nil
}

func smudgeCommand(cmd *cobra.Command, args []string) {
	requireStdin("This command should be run by the Git 'smudge' filter")
	lfs.InstallHooks(false)

	if !smudgeSkip && cfg.Os.Bool("GIT_LFS_SKIP_SMUDGE", false) {
		smudgeSkip = true
	}
	filter := filepathfilter.New(cfg.FetchIncludePaths(), cfg.FetchExcludePaths())

	if n, err := smudge(os.Stdout, os.Stdin, smudgeFilename(args), smudgeSkip, filter); err != nil {
		if errors.IsNotAPointerError(err) {
			fmt.Fprintln(os.Stderr, err.Error())
		} else {
			Error(err.Error())
		}
	} else if possiblyMalformedObjectSize(n) {
		fmt.Fprintln(os.Stderr, "Possibly malformed smudge on Windows: see `git lfs help smudge` for more info.")
	}
}

func smudgeFilename(args []string) string {
	if len(args) > 0 {
		return args[0]
	}
	return "<unknown file>"
}

func possiblyMalformedObjectSize(n int64) bool {
	return n > 4*humanize.Gigabyte
}

func init() {
	RegisterCommand("smudge", smudgeCommand, func(cmd *cobra.Command) {
		cmd.Flags().BoolVarP(&smudgeSkip, "skip", "s", false, "")
	})
}
<commit_msg>commands/smudge: teach 'delayedSmudge' only to transfer uncached files<commit_after>package commands

import (
	"fmt"
	"io"
	"os"

	"github.com/git-lfs/git-lfs/errors"
	"github.com/git-lfs/git-lfs/filepathfilter"
	"github.com/git-lfs/git-lfs/git"
	"github.com/git-lfs/git-lfs/lfs"
	"github.com/git-lfs/git-lfs/localstorage"
	"github.com/git-lfs/git-lfs/tools"
	"github.com/git-lfs/git-lfs/tools/humanize"
	"github.com/git-lfs/git-lfs/tq"
	"github.com/spf13/cobra"
)

var (
	// smudgeSkip is a command-line flag belonging to the "git-lfs smudge"
	// command specifying whether to skip the smudge process.
	smudgeSkip = false
)

// delayedSmudge performs a 'delayed' smudge, adding the LFS pointer to the
// `*tq.TransferQueue` "q" if the file is not present locally, passes the given
// filepathfilter, and is not skipped. If the pointer is malformed, or already
// exists, it streams the contents to be written into the working copy to "to".
//
// delayedSmudge returns the number of bytes written, whether the checkout was
// delayed, the *lfs.Pointer that was smudged, and an error, if one occurred.
func delayedSmudge(s *git.FilterProcessScanner, to io.Writer, from io.Reader, q *tq.TransferQueue, filename string, skip bool, filter *filepathfilter.Filter) (int64, bool, *lfs.Pointer, error) {
	ptr, pbuf, perr := lfs.DecodeFrom(from)
	if perr != nil {
		// Write 'statusFromErr(nil)', even though 'perr != nil', since
		// we are about to write non-delayed smudged contents to "to".
		if err := s.WriteStatus(statusFromErr(nil)); err != nil {
			return 0, false, nil, err
		}

		n, err := tools.Spool(to, pbuf, localstorage.Objects().TempDir)
		if err != nil {
			return n, false, nil, errors.Wrap(err, perr.Error())
		}

		if n != 0 {
			return 0, false, nil, errors.NewNotAPointerError(errors.Errorf(
				"Unable to parse pointer at: %q", filename,
			))
		}
		return 0, false, nil, nil
	}

	lfs.LinkOrCopyFromReference(ptr.Oid, ptr.Size)

	path, err := lfs.LocalMediaPath(ptr.Oid)
	if err != nil {
		return 0, false, nil, err
	}

	if !skip && filter.Allows(filename) {
		if _, statErr := os.Stat(path); statErr != nil {
			q.Add(filename, path, ptr.Oid, ptr.Size)
			return 0, true, ptr, nil
		}
	}

	if err := s.WriteStatus(statusFromErr(nil)); err != nil {
		return 0, false, nil, err
	}

	n, err := ptr.Encode(to)
	return int64(n), false, ptr, err
}

// smudge smudges the given `*lfs.Pointer`, "ptr", and writes its objects
// contents to the `io.Writer`, "to".
//
// If the encoded LFS pointer is not parse-able as a pointer, the contents of
// that file will instead be spooled to a temporary location on disk and then
// copied out back to Git. If the pointer file is empty, an empty file will be
// written with no error.
//
// If the smudged object did not "pass" the include and exclude filterset, it
// will not be downloaded, and the object will remain a pointer on disk, as if
// the smudge filter had not been applied at all.
//
// Any errors encountered along the way will be returned immediately if they
// were non-fatal, otherwise execution will halt and the process will be
// terminated by using the `commands.Panic()` func.
func smudge(to io.Writer, from io.Reader, filename string, skip bool, filter *filepathfilter.Filter) (int64, error) {
	ptr, pbuf, perr := lfs.DecodeFrom(from)
	if perr != nil {
		n, err := tools.Spool(to, pbuf, localstorage.Objects().TempDir)
		if err != nil {
			return 0, errors.Wrap(err, perr.Error())
		}

		if n != 0 {
			return 0, errors.NewNotAPointerError(errors.Errorf(
				"Unable to parse pointer at: %q", filename,
			))
		}
		return 0, nil
	}

	lfs.LinkOrCopyFromReference(ptr.Oid, ptr.Size)
	cb, file, err := lfs.CopyCallbackFile("download", filename, 1, 1)
	if err != nil {
		return 0, err
	}

	download := !skip
	if download {
		download = filter.Allows(filename)
	}

	n, err := ptr.Smudge(to, filename, download, getTransferManifest(), cb)
	if file != nil {
		file.Close()
	}

	if err != nil {
		ptr.Encode(to)
		// Download declined error is ok to skip if we weren't requesting download
		if !(errors.IsDownloadDeclinedError(err) && !download) {
			var oid string = ptr.Oid
			if len(oid) >= 7 {
				oid = oid[:7]
			}

			LoggedError(err, "Error downloading object: %s (%s): %s", filename, oid, err)
			if !cfg.SkipDownloadErrors() {
				os.Exit(2)
			}
		}
	}

	return n, nil
}

func smudgeCommand(cmd *cobra.Command, args []string) {
	requireStdin("This command should be run by the Git 'smudge' filter")
	lfs.InstallHooks(false)

	if !smudgeSkip && cfg.Os.Bool("GIT_LFS_SKIP_SMUDGE", false) {
		smudgeSkip = true
	}
	filter := filepathfilter.New(cfg.FetchIncludePaths(), cfg.FetchExcludePaths())

	if n, err := smudge(os.Stdout, os.Stdin, smudgeFilename(args), smudgeSkip, filter); err != nil {
		if errors.IsNotAPointerError(err) {
			fmt.Fprintln(os.Stderr, err.Error())
		} else {
			Error(err.Error())
		}
	} else if possiblyMalformedObjectSize(n) {
		fmt.Fprintln(os.Stderr, "Possibly malformed smudge on Windows: see `git lfs help smudge` for more info.")
	}
}

func smudgeFilename(args []string) string {
	if len(args) > 0 {
		return args[0]
	}
	return "<unknown file>"
}

func possiblyMalformedObjectSize(n int64) bool {
	return n > 4*humanize.Gigabyte
}

func init() {
	RegisterCommand("smudge", smudgeCommand, func(cmd *cobra.Command) {
		cmd.Flags().BoolVarP(&smudgeSkip, "skip", "s", false, "")
	})
}
<|endoftext|>


<commit_before>// +build windows

package commands

import (
	"log"
	"os"

	"github.com/codegangsta/cli"
)

func Hijack(c *cli.Context) {
	log.Fatalln("Command not supported on Windows!")
	os.Exit(0)
}
<commit_msg>Use non zero exit code for hijack on windows<commit_after>// +build windows

package commands

import (
	"log"
	"os"

	"github.com/codegangsta/cli"
)

func Hijack(c *cli.Context) {
	log.Fatalln("Command not supported on Windows!")
	os.Exit(1)
}
<|endoftext|>


<commit_before>package main

import (
	"github.com/cwchiu/go-winapi"
	"os"
	"sync"
	"unsafe"
	"syscall"
	"errors"
)

const (
	WINDOW_CLASS = "dacko"
)

const (
	MOD_ALT = 1
	MOD_CONTROL = 2
	MOD_SHIFT = 4
	MOD_WIN = 8
)

const (
	SYNCHRONIZE = 0x00100000
)

const (
	CREATE_NEW_CONSOLE = 0x00000010
)

const (
	WAIT_OBJECT_0 = 0x00000000
	WAIT_ABANDONED = 0x00000080
	WAIT_timeout = 0x00000102
	WAIT_FAILED = 0xFFFFFFFF
)

var (
	user32, _         = syscall.LoadLibrary("user32.dll")
	registerHotkey, _ = syscall.GetProcAddress(user32, "RegisterHotKey")
	getWindowThreadProcessId, _ = syscall.GetProcAddress(user32, "GetWindowThreadProcessId")
	enumWindows, _ = syscall.GetProcAddress(user32, "EnumWindows")
	procReplyMessage = syscall.NewLazyDLL("user32.dll").NewProc("ReplyMessage")

	kernel32, _            = syscall.LoadLibrary("kernel32.dll")
	waitForSingleObject, _ = syscall.GetProcAddress(kernel32, "WaitForSingleObject")
	dwProcessId uint32 = 0
	zeroProcAttr syscall.ProcAttr
	zeroSysProcAttr syscall.SysProcAttr
	ForkLock sync.RWMutex
)

func RegisterHotKey(hWnd winapi.HWND, id int, fsModifiers, vk uint) bool {
	ret, _, _ := syscall.Syscall6(uintptr(registerHotkey), 4,
		uintptr(hWnd),
		uintptr(id),
		uintptr(fsModifiers),
		uintptr(vk),0,0)

	return ret != 0
}

func GetWindowThreadProcessId(hWnd syscall.Handle) (winapi.HANDLE, uint32) {
	var id int
	ret, _, _ := syscall.Syscall(uintptr(getWindowThreadProcessId), 2,
		uintptr(hWnd),
		uintptr(unsafe.Pointer(&id)),
		0)

	return winapi.HANDLE(ret), uint32(id)
}

func EnumWindows(f func(syscall.Handle, uintptr) uintptr, lParam uint32) bool {
	ret, _, _ := syscall.Syscall(uintptr(enumWindows), 2,
		uintptr(syscall.NewCallback(f)),
		uintptr(lParam),
		0)

	return ret != 0
}

func createProcess(attr *syscall.ProcAttr) (pid int, handle uintptr, err error) {
	if attr == nil {
		attr = &zeroProcAttr
	}
	sys := attr.Sys
	if sys == nil {
		sys = &zeroSysProcAttr
	}

	if len(attr.Files) > 3 {
		return 0, 0, syscall.EWINDOWS
	}
	if len(attr.Files) < 3 {
		return 0, 0, syscall.EINVAL
	}

	var cmdline string
	// Windows CreateProcess takes the command line as a single string:
	// use attr.CmdLine if set, else build the command line by escaping
	// and joining each argument with spaces
	if sys.CmdLine != "" {
		cmdline = sys.CmdLine
	} else {
		return 0, 0, errors.New("sys.CmdLine is empty")
	}

	var argvp *uint16
	if len(cmdline) != 0 {
		argvp, err = syscall.UTF16PtrFromString(cmdline)
		if err != nil {
			return 0, 0, err
		}
	}

	var dirp *uint16
	if len(attr.Dir) != 0 {
		dirp, err = syscall.UTF16PtrFromString(attr.Dir)
		if err != nil {
			return 0, 0, err
		}
	}

	// Acquire the fork lock so that no other threads
	// create new fds that are not yet close-on-exec
	// before we fork.
	ForkLock.Lock()
	defer ForkLock.Unlock()

	p, _ := syscall.GetCurrentProcess()
	fd := make([]syscall.Handle, len(attr.Files))
	for i := range attr.Files {
		if attr.Files[i] > 0 {
			err := syscall.DuplicateHandle(p, syscall.Handle(attr.Files[i]), p, &fd[i], 0, true, syscall.DUPLICATE_SAME_ACCESS)
			if err != nil {
				return 0, 0, err
			}
			defer syscall.CloseHandle(syscall.Handle(fd[i]))
		}
	}
	si := new(syscall.StartupInfo)
	si.Cb = uint32(unsafe.Sizeof(*si))
	si.Flags = syscall.STARTF_USESTDHANDLES

	pi := new(syscall.ProcessInformation)

	flags := sys.CreationFlags | syscall.CREATE_UNICODE_ENVIRONMENT
	err = syscall.CreateProcess(nil, argvp, nil, nil, false, flags, nil, dirp, si, pi)
	if err != nil {
		return 0, 0, err
	}
	defer syscall.CloseHandle(syscall.Handle(pi.Thread))

	return int(pi.ProcessId), uintptr(pi.Process), nil
}

func showErrorMessage(hWnd winapi.HWND, msg string) {
	s, _ := syscall.UTF16PtrFromString(msg)
	t, _ := syscall.UTF16PtrFromString(WINDOW_CLASS)
	winapi.MessageBox(hWnd, s, t, winapi.MB_ICONWARNING|winapi.MB_OK)
}

func run() int {
	hInstance := winapi.GetModuleHandle(nil)

	if registerWindowClass(hInstance) == 0 {
		showErrorMessage(0, "registerWindowClass failed")
		return 1
	}

	if err := initializeInstance(hInstance, winapi.SW_SHOW); err != nil {
		showErrorMessage(0, err.Error())
		return 1
	}

	var msg winapi.MSG
	for winapi.GetMessage(&msg, 0,0,0) != 0 {
		winapi.TranslateMessage(&msg)
		winapi.DispatchMessage(&msg)
	}

	finalizeInstance(hInstance)

	return int(msg.WParam)
}

func registerWindowClass(hInstance winapi.HINSTANCE) winapi.ATOM {
	var wc winapi.WNDCLASSEX

	wc.CbSize = uint32(unsafe.Sizeof(winapi.WNDCLASSEX{}))
	wc.Style = 0
	wc.LpfnWndProc = syscall.NewCallback(wndProc)
	wc.CbClsExtra = 0
	wc.CbWndExtra = 0
	wc.HInstance = hInstance
	wc.HIcon         = winapi.LoadIcon(hInstance, winapi.MAKEINTRESOURCE(132))
	wc.HCursor       = winapi.LoadCursor(0, winapi.MAKEINTRESOURCE(winapi.IDC_HAND))
	wc.HbrBackground = winapi.HBRUSH(winapi.GetStockObject(winapi.WHITE_BRUSH))
	wc.LpszMenuName  = nil
	wc.LpszClassName,_ = syscall.UTF16PtrFromString(WINDOW_CLASS)

	return winapi.RegisterClassEx(&wc)
}

func initializeInstance(hInstance winapi.HINSTANCE, nCmdShow int) error {
	pc, _ := syscall.UTF16PtrFromString(WINDOW_CLASS)
	pt, _ := syscall.UTF16PtrFromString(WINDOW_CLASS)

	hWnd := winapi.CreateWindowEx(
		winapi.WS_EX_TOOLWINDOW|winapi.WS_EX_TOPMOST|winapi.WS_EX_NOACTIVATE|winapi.WS_EX_LAYERED,
		pc,pt,winapi.WS_POPUP,
		0,
		0,
		0,
		0,
		0,0,hInstance,nil)

	if hWnd == 0 {
		return errors.New("CreateWindowEx failed")
	}

	winapi.ShowWindow(hWnd, int32(nCmdShow))
	winapi.UpdateWindow(hWnd)

	if (RegisterHotKey(hWnd, 0, MOD_ALT|MOD_CONTROL, uint(config.Key)) == false) {
		return errors.New("RegisterHotKey failed")
	}

	return nil
}

func finalizeInstance(hInstance winapi.HINSTANCE) error {
	return nil
}

func isProcessGone(id uint32) bool {
	hProcess,err := syscall.OpenProcess(SYNCHRONIZE, false, id);
	if err != nil {
		return true;
	}

	dwWait,err := syscall.WaitForSingleObject(hProcess,0);
	if err != nil {
		return false
	}

	err = syscall.CloseHandle(hProcess)
	if err != nil {
		return false
	}

	return(dwWait == WAIT_OBJECT_0);
}

func createChildProcess(id uint32) (uint32,error) {
	if(!isProcessGone(id)) {
		return id,nil;
	}

	var sys syscall.SysProcAttr
	var attr syscall.ProcAttr

	sys.CreationFlags = CREATE_NEW_CONSOLE
	sys.CmdLine       = config.Cmdline

	attr.Files = make([]uintptr, 3)
	attr.Sys = &sys
	if config.WorkDir == "" {
		attr.Dir = os.Getenv("USERPROFILE")
	} else {
		attr.Dir = config.WorkDir
	}

	pid, _, err := createProcess(&attr)
	if err != nil {
		return 0,err
	}

	return uint32(pid), nil
}

func toggleWindowbyProcID(hwnd syscall.Handle, lparam uintptr) uintptr {
	var pid uint32
	_, pid = GetWindowThreadProcessId(hwnd)

	if (uint32(lparam) == pid) {
		if (winapi.IsWindowVisible(winapi.HWND(hwnd))) {
			winapi.ShowWindow(winapi.HWND(hwnd), winapi.SW_MINIMIZE);
			winapi.ShowWindow(winapi.HWND(hwnd), winapi.SW_HIDE);
		} else {
			winapi.ShowWindow(winapi.HWND(hwnd), winapi.SW_SHOW);
			winapi.ShowWindow(winapi.HWND(hwnd), winapi.SW_RESTORE);
		}
		return 0
	}
	return 1
}

func wndProc(hWnd winapi.HWND, msg uint32, wParam uintptr, lParam uintptr) uintptr {
	switch msg {
	case winapi.WM_HOTKEY:
		var err error
		dwProcessId,err = createChildProcess(dwProcessId);
		if err != nil {
			procReplyMessage.Call(1)
			return 0
		}

		EnumWindows(toggleWindowbyProcID, dwProcessId)
		return 0
	case winapi.WM_DESTROY:
		winapi.PostQuitMessage(0)
	default:
		return winapi.DefWindowProc(hWnd, msg, wParam, lParam)
	}
	return 0
}
<commit_msg>Fix to crash process<commit_after>package main

import (
	"github.com/cwchiu/go-winapi"
	"os"
	"sync"
	"unsafe"
	"syscall"
	"errors"
)

const (
	WINDOW_CLASS = "dacko"
)

const (
	MOD_ALT = 1
	MOD_CONTROL = 2
	MOD_SHIFT = 4
	MOD_WIN = 8
)

const (
	SYNCHRONIZE = 0x00100000
)

const (
	CREATE_NEW_CONSOLE = 0x00000010
)

const (
	WAIT_OBJECT_0 = 0x00000000
	WAIT_ABANDONED = 0x00000080
	WAIT_timeout = 0x00000102
	WAIT_FAILED = 0xFFFFFFFF
)

var (
	user32, _         = syscall.LoadLibrary("user32.dll")
	registerHotkey, _ = syscall.GetProcAddress(user32, "RegisterHotKey")
	getWindowThreadProcessId, _ = syscall.GetProcAddress(user32, "GetWindowThreadProcessId")
	enumWindows, _ = syscall.GetProcAddress(user32, "EnumWindows")
	procReplyMessage = syscall.NewLazyDLL("user32.dll").NewProc("ReplyMessage")

	kernel32, _            = syscall.LoadLibrary("kernel32.dll")
	waitForSingleObject, _ = syscall.GetProcAddress(kernel32, "WaitForSingleObject")
	dwProcessId uint32 = 0
	zeroProcAttr syscall.ProcAttr
	zeroSysProcAttr syscall.SysProcAttr
	ForkLock sync.RWMutex
)

func RegisterHotKey(hWnd winapi.HWND, id int, fsModifiers, vk uint) bool {
	ret, _, _ := syscall.Syscall6(uintptr(registerHotkey), 4,
		uintptr(hWnd),
		uintptr(id),
		uintptr(fsModifiers),
		uintptr(vk),0,0)

	return ret != 0
}

func GetWindowThreadProcessId(hWnd syscall.Handle) (winapi.HANDLE, uint32) {
	var id int
	ret, _, _ := syscall.Syscall(uintptr(getWindowThreadProcessId), 2,
		uintptr(hWnd),
		uintptr(unsafe.Pointer(&id)),
		0)

	return winapi.HANDLE(ret), uint32(id)
}

func EnumWindows(f func(syscall.Handle, uintptr) uintptr, lParam uint32) bool {
	ret, _, _ := syscall.Syscall(uintptr(enumWindows), 2,
		uintptr(syscall.NewCallback(f)),
		uintptr(lParam),
		0)

	return ret != 0
}

func createProcess(attr *syscall.ProcAttr) (pid int, handle uintptr, err error) {
	if attr == nil {
		attr = &zeroProcAttr
	}
	sys := attr.Sys
	if sys == nil {
		sys = &zeroSysProcAttr
	}

	if len(attr.Files) > 3 {
		return 0, 0, syscall.EWINDOWS
	}
	if len(attr.Files) < 3 {
		return 0, 0, syscall.EINVAL
	}

	var cmdline string
	// Windows CreateProcess takes the command line as a single string:
	// use attr.CmdLine if set, else build the command line by escaping
	// and joining each argument with spaces
	if sys.CmdLine != "" {
		cmdline = sys.CmdLine
	} else {
		return 0, 0, errors.New("sys.CmdLine is empty")
	}

	var argvp *uint16
	if len(cmdline) != 0 {
		argvp, err = syscall.UTF16PtrFromString(cmdline)
		if err != nil {
			return 0, 0, err
		}
	}

	var dirp *uint16
	if len(attr.Dir) != 0 {
		dirp, err = syscall.UTF16PtrFromString(attr.Dir)
		if err != nil {
			return 0, 0, err
		}
	}

	// Acquire the fork lock so that no other threads
	// create new fds that are not yet close-on-exec
	// before we fork.
	ForkLock.Lock()
	defer ForkLock.Unlock()

	p, _ := syscall.GetCurrentProcess()
	fd := make([]syscall.Handle, len(attr.Files))
	for i := range attr.Files {
		if attr.Files[i] > 0 {
			err := syscall.DuplicateHandle(p, syscall.Handle(attr.Files[i]), p, &fd[i], 0, true, syscall.DUPLICATE_SAME_ACCESS)
			if err != nil {
				return 0, 0, err
			}
			defer syscall.CloseHandle(syscall.Handle(fd[i]))
		}
	}
	si := new(syscall.StartupInfo)
	si.Cb = uint32(unsafe.Sizeof(*si))

	pi := new(syscall.ProcessInformation)

	flags := sys.CreationFlags | syscall.CREATE_UNICODE_ENVIRONMENT
	err = syscall.CreateProcess(nil, argvp, nil, nil, false, flags, nil, dirp, si, pi)
	if err != nil {
		return 0, 0, err
	}
	defer syscall.CloseHandle(syscall.Handle(pi.Thread))

	return int(pi.ProcessId), uintptr(pi.Process), nil
}

func showErrorMessage(hWnd winapi.HWND, msg string) {
	s, _ := syscall.UTF16PtrFromString(msg)
	t, _ := syscall.UTF16PtrFromString(WINDOW_CLASS)
	winapi.MessageBox(hWnd, s, t, winapi.MB_ICONWARNING|winapi.MB_OK)
}

func run() int {
	hInstance := winapi.GetModuleHandle(nil)

	if registerWindowClass(hInstance) == 0 {
		showErrorMessage(0, "registerWindowClass failed")
		return 1
	}

	if err := initializeInstance(hInstance, winapi.SW_SHOW); err != nil {
		showErrorMessage(0, err.Error())
		return 1
	}

	var msg winapi.MSG
	for winapi.GetMessage(&msg, 0,0,0) != 0 {
		winapi.TranslateMessage(&msg)
		winapi.DispatchMessage(&msg)
	}

	finalizeInstance(hInstance)

	return int(msg.WParam)
}

func registerWindowClass(hInstance winapi.HINSTANCE) winapi.ATOM {
	var wc winapi.WNDCLASSEX

	wc.CbSize = uint32(unsafe.Sizeof(winapi.WNDCLASSEX{}))
	wc.Style = 0
	wc.LpfnWndProc = syscall.NewCallback(wndProc)
	wc.CbClsExtra = 0
	wc.CbWndExtra = 0
	wc.HInstance = hInstance
	wc.HIcon         = winapi.LoadIcon(hInstance, winapi.MAKEINTRESOURCE(132))
	wc.HCursor       = winapi.LoadCursor(0, winapi.MAKEINTRESOURCE(winapi.IDC_HAND))
	wc.HbrBackground = winapi.HBRUSH(winapi.GetStockObject(winapi.WHITE_BRUSH))
	wc.LpszMenuName  = nil
	wc.LpszClassName,_ = syscall.UTF16PtrFromString(WINDOW_CLASS)

	return winapi.RegisterClassEx(&wc)
}

func initializeInstance(hInstance winapi.HINSTANCE, nCmdShow int) error {
	pc, _ := syscall.UTF16PtrFromString(WINDOW_CLASS)
	pt, _ := syscall.UTF16PtrFromString(WINDOW_CLASS)

	hWnd := winapi.CreateWindowEx(
		winapi.WS_EX_TOOLWINDOW|winapi.WS_EX_TOPMOST|winapi.WS_EX_NOACTIVATE|winapi.WS_EX_LAYERED,
		pc,pt,winapi.WS_POPUP,
		0,
		0,
		0,
		0,
		0,0,hInstance,nil)

	if hWnd == 0 {
		return errors.New("CreateWindowEx failed")
	}

	winapi.ShowWindow(hWnd, int32(nCmdShow))
	winapi.UpdateWindow(hWnd)

	if (RegisterHotKey(hWnd, 0, MOD_ALT|MOD_CONTROL, uint(config.Key)) == false) {
		return errors.New("RegisterHotKey failed")
	}

	return nil
}

func finalizeInstance(hInstance winapi.HINSTANCE) error {
	return nil
}

func isProcessGone(id uint32) bool {
	hProcess,err := syscall.OpenProcess(SYNCHRONIZE, false, id);
	if err != nil {
		return true;
	}

	dwWait,err := syscall.WaitForSingleObject(hProcess,0);
	if err != nil {
		return false
	}

	err = syscall.CloseHandle(hProcess)
	if err != nil {
		return false
	}

	return(dwWait == WAIT_OBJECT_0);
}

func createChildProcess(id uint32) (uint32,error) {
	if(!isProcessGone(id)) {
		return id,nil;
	}

	var sys syscall.SysProcAttr
	var attr syscall.ProcAttr

	sys.CreationFlags = CREATE_NEW_CONSOLE
	sys.CmdLine       = config.Cmdline

	attr.Files = make([]uintptr, 3)
	attr.Sys = &sys
	if config.WorkDir == "" {
		attr.Dir = os.Getenv("USERPROFILE")
	} else {
		attr.Dir = config.WorkDir
	}

	pid, _, err := createProcess(&attr)
	if err != nil {
		return 0,err
	}

	return uint32(pid), nil
}

func toggleWindowbyProcID(hwnd syscall.Handle, lparam uintptr) uintptr {
	var pid uint32
	_, pid = GetWindowThreadProcessId(hwnd)

	if (uint32(lparam) == pid) {
		if (winapi.IsWindowVisible(winapi.HWND(hwnd))) {
			winapi.ShowWindow(winapi.HWND(hwnd), winapi.SW_MINIMIZE);
			winapi.ShowWindow(winapi.HWND(hwnd), winapi.SW_HIDE);
		} else {
			winapi.ShowWindow(winapi.HWND(hwnd), winapi.SW_SHOW);
			winapi.ShowWindow(winapi.HWND(hwnd), winapi.SW_RESTORE);
		}
		return 0
	}
	return 1
}

func wndProc(hWnd winapi.HWND, msg uint32, wParam uintptr, lParam uintptr) uintptr {
	switch msg {
	case winapi.WM_HOTKEY:
		var err error
		dwProcessId,err = createChildProcess(dwProcessId);
		if err != nil {
			procReplyMessage.Call(1)
			return 0
		}

		EnumWindows(toggleWindowbyProcID, dwProcessId)
		return 0
	case winapi.WM_DESTROY:
		winapi.PostQuitMessage(0)
	default:
		return winapi.DefWindowProc(hWnd, msg, wParam, lParam)
	}
	return 0
}
<|endoftext|>


<commit_before>package disruptor

// Cursors should be a party of the same backing array to keep them as close together as possible:
// https://news.ycombinator.com/item?id=7800825

type Wireup struct {
	capacity int64
	groups   [][]Consumer
	cursors  []*Cursor // backing array keeps cursors (with padding) in contiguous memory
}

func Configure(capacity int64) Wireup {
	return Wireup{capacity: capacity, cursors: []*Cursor{NewCursor()}}
}

func (this Wireup) WithConsumerGroup(consumers ...Consumer) Wireup {
	if len(consumers) == 0 {
		return this
	}

	target := make([]Consumer, len(consumers))
	copy(target, consumers)

	for i := 0; i < len(consumers); i++ {
		this.cursors = append(this.cursors, NewCursor())
	}

	this.groups = append(this.groups, target)
	return this
}

func (this Wireup) Build() Disruptor {
	var listeners []ListenCloser
	var upstream Barrier = this.cursors[0]
	written := this.cursors[0]
	cursorIndex := 1 // 0 index is reserved for the writer Cursor

	for groupIndex, group := range this.groups {
		groupReaders, groupBarrier := this.buildReaders(groupIndex, cursorIndex, written, upstream)
		listeners = append(listeners, groupReaders...)
		upstream = groupBarrier
		cursorIndex += len(group)
	}

	sequencer := NewSingleSequencer(written, upstream, this.capacity)
	return Disruptor{sequencer: sequencer, listeners: listeners}
}

func (this Wireup) buildReaders(consumerIndex, cursorIndex int, writeSequence *Cursor, upstream Barrier) ([]ListenCloser, Barrier) {
	var barrierCursors []*Cursor
	var listeners []ListenCloser

	for _, consumer := range this.groups[consumerIndex] {
		readSequence := this.cursors[cursorIndex]
		barrierCursors = append(barrierCursors, readSequence)
		reader := NewReader(readSequence, writeSequence, upstream, consumer)
		listeners = append(listeners, reader)
		cursorIndex++
	}

	if len(barrierCursors) == 0 {
		panic("no barriers")
	}

	if len(this.groups[consumerIndex]) == 1 {
		return listeners, barrierCursors[0]
	} else {
		return listeners, NewCompositeBarrier(barrierCursors)
	}
}
<commit_msg>Inlined comment.<commit_after>package disruptor

type Wireup struct {
	capacity int64
	groups   [][]Consumer
	cursors  []*Cursor // slice keeps cursors in contiguous memory https://news.ycombinator.com/item?id=7800825
}

func Configure(capacity int64) Wireup {
	return Wireup{capacity: capacity, cursors: []*Cursor{NewCursor()}}
}

func (this Wireup) WithConsumerGroup(consumers ...Consumer) Wireup {
	if len(consumers) == 0 {
		return this
	}

	target := make([]Consumer, len(consumers))
	copy(target, consumers)

	for i := 0; i < len(consumers); i++ {
		this.cursors = append(this.cursors, NewCursor())
	}

	this.groups = append(this.groups, target)
	return this
}

func (this Wireup) Build() Disruptor {
	var listeners []ListenCloser
	var upstream Barrier = this.cursors[0]
	written := this.cursors[0]
	cursorIndex := 1 // 0 index is reserved for the writer Cursor

	for groupIndex, group := range this.groups {
		groupReaders, groupBarrier := this.buildReaders(groupIndex, cursorIndex, written, upstream)
		listeners = append(listeners, groupReaders...)
		upstream = groupBarrier
		cursorIndex += len(group)
	}

	sequencer := NewSingleSequencer(written, upstream, this.capacity)
	return Disruptor{sequencer: sequencer, listeners: listeners}
}

func (this Wireup) buildReaders(consumerIndex, cursorIndex int, writeSequence *Cursor, upstream Barrier) ([]ListenCloser, Barrier) {
	var barrierCursors []*Cursor
	var listeners []ListenCloser

	for _, consumer := range this.groups[consumerIndex] {
		readSequence := this.cursors[cursorIndex]
		barrierCursors = append(barrierCursors, readSequence)
		reader := NewReader(readSequence, writeSequence, upstream, consumer)
		listeners = append(listeners, reader)
		cursorIndex++
	}

	if len(barrierCursors) == 0 {
		panic("no barriers")
	}

	if len(this.groups[consumerIndex]) == 1 {
		return listeners, barrierCursors[0]
	} else {
		return listeners, NewCompositeBarrier(barrierCursors)
	}
}
<|endoftext|>


<commit_before>package gototo

import (
	"encoding/json"
	"fmt"
	"log"
	"reflect"
	"runtime"
	"sync/atomic"
	"time"

	zmq "github.com/JeremyOT/gototo/internal/github.com/pebbe/zmq4"
	"github.com/mitchellh/mapstructure"
)

// WorkerFunction may be registered with a worker
type WorkerFunction func(interface{}) interface{}

// MarshalFunction converts the result of a WorkerFunction to a byte slice to be
// sent back to the caller.
type MarshalFunction func(interface{}) ([]byte, error)

// UnmarshalFunction converts a byte slice to a worker invocation payload.
type UnmarshalFunction func([]byte) (map[string]interface{}, error)

func unmarshalJSON(buf []byte) (data map[string]interface{}, err error) {
	err = json.Unmarshal(buf, &data)
	return
}

// Response is a general response container that will be returned with Success = false
// if any errors are encountered while attempting to call a type safe function.
type Response struct {
	Success bool        `json:"success"`
	Error   string      `json:"error,omitempty"`
	Result  interface{} `json:"result,omitempty"`
}

// Validator allows automatic validation when calling type safe functions. If implemented
// by the type passed as an argument, Type.Validate() will be called before calling the
// worker function. If an error is returned then that error will be sent in response to
// the worker request and the method will not be called.
type Validator interface {
	Validate() error
}

// RequestUnpacker allows custom translation of request objects before calling type safe functions.
// Before calling a type safe function, UnpackRequest() will be called on a new instance of the
// parameter type. Note that mapstructure will not be used if this interface is implemented.
type RequestUnpacker interface {
	UnpackRequest(interface{}) error
}

// Worker listens for requests and invokes registered goroutines when called.
type Worker struct {
	logger                    *log.Logger
	registeredWorkerFunctions map[string]WorkerFunction
	activeTimeout             time.Duration
	passiveTimeout            time.Duration
	quit                      chan int
	wait                      chan int
	address                   string
	maxWorkers                int32
	runningWorkers            int32
	marshal                   MarshalFunction
	unmarshal                 UnmarshalFunction
	convertTypeTagName        string
	convertTypeDecoderConfig  *mapstructure.DecoderConfig
}

// Create a new worker bound to address that will run at most count functions at a time.
// If count == 0, it will default to runtime.NumCPU().
func New(address string, count int) *Worker {
	if count == 0 {
		count = runtime.NumCPU()
	}
	return &Worker{registeredWorkerFunctions: make(map[string]WorkerFunction),
		activeTimeout:      time.Millisecond,
		passiveTimeout:     100 * time.Millisecond,
		quit:               make(chan int),
		wait:               make(chan int),
		maxWorkers:         int32(count),
		address:            address,
		marshal:            json.Marshal,
		unmarshal:          unmarshalJSON,
		convertTypeTagName: "json",
	}
}

func (w *Worker) Quit() {
	w.quit <- 1
}

func (w *Worker) Wait() {
	<-w.wait
}

func (w *Worker) SetLogger(l *log.Logger) {
	w.logger = l
}

func (w *Worker) SetMarshalFunction(marshal MarshalFunction) {
	w.marshal = marshal
}

func (w *Worker) SetUnmarshalFunction(unmarshal UnmarshalFunction) {
	w.unmarshal = unmarshal
}

// SetConvertTypeTagName sets the tag name to use to find field information when
// converting request parameters to custom types. By default this is "json"
func (w *Worker) SetConvertTypeTagName(tagName string) {
	w.convertTypeTagName = tagName
}

// SetConvertTypeDecoderConfig sets the mapstructure config to use when decoding.
// if set, it takes precidence over SetConvertTypeTagName
func (w *Worker) SetConvertTypeDecoderConfig(config *mapstructure.DecoderConfig) {
	w.convertTypeDecoderConfig = config
}

func (w *Worker) writeLog(message ...interface{}) {
	if w.logger != nil {
		w.logger.Println(message)
	} else {
		log.Println(message)
	}
}

func (w *Worker) panicLog(message ...interface{}) {
	if w.logger != nil {
		w.logger.Panicln(message)
	} else {
		log.Panicln(message)
	}
}

func (w *Worker) RegisterWorkerFunction(name string, workerFunction WorkerFunction) {
	w.registeredWorkerFunctions[name] = workerFunction
}

// MakeWorkerFunction creates a wrapper around a function that allows a function to be
// called in a type safe manner. The function is expected to take a pointer to a struct
// as its only argument and return one value. Internally, github.com/mitchellh/mapstructure
// is used to convert the input parameters to a struct. MakeWorkerFunction will panic
// if called with an incorrect type.
func (w *Worker) MakeWorkerFunction(workerFunction interface{}) WorkerFunction {
	function := reflect.ValueOf(workerFunction)
	functionType := function.Type()
	if functionType.Kind() != reflect.Func {
		w.panicLog(fmt.Sprintf("Attempt to convert invalid type %#v to worker function", functionType))
	}
	if functionType.IsVariadic() {
		w.panicLog(fmt.Sprintf("Attempt to convert variadic function %#v to worker function", function))
	}
	if functionType.NumIn() != 1 {
		w.panicLog(fmt.Sprintf("Worker functions must accept one and only one argument: %#v", function))
	}
	if functionType.NumOut() != 1 {
		w.panicLog(fmt.Sprintf("Worker functions must only return one and only one result: %#v", function))
	}
	inputType := functionType.In(0)
	if inputType.Kind() != reflect.Ptr {
		w.panicLog(fmt.Sprintf("Worker functions must take a pointer to a struct as their argument: %#v", function))
	}
	validate := inputType.Implements(reflect.TypeOf((*Validator)(nil)).Elem())
	customUnpack := inputType.Implements(reflect.TypeOf((*RequestUnpacker)(nil)).Elem())
	inputType = inputType.Elem()
	if inputType.Kind() != reflect.Struct {
		w.panicLog(fmt.Sprintf("Worker functions must take a pointer to a struct as their argument: %#v", function))
	}
	return func(input interface{}) (output interface{}) {
		inputValue := reflect.New(inputType)
		parameters := inputValue.Interface()
		if customUnpack {
			if err := parameters.(RequestUnpacker).UnpackRequest(input); err != nil {
				w.writeLog(fmt.Sprintf("Failed to convert parameters to type %v: %s", inputType, err))
				return &Response{Success: false, Error: fmt.Sprintf("Failed to convert parameters to type %v: %s", inputType, err)}
			}
		} else {
			var config *mapstructure.DecoderConfig
			if w.convertTypeDecoderConfig != nil {
				config = &mapstructure.DecoderConfig{
					Metadata:         w.convertTypeDecoderConfig.Metadata,
					Result:           parameters,
					TagName:          w.convertTypeDecoderConfig.TagName,
					ErrorUnused:      w.convertTypeDecoderConfig.ErrorUnused,
					ZeroFields:       w.convertTypeDecoderConfig.ZeroFields,
					WeaklyTypedInput: w.convertTypeDecoderConfig.WeaklyTypedInput,
					DecodeHook:       w.convertTypeDecoderConfig.DecodeHook,
				}
			} else {
				config = &mapstructure.DecoderConfig{
					Metadata: nil,
					Result:   parameters,
					TagName:  w.convertTypeTagName,
				}
			}
			decoder, err := mapstructure.NewDecoder(config)
			if err != nil {
				w.writeLog("Failed to construct decoder:", err)
				return &Response{Success: false, Error: fmt.Sprintf("Failed to construct decoder: %s", err)}
			}
			if err = decoder.Decode(input); err != nil {
				w.writeLog(fmt.Sprintf("Failed to convert parameters to type %v: %s", inputType, err))
				return &Response{Success: false, Error: fmt.Sprintf("Failed to convert parameters to type %v: %s", inputType, err)}
			}
		}
		if validate {
			if err := parameters.(Validator).Validate(); err != nil {
				w.writeLog("Validation failed:", err)
				return &Response{Success: false, Error: fmt.Sprintf("Validation failed: %s", err)}
			}
		}
		return function.Call([]reflect.Value{inputValue})[0].Interface()
	}
}

// Run a worker function and send the response to responseChannel
func (w *Worker) runFunction(responseChannel chan [][]byte, message [][]byte, parameters interface{}, workerFunction WorkerFunction) {
	response := workerFunction(parameters)
	responseData, _ := w.marshal(response)
	message[len(message)-1] = responseData
	responseChannel <- message
}

func (w *Worker) Run() {
	defer func() { close(w.wait) }()
	socket, err := zmq.NewSocket(zmq.ROUTER)
	if err != nil {
		log.Println("Failed to start worker:", err)
	}
	defer socket.Close()
	socket.SetRcvtimeo(w.passiveTimeout)
	socket.Bind(w.address)
	atomic.StoreInt32(&w.runningWorkers, 0)
	responseChannel := make(chan [][]byte)
	sendResponse := func(response [][]byte) {
		atomic.AddInt32(&w.runningWorkers, -1)
		socket.SendMessage(response)
		if atomic.LoadInt32(&w.runningWorkers) == 0 {
			socket.SetRcvtimeo(w.passiveTimeout)
		}
	}
	for {
		if atomic.LoadInt32(&w.runningWorkers) == w.maxWorkers {
			select {
			case response := <-responseChannel:
				sendResponse(response)
			case <-w.quit:
				return
			}
		}
		select {
		case <-w.quit:
			return
		case response := <-responseChannel:
			sendResponse(response)
			break
		default:
			message, err := socket.RecvMessageBytes(0)
			if err != nil {
				// Needed to yield to goroutines when GOMAXPROCS is 1.
				// Note: The 1.3 preemptive scheduler doesn't seem to work here,
				// so this is still required.
				runtime.Gosched()
				break
			}
			data, err := w.unmarshal(message[len(message)-1])
			if err != nil {
				w.writeLog("Received invalid message", err)
				break
			}
			workerFunction := w.registeredWorkerFunctions[data["method"].(string)]
			if workerFunction == nil {
				w.writeLog("Unregistered worker function:", data["method"].(string))
				break
			}
			if atomic.LoadInt32(&w.runningWorkers) == 0 {
				socket.SetRcvtimeo(w.activeTimeout)
			}
			atomic.AddInt32(&w.runningWorkers, 1)
			go w.runFunction(responseChannel, message, data["parameters"], workerFunction)
		}
	}
}
<commit_msg>more cleanup, handle panics<commit_after>package gototo

import (
	"encoding/json"
	"fmt"
	"log"
	"reflect"
	"runtime"
	"sync/atomic"
	"time"

	zmq "github.com/JeremyOT/gototo/internal/github.com/pebbe/zmq4"
	"github.com/mitchellh/mapstructure"
)

// WorkerFunction may be registered with a worker
type WorkerFunction func(interface{}) interface{}

// MarshalFunction converts the result of a WorkerFunction to a byte slice to be
// sent back to the caller.
type MarshalFunction func(interface{}) ([]byte, error)

// UnmarshalFunction converts a byte slice to a worker invocation payload.
type UnmarshalFunction func([]byte) (*Request, error)

func unmarshalJSON(buf []byte) (data *Request, err error) {
	err = json.Unmarshal(buf, &data)
	return
}

// Response is a general response container that will be returned with Success = false
// if any errors are encountered while attempting to call a type safe function.
type Response struct {
	Success bool        `json:"success"`
	Error   string      `json:"error,omitempty"`
	Result  interface{} `json:"result,omitempty"`
}

// Request is a general request container. It has a Method identifier and Parameters.
type Request struct {
	Method     string      `json:"method"`
	Parameters interface{} `json:"parameters"`
}

// Validator allows automatic validation when calling type safe functions. If implemented
// by the type passed as an argument, Type.Validate() will be called before calling the
// worker function. If an error is returned then that error will be sent in response to
// the worker request and the method will not be called.
type Validator interface {
	Validate() error
}

// RequestUnpacker allows custom translation of request objects before calling type safe functions.
// Before calling a type safe function, UnpackRequest() will be called on a new instance of the
// parameter type. Note that mapstructure will not be used if this interface is implemented.
type RequestUnpacker interface {
	UnpackRequest(interface{}) error
}

// Worker listens for requests and invokes registered goroutines when called.
type Worker struct {
	logger                    *log.Logger
	registeredWorkerFunctions map[string]WorkerFunction
	activeTimeout             time.Duration
	passiveTimeout            time.Duration
	quit                      chan struct{}
	wait                      chan struct{}
	address                   string
	maxWorkers                int32
	runningWorkers            int32
	marshal                   MarshalFunction
	unmarshal                 UnmarshalFunction
	convertTypeTagName        string
	convertTypeDecoderConfig  *mapstructure.DecoderConfig
}

// Create a new worker bound to address that will run at most count functions at a time.
// If count == 0, it will default to runtime.NumCPU().
func New(address string, count int) *Worker {
	if count == 0 {
		count = runtime.NumCPU()
	}
	return &Worker{registeredWorkerFunctions: make(map[string]WorkerFunction),
		activeTimeout:      time.Millisecond,
		passiveTimeout:     100 * time.Millisecond,
		maxWorkers:         int32(count),
		address:            address,
		marshal:            json.Marshal,
		unmarshal:          unmarshalJSON,
		convertTypeTagName: "json",
	}
}

func (w *Worker) Quit() {
	close(w.quit)
}

func (w *Worker) Wait() {
	<-w.wait
}

func (w *Worker) SetLogger(l *log.Logger) {
	w.logger = l
}

func (w *Worker) SetMarshalFunction(marshal MarshalFunction) {
	w.marshal = marshal
}

func (w *Worker) SetUnmarshalFunction(unmarshal UnmarshalFunction) {
	w.unmarshal = unmarshal
}

// SetConvertTypeTagName sets the tag name to use to find field information when
// converting request parameters to custom types. By default this is "json"
func (w *Worker) SetConvertTypeTagName(tagName string) {
	w.convertTypeTagName = tagName
}

// SetConvertTypeDecoderConfig sets the mapstructure config to use when decoding.
// if set, it takes precidence over SetConvertTypeTagName
func (w *Worker) SetConvertTypeDecoderConfig(config *mapstructure.DecoderConfig) {
	w.convertTypeDecoderConfig = config
}

func (w *Worker) writeLog(message ...interface{}) {
	if w.logger != nil {
		w.logger.Println(message)
	} else {
		log.Println(message)
	}
}

func (w *Worker) panicLog(message ...interface{}) {
	if w.logger != nil {
		w.logger.Panicln(message)
	} else {
		log.Panicln(message)
	}
}

func (w *Worker) RegisterWorkerFunction(name string, workerFunction WorkerFunction) {
	w.registeredWorkerFunctions[name] = workerFunction
}

// MakeWorkerFunction creates a wrapper around a function that allows a function to be
// called in a type safe manner. The function is expected to take a pointer to a struct
// as its only argument and return one value. Internally, github.com/mitchellh/mapstructure
// is used to convert the input parameters to a struct. MakeWorkerFunction will panic
// if called with an incorrect type.
func (w *Worker) MakeWorkerFunction(workerFunction interface{}) WorkerFunction {
	function := reflect.ValueOf(workerFunction)
	functionType := function.Type()
	if functionType.Kind() != reflect.Func {
		w.panicLog(fmt.Sprintf("Attempt to convert invalid type %#v to worker function", functionType))
	}
	if functionType.IsVariadic() {
		w.panicLog(fmt.Sprintf("Attempt to convert variadic function %#v to worker function", function))
	}
	if functionType.NumIn() != 1 {
		w.panicLog(fmt.Sprintf("Worker functions must accept one and only one argument: %#v", function))
	}
	if functionType.NumOut() != 1 {
		w.panicLog(fmt.Sprintf("Worker functions must only return one and only one result: %#v", function))
	}
	inputType := functionType.In(0)
	if inputType.Kind() != reflect.Ptr {
		w.panicLog(fmt.Sprintf("Worker functions must take a pointer to a struct as their argument: %#v", function))
	}
	validate := inputType.Implements(reflect.TypeOf((*Validator)(nil)).Elem())
	customUnpack := inputType.Implements(reflect.TypeOf((*RequestUnpacker)(nil)).Elem())
	inputType = inputType.Elem()
	if inputType.Kind() != reflect.Struct {
		w.panicLog(fmt.Sprintf("Worker functions must take a pointer to a struct as their argument: %#v", function))
	}
	return func(input interface{}) (output interface{}) {
		inputValue := reflect.New(inputType)
		parameters := inputValue.Interface()
		if customUnpack {
			if err := parameters.(RequestUnpacker).UnpackRequest(input); err != nil {
				w.writeLog(fmt.Sprintf("Failed to convert parameters to type %v: %s", inputType, err))
				return &Response{Success: false, Error: fmt.Sprintf("Failed to convert parameters to type %v: %s", inputType, err)}
			}
		} else {
			var config *mapstructure.DecoderConfig
			if w.convertTypeDecoderConfig != nil {
				config = &mapstructure.DecoderConfig{
					Metadata:         w.convertTypeDecoderConfig.Metadata,
					Result:           parameters,
					TagName:          w.convertTypeDecoderConfig.TagName,
					ErrorUnused:      w.convertTypeDecoderConfig.ErrorUnused,
					ZeroFields:       w.convertTypeDecoderConfig.ZeroFields,
					WeaklyTypedInput: w.convertTypeDecoderConfig.WeaklyTypedInput,
					DecodeHook:       w.convertTypeDecoderConfig.DecodeHook,
				}
			} else {
				config = &mapstructure.DecoderConfig{
					Metadata: nil,
					Result:   parameters,
					TagName:  w.convertTypeTagName,
				}
			}
			decoder, err := mapstructure.NewDecoder(config)
			if err != nil {
				w.writeLog("Failed to construct decoder:", err)
				return &Response{Success: false, Error: fmt.Sprintf("Failed to construct decoder: %s", err)}
			}
			if err = decoder.Decode(input); err != nil {
				w.writeLog(fmt.Sprintf("Failed to convert parameters to type %v: %s", inputType, err))
				return &Response{Success: false, Error: fmt.Sprintf("Failed to convert parameters to type %v: %s", inputType, err)}
			}
		}
		if validate {
			if err := parameters.(Validator).Validate(); err != nil {
				w.writeLog("Validation failed:", err)
				return &Response{Success: false, Error: fmt.Sprintf("Validation failed: %s", err)}
			}
		}
		return function.Call([]reflect.Value{inputValue})[0].Interface()
	}
}

func (w *Worker) handlePanic(responseChannel chan [][]byte, message [][]byte) {
	if r := recover(); r != nil {
		errString := fmt.Sprintf("Panic while invoking worker function: %#v", r)
		w.writeLog(errString)
		if responseData, err := w.marshal(&Response{Success: false, Error: errString}); err != nil {
			w.writeLog("Error encoding response:", err)
		} else {
			message[len(message)-1] = responseData
			responseChannel <- message
		}
	}
}

// Run a worker function and send the response to responseChannel
func (w *Worker) runFunction(responseChannel chan [][]byte, message [][]byte, parameters interface{}, workerFunction WorkerFunction) {
	defer w.handlePanic(responseChannel, message)
	response := workerFunction(parameters)
	if responseData, err := w.marshal(response); err != nil {
		panic(err.Error())
	} else {
		message[len(message)-1] = responseData
		responseChannel <- message
	}
}

func (w *Worker) Run() {
	w.wait = make(chan struct{})
	w.quit = make(chan struct{})
	defer close(w.wait)
	socket, err := zmq.NewSocket(zmq.ROUTER)
	if err != nil {
		log.Println("Failed to start worker:", err)
	}
	defer socket.Close()
	socket.SetRcvtimeo(w.passiveTimeout)
	socket.Bind(w.address)
	atomic.StoreInt32(&w.runningWorkers, 0)
	responseChannel := make(chan [][]byte)
	sendResponse := func(response [][]byte) {
		atomic.AddInt32(&w.runningWorkers, -1)
		socket.SendMessage(response)
		if atomic.LoadInt32(&w.runningWorkers) == 0 {
			socket.SetRcvtimeo(w.passiveTimeout)
		}
	}
	for {
		if atomic.LoadInt32(&w.runningWorkers) == w.maxWorkers {
			select {
			case response := <-responseChannel:
				sendResponse(response)
			case <-w.quit:
				return
			}
		}
		select {
		case <-w.quit:
			return
		case response := <-responseChannel:
			sendResponse(response)
			break
		default:
			message, err := socket.RecvMessageBytes(0)
			if err != nil {
				// Needed to yield to goroutines when GOMAXPROCS is 1.
				// Note: The 1.3 preemptive scheduler doesn't seem to work here,
				// so this is still required.
				runtime.Gosched()
				break
			}
			request, err := w.unmarshal(message[len(message)-1])
			if err != nil {
				w.writeLog("Received invalid message", err)
				break
			}
			workerFunction := w.registeredWorkerFunctions[request.Method]
			if workerFunction == nil {
				w.writeLog("Unregistered worker function:", request.Method)
				break
			}
			if atomic.LoadInt32(&w.runningWorkers) == 0 {
				socket.SetRcvtimeo(w.activeTimeout)
			}
			atomic.AddInt32(&w.runningWorkers, 1)
			go w.runFunction(responseChannel, message, request.Parameters, workerFunction)
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"doozer"
	"doozer/util"
	"flag"
	"fmt"
	"net"
	"os"
	"log"
)

var (
	listenAddr  = flag.String("l", "127.0.0.1:8046", "The address to bind to.")
	attachAddr  = flag.String("a", "", "The address of another node to attach to.")
	webAddr     = flag.String("w", ":8080", "Serve web requests on this address.")
	clusterName = flag.String("c", "local", "The non-empty cluster name.")
	showVersion = flag.Bool("v", false, "print doozerd's version string")
	pi          = flag.Float64("pulse", 1, "how often (in seconds) to set applied key")
	fd          = flag.Float64("fill", 1.5, "delay (in seconds) to fill unowned seqns")
	kt          = flag.Float64("timeout", 60, "timeout (in seconds) to kick inactive nodes")
)

func Usage() {
	fmt.Fprintf(os.Stderr, "Usage: %s [OPTIONS]\n", os.Args[0])
	fmt.Fprintf(os.Stderr, "\nOptions:\n")
	flag.PrintDefaults()
}

func main() {
	util.LogWriter = os.Stderr
	flag.Usage = Usage
	flag.Parse()

	if *showVersion {
		fmt.Println("doozerd", doozer.Version)
		return
	}

	if *listenAddr == "" {
		fmt.Fprintln(os.Stderr, "require a listen address")
		flag.Usage()
		os.Exit(1)
	}

	log.SetPrefix("DOOZER " + *listenAddr + " ")

	listener, err := net.Listen("tcp", *listenAddr)
	if err != nil {
		panic(err)
	}

	conn, err := net.ListenPacket("udp", *listenAddr)
	if err != nil {
		panic(err)
	}

	var wl net.Listener
	if *webAddr != "" {
		wl, err = net.Listen("tcp", *webAddr)
		if err != nil {
			panic(err)
		}
	}

	doozer.Main(*clusterName, *attachAddr, conn, listener, wl, ns(*pi), ns(*fd), ns(*kt))
}

func ns(x float64) int64 {
	return int64(x * 1e9)
}
<commit_msg>forget this<commit_after>package main

import (
	"doozer"
	"doozer/util"
	"flag"
	"fmt"
	"net"
	"os"
	"log"
)

var (
	listenAddr  = flag.String("l", "127.0.0.1:8046", "The address to bind to.")
	attachAddr  = flag.String("a", "", "The address of another node to attach to.")
	webAddr     = flag.String("w", ":8080", "Serve web requests on this address.")
	clusterName = flag.String("c", "local", "The non-empty cluster name.")
	showVersion = flag.Bool("v", false, "print doozerd's version string")
	pi          = flag.Float64("pulse", 1, "how often (in seconds) to set applied key")
	fd          = flag.Float64("fill", 1.5, "delay (in seconds) to fill unowned seqns")
	kt          = flag.Float64("timeout", 60, "timeout (in seconds) to kick inactive nodes")
)

func Usage() {
	fmt.Fprintf(os.Stderr, "Usage: %s [OPTIONS]\n", os.Args[0])
	fmt.Fprintf(os.Stderr, "\nOptions:\n")
	flag.PrintDefaults()
}

func main() {
	util.LogWriter = os.Stderr
	flag.Usage = Usage
	flag.Parse()

	if *showVersion {
		fmt.Println("doozerd", doozer.Version)
		return
	}

	if *listenAddr == "" {
		fmt.Fprintln(os.Stderr, "require a listen address")
		flag.Usage()
		os.Exit(1)
	}

	log.SetPrefix("DOOZER ")

	listener, err := net.Listen("tcp", *listenAddr)
	if err != nil {
		panic(err)
	}

	conn, err := net.ListenPacket("udp", *listenAddr)
	if err != nil {
		panic(err)
	}

	var wl net.Listener
	if *webAddr != "" {
		wl, err = net.Listen("tcp", *webAddr)
		if err != nil {
			panic(err)
		}
	}

	doozer.Main(*clusterName, *attachAddr, conn, listener, wl, ns(*pi), ns(*fd), ns(*kt))
}

func ns(x float64) int64 {
	return int64(x * 1e9)
}
<|endoftext|>


<commit_before>// Copyright 2017 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

import (
	"bytes"
	cmddwarf "cmd/internal/dwarf"
	"cmd/internal/objfile"
	"debug/dwarf"
	"internal/testenv"
	"io/ioutil"
	"os"
	"os/exec"
	"path"
	"path/filepath"
	"runtime"
	"strings"
	"testing"
)

func testDWARF(t *testing.T, buildmode string, expectDWARF bool, env ...string) {
	testenv.MustHaveCGO(t)
	testenv.MustHaveGoBuild(t)

	if runtime.GOOS == "plan9" {
		t.Skip("skipping on plan9; no DWARF symbol table in executables")
	}

	t.Parallel()

	out, err := exec.Command(testenv.GoToolPath(t), "list", "-f", "{{.Stale}}", "cmd/link").CombinedOutput()
	if err != nil {
		t.Fatalf("go list: %v\n%s", err, out)
	}
	if string(out) != "false\n" {
		if strings.HasPrefix(testenv.Builder(), "darwin-") {
			t.Skipf("cmd/link is spuriously stale on Darwin builders - see #33598")
		}
		t.Fatalf("cmd/link is stale - run go install cmd/link")
	}

	for _, prog := range []string{"testprog", "testprogcgo"} {
		prog := prog
		expectDWARF := expectDWARF
		if runtime.GOOS == "aix" && prog == "testprogcgo" {
			extld := os.Getenv("CC")
			if extld == "" {
				extld = "gcc"
			}
			expectDWARF, err = cmddwarf.IsDWARFEnabledOnAIXLd(extld)
			if err != nil {
				t.Fatal(err)
			}

		}

		t.Run(prog, func(t *testing.T) {
			t.Parallel()

			tmpDir, err := ioutil.TempDir("", "go-link-TestDWARF")
			if err != nil {
				t.Fatal(err)
			}
			defer os.RemoveAll(tmpDir)

			exe := filepath.Join(tmpDir, prog+".exe")
			dir := "../../runtime/testdata/" + prog
			cmd := exec.Command(testenv.GoToolPath(t), "build", "-o", exe)
			if buildmode != "" {
				cmd.Args = append(cmd.Args, "-buildmode", buildmode)
			}
			cmd.Args = append(cmd.Args, dir)
			if env != nil {
				cmd.Env = append(os.Environ(), env...)
				cmd.Env = append(cmd.Env, "CGO_CFLAGS=") // ensure CGO_CFLAGS does not contain any flags. Issue #35459
			}
			out, err := cmd.CombinedOutput()
			if err != nil {
				t.Fatalf("go build -o %v %v: %v\n%s", exe, dir, err, out)
			}

			if buildmode == "c-archive" {
				// Extract the archive and use the go.o object within.
				cmd := exec.Command("ar", "-x", exe)
				cmd.Dir = tmpDir
				if out, err := cmd.CombinedOutput(); err != nil {
					t.Fatalf("ar -x %s: %v\n%s", exe, err, out)
				}
				exe = filepath.Join(tmpDir, "go.o")
			}

			if runtime.GOOS == "darwin" {
				if _, err = exec.LookPath("symbols"); err == nil {
					// Ensure Apple's tooling can parse our object for symbols.
					out, err = exec.Command("symbols", exe).CombinedOutput()
					if err != nil {
						t.Fatalf("symbols %v: %v: %s", filepath.Base(exe), err, out)
					} else {
						if bytes.HasPrefix(out, []byte("Unable to find file")) {
							// This failure will cause the App Store to reject our binaries.
							t.Fatalf("symbols %v: failed to parse file", filepath.Base(exe))
						} else if bytes.Contains(out, []byte(", Empty]")) {
							t.Fatalf("symbols %v: parsed as empty", filepath.Base(exe))
						}
					}
				}
			}

			f, err := objfile.Open(exe)
			if err != nil {
				t.Fatal(err)
			}
			defer f.Close()

			syms, err := f.Symbols()
			if err != nil {
				t.Fatal(err)
			}

			var addr uint64
			for _, sym := range syms {
				if sym.Name == "main.main" {
					addr = sym.Addr
					break
				}
			}
			if addr == 0 {
				t.Fatal("cannot find main.main in symbols")
			}

			d, err := f.DWARF()
			if err != nil {
				if expectDWARF {
					t.Fatal(err)
				}
				return
			} else {
				if !expectDWARF {
					t.Fatal("unexpected DWARF section")
				}
			}

			// TODO: We'd like to use filepath.Join here.
			// Also related: golang.org/issue/19784.
			wantFile := path.Join(prog, "main.go")
			wantLine := 24
			r := d.Reader()
			entry, err := r.SeekPC(addr)
			if err != nil {
				t.Fatal(err)
			}
			lr, err := d.LineReader(entry)
			if err != nil {
				t.Fatal(err)
			}
			var line dwarf.LineEntry
			if err := lr.SeekPC(addr, &line); err == dwarf.ErrUnknownPC {
				t.Fatalf("did not find file:line for %#x (main.main)", addr)
			} else if err != nil {
				t.Fatal(err)
			}
			if !strings.HasSuffix(line.File.Name, wantFile) || line.Line != wantLine {
				t.Errorf("%#x is %s:%d, want %s:%d", addr, line.File.Name, line.Line, filepath.Join("...", wantFile), wantLine)
			}
		})
	}
}

func TestDWARF(t *testing.T) {
	testDWARF(t, "", true)
	if !testing.Short() {
		if runtime.GOOS == "windows" {
			t.Skip("skipping Windows/c-archive; see Issue 35512 for more.")
		}
		t.Run("c-archive", func(t *testing.T) {
			testDWARF(t, "c-archive", true)
		})
	}
}

func TestDWARFiOS(t *testing.T) {
	// Normally we run TestDWARF on native platform. But on iOS we don't have
	// go build, so we do this test with a cross build.
	// Only run this on darwin/amd64, where we can cross build for iOS.
	if testing.Short() {
		t.Skip("skipping in short mode")
	}
	if runtime.GOARCH != "amd64" || runtime.GOOS != "darwin" {
		t.Skip("skipping on non-darwin/amd64 platform")
	}
	if err := exec.Command("xcrun", "--help").Run(); err != nil {
		t.Skipf("error running xcrun, required for iOS cross build: %v", err)
	}
	// Check to see if the ios tools are installed. It's possible to have the command line tools
	// installed without the iOS sdk.
	if output, err := exec.Command("xcodebuild -showsdks").CombinedOutput(); err != nil {
		t.Skipf("error running xcodebuild, required for iOS cross build: %v", err)
	} else if !strings.Contains(string(output), "iOS SDK") {
		t.Skipf("iOS SDK not detected.")
	}
	cc := "CC=" + runtime.GOROOT() + "/misc/ios/clangwrap.sh"
	// iOS doesn't allow unmapped segments, so iOS executables don't have DWARF.
	testDWARF(t, "", false, cc, "CGO_ENABLED=1", "GOOS=darwin", "GOARCH=arm64")
	// However, c-archive iOS objects have embedded DWARF.
	testDWARF(t, "c-archive", true, cc, "CGO_ENABLED=1", "GOOS=darwin", "GOARCH=arm64")
}
<commit_msg>cmd/link: fix TestDWARFiOS<commit_after>// Copyright 2017 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

import (
	"bytes"
	cmddwarf "cmd/internal/dwarf"
	"cmd/internal/objfile"
	"debug/dwarf"
	"internal/testenv"
	"io/ioutil"
	"os"
	"os/exec"
	"path"
	"path/filepath"
	"runtime"
	"strings"
	"testing"
)

func testDWARF(t *testing.T, buildmode string, expectDWARF bool, env ...string) {
	testenv.MustHaveCGO(t)
	testenv.MustHaveGoBuild(t)

	if runtime.GOOS == "plan9" {
		t.Skip("skipping on plan9; no DWARF symbol table in executables")
	}

	t.Parallel()

	out, err := exec.Command(testenv.GoToolPath(t), "list", "-f", "{{.Stale}}", "cmd/link").CombinedOutput()
	if err != nil {
		t.Fatalf("go list: %v\n%s", err, out)
	}
	if string(out) != "false\n" {
		if strings.HasPrefix(testenv.Builder(), "darwin-") {
			t.Skipf("cmd/link is spuriously stale on Darwin builders - see #33598")
		}
		t.Fatalf("cmd/link is stale - run go install cmd/link")
	}

	for _, prog := range []string{"testprog", "testprogcgo"} {
		prog := prog
		expectDWARF := expectDWARF
		if runtime.GOOS == "aix" && prog == "testprogcgo" {
			extld := os.Getenv("CC")
			if extld == "" {
				extld = "gcc"
			}
			expectDWARF, err = cmddwarf.IsDWARFEnabledOnAIXLd(extld)
			if err != nil {
				t.Fatal(err)
			}

		}

		t.Run(prog, func(t *testing.T) {
			t.Parallel()

			tmpDir, err := ioutil.TempDir("", "go-link-TestDWARF")
			if err != nil {
				t.Fatal(err)
			}
			defer os.RemoveAll(tmpDir)

			exe := filepath.Join(tmpDir, prog+".exe")
			dir := "../../runtime/testdata/" + prog
			cmd := exec.Command(testenv.GoToolPath(t), "build", "-o", exe)
			if buildmode != "" {
				cmd.Args = append(cmd.Args, "-buildmode", buildmode)
			}
			cmd.Args = append(cmd.Args, dir)
			if env != nil {
				cmd.Env = append(os.Environ(), env...)
				cmd.Env = append(cmd.Env, "CGO_CFLAGS=") // ensure CGO_CFLAGS does not contain any flags. Issue #35459
			}
			out, err := cmd.CombinedOutput()
			if err != nil {
				t.Fatalf("go build -o %v %v: %v\n%s", exe, dir, err, out)
			}

			if buildmode == "c-archive" {
				// Extract the archive and use the go.o object within.
				cmd := exec.Command("ar", "-x", exe)
				cmd.Dir = tmpDir
				if out, err := cmd.CombinedOutput(); err != nil {
					t.Fatalf("ar -x %s: %v\n%s", exe, err, out)
				}
				exe = filepath.Join(tmpDir, "go.o")
			}

			if runtime.GOOS == "darwin" {
				if _, err = exec.LookPath("symbols"); err == nil {
					// Ensure Apple's tooling can parse our object for symbols.
					out, err = exec.Command("symbols", exe).CombinedOutput()
					if err != nil {
						t.Fatalf("symbols %v: %v: %s", filepath.Base(exe), err, out)
					} else {
						if bytes.HasPrefix(out, []byte("Unable to find file")) {
							// This failure will cause the App Store to reject our binaries.
							t.Fatalf("symbols %v: failed to parse file", filepath.Base(exe))
						} else if bytes.Contains(out, []byte(", Empty]")) {
							t.Fatalf("symbols %v: parsed as empty", filepath.Base(exe))
						}
					}
				}
			}

			f, err := objfile.Open(exe)
			if err != nil {
				t.Fatal(err)
			}
			defer f.Close()

			syms, err := f.Symbols()
			if err != nil {
				t.Fatal(err)
			}

			var addr uint64
			for _, sym := range syms {
				if sym.Name == "main.main" {
					addr = sym.Addr
					break
				}
			}
			if addr == 0 {
				t.Fatal("cannot find main.main in symbols")
			}

			d, err := f.DWARF()
			if err != nil {
				if expectDWARF {
					t.Fatal(err)
				}
				return
			} else {
				if !expectDWARF {
					t.Fatal("unexpected DWARF section")
				}
			}

			// TODO: We'd like to use filepath.Join here.
			// Also related: golang.org/issue/19784.
			wantFile := path.Join(prog, "main.go")
			wantLine := 24
			r := d.Reader()
			entry, err := r.SeekPC(addr)
			if err != nil {
				t.Fatal(err)
			}
			lr, err := d.LineReader(entry)
			if err != nil {
				t.Fatal(err)
			}
			var line dwarf.LineEntry
			if err := lr.SeekPC(addr, &line); err == dwarf.ErrUnknownPC {
				t.Fatalf("did not find file:line for %#x (main.main)", addr)
			} else if err != nil {
				t.Fatal(err)
			}
			if !strings.HasSuffix(line.File.Name, wantFile) || line.Line != wantLine {
				t.Errorf("%#x is %s:%d, want %s:%d", addr, line.File.Name, line.Line, filepath.Join("...", wantFile), wantLine)
			}
		})
	}
}

func TestDWARF(t *testing.T) {
	testDWARF(t, "", true)
	if !testing.Short() {
		if runtime.GOOS == "windows" {
			t.Skip("skipping Windows/c-archive; see Issue 35512 for more.")
		}
		t.Run("c-archive", func(t *testing.T) {
			testDWARF(t, "c-archive", true)
		})
	}
}

func TestDWARFiOS(t *testing.T) {
	// Normally we run TestDWARF on native platform. But on iOS we don't have
	// go build, so we do this test with a cross build.
	// Only run this on darwin/amd64, where we can cross build for iOS.
	if testing.Short() {
		t.Skip("skipping in short mode")
	}
	if runtime.GOARCH != "amd64" || runtime.GOOS != "darwin" {
		t.Skip("skipping on non-darwin/amd64 platform")
	}
	if err := exec.Command("xcrun", "--help").Run(); err != nil {
		t.Skipf("error running xcrun, required for iOS cross build: %v", err)
	}
	// Check to see if the ios tools are installed. It's possible to have the command line tools
	// installed without the iOS sdk.
	if output, err := exec.Command("xcodebuild", "-showsdks").CombinedOutput(); err != nil {
		t.Skipf("error running xcodebuild, required for iOS cross build: %v", err)
	} else if !strings.Contains(string(output), "iOS SDK") {
		t.Skipf("iOS SDK not detected.")
	}
	cc := "CC=" + runtime.GOROOT() + "/misc/ios/clangwrap.sh"
	// iOS doesn't allow unmapped segments, so iOS executables don't have DWARF.
	t.Run("exe", func(t *testing.T) {
		testDWARF(t, "", false, cc, "CGO_ENABLED=1", "GOOS=ios", "GOARCH=arm64")
	})
	// However, c-archive iOS objects have embedded DWARF.
	t.Run("c-archive", func(t *testing.T) {
		testDWARF(t, "c-archive", true, cc, "CGO_ENABLED=1", "GOOS=ios", "GOARCH=arm64")
	})
}
<|endoftext|>


<commit_before>package iobit

import (
	"encoding/binary"
	"errors"
	"io"
)

type Writer struct {
	cache uint64
	dst   io.Writer
	fill  uint
	err   error
}

type bigEndian struct{}
type littleEndian struct{}

var (
	ErrUnderflow = errors.New("bit underflow")
	BigEndian    bigEndian
	LittleEndian littleEndian
)

func NewWriter(dst io.Writer) *Writer {
	return &Writer{dst: dst}
}

func (w *Writer) flushCache(bits uint) {
	if w.fill+bits <= 64 {
		return
	}
	var data [4]uint8
	binary.BigEndian.PutUint32(data[:], uint32(w.cache>>32))
	w.cache <<= 32
	w.fill -= 32
	w.write(data[:])
}

func (w *Writer) writeCache(bits uint, val uint32) {
	u := uint64(val)
	u &= ^(^uint64(0) << bits)
	u <<= 64 - w.fill - bits
	w.cache |= u
	w.fill += bits
}

func (bigEndian) PutUint32(w *Writer, bits uint, val uint32) {
	w.flushCache(bits)
	w.writeCache(bits, val)
}

func (littleEndian) PutUint32(w *Writer, bits uint, val uint32) {
	w.flushCache(bits)
	for bits > 8 {
		w.writeCache(8, val)
		val >>= 8
		bits -= 8
	}
	w.writeCache(bits, val)
}

func (bigEndian) PutUint64(w *Writer, bits uint, val uint64) {
	if bits > 32 {
		BigEndian.PutUint32(w, bits-32, uint32(val>>32))
		bits = 32
		val &= 0xFFFFFFFF
	}
	BigEndian.PutUint32(w, bits, uint32(val))
}

func (littleEndian) PutUint64(w *Writer, bits uint, val uint64) {
	if bits > 32 {
		LittleEndian.PutUint32(w, bits-32, uint32(val&0xFFFFFFFF))
		bits = 32
		val >>= 32
	}
	LittleEndian.PutUint32(w, bits, uint32(val))
}

func (w *Writer) write(data []uint8) {
	if w.err == nil {
		_, w.err = w.dst.Write(data)
	}
}

func (w *Writer) Flush() error {
	var data [8]uint8
	idx := 0
	for w.fill >= 8 {
		data[idx] = uint8(w.cache >> 56)
		w.cache <<= 8
		w.fill -= 8
		idx++
	}
	if w.fill != 0 {
		w.err = ErrUnderflow
	}
	w.write(data[:])
	return w.err
}

func (w *Writer) Write(p []uint8) (int, error) {
	err := w.Flush()
	if err != nil {
		return 0, err
	}
	return w.dst.Write(p)
}
<commit_msg>writer: tweak error checking<commit_after>package iobit

import (
	"encoding/binary"
	"errors"
	"io"
)

type Writer struct {
	dst   io.Writer
	cache uint64
	fill  uint
	err   error
}

type bigEndian struct{}
type littleEndian struct{}

var (
	ErrUnderflow = errors.New("bit underflow")
	BigEndian    bigEndian
	LittleEndian littleEndian
)

func NewWriter(dst io.Writer) *Writer {
	return &Writer{dst: dst}
}

func (w *Writer) flushCache(bits uint) {
	if w.fill+bits <= 64 {
		return
	}
	var data [4]uint8
	binary.BigEndian.PutUint32(data[:], uint32(w.cache>>32))
	w.cache <<= 32
	w.fill -= 32
	w.write(data[:])
}

func (w *Writer) writeCache(bits uint, val uint32) {
	u := uint64(val)
	u &= ^(^uint64(0) << bits)
	u <<= 64 - w.fill - bits
	w.cache |= u
	w.fill += bits
}

func (bigEndian) PutUint32(w *Writer, bits uint, val uint32) {
	w.flushCache(bits)
	w.writeCache(bits, val)
}

func (littleEndian) PutUint32(w *Writer, bits uint, val uint32) {
	w.flushCache(bits)
	for bits > 8 {
		w.writeCache(8, val)
		val >>= 8
		bits -= 8
	}
	w.writeCache(bits, val)
}

func (bigEndian) PutUint64(w *Writer, bits uint, val uint64) {
	if bits > 32 {
		BigEndian.PutUint32(w, bits-32, uint32(val>>32))
		bits = 32
		val &= 0xFFFFFFFF
	}
	BigEndian.PutUint32(w, bits, uint32(val))
}

func (littleEndian) PutUint64(w *Writer, bits uint, val uint64) {
	if bits > 32 {
		LittleEndian.PutUint32(w, bits-32, uint32(val&0xFFFFFFFF))
		bits = 32
		val >>= 32
	}
	LittleEndian.PutUint32(w, bits, uint32(val))
}

func (w *Writer) write(src []uint8) (int, error) {
	n := 0
	if w.err == nil {
		n, w.err = w.dst.Write(src)
	}
	return n, w.err
}

func (w *Writer) Flush() error {
	var data [8]uint8
	idx := 0
	for w.fill >= 8 {
		data[idx] = uint8(w.cache >> 56)
		w.cache <<= 8
		w.fill -= 8
		idx++
	}
	if w.err == nil && w.fill != 0 {
		w.err = ErrUnderflow
	}
	w.write(data[:])
	return w.err
}

func (w *Writer) Write(p []uint8) (int, error) {
	w.Flush()
	return w.write(p)
}
<|endoftext|>