<commit_before>package torrent

import (
	"crypto/rand"
	"errors"
	"io"
	"sync"

	"github.com/cenkalti/rain/internal/announcer"
	"github.com/cenkalti/rain/internal/bitfield"
	"github.com/cenkalti/rain/internal/downloader"
	"github.com/cenkalti/rain/internal/logger"
	"github.com/cenkalti/rain/internal/magnet"
	"github.com/cenkalti/rain/internal/metainfo"
	"github.com/cenkalti/rain/internal/peerlist"
	"github.com/cenkalti/rain/internal/peermanager"
	"github.com/cenkalti/rain/internal/worker"
	"github.com/cenkalti/rain/resume"
	"github.com/cenkalti/rain/storage"
	"github.com/cenkalti/rain/storage/filestorage"
)

var (
	// Version of client. Set during build.
	Version = "0000" // zero means development version

	// http://www.bittorrent.org/beps/bep_0020.html
	peerIDPrefix = []byte("-RN" + Version + "-")
)

// Torrent connect to peers and downloads files from swarm.
type Torrent struct {
	peerID     [20]byte // unique id per torrent
	infoHash   [20]byte
	announce   string
	port       int           // listen for peer connections
	closed     bool          // true after Close() is called
	m          sync.Mutex    // protects running and closed state
	completeC  chan struct{} // downloader closes this channel when all pieces are downloaded
	workers    worker.Workers
	log        logger.Logger
	downloader *downloader.Downloader
}

// DownloadTorrent returns a new torrent by reading a metainfo file.
//
// Files are read from disk. If there are existing files, hash check will be done.
//
// Close must be called before discarding the torrent.
//
// Seeding continues after all files are downloaded.
//
// You should listen NotifyComplete and NotifyError channels after starting the torrent.
func DownloadTorrent(r io.Reader, port int, sto storage.Storage, res resume.DB) (*Torrent, error) {
	if res != nil {
		rspec, err := res.Read()
		if err != nil {
			return nil, err
		}
		if rspec != nil {
			return loadResumeSpec(rspec)
		}
	}
	m, err := metainfo.New(r)
	if err != nil {
		return nil, err
	}
	spec := &downloader.Spec{
		InfoHash: m.Info.Hash,
		Storage:  sto,
		Resume:   res,
		Info:     m.Info,
	}
	if res != nil {
		err = writeResume(res, spec, port, m.Info.Name, m.Announce)
		if err != nil {
			return nil, err
		}
	}
	return newTorrent(spec, port, m.Info.Name, m.Announce)
}

func DownloadMagnet(magnetLink string, port int, sto storage.Storage, res resume.DB) (*Torrent, error) {
	if res != nil {
		rspec, err := res.Read()
		if err != nil {
			return nil, err
		}
		if rspec != nil {
			return loadResumeSpec(rspec)
		}
	}
	m, err := magnet.New(magnetLink)
	if err != nil {
		return nil, err
	}
	spec := &downloader.Spec{
		InfoHash: m.InfoHash,
		Storage:  sto,
		Resume:   res,
	}
	if res != nil {
		err = writeResume(res, spec, port, m.Name, m.Trackers[0])
		if err != nil {
			return nil, err
		}
	}
	return newTorrent(spec, port, m.Name, m.Trackers[0])
}

func Resume(res resume.DB) (*Torrent, error) {
	spec, err := res.Read()
	if err != nil {
		return nil, err
	}
	if spec == nil {
		return nil, errors.New("no resume info")
	}
	return loadResumeSpec(spec)
}

func loadResumeSpec(spec *resume.Spec) (*Torrent, error) {
	var err error
	dspec := &downloader.Spec{}
	copy(dspec.InfoHash[:], spec.InfoHash)
	if len(spec.Info) > 0 {
		dspec.Info, err = metainfo.NewInfo(spec.Info)
		if err != nil {
			return nil, err
		}
		if len(spec.Bitfield) > 0 {
			dspec.Bitfield = bitfield.New(dspec.Info.NumPieces)
			copy(dspec.Bitfield.Bytes(), spec.Bitfield)
		}
	}
	switch spec.StorageType {
	case filestorage.StorageType:
		dspec.Storage = &filestorage.FileStorage{}
	default:
		return nil, errors.New("unknown storage type: " + spec.StorageType)
	}
	err = dspec.Storage.Load(spec.StorageArgs)
	if err != nil {
		return nil, err
	}
	return newTorrent(dspec, spec.Port, spec.Name, spec.Trackers[0])
}

func writeResume(res resume.DB, dspec *downloader.Spec, port int, name string, tracker string) error {
	rspec := &resume.Spec{
		InfoHash: dspec.InfoHash[:],
		Port:     port,
		Name:     name,
		// TODO save every tracker
		Trackers:    []string{tracker},
		StorageType: dspec.Storage.Type(),
		StorageArgs: dspec.Storage.Args(),
	}
	if dspec.Info != nil {
		rspec.Info = dspec.Info.Bytes
	}
	if dspec.Bitfield != nil {
		rspec.Bitfield = dspec.Bitfield.Bytes()
	}
	return res.Write(rspec)
}

// TODO pass every tracker
func newTorrent(spec *downloader.Spec, port int, name string, tracker string) (*Torrent, error) {
	logName := name
	if len(logName) > 8 {
		logName = logName[:8]
	}

	var peerID [20]byte
	copy(peerID[:], peerIDPrefix)
	_, err := rand.Read(peerID[len(peerIDPrefix):]) // nolint: gosec
	if err != nil {
		return nil, err
	}

	completeC := make(chan struct{})
	l := logger.New("download " + logName)

	t := &Torrent{
		peerID:   peerID,
		infoHash: spec.InfoHash,
		// TODO pass every tracker to downloader
		announce:   tracker,
		port:       port,
		log:        l,
		completeC:  completeC,
		downloader: downloader.New(spec, completeC, l),
	}

	// keep list of peer addresses to connect
	pl := peerlist.New()
	t.workers.Start(pl)

	// get peers from tracker
	an := announcer.New(t.announce, t, t.completeC, pl, t.log)
	t.workers.Start(an)

	// manage peer connections
	pm := peermanager.New(t.port, pl, t.peerID, t.infoHash, t.downloader.NewPeers(), t.log)
	t.workers.Start(pm)

	return t, nil
}

// Close this torrent and release all resources.
func (t *Torrent) Close() error {
	t.m.Lock()
	if t.closed {
		t.m.Unlock()
		return nil
	}
	t.closed = true
	t.m.Unlock()

	t.workers.Stop()
	t.downloader.Close()
	return nil
}

// Port returns the port number that the client is listening.
func (t *Torrent) Port() int {
	return t.port
}

// PeerID is unique per torrent.
func (t *Torrent) PeerID() [20]byte { return t.peerID }

// InfoHash identifies the torrent file that is being downloaded.
func (t *Torrent) InfoHash() [20]byte { return t.infoHash }

// NotifyComplete returns a channel that is closed once all pieces are downloaded successfully.
func (t *Torrent) NotifyComplete() <-chan struct{} { return t.completeC }

// NotifyError returns a new channel for waiting download errors.
//
// When error is sent to the channel, torrent is stopped automatically.
func (t *Torrent) NotifyError() <-chan error { return t.downloader.ErrC() }

type Stats struct {
	// Bytes that are downloaded and passed hash check.
	BytesComplete int64

	// BytesLeft is the number of bytes that is needed to complete all missing pieces.
	BytesIncomplete int64

	// BytesTotal is the number of total bytes of files in torrent.
	//
	// BytesTotal = BytesComplete + BytesIncomplete
	BytesTotal int64

	// BytesDownloaded is the number of bytes downloaded from swarm.
	// Because some pieces may be downloaded more than once, this number may be greater than BytesCompleted returns.
	// BytesDownloaded int64

	// BytesUploaded is the number of bytes uploaded to the swarm.
	// BytesUploaded   int64
}

func (t *Torrent) Stats() *Stats {
	t.m.Lock()
	defer t.m.Unlock()

	if t.closed {
		return nil
	}

	ds := t.downloader.Stats()
	return &Stats{
		BytesComplete:   ds.BytesComplete,
		BytesIncomplete: ds.BytesIncomplete,
		BytesTotal:      ds.BytesTotal,
	}
}

func (t *Torrent) BytesDownloaded() int64 { return t.Stats().BytesComplete } // TODO not the same thing
func (t *Torrent) BytesUploaded() int64   { return 0 }                       // TODO implememnt
func (t *Torrent) BytesLeft() int64       { return t.Stats().BytesIncomplete }
<commit_msg>check hashes of resume files<commit_after>package torrent

import (
	"bytes"
	"crypto/rand"
	"errors"
	"io"
	"sync"

	"github.com/cenkalti/rain/internal/announcer"
	"github.com/cenkalti/rain/internal/bitfield"
	"github.com/cenkalti/rain/internal/downloader"
	"github.com/cenkalti/rain/internal/logger"
	"github.com/cenkalti/rain/internal/magnet"
	"github.com/cenkalti/rain/internal/metainfo"
	"github.com/cenkalti/rain/internal/peerlist"
	"github.com/cenkalti/rain/internal/peermanager"
	"github.com/cenkalti/rain/internal/worker"
	"github.com/cenkalti/rain/resume"
	"github.com/cenkalti/rain/storage"
	"github.com/cenkalti/rain/storage/filestorage"
)

var (
	// Version of client. Set during build.
	Version = "0000" // zero means development version

	// http://www.bittorrent.org/beps/bep_0020.html
	peerIDPrefix = []byte("-RN" + Version + "-")

	errInvalidResumeFile = errors.New("invalid resume file (info hashes does not match)")
)

// Torrent connect to peers and downloads files from swarm.
type Torrent struct {
	peerID     [20]byte // unique id per torrent
	infoHash   [20]byte
	announce   string
	port       int           // listen for peer connections
	closed     bool          // true after Close() is called
	m          sync.Mutex    // protects running and closed state
	completeC  chan struct{} // downloader closes this channel when all pieces are downloaded
	workers    worker.Workers
	log        logger.Logger
	downloader *downloader.Downloader
}

// DownloadTorrent returns a new torrent by reading a metainfo file.
//
// Files are read from disk. If there are existing files, hash check will be done.
//
// Close must be called before discarding the torrent.
//
// Seeding continues after all files are downloaded.
//
// You should listen NotifyComplete and NotifyError channels after starting the torrent.
func DownloadTorrent(r io.Reader, port int, sto storage.Storage, res resume.DB) (*Torrent, error) {
	m, err := metainfo.New(r)
	if err != nil {
		return nil, err
	}
	if res != nil {
		rspec, err2 := res.Read()
		if err2 != nil {
			return nil, err2
		}
		if !bytes.Equal(rspec.InfoHash, m.Info.Hash[:]) {
			return nil, errInvalidResumeFile
		}
		if rspec != nil {
			return loadResumeSpec(rspec)
		}
	}
	spec := &downloader.Spec{
		InfoHash: m.Info.Hash,
		Storage:  sto,
		Resume:   res,
		Info:     m.Info,
	}
	if res != nil {
		err = writeResume(res, spec, port, m.Info.Name, m.Announce)
		if err != nil {
			return nil, err
		}
	}
	return newTorrent(spec, port, m.Info.Name, m.Announce)
}

func DownloadMagnet(magnetLink string, port int, sto storage.Storage, res resume.DB) (*Torrent, error) {
	m, err := magnet.New(magnetLink)
	if err != nil {
		return nil, err
	}
	if res != nil {
		rspec, err2 := res.Read()
		if err2 != nil {
			return nil, err2
		}
		if !bytes.Equal(rspec.InfoHash, m.InfoHash[:]) {
			return nil, errInvalidResumeFile
		}
		if rspec != nil {
			return loadResumeSpec(rspec)
		}
	}
	spec := &downloader.Spec{
		InfoHash: m.InfoHash,
		Storage:  sto,
		Resume:   res,
	}
	if res != nil {
		err = writeResume(res, spec, port, m.Name, m.Trackers[0])
		if err != nil {
			return nil, err
		}
	}
	return newTorrent(spec, port, m.Name, m.Trackers[0])
}

func Resume(res resume.DB) (*Torrent, error) {
	spec, err := res.Read()
	if err != nil {
		return nil, err
	}
	if spec == nil {
		return nil, errors.New("no resume info")
	}
	return loadResumeSpec(spec)
}

func loadResumeSpec(spec *resume.Spec) (*Torrent, error) {
	var err error
	dspec := &downloader.Spec{}
	copy(dspec.InfoHash[:], spec.InfoHash)
	if len(spec.Info) > 0 {
		dspec.Info, err = metainfo.NewInfo(spec.Info)
		if err != nil {
			return nil, err
		}
		if len(spec.Bitfield) > 0 {
			dspec.Bitfield = bitfield.New(dspec.Info.NumPieces)
			copy(dspec.Bitfield.Bytes(), spec.Bitfield)
		}
	}
	switch spec.StorageType {
	case filestorage.StorageType:
		dspec.Storage = &filestorage.FileStorage{}
	default:
		return nil, errors.New("unknown storage type: " + spec.StorageType)
	}
	err = dspec.Storage.Load(spec.StorageArgs)
	if err != nil {
		return nil, err
	}
	return newTorrent(dspec, spec.Port, spec.Name, spec.Trackers[0])
}

