<commit_before>package par2

import (
	"bufio"
	"encoding/binary"
	"fmt"
	"io"
	"hash"
	"crypto/md5"
	"bytes"
	"errors"
	"math/big"
	"os"
)

type Fileset struct {
	setID [16]byte
	slicelen uint64
	complete bool
	files map[[16]byte]*File
	checksums map[[16]byte]chksum
}

type File struct {
	Name string
	md5 [16]byte
	md5_16k [16]byte
	length uint64
	checksums [][16]byte
}

func (f *File) numBlocks(fset *Fileset) int {
	blockcount := int(f.length / fset.slicelen)
	if f.length % fset.slicelen != 0 {
		blockcount++
	}
	return blockcount
}

type chksum struct {
	*File
	blockno int
}

// NewFileset reads r and returns a Fileset that can be used for verification and recovery of the files.
func NewFileset(r io.Reader) (*Fileset, error) {
	fset := &Fileset{}
	fset.files = make(map[[16]byte]*File)
	fset.checksums = make(map[[16]byte]chksum)
	bufr := bufio.NewReader(r)
	for {
		hdr, err := readHeader(bufr)
		if err == io.EOF {
			break
		} else if err != nil {
			return nil, err
		}
		if fset.setID == ([16]byte{}) {
			fset.setID = hdr.setID
		} else if hdr.setID != fset.setID {
			// this is weird and shouldn't happen
			return nil, errors.New("mismatched set ID in one file")
		}
		switch hdr.typ {
		case typeFileDesc:
			f, id, _ := readFileDesc(hdr, bufr)
			if f == nil {
				continue
			}
			fi, ok := fset.files[id]
			if ok {
				if fi.Name == "" {
					// file was discovered through some means
					// that didn't include the file info, usually IFSC
					// pkt. Fill in the information now that we know what file this is.
					f.checksums = fi.checksums
					fset.files[id] = f
				}
			} else {
				fset.files[id] = f
			}
		case typeIFSC:
			chksums, id, _ := readIFSC(hdr, bufr)
			if chksums == nil {
				continue
			}
			fi, ok := fset.files[id]
			if !ok {
				fi = new(File)
				fset.files[id] = fi
			}
			if fi.checksums == nil {
				fi.checksums = chksums
			}
			for i, chk := range fi.checksums {
				fset.checksums[chk] = chksum{
					File: fi,
					blockno: i,
				}
			}
		case typeMain:
			slicelen, ids, _ := readMain(hdr, bufr)
			for _, id := range ids {
				if _, ok := fset.files[id]; !ok {
					fset.files[id] = new(File)
				}
			}
			fset.slicelen = slicelen
		default:
		}
	}
	if !fset.CanVerify() {
		return fset, nil
	}
	return fset, nil
}

// CanVerify returns whether the current fileset can be
// used for verification.
func (f *Fileset) CanVerify() bool {
	if f.complete {
		return true
	}
	if f.slicelen == 0 {
		return false
	}
	for _, file := range f.files {
		if file.Name == "" || file.checksums == nil {
			return false
		}
	}
	f.complete = true
	return true
}

// Verify verifies the files at paths against the fileset.
// It returns a list of matches and how many blocks are needed in order to repair.
func (f *Fileset) Verify(paths []string) ([]*FileMatch, int) {
	if !f.complete {
		return nil, 0
	}
	files := make(map[[16]byte]*File, len(f.files))
	for _, v := range f.files {
		files[v.md5] = v
	}
	matches := make([]*FileMatch, 0, len(paths))
	blocksNeeded := 0
	for _, s := range paths {
		fm, blocksmissing := f.verifyfile(files, s)
		if fm != nil && fm.File != nil {
			delete(files, fm.File.md5)
			blocksNeeded += blocksmissing
			matches = append(matches, fm)
		}
	}
	for _, fi := range files {
		matches = append(matches, &FileMatch{ Err:ErrMissing, File: fi })
		blocksNeeded += fi.numBlocks(f)
	}
	return matches, blocksNeeded
}

var ErrMissing = errors.New("par2: file missing")

func (fset *Fileset) verifyfile(files map[[16]byte]*File, s string) (*FileMatch, int) {
	file, err := os.Open(s)
	if err != nil {
		return &FileMatch{Err: err}, 0
	}
	defer file.Close()

	chk := md5.New()
	_, err = io.Copy(chk, file)
	if err != nil {
		return &FileMatch{Err: err},0
	}
	var md5sum [16]byte
	chk.Sum(md5sum[:0])
	fi := files[md5sum]
	if fi != nil {
		return &FileMatch{ Path: s, File: fi },0
	}

	// ok, this file is either corrupted or not part of the recovery set.
	// Since we're downloading these files via yenc (hopefully), we can always
	// assume that they will have holes where they were corrupted.
	_, err = file.Seek(0, os.SEEK_SET)
	if err != nil {
		return &FileMatch{Err: err}, 0
	}
	partialMatch := &FileMatch{}
	for {
		mdchk := md5.New()
		_, err := io.CopyN(mdchk, file, int64(fset.slicelen))
		mdchk.Sum(md5sum[:0])
		if f, ok := fset.checksums[md5sum]; ok {
			if partialMatch.File == nil {
				// ok we have a match, init the block bitmap
				partialMatch.blocks = &big.Int{}
				partialMatch.File = f.File
				partialMatch.Path = s
			} else if partialMatch.File != f.File {
				// shit
				continue
			}
			partialMatch.blocks.SetBit(partialMatch.blocks, f.blockno, 1)
		}
		if err != nil {
			break
		}
	}
	if partialMatch.File == nil {
		// not part of the recovery set.
		return nil, 0
	}
	blockcount := partialMatch.File.numBlocks(fset)
	blocksmissing := 0
	for i := 0; i < blockcount; i++ {
		if partialMatch.blocks.Bit(i) == 0 {
			blocksmissing++
		}
	}
	return partialMatch, blocksmissing
}

type FileMatch struct {
	Err error
	Path string
	File *File
	blocks *big.Int
}

type hdr struct {
	length uint64
	hash   [16]byte
	setID  [16]byte
	typ  typ
	partialhash hash.Hash
}

func (h *hdr) String() string {
	return fmt.Sprintf("{\n\t%d\n\t%x\n\t%x\n\t%s\n}", h.length, h.hash, h.setID, h.typ)
}

func readHeader(r *bufio.Reader) (h hdr, err error) {
	err = findHeader(r)
	if err != nil {
		return h, err
	}
	defer func() {
		if err == io.EOF {
			err = io.ErrUnexpectedEOF
		}
	}()
	var buf [56]byte
	b := buf[:]
	_, err = io.ReadFull(r, b)
	if err != nil {
		return h, err
	}
	h.length, b = readint(b)
	// for convenience, turn the length into
	// length after the header has been read.
	h.length -= 64
	h.hash, b = readmd5(b)

	// create a partial hash so that we can create a hash of the
	// entire packet.
	h.partialhash = md5.New()
	h.partialhash.Write(b)

	h.setID, b = readmd5(b)

	// i know, it's not an md5sum
	typbuf, b := readmd5(b)
	switch typbuf {
	case magicMain:
		h.typ = typeMain
	case magicFiledesc:
		h.typ = typeFileDesc
	case magicIFSC:
		h.typ = typeIFSC
	case magicRecvSlic:
		h.typ = typeRecvSlic
	case magicCreator:
		h.typ = typeCreator
	default:
		h.typ = typeUnknown
	}

	return h, nil
}

// These are all the required types.
// Going to leave it at that
var (
	magicMain     = [16]byte{'P', 'A', 'R', ' ', '2', '.', '0', 0, 'M', 'a', 'i', 'n', 0, 0, 0, 0}
	magicFiledesc = [16]byte{'P', 'A', 'R', ' ', '2', '.', '0', 0, 'F', 'i', 'l', 'e', 'D', 'e', 's', 'c'}
	magicIFSC     = [16]byte{'P', 'A', 'R', ' ', '2', '.', '0', 0, 'I', 'F', 'S', 'C', 0, 0, 0, 0}
	magicRecvSlic = [16]byte{'P', 'A', 'R', ' ', '2', '.', '0', 0, 'R', 'e', 'c', 'v', 'S', 'l', 'i', 'c'}
	magicCreator  = [16]byte{'P', 'A', 'R', ' ', '2', '.', '0', 0, 'C', 'r', 'e', 'a', 't', 'o', 'r', 0}
)

// corruption might cause the header
// to become unaligned, so just search for the damn
// header.
func findHeader(r *bufio.Reader) error {
	i := 0
	str := "PAR2\x00PKT"
	for {
		//when completely matched
		if i == len(str) {
			return nil
		}
		c, err := r.ReadByte()
		if err != nil {
			return err
		}
		if str[i] == c {
			i++
			continue
		}
		i = 0
		if str[0] == c {
			i++
		}
	}
}

func readFileDesc(h hdr, r *bufio.Reader) (f *File, id [16]byte, err error) {
	buf, err := readPkt(h, r)
	if err != nil {
		return nil, id, nil
	}
	f = new(File)
	id, buf = readmd5(buf)
	f.md5, buf = readmd5(buf)
	f.md5_16k, buf = readmd5(buf)
	f.length, buf = readint(buf)

	// rest of block is name, trim 0 padding.
	i := bytes.LastIndex(buf, zero)
	if i != -1 {
		buf = buf[:i]
	}
	f.Name = string(buf)
	return f, id, nil
}

var zero = []byte{0}

func readIFSC(h hdr, r *bufio.Reader) (ss [][16]byte, id [16]byte, err error) {
	buf, err := readPkt(h, r)
	if err != nil {
		return nil, id, nil
	}
	id, buf = readmd5(buf)
	ss = make([][16]byte, 0, len(buf)/20)
	for len(buf) > 0 {
		var md5h [16]byte
		md5h, buf = readmd5(buf)
		// don't care about the crc, just that it gets consumed
		_, buf = readcrc(buf)
		ss = append(ss, md5h)
	}
	return ss, id, nil
}

func readMain(h hdr, r *bufio.Reader) (slicesize uint64, ids [][16]byte, err error) {
	buf, err := readPkt(h, r)
	if err != nil {
		return 0, nil, nil
	}
	slicesize, buf = readint(buf)
	numfiles, buf := readcrc(buf)
	ids = make([][16]byte, 0, numfiles)
	for i := uint32(0); i < numfiles; i++ {
		var nid [16]byte
		nid, buf = readmd5(buf)
		ids = append(ids, nid)
	}

	return slicesize, ids, nil
}

func readmd5(b []byte) ([16]byte, []byte) {
	var ret [16]byte
	copy(ret[:], b[:16])
	b = b[16:]
	return ret, b
}

func readint(b []byte) (uint64, []byte) {
	ret := binary.LittleEndian.Uint64(b)
	b = b[8:]
	return ret, b
}

func readcrc(b []byte) (uint32, []byte) {
	ret := binary.LittleEndian.Uint32(b)
	b = b[4:]
	return ret, b
}

// to be used for packets that can be read in a single go. Verifies
// the md5 sum of the packet as well
func readPkt(h hdr, r *bufio.Reader) ([]byte, error) {
	buf := make([]byte, h.length)
	_, err := io.ReadFull(r, buf)
	if err != nil {
		return nil, err
	}
	// check if the packet is valid
	h.partialhash.Write(buf)
	hash := h.partialhash.Sum(nil)
	for i, b := range h.hash {
		if hash[i] != b {
			return nil, errors.New("mismatch packet md5 and packet contents")
		}
	}
	return buf, nil
}
<commit_msg>gofmt<commit_after>package par2

import (
	"bufio"
	"bytes"
	"crypto/md5"
	"encoding/binary"
	"errors"
	"fmt"
	"hash"
	"io"
	"math/big"
	"os"
)

type Fileset struct {
	setID     [16]byte
	slicelen  uint64
	complete  bool
	files     map[[16]byte]*File
	checksums map[[16]byte]chksum
}

