<commit_before>package rain

import (
	"bytes"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"sync"
	"time"

	"github.com/cenkalti/log"
	"github.com/zeebo/bencode"

	"github.com/cenkalti/rain/internal/magnet"
	"github.com/cenkalti/rain/internal/protocol"
	"github.com/cenkalti/rain/internal/torrent"
	"github.com/cenkalti/rain/internal/tracker"
)

const (
	metadataPieceSize           = 16 * 1024
	concurrentMetadataDownloads = 1000
	metadataNetworkTimeout      = 2 * time.Minute
)

// Extension IDs
const (
	extensionHandshakeID = iota
	extensionMetadataID
)

// Metadata Extension Message Types
const (
	metadataRequest = iota
	metadataData
	metadataReject
)

type MetadataDownloader struct {
	magnet    *magnet.Magnet
	tracker   tracker.Tracker
	announceC chan *tracker.AnnounceResponse
	Result    chan *torrent.Info
	cancel    chan struct{}
	peers     map[tracker.Peer]struct{} // connecting or connected
	peersM    sync.Mutex
}

func NewMetadataDownloader(m *magnet.Magnet) (*MetadataDownloader, error) {
	if len(m.Trackers) == 0 {
		return nil, errors.New("magnet link does not contain a tracker")
	}
	c, err := newDummyClient()
	if err != nil {
		return nil, err
	}
	tr, err := tracker.New(m.Trackers[0], c)
	if err != nil {
		return nil, err
	}
	return &MetadataDownloader{
		magnet:    m,
		tracker:   tr,
		announceC: make(chan *tracker.AnnounceResponse),
		Result:    make(chan *torrent.Info, 1),
		cancel:    make(chan struct{}),
		peers:     make(map[tracker.Peer]struct{}),
	}, nil
}

func (m *MetadataDownloader) Run(announceInterval time.Duration) {
	t := emptyTransfer(m.magnet.InfoHash)
	events := make(chan tracker.Event)
	go m.tracker.Announce(&t, m.cancel, events, m.announceC)
	for {
		select {
		case resp := <-m.announceC:
			log.Infof("Seeders: %d Leechers: %d", resp.Seeders, resp.Leechers)
			for _, p := range resp.Peers {
				go m.worker(p)
			}
		case <-time.After(announceInterval):
			select {
			case events <- tracker.None:
			default:
			}
		}
	}
}

func (m *MetadataDownloader) worker(peer tracker.Peer) {
	// Do not open multiple connections to the same peer simultaneously.
	m.peersM.Lock()
	if _, ok := m.peers[peer]; ok {
		m.peersM.Unlock()
		return
	}
	m.peers[peer] = struct{}{}
	defer func() {
		m.peersM.Lock()
		delete(m.peers, peer)
		m.peersM.Unlock()
	}()
	m.peersM.Unlock()

	conn, err := net.DialTCP("tcp4", nil, peer.TCPAddr())
	if err != nil {
		log.Error(err)
		return
	}

	p := newPeer(conn)
	p.log.Debug("tcp connection is opened")

	info, err := downloadMetadataFromPeer(m.magnet, p)
	conn.Close()
	if err != nil {
		p.log.Error(err)
		return
	}

	select {
	case m.Result <- info:
		close(m.cancel) // will stop other workers
	case <-m.cancel:
		return
	}
}

func downloadMetadataFromPeer(m *magnet.Magnet, p *peer) (*torrent.Info, error) {
	err := p.conn.SetDeadline(time.Now().Add(metadataNetworkTimeout))
	if err != nil {
		return nil, err
	}

	peerID, err := generatePeerID()
	if err != nil {
		return nil, err
	}

	extensions := [8]byte{}
	extensions[5] |= 0x10 // BEP 10 Extension Protocol

	err = p.sendHandShake(m.InfoHash, peerID, extensions)
	if err != nil {
		p.log.Debug("cannot send BT handshake")
		return nil, err
	}
	p.log.Debug("sent BT handshake")

	ex, ih, err := p.readHandShake1()
	if err != nil {
		p.log.Debug("cannot read handshake part 1")
		return nil, err
	}
	if *ih != m.InfoHash {
		return nil, errors.New("unexpected info_hash")
	}
	if ex.Bytes()[5]&0x10 == 0 {
		return nil, errors.New("extension protocol is not supported by peer")
	}

	id, err := p.readHandShake2()
	if err != nil {
		p.log.Debug("cannot read handshake part 2")
		return nil, err
	}
	if *id == peerID {
		return nil, errors.New("rejected own connection: client")
	}

	p.log.Debug("BT handshake completed")

	// Extension Protocol Handshake
	d := &extensionHandshakeMessage{
		M: extensionMapping{
			UTMetadata: extensionMetadataID,
		},
	}

	err = p.sendExtensionHandshake(d)
	if err != nil {
		return nil, err
	}
	p.log.Debug("Sent extension handshake")

	var (
		v             extensionHandshakeMessage
		metadataBytes []byte
		numPieces     uint32
		lastPieceSize uint32
		remaining     uint32
	)

	for {
		err = p.conn.SetDeadline(time.Now().Add(metadataNetworkTimeout))
		if err != nil {
			return nil, err
		}

		var length uint32
		err = binary.Read(p.conn, binary.BigEndian, &length)
		if err != nil {
			return nil, err
		}
		if length == 0 { // keep-alive
			continue
		}

		var messageID protocol.MessageType
		err = binary.Read(p.conn, binary.BigEndian, &messageID)
		if err != nil {
			return nil, err
		}
		length--

		if messageID != protocol.Extension { // extension message id
			io.CopyN(ioutil.Discard, p.conn, int64(length))
			continue
		}
		p.log.Debugln("Read extension message")

		var extensionID uint8
		err = binary.Read(p.conn, binary.BigEndian, &extensionID)
		if err != nil {
			return nil, err
		}
		length--
		p.log.Debugln("LTEP message ID:", extensionID)

		switch extensionID {
		case extensionHandshakeID:
			payload := make([]byte, length)
			_, err = io.ReadFull(p.conn, payload)
			if err != nil {
				return nil, err
			}

			r := bytes.NewReader(payload)
			d := bencode.NewDecoder(r)
			err = d.Decode(&v)
			if err != nil {
				return nil, err
			}

			if v.M.UTMetadata == 0 {
				return nil, errors.New("ut_metadata extension is not supported")
			}

			if v.MetadataSize == 0 {
				return nil, errors.New("zero metadata size")
			}

			metadataBytes = make([]byte, v.MetadataSize)
			numPieces = v.MetadataSize / (metadataPieceSize)
			lastPieceSize = v.MetadataSize - (numPieces * metadataPieceSize)
			if lastPieceSize > 0 {
				numPieces++
			}
			remaining = numPieces
			p.log.Debugln("metadata has", numPieces, "pieces")

			// Send metadata piece requests.
			for i := uint32(0); i < numPieces; i++ {
				m := &metadataMessage{
					MessageType: metadataRequest,
					Piece:       i,
				}
				err = sendMetadataMessage(m, p, v.M.UTMetadata)
				if err != nil {
					return nil, err
				}
				p.log.Debugln("piece request sent", i)
			}
		case extensionMetadataID:
			payload := make([]byte, length)
			_, err = io.ReadFull(p.conn, payload)
			if err != nil {
				return nil, err
			}

			r := bytes.NewReader(payload)
			decoder := bencode.NewDecoder(r)

			in := make(map[string]uint32)
			err = decoder.Decode(&in)
			if err != nil {
				return nil, err
			}

			msgType, ok := in["msg_type"]
			if !ok {
				return nil, errors.New("no msg_type field in metadata message")
			}
			p.log.Debugln("msg_type:", msgType)

			i, ok := in["piece"]
			if !ok {
				return nil, errors.New("no piece field in metadata message")
			}
			if i >= numPieces {
				return nil, fmt.Errorf("metadata has %d pieces but peer sent piece #%d", numPieces, i)
			}

			switch msgType {
			case metadataRequest:
				req := &metadataMessage{
					MessageType: metadataReject,
					Piece:       i,
				}
				sendMetadataMessage(req, p, v.M.UTMetadata)
			case metadataData:
				var expectedSize uint32
				if i == numPieces-1 {
					expectedSize = lastPieceSize
				} else {
					expectedSize = metadataPieceSize
				}

				piece := payload[decoder.BytesParsed():]
				if uint32(len(piece)) != expectedSize {
					return nil, errors.New("received piece smaller than expected")
				}

				copy(metadataBytes[i*metadataPieceSize:], piece)

				remaining--
				if remaining == 0 {
					info, err := torrent.NewInfo(metadataBytes)
					if err != nil {
						return nil, err
					}
					if m.InfoHash != info.Hash {
						return nil, errors.New("invalid metadata received")
					}
					p.log.Info("peer has successfully sent the metadata")
					return info, nil
				}
			case metadataReject:
				return nil, errors.New("peer rejected our metadata request")
			}
		}
	}
}

func sendMetadataMessage(m *metadataMessage, p *peer, id uint8) error {
	var buf bytes.Buffer
	e := bencode.NewEncoder(&buf)
	err := e.Encode(m)
	if err != nil {
		return err
	}
	return p.sendExtensionMessage(id, buf.Bytes())
}

type extensionHandshakeMessage struct {
	M            extensionMapping `bencode:"m"`
	MetadataSize uint32           `bencode:"metadata_size,omitempty"`
}

type extensionMapping struct {
	UTMetadata uint8 `bencode:"ut_metadata"`
}

type metadataMessage struct {
	MessageType uint8  `bencode:"msg_type"`
	Piece       uint32 `bencode:"piece"`
}

type dummyClient struct {
	peerID protocol.PeerID
}

func newDummyClient() (*dummyClient, error) {
	var c dummyClient
	var err error
	c.peerID, err = generatePeerID()
	return &c, err
}

func (c *dummyClient) PeerID() protocol.PeerID { return c.peerID }
func (c *dummyClient) Port() uint16            { return 6881 }

// Required to make a fake announce to tracker to get peer list for metadata download.
type emptyTransfer protocol.InfoHash

func (t *emptyTransfer) InfoHash() protocol.InfoHash { return protocol.InfoHash(*t) }
func (t *emptyTransfer) Downloaded() int64           { return 0 }
func (t *emptyTransfer) Uploaded() int64             { return 0 }
func (t *emptyTransfer) Left() int64                 { return metadataPieceSize } // trackers don't accept 0
<commit_msg>remove unused const<commit_after>package rain

import (
	"bytes"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"sync"
	"time"

	"github.com/cenkalti/log"
	"github.com/zeebo/bencode"

	"github.com/cenkalti/rain/internal/magnet"
	"github.com/cenkalti/rain/internal/protocol"
	"github.com/cenkalti/rain/internal/torrent"
	"github.com/cenkalti/rain/internal/tracker"
)

const (
	metadataPieceSize      = 16 * 1024
	metadataNetworkTimeout = 2 * time.Minute
)

// Extension IDs
const (
	extensionHandshakeID = iota
	extensionMetadataID
)

// Metadata Extension Message Types
const (
	metadataRequest = iota
	metadataData
	metadataReject
)

type MetadataDownloader struct {
	magnet    *magnet.Magnet
	tracker   tracker.Tracker
	announceC chan *tracker.AnnounceResponse
	Result    chan *torrent.Info
	cancel    chan struct{}
	peers     map[tracker.Peer]struct{} // connecting or connected
	peersM    sync.Mutex
}

func NewMetadataDownloader(m *magnet.Magnet) (*MetadataDownloader, error) {
	if len(m.Trackers) == 0 {
		return nil, errors.New("magnet link does not contain a tracker")
	}
	c, err := newDummyClient()
	if err != nil {
		return nil, err
	}
	tr, err := tracker.New(m.Trackers[0], c)
	if err != nil {
		return nil, err
	}
	return &MetadataDownloader{
		magnet:    m,
		tracker:   tr,
		announceC: make(chan *tracker.AnnounceResponse),
		Result:    make(chan *torrent.Info, 1),
		cancel:    make(chan struct{}),
		peers:     make(map[tracker.Peer]struct{}),
	}, nil
}

func (m *MetadataDownloader) Run(announceInterval time.Duration) {
	t := emptyTransfer(m.magnet.InfoHash)
	events := make(chan tracker.Event)
	go m.tracker.Announce(&t, m.cancel, events, m.announceC)
	for {
		select {
		case resp := <-m.announceC:
			log.Infof("Seeders: %d Leechers: %d", resp.Seeders, resp.Leechers)
			for _, p := range resp.Peers {
				go m.worker(p)
			}
		case <-time.After(announceInterval):
			select {
			case events <- tracker.None:
			default:
			}
		}
	}
}

func (m *MetadataDownloader) worker(peer tracker.Peer) {
	// Do not open multiple connections to the same peer simultaneously.
	m.peersM.Lock()
	if _, ok := m.peers[peer]; ok {
		m.peersM.Unlock()
		return
	}
	m.peers[peer] = struct{}{}
	defer func() {
		m.peersM.Lock()
		delete(m.peers, peer)
		m.peersM.Unlock()
	}()
	m.peersM.Unlock()

	conn, err := net.DialTCP("tcp4", nil, peer.TCPAddr())
	if err != nil {
		log.Error(err)
		return
	}

	p := newPeer(conn)
	p.log.Debug("tcp connection is opened")

	info, err := downloadMetadataFromPeer(m.magnet, p)
	conn.Close()
	if err != nil {
		p.log.Error(err)
		return
	}

	select {
	case m.Result <- info:
		close(m.cancel) // will stop other workers
	case <-m.cancel:
		return
	}
}

