<commit_before>// Copyright 2009,2010, the 'gomongo' Authors.  All rights reserved.
// Use of this source code is governed by the New BSD License
// that can be found in the LICENSE file.

package mongo

import (
	"bytes"
	"container/vector"
	"encoding/binary"
	"io"
	"io/ioutil"
	"os"
)

const (
	_OP_REPLY        = 1
	_OP_MSG          = 1000
	_OP_UPDATE       = 2001
	_OP_INSERT       = 2002
	_OP_GET_BY_OID   = 2003
	_OP_QUERY        = 2004
	_OP_GET_MORE     = 2005
	_OP_DELETE       = 2006
	_OP_KILL_CURSORS = 2007
)

var last_req int32

type message interface {
	Bytes() []byte
	RequestID() int32
	OpCode() int32
}

func (c *Connection) writeMessage(m message) os.Error {
	body := m.Bytes()
	hb := header(int32(len(body)+16), m.RequestID(), 0, m.OpCode())
	msg := bytes.Add(hb, body)

	_, err := c.conn.Write(msg)

	last_req = m.RequestID()
	return err
}

func header(length, reqID, respTo, opCode int32) []byte {
	b := make([]byte, 16)
	binary.LittleEndian.PutUint32(b[0:4], uint32(length))
	binary.LittleEndian.PutUint32(b[4:8], uint32(reqID))
	binary.LittleEndian.PutUint32(b[8:12], uint32(respTo))
	binary.LittleEndian.PutUint32(b[12:16], uint32(opCode))
	return b
}

type replyMsg struct {
	responseTo     int32
	responseFlag   int32
	cursorID       int64
	startingFrom   int32
	numberReturned int32
	docs           *vector.Vector
}

func (c *Connection) readReply() (*replyMsg, os.Error) {
	size_bits, _ := ioutil.ReadAll(io.LimitReader(c.conn, 4))
	size := binary.LittleEndian.Uint32(size_bits)
	rest, _ := ioutil.ReadAll(io.LimitReader(c.conn, int64(size)-4))
	reply := parseReply(rest)
	return reply, nil
}

func parseReply(b []byte) *replyMsg {
	r := new(replyMsg)
	r.responseTo = int32(binary.LittleEndian.Uint32(b[4:8]))
	r.responseFlag = int32(binary.LittleEndian.Uint32(b[12:16]))
	r.cursorID = int64(binary.LittleEndian.Uint64(b[16:24]))
	r.startingFrom = int32(binary.LittleEndian.Uint32(b[24:28]))
	r.numberReturned = int32(binary.LittleEndian.Uint32(b[28:32]))
	r.docs = new(vector.Vector)

	if r.numberReturned > 0 {
		buf := bytes.NewBuffer(b[36:len(b)])
		for i := 0; int32(i) < r.numberReturned; i++ {
			var bson BSON
			bb := new(_BSONBuilder)
			bb.ptr = &bson
			bb.Object()
			Parse(buf, bb)
			r.docs.Push(bson)
			ioutil.ReadAll(io.LimitReader(buf, 4))
		}
	}

	return r
}

type deleteMsg struct {
	fullCollectionName string
	selector           BSON
	requestID          int32
}

func (d *deleteMsg) OpCode() int32    { return _OP_DELETE }
func (d *deleteMsg) RequestID() int32 { return d.requestID }
func (d *deleteMsg) Bytes() []byte {
	zero := make([]byte, 4)
	buf := bytes.NewBuffer(zero)
	buf.WriteString(d.fullCollectionName)
	buf.WriteByte(0)
	buf.Write(zero)
	buf.Write(d.selector.Bytes())
	return buf.Bytes()

}

type getMoreMsg struct {
	fullCollectionName string
	numberToReturn     int32
	cursorID           int64
	requestID          int32
}

func (g *getMoreMsg) OpCode() int32    { return _OP_GET_MORE }
func (g *getMoreMsg) RequestID() int32 { return g.requestID }
func (g *getMoreMsg) Bytes() []byte {
	buf := bytes.NewBuffer(make([]byte, 4))
	buf.WriteString(g.fullCollectionName)
	buf.WriteByte(0)

	b := make([]byte, 4)
	binary.LittleEndian.PutUint32(b, uint32(g.numberToReturn))
	buf.Write(b)

	b = make([]byte, 8)
	binary.LittleEndian.PutUint64(b, uint64(g.cursorID))
	buf.Write(b)

	return buf.Bytes()
}

type insertMsg struct {
	fullCollectionName string
	doc                BSON
	requestID          int32
}

func (i *insertMsg) OpCode() int32    { return _OP_INSERT }
func (i *insertMsg) RequestID() int32 { return i.requestID }
func (i *insertMsg) Bytes() []byte {
	buf := bytes.NewBuffer(make([]byte, 4))
	buf.WriteString(i.fullCollectionName)
	buf.WriteByte(0)
	buf.Write(i.doc.Bytes())
	return buf.Bytes()
}

type killMsg struct {
	numberOfCursorIDs int32
	cursorIDs         []int64
	requestID         int32
}

func (k *killMsg) OpCode() int32    { return _OP_KILL_CURSORS }
func (k *killMsg) RequestID() int32 { return k.requestID }
func (k *killMsg) Bytes() []byte {
	buf := bytes.NewBuffer(make([]byte, 4))

	b := make([]byte, 4)
	binary.LittleEndian.PutUint32(b, uint32(k.numberOfCursorIDs))
	buf.Write(b)

	b = make([]byte, 8)
	for _, id := range k.cursorIDs {
		binary.LittleEndian.PutUint64(b, uint64(id))
		buf.Write(b)
	}

	return buf.Bytes()
}

type queryMsg struct {
	opts               int32
	fullCollectionName string
	numberToSkip       int32
	numberToReturn     int32
	query              BSON
	requestID          int32
}

func (q *queryMsg) OpCode() int32    { return _OP_QUERY }
func (q *queryMsg) RequestID() int32 { return q.requestID }
func (q *queryMsg) Bytes() []byte {
	b := make([]byte, 4)
	binary.LittleEndian.PutUint32(b, uint32(q.opts))

	buf := bytes.NewBuffer(b)
	buf.WriteString(q.fullCollectionName)
	buf.WriteByte(0)

	binary.LittleEndian.PutUint32(b, uint32(q.numberToSkip))
	buf.Write(b)

	binary.LittleEndian.PutUint32(b, uint32(q.numberToReturn))
	buf.Write(b)

	buf.Write(q.query.Bytes())
	return buf.Bytes()
}

type updateMsg struct {
	fullCollectionName string
	flags              int32
	selector, document BSON
	requestID          int32
}

func (u *updateMsg) OpCode() int32    { return _OP_UPDATE }
func (u *updateMsg) RequestID() int32 { return u.requestID }
func (u *updateMsg) Bytes() []byte {
	buf := bytes.NewBuffer(make([]byte, 4))
	buf.WriteString(u.fullCollectionName)
	buf.WriteByte(0)

	b := make([]byte, 4)
	binary.LittleEndian.PutUint32(b, uint32(u.flags))
	buf.Write(b)

	buf.Write(u.selector.Bytes())
	buf.Write(u.document.Bytes())

	return buf.Bytes()
}

<commit_msg>adding comments<commit_after>// Copyright 2009,2010, the 'gomongo' Authors.  All rights reserved.
// Use of this source code is governed by the New BSD License
// that can be found in the LICENSE file.

/* Mongo Wire Protocol

http://www.mongodb.org/display/DOCS/Mongo+Wire+Protocol
*/

package mongo

import (
	"bytes"
	"container/vector"
	"encoding/binary"
	"io"
	"io/ioutil"
	"os"
)

// Request Opcodes
const (
	_OP_REPLY        = 1    // Reply to a client request. responseTo is set
	_OP_MSG          = 1000 // generic msg command followed by a string
	_OP_UPDATE       = 2001 // update document
	_OP_INSERT       = 2002 // insert new document
	_RESERVED        = 2003 // formerly used for _OP_GET_BY_OID
	_OP_QUERY        = 2004 // query a collection
	_OP_GET_MORE     = 2005 // Get more data from a query. See Cursors
	_OP_DELETE       = 2006 // Delete documents
	_OP_KILL_CURSORS = 2007 // Tell database client is done with a cursor
)

var last_req int32

// *** Standard Message Header
// ***

func header(length, reqID, respTo, opCode int32) []byte {
	b := make([]byte, 16)
	binary.LittleEndian.PutUint32(b[0:4], uint32(length))
	binary.LittleEndian.PutUint32(b[4:8], uint32(reqID))
	binary.LittleEndian.PutUint32(b[8:12], uint32(respTo))
	binary.LittleEndian.PutUint32(b[12:16], uint32(opCode))
	return b
}

// *** Interface
// ***

type message interface {
	Bytes() []byte
	RequestID() int32
	OpCode() int32
}

// *** Client Request Messages
// ***

type deleteMsg struct {
	fullCollectionName string
	selector           BSON
	requestID          int32
}

func (d *deleteMsg) OpCode() int32    { return _OP_DELETE }
func (d *deleteMsg) RequestID() int32 { return d.requestID }

func (d *deleteMsg) Bytes() []byte {
	zero := make([]byte, 4)
	buf := bytes.NewBuffer(zero)
	buf.WriteString(d.fullCollectionName)
	buf.WriteByte(0)
	buf.Write(zero)
	buf.Write(d.selector.Bytes())
	return buf.Bytes()

}

// ***

type getMoreMsg struct {
	fullCollectionName string
	numberToReturn     int32
	cursorID           int64
	requestID          int32
}

func (g *getMoreMsg) OpCode() int32    { return _OP_GET_MORE }
func (g *getMoreMsg) RequestID() int32 { return g.requestID }

func (g *getMoreMsg) Bytes() []byte {
	buf := bytes.NewBuffer(make([]byte, 4))
	buf.WriteString(g.fullCollectionName)
	buf.WriteByte(0)

	b := make([]byte, 4)
	binary.LittleEndian.PutUint32(b, uint32(g.numberToReturn))
	buf.Write(b)

	b = make([]byte, 8)
	binary.LittleEndian.PutUint64(b, uint64(g.cursorID))
	buf.Write(b)

	return buf.Bytes()
}

// ***

type insertMsg struct {
	fullCollectionName string
	doc                BSON
	requestID          int32
}

func (i *insertMsg) OpCode() int32    { return _OP_INSERT }
func (i *insertMsg) RequestID() int32 { return i.requestID }

func (i *insertMsg) Bytes() []byte {
	buf := bytes.NewBuffer(make([]byte, 4))
	buf.WriteString(i.fullCollectionName)
	buf.WriteByte(0)
	buf.Write(i.doc.Bytes())
	return buf.Bytes()
}

// ***

type killMsg struct {
	numberOfCursorIDs int32
	cursorIDs         []int64
	requestID         int32
}

func (k *killMsg) OpCode() int32    { return _OP_KILL_CURSORS }
func (k *killMsg) RequestID() int32 { return k.requestID }

func (k *killMsg) Bytes() []byte {
	buf := bytes.NewBuffer(make([]byte, 4))

	b := make([]byte, 4)
	binary.LittleEndian.PutUint32(b, uint32(k.numberOfCursorIDs))
	buf.Write(b)

	b = make([]byte, 8)
	for _, id := range k.cursorIDs {
		binary.LittleEndian.PutUint64(b, uint64(id))
		buf.Write(b)
	}

	return buf.Bytes()
}

// ***

type queryMsg struct {
	opts               int32
	fullCollectionName string
	numberToSkip       int32
	numberToReturn     int32
	query              BSON
	requestID          int32
}

func (q *queryMsg) OpCode() int32    { return _OP_QUERY }
func (q *queryMsg) RequestID() int32 { return q.requestID }

func (q *queryMsg) Bytes() []byte {
	b := make([]byte, 4)
	binary.LittleEndian.PutUint32(b, uint32(q.opts))

	buf := bytes.NewBuffer(b)
	buf.WriteString(q.fullCollectionName)
	buf.WriteByte(0)

	binary.LittleEndian.PutUint32(b, uint32(q.numberToSkip))
	buf.Write(b)

	binary.LittleEndian.PutUint32(b, uint32(q.numberToReturn))
	buf.Write(b)

	buf.Write(q.query.Bytes())
	return buf.Bytes()
}

// ***

type updateMsg struct {
	fullCollectionName string
	flags              int32
	selector, document BSON
	requestID          int32
}

func (u *updateMsg) OpCode() int32    { return _OP_UPDATE }
func (u *updateMsg) RequestID() int32 { return u.requestID }

func (u *updateMsg) Bytes() []byte {
	buf := bytes.NewBuffer(make([]byte, 4))
	buf.WriteString(u.fullCollectionName)
	buf.WriteByte(0)

	b := make([]byte, 4)
	binary.LittleEndian.PutUint32(b, uint32(u.flags))
	buf.Write(b)

	buf.Write(u.selector.Bytes())
	buf.Write(u.document.Bytes())

	return buf.Bytes()
}

// *** Database Response Message
// ***

type replyMsg struct {
	responseTo     int32
	responseFlag   int32
	cursorID       int64
	startingFrom   int32
	numberReturned int32
	docs           *vector.Vector
}

func (c *Connection) readReply() (*replyMsg, os.Error) {
	size_bits, _ := ioutil.ReadAll(io.LimitReader(c.conn, 4))
	size := binary.LittleEndian.Uint32(size_bits)
	rest, _ := ioutil.ReadAll(io.LimitReader(c.conn, int64(size)-4))
	reply := parseReply(rest)
	return reply, nil
}