type File struct {
	Name      string
	md5       [16]byte
	length    uint64
	checksums [][16]byte
}

func (f *File) numBlocks(fset *Fileset) int {
	blockcount := int(f.length / fset.slicelen)
	if f.length%fset.slicelen != 0 {
		blockcount++
	}
	return blockcount
}

type chksum struct {
	*File
	blockno int
}

// NewFileset reads r and returns a Fileset that can be used for verification and recovery of the files.
func NewFileset(r io.Reader) (*Fileset, error) {
	fset := &Fileset{}
	fset.files = make(map[[16]byte]*File)
	fset.checksums = make(map[[16]byte]chksum)
	bufr := bufio.NewReader(r)
	for {
		hdr, err := readHeader(bufr)
		if err == io.EOF {
			break
		} else if err != nil {
			return nil, err
		}
		if fset.setID == ([16]byte{}) {
			fset.setID = hdr.setID
		} else if hdr.setID != fset.setID {
			// this is weird and shouldn't happen
			return nil, errors.New("mismatched set ID in one file")
		}
		switch hdr.typ {
		case typeFileDesc:
			f, id, _ := readFileDesc(hdr, bufr)
			if f == nil {
				continue
			}
			fi, ok := fset.files[id]
			if ok {
				if fi.Name == "" {
					// file was discovered through some means
					// that didn't include the file info, usually IFSC
					// pkt. Fill in the information now that we know what file this is.
					f.checksums = fi.checksums
					fset.files[id] = f
				}
			} else {
				fset.files[id] = f
			}
		case typeIFSC:
			chksums, id, _ := readIFSC(hdr, bufr)
			if chksums == nil {
				continue
			}
			fi, ok := fset.files[id]
			if !ok {
				fi = new(File)
				fset.files[id] = fi
			}
			if fi.checksums == nil {
				fi.checksums = chksums
			}
			for i, chk := range fi.checksums {
				fset.checksums[chk] = chksum{
					File:    fi,
					blockno: i,
				}
			}
		case typeMain:
			slicelen, ids, _ := readMain(hdr, bufr)
			for _, id := range ids {
				if _, ok := fset.files[id]; !ok {
					fset.files[id] = new(File)
				}
			}
			fset.slicelen = slicelen
		default:
		}
	}
	if !fset.CanVerify() {
		return fset, nil
	}
	return fset, nil
}

// CanVerify returns whether the current fileset can be
// used for verification.
func (f *Fileset) CanVerify() bool {
	if f.complete {
		return true
	}
	if f.slicelen == 0 {
		return false
	}
	for _, file := range f.files {
		if file.Name == "" || file.checksums == nil {
			return false
		}
	}
	f.complete = true
	return true
}

// Verify verifies the files at paths against the fileset.
// It returns a list of matches and how many blocks are needed in order to repair.
func (f *Fileset) Verify(paths []string) ([]*FileMatch, int) {
	if !f.complete {
		return nil, 0
	}
	files := make(map[[16]byte]*File, len(f.files))
	for _, v := range f.files {
		files[v.md5] = v
	}
	matches := make([]*FileMatch, 0, len(paths))
	blocksNeeded := 0
	for _, s := range paths {
		fm, blocksmissing := f.verifyfile(files, s)
		if fm != nil && fm.File != nil {
			delete(files, fm.File.md5)
			blocksNeeded += blocksmissing
			matches = append(matches, fm)
		}
	}
	for _, fi := range files {
		matches = append(matches, &FileMatch{Err: ErrMissing, File: fi})
		blocksNeeded += fi.numBlocks(f)
	}
	return matches, blocksNeeded
}

var ErrMissing = errors.New("par2: file missing")

func (fset *Fileset) verifyfile(files map[[16]byte]*File, s string) (*FileMatch, int) {
	file, err := os.Open(s)
	if err != nil {
		return &FileMatch{Err: err}, 0
	}
	defer file.Close()

	chk := md5.New()
	_, err = io.Copy(chk, file)
	if err != nil {
		return &FileMatch{Err: err}, 0
	}
	var md5sum [16]byte
	chk.Sum(md5sum[:0])
	fi := files[md5sum]
	if fi != nil {
		return &FileMatch{Path: s, File: fi}, 0
	}

	// ok, this file is either corrupted or not part of the recovery set.
	// Since we're downloading these files via yenc (hopefully), we can always
	// assume that they will have holes where they were corrupted.
	_, err = file.Seek(0, os.SEEK_SET)
	if err != nil {
		return &FileMatch{Err: err}, 0
	}
	partialMatch := &FileMatch{}
	for {
		mdchk := md5.New()
		_, err := io.CopyN(mdchk, file, int64(fset.slicelen))
		mdchk.Sum(md5sum[:0])
		if f, ok := fset.checksums[md5sum]; ok {
			if partialMatch.File == nil {
				// ok we have a match, init the block bitmap
				partialMatch.blocks = &big.Int{}
				partialMatch.File = f.File
				partialMatch.Path = s
			} else if partialMatch.File != f.File {
				// shit
				continue
			}
			partialMatch.blocks.SetBit(partialMatch.blocks, f.blockno, 1)
		}
		if err != nil {
			break
		}
	}
	if partialMatch.File == nil {
		// not part of the recovery set.
		return nil, 0
	}
	blockcount := partialMatch.File.numBlocks(fset)
	blocksmissing := 0
	for i := 0; i < blockcount; i++ {
		if partialMatch.blocks.Bit(i) == 0 {
			blocksmissing++
		}
	}
	return partialMatch, blocksmissing
}

type FileMatch struct {
	Err    error
	Path   string
	File   *File
	blocks *big.Int
}

type hdr struct {
	length      uint64
	hash        [16]byte
	setID       [16]byte
	typ         typ
	partialhash hash.Hash
}

func (h *hdr) String() string {
	return fmt.Sprintf("{\n\t%d\n\t%x\n\t%x\n\t%s\n}", h.length, h.hash, h.setID, h.typ)
}

func readHeader(r *bufio.Reader) (h hdr, err error) {
	err = findHeader(r)
	if err != nil {
		return h, err
	}
	defer func() {
		if err == io.EOF {
			err = io.ErrUnexpectedEOF
		}
	}()
	var buf [56]byte
	b := buf[:]
	_, err = io.ReadFull(r, b)
	if err != nil {
		return h, err
	}
	h.length, b = readint(b)
	// for convenience, turn the length into
	// length after the header has been read.
	h.length -= 64
	h.hash, b = readmd5(b)

	// create a partial hash so that we can create a hash of the
	// entire packet.
	h.partialhash = md5.New()
	h.partialhash.Write(b)

	h.setID, b = readmd5(b)

	// i know, it's not an md5sum
	typbuf, b := readmd5(b)
	switch typbuf {
	case magicMain:
		h.typ = typeMain
	case magicFiledesc:
		h.typ = typeFileDesc
	case magicIFSC:
		h.typ = typeIFSC
	case magicRecvSlic:
		h.typ = typeRecvSlic
	case magicCreator:
		h.typ = typeCreator
	default:
		h.typ = typeUnknown
	}

	return h, nil
}

// These are all the required types.
// Going to leave it at that
var (
	magicMain     = [16]byte{'P', 'A', 'R', ' ', '2', '.', '0', 0, 'M', 'a', 'i', 'n', 0, 0, 0, 0}
	magicFiledesc = [16]byte{'P', 'A', 'R', ' ', '2', '.', '0', 0, 'F', 'i', 'l', 'e', 'D', 'e', 's', 'c'}
	magicIFSC     = [16]byte{'P', 'A', 'R', ' ', '2', '.', '0', 0, 'I', 'F', 'S', 'C', 0, 0, 0, 0}
	magicRecvSlic = [16]byte{'P', 'A', 'R', ' ', '2', '.', '0', 0, 'R', 'e', 'c', 'v', 'S', 'l', 'i', 'c'}
	magicCreator  = [16]byte{'P', 'A', 'R', ' ', '2', '.', '0', 0, 'C', 'r', 'e', 'a', 't', 'o', 'r', 0}
)

// corruption might cause the header
// to become unaligned, so just search for the damn
// header.
func findHeader(r *bufio.Reader) error {
	i := 0
	str := "PAR2\x00PKT"
	for {
		//when completely matched
		if i == len(str) {
			return nil
		}
		c, err := r.ReadByte()
		if err != nil {
			return err
		}
		if str[i] == c {
			i++
			continue
		}
		i = 0
		if str[0] == c {
			i++
		}
	}
}

func readFileDesc(h hdr, r *bufio.Reader) (f *File, id [16]byte, err error) {
	buf, err := readPkt(h, r)
	if err != nil {
		return nil, id, nil
	}
	f = new(File)
	id, buf = readmd5(buf)
	f.md5, buf = readmd5(buf)
	_, buf = readmd5(buf)
	f.length, buf = readint(buf)

	// rest of block is name, trim 0 padding.
	i := bytes.LastIndex(buf, zero)
	if i != -1 {
		buf = buf[:i]
	}
	f.Name = string(buf)
	return f, id, nil
}

var zero = []byte{0}

func readIFSC(h hdr, r *bufio.Reader) (ss [][16]byte, id [16]byte, err error) {
	buf, err := readPkt(h, r)
	if err != nil {
		return nil, id, nil
	}
	id, buf = readmd5(buf)
	ss = make([][16]byte, 0, len(buf)/20)
	for len(buf) > 0 {
		var md5h [16]byte
		md5h, buf = readmd5(buf)
		// don't care about the crc, just that it gets consumed
		_, buf = readcrc(buf)
		ss = append(ss, md5h)
	}
	return ss, id, nil
}

func readMain(h hdr, r *bufio.Reader) (slicesize uint64, ids [][16]byte, err error) {
	buf, err := readPkt(h, r)
	if err != nil {
		return 0, nil, nil
	}
	slicesize, buf = readint(buf)
	numfiles, buf := readcrc(buf)
	ids = make([][16]byte, 0, numfiles)
	for i := uint32(0); i < numfiles; i++ {
		var nid [16]byte
		nid, buf = readmd5(buf)
		ids = append(ids, nid)
	}

	return slicesize, ids, nil
}

func readmd5(b []byte) ([16]byte, []byte) {
	var ret [16]byte
	copy(ret[:], b[:16])
	b = b[16:]
	return ret, b
}

func readint(b []byte) (uint64, []byte) {
	ret := binary.LittleEndian.Uint64(b)
	b = b[8:]
	return ret, b
}

func readcrc(b []byte) (uint32, []byte) {
	ret := binary.LittleEndian.Uint32(b)
	b = b[4:]
	return ret, b
}

// to be used for packets that can be read in a single go. Verifies
// the md5 sum of the packet as well
func readPkt(h hdr, r *bufio.Reader) ([]byte, error) {
	buf := make([]byte, h.length)
	_, err := io.ReadFull(r, buf)
	if err != nil {
		return nil, err
	}
	// check if the packet is valid
	h.partialhash.Write(buf)
	hash := h.partialhash.Sum(nil)
	for i, b := range h.hash {
		if hash[i] != b {
			return nil, errors.New("mismatch packet md5 and packet contents")
		}
	}
	return buf, nil
}
<|endoftext|>


<commit_before>package qp

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

// ErrPayloadTooShort indicates that the message was not complete.
var ErrPayloadTooShort = errors.New("payload too short")

// Default is the protocol used by the raw Encode and Decode functions.
var Default = NineP2000

// Protocol defines a protocol message encoder/decoder
type Protocol interface {
	Decode(r io.Reader) (Message, error)
	Encode(w io.Writer, m Message) error
}

// MessageType is the type of the contained message.
type MessageType byte

// Message is an interface describing an item that can encode itself to a
// writer, decode itself from a reader, and inform how large the encoded form
// would be at the current time. It is also capable of getting/setting the
// message tag, which is merely a convenience feature to save a type assert
// for access to the tag.
type Message interface {
	encoding.BinaryUnmarshaler
	encoding.BinaryMarshaler
	GetTag() Tag
}

