<commit_before>package router

import (
	"bytes"
	"encoding/gob"
	"fmt"
	"io"
	"sync"
)

type Peers struct {
	sync.RWMutex
	table map[PeerName]*Peer
	onGC  func(*Peer)
}

func NewPeers(onGC func(*Peer)) *Peers {
	return &Peers{table: make(map[PeerName]*Peer), onGC: onGC}
}

func (peers *Peers) FetchWithDefault(peer *Peer) *Peer {
	peers.RLock()
	res, found := peers.fetchAlias(peer)
	peers.RUnlock()
	if found {
		return res
	}
	peers.Lock()
	defer peers.Unlock()
	res, found = peers.fetchAlias(peer)
	if found {
		return res
	}
	peers.table[peer.Name] = peer
	peer.IncrementLocalRefCount()
	return peer
}

func (peers *Peers) Fetch(name PeerName) (*Peer, bool) {
	peers.RLock()
	defer peers.RUnlock()
	peer, found := peers.table[name]
	return peer, found // GRRR, why can't I inline this!?
}

func (peers *Peers) ForEach(fun func(PeerName, *Peer)) {
	peers.RLock()
	defer peers.RUnlock()
	for name, peer := range peers.table {
		fun(name, peer)
	}
}

func (peers *Peers) ApplyUpdate(update []byte, router *Router) ([]byte, error) {
	peers.Lock()

	newPeers, decodedUpdate, decodedConns, err := peers.decodeUpdate(update, router)
	if err != nil {
		peers.Unlock()
		return nil, err
	}

	// By this point, we know the update doesn't refer to any peers we
	// have no knowledge of. We can now apply the update. Start by
	// adding in any new peers into the cache.
	for name, newPeer := range newPeers {
		peers.table[name] = newPeer
	}

	// Now apply the updates
	newUpdate := peers.applyUpdate(decodedUpdate, decodedConns, router)

	for _, peerRemoved := range peers.garbageCollect(router.Ourself) {
		delete(newUpdate, peerRemoved.Name)
	}

	// Don't need to hold peers lock any longer
	peers.Unlock()

	if len(newUpdate) > 0 {
		router.ConnectionMaker.Refresh()
	}

	return encodePeersMap(newUpdate), nil
}

func (peers *Peers) EncodeAllPeers() []byte {
	peers.RLock()
	defer peers.RUnlock()
	return encodePeersMap(peers.table)
}

func EncodePeers(peers ...*Peer) []byte {
	buf := new(bytes.Buffer)
	enc := gob.NewEncoder(buf)
	for _, peer := range peers {
		peer.encodePeer(enc)
	}
	return buf.Bytes()
}

func (peers *Peers) GarbageCollect(router *Router) []*Peer {
	peers.Lock()
	defer peers.Unlock()
	return peers.garbageCollect(router.Ourself)
}

func (peers *Peers) String() string {
	var buf bytes.Buffer
	peers.ForEach(func(name PeerName, peer *Peer) {
		buf.WriteString(fmt.Sprint(peer, "\n"))
		peer.ForEachConnection(func(remoteName PeerName, conn Connection) {
			buf.WriteString(fmt.Sprintf("   -> %v [%v]\n", remoteName, conn.RemoteTCPAddr()))
		})
	})
	return buf.String()
}

func (peers *Peers) fetchAlias(peer *Peer) (*Peer, bool) {
	if existingPeer, found := peers.table[peer.Name]; found {
		if existingPeer.UID == peer.UID {
			existingPeer.IncrementLocalRefCount()
			return existingPeer, true
		} else {
			return nil, true
		}
	}
	return nil, false
}

func (peers *Peers) garbageCollect(ourself *Peer) []*Peer {
	removed := []*Peer{}
	for name, peer := range peers.table {
		found, _ := ourself.Routes(peer, false)
		if !found && !peer.IsLocallyReferenced() {
			peers.onGC(peer)
			delete(peers.table, name)
			ourself.Router.Macs.Delete(peer)
			removed = append(removed, peer)
		}
	}
	return removed
}

func encodePeersMap(peers map[PeerName]*Peer) []byte {
	buf := new(bytes.Buffer)
	enc := gob.NewEncoder(buf)
	for _, peer := range peers {
		peer.encodePeer(enc)
	}
	return buf.Bytes()
}

func (peers *Peers) decodeUpdate(update []byte, router *Router) (newPeers map[PeerName]*Peer, decodedUpdate []*Peer, decodedConns [][]byte, err error) {
	newPeers = make(map[PeerName]*Peer)
	decodedUpdate = []*Peer{}
	decodedConns = [][]byte{}

	updateBuf := new(bytes.Buffer)
	updateBuf.Write(update)
	decoder := gob.NewDecoder(updateBuf)

	for {
		nameByte, uid, version, connsBuf, decErr := decodePeerNoConns(decoder)
		if decErr == io.EOF {
			break
		} else if decErr != nil {
			err = decErr
			return
		}
		name := PeerNameFromBin(nameByte)
		newPeer := NewPeer(name, uid, version, router)
		decodedUpdate = append(decodedUpdate, newPeer)
		decodedConns = append(decodedConns, connsBuf)
		existingPeer, found := peers.table[name]
		if !found {
			newPeers[name] = newPeer
		} else if existingPeer.UID != newPeer.UID {
			err = NameCollisionError{Name: newPeer.Name}
			return
		}
	}

	unknownPeers := false
	for _, connsBuf := range decodedConns {
		connsIterator(connsBuf, func(remoteNameByte []byte, _ string) {
			remoteName := PeerNameFromBin(remoteNameByte)
			if _, found := newPeers[remoteName]; found {
				return
			}
			if _, found := peers.table[remoteName]; found {
				return
			}
			// Update refers to a peer which we have no knowledge
			// of. Thus we can't apply the update. Abort.
			unknownPeers = true
		})
	}
	if unknownPeers {
		err = UnknownPeersError{}
		return
	}
	return
}

func (peers *Peers) applyUpdate(decodedUpdate []*Peer, decodedConns [][]byte, router *Router) map[PeerName]*Peer {
	newUpdate := make(map[PeerName]*Peer)
	for idx, newPeer := range decodedUpdate {
		connsBuf := decodedConns[idx]
		name := newPeer.Name
		// guaranteed to find peer in the peers.table
		peer := peers.table[name]
		if peer != newPeer {
			if peer.Version() > newPeer.Version() {
				// we know more about this one than the update. If
				// peer is ourself, this is slightly racey (further
				// changes could occur to ourself in the mean
				// time). But it doesn't matter as we know that we're
				// already > newPeer.Version() and that's not going to
				// change.
				newUpdate[name] = peer
				continue
			} else if peer == router.Ourself {
				// nobody but us updates us
				continue
			} else if peer.Version() == newPeer.Version() {
				// implication is that connections are equal too
				continue
			}
		}
		// If we're here, either it was a new peer, or the update has
		// more info about the peer than we do. Either case, we need
		// to set version and conns and include the updated peer in
		// the outgoing update.

		// Can peer have been updated by anyone else in the mean time?
		// No - we know that peer is not ourself, so the only prospect
		// for an update would be someone else calling
		// router.Peers.ApplyUpdate. But ApplyUpdate takes the Lock on
		// the router.Peers, so there can be no race here.
		conns := readConnsMap(peer, connsBuf, peers.table)
		peer.SetVersionAndConnections(newPeer.Version(), conns)
		newUpdate[name] = peer
	}
	return newUpdate
}

func (peer *Peer) encodePeer(enc *gob.Encoder) {
	peer.RLock()
	defer peer.RUnlock()

	checkFatal(enc.Encode(peer.NameByte))
	checkFatal(enc.Encode(peer.UID))
	checkFatal(enc.Encode(peer.version))

	connsBuf := new(bytes.Buffer)
	connsEnc := gob.NewEncoder(connsBuf)
	for _, conn := range peer.connections {
		// DANGER holding rlock on peer, going to take rlock on conn
		if !conn.Established() {
			continue
		}
		checkFatal(connsEnc.Encode(conn.Remote().NameByte))
		checkFatal(connsEnc.Encode(conn.RemoteTCPAddr()))
	}
	checkFatal(enc.Encode(connsBuf.Bytes()))
}

func decodePeerNoConns(dec *gob.Decoder) (nameByte []byte, uid uint64, version uint64, conns []byte, err error) {
	err = dec.Decode(&nameByte)
	if err != nil {
		return
	}
	err = dec.Decode(&uid)
	if err != nil {
		return
	}
	err = dec.Decode(&version)
	if err != nil {
		return
	}
	err = dec.Decode(&conns)
	if err == io.EOF {
		err = nil
	}
	return
}

func connsIterator(input []byte, fun func([]byte, string)) {
	buf := new(bytes.Buffer)
	buf.Write(input)
	dec := gob.NewDecoder(buf)
	for {
		var nameByte []byte
		err := dec.Decode(&nameByte)
		if err == io.EOF {
			return
		}
		checkFatal(err)
		var foundAt string
		checkFatal(dec.Decode(&foundAt))
		fun(nameByte, string(foundAt))
	}
}

func readConnsMap(peer *Peer, buf []byte, table map[PeerName]*Peer) map[PeerName]Connection {
	conns := make(map[PeerName]Connection)
	connsIterator(buf, func(nameByte []byte, remoteTCPAddr string) {
		name := PeerNameFromBin(nameByte)
		remotePeer := table[name]
		conn := NewRemoteConnection(peer, remotePeer, remoteTCPAddr)
		conns[name] = conn
	})
	return conns
}
<commit_msg>refactor: reduce dependency on peer.Router now the only code depending on it is Peer itself<commit_after>package router

import (
	"bytes"
	"encoding/gob"
	"fmt"
	"io"
	"sync"
)

type Peers struct {
	sync.RWMutex
	table map[PeerName]*Peer
	onGC  func(*Peer)
}

func NewPeers(onGC func(*Peer)) *Peers {
	return &Peers{table: make(map[PeerName]*Peer), onGC: onGC}
}

func (peers *Peers) FetchWithDefault(peer *Peer) *Peer {
	peers.RLock()
	res, found := peers.fetchAlias(peer)
	peers.RUnlock()
	if found {
		return res
	}
	peers.Lock()
	defer peers.Unlock()
	res, found = peers.fetchAlias(peer)
	if found {
		return res
	}
	peers.table[peer.Name] = peer
	peer.IncrementLocalRefCount()
	return peer
}

func (peers *Peers) Fetch(name PeerName) (*Peer, bool) {
	peers.RLock()
	defer peers.RUnlock()
	peer, found := peers.table[name]
	return peer, found // GRRR, why can't I inline this!?
}

func (peers *Peers) ForEach(fun func(PeerName, *Peer)) {
	peers.RLock()
	defer peers.RUnlock()
	for name, peer := range peers.table {
		fun(name, peer)
	}
}

func (peers *Peers) ApplyUpdate(update []byte, router *Router) ([]byte, error) {
	peers.Lock()

	newPeers, decodedUpdate, decodedConns, err := peers.decodeUpdate(update, router)
	if err != nil {
		peers.Unlock()
		return nil, err
	}

	// By this point, we know the update doesn't refer to any peers we
	// have no knowledge of. We can now apply the update. Start by
	// adding in any new peers into the cache.
	for name, newPeer := range newPeers {
		peers.table[name] = newPeer
	}

	// Now apply the updates
	newUpdate := peers.applyUpdate(decodedUpdate, decodedConns, router)

	for _, peerRemoved := range peers.garbageCollect(router.Ourself, router.Macs) {
		delete(newUpdate, peerRemoved.Name)
	}

	// Don't need to hold peers lock any longer
	peers.Unlock()

	if len(newUpdate) > 0 {
		router.ConnectionMaker.Refresh()
	}

	return encodePeersMap(newUpdate), nil
}

func (peers *Peers) EncodeAllPeers() []byte {
	peers.RLock()
	defer peers.RUnlock()
	return encodePeersMap(peers.table)
}

func EncodePeers(peers ...*Peer) []byte {
	buf := new(bytes.Buffer)
	enc := gob.NewEncoder(buf)
	for _, peer := range peers {
		peer.encodePeer(enc)
	}
	return buf.Bytes()
}

func (peers *Peers) GarbageCollect(router *Router) []*Peer {
	peers.Lock()
	defer peers.Unlock()
	return peers.garbageCollect(router.Ourself, router.Macs)
}

func (peers *Peers) String() string {
	var buf bytes.Buffer
	peers.ForEach(func(name PeerName, peer *Peer) {
		buf.WriteString(fmt.Sprint(peer, "\n"))
		peer.ForEachConnection(func(remoteName PeerName, conn Connection) {
			buf.WriteString(fmt.Sprintf("   -> %v [%v]\n", remoteName, conn.RemoteTCPAddr()))
		})
	})
	return buf.String()
}

func (peers *Peers) fetchAlias(peer *Peer) (*Peer, bool) {
	if existingPeer, found := peers.table[peer.Name]; found {
		if existingPeer.UID == peer.UID {
			existingPeer.IncrementLocalRefCount()
			return existingPeer, true
		} else {
			return nil, true
		}
	}
	return nil, false
}

func (peers *Peers) garbageCollect(ourself *Peer, macs *MacCache) []*Peer {
	removed := []*Peer{}
	for name, peer := range peers.table {
		found, _ := ourself.Routes(peer, false)
		if !found && !peer.IsLocallyReferenced() {
			peers.onGC(peer)
			delete(peers.table, name)
			macs.Delete(peer)
			removed = append(removed, peer)
		}
	}
	return removed
}

func encodePeersMap(peers map[PeerName]*Peer) []byte {
	buf := new(bytes.Buffer)
	enc := gob.NewEncoder(buf)
	for _, peer := range peers {
		peer.encodePeer(enc)
	}
	return buf.Bytes()
}