func parseReply(b []byte) *replyMsg {
	r := new(replyMsg)
	r.responseTo = int32(binary.LittleEndian.Uint32(b[4:8]))
	r.responseFlag = int32(binary.LittleEndian.Uint32(b[12:16]))
	r.cursorID = int64(binary.LittleEndian.Uint64(b[16:24]))
	r.startingFrom = int32(binary.LittleEndian.Uint32(b[24:28]))
	r.numberReturned = int32(binary.LittleEndian.Uint32(b[28:32]))
	r.docs = new(vector.Vector)

	if r.numberReturned > 0 {
		buf := bytes.NewBuffer(b[36:len(b)])
		for i := 0; int32(i) < r.numberReturned; i++ {
			var bson BSON
			bb := new(_BSONBuilder)
			bb.ptr = &bson
			bb.Object()
			Parse(buf, bb)
			r.docs.Push(bson)
			ioutil.ReadAll(io.LimitReader(buf, 4))
		}
	}

	return r
}

<|endoftext|>


<commit_before>package bytefmt_test

import (
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	. "github.com/pivotal-golang/bytefmt"
)

var _ = Describe("bytefmt", func() {

	Context("ByteSize", func() {
		It("Prints in the largest possible unit", func() {
			Expect(ByteSize(10 * TERABYTE)).To(Equal("10T"))
			Expect(ByteSize(uint64(10.5 * TERABYTE))).To(Equal("10.5T"))

			Expect(ByteSize(10 * GIGABYTE)).To(Equal("10G"))
			Expect(ByteSize(uint64(10.5 * GIGABYTE))).To(Equal("10.5G"))

			Expect(ByteSize(100 * MEGABYTE)).To(Equal("100M"))
			Expect(ByteSize(uint64(100.5 * MEGABYTE))).To(Equal("100.5M"))

			Expect(ByteSize(100 * KILOBYTE)).To(Equal("100K"))
			Expect(ByteSize(uint64(100.5 * KILOBYTE))).To(Equal("100.5K"))

			Expect(ByteSize(1)).To(Equal("1B"))
		})

		It("prints '0' for zero bytes", func() {
			Expect(ByteSize(0)).To(Equal("0"))
		})
	})

	Context("ToMegabytes", func() {
		It("parses byte amounts with short units (e.g. M, G)", func() {
			var (
				megabytes uint64
				err       error
			)

			megabytes, err = ToMegabytes("5B")
			Expect(megabytes).To(Equal(uint64(0)))
			Expect(err).NotTo(HaveOccurred())

			megabytes, err = ToMegabytes("5K")
			Expect(megabytes).To(Equal(uint64(0)))
			Expect(err).NotTo(HaveOccurred())

			megabytes, err = ToMegabytes("5M")
			Expect(megabytes).To(Equal(uint64(5)))
			Expect(err).NotTo(HaveOccurred())

			megabytes, err = ToMegabytes("5m")
			Expect(megabytes).To(Equal(uint64(5)))
			Expect(err).NotTo(HaveOccurred())

			megabytes, err = ToMegabytes("2G")
			Expect(megabytes).To(Equal(uint64(2 * 1024)))
			Expect(err).NotTo(HaveOccurred())

			megabytes, err = ToMegabytes("3T")
			Expect(megabytes).To(Equal(uint64(3 * 1024 * 1024)))
			Expect(err).NotTo(HaveOccurred())
		})

		It("parses byte amounts with long units (e.g MB, GB)", func() {
			var (
				megabytes uint64
				err       error
			)

			megabytes, err = ToMegabytes("5MB")
			Expect(megabytes).To(Equal(uint64(5)))
			Expect(err).NotTo(HaveOccurred())

			megabytes, err = ToMegabytes("5mb")
			Expect(megabytes).To(Equal(uint64(5)))
			Expect(err).NotTo(HaveOccurred())

			megabytes, err = ToMegabytes("2GB")
			Expect(megabytes).To(Equal(uint64(2 * 1024)))
			Expect(err).NotTo(HaveOccurred())

			megabytes, err = ToMegabytes("3TB")
			Expect(megabytes).To(Equal(uint64(3 * 1024 * 1024)))
			Expect(err).NotTo(HaveOccurred())
		})

		It("returns an error when the unit is missing", func() {
			_, err := ToMegabytes("5")
			Expect(err).To(HaveOccurred())
			Expect(err.Error()).To(ContainSubstring("unit of measurement"))
		})

		It("returns an error when the unit is unrecognized", func() {
			_, err := ToMegabytes("5MBB")
			Expect(err).To(HaveOccurred())
			Expect(err.Error()).To(ContainSubstring("unit of measurement"))

			_, err = ToMegabytes("5BB")
			Expect(err).To(HaveOccurred())
		})

		It("allows whitespace before and after the value", func() {
			megabytes, err := ToMegabytes("\t\n\r 5MB ")
			Expect(megabytes).To(Equal(uint64(5)))
			Expect(err).NotTo(HaveOccurred())
		})

		It("returns an error for negative values", func() {
			_, err := ToMegabytes("-5MB")
			Expect(err).To(HaveOccurred())
			Expect(err.Error()).To(ContainSubstring("unit of measurement"))
		})

		It("returns an error for zero values", func() {
			_, err := ToMegabytes("0TB")
			Expect(err).To(HaveOccurred())
			Expect(err.Error()).To(ContainSubstring("unit of measurement"))
		})
	})

	Context("ToBytes", func() {
		It("parses byte amounts with short units (e.g. M, G)", func() {
			var (
				bytes uint64
				err   error
			)

			bytes, err = ToBytes("5B")
			Expect(bytes).To(Equal(uint64(5)))
			Expect(err).NotTo(HaveOccurred())

			bytes, err = ToBytes("5K")
			Expect(bytes).To(Equal(uint64(5 * KILOBYTE)))
			Expect(err).NotTo(HaveOccurred())

			bytes, err = ToBytes("5M")
			Expect(bytes).To(Equal(uint64(5 * MEGABYTE)))
			Expect(err).NotTo(HaveOccurred())

			bytes, err = ToBytes("5m")
			Expect(bytes).To(Equal(uint64(5 * MEGABYTE)))
			Expect(err).NotTo(HaveOccurred())

			bytes, err = ToBytes("2G")
			Expect(bytes).To(Equal(uint64(2 * GIGABYTE)))
			Expect(err).NotTo(HaveOccurred())

			bytes, err = ToBytes("3T")
			Expect(bytes).To(Equal(uint64(3 * TERABYTE)))
			Expect(err).NotTo(HaveOccurred())
		})

		It("parses byte amounts with long units (e.g MB, GB)", func() {
			var (
				bytes uint64
				err   error
			)

			bytes, err = ToBytes("5MB")
			Expect(bytes).To(Equal(uint64(5 * MEGABYTE)))
			Expect(err).NotTo(HaveOccurred())

			bytes, err = ToBytes("5mb")
			Expect(bytes).To(Equal(uint64(5 * MEGABYTE)))
			Expect(err).NotTo(HaveOccurred())

			bytes, err = ToBytes("2GB")
			Expect(bytes).To(Equal(uint64(2 * GIGABYTE)))
			Expect(err).NotTo(HaveOccurred())

			bytes, err = ToBytes("3TB")
			Expect(bytes).To(Equal(uint64(3 * TERABYTE)))
			Expect(err).NotTo(HaveOccurred())
		})

		It("returns an error when the unit is missing", func() {
			_, err := ToBytes("5")
			Expect(err).To(HaveOccurred())
			Expect(err.Error()).To(ContainSubstring("unit of measurement"))
		})

		It("returns an error when the unit is unrecognized", func() {
			_, err := ToBytes("5MBB")
			Expect(err).To(HaveOccurred())
			Expect(err.Error()).To(ContainSubstring("unit of measurement"))

			_, err = ToBytes("5BB")
			Expect(err).To(HaveOccurred())
		})

		It("allows whitespace before and after the value", func() {
			bytes, err := ToBytes("\t\n\r 5MB ")
			Expect(bytes).To(Equal(uint64(5 * MEGABYTE)))
			Expect(err).NotTo(HaveOccurred())
		})

		It("returns an error for negative values", func() {
			_, err := ToBytes("-5MB")
			Expect(err).To(HaveOccurred())
			Expect(err.Error()).To(ContainSubstring("unit of measurement"))
		})

		It("returns an error for zero values", func() {
			_, err := ToBytes("0TB")
			Expect(err).To(HaveOccurred())
			Expect(err.Error()).To(ContainSubstring("unit of measurement"))
		})
	})
})
<commit_msg>correctly update import location for bytefmt<commit_after>package bytefmt_test