func downloadMetadataFromPeer(m *magnet.Magnet, p *peer) (*torrent.Info, error) {
	err := p.conn.SetDeadline(time.Now().Add(metadataNetworkTimeout))
	if err != nil {
		return nil, err
	}

	peerID, err := generatePeerID()
	if err != nil {
		return nil, err
	}

	extensions := [8]byte{}
	extensions[5] |= 0x10 // BEP 10 Extension Protocol

	err = p.sendHandShake(m.InfoHash, peerID, extensions)
	if err != nil {
		p.log.Debug("cannot send BT handshake")
		return nil, err
	}
	p.log.Debug("sent BT handshake")

	ex, ih, err := p.readHandShake1()
	if err != nil {
		p.log.Debug("cannot read handshake part 1")
		return nil, err
	}
	if *ih != m.InfoHash {
		return nil, errors.New("unexpected info_hash")
	}
	if ex.Bytes()[5]&0x10 == 0 {
		return nil, errors.New("extension protocol is not supported by peer")
	}

	id, err := p.readHandShake2()
	if err != nil {
		p.log.Debug("cannot read handshake part 2")
		return nil, err
	}
	if *id == peerID {
		return nil, errors.New("rejected own connection: client")
	}

	p.log.Debug("BT handshake completed")

	// Extension Protocol Handshake
	d := &extensionHandshakeMessage{
		M: extensionMapping{
			UTMetadata: extensionMetadataID,
		},
	}

	err = p.sendExtensionHandshake(d)
	if err != nil {
		return nil, err
	}
	p.log.Debug("Sent extension handshake")

	var (
		v             extensionHandshakeMessage
		metadataBytes []byte
		numPieces     uint32
		lastPieceSize uint32
		remaining     uint32
	)

	for {
		err = p.conn.SetDeadline(time.Now().Add(metadataNetworkTimeout))
		if err != nil {
			return nil, err
		}

		var length uint32
		err = binary.Read(p.conn, binary.BigEndian, &length)
		if err != nil {
			return nil, err
		}
		if length == 0 { // keep-alive
			continue
		}

		var messageID protocol.MessageType
		err = binary.Read(p.conn, binary.BigEndian, &messageID)
		if err != nil {
			return nil, err
		}
		length--

		if messageID != protocol.Extension { // extension message id
			io.CopyN(ioutil.Discard, p.conn, int64(length))
			continue
		}
		p.log.Debugln("Read extension message")

		var extensionID uint8
		err = binary.Read(p.conn, binary.BigEndian, &extensionID)
		if err != nil {
			return nil, err
		}
		length--
		p.log.Debugln("LTEP message ID:", extensionID)

		switch extensionID {
		case extensionHandshakeID:
			payload := make([]byte, length)
			_, err = io.ReadFull(p.conn, payload)
			if err != nil {
				return nil, err
			}

			r := bytes.NewReader(payload)
			d := bencode.NewDecoder(r)
			err = d.Decode(&v)
			if err != nil {
				return nil, err
			}

			if v.M.UTMetadata == 0 {
				return nil, errors.New("ut_metadata extension is not supported")
			}

			if v.MetadataSize == 0 {
				return nil, errors.New("zero metadata size")
			}

			metadataBytes = make([]byte, v.MetadataSize)
			numPieces = v.MetadataSize / (metadataPieceSize)
			lastPieceSize = v.MetadataSize - (numPieces * metadataPieceSize)
			if lastPieceSize > 0 {
				numPieces++
			}
			remaining = numPieces
			p.log.Debugln("metadata has", numPieces, "pieces")

			// Send metadata piece requests.
			for i := uint32(0); i < numPieces; i++ {
				m := &metadataMessage{
					MessageType: metadataRequest,
					Piece:       i,
				}
				err = sendMetadataMessage(m, p, v.M.UTMetadata)
				if err != nil {
					return nil, err
				}
				p.log.Debugln("piece request sent", i)
			}
		case extensionMetadataID:
			payload := make([]byte, length)
			_, err = io.ReadFull(p.conn, payload)
			if err != nil {
				return nil, err
			}

			r := bytes.NewReader(payload)
			decoder := bencode.NewDecoder(r)

			in := make(map[string]uint32)
			err = decoder.Decode(&in)
			if err != nil {
				return nil, err
			}

			msgType, ok := in["msg_type"]
			if !ok {
				return nil, errors.New("no msg_type field in metadata message")
			}
			p.log.Debugln("msg_type:", msgType)

			i, ok := in["piece"]
			if !ok {
				return nil, errors.New("no piece field in metadata message")
			}
			if i >= numPieces {
				return nil, fmt.Errorf("metadata has %d pieces but peer sent piece #%d", numPieces, i)
			}

			switch msgType {
			case metadataRequest:
				req := &metadataMessage{
					MessageType: metadataReject,
					Piece:       i,
				}
				sendMetadataMessage(req, p, v.M.UTMetadata)
			case metadataData:
				var expectedSize uint32
				if i == numPieces-1 {
					expectedSize = lastPieceSize
				} else {
					expectedSize = metadataPieceSize
				}

				piece := payload[decoder.BytesParsed():]
				if uint32(len(piece)) != expectedSize {
					return nil, errors.New("received piece smaller than expected")
				}

				copy(metadataBytes[i*metadataPieceSize:], piece)

				remaining--
				if remaining == 0 {
					info, err := torrent.NewInfo(metadataBytes)
					if err != nil {
						return nil, err
					}
					if m.InfoHash != info.Hash {
						return nil, errors.New("invalid metadata received")
					}
					p.log.Info("peer has successfully sent the metadata")
					return info, nil
				}
			case metadataReject:
				return nil, errors.New("peer rejected our metadata request")
			}
		}
	}
}

func sendMetadataMessage(m *metadataMessage, p *peer, id uint8) error {
	var buf bytes.Buffer
	e := bencode.NewEncoder(&buf)
	err := e.Encode(m)
	if err != nil {
		return err
	}
	return p.sendExtensionMessage(id, buf.Bytes())
}

type extensionHandshakeMessage struct {
	M            extensionMapping `bencode:"m"`
	MetadataSize uint32           `bencode:"metadata_size,omitempty"`
}

type extensionMapping struct {
	UTMetadata uint8 `bencode:"ut_metadata"`
}

type metadataMessage struct {
	MessageType uint8  `bencode:"msg_type"`
	Piece       uint32 `bencode:"piece"`
}

type dummyClient struct {
	peerID protocol.PeerID
}

func newDummyClient() (*dummyClient, error) {
	var c dummyClient
	var err error
	c.peerID, err = generatePeerID()
	return &c, err
}

func (c *dummyClient) PeerID() protocol.PeerID { return c.peerID }
func (c *dummyClient) Port() uint16            { return 6881 }

// Required to make a fake announce to tracker to get peer list for metadata download.
type emptyTransfer protocol.InfoHash

func (t *emptyTransfer) InfoHash() protocol.InfoHash { return protocol.InfoHash(*t) }
func (t *emptyTransfer) Downloaded() int64           { return 0 }
func (t *emptyTransfer) Uploaded() int64             { return 0 }
func (t *emptyTransfer) Left() int64                 { return metadataPieceSize } // trackers don't accept 0
<|endoftext|>


<commit_before>package rain

import (
	"bytes"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"sync"
	"time"

	"github.com/cenkalti/log"
	"github.com/zeebo/bencode"

	"github.com/cenkalti/rain/internal/connection"
	"github.com/cenkalti/rain/internal/magnet"
	"github.com/cenkalti/rain/internal/protocol"
	"github.com/cenkalti/rain/internal/torrent"
	"github.com/cenkalti/rain/internal/tracker"
)

const (
	metadataPieceSize      = 16 * 1024
	metadataNetworkTimeout = 2 * time.Minute
)

// Extension IDs
const (
	extensionHandshakeID = iota
	extensionMetadataID
)

// Metadata Extension Message Types
const (
	metadataRequest = iota
	metadataData
	metadataReject
)

type MetadataDownloader struct {
	magnet    *magnet.Magnet
	tracker   tracker.Tracker // TODO support multiple trackers
	announceC chan *tracker.AnnounceResponse
	Result    chan *torrent.Info
	cancel    chan struct{}
	peers     map[tracker.Peer]struct{} // connecting or connected
	peersM    sync.Mutex
}

func NewMetadataDownloader(m *magnet.Magnet) (*MetadataDownloader, error) {
	if len(m.Trackers) == 0 {
		return nil, errors.New("magnet link does not contain a tracker")
	}
	c, err := newDummyClient()
	if err != nil {
		return nil, err
	}
	tr, err := tracker.New(m.Trackers[0], c)
	if err != nil {
		return nil, err
	}
	return &MetadataDownloader{
		magnet:    m,
		tracker:   tr,
		announceC: make(chan *tracker.AnnounceResponse),
		Result:    make(chan *torrent.Info, 1),
		cancel:    make(chan struct{}),
		peers:     make(map[tracker.Peer]struct{}),
	}, nil
}

func (m *MetadataDownloader) Run(announceInterval time.Duration) {
	t := emptyTransfer(m.magnet.InfoHash)
	events := make(chan tracker.Event)
	go tracker.AnnouncePeriodically(m.tracker, &t, m.cancel, tracker.None, events, m.announceC)
	for {
		select {
		case resp := <-m.announceC:
			log.Infof("Seeders: %d Leechers: %d", resp.Seeders, resp.Leechers)
			for _, p := range resp.Peers {
				go m.worker(p)
			}
		case <-time.After(announceInterval):
			select {
			case events <- tracker.None:
			default:
			}
		}
	}
}

func (m *MetadataDownloader) worker(peer tracker.Peer) {
	// Do not open multiple connections to the same peer simultaneously.
	m.peersM.Lock()
	if _, ok := m.peers[peer]; ok {
		m.peersM.Unlock()
		return
	}
	m.peers[peer] = struct{}{}
	defer func() {
		m.peersM.Lock()
		delete(m.peers, peer)
		m.peersM.Unlock()
	}()
	m.peersM.Unlock()

	ourID, err := generatePeerID()
	if err != nil {
		panic(err)
	}

	ourExtensions := [8]byte{}
	ourExtensions[5] |= 0x10 // BEP 10 Extension Protocol
	conn, _, peerExtensions, _, err := connection.Dial(peer.TCPAddr(), true, false, ourExtensions, m.magnet.InfoHash, ourID)
	if err != nil {
		log.Error(err)
		return
	}
	defer conn.Close()

	if peerExtensions[5]&0x10 == 0 {
		log.Debug("Peer does not support extension protocol")
		return
	}

	p := newPeer(conn)

	info, err := downloadMetadataFromPeer(m.magnet, p)
	conn.Close()
	if err != nil {
		p.log.Error(err)
		return
	}

	select {
	case m.Result <- info:
		close(m.cancel) // will stop other workers
	case <-m.cancel:
		return
	}
}

func downloadMetadataFromPeer(m *magnet.Magnet, p *peer) (*torrent.Info, error) {
	d := &extensionHandshakeMessage{
		M: extensionMapping{
			UTMetadata: extensionMetadataID,
		},
	}

	err := p.sendExtensionHandshake(d)
	if err != nil {
		return nil, err
	}
	p.log.Debug("Sent extension handshake")

	var (
		v             extensionHandshakeMessage
		metadataBytes []byte
		numPieces     uint32
		lastPieceSize uint32
		remaining     uint32
	)

	for {
		err = p.conn.SetDeadline(time.Now().Add(metadataNetworkTimeout))
		if err != nil {
			return nil, err
		}

		p.log.Debug("Reading peer message...")
		var length uint32
		err = binary.Read(p.conn, binary.BigEndian, &length)
		if err != nil {
			return nil, err
		}
		if length == 0 { // keep-alive
			continue
		}
		p.log.Debugf("Next message length: %d", length)

		var messageID protocol.MessageType
		err = binary.Read(p.conn, binary.BigEndian, &messageID)
		if err != nil {
			return nil, err
		}
		p.log.Debugf("messageID: %s\n", messageID)
		length--

		if messageID != protocol.Extension { // extension message id
			io.CopyN(ioutil.Discard, p.conn, int64(length))
			p.log.Debugf("Discarded %d bytes", length)
			continue
		}
		p.log.Debugln("Read extension message")

		var extensionID uint8
		err = binary.Read(p.conn, binary.BigEndian, &extensionID)
		if err != nil {
			return nil, err
		}
		length--
		p.log.Debugln("LTEP message ID:", extensionID)

		switch extensionID {
		case extensionHandshakeID:
			payload := make([]byte, length)
			_, err = io.ReadFull(p.conn, payload)
			if err != nil {
				return nil, err
			}

			r := bytes.NewReader(payload)
			d := bencode.NewDecoder(r)
			err = d.Decode(&v)
			if err != nil {
				return nil, err
			}

			if v.M.UTMetadata == 0 {
				return nil, errors.New("ut_metadata extension is not supported")
			}

			if v.MetadataSize == 0 {
				return nil, errors.New("zero metadata size")
			}
			p.log.Infoln("Metadata size:", v.MetadataSize, "bytes")

			metadataBytes = make([]byte, v.MetadataSize)
			numPieces = v.MetadataSize / (metadataPieceSize)
			lastPieceSize = v.MetadataSize - (numPieces * metadataPieceSize)
			if lastPieceSize > 0 {
				numPieces++
			}
			remaining = numPieces
			p.log.Infoln("Metadata has", numPieces, "piece(s)")
			if numPieces == 1 {
				lastPieceSize = v.MetadataSize
			}

			// Send metadata piece requests.
			for i := uint32(0); i < numPieces; i++ {
				m := &metadataMessage{
					MessageType: metadataRequest,
					Piece:       i,
				}
				err = sendMetadataMessage(m, p, v.M.UTMetadata)
				if err != nil {
					return nil, err
				}
				p.log.Debugln("piece request sent", i)
			}
		case extensionMetadataID:
			payload := make([]byte, length)
			_, err = io.ReadFull(p.conn, payload)
			if err != nil {
				return nil, err
			}

			r := bytes.NewReader(payload)
			decoder := bencode.NewDecoder(r)

			in := make(map[string]uint32)
			err = decoder.Decode(&in)
			if err != nil {
				return nil, err
			}

			msgType, ok := in["msg_type"]
			if !ok {
				return nil, errors.New("no msg_type field in metadata message")
			}
			p.log.Debugln("msg_type:", msgType)

			i, ok := in["piece"]
			if !ok {
				return nil, errors.New("no piece field in metadata message")
			}
			if i >= numPieces {
				return nil, fmt.Errorf("metadata has %d pieces but peer sent piece #%d", numPieces, i)
			}

			switch msgType {
			case metadataRequest:
				req := &metadataMessage{
					MessageType: metadataReject,
					Piece:       i,
				}
				sendMetadataMessage(req, p, v.M.UTMetadata)
			case metadataData:
				var expectedSize uint32
				if i == numPieces-1 {
					expectedSize = lastPieceSize
				} else {
					expectedSize = metadataPieceSize
				}

				piece := payload[decoder.BytesParsed():]
				if uint32(len(piece)) != expectedSize {
					return nil, fmt.Errorf("received piece smaller than expected (%d/%d)", len(piece), expectedSize)
				}

				copy(metadataBytes[i*metadataPieceSize:], piece)

				remaining--
				if remaining == 0 {
					info, err := torrent.NewInfo(metadataBytes)
					if err != nil {
						return nil, err
					}
					if m.InfoHash != info.Hash {
						return nil, errors.New("invalid metadata received")
					}
					p.log.Info("peer has successfully sent the metadata")
					return info, nil
				}
			case metadataReject:
				return nil, errors.New("peer rejected our metadata request")
			}
		}
	}
}