func (peers *Peers) decodeUpdate(update []byte, router *Router) (newPeers map[PeerName]*Peer, decodedUpdate []*Peer, decodedConns [][]byte, err error) {
	newPeers = make(map[PeerName]*Peer)
	decodedUpdate = []*Peer{}
	decodedConns = [][]byte{}

	updateBuf := new(bytes.Buffer)
	updateBuf.Write(update)
	decoder := gob.NewDecoder(updateBuf)

	for {
		nameByte, uid, version, connsBuf, decErr := decodePeerNoConns(decoder)
		if decErr == io.EOF {
			break
		} else if decErr != nil {
			err = decErr
			return
		}
		name := PeerNameFromBin(nameByte)
		newPeer := NewPeer(name, uid, version, router)
		decodedUpdate = append(decodedUpdate, newPeer)
		decodedConns = append(decodedConns, connsBuf)
		existingPeer, found := peers.table[name]
		if !found {
			newPeers[name] = newPeer
		} else if existingPeer.UID != newPeer.UID {
			err = NameCollisionError{Name: newPeer.Name}
			return
		}
	}

	unknownPeers := false
	for _, connsBuf := range decodedConns {
		connsIterator(connsBuf, func(remoteNameByte []byte, _ string) {
			remoteName := PeerNameFromBin(remoteNameByte)
			if _, found := newPeers[remoteName]; found {
				return
			}
			if _, found := peers.table[remoteName]; found {
				return
			}
			// Update refers to a peer which we have no knowledge
			// of. Thus we can't apply the update. Abort.
			unknownPeers = true
		})
	}
	if unknownPeers {
		err = UnknownPeersError{}
		return
	}
	return
}

func (peers *Peers) applyUpdate(decodedUpdate []*Peer, decodedConns [][]byte, router *Router) map[PeerName]*Peer {
	newUpdate := make(map[PeerName]*Peer)
	for idx, newPeer := range decodedUpdate {
		connsBuf := decodedConns[idx]
		name := newPeer.Name
		// guaranteed to find peer in the peers.table
		peer := peers.table[name]
		if peer != newPeer {
			if peer.Version() > newPeer.Version() {
				// we know more about this one than the update. If
				// peer is ourself, this is slightly racey (further
				// changes could occur to ourself in the mean
				// time). But it doesn't matter as we know that we're
				// already > newPeer.Version() and that's not going to
				// change.
				newUpdate[name] = peer
				continue
			} else if peer == router.Ourself {
				// nobody but us updates us
				continue
			} else if peer.Version() == newPeer.Version() {
				// implication is that connections are equal too
				continue
			}
		}
		// If we're here, either it was a new peer, or the update has
		// more info about the peer than we do. Either case, we need
		// to set version and conns and include the updated peer in
		// the outgoing update.

		// Can peer have been updated by anyone else in the mean time?
		// No - we know that peer is not ourself, so the only prospect
		// for an update would be someone else calling
		// router.Peers.ApplyUpdate. But ApplyUpdate takes the Lock on
		// the router.Peers, so there can be no race here.
		conns := readConnsMap(peer, connsBuf, peers.table)
		peer.SetVersionAndConnections(newPeer.Version(), conns)
		newUpdate[name] = peer
	}
	return newUpdate
}

func (peer *Peer) encodePeer(enc *gob.Encoder) {
	peer.RLock()
	defer peer.RUnlock()

	checkFatal(enc.Encode(peer.NameByte))
	checkFatal(enc.Encode(peer.UID))
	checkFatal(enc.Encode(peer.version))

	connsBuf := new(bytes.Buffer)
	connsEnc := gob.NewEncoder(connsBuf)
	for _, conn := range peer.connections {
		// DANGER holding rlock on peer, going to take rlock on conn
		if !conn.Established() {
			continue
		}
		checkFatal(connsEnc.Encode(conn.Remote().NameByte))
		checkFatal(connsEnc.Encode(conn.RemoteTCPAddr()))
	}
	checkFatal(enc.Encode(connsBuf.Bytes()))
}

func decodePeerNoConns(dec *gob.Decoder) (nameByte []byte, uid uint64, version uint64, conns []byte, err error) {
	err = dec.Decode(&nameByte)
	if err != nil {
		return
	}
	err = dec.Decode(&uid)
	if err != nil {
		return
	}
	err = dec.Decode(&version)
	if err != nil {
		return
	}
	err = dec.Decode(&conns)
	if err == io.EOF {
		err = nil
	}
	return
}

func connsIterator(input []byte, fun func([]byte, string)) {
	buf := new(bytes.Buffer)
	buf.Write(input)
	dec := gob.NewDecoder(buf)
	for {
		var nameByte []byte
		err := dec.Decode(&nameByte)
		if err == io.EOF {
			return
		}
		checkFatal(err)
		var foundAt string
		checkFatal(dec.Decode(&foundAt))
		fun(nameByte, string(foundAt))
	}
}

func readConnsMap(peer *Peer, buf []byte, table map[PeerName]*Peer) map[PeerName]Connection {
	conns := make(map[PeerName]Connection)
	connsIterator(buf, func(nameByte []byte, remoteTCPAddr string) {
		name := PeerNameFromBin(nameByte)
		remotePeer := table[name]
		conn := NewRemoteConnection(peer, remotePeer, remoteTCPAddr)
		conns[name] = conn
	})
	return conns
}
<|endoftext|>


<commit_before>package debug_test

import (
	"net/http"
	"os"
	"path"
	"path/filepath"

	. "github.com/onsi/ginkgo"
	"github.com/onsi/ginkgo/reporters"
	. "github.com/onsi/gomega"
	"github.com/pivotal-cf/cf-redis-broker/brokerconfig"
	"github.com/pivotal-cf/cf-redis-broker/debug"
	"github.com/pivotal-cf/cf-redis-broker/integration/helpers"
	"github.com/pivotal-cf/cf-redis-broker/redis"
	"github.com/pivotal-golang/lager"

	"testing"
)

func TestDebug(t *testing.T) {
	RegisterFailHandler(Fail)
	junitReporter := reporters.NewJUnitReporter("junit_debug.xml")
	RunSpecsWithDefaultAndCustomReporters(t, "Debug Suite", []Reporter{junitReporter})
}

var dirs []string

var _ = BeforeSuite(func() {
	helpers.ResetTestDirs()

	dirs = []string{"/tmp/to/redis", "/tmp/redis/data/directory", "/tmp/redis/log/directory"}
	for _, dir := range dirs {
		err := os.MkdirAll(dir, 0755)
		Ω(err).ShouldNot(HaveOccurred())
	}
	_, err := os.Create("/tmp/to/redis/config.conf")
	Ω(err).ShouldNot(HaveOccurred())
	path, err := filepath.Abs(path.Join("..", "brokerconfig", "assets", "test_config.yml"))
	Ω(err).ToNot(HaveOccurred())
	config, err := brokerconfig.ParseConfig(path)
	Ω(err).NotTo(HaveOccurred())

	logger := lager.NewLogger("foo")
	repo, err := redis.NewRemoteRepository(&redis.RemoteAgentClient{}, config, logger)
	Ω(err).NotTo(HaveOccurred())

	handler := debug.NewHandler(repo)

	http.HandleFunc("/debug", handler)
	go func() {
		defer GinkgoRecover()
		err := http.ListenAndServe("localhost:3000", nil)
		Expect(err).ToNot(HaveOccurred())
	}()
})

var _ = AfterSuite(func() {
	for _, dir := range dirs {
		err := os.RemoveAll(dir)
		Ω(err).ShouldNot(HaveOccurred())
	}
})
<commit_msg>Wait for debug handler before testing<commit_after>package debug_test

import (
	"net/http"
	"os"
	"path"
	"path/filepath"
	"time"

	. "github.com/onsi/ginkgo"
	"github.com/onsi/ginkgo/reporters"
	. "github.com/onsi/gomega"
	"github.com/pivotal-cf/cf-redis-broker/brokerconfig"
	"github.com/pivotal-cf/cf-redis-broker/debug"
	"github.com/pivotal-cf/cf-redis-broker/integration/helpers"
	"github.com/pivotal-cf/cf-redis-broker/redis"
	"github.com/pivotal-golang/lager"

	"testing"
)

func TestDebug(t *testing.T) {
	RegisterFailHandler(Fail)
	junitReporter := reporters.NewJUnitReporter("junit_debug.xml")
	RunSpecsWithDefaultAndCustomReporters(t, "Debug Suite", []Reporter{junitReporter})
}

var dirs []string

var _ = BeforeSuite(func() {
	helpers.ResetTestDirs()

	dirs = []string{"/tmp/to/redis", "/tmp/redis/data/directory", "/tmp/redis/log/directory"}
	for _, dir := range dirs {
		err := os.MkdirAll(dir, 0755)
		Ω(err).ShouldNot(HaveOccurred())
	}
	_, err := os.Create("/tmp/to/redis/config.conf")
	Ω(err).ShouldNot(HaveOccurred())
	path, err := filepath.Abs(path.Join("..", "brokerconfig", "assets", "test_config.yml"))
	Ω(err).ToNot(HaveOccurred())
	config, err := brokerconfig.ParseConfig(path)
	Ω(err).NotTo(HaveOccurred())

	logger := lager.NewLogger("foo")
	repo, err := redis.NewRemoteRepository(&redis.RemoteAgentClient{}, config, logger)
	Ω(err).NotTo(HaveOccurred())

	handler := debug.NewHandler(repo)

	http.HandleFunc("/debug", handler)
	go func() {
		defer GinkgoRecover()
		err := http.ListenAndServe("localhost:3000", nil)
		Expect(err).ToNot(HaveOccurred())
	}()

	client := http.Client{}
	for i := 0; i < 10; i++ {
		_, err = client.Get("http://localhost:3000/debug")
		if err == nil {
			break
		}
		time.Sleep(time.Second * 1)
		if i == 9 {
			Fail("Timed out waiting for debug handler setup")
		}
	}
})

var _ = AfterSuite(func() {
	for _, dir := range dirs {
		err := os.RemoveAll(dir)
		Ω(err).ShouldNot(HaveOccurred())
	}
})
<|endoftext|>


<commit_before>package sshclient

//Handle all ssh connection and run command

import (
	"bytes"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"
	"sync"

	"golang.org/x/crypto/ssh"
)

type SSHAuthTypeEnum int

const (
	//Type authentication login in ssh, it can using password or using public-private key
	SSHAuthType_Password SSHAuthTypeEnum = iota
	SSHAuthType_Certificate
)

type SshSetting struct {
	//Setting information for ssh connection
	SSHHost        string
	SSHUser        string
	SSHPassword    string
	SSHKeyLocation string
	SSHAuthType    SSHAuthTypeEnum
}

/*
Parsing private key certicate using for connection over ssh
*/
func PublicKeyFile(file string) ssh.AuthMethod {
	buffer, err := ioutil.ReadFile(file)
	if err != nil {
		return nil
	}

	key, err := ssh.ParsePrivateKey(buffer)
	if err != nil {
		return nil
	}
	return ssh.PublicKeys(key)
}

/*
Build connection ssh client to ssh server
*/
func (S *SshSetting) Connect() (*ssh.Client, error) {
	var (
		cfg *ssh.ClientConfig
	)

	if S.SSHAuthType == SSHAuthType_Certificate {
		cfg = &ssh.ClientConfig{
			User:            S.SSHUser,
			HostKeyCallback: ssh.InsecureIgnoreHostKey(),
			Auth: []ssh.AuthMethod{
				PublicKeyFile(S.SSHKeyLocation),
			},
		}
	} else {
		cfg = &ssh.ClientConfig{
			User:            S.SSHUser,
			HostKeyCallback: ssh.InsecureIgnoreHostKey(),
			Auth: []ssh.AuthMethod{
				ssh.Password(S.SSHPassword),
			},
		}
	}

	client, e := ssh.Dial("tcp", S.SSHHost, cfg)
	return client, e

}

/*
Handle input and output into terminal using channel
*/
func TermInOut(w io.Writer, r io.Reader) (chan<- string, <-chan string) {
	in := make(chan string, 1)
	out := make(chan string, 1)
	var wg sync.WaitGroup
	wg.Add(1)
	go func() {
		for cmd := range in {
			wg.Add(1)
			w.Write([]byte(cmd + "\n"))
			wg.Wait()
		}
	}()
	go func() {
		var (
			buf [1024 * 1024]byte
			t   int
		)
		for {
			n, err := r.Read(buf[t:])
			if err != nil {
				close(in)
				close(out)
				return
			}
			t += n
			if buf[t-2] == '$' {
				out <- string(buf[:t])
				t = 0
				wg.Done()
			}
		}
	}()
	return in, out
}

/*
Create new session
*/
func (s *SshSetting) NewSession() (*ssh.Client, *ssh.Session, error) {
	c, e := s.Connect()
	if e != nil {
		e = fmt.Errorf("Unable to connect: %s", e.Error())
		return c, nil, e
	}

	Ses, e := c.NewSession()
	if e != nil {
		e = fmt.Errorf("Unable to start new session: %s", e.Error())
		return c, Ses, e
	}

	return c, Ses, e
}

/*
Build connection and run ssh script, catch the output or give error message if any
*/
func (S *SshSetting) RunCommandSsh(cmds ...string) (string, error) {
	var (
		res string
		err error
	)

	c, Ses, e := S.NewSession()
	if e != nil {
		err = fmt.Errorf("Unable to connect: %s", e.Error())
		return res, err
	}
	defer c.Close()
	defer Ses.Close()

	modes := ssh.TerminalModes{
		ssh.ECHO:          0,
		ssh.TTY_OP_ISPEED: 14400,
		ssh.TTY_OP_OSPEED: 14400,
	}

	if e = Ses.RequestPty("xterm", 80, 40, modes); e != nil {
		err = fmt.Errorf("Unable to start term: %s", e.Error())
		return res, err
	}

	w, _ := Ses.StdinPipe()
	r, _ := Ses.StdoutPipe()

	in, out := TermInOut(w, r)
	if e = Ses.Shell(); e != nil {
		err = fmt.Errorf("Unable to start shell: %s", e.Error())
		return res, err
	}
	<-out

	cmds = append(cmds, "exit")
	cmdtemp := ""

	for _, cmd := range cmds {
		in <- cmd
		outs := strings.Split(<-out, "\n")
		if len(outs) > 1 {
			outtemp := strings.Trim(strings.Join(outs[:len(outs)-1], "\n"), " ")
			res = res + "Output of " + cmdtemp + " : " + outtemp
		}
		cmdtemp = cmd
	}
	Ses.Wait()

	return res, err
}

type RunCommandResult struct {
	CMD    string
	Output string
}