// Write write all the provided data unless and io error occurs.
func write(w io.Writer, b []byte) error {
	var (
		written int
		err     error
		l       = len(b)
	)
	for written < l {
		written, err = w.Write(b[written:])
		if err != nil {
			return err
		}
	}

	return nil
}

// DecodeHdr reads 5 bytes and returns the decoded size and message type. It
// may return an error if reading from the Reader fails.
func DecodeHdr(r io.Reader) (uint32, MessageType, error) {
	var (
		n    int
		size uint32
		mt   MessageType
		err  error
	)

	b := make([]byte, 5)
	n, err = io.ReadFull(r, b)
	if n < 5 {
		return 0, 0, err
	}
	size = binary.LittleEndian.Uint32(b[0:4])
	mt = MessageType(b[4])
	return size, mt, err
}

// Codec encodes/decodes messages using the provided message type <-> message
// conversion.
type Codec struct {
	M2MT func(Message) (MessageType, error)
	MT2M func(MessageType) (Message, error)
}

// Decode decodes an entire message, including header, and returns the message.
// It may return an error if reading from the Reader fails, or if a message
// tries to consume more data than the size of the header indicated, making the
// message invalid.
func (c *Codec) Decode(r io.Reader) (Message, error) {
	var (
		size uint32
		mt   MessageType
		err  error
	)
	if size, mt, err = DecodeHdr(r); err != nil {
		return nil, err
	}

	size -= HeaderSize

	b := make([]byte, size)
	n, err := io.ReadFull(r, b)
	if err != nil {
		return nil, err
	}
	if n != int(size) {
		return nil, errors.New("short read")
	}

	m, err := c.MT2M(mt)
	if err != nil {
		return nil, err
	}

	if err = m.UnmarshalBinary(b); err != nil {
		return nil, err
	}
	return m, nil
}

// Encode write a header and message to the provided writer. It returns an
// error if writing failed.
func (c *Codec) Encode(w io.Writer, m Message) error {
	var err error
	var mt MessageType
	if mt, err = c.M2MT(m); err != nil {
		return err
	}

	var b []byte
	if b, err = m.MarshalBinary(); err != nil {
		return err
	}

	h := make([]byte, 5)
	binary.LittleEndian.PutUint32(h[0:4], uint32(len(b)+HeaderSize))
	h[4] = byte(mt)

	if err = write(w, h); err != nil {
		return err
	}

	if err = write(w, b); err != nil {
		return err
	}

	return nil
}

// Decode is a convenience function for calling decode on the default
// protocol.
func Decode(r io.Reader) (Message, error) {
	return Default.Decode(r)
}

// Encode is a convenience function for calling encode on the default
// protocol.
func Encode(w io.Writer, d Message) error {
	return Default.Encode(w, d)
}
<commit_msg>Add new decoder<commit_after>package qp

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

const (
	DefaultBufSize = 16384
	DefaultMinBuf  = 2048
)

// ErrPayloadTooShort indicates that the message was not complete.
var ErrPayloadTooShort = errors.New("payload too short")

// Default is the protocol used by the raw Encode and Decode functions.
var Default = NineP2000

// Protocol defines a protocol message encoder/decoder
type Protocol interface {
	Decode(r io.Reader) (Message, error)
	Encode(w io.Writer, m Message) error
	MessageType(Message) (MessageType, error)
	Message(MessageType) (Message, error)
}

// MessageType is the type of the contained message.
type MessageType byte

// Message is an interface describing an item that can encode itself to a
// writer, decode itself from a reader, and inform how large the encoded form
// would be at the current time. It is also capable of getting/setting the
// message tag, which is merely a convenience feature to save a type assert
// for access to the tag.
type Message interface {
	encoding.BinaryUnmarshaler
	encoding.BinaryMarshaler
	GetTag() Tag
}

// Write write all the provided data unless and io error occurs.
func write(w io.Writer, b []byte) error {
	var (
		written int
		err     error
		l       = len(b)
	)
	for written < l {
		written, err = w.Write(b[written:])
		if err != nil {
			return err
		}
	}

	return nil
}

// DecodeHdr reads 5 bytes and returns the decoded size and message type. It
// may return an error if reading from the Reader fails.
func DecodeHdr(r io.Reader) (uint32, MessageType, error) {
	var (
		n    int
		size uint32
		mt   MessageType
		err  error
	)

	b := make([]byte, 5)
	n, err = io.ReadFull(r, b)
	if n < 5 {
		return 0, 0, err
	}
	size = binary.LittleEndian.Uint32(b[0:4])
	mt = MessageType(b[4])
	return size, mt, err
}

// Codec encodes/decodes messages using the provided message type <-> message
// conversion.
type Codec struct {
	M2MT func(Message) (MessageType, error)
	MT2M func(MessageType) (Message, error)
}

// Decode decodes an entire message, including header, and returns the message.
// It may return an error if reading from the Reader fails, or if a message
// tries to consume more data than the size of the header indicated, making the
// message invalid.
func (c *Codec) Decode(r io.Reader) (Message, error) {
	var (
		size uint32
		mt   MessageType
		err  error
	)
	if size, mt, err = DecodeHdr(r); err != nil {
		return nil, err
	}

	size -= HeaderSize

	b := make([]byte, size)
	n, err := io.ReadFull(r, b)
	if err != nil {
		return nil, err
	}
	if n != int(size) {
		return nil, errors.New("short read")
	}

	m, err := c.MT2M(mt)
	if err != nil {
		return nil, err
	}

	if err = m.UnmarshalBinary(b); err != nil {
		return nil, err
	}
	return m, nil
}

// Encode write a header and message to the provided writer. It returns an
// error if writing failed.
func (c *Codec) Encode(w io.Writer, m Message) error {
	var err error
	var mt MessageType
	if mt, err = c.M2MT(m); err != nil {
		return err
	}

	var b []byte
	if b, err = m.MarshalBinary(); err != nil {
		return err
	}

	h := make([]byte, 5)
	binary.LittleEndian.PutUint32(h[0:4], uint32(len(b)+HeaderSize))
	h[4] = byte(mt)

	if err = write(w, h); err != nil {
		return err
	}

	if err = write(w, b); err != nil {
		return err
	}

	return nil
}

func (c *Codec) MessageType(m Message) (MessageType, error) {
	return c.M2MT(m)
}

func (c *Codec) Message(mt MessageType) (Message, error) {
	return c.MT2M(mt)
}

// Decode is a convenience function for calling decode on the default
// protocol.
func Decode(r io.Reader) (Message, error) {
	return Default.Decode(r)
}

// Encode is a convenience function for calling encode on the default
// protocol.
func Encode(w io.Writer, d Message) error {
	return Default.Encode(w, d)
}

// Decoder implements a decoding loop that calls a callback for every decoded
// message. Replacing the protocol codec from a callback is valid and well-
// defined, with the result being that the next message will be decoded with
// that codec. Replacing the reader is also well-defined if the callback is
// certain that the server have not sent and will not send any further
// messages before the replace have occured. Otherwise, a partial message may
// be left in the buffer, breaking the Decoder.
type Decoder struct {
	Proto    Protocol
	Callback func(m Message) error
	Reader   io.Reader
	Stopped  bool
	BufSize  int
	MinBuf   int
}

// If the callback returns an error, the reader returns an error, the message
// type is invalid or the message fails to decode, the loop exits with an
// error.
func (d *Decoder) Run() error {
	d.Stopped = false

	if d.BufSize == 0 {
		d.BufSize = DefaultBufSize
	}

	if d.MinBuf == 0 {
		d.MinBuf = DefaultMinBuf
	}

	var (
		// total is the count of bytes in the buffer.
		total uint32

		// needed is the amount of bytes missing.
		needed int = HeaderSize

		// size is the decoded message body size, not including message header.
		size uint32

		// ptr is the start index at the buffer at the current time.
		ptr uint32

		// m is the decoded message.
		m Message

		// buf is the rading buffer
		buf = make([]byte, d.BufSize)
	)

	for !d.Stopped {
		n, err := d.Reader.Read(buf[total:])
		if err != nil {
			return err
		}

		total += uint32(n)
		needed -= n

		// Handle the data we got
		for needed <= 0 {
			if m == nil { // Read a header if no message struct is set.
				size = binary.LittleEndian.Uint32(buf[ptr:ptr+4]) - HeaderSize
				mt := MessageType(buf[ptr+4])

				// Update message body size, missing bytes and the current ptr.
				needed += int(size)
				ptr += HeaderSize

				// We try to fetch the message struct immediately - better to fail
				// early rather than late.
				if m, err = d.Proto.Message(mt); err != nil {
					return err
				}

			} else { // Otherwise, read a body for the message.
				if err = m.UnmarshalBinary(buf[ptr : ptr+size]); err != nil {
					return err
				}
				if err = d.Callback(m); err != nil {
					return err
				}
				m = nil

				needed += HeaderSize
				ptr += size
				size = 0
			}
		}

		// Buffer checks and reset
		l := len(buf)
		remaining := l - int(total)
		if -needed > l {
			// Okay, we need to scale our buffer, because the total size of the
			// buffer is not sufficient. The normal 9P behaviour would be to
			// fail, as the version request specifically sets a maximum message
			// size that should be equivalent to the buffer size.
			for -needed > l {
				l *= 2
			}

			newbuf := make([]byte, l)
			copy(newbuf, buf[ptr:total])
			total -= ptr
			ptr = 0
			buf = newbuf
		} else if needed > remaining || remaining < d.MinBuf {
			// The remaining part of the buffer is smaller than what we need, or
			// smaller than the minimum hint - time for a cleaning.
			copy(buf, buf[ptr:total])
			total -= ptr
			ptr = 0
		}

	}
	return nil
}

func (d *Decoder) Stop() {
	d.Stopped = true
}
<|endoftext|>


<commit_before>package garden

import (
	"io"
	"time"
)

//go:generate counterfeiter . Container

type Container interface {
	Handle() string

	// Stop stops a container.
	//
	// If kill is false, garden stops a container by sending the processes running inside it the SIGTERM signal.
	// It then waits for the processes to terminate before returning a response.
	// If one or more processes do not terminate within 10 seconds,
	// garden sends these processes the SIGKILL signal, killing them ungracefully.
	//
	// If kill is true, garden stops a container by sending the processing running inside it a SIGKILL signal.
	//
	// It is possible to copy files in to and out of a stopped container.
	// It is only when a container is destroyed that its filesystem is cleaned up.
	//
	// Errors:
	// * None.
	Stop(kill bool) error

	// Returns information about a container.
	Info() (ContainerInfo, error)

	// StreamIn streams data into a file in a container.
	//
	// Errors:
	// *  TODO.
	StreamIn(spec StreamInSpec) error

	// StreamOut streams a file out of a container.
	//
	// Errors:
	// * TODO.
	StreamOut(spec StreamOutSpec) (io.ReadCloser, error)

	// Returns the current bandwidth limits set for the container.
	CurrentBandwidthLimits() (BandwidthLimits, error)

	// Returns the current CPU limts set for the container.
	CurrentCPULimits() (CPULimits, error)

	// Returns the current disk limts set for the container.
	CurrentDiskLimits() (DiskLimits, error)

	// Returns the current memory limts set for the container.
	CurrentMemoryLimits() (MemoryLimits, error)

	// Map a port on the host to a port in the container so that traffic to the
	// host port is forwarded to the container port. This is deprecated in
	// favour of passing NetIn configuration in the ContainerSpec at creation
	// time.
	//
	// If a host port is not given, a port will be acquired from the server's port
	// pool.
	//
	// If a container port is not given, the port will be the same as the
	// container port.
	//
	// The resulting host and container ports are returned in that order.
	//
	// Errors:
	// * When no port can be acquired from the server's port pool.
	NetIn(hostPort, containerPort uint32) (uint32, uint32, error)

	// Whitelist outbound network traffic. This is deprecated in favour of passing
	// NetOut configuration in the ContainerSpec at creation time.
	//
	// If the configuration directive deny_networks is not used,
	// all networks are already whitelisted and this command is effectively a no-op.
	//
	// Later NetOut calls take precedence over earlier calls, which is
	// significant only in relation to logging.
	//
	// Errors:
	// * An error is returned if the NetOut call fails.
	NetOut(netOutRule NetOutRule) error

	// A Bulk call for NetOut. This is deprecated in favour of passing
	// NetOut configuration in the ContainerSpec at creation time.
	//
	// Errors:
	// * An error is returned if any of the NetOut calls fail.
	BulkNetOut(netOutRules []NetOutRule) error

	// Run a script inside a container.
	//
	// The root user will be mapped to a non-root UID in the host unless the container (not this process) was created with 'privileged' true.
	//
	// Errors:
	// * TODO.
	Run(ProcessSpec, ProcessIO) (Process, error)

	// Attach starts streaming the output back to the client from a specified process.
	//
	// Errors:
	// * processID does not refer to a running process.
	Attach(processID string, io ProcessIO) (Process, error)

	// Metrics returns the current set of metrics for a container
	Metrics() (Metrics, error)

	// Sets the grace time.
	SetGraceTime(graceTime time.Duration) error

	// Properties returns the current set of properties
	Properties() (Properties, error)

	// Property returns the value of the property with the specified name.
	//
	// Errors:
	// * When the property does not exist on the container.
	Property(name string) (string, error)

	// Set a named property on a container to a specified value.
	//
	// Errors:
	// * None.
	SetProperty(name string, value string) error

	// Remove a property with the specified name from a container.
	//
	// Errors:
	// * None.
	RemoveProperty(name string) error
}