func sendMetadataMessage(m *metadataMessage, p *peer, id uint8) error {
	var buf bytes.Buffer
	e := bencode.NewEncoder(&buf)
	err := e.Encode(m)
	if err != nil {
		return err
	}
	return p.sendExtensionMessage(id, buf.Bytes())
}

type extensionHandshakeMessage struct {
	M            extensionMapping `bencode:"m"`
	MetadataSize uint32           `bencode:"metadata_size,omitempty"`
}

type extensionMapping struct {
	UTMetadata uint8 `bencode:"ut_metadata"`
}

type metadataMessage struct {
	MessageType uint8  `bencode:"msg_type"`
	Piece       uint32 `bencode:"piece"`
}

type dummyClient struct {
	peerID protocol.PeerID
}

func newDummyClient() (*dummyClient, error) {
	var c dummyClient
	var err error
	c.peerID, err = generatePeerID()
	return &c, err
}

func (c *dummyClient) PeerID() protocol.PeerID { return c.peerID }
func (c *dummyClient) Port() uint16            { return 6881 }

// Required to make a fake announce to tracker to get peer list for metadata download.
type emptyTransfer protocol.InfoHash

func (t *emptyTransfer) InfoHash() protocol.InfoHash { return protocol.InfoHash(*t) }
func (t *emptyTransfer) Downloaded() int64           { return 0 }
func (t *emptyTransfer) Uploaded() int64             { return 0 }
func (t *emptyTransfer) Left() int64                 { return metadataPieceSize } // trackers don't accept 0
<commit_msg>announce to multiple trackers<commit_after>package rain

import (
	"bytes"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"sync"
	"time"

	"github.com/cenkalti/log"
	"github.com/zeebo/bencode"

	"github.com/cenkalti/rain/internal/connection"
	"github.com/cenkalti/rain/internal/magnet"
	"github.com/cenkalti/rain/internal/protocol"
	"github.com/cenkalti/rain/internal/torrent"
	"github.com/cenkalti/rain/internal/tracker"
)

const (
	metadataPieceSize      = 16 * 1024
	metadataNetworkTimeout = 2 * time.Minute
)

// Extension IDs
const (
	extensionHandshakeID = iota
	extensionMetadataID
)

// Metadata Extension Message Types
const (
	metadataRequest = iota
	metadataData
	metadataReject
)

type MetadataDownloader struct {
	magnet    *magnet.Magnet
	trackers  []tracker.Tracker
	announceC chan *tracker.AnnounceResponse
	Result    chan *torrent.Info
	cancel    chan struct{}
	peers     map[tracker.Peer]struct{} // connecting or connected
	peersM    sync.Mutex
}

func NewMetadataDownloader(m *magnet.Magnet) (*MetadataDownloader, error) {
	if len(m.Trackers) == 0 {
		return nil, errors.New("magnet link does not contain a tracker")
	}
	c, err := newDummyClient()
	if err != nil {
		return nil, err
	}
	trackers := make([]tracker.Tracker, 0, len(m.Trackers))
	for _, s := range m.Trackers {
		tr, err := tracker.New(s, c)
		if err != nil {
			log.Error(err)
			continue
		}
		trackers = append(trackers, tr)
	}
	return &MetadataDownloader{
		magnet:    m,
		trackers:  trackers,
		announceC: make(chan *tracker.AnnounceResponse),
		Result:    make(chan *torrent.Info, 1),
		cancel:    make(chan struct{}),
		peers:     make(map[tracker.Peer]struct{}),
	}, nil
}

func (m *MetadataDownloader) Run(announceInterval time.Duration) {
	t := emptyTransfer(m.magnet.InfoHash)
	events := make(chan tracker.Event)
	for _, tr := range m.trackers {
		go tracker.AnnouncePeriodically(tr, &t, m.cancel, tracker.None, events, m.announceC)
	}
	for {
		select {
		case resp := <-m.announceC:
			log.Infof("Seeders: %d Leechers: %d", resp.Seeders, resp.Leechers)
			for _, p := range resp.Peers {
				go m.worker(p)
			}
		case <-time.After(announceInterval):
			select {
			case events <- tracker.None:
			default:
			}
		}
	}
}

func (m *MetadataDownloader) worker(peer tracker.Peer) {
	// Do not open multiple connections to the same peer simultaneously.
	m.peersM.Lock()
	if _, ok := m.peers[peer]; ok {
		m.peersM.Unlock()
		return
	}
	m.peers[peer] = struct{}{}
	defer func() {
		m.peersM.Lock()
		delete(m.peers, peer)
		m.peersM.Unlock()
	}()
	m.peersM.Unlock()

	ourID, err := generatePeerID()
	if err != nil {
		panic(err)
	}

	ourExtensions := [8]byte{}
	ourExtensions[5] |= 0x10 // BEP 10 Extension Protocol
	conn, _, peerExtensions, _, err := connection.Dial(peer.TCPAddr(), true, false, ourExtensions, m.magnet.InfoHash, ourID)
	if err != nil {
		log.Error(err)
		return
	}
	defer conn.Close()

	if peerExtensions[5]&0x10 == 0 {
		log.Debug("Peer does not support extension protocol")
		return
	}

	p := newPeer(conn)

	info, err := downloadMetadataFromPeer(m.magnet, p)
	conn.Close()
	if err != nil {
		p.log.Error(err)
		return
	}

	select {
	case m.Result <- info:
		close(m.cancel) // will stop other workers
	case <-m.cancel:
		return
	}
}

func downloadMetadataFromPeer(m *magnet.Magnet, p *peer) (*torrent.Info, error) {
	d := &extensionHandshakeMessage{
		M: extensionMapping{
			UTMetadata: extensionMetadataID,
		},
	}

	err := p.sendExtensionHandshake(d)
	if err != nil {
		return nil, err
	}
	p.log.Debug("Sent extension handshake")

	var (
		v             extensionHandshakeMessage
		metadataBytes []byte
		numPieces     uint32
		lastPieceSize uint32
		remaining     uint32
	)

	for {
		err = p.conn.SetDeadline(time.Now().Add(metadataNetworkTimeout))
		if err != nil {
			return nil, err
		}

		p.log.Debug("Reading peer message...")
		var length uint32
		err = binary.Read(p.conn, binary.BigEndian, &length)
		if err != nil {
			return nil, err
		}
		if length == 0 { // keep-alive
			continue
		}
		p.log.Debugf("Next message length: %d", length)

		var messageID protocol.MessageType
		err = binary.Read(p.conn, binary.BigEndian, &messageID)
		if err != nil {
			return nil, err
		}
		p.log.Debugf("messageID: %s\n", messageID)
		length--

		if messageID != protocol.Extension { // extension message id
			io.CopyN(ioutil.Discard, p.conn, int64(length))
			p.log.Debugf("Discarded %d bytes", length)
			continue
		}
		p.log.Debugln("Read extension message")

		var extensionID uint8
		err = binary.Read(p.conn, binary.BigEndian, &extensionID)
		if err != nil {
			return nil, err
		}
		length--
		p.log.Debugln("LTEP message ID:", extensionID)

		switch extensionID {
		case extensionHandshakeID:
			payload := make([]byte, length)
			_, err = io.ReadFull(p.conn, payload)
			if err != nil {
				return nil, err
			}

			r := bytes.NewReader(payload)
			d := bencode.NewDecoder(r)
			err = d.Decode(&v)
			if err != nil {
				return nil, err
			}

			if v.M.UTMetadata == 0 {
				return nil, errors.New("ut_metadata extension is not supported")
			}

			if v.MetadataSize == 0 {
				return nil, errors.New("zero metadata size")
			}
			p.log.Infoln("Metadata size:", v.MetadataSize, "bytes")

			metadataBytes = make([]byte, v.MetadataSize)
			numPieces = v.MetadataSize / (metadataPieceSize)
			lastPieceSize = v.MetadataSize - (numPieces * metadataPieceSize)
			if lastPieceSize > 0 {
				numPieces++
			}
			remaining = numPieces
			p.log.Infoln("Metadata has", numPieces, "piece(s)")
			if numPieces == 1 {
				lastPieceSize = v.MetadataSize
			}

			// Send metadata piece requests.
			for i := uint32(0); i < numPieces; i++ {
				m := &metadataMessage{
					MessageType: metadataRequest,
					Piece:       i,
				}
				err = sendMetadataMessage(m, p, v.M.UTMetadata)
				if err != nil {
					return nil, err
				}
				p.log.Debugln("piece request sent", i)
			}
		case extensionMetadataID:
			payload := make([]byte, length)
			_, err = io.ReadFull(p.conn, payload)
			if err != nil {
				return nil, err
			}

			r := bytes.NewReader(payload)
			decoder := bencode.NewDecoder(r)

			in := make(map[string]uint32)
			err = decoder.Decode(&in)
			if err != nil {
				return nil, err
			}

			msgType, ok := in["msg_type"]
			if !ok {
				return nil, errors.New("no msg_type field in metadata message")
			}
			p.log.Debugln("msg_type:", msgType)

			i, ok := in["piece"]
			if !ok {
				return nil, errors.New("no piece field in metadata message")
			}
			if i >= numPieces {
				return nil, fmt.Errorf("metadata has %d pieces but peer sent piece #%d", numPieces, i)
			}

			switch msgType {
			case metadataRequest:
				req := &metadataMessage{
					MessageType: metadataReject,
					Piece:       i,
				}
				sendMetadataMessage(req, p, v.M.UTMetadata)
			case metadataData:
				var expectedSize uint32
				if i == numPieces-1 {
					expectedSize = lastPieceSize
				} else {
					expectedSize = metadataPieceSize
				}

				piece := payload[decoder.BytesParsed():]
				if uint32(len(piece)) != expectedSize {
					return nil, fmt.Errorf("received piece smaller than expected (%d/%d)", len(piece), expectedSize)
				}

				copy(metadataBytes[i*metadataPieceSize:], piece)

				remaining--
				if remaining == 0 {
					info, err := torrent.NewInfo(metadataBytes)
					if err != nil {
						return nil, err
					}
					if m.InfoHash != info.Hash {
						return nil, errors.New("invalid metadata received")
					}
					p.log.Info("peer has successfully sent the metadata")
					return info, nil
				}
			case metadataReject:
				return nil, errors.New("peer rejected our metadata request")
			}
		}
	}
}

func sendMetadataMessage(m *metadataMessage, p *peer, id uint8) error {
	var buf bytes.Buffer
	e := bencode.NewEncoder(&buf)
	err := e.Encode(m)
	if err != nil {
		return err
	}
	return p.sendExtensionMessage(id, buf.Bytes())
}

type extensionHandshakeMessage struct {
	M            extensionMapping `bencode:"m"`
	MetadataSize uint32           `bencode:"metadata_size,omitempty"`
}

type extensionMapping struct {
	UTMetadata uint8 `bencode:"ut_metadata"`
}

type metadataMessage struct {
	MessageType uint8  `bencode:"msg_type"`
	Piece       uint32 `bencode:"piece"`
}

type dummyClient struct {
	peerID protocol.PeerID
}

func newDummyClient() (*dummyClient, error) {
	var c dummyClient
	var err error
	c.peerID, err = generatePeerID()
	return &c, err
}

func (c *dummyClient) PeerID() protocol.PeerID { return c.peerID }
func (c *dummyClient) Port() uint16            { return 6881 }