func writeResume(res resume.DB, dspec *downloader.Spec, port int, name string, tracker string) error {
	rspec := &resume.Spec{
		InfoHash: dspec.InfoHash[:],
		Port:     port,
		Name:     name,
		// TODO save every tracker
		Trackers:    []string{tracker},
		StorageType: dspec.Storage.Type(),
		StorageArgs: dspec.Storage.Args(),
	}
	if dspec.Info != nil {
		rspec.Info = dspec.Info.Bytes
	}
	if dspec.Bitfield != nil {
		rspec.Bitfield = dspec.Bitfield.Bytes()
	}
	return res.Write(rspec)
}

// TODO pass every tracker
func newTorrent(spec *downloader.Spec, port int, name string, tracker string) (*Torrent, error) {
	logName := name
	if len(logName) > 8 {
		logName = logName[:8]
	}

	var peerID [20]byte
	copy(peerID[:], peerIDPrefix)
	_, err := rand.Read(peerID[len(peerIDPrefix):]) // nolint: gosec
	if err != nil {
		return nil, err
	}

	completeC := make(chan struct{})
	l := logger.New("download " + logName)

	t := &Torrent{
		peerID:   peerID,
		infoHash: spec.InfoHash,
		// TODO pass every tracker to downloader
		announce:   tracker,
		port:       port,
		log:        l,
		completeC:  completeC,
		downloader: downloader.New(spec, completeC, l),
	}

	// keep list of peer addresses to connect
	pl := peerlist.New()
	t.workers.Start(pl)

	// get peers from tracker
	an := announcer.New(t.announce, t, t.completeC, pl, t.log)
	t.workers.Start(an)

	// manage peer connections
	pm := peermanager.New(t.port, pl, t.peerID, t.infoHash, t.downloader.NewPeers(), t.log)
	t.workers.Start(pm)

	return t, nil
}

// Close this torrent and release all resources.
func (t *Torrent) Close() error {
	t.m.Lock()
	if t.closed {
		t.m.Unlock()
		return nil
	}
	t.closed = true
	t.m.Unlock()

	t.workers.Stop()
	t.downloader.Close()
	return nil
}

// Port returns the port number that the client is listening.
func (t *Torrent) Port() int {
	return t.port
}

// PeerID is unique per torrent.
func (t *Torrent) PeerID() [20]byte { return t.peerID }

// InfoHash identifies the torrent file that is being downloaded.
func (t *Torrent) InfoHash() [20]byte { return t.infoHash }

// NotifyComplete returns a channel that is closed once all pieces are downloaded successfully.
func (t *Torrent) NotifyComplete() <-chan struct{} { return t.completeC }

// NotifyError returns a new channel for waiting download errors.
//
// When error is sent to the channel, torrent is stopped automatically.
func (t *Torrent) NotifyError() <-chan error { return t.downloader.ErrC() }

type Stats struct {
	// Bytes that are downloaded and passed hash check.
	BytesComplete int64

	// BytesLeft is the number of bytes that is needed to complete all missing pieces.
	BytesIncomplete int64

	// BytesTotal is the number of total bytes of files in torrent.
	//
	// BytesTotal = BytesComplete + BytesIncomplete
	BytesTotal int64

	// BytesDownloaded is the number of bytes downloaded from swarm.
	// Because some pieces may be downloaded more than once, this number may be greater than BytesCompleted returns.
	// BytesDownloaded int64

	// BytesUploaded is the number of bytes uploaded to the swarm.
	// BytesUploaded   int64
}

func (t *Torrent) Stats() *Stats {
	t.m.Lock()
	defer t.m.Unlock()

	if t.closed {
		return nil
	}

	ds := t.downloader.Stats()
	return &Stats{
		BytesComplete:   ds.BytesComplete,
		BytesIncomplete: ds.BytesIncomplete,
		BytesTotal:      ds.BytesTotal,
	}
}

func (t *Torrent) BytesDownloaded() int64 { return t.Stats().BytesComplete } // TODO not the same thing
func (t *Torrent) BytesUploaded() int64   { return 0 }                       // TODO implememnt
func (t *Torrent) BytesLeft() int64       { return t.Stats().BytesIncomplete }
<|endoftext|>


<commit_before>package starbound

import (
	"bytes"
	"io"
)

const (
	BTreeDB5HeaderSize = 512
)

var (
	BlockFree  = []byte("FF")
	BlockIndex = []byte("II")
	BlockLeaf  = []byte("LL")

	HeaderSignature = []byte("BTreeDB5")
)

func NewBTreeDB5(r io.ReaderAt) (db *BTreeDB5, err error) {
	db = &BTreeDB5{r: r}
	header := make([]byte, 67)
	n, err := r.ReadAt(header, 0)
	if n != len(header) || err != nil {
		return nil, ErrInvalidHeader
	}
	if !bytes.Equal(header[:8], HeaderSignature) {
		return nil, ErrInvalidHeader
	}
	db.BlockSize = getInt(header, 8)
	db.Name = string(bytes.TrimRight(header[12:28], "\x00"))
	db.KeySize = getInt(header, 28)
	db.Swap = (header[32] == 1)
	db.freeBlock1 = getInt(header, 33)
	// Skip 3 bytes...
	db.unknown1 = getInt(header, 40)
	// Skip 1 byte...
	db.rootBlock1 = getInt(header, 45)
	db.rootBlock1IsLeaf = (header[49] == 1)
	db.freeBlock2 = getInt(header, 50)
	// Skip 3 bytes...
	db.unknown2 = getInt(header, 57)
	// Skip 1 byte...
	db.rootBlock2 = getInt(header, 62)
	db.rootBlock2IsLeaf = (header[66] == 1)
	return
}

type BTreeDB5 struct {
	Name      string
	BlockSize int
	KeySize   int
	Swap      bool

	r io.ReaderAt

	freeBlock1, freeBlock2 int
	rootBlock1, rootBlock2 int
	rootBlock1IsLeaf       bool
	rootBlock2IsLeaf       bool
	unknown1, unknown2     int
}

func (db *BTreeDB5) FreeBlock() int {
	if !db.Swap {
		return db.freeBlock1
	} else {
		return db.freeBlock2
	}
}

func (db *BTreeDB5) Get(key []byte) (data []byte, err error) {
	if len(key) != db.KeySize {
		return nil, ErrInvalidKeyLength
	}
	bufSize := 11
	if db.KeySize > bufSize {
		bufSize = db.KeySize
	}
	buf := make([]byte, bufSize)
	block := db.RootBlock()
	offset := db.blockOffset(block)
	entrySize := db.KeySize + 4
	// Traverse the B-tree until we reach a leaf.
	for {
		if _, err = db.r.ReadAt(buf[:11], offset); err != nil {
			return
		}
		if !bytes.Equal(buf[:2], BlockIndex) {
			break
		}
		offset += 11
		// Binary search for the key.
		lo, hi := 0, getInt(buf, 3)
		block = getInt(buf, 7)
		for lo < hi {
			mid := (lo + hi) / 2
			if _, err = db.r.ReadAt(buf[:db.KeySize], offset+int64(entrySize*mid)); err != nil {
				return
			}
			if bytes.Compare(key, buf[:db.KeySize]) < 0 {
				hi = mid
			} else {
				lo = mid + 1
			}
		}
		if lo > 0 {
			// A candidate leaf/index was found in the current index. Get the block index.
			db.r.ReadAt(buf[:4], offset+int64(entrySize*(lo-1)+db.KeySize))
			block = getInt(buf, 0)
		}
		offset = db.blockOffset(block)
	}
	// Scan leaves for the key, then read the data.
	r := NewLeafReader(db, block)
	if _, err = r.Read(buf[:4]); err != nil {
		return
	}
	keyCount := getInt(buf, 0)
	for i := 0; i < keyCount; i += 1 {
		if _, err = r.Read(buf[:db.KeySize]); err != nil {
			return
		}
		var n int64
		if n, err = ReadVarint(r); err != nil {
			return
		}
		// TODO: Allow skipping without reading.
		temp := make([]byte, n)
		if _, err = io.ReadFull(r, temp); err != nil {
			return
		}
		if bytes.Equal(buf[:db.KeySize], key) {
			return temp, nil
		}
	}
	return nil, ErrKeyNotFound
}

func (db *BTreeDB5) RootBlock() int {
	if !db.Swap {
		return db.rootBlock1
	} else {
		return db.rootBlock2
	}
}

func (db *BTreeDB5) blockOffset(block int) int64 {
	return BTreeDB5HeaderSize + int64(block*db.BlockSize)
}

func NewLeafReader(db *BTreeDB5, block int) *LeafReader {
	return &LeafReader{
		db:  db,
		cur: db.blockOffset(block),
	}
}

type LeafReader struct {
	db       *BTreeDB5
	cur, end int64
}

func (l *LeafReader) Read(p []byte) (n int, err error) {
	buf := make([]byte, 4)
	if l.end == 0 {
		if _, err = l.db.r.ReadAt(buf[:2], l.cur); err != nil {
			return
		}
		if !bytes.Equal(buf[:2], BlockLeaf) {
			return 0, ErrDidNotReachLeaf
		}
		l.end = l.cur + int64(l.db.BlockSize-4)
		l.cur += 2
	}
	want := int64(len(p))
	if l.cur+want > l.end {
		want = l.end - l.cur
	}
	n, err = l.db.r.ReadAt(p[:want], l.cur)
	l.cur += int64(n)
	if l.cur == l.end {
		l.db.r.ReadAt(buf, l.cur)
		l.cur = l.db.blockOffset(getInt(buf, 0))
		l.end = 0
	}
	return
}
<commit_msg>Reduce allocations<commit_after>package starbound

import (
	"bytes"
	"io"
)

const (
	BTreeDB5HeaderSize = 512
)

var (
	BlockFree  = []byte("FF")
	BlockIndex = []byte("II")
	BlockLeaf  = []byte("LL")

	HeaderSignature = []byte("BTreeDB5")
)

func NewBTreeDB5(r io.ReaderAt) (db *BTreeDB5, err error) {
	db = &BTreeDB5{r: r}
	header := make([]byte, 67)
	n, err := r.ReadAt(header, 0)
	if n != len(header) || err != nil {
		return nil, ErrInvalidHeader
	}
	if !bytes.Equal(header[:8], HeaderSignature) {
		return nil, ErrInvalidHeader
	}
	db.BlockSize = getInt(header, 8)
	db.Name = string(bytes.TrimRight(header[12:28], "\x00"))
	db.KeySize = getInt(header, 28)
	db.Swap = (header[32] == 1)
	db.freeBlock1 = getInt(header, 33)
	// Skip 3 bytes...
	db.unknown1 = getInt(header, 40)
	// Skip 1 byte...
	db.rootBlock1 = getInt(header, 45)
	db.rootBlock1IsLeaf = (header[49] == 1)
	db.freeBlock2 = getInt(header, 50)
	// Skip 3 bytes...
	db.unknown2 = getInt(header, 57)
	// Skip 1 byte...
	db.rootBlock2 = getInt(header, 62)
	db.rootBlock2IsLeaf = (header[66] == 1)
	return
}

type BTreeDB5 struct {
	Name      string
	BlockSize int
	KeySize   int
	Swap      bool

	r io.ReaderAt

	freeBlock1, freeBlock2 int
	rootBlock1, rootBlock2 int
	rootBlock1IsLeaf       bool
	rootBlock2IsLeaf       bool
	unknown1, unknown2     int
}

func (db *BTreeDB5) FreeBlock() int {
	if !db.Swap {
		return db.freeBlock1
	} else {
		return db.freeBlock2
	}
}