/*
Build connection and run ssh script, catch the output or give error message if any
*/
func (s *SshSetting) RunCommandSshAsMap(cmds ...string) ([]RunCommandResult, error) {
	result := []RunCommandResult{}
	client, sess, err := s.NewSession()
	if err != nil {
		return result, fmt.Errorf("Unable to connect: %s", err.Error())
	}
	defer client.Close()
	defer sess.Close()

	modes := ssh.TerminalModes{
		ssh.ECHO:          0,
		ssh.TTY_OP_ISPEED: 14400,
		ssh.TTY_OP_OSPEED: 14400,
	}

	if err = sess.RequestPty("xterm", 80, 40, modes); err != nil {
		return result, fmt.Errorf("Unable to start term: %s", err.Error())
	}

	writer, _ := sess.StdinPipe()
	reader, _ := sess.StdoutPipe()

	in, out := TermInOut(writer, reader)
	if err = sess.Shell(); err != nil {
		return result, fmt.Errorf("Unable to start shell: %s", err.Error())
	}
	<-out

	cmds = append(cmds, "exit")

	for _, cmd := range cmds {
		in <- cmd
		if cmd != "exit" {
			res := <-out
			res = strings.Split(res, `]0;`)[0]
			res = strings.Split(res, string(0x1b))[0]
			res = strings.TrimSpace(res)
			result = append(result, RunCommandResult{cmd, res})
		}
	}
	sess.Wait()

	return result, nil
}

/*
Run single command, get the output
*/
func (s *SshSetting) GetOutputCommandSsh(cmd string) (string, error) {
	c, Ses, e := s.NewSession()
	if e != nil {
		e = fmt.Errorf("Unable to connect: %s", e.Error())
		return "", e
	}
	defer c.Close()
	defer Ses.Close()

	var out bytes.Buffer
	Ses.Stdout = &out
	var err bytes.Buffer
	Ses.Stderr = &err
	if e := Ses.Run(cmd); e != nil {
		return "", errors.New(fmt.Sprintf("%s. %s", e.Error(), err.String()))
	}

	return out.String(), nil
}

// Copy file adopted from https://github.com/tmc/scp/blob/master/scp.go
func (S *SshSetting) SshCopyByPath(filePath, destinationPath string) error {
	var (
		err error
	)

	f, err := os.Open(filePath)
	if err != nil {
		return err
	}

	defer f.Close()
	s, err := f.Stat()
	if err != nil {
		return err
	}

	err = S.SshCopyByFile(f, s.Size(), s.Mode().Perm(), filepath.Base(f.Name()), destinationPath)

	return nil
}

func (S *SshSetting) SshCopyByFile(content io.Reader, size int64, perm os.FileMode, filename string, destination string) error {
	var (
		err error
	)

	c, Ses, e := S.NewSession()
	if e != nil {
		err = fmt.Errorf("Unable to connect: %s", e.Error())
		return err
	}
	defer c.Close()
	defer Ses.Close()

	go func() {
		w, _ := Ses.StdinPipe()
		defer w.Close()
		fmt.Fprintf(w, "C%#o %d %s\n", perm, size, filename)
		io.Copy(w, content)
		fmt.Fprint(w, "\x00")
	}()

	cmd := fmt.Sprintf("scp -t %s", destination)

	if err = Ses.Run(cmd); err != nil {
		return err
	}

	return nil
}

func (S *SshSetting) SshGetFile(path string) (res bytes.Buffer, e error) {
	c, Ses, e := S.NewSession()
	if e != nil {
		e = fmt.Errorf("Unable to connect: %s", e.Error())
		return
	}
	defer c.Close()
	defer Ses.Close()

	cmd := fmt.Sprintf(CAT, path)

	Ses.Stdout = &res
	if e = Ses.Run(cmd); e != nil {
		return
	}

	return
}
<commit_msg>No password<commit_after>package sshclient

//Handle all ssh connection and run command

import (
	"bytes"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"
	"sync"

	"golang.org/x/crypto/ssh"
)

type SSHAuthTypeEnum int

const (
	//Type authentication login in ssh, it can using password or using public-private key
	SSHAuthType_Password SSHAuthTypeEnum = iota
	SSHAuthType_Certificate
	SSHAuthType_NoPassword
)

type SshSetting struct {
	//Setting information for ssh connection
	SSHHost        string
	SSHUser        string
	SSHPassword    string
	SSHKeyLocation string
	SSHAuthType    SSHAuthTypeEnum
}

/*
Parsing private key certicate using for connection over ssh
*/
func PublicKeyFile(file string) ssh.AuthMethod {
	buffer, err := ioutil.ReadFile(file)
	if err != nil {
		return nil
	}

	key, err := ssh.ParsePrivateKey(buffer)
	if err != nil {
		return nil
	}
	return ssh.PublicKeys(key)
}

/*
Build connection ssh client to ssh server
*/
func (S *SshSetting) Connect() (*ssh.Client, error) {
	var (
		cfg *ssh.ClientConfig
	)

	if S.SSHAuthType == SSHAuthType_Certificate {
		cfg = &ssh.ClientConfig{
			User:            S.SSHUser,
			HostKeyCallback: ssh.InsecureIgnoreHostKey(),
			Auth: []ssh.AuthMethod{
				PublicKeyFile(S.SSHKeyLocation),
			},
		}
	} else if S.SSHAuthType == SSHAuthType_NoPassword {
		cfg = &ssh.ClientConfig{
			User:            S.SSHUser,
			HostKeyCallback: ssh.InsecureIgnoreHostKey(),
		}
	} else {
		cfg = &ssh.ClientConfig{
			User:            S.SSHUser,
			HostKeyCallback: ssh.InsecureIgnoreHostKey(),
			Auth: []ssh.AuthMethod{
				ssh.Password(S.SSHPassword),
			},
		}
	}

	client, e := ssh.Dial("tcp", S.SSHHost, cfg)
	return client, e

}

/*
Handle input and output into terminal using channel
*/
func TermInOut(w io.Writer, r io.Reader) (chan<- string, <-chan string) {
	in := make(chan string, 1)
	out := make(chan string, 1)
	var wg sync.WaitGroup
	wg.Add(1)
	go func() {
		for cmd := range in {
			wg.Add(1)
			w.Write([]byte(cmd + "\n"))
			wg.Wait()
		}
	}()
	go func() {
		var (
			buf [1024 * 1024]byte
			t   int
		)
		for {
			n, err := r.Read(buf[t:])
			if err != nil {
				close(in)
				close(out)
				return
			}
			t += n
			if buf[t-2] == '$' {
				out <- string(buf[:t])
				t = 0
				wg.Done()
			}
		}
	}()
	return in, out
}

/*
Create new session
*/
func (s *SshSetting) NewSession() (*ssh.Client, *ssh.Session, error) {
	c, e := s.Connect()
	if e != nil {
		e = fmt.Errorf("Unable to connect: %s", e.Error())
		return c, nil, e
	}

	Ses, e := c.NewSession()
	if e != nil {
		e = fmt.Errorf("Unable to start new session: %s", e.Error())
		return c, Ses, e
	}

	return c, Ses, e
}

/*
Build connection and run ssh script, catch the output or give error message if any
*/
func (S *SshSetting) RunCommandSsh(cmds ...string) (string, error) {
	var (
		res string
		err error
	)

	c, Ses, e := S.NewSession()
	if e != nil {
		err = fmt.Errorf("Unable to connect: %s", e.Error())
		return res, err
	}
	defer c.Close()
	defer Ses.Close()

	modes := ssh.TerminalModes{
		ssh.ECHO:          0,
		ssh.TTY_OP_ISPEED: 14400,
		ssh.TTY_OP_OSPEED: 14400,
	}

	if e = Ses.RequestPty("xterm", 80, 40, modes); e != nil {
		err = fmt.Errorf("Unable to start term: %s", e.Error())
		return res, err
	}

	w, _ := Ses.StdinPipe()
	r, _ := Ses.StdoutPipe()

	in, out := TermInOut(w, r)
	if e = Ses.Shell(); e != nil {
		err = fmt.Errorf("Unable to start shell: %s", e.Error())
		return res, err
	}
	<-out

	cmds = append(cmds, "exit")
	cmdtemp := ""

	for _, cmd := range cmds {
		in <- cmd
		outs := strings.Split(<-out, "\n")
		if len(outs) > 1 {
			outtemp := strings.Trim(strings.Join(outs[:len(outs)-1], "\n"), " ")
			res = res + "Output of " + cmdtemp + " : " + outtemp
		}
		cmdtemp = cmd
	}
	Ses.Wait()

	return res, err
}

type RunCommandResult struct {
	CMD    string
	Output string
}

/*
Build connection and run ssh script, catch the output or give error message if any
*/
func (s *SshSetting) RunCommandSshAsMap(cmds ...string) ([]RunCommandResult, error) {
	result := []RunCommandResult{}
	client, sess, err := s.NewSession()
	if err != nil {
		return result, fmt.Errorf("Unable to connect: %s", err.Error())
	}
	defer client.Close()
	defer sess.Close()

	modes := ssh.TerminalModes{
		ssh.ECHO:          0,
		ssh.TTY_OP_ISPEED: 14400,
		ssh.TTY_OP_OSPEED: 14400,
	}

	if err = sess.RequestPty("xterm", 80, 40, modes); err != nil {
		return result, fmt.Errorf("Unable to start term: %s", err.Error())
	}

	writer, _ := sess.StdinPipe()
	reader, _ := sess.StdoutPipe()

	in, out := TermInOut(writer, reader)
	if err = sess.Shell(); err != nil {
		return result, fmt.Errorf("Unable to start shell: %s", err.Error())
	}
	<-out

	cmds = append(cmds, "exit")

	for _, cmd := range cmds {
		in <- cmd
		if cmd != "exit" {
			res := <-out
			res = strings.Split(res, `]0;`)[0]
			res = strings.Split(res, string(0x1b))[0]
			res = strings.TrimSpace(res)
			result = append(result, RunCommandResult{cmd, res})
		}
	}
	sess.Wait()

	return result, nil
}

/*
Run single command, get the output
*/
func (s *SshSetting) GetOutputCommandSsh(cmd string) (string, error) {
	c, Ses, e := s.NewSession()
	if e != nil {
		e = fmt.Errorf("Unable to connect: %s", e.Error())
		return "", e
	}
	defer c.Close()
	defer Ses.Close()

	var out bytes.Buffer
	Ses.Stdout = &out
	var err bytes.Buffer
	Ses.Stderr = &err
	if e := Ses.Run(cmd); e != nil {
		return "", errors.New(fmt.Sprintf("%s. %s", e.Error(), err.String()))
	}

	return out.String(), nil
}

// Copy file adopted from https://github.com/tmc/scp/blob/master/scp.go
func (S *SshSetting) SshCopyByPath(filePath, destinationPath string) error {
	var (
		err error
	)

	f, err := os.Open(filePath)
	if err != nil {
		return err
	}

	defer f.Close()
	s, err := f.Stat()
	if err != nil {
		return err
	}

	err = S.SshCopyByFile(f, s.Size(), s.Mode().Perm(), filepath.Base(f.Name()), destinationPath)

	return nil
}

func (S *SshSetting) SshCopyByFile(content io.Reader, size int64, perm os.FileMode, filename string, destination string) error {
	var (
		err error
	)

	c, Ses, e := S.NewSession()
	if e != nil {
		err = fmt.Errorf("Unable to connect: %s", e.Error())
		return err
	}
	defer c.Close()
	defer Ses.Close()

	go func() {
		w, _ := Ses.StdinPipe()
		defer w.Close()
		fmt.Fprintf(w, "C%#o %d %s\n", perm, size, filename)
		io.Copy(w, content)
		fmt.Fprint(w, "\x00")
	}()

	cmd := fmt.Sprintf("scp -t %s", destination)

	if err = Ses.Run(cmd); err != nil {
		return err
	}

	return nil
}

func (S *SshSetting) SshGetFile(path string) (res bytes.Buffer, e error) {
	c, Ses, e := S.NewSession()
	if e != nil {
		e = fmt.Errorf("Unable to connect: %s", e.Error())
		return
	}
	defer c.Close()
	defer Ses.Close()

	cmd := fmt.Sprintf(CAT, path)

	Ses.Stdout = &res
	if e = Ses.Run(cmd); e != nil {
		return
	}

	return
}
<|endoftext|>


<commit_before>package pomsg

import (
	"bytes"
	"fmt"

	"github.com/robfig/soy/ast"
)

// Validate checks if the given message is representable in a PO file.
// A MsgNode must be validated before trying to caculate its msgid or msgid_plural
//
// Rules:
//  - If a message contains a plural, it must be the sole child.
//  - A plural contains exactly {case 1} and {default} cases.
func Validate(n *ast.MsgNode) error {
	for i, child := range n.Body.Children() {
		if n, ok := child.(*ast.MsgPluralNode); ok {
			if i != 0 {
				return fmt.Errorf("plural node must be the sole child")
			}
			if len(n.Cases) != 1 || n.Cases[0].Value != 1 {
				return fmt.Errorf("PO requires two plural cases [1, default]. found %v", n.Cases)
			}
		}
	}
	return nil
}

// MsgId returns the msgid for the given msg node.
func Msgid(n *ast.MsgNode) string {
	return msgidn(n, true)
}

// MsgPlural returns the msgid_plural for the given message.
func MsgidPlural(n *ast.MsgNode) string {
	return msgidn(n, false)
}

func msgidn(n *ast.MsgNode, singular bool) string {
	var body = n.Body
	var children = body.Children()
	if len(children) == 0 {
		return ""
	}
	if pluralNode, ok := children[0].(*ast.MsgPluralNode); ok {
		body = pluralCase(pluralNode, singular)
	} else if !singular {
		return ""
	}
	var buf bytes.Buffer
	for _, child := range body.Children() {
		writeph(&buf, child)
	}
	return buf.String()
}

// pluralCase returns the singular or plural message body.
func pluralCase(n *ast.MsgPluralNode, singular bool) ast.ParentNode {
	if singular {
		return n.Cases[0].Body
	}
	return n.Default
}

// writeph writes the placeholder string for the given node to the given buffer.
func writeph(buf *bytes.Buffer, child ast.Node) {
	switch child := child.(type) {
	case *ast.RawTextNode:
		buf.Write(child.Text)
	case *ast.MsgPlaceholderNode:
		buf.WriteString("{" + child.Name + "}")
	}
}
<commit_msg>Fix function comments based on best practices from Effective Go<commit_after>package pomsg