import (
	. "code.cloudfoundry.org/bytefmt"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

var _ = Describe("bytefmt", func() {

	Context("ByteSize", func() {
		It("Prints in the largest possible unit", func() {
			Expect(ByteSize(10 * TERABYTE)).To(Equal("10T"))
			Expect(ByteSize(uint64(10.5 * TERABYTE))).To(Equal("10.5T"))

			Expect(ByteSize(10 * GIGABYTE)).To(Equal("10G"))
			Expect(ByteSize(uint64(10.5 * GIGABYTE))).To(Equal("10.5G"))

			Expect(ByteSize(100 * MEGABYTE)).To(Equal("100M"))
			Expect(ByteSize(uint64(100.5 * MEGABYTE))).To(Equal("100.5M"))

			Expect(ByteSize(100 * KILOBYTE)).To(Equal("100K"))
			Expect(ByteSize(uint64(100.5 * KILOBYTE))).To(Equal("100.5K"))

			Expect(ByteSize(1)).To(Equal("1B"))
		})

		It("prints '0' for zero bytes", func() {
			Expect(ByteSize(0)).To(Equal("0"))
		})
	})

	Context("ToMegabytes", func() {
		It("parses byte amounts with short units (e.g. M, G)", func() {
			var (
				megabytes uint64
				err       error
			)

			megabytes, err = ToMegabytes("5B")
			Expect(megabytes).To(Equal(uint64(0)))
			Expect(err).NotTo(HaveOccurred())

			megabytes, err = ToMegabytes("5K")
			Expect(megabytes).To(Equal(uint64(0)))
			Expect(err).NotTo(HaveOccurred())

			megabytes, err = ToMegabytes("5M")
			Expect(megabytes).To(Equal(uint64(5)))
			Expect(err).NotTo(HaveOccurred())

			megabytes, err = ToMegabytes("5m")
			Expect(megabytes).To(Equal(uint64(5)))
			Expect(err).NotTo(HaveOccurred())

			megabytes, err = ToMegabytes("2G")
			Expect(megabytes).To(Equal(uint64(2 * 1024)))
			Expect(err).NotTo(HaveOccurred())

			megabytes, err = ToMegabytes("3T")
			Expect(megabytes).To(Equal(uint64(3 * 1024 * 1024)))
			Expect(err).NotTo(HaveOccurred())
		})

		It("parses byte amounts with long units (e.g MB, GB)", func() {
			var (
				megabytes uint64
				err       error
			)

			megabytes, err = ToMegabytes("5MB")
			Expect(megabytes).To(Equal(uint64(5)))
			Expect(err).NotTo(HaveOccurred())

			megabytes, err = ToMegabytes("5mb")
			Expect(megabytes).To(Equal(uint64(5)))
			Expect(err).NotTo(HaveOccurred())

			megabytes, err = ToMegabytes("2GB")
			Expect(megabytes).To(Equal(uint64(2 * 1024)))
			Expect(err).NotTo(HaveOccurred())

			megabytes, err = ToMegabytes("3TB")
			Expect(megabytes).To(Equal(uint64(3 * 1024 * 1024)))
			Expect(err).NotTo(HaveOccurred())
		})

		It("returns an error when the unit is missing", func() {
			_, err := ToMegabytes("5")
			Expect(err).To(HaveOccurred())
			Expect(err.Error()).To(ContainSubstring("unit of measurement"))
		})

		It("returns an error when the unit is unrecognized", func() {
			_, err := ToMegabytes("5MBB")
			Expect(err).To(HaveOccurred())
			Expect(err.Error()).To(ContainSubstring("unit of measurement"))

			_, err = ToMegabytes("5BB")
			Expect(err).To(HaveOccurred())
		})

		It("allows whitespace before and after the value", func() {
			megabytes, err := ToMegabytes("\t\n\r 5MB ")
			Expect(megabytes).To(Equal(uint64(5)))
			Expect(err).NotTo(HaveOccurred())
		})

		It("returns an error for negative values", func() {
			_, err := ToMegabytes("-5MB")
			Expect(err).To(HaveOccurred())
			Expect(err.Error()).To(ContainSubstring("unit of measurement"))
		})

		It("returns an error for zero values", func() {
			_, err := ToMegabytes("0TB")
			Expect(err).To(HaveOccurred())
			Expect(err.Error()).To(ContainSubstring("unit of measurement"))
		})
	})

	Context("ToBytes", func() {
		It("parses byte amounts with short units (e.g. M, G)", func() {
			var (
				bytes uint64
				err   error
			)

			bytes, err = ToBytes("5B")
			Expect(bytes).To(Equal(uint64(5)))
			Expect(err).NotTo(HaveOccurred())

			bytes, err = ToBytes("5K")
			Expect(bytes).To(Equal(uint64(5 * KILOBYTE)))
			Expect(err).NotTo(HaveOccurred())

			bytes, err = ToBytes("5M")
			Expect(bytes).To(Equal(uint64(5 * MEGABYTE)))
			Expect(err).NotTo(HaveOccurred())

			bytes, err = ToBytes("5m")
			Expect(bytes).To(Equal(uint64(5 * MEGABYTE)))
			Expect(err).NotTo(HaveOccurred())

			bytes, err = ToBytes("2G")
			Expect(bytes).To(Equal(uint64(2 * GIGABYTE)))
			Expect(err).NotTo(HaveOccurred())

			bytes, err = ToBytes("3T")
			Expect(bytes).To(Equal(uint64(3 * TERABYTE)))
			Expect(err).NotTo(HaveOccurred())
		})

		It("parses byte amounts with long units (e.g MB, GB)", func() {
			var (
				bytes uint64
				err   error
			)

			bytes, err = ToBytes("5MB")
			Expect(bytes).To(Equal(uint64(5 * MEGABYTE)))
			Expect(err).NotTo(HaveOccurred())

			bytes, err = ToBytes("5mb")
			Expect(bytes).To(Equal(uint64(5 * MEGABYTE)))
			Expect(err).NotTo(HaveOccurred())

			bytes, err = ToBytes("2GB")
			Expect(bytes).To(Equal(uint64(2 * GIGABYTE)))
			Expect(err).NotTo(HaveOccurred())

			bytes, err = ToBytes("3TB")
			Expect(bytes).To(Equal(uint64(3 * TERABYTE)))
			Expect(err).NotTo(HaveOccurred())
		})

		It("returns an error when the unit is missing", func() {
			_, err := ToBytes("5")
			Expect(err).To(HaveOccurred())
			Expect(err.Error()).To(ContainSubstring("unit of measurement"))
		})

		It("returns an error when the unit is unrecognized", func() {
			_, err := ToBytes("5MBB")
			Expect(err).To(HaveOccurred())
			Expect(err.Error()).To(ContainSubstring("unit of measurement"))

			_, err = ToBytes("5BB")
			Expect(err).To(HaveOccurred())
		})

		It("allows whitespace before and after the value", func() {
			bytes, err := ToBytes("\t\n\r 5MB ")
			Expect(bytes).To(Equal(uint64(5 * MEGABYTE)))
			Expect(err).NotTo(HaveOccurred())
		})

		It("returns an error for negative values", func() {
			_, err := ToBytes("-5MB")
			Expect(err).To(HaveOccurred())
			Expect(err.Error()).To(ContainSubstring("unit of measurement"))
		})

		It("returns an error for zero values", func() {
			_, err := ToBytes("0TB")
			Expect(err).To(HaveOccurred())
			Expect(err.Error()).To(ContainSubstring("unit of measurement"))
		})
	})
})
<|endoftext|>


<commit_before>package cache

import (
	"context"
	"sync"

	"github.com/containerd/containerd/mount"
	"github.com/moby/buildkit/cache/metadata"
	"github.com/moby/buildkit/identity"
	"github.com/moby/buildkit/util/flightcontrol"
	"github.com/pkg/errors"
	"github.com/sirupsen/logrus"
)

// Ref is a reference to cacheable objects.
type Ref interface {
	Mountable
	ID() string
	Release(context.Context) error
	Size(ctx context.Context) (int64, error)
	Metadata() *metadata.StorageItem
}

type ImmutableRef interface {
	Ref
	Parent() ImmutableRef
	Finalize(ctx context.Context) error // Make sure reference is flushed to driver
}

type MutableRef interface {
	Ref
	Commit(context.Context) (ImmutableRef, error)
}

type Mountable interface {
	Mount(ctx context.Context, readonly bool) ([]mount.Mount, error)
}

type cacheRecord struct {
	cm *cacheManager
	mu *sync.Mutex // the mutex is shared by records sharing data

	mutable bool
	refs    map[Mountable]struct{}
	parent  ImmutableRef
	md      *metadata.StorageItem

	// dead means record is marked as deleted
	dead bool

	view      string
	viewMount []mount.Mount

	sizeG flightcontrol.Group

	// these are filled if multiple refs point to same data
	equalMutable   *mutableRef
	equalImmutable *immutableRef
}

// hold ref lock before calling
func (cr *cacheRecord) ref() *immutableRef {
	ref := &immutableRef{cacheRecord: cr}
	cr.refs[ref] = struct{}{}
	return ref
}

// hold ref lock before calling
func (cr *cacheRecord) mref() *mutableRef {
	ref := &mutableRef{cacheRecord: cr}
	cr.refs[ref] = struct{}{}
	return ref
}

// hold ref lock before calling
func (cr *cacheRecord) isDead() bool {
	return cr.dead || (cr.equalImmutable != nil && cr.equalImmutable.dead) || (cr.equalMutable != nil && cr.equalMutable.dead)
}

func (cr *cacheRecord) Size(ctx context.Context) (int64, error) {
	// this expects that usage() is implemented lazily
	s, err := cr.sizeG.Do(ctx, cr.ID(), func(ctx context.Context) (interface{}, error) {
		cr.mu.Lock()
		s := getSize(cr.md)
		if s != sizeUnknown {
			cr.mu.Unlock()
			return s, nil
		}
		driverID := cr.ID()
		if cr.equalMutable != nil {
			driverID = cr.equalMutable.ID()
		}
		cr.mu.Unlock()
		usage, err := cr.cm.ManagerOpt.Snapshotter.Usage(ctx, driverID)
		if err != nil {
			if cr.isDead() {
				cr.mu.Unlock()
				return int64(0), nil
			}
			return s, errors.Wrapf(err, "failed to get usage for %s", cr.ID())
		}
		cr.mu.Lock()
		setSize(cr.md, usage.Size)
		if err := cr.md.Commit(); err != nil {
			return s, err
		}
		cr.mu.Unlock()
		return usage.Size, nil
	})
	return s.(int64), err
}

func (cr *cacheRecord) Parent() ImmutableRef {
	if cr.parent == nil {
		return nil
	}
	p := cr.parent.(*immutableRef)
	p.mu.Lock()
	defer p.mu.Unlock()
	return p.ref()
}

func (cr *cacheRecord) Mount(ctx context.Context, readonly bool) ([]mount.Mount, error) {
	cr.mu.Lock()
	defer cr.mu.Unlock()

	if cr.mutable {
		m, err := cr.cm.Snapshotter.Mounts(ctx, cr.ID())
		if err != nil {
			return nil, errors.Wrapf(err, "failed to mount %s", cr.ID())
		}
		if readonly {
			m = setReadonly(m)
		}
		return m, nil
	}

	if cr.equalMutable != nil && readonly {
		m, err := cr.cm.Snapshotter.Mounts(ctx, cr.equalMutable.ID())
		if err != nil {
			return nil, errors.Wrapf(err, "failed to mount %s", cr.equalMutable.ID())
		}
		return setReadonly(m), nil
	}

	if err := cr.finalize(ctx); err != nil {
		return nil, err
	}
	if cr.viewMount == nil { // TODO: handle this better
		cr.view = identity.NewID()
		m, err := cr.cm.Snapshotter.View(ctx, cr.view, cr.ID())
		if err != nil {
			cr.view = ""
			return nil, errors.Wrapf(err, "failed to mount %s", cr.ID())
		}
		cr.viewMount = m
	}
	return cr.viewMount, nil
}

// call when holding the manager lock
func (cr *cacheRecord) remove(ctx context.Context, removeSnapshot bool) error {
	delete(cr.cm.records, cr.ID())
	if cr.parent != nil {
		if err := cr.parent.(*immutableRef).release(ctx); err != nil {
			return err
		}
	}
	if removeSnapshot {
		if err := cr.cm.Snapshotter.Remove(ctx, cr.ID()); err != nil {
			return err
		}
	}
	if err := cr.cm.md.Clear(cr.ID()); err != nil {
		return err
	}
	return nil
}

func (cr *cacheRecord) ID() string {
	return cr.md.ID()
}

type immutableRef struct {
	*cacheRecord
}

type mutableRef struct {
	*cacheRecord
}

func (sr *immutableRef) Release(ctx context.Context) error {
	sr.cm.mu.Lock()
	defer sr.cm.mu.Unlock()

	sr.mu.Lock()
	defer sr.mu.Unlock()

	return sr.release(ctx)
}

func (sr *immutableRef) release(ctx context.Context) error {
	updateLastUsed(sr.md)

	delete(sr.refs, sr)

	if len(sr.refs) == 0 {
		if sr.viewMount != nil { // TODO: release viewMount earlier if possible
			if err := sr.cm.Snapshotter.Remove(ctx, sr.view); err != nil {
				return err
			}
			sr.view = ""
			sr.viewMount = nil
		}

		if sr.equalMutable != nil {
			sr.equalMutable.release(ctx)
		}
		// go sr.cm.GC()
	}

	return nil
}

func (sr *immutableRef) Finalize(ctx context.Context) error {
	sr.mu.Lock()
	defer sr.mu.Unlock()

	return sr.finalize(ctx)
}

func (cr *cacheRecord) Metadata() *metadata.StorageItem {
	return cr.md
}

func (cr *cacheRecord) finalize(ctx context.Context) error {
	mutable := cr.equalMutable
	if mutable == nil {
		return nil
	}
	err := cr.cm.Snapshotter.Commit(ctx, cr.ID(), mutable.ID())
	if err != nil {
		return errors.Wrapf(err, "failed to commit %s", mutable.ID())
	}
	mutable.dead = true
	go func() {
		cr.cm.mu.Lock()
		defer cr.cm.mu.Unlock()
		if err := mutable.remove(context.TODO(), false); err != nil {
			logrus.Error(err)
		}
	}()
	cr.equalMutable = nil
	clearEqualMutable(cr.md)
	return cr.md.Commit()
}

func (sr *mutableRef) commit(ctx context.Context) (ImmutableRef, error) {
	if !sr.mutable || len(sr.refs) == 0 {
		return nil, errors.Wrapf(errInvalid, "invalid mutable ref")
	}

	id := identity.NewID()
	md, _ := sr.cm.md.Get(id)

	rec := &cacheRecord{
		mu:           sr.mu,
		cm:           sr.cm,
		parent:       sr.Parent(),
		equalMutable: sr,
		refs:         make(map[Mountable]struct{}),
		md:           md,
	}

	if descr := GetDescription(sr.md); descr != "" {
		if err := queueDescription(md, descr); err != nil {
			return nil, err
		}
	}

	if err := initializeMetadata(rec); err != nil {
		return nil, err
	}

	sr.cm.records[id] = rec

	if err := sr.md.Commit(); err != nil {
		return nil, err
	}

	setSize(md, sizeUnknown)
	setEqualMutable(md, sr.ID())
	if err := md.Commit(); err != nil {
		return nil, err
	}

	ref := rec.ref()
	sr.equalImmutable = ref
	return ref, nil
}

func (sr *mutableRef) Commit(ctx context.Context) (ImmutableRef, error) {
	sr.cm.mu.Lock()
	defer sr.cm.mu.Unlock()

	sr.mu.Lock()
	defer sr.mu.Unlock()

	return sr.commit(ctx)
}

func (sr *mutableRef) Release(ctx context.Context) error {
	sr.cm.mu.Lock()
	defer sr.cm.mu.Unlock()

	sr.mu.Lock()
	defer sr.mu.Unlock()

	return sr.release(ctx)
}

func (sr *mutableRef) release(ctx context.Context) error {
	delete(sr.refs, sr)
	updateLastUsed(sr.md)
	if getCachePolicy(sr.md) != cachePolicyRetain {
		if sr.equalImmutable != nil {
			if getCachePolicy(sr.equalImmutable.md) == cachePolicyRetain {
				return nil
			}
			if err := sr.equalImmutable.remove(ctx, false); err != nil {
				return err
			}
		}
		if sr.parent != nil {
			if err := sr.parent.(*immutableRef).release(ctx); err != nil {
				return err
			}
		}
		return sr.remove(ctx, true)
	}
	return nil
}

func setReadonly(mounts []mount.Mount) []mount.Mount {
	for i, m := range mounts {
		opts := make([]string, 0, len(m.Options))
		for _, opt := range m.Options {
			if opt != "rw" {
				opts = append(opts, opt)
			}
		}
		opts = append(opts, "ro")
		mounts[i].Options = opts
	}
	return mounts
}
<commit_msg>cache: fix refs unlocking bug<commit_after>package cache

import (
	"context"
	"sync"

	"github.com/containerd/containerd/mount"
	"github.com/moby/buildkit/cache/metadata"
	"github.com/moby/buildkit/identity"
	"github.com/moby/buildkit/util/flightcontrol"
	"github.com/pkg/errors"
	"github.com/sirupsen/logrus"
)

// Ref is a reference to cacheable objects.
type Ref interface {
	Mountable
	ID() string
	Release(context.Context) error
	Size(ctx context.Context) (int64, error)
	Metadata() *metadata.StorageItem
}

type ImmutableRef interface {
	Ref
	Parent() ImmutableRef
	Finalize(ctx context.Context) error // Make sure reference is flushed to driver
}

type MutableRef interface {
	Ref
	Commit(context.Context) (ImmutableRef, error)
}

type Mountable interface {
	Mount(ctx context.Context, readonly bool) ([]mount.Mount, error)
}

type cacheRecord struct {
	cm *cacheManager
	mu *sync.Mutex // the mutex is shared by records sharing data

	mutable bool
	refs    map[Mountable]struct{}
	parent  ImmutableRef
	md      *metadata.StorageItem

	// dead means record is marked as deleted
	dead bool

	view      string
	viewMount []mount.Mount

	sizeG flightcontrol.Group

	// these are filled if multiple refs point to same data
	equalMutable   *mutableRef
	equalImmutable *immutableRef
}

// hold ref lock before calling
func (cr *cacheRecord) ref() *immutableRef {
	ref := &immutableRef{cacheRecord: cr}
	cr.refs[ref] = struct{}{}
	return ref
}

// hold ref lock before calling
func (cr *cacheRecord) mref() *mutableRef {
	ref := &mutableRef{cacheRecord: cr}
	cr.refs[ref] = struct{}{}
	return ref
}

// hold ref lock before calling
func (cr *cacheRecord) isDead() bool {
	return cr.dead || (cr.equalImmutable != nil && cr.equalImmutable.dead) || (cr.equalMutable != nil && cr.equalMutable.dead)
}

func (cr *cacheRecord) Size(ctx context.Context) (int64, error) {
	// this expects that usage() is implemented lazily
	s, err := cr.sizeG.Do(ctx, cr.ID(), func(ctx context.Context) (interface{}, error) {
		cr.mu.Lock()
		s := getSize(cr.md)
		if s != sizeUnknown {
			cr.mu.Unlock()
			return s, nil
		}
		driverID := cr.ID()
		if cr.equalMutable != nil {
			driverID = cr.equalMutable.ID()
		}
		cr.mu.Unlock()
		usage, err := cr.cm.ManagerOpt.Snapshotter.Usage(ctx, driverID)
		if err != nil {
			cr.mu.Lock()
			isDead := cr.isDead()
			cr.mu.Unlock()
			if isDead {
				return int64(0), nil
			}
			return s, errors.Wrapf(err, "failed to get usage for %s", cr.ID())
		}
		cr.mu.Lock()
		setSize(cr.md, usage.Size)
		if err := cr.md.Commit(); err != nil {
			cr.mu.Unlock()
			return s, err
		}
		cr.mu.Unlock()
		return usage.Size, nil
	})
	return s.(int64), err
}

func (cr *cacheRecord) Parent() ImmutableRef {
	if cr.parent == nil {
		return nil
	}
	p := cr.parent.(*immutableRef)
	p.mu.Lock()
	defer p.mu.Unlock()
	return p.ref()
}

func (cr *cacheRecord) Mount(ctx context.Context, readonly bool) ([]mount.Mount, error) {
	cr.mu.Lock()
	defer cr.mu.Unlock()

	if cr.mutable {
		m, err := cr.cm.Snapshotter.Mounts(ctx, cr.ID())
		if err != nil {
			return nil, errors.Wrapf(err, "failed to mount %s", cr.ID())
		}
		if readonly {
			m = setReadonly(m)
		}
		return m, nil
	}

	if cr.equalMutable != nil && readonly {
		m, err := cr.cm.Snapshotter.Mounts(ctx, cr.equalMutable.ID())
		if err != nil {
			return nil, errors.Wrapf(err, "failed to mount %s", cr.equalMutable.ID())
		}
		return setReadonly(m), nil
	}

	if err := cr.finalize(ctx); err != nil {
		return nil, err
	}
	if cr.viewMount == nil { // TODO: handle this better
		cr.view = identity.NewID()
		m, err := cr.cm.Snapshotter.View(ctx, cr.view, cr.ID())
		if err != nil {
			cr.view = ""
			return nil, errors.Wrapf(err, "failed to mount %s", cr.ID())
		}
		cr.viewMount = m
	}
	return cr.viewMount, nil
}

// call when holding the manager lock
func (cr *cacheRecord) remove(ctx context.Context, removeSnapshot bool) error {
	delete(cr.cm.records, cr.ID())
	if cr.parent != nil {
		if err := cr.parent.(*immutableRef).release(ctx); err != nil {
			return err
		}
	}
	if removeSnapshot {
		if err := cr.cm.Snapshotter.Remove(ctx, cr.ID()); err != nil {
			return err
		}
	}
	if err := cr.cm.md.Clear(cr.ID()); err != nil {
		return err
	}
	return nil
}

func (cr *cacheRecord) ID() string {
	return cr.md.ID()
}

type immutableRef struct {
	*cacheRecord
}

type mutableRef struct {
	*cacheRecord
}

func (sr *immutableRef) Release(ctx context.Context) error {
	sr.cm.mu.Lock()
	defer sr.cm.mu.Unlock()

	sr.mu.Lock()
	defer sr.mu.Unlock()

	return sr.release(ctx)
}

func (sr *immutableRef) release(ctx context.Context) error {
	updateLastUsed(sr.md)

	delete(sr.refs, sr)

	if len(sr.refs) == 0 {
		if sr.viewMount != nil { // TODO: release viewMount earlier if possible
			if err := sr.cm.Snapshotter.Remove(ctx, sr.view); err != nil {
				return err
			}
			sr.view = ""
			sr.viewMount = nil
		}

		if sr.equalMutable != nil {
			sr.equalMutable.release(ctx)
		}
		// go sr.cm.GC()
	}

	return nil
}

func (sr *immutableRef) Finalize(ctx context.Context) error {
	sr.mu.Lock()
	defer sr.mu.Unlock()

	return sr.finalize(ctx)
}

func (cr *cacheRecord) Metadata() *metadata.StorageItem {
	return cr.md
}

func (cr *cacheRecord) finalize(ctx context.Context) error {
	mutable := cr.equalMutable
	if mutable == nil {
		return nil
	}
	err := cr.cm.Snapshotter.Commit(ctx, cr.ID(), mutable.ID())
	if err != nil {
		return errors.Wrapf(err, "failed to commit %s", mutable.ID())
	}
	mutable.dead = true
	go func() {
		cr.cm.mu.Lock()
		defer cr.cm.mu.Unlock()
		if err := mutable.remove(context.TODO(), false); err != nil {
			logrus.Error(err)
		}
	}()
	cr.equalMutable = nil
	clearEqualMutable(cr.md)
	return cr.md.Commit()
}

func (sr *mutableRef) commit(ctx context.Context) (ImmutableRef, error) {
	if !sr.mutable || len(sr.refs) == 0 {
		return nil, errors.Wrapf(errInvalid, "invalid mutable ref")
	}

	id := identity.NewID()
	md, _ := sr.cm.md.Get(id)

	rec := &cacheRecord{
		mu:           sr.mu,
		cm:           sr.cm,
		parent:       sr.Parent(),
		equalMutable: sr,
		refs:         make(map[Mountable]struct{}),
		md:           md,
	}

	if descr := GetDescription(sr.md); descr != "" {
		if err := queueDescription(md, descr); err != nil {
			return nil, err
		}
	}

	if err := initializeMetadata(rec); err != nil {
		return nil, err
	}

	sr.cm.records[id] = rec

	if err := sr.md.Commit(); err != nil {
		return nil, err
	}

	setSize(md, sizeUnknown)
	setEqualMutable(md, sr.ID())
	if err := md.Commit(); err != nil {
		return nil, err
	}

	ref := rec.ref()
	sr.equalImmutable = ref
	return ref, nil
}

func (sr *mutableRef) Commit(ctx context.Context) (ImmutableRef, error) {
	sr.cm.mu.Lock()
	defer sr.cm.mu.Unlock()

	sr.mu.Lock()
	defer sr.mu.Unlock()

	return sr.commit(ctx)
}

func (sr *mutableRef) Release(ctx context.Context) error {
	sr.cm.mu.Lock()
	defer sr.cm.mu.Unlock()

	sr.mu.Lock()
	defer sr.mu.Unlock()

	return sr.release(ctx)
}

func (sr *mutableRef) release(ctx context.Context) error {
	delete(sr.refs, sr)
	updateLastUsed(sr.md)
	if getCachePolicy(sr.md) != cachePolicyRetain {
		if sr.equalImmutable != nil {
			if getCachePolicy(sr.equalImmutable.md) == cachePolicyRetain {
				return nil
			}
			if err := sr.equalImmutable.remove(ctx, false); err != nil {
				return err
			}
		}
		if sr.parent != nil {
			if err := sr.parent.(*immutableRef).release(ctx); err != nil {
				return err
			}
		}
		return sr.remove(ctx, true)
	}
	return nil
}

func setReadonly(mounts []mount.Mount) []mount.Mount {
	for i, m := range mounts {
		opts := make([]string, 0, len(m.Options))
		for _, opt := range m.Options {
			if opt != "rw" {
				opts = append(opts, opt)
			}
		}
		opts = append(opts, "ro")
		mounts[i].Options = opts
	}
	return mounts
}
<|endoftext|>


<commit_before>package hkit

import (
	"errors"
	"net/http"
)

type HTTPCache struct {
	cacheService CacheService
}

var (
	errHeaderAlreadyLocked = errors.New("Header has already been locked")
	ErrAlreadyWritten      = errors.New("Response has already been written")
)

// HTTPResponseInterceptor will capture actions on an underlaying
// ResponseWriter. IMPORTANT NOTE: will only capture header state until
// the first call to WriteHeader or Write. Therefore no support for
// trailers etc.
// Behavior may diverge from the implementors in the "net/http", so no
// crazily protective behaviors should be expected from this wrapper.
// TL;DR: you probably have your own cache wrapper if you are writing
// e.g. trailer headers all over your handlers.
type HTTPResponseInterceptor struct {
	tempHeader   http.Header
	lockedHeader http.Header
	writtenBytes []byte
	writtenCode  int
}

func NewHTTPResponseInterceptor(w http.ResponseWriter) *HTTPResponseInterceptor {
	interceptor := &HTTPResponseInterceptor{
		tempHeader:   cloneHeader(w.Header()),
		lockedHeader: nil,
		writtenBytes: nil,
		writtenCode:  -1,
	}
	return interceptor
}

func (w *HTTPResponseInterceptor) Write(bytes []byte) (int, error) {
	w.lockHeader()
	if w.writtenBytes != nil {
		return 0, ErrAlreadyWritten
	}
	w.writtenBytes = bytes
	return w.Write(bytes)
}

func (w *HTTPResponseInterceptor) WriteHeader(code int) {
	if err := w.lockHeader(); err != nil {
		return
	}
	w.writtenCode = code
	w.WriteHeader(code)
}

func (w *HTTPResponseInterceptor) Header() http.Header {
	return w.tempHeader
}

func (w *HTTPResponseInterceptor) lockHeader() error {
	if w.lockedHeader != nil {
		return errHeaderAlreadyLocked
	}
	w.lockedHeader = cloneHeader(w.tempHeader)
	overwriteHeader(w.Header(), w.lockedHeader)
	return nil
}

func (w *HTTPResponseInterceptor) Replay(writer http.ResponseWriter) {
	if w.lockedHeader != nil {
		overwriteHeader(writer.Header(), w.lockedHeader)
	}
	if w.writtenCode != -1 {
		writer.WriteHeader(w.writtenCode)
	}
	if w.writtenBytes != nil {
		writer.Write(w.writtenBytes)
	}
}

func cloneHeader(h http.Header) http.Header {
	h2 := make(http.Header, len(h))
	for k, vv := range h {
		vv2 := make([]string, len(vv))
		copy(vv2, vv)
		h2[k] = vv2
	}
	return h2
}

func overwriteHeader(lhs, rhs http.Header) {
	keys := make([]string, 0)
	for k, _ := range lhs {
		keys = append(keys, k)
	}
	for _, k := range keys {
		lhs.Del(k)
	}

	for k, v := range rhs {
		vv := make([]string, len(v))
		copy(vv, v)
		lhs[k] = vv
	}
}
<commit_msg>Fix 80 character line limit in cache_http.go<commit_after>package hkit

import (
	"errors"
	"net/http"
)

type HTTPCache struct {
	cacheService CacheService
}

var (
	errHeaderAlreadyLocked = errors.New("Header has already been locked")
	ErrAlreadyWritten      = errors.New("Response has already been written")
)

// HTTPResponseInterceptor will capture actions on an underlaying
// ResponseWriter. IMPORTANT NOTE: will only capture header state until
// the first call to WriteHeader or Write. Therefore no support for
// trailers etc.
// Behavior may diverge from the implementors in the "net/http", so no
// crazily protective behaviors should be expected from this wrapper.
// TL;DR: you probably have your own cache wrapper if you are writing
// e.g. trailer headers all over your handlers.
type HTTPResponseInterceptor struct {
	tempHeader   http.Header
	lockedHeader http.Header
	writtenBytes []byte
	writtenCode  int
}

func NewHTTPResponseInterceptor(
	w http.ResponseWriter) *HTTPResponseInterceptor {

	interceptor := &HTTPResponseInterceptor{
		tempHeader:   cloneHeader(w.Header()),
		lockedHeader: nil,
		writtenBytes: nil,
		writtenCode:  -1,
	}
	return interceptor
}

func (w *HTTPResponseInterceptor) Write(bytes []byte) (int, error) {
	w.lockHeader()
	if w.writtenBytes != nil {
		return 0, ErrAlreadyWritten
	}
	w.writtenBytes = bytes
	return w.Write(bytes)
}

func (w *HTTPResponseInterceptor) WriteHeader(code int) {
	if err := w.lockHeader(); err != nil {
		return
	}
	w.writtenCode = code
	w.WriteHeader(code)
}

func (w *HTTPResponseInterceptor) Header() http.Header {
	return w.tempHeader
}

func (w *HTTPResponseInterceptor) lockHeader() error {
	if w.lockedHeader != nil {
		return errHeaderAlreadyLocked
	}
	w.lockedHeader = cloneHeader(w.tempHeader)
	overwriteHeader(w.Header(), w.lockedHeader)
	return nil
}

func (w *HTTPResponseInterceptor) Replay(writer http.ResponseWriter) {
	if w.lockedHeader != nil {
		overwriteHeader(writer.Header(), w.lockedHeader)
	}
	if w.writtenCode != -1 {
		writer.WriteHeader(w.writtenCode)
	}
	if w.writtenBytes != nil {
		writer.Write(w.writtenBytes)
	}
}

func cloneHeader(h http.Header) http.Header {
	h2 := make(http.Header, len(h))
	for k, vv := range h {
		vv2 := make([]string, len(vv))
		copy(vv2, vv)
		h2[k] = vv2
	}
	return h2
}

func overwriteHeader(lhs, rhs http.Header) {
	keys := make([]string, 0)
	for k, _ := range lhs {
		keys = append(keys, k)
	}
	for _, k := range keys {
		lhs.Del(k)
	}

	for k, v := range rhs {
		vv := make([]string, len(v))
		copy(vv, v)
		lhs[k] = vv
	}
}
<|endoftext|>


<commit_before>package cache

import (
	"encoding/json"
	"testing"
)

type cValue struct {
	size int
}

func (c *cValue) Size() int {
	return len(string(c.size))
}

func testInitialState(t *testing.T) {
	cache := LRUCache(10)
	l, s, capacity := cache.CacheStats()
	if l != 0 {
		t.Errorf("length = %v, want 0", l)
	}
	if s != 0 {
		t.Errorf("size = %v, want 0", s)
	}
	if capacity != 10 {
		t.Errorf("capacity = %v, want 10", capacity)
	}
}

func testInsert(t *testing.T) {
	cache := LRUCache(100)
	data := json.Marshal(&cValue{0})
	cache.Insert("key", string(data))

	v, ok := cache.Get("key")
	if !ok || v(*cValue) != data {
		t.Errorf("Cache has incorrect value: %v != %v", data, v)
	}
}

func testLRUEviction(t *testing.T) {
	size := uint64(3)
	cache := LRUCache(size)

	cache.Insert("key1", &cValue{1})
	cache.Insert("key2", &cValue{1})
	cache.Insert("key3", &cValue{1})
	// lru: [key3, key2, key1]

	// Look up the elements. This will rearrange the LRU ordering.
	cache.Get("key3")
	cache.Get("key2")
	cache.Get("key1")
	// lru: [key1, key2, key3]

	cache.Insert("key0", &cValue{1})
	// lru: [key0, key1, key2]

	// The least recently used one should have been evicted.
	if _, ok := cache.Get("key3"); ok {
		t.Error("Least recently used element was not evicted.")
	}
}

func testCapacityIsObeyed(t *testing.T) {
	size := uint64(3)
	cache := LRUCache(size)
	value := &cValue{1}

	// Insert up to the cache's capacity.
	cache.Insert("key1", value)
	cache.Insert("key2", value)
	cache.Insert("key3", value)
	if _, sz, _, _ := cache.CacheStats(); sz != size {
		t.Errorf("cache.Size() = %v, expected %v", sz, size)
	}
	// Insert one more; something should be evicted to make room.
	cache.Insert("key4", value)
	if _, sz, _ := cache.CacheStats(); sz != size {
		t.Errorf("post-evict cache.Size() = %v, expected %v", sz, size)
	}
}

func testInsertUpdatesSize(t *testing.T) {
	cache := LRUCache(100)
	emptyValue := &cValue{0}
	key := "key1"
	cache.Insert(key, emptyValue)
	if _, sz, _ := cache.CacheStats(); sz != 0 {
		t.Errorf("cache.Size() = %v, expected 0", sz)
	}
	someValue := &cValue{20}
	key = "key2"
	cache.Insert(key, someValue)
	if _, sz, _ := cache.CacheStats(); sz != 20 {
		t.Errorf("cache.Size() = %v, expected 20", sz)
	}
}

func testInsertWithOldKeyUpdatesValue(t *testing.T) {
	cache := LRUCache(100)
	emptyValue := &cValue{0}
	key := "key1"
	cache.Insert(key, emptyValue)
	someValue := &cValue{20}
	cache.Insert(key, someValue)

	v, ok := cache.Get(key)
	if !ok || v.(*cValue) != someValue {
		t.Errorf("Cache has incorrect value: %v != %v", someValue, v)
	}
}
<commit_msg>working test cases<commit_after>//TODO: fix cachestats test cases

package cache

import (
	"encoding/json"
	"testing"
)

type cValue struct {
	Size int
}

func (c *cValue) getSize() int {
	return len(string(c.Size))
}

func testInitialState(t *testing.T) {
	cache := LRUCache(10)
	l, s, capacity := cache.CacheStats()
	if l != 0 {
		t.Errorf("length = %v, want 0", l)
	}
	if s != 0 {
		t.Errorf("Size = %v, want 0", s)
	}
	if capacity != 10 {
		t.Errorf("capacity = %v, want 10", capacity)
	}
}

func testInsert(t *testing.T) {
	cache := LRUCache(100)
	data, _ := json.Marshal(&cValue{Size: 0})
	cache.Insert("key", string(data))

	v, ok := cache.Get("key")
	if !ok || v != string(data) {
		t.Errorf("Cache has incorrect value: %v != %v", data, v)
	}
}

func testLRUEviction(t *testing.T) {
	Size := uint64(3)
	cache := LRUCache(Size)

	data, _ := json.Marshal(&cValue{Size: 1})

	cache.Insert("key1", string(data))
	cache.Insert("key2", string(data))
	cache.Insert("key3", string(data))
	// lru: [key3, key2, key1]

	// Look up the elements. This will rearrange the LRU ordering.
	cache.Get("key3")
	cache.Get("key2")
	cache.Get("key1")
	// lru: [key1, key2, key3]

	cache.Insert("key0", string(data))
	// lru: [key0, key1, key2]

	// The least recently used one should have been evicted.
	if _, ok := cache.Get("key3"); ok {
		t.Error("Least recently used element was not evicted.")
	}
}

/*
func testCapacityIsObeyed(t *testing.T) {
	Size := uint64(3)
	cache := LRUCache(Size)
    value,_ := json.Marshal(&cValue{Size: 1})

	// Insert up to the cache's capacity.
	cache.Insert("key1", string(value))
	cache.Insert("key2", string(value))
	cache.Insert("key3", string(value))
	if _, sz, _, _ := cache.CacheStats(); sz != Size {
		t.Errorf("cache.getSize() = %v, expected %v", sz, Size)
	}
	// Insert one more; something should be evicted to make room.
	cache.Insert("key4", string(value))
	if _, sz, _ := cache.CacheStats(); sz != Size {
		t.Errorf("post-evict cache.getSize() = %v, expected %v", sz, Size)
	}
}

func testInsertUpdatesgetSize(t *testing.T) {
	cache := LRUCache(100)
    emptyValue, _ := json.Marshal(&cValue{Size: 0})
	key := "key1"
	cache.Insert(key, string(emptyValue))
	if _, sz, _ := cache.CacheStats(); sz != 0 {
		t.Errorf("cache.getSize() = %v, expected 0", sz)
	}
    someValue, _ := json.Marshal(&cValue{Size: 20})
	key = "key2"
	cache.Insert(key, string(someValue))
	if _, sz, _ := cache.CacheStats(); sz != 20 {
		t.Errorf("cache.getSize() = %v, expected 20", sz)
	}
}
*/

func testInsertWithOldKeyUpdatesValue(t *testing.T) {
	cache := LRUCache(100)
	emptyValue, _ := json.Marshal(&cValue{Size: 0})
	key := "key1"
	cache.Insert(key, string(emptyValue))
	someValue, _ := json.Marshal(&cValue{Size: 20})
	cache.Insert(key, string(someValue))

	v, ok := cache.Get(key)
	if !ok || v != string(someValue) {
		t.Errorf("Cache has incorrect value: %v != %v", someValue, v)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"github.com/codegangsta/cli"
	"github.com/dustin/go-humanize"
	. "github.com/francoishill/golang-web-dry/errors/checkerror"
	"github.com/francoishill/golang-web-dry/zip/ziputils"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
	"os"
	"strings"
	"time"
)

type Logger interface {
	Debug(msg string, args ...interface{})
	Info(msg string, args ...interface{})
	Error(msg string, args ...interface{})
}

type defaultLogger struct {
	d *log.Logger
	i *log.Logger
	e *log.Logger
}

func (l *defaultLogger) Debug(msg string, args ...interface{}) {
	l.d.Println(fmt.Sprintf(msg, args...))
}

func (l *defaultLogger) Info(msg string, args ...interface{}) {
	l.i.Println(fmt.Sprintf(msg, args...))
}

func (l *defaultLogger) Error(msg string, args ...interface{}) {
	l.e.Println(fmt.Sprintf(msg, args...))
}

type cliExtendedContext struct {
	*cli.Context
}

func (c *cliExtendedContext) RequireGlobalString(flagName string) string {
	val := c.GlobalString(flagName)
	if strings.TrimSpace(val) == "" {
		panic("Flag '" + flagName + "' is empty")
	}
	return val
}

type timer struct {
	logger    Logger
	startTime time.Time
}

func (t *timer) printDuration() {
	duration := time.Now().Sub(t.startTime)
	t.logger.Debug("Duration %s", duration.String())
}

type appContext struct {
	logger Logger
}

func (a *appContext) getFileSize(file *os.File) int64 {
	fi, err := file.Stat()
	CheckError(err)
	return fi.Size()
}

func checkServerResponse(resp *http.Response) {
	if resp.StatusCode != http.StatusOK {
		b, e := ioutil.ReadAll(resp.Body)
		if e != nil {
			panic(fmt.Sprintf("The server returned status code %d but could not read response body. Error: %s", e.Error()))
		}
		panic(fmt.Sprintf("Server status code %d with response %s", resp.StatusCode, string(b)))
	}
}

func (a *appContext) download(serverUrl, localPath, remotePath, dirFileFilterPattern string) {
	defer (&timer{a.logger, time.Now()}).printDuration()

	var fileFilterQueryPart = ""
	if dirFileFilterPattern != "" {
		fileFilterQueryPart = "&filefilter=" + url.QueryEscape(dirFileFilterPattern)
	}

	resp, err := http.Get(serverUrl + "?path=" + url.QueryEscape(remotePath) + fileFilterQueryPart)
	CheckError(err)
	defer resp.Body.Close()

	checkServerResponse(resp)

	ziputils.SaveTarReaderToPath(a.logger, resp.Body, localPath)
}

func (a *appContext) uploadFile(serverUrl, localPath, remotePath string) {
	defer (&timer{a.logger, time.Now()}).printDuration()

	file, err := os.OpenFile(localPath, 0, 0600)
	CheckError(err)

	a.logger.Debug("Now starting to upload local file '%s' of size %s to remote path '%s'", localPath, humanize.IBytes(uint64(a.getFileSize(file))), remotePath)
	resp, err := http.Post(serverUrl+"?path="+url.QueryEscape(remotePath), "application/octet-stream", file)
	CheckError(err)

	checkServerResponse(resp)
}

func (a *appContext) uploadDirectory(serverUrl, localPath, remotePath, dirFileFilterPattern string) {
	defer (&timer{a.logger, time.Now()}).printDuration()

	a.logger.Debug("Now starting to upload local directory '%s' to remote '%s", localPath, remotePath)
	checkResponseFunc := checkServerResponse
	walkContext := ziputils.NewDirWalkContext(dirFileFilterPattern)
	ziputils.UploadDirectoryToUrl(a.logger, serverUrl+"?dir="+url.QueryEscape(remotePath), "application/octet-stream", localPath, walkContext, checkResponseFunc)
}

func (a *appContext) isDir(path string) bool {
	p, err := os.Open(path)
	CheckError(err)
	defer p.Close()
	info, err := p.Stat()
	CheckError(err)
	return info.IsDir()
}

func (a *appContext) upload(serverUrl, localPath, remotePath, dirFileFilterPattern string) {
	if a.isDir(localPath) {
		a.uploadDirectory(serverUrl, localPath, remotePath, dirFileFilterPattern)
	} else {
		a.uploadFile(serverUrl, localPath, remotePath)
	}
}

func (a *appContext) delete(serverUrl, remotePath, dirFileFilterPattern string) {
	defer (&timer{a.logger, time.Now()}).printDuration()

	var fileFilterQueryPart = ""
	if dirFileFilterPattern != "" {
		fileFilterQueryPart = "&filefilter=" + url.QueryEscape(dirFileFilterPattern)
	}

	req, err := http.NewRequest("DELETE", serverUrl+"?path="+url.QueryEscape(remotePath)+fileFilterQueryPart, nil)
	CheckError(err)

	resp, err := http.DefaultClient.Do(req)
	CheckError(err)

	checkServerResponse(resp)
}

func (a *appContext) writeStatsToStdout(serverUrl, remotePath string) {
	defer (&timer{a.logger, time.Now()}).printDuration()

	resp, err := http.Head(serverUrl + "?path=" + url.QueryEscape(remotePath))
	CheckError(err)
	defer resp.Body.Close()

	checkServerResponse(resp)

	if exists := resp.Header.Get("EXISTS"); exists == "" {
		panic("Could not find 'EXISTS' header")
	} else {
		a.logger.Info("STATS_EXISTS=%s", exists)
		if exists == "0" {
			return
		}
	}

	if isDir := resp.Header.Get("IS_DIR"); isDir == "" {
		panic("Could not find 'IS_DIR' header")
	} else {
		a.logger.Info("STATS_IS_DIR=%s", isDir)
	}
}

func (a *appContext) move(serverUrl, remotePath, newPath string) {
	defer (&timer{a.logger, time.Now()}).printDuration()

	escapedOldPath := url.QueryEscape(remotePath)
	escapedNewPath := url.QueryEscape(newPath)
	req, err := http.NewRequest("PUT", serverUrl+"?action=move&path="+escapedOldPath+"&newpath="+escapedNewPath, nil)
	CheckError(err)

	resp, err := http.DefaultClient.Do(req)
	CheckError(err)
	checkServerResponse(resp)
}

func (a *appContext) MainAction(c *cli.Context) {
	c2 := &cliExtendedContext{c}

	mode := c2.RequireGlobalString("mode")

	serverUrl := c2.RequireGlobalString("serverurl")
	remotePath := c2.RequireGlobalString("remotepath")

	switch mode {
	case "DOWNLOAD":
		localPath := c2.RequireGlobalString("localpath")
		dirFileFilterPattern := c.GlobalString("filefilter") //Not required
		a.download(serverUrl, localPath, remotePath, dirFileFilterPattern)
		break
	case "UPLOAD":
		localPath := c2.RequireGlobalString("localpath")
		dirFileFilterPattern := c.GlobalString("filefilter") //Not required
		a.upload(serverUrl, localPath, remotePath, dirFileFilterPattern)
		break
	case "DELETE":
		dirFileFilterPattern := c.GlobalString("filefilter") //Not required
		a.delete(serverUrl, remotePath, dirFileFilterPattern)
		break
	case "STATS":
		a.writeStatsToStdout(serverUrl, remotePath)
		break
	case "MOVE":
		newPath := c.GlobalString("newpath") //Not required
		a.move(serverUrl, remotePath, newPath)
		break
	default:
		panic("Unknown mode '" + mode + "'")
	}
}

func main() {
	logger := &defaultLogger{
		d: log.New(os.Stdout, "[D] ", log.Ldate|log.Ltime|log.Lshortfile),
		i: log.New(os.Stdout, "[I] ", log.Ldate|log.Ltime|log.Lshortfile),
		e: log.New(os.Stderr, "[E] ", log.Ldate|log.Ltime|log.Lshortfile),
	}
	context := &appContext{logger}

	defer func() {
		if r := recover(); r != nil {
			logger.Error("%+v", r)
		}
	}()

	app := cli.NewApp()
	app.Name = "copyclient"
	app.Usage = "A http client to copy files to a server"
	app.Action = context.MainAction
	app.Flags = []cli.Flag{
		cli.StringFlag{
			Name:  "mode,m",
			Value: "",
			Usage: "The mode of the action (UPLOADFILE, DOWNLOADFILE, UPLOADFOLDER, DOWNLOADFOLDER)",
		},
		cli.StringFlag{
			Name:  "serverurl,s",
			Value: "",
			Usage: "Url to the file server",
		},
		cli.StringFlag{
			Name:  "localpath,l",
			Value: "",
			Usage: "The full absolute LOCAL path (file/folder)",
		},
		cli.StringFlag{
			Name:  "remotepath,r",
			Value: "",
			Usage: "The full absolute REMOTE path (file/folder)",
		},
		cli.StringFlag{
			Name:  "filefilter,ff",
			Value: "",
			Usage: "The golang filepath filter pattern (for file base name), see http://golang.org/pkg/path/filepath/#Match",
		},
		cli.StringFlag{
			Name:  "newpath,np",
			Value: "",
			Usage: "The new path, this is only currently applicable to the 'MOVE' method.",
		},
	}
	app.Run(os.Args)
}
<commit_msg>Use `os.Exit` on panic recovery<commit_after>package main

import (
	"fmt"
	"github.com/codegangsta/cli"
	"github.com/dustin/go-humanize"
	. "github.com/francoishill/golang-web-dry/errors/checkerror"
	"github.com/francoishill/golang-web-dry/zip/ziputils"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
	"os"
	"strings"
	"time"
)

type Logger interface {
	Debug(msg string, args ...interface{})
	Info(msg string, args ...interface{})
	Error(msg string, args ...interface{})
}

type defaultLogger struct {
	d *log.Logger
	i *log.Logger
	e *log.Logger
}

func (l *defaultLogger) Debug(msg string, args ...interface{}) {
	l.d.Println(fmt.Sprintf(msg, args...))
}

func (l *defaultLogger) Info(msg string, args ...interface{}) {
	l.i.Println(fmt.Sprintf(msg, args...))
}

func (l *defaultLogger) Error(msg string, args ...interface{}) {
	l.e.Println(fmt.Sprintf(msg, args...))
}

type cliExtendedContext struct {
	*cli.Context
}

func (c *cliExtendedContext) RequireGlobalString(flagName string) string {
	val := c.GlobalString(flagName)
	if strings.TrimSpace(val) == "" {
		panic("Flag '" + flagName + "' is empty")
	}
	return val
}

type timer struct {
	logger    Logger
	startTime time.Time
}

func (t *timer) printDuration() {
	duration := time.Now().Sub(t.startTime)
	t.logger.Debug("Duration %s", duration.String())
}

type appContext struct {
	logger Logger
}

func (a *appContext) getFileSize(file *os.File) int64 {
	fi, err := file.Stat()
	CheckError(err)
	return fi.Size()
}

func checkServerResponse(resp *http.Response) {
	if resp.StatusCode != http.StatusOK {
		b, e := ioutil.ReadAll(resp.Body)
		if e != nil {
			panic(fmt.Sprintf("The server returned status code %d but could not read response body. Error: %s", e.Error()))
		}
		panic(fmt.Sprintf("Server status code %d with response %s", resp.StatusCode, string(b)))
	}
}

func (a *appContext) download(serverUrl, localPath, remotePath, dirFileFilterPattern string) {
	defer (&timer{a.logger, time.Now()}).printDuration()

	var fileFilterQueryPart = ""
	if dirFileFilterPattern != "" {
		fileFilterQueryPart = "&filefilter=" + url.QueryEscape(dirFileFilterPattern)
	}

	resp, err := http.Get(serverUrl + "?path=" + url.QueryEscape(remotePath) + fileFilterQueryPart)
	CheckError(err)
	defer resp.Body.Close()

	checkServerResponse(resp)

	ziputils.SaveTarReaderToPath(a.logger, resp.Body, localPath)
}

func (a *appContext) uploadFile(serverUrl, localPath, remotePath string) {
	defer (&timer{a.logger, time.Now()}).printDuration()

	file, err := os.OpenFile(localPath, 0, 0600)
	CheckError(err)

	a.logger.Debug("Now starting to upload local file '%s' of size %s to remote path '%s'", localPath, humanize.IBytes(uint64(a.getFileSize(file))), remotePath)
	resp, err := http.Post(serverUrl+"?path="+url.QueryEscape(remotePath), "application/octet-stream", file)
	CheckError(err)

	checkServerResponse(resp)
}

func (a *appContext) uploadDirectory(serverUrl, localPath, remotePath, dirFileFilterPattern string) {
	defer (&timer{a.logger, time.Now()}).printDuration()

	a.logger.Debug("Now starting to upload local directory '%s' to remote '%s", localPath, remotePath)
	checkResponseFunc := checkServerResponse
	walkContext := ziputils.NewDirWalkContext(dirFileFilterPattern)
	ziputils.UploadDirectoryToUrl(a.logger, serverUrl+"?dir="+url.QueryEscape(remotePath), "application/octet-stream", localPath, walkContext, checkResponseFunc)
}

func (a *appContext) isDir(path string) bool {
	p, err := os.Open(path)
	CheckError(err)
	defer p.Close()
	info, err := p.Stat()
	CheckError(err)
	return info.IsDir()
}

func (a *appContext) upload(serverUrl, localPath, remotePath, dirFileFilterPattern string) {
	if a.isDir(localPath) {
		a.uploadDirectory(serverUrl, localPath, remotePath, dirFileFilterPattern)
	} else {
		a.uploadFile(serverUrl, localPath, remotePath)
	}
}

func (a *appContext) delete(serverUrl, remotePath, dirFileFilterPattern string) {
	defer (&timer{a.logger, time.Now()}).printDuration()

	var fileFilterQueryPart = ""
	if dirFileFilterPattern != "" {
		fileFilterQueryPart = "&filefilter=" + url.QueryEscape(dirFileFilterPattern)
	}

	req, err := http.NewRequest("DELETE", serverUrl+"?path="+url.QueryEscape(remotePath)+fileFilterQueryPart, nil)
	CheckError(err)

	resp, err := http.DefaultClient.Do(req)
	CheckError(err)

	checkServerResponse(resp)
}

func (a *appContext) writeStatsToStdout(serverUrl, remotePath string) {
	defer (&timer{a.logger, time.Now()}).printDuration()

	resp, err := http.Head(serverUrl + "?path=" + url.QueryEscape(remotePath))
	CheckError(err)
	defer resp.Body.Close()

	checkServerResponse(resp)

	if exists := resp.Header.Get("EXISTS"); exists == "" {
		panic("Could not find 'EXISTS' header")
	} else {
		a.logger.Info("STATS_EXISTS=%s", exists)
		if exists == "0" {
			return
		}
	}

	if isDir := resp.Header.Get("IS_DIR"); isDir == "" {
		panic("Could not find 'IS_DIR' header")
	} else {
		a.logger.Info("STATS_IS_DIR=%s", isDir)
	}
}

func (a *appContext) move(serverUrl, remotePath, newPath string) {
	defer (&timer{a.logger, time.Now()}).printDuration()

	escapedOldPath := url.QueryEscape(remotePath)
	escapedNewPath := url.QueryEscape(newPath)
	req, err := http.NewRequest("PUT", serverUrl+"?action=move&path="+escapedOldPath+"&newpath="+escapedNewPath, nil)
	CheckError(err)

	resp, err := http.DefaultClient.Do(req)
	CheckError(err)
	checkServerResponse(resp)
}

func (a *appContext) MainAction(c *cli.Context) {
	c2 := &cliExtendedContext{c}

	mode := c2.RequireGlobalString("mode")

	serverUrl := c2.RequireGlobalString("serverurl")
	remotePath := c2.RequireGlobalString("remotepath")

	switch mode {
	case "DOWNLOAD":
		localPath := c2.RequireGlobalString("localpath")
		dirFileFilterPattern := c.GlobalString("filefilter") //Not required
		a.download(serverUrl, localPath, remotePath, dirFileFilterPattern)
		break
	case "UPLOAD":
		localPath := c2.RequireGlobalString("localpath")
		dirFileFilterPattern := c.GlobalString("filefilter") //Not required
		a.upload(serverUrl, localPath, remotePath, dirFileFilterPattern)
		break
	case "DELETE":
		dirFileFilterPattern := c.GlobalString("filefilter") //Not required
		a.delete(serverUrl, remotePath, dirFileFilterPattern)
		break
	case "STATS":
		a.writeStatsToStdout(serverUrl, remotePath)
		break
	case "MOVE":
		newPath := c.GlobalString("newpath") //Not required
		a.move(serverUrl, remotePath, newPath)
		break
	default:
		panic("Unknown mode '" + mode + "'")
	}
}

func main() {
	logger := &defaultLogger{
		d: log.New(os.Stdout, "[D] ", log.Ldate|log.Ltime|log.Lshortfile),
		i: log.New(os.Stdout, "[I] ", log.Ldate|log.Ltime|log.Lshortfile),
		e: log.New(os.Stderr, "[E] ", log.Ldate|log.Ltime|log.Lshortfile),
	}
	context := &appContext{logger}

	defer func() {
		if r := recover(); r != nil {
			logger.Error("%+v", r)
			os.Exit(2)
		}
	}()

	app := cli.NewApp()
	app.Name = "copyclient"
	app.Usage = "A http client to copy files to a server"
	app.Action = context.MainAction
	app.Flags = []cli.Flag{
		cli.StringFlag{
			Name:  "mode,m",
			Value: "",
			Usage: "The mode of the action (UPLOADFILE, DOWNLOADFILE, UPLOADFOLDER, DOWNLOADFOLDER)",
		},
		cli.StringFlag{
			Name:  "serverurl,s",
			Value: "",
			Usage: "Url to the file server",
		},
		cli.StringFlag{
			Name:  "localpath,l",
			Value: "",
			Usage: "The full absolute LOCAL path (file/folder)",
		},
		cli.StringFlag{
			Name:  "remotepath,r",
			Value: "",
			Usage: "The full absolute REMOTE path (file/folder)",
		},
		cli.StringFlag{
			Name:  "filefilter,ff",
			Value: "",
			Usage: "The golang filepath filter pattern (for file base name), see http://golang.org/pkg/path/filepath/#Match",
		},
		cli.StringFlag{
			Name:  "newpath,np",
			Value: "",
			Usage: "The new path, this is only currently applicable to the 'MOVE' method.",
		},
	}
	app.Run(os.Args)
}
<|endoftext|>


<commit_before>package main

import (
	"log"
	"net/http"

	"github.com/boltdb/bolt"
	"github.com/gorilla/mux"
)

func main() {
	r := mux.NewRouter()
	apiR := r.PathPrefix("/api/").Subrouter()

	db, err := bolt.Open("records.bolt", 0600, &bolt.Options{Timeout: 1})
	if err != nil {
		log.Fatalf("Failed to open bolt database, err: %s", err)
	}

	gprdb, err := NewPreRegBoltDb(db)
	if err != nil {
		log.Fatalf("Failed to get group preregistration database started, err %s", err)
	}

	ces := NewConfirmationEmailService("registration.cubjamboree.ca", "no-reply@cubjamboree.ca", "info@cubjamboree.ca", NewLocalMailder("localhost:25"), gprdb)

	NewGroupPreRegistrationHandler(apiR, gprdb, ces)

	http.Handle("/api/", r)
	otherFiles := http.FileServer(http.Dir("../app"))
	http.Handle("/app.css", otherFiles)
	http.Handle("/app.js", otherFiles)
	http.Handle("/components/", otherFiles)
	http.Handle("/views/", otherFiles)
	http.Handle("/bower_components/", otherFiles)
	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		http.ServeFile(w, r, "../app/index.html")
	})
	panic(http.ListenAndServe(":8080", nil))
}
<commit_msg>Add quick and dirty request logger.<commit_after>package main