func (db *BTreeDB5) Get(key []byte) (data []byte, err error) {
	if len(key) != db.KeySize {
		return nil, ErrInvalidKeyLength
	}
	bufSize := 11
	if db.KeySize > bufSize {
		bufSize = db.KeySize
	}
	buf := make([]byte, bufSize)
	bufBlock := buf[:4]
	bufHead := buf[:11]
	bufKey := buf[:db.KeySize]
	bufType := buf[:2]
	block := db.RootBlock()
	offset := db.blockOffset(block)
	entrySize := db.KeySize + 4
	// Traverse the B-tree until we reach a leaf.
	for {
		if _, err = db.r.ReadAt(bufHead, offset); err != nil {
			return
		}
		if !bytes.Equal(bufType, BlockIndex) {
			break
		}
		offset += 11
		// Binary search for the key.
		lo, hi := 0, getInt(buf, 3)
		block = getInt(buf, 7)
		for lo < hi {
			mid := (lo + hi) / 2
			if _, err = db.r.ReadAt(bufKey, offset+int64(entrySize*mid)); err != nil {
				return
			}
			if bytes.Compare(key, bufKey) < 0 {
				hi = mid
			} else {
				lo = mid + 1
			}
		}
		if lo > 0 {
			// A candidate leaf/index was found in the current index. Get the block index.
			db.r.ReadAt(bufBlock, offset+int64(entrySize*(lo-1)+db.KeySize))
			block = getInt(buf, 0)
		}
		offset = db.blockOffset(block)
	}
	// Scan leaves for the key, then read the data.
	r := NewLeafReader(db, block)
	if _, err = r.Read(bufBlock); err != nil {
		return
	}
	keyCount := getInt(buf, 0)
	for i := 0; i < keyCount; i += 1 {
		if _, err = r.Read(bufKey); err != nil {
			return
		}
		var n int64
		if n, err = ReadVarint(r); err != nil {
			return
		}
		// TODO: Allow skipping without reading.
		temp := make([]byte, n)
		if _, err = io.ReadFull(r, temp); err != nil {
			return
		}
		if bytes.Equal(bufKey, key) {
			return temp, nil
		}
	}
	return nil, ErrKeyNotFound
}

func (db *BTreeDB5) RootBlock() int {
	if !db.Swap {
		return db.rootBlock1
	} else {
		return db.rootBlock2
	}
}

func (db *BTreeDB5) blockOffset(block int) int64 {
	return BTreeDB5HeaderSize + int64(block*db.BlockSize)
}

func NewLeafReader(db *BTreeDB5, block int) *LeafReader {
	return &LeafReader{
		db:  db,
		cur: db.blockOffset(block),
	}
}

type LeafReader struct {
	db       *BTreeDB5
	cur, end int64
}

func (l *LeafReader) Read(p []byte) (n int, err error) {
	buf := make([]byte, 4)
	if l.end == 0 {
		if _, err = l.db.r.ReadAt(buf[:2], l.cur); err != nil {
			return
		}
		if !bytes.Equal(buf[:2], BlockLeaf) {
			return 0, ErrDidNotReachLeaf
		}
		l.end = l.cur + int64(l.db.BlockSize-4)
		l.cur += 2
	}
	want := int64(len(p))
	if l.cur+want > l.end {
		want = l.end - l.cur
	}
	n, err = l.db.r.ReadAt(p[:want], l.cur)
	l.cur += int64(n)
	if l.cur == l.end {
		l.db.r.ReadAt(buf, l.cur)
		l.cur = l.db.blockOffset(getInt(buf, 0))
		l.end = 0
	}
	return
}
<|endoftext|>


<commit_before>package remote

import (
	"bytes"
	"crypto/md5"
	"encoding/base64"
	"fmt"
	"io"
	"log"
	"net/http"
	"net/url"
	"os"
	"path"
	"strings"

	"github.com/hashicorp/go-retryablehttp"
	"github.com/hashicorp/terraform/terraform"
)

const (
	// defaultAtlasServer is used when no address is given
	defaultAtlasServer = "https://atlas.hashicorp.com/"
)

func atlasFactory(conf map[string]string) (Client, error) {
	var client AtlasClient

	server, ok := conf["address"]
	if !ok || server == "" {
		server = defaultAtlasServer
	}

	url, err := url.Parse(server)
	if err != nil {
		return nil, err
	}

	token, ok := conf["access_token"]
	if token == "" {
		token = os.Getenv("ATLAS_TOKEN")
		ok = true
	}
	if !ok || token == "" {
		return nil, fmt.Errorf(
			"missing 'access_token' configuration or ATLAS_TOKEN environmental variable")
	}

	name, ok := conf["name"]
	if !ok || name == "" {
		return nil, fmt.Errorf("missing 'name' configuration")
	}

	parts := strings.Split(name, "/")
	if len(parts) != 2 {
		return nil, fmt.Errorf("malformed name '%s', expected format '<account>/<name>'", name)
	}

	// If it exists, add the `ATLAS_RUN_ID` environment
	// variable as a param, which is injected during Atlas Terraform
	// runs. This is completely optional.
	client.RunId = os.Getenv("ATLAS_RUN_ID")

	client.Server = server
	client.ServerURL = url
	client.AccessToken = token
	client.User = parts[0]
	client.Name = parts[1]

	return &client, nil
}

// AtlasClient implements the Client interface for an Atlas compatible server.
type AtlasClient struct {
	Server      string
	ServerURL   *url.URL
	User        string
	Name        string
	AccessToken string
	RunId       string
	HTTPClient  *retryablehttp.Client

	conflictHandlingAttempted bool
}

func (c *AtlasClient) Get() (*Payload, error) {
	// Make the HTTP request
	req, err := retryablehttp.NewRequest("GET", c.url().String(), nil)
	if err != nil {
		return nil, fmt.Errorf("Failed to make HTTP request: %v", err)
	}

	// Request the url
	client := c.http()
	resp, err := client.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	// Handle the common status codes
	switch resp.StatusCode {
	case http.StatusOK:
		// Handled after
	case http.StatusNoContent:
		return nil, nil
	case http.StatusNotFound:
		return nil, nil
	case http.StatusUnauthorized:
		return nil, fmt.Errorf("HTTP remote state endpoint requires auth")
	case http.StatusForbidden:
		return nil, fmt.Errorf("HTTP remote state endpoint invalid auth")
	case http.StatusInternalServerError:
		return nil, fmt.Errorf("HTTP remote state internal server error")
	default:
		return nil, fmt.Errorf(
			"Unexpected HTTP response code: %d\n\nBody: %s",
			resp.StatusCode, c.readBody(resp.Body))
	}

	// Read in the body
	buf := bytes.NewBuffer(nil)
	if _, err := io.Copy(buf, resp.Body); err != nil {
		return nil, fmt.Errorf("Failed to read remote state: %v", err)
	}

	// Create the payload
	payload := &Payload{
		Data: buf.Bytes(),
	}

	if len(payload.Data) == 0 {
		return nil, nil
	}

	// Check for the MD5
	if raw := resp.Header.Get("Content-MD5"); raw != "" {
		md5, err := base64.StdEncoding.DecodeString(raw)
		if err != nil {
			return nil, fmt.Errorf("Failed to decode Content-MD5 '%s': %v", raw, err)
		}

		payload.MD5 = md5
	} else {
		// Generate the MD5
		hash := md5.Sum(payload.Data)
		payload.MD5 = hash[:]
	}

	return payload, nil
}

func (c *AtlasClient) Put(state []byte) error {
	// Get the target URL
	base := c.url()

	// Generate the MD5
	hash := md5.Sum(state)
	b64 := base64.StdEncoding.EncodeToString(hash[:])

	// Make the HTTP client and request
	req, err := retryablehttp.NewRequest("PUT", base.String(), bytes.NewReader(state))
	if err != nil {
		return fmt.Errorf("Failed to make HTTP request: %v", err)
	}

	// Prepare the request
	req.Header.Set("Content-MD5", b64)
	req.Header.Set("Content-Type", "application/json")
	req.ContentLength = int64(len(state))

	// Make the request
	client := c.http()
	resp, err := client.Do(req)
	if err != nil {
		return fmt.Errorf("Failed to upload state: %v", err)
	}
	defer resp.Body.Close()

	// Handle the error codes
	switch resp.StatusCode {
	case http.StatusOK:
		return nil
	case http.StatusConflict:
		return c.handleConflict(c.readBody(resp.Body), state)
	default:
		return fmt.Errorf(
			"HTTP error: %d\n\nBody: %s",
			resp.StatusCode, c.readBody(resp.Body))
	}
}

func (c *AtlasClient) Delete() error {
	// Make the HTTP request
	req, err := retryablehttp.NewRequest("DELETE", c.url().String(), nil)
	if err != nil {
		return fmt.Errorf("Failed to make HTTP request: %v", err)
	}

	// Make the request
	client := c.http()
	resp, err := client.Do(req)
	if err != nil {
		return fmt.Errorf("Failed to delete state: %v", err)
	}
	defer resp.Body.Close()

	// Handle the error codes
	switch resp.StatusCode {
	case http.StatusOK:
		return nil
	case http.StatusNoContent:
		return nil
	case http.StatusNotFound:
		return nil
	default:
		return fmt.Errorf(
			"HTTP error: %d\n\nBody: %s",
			resp.StatusCode, c.readBody(resp.Body))
	}
}

func (c *AtlasClient) readBody(b io.Reader) string {
	var buf bytes.Buffer
	if _, err := io.Copy(&buf, b); err != nil {
		return fmt.Sprintf("Error reading body: %s", err)
	}

	result := buf.String()
	if result == "" {
		result = "<empty>"
	}

	return result
}

func (c *AtlasClient) url() *url.URL {
	values := url.Values{}

	values.Add("atlas_run_id", c.RunId)
	values.Add("access_token", c.AccessToken)

	return &url.URL{
		Scheme:   c.ServerURL.Scheme,
		Host:     c.ServerURL.Host,
		Path:     path.Join("api/v1/terraform/state", c.User, c.Name),
		RawQuery: values.Encode(),
	}
}

func (c *AtlasClient) http() *retryablehttp.Client {
	if c.HTTPClient != nil {
		return c.HTTPClient
	}
	return retryablehttp.NewClient()
}

// Atlas returns an HTTP 409 - Conflict if the pushed state reports the same
// Serial number but the checksum of the raw content differs. This can
// sometimes happen when Terraform changes state representation internally
// between versions in a way that's semantically neutral but affects the JSON
// output and therefore the checksum.
//
// Here we detect and handle this situation by ticking the serial and retrying
// iff for the previous state and the proposed state:
//
//   * the serials match
//   * the parsed states are Equal (semantically equivalent)
//
// In other words, in this situation Terraform can override Atlas's detected
// conflict by asserting that the state it is pushing is indeed correct.
func (c *AtlasClient) handleConflict(msg string, state []byte) error {
	log.Printf("[DEBUG] Handling Atlas conflict response: %s", msg)

	if c.conflictHandlingAttempted {
		log.Printf("[DEBUG] Already attempted conflict resolution; returning conflict.")
	} else {
		c.conflictHandlingAttempted = true
		log.Printf("[DEBUG] Atlas reported conflict, checking for equivalent states.")

		payload, err := c.Get()
		if err != nil {
			return conflictHandlingError(err)
		}

		currentState, err := terraform.ReadState(bytes.NewReader(payload.Data))
		if err != nil {
			return conflictHandlingError(err)
		}

		proposedState, err := terraform.ReadState(bytes.NewReader(state))
		if err != nil {
			return conflictHandlingError(err)
		}

		if statesAreEquivalent(currentState, proposedState) {
			log.Printf("[DEBUG] States are equivalent, incrementing serial and retrying.")
			proposedState.Serial++
			var buf bytes.Buffer
			if err := terraform.WriteState(proposedState, &buf); err != nil {
				return conflictHandlingError(err)
			}
			return c.Put(buf.Bytes())
		} else {
			log.Printf("[DEBUG] States are not equivalent, returning conflict.")
		}
	}

	return fmt.Errorf(
		"Atlas detected a remote state conflict.\n\nMessage: %s", msg)
}

func conflictHandlingError(err error) error {
	return fmt.Errorf(
		"Error while handling a conflict response from Atlas: %s", err)
}

func statesAreEquivalent(current, proposed *terraform.State) bool {
	return current.Serial == proposed.Serial && current.Equal(proposed)
}
<commit_msg>state/remote/atlas: Use go-rootcerts for certificate loading<commit_after>package remote

import (
	"bytes"
	"crypto/md5"
	"crypto/tls"
	"encoding/base64"
	"fmt"
	"io"
	"log"
	"net/http"
	"net/url"
	"os"
	"path"
	"strings"

	"github.com/hashicorp/go-cleanhttp"
	"github.com/hashicorp/go-retryablehttp"
	"github.com/hashicorp/go-rootcerts"
	"github.com/hashicorp/terraform/terraform"
)

const (
	// defaultAtlasServer is used when no address is given
	defaultAtlasServer = "https://atlas.hashicorp.com/"
)

func atlasFactory(conf map[string]string) (Client, error) {
	var client AtlasClient

	server, ok := conf["address"]
	if !ok || server == "" {
		server = defaultAtlasServer
	}

	url, err := url.Parse(server)
	if err != nil {
		return nil, err
	}

	token, ok := conf["access_token"]
	if token == "" {
		token = os.Getenv("ATLAS_TOKEN")
		ok = true
	}
	if !ok || token == "" {
		return nil, fmt.Errorf(
			"missing 'access_token' configuration or ATLAS_TOKEN environmental variable")
	}

	name, ok := conf["name"]
	if !ok || name == "" {
		return nil, fmt.Errorf("missing 'name' configuration")
	}

	parts := strings.Split(name, "/")
	if len(parts) != 2 {
		return nil, fmt.Errorf("malformed name '%s', expected format '<account>/<name>'", name)
	}

	// If it exists, add the `ATLAS_RUN_ID` environment
	// variable as a param, which is injected during Atlas Terraform
	// runs. This is completely optional.
	client.RunId = os.Getenv("ATLAS_RUN_ID")

	client.Server = server
	client.ServerURL = url
	client.AccessToken = token
	client.User = parts[0]
	client.Name = parts[1]

	return &client, nil
}

