<commit_before>package zoo

import (
	"fmt"
	"path"
	"strconv"
	"strings"
	"time"

	"github.com/gotgo/fw/me"
	"github.com/samuel/go-zookeeper/zk"
)

const connectTimeout = 3 * time.Second

func NewKafkaKeeper(hosts []string, c *TopicConsumer, s *KafkaState) *KafkaKeeper {
	return &KafkaKeeper{
		acl:      zk.WorldACL(zk.PermAll),
		hosts:    hosts,
		Consumer: c,
		State:    s,
	}
}

// KafkaKeeper - Used to manage consumer per topic, per application consuming
// > /{TopicConsumer.Root}/{topic}/apps/{app}/partitions
//														/{partition}/consumed/{offset}
type KafkaKeeper struct {
	acl      []zk.ACL
	hosts    []string
	Consumer *TopicConsumer
	State    *KafkaState
}

func TestConnect(hosts []string) error {
	if conn, _, err := zk.Connect(hosts, connectTimeout); err != nil {
		return err
	} else {
		conn.Close()
		return nil
	}
}

func (k *KafkaKeeper) connect() *zk.Conn {
	conn, _, err := zk.Connect(k.hosts, connectTimeout)
	if err != nil {
		panic(me.Err(err, "failed to connect to zookeeper nodes"))
	}
	return conn
}

func (z *KafkaKeeper) ensureExists(c *zk.Conn, p string, data string) error {
	parts := strings.Split(p, "/")
	current := "/"
	for _, part := range parts {
		if len(part) == 0 {
			continue
		}

		current := path.Join(current, part)
		fmt.Printf("cheking for %s", current)
		exists, _, err := c.Exists(current)
		if err != nil {
			return me.Err(err, "error checking if path exists: "+current)
		}
		if !exists {
			const flags = 0
			fmt.Printf("creating %s", current)
			fmt.Println()
			if _, err := c.Create(current, []byte(data), flags, z.acl); err != nil {
				return me.Err(err, "error creating path: "+current)
			}
		}
	}

	return nil
}

func (z *KafkaKeeper) get(c *zk.Conn, path string) (string, int32, error) {
	bts, stat, err := c.Get(path)
	if err != nil {
		return "", 0, me.Err(err, "failed to get path: "+path)
	}
	return fmt.Sprintf("%s", bts), stat.Version, nil
}

func (z *KafkaKeeper) set(c *zk.Conn, path, data string, version int32) error {
	if _, err := c.Set(path, []byte(data), version); err != nil {
		return me.Err(err, "failed to set path: "+path)
	}
	return nil
}

func (z *KafkaKeeper) getChildren(c *zk.Conn, path string) ([]string, error) {
	children, _, err := c.Children(path)
	if err != nil {
		return nil, me.Err(err, "failed to get children for path: "+path)
	}
	return children, nil
}

func (z *KafkaKeeper) ensureSetup(c *zk.Conn) error {
	tc := z.Consumer

	//  /kafka-topics
	if err := z.ensureExists(c, tc.Root, ""); err != nil {
		return err
	}

	//  /kafka-topics/{topic}/apps/{app}/partitions/"
	if err := z.ensureExists(c, tc.PartitionsPath(), ""); err != nil {
		return err
	}

	return nil
}

func (z *KafkaKeeper) GetOffsets() (map[int]int64, error) {
	conn := z.connect()
	defer conn.Close()

	err := z.ensureSetup(conn)
	if err != nil {
		return nil, me.Err(err, "failed to get offsets")
	}

	// get offical partition data from kafka settings
	partitions, err := z.getChildren(conn, z.State.PartitionsPath())
	if err != nil {
		return nil, me.Err(err, "failed to get kafka partition info")
	}

	offsets := make(map[int]int64)
	for _, p := range partitions {
		//{partition}/consumed/{offset}
		err := z.ensureExists(conn, path.Join(z.Consumer.PartitionsPath(), p, "consumed"), "0")
		if err != nil {
			return nil, err
		}
		offset, _, err := z.get(conn, path.Join(z.Consumer.PartitionsPath(), p, "consumed"))
		if err != nil {
			return nil, err
		}
		k, _ := strconv.Atoi(p)
		v, _ := strconv.ParseInt(offset, 10, 0)
		offsets[k] = v
	}
	return offsets, nil
}

func (z *KafkaKeeper) SetOffset(partition int32, offset int64) error {
	conn := z.connect()
	defer conn.Close()

	err := z.ensureSetup(conn)
	if err != nil {
		return me.Err(err, fmt.Sprintf("failed to set offset on partition: %d", partition))
	}
	path := path.Join(z.Consumer.PartitionsPath(), fmt.Sprintf("%d", partition), "consumed")

	//force version for now
	_, version, err := z.get(conn, path)
	if err != nil {
		return err
	}

	err = z.set(conn, path, fmt.Sprintf("%d", offset), version)
	if err != nil {
		return me.Err(err, "failed to set value for path: "+path)
	}
	return err
}
<commit_msg>path creation<commit_after>package zoo

import (
	"fmt"
	"path"
	"strconv"
	"strings"
	"time"

	"github.com/gotgo/fw/me"
	"github.com/samuel/go-zookeeper/zk"
)

const connectTimeout = 3 * time.Second

func NewKafkaKeeper(hosts []string, c *TopicConsumer, s *KafkaState) *KafkaKeeper {
	return &KafkaKeeper{
		acl:      zk.WorldACL(zk.PermAll),
		hosts:    hosts,
		Consumer: c,
		State:    s,
	}
}

// KafkaKeeper - Used to manage consumer per topic, per application consuming
// > /{TopicConsumer.Root}/{topic}/apps/{app}/partitions
//														/{partition}/consumed/{offset}
type KafkaKeeper struct {
	acl      []zk.ACL
	hosts    []string
	Consumer *TopicConsumer
	State    *KafkaState
}

func TestConnect(hosts []string) error {
	if conn, _, err := zk.Connect(hosts, connectTimeout); err != nil {
		return err
	} else {
		conn.Close()
		return nil
	}
}

func (k *KafkaKeeper) connect() *zk.Conn {
	conn, _, err := zk.Connect(k.hosts, connectTimeout)
	if err != nil {
		panic(me.Err(err, "failed to connect to zookeeper nodes"))
	}
	return conn
}

func (z *KafkaKeeper) ensureExists(c *zk.Conn, p string, data string) error {
	parts := strings.Split(p, "/")
	current := "/"
	for _, part := range parts {
		if len(part) == 0 {
			continue
		}

		current = path.Join(current, part)
		fmt.Printf("cheking for %s", current)
		exists, _, err := c.Exists(current)
		if err != nil {
			return me.Err(err, "error checking if path exists: "+current)
		}
		if !exists {
			const flags = 0
			fmt.Printf("creating %s", current)
			fmt.Println()
			if _, err := c.Create(current, []byte(data), flags, z.acl); err != nil {
				return me.Err(err, "error creating path: "+current)
			}
		}
	}

	return nil
}

func (z *KafkaKeeper) get(c *zk.Conn, path string) (string, int32, error) {
	bts, stat, err := c.Get(path)
	if err != nil {
		return "", 0, me.Err(err, "failed to get path: "+path)
	}
	return fmt.Sprintf("%s", bts), stat.Version, nil
}

func (z *KafkaKeeper) set(c *zk.Conn, path, data string, version int32) error {
	if _, err := c.Set(path, []byte(data), version); err != nil {
		return me.Err(err, "failed to set path: "+path)
	}
	return nil
}

func (z *KafkaKeeper) getChildren(c *zk.Conn, path string) ([]string, error) {
	children, _, err := c.Children(path)
	if err != nil {
		return nil, me.Err(err, "failed to get children for path: "+path)
	}
	return children, nil
}

func (z *KafkaKeeper) ensureSetup(c *zk.Conn) error {
	tc := z.Consumer

	//  /kafka-topics
	if err := z.ensureExists(c, tc.Root, ""); err != nil {
		return err
	}

	//  /kafka-topics/{topic}/apps/{app}/partitions/"
	if err := z.ensureExists(c, tc.PartitionsPath(), ""); err != nil {
		return err
	}

	return nil
}

func (z *KafkaKeeper) GetOffsets() (map[int]int64, error) {
	conn := z.connect()
	defer conn.Close()

	err := z.ensureSetup(conn)
	if err != nil {
		return nil, me.Err(err, "failed to get offsets")
	}

	// get offical partition data from kafka settings
	partitions, err := z.getChildren(conn, z.State.PartitionsPath())
	if err != nil {
		return nil, me.Err(err, "failed to get kafka partition info")
	}

	offsets := make(map[int]int64)
	for _, p := range partitions {
		//{partition}/consumed/{offset}
		err := z.ensureExists(conn, path.Join(z.Consumer.PartitionsPath(), p, "consumed"), "0")
		if err != nil {
			return nil, err
		}
		offset, _, err := z.get(conn, path.Join(z.Consumer.PartitionsPath(), p, "consumed"))
		if err != nil {
			return nil, err
		}
		k, _ := strconv.Atoi(p)
		v, _ := strconv.ParseInt(offset, 10, 0)
		offsets[k] = v
	}
	return offsets, nil
}

func (z *KafkaKeeper) SetOffset(partition int32, offset int64) error {
	conn := z.connect()
	defer conn.Close()

	err := z.ensureSetup(conn)
	if err != nil {
		return me.Err(err, fmt.Sprintf("failed to set offset on partition: %d", partition))
	}
	path := path.Join(z.Consumer.PartitionsPath(), fmt.Sprintf("%d", partition), "consumed")

	//force version for now
	_, version, err := z.get(conn, path)
	if err != nil {
		return err
	}

	err = z.set(conn, path, fmt.Sprintf("%d", offset), version)
	if err != nil {
		return me.Err(err, "failed to set value for path: "+path)
	}
	return err
}
<|endoftext|>


<commit_before>package ssh

import (
	"bytes"
	"crypto/rand"
	"encoding/binary"
	"errors"
	"math/big"
	"net"
)

type Conn struct {
	// Underlying network connection
	conn net.Conn

	config *Config

	// Key information
	macLength uint32

	// Log for ZGrab output
	handshakeLog HandshakeLog

	currentCipher cipher

	kexAlgorithm     string
	hostKeyAlgorithm string
}

type sshPayload interface {
	MsgType() byte
	Marshal() ([]byte, error)
	Unmarshal([]byte) bool
}

func (c *Conn) ClientHandshake() error {
	clientProtocol := MakeZGrabProtocolAgreement()
	clientProtocolBytes := clientProtocol.Marshal()
	c.conn.Write(clientProtocolBytes)

	buf := make([]byte, 1024)
	protocolDone := false
	protocolRead := 0
	for !protocolDone && protocolRead < 1024 {
		n, err := c.conn.Read(buf[protocolRead : protocolRead+1])
		protocolRead += n
		if err != nil {
			return err
		}
		if protocolRead < 2 {
			continue
		}
		if bytes.Equal(buf[protocolRead-2:protocolRead], []byte("\r\n")) {
			protocolDone = true
		}
	}
	serverProtocol := ProtocolAgreement{
		RawBanner: string(buf[0:protocolRead]),
	}
	serverProtocol.ParseRawBanner()
	c.handshakeLog.ServerProtocol = &serverProtocol

	// See if it matches????

	//
	ckxi, err := GenerateKeyExchangeInit(c.config)
	if err != nil {
		return err
	}
	if err = c.writePacket(ckxi); err != nil {
		return err
	}

	// Read the key options
	serverKex := new(KeyExchangeInit)
	err = c.readPacket(serverKex)
	if err != nil {
		return err
	}

	c.handshakeLog.ServerKeyExchangeInit = serverKex

	if c.kexAlgorithm, err = chooseAlgorithm(ckxi.KexAlgorithms, serverKex.KexAlgorithms); err != nil {
		return err
	}

	c.handshakeLog.Algorithms = new(AlgorithmSelection)
	c.handshakeLog.Algorithms.KexAlgorithm = c.kexAlgorithm

	if c.hostKeyAlgorithm, err = chooseAlgorithm(ckxi.HostKeyAlgorithms, serverKex.HostKeyAlgorithms); err != nil {
		return err
	}
	c.handshakeLog.Algorithms.HostKeyAlgorithm = c.hostKeyAlgorithm

	switch c.kexAlgorithm {
	case KEX_DH_GROUP1_SHA1:
		if err := c.dhGroup1Kex(); err != nil {
			return err
		}
	case KEX_DH_GROUP14_SHA1:
		if err := c.dhGroup14Kex(); err != nil {
			return err
		}
	case KEX_DH_SHA1, KEX_DH_SHA256:
		if err := c.dhGroupExchange(); err != nil {
			return err
		}
	default:
		return errors.New("unimplemented kex method")
	}
	return nil
}

func (c *Conn) HandshakeLog() *HandshakeLog {
	return &c.handshakeLog
}