import (
	"log"
	"net/http"
	"time"

	"github.com/boltdb/bolt"
	"github.com/gorilla/mux"
)

type requestLogger struct {
	H http.Handler
}

type wWrapper struct {
	code int
	http.ResponseWriter
}

func (w *wWrapper) WriteHeader(code int) {
	w.code = code
	w.ResponseWriter.WriteHeader(code)
}

func (h *requestLogger) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	start := time.Now()
	wrappedW := &wWrapper{
		ResponseWriter: w,
		code: http.StatusOK, // Default code
	}
	h.H.ServeHTTP(wrappedW, r)
	duration := time.Now().Sub(start)
	log.Printf("Handled request for url %s, code %v, took %s seconds", r.URL, wrappedW.code, duration)
}

func main() {
	r := mux.NewRouter()
	apiR := r.PathPrefix("/api/").Subrouter()

	db, err := bolt.Open("records.bolt", 0600, &bolt.Options{Timeout: 1})
	if err != nil {
		log.Fatalf("Failed to open bolt database, err: %s", err)
	}

	gprdb, err := NewPreRegBoltDb(db)
	if err != nil {
		log.Fatalf("Failed to get group preregistration database started, err %s", err)
	}

	ces := NewConfirmationEmailService("registration.cubjamboree.ca", "no-reply@cubjamboree.ca", "info@cubjamboree.ca", NewLocalMailder("localhost:25"), gprdb)

	NewGroupPreRegistrationHandler(apiR, gprdb, ces)

	http.Handle("/api/", r)
	otherFiles := http.FileServer(http.Dir("../app"))
	http.Handle("/app.css", otherFiles)
	http.Handle("/app.js", otherFiles)
	http.Handle("/components/", otherFiles)
	http.Handle("/views/", otherFiles)
	http.Handle("/bower_components/", otherFiles)
	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		http.ServeFile(w, r, "../app/index.html")
	})
	panic(http.ListenAndServe(":8080", &requestLogger{http.DefaultServeMux}))
}
<|endoftext|>