// AtlasClient implements the Client interface for an Atlas compatible server.
type AtlasClient struct {
	Server      string
	ServerURL   *url.URL
	User        string
	Name        string
	AccessToken string
	RunId       string
	HTTPClient  *retryablehttp.Client

	conflictHandlingAttempted bool
}

func (c *AtlasClient) Get() (*Payload, error) {
	// Make the HTTP request
	req, err := retryablehttp.NewRequest("GET", c.url().String(), nil)
	if err != nil {
		return nil, fmt.Errorf("Failed to make HTTP request: %v", err)
	}

	// Request the url
	client, err := c.http()
	if err != nil {
		return nil, err
	}
	resp, err := client.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	// Handle the common status codes
	switch resp.StatusCode {
	case http.StatusOK:
		// Handled after
	case http.StatusNoContent:
		return nil, nil
	case http.StatusNotFound:
		return nil, nil
	case http.StatusUnauthorized:
		return nil, fmt.Errorf("HTTP remote state endpoint requires auth")
	case http.StatusForbidden:
		return nil, fmt.Errorf("HTTP remote state endpoint invalid auth")
	case http.StatusInternalServerError:
		return nil, fmt.Errorf("HTTP remote state internal server error")
	default:
		return nil, fmt.Errorf(
			"Unexpected HTTP response code: %d\n\nBody: %s",
			resp.StatusCode, c.readBody(resp.Body))
	}

	// Read in the body
	buf := bytes.NewBuffer(nil)
	if _, err := io.Copy(buf, resp.Body); err != nil {
		return nil, fmt.Errorf("Failed to read remote state: %v", err)
	}

	// Create the payload
	payload := &Payload{
		Data: buf.Bytes(),
	}

	if len(payload.Data) == 0 {
		return nil, nil
	}

	// Check for the MD5
	if raw := resp.Header.Get("Content-MD5"); raw != "" {
		md5, err := base64.StdEncoding.DecodeString(raw)
		if err != nil {
			return nil, fmt.Errorf("Failed to decode Content-MD5 '%s': %v", raw, err)
		}

		payload.MD5 = md5
	} else {
		// Generate the MD5
		hash := md5.Sum(payload.Data)
		payload.MD5 = hash[:]
	}

	return payload, nil
}

func (c *AtlasClient) Put(state []byte) error {
	// Get the target URL
	base := c.url()

	// Generate the MD5
	hash := md5.Sum(state)
	b64 := base64.StdEncoding.EncodeToString(hash[:])

	// Make the HTTP client and request
	req, err := retryablehttp.NewRequest("PUT", base.String(), bytes.NewReader(state))
	if err != nil {
		return fmt.Errorf("Failed to make HTTP request: %v", err)
	}

	// Prepare the request
	req.Header.Set("Content-MD5", b64)
	req.Header.Set("Content-Type", "application/json")
	req.ContentLength = int64(len(state))

	// Make the request
	client, err := c.http()
	if err != nil {
		return err
	}
	resp, err := client.Do(req)
	if err != nil {
		return fmt.Errorf("Failed to upload state: %v", err)
	}
	defer resp.Body.Close()

	// Handle the error codes
	switch resp.StatusCode {
	case http.StatusOK:
		return nil
	case http.StatusConflict:
		return c.handleConflict(c.readBody(resp.Body), state)
	default:
		return fmt.Errorf(
			"HTTP error: %d\n\nBody: %s",
			resp.StatusCode, c.readBody(resp.Body))
	}
}

func (c *AtlasClient) Delete() error {
	// Make the HTTP request
	req, err := retryablehttp.NewRequest("DELETE", c.url().String(), nil)
	if err != nil {
		return fmt.Errorf("Failed to make HTTP request: %v", err)
	}

	// Make the request
	client, err := c.http()
	if err != nil {
		return err
	}
	resp, err := client.Do(req)
	if err != nil {
		return fmt.Errorf("Failed to delete state: %v", err)
	}
	defer resp.Body.Close()

	// Handle the error codes
	switch resp.StatusCode {
	case http.StatusOK:
		return nil
	case http.StatusNoContent:
		return nil
	case http.StatusNotFound:
		return nil
	default:
		return fmt.Errorf(
			"HTTP error: %d\n\nBody: %s",
			resp.StatusCode, c.readBody(resp.Body))
	}
}

func (c *AtlasClient) readBody(b io.Reader) string {
	var buf bytes.Buffer
	if _, err := io.Copy(&buf, b); err != nil {
		return fmt.Sprintf("Error reading body: %s", err)
	}

	result := buf.String()
	if result == "" {
		result = "<empty>"
	}

	return result
}

func (c *AtlasClient) url() *url.URL {
	values := url.Values{}

	values.Add("atlas_run_id", c.RunId)
	values.Add("access_token", c.AccessToken)

	return &url.URL{
		Scheme:   c.ServerURL.Scheme,
		Host:     c.ServerURL.Host,
		Path:     path.Join("api/v1/terraform/state", c.User, c.Name),
		RawQuery: values.Encode(),
	}
}

func (c *AtlasClient) http() (*retryablehttp.Client, error) {
	if c.HTTPClient != nil {
		return c.HTTPClient, nil
	}
	tlsConfig := &tls.Config{}
	err := rootcerts.ConfigureTLS(tlsConfig, &rootcerts.Config{
		CAFile: os.Getenv("ATLAS_CAFILE"),
		CAPath: os.Getenv("ATLAS_CAPATH"),
	})
	if err != nil {
		return nil, err
	}
	rc := retryablehttp.NewClient()
	t := cleanhttp.DefaultTransport()
	t.TLSClientConfig = tlsConfig
	rc.HTTPClient.Transport = t
	return rc, nil
}

// Atlas returns an HTTP 409 - Conflict if the pushed state reports the same
// Serial number but the checksum of the raw content differs. This can
// sometimes happen when Terraform changes state representation internally
// between versions in a way that's semantically neutral but affects the JSON
// output and therefore the checksum.
//
// Here we detect and handle this situation by ticking the serial and retrying
// iff for the previous state and the proposed state:
//
//   * the serials match
//   * the parsed states are Equal (semantically equivalent)
//
// In other words, in this situation Terraform can override Atlas's detected
// conflict by asserting that the state it is pushing is indeed correct.
func (c *AtlasClient) handleConflict(msg string, state []byte) error {
	log.Printf("[DEBUG] Handling Atlas conflict response: %s", msg)

	if c.conflictHandlingAttempted {
		log.Printf("[DEBUG] Already attempted conflict resolution; returning conflict.")
	} else {
		c.conflictHandlingAttempted = true
		log.Printf("[DEBUG] Atlas reported conflict, checking for equivalent states.")

		payload, err := c.Get()
		if err != nil {
			return conflictHandlingError(err)
		}

		currentState, err := terraform.ReadState(bytes.NewReader(payload.Data))
		if err != nil {
			return conflictHandlingError(err)
		}

		proposedState, err := terraform.ReadState(bytes.NewReader(state))
		if err != nil {
			return conflictHandlingError(err)
		}

		if statesAreEquivalent(currentState, proposedState) {
			log.Printf("[DEBUG] States are equivalent, incrementing serial and retrying.")
			proposedState.Serial++
			var buf bytes.Buffer
			if err := terraform.WriteState(proposedState, &buf); err != nil {
				return conflictHandlingError(err)
			}
			return c.Put(buf.Bytes())
		} else {
			log.Printf("[DEBUG] States are not equivalent, returning conflict.")
		}
	}

	return fmt.Errorf(
		"Atlas detected a remote state conflict.\n\nMessage: %s", msg)
}

func conflictHandlingError(err error) error {
	return fmt.Errorf(
		"Error while handling a conflict response from Atlas: %s", err)
}

func statesAreEquivalent(current, proposed *terraform.State) bool {
	return current.Serial == proposed.Serial && current.Equal(proposed)
}
<|endoftext|>


<commit_before>package atom

import (
	"encoding/binary"
	"fmt"
	"io"
	"time"
)

type AtomType interface {
	Size() uint32
	Type() string
	Print()
}

type Atom struct {
	size     uint32
	atype    string
	children []*AtomType
	Print    func()
	Type     func()
}

var atom_parsers map[string]func(Atom, io.Reader) *AtomType

// ftyp
type ftypAtom struct {
	Atom
	major_brand       string
	minor_version     string
	compatible_brands []string
}

func (fa *ftypAtom) Print() {

	fmt.Printf("size: %d\n", fa.size)
	fmt.Printf("type: %s\n", fa.atype)
	fmt.Printf("Major_brand: %s\n", fa.major_brand)
	fmt.Printf("Minor_version: %s\n", fa.minor_version)

	for _, v := range fa.compatible_brands {
		fmt.Printf("Compatible_brands: %s\n", v)
	}
}

func (fa *ftypAtom) Size() uint32 {
	return fa.size
}

func (fa *ftypAtom) Type() string {
	return fa.atype
}

func parse_ftyp(a Atom, r io.Reader) *AtomType {

	var ftyp ftypAtom

	ftyp.size = a.size
	ftyp.atype = a.atype

	buf := make([]byte, 4)

	r.Read(buf)
	ftyp.major_brand = string(buf)

	r.Read(buf)
	ftyp.minor_version = string(buf)

	remain_size := ftyp.size - 16

	for i := uint32(0); i < remain_size; i += 4 {
		r.Read(buf)
		ftyp.compatible_brands = append(ftyp.compatible_brands, string(buf))
	}

	var ret AtomType = &ftyp

	return &ret
}

// moov
type moovAtom struct {
	Atom
}

func (ma *moovAtom) Print() {
	fmt.Printf("size: %d\n", ma.size)
	fmt.Printf("type: %s\n", ma.atype)
}

func (ma *moovAtom) Size() uint32 {
	return ma.size
}

func (ma *moovAtom) Type() string {
	return ma.atype
}

func parse_moov(a Atom, r io.Reader) *AtomType {

	var moov moovAtom

	moov.size = a.size
	moov.atype = a.atype

	moov.children = Parse_atom(r)

	var ret AtomType = &moov

	return &ret
}

func (a *Atom) Parse(r io.Reader) {

}

// mvhd
type mvhdAtom struct {
	Atom
	version            uint32
	flags              uint32
	creation_time      time.Time
	modification_time  time.Time
	time_scale         uint32
	duration           uint32
	preferred_rate     uint32
	matrix_structure   [3][3]uint32
	preview_time       time.Time
	preview_duration   uint32
	poster_time        time.Time
	selection_time     time.Time
	selection_duration uint32
	current_time       time.Time
	next_track_id      uint32
}

func (ma *mvhdAtom) Print() {
	fmt.Printf("size: %d\n", ma.size)
	fmt.Printf("type: %s\n", ma.atype)
	fmt.Printf("version: %d\n", ma.version)
	fmt.Printf("flags: %x\n", ma.flags)
}

func (ma *mvhdAtom) Size() uint32 {
	return ma.size
}

func (ma *mvhdAtom) Type() string {
	return ma.atype
}

func parse_mvhd(a Atom, r io.Reader) *AtomType {
	var mvhd mvhdAtom

	mvhd.size = a.size
	mvhd.atype = a.atype

	var tmp uint32
	binary.Read(r, binary.LittleEndian, &tmp)
	mvhd.version = (tmp >> 24) & 0xff
	mvhd.flags = tmp & 0xffffff

	fmt.Printf("(%s) %d bytes were ignored\n", a.atype, a.size-12)

	buf := make([]byte, a.size-12)
	r.Read(buf)

	var ret AtomType = &mvhd

	return &ret
}

// free
type freeAtom struct {
	Atom
}

func (fa *freeAtom) Print() {
	fmt.Printf("size: %d\n", fa.size)
	fmt.Printf("type: %s\n", fa.atype)
}

func (fa *freeAtom) Size() uint32 {
	return fa.size
}

func (fa *freeAtom) Type() string {
	return fa.atype
}

func parse_free(a Atom, r io.Reader) *AtomType {
	var free freeAtom

	free.size = a.size
	free.atype = a.atype

	var ret AtomType = &free

	return &ret
}

// general
type generalAtom struct {
	Atom
}

func (ga *generalAtom) Print() {
	fmt.Printf("size: %d\n", ga.size)
	fmt.Printf("type: %s\n", ga.atype)
}

func (ga *generalAtom) Size() uint32 {
	return ga.size
}