func (c *Conn) readPacket(expected sshPayload) error {
	// Make a buffer of max packet size
	buf := make([]byte, 35001)
	totalRead, err := c.conn.Read(buf[0:4])
	if err != nil {
		return err
	}
	for totalRead < 4 {
		n, err := c.conn.Read(buf[totalRead:4])
		totalRead += n
		if err != nil {
			return err
		}
	}
	var p packet
	p.packetLength = binary.BigEndian.Uint32(buf[0:4])
	if p.packetLength > 35000 {
		return errLongPacket
	}
	totalLength := expectedLength(p.packetLength, c.macLength)
	for totalRead < totalLength {
		n, err := c.conn.Read(buf[totalRead:totalLength])
		totalRead += n
		if err != nil {
			return err
		}
	}
	// Fill out the rest of the packet
	b := buf[4:totalLength]

	// Read padding length
	if len(b) < 1 {
		return errShortPacket
	}
	p.paddingLength = b[0]
	b = b[1:]
	if uint32(p.paddingLength) > p.packetLength-1 {
		return errInvalidPadding
	}

	// Read the payload
	payloadLength := p.packetLength - uint32(p.paddingLength) - 1
	p.msgType = b[0]
	p.payload = b[1:payloadLength]
	b = b[payloadLength:]

	// Read the padding
	p.padding = b[0:p.paddingLength]
	b = b[p.paddingLength:]

	// Read the MAC, if applicable
	if uint32(len(b)) != c.macLength {
		return errShortPacket
	}

	if c.macLength > 0 {
		p.mac = b[0:c.macLength]
	}
	if len(p.payload) < 1 {
		return errShortPacket
	}
	if p.msgType != expected.MsgType() {
		return errUnexpectedMessage
	}
	if ok := expected.Unmarshal(p.payload); !ok {
		return errors.New("could not unmarshal")
	}

	return nil
}

func (c *Conn) writePacket(payload sshPayload) error {
	payloadBytes, err := payload.Marshal()
	msgType := payload.MsgType()
	if err != nil {
		return err
	}
	if len(payloadBytes) > 32768 {
		return errLongPacket
	}
	paddingLen := 8 - ((4 + 1 + 1 + len(payloadBytes)) % 8)
	if paddingLen < 4 {
		paddingLen += 8
	}
	paddingBytes := make([]byte, paddingLen)
	if len(paddingBytes) > 255 {
		return errInvalidPadding
	}
	pkt := packet{
		packetLength: uint32(2 + len(payloadBytes) + len(paddingBytes)),
		msgType:      msgType,
		payload:      payloadBytes,
		padding:      paddingBytes,
		mac:          []byte{},
	}
	out := make([]byte, 4+1+1+len(pkt.payload)+len(pkt.padding))
	binary.BigEndian.PutUint32(out, pkt.packetLength)
	out[4] = byte(len(pkt.padding))
	out[5] = pkt.msgType
	copy(out[6:], pkt.payload)
	copy(out[6+len(pkt.payload):], pkt.padding)

	written := 0
	for written < len(out) {
		n, err := c.conn.Write(out[written:])
		written += n
		if err != nil {
			return err
		}
	}
	written = 0
	mac := make([]byte, 0)
	for written < len(mac) {
		n, err := c.conn.Write(mac[written:])
		written += n
		if err != nil {
			return err
		}
	}
	return nil
}

func (c *Conn) dhGroupExchange() error {
	gexRequest := new(KeyExchangeDHGroupRequest)
	gexRequest.Min = 1024
	gexRequest.Preferred = 2048
	gexRequest.Max = 4096
	if err := c.writePacket(gexRequest); err != nil {
		return err
	}
	gexParams := new(KeyExchangeDHGroupParameters)
	if err := c.readPacket(gexParams); err != nil {
		return err
	}
	c.handshakeLog.KexDHGroupParams = gexParams

	gexInit := new(KeyExchangeDHGroupInit)
	g := big.NewInt(0).SetBytes(gexParams.Generator.Bytes())
	p := big.NewInt(0).SetBytes(gexParams.Prime.Bytes())
	order := big.NewInt(0)
	order.Sub(p, big.NewInt(1))
	x, err := rand.Int(c.config.getRandom(), order)
	if err != nil {
		return err
	}
	gexInit.E.Exp(g, x, p)
	if err = c.writePacket(gexInit); err != nil {
		return err
	}
	gexReply := new(KeyExchangeDHGroupReply)
	if err = c.readPacket(gexReply); err != nil {
		return err
	}
	c.handshakeLog.KexDHGroupReply = gexReply
	return nil
}

func (c *Conn) dhExchange(params *DHParams) error {
	x, err := rand.Int(c.config.getRandom(), params.order)
	if err != nil {
		return err
	}
	dhi := new(KeyExchangeDHInit)
	E := big.NewInt(0)
	E.Exp(params.Generator, x, params.Prime)
	dhi.E.Set(E)
	c.writePacket(dhi)
	dhReply := new(KeyExchangeDHInitReply)
	if err = c.readPacket(dhReply); err != nil {
		return err
	}

	c.handshakeLog.DHReply = dhReply
	return nil
}

func (c *Conn) dhGroup1Kex() error {
	return c.dhExchange(&dhOakleyGroup2)
}

func (c *Conn) dhGroup14Kex() error {
	return c.dhExchange(&dhOakleyGroup14)
}
<commit_msg>SSH - Update DH GEX request to use OpenSSH sizes<commit_after>package ssh

import (
	"bytes"
	"crypto/rand"
	"encoding/binary"
	"errors"
	"math/big"
	"net"
)

type Conn struct {
	// Underlying network connection
	conn net.Conn

	config *Config

	// Key information
	macLength uint32

	// Log for ZGrab output
	handshakeLog HandshakeLog

	currentCipher cipher

	kexAlgorithm     string
	hostKeyAlgorithm string
}

type sshPayload interface {
	MsgType() byte
	Marshal() ([]byte, error)
	Unmarshal([]byte) bool
}

func (c *Conn) ClientHandshake() error {
	clientProtocol := MakeZGrabProtocolAgreement()
	clientProtocolBytes := clientProtocol.Marshal()
	c.conn.Write(clientProtocolBytes)

	buf := make([]byte, 1024)
	protocolDone := false
	protocolRead := 0
	for !protocolDone && protocolRead < 1024 {
		n, err := c.conn.Read(buf[protocolRead : protocolRead+1])
		protocolRead += n
		if err != nil {
			return err
		}
		if protocolRead < 2 {
			continue
		}
		if bytes.Equal(buf[protocolRead-2:protocolRead], []byte("\r\n")) {
			protocolDone = true
		}
	}
	serverProtocol := ProtocolAgreement{
		RawBanner: string(buf[0:protocolRead]),
	}
	serverProtocol.ParseRawBanner()
	c.handshakeLog.ServerProtocol = &serverProtocol

	// See if it matches????

	//
	ckxi, err := GenerateKeyExchangeInit(c.config)
	if err != nil {
		return err
	}
	if err = c.writePacket(ckxi); err != nil {
		return err
	}

	// Read the key options
	serverKex := new(KeyExchangeInit)
	err = c.readPacket(serverKex)
	if err != nil {
		return err
	}

	c.handshakeLog.ServerKeyExchangeInit = serverKex

	if c.kexAlgorithm, err = chooseAlgorithm(ckxi.KexAlgorithms, serverKex.KexAlgorithms); err != nil {
		return err
	}

	c.handshakeLog.Algorithms = new(AlgorithmSelection)
	c.handshakeLog.Algorithms.KexAlgorithm = c.kexAlgorithm

	if c.hostKeyAlgorithm, err = chooseAlgorithm(ckxi.HostKeyAlgorithms, serverKex.HostKeyAlgorithms); err != nil {
		return err
	}
	c.handshakeLog.Algorithms.HostKeyAlgorithm = c.hostKeyAlgorithm

	switch c.kexAlgorithm {
	case KEX_DH_GROUP1_SHA1:
		if err := c.dhGroup1Kex(); err != nil {
			return err
		}
	case KEX_DH_GROUP14_SHA1:
		if err := c.dhGroup14Kex(); err != nil {
			return err
		}
	case KEX_DH_SHA1, KEX_DH_SHA256:
		if err := c.dhGroupExchange(); err != nil {
			return err
		}
	default:
		return errors.New("unimplemented kex method")
	}
	return nil
}

func (c *Conn) HandshakeLog() *HandshakeLog {
	return &c.handshakeLog
}

func (c *Conn) readPacket(expected sshPayload) error {
	// Make a buffer of max packet size
	buf := make([]byte, 35001)
	totalRead, err := c.conn.Read(buf[0:4])
	if err != nil {
		return err
	}
	for totalRead < 4 {
		n, err := c.conn.Read(buf[totalRead:4])
		totalRead += n
		if err != nil {
			return err
		}
	}
	var p packet
	p.packetLength = binary.BigEndian.Uint32(buf[0:4])
	if p.packetLength > 35000 {
		return errLongPacket
	}
	totalLength := expectedLength(p.packetLength, c.macLength)
	for totalRead < totalLength {
		n, err := c.conn.Read(buf[totalRead:totalLength])
		totalRead += n
		if err != nil {
			return err
		}
	}
	// Fill out the rest of the packet
	b := buf[4:totalLength]

	// Read padding length
	if len(b) < 1 {
		return errShortPacket
	}
	p.paddingLength = b[0]
	b = b[1:]
	if uint32(p.paddingLength) > p.packetLength-1 {
		return errInvalidPadding
	}

	// Read the payload
	payloadLength := p.packetLength - uint32(p.paddingLength) - 1
	p.msgType = b[0]
	p.payload = b[1:payloadLength]
	b = b[payloadLength:]

	// Read the padding
	p.padding = b[0:p.paddingLength]
	b = b[p.paddingLength:]

	// Read the MAC, if applicable
	if uint32(len(b)) != c.macLength {
		return errShortPacket
	}

	if c.macLength > 0 {
		p.mac = b[0:c.macLength]
	}
	if len(p.payload) < 1 {
		return errShortPacket
	}
	if p.msgType != expected.MsgType() {
		return errUnexpectedMessage
	}
	if ok := expected.Unmarshal(p.payload); !ok {
		return errors.New("could not unmarshal")
	}

	return nil
}

func (c *Conn) writePacket(payload sshPayload) error {
	payloadBytes, err := payload.Marshal()
	msgType := payload.MsgType()
	if err != nil {
		return err
	}
	if len(payloadBytes) > 32768 {
		return errLongPacket
	}
	paddingLen := 8 - ((4 + 1 + 1 + len(payloadBytes)) % 8)
	if paddingLen < 4 {
		paddingLen += 8
	}
	paddingBytes := make([]byte, paddingLen)
	if len(paddingBytes) > 255 {
		return errInvalidPadding
	}
	pkt := packet{
		packetLength: uint32(2 + len(payloadBytes) + len(paddingBytes)),
		msgType:      msgType,
		payload:      payloadBytes,
		padding:      paddingBytes,
		mac:          []byte{},
	}
	out := make([]byte, 4+1+1+len(pkt.payload)+len(pkt.padding))
	binary.BigEndian.PutUint32(out, pkt.packetLength)
	out[4] = byte(len(pkt.padding))
	out[5] = pkt.msgType
	copy(out[6:], pkt.payload)
	copy(out[6+len(pkt.payload):], pkt.padding)

	written := 0
	for written < len(out) {
		n, err := c.conn.Write(out[written:])
		written += n
		if err != nil {
			return err
		}
	}
	written = 0
	mac := make([]byte, 0)
	for written < len(mac) {
		n, err := c.conn.Write(mac[written:])
		written += n
		if err != nil {
			return err
		}
	}
	return nil
}

func (c *Conn) dhGroupExchange() error {
	gexRequest := new(KeyExchangeDHGroupRequest)
	gexRequest.Min = 1024
	gexRequest.Preferred = 3072
	gexRequest.Max = 8192
	if err := c.writePacket(gexRequest); err != nil {
		return err
	}
	gexParams := new(KeyExchangeDHGroupParameters)
	if err := c.readPacket(gexParams); err != nil {
		return err
	}
	c.handshakeLog.KexDHGroupParams = gexParams

	gexInit := new(KeyExchangeDHGroupInit)
	g := big.NewInt(0).SetBytes(gexParams.Generator.Bytes())
	p := big.NewInt(0).SetBytes(gexParams.Prime.Bytes())
	order := big.NewInt(0)
	order.Sub(p, big.NewInt(1))
	x, err := rand.Int(c.config.getRandom(), order)
	if err != nil {
		return err
	}
	gexInit.E.Exp(g, x, p)
	if err = c.writePacket(gexInit); err != nil {
		return err
	}
	gexReply := new(KeyExchangeDHGroupReply)
	if err = c.readPacket(gexReply); err != nil {
		return err
	}
	c.handshakeLog.KexDHGroupReply = gexReply
	return nil
}

func (c *Conn) dhExchange(params *DHParams) error {
	x, err := rand.Int(c.config.getRandom(), params.order)
	if err != nil {
		return err
	}
	dhi := new(KeyExchangeDHInit)
	E := big.NewInt(0)
	E.Exp(params.Generator, x, params.Prime)
	dhi.E.Set(E)
	c.writePacket(dhi)
	dhReply := new(KeyExchangeDHInitReply)
	if err = c.readPacket(dhReply); err != nil {
		return err
	}

	c.handshakeLog.DHReply = dhReply
	return nil
}

func (c *Conn) dhGroup1Kex() error {
	return c.dhExchange(&dhOakleyGroup2)
}

func (c *Conn) dhGroup14Kex() error {
	return c.dhExchange(&dhOakleyGroup14)
}
<|endoftext|>


<commit_before>package main

import (
    "errors"
    "fmt"
    "net"
    "time"
//    "syscall"
//"golang.org/x/net/ipv4"
)

type IP_Reader struct {
    incomingPackets chan []byte
    nr              *Network_Reader
    protocol        uint8
    ip              string
    fragBuf         map[string](chan []byte)
    //fragQuit        map[string](chan bool)
}

func (nr *Network_Reader) NewIP_Reader(ip string, protocol uint8) (*IP_Reader, error) {
    c, err := nr.bind(ip, protocol)
    if err != nil {
        return nil, err
    }

    return &IP_Reader{
        incomingPackets: c,
        nr:              nr,
        protocol:        protocol,
        ip:              ip,
        fragBuf:         make(map[string](chan []byte)),
        //fragQuit:        make(map[string](chan bool)),
    }, nil
}