import (
	"bytes"
	"fmt"

	"github.com/robfig/soy/ast"
)

// Validate checks if the given message is representable in a PO file.
// A MsgNode must be validated before trying to caculate its msgid or msgid_plural
//
// Rules:
//  - If a message contains a plural, it must be the sole child.
//  - A plural contains exactly {case 1} and {default} cases.
func Validate(n *ast.MsgNode) error {
	for i, child := range n.Body.Children() {
		if n, ok := child.(*ast.MsgPluralNode); ok {
			if i != 0 {
				return fmt.Errorf("plural node must be the sole child")
			}
			if len(n.Cases) != 1 || n.Cases[0].Value != 1 {
				return fmt.Errorf("PO requires two plural cases [1, default]. found %v", n.Cases)
			}
		}
	}
	return nil
}

// MsgId returns the msgid for the given msg node.
func Msgid(n *ast.MsgNode) string {
	return msgidn(n, true)
}

// MsgidPlural returns the msgid_plural for the given message.
func MsgidPlural(n *ast.MsgNode) string {
	return msgidn(n, false)
}

func msgidn(n *ast.MsgNode, singular bool) string {
	var body = n.Body
	var children = body.Children()
	if len(children) == 0 {
		return ""
	}
	if pluralNode, ok := children[0].(*ast.MsgPluralNode); ok {
		body = pluralCase(pluralNode, singular)
	} else if !singular {
		return ""
	}
	var buf bytes.Buffer
	for _, child := range body.Children() {
		writeph(&buf, child)
	}
	return buf.String()
}

// pluralCase returns the singular or plural message body.
func pluralCase(n *ast.MsgPluralNode, singular bool) ast.ParentNode {
	if singular {
		return n.Cases[0].Body
	}
	return n.Default
}

// writeph writes the placeholder string for the given node to the given buffer.
func writeph(buf *bytes.Buffer, child ast.Node) {
	switch child := child.(type) {
	case *ast.RawTextNode:
		buf.Write(child.Text)
	case *ast.MsgPlaceholderNode:
		buf.WriteString("{" + child.Name + "}")
	}
}
<|endoftext|>


<commit_before>package btrfs

import "sort"

/*
#cgo LDFLAGS: -lbtrfs

#include <stddef.h>
#include <btrfs/ioctl.h>
#include "btrfs.h"

// Required because Go has struct casting rules for negative numbers
const __u64 u64_BTRFS_LAST_FREE_OBJECTID = (__u64)BTRFS_LAST_FREE_OBJECTID;
const __u64 negative_one = (__u64)-1;

static char* get_name_btrfs_ioctl_vol_args_v2(struct btrfs_ioctl_vol_args_v2* btrfs_struct) {
	return btrfs_struct->name;
}
*/
import "C"

import (
	"fmt"
	"os"
	"path/filepath"
	"syscall"
	"unsafe"

	"github.com/pkg/errors"
)

// IsSubvolume returns nil if the path is a valid subvolume. An error is
// returned if the path does not exist or the path is not a valid subvolume.
func IsSubvolume(path string) error {
	fi, err := os.Lstat(path)
	if err != nil {
		return err
	}

	if err := isFileInfoSubvol(fi); err != nil {
		return err
	}

	var statfs syscall.Statfs_t
	if err := syscall.Statfs(path, &statfs); err != nil {
		return err
	}

	return isStatfsSubvol(&statfs)
}

// SubvolInfo returns information about the subvolume at the provided path.
func SubvolInfo(path string) (info Info, err error) {
	path, err = filepath.EvalSymlinks(path)
	if err != nil {
		return info, err
	}

	fp, err := openSubvolDir(path)
	if err != nil {
		return info, err
	}
	defer fp.Close()

	id, err := subvolID(fp.Fd())
	if err != nil {
		return info, err
	}

	subvolsByID, err := subvolMap(path)
	if err != nil {
		return info, err
	}

	if info, ok := subvolsByID[id]; ok {
		return info, nil
	}

	return info, errors.Errorf("%q not found", path)
}

func subvolMap(path string) (map[uint64]Info, error) {
	fp, err := openSubvolDir(path)
	if err != nil {
		return nil, err
	}
	defer fp.Close()

	var args C.struct_btrfs_ioctl_search_args

	args.key.tree_id = C.BTRFS_ROOT_TREE_OBJECTID
	args.key.min_type = C.BTRFS_ROOT_ITEM_KEY
	args.key.max_type = C.BTRFS_ROOT_BACKREF_KEY
	args.key.min_objectid = C.BTRFS_FS_TREE_OBJECTID
	args.key.max_objectid = C.u64_BTRFS_LAST_FREE_OBJECTID
	args.key.max_offset = C.negative_one
	args.key.max_transid = C.negative_one

	subvolsByID := map[uint64]Info{}

	for {
		args.key.nr_items = 4096
		if err := ioctl(fp.Fd(), C.BTRFS_IOC_TREE_SEARCH, uintptr(unsafe.Pointer(&args))); err != nil {
			return nil, err
		}

		if args.key.nr_items == 0 {
			break
		}

		var (
			sh     C.struct_btrfs_ioctl_search_header
			shSize = unsafe.Sizeof(sh)
			buf    = (*[1<<31 - 1]byte)(unsafe.Pointer(&args.buf[0]))[:C.BTRFS_SEARCH_ARGS_BUFSIZE]
		)

		for i := 0; i < int(args.key.nr_items); i++ {
			sh = (*(*C.struct_btrfs_ioctl_search_header)(unsafe.Pointer(&buf[0])))
			buf = buf[shSize:]

			info := subvolsByID[uint64(sh.objectid)]
			info.ID = uint64(sh.objectid)

			if sh._type == C.BTRFS_ROOT_BACKREF_KEY {
				rr := (*(*C.struct_btrfs_root_ref)(unsafe.Pointer(&buf[0])))

				// This branch processes the backrefs from the root object. We
				// get an entry of the objectid, with name, but the parent is
				// the offset.

				nname := C.btrfs_stack_root_ref_name_len(&rr)
				name := string(buf[C.sizeof_struct_btrfs_root_ref : C.sizeof_struct_btrfs_root_ref+uintptr(nname)])

				info.ID = uint64(sh.objectid)
				info.ParentID = uint64(sh.offset)
				info.Name = name
				info.DirID = uint64(C.btrfs_stack_root_ref_dirid(&rr))

				subvolsByID[uint64(sh.objectid)] = info
			} else if sh._type == C.BTRFS_ROOT_ITEM_KEY &&
				(sh.objectid >= C.BTRFS_ROOT_ITEM_KEY ||
					sh.objectid == C.BTRFS_FS_TREE_OBJECTID) {

				var (
					ri  = (*C.struct_btrfs_root_item)(unsafe.Pointer(&buf[0]))
					gri C.struct_gosafe_btrfs_root_item
				)

				C.unpack_root_item(&gri, ri)

				if gri.flags&C.BTRFS_ROOT_SUBVOL_RDONLY != 0 {
					info.Readonly = true
				}

				// in this case, the offset is the actual offset.
				info.Offset = uint64(sh.offset)

				info.UUID = uuidString(&gri.uuid)
				info.ParentUUID = uuidString(&gri.parent_uuid)
				info.ReceivedUUID = uuidString(&gri.received_uuid)

				info.Generation = uint64(gri.gen)
				info.OriginalGeneration = uint64(gri.ogen)

				subvolsByID[uint64(sh.objectid)] = info
			}

			args.key.min_objectid = sh.objectid
			args.key.min_offset = sh.offset
			args.key.min_type = sh._type //  this is very questionable.

			buf = buf[sh.len:]
		}

		args.key.min_offset++
		if args.key.min_offset == 0 {
			args.key.min_type++
		} else {
			continue
		}

		if args.key.min_type > C.BTRFS_ROOT_BACKREF_KEY {
			args.key.min_type = C.BTRFS_ROOT_ITEM_KEY
			args.key.min_objectid++
		} else {
			continue
		}

		if args.key.min_objectid > args.key.max_objectid {
			break
		}
	}

	mnt, err := findMountPoint(path)
	if err != nil {
		return nil, err
	}

	for _, sv := range subvolsByID {
		path := sv.Name
		parentID := sv.ParentID

		for parentID != 0 {
			parent, ok := subvolsByID[parentID]
			if !ok {
				break
			}

			parentID = parent.ParentID
			path = filepath.Join(parent.Name, path)
		}

		sv.Path = filepath.Join(mnt, path)
	}
	return subvolsByID, nil
}

// SubvolList will return the information for all subvolumes corresponding to
// the provided path.
func SubvolList(path string) ([]Info, error) {
	subvolsByID, err := subvolMap(path)
	if err != nil {
		return nil, err
	}

	subvols := make([]Info, 0, len(subvolsByID))
	for _, sv := range subvolsByID {
		subvols = append(subvols, sv)
	}

	sort.Sort(infosByID(subvols))

	return subvols, nil
}

// SubvolCreate creates a subvolume at the provided path.
func SubvolCreate(path string) error {
	dir, name := filepath.Split(path)

	fp, err := os.Open(dir)
	if err != nil {
		return err
	}
	defer fp.Close()

	var args C.struct_btrfs_ioctl_vol_args
	args.fd = C.__s64(fp.Fd())

	if len(name) > C.BTRFS_PATH_NAME_MAX {
		return errors.Errorf("%q too long for subvolume", name)
	}
	nameptr := (*[1<<31 - 1]byte)(unsafe.Pointer(&args.name[0]))
	copy(nameptr[:C.BTRFS_PATH_NAME_MAX], []byte(name))

	if err := ioctl(fp.Fd(), C.BTRFS_IOC_SUBVOL_CREATE, uintptr(unsafe.Pointer(&args))); err != nil {
		return errors.Wrap(err, "btrfs subvolume create failed")
	}

	return nil
}

// SubvolSnapshot creates a snapshot in dst from src. If readonly is true, the
// snapshot will be readonly.
func SubvolSnapshot(dst, src string, readonly bool) error {
	dstdir, dstname := filepath.Split(dst)

	dstfp, err := openSubvolDir(dstdir)
	if err != nil {
		return errors.Wrapf(err, "opening snapshot desination subvolume failed")
	}
	defer dstfp.Close()

	srcfp, err := openSubvolDir(src)
	if err != nil {
		return errors.Wrapf(err, "opening snapshot source subvolume failed")
	}
	defer srcfp.Close()

	// dstdir is the ioctl arg, wile srcdir gets set on the args
	var args C.struct_btrfs_ioctl_vol_args_v2
	args.fd = C.__s64(srcfp.Fd())
	name := C.get_name_btrfs_ioctl_vol_args_v2(&args)

	if len(dstname) > C.BTRFS_SUBVOL_NAME_MAX {
		return errors.Errorf("%q too long for subvolume", dstname)
	}

	nameptr := (*[1<<31 - 1]byte)(unsafe.Pointer(name))
	copy(nameptr[:C.BTRFS_SUBVOL_NAME_MAX], []byte(dstname))

	if readonly {
		args.flags |= C.BTRFS_SUBVOL_RDONLY
	}

	if err := ioctl(dstfp.Fd(), C.BTRFS_IOC_SNAP_CREATE_V2, uintptr(unsafe.Pointer(&args))); err != nil {
		return errors.Wrapf(err, "snapshot create failed")
	}

	return nil
}

// SubvolDelete deletes the subvolumes under the given path.
func SubvolDelete(path string) error {
	fmt.Println("delete", path)
	dir, name := filepath.Split(path)
	fp, err := openSubvolDir(dir)
	if err != nil {
		return errors.Wrapf(err, "failed opening %v", path)
	}
	defer fp.Close()

	// remove child subvolumes
	if err := filepath.Walk(path, func(p string, fi os.FileInfo, err error) error {
		if err != nil {
			if os.IsNotExist(err) || p == path {
				return nil
			}

			return errors.Wrapf(err, "failed walking subvolume %v", p)
		}

		if !fi.IsDir() {
			return nil // just ignore it!
		}

		if p == path {
			return nil
		}

		if err := isFileInfoSubvol(fi); err != nil {
			return nil
		}

		if err := SubvolDelete(p); err != nil {
			return errors.Wrapf(err, "recursive delete of %v failed", p)
		}

		return filepath.SkipDir // children get walked by call above.
	}); err != nil {
		return err
	}

	var args C.struct_btrfs_ioctl_vol_args
	if len(name) > C.BTRFS_SUBVOL_NAME_MAX {
		return errors.Errorf("%q too long for subvolume", name)
	}

	nameptr := (*[1<<31 - 1]byte)(unsafe.Pointer(&args.name[0]))
	copy(nameptr[:C.BTRFS_SUBVOL_NAME_MAX], []byte(name))

	if err := ioctl(fp.Fd(), C.BTRFS_IOC_SNAP_DESTROY, uintptr(unsafe.Pointer(&args))); err != nil {
		return errors.Wrapf(err, "failed removing subvolume %v", path)
	}

	return nil
}

func openSubvolDir(path string) (*os.File, error) {
	fp, err := os.Open(path)
	if err != nil {
		return nil, errors.Wrapf(err, "opening %v as subvolume failed", path)
	}

	return fp, nil
}

func isStatfsSubvol(statfs *syscall.Statfs_t) error {
	if int64(statfs.Type) != int64(C.BTRFS_SUPER_MAGIC) {
		return errors.Errorf("not a btrfs filesystem")
	}

	return nil
}

func isFileInfoSubvol(fi os.FileInfo) error {
	if !fi.IsDir() {
		errors.Errorf("must be a directory")
	}

	stat := fi.Sys().(*syscall.Stat_t)

	if stat.Ino != C.BTRFS_FIRST_FREE_OBJECTID {
		return errors.Errorf("incorrect inode type")
	}

	return nil
}
<commit_msg>Add function to get subvolume ID without all info<commit_after>package btrfs

import "sort"

/*
#cgo LDFLAGS: -lbtrfs

#include <stddef.h>
#include <btrfs/ioctl.h>
#include "btrfs.h"

// Required because Go has struct casting rules for negative numbers
const __u64 u64_BTRFS_LAST_FREE_OBJECTID = (__u64)BTRFS_LAST_FREE_OBJECTID;
const __u64 negative_one = (__u64)-1;

static char* get_name_btrfs_ioctl_vol_args_v2(struct btrfs_ioctl_vol_args_v2* btrfs_struct) {
	return btrfs_struct->name;
}
*/
import "C"