func (ga *generalAtom) Type() string {
	return ga.atype
}

func parse_general(a Atom, r io.Reader) *AtomType {

	var ga generalAtom

	ga.size = a.size
	ga.atype = a.atype

	// skip
	fmt.Printf("(%s) %d bytes were ignored\n", a.atype, a.size-8)
	buf := make([]byte, a.size-8)
	r.Read(buf)
	//

	var ret AtomType = &ga

	return &ret

}

func Parse_atom(r io.Reader) []*AtomType {

	var atom Atom

	var atoms = make([]*AtomType, 0)

	buf := make([]byte, 4)

	for binary.Read(r, binary.BigEndian, &atom.size) == nil {

		r.Read(buf)
		atom.atype = string(buf)

		if atom_parsers[atom.atype] != nil {

			mp4 := *atom_parsers[atom.atype](atom, r)
			mp4.Print()

			atoms = append(atoms, &mp4)

		} else {

			mp4 := *parse_general(atom, r)
			mp4.Print()

			atoms = append(atoms, &mp4)
		}
	}

	return atoms
}

func init() {

	atom_parsers = map[string]func(Atom, io.Reader) *AtomType{
		"ftyp": parse_ftyp,
		"moov": parse_moov,
		"mvhd": parse_mvhd,
		"free": parse_free,
	}
}
<commit_msg>improve code by using map<commit_after>package atom

import (
	"encoding/binary"
	"fmt"
	"io"
	//	"time"
)

type Atom struct {
	size     uint32
	atype    string
	children []*Atom
	elements map[string]interface{}
}

var atom_parsers map[string]func(*Atom, io.Reader) *Atom

// ftyp
func parse_ftyp(a *Atom, r io.Reader) *Atom {

	buf := make([]byte, 4)
	el := make(map[string]interface{})

	r.Read(buf)
	el["major_brand"] = string(buf)

	r.Read(buf)
	el["minor_version"] = string(buf)

	remain_size := a.size - 16
	tmp := make([]string, remain_size/4+1)

	for i := uint32(0); i < remain_size; i += 4 {
		r.Read(buf)
		tmp = append(tmp, string(buf))
	}
	el["compatible_brands"] = tmp

	a.elements = el

	return a
}

// moov
func parse_moov(a *Atom, r io.Reader) *Atom {

	a.children = Parse_atom(r)

	return a
}

// mvhd
func parse_mvhd(a *Atom, r io.Reader) *Atom {

	el := make(map[string]interface{})
	var tmp uint32

	binary.Read(r, binary.LittleEndian, &tmp)
	el["version"] = (tmp >> 24) & 0xff
	el["flags"] = tmp & 0xffffff

	fmt.Printf("(%s) %d bytes were ignored\n", a.atype, a.size-12)

	buf := make([]byte, a.size-12)
	r.Read(buf)

	a.elements = el

	return a
}

// free
func parse_free(a *Atom, r io.Reader) *Atom {

	fmt.Printf("(%s) %d bytes were ignored\n", a.atype, a.size-8)
	buf := make([]byte, a.size-8)
	r.Read(buf)

	return a
}

// general
func parse_general(a *Atom, r io.Reader) *Atom {

	// skip
	fmt.Printf("(%s) %d bytes were ignored\n", a.atype, a.size-8)
	buf := make([]byte, a.size-8)
	r.Read(buf)
	//

	return a
}

func Print_atom(a *Atom) {

	fmt.Printf("type: %s\n", a.atype)
	fmt.Printf("size: %d\n", a.size)

	for k, v := range a.elements {
		fmt.Printf("%s: %v\n", k, v)
	}

	fmt.Printf("\n")
}

func Parse_atom(r io.Reader) []*Atom {

	var atom Atom

	var atoms = make([]*Atom, 0)

	buf := make([]byte, 4)

	for binary.Read(r, binary.BigEndian, &atom.size) == nil {

		var mp4 *Atom

		r.Read(buf)
		atom.atype = string(buf)

		if atom_parsers[atom.atype] != nil {
			mp4 = atom_parsers[atom.atype](&atom, r)
		} else {
			mp4 = parse_general(&atom, r)
		}

		Print_atom(mp4)
		atoms = append(atoms, mp4)
	}

	return atoms
}

func init() {

	atom_parsers = map[string]func(*Atom, io.Reader) *Atom{
		"ftyp": parse_ftyp,
		"moov": parse_moov,
		"mvhd": parse_mvhd,
		"free": parse_free,
	}
}
<|endoftext|>


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

package main

import (
	"bufio"
	"flag"
	"fmt"
	"io"
	"log"
	"os"
	"path/filepath"
	"runtime"
	"sync"
	"time"

	"github.com/richardlehane/siegfried"
	"github.com/richardlehane/siegfried/config"
	"github.com/richardlehane/siegfried/pkg/core/siegreader"

	/*// Uncomment to build with profiler
	"net/http"
	_ "net/http/pprof"
	*/)

const PROCS = -1

// flags
var (
	update    = flag.Bool("update", false, "update or install the default signature file")
	version   = flag.Bool("version", false, "display version information")
	logf      = flag.String("log", "error", "log errors, warnings, debug or slow output, knowns or unknowns to stderr or stdout e.g. -log error,warn,unknown,stdout")
	nr        = flag.Bool("nr", false, "prevent automatic directory recursion")
	csvo      = flag.Bool("csv", false, "CSV output format")
	jsono     = flag.Bool("json", false, "JSON output format")
	droido    = flag.Bool("droid", false, "DROID CSV output format")
	sig       = flag.String("sig", config.SignatureBase(), "set the signature file")
	home      = flag.String("home", config.Home(), "override the default home directory")
	serve     = flag.String("serve", "", "start siegfried server e.g. -serve localhost:5138")
	multi     = flag.Int("multi", 1, "set number of file ID processes")
	archive   = flag.Bool("z", false, "scan archive formats (zip, tar, gzip, warc, arc)")
	hashf     = flag.String("hash", "", "calculate file checksum with hash algorithm; options "+hashChoices)
	throttlef = flag.Duration("throttle", 0, "set a time to wait between scanning files e.g. 50ms")
)

var throttle *time.Ticker

func writeError(w writer, path string, sz int64, mod string, err error) {
	w.writeFile(path, sz, mod, nil, err, nil)
	// log the error too
	lg.set(path)
	lg.err(err)
	lg.reset()
}

type res struct {
	path string
	sz   int64
	mod  string
	c    iterableID
	err  error
}

func printer(w writer, resc chan chan res, wg *sync.WaitGroup) {
	for rr := range resc {
		r := <-rr
		w.writeFile(r.path, r.sz, r.mod, nil, r.err, r.c)
		wg.Done()
	}
}

func multiIdentifyP(w writer, s *siegfried.Siegfried, r string, norecurse bool) error {
	wg := &sync.WaitGroup{}
	runtime.GOMAXPROCS(PROCS)
	resc := make(chan chan res, *multi)
	go printer(w, resc, wg)
	wf := func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return fmt.Errorf("walking %s; got %v", path, err)
		}
		if info.IsDir() {
			if norecurse && path != r {
				return filepath.SkipDir
			}
			if *droido {
				wg.Add(1)
				rchan := make(chan res, 1)
				resc <- rchan
				go func() {
					rchan <- res{path, -1, info.ModTime().String(), nil, nil} // write directory with a -1 size for droid output only
				}()
			}
			return nil
		}
		wg.Add(1)
		rchan := make(chan res, 1)
		resc <- rchan
		go func() {
			f, err := os.Open(path)
			if err != nil {
				rchan <- res{path, 0, "", nil, err.(*os.PathError).Err} // return summary error only
				return
			}
			c, err := s.Identify(f, path, "")
			if c == nil {
				f.Close()
				rchan <- res{path, 0, "", nil, err}
				return
			}
			ids := makeIdSlice(idChan(c))
			f.Close()
			rchan <- res{path, info.Size(), info.ModTime().Format(time.RFC3339), ids, err}
		}()
		return nil
	}
	err := filepath.Walk(r, wf)
	wg.Wait()
	close(resc)
	return err
}

// multiIdentifyS() defined in longpath.go and longpath_windows.go

func identifyFile(w writer, s *siegfried.Siegfried, path string, sz int64, mod string) {
	f, err := os.Open(path)
	if err != nil {
		f, err = retryOpen(path, err) // retry open in case is a windows long path error
		if err != nil {
			writeError(w, path, sz, mod, err.(*os.PathError).Err) // write summary error
			return
		}
	}
	identifyRdr(w, s, f, sz, path, "", mod)
	f.Close()
}

func identifyRdr(w writer, s *siegfried.Siegfried, r io.Reader, sz int64, path, mime, mod string) {
	lg.set(path)
	c, err := s.Identify(r, path, mime)
	lg.err(err)
	if c == nil {
		w.writeFile(path, sz, mod, nil, err, nil)
		lg.reset()
		return
	}
	var b *siegreader.Buffer
	var cs []byte
	if checksum != nil {
		b = s.Buffer()
		var i int64
		l := checksum.BlockSize()
		for ; ; i += int64(l) {
			buf, _ := b.Slice(i, l)
			if buf == nil {
				break
			}
			checksum.Write(buf)
		}
		cs = checksum.Sum(nil)
		checksum.Reset()
	}
	a := w.writeFile(path, sz, mod, cs, err, idChan(c))
	lg.reset()
	if !*archive || a == config.None {
		return
	}
	var d decompressor
	if b == nil {
		b = s.Buffer()
	}
	switch a {
	case config.Zip:
		d, err = newZip(siegreader.ReaderFrom(b), path, sz)
	case config.Gzip:
		d, err = newGzip(b, path)
	case config.Tar:
		d, err = newTar(siegreader.ReaderFrom(b), path)
	case config.ARC:
		d, err = newARC(siegreader.ReaderFrom(b), path)
	case config.WARC:
		d, err = newWARC(siegreader.ReaderFrom(b), path)
	}
	if err != nil {
		writeError(w, path, sz, mod, fmt.Errorf("failed to decompress, got: %v", err))
		return
	}
	for err = d.next(); err == nil; err = d.next() {
		if *droido {
			for _, v := range d.dirs() {
				w.writeFile(v, -1, "", nil, nil, nil)
			}
		}
		identifyRdr(w, s, d.reader(), d.size(), d.path(), d.mime(), d.mod())
	}
}

func main() {

	flag.Parse()

	/*//UNCOMMENT TO RUN PROFILER
	go func() {
		log.Println(http.ListenAndServe("localhost:6060", nil))
	}()*/

	if *version {
		version := config.Version()
		fmt.Printf("siegfried %d.%d.%d\n", version[0], version[1], version[2])
		s, err := siegfried.Load(config.Signature())
		if err != nil {
			fmt.Println(err)
			return
		}
		fmt.Print(s)
		return
	}

	if *home != config.Home() {
		config.SetHome(*home)
	}

	if *sig != config.SignatureBase() {
		config.SetSignature(*sig)
	}

	if *update {
		msg, err := updateSigs()
		if err != nil {
			log.Fatalf("[FATAL] failed to update signature file, %v", err)
		}
		fmt.Println(msg)
		return
	}

	// during parallel scanning or in server mode, unsafe to access the last read buffer - so can't unzip or hash
	if *multi > 1 || *serve != "" {
		if *archive {
			log.Fatalln("[FATAL] cannot scan archive formats when running in parallel or server mode")
		}
		if *hashf != "" {
			log.Fatalln("[FATAL] cannot calculate file checksum when running in parallel or server mode")
		}
	}

	if *logf != "" {
		if *multi > 1 && *logf != "error" {
			log.Fatalln("[FATAL] cannot log in parallel mode")
		}
		if err := newLogger(*logf); err != nil {
			log.Fatalln(err)
		}
	}

	if err := setHash(); err != nil {
		log.Fatal(err)
	}

	if *serve != "" || *fprflag {
		s, err := siegfried.Load(config.Signature())
		if err != nil {
			log.Fatalf("[FATAL] error loading signature file, got: %v", err)
		}
		if *serve != "" {
			log.Printf("Starting server at %s. Use CTRL-C to quit.\n", *serve)
			listen(*serve, s)
			return
		}
		log.Printf("FPR server started at %s. Use CTRL-C to quit.\n", config.Fpr())
		serveFpr(config.Fpr(), s)
		return
	}

	if flag.NArg() != 1 {
		log.Fatalln("[FATAL] expecting a single file or directory argument")
	}

	s, err := siegfried.Load(config.Signature())
	if err != nil {
		log.Fatalf("[FATAL] error loading signature file, got: %v", err)
	}

	var w writer
	switch {
	case *csvo:
		w = newCSV(os.Stdout)
	case *jsono:
		w = newJSON(os.Stdout)
	case *droido:
		w = newDroid(os.Stdout)
	default:
		w = newYAML(os.Stdout)
	}

	if lg != nil && lg.w == os.Stdout {
		w = logWriter{}
	}

	// support reading list files from stdin
	if flag.Arg(0) == "-" {
		w.writeHead(s)
		scanner := bufio.NewScanner(os.Stdin)
		for scanner.Scan() {
			info, err := os.Stat(scanner.Text())
			if err != nil {
				info, err = retryStat(scanner.Text(), err)
			}
			if err != nil || info.IsDir() {
				writeError(w, scanner.Text(), 0, "", fmt.Errorf("failed to identify %s (in scanning mode, inputs must all be files and not directories), got: %v", scanner.Text(), err))
			} else {
				identifyFile(w, s, scanner.Text(), info.Size(), info.ModTime().Format(time.RFC3339))
			}
		}
		w.writeTail()
		lg.printElapsed()
		os.Exit(0)
	}

	info, err := os.Stat(flag.Arg(0))
	if err != nil {
		info, err = retryStat(flag.Arg(0), err)
		if err != nil {
			log.Fatalf("[FATAL] cannot get info for %v, got: %v", flag.Arg(0), err)
		}
	}

	if info.IsDir() {
		w.writeHead(s)
		if *multi > 16 {
			*multi = 16
		}
		if *multi > 1 {
			err = multiIdentifyP(w, s, flag.Arg(0), *nr)
		} else {
			if *throttlef != 0 {
				throttle = time.NewTicker(*throttlef)
				defer throttle.Stop()
			}
			err = multiIdentifyS(w, s, flag.Arg(0), "", *nr)
		}
		w.writeTail()
		if err != nil {
			log.Fatalf("[FATAL] %v\n", err)
		}
		lg.printElapsed()
		os.Exit(0)
	}
	w.writeHead(s)
	identifyFile(w, s, flag.Arg(0), info.Size(), info.ModTime().Format(time.RFC3339))
	w.writeTail()
	lg.printElapsed()
	os.Exit(0)
}
<commit_msg>limit DROID to single PRONOM IDs<commit_after>// Copyright 2014 Richard Lehane. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package main