func slicePacket(b []byte) (hrd, payload []byte) {
    hdrLen := int(b[0]&0x0f) * 4
    //fmt.Println("HdrLen: ", hdrLen)
    return b[:hdrLen], b[hdrLen:]
}

func fragmentAssembler(in <-chan []byte, quit <-chan bool, finished chan<- []byte, done chan<- bool) {
    payload := make([]byte, 0)
    extraFrags := make(map[uint64]([]byte))
    recvLast := false

    for {
        select {
        case <- quit:
            fmt.Println("quitting upon force quit")
            done <- true
            return
        case frag := <-in:
            fmt.Println("got a fragment packet. len:", len(frag))
            hdr, p := slicePacket(frag)
            //offset := 8 * (uint64(hdr[6]&0x1f)<<8 + uint64(hdr[7]))
            //fmt.Println("RECEIVED FRAG")
            //fmt.Println("Offset:", offset)
            //fmt.Println(len(payload))

            // add to map
            offset := 8 * (uint64(hdr[6]) & 0x1F + uint64(hdr[7]))
            fmt.Println("Offset:", offset)
            extraFrags[offset] = p

            // check more fragments flag
            if (hdr[6]>>5)&0x01 == 0 {
                recvLast = true
            }

            // add to payload
            fmt.Println("Begin to add to the payload")
            for {
                if storedFrag, found := extraFrags[uint64(len(payload))]; found {
                    fmt.Println("New Payload Len: ", len(payload))
                    delete(extraFrags, uint64(len(payload)))
                    payload = append(payload, storedFrag...)
                } else {
                    break
                }
            }
            fmt.Println("Finished add to the payload")

            // deal with the payload
            if recvLast && len(extraFrags) == 0 {
				fmt.Println("Done")
                // correct the header
                fullPacketHdr := hdr
                totalLen := uint16(fullPacketHdr[0]&0x0F)*4 + uint16(len(payload))
                fullPacketHdr[2] = byte(totalLen >> 8)
                fullPacketHdr[3] = byte(totalLen)
                fullPacketHdr[6] = 0
                fullPacketHdr[7] = 0

                // update the checksum
                check := calculateChecksum(fullPacketHdr[:20])
                fullPacketHdr[10] = byte(check >> 8)
                fullPacketHdr[11] = byte(check)

                // send the packet back into processing
                go func() {
                    finished <- append(fullPacketHdr, payload...)
                    //fmt.Println("FINISHED")
                }()
                fmt.Println("Just wrote back in")
                done <- true
                return // from goroutine
            }
			fmt.Println("Looping")
        default:
        // make the timeout actually have a chance of being hit
        }
    }

    // drop the packet upon timeout
    fmt.Println(errors.New("Fragments took too long, packet dropped"))
    return
}

func killFragmentAssembler(quit chan<- bool, done <-chan bool, bufID string) {
    // sends quit to the assembler if it doesn't send done
    t := time.Now()
    finished := false
	for time.Since(t).Seconds() <= FRAGMENT_TIMEOUT {
        select {
        case <-done:
            fmt.Println("Recieved done msg.")
            return
        default:
        }
    }
    if !finished {
        fmt.Println("Force quitting")
        quit <- true
        <-done // wait until it is done
    }
    fmt.Println("Frag Assemble Ended, finished:", finished)
    // TODO: clean the buffer for bufID
}

func (ipr *IP_Reader) ReadFrom() (ip string, b, payload []byte, e error) {
    //fmt.Println("STARTING READ")
    b = <-ipr.incomingPackets
    //fmt.Println("RAW READ COMPLETED")
    //fmt.Println("Read Length: ", len(b))
    //fmt.Print(".")
    //fmt.Println("Full Read Data: ", b)

    hdr, p := slicePacket(b)

    // extract source IP and protocol
    ip = net.IPv4(hdr[12], hdr[13], hdr[14], hdr[15]).String()
    proto := uint8(hdr[9])
    if !((ipr.ip == ip || ipr.ip == "*") && ipr.protocol == proto) {
        fmt.Println("Not interested in packet: dropping.")
        return ipr.ReadFrom()
    }

    // verify checksum
    if verifyChecksum(hdr) != 0 {
        //fmt.Println("Header checksum verification failed. Packet dropped.")
        //fmt.Println("Wrong header: ", hdr)
        //fmt.Println("Payload (dropped): ", p)
        fmt.Println("Header checksum incorrect, packet dropped")
        return ipr.ReadFrom() // return another packet instead
    }

    packetOffset := uint16(hdr[6]&0x1f)<<8 + uint16(hdr[7])
    //fmt.Println("PACK OFF", packetOffset, "HEADER FLAGS", (hdr[6] >> 5))
    if ((hdr[6]>>5)&0x01 == 0) && (packetOffset == 0) {
        // not a fragment
        //fmt.Println("Payload Length: ", len(p))
        //fmt.Println("Full payload: ", p)
        //fmt.Println("PACKET COMPLETELY READ")
        return ip, b, p, nil
    } else {
        // is a fragment
        bufID := string([]byte{hdr[12], hdr[13], hdr[14], hdr[15], // the source IP
            hdr[16], hdr[17], hdr[18], hdr[19], // the destination IP
            hdr[9],         // protocol
            hdr[4], hdr[5], // identification
        })
        fmt.Printf("rcv a fragment-bufId: %x, len: %d\n", bufID, len(b))

        if c, ok := ipr.fragBuf[bufID]; ok {
            // the fragment has already started
            fmt.Printf("sending to already existing assembler %x\n", bufID)
            go func() { c <- b }()
        } else {
            // create the fragment buffer and quit
            fmt.Printf("creating a new buffer for %x\n", bufID)
            ipr.fragBuf[bufID] = make(chan []byte)
            quit := make(chan bool)
            done := make(chan bool)

            // create the packet assembler in a goroutine to allow the program to continue
            go fragmentAssembler(ipr.fragBuf[bufID], quit, ipr.incomingPackets, done)
            go killFragmentAssembler(quit, done, bufID)

            // send in the first fragment
            ipr.fragBuf[bufID] <- b
        }

        // after dealing with the fragment, try reading again
        //fmt.Println("RECURSE")
        return ipr.ReadFrom()
    }
}

func (ipr *IP_Reader) Close() error {
    return ipr.nr.unbind(ipr.ip, ipr.protocol)
}

/* h := &ipv4.Header{
	Version:  ipv4.Version,      // protocol version
	Len:      20,                // header length
	TOS:      0,                 // type-of-service (0 is everything normal)
	TotalLen: len(x) + 20,       // packet total length (octets)
	ID:       0,                 // identification
	Flags:    ipv4.DontFragment, // flags
	FragOff:  0,                 // fragment offset
	TTL:      8,                 // time-to-live (maximum lifespan in seconds)
	Protocol: 17,                // next protocol (17 is UDP)
	Checksum: 0,                 // checksum (apparently autocomputed)
	//Src:    net.IPv4(127, 0, 0, 1), // source address, apparently done automatically
	Dst: net.ParseIP(c.manager.ipAddress), // destination address
	//Options                         // options, extension headers
}
*/
<commit_msg>Minor bug fix to solve the problem<commit_after>package main

import (
    "errors"
    "fmt"
    "net"
    "time"
//    "syscall"
//"golang.org/x/net/ipv4"
)

type IP_Reader struct {
    incomingPackets chan []byte
    nr              *Network_Reader
    protocol        uint8
    ip              string
    fragBuf         map[string](chan []byte)
    //fragQuit        map[string](chan bool)
}

func (nr *Network_Reader) NewIP_Reader(ip string, protocol uint8) (*IP_Reader, error) {
    c, err := nr.bind(ip, protocol)
    if err != nil {
        return nil, err
    }

    return &IP_Reader{
        incomingPackets: c,
        nr:              nr,
        protocol:        protocol,
        ip:              ip,
        fragBuf:         make(map[string](chan []byte)),
        //fragQuit:        make(map[string](chan bool)),
    }, nil
}

func slicePacket(b []byte) (hrd, payload []byte) {
    hdrLen := int(b[0]&0x0f) * 4
    //fmt.Println("HdrLen: ", hdrLen)
    return b[:hdrLen], b[hdrLen:]
}

func fragmentAssembler(in <-chan []byte, quit <-chan bool, finished chan<- []byte, done chan<- bool) {
    payload := make([]byte, 0)
    extraFrags := make(map[uint64]([]byte))
    recvLast := false

    for {
        select {
        case <- quit:
            fmt.Println("quitting upon force quit")
            done <- true
            return
        case frag := <-in:
            fmt.Println("got a fragment packet. len:", len(frag))
            hdr, p := slicePacket(frag)
            //offset := 8 * (uint64(hdr[6]&0x1f)<<8 + uint64(hdr[7]))
            //fmt.Println("RECEIVED FRAG")
            //fmt.Println("Offset:", offset)
            //fmt.Println(len(payload))

            // add to map
            offset := 8 * (uint64(hdr[6] & 0x1F) << 8 + uint64(hdr[7]))
            fmt.Println("Offset:", offset)
            extraFrags[offset] = p

            // check more fragments flag
            if (hdr[6]>>5)&0x01 == 0 {
                recvLast = true
            }

            // add to payload
            fmt.Println("Begin to add to the payload")
            for {
                if storedFrag, found := extraFrags[uint64(len(payload))]; found {
                    fmt.Println("New Payload Len: ", len(payload))
                    delete(extraFrags, uint64(len(payload)))
                    payload = append(payload, storedFrag...)
                } else {
                    break
                }
            }
            fmt.Println("Finished add to the payload")

            // deal with the payload
            if recvLast && len(extraFrags) == 0 {
				fmt.Println("Done")
                // correct the header
                fullPacketHdr := hdr
                totalLen := uint16(fullPacketHdr[0]&0x0F)*4 + uint16(len(payload))
                fullPacketHdr[2] = byte(totalLen >> 8)
                fullPacketHdr[3] = byte(totalLen)
                fullPacketHdr[6] = 0
                fullPacketHdr[7] = 0

                // update the checksum
                check := calculateChecksum(fullPacketHdr[:20])
                fullPacketHdr[10] = byte(check >> 8)
                fullPacketHdr[11] = byte(check)

                // send the packet back into processing
                go func() {
                    finished <- append(fullPacketHdr, payload...)
                    //fmt.Println("FINISHED")
                }()
                fmt.Println("Just wrote back in")
                done <- true
                return // from goroutine
            }
			fmt.Println("Looping")
        default:
        // make the timeout actually have a chance of being hit
        }
    }

    // drop the packet upon timeout
    fmt.Println(errors.New("Fragments took too long, packet dropped"))
    return
}

func killFragmentAssembler(quit chan<- bool, done <-chan bool, bufID string) {
    // sends quit to the assembler if it doesn't send done
    t := time.Now()
    finished := false
	for time.Since(t).Seconds() <= FRAGMENT_TIMEOUT {
        select {
        case <-done:
            fmt.Println("Recieved done msg.")
            return
        default:
        }
    }
    if !finished {
        fmt.Println("Force quitting")
        quit <- true
        <-done // wait until it is done
    }
    fmt.Println("Frag Assemble Ended, finished:", finished)
    // TODO: clean the buffer for bufID
}

func (ipr *IP_Reader) ReadFrom() (ip string, b, payload []byte, e error) {
    //fmt.Println("STARTING READ")
    b = <-ipr.incomingPackets
    //fmt.Println("RAW READ COMPLETED")
    //fmt.Println("Read Length: ", len(b))
    //fmt.Print(".")
    //fmt.Println("Full Read Data: ", b)

    hdr, p := slicePacket(b)

    // extract source IP and protocol
    ip = net.IPv4(hdr[12], hdr[13], hdr[14], hdr[15]).String()
    proto := uint8(hdr[9])
    if !((ipr.ip == ip || ipr.ip == "*") && ipr.protocol == proto) {
        fmt.Println("Not interested in packet: dropping.")
        return ipr.ReadFrom()
    }

    // verify checksum
    if verifyChecksum(hdr) != 0 {
        //fmt.Println("Header checksum verification failed. Packet dropped.")
        //fmt.Println("Wrong header: ", hdr)
        //fmt.Println("Payload (dropped): ", p)
        fmt.Println("Header checksum incorrect, packet dropped")
        return ipr.ReadFrom() // return another packet instead
    }

    packetOffset := uint16(hdr[6]&0x1f)<<8 + uint16(hdr[7])
    //fmt.Println("PACK OFF", packetOffset, "HEADER FLAGS", (hdr[6] >> 5))
    if ((hdr[6]>>5)&0x01 == 0) && (packetOffset == 0) {
        // not a fragment
        //fmt.Println("Payload Length: ", len(p))
        //fmt.Println("Full payload: ", p)
        //fmt.Println("PACKET COMPLETELY READ")
        return ip, b, p, nil
    } else {
        // is a fragment
        bufID := string([]byte{hdr[12], hdr[13], hdr[14], hdr[15], // the source IP
            hdr[16], hdr[17], hdr[18], hdr[19], // the destination IP
            hdr[9],         // protocol
            hdr[4], hdr[5], // identification
        })
        fmt.Printf("rcv a fragment-bufId: %x, len: %d\n", bufID, len(b))

        if c, ok := ipr.fragBuf[bufID]; ok {
            // the fragment has already started
            fmt.Printf("sending to already existing assembler %x\n", bufID)
            go func() { c <- b }()
        } else {
            // create the fragment buffer and quit
            fmt.Printf("creating a new buffer for %x\n", bufID)
            ipr.fragBuf[bufID] = make(chan []byte)
            quit := make(chan bool)
            done := make(chan bool)

            // create the packet assembler in a goroutine to allow the program to continue
            go fragmentAssembler(ipr.fragBuf[bufID], quit, ipr.incomingPackets, done)
            go killFragmentAssembler(quit, done, bufID)

            // send in the first fragment
            ipr.fragBuf[bufID] <- b
        }

        // after dealing with the fragment, try reading again
        //fmt.Println("RECURSE")
        return ipr.ReadFrom()
    }
}