import (
	"fmt"
	"os"
	"path/filepath"
	"syscall"
	"unsafe"

	"github.com/pkg/errors"
)

// IsSubvolume returns nil if the path is a valid subvolume. An error is
// returned if the path does not exist or the path is not a valid subvolume.
func IsSubvolume(path string) error {
	fi, err := os.Lstat(path)
	if err != nil {
		return err
	}

	if err := isFileInfoSubvol(fi); err != nil {
		return err
	}

	var statfs syscall.Statfs_t
	if err := syscall.Statfs(path, &statfs); err != nil {
		return err
	}

	return isStatfsSubvol(&statfs)
}

// SubvolID returns the subvolume ID for the provided path
func SubvolID(path string) (uint64, error) {
	fp, err := openSubvolDir(path)
	if err != nil {
		return 0, err
	}
	defer fp.Close()

	return subvolID(fp.Fd())
}

// SubvolInfo returns information about the subvolume at the provided path.
func SubvolInfo(path string) (info Info, err error) {
	path, err = filepath.EvalSymlinks(path)
	if err != nil {
		return info, err
	}

	fp, err := openSubvolDir(path)
	if err != nil {
		return info, err
	}
	defer fp.Close()

	id, err := subvolID(fp.Fd())
	if err != nil {
		return info, err
	}

	subvolsByID, err := subvolMap(path)
	if err != nil {
		return info, err
	}

	if info, ok := subvolsByID[id]; ok {
		return info, nil
	}

	return info, errors.Errorf("%q not found", path)
}

func subvolMap(path string) (map[uint64]Info, error) {
	fp, err := openSubvolDir(path)
	if err != nil {
		return nil, err
	}
	defer fp.Close()

	var args C.struct_btrfs_ioctl_search_args

	args.key.tree_id = C.BTRFS_ROOT_TREE_OBJECTID
	args.key.min_type = C.BTRFS_ROOT_ITEM_KEY
	args.key.max_type = C.BTRFS_ROOT_BACKREF_KEY
	args.key.min_objectid = C.BTRFS_FS_TREE_OBJECTID
	args.key.max_objectid = C.u64_BTRFS_LAST_FREE_OBJECTID
	args.key.max_offset = C.negative_one
	args.key.max_transid = C.negative_one

	subvolsByID := map[uint64]Info{}

	for {
		args.key.nr_items = 4096
		if err := ioctl(fp.Fd(), C.BTRFS_IOC_TREE_SEARCH, uintptr(unsafe.Pointer(&args))); err != nil {
			return nil, err
		}

		if args.key.nr_items == 0 {
			break
		}

		var (
			sh     C.struct_btrfs_ioctl_search_header
			shSize = unsafe.Sizeof(sh)
			buf    = (*[1<<31 - 1]byte)(unsafe.Pointer(&args.buf[0]))[:C.BTRFS_SEARCH_ARGS_BUFSIZE]
		)

		for i := 0; i < int(args.key.nr_items); i++ {
			sh = (*(*C.struct_btrfs_ioctl_search_header)(unsafe.Pointer(&buf[0])))
			buf = buf[shSize:]

			info := subvolsByID[uint64(sh.objectid)]
			info.ID = uint64(sh.objectid)

			if sh._type == C.BTRFS_ROOT_BACKREF_KEY {
				rr := (*(*C.struct_btrfs_root_ref)(unsafe.Pointer(&buf[0])))

				// This branch processes the backrefs from the root object. We
				// get an entry of the objectid, with name, but the parent is
				// the offset.

				nname := C.btrfs_stack_root_ref_name_len(&rr)
				name := string(buf[C.sizeof_struct_btrfs_root_ref : C.sizeof_struct_btrfs_root_ref+uintptr(nname)])

				info.ID = uint64(sh.objectid)
				info.ParentID = uint64(sh.offset)
				info.Name = name
				info.DirID = uint64(C.btrfs_stack_root_ref_dirid(&rr))

				subvolsByID[uint64(sh.objectid)] = info
			} else if sh._type == C.BTRFS_ROOT_ITEM_KEY &&
				(sh.objectid >= C.BTRFS_ROOT_ITEM_KEY ||
					sh.objectid == C.BTRFS_FS_TREE_OBJECTID) {

				var (
					ri  = (*C.struct_btrfs_root_item)(unsafe.Pointer(&buf[0]))
					gri C.struct_gosafe_btrfs_root_item
				)

				C.unpack_root_item(&gri, ri)

				if gri.flags&C.BTRFS_ROOT_SUBVOL_RDONLY != 0 {
					info.Readonly = true
				}

				// in this case, the offset is the actual offset.
				info.Offset = uint64(sh.offset)

				info.UUID = uuidString(&gri.uuid)
				info.ParentUUID = uuidString(&gri.parent_uuid)
				info.ReceivedUUID = uuidString(&gri.received_uuid)

				info.Generation = uint64(gri.gen)
				info.OriginalGeneration = uint64(gri.ogen)

				subvolsByID[uint64(sh.objectid)] = info
			}

			args.key.min_objectid = sh.objectid
			args.key.min_offset = sh.offset
			args.key.min_type = sh._type //  this is very questionable.

			buf = buf[sh.len:]
		}

		args.key.min_offset++
		if args.key.min_offset == 0 {
			args.key.min_type++
		} else {
			continue
		}

		if args.key.min_type > C.BTRFS_ROOT_BACKREF_KEY {
			args.key.min_type = C.BTRFS_ROOT_ITEM_KEY
			args.key.min_objectid++
		} else {
			continue
		}

		if args.key.min_objectid > args.key.max_objectid {
			break
		}
	}

	mnt, err := findMountPoint(path)
	if err != nil {
		return nil, err
	}

	for _, sv := range subvolsByID {
		path := sv.Name
		parentID := sv.ParentID

		for parentID != 0 {
			parent, ok := subvolsByID[parentID]
			if !ok {
				break
			}

			parentID = parent.ParentID
			path = filepath.Join(parent.Name, path)
		}

		sv.Path = filepath.Join(mnt, path)
	}
	return subvolsByID, nil
}

// SubvolList will return the information for all subvolumes corresponding to
// the provided path.
func SubvolList(path string) ([]Info, error) {
	subvolsByID, err := subvolMap(path)
	if err != nil {
		return nil, err
	}

	subvols := make([]Info, 0, len(subvolsByID))
	for _, sv := range subvolsByID {
		subvols = append(subvols, sv)
	}

	sort.Sort(infosByID(subvols))

	return subvols, nil
}

// SubvolCreate creates a subvolume at the provided path.
func SubvolCreate(path string) error {
	dir, name := filepath.Split(path)

	fp, err := os.Open(dir)
	if err != nil {
		return err
	}
	defer fp.Close()

	var args C.struct_btrfs_ioctl_vol_args
	args.fd = C.__s64(fp.Fd())

	if len(name) > C.BTRFS_PATH_NAME_MAX {
		return errors.Errorf("%q too long for subvolume", name)
	}
	nameptr := (*[1<<31 - 1]byte)(unsafe.Pointer(&args.name[0]))
	copy(nameptr[:C.BTRFS_PATH_NAME_MAX], []byte(name))

	if err := ioctl(fp.Fd(), C.BTRFS_IOC_SUBVOL_CREATE, uintptr(unsafe.Pointer(&args))); err != nil {
		return errors.Wrap(err, "btrfs subvolume create failed")
	}

	return nil
}

// SubvolSnapshot creates a snapshot in dst from src. If readonly is true, the
// snapshot will be readonly.
func SubvolSnapshot(dst, src string, readonly bool) error {
	dstdir, dstname := filepath.Split(dst)

	dstfp, err := openSubvolDir(dstdir)
	if err != nil {
		return errors.Wrapf(err, "opening snapshot desination subvolume failed")
	}
	defer dstfp.Close()

	srcfp, err := openSubvolDir(src)
	if err != nil {
		return errors.Wrapf(err, "opening snapshot source subvolume failed")
	}
	defer srcfp.Close()

	// dstdir is the ioctl arg, wile srcdir gets set on the args
	var args C.struct_btrfs_ioctl_vol_args_v2
	args.fd = C.__s64(srcfp.Fd())
	name := C.get_name_btrfs_ioctl_vol_args_v2(&args)

	if len(dstname) > C.BTRFS_SUBVOL_NAME_MAX {
		return errors.Errorf("%q too long for subvolume", dstname)
	}

	nameptr := (*[1<<31 - 1]byte)(unsafe.Pointer(name))
	copy(nameptr[:C.BTRFS_SUBVOL_NAME_MAX], []byte(dstname))

	if readonly {
		args.flags |= C.BTRFS_SUBVOL_RDONLY
	}

	if err := ioctl(dstfp.Fd(), C.BTRFS_IOC_SNAP_CREATE_V2, uintptr(unsafe.Pointer(&args))); err != nil {
		return errors.Wrapf(err, "snapshot create failed")
	}

	return nil
}

// SubvolDelete deletes the subvolumes under the given path.
func SubvolDelete(path string) error {
	fmt.Println("delete", path)
	dir, name := filepath.Split(path)
	fp, err := openSubvolDir(dir)
	if err != nil {
		return errors.Wrapf(err, "failed opening %v", path)
	}
	defer fp.Close()

	// remove child subvolumes
	if err := filepath.Walk(path, func(p string, fi os.FileInfo, err error) error {
		if err != nil {
			if os.IsNotExist(err) || p == path {
				return nil
			}

			return errors.Wrapf(err, "failed walking subvolume %v", p)
		}

		if !fi.IsDir() {
			return nil // just ignore it!
		}

		if p == path {
			return nil
		}

		if err := isFileInfoSubvol(fi); err != nil {
			return nil
		}

		if err := SubvolDelete(p); err != nil {
			return errors.Wrapf(err, "recursive delete of %v failed", p)
		}

		return filepath.SkipDir // children get walked by call above.
	}); err != nil {
		return err
	}

	var args C.struct_btrfs_ioctl_vol_args
	if len(name) > C.BTRFS_SUBVOL_NAME_MAX {
		return errors.Errorf("%q too long for subvolume", name)
	}

	nameptr := (*[1<<31 - 1]byte)(unsafe.Pointer(&args.name[0]))
	copy(nameptr[:C.BTRFS_SUBVOL_NAME_MAX], []byte(name))

	if err := ioctl(fp.Fd(), C.BTRFS_IOC_SNAP_DESTROY, uintptr(unsafe.Pointer(&args))); err != nil {
		return errors.Wrapf(err, "failed removing subvolume %v", path)
	}

	return nil
}

func openSubvolDir(path string) (*os.File, error) {
	fp, err := os.Open(path)
	if err != nil {
		return nil, errors.Wrapf(err, "opening %v as subvolume failed", path)
	}

	return fp, nil
}

func isStatfsSubvol(statfs *syscall.Statfs_t) error {
	if int64(statfs.Type) != int64(C.BTRFS_SUPER_MAGIC) {
		return errors.Errorf("not a btrfs filesystem")
	}

	return nil
}

func isFileInfoSubvol(fi os.FileInfo) error {
	if !fi.IsDir() {
		errors.Errorf("must be a directory")
	}

	stat := fi.Sys().(*syscall.Stat_t)

	if stat.Ino != C.BTRFS_FIRST_FREE_OBJECTID {
		return errors.Errorf("incorrect inode type")
	}

	return nil
}
<|endoftext|>


<commit_before>package bufit

import (
	"container/heap"
	"io"
	"sync"
	"sync/atomic"
)

// Reader provides an io.Reader whose methods MUST be concurrent-safe
// with the methods of the Writer from which it was generated.
type Reader interface {

	// Len returns the unread # of bytes in this Reader
	Len() int

	// Discard drops the next n bytes from the Reader, as if it were Read()
	// it returns the # of bytes actually dropped. It may return io.EOF
	// if all remaining bytes have been discarded.
	Discard(int) (int, error)

	// Read bytes into the provided buffer.
	io.Reader
}

// Writer accepts bytes and generates Readers who consume those bytes.
// Generated Readers methods must be concurrent-safe with the Writers methods.
type Writer interface {

	// Len returns the # of bytes buffered for Readers
	Len() int

	// Discard drops the next n buffered bytes. It returns the actual number of
	// bytes dropped and may return io.EOF if all remaining bytes have been
	// discarded.
	Discard(int) (int, error)

	// NextReader returns a Reader which reads a "snapshot" of the current written bytes
	// (excluding discarded bytes). The Reader should work independently of the Writer
	// and therefore be concurrent safe with operations on the Writer.
	NextReader() Reader

	// Write writes the given bytes into the Writer's underlying buffer. Which will
	// be available for reading using NextReader() to grab a snapshot of the current
	// written bytes.
	io.Writer
}

// Buffer is used to provide multiple readers with access to a shared buffer.
// Readers may join/leave at any time, however a joining reader will only
// see whats currently in the buffer onwards. Data is evicted from the buffer
// once all active readers have read that section.
type Buffer struct {
	mu   sync.RWMutex
	cond *sync.Cond
	off  int
	rh   readerHeap
	buf  Writer
	life
}

type life struct {
	state int32
}

func (l *life) alive() bool { return atomic.LoadInt32(&l.state) == 0 }
func (l *life) kill()       { atomic.AddInt32(&l.state, 1) }

func (b *Buffer) fetch(r *reader) {
	b.mu.Lock()
	defer b.mu.Unlock()

	for r.off+r.chunk-b.off == b.buf.Len() && b.alive() && r.alive() {
		b.cond.Wait()
	}

	if !r.alive() {
		return
	}

	r.off += r.chunk
	diff := r.off - b.off
	r.data = b.buf.NextReader()
	r.data.Discard(diff)
	r.chunk = r.data.Len()

	heap.Fix(&b.rh, r.i)
	b.shift()
}

func (b *Buffer) drop(r *reader) {
	b.mu.Lock()
	defer b.mu.Unlock()
	heap.Remove(&b.rh, r.i)
	b.shift()
}

func (b *Buffer) shift() {
	if b.rh.Len() == 0 {
		return
	}

	if diff := b.rh.Peek().off - b.off; diff > 0 {
		b.buf.Discard(diff)
		b.off += diff
	}
}