import (
	"bufio"
	"flag"
	"fmt"
	"io"
	"log"
	"os"
	"path/filepath"
	"runtime"
	"sync"
	"time"

	"github.com/richardlehane/siegfried"
	"github.com/richardlehane/siegfried/config"
	"github.com/richardlehane/siegfried/pkg/core/siegreader"

	/*// Uncomment to build with profiler
	"net/http"
	_ "net/http/pprof"
	*/)

const PROCS = -1

// flags
var (
	update    = flag.Bool("update", false, "update or install the default signature file")
	version   = flag.Bool("version", false, "display version information")
	logf      = flag.String("log", "error", "log errors, warnings, debug or slow output, knowns or unknowns to stderr or stdout e.g. -log error,warn,unknown,stdout")
	nr        = flag.Bool("nr", false, "prevent automatic directory recursion")
	csvo      = flag.Bool("csv", false, "CSV output format")
	jsono     = flag.Bool("json", false, "JSON output format")
	droido    = flag.Bool("droid", false, "DROID CSV output format")
	sig       = flag.String("sig", config.SignatureBase(), "set the signature file")
	home      = flag.String("home", config.Home(), "override the default home directory")
	serve     = flag.String("serve", "", "start siegfried server e.g. -serve localhost:5138")
	multi     = flag.Int("multi", 1, "set number of file ID processes")
	archive   = flag.Bool("z", false, "scan archive formats (zip, tar, gzip, warc, arc)")
	hashf     = flag.String("hash", "", "calculate file checksum with hash algorithm; options "+hashChoices)
	throttlef = flag.Duration("throttle", 0, "set a time to wait between scanning files e.g. 50ms")
)

var throttle *time.Ticker

func writeError(w writer, path string, sz int64, mod string, err error) {
	w.writeFile(path, sz, mod, nil, err, nil)
	// log the error too
	lg.set(path)
	lg.err(err)
	lg.reset()
}

type res struct {
	path string
	sz   int64
	mod  string
	c    iterableID
	err  error
}

func printer(w writer, resc chan chan res, wg *sync.WaitGroup) {
	for rr := range resc {
		r := <-rr
		w.writeFile(r.path, r.sz, r.mod, nil, r.err, r.c)
		wg.Done()
	}
}

func multiIdentifyP(w writer, s *siegfried.Siegfried, r string, norecurse bool) error {
	wg := &sync.WaitGroup{}
	runtime.GOMAXPROCS(PROCS)
	resc := make(chan chan res, *multi)
	go printer(w, resc, wg)
	wf := func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return fmt.Errorf("walking %s; got %v", path, err)
		}
		if info.IsDir() {
			if norecurse && path != r {
				return filepath.SkipDir
			}
			if *droido {
				wg.Add(1)
				rchan := make(chan res, 1)
				resc <- rchan
				go func() {
					rchan <- res{path, -1, info.ModTime().String(), nil, nil} // write directory with a -1 size for droid output only
				}()
			}
			return nil
		}
		wg.Add(1)
		rchan := make(chan res, 1)
		resc <- rchan
		go func() {
			f, err := os.Open(path)
			if err != nil {
				rchan <- res{path, 0, "", nil, err.(*os.PathError).Err} // return summary error only
				return
			}
			c, err := s.Identify(f, path, "")
			if c == nil {
				f.Close()
				rchan <- res{path, 0, "", nil, err}
				return
			}
			ids := makeIdSlice(idChan(c))
			f.Close()
			rchan <- res{path, info.Size(), info.ModTime().Format(time.RFC3339), ids, err}
		}()
		return nil
	}
	err := filepath.Walk(r, wf)
	wg.Wait()
	close(resc)
	return err
}

// multiIdentifyS() defined in longpath.go and longpath_windows.go

func identifyFile(w writer, s *siegfried.Siegfried, path string, sz int64, mod string) {
	f, err := os.Open(path)
	if err != nil {
		f, err = retryOpen(path, err) // retry open in case is a windows long path error
		if err != nil {
			writeError(w, path, sz, mod, err.(*os.PathError).Err) // write summary error
			return
		}
	}
	identifyRdr(w, s, f, sz, path, "", mod)
	f.Close()
}

func identifyRdr(w writer, s *siegfried.Siegfried, r io.Reader, sz int64, path, mime, mod string) {
	lg.set(path)
	c, err := s.Identify(r, path, mime)
	lg.err(err)
	if c == nil {
		w.writeFile(path, sz, mod, nil, err, nil)
		lg.reset()
		return
	}
	var b *siegreader.Buffer
	var cs []byte
	if checksum != nil {
		b = s.Buffer()
		var i int64
		l := checksum.BlockSize()
		for ; ; i += int64(l) {
			buf, _ := b.Slice(i, l)
			if buf == nil {
				break
			}
			checksum.Write(buf)
		}
		cs = checksum.Sum(nil)
		checksum.Reset()
	}
	a := w.writeFile(path, sz, mod, cs, err, idChan(c))
	lg.reset()
	if !*archive || a == config.None {
		return
	}
	var d decompressor
	if b == nil {
		b = s.Buffer()
	}
	switch a {
	case config.Zip:
		d, err = newZip(siegreader.ReaderFrom(b), path, sz)
	case config.Gzip:
		d, err = newGzip(b, path)
	case config.Tar:
		d, err = newTar(siegreader.ReaderFrom(b), path)
	case config.ARC:
		d, err = newARC(siegreader.ReaderFrom(b), path)
	case config.WARC:
		d, err = newWARC(siegreader.ReaderFrom(b), path)
	}
	if err != nil {
		writeError(w, path, sz, mod, fmt.Errorf("failed to decompress, got: %v", err))
		return
	}
	for err = d.next(); err == nil; err = d.next() {
		if *droido {
			for _, v := range d.dirs() {
				w.writeFile(v, -1, "", nil, nil, nil)
			}
		}
		identifyRdr(w, s, d.reader(), d.size(), d.path(), d.mime(), d.mod())
	}
}

func main() {

	flag.Parse()

	/*//UNCOMMENT TO RUN PROFILER
	go func() {
		log.Println(http.ListenAndServe("localhost:6060", nil))
	}()*/

	if *version {
		version := config.Version()
		fmt.Printf("siegfried %d.%d.%d\n", version[0], version[1], version[2])
		s, err := siegfried.Load(config.Signature())
		if err != nil {
			fmt.Println(err)
			return
		}
		fmt.Print(s)
		return
	}

	if *home != config.Home() {
		config.SetHome(*home)
	}

	if *sig != config.SignatureBase() {
		config.SetSignature(*sig)
	}

	if *update {
		msg, err := updateSigs()
		if err != nil {
			log.Fatalf("[FATAL] failed to update signature file, %v", err)
		}
		fmt.Println(msg)
		return
	}

	// during parallel scanning or in server mode, unsafe to access the last read buffer - so can't unzip or hash
	if *multi > 1 || *serve != "" {
		if *archive {
			log.Fatalln("[FATAL] cannot scan archive formats when running in parallel or server mode")
		}
		if *hashf != "" {
			log.Fatalln("[FATAL] cannot calculate file checksum when running in parallel or server mode")
		}
	}

	if *logf != "" {
		if *multi > 1 && *logf != "error" {
			log.Fatalln("[FATAL] cannot log in parallel mode")
		}
		if err := newLogger(*logf); err != nil {
			log.Fatalln(err)
		}
	}

	if err := setHash(); err != nil {
		log.Fatal(err)
	}

	if *serve != "" || *fprflag {
		s, err := siegfried.Load(config.Signature())
		if err != nil {
			log.Fatalf("[FATAL] error loading signature file, got: %v", err)
		}
		if *serve != "" {
			log.Printf("Starting server at %s. Use CTRL-C to quit.\n", *serve)
			listen(*serve, s)
			return
		}
		log.Printf("FPR server started at %s. Use CTRL-C to quit.\n", config.Fpr())
		serveFpr(config.Fpr(), s)
		return
	}

	if flag.NArg() != 1 {
		log.Fatalln("[FATAL] expecting a single file or directory argument")
	}

	s, err := siegfried.Load(config.Signature())
	if err != nil {
		log.Fatalf("[FATAL] error loading signature file, got: %v", err)
	}

	var w writer
	switch {
	case *csvo:
		w = newCSV(os.Stdout)
	case *jsono:
		w = newJSON(os.Stdout)
	case *droido:
		w = newDroid(os.Stdout)
		if len(s.Fields()) != 1 || len(s.Fields()[0]) != 7 {
			log.Fatalln("[FATAL] DROID output is limited to signature files with a single PRONOM identifier")
		}
	default:
		w = newYAML(os.Stdout)
	}

	if lg != nil && lg.w == os.Stdout {
		w = logWriter{}
	}

	// support reading list files from stdin
	if flag.Arg(0) == "-" {
		w.writeHead(s)
		scanner := bufio.NewScanner(os.Stdin)
		for scanner.Scan() {
			info, err := os.Stat(scanner.Text())
			if err != nil {
				info, err = retryStat(scanner.Text(), err)
			}
			if err != nil || info.IsDir() {
				writeError(w, scanner.Text(), 0, "", fmt.Errorf("failed to identify %s (in scanning mode, inputs must all be files and not directories), got: %v", scanner.Text(), err))
			} else {
				identifyFile(w, s, scanner.Text(), info.Size(), info.ModTime().Format(time.RFC3339))
			}
		}
		w.writeTail()
		lg.printElapsed()
		os.Exit(0)
	}

	info, err := os.Stat(flag.Arg(0))
	if err != nil {
		info, err = retryStat(flag.Arg(0), err)
		if err != nil {
			log.Fatalf("[FATAL] cannot get info for %v, got: %v", flag.Arg(0), err)
		}
	}

	if info.IsDir() {
		w.writeHead(s)
		if *multi > 16 {
			*multi = 16
		}
		if *multi > 1 {
			err = multiIdentifyP(w, s, flag.Arg(0), *nr)
		} else {
			if *throttlef != 0 {
				throttle = time.NewTicker(*throttlef)
				defer throttle.Stop()
			}
			err = multiIdentifyS(w, s, flag.Arg(0), "", *nr)
		}
		w.writeTail()
		if err != nil {
			log.Fatalf("[FATAL] %v\n", err)
		}
		lg.printElapsed()
		os.Exit(0)
	}
	w.writeHead(s)
	identifyFile(w, s, flag.Arg(0), info.Size(), info.ModTime().Format(time.RFC3339))
	w.writeTail()
	lg.printElapsed()
	os.Exit(0)
}
<|endoftext|>


<commit_before>package wca

const (
	eConsole         = 0x0
	eMultimedia      = 0x1
	eCommunications  = 0x2
	ERole_enum_count = 0x3
)

const (
	DELETE       = 0x00010000
	READ_CONTROL = 0x00020000
	SYNCHRONIZE  = 0x00100000
	WRITE_DAC    = 0x00040000
	WRITE_OWNER  = 0x00080000
)