// Required to make a fake announce to tracker to get peer list for metadata download.
type emptyTransfer protocol.InfoHash

func (t *emptyTransfer) InfoHash() protocol.InfoHash { return protocol.InfoHash(*t) }
func (t *emptyTransfer) Downloaded() int64           { return 0 }
func (t *emptyTransfer) Uploaded() int64             { return 0 }
func (t *emptyTransfer) Left() int64                 { return metadataPieceSize } // trackers don't accept 0
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"strings"

	"github.com/spf13/cobra"

	"github.com/lxc/lxd/client"
	"github.com/lxc/lxd/lxc/config"
	"github.com/lxc/lxd/lxc/utils"
	"github.com/lxc/lxd/shared"
	"github.com/lxc/lxd/shared/api"
	cli "github.com/lxc/lxd/shared/cmd"
	"github.com/lxc/lxd/shared/i18n"
)

type cmdCopy struct {
	global *cmdGlobal

	flagNoProfiles    bool
	flagProfile       []string
	flagConfig        []string
	flagDevice        []string
	flagEphemeral     bool
	flagContainerOnly bool
	flagMode          string
	flagStateless     bool
	flagStorage       string
	flagTarget        string
	flagTargetProject string
	flagRefresh       bool
}

func (c *cmdCopy) Command() *cobra.Command {
	cmd := &cobra.Command{}
	cmd.Use = i18n.G("copy [<remote>:]<source>[/<snapshot>] [[<remote>:]<destination>]")
	cmd.Aliases = []string{"cp"}
	cmd.Short = i18n.G("Copy containers within or in between LXD instances")
	cmd.Long = cli.FormatSection(i18n.G("Description"), i18n.G(
		`Copy containers within or in between LXD instances`))

	cmd.RunE = c.Run
	cmd.Flags().StringArrayVarP(&c.flagConfig, "config", "c", nil, i18n.G("Config key/value to apply to the new container")+"``")
	cmd.Flags().StringArrayVarP(&c.flagDevice, "device", "d", nil, i18n.G("New key/value to apply to a specific device")+"``")
	cmd.Flags().StringArrayVarP(&c.flagProfile, "profile", "p", nil, i18n.G("Profile to apply to the new container")+"``")
	cmd.Flags().BoolVarP(&c.flagEphemeral, "ephemeral", "e", false, i18n.G("Ephemeral container"))
	cmd.Flags().StringVar(&c.flagMode, "mode", "pull", i18n.G("Transfer mode. One of pull (default), push or relay")+"``")
	cmd.Flags().BoolVar(&c.flagContainerOnly, "container-only", false, i18n.G("Copy the container without its snapshots"))
	cmd.Flags().BoolVar(&c.flagStateless, "stateless", false, i18n.G("Copy a stateful container stateless"))
	cmd.Flags().StringVarP(&c.flagStorage, "storage", "s", "", i18n.G("Storage pool name")+"``")
	cmd.Flags().StringVar(&c.flagTarget, "target", "", i18n.G("Cluster member name")+"``")
	cmd.Flags().StringVar(&c.flagTargetProject, "target-project", "", i18n.G("Copy to a project different from the source")+"``")
	cmd.Flags().BoolVar(&c.flagNoProfiles, "no-profiles", false, i18n.G("Create the container with no profiles applied"))
	cmd.Flags().BoolVar(&c.flagRefresh, "refresh", false, i18n.G("Perform an incremental copy"))

	return cmd
}

func (c *cmdCopy) copyContainer(conf *config.Config, sourceResource string,
	destResource string, keepVolatile bool, ephemeral int, stateful bool,
	containerOnly bool, mode string, pool string) error {
	// Parse the source
	sourceRemote, sourceName, err := conf.ParseRemote(sourceResource)
	if err != nil {
		return err
	}

	// Parse the destination
	destRemote, destName, err := conf.ParseRemote(destResource)
	if err != nil {
		return err
	}

	// Make sure we have a container or snapshot name
	if sourceName == "" {
		return fmt.Errorf(i18n.G("You must specify a source container name"))
	}

	// Check that a destination container was specified, if --target is passed.
	if destName == "" && c.flagTarget != "" {
		return fmt.Errorf(i18n.G("You must specify a destination container name when using --target"))
	}

	// If no destination name was provided, use the same as the source
	if destName == "" && destResource != "" {
		destName = sourceName
	}

	// Connect to the source host
	source, err := conf.GetContainerServer(sourceRemote)
	if err != nil {
		return err
	}

	// Connect to the destination host
	var dest lxd.ContainerServer
	if sourceRemote == destRemote {
		// Source and destination are the same
		dest = source
	} else {
		// Destination is different, connect to it
		dest, err = conf.GetContainerServer(destRemote)
		if err != nil {
			return err
		}
	}

	// Project copies
	if c.flagTargetProject != "" {
		dest = dest.UseProject(c.flagTargetProject)
	}

	// Confirm that --target is only used with a cluster
	if c.flagTarget != "" && !dest.IsClustered() {
		return fmt.Errorf(i18n.G("To use --target, the destination remote must be a cluster"))
	}

	// Parse the config overrides
	configMap := map[string]string{}
	for _, entry := range c.flagConfig {
		if !strings.Contains(entry, "=") {
			return fmt.Errorf(i18n.G("Bad key=value pair: %s"), entry)
		}

		fields := strings.SplitN(entry, "=", 2)
		configMap[fields[0]] = fields[1]
	}

	// Parse the device overrides
	deviceMap := map[string]map[string]string{}
	for _, entry := range c.flagDevice {
		if !strings.Contains(entry, "=") || !strings.Contains(entry, ",") {
			return fmt.Errorf(i18n.G("Bad syntax, expecting <device>,<key>=<value>: %s"), entry)
		}

		deviceFields := strings.SplitN(entry, ",", 2)
		keyFields := strings.SplitN(deviceFields[1], "=", 2)

		if deviceMap[deviceFields[0]] == nil {
			deviceMap[deviceFields[0]] = map[string]string{}
		}

		deviceMap[deviceFields[0]][keyFields[0]] = keyFields[1]
	}

	var op lxd.RemoteOperation
	var writable api.ContainerPut

	if shared.IsSnapshot(sourceName) {
		if containerOnly {
			return fmt.Errorf(i18n.G("--container-only can't be passed when the source is a snapshot"))
		}

		// Prepare the container creation request
		args := lxd.ContainerSnapshotCopyArgs{
			Name: destName,
			Mode: mode,
			Live: stateful,
		}

		if c.flagRefresh {
			return fmt.Errorf(i18n.G("--refresh can only be used with containers"))
		}

		// Copy of a snapshot into a new container
		srcFields := strings.SplitN(sourceName, shared.SnapshotDelimiter, 2)
		entry, _, err := source.GetContainerSnapshot(srcFields[0], srcFields[1])
		if err != nil {
			return err
		}

		// Allow adding additional profiles
		if c.flagProfile != nil {
			entry.Profiles = append(entry.Profiles, c.flagProfile...)
		} else if c.flagNoProfiles {
			entry.Profiles = []string{}
		}

		// Allow setting additional config keys
		if configMap != nil {
			for key, value := range configMap {
				entry.Config[key] = value
			}
		}

		// Allow setting device overrides
		if deviceMap != nil {
			for k, m := range deviceMap {
				if entry.Devices[k] == nil {
					entry.Devices[k] = m
					continue
				}

				for key, value := range m {
					entry.Devices[k][key] = value
				}
			}
		}

		// Allow overriding the ephemeral status
		if ephemeral == 1 {
			entry.Ephemeral = true
		} else if ephemeral == 0 {
			entry.Ephemeral = false
		}

		rootDiskDeviceKey, _, _ := shared.GetRootDiskDevice(entry.Devices)
		if err != nil {
			return err
		}

		if rootDiskDeviceKey != "" && pool != "" {
			entry.Devices[rootDiskDeviceKey]["pool"] = pool
		} else if pool != "" {
			entry.Devices["root"] = map[string]string{
				"type": "disk",
				"path": "/",
				"pool": pool,
			}
		}

		// Strip the volatile keys if requested
		if !keepVolatile {
			for k := range entry.Config {
				if k == "volatile.base_image" {
					continue
				}

				if strings.HasPrefix(k, "volatile") {
					delete(entry.Config, k)
				}
			}
		}

		// Do the actual copy
		if c.flagTarget != "" {
			dest = dest.UseTarget(c.flagTarget)
		}

		op, err = dest.CopyContainerSnapshot(source, srcFields[0], *entry, &args)
		if err != nil {
			return err
		}
	} else {
		// Prepare the container creation request
		args := lxd.ContainerCopyArgs{
			Name:          destName,
			Live:          stateful,
			ContainerOnly: containerOnly,
			Mode:          mode,
			Refresh:       c.flagRefresh,
		}

		// Copy of a container into a new container
		entry, _, err := source.GetContainer(sourceName)
		if err != nil {
			return err
		}

		// Allow adding additional profiles
		if c.flagProfile != nil {
			entry.Profiles = append(entry.Profiles, c.flagProfile...)
		} else if c.flagNoProfiles {
			entry.Profiles = []string{}
		}

		// Allow setting additional config keys
		if configMap != nil {
			for key, value := range configMap {
				entry.Config[key] = value
			}
		}

		// Allow setting device overrides
		if deviceMap != nil {
			for k, m := range deviceMap {
				if entry.Devices[k] == nil {
					entry.Devices[k] = m
					continue
				}

				for key, value := range m {
					entry.Devices[k][key] = value
				}
			}
		}

		// Allow overriding the ephemeral status
		if ephemeral == 1 {
			entry.Ephemeral = true
		} else if ephemeral == 0 {
			entry.Ephemeral = false
		}

		rootDiskDeviceKey, _, _ := shared.GetRootDiskDevice(entry.Devices)
		if err != nil {
			return err
		}

		if rootDiskDeviceKey != "" && pool != "" {
			entry.Devices[rootDiskDeviceKey]["pool"] = pool
		} else if pool != "" {
			entry.Devices["root"] = map[string]string{
				"type": "disk",
				"path": "/",
				"pool": pool,
			}
		}

		// Strip the volatile keys if requested
		if !keepVolatile {
			for k := range entry.Config {
				if k == "volatile.base_image" {
					continue
				}

				if strings.HasPrefix(k, "volatile") {
					delete(entry.Config, k)
				}
			}
		}

		// Do the actual copy
		if c.flagTarget != "" {
			dest = dest.UseTarget(c.flagTarget)
		}

		op, err = dest.CopyContainer(source, *entry, &args)
		if err != nil {
			return err
		}

		writable = entry.Writable()
	}

	// Watch the background operation
	progress := utils.ProgressRenderer{
		Format: i18n.G("Transferring container: %s"),
		Quiet:  c.global.flagQuiet,
	}

	_, err = op.AddHandler(progress.UpdateOp)
	if err != nil {
		progress.Done("")
		return err
	}

	// Wait for the copy to complete
	err = utils.CancelableWait(op, &progress)
	if err != nil {
		progress.Done("")
		return err
	}
	progress.Done("")

	if c.flagRefresh {
		_, etag, err := dest.GetContainer(destName)
		if err != nil {
			return fmt.Errorf("Failed to refresh target container '%s': %v", destName, err)
		}

		op, err := dest.UpdateContainer(destName, writable, etag)
		if err != nil {
			return err
		}

		// Watch the background operation
		progress := utils.ProgressRenderer{
			Format: i18n.G("Refreshing container: %s"),
			Quiet:  c.global.flagQuiet,
		}

		_, err = op.AddHandler(progress.UpdateOp)
		if err != nil {
			progress.Done("")
			return err
		}

		// Wait for the copy to complete
		err = utils.CancelableWait(op, &progress)
		if err != nil {
			progress.Done("")
			return err
		}
		progress.Done("")
	}

	// If choosing a random name, show it to the user
	if destResource == "" {
		// Get the successful operation data
		opInfo, err := op.GetTarget()
		if err != nil {
			return err
		}

		// Extract the list of affected containers
		containers, ok := opInfo.Resources["containers"]
		if !ok || len(containers) != 1 {
			return fmt.Errorf(i18n.G("Failed to get the new container name"))
		}

		// Extract the name of the container
		fields := strings.Split(containers[0], "/")
		fmt.Printf(i18n.G("Container name is: %s")+"\n", fields[len(fields)-1])
	}

	return nil
}

func (c *cmdCopy) Run(cmd *cobra.Command, args []string) error {
	conf := c.global.conf

	// Sanity checks
	exit, err := c.global.CheckArgs(cmd, args, 1, 2)
	if exit {
		return err
	}

	// For copies, default to non-ephemeral and allow override (move uses -1)
	ephem := 0
	if c.flagEphemeral {
		ephem = 1
	}

	// Parse the mode
	mode := "pull"
	if c.flagMode != "" {
		mode = c.flagMode
	}

	stateful := !c.flagStateless && !c.flagRefresh

	// If not target name is specified, one will be chosed by the server
	if len(args) < 2 {
		return c.copyContainer(conf, args[0], "", false, ephem,
			stateful, c.flagContainerOnly, mode, c.flagStorage)
	}

	// Normal copy with a pre-determined name
	return c.copyContainer(conf, args[0], args[1], false, ephem,
		stateful, c.flagContainerOnly, mode, c.flagStorage)
}
<commit_msg>lxc/copy: Don't strip volatile keys on refresh<commit_after>package main