// NextReader returns a new ReadCloser for this shared buffer.
// Read/Close are safe to call concurrently with the buffers Write/Close methods.
// Read calls will block if the Buffer is not Closed and contains no data.
func (b *Buffer) NextReader() io.ReadCloser {
	b.mu.Lock()
	defer b.mu.Unlock()
	r := &reader{
		buf:   b,
		chunk: b.buf.Len(),
		off:   b.off,
		data:  b.buf.NextReader(),
	}
	heap.Push(&b.rh, r)
	return r
}

// Write appends the given data to the buffer. All active readers will
// see this write.
func (b *Buffer) Write(p []byte) (int, error) {
	b.mu.Lock()
	defer b.cond.Broadcast()
	defer b.mu.Unlock()
	return b.buf.Write(p)
}

// Close marks the buffer as complete. Readers will return io.EOF instead of blocking
// when they reach the end of the buffer.
func (b *Buffer) Close() error {
	defer b.cond.Broadcast()
	b.kill()
	return nil
}

// NewBuffer creates and returns a new Buffer backed by the passed Writer
func NewBuffer(w Writer) *Buffer {
	buf := Buffer{
		buf: w,
	}
	buf.cond = sync.NewCond(&buf.mu)
	return &buf
}

// New creates and returns a new Buffer
func New() *Buffer {
	return NewBuffer(newWriter(nil))
}
<commit_msg>fixed doc<commit_after>package bufit

import (
	"container/heap"
	"io"
	"sync"
	"sync/atomic"
)

// Reader provides an io.Reader whose methods MUST be concurrent-safe
// with the Write method of the Writer from which it was generated.
type Reader interface {

	// Len returns the unread # of bytes in this Reader
	Len() int

	// Discard drops the next n bytes from the Reader, as if it were Read()
	// it returns the # of bytes actually dropped. It may return io.EOF
	// if all remaining bytes have been discarded.
	Discard(int) (int, error)

	// Read bytes into the provided buffer.
	io.Reader
}

// Writer accepts bytes and generates Readers who consume those bytes.
// Generated Readers methods must be concurrent-safe with the Write method.
type Writer interface {

	// Len returns the # of bytes buffered for Readers
	Len() int

	// Discard drops the next n buffered bytes. It returns the actual number of
	// bytes dropped and may return io.EOF if all remaining bytes have been
	// discarded.
	Discard(int) (int, error)

	// NextReader returns a Reader which reads a "snapshot" of the current written bytes
	// (excluding discarded bytes). The Reader should work independently of the Writer
	// and be concurrent safe with the Write method on the Writer.
	NextReader() Reader

	// Write writes the given bytes into the Writer's underlying buffer. Which will
	// be available for reading using NextReader() to grab a snapshot of the current
	// written bytes.
	io.Writer
}

// Buffer is used to provide multiple readers with access to a shared buffer.
// Readers may join/leave at any time, however a joining reader will only
// see whats currently in the buffer onwards. Data is evicted from the buffer
// once all active readers have read that section.
type Buffer struct {
	mu   sync.RWMutex
	cond *sync.Cond
	off  int
	rh   readerHeap
	buf  Writer
	life
}

type life struct {
	state int32
}

func (l *life) alive() bool { return atomic.LoadInt32(&l.state) == 0 }
func (l *life) kill()       { atomic.AddInt32(&l.state, 1) }

func (b *Buffer) fetch(r *reader) {
	b.mu.Lock()
	defer b.mu.Unlock()

	for r.off+r.chunk-b.off == b.buf.Len() && b.alive() && r.alive() {
		b.cond.Wait()
	}

	if !r.alive() {
		return
	}

	r.off += r.chunk
	diff := r.off - b.off
	r.data = b.buf.NextReader()
	r.data.Discard(diff)
	r.chunk = r.data.Len()

	heap.Fix(&b.rh, r.i)
	b.shift()
}

func (b *Buffer) drop(r *reader) {
	b.mu.Lock()
	defer b.mu.Unlock()
	heap.Remove(&b.rh, r.i)
	b.shift()
}

func (b *Buffer) shift() {
	if b.rh.Len() == 0 {
		return
	}

	if diff := b.rh.Peek().off - b.off; diff > 0 {
		b.buf.Discard(diff)
		b.off += diff
	}
}

// NextReader returns a new ReadCloser for this shared buffer.
// Read/Close are safe to call concurrently with the buffers Write/Close methods.
// Read calls will block if the Buffer is not Closed and contains no data.
func (b *Buffer) NextReader() io.ReadCloser {
	b.mu.Lock()
	defer b.mu.Unlock()
	r := &reader{
		buf:   b,
		chunk: b.buf.Len(),
		off:   b.off,
		data:  b.buf.NextReader(),
	}
	heap.Push(&b.rh, r)
	return r
}

// Write appends the given data to the buffer. All active readers will
// see this write.
func (b *Buffer) Write(p []byte) (int, error) {
	b.mu.Lock()
	defer b.cond.Broadcast()
	defer b.mu.Unlock()
	return b.buf.Write(p)
}

// Close marks the buffer as complete. Readers will return io.EOF instead of blocking
// when they reach the end of the buffer.
func (b *Buffer) Close() error {
	defer b.cond.Broadcast()
	b.kill()
	return nil
}

// NewBuffer creates and returns a new Buffer backed by the passed Writer
func NewBuffer(w Writer) *Buffer {
	buf := Buffer{
		buf: w,
	}
	buf.cond = sync.NewCond(&buf.mu)
	return &buf
}

// New creates and returns a new Buffer
func New() *Buffer {
	return NewBuffer(newWriter(nil))
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/hex"
	"flag"
	"fmt"
	"log"
	"os"
	"path/filepath"

	"github.com/andres-erbsen/chatterbox/client/daemon"
)

type hex32Byte [32]byte

func (h *hex32Byte) String() string {
	return hex.EncodeToString(h[:])
}
func (h *hex32Byte) Set(value string) error {
	if len(value) != 2*32 {
		return fmt.Errorf("Server pubkey must be 64 hex digits long, got %d", len(value))
	}
	_, err := hex.Decode(h[:], []byte(value))
	if err != nil {
		return err
	}
	return err
}

func main() {
	dename := flag.String("dename", "", "Your dename username.")
	var serverTransportPubkey [32]byte
	if err := ((*hex32Byte)(&serverTransportPubkey)).Set("70eb7fb3e6c85c006ba76b48208ccf75f99eb6ec98dffb4292388369cb197e30"); err != nil {
		panic(err)
	}
	flag.Var((*hex32Byte)(&serverTransportPubkey), "server-pubkey", "The TCP port which the server listens on. Note that people sending you mesages expct to be able to reach your home server at port 1984.")
	serverAddress := flag.String("server-host", "chatterbox.xvm.mit.edu", "The IP address or hostname on which your (prospective) home server server can be reached")
	serverPort := flag.Int("server-port", 1984, "The TCP port which the server listens on.")
	dir := flag.String("account-directory", "", "Dedicated directory for the account.")
	torAddress := flag.Bool("tor-address", "127.0.0.1:9050", "Address of the local TOR proxy.")
	flag.Parse()

	if *dename == "" || serverTransportPubkey == [32]byte{} || *serverAddress == "" {
		flag.Usage()
		os.Exit(2)
	}
	if *dir == "" {
		*dir = filepath.Join(os.Getenv("HOME"), ".chatterbox", *dename)
	}

	if err := daemon.Init(*dir, *dename, *serverAddress, *serverPort, &serverTransportPubkey, *torAddress); err != nil {
		log.Fatal(err)
	}
	fmt.Printf("Account initialization done.\n"+
		"You may use the following command to link this account with your dename profile:.\n"+
		"torify dnmgr set '%s' 1984 < %s/.daemon/chatterbox-profile.pb\n", *dename, *dir)
}
<commit_msg>fix some more build...<commit_after>package main

import (
	"encoding/hex"
	"flag"
	"fmt"
	"log"
	"os"
	"path/filepath"

	"github.com/andres-erbsen/chatterbox/client/daemon"
)

type hex32Byte [32]byte

func (h *hex32Byte) String() string {
	return hex.EncodeToString(h[:])
}
func (h *hex32Byte) Set(value string) error {
	if len(value) != 2*32 {
		return fmt.Errorf("Server pubkey must be 64 hex digits long, got %d", len(value))
	}
	_, err := hex.Decode(h[:], []byte(value))
	if err != nil {
		return err
	}
	return err
}

func main() {
	dename := flag.String("dename", "", "Your dename username.")
	var serverTransportPubkey [32]byte
	if err := ((*hex32Byte)(&serverTransportPubkey)).Set("70eb7fb3e6c85c006ba76b48208ccf75f99eb6ec98dffb4292388369cb197e30"); err != nil {
		panic(err)
	}
	flag.Var((*hex32Byte)(&serverTransportPubkey), "server-pubkey", "The TCP port which the server listens on. Note that people sending you mesages expct to be able to reach your home server at port 1984.")
	serverAddress := flag.String("server-host", "chatterbox.xvm.mit.edu", "The IP address or hostname on which your (prospective) home server server can be reached")
	serverPort := flag.Int("server-port", 1984, "The TCP port which the server listens on.")
	dir := flag.String("account-directory", "", "Dedicated directory for the account.")
	torAddress := flag.String("tor-address", "127.0.0.1:9050", "Address of the local TOR proxy.")
	flag.Parse()

	if *dename == "" || serverTransportPubkey == [32]byte{} || *serverAddress == "" {
		flag.Usage()
		os.Exit(2)
	}
	if *dir == "" {
		*dir = filepath.Join(os.Getenv("HOME"), ".chatterbox", *dename)
	}

	if err := daemon.Init(*dir, *dename, *serverAddress, *serverPort, &serverTransportPubkey, *torAddress); err != nil {
		log.Fatal(err)
	}
	fmt.Printf("Account initialization done.\n"+
		"You may use the following command to link this account with your dename profile:.\n"+
		"torify dnmgr set '%s' 1984 < %s/.daemon/chatterbox-profile.pb\n", *dename, *dir)
}
<|endoftext|>


<commit_before>// Copyright 2016 Albert Nigmatzianov. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.

package id3v2

import (
	"bytes"
	"io"
	"io/ioutil"
	"os"

	"github.com/bogem/id3v2/bbpool"
	"github.com/bogem/id3v2/util"
)

// Tag stores all frames of opened file.
type Tag struct {
	framesCoords map[string][]frameCoordinates
	frames       map[string]Framer
	sequences    map[string]sequencer
	ids          map[string]string

	file         *os.File
	originalSize int64
}

func (t *Tag) AddFrame(id string, f Framer) {
	if addFunc := t.findSpecificAddFunction(id); addFunc != nil {
		addFunc(f)
	} else {
		t.frames[id] = f
	}
}

// ID returns ID3v2.3 or ID3v2.4 (in appropriate to version of Tag) frame ID
// from given description.
// For example, ID("Mood") will return "TMOO".
func (t Tag) ID(description string) string {
	return t.ids[description]
}

func (t *Tag) findSpecificAddFunction(id string) func(Framer) {
	switch id {
	case t.ID("Attached picture"):
		return t.AddAttachedPicture
	case t.ID("Comments"):
		return t.AddCommentFrame
	case t.ID("Unsynchronised lyrics/text transcription"):
		return t.AddUnsynchronisedLyricsFrame
	}
	return nil
}

func (t *Tag) AddAttachedPicture(f Framer) {
	id := t.ID("Attached picture")
	t.checkExistenceOfSequence(id, newPictureSequence)
	t.addFrameToSequence(id, f)
}

func (t *Tag) AddCommentFrame(f Framer) {
	id := t.ID("Comments")
	t.checkExistenceOfSequence(id, newCommentSequence)
	t.addFrameToSequence(id, f)
}

func (t *Tag) AddUnsynchronisedLyricsFrame(f Framer) {
	id := t.ID("Unsynchronised lyrics/text transcription")
	t.checkExistenceOfSequence(id, newUSLFSequence)
	t.addFrameToSequence(id, f)
}

func (t *Tag) checkExistenceOfSequence(id string, newSequence func() sequencer) {
	if t.sequences[id] == nil {
		t.sequences[id] = newSequence()
	}
}

func (t *Tag) addFrameToSequence(id string, f Framer) {
	t.sequences[id].AddFrame(f)
}

// AllFrames returns map, that contains all frames in tag, that could be parsed.
// The key of this map is an ID of frame and value is an array of frames.
func (t *Tag) AllFrames() map[string][]Framer {
	t.parseAllFramesCoords()

	frames := make(map[string][]Framer)

	// Add frames from t.frames
	for id, frame := range t.frames {
		frames[id] = append(frames[id], frame)
	}

	// Add frames from t.sequences
	for id, sequence := range t.sequences {
		frames[id] = append(frames[id], sequence.Frames()...)
	}

	return frames
}

// GetLastFrame returns last frame from slice, that is returned from GetFrames function.
// GetLastFrame is suitable for frames, that can be only one in whole tag.
// For example, for text frames.
//
// Example of usage:
//	bpmFramer := t.GetLastFrame(t.ID("BPM"))
//	bpm, ok := bpmFramer.(id3v2.TextFrame)
//	if !ok {
//		log.Fatal("Couldn't assert bpm frame")
//	}
//	fmt.Println(bpm.Text)
func (t *Tag) GetLastFrame(id string) Framer {
	fs := t.GetFrames(id)
	if len(fs) == 0 || fs == nil {
		return nil
	}
	return fs[len(fs)-1]
}

// GetFrames returns frames with corresponding id.
//
// Example of usage:
//	pictures := tag.GetFrames(tag.ID("Attached picture"))
//	for _, f := range pictures {
//		pic, ok := f.(id3v2.PictureFrame)
//		if !ok {
//			log.Fatal("Couldn't assert picture frame")
//		}
//
//		// Do some operations with picture frame:
//		fmt.Println(pic.Description) // For example, print description of picture frame
//		image, err := ioutil.ReadAll(pic.Picture) // Or read a picture from picture frame
//		if err != nil {
//			log.Fatal("Error while reading a picture from picture frame: ", err)
//		}
//	}
func (t *Tag) GetFrames(id string) []Framer {
	// If frames with id didn't parsed yet, parse them
	if _, exists := t.framesCoords[id]; exists {
		t.parseFramesCoordsWithID(id)
	}

	if f, exists := t.frames[id]; exists {
		return []Framer{f}
	}

	if s, exists := t.sequences[id]; exists {
		return s.Frames()
	}

	return nil
}