<commit_before>package relay

import (
	"errors"
	"fmt"
	"io"
	"os"
	"path"
	"regexp"
	"strings"

	"github.com/fsouza/go-dockerclient"
)

// ErrInvalidRepo is returned by parseRepo when the repo is not a valid repo.
var ErrInvalidRepo = errors.New("registry: not a valid docker repo")

// RepoPattern matches strings like `quay.io/remind101/acme-inc:latest` or `remind101/acme-inc:latest`
var RepoPattern = regexp.MustCompile(`(\S+\/)?(\S+\/\S+):(\S+)`)

// ContainerManager defines an interface for managing containers.
type ContainerManager interface {
	Pull(*Container) error
	Create(*Container) error
	Attach(string, io.Reader, io.Writer) error
	Start(string) error
	Wait(string) (int, error)
	Stop(string) error
	Remove(string) error
}

// newContainerManager returns a ContainerManager based on the given options.
func newContainerManager(options DockerOptions) (manager ContainerManager) {
	var err error

	if options.Socket == "fake" {
		manager = &fakeManager{}
	} else {
		manager, err = NewDockerManager(options.Socket, options.CertPath, options.Auth)
		if err != nil {
			panic(err)
		}
	}
	return manager
}

type fakeManager struct {
}

func (f *fakeManager) Pull(c *Container) error {
	return nil
}