func (ipr *IP_Reader) Close() error {
    return ipr.nr.unbind(ipr.ip, ipr.protocol)
}

/* h := &ipv4.Header{
	Version:  ipv4.Version,      // protocol version
	Len:      20,                // header length
	TOS:      0,                 // type-of-service (0 is everything normal)
	TotalLen: len(x) + 20,       // packet total length (octets)
	ID:       0,                 // identification
	Flags:    ipv4.DontFragment, // flags
	FragOff:  0,                 // fragment offset
	TTL:      8,                 // time-to-live (maximum lifespan in seconds)
	Protocol: 17,                // next protocol (17 is UDP)
	Checksum: 0,                 // checksum (apparently autocomputed)
	//Src:    net.IPv4(127, 0, 0, 1), // source address, apparently done automatically
	Dst: net.ParseIP(c.manager.ipAddress), // destination address
	//Options                         // options, extension headers
}
*/
<|endoftext|>


<commit_before>package torrent

import "time"

type Config struct {
	// Number of unchoked peers.
	UnchokedPeers int
	// Number of optimistic unchoked peers.
	OptimisticUnchokedPeers int
	// Max number of blocks requested from a peer but not received yet
	RequestQueueLength int
	// Time to wait for a requested block to be received before marking peer as snubbed
	RequestTimeout time.Duration
	// Max number of running downloads on piece in endgame mode, snubbed and choed peers don't count
	EndgameParallelDownloadsPerPiece int
	// Max number of outgoing connections to dial
	MaxPeerDial int
	// Max number of incoming connections to accept
	MaxPeerAccept int
	// Running piece downloads, snubbed and choked peers don't count
	ParallelPieceDownloads int
	// Running metadata downloads, snubbed peers don't count
	ParallelMetadataDownloads int
	// Time to wait for TCP connection to open.
	PeerConnectTimeout time.Duration
	// Time to wait for BitTorrent handshake to complete.
	PeerHandshakeTimeout time.Duration
	// When peer has started to send piece block, if it does not send any bytes in PieceTimeout, the connection is closed.
	PieceTimeout time.Duration
	// Buffer size for messages read from a single peer
	PeerReadBufferSize int
	// When the client want to connect a peer, first it tries to do encrypted handshake.
	// If it does not work, it connects to same peer again and does unencrypted handshake.
	// This behavior can be changed via this variable.
	DisableOutgoingEncryption bool
	// Dial only encrypted connections.
	ForceOutgoingEncryption bool
	// Do not accept unencrypted connections.
	ForceIncomingEncryption bool
	// Number of peer addresses to request in announce request.
	TrackerNumWant int
	// Time to wait for announcing stopped event.
	// Stopped event is sent to the tracker when torrent is stopped.
	TrackerStopTimeout time.Duration
	// When the client needs new peer addresses to connect, it ask to the tracker.
	// To prevent spamming the tracker an interval is set to wait before the next announce.
	MinAnnounceInterval time.Duration
	// Total time to wait for response to be read.
	// This includes ConnectTimeout and TLSHandshakeTimeout.
	HTTPTrackerTimeout time.Duration
	// DHT announce interval
	DHTAnnounceInterval time.Duration
	// Minimum announce interval when announcing to DHT.
	MinDHTAnnounceInterval time.Duration
	// Bitfield is saved to disk for fast resume without hash checking.
	// There is an interval to keep IO lower.
	BitfieldWriteInterval time.Duration
	// Number of bytes to read when a piece is requested by a peer.
	PieceReadSize int64
	// Number of cached bytes for piece read requests.
	PieceCacheSize int64
	// Read bytes for a piece part expires after duration.
	PieceCacheTTL time.Duration
}

var DefaultConfig = Config{
	UnchokedPeers:                    3,
	OptimisticUnchokedPeers:          1,
	RequestQueueLength:               50,
	RequestTimeout:                   20 * time.Second,
	EndgameParallelDownloadsPerPiece: 2,
	MaxPeerDial:                      25,
	MaxPeerAccept:                    25,
	ParallelPieceDownloads:           50,
	ParallelMetadataDownloads:        2,
	PeerConnectTimeout:               5 * time.Second,
	PeerHandshakeTimeout:             10 * time.Second,
	PieceTimeout:                     30 * time.Second,
	PeerReadBufferSize:               13 + 16*1024,
	TrackerNumWant:                   100,
	TrackerStopTimeout:               5 * time.Second,
	MinAnnounceInterval:              time.Minute,
	HTTPTrackerTimeout:               10 * time.Second,
	DHTAnnounceInterval:              30 * time.Minute,
	MinDHTAnnounceInterval:           time.Minute,
	BitfieldWriteInterval:            time.Minute,
	PieceReadSize:                    256 * 1024,
	PieceCacheSize:                   50 * 256 * 1024,
	PieceCacheTTL:                    5 * time.Minute,
}
<commit_msg>adjust default config values<commit_after>package torrent

import "time"

type Config struct {
	// Number of unchoked peers.
	UnchokedPeers int
	// Number of optimistic unchoked peers.
	OptimisticUnchokedPeers int
	// Max number of blocks requested from a peer but not received yet
	RequestQueueLength int
	// Time to wait for a requested block to be received before marking peer as snubbed
	RequestTimeout time.Duration
	// Max number of running downloads on piece in endgame mode, snubbed and choed peers don't count
	EndgameParallelDownloadsPerPiece int
	// Max number of outgoing connections to dial
	MaxPeerDial int
	// Max number of incoming connections to accept
	MaxPeerAccept int
	// Running piece downloads, snubbed and choked peers don't count
	ParallelPieceDownloads int
	// Running metadata downloads, snubbed peers don't count
	ParallelMetadataDownloads int
	// Time to wait for TCP connection to open.
	PeerConnectTimeout time.Duration
	// Time to wait for BitTorrent handshake to complete.
	PeerHandshakeTimeout time.Duration
	// When peer has started to send piece block, if it does not send any bytes in PieceTimeout, the connection is closed.
	PieceTimeout time.Duration
	// Buffer size for messages read from a single peer
	PeerReadBufferSize int
	// When the client want to connect a peer, first it tries to do encrypted handshake.
	// If it does not work, it connects to same peer again and does unencrypted handshake.
	// This behavior can be changed via this variable.
	DisableOutgoingEncryption bool
	// Dial only encrypted connections.
	ForceOutgoingEncryption bool
	// Do not accept unencrypted connections.
	ForceIncomingEncryption bool
	// Number of peer addresses to request in announce request.
	TrackerNumWant int
	// Time to wait for announcing stopped event.
	// Stopped event is sent to the tracker when torrent is stopped.
	TrackerStopTimeout time.Duration
	// When the client needs new peer addresses to connect, it ask to the tracker.
	// To prevent spamming the tracker an interval is set to wait before the next announce.
	MinAnnounceInterval time.Duration
	// Total time to wait for response to be read.
	// This includes ConnectTimeout and TLSHandshakeTimeout.
	HTTPTrackerTimeout time.Duration
	// DHT announce interval
	DHTAnnounceInterval time.Duration
	// Minimum announce interval when announcing to DHT.
	MinDHTAnnounceInterval time.Duration
	// Bitfield is saved to disk for fast resume without hash checking.
	// There is an interval to keep IO lower.
	BitfieldWriteInterval time.Duration
	// Number of bytes to read when a piece is requested by a peer.
	PieceReadSize int64
	// Number of cached bytes for piece read requests.
	PieceCacheSize int64
	// Read bytes for a piece part expires after duration.
	PieceCacheTTL time.Duration
}

var DefaultConfig = Config{
	UnchokedPeers:                    3,
	OptimisticUnchokedPeers:          1,
	RequestQueueLength:               50,
	RequestTimeout:                   20 * time.Second,
	EndgameParallelDownloadsPerPiece: 2,
	MaxPeerDial:                      20,
	MaxPeerAccept:                    5,
	ParallelPieceDownloads:           10,
	ParallelMetadataDownloads:        2,
	PeerConnectTimeout:               5 * time.Second,
	PeerHandshakeTimeout:             10 * time.Second,
	PieceTimeout:                     30 * time.Second,
	PeerReadBufferSize:               13 + 16*1024,
	TrackerNumWant:                   100,
	TrackerStopTimeout:               5 * time.Second,
	MinAnnounceInterval:              time.Minute,
	HTTPTrackerTimeout:               10 * time.Second,
	DHTAnnounceInterval:              30 * time.Minute,
	MinDHTAnnounceInterval:           time.Minute,
	BitfieldWriteInterval:            time.Minute,
	PieceReadSize:                    256 * 1024,
	PieceCacheSize:                   50 * 256 * 1024,
	PieceCacheTTL:                    5 * time.Minute,
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"github.com/keybase/go/libkb/engine"
	"github.com/keybase/protocol/go"
	"github.com/maxtaco/go-framed-msgpack-rpc/rpc2"
)

type RemoteBaseIdentifyUI struct {
	sessionId int
	username  string
	uicli     keybase_1.IdentifyUiClient
	logcli    keybase_1.LogUiClient
}

type RemoteSelfIdentifyUI struct {
	RemoteBaseIdentifyUI
}

type IdentifyHandler struct {
	BaseHandler
}

func NewIdentifyHandler(xp *rpc2.Transport) *IdentifyHandler {
	return &IdentifyHandler{BaseHandler{xp: xp}}
}

func (h *IdentifyHandler) Identify(arg keybase_1.IdentifyArg) (keybase_1.IdentifyRes, error) {
	iarg := engine.ImportIdentifyArg(arg)
	res, err := h.identify(iarg, true)
	if err != nil {
		return keybase_1.IdentifyRes{}, err
	}
	return *(res.Export()), nil
}

func (h *IdentifyHandler) IdentifyDefault(username string) (keybase_1.IdentifyRes, error) {
	arg := engine.IdentifyArgPrime{User: username}
	res, err := h.identify(arg, true)
	if err != nil {
		return keybase_1.IdentifyRes{}, err
	}
	return *(res.Export()), nil
}

func (h *IdentifyHandler) identify(iarg engine.IdentifyArgPrime, doInteractive bool) (*engine.IdentifyRes, error) {
	sessionId := nextSessionId()
	iarg.LogUI = h.getLogUI(sessionId)

	eng := engine.NewIdentifyEng(&iarg, NewRemoteIdentifyUI(sessionId, iarg.User, h.cli))
	return eng.Run()
}

var (
	__sessionId = 0
)

func nextSessionId() int {
	ret := __sessionId
	__sessionId++
	return ret
}

func NewRemoteSelfIdentifyUI(sessionId int, username string, c *rpc2.Client) *RemoteSelfIdentifyUI {
	fmt.Printf("************** NewRemoteSelfIdentifyUI\n")
	return &RemoteSelfIdentifyUI{RemoteBaseIdentifyUI{
		sessionId: sessionId,
		username:  username,
		uicli:     keybase_1.IdentifyUiClient{c},
	}}
}

func (u *RemoteBaseIdentifyUI) FinishWebProofCheck(p keybase_1.RemoteProof, lcr keybase_1.LinkCheckResult) {
	u.uicli.FinishWebProofCheck(keybase_1.FinishWebProofCheckArg{
		SessionId: u.sessionId,
		Rp:        p,
		Lcr:       lcr,
	})
	return
}

func (u *RemoteBaseIdentifyUI) FinishSocialProofCheck(p keybase_1.RemoteProof, lcr keybase_1.LinkCheckResult) {
	u.uicli.FinishSocialProofCheck(keybase_1.FinishSocialProofCheckArg{
		SessionId: u.sessionId,
		Rp:        p,
		Lcr:       lcr,
	})
	return
}

func (u *RemoteBaseIdentifyUI) FinishAndPrompt(io *keybase_1.IdentifyOutcome) (keybase_1.FinishAndPromptRes, error) {
	return u.uicli.FinishAndPrompt(keybase_1.FinishAndPromptArg{SessionId: u.sessionId, Outcome: *io})
}

func (u *RemoteBaseIdentifyUI) DisplayCryptocurrency(c keybase_1.Cryptocurrency) {
	u.uicli.DisplayCryptocurrency(keybase_1.DisplayCryptocurrencyArg{SessionId: u.sessionId, C: c})
	return
}

func (u *RemoteBaseIdentifyUI) DisplayKey(k keybase_1.FOKID, d *keybase_1.TrackDiff) {
	u.uicli.DisplayKey(keybase_1.DisplayKeyArg{SessionId: u.sessionId, Fokid: k, Diff: d})
	return
}

func (u *RemoteBaseIdentifyUI) ReportLastTrack(t *keybase_1.TrackSummary) {
	u.uicli.ReportLastTrack(keybase_1.ReportLastTrackArg{SessionId: u.sessionId, Track: t})
	return
}

func (u *RemoteBaseIdentifyUI) DisplayTrackStatement(s string) error {
	return u.uicli.DisplayTrackStatement(keybase_1.DisplayTrackStatementArg{Stmt: s, SessionId: u.sessionId})
	// return
}

func (u *RemoteBaseIdentifyUI) LaunchNetworkChecks(id *keybase_1.Identity) {
	u.uicli.LaunchNetworkChecks(keybase_1.LaunchNetworkChecksArg{
		SessionId: u.sessionId,
		Id:        *id,
	})
	return
}

type RemoteIdentifyUI struct {
	RemoteBaseIdentifyUI
}