import (
	"fmt"
	"strings"

	"github.com/spf13/cobra"

	"github.com/lxc/lxd/client"
	"github.com/lxc/lxd/lxc/config"
	"github.com/lxc/lxd/lxc/utils"
	"github.com/lxc/lxd/shared"
	"github.com/lxc/lxd/shared/api"
	cli "github.com/lxc/lxd/shared/cmd"
	"github.com/lxc/lxd/shared/i18n"
)

type cmdCopy struct {
	global *cmdGlobal

	flagNoProfiles    bool
	flagProfile       []string
	flagConfig        []string
	flagDevice        []string
	flagEphemeral     bool
	flagContainerOnly bool
	flagMode          string
	flagStateless     bool
	flagStorage       string
	flagTarget        string
	flagTargetProject string
	flagRefresh       bool
}

func (c *cmdCopy) Command() *cobra.Command {
	cmd := &cobra.Command{}
	cmd.Use = i18n.G("copy [<remote>:]<source>[/<snapshot>] [[<remote>:]<destination>]")
	cmd.Aliases = []string{"cp"}
	cmd.Short = i18n.G("Copy containers within or in between LXD instances")
	cmd.Long = cli.FormatSection(i18n.G("Description"), i18n.G(
		`Copy containers within or in between LXD instances`))

	cmd.RunE = c.Run
	cmd.Flags().StringArrayVarP(&c.flagConfig, "config", "c", nil, i18n.G("Config key/value to apply to the new container")+"``")
	cmd.Flags().StringArrayVarP(&c.flagDevice, "device", "d", nil, i18n.G("New key/value to apply to a specific device")+"``")
	cmd.Flags().StringArrayVarP(&c.flagProfile, "profile", "p", nil, i18n.G("Profile to apply to the new container")+"``")
	cmd.Flags().BoolVarP(&c.flagEphemeral, "ephemeral", "e", false, i18n.G("Ephemeral container"))
	cmd.Flags().StringVar(&c.flagMode, "mode", "pull", i18n.G("Transfer mode. One of pull (default), push or relay")+"``")
	cmd.Flags().BoolVar(&c.flagContainerOnly, "container-only", false, i18n.G("Copy the container without its snapshots"))
	cmd.Flags().BoolVar(&c.flagStateless, "stateless", false, i18n.G("Copy a stateful container stateless"))
	cmd.Flags().StringVarP(&c.flagStorage, "storage", "s", "", i18n.G("Storage pool name")+"``")
	cmd.Flags().StringVar(&c.flagTarget, "target", "", i18n.G("Cluster member name")+"``")
	cmd.Flags().StringVar(&c.flagTargetProject, "target-project", "", i18n.G("Copy to a project different from the source")+"``")
	cmd.Flags().BoolVar(&c.flagNoProfiles, "no-profiles", false, i18n.G("Create the container with no profiles applied"))
	cmd.Flags().BoolVar(&c.flagRefresh, "refresh", false, i18n.G("Perform an incremental copy"))

	return cmd
}

func (c *cmdCopy) copyContainer(conf *config.Config, sourceResource string,
	destResource string, keepVolatile bool, ephemeral int, stateful bool,
	containerOnly bool, mode string, pool string) error {
	// Parse the source
	sourceRemote, sourceName, err := conf.ParseRemote(sourceResource)
	if err != nil {
		return err
	}

	// Parse the destination
	destRemote, destName, err := conf.ParseRemote(destResource)
	if err != nil {
		return err
	}

	// Make sure we have a container or snapshot name
	if sourceName == "" {
		return fmt.Errorf(i18n.G("You must specify a source container name"))
	}

	// Check that a destination container was specified, if --target is passed.
	if destName == "" && c.flagTarget != "" {
		return fmt.Errorf(i18n.G("You must specify a destination container name when using --target"))
	}

	// If no destination name was provided, use the same as the source
	if destName == "" && destResource != "" {
		destName = sourceName
	}

	// Connect to the source host
	source, err := conf.GetContainerServer(sourceRemote)
	if err != nil {
		return err
	}

	// Connect to the destination host
	var dest lxd.ContainerServer
	if sourceRemote == destRemote {
		// Source and destination are the same
		dest = source
	} else {
		// Destination is different, connect to it
		dest, err = conf.GetContainerServer(destRemote)
		if err != nil {
			return err
		}
	}

	// Project copies
	if c.flagTargetProject != "" {
		dest = dest.UseProject(c.flagTargetProject)
	}

	// Confirm that --target is only used with a cluster
	if c.flagTarget != "" && !dest.IsClustered() {
		return fmt.Errorf(i18n.G("To use --target, the destination remote must be a cluster"))
	}

	// Parse the config overrides
	configMap := map[string]string{}
	for _, entry := range c.flagConfig {
		if !strings.Contains(entry, "=") {
			return fmt.Errorf(i18n.G("Bad key=value pair: %s"), entry)
		}

		fields := strings.SplitN(entry, "=", 2)
		configMap[fields[0]] = fields[1]
	}

	// Parse the device overrides
	deviceMap := map[string]map[string]string{}
	for _, entry := range c.flagDevice {
		if !strings.Contains(entry, "=") || !strings.Contains(entry, ",") {
			return fmt.Errorf(i18n.G("Bad syntax, expecting <device>,<key>=<value>: %s"), entry)
		}

		deviceFields := strings.SplitN(entry, ",", 2)
		keyFields := strings.SplitN(deviceFields[1], "=", 2)

		if deviceMap[deviceFields[0]] == nil {
			deviceMap[deviceFields[0]] = map[string]string{}
		}

		deviceMap[deviceFields[0]][keyFields[0]] = keyFields[1]
	}

	var op lxd.RemoteOperation
	var writable api.ContainerPut

	if shared.IsSnapshot(sourceName) {
		if containerOnly {
			return fmt.Errorf(i18n.G("--container-only can't be passed when the source is a snapshot"))
		}

		// Prepare the container creation request
		args := lxd.ContainerSnapshotCopyArgs{
			Name: destName,
			Mode: mode,
			Live: stateful,
		}

		if c.flagRefresh {
			return fmt.Errorf(i18n.G("--refresh can only be used with containers"))
		}

		// Copy of a snapshot into a new container
		srcFields := strings.SplitN(sourceName, shared.SnapshotDelimiter, 2)
		entry, _, err := source.GetContainerSnapshot(srcFields[0], srcFields[1])
		if err != nil {
			return err
		}

		// Allow adding additional profiles
		if c.flagProfile != nil {
			entry.Profiles = append(entry.Profiles, c.flagProfile...)
		} else if c.flagNoProfiles {
			entry.Profiles = []string{}
		}

		// Allow setting additional config keys
		if configMap != nil {
			for key, value := range configMap {
				entry.Config[key] = value
			}
		}

		// Allow setting device overrides
		if deviceMap != nil {
			for k, m := range deviceMap {
				if entry.Devices[k] == nil {
					entry.Devices[k] = m
					continue
				}

				for key, value := range m {
					entry.Devices[k][key] = value
				}
			}
		}

		// Allow overriding the ephemeral status
		if ephemeral == 1 {
			entry.Ephemeral = true
		} else if ephemeral == 0 {
			entry.Ephemeral = false
		}

		rootDiskDeviceKey, _, _ := shared.GetRootDiskDevice(entry.Devices)
		if err != nil {
			return err
		}

		if rootDiskDeviceKey != "" && pool != "" {
			entry.Devices[rootDiskDeviceKey]["pool"] = pool
		} else if pool != "" {
			entry.Devices["root"] = map[string]string{
				"type": "disk",
				"path": "/",
				"pool": pool,
			}
		}

		// Strip the volatile keys if requested
		if !keepVolatile {
			for k := range entry.Config {
				if k == "volatile.base_image" {
					continue
				}

				if strings.HasPrefix(k, "volatile") {
					delete(entry.Config, k)
				}
			}
		}

		// Do the actual copy
		if c.flagTarget != "" {
			dest = dest.UseTarget(c.flagTarget)
		}

		op, err = dest.CopyContainerSnapshot(source, srcFields[0], *entry, &args)
		if err != nil {
			return err
		}
	} else {
		// Prepare the container creation request
		args := lxd.ContainerCopyArgs{
			Name:          destName,
			Live:          stateful,
			ContainerOnly: containerOnly,
			Mode:          mode,
			Refresh:       c.flagRefresh,
		}

		// Copy of a container into a new container
		entry, _, err := source.GetContainer(sourceName)
		if err != nil {
			return err
		}

		// Allow adding additional profiles
		if c.flagProfile != nil {
			entry.Profiles = append(entry.Profiles, c.flagProfile...)
		} else if c.flagNoProfiles {
			entry.Profiles = []string{}
		}

		// Allow setting additional config keys
		if configMap != nil {
			for key, value := range configMap {
				entry.Config[key] = value
			}
		}

		// Allow setting device overrides
		if deviceMap != nil {
			for k, m := range deviceMap {
				if entry.Devices[k] == nil {
					entry.Devices[k] = m
					continue
				}

				for key, value := range m {
					entry.Devices[k][key] = value
				}
			}
		}

		// Allow overriding the ephemeral status
		if ephemeral == 1 {
			entry.Ephemeral = true
		} else if ephemeral == 0 {
			entry.Ephemeral = false
		}

		rootDiskDeviceKey, _, _ := shared.GetRootDiskDevice(entry.Devices)
		if err != nil {
			return err
		}

		if rootDiskDeviceKey != "" && pool != "" {
			entry.Devices[rootDiskDeviceKey]["pool"] = pool
		} else if pool != "" {
			entry.Devices["root"] = map[string]string{
				"type": "disk",
				"path": "/",
				"pool": pool,
			}
		}

		// Strip the volatile keys if requested
		if !keepVolatile {
			for k := range entry.Config {
				if k == "volatile.base_image" {
					continue
				}

				if strings.HasPrefix(k, "volatile") {
					delete(entry.Config, k)
				}
			}
		}

		// Do the actual copy
		if c.flagTarget != "" {
			dest = dest.UseTarget(c.flagTarget)
		}

		op, err = dest.CopyContainer(source, *entry, &args)
		if err != nil {
			return err
		}

		writable = entry.Writable()
	}

	// Watch the background operation
	progress := utils.ProgressRenderer{
		Format: i18n.G("Transferring container: %s"),
		Quiet:  c.global.flagQuiet,
	}

	_, err = op.AddHandler(progress.UpdateOp)
	if err != nil {
		progress.Done("")
		return err
	}

	// Wait for the copy to complete
	err = utils.CancelableWait(op, &progress)
	if err != nil {
		progress.Done("")
		return err
	}
	progress.Done("")

	if c.flagRefresh {
		_, etag, err := dest.GetContainer(destName)
		if err != nil {
			return fmt.Errorf("Failed to refresh target container '%s': %v", destName, err)
		}

		op, err := dest.UpdateContainer(destName, writable, etag)
		if err != nil {
			return err
		}

		// Watch the background operation
		progress := utils.ProgressRenderer{
			Format: i18n.G("Refreshing container: %s"),
			Quiet:  c.global.flagQuiet,
		}

		_, err = op.AddHandler(progress.UpdateOp)
		if err != nil {
			progress.Done("")
			return err
		}

		// Wait for the copy to complete
		err = utils.CancelableWait(op, &progress)
		if err != nil {
			progress.Done("")
			return err
		}
		progress.Done("")
	}

	// If choosing a random name, show it to the user
	if destResource == "" {
		// Get the successful operation data
		opInfo, err := op.GetTarget()
		if err != nil {
			return err
		}

		// Extract the list of affected containers
		containers, ok := opInfo.Resources["containers"]
		if !ok || len(containers) != 1 {
			return fmt.Errorf(i18n.G("Failed to get the new container name"))
		}

		// Extract the name of the container
		fields := strings.Split(containers[0], "/")
		fmt.Printf(i18n.G("Container name is: %s")+"\n", fields[len(fields)-1])
	}

	return nil
}

func (c *cmdCopy) Run(cmd *cobra.Command, args []string) error {
	conf := c.global.conf

	// Sanity checks
	exit, err := c.global.CheckArgs(cmd, args, 1, 2)
	if exit {
		return err
	}

	// For copies, default to non-ephemeral and allow override (move uses -1)
	ephem := 0
	if c.flagEphemeral {
		ephem = 1
	}

	// Parse the mode
	mode := "pull"
	if c.flagMode != "" {
		mode = c.flagMode
	}

	stateful := !c.flagStateless && !c.flagRefresh
	keepVolatile := c.flagRefresh

	// If not target name is specified, one will be chosed by the server
	if len(args) < 2 {
		return c.copyContainer(conf, args[0], "", keepVolatile, ephem,
			stateful, c.flagContainerOnly, mode, c.flagStorage)
	}

	// Normal copy with a pre-determined name
	return c.copyContainer(conf, args[0], args[1], keepVolatile, ephem,
		stateful, c.flagContainerOnly, mode, c.flagStorage)
}
<|endoftext|>


<commit_before>/*
 * lxc_test.go: Go bindings for lxc
 *
 * Copyright © 2013, S.Çağlar Onur
 *
 * Authors:
 * S.Çağlar Onur <caglar@10ur.org>
 *
 * This library is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License version 2, as
 * published by the Free Software Foundation.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License along
 * with this program; if not, write to the Free Software Foundation, Inc.,
 * 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
 */

package lxc

import (
	"fmt"
	"testing"
)