func (f *fakeManager) Create(c *Container) error {
	return nil
}

func (f *fakeManager) Attach(name string, input io.Reader, output io.Writer) error {
	return nil
}

func (f *fakeManager) Start(name string) error {
	return nil
}

func (f *fakeManager) Wait(name string) (int, error) {
	return 0, nil
}

func (f *fakeManager) Stop(name string) error {
	return nil
}

func (f *fakeManager) Remove(name string) error {
	return nil
}

type dockerManager struct {
	client *docker.Client
	auth   *docker.AuthConfigurations
}

func NewDockerManager(socket, certPath string, auth *docker.AuthConfigurations) (*dockerManager, error) {
	var err error
	var dc *docker.Client

	switch {
	case certPath != "":
		cert := path.Join(certPath, "cert.pem")
		key := path.Join(certPath, "key.pem")
		ca := ""
		dc, err = docker.NewTLSClient(socket, cert, key, ca)
	case socket != "":
		dc, err = docker.NewClient(socket)
	default:
		return nil, errors.New("NewDockerManager needs a socket or a certPath")
	}

	if err != nil {
		return nil, err
	}

	return &dockerManager{
		client: dc,
		auth:   auth,
	}, nil
}

func (d *dockerManager) Pull(c *Container) error {
	return d.pullImage(c.Image)
}