// ProcessSpec contains parameters for running a script inside a container.
type ProcessSpec struct {
	// ID for the process. If empty, an ID will be generated.
	ID string `json:"id,omitempty"`

	// Path to command to execute.
	Path string `json:"path,omitempty"`

	// Arguments to pass to command.
	Args []string `json:"args,omitempty"`

	// Environment variables.
	Env []string `json:"env,omitempty"`

	// Working directory (default: home directory).
	Dir string `json:"dir,omitempty"`

	// The name of a user in the container to run the process as.
	// This must either be a username, or uid:gid.
	User string `json:"user,omitempty"`

	// Resource limits
	Limits ResourceLimits `json:"rlimits,omitempty"`

	// Limits to be applied to the newly created process
	OverrideContainerLimits *ProcessLimits `json:"limits,omitempty"`

	// Execute with a TTY for stdio.
	TTY *TTYSpec `json:"tty,omitempty"`

	// Execute process in own root filesystem, different from the other processes
	// in the container.
	Image ImageRef `json:"image,omitempty"`
}

type TTYSpec struct {
	WindowSize *WindowSize `json:"window_size,omitempty"`
}

type WindowSize struct {
	Columns int `json:"columns,omitempty"`
	Rows    int `json:"rows,omitempty"`
}

type ProcessIO struct {
	Stdin  io.Reader
	Stdout io.Writer
	Stderr io.Writer
}

//go:generate counterfeiter . Process

type Process interface {
	ID() string
	Wait() (int, error)
	SetTTY(TTYSpec) error
	Signal(Signal) error
}

type Signal int

const (
	SignalTerminate Signal = iota
	SignalKill
)

type PortMapping struct {
	HostPort      uint32
	ContainerPort uint32
}

type StreamInSpec struct {
	Path      string
	User      string
	TarStream io.Reader
}

type StreamOutSpec struct {
	Path string
	User string
}

// ContainerInfo holds information about a container.
type ContainerInfo struct {
	State         string        // Either "active" or "stopped".
	Events        []string      // List of events that occurred for the container. It currently includes only "oom" (Out Of Memory) event if it occurred.
	HostIP        string        // The IP address of the gateway which controls the host side of the container's virtual ethernet pair.
	ContainerIP   string        // The IP address of the container side of the container's virtual ethernet pair.
	ExternalIP    string        //
	ContainerPath string        // The path to the directory holding the container's files (both its control scripts and filesystem).
	ProcessIDs    []string      // List of running processes.
	Properties    Properties    // List of properties defined for the container.
	MappedPorts   []PortMapping //
}

type ContainerInfoEntry struct {
	Info ContainerInfo
	Err  *Error
}

type Metrics struct {
	MemoryStat  ContainerMemoryStat
	CPUStat     ContainerCPUStat
	DiskStat    ContainerDiskStat
	NetworkStat ContainerNetworkStat
}

type ContainerMetricsEntry struct {
	Metrics Metrics
	Err     *Error
}

type ContainerMemoryStat struct {
	ActiveAnon              uint64 `json:"active_anon"`
	ActiveFile              uint64 `json:"active_file"`
	Cache                   uint64 `json:"cache"`
	HierarchicalMemoryLimit uint64 `json:"hierarchical_memory_limit"`
	InactiveAnon            uint64 `json:"inactive_anon"`
	InactiveFile            uint64 `json:"inactive_file"`
	MappedFile              uint64 `json:"mapped_file"`
	Pgfault                 uint64 `json:"pgfault"`
	Pgmajfault              uint64 `json:"pgmajfault"`
	Pgpgin                  uint64 `json:"pgpgin"`
	Pgpgout                 uint64 `json:"pgpgout"`
	Rss                     uint64 `json:"rss"`
	TotalActiveAnon         uint64 `json:"total_active_anon"`
	TotalActiveFile         uint64 `json:"total_active_file"`
	TotalCache              uint64 `json:"total_cache"`
	TotalInactiveAnon       uint64 `json:"total_inactive_anon"`
	TotalInactiveFile       uint64 `json:"total_inactive_file"`
	TotalMappedFile         uint64 `json:"total_mapped_file"`
	TotalPgfault            uint64 `json:"total_pgfault"`
	TotalPgmajfault         uint64 `json:"total_pgmajfault"`
	TotalPgpgin             uint64 `json:"total_pgpgin"`
	TotalPgpgout            uint64 `json:"total_pgpgout"`
	TotalRss                uint64 `json:"total_rss"`
	TotalUnevictable        uint64 `json:"total_unevictable"`
	Unevictable             uint64 `json:"unevictable"`
	Swap                    uint64 `json:"swap"`
	HierarchicalMemswLimit  uint64 `json:"hierarchical_memsw_limit"`
	TotalSwap               uint64 `json:"total_swap"`
	// A memory usage total which reports memory usage in the same way that limits are enforced.
	// This value includes memory consumed by nested containers.
	TotalUsageTowardLimit uint64
}

type ContainerCPUStat struct {
	Usage  uint64
	User   uint64
	System uint64
}

type ContainerDiskStat struct {
	TotalBytesUsed      uint64
	TotalInodesUsed     uint64
	ExclusiveBytesUsed  uint64
	ExclusiveInodesUsed uint64
}

type ContainerBandwidthStat struct {
	InRate   uint64
	InBurst  uint64
	OutRate  uint64
	OutBurst uint64
}

type ContainerNetworkStat struct {
	RxBytes uint64
	TxBytes uint64
}

type BandwidthLimits struct {
	RateInBytesPerSecond      uint64 `json:"rate,omitempty"`
	BurstRateInBytesPerSecond uint64 `json:"burst,omitempty"`
}

type ProcessLimits struct {
	CPU    CPULimits    `json:"cpu_limits,omitempty"`
	Memory MemoryLimits `json:"memory_limits,omitempty"`
}

type DiskLimits struct {
	InodeSoft uint64 `json:"inode_soft,omitempty"`
	InodeHard uint64 `json:"inode_hard,omitempty"`

	ByteSoft uint64 `json:"byte_soft,omitempty"`
	ByteHard uint64 `json:"byte_hard,omitempty"`

	Scope DiskLimitScope `json:"scope,omitempty"`
}

type MemoryLimits struct {
	//	Memory usage limit in bytes.
	LimitInBytes uint64 `json:"limit_in_bytes,omitempty"`
}

type CPULimits struct {
	LimitInShares uint64 `json:"limit_in_shares,omitempty"`
}

type PidLimits struct {
	// Limits the number of pids a container may create before new forks or clones are disallowed to processes in the container.
	// Note: this may only be enforced when a process attempts to fork, so it does not guarantee that a new container.Run(ProcessSpec)
	// will not succeed even if the limit has been exceeded, but the process will not be able to spawn further processes or threads.
	Max uint64 `json:"max,omitempty"`
}

// Resource limits.
//
// Please refer to the manual page of getrlimit for a description of the individual fields:
// http://www.kernel.org/doc/man-pages/online/pages/man2/getrlimit.2.html
type ResourceLimits struct {
	As         *uint64 `json:"as,omitempty"`
	Core       *uint64 `json:"core,omitempty"`
	Cpu        *uint64 `json:"cpu,omitempty"`
	Data       *uint64 `json:"data,omitempty"`
	Fsize      *uint64 `json:"fsize,omitempty"`
	Locks      *uint64 `json:"locks,omitempty"`
	Memlock    *uint64 `json:"memlock,omitempty"`
	Msgqueue   *uint64 `json:"msgqueue,omitempty"`
	Nice       *uint64 `json:"nice,omitempty"`
	Nofile     *uint64 `json:"nofile,omitempty"`
	Nproc      *uint64 `json:"nproc,omitempty"`
	Rss        *uint64 `json:"rss,omitempty"`
	Rtprio     *uint64 `json:"rtprio,omitempty"`
	Sigpending *uint64 `json:"sigpending,omitempty"`
	Stack      *uint64 `json:"stack,omitempty"`
}

type DiskLimitScope uint8

const DiskLimitScopeTotal DiskLimitScope = 0
const DiskLimitScopeExclusive DiskLimitScope = 1
<commit_msg>Add bindmounts to process spec<commit_after>package garden

import (
	"io"
	"time"
)

//go:generate counterfeiter . Container

type Container interface {
	Handle() string

	// Stop stops a container.
	//
	// If kill is false, garden stops a container by sending the processes running inside it the SIGTERM signal.
	// It then waits for the processes to terminate before returning a response.
	// If one or more processes do not terminate within 10 seconds,
	// garden sends these processes the SIGKILL signal, killing them ungracefully.
	//
	// If kill is true, garden stops a container by sending the processing running inside it a SIGKILL signal.
	//
	// It is possible to copy files in to and out of a stopped container.
	// It is only when a container is destroyed that its filesystem is cleaned up.
	//
	// Errors:
	// * None.
	Stop(kill bool) error

	// Returns information about a container.
	Info() (ContainerInfo, error)

	// StreamIn streams data into a file in a container.
	//
	// Errors:
	// *  TODO.
	StreamIn(spec StreamInSpec) error

	// StreamOut streams a file out of a container.
	//
	// Errors:
	// * TODO.
	StreamOut(spec StreamOutSpec) (io.ReadCloser, error)

	// Returns the current bandwidth limits set for the container.
	CurrentBandwidthLimits() (BandwidthLimits, error)

	// Returns the current CPU limts set for the container.
	CurrentCPULimits() (CPULimits, error)

	// Returns the current disk limts set for the container.
	CurrentDiskLimits() (DiskLimits, error)

	// Returns the current memory limts set for the container.
	CurrentMemoryLimits() (MemoryLimits, error)

	// Map a port on the host to a port in the container so that traffic to the
	// host port is forwarded to the container port. This is deprecated in
	// favour of passing NetIn configuration in the ContainerSpec at creation
	// time.
	//
	// If a host port is not given, a port will be acquired from the server's port
	// pool.
	//
	// If a container port is not given, the port will be the same as the
	// container port.
	//
	// The resulting host and container ports are returned in that order.
	//
	// Errors:
	// * When no port can be acquired from the server's port pool.
	NetIn(hostPort, containerPort uint32) (uint32, uint32, error)

	// Whitelist outbound network traffic. This is deprecated in favour of passing
	// NetOut configuration in the ContainerSpec at creation time.
	//
	// If the configuration directive deny_networks is not used,
	// all networks are already whitelisted and this command is effectively a no-op.
	//
	// Later NetOut calls take precedence over earlier calls, which is
	// significant only in relation to logging.
	//
	// Errors:
	// * An error is returned if the NetOut call fails.
	NetOut(netOutRule NetOutRule) error

	// A Bulk call for NetOut. This is deprecated in favour of passing
	// NetOut configuration in the ContainerSpec at creation time.
	//
	// Errors:
	// * An error is returned if any of the NetOut calls fail.
	BulkNetOut(netOutRules []NetOutRule) error

	// Run a script inside a container.
	//
	// The root user will be mapped to a non-root UID in the host unless the container (not this process) was created with 'privileged' true.
	//
	// Errors:
	// * TODO.
	Run(ProcessSpec, ProcessIO) (Process, error)

	// Attach starts streaming the output back to the client from a specified process.
	//
	// Errors:
	// * processID does not refer to a running process.
	Attach(processID string, io ProcessIO) (Process, error)

	// Metrics returns the current set of metrics for a container
	Metrics() (Metrics, error)

	// Sets the grace time.
	SetGraceTime(graceTime time.Duration) error

	// Properties returns the current set of properties
	Properties() (Properties, error)

	// Property returns the value of the property with the specified name.
	//
	// Errors:
	// * When the property does not exist on the container.
	Property(name string) (string, error)

	// Set a named property on a container to a specified value.
	//
	// Errors:
	// * None.
	SetProperty(name string, value string) error

	// Remove a property with the specified name from a container.
	//
	// Errors:
	// * None.
	RemoveProperty(name string) error
}