func NewRemoteIdentifyUI(sessionId int, username string, c *rpc2.Client) *RemoteIdentifyUI {
	return &RemoteIdentifyUI{RemoteBaseIdentifyUI{
		sessionId: sessionId,
		username:  username,
		uicli:     keybase_1.IdentifyUiClient{c},
		logcli:    keybase_1.LogUiClient{c},
	}}
}

func (u *RemoteBaseIdentifyUI) Start() {
	u.logcli.Log(keybase_1.LogArg{SessionId: u.sessionId, Level: keybase_1.LogLevel_INFO, Text: keybase_1.Text{Data: "Identifying " + u.username}})
	return
}

func (u *RemoteBaseIdentifyUI) SetUsername(username string) {
	u.username = username
}
<commit_msg>debug cleanup<commit_after>package main

import (
	"github.com/keybase/go/libkb/engine"
	"github.com/keybase/protocol/go"
	"github.com/maxtaco/go-framed-msgpack-rpc/rpc2"
)

type RemoteBaseIdentifyUI struct {
	sessionId int
	username  string
	uicli     keybase_1.IdentifyUiClient
	logcli    keybase_1.LogUiClient
}

type RemoteSelfIdentifyUI struct {
	RemoteBaseIdentifyUI
}

type IdentifyHandler struct {
	BaseHandler
}

func NewIdentifyHandler(xp *rpc2.Transport) *IdentifyHandler {
	return &IdentifyHandler{BaseHandler{xp: xp}}
}

func (h *IdentifyHandler) Identify(arg keybase_1.IdentifyArg) (keybase_1.IdentifyRes, error) {
	iarg := engine.ImportIdentifyArg(arg)
	res, err := h.identify(iarg, true)
	if err != nil {
		return keybase_1.IdentifyRes{}, err
	}
	return *(res.Export()), nil
}

func (h *IdentifyHandler) IdentifyDefault(username string) (keybase_1.IdentifyRes, error) {
	arg := engine.IdentifyArgPrime{User: username}
	res, err := h.identify(arg, true)
	if err != nil {
		return keybase_1.IdentifyRes{}, err
	}
	return *(res.Export()), nil
}

func (h *IdentifyHandler) identify(iarg engine.IdentifyArgPrime, doInteractive bool) (*engine.IdentifyRes, error) {
	sessionId := nextSessionId()
	iarg.LogUI = h.getLogUI(sessionId)

	eng := engine.NewIdentifyEng(&iarg, NewRemoteIdentifyUI(sessionId, iarg.User, h.cli))
	return eng.Run()
}

var (
	__sessionId = 0
)

func nextSessionId() int {
	ret := __sessionId
	__sessionId++
	return ret
}

func NewRemoteSelfIdentifyUI(sessionId int, username string, c *rpc2.Client) *RemoteSelfIdentifyUI {
	return &RemoteSelfIdentifyUI{RemoteBaseIdentifyUI{
		sessionId: sessionId,
		username:  username,
		uicli:     keybase_1.IdentifyUiClient{c},
	}}
}

func (u *RemoteBaseIdentifyUI) FinishWebProofCheck(p keybase_1.RemoteProof, lcr keybase_1.LinkCheckResult) {
	u.uicli.FinishWebProofCheck(keybase_1.FinishWebProofCheckArg{
		SessionId: u.sessionId,
		Rp:        p,
		Lcr:       lcr,
	})
	return
}

func (u *RemoteBaseIdentifyUI) FinishSocialProofCheck(p keybase_1.RemoteProof, lcr keybase_1.LinkCheckResult) {
	u.uicli.FinishSocialProofCheck(keybase_1.FinishSocialProofCheckArg{
		SessionId: u.sessionId,
		Rp:        p,
		Lcr:       lcr,
	})
	return
}

func (u *RemoteBaseIdentifyUI) FinishAndPrompt(io *keybase_1.IdentifyOutcome) (keybase_1.FinishAndPromptRes, error) {
	return u.uicli.FinishAndPrompt(keybase_1.FinishAndPromptArg{SessionId: u.sessionId, Outcome: *io})
}

func (u *RemoteBaseIdentifyUI) DisplayCryptocurrency(c keybase_1.Cryptocurrency) {
	u.uicli.DisplayCryptocurrency(keybase_1.DisplayCryptocurrencyArg{SessionId: u.sessionId, C: c})
	return
}

func (u *RemoteBaseIdentifyUI) DisplayKey(k keybase_1.FOKID, d *keybase_1.TrackDiff) {
	u.uicli.DisplayKey(keybase_1.DisplayKeyArg{SessionId: u.sessionId, Fokid: k, Diff: d})
	return
}

func (u *RemoteBaseIdentifyUI) ReportLastTrack(t *keybase_1.TrackSummary) {
	u.uicli.ReportLastTrack(keybase_1.ReportLastTrackArg{SessionId: u.sessionId, Track: t})
	return
}

func (u *RemoteBaseIdentifyUI) DisplayTrackStatement(s string) error {
	return u.uicli.DisplayTrackStatement(keybase_1.DisplayTrackStatementArg{Stmt: s, SessionId: u.sessionId})
	// return
}

func (u *RemoteBaseIdentifyUI) LaunchNetworkChecks(id *keybase_1.Identity) {
	u.uicli.LaunchNetworkChecks(keybase_1.LaunchNetworkChecksArg{
		SessionId: u.sessionId,
		Id:        *id,
	})
	return
}

type RemoteIdentifyUI struct {
	RemoteBaseIdentifyUI
}

func NewRemoteIdentifyUI(sessionId int, username string, c *rpc2.Client) *RemoteIdentifyUI {
	return &RemoteIdentifyUI{RemoteBaseIdentifyUI{
		sessionId: sessionId,
		username:  username,
		uicli:     keybase_1.IdentifyUiClient{c},
		logcli:    keybase_1.LogUiClient{c},
	}}
}

func (u *RemoteBaseIdentifyUI) Start() {
	u.logcli.Log(keybase_1.LogArg{SessionId: u.sessionId, Level: keybase_1.LogLevel_INFO, Text: keybase_1.Text{Data: "Identifying " + u.username}})
	return
}

func (u *RemoteBaseIdentifyUI) SetUsername(username string) {
	u.username = username
}
<|endoftext|>


<commit_before>package transfer

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

	"github.com/github/git-lfs/api"
	"github.com/github/git-lfs/config"
	"github.com/github/git-lfs/errutil"
	"github.com/github/git-lfs/httputil"
	"github.com/github/git-lfs/tools"

	"github.com/github/git-lfs/progress"
)

// Base implementation of basic all-or-nothing HTTP upload / download adapter
type basicAdapter struct {
	direction Direction
	jobChan   chan *Transfer
	cb        progress.CopyCallback
	outChan   chan TransferResult
	// WaitGroup to sync the completion of all workers
	workerWait sync.WaitGroup
	// WaitGroup to serialise the first transfer response to perform login if needed
	authWait sync.WaitGroup
}

func newBasicAdapter(d Direction) *basicAdapter {
	return &basicAdapter{
		direction: d,
	}
}

func (a *basicAdapter) Direction() Direction {
	return a.direction
}

func (a *basicAdapter) Name() string {
	return "basic"
}

func (a *basicAdapter) Begin(cb progress.CopyCallback, completion chan TransferResult) error {
	a.cb = cb
	a.outChan = completion
	a.jobChan = make(chan *Transfer, 100)

	numworkers := config.Config.ConcurrentTransfers()
	a.workerWait.Add(numworkers)
	a.authWait.Add(1)
	for i := 0; i < numworkers; i++ {
		go a.worker(i)
	}
	return nil
}

func (a *basicAdapter) Add(t *Transfer) {
	a.jobChan <- t
}

func (a *basicAdapter) End() {
	a.jobChan.Close()
	// wait for all transfers to complete
	a.workerWait.Wait()
}

func (a *basicAdapter) ClearTempStorage() error {
	// TODO @sinbad
	return nil
}

// worker function, many of these run per adapter
func (a *basicAdapter) worker(workerNum int) {

	isFirstWorker := workerNum == 0
	signalAuthOnResponse := isFirstWorker

	for t := range a.jobChan {
		if !isFirstWorker {
			// First worker is the only one allowed to start immediately
			// The rest wait until successful response from 1st worker to
			// make sure only 1 login prompt is presented if necessary
			a.authWait.Wait()
		}
		var err error
		switch a.Direction() {
		case Download:
			err = a.download(t, signalAuthOnResponse)
		case Upload:
			err = a.upload(t, signalAuthOnResponse)
		}

		res := TransferResult{t, err}
		a.outChan <- res

		signalAuthOnResponse = false
	}
	a.workerWait.Done()
}

func (a *basicAdapter) tempDir() string {
	return filepath.Join(os.TempDir(), "git-lfs-basic")
}

func (a *basicAdapter) download(t *Transfer, signalAuthOnResponse bool) error {
	rel, ok := t.Object.Rel("download")
	if !ok {
		return errors.New("Object not found on the server.")
	}

	req, err := httputil.NewHttpRequest("GET", rel.Href, rel.Header)
	if err != nil {
		return err
	}

	res, err := httputil.DoHttpRequest(req, true)
	if err != nil {
		return errutil.NewRetriableError(err)
	}
	httputil.LogTransfer("lfs.data.download", res)
	defer res.Body.Close()

	// Signal auth OK on success response, before starting download to free up
	// other workers immediately
	if signalAuthOnResponse {
		a.authWait.Done()
	}

	// Now do transfer of content
	f, err := ioutil.TempFile(a.tempDir(), t.Object.Oid+"-")
	if err != nil {
		return fmt.Errorf("cannot create temp file: %v", err)
	}

	defer func() {
		if err != nil {
			// Don't leave the temp file lying around on error.
			_ = os.Remove(f.Name()) // yes, ignore the error, not much we can do about it.
		}
	}()

	hasher := tools.NewHashingReader(res.Body)

	// ensure we always close f. Note that this does not conflict with  the
	// close below, as close is idempotent.
	defer f.Close()
	tempfilename := f.Name()
	written, err := tools.CopyWithCallback(f, hasher, res.ContentLength, a.cb)
	if err != nil {
		return fmt.Errorf("cannot write data to tempfile %q: %v", tempfilename, err)
	}
	if err := f.Close(); err != nil {
		return fmt.Errorf("can't close tempfile %q: %v", tempfilename, err)
	}

	if actual := hasher.Hash(); actual != t.Object.Oid {
		return fmt.Errorf("Expected OID %s, got %s after %d bytes written", t.Object.Oid, actual, written)
	}

	return tools.RenameFileCopyPermissions(tempfilename, t.Path)

}
func (a *basicAdapter) upload(t *Transfer, signalAuthOnResponse bool) error {
	rel, ok := t.Object.Rel("upload")
	if !ok {
		return fmt.Errorf("No upload action for this object.")
	}

	req, err := httputil.NewHttpRequest("PUT", rel.Href, rel.Header)
	if err != nil {
		return err
	}

	if len(req.Header.Get("Content-Type")) == 0 {
		req.Header.Set("Content-Type", "application/octet-stream")
	}

	if req.Header.Get("Transfer-Encoding") == "chunked" {
		req.TransferEncoding = []string{"chunked"}
	} else {
		req.Header.Set("Content-Length", strconv.FormatInt(t.Object.Size, 10))
	}

	req.ContentLength = t.Object.Size

	f, err := os.OpenFile(t.Path, os.O_RDONLY, 0644)
	if err != nil {
		return errutil.Error(err)
	}
	defer f.Close()

	// Ensure progress callbacks made while uploading
	var reader io.Reader
	reader = &progress.CallbackReader{
		C:         a.cb,
		TotalSize: t.Object.Size,
		Reader:    f,
	}

	if signalAuthOnResponse {
		// Signal auth was ok on first read; this frees up other workers to start
		reader = newStartCallbackReader(reader, func(*startCallbackReader) {
			a.authWait.Done()
		})
	}

	req.Body = ioutil.NopCloser(reader)

	res, err := httputil.DoHttpRequest(req, true)
	if err != nil {
		return errutil.NewRetriableError(err)
	}
	httputil.LogTransfer("lfs.data.upload", res)

	// A status code of 403 likely means that an authentication token for the
	// upload has expired. This can be safely retried.
	if res.StatusCode == 403 {
		return errutil.NewRetriableError(err)
	}

	if res.StatusCode > 299 {
		return errutil.Errorf(nil, "Invalid status for %s: %d", httputil.TraceHttpReq(req), res.StatusCode)
	}

	io.Copy(ioutil.Discard, res.Body)
	res.Body.Close()

	return api.VerifyUpload(t.Object)
}

// startCallbackReader is a reader wrapper which calls a function as soon as the
// first Read() call is made. This callback is only made once
type startCallbackReader struct {
	r      io.Reader
	cb     func(*startCallbackReader)
	cbDone bool
}

func (s *startCallbackReader) Read(p []byte) (n int, err error) {
	if !s.cbDone && s.cb != nil {
		s.cb(s)
		s.cbDone = true
	}
	return s.r.Read(p)
}
func newStartCallbackReader(r io.Reader, cb func(*startCallbackReader)) *startCallbackReader {
	return &startCallbackReader{r, cb, false}
}

func init() {
	ul := newBasicAdapter(Upload)
	RegisterAdapter(ul)
	dl := newBasicAdapter(Download)
	RegisterAdapter(dl)
}
<commit_msg>ClearTempStorage implementation not needed<commit_after>package transfer

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

	"github.com/github/git-lfs/api"
	"github.com/github/git-lfs/config"
	"github.com/github/git-lfs/errutil"
	"github.com/github/git-lfs/httputil"
	"github.com/github/git-lfs/tools"

	"github.com/github/git-lfs/progress"
)