func (d *dockerManager) Create(c *Container) error {
	env := []string{}
	for k, v := range c.Env {
		env = append(env, fmt.Sprintf("%s=%s", k, v))
	}

	opts := docker.CreateContainerOptions{
		Name: c.Name,
		Config: &docker.Config{
			Tty:          c.Attach,
			AttachStdin:  c.Attach,
			AttachStdout: c.Attach,
			AttachStderr: c.Attach,
			OpenStdin:    c.Attach,

			Image: c.Image,
			Cmd:   []string{c.Command},
			Env:   env,
		},
		HostConfig: &docker.HostConfig{},
	}
	_, err := d.client.CreateContainer(opts)
	return err
}

func (d *dockerManager) Attach(name string, input io.Reader, output io.Writer) error {
	opts := docker.AttachToContainerOptions{
		Container:    name,
		InputStream:  input,
		OutputStream: output,
		ErrorStream:  output,
		Logs:         true,
		Stream:       true,
		Stdin:        true,
		Stdout:       true,
		Stderr:       true,
		RawTerminal:  true,
	}

	return d.client.AttachToContainer(opts)
}

func (d *dockerManager) Start(name string) error {
	return d.client.StartContainer(name, nil)
}

func (d *dockerManager) Wait(name string) (int, error) {
	return d.client.WaitContainer(name)
}

func (d *dockerManager) Stop(name string) error {
	return d.client.StopContainer(name, 10)
}

func (d *dockerManager) Remove(name string) error {
	return d.client.RemoveContainer(docker.RemoveContainerOptions{
		ID:    name,
		Force: true,
	})
}

func (d *dockerManager) pullImage(image string) error {
	var a docker.AuthConfiguration

	reg, repo, tag, err := parseRepo(image)
	if err != nil {
		return err
	}

	if reg == "" {
		reg = "https://index.docker.io/v1/"
	}

	if c, ok := d.auth.Configs[reg]; ok {
		a = c
	}

	return d.client.PullImage(docker.PullImageOptions{
		Repository:   repo,
		Tag:          tag,
		OutputStream: os.Stdout,
	}, a)
}

// parseRepo splits a full docker repo into registry, repo and tag segments.
//
// Examples:
//
//     quay.io/remind101/acme-inc:latest # => registry: "quay.io", repo: "remind101/acme-inc", tag: "latest"
//     remind101/acme-inc:latest         # => registry: "", repo: "remind101/acme-inc", tag: "latest"
func parseRepo(fullRepo string) (registry string, repo string, tag string, err error) {
	m := RepoPattern.FindStringSubmatch(fullRepo)
	if len(m) == 0 {
		return "", "", "", ErrInvalidRepo
	}

	// Registy subpattern was matched.
	if len(m) == 4 {
		return strings.TrimRight(m[1], "/"), m[2], m[3], nil
	}

	return "", m[1], m[2], nil
}
<commit_msg>Fix hk run with commands with spaces<commit_after>package relay

import (
	"errors"
	"fmt"
	"io"
	"os"
	"path"
	"regexp"
	"strings"

	"github.com/fsouza/go-dockerclient"
)

// ErrInvalidRepo is returned by parseRepo when the repo is not a valid repo.
var ErrInvalidRepo = errors.New("registry: not a valid docker repo")

// RepoPattern matches strings like `quay.io/remind101/acme-inc:latest` or `remind101/acme-inc:latest`
var RepoPattern = regexp.MustCompile(`(\S+\/)?(\S+\/\S+):(\S+)`)

// ContainerManager defines an interface for managing containers.
type ContainerManager interface {
	Pull(*Container) error
	Create(*Container) error
	Attach(string, io.Reader, io.Writer) error
	Start(string) error
	Wait(string) (int, error)
	Stop(string) error
	Remove(string) error
}

// newContainerManager returns a ContainerManager based on the given options.
func newContainerManager(options DockerOptions) (manager ContainerManager) {
	var err error

	if options.Socket == "fake" {
		manager = &fakeManager{}
	} else {
		manager, err = NewDockerManager(options.Socket, options.CertPath, options.Auth)
		if err != nil {
			panic(err)
		}
	}
	return manager
}

type fakeManager struct {
}

func (f *fakeManager) Pull(c *Container) error {
	return nil
}

func (f *fakeManager) Create(c *Container) error {
	return nil
}

func (f *fakeManager) Attach(name string, input io.Reader, output io.Writer) error {
	return nil
}

func (f *fakeManager) Start(name string) error {
	return nil
}

func (f *fakeManager) Wait(name string) (int, error) {
	return 0, nil
}

func (f *fakeManager) Stop(name string) error {
	return nil
}

func (f *fakeManager) Remove(name string) error {
	return nil
}

type dockerManager struct {
	client *docker.Client
	auth   *docker.AuthConfigurations
}

func NewDockerManager(socket, certPath string, auth *docker.AuthConfigurations) (*dockerManager, error) {
	var err error
	var dc *docker.Client

	switch {
	case certPath != "":
		cert := path.Join(certPath, "cert.pem")
		key := path.Join(certPath, "key.pem")
		ca := ""
		dc, err = docker.NewTLSClient(socket, cert, key, ca)
	case socket != "":
		dc, err = docker.NewClient(socket)
	default:
		return nil, errors.New("NewDockerManager needs a socket or a certPath")
	}

	if err != nil {
		return nil, err
	}

	return &dockerManager{
		client: dc,
		auth:   auth,
	}, nil
}

func (d *dockerManager) Pull(c *Container) error {
	return d.pullImage(c.Image)
}

func (d *dockerManager) Create(c *Container) error {
	env := []string{}
	for k, v := range c.Env {
		env = append(env, fmt.Sprintf("%s=%s", k, v))
	}

	opts := docker.CreateContainerOptions{
		Name: c.Name,
		Config: &docker.Config{
			Tty:          c.Attach,
			AttachStdin:  c.Attach,
			AttachStdout: c.Attach,
			AttachStderr: c.Attach,
			OpenStdin:    c.Attach,

			Image: c.Image,
			Cmd:   strings.Split(c.Command, " "),
			Env:   env,
		},
		HostConfig: &docker.HostConfig{},
	}
	_, err := d.client.CreateContainer(opts)
	return err
}

func (d *dockerManager) Attach(name string, input io.Reader, output io.Writer) error {
	opts := docker.AttachToContainerOptions{
		Container:    name,
		InputStream:  input,
		OutputStream: output,
		ErrorStream:  output,
		Logs:         true,
		Stream:       true,
		Stdin:        true,
		Stdout:       true,
		Stderr:       true,
		RawTerminal:  true,
	}

	return d.client.AttachToContainer(opts)
}

func (d *dockerManager) Start(name string) error {
	return d.client.StartContainer(name, nil)
}

func (d *dockerManager) Wait(name string) (int, error) {
	return d.client.WaitContainer(name)
}

func (d *dockerManager) Stop(name string) error {
	return d.client.StopContainer(name, 10)
}

func (d *dockerManager) Remove(name string) error {
	return d.client.RemoveContainer(docker.RemoveContainerOptions{
		ID:    name,
		Force: true,
	})
}

func (d *dockerManager) pullImage(image string) error {
	var a docker.AuthConfiguration

	reg, repo, tag, err := parseRepo(image)
	if err != nil {
		return err
	}

	if reg == "" {
		reg = "https://index.docker.io/v1/"
	}

	if c, ok := d.auth.Configs[reg]; ok {
		a = c
	}

	return d.client.PullImage(docker.PullImageOptions{
		Repository:   repo,
		Tag:          tag,
		OutputStream: os.Stdout,
	}, a)
}

// parseRepo splits a full docker repo into registry, repo and tag segments.
//
// Examples:
//
//     quay.io/remind101/acme-inc:latest # => registry: "quay.io", repo: "remind101/acme-inc", tag: "latest"
//     remind101/acme-inc:latest         # => registry: "", repo: "remind101/acme-inc", tag: "latest"
func parseRepo(fullRepo string) (registry string, repo string, tag string, err error) {
	m := RepoPattern.FindStringSubmatch(fullRepo)
	if len(m) == 0 {
		return "", "", "", ErrInvalidRepo
	}

	// Registy subpattern was matched.
	if len(m) == 4 {
		return strings.TrimRight(m[1], "/"), m[2], m[3], nil
	}

	return "", m[1], m[2], nil
}
<|endoftext|>


<commit_before>package archive