// ProcessSpec contains parameters for running a script inside a container.
type ProcessSpec struct {
	// ID for the process. If empty, an ID will be generated.
	ID string `json:"id,omitempty"`

	// Path to command to execute.
	Path string `json:"path,omitempty"`

	// Arguments to pass to command.
	Args []string `json:"args,omitempty"`

	// Environment variables.
	Env []string `json:"env,omitempty"`

	// Working directory (default: home directory).
	Dir string `json:"dir,omitempty"`

	// The name of a user in the container to run the process as.
	// This must either be a username, or uid:gid.
	User string `json:"user,omitempty"`

	// Resource limits
	Limits ResourceLimits `json:"rlimits,omitempty"`

	// Limits to be applied to the newly created process
	OverrideContainerLimits *ProcessLimits `json:"limits,omitempty"`

	// Execute with a TTY for stdio.
	TTY *TTYSpec `json:"tty,omitempty"`

	// Execute process in own root filesystem, different from the other processes
	// in the container.
	Image ImageRef `json:"image,omitempty"`

	// Bind mounts to be applied to the process's filesystem
	// An error is returned if ProcessSpec.Image is not also set.
	BindMounts []BindMount `json:"bind_mounts,omitempty"`
}

type TTYSpec struct {
	WindowSize *WindowSize `json:"window_size,omitempty"`
}

type WindowSize struct {
	Columns int `json:"columns,omitempty"`
	Rows    int `json:"rows,omitempty"`
}

type ProcessIO struct {
	Stdin  io.Reader
	Stdout io.Writer
	Stderr io.Writer
}

//go:generate counterfeiter . Process

type Process interface {
	ID() string
	Wait() (int, error)
	SetTTY(TTYSpec) error
	Signal(Signal) error
}

type Signal int

const (
	SignalTerminate Signal = iota
	SignalKill
)

type PortMapping struct {
	HostPort      uint32
	ContainerPort uint32
}

type StreamInSpec struct {
	Path      string
	User      string
	TarStream io.Reader
}

type StreamOutSpec struct {
	Path string
	User string
}

// ContainerInfo holds information about a container.
type ContainerInfo struct {
	State         string        // Either "active" or "stopped".
	Events        []string      // List of events that occurred for the container. It currently includes only "oom" (Out Of Memory) event if it occurred.
	HostIP        string        // The IP address of the gateway which controls the host side of the container's virtual ethernet pair.
	ContainerIP   string        // The IP address of the container side of the container's virtual ethernet pair.
	ExternalIP    string        //
	ContainerPath string        // The path to the directory holding the container's files (both its control scripts and filesystem).
	ProcessIDs    []string      // List of running processes.
	Properties    Properties    // List of properties defined for the container.
	MappedPorts   []PortMapping //
}

type ContainerInfoEntry struct {
	Info ContainerInfo
	Err  *Error
}

type Metrics struct {
	MemoryStat  ContainerMemoryStat
	CPUStat     ContainerCPUStat
	DiskStat    ContainerDiskStat
	NetworkStat ContainerNetworkStat
}

type ContainerMetricsEntry struct {
	Metrics Metrics
	Err     *Error
}

type ContainerMemoryStat struct {
	ActiveAnon              uint64 `json:"active_anon"`
	ActiveFile              uint64 `json:"active_file"`
	Cache                   uint64 `json:"cache"`
	HierarchicalMemoryLimit uint64 `json:"hierarchical_memory_limit"`
	InactiveAnon            uint64 `json:"inactive_anon"`
	InactiveFile            uint64 `json:"inactive_file"`
	MappedFile              uint64 `json:"mapped_file"`
	Pgfault                 uint64 `json:"pgfault"`
	Pgmajfault              uint64 `json:"pgmajfault"`
	Pgpgin                  uint64 `json:"pgpgin"`
	Pgpgout                 uint64 `json:"pgpgout"`
	Rss                     uint64 `json:"rss"`
	TotalActiveAnon         uint64 `json:"total_active_anon"`
	TotalActiveFile         uint64 `json:"total_active_file"`
	TotalCache              uint64 `json:"total_cache"`
	TotalInactiveAnon       uint64 `json:"total_inactive_anon"`
	TotalInactiveFile       uint64 `json:"total_inactive_file"`
	TotalMappedFile         uint64 `json:"total_mapped_file"`
	TotalPgfault            uint64 `json:"total_pgfault"`
	TotalPgmajfault         uint64 `json:"total_pgmajfault"`
	TotalPgpgin             uint64 `json:"total_pgpgin"`
	TotalPgpgout            uint64 `json:"total_pgpgout"`
	TotalRss                uint64 `json:"total_rss"`
	TotalUnevictable        uint64 `json:"total_unevictable"`
	Unevictable             uint64 `json:"unevictable"`
	Swap                    uint64 `json:"swap"`
	HierarchicalMemswLimit  uint64 `json:"hierarchical_memsw_limit"`
	TotalSwap               uint64 `json:"total_swap"`
	// A memory usage total which reports memory usage in the same way that limits are enforced.
	// This value includes memory consumed by nested containers.
	TotalUsageTowardLimit uint64
}

type ContainerCPUStat struct {
	Usage  uint64
	User   uint64
	System uint64
}

type ContainerDiskStat struct {
	TotalBytesUsed      uint64
	TotalInodesUsed     uint64
	ExclusiveBytesUsed  uint64
	ExclusiveInodesUsed uint64
}

type ContainerBandwidthStat struct {
	InRate   uint64
	InBurst  uint64
	OutRate  uint64
	OutBurst uint64
}

type ContainerNetworkStat struct {
	RxBytes uint64
	TxBytes uint64
}

type BandwidthLimits struct {
	RateInBytesPerSecond      uint64 `json:"rate,omitempty"`
	BurstRateInBytesPerSecond uint64 `json:"burst,omitempty"`
}

type ProcessLimits struct {
	CPU    CPULimits    `json:"cpu_limits,omitempty"`
	Memory MemoryLimits `json:"memory_limits,omitempty"`
}

type DiskLimits struct {
	InodeSoft uint64 `json:"inode_soft,omitempty"`
	InodeHard uint64 `json:"inode_hard,omitempty"`

	ByteSoft uint64 `json:"byte_soft,omitempty"`
	ByteHard uint64 `json:"byte_hard,omitempty"`

	Scope DiskLimitScope `json:"scope,omitempty"`
}

type MemoryLimits struct {
	//	Memory usage limit in bytes.
	LimitInBytes uint64 `json:"limit_in_bytes,omitempty"`
}

type CPULimits struct {
	LimitInShares uint64 `json:"limit_in_shares,omitempty"`
}

type PidLimits struct {
	// Limits the number of pids a container may create before new forks or clones are disallowed to processes in the container.
	// Note: this may only be enforced when a process attempts to fork, so it does not guarantee that a new container.Run(ProcessSpec)
	// will not succeed even if the limit has been exceeded, but the process will not be able to spawn further processes or threads.
	Max uint64 `json:"max,omitempty"`
}

// Resource limits.
//
// Please refer to the manual page of getrlimit for a description of the individual fields:
// http://www.kernel.org/doc/man-pages/online/pages/man2/getrlimit.2.html
type ResourceLimits struct {
	As         *uint64 `json:"as,omitempty"`
	Core       *uint64 `json:"core,omitempty"`
	Cpu        *uint64 `json:"cpu,omitempty"`
	Data       *uint64 `json:"data,omitempty"`
	Fsize      *uint64 `json:"fsize,omitempty"`
	Locks      *uint64 `json:"locks,omitempty"`
	Memlock    *uint64 `json:"memlock,omitempty"`
	Msgqueue   *uint64 `json:"msgqueue,omitempty"`
	Nice       *uint64 `json:"nice,omitempty"`
	Nofile     *uint64 `json:"nofile,omitempty"`
	Nproc      *uint64 `json:"nproc,omitempty"`
	Rss        *uint64 `json:"rss,omitempty"`
	Rtprio     *uint64 `json:"rtprio,omitempty"`
	Sigpending *uint64 `json:"sigpending,omitempty"`
	Stack      *uint64 `json:"stack,omitempty"`
}

type DiskLimitScope uint8

const DiskLimitScopeTotal DiskLimitScope = 0
const DiskLimitScopeExclusive DiskLimitScope = 1
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"context"
	"errors"
	"fmt"
	"log"
	"strings"

	"github.com/docker/docker/api/types"
	"github.com/docker/docker/api/types/container"
	"github.com/docker/docker/api/types/strslice"
	"github.com/docker/docker/client"
	"github.com/docker/go-connections/nat"
)

func CreateContainer(ctx context.Context, s serviceConfig, networkName, formattedImageName string) (string, error) {
	cli, err := client.NewEnvClient()
	if err != nil {
		return "", &popagateError{
			originalErr: err,
			newErr:      errors.New("unable to intialize docker client")}
	}
	defer cli.Close()

	// 2.1 make labels
	labelsMap := make(map[string]string)
	if len(s.Labels) > 0 {
		for _, v := range s.Labels {
			onelabel := fomatLabels(v)
			labelsMap[onelabel[0]] = onelabel[1]
		}
	}
	//2.2 make ports
	portsMap := make(map[nat.Port]struct{})
	portBindingMap := make(map[nat.Port][]nat.PortBinding)
	if len(s.Ports) > 0 {
		for _, v := range s.Ports {
			oneport := fomatPorts(v)
			hostport := oneport[0]
			containerport := oneport[1]
			port, err := nat.NewPort("tcp", containerport)
			myPortBinding := nat.PortBinding{HostPort: hostport}
			if err != nil {
				log.Println(err, "unable to create a nat.Port")
			}
			portsMap[port] = emptyStruct{}
			portBindingMap[port] = []nat.PortBinding{myPortBinding}
		}
	}
	//2.3 create command
	cmd := strslice.StrSlice{}
	if s.Command != "" {
		sliceCommand := strings.Fields(s.Command)
		cmd = strslice.StrSlice(sliceCommand)
	}
	//2.4 create restart policy
	restartPolicy := container.RestartPolicy{}
	if s.Restart != "" {
		// you cannot set MaximumRetryCount for the following restart policies;
		// always, no, unless-stopped
		if s.Restart == "on-failure" {
			restartPolicy = container.RestartPolicy{Name: s.Restart, MaximumRetryCount: 3}
		} else {
			restartPolicy = container.RestartPolicy{Name: s.Restart}
		}

	}
	//2.5 build image
	imageName := s.Image
	if s.Build != (buildstruct{}) {
		imageName, err = BuildDockerImage(ctx, s.Build.Dockerfile)
		if err != nil {
			return "", &popagateError{originalErr: err}
		}
	}

	//2.6 add volumes
	volume := make(map[string]struct{})
	binds := []string{}
	if len(s.Volumes) > 0 {
		vol := fomatServiceVolumes(s.Volumes[0])
		volume[vol[1]] = emptyStruct{}
		// TODO: handle other read/write modes
		whatToBind := "meli_" + vol[0] + ":" + vol[1] + ":rw"
		binds = append(binds, whatToBind)
	}

	// TODO: we should skip creating the container again if already exists
	// instead of creating a uniquely named container name
	containerCreateResp, err := cli.ContainerCreate(
		ctx,
		&container.Config{
			Image:        imageName,
			Labels:       labelsMap,
			Env:          s.Environment,
			ExposedPorts: portsMap,
			Cmd:          cmd,
			Volumes:      volume},
		&container.HostConfig{
			PublishAllPorts: false,
			PortBindings:    portBindingMap,
			NetworkMode:     container.NetworkMode(networkName),
			RestartPolicy:   restartPolicy,
			Binds:           binds},
		nil,
		formattedImageName)
	if err != nil {
		if err != nil {
			return "", &popagateError{
				originalErr: err,
				newErr:      errors.New("unable to create container")}
		}
	}

	return containerCreateResp.ID, nil
}