func TestAll(t *testing.T) {
	z := NewContainer("rubik")

	fmt.Printf("Config file:%+v\n", z.ConfigFileName())
	fmt.Printf("Daemonize: %+v\n", z.Daemonize())
	fmt.Printf("Init PID: %+v\n", z.InitPID())
	fmt.Printf("Defined: %+v\n", z.Defined())
	fmt.Printf("Running: %+v\n", z.Running())
	fmt.Printf("State: %+v\n", z.State())
	z.SetDaemonize()
	fmt.Printf("Daemonize: %+v\n", z.Daemonize())

	if !z.Defined() {
		fmt.Printf("Creating rubik container...\n")
		fmt.Printf("Create: %+v\n", z.Create("ubuntu", []string{"amd64", "quantal"}))
	} else {
		fmt.Printf("Starting rubik container...\n\n")
		fmt.Printf("Start: %+v\n", z.Start(false, nil))
		fmt.Printf("State: %+v\n", z.State())
		fmt.Printf("Init PID: %+v\n", z.InitPID())
		fmt.Printf("Freeze: %+v\n", z.Freeze())
		fmt.Printf("State: %+v\n", z.State())
		fmt.Printf("Unfreeze: %+v\n", z.Unfreeze())
		fmt.Printf("State: %+v\n", z.State())
	}

	if z.Running() {
		fmt.Printf("Shutdown: %+v\n", z.Shutdown(30))
		fmt.Printf("State: %+v\n", z.State())
		fmt.Printf("Stop: %+v\n", z.Stop())
		fmt.Printf("State: %+v\n", z.State())
	}
	fmt.Printf("Destroy: %+v\n", z.Destroy())
}
<commit_msg>call Destroy after Shutdown/Stop<commit_after>/*
 * lxc_test.go: Go bindings for lxc
 *
 * Copyright © 2013, S.Çağlar Onur
 *
 * Authors:
 * S.Çağlar Onur <caglar@10ur.org>
 *
 * This library is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License version 2, as
 * published by the Free Software Foundation.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License along
 * with this program; if not, write to the Free Software Foundation, Inc.,
 * 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
 */

package lxc

import (
	"fmt"
	"testing"
)

func TestAll(t *testing.T) {
	z := NewContainer("rubik")

	fmt.Printf("Config file:%+v\n", z.ConfigFileName())
	fmt.Printf("Daemonize: %+v\n", z.Daemonize())
	fmt.Printf("Init PID: %+v\n", z.InitPID())
	fmt.Printf("Defined: %+v\n", z.Defined())
	fmt.Printf("Running: %+v\n", z.Running())
	fmt.Printf("State: %+v\n", z.State())
	z.SetDaemonize()
	fmt.Printf("Daemonize: %+v\n", z.Daemonize())

	if !z.Defined() {
		fmt.Printf("Creating rubik container...\n")
		fmt.Printf("Create: %+v\n", z.Create("ubuntu", []string{"amd64", "quantal"}))
	} else {
		fmt.Printf("Starting rubik container...\n\n")
		fmt.Printf("Start: %+v\n", z.Start(false, nil))
		fmt.Printf("State: %+v\n", z.State())
		fmt.Printf("Init PID: %+v\n", z.InitPID())
		fmt.Printf("Freeze: %+v\n", z.Freeze())
		fmt.Printf("State: %+v\n", z.State())
		fmt.Printf("Unfreeze: %+v\n", z.Unfreeze())
		fmt.Printf("State: %+v\n", z.State())
	}

	if z.Running() {
		fmt.Printf("Shutdown: %+v\n", z.Shutdown(30))
		fmt.Printf("State: %+v\n", z.State())
		fmt.Printf("Stop: %+v\n", z.Stop())
		fmt.Printf("State: %+v\n", z.State())
		fmt.Printf("Destroy: %+v\n", z.Destroy())
	}
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"os"
	"os/signal"
	"strconv"
	"strings"
	"syscall"

	"github.com/chai2010/gettext-go/gettext"
	"github.com/gorilla/websocket"
	"golang.org/x/crypto/ssh/terminal"

	"github.com/lxc/lxd"
	"github.com/lxc/lxd/shared"
	"github.com/lxc/lxd/shared/gnuflag"
)

type execCmd struct{}

func (c *execCmd) showByDefault() bool {
	return true
}

func (c *execCmd) usage() string {
	return gettext.Gettext(
		`Execute the specified command in a container.

lxc exec [remote:]container [--env EDITOR=/usr/bin/vim]... <command>`)
}

type envFlag []string

func (f *envFlag) String() string {
	return fmt.Sprint(*f)
}

func (f *envFlag) Set(value string) error {
	if f == nil {
		*f = make(envFlag, 1)
	} else {
		*f = append(*f, value)
	}
	return nil
}

var envArgs envFlag

func (c *execCmd) flags() {
	gnuflag.Var(&envArgs, "env", gettext.Gettext("An environment variable of the form HOME=/home/foo"))
}

func controlSocketHandler(c *lxd.Client, control *websocket.Conn) {
	for {
		width, height, err := terminal.GetSize(syscall.Stdout)
		if err != nil {
			continue
		}

		shared.Debugf("Window size is now: %dx%d", width, height)

		w, err := control.NextWriter(websocket.TextMessage)
		if err != nil {
			shared.Debugf("Got error getting next writer %s", err)
			break
		}

		msg := shared.ContainerExecControl{}
		msg.Command = "window-resize"
		msg.Args = make(map[string]string)
		msg.Args["width"] = strconv.Itoa(width)
		msg.Args["height"] = strconv.Itoa(height)

		buf, err := json.Marshal(msg)
		if err != nil {
			shared.Debugf("Failed to convert to json %s", err)
			break
		}
		_, err = w.Write(buf)

		w.Close()
		if err != nil {
			shared.Debugf("Got err writing %s", err)
			break
		}

		ch := make(chan os.Signal)
		signal.Notify(ch, syscall.SIGWINCH)
		sig := <-ch

		shared.Debugf("Received '%s signal', updating window geometry.", sig)
	}

	closeMsg := websocket.FormatCloseMessage(websocket.CloseNormalClosure, "")
	control.WriteMessage(websocket.CloseMessage, closeMsg)
}

func (c *execCmd) run(config *lxd.Config, args []string) error {
	if len(args) < 2 {
		return errArgs
	}

	remote, name := config.ParseRemoteAndContainer(args[0])
	d, err := lxd.NewClient(config, remote)
	if err != nil {
		return err
	}

	env := map[string]string{"HOME": "/root", "USER": "root"}
	myEnv := os.Environ()
	for _, ent := range myEnv {
		if strings.HasPrefix(ent, "TERM=") {
			env["TERM"] = ent[len("TERM="):]
		}
	}

	for _, arg := range envArgs {
		pieces := strings.SplitN(arg, "=", 2)
		value := ""
		if len(pieces) > 1 {
			value = pieces[1]
		}
		env[pieces[0]] = value
	}

	cfd := syscall.Stdout
	var oldttystate *terminal.State
	interactive := terminal.IsTerminal(cfd)
	if interactive {
		oldttystate, err = terminal.MakeRaw(cfd)
		if err != nil {
			return err
		}
		defer terminal.Restore(cfd, oldttystate)
	}

	handler := controlSocketHandler
	if !interactive {
		handler = nil
	}

	ret, err := d.Exec(name, args[1:], env, os.Stdin, os.Stdout, os.Stderr, handler)
	if err != nil {
		return err
	}

	if oldttystate != nil {
		/* A bit of a special case here: we want to exit with the same code as
		 * the process inside the container, so we explicitly exit here
		 * instead of returning an error.
		 *
		 * Additionally, since os.Exit() exits without running deferred
		 * functions, we restore the terminal explicitly.
		 */
		terminal.Restore(cfd, oldttystate)
	}

	/* we get the result of waitpid() here so we need to transform it */
	os.Exit(ret >> 8)
	return fmt.Errorf(gettext.Gettext("unreachable return reached"))
}
<commit_msg>Fix non-interactive mode detection<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"os"
	"os/signal"
	"strconv"
	"strings"
	"syscall"

	"github.com/chai2010/gettext-go/gettext"
	"github.com/gorilla/websocket"
	"golang.org/x/crypto/ssh/terminal"

	"github.com/lxc/lxd"
	"github.com/lxc/lxd/shared"
	"github.com/lxc/lxd/shared/gnuflag"
)

type execCmd struct{}

func (c *execCmd) showByDefault() bool {
	return true
}

func (c *execCmd) usage() string {
	return gettext.Gettext(
		`Execute the specified command in a container.

lxc exec [remote:]container [--env EDITOR=/usr/bin/vim]... <command>`)
}

type envFlag []string

func (f *envFlag) String() string {
	return fmt.Sprint(*f)
}

func (f *envFlag) Set(value string) error {
	if f == nil {
		*f = make(envFlag, 1)
	} else {
		*f = append(*f, value)
	}
	return nil
}

var envArgs envFlag

func (c *execCmd) flags() {
	gnuflag.Var(&envArgs, "env", gettext.Gettext("An environment variable of the form HOME=/home/foo"))
}

func controlSocketHandler(c *lxd.Client, control *websocket.Conn) {
	for {
		width, height, err := terminal.GetSize(syscall.Stdout)
		if err != nil {
			continue
		}

		shared.Debugf("Window size is now: %dx%d", width, height)

		w, err := control.NextWriter(websocket.TextMessage)
		if err != nil {
			shared.Debugf("Got error getting next writer %s", err)
			break
		}

		msg := shared.ContainerExecControl{}
		msg.Command = "window-resize"
		msg.Args = make(map[string]string)
		msg.Args["width"] = strconv.Itoa(width)
		msg.Args["height"] = strconv.Itoa(height)

		buf, err := json.Marshal(msg)
		if err != nil {
			shared.Debugf("Failed to convert to json %s", err)
			break
		}
		_, err = w.Write(buf)

		w.Close()
		if err != nil {
			shared.Debugf("Got err writing %s", err)
			break
		}

		ch := make(chan os.Signal)
		signal.Notify(ch, syscall.SIGWINCH)
		sig := <-ch

		shared.Debugf("Received '%s signal', updating window geometry.", sig)
	}

	closeMsg := websocket.FormatCloseMessage(websocket.CloseNormalClosure, "")
	control.WriteMessage(websocket.CloseMessage, closeMsg)
}

func (c *execCmd) run(config *lxd.Config, args []string) error {
	if len(args) < 2 {
		return errArgs
	}

	remote, name := config.ParseRemoteAndContainer(args[0])
	d, err := lxd.NewClient(config, remote)
	if err != nil {
		return err
	}

	env := map[string]string{"HOME": "/root", "USER": "root"}
	myEnv := os.Environ()
	for _, ent := range myEnv {
		if strings.HasPrefix(ent, "TERM=") {
			env["TERM"] = ent[len("TERM="):]
		}
	}

	for _, arg := range envArgs {
		pieces := strings.SplitN(arg, "=", 2)
		value := ""
		if len(pieces) > 1 {
			value = pieces[1]
		}
		env[pieces[0]] = value
	}

	cfd := syscall.Stdin
	var oldttystate *terminal.State
	interactive := terminal.IsTerminal(cfd)
	if interactive {
		oldttystate, err = terminal.MakeRaw(cfd)
		if err != nil {
			return err
		}
		defer terminal.Restore(cfd, oldttystate)
	}

	handler := controlSocketHandler
	if !interactive {
		handler = nil
	}

	ret, err := d.Exec(name, args[1:], env, os.Stdin, os.Stdout, os.Stderr, handler)
	if err != nil {
		return err
	}

	if oldttystate != nil {
		/* A bit of a special case here: we want to exit with the same code as
		 * the process inside the container, so we explicitly exit here
		 * instead of returning an error.
		 *
		 * Additionally, since os.Exit() exits without running deferred
		 * functions, we restore the terminal explicitly.
		 */
		terminal.Restore(cfd, oldttystate)
	}

	/* we get the result of waitpid() here so we need to transform it */
	os.Exit(ret >> 8)
	return fmt.Errorf(gettext.Gettext("unreachable return reached"))
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"bytes"
	"fmt"
	"os"
	"sort"
	"strings"

	"github.com/lxc/lxd"
	"github.com/lxc/lxd/shared/gnuflag"
	"github.com/lxc/lxd/shared/i18n"
)

type helpCmd struct {
	showAll bool
}

func (c *helpCmd) showByDefault() bool {
	return true
}

func (c *helpCmd) usage() string {
	return i18n.G(
		`Presents details on how to use LXD.

lxd help [--all]`)
}

func (c *helpCmd) flags() {
	gnuflag.BoolVar(&c.showAll, "all", false, i18n.G("Show all commands (not just interesting ones)"))
}

func (c *helpCmd) run(_ *lxd.Config, args []string) error {
	if len(args) > 0 {
		for _, name := range args {
			cmd, ok := commands[name]
			if !ok {
				fmt.Fprintf(os.Stderr, i18n.G("error: unknown command: %s")+"\n", name)
			} else {
				fmt.Fprintf(os.Stdout, cmd.usage()+"\n")
			}
		}
		return nil
	}

	fmt.Println(i18n.G("Usage: lxc [subcommand] [options]"))
	fmt.Println(i18n.G("Available commands:"))
	var names []string
	for name := range commands {
		names = append(names, name)
	}
	sort.Strings(names)
	for _, name := range names {
		cmd := commands[name]
		if c.showAll || cmd.showByDefault() {
			fmt.Printf("\t%-10s - %s\n", name, c.summaryLine(cmd.usage()))
		}
	}
	if !c.showAll {
		fmt.Println()
		fmt.Println(i18n.G("Options:"))
		fmt.Println("  --all              " + i18n.G("Print less common commands."))
		fmt.Println("  --debug            " + i18n.G("Print debug information."))
		fmt.Println("  --verbose          " + i18n.G("Print verbose information."))
		fmt.Println("  --version          " + i18n.G("Show client version."))
		fmt.Println()
		fmt.Println(i18n.G("Environment:"))
		fmt.Println("  LXD_CONF           " + i18n.G("Path to an alternate client configuration directory."))
		fmt.Println("  LXD_DIR            " + i18n.G("Path to an alternate server directory."))
	}
	return nil
}