// GetTextFrame returns text frame with corresponding id.
func (t Tag) GetTextFrame(id string) TextFrame {
	f := t.GetLastFrame(id)
	if f == nil {
		return TextFrame{}
	}
	tf := f.(TextFrame)
	return tf
}

func (t Tag) Title() string {
	f := t.GetTextFrame(t.ID("Title/Songname/Content description"))
	return f.Text
}

func (t *Tag) SetTitle(title string) {
	t.AddFrame(t.ID("Title/Songname/Content description"), TextFrame{Encoding: ENUTF8, Text: title})
}

func (t Tag) Artist() string {
	f := t.GetTextFrame(t.ID("Lead artist/Lead performer/Soloist/Performing group"))
	return f.Text
}

func (t *Tag) SetArtist(artist string) {
	t.AddFrame(t.ID("Lead artist/Lead performer/Soloist/Performing group"), TextFrame{Encoding: ENUTF8, Text: artist})
}

func (t Tag) Album() string {
	f := t.GetTextFrame(t.ID("Album/Movie/Show title"))
	return f.Text
}

func (t *Tag) SetAlbum(album string) {
	t.AddFrame(t.ID("Album/Movie/Show title"), TextFrame{Encoding: ENUTF8, Text: album})
}

func (t Tag) Year() string {
	f := t.GetTextFrame(t.ID("Recording time"))
	return f.Text
}

func (t *Tag) SetYear(year string) {
	t.AddFrame(t.ID("Recording time"), TextFrame{Encoding: ENUTF8, Text: year})
}

func (t Tag) Genre() string {
	f := t.GetTextFrame(t.ID("Content type"))
	return f.Text
}

func (t *Tag) SetGenre(genre string) {
	t.AddFrame(t.ID("Content type"), TextFrame{Encoding: ENUTF8, Text: genre})
}

// Save writes tag to the file.
func (t *Tag) Save() error {
	// Forming new frames
	frames := t.formAllFrames()

	// Forming size of new frames
	framesSize := util.FormSize(int64(len(frames)))

	// Creating a temp file for mp3 file, which will contain new tag
	newFile, err := ioutil.TempFile("", "")
	if err != nil {
		return err
	}

	// Writing to new file new tag header
	if _, err = newFile.Write(formTagHeader(framesSize)); err != nil {
		return err
	}

	// Writing to new file new frames
	if _, err = newFile.Write(frames); err != nil {
		return err
	}

	// Seeking to a music part of mp3
	originalFile := t.file
	defer originalFile.Close()
	if _, err = originalFile.Seek(t.originalSize, os.SEEK_SET); err != nil {
		return err
	}

	// Writing to new file the music part
	if _, err = io.Copy(newFile, originalFile); err != nil {
		return err
	}

	// Getting original file mode
	originalFileStat, err := originalFile.Stat()
	if err != nil {
		return err
	}
	originalFileMode := originalFileStat.Mode()

	// Setting new file mode
	if err = newFile.Chmod(originalFileMode); err != nil {
		return err
	}

	// Replacing original file with new file
	if err = os.Rename(newFile.Name(), originalFile.Name()); err != nil {
		return err
	}
	t.file = newFile

	return nil
}

// Close closes the tag's file, rendering it unusable for I/O.
// It returns an error, if any.
func (t *Tag) Close() error {
	return t.file.Close()
}

func (t Tag) formAllFrames() []byte {
	framesBuffer := bbpool.Get()
	defer bbpool.Put(framesBuffer)

	t.writeFrames(framesBuffer)

	return framesBuffer.Bytes()
}

func (t Tag) writeFrames(w io.Writer) {
	for id, frames := range t.AllFrames() {
		for _, f := range frames {
			w.Write(formFrame(id, f))
		}
	}
}

func formFrame(id string, frame Framer) []byte {
	if id == "" {
		panic("there is blank ID in frames")
	}

	frameBuffer := bbpool.Get()
	defer bbpool.Put(frameBuffer)

	frameBody := frame.Body()
	writeFrameHeader(frameBuffer, id, int64(len(frameBody)))
	frameBuffer.Write(frameBody)

	return frameBuffer.Bytes()
}

func writeFrameHeader(buf *bytes.Buffer, id string, frameSize int64) {
	buf.WriteString(id)
	buf.Write(util.FormSize(frameSize))
	buf.Write([]byte{0, 0})
}
<commit_msg>Add version property to Tag<commit_after>// Copyright 2016 Albert Nigmatzianov. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.

package id3v2

import (
	"bytes"
	"io"
	"io/ioutil"
	"os"

	"github.com/bogem/id3v2/bbpool"
	"github.com/bogem/id3v2/util"
)

// Tag stores all frames of opened file.
type Tag struct {
	framesCoords map[string][]frameCoordinates
	frames       map[string]Framer
	sequences    map[string]sequencer
	ids          map[string]string

	file         *os.File
	originalSize int64
	version      byte
}

func (t *Tag) AddFrame(id string, f Framer) {
	if addFunc := t.findSpecificAddFunction(id); addFunc != nil {
		addFunc(f)
	} else {
		t.frames[id] = f
	}
}

// ID returns ID3v2.3 or ID3v2.4 (in appropriate to version of Tag) frame ID
// from given description.
// For example, ID("Mood") will return "TMOO".
func (t Tag) ID(description string) string {
	return t.ids[description]
}

func (t *Tag) findSpecificAddFunction(id string) func(Framer) {
	switch id {
	case t.ID("Attached picture"):
		return t.AddAttachedPicture
	case t.ID("Comments"):
		return t.AddCommentFrame
	case t.ID("Unsynchronised lyrics/text transcription"):
		return t.AddUnsynchronisedLyricsFrame
	}
	return nil
}

func (t *Tag) AddAttachedPicture(f Framer) {
	id := t.ID("Attached picture")
	t.checkExistenceOfSequence(id, newPictureSequence)
	t.addFrameToSequence(id, f)
}

func (t *Tag) AddCommentFrame(f Framer) {
	id := t.ID("Comments")
	t.checkExistenceOfSequence(id, newCommentSequence)
	t.addFrameToSequence(id, f)
}

func (t *Tag) AddUnsynchronisedLyricsFrame(f Framer) {
	id := t.ID("Unsynchronised lyrics/text transcription")
	t.checkExistenceOfSequence(id, newUSLFSequence)
	t.addFrameToSequence(id, f)
}

func (t *Tag) checkExistenceOfSequence(id string, newSequence func() sequencer) {
	if t.sequences[id] == nil {
		t.sequences[id] = newSequence()
	}
}

func (t *Tag) addFrameToSequence(id string, f Framer) {
	t.sequences[id].AddFrame(f)
}

// AllFrames returns map, that contains all frames in tag, that could be parsed.
// The key of this map is an ID of frame and value is an array of frames.
func (t *Tag) AllFrames() map[string][]Framer {
	t.parseAllFramesCoords()

	frames := make(map[string][]Framer)

	// Add frames from t.frames
	for id, frame := range t.frames {
		frames[id] = append(frames[id], frame)
	}

	// Add frames from t.sequences
	for id, sequence := range t.sequences {
		frames[id] = append(frames[id], sequence.Frames()...)
	}

	return frames
}

// GetLastFrame returns last frame from slice, that is returned from GetFrames function.
// GetLastFrame is suitable for frames, that can be only one in whole tag.
// For example, for text frames.
//
// Example of usage:
//	bpmFramer := t.GetLastFrame(t.ID("BPM"))
//	bpm, ok := bpmFramer.(id3v2.TextFrame)
//	if !ok {
//		log.Fatal("Couldn't assert bpm frame")
//	}
//	fmt.Println(bpm.Text)
func (t *Tag) GetLastFrame(id string) Framer {
	fs := t.GetFrames(id)
	if len(fs) == 0 || fs == nil {
		return nil
	}
	return fs[len(fs)-1]
}

// GetFrames returns frames with corresponding id.
//
// Example of usage:
//	pictures := tag.GetFrames(tag.ID("Attached picture"))
//	for _, f := range pictures {
//		pic, ok := f.(id3v2.PictureFrame)
//		if !ok {
//			log.Fatal("Couldn't assert picture frame")
//		}
//
//		// Do some operations with picture frame:
//		fmt.Println(pic.Description) // For example, print description of picture frame
//		image, err := ioutil.ReadAll(pic.Picture) // Or read a picture from picture frame
//		if err != nil {
//			log.Fatal("Error while reading a picture from picture frame: ", err)
//		}
//	}
func (t *Tag) GetFrames(id string) []Framer {
	// If frames with id didn't parsed yet, parse them
	if _, exists := t.framesCoords[id]; exists {
		t.parseFramesCoordsWithID(id)
	}

	if f, exists := t.frames[id]; exists {
		return []Framer{f}
	}

	if s, exists := t.sequences[id]; exists {
		return s.Frames()
	}

	return nil
}

// GetTextFrame returns text frame with corresponding id.
func (t Tag) GetTextFrame(id string) TextFrame {
	f := t.GetLastFrame(id)
	if f == nil {
		return TextFrame{}
	}
	tf := f.(TextFrame)
	return tf
}

func (t Tag) Title() string {
	f := t.GetTextFrame(t.ID("Title/Songname/Content description"))
	return f.Text
}

func (t *Tag) SetTitle(title string) {
	t.AddFrame(t.ID("Title/Songname/Content description"), TextFrame{Encoding: ENUTF8, Text: title})
}

func (t Tag) Artist() string {
	f := t.GetTextFrame(t.ID("Lead artist/Lead performer/Soloist/Performing group"))
	return f.Text
}

func (t *Tag) SetArtist(artist string) {
	t.AddFrame(t.ID("Lead artist/Lead performer/Soloist/Performing group"), TextFrame{Encoding: ENUTF8, Text: artist})
}

func (t Tag) Album() string {
	f := t.GetTextFrame(t.ID("Album/Movie/Show title"))
	return f.Text
}

func (t *Tag) SetAlbum(album string) {
	t.AddFrame(t.ID("Album/Movie/Show title"), TextFrame{Encoding: ENUTF8, Text: album})
}

func (t Tag) Year() string {
	f := t.GetTextFrame(t.ID("Recording time"))
	return f.Text
}

func (t *Tag) SetYear(year string) {
	t.AddFrame(t.ID("Recording time"), TextFrame{Encoding: ENUTF8, Text: year})
}

func (t Tag) Genre() string {
	f := t.GetTextFrame(t.ID("Content type"))
	return f.Text
}

func (t *Tag) SetGenre(genre string) {
	t.AddFrame(t.ID("Content type"), TextFrame{Encoding: ENUTF8, Text: genre})
}

// Save writes tag to the file.
func (t *Tag) Save() error {
	// Forming new frames
	frames := t.formAllFrames()

	// Forming size of new frames
	framesSize := util.FormSize(int64(len(frames)))

	// Creating a temp file for mp3 file, which will contain new tag
	newFile, err := ioutil.TempFile("", "")
	if err != nil {
		return err
	}

	// Writing to new file new tag header
	if _, err = newFile.Write(formTagHeader(framesSize)); err != nil {
		return err
	}

	// Writing to new file new frames
	if _, err = newFile.Write(frames); err != nil {
		return err
	}

	// Seeking to a music part of mp3
	originalFile := t.file
	defer originalFile.Close()
	if _, err = originalFile.Seek(t.originalSize, os.SEEK_SET); err != nil {
		return err
	}

	// Writing to new file the music part
	if _, err = io.Copy(newFile, originalFile); err != nil {
		return err
	}

	// Getting original file mode
	originalFileStat, err := originalFile.Stat()
	if err != nil {
		return err
	}
	originalFileMode := originalFileStat.Mode()

	// Setting new file mode
	if err = newFile.Chmod(originalFileMode); err != nil {
		return err
	}

	// Replacing original file with new file
	if err = os.Rename(newFile.Name(), originalFile.Name()); err != nil {
		return err
	}
	t.file = newFile

	return nil
}

// Close closes the tag's file, rendering it unusable for I/O.
// It returns an error, if any.
func (t *Tag) Close() error {
	return t.file.Close()
}

func (t Tag) formAllFrames() []byte {
	framesBuffer := bbpool.Get()
	defer bbpool.Put(framesBuffer)

	t.writeFrames(framesBuffer)

	return framesBuffer.Bytes()
}

func (t Tag) writeFrames(w io.Writer) {
	for id, frames := range t.AllFrames() {
		for _, f := range frames {
			w.Write(formFrame(id, f))
		}
	}
}

func formFrame(id string, frame Framer) []byte {
	if id == "" {
		panic("there is blank ID in frames")
	}

	frameBuffer := bbpool.Get()
	defer bbpool.Put(frameBuffer)

	frameBody := frame.Body()
	writeFrameHeader(frameBuffer, id, int64(len(frameBody)))
	frameBuffer.Write(frameBody)

	return frameBuffer.Bytes()
}

func writeFrameHeader(buf *bytes.Buffer, id string, frameSize int64) {
	buf.WriteString(id)
	buf.Write(util.FormSize(frameSize))
	buf.Write([]byte{0, 0})
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"bytes"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"text/template"

	"github.com/fatih/color"
)

func check(e error) {
	if e != nil {
		log.Fatal(e)
	}
}

var (
	red           = color.RedString
	blue          = color.BlueString
	pathRe        = regexp.MustCompile(`^(?:\x1b\[[^m]+m)?([^\x1b]+).*`)
	lineNumberRe  = regexp.MustCompile(`^(?:\x1b\[[^m]+m)?(\d+)(?:\x1b\[0m\x1b\[K)?:.*`)
	cleanFilename = regexp.MustCompile(`([ \(\)\[\]\<\>])`)
)

type AliasFile struct {
	filename string
	fmtStr   string
	buf      bytes.Buffer
	writer   *bufio.Writer
}