import (
	"fmt"
	"io"
	"os"
	"path"
	"time"

	"golang.org/x/crypto/openpgp"

	"pault.ag/go/blobstore"
	"pault.ag/go/debian/control"
	"pault.ag/go/debian/dependency"
	"pault.ag/go/debian/transput"
)

// New {{{

func New(path string) (*Archive, error) {
	store, err := blobstore.Load(path)
	if err != nil {
		return nil, err
	}

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

	return &Archive{
		store: *store,
	}, nil
}

func NewWritable(path string, keyring string, keyid uint64) (*Archive, error) {
	archive, err := New(path)
	if err != nil {
		return nil, err
	}

	fd, err := os.Open(keyring)
	if err != nil {
		return nil, err
	}

	el, err := openpgp.ReadKeyRing(fd)
	if err != nil {
		return nil, err
	}

	keys := el.KeysById(keyid)

	if len(keys) == 0 {
		return nil, fmt.Errorf("No keys matched that key ID")
	}

	if len(keys) != 1 {
		return nil, fmt.Errorf("Too many keys matched that key ID")
	}

	archive.signingKey = keys[0].Entity

	return archive, err
}

// }}}

// Archive magic {{{

type Archive struct {
	store      blobstore.Store
	signingKey *openpgp.Entity
}

func (a Archive) Suite(name string) (*Suite, error) {
	/* Get the Release / InRelease */
	inRelease := Release{}
	components := map[string]*Component{}

	fd, err := a.store.OpenPath(path.Join("dists", name, "InRelease"))
	if err == nil {
		defer fd.Close()
		if err := control.Unmarshal(&inRelease, fd); err != nil {
			return nil, err
		}

		for _, name := range inRelease.Components {
			components[name] = &Component{Packages: []Package{}}
		}
	}

	suite := Suite{
		Name: name,

		release:    inRelease,
		Components: components,
	}

	suite.features.Hashes = []string{"sha256", "sha1"}

	return &suite, nil
}

func (a Archive) encodeHashedBySuite(path string, suite Suite, data interface{}) (*blobstore.Object, []control.FileHash, error) {

	hashers := []*transput.Hasher{}
	for _, algorithm := range suite.features.Hashes {
		hasher, err := transput.NewHasher(algorithm)
		if err != nil {
			return nil, nil, err
		}
		hashers = append(hashers, hasher)
	}

	return a.encodeHashed(path, hashers, data)
}

func (a Archive) encodeHashed(path string, hashers []*transput.Hasher, data interface{}) (*blobstore.Object, []control.FileHash, error) {

	writers := []io.Writer{}
	for _, hasher := range hashers {
		writers = append(writers, hasher)
	}

	obj, err := a.encode(data, io.MultiWriter(writers...))
	if err != nil {
		return nil, nil, err
	}

	fileHashs := []control.FileHash{}
	for _, hasher := range hashers {
		fileHashs = append(fileHashs, control.FileHashFromHasher(path, *hasher))
	}

	return obj, fileHashs, nil
}

func (a Archive) encode(data interface{}, tee io.Writer) (*blobstore.Object, error) {
	writer, err := a.store.Create()
	if err != nil {
		return nil, err
	}
	defer writer.Close()

	var target io.Writer = writer
	if tee != nil {
		target = io.MultiWriter(writer, tee)
	}

	encoder, err := control.NewEncoder(target)
	if err != nil {
		return nil, err
	}

	if err := encoder.Encode(data); err != nil {
		return nil, err
	}

	obj, err := a.store.Commit(*writer)
	if err != nil {
		return nil, err
	}

	return obj, nil
}

func (a Archive) Engross(suite Suite) (map[string]blobstore.Object, error) {
	files := map[string]blobstore.Object{}

	release := Release{
		Description:   "",
		Origin:        "",
		Label:         "",
		Version:       "",
		Suite:         suite.Name,
		Codename:      "",
		Components:    suite.ComponenetNames(),
		Architectures: suite.Arches(),
		Date:          time.Now().Format(time.RFC1123Z),
		SHA256:        []control.SHA256FileHash{},
		SHA1:          []control.SHA1FileHash{},
		SHA512:        []control.SHA512FileHash{},
		MD5Sum:        []control.MD5FileHash{},
	}

	for name, component := range suite.Components {
		for arch, pkgs := range component.ByArch() {
			filePath := path.Join("dists", suite.Name, name,
				fmt.Sprintf("binary-%s", arch), "Packages")

			obj, hashes, err := a.encodeHashedBySuite(filePath, suite, pkgs)
			if err != nil {
				return nil, err
			}

			for _, hash := range hashes {
				if err := release.AddHash(hash); err != nil {
					return nil, err
				}
			}

			files[filePath] = *obj
		}
	}

	filePath := path.Join("dists", suite.Name, "Release")
	obj, sig, err := a.encodeSigned(release)
	if err != nil {
		return nil, err
	}
	files[filePath] = *obj
	files[fmt.Sprintf("%s.gpg", filePath)] = *sig

	return files, nil
}

func (a Archive) Link(blobs map[string]blobstore.Object) error {
	for p, obj := range blobs {
		if err := a.store.Link(obj, p); err != nil {
			return err
		}
	}
	return nil
}

func (a Archive) Decruft() error {
	return a.store.GC(blobstore.DumbGarbageCollector{})
}

// }}}

// Suite magic {{{

type Suite struct {
	Name string

	release    Release
	Components map[string]*Component

	features struct {
		Hashes []string
	}
}

func (s Suite) Arches() []dependency.Arch {
	ret := map[dependency.Arch]bool{}
	for _, component := range s.Components {
		for _, arch := range component.Arches() {
			ret[arch] = true
		}
	}
	r := []dependency.Arch{}
	for arch, _ := range ret {
		r = append(r, arch)
	}
	return r
}

func (s Suite) ComponenetNames() []string {
	ret := []string{}
	for name, _ := range s.Components {
		ret = append(ret, name)
	}
	return ret
}

func (s Suite) Add(name string, pkg Package) {
	if _, ok := s.Components[name]; !ok {
		s.Components[name] = &Component{Packages: []Package{}}
	}
	s.Components[name].Add(pkg)
}

// }}}

// Component magic {{{

type Component struct {
	Packages []Package
}

func (c *Component) ByArch() map[dependency.Arch][]Package {
	ret := map[dependency.Arch][]Package{}

	for _, pkg := range c.Packages {
		packages := ret[pkg.Architecture]
		ret[pkg.Architecture] = append(packages, pkg)
	}

	return ret
}

func (c *Component) Arches() []dependency.Arch {
	ret := []dependency.Arch{}
	for _, pkg := range c.Packages {
		ret = append(ret, pkg.Architecture)
	}
	return ret
}

func (c *Component) Add(p Package) {
	c.Packages = append(c.Packages, p)
}

// }}}

// vim: foldmethod=marker
<commit_msg>update paths<commit_after>package archive

import (
	"fmt"
	"io"
	"os"
	"path"
	"time"

	"golang.org/x/crypto/openpgp"

	"pault.ag/go/blobstore"
	"pault.ag/go/debian/control"
	"pault.ag/go/debian/dependency"
	"pault.ag/go/debian/transput"
)

// New {{{

func New(path string) (*Archive, error) {
	store, err := blobstore.Load(path)
	if err != nil {
		return nil, err
	}

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

	return &Archive{
		store: *store,
	}, nil
}

func NewWritable(path string, keyring string, keyid uint64) (*Archive, error) {
	archive, err := New(path)
	if err != nil {
		return nil, err
	}

	fd, err := os.Open(keyring)
	if err != nil {
		return nil, err
	}

	el, err := openpgp.ReadKeyRing(fd)
	if err != nil {
		return nil, err
	}

	keys := el.KeysById(keyid)

	if len(keys) == 0 {
		return nil, fmt.Errorf("No keys matched that key ID")
	}

	if len(keys) != 1 {
		return nil, fmt.Errorf("Too many keys matched that key ID")
	}

	archive.signingKey = keys[0].Entity

	return archive, err
}

// }}}

// Archive magic {{{

type Archive struct {
	store      blobstore.Store
	signingKey *openpgp.Entity
}

func (a Archive) Suite(name string) (*Suite, error) {
	/* Get the Release / InRelease */
	inRelease := Release{}
	components := map[string]*Component{}

	fd, err := a.store.OpenPath(path.Join("dists", name, "InRelease"))
	if err == nil {
		defer fd.Close()
		if err := control.Unmarshal(&inRelease, fd); err != nil {
			return nil, err
		}

		for _, name := range inRelease.Components {
			components[name] = &Component{Packages: []Package{}}
		}
	}

	suite := Suite{
		Name: name,

		release:    inRelease,
		Components: components,
	}

	suite.features.Hashes = []string{"sha256", "sha1"}

	return &suite, nil
}

func (a Archive) encodeHashedBySuite(path string, suite Suite, data interface{}) (*blobstore.Object, []control.FileHash, error) {

	hashers := []*transput.Hasher{}
	for _, algorithm := range suite.features.Hashes {
		hasher, err := transput.NewHasher(algorithm)
		if err != nil {
			return nil, nil, err
		}
		hashers = append(hashers, hasher)
	}

	return a.encodeHashed(path, hashers, data)
}

func (a Archive) encodeHashed(path string, hashers []*transput.Hasher, data interface{}) (*blobstore.Object, []control.FileHash, error) {

	writers := []io.Writer{}
	for _, hasher := range hashers {
		writers = append(writers, hasher)
	}

	obj, err := a.encode(data, io.MultiWriter(writers...))
	if err != nil {
		return nil, nil, err
	}

	fileHashs := []control.FileHash{}
	for _, hasher := range hashers {
		fileHashs = append(fileHashs, control.FileHashFromHasher(path, *hasher))
	}

	return obj, fileHashs, nil
}

func (a Archive) encode(data interface{}, tee io.Writer) (*blobstore.Object, error) {
	writer, err := a.store.Create()
	if err != nil {
		return nil, err
	}
	defer writer.Close()

	var target io.Writer = writer
	if tee != nil {
		target = io.MultiWriter(writer, tee)
	}

	encoder, err := control.NewEncoder(target)
	if err != nil {
		return nil, err
	}

	if err := encoder.Encode(data); err != nil {
		return nil, err
	}

	obj, err := a.store.Commit(*writer)
	if err != nil {
		return nil, err
	}

	return obj, nil
}

func (a Archive) Engross(suite Suite) (map[string]blobstore.Object, error) {
	files := map[string]blobstore.Object{}

	release := Release{
		Description:   "",
		Origin:        "",
		Label:         "",
		Version:       "",
		Suite:         suite.Name,
		Codename:      "",
		Components:    suite.ComponenetNames(),
		Architectures: suite.Arches(),
		Date:          time.Now().Format(time.RFC1123Z),
		SHA256:        []control.SHA256FileHash{},
		SHA1:          []control.SHA1FileHash{},
		SHA512:        []control.SHA512FileHash{},
		MD5Sum:        []control.MD5FileHash{},
	}

	for name, component := range suite.Components {
		for arch, pkgs := range component.ByArch() {
			suitePath := path.Join(name, fmt.Sprintf("binary-%s", arch),
				"Packages")
			filePath := path.Join("dists", suite.Name, suitePath)

			obj, hashes, err := a.encodeHashedBySuite(suitePath, suite, pkgs)
			if err != nil {
				return nil, err
			}

			for _, hash := range hashes {
				if err := release.AddHash(hash); err != nil {
					return nil, err
				}
			}

			files[filePath] = *obj
		}
	}

	filePath := path.Join("dists", suite.Name, "Release")
	obj, sig, err := a.encodeSigned(release)
	if err != nil {
		return nil, err
	}
	files[filePath] = *obj
	files[fmt.Sprintf("%s.gpg", filePath)] = *sig

	return files, nil
}

func (a Archive) Link(blobs map[string]blobstore.Object) error {
	for p, obj := range blobs {
		if err := a.store.Link(obj, p); err != nil {
			return err
		}
	}
	return nil
}

func (a Archive) Decruft() error {
	return a.store.GC(blobstore.DumbGarbageCollector{})
}

// }}}

// Suite magic {{{

type Suite struct {
	Name string

	release    Release
	Components map[string]*Component

	features struct {
		Hashes []string
	}
}

func (s Suite) Arches() []dependency.Arch {
	ret := map[dependency.Arch]bool{}
	for _, component := range s.Components {
		for _, arch := range component.Arches() {
			ret[arch] = true
		}
	}
	r := []dependency.Arch{}
	for arch, _ := range ret {
		r = append(r, arch)
	}
	return r
}

func (s Suite) ComponenetNames() []string {
	ret := []string{}
	for name, _ := range s.Components {
		ret = append(ret, name)
	}
	return ret
}

func (s Suite) Add(name string, pkg Package) {
	if _, ok := s.Components[name]; !ok {
		s.Components[name] = &Component{Packages: []Package{}}
	}
	s.Components[name].Add(pkg)
}

// }}}

// Component magic {{{

type Component struct {
	Packages []Package
}

func (c *Component) ByArch() map[dependency.Arch][]Package {
	ret := map[dependency.Arch][]Package{}

	for _, pkg := range c.Packages {
		packages := ret[pkg.Architecture]
		ret[pkg.Architecture] = append(packages, pkg)
	}

	return ret
}

func (c *Component) Arches() []dependency.Arch {
	ret := []dependency.Arch{}
	for _, pkg := range c.Packages {
		ret = append(ret, pkg.Architecture)
	}
	return ret
}

func (c *Component) Add(p Package) {
	c.Packages = append(c.Packages, p)
}

// }}}

// vim: foldmethod=marker
<|endoftext|>