// summaryLine returns the first line of the help text. Conventionally, this
// should be a one-line command summary, potentially followed by a longer
// explanation.
func (c *helpCmd) summaryLine(usage string) string {
	usage = strings.TrimSpace(usage)
	s := bufio.NewScanner(bytes.NewBufferString(usage))
	if s.Scan() {
		if len(s.Text()) > 1 {
			return s.Text()
		}
	}
	return i18n.G("Missing summary.")
}
<commit_msg>help: Update help<commit_after>package main

import (
	"bufio"
	"bytes"
	"fmt"
	"os"
	"sort"
	"strings"

	"github.com/lxc/lxd"
	"github.com/lxc/lxd/shared/gnuflag"
	"github.com/lxc/lxd/shared/i18n"
)

type helpCmd struct {
	showAll bool
}

func (c *helpCmd) showByDefault() bool {
	return true
}

func (c *helpCmd) usage() string {
	return i18n.G(
		`Help page for the LXD client.

lxc help [--all]`)
}

func (c *helpCmd) flags() {
	gnuflag.BoolVar(&c.showAll, "all", false, i18n.G("Show all commands (not just interesting ones)"))
}

func (c *helpCmd) run(_ *lxd.Config, args []string) error {
	if len(args) > 0 {
		for _, name := range args {
			cmd, ok := commands[name]
			if !ok {
				fmt.Fprintf(os.Stderr, i18n.G("error: unknown command: %s")+"\n", name)
			} else {
				fmt.Fprintf(os.Stdout, cmd.usage()+"\n")
			}
		}
		return nil
	}

	fmt.Println(i18n.G("Usage: lxc <command> [options]"))
	fmt.Println(i18n.G("Available commands:"))
	var names []string
	for name := range commands {
		names = append(names, name)
	}
	sort.Strings(names)
	for _, name := range names {
		cmd := commands[name]
		if c.showAll || cmd.showByDefault() {
			fmt.Printf("\t%-10s - %s\n", name, c.summaryLine(cmd.usage()))
		}
	}
	if !c.showAll {
		fmt.Println()
		fmt.Println(i18n.G("Options:"))
		fmt.Println("  --all              " + i18n.G("Print less common commands"))
		fmt.Println("  --debug            " + i18n.G("Print debug information"))
		fmt.Println("  --verbose          " + i18n.G("Print verbose information"))
		fmt.Println("  --version          " + i18n.G("Show client version"))
		fmt.Println()
		fmt.Println(i18n.G("Environment:"))
		fmt.Println("  LXD_CONF           " + i18n.G("Path to an alternate client configuration directory"))
		fmt.Println("  LXD_DIR            " + i18n.G("Path to an alternate server directory"))
	}
	return nil
}

// summaryLine returns the first line of the help text. Conventionally, this
// should be a one-line command summary, potentially followed by a longer
// explanation.
func (c *helpCmd) summaryLine(usage string) string {
	usage = strings.TrimSpace(usage)
	s := bufio.NewScanner(bytes.NewBufferString(usage))
	if s.Scan() {
		if len(s.Text()) > 1 {
			return s.Text()
		}
	}
	return i18n.G("Missing summary.")
}
<|endoftext|>


<commit_before>// Copyright (c) 2017 David R. Jenni. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Command wc implements word count as MapReduce job.
package main

import (
	"bufio"
	"bytes"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"runtime/pprof"
	"strconv"
	"strings"

	"github.com/davidrjenni/lib/mr"
)

type wordCount struct{}

func (w wordCount) Map(key, value string, out chan<- mr.Tuple) {
	s := bufio.NewScanner(strings.NewReader(value))
	s.Split(bufio.ScanWords)
	for s.Scan() {
		out <- mr.Tuple{s.Text(), "1"}
	}
	if err := s.Err(); err != nil {
		log.Println("map error: ", err)
	}
}

func (w wordCount) Reduce(key string, values []string, out chan<- mr.Tuple) {
	c := 0
	for _, v := range values {
		n, err := strconv.Atoi(v)
		if err != nil {
			log.Println("reduce error: ", err)
		} else {
			c += n
		}
	}
	out <- mr.Tuple{key, strconv.Itoa(c)}
}

func main() {
	log.SetFlags(0)
	log.SetPrefix("wc: ")

	var (
		input      = flag.String("input", "", "input text file")
		cpuprofile = flag.String("cpu", "", "cpu profile output")
	)

	flag.Parse()
	if *input == "" {
		flag.Usage()
		os.Exit(1)
	}

	b, err := ioutil.ReadFile(*input)
	if err != nil {
		log.Fatal("read input file: ", err)
	}

	if *cpuprofile != "" {
		f, err := os.Create(*cpuprofile)
		if err != nil {
			log.Fatal("create file: ", err)
		}
		pprof.StartCPUProfile(f)
		defer pprof.StopCPUProfile()
	}

	s := bufio.NewScanner(bytes.NewReader(b))
	s.Split(bufio.ScanLines)
	var values []string
	for s.Scan() {
		values = append(values, s.Text())
	}

	tuples := mr.Run(wordCount{}, values)
	for t := range tuples {
		fmt.Println(t.First, ":", t.Second)
	}
}
<commit_msg>mr/wc: keyify Tuple literals.<commit_after>// Copyright (c) 2017 David R. Jenni. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Command wc implements word count as MapReduce job.
package main

import (
	"bufio"
	"bytes"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"runtime/pprof"
	"strconv"
	"strings"

	"github.com/davidrjenni/lib/mr"
)

type wordCount struct{}

func (w wordCount) Map(key, value string, out chan<- mr.Tuple) {
	s := bufio.NewScanner(strings.NewReader(value))
	s.Split(bufio.ScanWords)
	for s.Scan() {
		out <- mr.Tuple{
			First:  s.Text(),
			Second: "1",
		}
	}
	if err := s.Err(); err != nil {
		log.Println("map error: ", err)
	}
}

func (w wordCount) Reduce(key string, values []string, out chan<- mr.Tuple) {
	c := 0
	for _, v := range values {
		n, err := strconv.Atoi(v)
		if err != nil {
			log.Println("reduce error: ", err)
		} else {
			c += n
		}
	}
	out <- mr.Tuple{
		First:  key,
		Second: strconv.Itoa(c),
	}
}

func main() {
	log.SetFlags(0)
	log.SetPrefix("wc: ")

	var (
		input      = flag.String("input", "", "input text file")
		cpuprofile = flag.String("cpu", "", "cpu profile output")
	)

	flag.Parse()
	if *input == "" {
		flag.Usage()
		os.Exit(1)
	}

	b, err := ioutil.ReadFile(*input)
	if err != nil {
		log.Fatal("read input file: ", err)
	}

	if *cpuprofile != "" {
		f, err := os.Create(*cpuprofile)
		if err != nil {
			log.Fatal("create file: ", err)
		}
		pprof.StartCPUProfile(f)
		defer pprof.StopCPUProfile()
	}

	s := bufio.NewScanner(bytes.NewReader(b))
	s.Split(bufio.ScanLines)
	var values []string
	for s.Scan() {
		values = append(values, s.Text())
	}

	tuples := mr.Run(wordCount{}, values)
	for t := range tuples {
		fmt.Println(t.First, ":", t.Second)
	}
}
<|endoftext|>


<commit_before>package mustache

import (
    "bytes"
    "container/vector"
    "fmt"
    "io"
    "io/ioutil"
    "os"
    "path"
    "reflect"
    "strings"
)

type textElement struct {
    text []byte
}

type varElement struct {
    name string
}

type sectionElement struct {
    name      string
    startline int
    elems     *vector.Vector
}

type template struct {
    data    string
    otag    string
    ctag    string
    p       int
    curline int
    dir     string
    elems   *vector.Vector
}

type parseError struct {
    line    int
    message string
}

func (p parseError) String() string { return fmt.Sprintf("line %d: %s", p.line, p.message) }

func (tmpl *template) readString(s string) (string, os.Error) {
    i := tmpl.p
    newlines := 0
    for true {
        //are we at the end of the string?
        if i+len(s) > len(tmpl.data) {
            return tmpl.data[tmpl.p:], os.EOF
        }

        if tmpl.data[i] == '\n' {
            newlines++
        }

        if tmpl.data[i] != s[0] {
            i++
            continue
        }

        match := true
        for j := 1; j < len(s); j++ {
            if s[j] != tmpl.data[i+j] {
                match = false
                break
            }
        }

        if match {
            e := i + len(s)
            text := tmpl.data[tmpl.p:e]
            tmpl.p = e

            tmpl.curline += newlines
            return text, nil
        } else {
            i++
        }
    }

    //should never be here
    return "", nil
}

func (tmpl *template) parsePartial(name string) (*template, os.Error) {
    filename := path.Join(tmpl.dir, name+".mustache")

    partial, err := ParseFile(filename)

    if err != nil {
        return nil, err
    }

    return partial, nil
}

func (tmpl *template) parseSection(section *sectionElement) os.Error {
    for {
        text, err := tmpl.readString(tmpl.otag)

        if err == os.EOF {
            return parseError{section.startline, "Section " + section.name + " has no closing tag"}
        }

        // put text into an item
        text = text[0 : len(text)-len(tmpl.otag)]
        section.elems.Push(&textElement{strings.Bytes(text)})

        text, err = tmpl.readString(tmpl.ctag)
        if err == os.EOF {
            //put the remaining text in a block
            return parseError{tmpl.curline, "unmatched open tag"}
        }

        //trim the close tag off the text
        tag := strings.TrimSpace(text[0 : len(text)-len(tmpl.ctag)])
        if len(tag) == 0 {
            return parseError{tmpl.curline, "empty tag"}
        }
        switch tag[0] {
        case '!':
            //ignore comment
            break
        case '#':
            name := strings.TrimSpace(tag[1:])

            if len(tmpl.data) > tmpl.p && tmpl.data[tmpl.p] == '\n' {
                tmpl.p += 1
            } else if len(tmpl.data) > tmpl.p+1 && tmpl.data[tmpl.p] == '\r' && tmpl.data[tmpl.p+1] == '\n' {
                tmpl.p += 2
            }

            se := sectionElement{name, tmpl.curline, new(vector.Vector)}
            err := tmpl.parseSection(&se)
            if err != nil {
                return err
            }
            section.elems.Push(&se)
        case '/':
            name := strings.TrimSpace(tag[1:])
            if name != section.name {
                return parseError{tmpl.curline, "interleaved closing tag: " + name}
            } else {
                return nil
            }
        case '>':
            name := strings.TrimSpace(tag[1:])
            partial, err := tmpl.parsePartial(name)
            if err != nil {
                return err
            }
            tmpl.elems.Push(partial)
        case '=':
            if tag[len(tag)-1] != '=' {
                panicln("Invalid meta tag")
            }
            tag = strings.TrimSpace(tag[1 : len(tag)-1])
            newtags := strings.Split(tag, " ", 0)
            if len(newtags) == 2 {
                tmpl.otag = newtags[0]
                tmpl.ctag = newtags[1]
            }
        default:
            section.elems.Push(&varElement{tag})
        }
    }

    return nil
}

func (tmpl *template) parse() os.Error {
    for {
        text, err := tmpl.readString(tmpl.otag)

        if err == os.EOF {
            //put the remaining text in a block
            tmpl.elems.Push(&textElement{strings.Bytes(text)})
            return nil
        }

        // put text into an item
        text = text[0 : len(text)-len(tmpl.otag)]
        tmpl.elems.Push(&textElement{strings.Bytes(text)})

        text, err = tmpl.readString(tmpl.ctag)
        if err == os.EOF {
            //put the remaining text in a block
            return parseError{tmpl.curline, "unmatched open tag"}
        }

        //trim the close tag off the text
        tag := strings.TrimSpace(text[0 : len(text)-len(tmpl.ctag)])
        if len(tag) == 0 {
            return parseError{tmpl.curline, "empty tag"}
        }
        switch tag[0] {
        case '!':
            //ignore comment
            break
        case '#':
            name := strings.TrimSpace(tag[1:])

            if len(tmpl.data) > tmpl.p && tmpl.data[tmpl.p] == '\n' {
                tmpl.p += 1
            } else if len(tmpl.data) > tmpl.p+1 && tmpl.data[tmpl.p] == '\r' && tmpl.data[tmpl.p+1] == '\n' {
                tmpl.p += 2
            }

            se := sectionElement{name, tmpl.curline, new(vector.Vector)}
            err := tmpl.parseSection(&se)
            if err != nil {
                return err
            }
            tmpl.elems.Push(&se)
        case '/':
            return parseError{tmpl.curline, "unmatched close tag"}
        case '>':
            name := strings.TrimSpace(tag[1:])
            partial, err := tmpl.parsePartial(name)
            if err != nil {
                return err
            }
            tmpl.elems.Push(partial)
        case '=':
            if tag[len(tag)-1] != '=' {
                panicln("Invalid meta tag")
            }
            tag = strings.TrimSpace(tag[1 : len(tag)-1])
            newtags := strings.Split(tag, " ", 0)
            if len(newtags) == 2 {
                tmpl.otag = newtags[0]
                tmpl.ctag = newtags[1]
            }
        default:
            tmpl.elems.Push(&varElement{tag})
        }
    }

    return nil
}