func NewAliasFile() *AliasFile {
	aliasFilename := os.Getenv("TAG_ALIAS_FILE")
	if len(aliasFilename) == 0 {
		aliasFilename = "/tmp/tag_aliases"
	}

	aliasPrefix := os.Getenv("TAG_ALIAS_PREFIX")
	if len(aliasPrefix) == 0 {
		aliasPrefix = "e"
	}

	aliasCmdFmtString := os.Getenv("TAG_CMD_FMT_STRING")
	if len(aliasCmdFmtString) == 0 {
		aliasCmdFmtString = "vim {{.Filename}} +{{.LineNumber}}"
	}

	a := &AliasFile{
		fmtStr:   "alias " + aliasPrefix + "{{.MatchIndex}}='" + aliasCmdFmtString + "'\n",
		filename: aliasFilename,
	}
	a.writer = bufio.NewWriter(&a.buf)
	return a
}

func (a *AliasFile) WriteAlias(index int, filename, linenum string) {
	t := template.Must(template.New("alias").Parse(a.fmtStr))

	filename = cleanFilename.ReplaceAllString(filename, "\\$1")

	aliasVars := struct {
		MatchIndex int
		Filename   string
		LineNumber string
	}{index, filename, linenum}

	err := t.Execute(a.writer, aliasVars)
	check(err)
}

func (a *AliasFile) WriteFile() {
	err := a.writer.Flush()
	check(err)

	err = ioutil.WriteFile(a.filename, a.buf.Bytes(), 0644)
	check(err)
}

func optionIndex(args []string, option string) int {
	for i := len(args) - 1; i >= 0; i-- {
		if args[i] == option {
			return i
		}
	}
	return -1
}

func tagPrefix(aliasIndex int) string {
	return blue("[") + red("%d", aliasIndex) + blue("]")
}

// Modified from: https://golang.org/src/bufio/scan.go?s=10982:11013
// dropCR drops a terminal \r from the data.
func dropCR(data []byte) []byte {
	if len(data) > 0 && data[len(data)-1] == '\r' {
		return data[0 : len(data)-1]
	}
	return data
}

// Modified from: https://golang.org/src/bufio/scan.go?s=11500:11578
func scanLinesAndNulls(data []byte, atEOF bool) (advance int, token []byte, err error) {
	if atEOF && len(data) == 0 {
		return 0, nil, nil
	}
	if i := bytes.IndexAny(data, "\n\x00"); i >= 0 {
		// We have a full newline-terminated line or a null-terminated sequence.
		return i + 1, dropCR(data[0:i]), nil
	}
	// If we're at EOF, we have a final, non-terminated line. Return it.
	if atEOF {
		return len(data), dropCR(data), nil
	}
	// Request more data.
	return 0, nil, nil
}

func generateTags(cmd *exec.Cmd) int {
	color.NoColor = (optionIndex(cmd.Args, "--nocolor") > 0)
	cmd.Args = append(cmd.Args, "--null")

	stdout, err := cmd.StdoutPipe()
	check(err)

	scanner := bufio.NewScanner(stdout)
	scanner.Split(scanLinesAndNulls)

	var (
		line      []byte
		curPath   string
		groupIdxs []int
	)

	aliasFile := NewAliasFile()
	defer aliasFile.WriteFile()

	aliasIndex := 1

	err = cmd.Start()
	check(err)

	for scanner.Scan() {
		line = scanner.Bytes()
		if groupIdxs = lineNumberRe.FindSubmatchIndex(line); len(groupIdxs) > 0 {
			// Extract and tagged match
			aliasFile.WriteAlias(aliasIndex, curPath, string(line[groupIdxs[2]:groupIdxs[3]]))
			fmt.Printf("%s %s\n", tagPrefix(aliasIndex), string(line))
			aliasIndex++
		} else if groupIdxs = pathRe.FindSubmatchIndex(line); len(groupIdxs) > 0 {
			// Extract and print path
			curPath = string(line[groupIdxs[2]:groupIdxs[3]])
			curPath, err = filepath.Abs(curPath)
			check(err)
			fmt.Println(string(line[:groupIdxs[1]]))
		} else {
			fmt.Println(string(line))
		}
	}

	err = cmd.Wait()
	if err != nil {
		return 1
	}
	return 0
}

func passThrough(cmd *exec.Cmd) int {
	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout

	err := cmd.Run()
	if err != nil {
		return 1
	}
	return 0
}

func isatty(f *os.File) bool {
	stat, err := f.Stat()
	check(err)
	return stat.Mode()&os.ModeCharDevice != 0
}

func main() {
	noTag := false
	var tagArgs []string

	switch i := optionIndex(os.Args, "--notag"); {
	case i > 0:
		os.Args = append(os.Args[:i], os.Args[i+1:]...)
		fallthrough
	case len(os.Args) == 1:
		noTag = true
	default:
		tagArgs = []string{"--group", "--color"}
	}

	/* From ag src/options.c:
	 * If we're not outputting to a terminal. change output to:
	 * turn off colors
	 * print filenames on every line
	 */
	args := os.Args[1:]
	if isatty(os.Stdout) {
		args = append(tagArgs, args...)
	}

	cmd := exec.Command("ag", args...)
	cmd.Stderr = os.Stderr

	if noTag || !isatty(os.Stdin) || !isatty(os.Stdout) {
		// Data being piped from stdin
		os.Exit(passThrough(cmd))
	}

	os.Exit(generateTags(cmd))
}
<commit_msg>Implement column number parsing<commit_after>package main

import (
	"bufio"
	"bytes"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"text/template"

	"github.com/fatih/color"
)

func check(e error) {
	if e != nil {
		log.Fatal(e)
	}
}

var (
	red           = color.RedString
	blue          = color.BlueString
	pathRe        = regexp.MustCompile(`^(?:\x1b\[[^m]+m)?([^\x1b]+).*`)
	lineNumberRe  = regexp.MustCompile(`^(?:\x1b\[[^m]+m)?(\d+)(?:\x1b\[0m\x1b\[K)?:(\d+):.*`)
	cleanFilename = regexp.MustCompile(`([ \(\)\[\]\<\>])`)
)

type AliasFile struct {
	filename string
	fmtStr   string
	buf      bytes.Buffer
	writer   *bufio.Writer
}

func NewAliasFile() *AliasFile {
	aliasFilename := os.Getenv("TAG_ALIAS_FILE")
	if len(aliasFilename) == 0 {
		aliasFilename = "/tmp/tag_aliases"
	}

	aliasPrefix := os.Getenv("TAG_ALIAS_PREFIX")
	if len(aliasPrefix) == 0 {
		aliasPrefix = "e"
	}

	aliasCmdFmtString := os.Getenv("TAG_CMD_FMT_STRING")
	if len(aliasCmdFmtString) == 0 {
		aliasCmdFmtString = "vim {{.Filename}} \"+call cursor({{.LineNumber}}, {{.ColumnNumber}})\""
	}

	a := &AliasFile{
		fmtStr:   "alias " + aliasPrefix + "{{.MatchIndex}}='" + aliasCmdFmtString + "'\n",
		filename: aliasFilename,
	}
	a.writer = bufio.NewWriter(&a.buf)
	return a
}

func (a *AliasFile) WriteAlias(index int, filename, linenum string, colnum string) {
	t := template.Must(template.New("alias").Parse(a.fmtStr))

	filename = cleanFilename.ReplaceAllString(filename, "\\$1")

	aliasVars := struct {
		MatchIndex   int
		Filename     string
		LineNumber   string
		ColumnNumber string
	}{index, filename, linenum, colnum}

	err := t.Execute(a.writer, aliasVars)
	check(err)
}

func (a *AliasFile) WriteFile() {
	err := a.writer.Flush()
	check(err)

	err = ioutil.WriteFile(a.filename, a.buf.Bytes(), 0644)
	check(err)
}

func optionIndex(args []string, option string) int {
	for i := len(args) - 1; i >= 0; i-- {
		if args[i] == option {
			return i
		}
	}
	return -1
}

func tagPrefix(aliasIndex int) string {
	return blue("[") + red("%d", aliasIndex) + blue("]")
}

// Modified from: https://golang.org/src/bufio/scan.go?s=10982:11013
// dropCR drops a terminal \r from the data.
func dropCR(data []byte) []byte {
	if len(data) > 0 && data[len(data)-1] == '\r' {
		return data[0 : len(data)-1]
	}
	return data
}

// Modified from: https://golang.org/src/bufio/scan.go?s=11500:11578
func scanLinesAndNulls(data []byte, atEOF bool) (advance int, token []byte, err error) {
	if atEOF && len(data) == 0 {
		return 0, nil, nil
	}
	if i := bytes.IndexAny(data, "\n\x00"); i >= 0 {
		// We have a full newline-terminated line or a null-terminated sequence.
		return i + 1, dropCR(data[0:i]), nil
	}
	// If we're at EOF, we have a final, non-terminated line. Return it.
	if atEOF {
		return len(data), dropCR(data), nil
	}
	// Request more data.
	return 0, nil, nil
}

func generateTags(cmd *exec.Cmd) int {
	color.NoColor = (optionIndex(cmd.Args, "--nocolor") > 0)
	cmd.Args = append(cmd.Args, "--null")

	stdout, err := cmd.StdoutPipe()
	check(err)

	scanner := bufio.NewScanner(stdout)
	scanner.Split(scanLinesAndNulls)

	var (
		line      []byte
		curPath   string
		groupIdxs []int
	)

	aliasFile := NewAliasFile()
	defer aliasFile.WriteFile()

	aliasIndex := 1

	err = cmd.Start()
	check(err)

	for scanner.Scan() {
		line = scanner.Bytes()
		if groupIdxs = lineNumberRe.FindSubmatchIndex(line); len(groupIdxs) > 0 {
			// Extract and tagged match
			aliasFile.WriteAlias(aliasIndex, curPath, string(line[groupIdxs[2]:groupIdxs[3]]), string(line[groupIdxs[4]:groupIdxs[5]]))
			fmt.Printf("%s %s\n", tagPrefix(aliasIndex), string(line))
			aliasIndex++
		} else if groupIdxs = pathRe.FindSubmatchIndex(line); len(groupIdxs) > 0 {
			// Extract and print path
			curPath = string(line[groupIdxs[2]:groupIdxs[3]])
			curPath, err = filepath.Abs(curPath)
			check(err)
			fmt.Println(string(line[:groupIdxs[1]]))
		} else {
			fmt.Println(string(line))
		}
	}

	err = cmd.Wait()
	if err != nil {
		return 1
	}
	return 0
}

func passThrough(cmd *exec.Cmd) int {
	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout

	err := cmd.Run()
	if err != nil {
		return 1
	}
	return 0
}

func isatty(f *os.File) bool {
	stat, err := f.Stat()
	check(err)
	return stat.Mode()&os.ModeCharDevice != 0
}

func main() {
	noTag := false
	var tagArgs []string

	switch i := optionIndex(os.Args, "--notag"); {
	case i > 0:
		os.Args = append(os.Args[:i], os.Args[i+1:]...)
		fallthrough
	case len(os.Args) == 1:
		noTag = true
	default:
		tagArgs = []string{"--group", "--color", "--column"}
	}

	/* From ag src/options.c:
	 * If we're not outputting to a terminal. change output to:
	 * turn off colors
	 * print filenames on every line
	 */
	args := os.Args[1:]
	if isatty(os.Stdout) {
		args = append(tagArgs, args...)
	}

	cmd := exec.Command("ag", args...)
	cmd.Stderr = os.Stderr

	if noTag || !isatty(os.Stdin) || !isatty(os.Stdout) {
		// Data being piped from stdin
		os.Exit(passThrough(cmd))
	}

	os.Exit(generateTags(cmd))
}
<|endoftext|>


<commit_before>package stackup

import (
	"fmt"
	"io"

	"os/exec"
)

// Copying dirs/files over SSH using TAR.
// tar -C . -cvzf - <dirs/files> | ssh <host> "tar -C <dst_dir> -xvzf -"

// RemoteTarCommand returns command to be run on remote SSH host
// to properly receive the created TAR stream.
// TODO: Check for relative directory.
func RemoteTarCommand(dir string) string {
	return fmt.Sprintf("tar -C '%s' -xvzf -", dir)
}

func LocalTarCommand(path string) string {
	return fmt.Sprintf("tar -C '.' -cvzf - %s", path)
}

// NewTarStreamReader creates a tar stream reader from a local path.
// TODO: Refactor. Use "archive/tar" instead.
func NewTarStreamReader(path, env string) io.Reader {
	// // Dumb way how to check if the "path" exist
	// _, err := os.Stat(path)
	// if err != nil {
	// 	log.Fatal(err)
	// }

	cmd := exec.Command("bash", "-c", env+LocalTarCommand(path))
	//cmd := exec.Command("tar", "-C", ".", "-cvzf", "-", path)

	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return nil
	}

	stderr, err := cmd.StderrPipe()
	if err != nil {
		return nil
	}

	output := io.MultiReader(stdout, stderr)

	if err := cmd.Start(); err != nil {
		return nil
	}

	return output
}
<commit_msg>Allow remote TAR to accept ENV VARs<commit_after>package stackup

import (
	"fmt"
	"io"

	"os/exec"
)

// Copying dirs/files over SSH using TAR.
// tar -C . -cvzf - <dirs/files> | ssh <host> "tar -C <dst_dir> -xvzf -"

// RemoteTarCommand returns command to be run on remote SSH host
// to properly receive the created TAR stream.
// TODO: Check for relative directory.
func RemoteTarCommand(dir string) string {
	return fmt.Sprintf("tar -C \"%s\" -xvzf -", dir)
}

func LocalTarCommand(path string) string {
	return fmt.Sprintf("tar -C '.' -cvzf - %s", path)
}

// NewTarStreamReader creates a tar stream reader from a local path.
// TODO: Refactor. Use "archive/tar" instead.
func NewTarStreamReader(path, env string) io.Reader {
	// // Dumb way how to check if the "path" exist
	// _, err := os.Stat(path)
	// if err != nil {
	// 	log.Fatal(err)
	// }

	cmd := exec.Command("bash", "-c", env+LocalTarCommand(path))
	//cmd := exec.Command("tar", "-C", ".", "-cvzf", "-", path)

	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return nil
	}

	stderr, err := cmd.StderrPipe()
	if err != nil {
		return nil
	}

	output := io.MultiReader(stdout, stderr)

	if err := cmd.Start(); err != nil {
		return nil
	}

	return output
}
<|endoftext|>