func ContainerStart(ctx context.Context, containerId string) error {
	cli, err := client.NewEnvClient()
	if err != nil {
		return &popagateError{
			originalErr: err,
			newErr:      errors.New("unable to intialize docker client")}
	}
	defer cli.Close()

	err = cli.ContainerStart(
		ctx,
		containerId,
		types.ContainerStartOptions{})
	if err != nil {
		return &popagateError{
			originalErr: err,
			newErr:      fmt.Errorf("unable to start container %s", containerId)}

	}
	return nil
}

func ContainerLogs(ctx context.Context, containerId string) error {
	cli, err := client.NewEnvClient()
	if err != nil {
		return &popagateError{
			originalErr: err,
			newErr:      errors.New("unable to intialize docker client")}
	}
	defer cli.Close()

	containerLogResp, err := cli.ContainerLogs(
		ctx,
		containerId,
		types.ContainerLogsOptions{
			ShowStdout: true,
			ShowStderr: true,
			Timestamps: true})
	if err != nil {
		if err != nil {
			return &popagateError{
				originalErr: err,
				newErr:      fmt.Errorf("unable to get container logs %s", containerId)}
		}
	}
	defer containerLogResp.Close()

	scanner := bufio.NewScanner(containerLogResp)
	for scanner.Scan() {
		output := strings.Replace(scanner.Text(), "u003e", ">", -1)
		log.Println(output)
	}
	err = scanner.Err()
	if err != nil {
		log.Println(err, "error in scanning")
	}
	return nil
}
<commit_msg>tail container logs<commit_after>package main

import (
	"bufio"
	"context"
	"errors"
	"fmt"
	"log"
	"strings"

	"github.com/docker/docker/api/types"
	"github.com/docker/docker/api/types/container"
	"github.com/docker/docker/api/types/strslice"
	"github.com/docker/docker/client"
	"github.com/docker/go-connections/nat"
)

func CreateContainer(ctx context.Context, s serviceConfig, networkName, formattedImageName string) (string, error) {
	cli, err := client.NewEnvClient()
	if err != nil {
		return "", &popagateError{
			originalErr: err,
			newErr:      errors.New("unable to intialize docker client")}
	}
	defer cli.Close()

	// 2.1 make labels
	labelsMap := make(map[string]string)
	if len(s.Labels) > 0 {
		for _, v := range s.Labels {
			onelabel := fomatLabels(v)
			labelsMap[onelabel[0]] = onelabel[1]
		}
	}
	//2.2 make ports
	portsMap := make(map[nat.Port]struct{})
	portBindingMap := make(map[nat.Port][]nat.PortBinding)
	if len(s.Ports) > 0 {
		for _, v := range s.Ports {
			oneport := fomatPorts(v)
			hostport := oneport[0]
			containerport := oneport[1]
			port, err := nat.NewPort("tcp", containerport)
			myPortBinding := nat.PortBinding{HostPort: hostport}
			if err != nil {
				log.Println(err, "unable to create a nat.Port")
			}
			portsMap[port] = emptyStruct{}
			portBindingMap[port] = []nat.PortBinding{myPortBinding}
		}
	}
	//2.3 create command
	cmd := strslice.StrSlice{}
	if s.Command != "" {
		sliceCommand := strings.Fields(s.Command)
		cmd = strslice.StrSlice(sliceCommand)
	}
	//2.4 create restart policy
	restartPolicy := container.RestartPolicy{}
	if s.Restart != "" {
		// you cannot set MaximumRetryCount for the following restart policies;
		// always, no, unless-stopped
		if s.Restart == "on-failure" {
			restartPolicy = container.RestartPolicy{Name: s.Restart, MaximumRetryCount: 3}
		} else {
			restartPolicy = container.RestartPolicy{Name: s.Restart}
		}

	}
	//2.5 build image
	imageName := s.Image
	if s.Build != (buildstruct{}) {
		imageName, err = BuildDockerImage(ctx, s.Build.Dockerfile)
		if err != nil {
			return "", &popagateError{originalErr: err}
		}
	}

	//2.6 add volumes
	volume := make(map[string]struct{})
	binds := []string{}
	if len(s.Volumes) > 0 {
		vol := fomatServiceVolumes(s.Volumes[0])
		volume[vol[1]] = emptyStruct{}
		// TODO: handle other read/write modes
		whatToBind := "meli_" + vol[0] + ":" + vol[1] + ":rw"
		binds = append(binds, whatToBind)
	}

	// TODO: we should skip creating the container again if already exists
	// instead of creating a uniquely named container name
	containerCreateResp, err := cli.ContainerCreate(
		ctx,
		&container.Config{
			Image:        imageName,
			Labels:       labelsMap,
			Env:          s.Environment,
			ExposedPorts: portsMap,
			Cmd:          cmd,
			Volumes:      volume},
		&container.HostConfig{
			PublishAllPorts: false,
			PortBindings:    portBindingMap,
			NetworkMode:     container.NetworkMode(networkName),
			RestartPolicy:   restartPolicy,
			Binds:           binds},
		nil,
		formattedImageName)
	if err != nil {
		if err != nil {
			return "", &popagateError{
				originalErr: err,
				newErr:      errors.New("unable to create container")}
		}
	}

	return containerCreateResp.ID, nil
}

func ContainerStart(ctx context.Context, containerId string) error {
	cli, err := client.NewEnvClient()
	if err != nil {
		return &popagateError{
			originalErr: err,
			newErr:      errors.New("unable to intialize docker client")}
	}
	defer cli.Close()

	err = cli.ContainerStart(
		ctx,
		containerId,
		types.ContainerStartOptions{})
	if err != nil {
		return &popagateError{
			originalErr: err,
			newErr:      fmt.Errorf("unable to start container %s", containerId)}

	}
	return nil
}

func ContainerLogs(ctx context.Context, containerId string) error {
	cli, err := client.NewEnvClient()
	if err != nil {
		return &popagateError{
			originalErr: err,
			newErr:      errors.New("unable to intialize docker client")}
	}
	defer cli.Close()

	containerLogResp, err := cli.ContainerLogs(
		ctx,
		containerId,
		types.ContainerLogsOptions{
			ShowStdout: true,
			ShowStderr: true,
			Timestamps: true,
			Follow:     true,
			Details:    true,
			Tail:       "all"})

	if err != nil {
		if err != nil {
			return &popagateError{
				originalErr: err,
				newErr:      fmt.Errorf("unable to get container logs %s", containerId)}
		}
	}
	defer containerLogResp.Close()

	scanner := bufio.NewScanner(containerLogResp)
	for scanner.Scan() {
		output := strings.Replace(scanner.Text(), "u003e", ">", -1)
		log.Println(output)
	}
	err = scanner.Err()
	if err != nil {
		log.Println(err, "error in scanning")
	}
	return nil
}
<|endoftext|>


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

package common_test

import (
	"fmt"
	"os"
	"time"

	"github.com/juju/loggo"
	"github.com/juju/testing"
	jc "github.com/juju/testing/checkers"
	gc "launchpad.net/gocheck"

	"github.com/juju/juju/constraints"
	"github.com/juju/juju/environs"
	"github.com/juju/juju/environs/cloudinit"
	"github.com/juju/juju/environs/config"
	"github.com/juju/juju/environs/storage"
	envtesting "github.com/juju/juju/environs/testing"
	"github.com/juju/juju/instance"
	"github.com/juju/juju/network"
	"github.com/juju/juju/provider/common"
	coretesting "github.com/juju/juju/testing"
	"github.com/juju/juju/tools"
	"github.com/juju/juju/utils/ssh"
)

type BootstrapSuite struct {
	coretesting.FakeJujuHomeSuite
	envtesting.ToolsFixture
}

var _ = gc.Suite(&BootstrapSuite{})

type cleaner interface {
	AddCleanup(testing.CleanupFunc)
}

func (s *BootstrapSuite) SetUpTest(c *gc.C) {
	s.FakeJujuHomeSuite.SetUpTest(c)
	s.ToolsFixture.SetUpTest(c)
	s.PatchValue(common.ConnectSSH, func(_ ssh.Client, host, checkHostScript string) error {
		return fmt.Errorf("mock connection failure to %s", host)
	})
}

func (s *BootstrapSuite) TearDownTest(c *gc.C) {
	s.ToolsFixture.TearDownTest(c)
	s.FakeJujuHomeSuite.TearDownTest(c)
}

func newStorage(suite cleaner, c *gc.C) storage.Storage {
	closer, stor, _ := envtesting.CreateLocalTestStorage(c)
	suite.AddCleanup(func(*gc.C) { closer.Close() })
	envtesting.UploadFakeTools(c, stor)
	return stor
}

func minimalConfig(c *gc.C) *config.Config {
	attrs := map[string]interface{}{
		"name":            "whatever",
		"type":            "anything, really",
		"ca-cert":         coretesting.CACert,
		"ca-private-key":  coretesting.CAKey,
		"authorized-keys": coretesting.FakeAuthKeys,
	}
	cfg, err := config.New(config.UseDefaults, attrs)
	c.Assert(err, gc.IsNil)
	return cfg
}

func configGetter(c *gc.C) configFunc {
	cfg := minimalConfig(c)
	return func() *config.Config { return cfg }
}

func (s *BootstrapSuite) TestCannotStartInstance(c *gc.C) {
	checkPlacement := "directive"
	checkCons := constraints.MustParse("mem=8G")

	startInstance := func(
		placement string, cons constraints.Value, _ []string, possibleTools tools.List, mcfg *cloudinit.MachineConfig,
	) (
		instance.Instance, *instance.HardwareCharacteristics, []network.Info, error,
	) {
		c.Assert(placement, gc.DeepEquals, checkPlacement)
		c.Assert(cons, gc.DeepEquals, checkCons)
		c.Assert(mcfg, gc.DeepEquals, environs.NewBootstrapMachineConfig(cons, mcfg.SystemPrivateSSHKey))
		return nil, nil, nil, fmt.Errorf("meh, not started")
	}

	env := &mockEnviron{
		storage:       newStorage(s, c),
		startInstance: startInstance,
		config:        configGetter(c),
	}

	ctx := coretesting.Context(c)
	_, _, _, err := common.Bootstrap(ctx, env, environs.BootstrapParams{
		Constraints: checkCons,
		Placement:   checkPlacement,
	})
	c.Assert(err, gc.ErrorMatches, "cannot start bootstrap instance: meh, not started")
}