func lookup(context reflect.Value, name string) reflect.Value {
    //if the context is an interface, get the actual value
    if iface, ok := context.(*reflect.InterfaceValue); ok && !iface.IsNil() {
        context = iface.Elem()
    }

    //the context may be a pointer, so do an indirect
    contextInd := reflect.Indirect(context)

    var ret reflect.Value = nil

    switch val := contextInd.(type) {
    case *reflect.MapValue:
        ret = val.Elem(reflect.NewValue(name))
    case *reflect.StructValue:
        ret = val.FieldByName(name)
    }

    //if the lookup value is an interface, return the actual value
    if iface, ok := ret.(*reflect.InterfaceValue); ok && !iface.IsNil() {
        ret = iface.Elem()
    }

    return ret
}

func renderSection(section *sectionElement, context reflect.Value, buf io.Writer) {
    value := lookup(context, section.name)

    valueInd := reflect.Indirect(value)

    var contexts = new(vector.Vector)

    switch val := valueInd.(type) {
    case *reflect.BoolValue:
        if !val.Get() {
            return
        } else {
            contexts.Push(context)
        }
    case *reflect.SliceValue:
        for i := 0; i < val.Len(); i++ {
            contexts.Push(val.Elem(i))
        }
    case *reflect.ArrayValue:
        for i := 0; i < val.Len(); i++ {
            contexts.Push(val.Elem(i))
        }
    default:
        contexts.Push(context)
    }

    //by default we execute the section
    for j := 0; j < contexts.Len(); j++ {
        ctx := contexts.At(j).(reflect.Value)
        for i := 0; i < section.elems.Len(); i++ {
            renderElement(section.elems.At(i), ctx, buf)
        }
    }
}

func renderElement(element interface{}, context reflect.Value, buf io.Writer) {

    switch elem := element.(type) {
    case *textElement:
        buf.Write(elem.text)
    case *varElement:
        val := lookup(context, elem.name)
        if val != nil {
            fmt.Fprint(buf, val.Interface())
        }
    case *sectionElement:
        renderSection(elem, context, buf)
    case *template:
        elem.renderTemplate(context, buf)
    }
}

func (tmpl *template) renderTemplate(context reflect.Value, buf io.Writer) {
    for i := 0; i < tmpl.elems.Len(); i++ {
        renderElement(tmpl.elems.At(i), context, buf)
    }
}

func (tmpl *template) Render(context interface{}, buf io.Writer) {
    val := reflect.NewValue(context)
    tmpl.renderTemplate(val, buf)
}

func ParseString(data string) (*template, os.Error) {
    cwd := os.Getenv("CWD")
    tmpl := template{data, "{{", "}}", 0, 1, cwd, new(vector.Vector)}
    err := tmpl.parse()

    if err != nil {
        return nil, err
    }

    return &tmpl, err
}

func ParseFile(filename string) (*template, os.Error) {
    data, err := ioutil.ReadFile(filename)

    if err != nil {
        return nil, err
    }

    dirname, _ := path.Split(filename)

    tmpl := template{string(data), "{{", "}}", 0, 1, dirname, new(vector.Vector)}
    err = tmpl.parse()

    if err != nil {
        return nil, err
    }

    return &tmpl, nil
}

func Render(data string, context interface{}) (string, os.Error) {
    tmpl, err := ParseString(data)

    if err != nil {
        return "", err
    }

    var buf bytes.Buffer
    tmpl.Render(context, &buf)

    return buf.String(), nil
}

func RenderFile(filename string, context interface{}) (string, os.Error) {
    tmpl, err := ParseFile(filename)

    if err != nil {
        return "", err
    }

    var buf bytes.Buffer
    tmpl.Render(context, &buf)

    return buf.String(), nil
}
<commit_msg>strings.Bytes -> []byte<commit_after>package mustache

import (
    "bytes"
    "container/vector"
    "fmt"
    "io"
    "io/ioutil"
    "os"
    "path"
    "reflect"
    "strings"
)

type textElement struct {
    text []byte
}

type varElement struct {
    name string
}

type sectionElement struct {
    name      string
    startline int
    elems     *vector.Vector
}

type template struct {
    data    string
    otag    string
    ctag    string
    p       int
    curline int
    dir     string
    elems   *vector.Vector
}

type parseError struct {
    line    int
    message string
}

func (p parseError) String() string { return fmt.Sprintf("line %d: %s", p.line, p.message) }

func (tmpl *template) readString(s string) (string, os.Error) {
    i := tmpl.p
    newlines := 0
    for true {
        //are we at the end of the string?
        if i+len(s) > len(tmpl.data) {
            return tmpl.data[tmpl.p:], os.EOF
        }

        if tmpl.data[i] == '\n' {
            newlines++
        }

        if tmpl.data[i] != s[0] {
            i++
            continue
        }

        match := true
        for j := 1; j < len(s); j++ {
            if s[j] != tmpl.data[i+j] {
                match = false
                break
            }
        }

        if match {
            e := i + len(s)
            text := tmpl.data[tmpl.p:e]
            tmpl.p = e

            tmpl.curline += newlines
            return text, nil
        } else {
            i++
        }
    }

    //should never be here
    return "", nil
}

func (tmpl *template) parsePartial(name string) (*template, os.Error) {
    filename := path.Join(tmpl.dir, name+".mustache")

    partial, err := ParseFile(filename)

    if err != nil {
        return nil, err
    }

    return partial, nil
}

func (tmpl *template) parseSection(section *sectionElement) os.Error {
    for {
        text, err := tmpl.readString(tmpl.otag)

        if err == os.EOF {
            return parseError{section.startline, "Section " + section.name + " has no closing tag"}
        }

        // put text into an item
        text = text[0 : len(text)-len(tmpl.otag)]
        section.elems.Push(&textElement{[]byte(text)})

        text, err = tmpl.readString(tmpl.ctag)
        if err == os.EOF {
            //put the remaining text in a block
            return parseError{tmpl.curline, "unmatched open tag"}
        }

        //trim the close tag off the text
        tag := strings.TrimSpace(text[0 : len(text)-len(tmpl.ctag)])
        if len(tag) == 0 {
            return parseError{tmpl.curline, "empty tag"}
        }
        switch tag[0] {
        case '!':
            //ignore comment
            break
        case '#':
            name := strings.TrimSpace(tag[1:])

            //ignore the newline when a section starts
            if len(tmpl.data) > tmpl.p && tmpl.data[tmpl.p] == '\n' {
                tmpl.p += 1
            } else if len(tmpl.data) > tmpl.p+1 && tmpl.data[tmpl.p] == '\r' && tmpl.data[tmpl.p+1] == '\n' {
                tmpl.p += 2
            }

            se := sectionElement{name, tmpl.curline, new(vector.Vector)}
            err := tmpl.parseSection(&se)
            if err != nil {
                return err
            }
            section.elems.Push(&se)
        case '/':
            name := strings.TrimSpace(tag[1:])
            if name != section.name {
                return parseError{tmpl.curline, "interleaved closing tag: " + name}
            } else {
                return nil
            }
        case '>':
            name := strings.TrimSpace(tag[1:])
            partial, err := tmpl.parsePartial(name)
            if err != nil {
                return err
            }
            tmpl.elems.Push(partial)
        case '=':
            if tag[len(tag)-1] != '=' {
                panicln("Invalid meta tag")
            }
            tag = strings.TrimSpace(tag[1 : len(tag)-1])
            newtags := strings.Split(tag, " ", 0)
            if len(newtags) == 2 {
                tmpl.otag = newtags[0]
                tmpl.ctag = newtags[1]
            }
        default:
            section.elems.Push(&varElement{tag})
        }
    }

    return nil
}

func (tmpl *template) parse() os.Error {
    for {
        text, err := tmpl.readString(tmpl.otag)

        if err == os.EOF {
            //put the remaining text in a block
            tmpl.elems.Push(&textElement{[]byte(text)})
            return nil
        }

        // put text into an item
        text = text[0 : len(text)-len(tmpl.otag)]
        tmpl.elems.Push(&textElement{[]byte(text)})

        text, err = tmpl.readString(tmpl.ctag)
        if err == os.EOF {
            //put the remaining text in a block
            return parseError{tmpl.curline, "unmatched open tag"}
        }

        //trim the close tag off the text
        tag := strings.TrimSpace(text[0 : len(text)-len(tmpl.ctag)])
        if len(tag) == 0 {
            return parseError{tmpl.curline, "empty tag"}
        }
        switch tag[0] {
        case '!':
            //ignore comment
            break
        case '#':
            name := strings.TrimSpace(tag[1:])

            if len(tmpl.data) > tmpl.p && tmpl.data[tmpl.p] == '\n' {
                tmpl.p += 1
            } else if len(tmpl.data) > tmpl.p+1 && tmpl.data[tmpl.p] == '\r' && tmpl.data[tmpl.p+1] == '\n' {
                tmpl.p += 2
            }

            se := sectionElement{name, tmpl.curline, new(vector.Vector)}
            err := tmpl.parseSection(&se)
            if err != nil {
                return err
            }
            tmpl.elems.Push(&se)
        case '/':
            return parseError{tmpl.curline, "unmatched close tag"}
        case '>':
            name := strings.TrimSpace(tag[1:])
            partial, err := tmpl.parsePartial(name)
            if err != nil {
                return err
            }
            tmpl.elems.Push(partial)
        case '=':
            if tag[len(tag)-1] != '=' {
                panicln("Invalid meta tag")
            }
            tag = strings.TrimSpace(tag[1 : len(tag)-1])
            newtags := strings.Split(tag, " ", 0)
            if len(newtags) == 2 {
                tmpl.otag = newtags[0]
                tmpl.ctag = newtags[1]
            }
        default:
            tmpl.elems.Push(&varElement{tag})
        }
    }

    return nil
}

func lookup(context reflect.Value, name string) reflect.Value {
    //if the context is an interface, get the actual value
    if iface, ok := context.(*reflect.InterfaceValue); ok && !iface.IsNil() {
        context = iface.Elem()
    }

    //the context may be a pointer, so do an indirect
    contextInd := reflect.Indirect(context)

    var ret reflect.Value = nil

    switch val := contextInd.(type) {
    case *reflect.MapValue:
        ret = val.Elem(reflect.NewValue(name))
    case *reflect.StructValue:
        ret = val.FieldByName(name)
    }

    //if the lookup value is an interface, return the actual value
    if iface, ok := ret.(*reflect.InterfaceValue); ok && !iface.IsNil() {
        ret = iface.Elem()
    }

    return ret
}

func renderSection(section *sectionElement, context reflect.Value, buf io.Writer) {
    value := lookup(context, section.name)

    valueInd := reflect.Indirect(value)

    var contexts = new(vector.Vector)

    switch val := valueInd.(type) {
    case *reflect.BoolValue:
        if !val.Get() {
            return
        } else {
            contexts.Push(context)
        }
    case *reflect.SliceValue:
        for i := 0; i < val.Len(); i++ {
            contexts.Push(val.Elem(i))
        }
    case *reflect.ArrayValue:
        for i := 0; i < val.Len(); i++ {
            contexts.Push(val.Elem(i))
        }
    default:
        contexts.Push(context)
    }

    //by default we execute the section
    for j := 0; j < contexts.Len(); j++ {
        ctx := contexts.At(j).(reflect.Value)
        for i := 0; i < section.elems.Len(); i++ {
            renderElement(section.elems.At(i), ctx, buf)
        }
    }
}

func renderElement(element interface{}, context reflect.Value, buf io.Writer) {

    switch elem := element.(type) {
    case *textElement:
        buf.Write(elem.text)
    case *varElement:
        val := lookup(context, elem.name)
        if val != nil {
            fmt.Fprint(buf, val.Interface())
        }
    case *sectionElement:
        renderSection(elem, context, buf)
    case *template:
        elem.renderTemplate(context, buf)
    }
}

func (tmpl *template) renderTemplate(context reflect.Value, buf io.Writer) {
    for i := 0; i < tmpl.elems.Len(); i++ {
        renderElement(tmpl.elems.At(i), context, buf)
    }
}

func (tmpl *template) Render(context interface{}, buf io.Writer) {
    val := reflect.NewValue(context)
    tmpl.renderTemplate(val, buf)
}

func ParseString(data string) (*template, os.Error) {
    cwd := os.Getenv("CWD")
    tmpl := template{data, "{{", "}}", 0, 1, cwd, new(vector.Vector)}
    err := tmpl.parse()

    if err != nil {
        return nil, err
    }

    return &tmpl, err
}

func ParseFile(filename string) (*template, os.Error) {
    data, err := ioutil.ReadFile(filename)

    if err != nil {
        return nil, err
    }

    dirname, _ := path.Split(filename)

    tmpl := template{string(data), "{{", "}}", 0, 1, dirname, new(vector.Vector)}
    err = tmpl.parse()

    if err != nil {
        return nil, err
    }

    return &tmpl, nil
}

func Render(data string, context interface{}) (string, os.Error) {
    tmpl, err := ParseString(data)

    if err != nil {
        return "", err
    }

    var buf bytes.Buffer
    tmpl.Render(context, &buf)

    return buf.String(), nil
}

func RenderFile(filename string, context interface{}) (string, os.Error) {
    tmpl, err := ParseFile(filename)

    if err != nil {
        return "", err
    }

    var buf bytes.Buffer
    tmpl.Render(context, &buf)

    return buf.String(), nil
}
<|endoftext|>