const (
	EVENT_ALL_ACCESS   = 0x1F0003
	EVENT_MODIFY_STATE = 0x0002
)

const (
	CREATE_EVENT_INITIAL_SET  = 0x00000002
	CREATE_EVENT_MANUAL_RESET = 0x00000001
)

// The following constants are defined in Audioclient.h.
const (
	AUDCLNT_BUFFERFLAGS_DATA_DISCONTINUITY = 0x1
	AUDCLNT_BUFFERFLAGS_SILENT             = 0x2
	AUDCLNT_BUFFERFLAGS_TIMESTAMP_ERROR    = 0x4
)

const (
	AUDCLNT_STREAMFLAGS_CROSSPROCESS        = 0x00010000
	AUDCLNT_STREAMFLAGS_LOOPBACK            = 0x00020000
	AUDCLNT_STREAMFLAGS_EVENTCALLBACK       = 0x00040000
	AUDCLNT_STREAMFLAGS_NOPERSIST           = 0x00080000
	AUDCLNT_STREAMFLAGS_RATEADJUST          = 0x00100000
	AUDCLNT_STREAMFLAGS_AUTOCONVERTPCM      = 0x80000000
	AUDCLNT_STREAMFLAGS_SRC_DEFAULT_QUALITY = 0x08000000
)

// the following AUDCLNT_SHAREMODE enumeration values are defined in audiosessiontypes.h.
const (
	AUDCLNT_SHAREMODE_SHARED    = 0x0
	AUDCLNT_SHAREMODE_EXCLUSIVE = 0x1
)

const (
	ENDPOINT_SYSFX_ENABLED  = 0x00000000
	ENDPOINT_SYSFX_DISABLED = 0x00000001
)

const (
	DEVICE_STATE_ACTIVE     = 0x00000001
	DEVICE_STATE_DISABLED   = 0x00000002
	DEVICE_STATE_NOTPRESENT = 0x00000004
	DEVICE_STATE_UNPLUGGED  = 0x00000008
	DEVICE_STATEMASK_ALL    = 0x0000000F
)

const (
	ERender              = 0x0
	ECapture             = 0x1
	EAll                 = 0x2
	EDataFlow_enum_count = 0x3
)

const (
	STGM_READ       = 0x0
	STGM_WRITE      = 0x1
	STGM_READ_WRITE = 0x2
)

const (
	CLSCTX_INPROC_SERVER          = 0x1
	CLSCTX_INPROC_HANDLER         = 0x2
	CLSCTX_LOCAL_SERVER           = 0x4
	CLSCTX_INPROC_SERVER16        = 0x8
	CLSCTX_REMOTE_SERVER          = 0x10
	CLSCTX_INPROC_HANDLER16       = 0x20
	CLSCTX_RESERVED1              = 0x40
	CLSCTX_RESERVED2              = 0x80
	CLSCTX_RESERVED3              = 0x100
	CLSCTX_RESERVED4              = 0x200
	CLSCTX_NO_CODE_DOWNLOAD       = 0x400
	CLSCTX_RESERVED5              = 0x800
	CLSCTX_NO_CUSTOM_MARSHAL      = 0x1000
	CLSCTX_ENABLE_CODE_DOWNLOAD   = 0x2000
	CLSCTX_NO_FAILURE_LOG         = 0x4000
	CLSCTX_DISABLE_AAA            = 0x8000
	CLSCTX_ENABLE_AAA             = 0x10000
	CLSCTX_FROM_DEFAULT_CONTEXT   = 0x20000
	CLSCTX_ACTIVATE_32_BIT_SERVER = 0x40000
	CLSCTX_ACTIVATE_64_BIT_SERVER = 0x80000
	CLSCTX_ENABLE_CLOAKING        = 0x100000
	CLSCTX_APPCONTAINER           = 0x400000
	CLSCTX_ACTIVATE_AAA_AS_IU     = 0x800000
	CLSCTX_PS_DLL                 = 0x80000000
)
<commit_msg>Add constant<commit_after>package wca

const (
	EConsole         = 0x0
	EMultimedia      = 0x1
	ECommunications  = 0x2
	ERole_enum_count = 0x3
)

const (
	DELETE       = 0x00010000
	READ_CONTROL = 0x00020000
	SYNCHRONIZE  = 0x00100000
	WRITE_DAC    = 0x00040000
	WRITE_OWNER  = 0x00080000
)

const (
	EVENT_ALL_ACCESS   = 0x1F0003
	EVENT_MODIFY_STATE = 0x0002
)

const (
	CREATE_EVENT_INITIAL_SET  = 0x00000002
	CREATE_EVENT_MANUAL_RESET = 0x00000001
)

// The following constants are defined in Audioclient.h.
const (
	AUDCLNT_BUFFERFLAGS_DATA_DISCONTINUITY = 0x1
	AUDCLNT_BUFFERFLAGS_SILENT             = 0x2
	AUDCLNT_BUFFERFLAGS_TIMESTAMP_ERROR    = 0x4
)

const (
	AUDCLNT_STREAMFLAGS_CROSSPROCESS        = 0x00010000
	AUDCLNT_STREAMFLAGS_LOOPBACK            = 0x00020000
	AUDCLNT_STREAMFLAGS_EVENTCALLBACK       = 0x00040000
	AUDCLNT_STREAMFLAGS_NOPERSIST           = 0x00080000
	AUDCLNT_STREAMFLAGS_RATEADJUST          = 0x00100000
	AUDCLNT_STREAMFLAGS_AUTOCONVERTPCM      = 0x80000000
	AUDCLNT_STREAMFLAGS_SRC_DEFAULT_QUALITY = 0x08000000
)

// the following AUDCLNT_SHAREMODE enumeration values are defined in audiosessiontypes.h.
const (
	AUDCLNT_SHAREMODE_SHARED    = 0x0
	AUDCLNT_SHAREMODE_EXCLUSIVE = 0x1
)

const (
	ENDPOINT_SYSFX_ENABLED  = 0x00000000
	ENDPOINT_SYSFX_DISABLED = 0x00000001
)

const (
	DEVICE_STATE_ACTIVE     = 0x00000001
	DEVICE_STATE_DISABLED   = 0x00000002
	DEVICE_STATE_NOTPRESENT = 0x00000004
	DEVICE_STATE_UNPLUGGED  = 0x00000008
	DEVICE_STATEMASK_ALL    = 0x0000000F
)

const (
	ERender              = 0x0
	ECapture             = 0x1
	EAll                 = 0x2
	EDataFlow_enum_count = 0x3
)

const (
	STGM_READ       = 0x0
	STGM_WRITE      = 0x1
	STGM_READ_WRITE = 0x2
)

const (
	CLSCTX_INPROC_SERVER          = 0x1
	CLSCTX_INPROC_HANDLER         = 0x2
	CLSCTX_LOCAL_SERVER           = 0x4
	CLSCTX_INPROC_SERVER16        = 0x8
	CLSCTX_REMOTE_SERVER          = 0x10
	CLSCTX_INPROC_HANDLER16       = 0x20
	CLSCTX_RESERVED1              = 0x40
	CLSCTX_RESERVED2              = 0x80
	CLSCTX_RESERVED3              = 0x100
	CLSCTX_RESERVED4              = 0x200
	CLSCTX_NO_CODE_DOWNLOAD       = 0x400
	CLSCTX_RESERVED5              = 0x800
	CLSCTX_NO_CUSTOM_MARSHAL      = 0x1000
	CLSCTX_ENABLE_CODE_DOWNLOAD   = 0x2000
	CLSCTX_NO_FAILURE_LOG         = 0x4000
	CLSCTX_DISABLE_AAA            = 0x8000
	CLSCTX_ENABLE_AAA             = 0x10000
	CLSCTX_FROM_DEFAULT_CONTEXT   = 0x20000
	CLSCTX_ACTIVATE_32_BIT_SERVER = 0x40000
	CLSCTX_ACTIVATE_64_BIT_SERVER = 0x80000
	CLSCTX_ENABLE_CLOAKING        = 0x100000
	CLSCTX_APPCONTAINER           = 0x400000
	CLSCTX_ACTIVATE_AAA_AS_IU     = 0x800000
	CLSCTX_PS_DLL                 = 0x80000000
	CLSCTX_ALL                    = CLSCTX_INPROC_SERVER | CLSCTX_INPROC_HANDLER | CLSCTX_LOCAL_SERVER | CLSCTX_REMOTE_SERVER
)
<|endoftext|>


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

package badger

import (
	"sync"

	"github.com/dgraph-io/badger/v2/y"
	"github.com/pkg/errors"
)

// WriteBatch holds the necessary info to perform batched writes.
type WriteBatch struct {
	sync.Mutex
	txn      *Txn
	db       *DB
	throttle *y.Throttle
	err      error
	commitTs uint64
}

// NewWriteBatch creates a new WriteBatch. This provides a way to conveniently do a lot of writes,
// batching them up as tightly as possible in a single transaction and using callbacks to avoid
// waiting for them to commit, thus achieving good performance. This API hides away the logic of
// creating and committing transactions. Due to the nature of SSI guaratees provided by Badger,
// blind writes can never encounter transaction conflicts (ErrConflict).
func (db *DB) NewWriteBatch() *WriteBatch {
	if db.opt.managedTxns {
		panic("cannot use NewWriteBatch in managed mode. Use NewWriteBatchAt instead")
	}
	return db.newWriteBatch()
}

func (db *DB) newWriteBatch() *WriteBatch {
	return &WriteBatch{
		db:       db,
		txn:      db.newTransaction(true, true),
		throttle: y.NewThrottle(16),
	}
}

// SetMaxPendingTxns sets a limit on maximum number of pending transactions while writing batches.
// This function should be called before using WriteBatch. Default value of MaxPendingTxns is
// 16 to minimise memory usage.
func (wb *WriteBatch) SetMaxPendingTxns(max int) {
	wb.throttle = y.NewThrottle(max)
}

// Cancel function must be called if there's a chance that Flush might not get
// called. If neither Flush or Cancel is called, the transaction oracle would
// never get a chance to clear out the row commit timestamp map, thus causing an
// unbounded memory consumption. Typically, you can call Cancel as a defer
// statement right after NewWriteBatch is called.
//
// Note that any committed writes would still go through despite calling Cancel.
func (wb *WriteBatch) Cancel() {
	if err := wb.throttle.Finish(); err != nil {
		wb.db.opt.Errorf("WatchBatch.Cancel error while finishing: %v", err)
	}
	wb.txn.Discard()
}

func (wb *WriteBatch) callback(err error) {
	// sync.WaitGroup is thread-safe, so it doesn't need to be run inside wb.Lock.
	defer wb.throttle.Done(err)
	if err == nil {
		return
	}

	wb.Lock()
	defer wb.Unlock()
	if wb.err != nil {
		return
	}
	wb.err = err
}

// SetEntryAt is the equivalent of Txn.SetEntry but it also allows setting version for the entry.
// SetEntryAt can be used only in managed mode.
func (wb *WriteBatch) SetEntryAt(e *Entry, ts uint64) error {
	if !wb.db.opt.managedTxns {
		return errors.New("SetEntryAt can only be used in managed mode. Use SetEntry instead")
	}
	e.version = ts
	return wb.SetEntry(e)
}

// SetEntry is the equivalent of Txn.SetEntry.
func (wb *WriteBatch) SetEntry(e *Entry) error {
	wb.Lock()
	defer wb.Unlock()

	if err := wb.txn.SetEntry(e); err != ErrTxnTooBig {
		return err
	}
	// Txn has reached it's zenith. Commit now.
	if cerr := wb.commit(); cerr != nil {
		return cerr
	}
	// This time the error must not be ErrTxnTooBig, otherwise, we make the
	// error permanent.
	if err := wb.txn.SetEntry(e); err != nil {
		wb.err = err
		return err
	}
	return nil
}

// Set is equivalent of Txn.Set().
func (wb *WriteBatch) Set(k, v []byte) error {
	e := &Entry{Key: k, Value: v}
	return wb.SetEntry(e)
}

// DeleteAt is equivalent of Txn.Delete but accepts a delete timestamp.
func (wb *WriteBatch) DeleteAt(k []byte, ts uint64) error {
	e := Entry{Key: k, meta: bitDelete, version: ts}
	return wb.SetEntry(&e)
}

// Delete is equivalent of Txn.Delete.
func (wb *WriteBatch) Delete(k []byte) error {
	wb.Lock()
	defer wb.Unlock()

	if err := wb.txn.Delete(k); err != ErrTxnTooBig {
		return err
	}
	if err := wb.commit(); err != nil {
		return err
	}
	if err := wb.txn.Delete(k); err != nil {
		wb.err = err
		return err
	}
	return nil
}

// Caller to commit must hold a write lock.
func (wb *WriteBatch) commit() error {
	if wb.err != nil {
		return wb.err
	}
	if err := wb.throttle.Do(); err != nil {
		return err
	}
	wb.txn.CommitWith(wb.callback)
	wb.txn = wb.db.newTransaction(true, true)
	wb.txn.readTs = 0 // We're not reading anything.
	wb.txn.commitTs = wb.commitTs
	return wb.err
}