func (s *BootstrapSuite) TestCannotRecordStartedInstance(c *gc.C) {
	innerStorage := newStorage(s, c)
	stor := &mockStorage{Storage: innerStorage}

	startInstance := func(
		_ string, _ constraints.Value, _ []string, _ tools.List, _ *cloudinit.MachineConfig,
	) (
		instance.Instance, *instance.HardwareCharacteristics, []network.Info, error,
	) {
		stor.putErr = fmt.Errorf("suddenly a wild blah")
		return &mockInstance{id: "i-blah"}, nil, nil, nil
	}

	var stopped []instance.Id
	stopInstances := func(ids []instance.Id) error {
		stopped = append(stopped, ids...)
		return nil
	}

	env := &mockEnviron{
		storage:       stor,
		startInstance: startInstance,
		stopInstances: stopInstances,
		config:        configGetter(c),
	}

	ctx := coretesting.Context(c)
	_, _, _, err := common.Bootstrap(ctx, env, environs.BootstrapParams{})
	c.Assert(err, gc.ErrorMatches, "cannot save state: suddenly a wild blah")
	c.Assert(stopped, gc.HasLen, 1)
	c.Assert(stopped[0], gc.Equals, instance.Id("i-blah"))
}

func (s *BootstrapSuite) TestCannotRecordThenCannotStop(c *gc.C) {
	innerStorage := newStorage(s, c)
	stor := &mockStorage{Storage: innerStorage}

	startInstance := func(
		_ string, _ constraints.Value, _ []string, _ tools.List, _ *cloudinit.MachineConfig,
	) (
		instance.Instance, *instance.HardwareCharacteristics, []network.Info, error,
	) {
		stor.putErr = fmt.Errorf("suddenly a wild blah")
		return &mockInstance{id: "i-blah"}, nil, nil, nil
	}

	var stopped []instance.Id
	stopInstances := func(instances []instance.Id) error {
		stopped = append(stopped, instances...)
		return fmt.Errorf("bork bork borken")
	}

	var tw loggo.TestWriter
	c.Assert(loggo.RegisterWriter("bootstrap-tester", &tw, loggo.DEBUG), gc.IsNil)
	defer loggo.RemoveWriter("bootstrap-tester")

	env := &mockEnviron{
		storage:       stor,
		startInstance: startInstance,
		stopInstances: stopInstances,
		config:        configGetter(c),
	}

	ctx := coretesting.Context(c)
	_, _, _, err := common.Bootstrap(ctx, env, environs.BootstrapParams{})
	c.Assert(err, gc.ErrorMatches, "cannot save state: suddenly a wild blah")
	c.Assert(stopped, gc.HasLen, 1)
	c.Assert(stopped[0], gc.Equals, instance.Id("i-blah"))
	c.Assert(tw.Log(), jc.LogMatches, []jc.SimpleMessage{{
		loggo.ERROR, `cannot stop failed bootstrap instance "i-blah": bork bork borken`,
	}})
}

func (s *BootstrapSuite) TestSuccess(c *gc.C) {
	stor := newStorage(s, c)
	checkInstanceId := "i-success"
	checkHardware := instance.MustParseHardware("arch=ppc64el mem=2T")

	startInstance := func(
		_ string, _ constraints.Value, _ []string, _ tools.List, mcfg *cloudinit.MachineConfig,
	) (
		instance.Instance, *instance.HardwareCharacteristics, []network.Info, error,
	) {
		return &mockInstance{id: checkInstanceId}, &checkHardware, nil, nil
	}
	var mocksConfig = minimalConfig(c)
	var getConfigCalled int
	getConfig := func() *config.Config {
		getConfigCalled++
		return mocksConfig
	}
	setConfig := func(c *config.Config) error {
		mocksConfig = c
		return nil
	}

	env := &mockEnviron{
		storage:       stor,
		startInstance: startInstance,
		config:        getConfig,
		setConfig:     setConfig,
	}
	ctx := coretesting.Context(c)
	arch, series, _, err := common.Bootstrap(ctx, env, environs.BootstrapParams{})
	c.Assert(err, gc.IsNil)
	c.Assert(arch, gc.Equals, "ppc64el") // based on hardware characteristics
	c.Assert(series, gc.Equals, "trusty")
}

type neverRefreshes struct {
}

func (neverRefreshes) Refresh() error {
	return nil
}

type neverAddresses struct {
	neverRefreshes
}

func (neverAddresses) Addresses() ([]network.Address, error) {
	return nil, nil
}

var testSSHTimeout = config.SSHTimeoutOpts{
	Timeout:        coretesting.ShortWait,
	RetryDelay:     1 * time.Millisecond,
	AddressesDelay: 1 * time.Millisecond,
}

func (s *BootstrapSuite) TestWaitSSHTimesOutWaitingForAddresses(c *gc.C) {
	ctx := coretesting.Context(c)
	_, err := common.WaitSSH(ctx, nil, ssh.DefaultClient, "/bin/true", neverAddresses{}, testSSHTimeout)
	c.Check(err, gc.ErrorMatches, `waited for `+testSSHTimeout.Timeout.String()+` without getting any addresses`)
	c.Check(coretesting.Stderr(ctx), gc.Matches, "Waiting for address\n")
}

func (s *BootstrapSuite) TestWaitSSHKilledWaitingForAddresses(c *gc.C) {
	ctx := coretesting.Context(c)
	interrupted := make(chan os.Signal, 1)
	interrupted <- os.Interrupt
	_, err := common.WaitSSH(ctx, interrupted, ssh.DefaultClient, "/bin/true", neverAddresses{}, testSSHTimeout)
	c.Check(err, gc.ErrorMatches, "interrupted")
	c.Check(coretesting.Stderr(ctx), gc.Matches, "Waiting for address\n")
}

type brokenAddresses struct {
	neverRefreshes
}

func (brokenAddresses) Addresses() ([]network.Address, error) {
	return nil, fmt.Errorf("Addresses will never work")
}

func (s *BootstrapSuite) TestWaitSSHStopsOnBadError(c *gc.C) {
	ctx := coretesting.Context(c)
	_, err := common.WaitSSH(ctx, nil, ssh.DefaultClient, "/bin/true", brokenAddresses{}, testSSHTimeout)
	c.Check(err, gc.ErrorMatches, "getting addresses: Addresses will never work")
	c.Check(coretesting.Stderr(ctx), gc.Equals, "Waiting for address\n")
}

type neverOpensPort struct {
	neverRefreshes
	addr string
}

func (n *neverOpensPort) Addresses() ([]network.Address, error) {
	return network.NewAddresses(n.addr), nil
}

func (s *BootstrapSuite) TestWaitSSHTimesOutWaitingForDial(c *gc.C) {
	ctx := coretesting.Context(c)
	// 0.x.y.z addresses are always invalid
	_, err := common.WaitSSH(ctx, nil, ssh.DefaultClient, "/bin/true", &neverOpensPort{addr: "0.1.2.3"}, testSSHTimeout)
	c.Check(err, gc.ErrorMatches,
		`waited for `+testSSHTimeout.Timeout.String()+` without being able to connect: mock connection failure to 0.1.2.3`)
	c.Check(coretesting.Stderr(ctx), gc.Matches,
		"Waiting for address\n"+
			"(Attempting to connect to 0.1.2.3:22\n)+")
}

type interruptOnDial struct {
	neverRefreshes
	name        string
	interrupted chan os.Signal
	returned    bool
}

func (i *interruptOnDial) Addresses() ([]network.Address, error) {
	// kill the tomb the second time Addresses is called
	if !i.returned {
		i.returned = true
	} else {
		i.interrupted <- os.Interrupt
	}
	return []network.Address{network.NewAddress(i.name, network.ScopeUnknown)}, nil
}

func (s *BootstrapSuite) TestWaitSSHKilledWaitingForDial(c *gc.C) {
	ctx := coretesting.Context(c)
	timeout := testSSHTimeout
	timeout.Timeout = 1 * time.Minute
	interrupted := make(chan os.Signal, 1)
	_, err := common.WaitSSH(ctx, interrupted, ssh.DefaultClient, "", &interruptOnDial{name: "0.1.2.3", interrupted: interrupted}, timeout)
	c.Check(err, gc.ErrorMatches, "interrupted")
	// Exact timing is imprecise but it should have tried a few times before being killed
	c.Check(coretesting.Stderr(ctx), gc.Matches,
		"Waiting for address\n"+
			"(Attempting to connect to 0.1.2.3:22\n)+")
}

type addressesChange struct {
	addrs [][]string
}

func (ac *addressesChange) Refresh() error {
	if len(ac.addrs) > 1 {
		ac.addrs = ac.addrs[1:]
	}
	return nil
}

func (ac *addressesChange) Addresses() ([]network.Address, error) {
	var addrs []network.Address
	for _, addr := range ac.addrs[0] {
		addrs = append(addrs, network.NewAddress(addr, network.ScopeUnknown))
	}
	return addrs, nil
}

func (s *BootstrapSuite) TestWaitSSHRefreshAddresses(c *gc.C) {
	ctx := coretesting.Context(c)
	_, err := common.WaitSSH(ctx, nil, ssh.DefaultClient, "", &addressesChange{addrs: [][]string{
		nil,
		nil,
		[]string{"0.1.2.3"},
		[]string{"0.1.2.3"},
		nil,
		[]string{"0.1.2.4"},
	}}, testSSHTimeout)
	// Not necessarily the last one in the list, due to scheduling.
	c.Check(err, gc.ErrorMatches,
		`waited for `+testSSHTimeout.Timeout.String()+` without being able to connect: mock connection failure to 0.1.2.[34]`)
	stderr := coretesting.Stderr(ctx)
	c.Check(stderr, gc.Matches,
		"Waiting for address\n"+
			"(.|\n)*(Attempting to connect to 0.1.2.3:22\n)+(.|\n)*")
	c.Check(stderr, gc.Matches,
		"Waiting for address\n"+
			"(.|\n)*(Attempting to connect to 0.1.2.4:22\n)+(.|\n)*")
}
<commit_msg>provider/common: fix test<commit_after>// Copyright 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package common_test

import (
	"fmt"
	"os"
	"time"

	"github.com/juju/loggo"
	"github.com/juju/testing"
	jc "github.com/juju/testing/checkers"
	gc "launchpad.net/gocheck"

	"github.com/juju/juju/constraints"
	"github.com/juju/juju/environs"
	"github.com/juju/juju/environs/cloudinit"
	"github.com/juju/juju/environs/config"
	"github.com/juju/juju/environs/storage"
	envtesting "github.com/juju/juju/environs/testing"
	"github.com/juju/juju/instance"
	"github.com/juju/juju/network"
	"github.com/juju/juju/provider/common"
	coretesting "github.com/juju/juju/testing"
	"github.com/juju/juju/tools"
	"github.com/juju/juju/utils/ssh"
	"github.com/juju/juju/version"
)

type BootstrapSuite struct {
	coretesting.FakeJujuHomeSuite
	envtesting.ToolsFixture
}

var _ = gc.Suite(&BootstrapSuite{})

type cleaner interface {
	AddCleanup(testing.CleanupFunc)
}

func (s *BootstrapSuite) SetUpTest(c *gc.C) {
	s.FakeJujuHomeSuite.SetUpTest(c)
	s.ToolsFixture.SetUpTest(c)
	s.PatchValue(common.ConnectSSH, func(_ ssh.Client, host, checkHostScript string) error {
		return fmt.Errorf("mock connection failure to %s", host)
	})
}

func (s *BootstrapSuite) TearDownTest(c *gc.C) {
	s.ToolsFixture.TearDownTest(c)
	s.FakeJujuHomeSuite.TearDownTest(c)
}

func newStorage(suite cleaner, c *gc.C) storage.Storage {
	closer, stor, _ := envtesting.CreateLocalTestStorage(c)
	suite.AddCleanup(func(*gc.C) { closer.Close() })
	envtesting.UploadFakeTools(c, stor)
	return stor
}