// Base implementation of basic all-or-nothing HTTP upload / download adapter
type basicAdapter struct {
	direction Direction
	jobChan   chan *Transfer
	cb        progress.CopyCallback
	outChan   chan TransferResult
	// WaitGroup to sync the completion of all workers
	workerWait sync.WaitGroup
	// WaitGroup to serialise the first transfer response to perform login if needed
	authWait sync.WaitGroup
}

func newBasicAdapter(d Direction) *basicAdapter {
	return &basicAdapter{
		direction: d,
	}
}

func (a *basicAdapter) Direction() Direction {
	return a.direction
}

func (a *basicAdapter) Name() string {
	return "basic"
}

func (a *basicAdapter) Begin(cb progress.CopyCallback, completion chan TransferResult) error {
	a.cb = cb
	a.outChan = completion
	a.jobChan = make(chan *Transfer, 100)

	numworkers := config.Config.ConcurrentTransfers()
	a.workerWait.Add(numworkers)
	a.authWait.Add(1)
	for i := 0; i < numworkers; i++ {
		go a.worker(i)
	}
	return nil
}

func (a *basicAdapter) Add(t *Transfer) {
	a.jobChan <- t
}

func (a *basicAdapter) End() {
	a.jobChan.Close()
	// wait for all transfers to complete
	a.workerWait.Wait()
}

func (a *basicAdapter) ClearTempStorage() error {
	// Not actually necessary for this adapter, the only temp files are deleted
	// immediately via defer
	return nil
}

// worker function, many of these run per adapter
func (a *basicAdapter) worker(workerNum int) {

	isFirstWorker := workerNum == 0
	signalAuthOnResponse := isFirstWorker

	for t := range a.jobChan {
		if !isFirstWorker {
			// First worker is the only one allowed to start immediately
			// The rest wait until successful response from 1st worker to
			// make sure only 1 login prompt is presented if necessary
			a.authWait.Wait()
		}
		var err error
		switch a.Direction() {
		case Download:
			err = a.download(t, signalAuthOnResponse)
		case Upload:
			err = a.upload(t, signalAuthOnResponse)
		}

		res := TransferResult{t, err}
		a.outChan <- res

		signalAuthOnResponse = false
	}
	a.workerWait.Done()
}

func (a *basicAdapter) tempDir() string {
	return filepath.Join(os.TempDir(), "git-lfs-basic")
}

func (a *basicAdapter) download(t *Transfer, signalAuthOnResponse bool) error {
	rel, ok := t.Object.Rel("download")
	if !ok {
		return errors.New("Object not found on the server.")
	}

	req, err := httputil.NewHttpRequest("GET", rel.Href, rel.Header)
	if err != nil {
		return err
	}

	res, err := httputil.DoHttpRequest(req, true)
	if err != nil {
		return errutil.NewRetriableError(err)
	}
	httputil.LogTransfer("lfs.data.download", res)
	defer res.Body.Close()

	// Signal auth OK on success response, before starting download to free up
	// other workers immediately
	if signalAuthOnResponse {
		a.authWait.Done()
	}

	// Now do transfer of content
	f, err := ioutil.TempFile(a.tempDir(), t.Object.Oid+"-")
	if err != nil {
		return fmt.Errorf("cannot create temp file: %v", err)
	}

	defer func() {
		if err != nil {
			// Don't leave the temp file lying around on error.
			_ = os.Remove(f.Name()) // yes, ignore the error, not much we can do about it.
		}
	}()

	hasher := tools.NewHashingReader(res.Body)

	// ensure we always close f. Note that this does not conflict with  the
	// close below, as close is idempotent.
	defer f.Close()
	tempfilename := f.Name()
	written, err := tools.CopyWithCallback(f, hasher, res.ContentLength, a.cb)
	if err != nil {
		return fmt.Errorf("cannot write data to tempfile %q: %v", tempfilename, err)
	}
	if err := f.Close(); err != nil {
		return fmt.Errorf("can't close tempfile %q: %v", tempfilename, err)
	}

	if actual := hasher.Hash(); actual != t.Object.Oid {
		return fmt.Errorf("Expected OID %s, got %s after %d bytes written", t.Object.Oid, actual, written)
	}

	return tools.RenameFileCopyPermissions(tempfilename, t.Path)

}
func (a *basicAdapter) upload(t *Transfer, signalAuthOnResponse bool) error {
	rel, ok := t.Object.Rel("upload")
	if !ok {
		return fmt.Errorf("No upload action for this object.")
	}

	req, err := httputil.NewHttpRequest("PUT", rel.Href, rel.Header)
	if err != nil {
		return err
	}

	if len(req.Header.Get("Content-Type")) == 0 {
		req.Header.Set("Content-Type", "application/octet-stream")
	}

	if req.Header.Get("Transfer-Encoding") == "chunked" {
		req.TransferEncoding = []string{"chunked"}
	} else {
		req.Header.Set("Content-Length", strconv.FormatInt(t.Object.Size, 10))
	}

	req.ContentLength = t.Object.Size

	f, err := os.OpenFile(t.Path, os.O_RDONLY, 0644)
	if err != nil {
		return errutil.Error(err)
	}
	defer f.Close()

	// Ensure progress callbacks made while uploading
	var reader io.Reader
	reader = &progress.CallbackReader{
		C:         a.cb,
		TotalSize: t.Object.Size,
		Reader:    f,
	}

	if signalAuthOnResponse {
		// Signal auth was ok on first read; this frees up other workers to start
		reader = newStartCallbackReader(reader, func(*startCallbackReader) {
			a.authWait.Done()
		})
	}

	req.Body = ioutil.NopCloser(reader)

	res, err := httputil.DoHttpRequest(req, true)
	if err != nil {
		return errutil.NewRetriableError(err)
	}
	httputil.LogTransfer("lfs.data.upload", res)

	// A status code of 403 likely means that an authentication token for the
	// upload has expired. This can be safely retried.
	if res.StatusCode == 403 {
		return errutil.NewRetriableError(err)
	}

	if res.StatusCode > 299 {
		return errutil.Errorf(nil, "Invalid status for %s: %d", httputil.TraceHttpReq(req), res.StatusCode)
	}

	io.Copy(ioutil.Discard, res.Body)
	res.Body.Close()

	return api.VerifyUpload(t.Object)
}

// startCallbackReader is a reader wrapper which calls a function as soon as the
// first Read() call is made. This callback is only made once
type startCallbackReader struct {
	r      io.Reader
	cb     func(*startCallbackReader)
	cbDone bool
}

func (s *startCallbackReader) Read(p []byte) (n int, err error) {
	if !s.cbDone && s.cb != nil {
		s.cb(s)
		s.cbDone = true
	}
	return s.r.Read(p)
}
func newStartCallbackReader(r io.Reader, cb func(*startCallbackReader)) *startCallbackReader {
	return &startCallbackReader{r, cb, false}
}

func init() {
	ul := newBasicAdapter(Upload)
	RegisterAdapter(ul)
	dl := newBasicAdapter(Download)
	RegisterAdapter(dl)
}
<|endoftext|>


<commit_before>package queue

import (
	"context"
	"log"
)

// Worker instantiates and calls MessageWorker for every Message in Queue
type Worker struct {
	errChan chan<- error
	queue   *Queue
	factory MessageWorkerFactory
}

// NewWorker returns a worker for a given queue with error channel. The
// MessageWorkerFactory is itself wrapped in a messageWorker for proper
// error handling etc.
func NewWorker(errc chan<- error, queue *Queue, factory MessageWorkerFactory) *Worker {
	return &Worker{
		errChan: errc,
		queue:   queue,
		factory: newMessageWorker(factory),
	}
}

// String returns the name of the worker queue
func (w *Worker) String() string {
	return w.queue.String()
}

// Work performs consumption of messages in the worker's Queue
func (w *Worker) Work(ctx context.Context) error {
	msgs, err := w.queue.Consume()
	if err != nil {
		return err
	}

	// Keep consuming messages until context is cancelled
	for {
		select {
		case <-ctx.Done():
			// Context canceled, stop processing messages
			log.Printf("Stopping worker %s: %s", w, ctx.Err())
			return ctx.Err()
		case msg := <-msgs:
			worker := w.factory(&msg)
			err = worker.Work(ctx)
			if err != nil {
				w.errChan <- err
			}
		}
	}
}
<commit_msg>Fix naming issue.<commit_after>package queue

import (
	"context"
	"log"
)

// Worker instantiates and calls MessageWorker for every Message in Queue
type Worker struct {
	errChan chan<- error
	queue   *Queue
	factory MessageWorkerFactory
}

// NewWorker returns a worker for a given queue with error channel. The
// MessageWorkerFactory is itself wrapped in a messageWorker for proper
// error handling etc.
func NewWorker(errc chan<- error, queue *Queue, factory MessageWorkerFactory) *Worker {
	return &Worker{
		errChan: errc,
		queue:   queue,
		factory: newMessageWorker(factory),
	}
}

// String returns the name of the worker queue
func (w *Worker) String() string {
	return w.queue.String()
}

// Work performs consumption of messages in the worker's Queue
func (w *Worker) Work(ctx context.Context) error {
	msgs, err := w.queue.Consume()
	if err != nil {
		return err
	}

	// Keep consuming messages until context is cancelled
	for {
		select {
		case <-ctx.Done():
			// Context canceled, stop processing messages
			log.Printf("Stopping worker %s: %s", w, ctx.Err())
			return ctx.Err()
		case msg := <-msgs:
			msgWorker := w.factory(&msg)
			err = msgWorker.Work(ctx)
			if err != nil {
				w.errChan <- err
			}
		}
	}
}
<|endoftext|>


<commit_before>// Copyright 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package provisioner_test

import (
	. "launchpad.net/gocheck"
	"launchpad.net/golxc"
	"launchpad.net/juju-core/constraints"
	"launchpad.net/juju-core/container/lxc"
	"launchpad.net/juju-core/container/lxc/mock"
	"launchpad.net/juju-core/instance"
	jujutesting "launchpad.net/juju-core/juju/testing"
	"launchpad.net/juju-core/state"
	"launchpad.net/juju-core/testing"
	"launchpad.net/juju-core/version"
	"launchpad.net/juju-core/worker/provisioner"
)

type lxcBrokerSuite struct {
	testing.LoggingSuite
	golxc              golxc.ContainerFactory
	broker             provisioner.Broker
	containerDir       string
	removedDir         string
	lxcDir             string
	oldContainerDir    string
	oldRemovedDir      string
	oldLxcContainerDir string
}

var _ = Suite(&lxcBrokerSuite{})

func (s *lxcBrokerSuite) SetUpSuite(c *C) {
	s.LoggingSuite.SetUpSuite(c)
}

func (s *lxcBrokerSuite) TearDownSuite(c *C) {
	s.LoggingSuite.TearDownSuite(c)
}

func (s *lxcBrokerSuite) SetUpTest(c *C) {
	s.LoggingSuite.SetUpTest(c)
	s.containerDir = c.MkDir()
	s.oldContainerDir = lxc.SetContainerDir(s.containerDir)
	s.removedDir = c.MkDir()
	s.oldRemovedDir = lxc.SetRemovedContainerDir(s.removedDir)
	s.lxcDir = c.MkDir()
	s.oldLxcContainerDir = lxc.SetLxcContainerDir(s.lxcDir)
	s.golxc = mock.MockFactory()
	tools := &state.Tools{
		Binary: version.MustParseBinary("2.3.4-foo-bar"),
		URL:    "http://tools.example.com/2.3.4-foo-bar.tgz",
	}
	s.broker = provisioner.NewLxcBroker(s.golxc, testing.EnvironConfig(c), tools)
}

func (s *lxcBrokerSuite) TearDownTest(c *C) {
	lxc.SetContainerDir(s.oldContainerDir)
	lxc.SetLxcContainerDir(s.oldLxcContainerDir)
	lxc.SetRemovedContainerDir(s.oldRemovedDir)
	s.LoggingSuite.TearDownTest(c)
}

func (s *lxcBrokerSuite) startInstance(c *C, machineId string) instance.Instance {
	stateInfo := jujutesting.FakeStateInfo(machineId)
	apiInfo := jujutesting.FakeAPIInfo(machineId)

	series := "series"
	nonce := "fake-nonce"
	cons := constraints.Value{}
	lxc, err := s.broker.StartInstance(machineId, nonce, series, cons, stateInfo, apiInfo)
	c.Assert(err, IsNil)
	return lxc
}

func (s *lxcBrokerSuite) TestStartInstance(c *C) {
	machineId := "1/lxc/0"
	lxc := s.startInstance(c, machineId)
	c.Assert(lxc.Id(), Equals, instance.Id("machine-1-lxc-0"))
}

func (s *lxcBrokerSuite) TestStopInstance(c *C) {
	lxc0 := s.startInstance(c, "1/lxc/0")
	lxc1 := s.startInstance(c, "1/lxc/1")
	lxc2 := s.startInstance(c, "1/lxc/2")
	err := s.broker.StopInstances([]instance.Instance{lxc0, lxc1, lxc2})
	c.Assert(err, IsNil)
}
<commit_msg>Check list.<commit_after>// Copyright 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package provisioner_test

import (
	. "launchpad.net/gocheck"
	"launchpad.net/golxc"
	"launchpad.net/juju-core/constraints"
	"launchpad.net/juju-core/container/lxc"
	"launchpad.net/juju-core/container/lxc/mock"
	"launchpad.net/juju-core/instance"
	jujutesting "launchpad.net/juju-core/juju/testing"
	"launchpad.net/juju-core/state"
	"launchpad.net/juju-core/testing"
	"launchpad.net/juju-core/version"
	"launchpad.net/juju-core/worker/provisioner"
)

type lxcBrokerSuite struct {
	testing.LoggingSuite
	golxc              golxc.ContainerFactory
	broker             provisioner.Broker
	containerDir       string
	removedDir         string
	lxcDir             string
	oldContainerDir    string
	oldRemovedDir      string
	oldLxcContainerDir string
}

