<commit_before>package gocql

import (
	"context"
	crand "crypto/rand"
	"errors"
	"fmt"
	"math/rand"
	"net"
	"os"
	"regexp"
	"strconv"
	"sync"
	"sync/atomic"
	"time"
)

var (
	randr    *rand.Rand
	mutRandr sync.Mutex
)

func init() {
	b := make([]byte, 4)
	if _, err := crand.Read(b); err != nil {
		panic(fmt.Sprintf("unable to seed random number generator: %v", err))
	}

	randr = rand.New(rand.NewSource(int64(readInt(b))))
}

// Ensure that the atomic variable is aligned to a 64bit boundary
// so that atomic operations can be applied on 32bit architectures.
type controlConn struct {
	started      int32
	reconnecting int32

	session *Session
	conn    atomic.Value

	retry RetryPolicy

	quit chan struct{}
}

func createControlConn(session *Session) *controlConn {
	control := &controlConn{
		session: session,
		quit:    make(chan struct{}),
		retry:   &SimpleRetryPolicy{NumRetries: 3},
	}

	control.conn.Store((*connHost)(nil))

	return control
}

func (c *controlConn) heartBeat() {
	if !atomic.CompareAndSwapInt32(&c.started, 0, 1) {
		return
	}

	sleepTime := 1 * time.Second
	timer := time.NewTimer(sleepTime)
	defer timer.Stop()

	for {
		timer.Reset(sleepTime)

		select {
		case <-c.quit:
			return
		case <-timer.C:
		}

		resp, err := c.writeFrame(&writeOptionsFrame{})
		if err != nil {
			goto reconn
		}

		switch resp.(type) {
		case *supportedFrame:
			// Everything ok
			sleepTime = 5 * time.Second
			continue
		case error:
			goto reconn
		default:
			panic(fmt.Sprintf("gocql: unknown frame in response to options: %T", resp))
		}

	reconn:
		// try to connect a bit faster
		sleepTime = 1 * time.Second
		c.reconnect(true)
		continue
	}
}

var hostLookupPreferV4 = os.Getenv("GOCQL_HOST_LOOKUP_PREFER_V4") == "true"

func hostInfo(addr string, defaultPort int) ([]*HostInfo, error) {
	var port int
	host, portStr, err := net.SplitHostPort(addr)
	if err != nil {
		host = addr
		port = defaultPort
	} else {
		port, err = strconv.Atoi(portStr)
		if err != nil {
			return nil, err
		}
	}

	var hosts []*HostInfo

	// Check if host is a literal IP address
	if ip := net.ParseIP(host); ip != nil {
		hosts = append(hosts, &HostInfo{connectAddress: ip, port: port})
		return hosts, nil
	}

	// Look up host in DNS
	ips, err := LookupIP(host)
	if err != nil {
		return nil, err
	} else if len(ips) == 0 {
		return nil, fmt.Errorf("No IP's returned from DNS lookup for %q", addr)
	}

	// Filter to v4 addresses if any present
	if hostLookupPreferV4 {
		var preferredIPs []net.IP
		for _, v := range ips {
			if v4 := v.To4(); v4 != nil {
				preferredIPs = append(preferredIPs, v4)
			}
		}
		if len(preferredIPs) != 0 {
			ips = preferredIPs
		}
	}

	for _, ip := range ips {
		hosts = append(hosts, &HostInfo{connectAddress: ip, port: port})
	}

	return hosts, nil
}

func shuffleHosts(hosts []*HostInfo) []*HostInfo {
	mutRandr.Lock()
	perm := randr.Perm(len(hosts))
	mutRandr.Unlock()
	shuffled := make([]*HostInfo, len(hosts))

	for i, host := range hosts {
		shuffled[perm[i]] = host
	}

	return shuffled
}

func (c *controlConn) shuffleDial(endpoints []*HostInfo) (*Conn, error) {
	// shuffle endpoints so not all drivers will connect to the same initial
	// node.
	shuffled := shuffleHosts(endpoints)

	cfg := *c.session.connCfg
	cfg.disableCoalesce = true

	var err error
	for _, host := range shuffled {
		var conn *Conn
		conn, err = c.session.dial(host, &cfg, c)
		if err == nil {
			return conn, nil
		}

		Logger.Printf("gocql: unable to dial control conn %v: %v\n", host.ConnectAddress(), err)
	}

	return nil, err
}

// this is going to be version dependant and a nightmare to maintain :(
var protocolSupportRe = regexp.MustCompile(`the lowest supported version is \d+ and the greatest is (\d+)$`)

func parseProtocolFromError(err error) int {
	// I really wish this had the actual info in the error frame...
	matches := protocolSupportRe.FindAllStringSubmatch(err.Error(), -1)
	if len(matches) != 1 || len(matches[0]) != 2 {
		if verr, ok := err.(*protocolError); ok {
			return int(verr.frame.Header().version.version())
		}
		return 0
	}

	max, err := strconv.Atoi(matches[0][1])
	if err != nil {
		return 0
	}

	return max
}

func (c *controlConn) discoverProtocol(hosts []*HostInfo) (int, error) {
	hosts = shuffleHosts(hosts)

	connCfg := *c.session.connCfg
	connCfg.ProtoVersion = 4 // TODO: define maxProtocol

	handler := connErrorHandlerFn(func(c *Conn, err error, closed bool) {
		// we should never get here, but if we do it means we connected to a
		// host successfully which means our attempted protocol version worked
		if !closed {
			c.Close()
		}
	})

	var err error
	for _, host := range hosts {
		var conn *Conn
		conn, err = c.session.dial(host, &connCfg, handler)
		if conn != nil {
			conn.Close()
		}

		if err == nil {
			return connCfg.ProtoVersion, nil
		}

		if proto := parseProtocolFromError(err); proto > 0 {
			return proto, nil
		}
	}

	return 0, err
}

func (c *controlConn) connect(hosts []*HostInfo) error {
	if len(hosts) == 0 {
		return errors.New("control: no endpoints specified")
	}

	conn, err := c.shuffleDial(hosts)
	if err != nil {
		return fmt.Errorf("control: unable to connect to initial hosts: %v", err)
	}

	if err := c.setupConn(conn); err != nil {
		conn.Close()
		return fmt.Errorf("control: unable to setup connection: %v", err)
	}

	// we could fetch the initial ring here and update initial host data. So that
	// when we return from here we have a ring topology ready to go.

	go c.heartBeat()

	return nil
}

type connHost struct {
	conn *Conn
	host *HostInfo
}

func (c *controlConn) setupConn(conn *Conn) error {
	if err := c.registerEvents(conn); err != nil {
		conn.Close()
		return err
	}

	// TODO(zariel): do we need to fetch host info everytime
	// the control conn connects? Surely we have it cached?
	host, err := conn.localHostInfo(context.TODO())
	if err != nil {
		return err
	}

	ch := &connHost{
		conn: conn,
		host: host,
	}

	c.conn.Store(ch)
	c.session.handleNodeUp(host.ConnectAddress(), host.Port(), false)

	return nil
}

func (c *controlConn) registerEvents(conn *Conn) error {
	var events []string

	if !c.session.cfg.Events.DisableTopologyEvents {
		events = append(events, "TOPOLOGY_CHANGE")
	}
	if !c.session.cfg.Events.DisableNodeStatusEvents {
		events = append(events, "STATUS_CHANGE")
	}
	if !c.session.cfg.Events.DisableSchemaEvents {
		events = append(events, "SCHEMA_CHANGE")
	}

	if len(events) == 0 {
		return nil
	}

	framer, err := conn.exec(context.Background(),
		&writeRegisterFrame{
			events: events,
		}, nil)
	if err != nil {
		return err
	}

	frame, err := framer.parseFrame()
	if err != nil {
		return err
	} else if _, ok := frame.(*readyFrame); !ok {
		return fmt.Errorf("unexpected frame in response to register: got %T: %v\n", frame, frame)
	}

	return nil
}

func (c *controlConn) reconnect(refreshring bool) {
	if !atomic.CompareAndSwapInt32(&c.reconnecting, 0, 1) {
		return
	}
	defer atomic.StoreInt32(&c.reconnecting, 0)
	// TODO: simplify this function, use session.ring to get hosts instead of the
	// connection pool

	var host *HostInfo
	ch := c.getConn()
	if ch != nil {
		host = ch.host
		ch.conn.Close()
	}

	var newConn *Conn
	if host != nil {
		// try to connect to the old host
		conn, err := c.session.connect(host, c)
		if err != nil {
			// host is dead
			// TODO: this is replicated in a few places
			if c.session.cfg.ConvictionPolicy.AddFailure(err, host) {
				c.session.handleNodeDown(host.ConnectAddress(), host.Port())
			}
		} else {
			newConn = conn
		}
	}

	// TODO: should have our own round-robin for hosts so that we can try each
	// in succession and guarantee that we get a different host each time.
	if newConn == nil {
		host := c.session.ring.rrHost()
		if host == nil {
			c.connect(c.session.ring.endpoints)
			return
		}

		var err error
		newConn, err = c.session.connect(host, c)
		if err != nil {
			// TODO: add log handler for things like this
			return
		}
	}

	if err := c.setupConn(newConn); err != nil {
		newConn.Close()
		Logger.Printf("gocql: control unable to register events: %v\n", err)
		return
	}

	if refreshring {
		c.session.hostSource.refreshRing()
	}
}

func (c *controlConn) HandleError(conn *Conn, err error, closed bool) {
	if !closed {
		return
	}

	oldConn := c.getConn()
	if oldConn.conn != conn {
		return
	}

	c.reconnect(false)
}

func (c *controlConn) getConn() *connHost {
	return c.conn.Load().(*connHost)
}

func (c *controlConn) writeFrame(w frameWriter) (frame, error) {
	ch := c.getConn()
	if ch == nil {
		return nil, errNoControl
	}

	framer, err := ch.conn.exec(context.Background(), w, nil)
	if err != nil {
		return nil, err
	}

	return framer.parseFrame()
}

func (c *controlConn) withConnHost(fn func(*connHost) *Iter) *Iter {
	const maxConnectAttempts = 5
	connectAttempts := 0

	for i := 0; i < maxConnectAttempts; i++ {
		ch := c.getConn()
		if ch == nil {
			if connectAttempts > maxConnectAttempts {
				break
			}

			connectAttempts++

			c.reconnect(false)
			continue
		}

		return fn(ch)
	}

	return &Iter{err: errNoControl}
}

func (c *controlConn) withConn(fn func(*Conn) *Iter) *Iter {
	return c.withConnHost(func(ch *connHost) *Iter {
		return fn(ch.conn)
	})
}

// query will return nil if the connection is closed or nil
func (c *controlConn) query(statement string, values ...interface{}) (iter *Iter) {
	q := c.session.Query(statement, values...).Consistency(One).RoutingKey([]byte{}).Trace(nil)

	for {
		iter = c.withConn(func(conn *Conn) *Iter {
			return conn.executeQuery(context.TODO(), q)
		})

		if gocqlDebug && iter.err != nil {
			Logger.Printf("control: error executing %q: %v\n", statement, iter.err)
		}

		q.AddAttempts(1, c.getConn().host)
		if iter.err == nil || !c.retry.Attempt(q) {
			break
		}
	}

	return
}

func (c *controlConn) awaitSchemaAgreement() error {
	return c.withConn(func(conn *Conn) *Iter {
		return &Iter{err: conn.awaitSchemaAgreement(context.TODO())}
	}).err
}

func (c *controlConn) close() {
	if atomic.CompareAndSwapInt32(&c.started, 1, -1) {
		c.quit <- struct{}{}
	}

	ch := c.getConn()
	if ch != nil {
		ch.conn.Close()
	}
}