// Flush must be called at the end to ensure that any pending writes get committed to Badger. Flush
// returns any error stored by WriteBatch.
func (wb *WriteBatch) Flush() error {
	wb.Lock()
	_ = wb.commit()
	wb.txn.Discard()
	wb.Unlock()

	if err := wb.throttle.Finish(); err != nil {
		return err
	}

	return wb.err
}

// Error returns any errors encountered so far. No commits would be run once an error is detected.
func (wb *WriteBatch) Error() error {
	wb.Lock()
	defer wb.Unlock()
	return wb.err
}
<commit_msg>Add Write method to batch write (#1321)<commit_after>/*
 * Copyright 2018 Dgraph Labs, Inc. and Contributors
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package badger

import (
	"sync"

	"github.com/dgraph-io/badger/v2/pb"
	"github.com/dgraph-io/badger/v2/y"
	"github.com/pkg/errors"
)

// WriteBatch holds the necessary info to perform batched writes.
type WriteBatch struct {
	sync.Mutex
	txn      *Txn
	db       *DB
	throttle *y.Throttle
	err      error
	commitTs uint64
}

// NewWriteBatch creates a new WriteBatch. This provides a way to conveniently do a lot of writes,
// batching them up as tightly as possible in a single transaction and using callbacks to avoid
// waiting for them to commit, thus achieving good performance. This API hides away the logic of
// creating and committing transactions. Due to the nature of SSI guaratees provided by Badger,
// blind writes can never encounter transaction conflicts (ErrConflict).
func (db *DB) NewWriteBatch() *WriteBatch {
	if db.opt.managedTxns {
		panic("cannot use NewWriteBatch in managed mode. Use NewWriteBatchAt instead")
	}
	return db.newWriteBatch()
}

func (db *DB) newWriteBatch() *WriteBatch {
	return &WriteBatch{
		db:       db,
		txn:      db.newTransaction(true, true),
		throttle: y.NewThrottle(16),
	}
}

// SetMaxPendingTxns sets a limit on maximum number of pending transactions while writing batches.
// This function should be called before using WriteBatch. Default value of MaxPendingTxns is
// 16 to minimise memory usage.
func (wb *WriteBatch) SetMaxPendingTxns(max int) {
	wb.throttle = y.NewThrottle(max)
}

// Cancel function must be called if there's a chance that Flush might not get
// called. If neither Flush or Cancel is called, the transaction oracle would
// never get a chance to clear out the row commit timestamp map, thus causing an
// unbounded memory consumption. Typically, you can call Cancel as a defer
// statement right after NewWriteBatch is called.
//
// Note that any committed writes would still go through despite calling Cancel.
func (wb *WriteBatch) Cancel() {
	if err := wb.throttle.Finish(); err != nil {
		wb.db.opt.Errorf("WatchBatch.Cancel error while finishing: %v", err)
	}
	wb.txn.Discard()
}

func (wb *WriteBatch) callback(err error) {
	// sync.WaitGroup is thread-safe, so it doesn't need to be run inside wb.Lock.
	defer wb.throttle.Done(err)
	if err == nil {
		return
	}

	wb.Lock()
	defer wb.Unlock()
	if wb.err != nil {
		return
	}
	wb.err = err
}

func (wb *WriteBatch) Write(kvList *pb.KVList) error {
	wb.Lock()
	defer wb.Unlock()
	for _, kv := range kvList.Kv {
		e := Entry{Key: kv.Key, Value: kv.Value}
		if len(kv.UserMeta) > 0 {
			e.UserMeta = kv.UserMeta[0]
		}
		y.AssertTrue(kv.Version != 0)
		e.version = kv.Version
		if err := wb.handleEntry(&e); err != nil {
			return err
		}
	}
	return nil
}

// SetEntryAt is the equivalent of Txn.SetEntry but it also allows setting version for the entry.
// SetEntryAt can be used only in managed mode.
func (wb *WriteBatch) SetEntryAt(e *Entry, ts uint64) error {
	if !wb.db.opt.managedTxns {
		return errors.New("SetEntryAt can only be used in managed mode. Use SetEntry instead")
	}
	e.version = ts
	return wb.SetEntry(e)
}

// Should be called with lock acquired.
func (wb *WriteBatch) handleEntry(e *Entry) error {
	if err := wb.txn.SetEntry(e); err != ErrTxnTooBig {
		return err
	}
	// Txn has reached it's zenith. Commit now.
	if cerr := wb.commit(); cerr != nil {
		return cerr
	}
	// This time the error must not be ErrTxnTooBig, otherwise, we make the
	// error permanent.
	if err := wb.txn.SetEntry(e); err != nil {
		wb.err = err
		return err
	}
	return nil
}

// SetEntry is the equivalent of Txn.SetEntry.
func (wb *WriteBatch) SetEntry(e *Entry) error {
	wb.Lock()
	defer wb.Unlock()
	return wb.handleEntry(e)
}

// Set is equivalent of Txn.Set().
func (wb *WriteBatch) Set(k, v []byte) error {
	e := &Entry{Key: k, Value: v}
	return wb.SetEntry(e)
}

// DeleteAt is equivalent of Txn.Delete but accepts a delete timestamp.
func (wb *WriteBatch) DeleteAt(k []byte, ts uint64) error {
	e := Entry{Key: k, meta: bitDelete, version: ts}
	return wb.SetEntry(&e)
}

// Delete is equivalent of Txn.Delete.
func (wb *WriteBatch) Delete(k []byte) error {
	wb.Lock()
	defer wb.Unlock()

	if err := wb.txn.Delete(k); err != ErrTxnTooBig {
		return err
	}
	if err := wb.commit(); err != nil {
		return err
	}
	if err := wb.txn.Delete(k); err != nil {
		wb.err = err
		return err
	}
	return nil
}

// Caller to commit must hold a write lock.
func (wb *WriteBatch) commit() error {
	if wb.err != nil {
		return wb.err
	}
	if err := wb.throttle.Do(); err != nil {
		return err
	}
	wb.txn.CommitWith(wb.callback)
	wb.txn = wb.db.newTransaction(true, true)
	wb.txn.readTs = 0 // We're not reading anything.
	wb.txn.commitTs = wb.commitTs
	return wb.err
}

// Flush must be called at the end to ensure that any pending writes get committed to Badger. Flush
// returns any error stored by WriteBatch.
func (wb *WriteBatch) Flush() error {
	wb.Lock()
	_ = wb.commit()
	wb.txn.Discard()
	wb.Unlock()

	if err := wb.throttle.Finish(); err != nil {
		return err
	}

	return wb.err
}

// Error returns any errors encountered so far. No commits would be run once an error is detected.
func (wb *WriteBatch) Error() error {
	wb.Lock()
	defer wb.Unlock()
	return wb.err
}
<|endoftext|>


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

package csdb

import (
	"github.com/corestoreio/csfw/storage/dbr"
	"github.com/juju/errgo"
)

// Table represents a table from the database
type Table struct {
	// Name is the table name
	Name string
	// Columns all table columns
	Columns Columns
	// CountPK number of primary keys
	CountPK int
	// CountUnique number of unique keys
	CountUnique int

	// internal caches
	fieldsPK  []string // all PK column field
	fieldsUNI []string // all unique key column field
	fields    []string // all other non-pk column field
}

// NewTable initializes a new table structure
func NewTable(n string, cs ...Column) *Table {
	ts := &Table{
		Name:    n,
		Columns: Columns(cs),
	}
	return ts.update()
}

// update recalculates the internal cached fields
func (ts *Table) update() *Table {
	ts.fieldsPK = ts.Columns.PrimaryKeys().FieldNames()
	ts.fieldsUNI = ts.Columns.UniqueKeys().FieldNames()
	ts.fields = ts.Columns.ColumnsNoPK().FieldNames()
	ts.CountPK = ts.Columns.PrimaryKeys().Len()
	ts.CountUnique = ts.Columns.UniqueKeys().Len()
	return ts
}

// Load reads the column information from the DB. @todo
func (ts *Table) LoadColumns(dbrSess dbr.SessionRunner) (err error) {
	ts.Columns, err = GetColumns(dbrSess, ts.Name)
	ts.update()
	return errgo.Mask(err)
}

// remove this once the ALIAS via []string is implemented in DBR
func (ts *Table) TableAliasQuote(alias string) string {
	return "`" + ts.Name + "` AS `" + alias + "`"
}

// ColumnAliasQuote prefixes non-id columns with an alias and puts quotes around them. Returns a copy.
func (ts *Table) ColumnAliasQuote(alias string) []string {
	return dbr.TableColumnQuote(alias, append([]string(nil), ts.fields...)...)
}

// AllColumnAliasQuote prefixes all columns with an alias and puts quotes around them. Returns a copy.
func (ts *Table) AllColumnAliasQuote(alias string) []string {
	c := append([]string(nil), ts.fieldsPK...)
	return dbr.TableColumnQuote(alias, append(c, ts.fields...)...)
}

// In checks if column name n is a column of this table
func (ts *Table) In(n string) bool {
	for _, c := range ts.fieldsPK {
		if c == n {
			return true
		}
	}
	for _, c := range ts.fields {
		if c == n {
			return true
		}
	}
	return false
}

// Select generates a SELECT * FROM tableName statement
func (ts *Table) Select(dbrSess dbr.SessionRunner) (*dbr.SelectBuilder, error) {
	if ts == nil {
		return nil, ErrTableNotFound
	}
	return dbrSess.
		Select(ts.AllColumnAliasQuote("main_table")...).
		From(ts.Name, "main_table"), nil
}
<commit_msg>storage/csdb: Initial idea for tiny SQL abstraction<commit_after>// Copyright 2015, Cyrill @ Schumacher.fm and the CoreStore contributors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package csdb

import (
	"github.com/corestoreio/csfw/storage/dbr"
	"github.com/juju/errgo"
)

// Table represents a table from the database
type Table struct {
	// Name is the table name
	Name string
	// Columns all table columns
	Columns Columns
	// CountPK number of primary keys
	CountPK int
	// CountUnique number of unique keys
	CountUnique int

	// internal caches
	fieldsPK  []string // all PK column field
	fieldsUNI []string // all unique key column field
	fields    []string // all other non-pk column field
}

// NewTable initializes a new table structure
func NewTable(n string, cs ...Column) *Table {
	ts := &Table{
		Name:    n,
		Columns: Columns(cs),
	}
	return ts.update()
}

// update recalculates the internal cached fields
func (ts *Table) update() *Table {
	ts.fieldsPK = ts.Columns.PrimaryKeys().FieldNames()
	ts.fieldsUNI = ts.Columns.UniqueKeys().FieldNames()
	ts.fields = ts.Columns.ColumnsNoPK().FieldNames()
	ts.CountPK = ts.Columns.PrimaryKeys().Len()
	ts.CountUnique = ts.Columns.UniqueKeys().Len()
	return ts
}

// Load reads the column information from the DB. @todo
func (ts *Table) LoadColumns(dbrSess dbr.SessionRunner) (err error) {
	ts.Columns, err = GetColumns(dbrSess, ts.Name)
	ts.update()
	return errgo.Mask(err)
}

// remove this once the ALIAS via []string is implemented in DBR
func (ts *Table) TableAliasQuote(alias string) string {
	return "`" + ts.Name + "` AS `" + alias + "`"
}

// ColumnAliasQuote prefixes non-id columns with an alias and puts quotes around them. Returns a copy.
func (ts *Table) ColumnAliasQuote(alias string) []string {
	return dbr.TableColumnQuote(alias, append([]string(nil), ts.fields...)...)
}

// AllColumnAliasQuote prefixes all columns with an alias and puts quotes around them. Returns a copy.
func (ts *Table) AllColumnAliasQuote(alias string) []string {
	c := append([]string(nil), ts.fieldsPK...)
	return dbr.TableColumnQuote(alias, append(c, ts.fields...)...)
}

// In checks if column name n is a column of this table
func (ts *Table) In(n string) bool {
	for _, c := range ts.fieldsPK {
		if c == n {
			return true
		}
	}
	for _, c := range ts.fields {
		if c == n {
			return true
		}
	}
	return false
}

// Select generates a SELECT * FROM tableName statement
func (ts *Table) Select(dbrSess dbr.SessionRunner) (*dbr.SelectBuilder, error) {
	if ts == nil {
		return nil, ErrTableNotFound
	}
	return dbrSess.
		Select(ts.AllColumnAliasQuote("main_table")...).
		From(ts.Name, "main_table"), nil
}

func (ts *Table) Update() {}
func (ts *Table) Delete() {}
func (ts *Table) Insert() {}
func (ts *Table) Alter()  {}
func (ts *Table) Drop()   {}
func (ts *Table) Create() {}
<|endoftext|>