var _ = Suite(&lxcBrokerSuite{})

func (s *lxcBrokerSuite) SetUpSuite(c *C) {
	s.LoggingSuite.SetUpSuite(c)
}

func (s *lxcBrokerSuite) TearDownSuite(c *C) {
	s.LoggingSuite.TearDownSuite(c)
}

func (s *lxcBrokerSuite) SetUpTest(c *C) {
	s.LoggingSuite.SetUpTest(c)
	s.containerDir = c.MkDir()
	s.oldContainerDir = lxc.SetContainerDir(s.containerDir)
	s.removedDir = c.MkDir()
	s.oldRemovedDir = lxc.SetRemovedContainerDir(s.removedDir)
	s.lxcDir = c.MkDir()
	s.oldLxcContainerDir = lxc.SetLxcContainerDir(s.lxcDir)
	s.golxc = mock.MockFactory()
	tools := &state.Tools{
		Binary: version.MustParseBinary("2.3.4-foo-bar"),
		URL:    "http://tools.example.com/2.3.4-foo-bar.tgz",
	}
	s.broker = provisioner.NewLxcBroker(s.golxc, testing.EnvironConfig(c), tools)
}

func (s *lxcBrokerSuite) TearDownTest(c *C) {
	lxc.SetContainerDir(s.oldContainerDir)
	lxc.SetLxcContainerDir(s.oldLxcContainerDir)
	lxc.SetRemovedContainerDir(s.oldRemovedDir)
	s.LoggingSuite.TearDownTest(c)
}

func (s *lxcBrokerSuite) startInstance(c *C, machineId string) instance.Instance {
	stateInfo := jujutesting.FakeStateInfo(machineId)
	apiInfo := jujutesting.FakeAPIInfo(machineId)

	series := "series"
	nonce := "fake-nonce"
	cons := constraints.Value{}
	lxc, err := s.broker.StartInstance(machineId, nonce, series, cons, stateInfo, apiInfo)
	c.Assert(err, IsNil)
	return lxc
}

func (s *lxcBrokerSuite) TestStartInstance(c *C) {
	machineId := "1/lxc/0"
	lxc := s.startInstance(c, machineId)
	c.Assert(lxc.Id(), Equals, instance.Id("machine-1-lxc-0"))
}

func (s *lxcBrokerSuite) TestStopInstance(c *C) {
	lxc0 := s.startInstance(c, "1/lxc/0")
	lxc1 := s.startInstance(c, "1/lxc/1")
	lxc2 := s.startInstance(c, "1/lxc/2")
	err := s.broker.StopInstances([]instance.Instance{lxc0, lxc1, lxc2})
	c.Assert(err, IsNil)
}

func (s *lxcBrokerSuite) matchInstances(c *C, result []instance.Instance, expected ...instance.Instance) {
	resultMap := make(map[instance.Id]instance.Instance)
	for _, i := range result {
		resultMap[i.Id()] = i
	}

	expectedMap := make(map[instance.Id]instance.Instance)
	for _, i := range expected {
		expectedMap[i.Id()] = i
	}
	c.Assert(resultMap, DeepEquals, expectedMap)
}

func (s *lxcBrokerSuite) TestAllInstances(c *C) {
	lxc0 := s.startInstance(c, "1/lxc/0")
	lxc1 := s.startInstance(c, "1/lxc/1")
	results, err := s.broker.AllInstances()
	c.Assert(err, IsNil)
	s.matchInstances(c, results, lxc0, lxc1)
}
<|endoftext|>


<commit_before>// Copyright 2012, 2013, 2014 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package jujuc

import (
	"fmt"
	"strconv"
	"strings"
	"time"

	"github.com/juju/errors"
	"github.com/juju/names"
	"gopkg.in/juju/charm.v5"

	"github.com/juju/juju/apiserver/params"
	"github.com/juju/juju/network"
	"github.com/juju/juju/storage"
)

type RebootPriority int

const (
	// noop
	RebootSkip RebootPriority = iota
	// wait for current hook to finish before rebooting
	RebootAfterHook
	// reboot immediately, killing and requeueing the calling hook
	RebootNow
)

// Context is the interface that all hook helper commands
// depend on to interact with the rest of the system.
type Context interface {

	// UnitName returns the executing unit's name.
	UnitName() string

	// UnitStatus returns the executing unit's current status.
	UnitStatus() (*StatusInfo, error)

	// SetUnitStatus updates the unit's status.
	SetUnitStatus(StatusInfo) error

	// SetServiceStatus updates the status for the unit's service.
	SetServiceStatus(StatusInfo) error

	// ServiceStatus returns the executing unit's service status (including all units).
	ServiceStatus() (ServiceStatusInfo, error)

	// PublicAddress returns the executing unit's public address.
	PublicAddress() (string, bool)

	// PrivateAddress returns the executing unit's private address.
	PrivateAddress() (string, bool)

	// AvailabilityZone returns the executing unit's availablilty zone.
	AvailabilityZone() (string, bool)

	// OpenPorts marks the supplied port range for opening when the
	// executing unit's service is exposed.
	OpenPorts(protocol string, fromPort, toPort int) error

	// ClosePorts ensures the supplied port range is closed even when
	// the executing unit's service is exposed (unless it is opened
	// separately by a co- located unit).
	ClosePorts(protocol string, fromPort, toPort int) error

	// OpenedPorts returns all port ranges currently opened by this
	// unit on its assigned machine. The result is sorted first by
	// protocol, then by number.
	OpenedPorts() []network.PortRange

	// Config returns the current service configuration of the executing unit.
	ConfigSettings() (charm.Settings, error)

	// IsLeader returns true if the local unit is known to be leader for at
	// least the next 30s.
	IsLeader() (bool, error)

	// LeaderSettings returns the current leader settings. Once leader settings
	// have been read in a given context, they will not be updated other than
	// via successful calls to WriteLeaderSettings.
	LeaderSettings() (map[string]string, error)

	// WriteLeaderSettings writes the supplied settings directly to state, or
	// fails if the local unit is not the service's leader.
	WriteLeaderSettings(map[string]string) error

	// ActionParams returns the map of params passed with an Action.
	ActionParams() (map[string]interface{}, error)

	// UpdateActionResults inserts new values for use with action-set.
	// The results struct will be delivered to the state server upon
	// completion of the Action.
	UpdateActionResults(keys []string, value string) error

	// SetActionMessage sets a message for the Action.
	SetActionMessage(string) error

	// SetActionFailed sets a failure state for the Action.
	SetActionFailed() error

	// HookRelation returns the ContextRelation associated with the executing
	// hook if it was found, and whether it was found.
	HookRelation() (ContextRelation, bool)

	// RemoteUnitName returns the name of the remote unit the hook execution
	// is associated with if it was found, and whether it was found.
	RemoteUnitName() (string, bool)

	// Relation returns the relation with the supplied id if it was found, and
	// whether it was found.
	Relation(id int) (ContextRelation, bool)

	// RelationIds returns the ids of all relations the executing unit is
	// currently participating in.
	RelationIds() []int

	// OwnerTag returns the user tag of the service the executing
	// units belongs to.
	OwnerTag() string

	// AddMetric records a metric to return after hook execution.
	AddMetric(string, string, time.Time) error

	// RequestReboot will set the reboot flag to true on the machine agent
	RequestReboot(prio RebootPriority) error

	// Storage returns the ContextStorageAttachment with the supplied
	// tag if it was found, and whether it was found.
	Storage(names.StorageTag) (ContextStorageAttachment, bool)

	// HookStorageAttachment returns the storage attachment associated
	// the executing hook if it was found, and whether it was found.
	HookStorage() (ContextStorageAttachment, bool)

	// AddUnitStorage saves storage constraints in the context.
	AddUnitStorage(map[string]params.StorageConstraints)
}

// ContextRelation expresses the capabilities of a hook with respect to a relation.
type ContextRelation interface {

	// Id returns an integer which uniquely identifies the relation.
	Id() int

	// Name returns the name the locally executing charm assigned to this relation.
	Name() string

	// FakeId returns a string of the form "relation-name:123", which uniquely
	// identifies the relation to the hook. In reality, the identification
	// of the relation is the integer following the colon, but the composed
	// name is useful to humans observing it.
	FakeId() string

	// Settings allows read/write access to the local unit's settings in
	// this relation.
	Settings() (Settings, error)

	// UnitNames returns a list of the remote units in the relation.
	UnitNames() []string

	// ReadSettings returns the settings of any remote unit in the relation.
	ReadSettings(unit string) (params.Settings, error)
}

// ContextStorageAttachment expresses the capabilities of a hook with
// respect to a storage attachment.
type ContextStorageAttachment interface {

	// Tag returns a tag which uniquely identifies the storage attachment
	// in the context of the unit.
	Tag() names.StorageTag

	// Kind returns the kind of the storage.
	Kind() storage.StorageKind

	// Location returns the location of the storage: the mount point for
	// filesystem-kind stores, and the device path for block-kind stores.
	Location() string
}

// Settings is implemented by types that manipulate unit settings.
type Settings interface {
	Map() params.Settings
	Set(string, string)
	Delete(string)
}

// newRelationIdValue returns a gnuflag.Value for convenient parsing of relation
// ids in ctx.
func newRelationIdValue(ctx Context, result *int) *relationIdValue {
	v := &relationIdValue{result: result, ctx: ctx}
	id := -1
	if r, found := ctx.HookRelation(); found {
		id = r.Id()
		v.value = r.FakeId()
	}
	*result = id
	return v
}

// relationIdValue implements gnuflag.Value for use in relation commands.
type relationIdValue struct {
	result *int
	ctx    Context
	value  string
}

// String returns the current value.
func (v *relationIdValue) String() string {
	return v.value
}

// Set interprets value as a relation id, if possible, and returns an error
// if it is not known to the system. The parsed relation id will be written
// to v.result.
func (v *relationIdValue) Set(value string) error {
	trim := value
	if idx := strings.LastIndex(trim, ":"); idx != -1 {
		trim = trim[idx+1:]
	}
	id, err := strconv.Atoi(trim)
	if err != nil {
		return fmt.Errorf("invalid relation id")
	}
	if _, found := v.ctx.Relation(id); !found {
		return fmt.Errorf("unknown relation id")
	}
	*v.result = id
	v.value = value
	return nil
}

// newStorageIdValue returns a gnuflag.Value for convenient parsing of storage
// ids in ctx.
func newStorageIdValue(ctx Context, result *names.StorageTag) *storageIdValue {
	v := &storageIdValue{result: result, ctx: ctx}
	if s, found := ctx.HookStorage(); found {
		*v.result = s.Tag()
	}
	return v
}

// storageIdValue implements gnuflag.Value for use in storage commands.
type storageIdValue struct {
	result *names.StorageTag
	ctx    Context
}

// String returns the current value.
func (v *storageIdValue) String() string {
	if *v.result == (names.StorageTag{}) {
		return ""
	}
	return v.result.Id()
}

// Set interprets value as a storage id, if possible, and returns an error
// if it is not known to the system. The parsed storage id will be written
// to v.result.
func (v *storageIdValue) Set(value string) error {
	if !names.IsValidStorage(value) {
		return errors.Errorf("invalid storage ID %q", value)
	}
	tag := names.NewStorageTag(value)
	if _, found := v.ctx.Storage(tag); !found {
		return fmt.Errorf("unknown storage ID")
	}
	*v.result = tag
	return nil
}
<commit_msg>Split up the jujuc.Context interface by component.<commit_after>// Copyright 2012, 2013, 2014 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package jujuc

import (
	"fmt"
	"strconv"
	"strings"
	"time"

	"github.com/juju/errors"
	"github.com/juju/names"
	"gopkg.in/juju/charm.v5"

	"github.com/juju/juju/apiserver/params"
	"github.com/juju/juju/network"
	"github.com/juju/juju/storage"
)

type RebootPriority int

const (
	// noop
	RebootSkip RebootPriority = iota
	// wait for current hook to finish before rebooting
	RebootAfterHook
	// reboot immediately, killing and requeueing the calling hook
	RebootNow
)

// Context is the interface that all hook helper commands
// depend on to interact with the rest of the system.
type Context interface {
	HookContext
	relationHookContext
	actionHookContext
}

// HookContext represents the information and functionality that is
// common to all charm hooks.
type HookContext interface {
	ContextUnit
	ContextStatus
	ContextInstance
	ContextNetworking
	ContextLeadership
	ContextMetrics
	ContextStorage
	ContextRelations
}

// UnitHookContext is the context for a unit hook.
type UnitHookContext interface {
	HookContext
}

// RelationHookContext is the context for a relation hook.
type RelationHookContext interface {
	HookContext
	relationHookContext
}

type relationHookContext interface {
	// HookRelation returns the ContextRelation associated with the executing
	// hook if it was found, and whether it was found.
	HookRelation() (ContextRelation, bool)

	// RemoteUnitName returns the name of the remote unit the hook execution
	// is associated with if it was found, and whether it was found.
	RemoteUnitName() (string, bool)
}

// ActionHookContext is the context for an action hook.
type ActionHookContext interface {
	HookContext
	actionHookContext
}

type actionHookContext interface {
	// ActionParams returns the map of params passed with an Action.
	ActionParams() (map[string]interface{}, error)

	// UpdateActionResults inserts new values for use with action-set.
	// The results struct will be delivered to the state server upon
	// completion of the Action.
	UpdateActionResults(keys []string, value string) error

	// SetActionMessage sets a message for the Action.
	SetActionMessage(string) error

	// SetActionFailed sets a failure state for the Action.
	SetActionFailed() error
}