func minimalConfig(c *gc.C) *config.Config {
	attrs := map[string]interface{}{
		"name":            "whatever",
		"type":            "anything, really",
		"ca-cert":         coretesting.CACert,
		"ca-private-key":  coretesting.CAKey,
		"authorized-keys": coretesting.FakeAuthKeys,
	}
	cfg, err := config.New(config.UseDefaults, attrs)
	c.Assert(err, gc.IsNil)
	return cfg
}

func configGetter(c *gc.C) configFunc {
	cfg := minimalConfig(c)
	return func() *config.Config { return cfg }
}

func (s *BootstrapSuite) TestCannotStartInstance(c *gc.C) {
	checkPlacement := "directive"
	checkCons := constraints.MustParse("mem=8G")

	startInstance := func(
		placement string, cons constraints.Value, _ []string, possibleTools tools.List, mcfg *cloudinit.MachineConfig,
	) (
		instance.Instance, *instance.HardwareCharacteristics, []network.Info, error,
	) {
		c.Assert(placement, gc.DeepEquals, checkPlacement)
		c.Assert(cons, gc.DeepEquals, checkCons)
		c.Assert(mcfg, gc.DeepEquals, environs.NewBootstrapMachineConfig(cons, mcfg.SystemPrivateSSHKey))
		return nil, nil, nil, fmt.Errorf("meh, not started")
	}

	env := &mockEnviron{
		storage:       newStorage(s, c),
		startInstance: startInstance,
		config:        configGetter(c),
	}

	ctx := coretesting.Context(c)
	_, _, _, err := common.Bootstrap(ctx, env, environs.BootstrapParams{
		Constraints: checkCons,
		Placement:   checkPlacement,
	})
	c.Assert(err, gc.ErrorMatches, "cannot start bootstrap instance: meh, not started")
}

func (s *BootstrapSuite) TestCannotRecordStartedInstance(c *gc.C) {
	innerStorage := newStorage(s, c)
	stor := &mockStorage{Storage: innerStorage}

	startInstance := func(
		_ string, _ constraints.Value, _ []string, _ tools.List, _ *cloudinit.MachineConfig,
	) (
		instance.Instance, *instance.HardwareCharacteristics, []network.Info, error,
	) {
		stor.putErr = fmt.Errorf("suddenly a wild blah")
		return &mockInstance{id: "i-blah"}, nil, nil, nil
	}

	var stopped []instance.Id
	stopInstances := func(ids []instance.Id) error {
		stopped = append(stopped, ids...)
		return nil
	}

	env := &mockEnviron{
		storage:       stor,
		startInstance: startInstance,
		stopInstances: stopInstances,
		config:        configGetter(c),
	}

	ctx := coretesting.Context(c)
	_, _, _, err := common.Bootstrap(ctx, env, environs.BootstrapParams{})
	c.Assert(err, gc.ErrorMatches, "cannot save state: suddenly a wild blah")
	c.Assert(stopped, gc.HasLen, 1)
	c.Assert(stopped[0], gc.Equals, instance.Id("i-blah"))
}

func (s *BootstrapSuite) TestCannotRecordThenCannotStop(c *gc.C) {
	innerStorage := newStorage(s, c)
	stor := &mockStorage{Storage: innerStorage}

	startInstance := func(
		_ string, _ constraints.Value, _ []string, _ tools.List, _ *cloudinit.MachineConfig,
	) (
		instance.Instance, *instance.HardwareCharacteristics, []network.Info, error,
	) {
		stor.putErr = fmt.Errorf("suddenly a wild blah")
		return &mockInstance{id: "i-blah"}, nil, nil, nil
	}

	var stopped []instance.Id
	stopInstances := func(instances []instance.Id) error {
		stopped = append(stopped, instances...)
		return fmt.Errorf("bork bork borken")
	}

	var tw loggo.TestWriter
	c.Assert(loggo.RegisterWriter("bootstrap-tester", &tw, loggo.DEBUG), gc.IsNil)
	defer loggo.RemoveWriter("bootstrap-tester")

	env := &mockEnviron{
		storage:       stor,
		startInstance: startInstance,
		stopInstances: stopInstances,
		config:        configGetter(c),
	}

	ctx := coretesting.Context(c)
	_, _, _, err := common.Bootstrap(ctx, env, environs.BootstrapParams{})
	c.Assert(err, gc.ErrorMatches, "cannot save state: suddenly a wild blah")
	c.Assert(stopped, gc.HasLen, 1)
	c.Assert(stopped[0], gc.Equals, instance.Id("i-blah"))
	c.Assert(tw.Log(), jc.LogMatches, []jc.SimpleMessage{{
		loggo.ERROR, `cannot stop failed bootstrap instance "i-blah": bork bork borken`,
	}})
}

func (s *BootstrapSuite) TestSuccess(c *gc.C) {
	stor := newStorage(s, c)
	checkInstanceId := "i-success"
	checkHardware := instance.MustParseHardware("arch=ppc64el mem=2T")

	startInstance := func(
		_ string, _ constraints.Value, _ []string, _ tools.List, mcfg *cloudinit.MachineConfig,
	) (
		instance.Instance, *instance.HardwareCharacteristics, []network.Info, error,
	) {
		return &mockInstance{id: checkInstanceId}, &checkHardware, nil, nil
	}
	var mocksConfig = minimalConfig(c)
	var getConfigCalled int
	getConfig := func() *config.Config {
		getConfigCalled++
		return mocksConfig
	}
	setConfig := func(c *config.Config) error {
		mocksConfig = c
		return nil
	}

	env := &mockEnviron{
		storage:       stor,
		startInstance: startInstance,
		config:        getConfig,
		setConfig:     setConfig,
	}
	ctx := coretesting.Context(c)
	arch, series, _, err := common.Bootstrap(ctx, env, environs.BootstrapParams{})
	c.Assert(err, gc.IsNil)
	c.Assert(arch, gc.Equals, "ppc64el") // based on hardware characteristics
	c.Assert(series, gc.Equals, version.Current.Series)
}

type neverRefreshes struct {
}

func (neverRefreshes) Refresh() error {
	return nil
}

type neverAddresses struct {
	neverRefreshes
}

func (neverAddresses) Addresses() ([]network.Address, error) {
	return nil, nil
}

var testSSHTimeout = config.SSHTimeoutOpts{
	Timeout:        coretesting.ShortWait,
	RetryDelay:     1 * time.Millisecond,
	AddressesDelay: 1 * time.Millisecond,
}

func (s *BootstrapSuite) TestWaitSSHTimesOutWaitingForAddresses(c *gc.C) {
	ctx := coretesting.Context(c)
	_, err := common.WaitSSH(ctx, nil, ssh.DefaultClient, "/bin/true", neverAddresses{}, testSSHTimeout)
	c.Check(err, gc.ErrorMatches, `waited for `+testSSHTimeout.Timeout.String()+` without getting any addresses`)
	c.Check(coretesting.Stderr(ctx), gc.Matches, "Waiting for address\n")
}

func (s *BootstrapSuite) TestWaitSSHKilledWaitingForAddresses(c *gc.C) {
	ctx := coretesting.Context(c)
	interrupted := make(chan os.Signal, 1)
	interrupted <- os.Interrupt
	_, err := common.WaitSSH(ctx, interrupted, ssh.DefaultClient, "/bin/true", neverAddresses{}, testSSHTimeout)
	c.Check(err, gc.ErrorMatches, "interrupted")
	c.Check(coretesting.Stderr(ctx), gc.Matches, "Waiting for address\n")
}

type brokenAddresses struct {
	neverRefreshes
}

func (brokenAddresses) Addresses() ([]network.Address, error) {
	return nil, fmt.Errorf("Addresses will never work")
}

func (s *BootstrapSuite) TestWaitSSHStopsOnBadError(c *gc.C) {
	ctx := coretesting.Context(c)
	_, err := common.WaitSSH(ctx, nil, ssh.DefaultClient, "/bin/true", brokenAddresses{}, testSSHTimeout)
	c.Check(err, gc.ErrorMatches, "getting addresses: Addresses will never work")
	c.Check(coretesting.Stderr(ctx), gc.Equals, "Waiting for address\n")
}

type neverOpensPort struct {
	neverRefreshes
	addr string
}

func (n *neverOpensPort) Addresses() ([]network.Address, error) {
	return network.NewAddresses(n.addr), nil
}

func (s *BootstrapSuite) TestWaitSSHTimesOutWaitingForDial(c *gc.C) {
	ctx := coretesting.Context(c)
	// 0.x.y.z addresses are always invalid
	_, err := common.WaitSSH(ctx, nil, ssh.DefaultClient, "/bin/true", &neverOpensPort{addr: "0.1.2.3"}, testSSHTimeout)
	c.Check(err, gc.ErrorMatches,
		`waited for `+testSSHTimeout.Timeout.String()+` without being able to connect: mock connection failure to 0.1.2.3`)
	c.Check(coretesting.Stderr(ctx), gc.Matches,
		"Waiting for address\n"+
			"(Attempting to connect to 0.1.2.3:22\n)+")
}

type interruptOnDial struct {
	neverRefreshes
	name        string
	interrupted chan os.Signal
	returned    bool
}

func (i *interruptOnDial) Addresses() ([]network.Address, error) {
	// kill the tomb the second time Addresses is called
	if !i.returned {
		i.returned = true
	} else {
		i.interrupted <- os.Interrupt
	}
	return []network.Address{network.NewAddress(i.name, network.ScopeUnknown)}, nil
}

func (s *BootstrapSuite) TestWaitSSHKilledWaitingForDial(c *gc.C) {
	ctx := coretesting.Context(c)
	timeout := testSSHTimeout
	timeout.Timeout = 1 * time.Minute
	interrupted := make(chan os.Signal, 1)
	_, err := common.WaitSSH(ctx, interrupted, ssh.DefaultClient, "", &interruptOnDial{name: "0.1.2.3", interrupted: interrupted}, timeout)
	c.Check(err, gc.ErrorMatches, "interrupted")
	// Exact timing is imprecise but it should have tried a few times before being killed
	c.Check(coretesting.Stderr(ctx), gc.Matches,
		"Waiting for address\n"+
			"(Attempting to connect to 0.1.2.3:22\n)+")
}

type addressesChange struct {
	addrs [][]string
}

func (ac *addressesChange) Refresh() error {
	if len(ac.addrs) > 1 {
		ac.addrs = ac.addrs[1:]
	}
	return nil
}

func (ac *addressesChange) Addresses() ([]network.Address, error) {
	var addrs []network.Address
	for _, addr := range ac.addrs[0] {
		addrs = append(addrs, network.NewAddress(addr, network.ScopeUnknown))
	}
	return addrs, nil
}

func (s *BootstrapSuite) TestWaitSSHRefreshAddresses(c *gc.C) {
	ctx := coretesting.Context(c)
	_, err := common.WaitSSH(ctx, nil, ssh.DefaultClient, "", &addressesChange{addrs: [][]string{
		nil,
		nil,
		[]string{"0.1.2.3"},
		[]string{"0.1.2.3"},
		nil,
		[]string{"0.1.2.4"},
	}}, testSSHTimeout)
	// Not necessarily the last one in the list, due to scheduling.
	c.Check(err, gc.ErrorMatches,
		`waited for `+testSSHTimeout.Timeout.String()+` without being able to connect: mock connection failure to 0.1.2.[34]`)
	stderr := coretesting.Stderr(ctx)
	c.Check(stderr, gc.Matches,
		"Waiting for address\n"+
			"(.|\n)*(Attempting to connect to 0.1.2.3:22\n)+(.|\n)*")
	c.Check(stderr, gc.Matches,
		"Waiting for address\n"+
			"(.|\n)*(Attempting to connect to 0.1.2.4:22\n)+(.|\n)*")
}
<|endoftext|>