var errNoControl = errors.New("gocql: no control connection available")
<commit_msg>Add non-nil check to old connection in error handler (#1305)<commit_after>package gocql

import (
	"context"
	crand "crypto/rand"
	"errors"
	"fmt"
	"math/rand"
	"net"
	"os"
	"regexp"
	"strconv"
	"sync"
	"sync/atomic"
	"time"
)

var (
	randr    *rand.Rand
	mutRandr sync.Mutex
)

func init() {
	b := make([]byte, 4)
	if _, err := crand.Read(b); err != nil {
		panic(fmt.Sprintf("unable to seed random number generator: %v", err))
	}

	randr = rand.New(rand.NewSource(int64(readInt(b))))
}

// Ensure that the atomic variable is aligned to a 64bit boundary
// so that atomic operations can be applied on 32bit architectures.
type controlConn struct {
	started      int32
	reconnecting int32

	session *Session
	conn    atomic.Value

	retry RetryPolicy

	quit chan struct{}
}

func createControlConn(session *Session) *controlConn {
	control := &controlConn{
		session: session,
		quit:    make(chan struct{}),
		retry:   &SimpleRetryPolicy{NumRetries: 3},
	}

	control.conn.Store((*connHost)(nil))

	return control
}

func (c *controlConn) heartBeat() {
	if !atomic.CompareAndSwapInt32(&c.started, 0, 1) {
		return
	}

	sleepTime := 1 * time.Second
	timer := time.NewTimer(sleepTime)
	defer timer.Stop()

	for {
		timer.Reset(sleepTime)

		select {
		case <-c.quit:
			return
		case <-timer.C:
		}

		resp, err := c.writeFrame(&writeOptionsFrame{})
		if err != nil {
			goto reconn
		}

		switch resp.(type) {
		case *supportedFrame:
			// Everything ok
			sleepTime = 5 * time.Second
			continue
		case error:
			goto reconn
		default:
			panic(fmt.Sprintf("gocql: unknown frame in response to options: %T", resp))
		}

	reconn:
		// try to connect a bit faster
		sleepTime = 1 * time.Second
		c.reconnect(true)
		continue
	}
}

var hostLookupPreferV4 = os.Getenv("GOCQL_HOST_LOOKUP_PREFER_V4") == "true"

func hostInfo(addr string, defaultPort int) ([]*HostInfo, error) {
	var port int
	host, portStr, err := net.SplitHostPort(addr)
	if err != nil {
		host = addr
		port = defaultPort
	} else {
		port, err = strconv.Atoi(portStr)
		if err != nil {
			return nil, err
		}
	}

	var hosts []*HostInfo

	// Check if host is a literal IP address
	if ip := net.ParseIP(host); ip != nil {
		hosts = append(hosts, &HostInfo{connectAddress: ip, port: port})
		return hosts, nil
	}

	// Look up host in DNS
	ips, err := LookupIP(host)
	if err != nil {
		return nil, err
	} else if len(ips) == 0 {
		return nil, fmt.Errorf("No IP's returned from DNS lookup for %q", addr)
	}

	// Filter to v4 addresses if any present
	if hostLookupPreferV4 {
		var preferredIPs []net.IP
		for _, v := range ips {
			if v4 := v.To4(); v4 != nil {
				preferredIPs = append(preferredIPs, v4)
			}
		}
		if len(preferredIPs) != 0 {
			ips = preferredIPs
		}
	}

	for _, ip := range ips {
		hosts = append(hosts, &HostInfo{connectAddress: ip, port: port})
	}

	return hosts, nil
}

func shuffleHosts(hosts []*HostInfo) []*HostInfo {
	mutRandr.Lock()
	perm := randr.Perm(len(hosts))
	mutRandr.Unlock()
	shuffled := make([]*HostInfo, len(hosts))

	for i, host := range hosts {
		shuffled[perm[i]] = host
	}

	return shuffled
}

func (c *controlConn) shuffleDial(endpoints []*HostInfo) (*Conn, error) {
	// shuffle endpoints so not all drivers will connect to the same initial
	// node.
	shuffled := shuffleHosts(endpoints)

	cfg := *c.session.connCfg
	cfg.disableCoalesce = true

	var err error
	for _, host := range shuffled {
		var conn *Conn
		conn, err = c.session.dial(host, &cfg, c)
		if err == nil {
			return conn, nil
		}

		Logger.Printf("gocql: unable to dial control conn %v: %v\n", host.ConnectAddress(), err)
	}

	return nil, err
}

// this is going to be version dependant and a nightmare to maintain :(
var protocolSupportRe = regexp.MustCompile(`the lowest supported version is \d+ and the greatest is (\d+)$`)

func parseProtocolFromError(err error) int {
	// I really wish this had the actual info in the error frame...
	matches := protocolSupportRe.FindAllStringSubmatch(err.Error(), -1)
	if len(matches) != 1 || len(matches[0]) != 2 {
		if verr, ok := err.(*protocolError); ok {
			return int(verr.frame.Header().version.version())
		}
		return 0
	}

	max, err := strconv.Atoi(matches[0][1])
	if err != nil {
		return 0
	}

	return max
}

func (c *controlConn) discoverProtocol(hosts []*HostInfo) (int, error) {
	hosts = shuffleHosts(hosts)

	connCfg := *c.session.connCfg
	connCfg.ProtoVersion = 4 // TODO: define maxProtocol

	handler := connErrorHandlerFn(func(c *Conn, err error, closed bool) {
		// we should never get here, but if we do it means we connected to a
		// host successfully which means our attempted protocol version worked
		if !closed {
			c.Close()
		}
	})

	var err error
	for _, host := range hosts {
		var conn *Conn
		conn, err = c.session.dial(host, &connCfg, handler)
		if conn != nil {
			conn.Close()
		}

		if err == nil {
			return connCfg.ProtoVersion, nil
		}

		if proto := parseProtocolFromError(err); proto > 0 {
			return proto, nil
		}
	}

	return 0, err
}

func (c *controlConn) connect(hosts []*HostInfo) error {
	if len(hosts) == 0 {
		return errors.New("control: no endpoints specified")
	}

	conn, err := c.shuffleDial(hosts)
	if err != nil {
		return fmt.Errorf("control: unable to connect to initial hosts: %v", err)
	}

	if err := c.setupConn(conn); err != nil {
		conn.Close()
		return fmt.Errorf("control: unable to setup connection: %v", err)
	}

	// we could fetch the initial ring here and update initial host data. So that
	// when we return from here we have a ring topology ready to go.

	go c.heartBeat()

	return nil
}

type connHost struct {
	conn *Conn
	host *HostInfo
}

func (c *controlConn) setupConn(conn *Conn) error {
	if err := c.registerEvents(conn); err != nil {
		conn.Close()
		return err
	}

	// TODO(zariel): do we need to fetch host info everytime
	// the control conn connects? Surely we have it cached?
	host, err := conn.localHostInfo(context.TODO())
	if err != nil {
		return err
	}

	ch := &connHost{
		conn: conn,
		host: host,
	}

	c.conn.Store(ch)
	c.session.handleNodeUp(host.ConnectAddress(), host.Port(), false)

	return nil
}

func (c *controlConn) registerEvents(conn *Conn) error {
	var events []string

	if !c.session.cfg.Events.DisableTopologyEvents {
		events = append(events, "TOPOLOGY_CHANGE")
	}
	if !c.session.cfg.Events.DisableNodeStatusEvents {
		events = append(events, "STATUS_CHANGE")
	}
	if !c.session.cfg.Events.DisableSchemaEvents {
		events = append(events, "SCHEMA_CHANGE")
	}

	if len(events) == 0 {
		return nil
	}

	framer, err := conn.exec(context.Background(),
		&writeRegisterFrame{
			events: events,
		}, nil)
	if err != nil {
		return err
	}

	frame, err := framer.parseFrame()
	if err != nil {
		return err
	} else if _, ok := frame.(*readyFrame); !ok {
		return fmt.Errorf("unexpected frame in response to register: got %T: %v\n", frame, frame)
	}

	return nil
}

func (c *controlConn) reconnect(refreshring bool) {
	if !atomic.CompareAndSwapInt32(&c.reconnecting, 0, 1) {
		return
	}
	defer atomic.StoreInt32(&c.reconnecting, 0)
	// TODO: simplify this function, use session.ring to get hosts instead of the
	// connection pool

	var host *HostInfo
	ch := c.getConn()
	if ch != nil {
		host = ch.host
		ch.conn.Close()
	}

	var newConn *Conn
	if host != nil {
		// try to connect to the old host
		conn, err := c.session.connect(host, c)
		if err != nil {
			// host is dead
			// TODO: this is replicated in a few places
			if c.session.cfg.ConvictionPolicy.AddFailure(err, host) {
				c.session.handleNodeDown(host.ConnectAddress(), host.Port())
			}
		} else {
			newConn = conn
		}
	}

	// TODO: should have our own round-robin for hosts so that we can try each
	// in succession and guarantee that we get a different host each time.
	if newConn == nil {
		host := c.session.ring.rrHost()
		if host == nil {
			c.connect(c.session.ring.endpoints)
			return
		}

		var err error
		newConn, err = c.session.connect(host, c)
		if err != nil {
			// TODO: add log handler for things like this
			return
		}
	}

	if err := c.setupConn(newConn); err != nil {
		newConn.Close()
		Logger.Printf("gocql: control unable to register events: %v\n", err)
		return
	}

	if refreshring {
		c.session.hostSource.refreshRing()
	}
}

func (c *controlConn) HandleError(conn *Conn, err error, closed bool) {
	if !closed {
		return
	}

	oldConn := c.getConn()

	// If connection has long gone, and not been attempted for awhile,
	// it's possible to have oldConn as nil here (#1297).
	if oldConn != nil && oldConn.conn != conn {
		return
	}

	c.reconnect(false)
}

func (c *controlConn) getConn() *connHost {
	return c.conn.Load().(*connHost)
}

func (c *controlConn) writeFrame(w frameWriter) (frame, error) {
	ch := c.getConn()
	if ch == nil {
		return nil, errNoControl
	}

	framer, err := ch.conn.exec(context.Background(), w, nil)
	if err != nil {
		return nil, err
	}

	return framer.parseFrame()
}

func (c *controlConn) withConnHost(fn func(*connHost) *Iter) *Iter {
	const maxConnectAttempts = 5
	connectAttempts := 0

	for i := 0; i < maxConnectAttempts; i++ {
		ch := c.getConn()
		if ch == nil {
			if connectAttempts > maxConnectAttempts {
				break
			}

			connectAttempts++

			c.reconnect(false)
			continue
		}

		return fn(ch)
	}

	return &Iter{err: errNoControl}
}

func (c *controlConn) withConn(fn func(*Conn) *Iter) *Iter {
	return c.withConnHost(func(ch *connHost) *Iter {
		return fn(ch.conn)
	})
}

// query will return nil if the connection is closed or nil
func (c *controlConn) query(statement string, values ...interface{}) (iter *Iter) {
	q := c.session.Query(statement, values...).Consistency(One).RoutingKey([]byte{}).Trace(nil)

	for {
		iter = c.withConn(func(conn *Conn) *Iter {
			return conn.executeQuery(context.TODO(), q)
		})

		if gocqlDebug && iter.err != nil {
			Logger.Printf("control: error executing %q: %v\n", statement, iter.err)
		}

		q.AddAttempts(1, c.getConn().host)
		if iter.err == nil || !c.retry.Attempt(q) {
			break
		}
	}

	return
}

func (c *controlConn) awaitSchemaAgreement() error {
	return c.withConn(func(conn *Conn) *Iter {
		return &Iter{err: conn.awaitSchemaAgreement(context.TODO())}
	}).err
}

func (c *controlConn) close() {
	if atomic.CompareAndSwapInt32(&c.started, 1, -1) {
		c.quit <- struct{}{}
	}

	ch := c.getConn()
	if ch != nil {
		ch.conn.Close()
	}
}

var errNoControl = errors.New("gocql: no control connection available")
<|endoftext|>


<commit_before>package gocql

import (
	"context"
	crand "crypto/rand"
	"errors"
	"fmt"
	"math/rand"
	"net"
	"os"
	"regexp"
	"strconv"
	"sync"
	"sync/atomic"
	"time"
)

var (
	randr    *rand.Rand
	mutRandr sync.Mutex
)

func init() {
	b := make([]byte, 4)
	if _, err := crand.Read(b); err != nil {
		panic(fmt.Sprintf("unable to seed random number generator: %v", err))
	}

	randr = rand.New(rand.NewSource(int64(readInt(b))))
}

// Ensure that the atomic variable is aligned to a 64bit boundary
// so that atomic operations can be applied on 32bit architectures.
type controlConn struct {
	started      int32
	reconnecting int32

	session *Session
	conn    atomic.Value

	retry RetryPolicy

	quit chan struct{}
}

func createControlConn(session *Session) *controlConn {
	control := &controlConn{
		session: session,
		quit:    make(chan struct{}),
		retry:   &SimpleRetryPolicy{NumRetries: 3},
	}

	control.conn.Store((*connHost)(nil))

	return control
}

func (c *controlConn) heartBeat() {
	if !atomic.CompareAndSwapInt32(&c.started, 0, 1) {
		return
	}

	sleepTime := 1 * time.Second
	timer := time.NewTimer(sleepTime)
	defer timer.Stop()

	for {
		timer.Reset(sleepTime)

		select {
		case <-c.quit:
			return
		case <-timer.C:
		}

		resp, err := c.writeFrame(&writeOptionsFrame{})
		if err != nil {
			goto reconn
		}

		switch resp.(type) {
		case *supportedFrame:
			// Everything ok
			sleepTime = 5 * time.Second
			continue
		case error:
			goto reconn
		default:
			panic(fmt.Sprintf("gocql: unknown frame in response to options: %T", resp))
		}

	reconn:
		// try to connect a bit faster
		sleepTime = 1 * time.Second
		c.reconnect(true)
		continue
	}
}

var hostLookupPreferV4 = os.Getenv("GOCQL_HOST_LOOKUP_PREFER_V4") == "true"

func hostInfo(addr string, defaultPort int) ([]*HostInfo, error) {
	var port int
	host, portStr, err := net.SplitHostPort(addr)
	if err != nil {
		host = addr
		port = defaultPort
	} else {
		port, err = strconv.Atoi(portStr)
		if err != nil {
			return nil, err
		}
	}

	var hosts []*HostInfo

	// Check if host is a literal IP address
	if ip := net.ParseIP(host); ip != nil {
		hosts = append(hosts, &HostInfo{connectAddress: ip, port: port})
		return hosts, nil
	}

	// Look up host in DNS
	ips, err := LookupIP(host)
	if err != nil {
		return nil, err
	} else if len(ips) == 0 {
		return nil, fmt.Errorf("No IP's returned from DNS lookup for %q", addr)
	}

	// Filter to v4 addresses if any present
	if hostLookupPreferV4 {
		var preferredIPs []net.IP
		for _, v := range ips {
			if v4 := v.To4(); v4 != nil {
				preferredIPs = append(preferredIPs, v4)
			}
		}
		if len(preferredIPs) != 0 {
			ips = preferredIPs
		}
	}

	for _, ip := range ips {
		hosts = append(hosts, &HostInfo{connectAddress: ip, port: port})
	}

	return hosts, nil
}

func shuffleHosts(hosts []*HostInfo) []*HostInfo {
	mutRandr.Lock()
	perm := randr.Perm(len(hosts))
	mutRandr.Unlock()
	shuffled := make([]*HostInfo, len(hosts))

	for i, host := range hosts {
		shuffled[perm[i]] = host
	}

	return shuffled
}

func (c *controlConn) shuffleDial(endpoints []*HostInfo) (*Conn, error) {
	// shuffle endpoints so not all drivers will connect to the same initial
	// node.
	shuffled := shuffleHosts(endpoints)

	cfg := *c.session.connCfg
	cfg.disableCoalesce = true

	var err error
	for _, host := range shuffled {
		var conn *Conn
		c.session.dial(host, &cfg, c)
		conn, err = c.session.connect(host, c)
		if err == nil {
			return conn, nil
		}

		Logger.Printf("gocql: unable to dial control conn %v: %v\n", host.ConnectAddress(), err)
	}

	return nil, err
}

// this is going to be version dependant and a nightmare to maintain :(
var protocolSupportRe = regexp.MustCompile(`the lowest supported version is \d+ and the greatest is (\d+)$`)

func parseProtocolFromError(err error) int {
	// I really wish this had the actual info in the error frame...
	matches := protocolSupportRe.FindAllStringSubmatch(err.Error(), -1)
	if len(matches) != 1 || len(matches[0]) != 2 {
		if verr, ok := err.(*protocolError); ok {
			return int(verr.frame.Header().version.version())
		}
		return 0
	}

	max, err := strconv.Atoi(matches[0][1])
	if err != nil {
		return 0
	}

	return max
}

func (c *controlConn) discoverProtocol(hosts []*HostInfo) (int, error) {
	hosts = shuffleHosts(hosts)

	connCfg := *c.session.connCfg
	connCfg.ProtoVersion = 4 // TODO: define maxProtocol

	handler := connErrorHandlerFn(func(c *Conn, err error, closed bool) {
		// we should never get here, but if we do it means we connected to a
		// host successfully which means our attempted protocol version worked
		if !closed {
			c.Close()
		}
	})

	var err error
	for _, host := range hosts {
		var conn *Conn
		conn, err = c.session.dial(host, &connCfg, handler)
		if conn != nil {
			conn.Close()
		}

		if err == nil {
			return connCfg.ProtoVersion, nil
		}

		if proto := parseProtocolFromError(err); proto > 0 {
			return proto, nil
		}
	}

	return 0, err
}

func (c *controlConn) connect(hosts []*HostInfo) error {
	if len(hosts) == 0 {
		return errors.New("control: no endpoints specified")
	}

	conn, err := c.shuffleDial(hosts)
	if err != nil {
		return fmt.Errorf("control: unable to connect to initial hosts: %v", err)
	}

	if err := c.setupConn(conn); err != nil {
		conn.Close()
		return fmt.Errorf("control: unable to setup connection: %v", err)
	}

	// we could fetch the initial ring here and update initial host data. So that
	// when we return from here we have a ring topology ready to go.

	go c.heartBeat()

	return nil
}

type connHost struct {
	conn *Conn
	host *HostInfo
}

func (c *controlConn) setupConn(conn *Conn) error {
	if err := c.registerEvents(conn); err != nil {
		conn.Close()
		return err
	}

	// TODO(zariel): do we need to fetch host info everytime
	// the control conn connects? Surely we have it cached?
	host, err := conn.localHostInfo(context.TODO())
	if err != nil {
		return err
	}

	ch := &connHost{
		conn: conn,
		host: host,
	}

	c.conn.Store(ch)
	c.session.handleNodeUp(host.ConnectAddress(), host.Port(), false)

	return nil
}

func (c *controlConn) registerEvents(conn *Conn) error {
	var events []string

	if !c.session.cfg.Events.DisableTopologyEvents {
		events = append(events, "TOPOLOGY_CHANGE")
	}
	if !c.session.cfg.Events.DisableNodeStatusEvents {
		events = append(events, "STATUS_CHANGE")
	}
	if !c.session.cfg.Events.DisableSchemaEvents {
		events = append(events, "SCHEMA_CHANGE")
	}

	if len(events) == 0 {
		return nil
	}

	framer, err := conn.exec(context.Background(),
		&writeRegisterFrame{
			events: events,
		}, nil)
	if err != nil {
		return err
	}

	frame, err := framer.parseFrame()
	if err != nil {
		return err
	} else if _, ok := frame.(*readyFrame); !ok {
		return fmt.Errorf("unexpected frame in response to register: got %T: %v\n", frame, frame)
	}

	return nil
}

func (c *controlConn) reconnect(refreshring bool) {
	if !atomic.CompareAndSwapInt32(&c.reconnecting, 0, 1) {
		return
	}
	defer atomic.StoreInt32(&c.reconnecting, 0)
	// TODO: simplify this function, use session.ring to get hosts instead of the
	// connection pool

	var host *HostInfo
	ch := c.getConn()
	if ch != nil {
		host = ch.host
		ch.conn.Close()
	}

	var newConn *Conn
	if host != nil {
		// try to connect to the old host
		conn, err := c.session.connect(host, c)
		if err != nil {
			// host is dead
			// TODO: this is replicated in a few places
			if c.session.cfg.ConvictionPolicy.AddFailure(err, host) {
				c.session.handleNodeDown(host.ConnectAddress(), host.Port())
			}
		} else {
			newConn = conn
		}
	}

	// TODO: should have our own round-robin for hosts so that we can try each
	// in succession and guarantee that we get a different host each time.
	if newConn == nil {
		host := c.session.ring.rrHost()
		if host == nil {
			c.connect(c.session.ring.endpoints)
			return
		}

		var err error
		newConn, err = c.session.connect(host, c)
		if err != nil {
			// TODO: add log handler for things like this
			return
		}
	}

	if err := c.setupConn(newConn); err != nil {
		newConn.Close()
		Logger.Printf("gocql: control unable to register events: %v\n", err)
		return
	}

	if refreshring {
		c.session.hostSource.refreshRing()
	}
}

func (c *controlConn) HandleError(conn *Conn, err error, closed bool) {
	if !closed {
		return
	}

	oldConn := c.getConn()
	if oldConn.conn != conn {
		return
	}

	c.reconnect(false)
}

func (c *controlConn) getConn() *connHost {
	return c.conn.Load().(*connHost)
}

func (c *controlConn) writeFrame(w frameWriter) (frame, error) {
	ch := c.getConn()
	if ch == nil {
		return nil, errNoControl
	}

	framer, err := ch.conn.exec(context.Background(), w, nil)
	if err != nil {
		return nil, err
	}

	return framer.parseFrame()
}

func (c *controlConn) withConnHost(fn func(*connHost) *Iter) *Iter {
	const maxConnectAttempts = 5
	connectAttempts := 0

	for i := 0; i < maxConnectAttempts; i++ {
		ch := c.getConn()
		if ch == nil {
			if connectAttempts > maxConnectAttempts {
				break
			}

			connectAttempts++

			c.reconnect(false)
			continue
		}

		return fn(ch)
	}

	return &Iter{err: errNoControl}
}

func (c *controlConn) withConn(fn func(*Conn) *Iter) *Iter {
	return c.withConnHost(func(ch *connHost) *Iter {
		return fn(ch.conn)
	})
}

// query will return nil if the connection is closed or nil
func (c *controlConn) query(statement string, values ...interface{}) (iter *Iter) {
	q := c.session.Query(statement, values...).Consistency(One).RoutingKey([]byte{}).Trace(nil)

	for {
		iter = c.withConn(func(conn *Conn) *Iter {
			return conn.executeQuery(context.TODO(), q)
		})

		if gocqlDebug && iter.err != nil {
			Logger.Printf("control: error executing %q: %v\n", statement, iter.err)
		}

		q.AddAttempts(1, c.getConn().host)
		if iter.err == nil || !c.retry.Attempt(q) {
			break
		}
	}

	return
}

func (c *controlConn) awaitSchemaAgreement() error {
	return c.withConn(func(conn *Conn) *Iter {
		return &Iter{err: conn.awaitSchemaAgreement(context.TODO())}
	}).err
}

func (c *controlConn) close() {
	if atomic.CompareAndSwapInt32(&c.started, 1, -1) {
		c.quit <- struct{}{}
	}

	ch := c.getConn()
	if ch != nil {
		ch.conn.Close()
	}
}

var errNoControl = errors.New("gocql: no control connection available")
<commit_msg>Removing redundant connection + dependant goroutines (#1278)<commit_after>package gocql

import (
	"context"
	crand "crypto/rand"
	"errors"
	"fmt"
	"math/rand"
	"net"
	"os"
	"regexp"
	"strconv"
	"sync"
	"sync/atomic"
	"time"
)

var (
	randr    *rand.Rand
	mutRandr sync.Mutex
)

func init() {
	b := make([]byte, 4)
	if _, err := crand.Read(b); err != nil {
		panic(fmt.Sprintf("unable to seed random number generator: %v", err))
	}

	randr = rand.New(rand.NewSource(int64(readInt(b))))
}

// Ensure that the atomic variable is aligned to a 64bit boundary
// so that atomic operations can be applied on 32bit architectures.
type controlConn struct {
	started      int32
	reconnecting int32

	session *Session
	conn    atomic.Value

	retry RetryPolicy

	quit chan struct{}
}

func createControlConn(session *Session) *controlConn {
	control := &controlConn{
		session: session,
		quit:    make(chan struct{}),
		retry:   &SimpleRetryPolicy{NumRetries: 3},
	}

	control.conn.Store((*connHost)(nil))

	return control
}

func (c *controlConn) heartBeat() {
	if !atomic.CompareAndSwapInt32(&c.started, 0, 1) {
		return
	}

	sleepTime := 1 * time.Second
	timer := time.NewTimer(sleepTime)
	defer timer.Stop()

	for {
		timer.Reset(sleepTime)

		select {
		case <-c.quit:
			return
		case <-timer.C:
		}

		resp, err := c.writeFrame(&writeOptionsFrame{})
		if err != nil {
			goto reconn
		}

		switch resp.(type) {
		case *supportedFrame:
			// Everything ok
			sleepTime = 5 * time.Second
			continue
		case error:
			goto reconn
		default:
			panic(fmt.Sprintf("gocql: unknown frame in response to options: %T", resp))
		}

	reconn:
		// try to connect a bit faster
		sleepTime = 1 * time.Second
		c.reconnect(true)
		continue
	}
}

var hostLookupPreferV4 = os.Getenv("GOCQL_HOST_LOOKUP_PREFER_V4") == "true"

func hostInfo(addr string, defaultPort int) ([]*HostInfo, error) {
	var port int
	host, portStr, err := net.SplitHostPort(addr)
	if err != nil {
		host = addr
		port = defaultPort
	} else {
		port, err = strconv.Atoi(portStr)
		if err != nil {
			return nil, err
		}
	}

	var hosts []*HostInfo

	// Check if host is a literal IP address
	if ip := net.ParseIP(host); ip != nil {
		hosts = append(hosts, &HostInfo{connectAddress: ip, port: port})
		return hosts, nil
	}

	// Look up host in DNS
	ips, err := LookupIP(host)
	if err != nil {
		return nil, err
	} else if len(ips) == 0 {
		return nil, fmt.Errorf("No IP's returned from DNS lookup for %q", addr)
	}

	// Filter to v4 addresses if any present
	if hostLookupPreferV4 {
		var preferredIPs []net.IP
		for _, v := range ips {
			if v4 := v.To4(); v4 != nil {
				preferredIPs = append(preferredIPs, v4)
			}
		}
		if len(preferredIPs) != 0 {
			ips = preferredIPs
		}
	}

	for _, ip := range ips {
		hosts = append(hosts, &HostInfo{connectAddress: ip, port: port})
	}

	return hosts, nil
}

func shuffleHosts(hosts []*HostInfo) []*HostInfo {
	mutRandr.Lock()
	perm := randr.Perm(len(hosts))
	mutRandr.Unlock()
	shuffled := make([]*HostInfo, len(hosts))

	for i, host := range hosts {
		shuffled[perm[i]] = host
	}

	return shuffled
}

func (c *controlConn) shuffleDial(endpoints []*HostInfo) (*Conn, error) {
	// shuffle endpoints so not all drivers will connect to the same initial
	// node.
	shuffled := shuffleHosts(endpoints)

	cfg := *c.session.connCfg
	cfg.disableCoalesce = true

	var err error
	for _, host := range shuffled {
		var conn *Conn
		conn, err = c.session.dial(host, &cfg, c)
		if err == nil {
			return conn, nil
		}

		Logger.Printf("gocql: unable to dial control conn %v: %v\n", host.ConnectAddress(), err)
	}

	return nil, err
}

// this is going to be version dependant and a nightmare to maintain :(
var protocolSupportRe = regexp.MustCompile(`the lowest supported version is \d+ and the greatest is (\d+)$`)

func parseProtocolFromError(err error) int {
	// I really wish this had the actual info in the error frame...
	matches := protocolSupportRe.FindAllStringSubmatch(err.Error(), -1)
	if len(matches) != 1 || len(matches[0]) != 2 {
		if verr, ok := err.(*protocolError); ok {
			return int(verr.frame.Header().version.version())
		}
		return 0
	}

	max, err := strconv.Atoi(matches[0][1])
	if err != nil {
		return 0
	}

	return max
}

func (c *controlConn) discoverProtocol(hosts []*HostInfo) (int, error) {
	hosts = shuffleHosts(hosts)

	connCfg := *c.session.connCfg
	connCfg.ProtoVersion = 4 // TODO: define maxProtocol

	handler := connErrorHandlerFn(func(c *Conn, err error, closed bool) {
		// we should never get here, but if we do it means we connected to a
		// host successfully which means our attempted protocol version worked
		if !closed {
			c.Close()
		}
	})

	var err error
	for _, host := range hosts {
		var conn *Conn
		conn, err = c.session.dial(host, &connCfg, handler)
		if conn != nil {
			conn.Close()
		}

		if err == nil {
			return connCfg.ProtoVersion, nil
		}

		if proto := parseProtocolFromError(err); proto > 0 {
			return proto, nil
		}
	}

	return 0, err
}

func (c *controlConn) connect(hosts []*HostInfo) error {
	if len(hosts) == 0 {
		return errors.New("control: no endpoints specified")
	}

	conn, err := c.shuffleDial(hosts)
	if err != nil {
		return fmt.Errorf("control: unable to connect to initial hosts: %v", err)
	}

	if err := c.setupConn(conn); err != nil {
		conn.Close()
		return fmt.Errorf("control: unable to setup connection: %v", err)
	}

	// we could fetch the initial ring here and update initial host data. So that
	// when we return from here we have a ring topology ready to go.

	go c.heartBeat()

	return nil
}

type connHost struct {
	conn *Conn
	host *HostInfo
}

func (c *controlConn) setupConn(conn *Conn) error {
	if err := c.registerEvents(conn); err != nil {
		conn.Close()
		return err
	}

	// TODO(zariel): do we need to fetch host info everytime
	// the control conn connects? Surely we have it cached?
	host, err := conn.localHostInfo(context.TODO())
	if err != nil {
		return err
	}

	ch := &connHost{
		conn: conn,
		host: host,
	}

	c.conn.Store(ch)
	c.session.handleNodeUp(host.ConnectAddress(), host.Port(), false)

	return nil
}

func (c *controlConn) registerEvents(conn *Conn) error {
	var events []string

	if !c.session.cfg.Events.DisableTopologyEvents {
		events = append(events, "TOPOLOGY_CHANGE")
	}
	if !c.session.cfg.Events.DisableNodeStatusEvents {
		events = append(events, "STATUS_CHANGE")
	}
	if !c.session.cfg.Events.DisableSchemaEvents {
		events = append(events, "SCHEMA_CHANGE")
	}

	if len(events) == 0 {
		return nil
	}

	framer, err := conn.exec(context.Background(),
		&writeRegisterFrame{
			events: events,
		}, nil)
	if err != nil {
		return err
	}

	frame, err := framer.parseFrame()
	if err != nil {
		return err
	} else if _, ok := frame.(*readyFrame); !ok {
		return fmt.Errorf("unexpected frame in response to register: got %T: %v\n", frame, frame)
	}

	return nil
}

func (c *controlConn) reconnect(refreshring bool) {
	if !atomic.CompareAndSwapInt32(&c.reconnecting, 0, 1) {
		return
	}
	defer atomic.StoreInt32(&c.reconnecting, 0)
	// TODO: simplify this function, use session.ring to get hosts instead of the
	// connection pool

	var host *HostInfo
	ch := c.getConn()
	if ch != nil {
		host = ch.host
		ch.conn.Close()
	}

	var newConn *Conn
	if host != nil {
		// try to connect to the old host
		conn, err := c.session.connect(host, c)
		if err != nil {
			// host is dead
			// TODO: this is replicated in a few places
			if c.session.cfg.ConvictionPolicy.AddFailure(err, host) {
				c.session.handleNodeDown(host.ConnectAddress(), host.Port())
			}
		} else {
			newConn = conn
		}
	}

	// TODO: should have our own round-robin for hosts so that we can try each
	// in succession and guarantee that we get a different host each time.
	if newConn == nil {
		host := c.session.ring.rrHost()
		if host == nil {
			c.connect(c.session.ring.endpoints)
			return
		}

		var err error
		newConn, err = c.session.connect(host, c)
		if err != nil {
			// TODO: add log handler for things like this
			return
		}
	}

	if err := c.setupConn(newConn); err != nil {
		newConn.Close()
		Logger.Printf("gocql: control unable to register events: %v\n", err)
		return
	}

	if refreshring {
		c.session.hostSource.refreshRing()
	}
}

func (c *controlConn) HandleError(conn *Conn, err error, closed bool) {
	if !closed {
		return
	}

	oldConn := c.getConn()
	if oldConn.conn != conn {
		return
	}

	c.reconnect(false)
}

func (c *controlConn) getConn() *connHost {
	return c.conn.Load().(*connHost)
}

func (c *controlConn) writeFrame(w frameWriter) (frame, error) {
	ch := c.getConn()
	if ch == nil {
		return nil, errNoControl
	}

	framer, err := ch.conn.exec(context.Background(), w, nil)
	if err != nil {
		return nil, err
	}

	return framer.parseFrame()
}

func (c *controlConn) withConnHost(fn func(*connHost) *Iter) *Iter {
	const maxConnectAttempts = 5
	connectAttempts := 0

	for i := 0; i < maxConnectAttempts; i++ {
		ch := c.getConn()
		if ch == nil {
			if connectAttempts > maxConnectAttempts {
				break
			}

			connectAttempts++

			c.reconnect(false)
			continue
		}

		return fn(ch)
	}

	return &Iter{err: errNoControl}
}

func (c *controlConn) withConn(fn func(*Conn) *Iter) *Iter {
	return c.withConnHost(func(ch *connHost) *Iter {
		return fn(ch.conn)
	})
}

// query will return nil if the connection is closed or nil
func (c *controlConn) query(statement string, values ...interface{}) (iter *Iter) {
	q := c.session.Query(statement, values...).Consistency(One).RoutingKey([]byte{}).Trace(nil)

	for {
		iter = c.withConn(func(conn *Conn) *Iter {
			return conn.executeQuery(context.TODO(), q)
		})

		if gocqlDebug && iter.err != nil {
			Logger.Printf("control: error executing %q: %v\n", statement, iter.err)
		}

		q.AddAttempts(1, c.getConn().host)
		if iter.err == nil || !c.retry.Attempt(q) {
			break
		}
	}

	return
}

func (c *controlConn) awaitSchemaAgreement() error {
	return c.withConn(func(conn *Conn) *Iter {
		return &Iter{err: conn.awaitSchemaAgreement(context.TODO())}
	}).err
}

func (c *controlConn) close() {
	if atomic.CompareAndSwapInt32(&c.started, 1, -1) {
		c.quit <- struct{}{}
	}

	ch := c.getConn()
	if ch != nil {
		ch.conn.Close()
	}
}

var errNoControl = errors.New("gocql: no control connection available")
<|endoftext|>


<commit_before>package repl

import (
	"fmt"
	"io"
	"os"

	"github.com/chzyer/readline"
	"github.com/mattn/go-shellwords"
)

type (
	// REPL is a read-eval-print loop used to create a simple, minimalistic,
	// easy-to-use command line interface for masterkey.
	REPL struct {
		prompt          string
		commands        map[string]Command
		prefixCompleter *readline.PrefixCompleter
		input           io.Reader
		output          io.Writer
		rl              *readline.Instance
		stopfunc        func()
	}

	// Command is a command that can be registered with the REPL. It consists
	// of a name, an action that is run when the name is input to the REPL, and
	// a usage string.
	Command struct {
		Name   string
		Action ActionFunc
		Usage  string
	}

	// ActionFunc defines the signature of an action associated with a command.
	// Actions take on parameter, a slice of strings, representing the arguments
	// passed to the command. Actions should return a string representing the
	// result of the action, or an error if the action fails.
	ActionFunc func([]string) (string, error)
)

// New instantiates a new REPL using the provided `prompt`.
func New(prompt string) *REPL {
	return &REPL{
		commands:        make(map[string]Command),
		prefixCompleter: readline.NewPrefixCompleter(readline.PcItem("exit"), readline.PcItem("help")),
		prompt:          prompt,
		input:           os.Stdin,
		output:          os.Stdout,
	}
}

// OnStop registers a function to be called when the REPL stops.
func (r *REPL) OnStop(sf func()) {
	r.stopfunc = sf
}

// Usage returns the usage for every command in the REPL.
func (r *REPL) Usage() string {
	printstring := ""
	for _, command := range r.commands {
		printstring += command.Usage + "\n"
	}
	return printstring
}

// AddCommand registers the command provided in `cmd` with the REPL.
func (r *REPL) AddCommand(cmd Command) {
	r.commands[cmd.Name] = cmd

	var completers []readline.PrefixCompleterInterface
	for name := range r.commands {
		completers = append(completers, readline.PcItem(name))
	}

	r.prefixCompleter = readline.NewPrefixCompleter(completers...)
}

// eval evaluates a line that was input to the REPL.
func (r *REPL) eval(line string) (string, error) {
	args, err := shellwords.Parse(line)
	if err != nil {
		return "", err
	}
	command := args[0]

	if command == "help" {
		return r.Usage(), nil
	}

	if command == "exit" {
		return "", r.rl.Close()
	}

	if command == "clear" {
		_, err := readline.ClearScreen(r.output)
		if err != nil {
			return "", err
		}
		return "terminal cleared", nil
	}

	cmd, exists := r.commands[command]
	if !exists {
		return "", fmt.Errorf("command not recognized. Type `help` for a list of commands.")
	}

	res, err := cmd.Action(args[1:])
	if err != nil {
		return "", err
	}

	return res, nil
}

// Loop starts the Read-Eval-Print loop.
func (r *REPL) Loop() error {
	rl, err := readline.NewEx(&readline.Config{
		Prompt:       r.prompt,
		AutoComplete: r.prefixCompleter,
	})

	if err != nil {
		return err
	}
	defer rl.Close()
	r.rl = rl

	for {
		line, err := r.rl.Readline()
		if err != nil {
			if err == readline.ErrInterrupt && r.stopfunc != nil {
				r.stopfunc()
			}
			break
		}
		if line != "" {
			res, err := r.eval(line)
			if err != nil {
				fmt.Fprintln(r.output, err.Error())
				continue
			}
			fmt.Fprint(r.output, res)
		}
	}
	return nil
}
<commit_msg>register default commands in repl.New() instead of eval<commit_after>package repl

import (
	"fmt"
	"io"
	"os"

	"github.com/chzyer/readline"
	"github.com/mattn/go-shellwords"
)

type (
	// REPL is a read-eval-print loop used to create a simple, minimalistic,
	// easy-to-use command line interface for masterkey.
	REPL struct {
		prompt          string
		commands        map[string]Command
		prefixCompleter *readline.PrefixCompleter
		input           io.Reader
		output          io.Writer
		rl              *readline.Instance
		stopfunc        func()
	}

	// Command is a command that can be registered with the REPL. It consists
	// of a name, an action that is run when the name is input to the REPL, and
	// a usage string.
	Command struct {
		Name   string
		Action ActionFunc
		Usage  string
	}

	// ActionFunc defines the signature of an action associated with a command.
	// Actions take on parameter, a slice of strings, representing the arguments
	// passed to the command. Actions should return a string representing the
	// result of the action, or an error if the action fails.
	ActionFunc func([]string) (string, error)
)

// New instantiates a new REPL using the provided `prompt`.
func New(prompt string) *REPL {
	r := &REPL{
		commands: make(map[string]Command),
		prompt:   prompt,
		input:    os.Stdin,
		output:   os.Stdout,
	}

	// Add default commands clear, exit, and help
	r.AddCommand(Command{
		Name:  "help",
		Usage: "help: displays available commands and their usage",
		Action: func(args []string) (string, error) {
			return r.Usage(), nil
		},
	})

	r.AddCommand(Command{
		Name:  "exit",
		Usage: "exit: exit the interactive prompt",
		Action: func(args []string) (string, error) {
			return "exiting", r.rl.Close()
		},
	})

	r.AddCommand(Command{
		Name:  "clear",
		Usage: "clear: clear the terminal",
		Action: func(args []string) (string, error) {
			_, err := readline.ClearScreen(r.output)
			if err != nil {
				return "", err
			}
			return "cleared terminal", nil
		},
	})

	return r
}

// OnStop registers a function to be called when the REPL stops.
func (r *REPL) OnStop(sf func()) {
	r.stopfunc = sf
}

// Usage returns the usage for every command in the REPL.
func (r *REPL) Usage() string {
	printstring := ""
	for _, command := range r.commands {
		printstring += command.Usage + "\n"
	}
	return printstring
}

// AddCommand registers the command provided in `cmd` with the REPL.
func (r *REPL) AddCommand(cmd Command) {
	r.commands[cmd.Name] = cmd

	var completers []readline.PrefixCompleterInterface
	for name := range r.commands {
		completers = append(completers, readline.PcItem(name))
	}

	r.prefixCompleter = readline.NewPrefixCompleter(completers...)
}

// eval evaluates a line that was input to the REPL.
func (r *REPL) eval(line string) (string, error) {
	args, err := shellwords.Parse(line)
	if err != nil {
		return "", err
	}
	command := args[0]

	cmd, exists := r.commands[command]
	if !exists {
		return "", fmt.Errorf("command not recognized. Type `help` for a list of commands.")
	}

	res, err := cmd.Action(args[1:])
	if err != nil {
		return "", err
	}

	return res, nil
}

// Loop starts the Read-Eval-Print loop.
func (r *REPL) Loop() error {
	rl, err := readline.NewEx(&readline.Config{
		Prompt:       r.prompt,
		AutoComplete: r.prefixCompleter,
	})

	if err != nil {
		return err
	}
	defer rl.Close()
	r.rl = rl

	for {
		line, err := r.rl.Readline()
		if err != nil {
			if err == readline.ErrInterrupt && r.stopfunc != nil {
				r.stopfunc()
			}
			break
		}
		if line != "" {
			res, err := r.eval(line)
			if err != nil {
				fmt.Fprintln(r.output, err.Error())
				continue
			}
			fmt.Fprint(r.output, res)
		}
	}
	return nil
}
<|endoftext|>


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

package main

import (
	"os"
	"time"

	"github.com/codegangsta/cli"
	mountpkg "github.com/googlecloudplatform/gcsfuse/mount"
)

func newApp() (app *cli.App) {
	app = &cli.App{
		Name:          "gcsfuse",
		Usage:         "Mount a GCS bucket locally",
		ArgumentUsage: "bucket mountpoint",
		HideHelp:      true,
		HideVersion:   true,
		Writer:        os.Stderr,
		Flags: []cli.Flag{

			cli.BoolFlag{
				Name:  "help, h",
				Usage: "Print this help text and exit successfuly.",
			},

			/////////////////////////
			// File system
			/////////////////////////

			cli.StringSliceFlag{
				Name:  "o",
				Usage: "Additional system-specific mount options. Be careful!",
			},

			cli.IntFlag{
				Name:        "dir-mode",
				Value:       0755,
				Usage:       "Permissions bits for directories. (default: 0755)",
				HideDefault: true,
			},

			cli.IntFlag{
				Name:        "file-mode",
				Value:       0644,
				Usage:       "Permission bits for files (default: 0644)",
				HideDefault: true,
			},

			cli.IntFlag{
				Name:        "uid",
				Value:       -1,
				HideDefault: true,
				Usage:       "UID owner of all inodes.",
			},

			cli.IntFlag{
				Name:        "gid",
				Value:       -1,
				HideDefault: true,
				Usage:       "GID owner of all inodes.",
			},

			cli.BoolFlag{
				Name: "implicit-dirs",
				Usage: "Implicitly define directories based on content. See" +
					"docs/semantics.md",
			},

			/////////////////////////
			// GCS
			/////////////////////////

			cli.StringFlag{
				Name:        "key-file",
				Value:       "",
				HideDefault: true,
				Usage: "Path to JSON key file for use with GCS. " +
					"(default: none, Google application default credentials used)",
			},

			cli.Float64Flag{
				Name:  "limit-bytes-per-sec",
				Value: -1,
				Usage: "Bandwidth limit for reading data, measured over a 30-second " +
					"window. (use -1 for no limit)",
			},

			cli.Float64Flag{
				Name:  "limit-ops-per-sec",
				Value: 5.0,
				Usage: "Operations per second limit, measured over a 30-second window " +
					"(use -1 for no limit)",
			},

			/////////////////////////
			// Tuning
			/////////////////////////

			cli.DurationFlag{
				Name:  "stat-cache-ttl",
				Value: time.Minute,
				Usage: "How long to cache StatObject results from GCS.",
			},

			cli.DurationFlag{
				Name:  "type-cache-ttl",
				Value: time.Minute,
				Usage: "How long to cache name -> file/dir mappings in directory " +
					"inodes.",
			},

			cli.IntFlag{
				Name:  "gcs-chunk-size",
				Value: 1 << 24,
				Usage: "Max chunk size for loading GCS objects.",
			},

			cli.StringFlag{
				Name:        "temp-dir",
				Value:       "",
				HideDefault: true,
				Usage: "Temporary directory for local GCS object copies. " +
					"(default: system default, likely /tmp)",
			},

			cli.IntFlag{
				Name:  "temp-dir-bytes",
				Value: 1 << 31,
				Usage: "Size limit of the temporary directory.",
			},

			/////////////////////////
			// Debugging
			/////////////////////////

			cli.BoolFlag{
				Name:  "debug_cpu_profile",
				Usage: "Write a 10-second CPU profile to /tmp on SIGHUP.",
			},

			cli.BoolFlag{
				Name:  "debug_fuse",
				Usage: "Enable fuse-related debugging output.",
			},

			cli.BoolFlag{
				Name:  "debug_gcs",
				Usage: "Print GCS request and timing information.",
			},

			cli.BoolFlag{
				Name:  "debug_http",
				Usage: "Dump HTTP requests and responses to/from GCS.",
			},

			cli.BoolFlag{
				Name:  "debug_invariants",
				Usage: "Panic when internal invariants are violated.",
			},

			cli.BoolFlag{
				Name:  "debug_mem_profile",
				Usage: "Write a 10-second memory profile to /tmp on SIGHUP.",
			},
		},
	}

	return
}

type flagStorage struct {
	// File system
	MountOptions map[string]string
	DirMode      os.FileMode
	FileMode     os.FileMode
	Uid          int64
	Gid          int64
	ImplicitDirs bool

	// GCS
	KeyFile                            string
	EgressBandwidthLimitBytesPerSecond float64
	OpRateLimitHz                      float64

	// Tuning
	StatCacheTTL time.Duration
	TypeCacheTTL time.Duration
	GCSChunkSize uint64
	TempDir      string
	TempDirLimit int64

	// Debugging
	DebugCPUProfile bool
	DebugFuse       bool
	DebugGCS        bool
	DebugHTTP       bool
	DebugInvariants bool
	DebugMemProfile bool
}

// Add the flags accepted by run to the supplied flag set, returning the
// variables into which the flags will parse.
func populateFlags(c *cli.Context) (flags *flagStorage) {
	flags = &flagStorage{
		// File system
		MountOptions: make(map[string]string),
		DirMode:      os.FileMode(c.Int("dir-mode")),
		FileMode:     os.FileMode(c.Int("file-mode")),
		Uid:          int64(c.Int("uid")),
		Gid:          int64(c.Int("gid")),

		// GCS,
		KeyFile: c.String("key-file"),
		EgressBandwidthLimitBytesPerSecond: c.Float64("limit-bytes-per-sec"),
		OpRateLimitHz:                      c.Float64("limit-ops-per-sec"),

		// Tuning,
		StatCacheTTL: c.Duration("stat-cache-ttl"),
		TypeCacheTTL: c.Duration("type-cache-ttl"),
		GCSChunkSize: uint64(c.Int("gcs-chunk-size")),
		TempDir:      c.String("temp-dir"),
		TempDirLimit: int64(c.Int("temp-dir-bytes")),
		ImplicitDirs: c.Bool("implicit-dirs"),

		// Debugging,
		DebugCPUProfile: c.Bool("debug_cpu_profile"),
		DebugFuse:       c.Bool("debug_fuse"),
		DebugGCS:        c.Bool("debug_gcs"),
		DebugHTTP:       c.Bool("debug_http"),
		DebugInvariants: c.Bool("debug_invariants"),
		DebugMemProfile: c.Bool("debug_mem_profile"),
	}

	// Handle the repeated "-o" flag.
	for _, o := range c.StringSlice("o") {
		mountpkg.ParseOptions(flags.MountOptions, o)
	}

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

package main

import (
	"os"
	"time"

	"github.com/codegangsta/cli"
	mountpkg "github.com/googlecloudplatform/gcsfuse/mount"
)

func newApp() (app *cli.App) {
	app = &cli.App{
		Name:  "gcsfuse",
		Usage: "Mount a GCS bucket locally",
		// TODO(jacobsa): ArgumentUsage: "bucket mountpoint",
		HideHelp:    true,
		HideVersion: true,
		Writer:      os.Stderr,
		Flags: []cli.Flag{

			cli.BoolFlag{
				Name:  "help, h",
				Usage: "Print this help text and exit successfuly.",
			},

			/////////////////////////
			// File system
			/////////////////////////

			cli.StringSliceFlag{
				Name:  "o",
				Usage: "Additional system-specific mount options. Be careful!",
			},

			cli.IntFlag{
				Name:  "dir-mode",
				Value: 0755,
				Usage: "Permissions bits for directories. (default: 0755)",
				// TODO(jacobsa): HideDefault: true,
			},

			cli.IntFlag{
				Name:  "file-mode",
				Value: 0644,
				Usage: "Permission bits for files (default: 0644)",
				// TODO(jacobsa): HideDefault: true,
			},

			cli.IntFlag{
				Name:  "uid",
				Value: -1,
				// TODO(jacobsa): HideDefault: true,
				Usage: "UID owner of all inodes.",
			},

			cli.IntFlag{
				Name:  "gid",
				Value: -1,
				// TODO(jacobsa): HideDefault: true,
				Usage: "GID owner of all inodes.",
			},

			cli.BoolFlag{
				Name: "implicit-dirs",
				Usage: "Implicitly define directories based on content. See" +
					"docs/semantics.md",
			},

			/////////////////////////
			// GCS
			/////////////////////////

			cli.StringFlag{
				Name:  "key-file",
				Value: "",
				// TODO(jacobsa): HideDefault: true,
				Usage: "Path to JSON key file for use with GCS. " +
					"(default: none, Google application default credentials used)",
			},

			cli.Float64Flag{
				Name:  "limit-bytes-per-sec",
				Value: -1,
				Usage: "Bandwidth limit for reading data, measured over a 30-second " +
					"window. (use -1 for no limit)",
			},

			cli.Float64Flag{
				Name:  "limit-ops-per-sec",
				Value: 5.0,
				Usage: "Operations per second limit, measured over a 30-second window " +
					"(use -1 for no limit)",
			},

			/////////////////////////
			// Tuning
			/////////////////////////

			cli.DurationFlag{
				Name:  "stat-cache-ttl",
				Value: time.Minute,
				Usage: "How long to cache StatObject results from GCS.",
			},

			cli.DurationFlag{
				Name:  "type-cache-ttl",
				Value: time.Minute,
				Usage: "How long to cache name -> file/dir mappings in directory " +
					"inodes.",
			},

			cli.IntFlag{
				Name:  "gcs-chunk-size",
				Value: 1 << 24,
				Usage: "Max chunk size for loading GCS objects.",
			},

			cli.StringFlag{
				Name:  "temp-dir",
				Value: "",
				// TODO(jacobsa): HideDefault: true,
				Usage: "Temporary directory for local GCS object copies. " +
					"(default: system default, likely /tmp)",
			},

			cli.IntFlag{
				Name:  "temp-dir-bytes",
				Value: 1 << 31,
				Usage: "Size limit of the temporary directory.",
			},

			/////////////////////////
			// Debugging
			/////////////////////////

			cli.BoolFlag{
				Name:  "debug_cpu_profile",
				Usage: "Write a 10-second CPU profile to /tmp on SIGHUP.",
			},

			cli.BoolFlag{
				Name:  "debug_fuse",
				Usage: "Enable fuse-related debugging output.",
			},

			cli.BoolFlag{
				Name:  "debug_gcs",
				Usage: "Print GCS request and timing information.",
			},

			cli.BoolFlag{
				Name:  "debug_http",
				Usage: "Dump HTTP requests and responses to/from GCS.",
			},

			cli.BoolFlag{
				Name:  "debug_invariants",
				Usage: "Panic when internal invariants are violated.",
			},

			cli.BoolFlag{
				Name:  "debug_mem_profile",
				Usage: "Write a 10-second memory profile to /tmp on SIGHUP.",
			},
		},
	}

	return
}

type flagStorage struct {
	// File system
	MountOptions map[string]string
	DirMode      os.FileMode
	FileMode     os.FileMode
	Uid          int64
	Gid          int64
	ImplicitDirs bool

	// GCS
	KeyFile                            string
	EgressBandwidthLimitBytesPerSecond float64
	OpRateLimitHz                      float64

	// Tuning
	StatCacheTTL time.Duration
	TypeCacheTTL time.Duration
	GCSChunkSize uint64
	TempDir      string
	TempDirLimit int64

	// Debugging
	DebugCPUProfile bool
	DebugFuse       bool
	DebugGCS        bool
	DebugHTTP       bool
	DebugInvariants bool
	DebugMemProfile bool
}

// Add the flags accepted by run to the supplied flag set, returning the
// variables into which the flags will parse.
func populateFlags(c *cli.Context) (flags *flagStorage) {
	flags = &flagStorage{
		// File system
		MountOptions: make(map[string]string),
		DirMode:      os.FileMode(c.Int("dir-mode")),
		FileMode:     os.FileMode(c.Int("file-mode")),
		Uid:          int64(c.Int("uid")),
		Gid:          int64(c.Int("gid")),

		// GCS,
		KeyFile: c.String("key-file"),
		EgressBandwidthLimitBytesPerSecond: c.Float64("limit-bytes-per-sec"),
		OpRateLimitHz:                      c.Float64("limit-ops-per-sec"),

		// Tuning,
		StatCacheTTL: c.Duration("stat-cache-ttl"),
		TypeCacheTTL: c.Duration("type-cache-ttl"),
		GCSChunkSize: uint64(c.Int("gcs-chunk-size")),
		TempDir:      c.String("temp-dir"),
		TempDirLimit: int64(c.Int("temp-dir-bytes")),
		ImplicitDirs: c.Bool("implicit-dirs"),

		// Debugging,
		DebugCPUProfile: c.Bool("debug_cpu_profile"),
		DebugFuse:       c.Bool("debug_fuse"),
		DebugGCS:        c.Bool("debug_gcs"),
		DebugHTTP:       c.Bool("debug_http"),
		DebugInvariants: c.Bool("debug_invariants"),
		DebugMemProfile: c.Bool("debug_mem_profile"),
	}

	// Handle the repeated "-o" flag.
	for _, o := range c.StringSlice("o") {
		mountpkg.ParseOptions(flags.MountOptions, o)
	}

	return
}
<|endoftext|>


<commit_before>package repo

import (
	"crypto/rand"
	"io"
	"os"
	"path/filepath"

	"code.google.com/p/go-uuid/uuid"
	log "github.com/Sirupsen/logrus"
	"github.com/disorganizer/brig/repo/config"
	"github.com/disorganizer/brig/repo/global"
	"github.com/disorganizer/brig/store"
	"github.com/disorganizer/brig/util"
	logutil "github.com/disorganizer/brig/util/log"
	ipfsconfig "github.com/ipfs/go-ipfs/repo/config"
	"github.com/ipfs/go-ipfs/repo/fsrepo"
)

// NewRepository creates a new repository at filesystem level
// and returns a Repository interface
func NewRepository(jid, pwd, folder string) (*Repository, error) {
	absFolderPath, err := filepath.Abs(folder)
	if err != nil {
		return nil, err
	}

	if _, err = os.Stat(absFolderPath); os.IsExist(err) {
		return nil, err
	}

	if err := createRepositoryTree(absFolderPath); err != nil {
		return nil, err
	}

	cfg := config.CreateDefaultConfig()
	minilockID, err := GenerateMinilockID(jid, pwd)
	if err != nil {
		return nil, err
	}

	configDefaults := map[string]interface{}{
		"repository.jid":  jid,
		"repository.uuid": uuid.NewRandom().String(),
		"repository.mid":  minilockID,
		"ipfs.path":       filepath.Join(absFolderPath, ".brig", "ipfs"),
	}

	for key, value := range configDefaults {
		if err = cfg.Set(key, value); err != nil {
			return nil, err
		}
	}

	configPath := filepath.Join(absFolderPath, ".brig", "config")
	if _, err := config.SaveConfig(configPath, cfg); err != nil {
		return nil, err
	}

	return LoadRepository(pwd, absFolderPath)
}

// CloneRepository clones a brig repository in a git like way
func CloneRepository() *Repository {
	return nil
}

// LoadRepository load a brig repository from a given folder.
func LoadRepository(pwd, folder string) (*Repository, error) {
	absFolderPath, err := filepath.Abs(folder)
	if err != nil {
		return nil, err
	}

	brigPath := filepath.Join(absFolderPath, ".brig")
	cfg, err := config.LoadConfig(filepath.Join(brigPath, "config"))
	if err != nil {
		return nil, err
	}

	configValues := map[string]string{
		"repository.jid":  "",
		"repository.mid":  "",
		"repository.uuid": "",
	}

	for key := range configValues {
		configValues[key], err = cfg.String(key)
		if err != nil {
			return nil, err
		}
	}

	// Init the global repo (similar to .gitconfig)
	globalRepo, err := global.New()
	if err != nil {
		return nil, err
	}

	globalRepo.AddRepo(global.RepoListEntry{
		UniqueID:   configValues["repository.uuid"],
		RepoPath:   folder,
		DaemonPort: 6666,
		IpfsPort:   4001,
	})

	store, err := store.Open(brigPath)
	if err != nil {
		return nil, err
	}

	repo := Repository{
		Jid:            configValues["repository.jid"],
		Mid:            configValues["repository.mid"],
		Folder:         absFolderPath,
		InternalFolder: brigPath,
		UniqueID:       configValues["repository.uuid"],
		Config:         cfg,
		globalRepo:     globalRepo,
		Store:          store,
		Password:       pwd,
	}

	return &repo, nil
}

func createRepositoryTree(absFolderPath string) error {
	if err := os.Mkdir(absFolderPath, 0755); err != nil {
		return err
	}

	brigPath := filepath.Join(absFolderPath, ".brig")
	if err := os.Mkdir(brigPath, 0755); err != nil {
		return err
	}

	ipfsPath := filepath.Join(brigPath, "ipfs")
	if err := os.Mkdir(ipfsPath, 0755); err != nil {
		return err
	}

	empties := []string{"index.bolt", "otr.key", "otr.buddies"}
	for _, empty := range empties {
		fullPath := filepath.Join(brigPath, empty)
		if err := util.Touch(fullPath); err != nil {
			return err
		}
	}

	// Make the key larger than needed:
	if err := createMasterKey(brigPath, 1024); err != nil {
		return err
	}

	return CreateIpfsRepo(ipfsPath)
}

func createMasterKey(brigPath string, keySize int) error {
	keyPath := filepath.Join(brigPath, "master.key")
	fd, err := os.OpenFile(keyPath, os.O_CREATE|os.O_WRONLY, 0755)
	if err != nil {
		return err
	}

	defer util.Closer(fd)

	if _, err := io.CopyN(fd, rand.Reader, int64(keySize/8)); err != nil {
		return err
	}

	return nil
}

func CreateIpfsRepo(ipfsRootPath string) error {
	logger := &logutil.Writer{Level: log.InfoLevel}
	cfg, err := ipfsconfig.Init(logger, 2048)
	if err != nil {
		return err
	}

	if err := fsrepo.Init(ipfsRootPath, cfg); err != nil {
		return err
	}

	return nil
}
<commit_msg>Remove code.google imports.<commit_after>package repo

import (
	"crypto/rand"
	"io"
	"os"
	"path/filepath"

	log "github.com/Sirupsen/logrus"
	"github.com/disorganizer/brig/repo/config"
	"github.com/disorganizer/brig/repo/global"
	"github.com/disorganizer/brig/store"
	"github.com/disorganizer/brig/util"
	logutil "github.com/disorganizer/brig/util/log"
	ipfsconfig "github.com/ipfs/go-ipfs/repo/config"
	"github.com/ipfs/go-ipfs/repo/fsrepo"
	"github.com/wayn3h0/go-uuid"
)

// NewRepository creates a new repository at filesystem level
// and returns a Repository interface
func NewRepository(jid, pwd, folder string) (*Repository, error) {
	absFolderPath, err := filepath.Abs(folder)
	if err != nil {
		return nil, err
	}

	if _, err = os.Stat(absFolderPath); os.IsExist(err) {
		return nil, err
	}

	if err := createRepositoryTree(absFolderPath); err != nil {
		return nil, err
	}

	cfg := config.CreateDefaultConfig()
	minilockID, err := GenerateMinilockID(jid, pwd)
	if err != nil {
		return nil, err
	}

	repoUUID, err := uuid.NewRandom()
	if err != nil {
		return nil, err
	}

	configDefaults := map[string]interface{}{
		"repository.jid":  jid,
		"repository.uuid": repoUUID.String(),
		"repository.mid":  minilockID,
		"ipfs.path":       filepath.Join(absFolderPath, ".brig", "ipfs"),
	}

	for key, value := range configDefaults {
		if err = cfg.Set(key, value); err != nil {
			return nil, err
		}
	}

	configPath := filepath.Join(absFolderPath, ".brig", "config")
	if _, err := config.SaveConfig(configPath, cfg); err != nil {
		return nil, err
	}

	return LoadRepository(pwd, absFolderPath)
}

// CloneRepository clones a brig repository in a git like way
func CloneRepository() *Repository {
	return nil
}

// LoadRepository load a brig repository from a given folder.
func LoadRepository(pwd, folder string) (*Repository, error) {
	absFolderPath, err := filepath.Abs(folder)
	if err != nil {
		return nil, err
	}

	brigPath := filepath.Join(absFolderPath, ".brig")
	cfg, err := config.LoadConfig(filepath.Join(brigPath, "config"))
	if err != nil {
		return nil, err
	}

	configValues := map[string]string{
		"repository.jid":  "",
		"repository.mid":  "",
		"repository.uuid": "",
	}

	for key := range configValues {
		configValues[key], err = cfg.String(key)
		if err != nil {
			return nil, err
		}
	}

	// Init the global repo (similar to .gitconfig)
	globalRepo, err := global.New()
	if err != nil {
		return nil, err
	}

	globalRepo.AddRepo(global.RepoListEntry{
		UniqueID:   configValues["repository.uuid"],
		RepoPath:   folder,
		DaemonPort: 6666,
		IpfsPort:   4001,
	})

	store, err := store.Open(brigPath)
	if err != nil {
		return nil, err
	}

	repo := Repository{
		Jid:            configValues["repository.jid"],
		Mid:            configValues["repository.mid"],
		Folder:         absFolderPath,
		InternalFolder: brigPath,
		UniqueID:       configValues["repository.uuid"],
		Config:         cfg,
		globalRepo:     globalRepo,
		Store:          store,
		Password:       pwd,
	}

	return &repo, nil
}

func createRepositoryTree(absFolderPath string) error {
	if err := os.Mkdir(absFolderPath, 0755); err != nil {
		return err
	}

	brigPath := filepath.Join(absFolderPath, ".brig")
	if err := os.Mkdir(brigPath, 0755); err != nil {
		return err
	}

	ipfsPath := filepath.Join(brigPath, "ipfs")
	if err := os.Mkdir(ipfsPath, 0755); err != nil {
		return err
	}

	empties := []string{"index.bolt", "otr.key", "otr.buddies"}
	for _, empty := range empties {
		fullPath := filepath.Join(brigPath, empty)
		if err := util.Touch(fullPath); err != nil {
			return err
		}
	}

	// Make the key larger than needed:
	if err := createMasterKey(brigPath, 1024); err != nil {
		return err
	}

	return CreateIpfsRepo(ipfsPath)
}

func createMasterKey(brigPath string, keySize int) error {
	keyPath := filepath.Join(brigPath, "master.key")
	fd, err := os.OpenFile(keyPath, os.O_CREATE|os.O_WRONLY, 0755)
	if err != nil {
		return err
	}

	defer util.Closer(fd)

	if _, err := io.CopyN(fd, rand.Reader, int64(keySize/8)); err != nil {
		return err
	}

	return nil
}

func CreateIpfsRepo(ipfsRootPath string) error {
	logger := &logutil.Writer{Level: log.InfoLevel}
	cfg, err := ipfsconfig.Init(logger, 2048)
	if err != nil {
		return err
	}

	if err := fsrepo.Init(ipfsRootPath, cfg); err != nil {
		return err
	}

	return nil
}
<|endoftext|>


<commit_before>package repo

import (
	"html/template"
	"log"
	"sort"

	"github.com/longkai/xiaolongtongxue.com/helper"
)

// Repo the documents repository.
type Repo interface {
	List(since string, size int) Docs
	Get(path string) (Doc, error)
	Del(path string)
	Put(path string)
	Post(path string)
	Batch(adds, mods, dels []string)
}

type entry struct {
	ready chan struct{}
	val   template.HTML
	err   error
}

func (e *entry) call(path string, r Renderer) {
	val, err := helper.Try(3, func() (interface{}, error) {
		return r.Render(path)
	})
	e.val, e.err = val.(template.HTML), err
	close(e.ready)
}

type listReq struct {
	path string
	size int
	resp chan Docs
}

type batchReq struct {
	adds, mods, dels []string
}

type getReq struct {
	path string
	resp chan getResp
}

type getResp struct {
	doc Doc
	err error
}

type reqs struct {
	get   chan getReq
	post  chan Docs
	list  chan listReq
	batch chan batchReq
}

// DocRepo documents repository implements.
type DocRepo struct {
	reqs

	dir       Dir
	renderer  Renderer
	processor Processor
	visitors  []Visitor

	docs  Docs              // More read, less write.
	index map[string]int    // Fast lookup.
	cache map[string]*entry // Rendering cache.
}

func (r *DocRepo) loop() {
	for {
		select {
		case req := <-r.reqs.get:
			r.get(req)
		case doc := <-r.reqs.post:
			r.post(doc)
		case req := <-r.reqs.list:
			r.list(req)
		case req := <-r.reqs.batch:
			r.batch(req)
		}
	}
}

func (r *DocRepo) batch(req batchReq) {
	rm := func(paths []string) {
		for _, path := range paths {
			if _, ok := r.index[path]; ok {
				delete(r.index, path)
				delete(r.cache, path)
			}
		}
	}
	// Re-process adds and mods.
	if plen := len(req.adds) + len(req.mods); plen > 0 {
		combine := make([]string, plen)
		copy(combine, req.adds)
		copy(combine[len(req.adds):], req.mods)
		go func() {
			docs := r.processor.Process(combine...)
			// Post process.
			for _, v := range r.visitors {
				v.Visit(docs, map[int]interface{}{
					Adds: req.adds,
					Mods: req.mods,
				})
			}
		}()
	}
	// Deletions.
	rm(req.dels)
	// Modifications.
	rm(req.mods)
	// Rearrangement, strip dels and mods, the order still remains.
	idx := 0
	tmp := make(Docs, 0, len(r.docs))
	for _, doc := range r.docs {
		// Pick those who are not deleted or modified.
		if _, ok := r.index[doc.URL]; ok {
			tmp = append(tmp, doc)
			r.index[doc.URL] = idx
			idx++
		}
	}
	r.docs = tmp
	// Hence, the time complexity is O(n).
}

func (r *DocRepo) get(req getReq) {
	i, ok := r.index[req.path]
	if !ok {
		go func() { req.resp <- getResp{Doc{}, NotFoundError(req.path)} }()
		return
	}

	doc := r.docs[i]

	// A hidden doc has no newer/older navigation.
	if !doc.Hide {
		if docs := r.docs.travel(i+1, 1, true, r.docs.filterHidden); docs.Len() > 0 {
			doc.Older = docs[0].URL
		}
		if docs := r.docs.travel(i-1, 1, false, r.docs.filterHidden); docs.Len() > 0 {
			doc.Newer = docs[0].URL
		}
	}

	e := r.cache[req.path]
	if e == nil {
		// Cache misses.
		e = &entry{ready: make(chan struct{})}
		go e.call(doc.Path, r.renderer)
		r.cache[req.path] = e
	}

	go func() {
		<-e.ready
		doc.Body = e.val
		req.resp <- getResp{doc, e.err}
	}()
}

func (r *DocRepo) post(docs Docs) {
	oldSize := r.docs.Len()
	// If a some of a newly doc has been existed already, replace them.
	// i.e., avoid duplication.
	for _, d := range docs {
		if i, ok := r.index[d.URL]; ok {
			log.Printf("replace %q with %q", r.docs[i].URL, d.URL)
			r.docs[i] = d          // Replace the old one.
			delete(r.cache, d.URL) // Clear its rendering cache, if any.
		} else {
			r.docs = append(r.docs, d) // Append the new one.
			// If there are multiple docs have a same URL,
			// the last one will be kept.
			r.index[d.URL] = r.docs.Len() - 1
		}
	}
	sort.Sort(r.docs)

	// Rebuild index.
	index := r.index
	for i, d := range r.docs {
		index[d.URL] = i
	}

	log.Printf("receive %d docs, len %d to %d", docs.Len(), oldSize, r.docs.Len())
}

func (r *DocRepo) list(req listReq) {
	i, ok := r.index[req.path]
	if !ok {
		i = 0 // If not found any match, start from 0.
	} else {
		i++ // Skip the current one.
	}

	res := r.docs.travel(i, req.size, true, r.docs.filterHidden)

	go func() { req.resp <- res }()
}

// List articles since a specific path, excluded.
func (r *DocRepo) List(since string, size int) Docs {
	resp := make(chan Docs)
	r.reqs.list <- listReq{since, size, resp}
	return <-resp
}

// Get a document for the path.
func (r *DocRepo) Get(path string) (Doc, error) {
	// Read only index, fast indexing without channel synchronization.
	// It's safe since lookup success, the go-routine will lookup again.
	// Hence, when lookup fail, maybe it's just removed or never exists.
	if _, ok := r.index[path]; !ok {
		return Doc{}, NotFoundError(path)
	}
	resp := make(chan getResp)
	r.reqs.get <- getReq{path, resp}
	v := <-resp
	return v.doc, v.err
}

// Del a document for the path.
func (r *DocRepo) Del(path string) {
	r.Batch(nil, nil, []string{path})
}

// Put revalidate a document.
func (r *DocRepo) Put(path string) {
	r.Batch(nil, []string{path}, nil)
}

// Post publish the path for documents.
// This method should be called when you start the application.
// Since it won't call the visitors let them do post process.
func (r *DocRepo) Post(path string) {
	log.Printf("post %s", path)
	r.processor.Process(path)
}

// Batch additions, modifications and deletions into a single request.
func (r *DocRepo) Batch(adds, mods, dels []string) {
	// Git only tracks files, hence, all of the slice are files path.

	// `git mv a b`: a deletion plus a addition, a and b is different.
	// `git rm`: deletion only.
	// `git add`: addition or modification.

	// Hence, the strategy is:
	// 1. deletion: just delete it from slice and rendering cache.
	// 2. modification: delete first then adding.
	// 3. addition: adding it.
	// 4. renaming: a delete and a addition.

	// What about a user modifies follows a renaming? A deletion and addition.

	// The key point: `adds`, `mods` and `dels` slice are distinct.
	// Therefore, order doesn't matter.

	log.Printf("Batch(%v, %v, %v)", adds, mods, dels)

	r.reqs.batch <- batchReq{adds, mods, dels}
}

// NewRepo create a new article repository.
func NewRepo(repoDir string, skipDirs, globDocs []string,
	user, repo string, vistors ...Visitor) Repo {
	dir := Dir(repoDir)

	p := &DocProcessor{dir: dir}
	p.scanner = &DocScanner{
		dir:      dir,
		skipDirs: skipDirs,
		globDocs: globDocs,
	}
	p.parser = &DocParser{}

	r := new(DocRepo)
	r.dir = dir
	r.cache = make(map[string]*entry)
	r.index = make(map[string]int)
	r.processor = p
	r.visitors = vistors

	r.reqs = reqs{
		list:  make(chan listReq),
		get:   make(chan getReq),
		post:  make(chan Docs),
		batch: make(chan batchReq),
	}

	r.renderer = NewRenderer(user, repo, dir)

	// Receive result asynchronously.
	p.callback = func(docs Docs) { r.reqs.post <- docs }

	go r.loop()
	go r.Post(repoDir)
	return r
}
<commit_msg>Fix dup git diffs<commit_after>package repo

import (
	"html/template"
	"log"
	"sort"

	"github.com/longkai/xiaolongtongxue.com/helper"
)

// Repo the documents repository.
type Repo interface {
	List(since string, size int) Docs
	Get(path string) (Doc, error)
	Del(path string)
	Put(path string)
	Post(path string)
	Batch(adds, mods, dels []string)
}

type entry struct {
	ready chan struct{}
	val   template.HTML
	err   error
}

func (e *entry) call(path string, r Renderer) {
	val, err := helper.Try(3, func() (interface{}, error) {
		return r.Render(path)
	})
	e.val, e.err = val.(template.HTML), err
	close(e.ready)
}

type listReq struct {
	path string
	size int
	resp chan Docs
}

type batchReq struct {
	adds, mods, dels []string
}

type getReq struct {
	path string
	resp chan getResp
}

type getResp struct {
	doc Doc
	err error
}

type reqs struct {
	get   chan getReq
	post  chan Docs
	list  chan listReq
	batch chan batchReq
}

// DocRepo documents repository implements.
type DocRepo struct {
	reqs

	dir       Dir
	renderer  Renderer
	processor Processor
	visitors  []Visitor

	docs  Docs              // More read, less write.
	index map[string]int    // Fast lookup.
	cache map[string]*entry // Rendering cache.
}

func (r *DocRepo) loop() {
	for {
		select {
		case req := <-r.reqs.get:
			r.get(req)
		case doc := <-r.reqs.post:
			r.post(doc)
		case req := <-r.reqs.list:
			r.list(req)
		case req := <-r.reqs.batch:
			r.batch(req)
		}
	}
}

func (r *DocRepo) batch(req batchReq) {
	rm := func(paths []string) {
		for _, path := range paths {
			if _, ok := r.index[path]; ok {
				delete(r.index, path)
				delete(r.cache, path)
			}
		}
	}
	// Re-process adds and mods.
	if plen := len(req.adds) + len(req.mods); plen > 0 {
		combine := make([]string, plen)
		copy(combine, req.adds)
		copy(combine[len(req.adds):], req.mods)
		go func() {
			docs := r.processor.Process(combine...)
			// Post process.
			for _, v := range r.visitors {
				v.Visit(docs, map[int]interface{}{
					Adds: req.adds,
					Mods: req.mods,
				})
			}
		}()
	}
	// Deletions.
	rm(req.dels)
	// Modifications.
	rm(req.mods)
	// Rearrangement, strip dels and mods, the order still remains.
	idx := 0
	tmp := make(Docs, 0, len(r.docs))
	for _, doc := range r.docs {
		// Pick those who are not deleted or modified.
		if _, ok := r.index[doc.URL]; ok {
			tmp = append(tmp, doc)
			r.index[doc.URL] = idx
			idx++
		}
	}
	r.docs = tmp
	// Hence, the time complexity is O(n).
}

func (r *DocRepo) get(req getReq) {
	i, ok := r.index[req.path]
	if !ok {
		go func() { req.resp <- getResp{Doc{}, NotFoundError(req.path)} }()
		return
	}

	doc := r.docs[i]

	// A hidden doc has no newer/older navigation.
	if !doc.Hide {
		if docs := r.docs.travel(i+1, 1, true, r.docs.filterHidden); docs.Len() > 0 {
			doc.Older = docs[0].URL
		}
		if docs := r.docs.travel(i-1, 1, false, r.docs.filterHidden); docs.Len() > 0 {
			doc.Newer = docs[0].URL
		}
	}

	e := r.cache[req.path]
	if e == nil {
		// Cache misses.
		e = &entry{ready: make(chan struct{})}
		go e.call(doc.Path, r.renderer)
		r.cache[req.path] = e
	}

	go func() {
		<-e.ready
		doc.Body = e.val
		req.resp <- getResp{doc, e.err}
	}()
}

func (r *DocRepo) post(docs Docs) {
	oldSize := r.docs.Len()
	// If a some of a newly doc has been existed already, replace them.
	// i.e., avoid duplication.
	for _, d := range docs {
		if i, ok := r.index[d.URL]; ok {
			log.Printf("replace %q with %q", r.docs[i].URL, d.URL)
			r.docs[i] = d          // Replace the old one.
			delete(r.cache, d.URL) // Clear its rendering cache, if any.
		} else {
			r.docs = append(r.docs, d) // Append the new one.
			// If there are multiple docs have a same URL,
			// the last one will be kept.
			r.index[d.URL] = r.docs.Len() - 1
		}
	}
	sort.Sort(r.docs)

	// Rebuild index.
	index := r.index
	for i, d := range r.docs {
		index[d.URL] = i
	}

	log.Printf("receive %d docs, len %d to %d", docs.Len(), oldSize, r.docs.Len())
}

func (r *DocRepo) list(req listReq) {
	i, ok := r.index[req.path]
	if !ok {
		i = 0 // If not found any match, start from 0.
	} else {
		i++ // Skip the current one.
	}

	res := r.docs.travel(i, req.size, true, r.docs.filterHidden)

	go func() { req.resp <- res }()
}

// List articles since a specific path, excluded.
func (r *DocRepo) List(since string, size int) Docs {
	resp := make(chan Docs)
	r.reqs.list <- listReq{since, size, resp}
	return <-resp
}

// Get a document for the path.
func (r *DocRepo) Get(path string) (Doc, error) {
	// Read only index, fast indexing without channel synchronization.
	// It's safe since lookup success, the go-routine will lookup again.
	// Hence, when lookup fail, maybe it's just removed or never exists.
	if _, ok := r.index[path]; !ok {
		return Doc{}, NotFoundError(path)
	}
	resp := make(chan getResp)
	r.reqs.get <- getReq{path, resp}
	v := <-resp
	return v.doc, v.err
}

// Del a document for the path.
func (r *DocRepo) Del(path string) {
	r.Batch(nil, nil, []string{path})
}

// Put revalidate a document.
func (r *DocRepo) Put(path string) {
	r.Batch(nil, []string{path}, nil)
}

// Post publish the path for documents.
// This method should be called when you start the application.
// Since it won't call the visitors let them do post process.
func (r *DocRepo) Post(path string) {
	log.Printf("post %s", path)
	r.processor.Process(path)
}

// Batch additions, modifications and deletions into a single request.
func (r *DocRepo) Batch(adds, mods, dels []string) {
	// Git only tracks files, hence, all of the slice are files path.

	// `git mv a b`: a deletion plus a addition, a and b is different.
	// `git rm`: deletion only.
	// `git add`: addition or modification.

	// Hence, the strategy is:
	// 1. deletion: just delete it from slice and rendering cache.
	// 2. modification: delete first then adding.
	// 3. addition: adding it.
	// 4. renaming: a delete and a addition.

	// What about a user modifies follows a renaming? A deletion and addition.

	// The key point: `adds`, `mods` and `dels` slice are distinct.
	// Therefore, order doesn't matter.

	log.Printf("Batch(%v, %v, %v)", filter(adds), filter(mods), filter(dels))

	r.reqs.batch <- batchReq{adds, mods, dels}
}

// NewRepo create a new article repository.
func NewRepo(repoDir string, skipDirs, globDocs []string,
	user, repo string, vistors ...Visitor) Repo {
	dir := Dir(repoDir)

	p := &DocProcessor{dir: dir}
	p.scanner = &DocScanner{
		dir:      dir,
		skipDirs: skipDirs,
		globDocs: globDocs,
	}
	p.parser = &DocParser{}

	r := new(DocRepo)
	r.dir = dir
	r.cache = make(map[string]*entry)
	r.index = make(map[string]int)
	r.processor = p
	r.visitors = vistors

	r.reqs = reqs{
		list:  make(chan listReq),
		get:   make(chan getReq),
		post:  make(chan Docs),
		batch: make(chan batchReq),
	}

	r.renderer = NewRenderer(user, repo, dir)

	// Receive result asynchronously.
	p.callback = func(docs Docs) { r.reqs.post <- docs }

	go r.loop()
	go r.Post(repoDir)
	return r
}
<|endoftext|>


<commit_before>package main

import (
  "bytes"
  "fmt"
  "net/http"
  "io/ioutil"
  "log"
  "encoding/json"
  "io"
  // "errors"
)

// STRUCTS

type Client struct {
  Url string
}

// http://docs.couchdb.org/en/latest/intro/api.html#server
type Server struct {
  Couchdb string
  Uuid string
  Vendor struct {
    Version string
    Name string
  }
  Version string
}

type Database struct {
  Url string
}

type DatabaseInfo struct {
  Name string `json:"db_name"`
  DocCount int `json:"doc_count"`
  DocDelCount int `json:"doc_del_count"`
}

type DbResponse struct {
  Ok bool
  Error string
  Reason string
}

type DocResponse struct {
  Ok bool
  Id string
  Rev string
  Error string
  Reason string
}

type Error struct {
  Method string
  Url string
  StatusCode int
  Type string `json:"error"`
  Reason string
}

// custom Error struct has to implement Error method
func (e *Error) Error() string {
  return fmt.Sprintf("CouchDB - %s %s, Status Code: %d, Error: %s, Reason: %s", e.Method, e.Url, e.StatusCode, e.Type, e.Reason)
}

// CLIENT OPERATIONS

/**
 * Get server information.
 */
func (c *Client) info() (*Server, error) {
  body, err := request("GET", c.Url, nil)
  if err != nil {
    return nil, err
  }
  var server *Server
  err = json.Unmarshal(body, &server)
  if err != nil {
    return nil, err
  }
  return server, nil
}

/**
 * Get all databases.
 */
func (c *Client) all() ([]string, error) {
  body, err := request("GET", c.Url + "_all_dbs", nil)
  if err != nil {
    return nil, err
  }
  var data []string
  return data, json.Unmarshal(body, &data)
}

/**
 * Get single database.
 */
func (c *Client) get(name string) (*DatabaseInfo, error) {
  body, err := request("GET", c.Url + name, nil)
  if err != nil {
    return nil, err
  }
  var dbInfo *DatabaseInfo
  return dbInfo, json.Unmarshal(body, &dbInfo)
}

/**
 * Create single database.
 */
func (c *Client) create(name string) (*DbResponse, error) {
  body, err := request("PUT", c.Url + name, nil)
  if err != nil {
    return nil, err
  }
  var DbResponse *DbResponse
  return DbResponse, json.Unmarshal(body, &DbResponse)
}

/**
 * Delete single database.
 */
func (c *Client) delete(name string) (*DbResponse, error) {
  body, err := request("DELETE", c.Url + name, nil)
  if err != nil {
    return nil, err
  }
  var DbResponse *DbResponse
  return DbResponse, json.Unmarshal(body, &DbResponse)
}

func (c *Client) use(name string) (Database) {
  return Database{c.Url + name + "/"}
}

// DATABASE OPERATIONS

/**
 * Head request.
 * http://docs.couchdb.org/en/latest/api/document/common.html#head--db-docid
 */
func (db *Database) head(id string) (*http.Response, error) {
  return http.Head(db.Url + id)
}

func (db *Database) get(id string) (map[string]interface{}, error) {
  body, err := request("GET", db.Url + id, nil)
  if err != nil {
    return nil, err
  }
  var data map[string]interface{}
  return data, json.Unmarshal(body, &data)
}

func (db *Database) put(id string, document interface{}) (*DocResponse, error) {
  data, err := marshal(document)
  if err != nil {
    return nil, err
  }
  body, err := request("PUT", db.Url + id, data)
  if err != nil {
    return nil, err
  }
  var res *DocResponse
  return res, json.Unmarshal(body, &res)
}

// FUNC MAIN

func main() {

  const url = "http://127.0.0.1:5984/"

  // create client
  client := &Client{url}

  // get server info
  // couch, err := client.info()
  // if err != nil {
  //   log.Fatal(err)
  // }
  // fmt.Println(couch.Vendor.Version)

  // get all dbs
  // res, err := client.all()
  // if err != nil {
  //   log.Fatal(err)
  // }
  // fmt.Println(res)

  // get db information
  // info, err := client.get("nice")
  // if err != nil {
  //   log.Fatal(err)
  // }
  // fmt.Println(info)

  // use db
  db := client.use("nice")

  // get document head
  // head, err := db.head("awesome")
  // if err != nil {
  //   log.Fatal(err)
  // }
  // fmt.Println(head.StatusCode)

  // get document
  // doc, err := db.get("awesome")
  // if err != nil {
  //   log.Fatal(err)
  // }
  // nested := doc["nested"].(map[string]interface{})
  // fmt.Println(nested["awesome"])

  // put document
  type MyDoc struct {
    Brand string `json:"brand"`
  }
  myDoc := MyDoc{"audi"}
  _, err := db.put("tight", myDoc)
  if err != nil {
    // fmt.Println(err.Type)
    log.Fatal(err)
  }

  //
  // // create db
  // status, err := client.create("nice")
  // if err != nil {
  //   log.Fatal(err)
  // }
  // fmt.Println(status)
  // fmt.Println(status.Ok)
  //
  // // delete database
  // status, err = client.delete("awesome")
  // if err != nil {
  //   log.Fatal(err)
  // }
  // fmt.Println(status)
}

// HELPER FUNCTIONS
func request(method, url string, data io.Reader) ([]byte, error) {
  client := &http.Client{}
  req, err := http.NewRequest(method, url, data)
  if err != nil {
    return nil, err
  }
  res, err := client.Do(req)
  if err != nil {
    return nil, err
  }
  defer res.Body.Close()
  body, err := ioutil.ReadAll(res.Body)
  if err != nil {
    return nil, err
  }
  // handle CouchDB http errors
  var error *Error
  err = json.Unmarshal(body, &error)
  if err != nil {
    return nil, err
  }
  if error.Type != "" && error.Reason != "" {
    error.Method = method
    error.Url = url
    error.StatusCode = res.StatusCode
    return nil, error
  }
  return body, nil
}

func marshal(v interface{}) (io.Reader, error) {
  json, err := json.Marshal(v)
  if err != nil {
    return nil, err
  }
  return bytes.NewReader(json), nil
}
<commit_msg>improve put function by using inout params<commit_after>package main

import (
  "bytes"
  "fmt"
  "net/http"
  "io/ioutil"
  "log"
  "encoding/json"
  "io"
  // "errors"
)

// STRUCTS

type Client struct {
  Url string
}

// http://docs.couchdb.org/en/latest/intro/api.html#server
type Server struct {
  Couchdb string
  Uuid string
  Vendor struct {
    Version string
    Name string
  }
  Version string
}

type Database struct {
  Url string
}

type DatabaseInfo struct {
  Name string `json:"db_name"`
  DocCount int `json:"doc_count"`
  DocDelCount int `json:"doc_del_count"`
}

type DbResponse struct {
  Ok bool
  Error string
  Reason string
}

type Error struct {
  Method string
  Url string
  StatusCode int
  Type string `json:"error"`
  Reason string
}

// custom Error struct has to implement Error method
func (e *Error) Error() string {
  return fmt.Sprintf("CouchDB - %s %s, Status Code: %d, Error: %s, Reason: %s", e.Method, e.Url, e.StatusCode, e.Type, e.Reason)
}

// leave out _rev when empty otherwise "Invalid rev format"
type Document struct {
  Id string `json:"_id"`
  Rev string `json:"_rev,omitempty"`
}

// CLIENT OPERATIONS

/**
 * Get server information.
 */
func (c *Client) info() (*Server, error) {
  body, err := request("GET", c.Url, nil)
  if err != nil {
    return nil, err
  }
  var server *Server
  err = json.Unmarshal(body, &server)
  if err != nil {
    return nil, err
  }
  return server, nil
}

/**
 * Get all databases.
 */
func (c *Client) all() ([]string, error) {
  body, err := request("GET", c.Url + "_all_dbs", nil)
  if err != nil {
    return nil, err
  }
  var data []string
  return data, json.Unmarshal(body, &data)
}

/**
 * Get single database.
 */
func (c *Client) get(name string) (*DatabaseInfo, error) {
  body, err := request("GET", c.Url + name, nil)
  if err != nil {
    return nil, err
  }
  var dbInfo *DatabaseInfo
  return dbInfo, json.Unmarshal(body, &dbInfo)
}

/**
 * Create single database.
 */
func (c *Client) create(name string) (*DbResponse, error) {
  body, err := request("PUT", c.Url + name, nil)
  if err != nil {
    return nil, err
  }
  var DbResponse *DbResponse
  return DbResponse, json.Unmarshal(body, &DbResponse)
}

/**
 * Delete single database.
 */
func (c *Client) delete(name string) (*DbResponse, error) {
  body, err := request("DELETE", c.Url + name, nil)
  if err != nil {
    return nil, err
  }
  var DbResponse *DbResponse
  return DbResponse, json.Unmarshal(body, &DbResponse)
}

func (c *Client) use(name string) (Database) {
  return Database{c.Url + name + "/"}
}

// DATABASE OPERATIONS

/**
 * Head request.
 * http://docs.couchdb.org/en/latest/api/document/common.html#head--db-docid
 */
func (db *Database) head(id string) (*http.Response, error) {
  return http.Head(db.Url + id)
}

func (db *Database) get(document interface{}, id string) error {
  body, err := request("GET", db.Url + id, nil)
  if err != nil {
    return err
  }
  return json.Unmarshal(body, &document)
}

func (db *Database) put(doc interface{}) error {
  res, err := json.Marshal(doc)
  if err != nil {
    return err
  }
  var document *Document
  err = json.Unmarshal(res, &document)
  if err != nil {
    return err
  }
  data := bytes.NewReader(res)
  if err != nil {
    return err
  }
  _, err = request("PUT", db.Url + document.Id, data)
  if err != nil {
    return err
  }
  return db.get(doc, document.Id)
}

// FUNC MAIN

func main() {

  const url = "http://127.0.0.1:5984/"

  // create client
  client := &Client{url}

  // get server info
  // couch, err := client.info()
  // if err != nil {
  //   log.Fatal(err)
  // }
  // fmt.Println(couch.Vendor.Version)

  // get all dbs
  // res, err := client.all()
  // if err != nil {
  //   log.Fatal(err)
  // }
  // fmt.Println(res)

  // get db information
  // info, err := client.get("nice")
  // if err != nil {
  //   log.Fatal(err)
  // }
  // fmt.Println(info)

  // use db
  db := client.use("nice")

  // get document head
  // head, err := db.head("awesome")
  // if err != nil {
  //   log.Fatal(err)
  // }
  // fmt.Println(head.StatusCode)

  type MyDoc struct {
    Document
    Brand string `json:"brand"`
    Name string
    Nested struct {
      Awesome string
    }
  }

  // get document
  var myDoc *MyDoc
  err := db.get(&myDoc, "awesome")
  if err != nil {
    log.Fatal(err)
  }
  fmt.Println(myDoc)

  myDoc.Name = "sour"
  err = db.put(&myDoc)
  if err != nil {
    log.Fatal(err)
  }
  fmt.Println(myDoc)

  // doc["foo"] = "bar"
  // nested := doc["nested"].(map[string]interface{})
  // fmt.Println(nested["awesome"])

  //
  // // create db
  // status, err := client.create("nice")
  // if err != nil {
  //   log.Fatal(err)
  // }
  // fmt.Println(status)
  // fmt.Println(status.Ok)
  //
  // // delete database
  // status, err = client.delete("awesome")
  // if err != nil {
  //   log.Fatal(err)
  // }
  // fmt.Println(status)
}

// HELPER FUNCTIONS
func request(method, url string, data io.Reader) ([]byte, error) {
  client := &http.Client{}
  req, err := http.NewRequest(method, url, data)
  if err != nil {
    return nil, err
  }
  res, err := client.Do(req)
  if err != nil {
    return nil, err
  }
  defer res.Body.Close()
  body, err := ioutil.ReadAll(res.Body)
  if err != nil {
    return nil, err
  }
  // handle CouchDB http errors
  var error *Error
  err = json.Unmarshal(body, &error)
  if err != nil {
    return nil, err
  }
  if error.Type != "" && error.Reason != "" {
    error.Method = method
    error.Url = url
    error.StatusCode = res.StatusCode
    return nil, error
  }
  return body, nil
}

// func marshal(v interface{}) (io.Reader, error) {
//   res, err := json.Marshal(v)
//   var document *Document
//   json.Unmarshal(res, &document)
//   fmt.Println(document.Id)
//   if err != nil {
//     return nil, err
//   }
//   return bytes.NewReader(res), nil
// }
<|endoftext|>


<commit_before>package repr

import (
	"testing"

	"github.com/stretchr/testify/assert"
)

type anotherStruct struct {
	A []int
}

type testStruct struct {
	S string
	I *int
	A anotherStruct
}

func TestReprEmptyArray(t *testing.T) {
	assert.Equal(t, "[]string{}", String([]string{}, OmitEmpty(false)))
}

func TestReprStringArray(t *testing.T) {
	assert.Equal(t, "[]string{\"a\", \"b\"}", String([]string{"a", "b"}))
}

func TestReprIntArray(t *testing.T) {
	assert.Equal(t, "[]int{1, 2}", String([]int{1, 2}))
}

func TestReprPointerToInt(t *testing.T) {
	pi := new(int)
	*pi = 13
	assert.Equal(t, `&13`, String(pi))
}

func TestReprChannel(t *testing.T) {
	ch := make(<-chan map[string]*testStruct, 1)
	assert.Equal(t, `make(<-chan map[string]*repr.testStruct, 1)`, String(ch))
}

func TestReprEmptyMap(t *testing.T) {
	assert.Equal(t, "map[string]bool{}", String(map[string]bool{}))
}

func TestReprMap(t *testing.T) {
	m := map[string]int{"a": 1}
	assert.Equal(t, "map[string]int{\"a\": 1}", String(m))
}

func TestReprStructWithIndent(t *testing.T) {
	pi := new(int)
	*pi = 13
	s := &testStruct{
		S: "String",
		I: pi,
		A: anotherStruct{
			A: []int{1, 2, 3},
		},
	}
	assert.Equal(t, `&repr.testStruct{
  S: "String",
  I: &13,
  A: repr.anotherStruct{
    A: []int{
      1,
      2,
      3,
    },
  },
}`, String(s, Indent("  ")))

}

func TestReprByteArray(t *testing.T) {
	b := []byte{1, 2, 3}
	assert.Equal(t, `[]uint8{1, 2, 3}`, String(b))
}

type privateTestStruct struct {
	a string
}

func TestReprPrivateField(t *testing.T) {
	s := privateTestStruct{"hello"}
	assert.Equal(t, `repr.privateTestStruct{a: "hello"}`, String(s))
}

type Enum int

func (e Enum) String() string {
	return "Value"
}

func TestEnum(t *testing.T) {
	v := Enum(1)
	s := String(v)
	assert.Equal(t, "repr.Enum(Value)", s)
}
<commit_msg>Fix tests!<commit_after>package repr

import (
	"testing"

	"github.com/stretchr/testify/assert"
)

type anotherStruct struct {
	A []int
}

type testStruct struct {
	S string
	I *int
	A anotherStruct
}

func TestReprEmptyArray(t *testing.T) {
	assert.Equal(t, "[]string{}", String([]string{}, OmitEmpty(false)))
}

func TestReprStringArray(t *testing.T) {
	assert.Equal(t, "[]string{\"a\", \"b\"}", String([]string{"a", "b"}))
}

func TestReprIntArray(t *testing.T) {
	assert.Equal(t, "[]int{1, 2}", String([]int{1, 2}))
}

func TestReprPointerToInt(t *testing.T) {
	pi := new(int)
	*pi = 13
	assert.Equal(t, `&13`, String(pi))
}

func TestReprChannel(t *testing.T) {
	ch := make(<-chan map[string]*testStruct, 1)
	assert.Equal(t, `make(<-chan map[string]*repr.testStruct, 1)`, String(ch))
}

func TestReprEmptyMap(t *testing.T) {
	assert.Equal(t, "map[string]bool{}", String(map[string]bool{}))
}

func TestReprMap(t *testing.T) {
	m := map[string]int{"a": 1}
	assert.Equal(t, "map[string]int{\"a\": 1}", String(m))
}

func TestReprStructWithIndent(t *testing.T) {
	pi := new(int)
	*pi = 13
	s := &testStruct{
		S: "String",
		I: pi,
		A: anotherStruct{
			A: []int{1, 2, 3},
		},
	}
	assert.Equal(t, `&repr.testStruct{
  S: "String",
  I: &13,
  A: repr.anotherStruct{
    A: []int{
      1,
      2,
      3,
    },
  },
}`, String(s, Indent("  ")))

}

func TestReprByteArray(t *testing.T) {
	b := []byte{1, 2, 3}
	assert.Equal(t, "[]byte(\"\\x01\\x02\\x03\")", String(b))
}

type privateTestStruct struct {
	a string
}

func TestReprPrivateField(t *testing.T) {
	s := privateTestStruct{"hello"}
	assert.Equal(t, `repr.privateTestStruct{a: "hello"}`, String(s))
}

type Enum int

func (e Enum) String() string {
	return "Value"
}

func TestEnum(t *testing.T) {
	v := Enum(1)
	s := String(v)
	assert.Equal(t, "repr.Enum(Value)", s)
}
<|endoftext|>