// ContextUnit is the part of a hook context related to the unit.
type ContextUnit interface {
	// UnitName returns the executing unit's name.
	UnitName() string

	// OwnerTag returns the user tag of the service the executing
	// units belongs to.
	OwnerTag() string

	// Config returns the current service configuration of the executing unit.
	ConfigSettings() (charm.Settings, error)
}

// ContextStatus is the part of a hook context related to the unit's status.
type ContextStatus interface {
	// UnitStatus returns the executing unit's current status.
	UnitStatus() (*StatusInfo, error)

	// SetUnitStatus updates the unit's status.
	SetUnitStatus(StatusInfo) error

	// ServiceStatus returns the executing unit's service status
	// (including all units).
	ServiceStatus() (ServiceStatusInfo, error)

	// SetServiceStatus updates the status for the unit's service.
	SetServiceStatus(StatusInfo) error
}

// ContextInstace is the part of a hook context related to the unit's intance.
type ContextInstance interface {
	// AvailabilityZone returns the executing unit's availablilty zone.
	AvailabilityZone() (string, bool)

	// RequestReboot will set the reboot flag to true on the machine agent
	RequestReboot(prio RebootPriority) error
}

// ContextNetworking is the part of a hook context related to network
// interface of the unit's instance.
type ContextNetworking interface {
	// PublicAddress returns the executing unit's public address.
	PublicAddress() (string, bool)

	// PrivateAddress returns the executing unit's private address.
	PrivateAddress() (string, bool)

	// OpenPorts marks the supplied port range for opening when the
	// executing unit's service is exposed.
	OpenPorts(protocol string, fromPort, toPort int) error

	// ClosePorts ensures the supplied port range is closed even when
	// the executing unit's service is exposed (unless it is opened
	// separately by a co- located unit).
	ClosePorts(protocol string, fromPort, toPort int) error

	// OpenedPorts returns all port ranges currently opened by this
	// unit on its assigned machine. The result is sorted first by
	// protocol, then by number.
	OpenedPorts() []network.PortRange
}

// ContextLeadership is the part of a hook context related to the
// unit leadership.
type ContextLeadership interface {
	// IsLeader returns true if the local unit is known to be leader for at
	// least the next 30s.
	IsLeader() (bool, error)

	// LeaderSettings returns the current leader settings. Once leader settings
	// have been read in a given context, they will not be updated other than
	// via successful calls to WriteLeaderSettings.
	LeaderSettings() (map[string]string, error)

	// WriteLeaderSettings writes the supplied settings directly to state, or
	// fails if the local unit is not the service's leader.
	WriteLeaderSettings(map[string]string) error
}

// ContextLeadership is the part of a hook context related to metrics.
type ContextMetrics interface {
	// AddMetric records a metric to return after hook execution.
	AddMetric(string, string, time.Time) error
}

// ContextStorage is the part of a hook context related to storage
// resources associated with the unit.
type ContextStorage interface {
	// Storage returns the ContextStorageAttachment with the supplied
	// tag if it was found, and whether it was found.
	Storage(names.StorageTag) (ContextStorageAttachment, bool)

	// HookStorage returns the storage attachment associated
	// the executing hook if it was found, and whether it was found.
	HookStorage() (ContextStorageAttachment, bool)

	// AddUnitStorage saves storage constraints in the context.
	AddUnitStorage(map[string]params.StorageConstraints)
}

// ContextRelations exposes the relations associated with the unit.
type ContextRelations interface {
	// Relation returns the relation with the supplied id if it was found, and
	// whether it was found.
	Relation(id int) (ContextRelation, bool)

	// RelationIds returns the ids of all relations the executing unit is
	// currently participating in.
	RelationIds() []int
}

// ContextRelation expresses the capabilities of a hook with respect to a relation.
type ContextRelation interface {

	// Id returns an integer which uniquely identifies the relation.
	Id() int

	// Name returns the name the locally executing charm assigned to this relation.
	Name() string

	// FakeId returns a string of the form "relation-name:123", which uniquely
	// identifies the relation to the hook. In reality, the identification
	// of the relation is the integer following the colon, but the composed
	// name is useful to humans observing it.
	FakeId() string

	// Settings allows read/write access to the local unit's settings in
	// this relation.
	Settings() (Settings, error)

	// UnitNames returns a list of the remote units in the relation.
	UnitNames() []string

	// ReadSettings returns the settings of any remote unit in the relation.
	ReadSettings(unit string) (params.Settings, error)
}

// ContextStorageAttachment expresses the capabilities of a hook with
// respect to a storage attachment.
type ContextStorageAttachment interface {

	// Tag returns a tag which uniquely identifies the storage attachment
	// in the context of the unit.
	Tag() names.StorageTag

	// Kind returns the kind of the storage.
	Kind() storage.StorageKind

	// Location returns the location of the storage: the mount point for
	// filesystem-kind stores, and the device path for block-kind stores.
	Location() string
}

// Settings is implemented by types that manipulate unit settings.
type Settings interface {
	Map() params.Settings
	Set(string, string)
	Delete(string)
}

// newRelationIdValue returns a gnuflag.Value for convenient parsing of relation
// ids in ctx.
func newRelationIdValue(ctx Context, result *int) *relationIdValue {
	v := &relationIdValue{result: result, ctx: ctx}
	id := -1
	if r, found := ctx.HookRelation(); found {
		id = r.Id()
		v.value = r.FakeId()
	}
	*result = id
	return v
}

// relationIdValue implements gnuflag.Value for use in relation commands.
type relationIdValue struct {
	result *int
	ctx    Context
	value  string
}

// String returns the current value.
func (v *relationIdValue) String() string {
	return v.value
}

// Set interprets value as a relation id, if possible, and returns an error
// if it is not known to the system. The parsed relation id will be written
// to v.result.
func (v *relationIdValue) Set(value string) error {
	trim := value
	if idx := strings.LastIndex(trim, ":"); idx != -1 {
		trim = trim[idx+1:]
	}
	id, err := strconv.Atoi(trim)
	if err != nil {
		return fmt.Errorf("invalid relation id")
	}
	if _, found := v.ctx.Relation(id); !found {
		return fmt.Errorf("unknown relation id")
	}
	*v.result = id
	v.value = value
	return nil
}

// newStorageIdValue returns a gnuflag.Value for convenient parsing of storage
// ids in ctx.
func newStorageIdValue(ctx Context, result *names.StorageTag) *storageIdValue {
	v := &storageIdValue{result: result, ctx: ctx}
	if s, found := ctx.HookStorage(); found {
		*v.result = s.Tag()
	}
	return v
}

// storageIdValue implements gnuflag.Value for use in storage commands.
type storageIdValue struct {
	result *names.StorageTag
	ctx    Context
}

// String returns the current value.
func (v *storageIdValue) String() string {
	if *v.result == (names.StorageTag{}) {
		return ""
	}
	return v.result.Id()
}

// Set interprets value as a storage id, if possible, and returns an error
// if it is not known to the system. The parsed storage id will be written
// to v.result.
func (v *storageIdValue) Set(value string) error {
	if !names.IsValidStorage(value) {
		return errors.Errorf("invalid storage ID %q", value)
	}
	tag := names.NewStorageTag(value)
	if _, found := v.ctx.Storage(tag); !found {
		return fmt.Errorf("unknown storage ID")
	}
	*v.result = tag
	return nil
}
<|endoftext|>


<commit_before>package table

import (
	"bytes"
	"compress/gzip"
	"gnd.la/encoding/binary"
	"io"
	"strings"
)

type Formula func(n int) int

type Table struct {
	translations map[string]Translation
	formula      Formula
}

func (t *Table) Singular(ctx string, msg string) string {
	k := Key(ctx, msg, "")
	if tr := t.translations[k]; len(tr) > 0 {
		return tr[0]
	}
	return msg
}

func (t *Table) Plural(ctx string, singular string, plural string, n int) string {
	k := Key(ctx, singular, plural)
	if tr := t.translations[k]; tr != nil {
		ii := t.formula(n)
		if ii < len(tr) {
			return tr[ii]
		}
	}
	if n == 1 {
		return singular
	}
	return plural
}

func (t *Table) Encode() (string, error) {
	var buf bytes.Buffer
	w, err := gzip.NewWriterLevel(&buf, gzip.BestCompression)
	if err != nil {
		return "", err
	}
	if err := binary.Write(w, binary.BigEndian, int32(len(t.translations))); err != nil {
		return "", err
	}
	for k, v := range t.translations {
		if err := writeString(w, k); err != nil {
			return "", err
		}
		if err := binary.Write(w, binary.BigEndian, int32(len(v))); err != nil {
			return "", err
		}
		for _, s := range v {
			if err := writeString(w, s); err != nil {
				return "", err
			}
		}
	}
	if err := w.Close(); err != nil {
		return "", err
	}
	return buf.String(), nil
}

func (t *Table) Update(other *Table) error {
	if other.formula != nil {
		t.formula = other.formula
	}
	for k, v := range other.translations {
		t.translations[k] = v
	}
	return nil
}

func Decode(data string) (*Table, error) {
	r, err := gzip.NewReader(strings.NewReader(data))
	if err != nil {
		return nil, err
	}
	defer r.Close()
	var count int32
	if err := binary.Read(r, binary.BigEndian, &count); err != nil {
		return nil, err
	}
	c := int(count)
	translations := make(map[string]Translation, c)
	for ii := 0; ii < c; ii++ {
		key, err := readString(r)
		if err != nil {
			return nil, err
		}
		var tcount int32
		if err := binary.Read(r, binary.BigEndian, &tcount); err != nil {
			return nil, err
		}
		tc := int(tcount)
		value := make([]string, tc)
		for jj := 0; jj < tc; jj++ {
			tr, err := readString(r)
			if err != nil {
				return nil, err
			}
			value[jj] = tr
		}
		translations[key] = value
	}
	return &Table{
		translations: translations,
	}, nil
}

func New(formula Formula, translations map[string]Translation) (*Table, error) {
	return &Table{
		translations: translations,
		formula:      formula,
	}, nil
}

func writeString(w io.Writer, s string) error {
	b := []byte(s)
	if err := binary.Write(w, binary.BigEndian, int32(len(b))); err != nil {
		return err
	}
	_, err := w.Write(b)
	return err
}

func readString(r io.Reader) (string, error) {
	var s int32
	if err := binary.Read(r, binary.BigEndian, &s); err != nil {
		return "", err
	}
	b := make([]byte, int(s))
	if _, err := r.Read(b); err != nil {
		return "", err
	}
	return string(b), nil
}
<commit_msg>Sort imports<commit_after>package table

import (
	"bytes"
	"compress/gzip"
	"fmt"
	"io"
	"strings"

	"gnd.la/encoding/binary"
)

type Formula func(n int) int

type Table struct {
	translations map[string]Translation
	formula      Formula
}

func (t *Table) Singular(ctx string, msg string) string {
	k := Key(ctx, msg, "")
	if tr := t.translations[k]; len(tr) > 0 {
		return tr[0]
	}
	return msg
}

func (t *Table) Plural(ctx string, singular string, plural string, n int) string {
	k := Key(ctx, singular, plural)
	if tr := t.translations[k]; tr != nil {
		ii := t.formula(n)
		if ii < len(tr) {
			return tr[ii]
		}
	}
	if n == 1 {
		return singular
	}
	return plural
}

func (t *Table) Encode() (string, error) {
	var buf bytes.Buffer
	w, err := gzip.NewWriterLevel(&buf, gzip.BestCompression)
	if err != nil {
		return "", err
	}
	if err := binary.Write(w, binary.BigEndian, int32(len(t.translations))); err != nil {
		return "", err
	}
	for k, v := range t.translations {
		if err := writeString(w, k); err != nil {
			return "", err
		}
		if err := binary.Write(w, binary.BigEndian, int32(len(v))); err != nil {
			return "", err
		}
		for _, s := range v {
			if err := writeString(w, s); err != nil {
				return "", err
			}
		}
	}
	if err := w.Close(); err != nil {
		return "", err
	}
	return buf.String(), nil
}

func (t *Table) Update(other *Table) error {
	if other.formula != nil {
		t.formula = other.formula
	}
	for k, v := range other.translations {
		t.translations[k] = v
	}
	return nil
}

func Decode(data string) (*Table, error) {
	r, err := gzip.NewReader(strings.NewReader(data))
	if err != nil {
		return nil, err
	}
	defer r.Close()
	var count int32
	if err := binary.Read(r, binary.BigEndian, &count); err != nil {
		return nil, err
	}
	c := int(count)
	translations := make(map[string]Translation, c)
	for ii := 0; ii < c; ii++ {
		key, err := readString(r)
		if err != nil {
			return nil, err
		}
		var tcount int32
		if err := binary.Read(r, binary.BigEndian, &tcount); err != nil {
			return nil, err
		}
		tc := int(tcount)
		value := make([]string, tc)
		for jj := 0; jj < tc; jj++ {
			tr, err := readString(r)
			if err != nil {
				return nil, err
			}
			value[jj] = tr
		}
		translations[key] = value
	}
	return &Table{
		translations: translations,
	}, nil
}

func New(formula Formula, translations map[string]Translation) (*Table, error) {
	return &Table{
		translations: translations,
		formula:      formula,
	}, nil
}

func writeString(w io.Writer, s string) error {
	b := []byte(s)
	if err := binary.Write(w, binary.BigEndian, int32(len(b))); err != nil {
		return err
	}
	_, err := w.Write(b)
	return err
}

func readString(r io.Reader) (string, error) {
	var s int32
	if err := binary.Read(r, binary.BigEndian, &s); err != nil {
		return "", err
	}
	b := make([]byte, int(s))
	if _, err := r.Read(b); err != nil {
		return "", err
	}
	return string(b), nil
}
<|endoftext|>