<commit_before>//  Copyright (c) 2015 Couchbase, Inc.
//  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 firestorm

import (
	"encoding/json"
	"fmt"
	"sync"
	"sync/atomic"
	"time"

	"github.com/blevesearch/bleve/document"
	"github.com/blevesearch/bleve/index"
	"github.com/blevesearch/bleve/index/store"
	"github.com/blevesearch/bleve/registry"
)

const Name = "firestorm"

var UnsafeBatchUseDetected = fmt.Errorf("bleve.Batch is NOT thread-safe, modification after execution detected")

type Firestorm struct {
	storeName        string
	storeConfig      map[string]interface{}
	store            store.KVStore
	compensator      *Compensator
	analysisQueue    *index.AnalysisQueue
	fieldCache       *index.FieldCache
	highDocNumber    uint64
	docCount         *uint64
	garbageCollector *GarbageCollector
	lookuper         *Lookuper
	dictUpdater      *DictUpdater
	stats            *indexStat
}

func NewFirestorm(storeName string, storeConfig map[string]interface{}, analysisQueue *index.AnalysisQueue) (index.Index, error) {
	initialCount := uint64(0)
	rv := Firestorm{
		storeName:     storeName,
		storeConfig:   storeConfig,
		compensator:   NewCompensator(),
		analysisQueue: analysisQueue,
		fieldCache:    index.NewFieldCache(),
		docCount:      &initialCount,
		highDocNumber: 0,
		stats:         &indexStat{},
	}
	rv.stats.f = &rv
	rv.garbageCollector = NewGarbageCollector(&rv)
	rv.lookuper = NewLookuper(&rv)
	rv.dictUpdater = NewDictUpdater(&rv)
	return &rv, nil
}

func (f *Firestorm) Open() (err error) {

	// open the kv store
	storeConstructor := registry.KVStoreConstructorByName(f.storeName)
	if storeConstructor == nil {
		err = index.ErrorUnknownStorageType
		return
	}

	// now open the store
	f.store, err = storeConstructor(&mergeOperator, f.storeConfig)
	if err != nil {
		return
	}

	// start a reader
	var kvreader store.KVReader
	kvreader, err = f.store.Reader()
	if err != nil {
		return
	}

	// assert correct version, and find out if this is new index
	var newIndex bool
	newIndex, err = f.checkVersion(kvreader)
	if err != nil {
		return
	}

	if !newIndex {
		// process existing index before opening
		err = f.warmup(kvreader)
		if err != nil {
			return
		}
	}

	err = kvreader.Close()
	if err != nil {
		return
	}

	if newIndex {
		// prepare a new index
		err = f.bootstrap()
		if err != nil {
			return
		}
	}

	// start the garbage collector
	f.garbageCollector.Start()

	// start the lookuper
	f.lookuper.Start()

	// start the dict updater
	f.dictUpdater.Start()

	return
}

func (f *Firestorm) Close() error {
	f.garbageCollector.Stop()
	f.lookuper.Stop()
	f.dictUpdater.Stop()
	return f.store.Close()
}

func (f *Firestorm) DocCount() (uint64, error) {
	count := atomic.LoadUint64(f.docCount)
	return count, nil

}

func (f *Firestorm) Update(doc *document.Document) (err error) {

	// assign this document a number
	doc.Number = atomic.AddUint64(&f.highDocNumber, 1)

	// do analysis before acquiring write lock
	analysisStart := time.Now()
	resultChan := make(chan *index.AnalysisResult)
	aw := index.NewAnalysisWork(f, doc, resultChan)

	// put the work on the queue
	go f.analysisQueue.Queue(aw)

	// wait for the result
	result := <-resultChan
	close(resultChan)
	atomic.AddUint64(&f.stats.analysisTime, uint64(time.Since(analysisStart)))

	// start a writer for this update
	indexStart := time.Now()
	var kvwriter store.KVWriter
	kvwriter, err = f.store.Writer()
	if err != nil {
		return
	}
	defer func() {
		if cerr := kvwriter.Close(); err == nil && cerr != nil {
			err = cerr
		}
	}()

	var dictionaryDeltas map[string]int64
	dictionaryDeltas, err = f.batchRows(kvwriter, result.Rows, nil)
	if err != nil {
		_ = kvwriter.Close()
		atomic.AddUint64(&f.stats.errors, 1)
		return
	}

	f.compensator.Mutate([]byte(doc.ID), doc.Number)
	f.lookuper.Notify(doc.Number, []byte(doc.ID))
	f.dictUpdater.NotifyBatch(dictionaryDeltas)

	atomic.AddUint64(&f.stats.indexTime, uint64(time.Since(indexStart)))
	return
}

func (f *Firestorm) Delete(id string) error {
	indexStart := time.Now()
	f.compensator.Mutate([]byte(id), 0)
	f.lookuper.Notify(0, []byte(id))
	atomic.AddUint64(&f.stats.indexTime, uint64(time.Since(indexStart)))
	return nil
}

func (f *Firestorm) batchRows(writer store.KVWriter, rows []index.IndexRow, deleteKeys [][]byte) (map[string]int64, error) {

	// prepare batch
	wb := writer.NewBatch()

	dictionaryDeltas := make(map[string]int64)
	for _, row := range rows {
		tfr, ok := row.(*TermFreqRow)
		if ok {
			if tfr.Field() != 0 {
				drk := tfr.DictionaryRowKey()
				dictionaryDeltas[string(drk)] += 1
			}
		}
		wb.Set(row.Key(), row.Value())
	}

	for _, dk := range deleteKeys {
		wb.Delete(dk)
	}

	// write out the batch
	err := writer.ExecuteBatch(wb)
	if err != nil {
		return nil, err
	}
	return dictionaryDeltas, nil
}

func (f *Firestorm) Batch(batch *index.Batch) (err error) {

	// acquire enough doc numbers for all updates in the batch
	// FIXME we actually waste doc numbers because deletes are in the
	// same map and we don't need numbers for them
	lastDocNumber := atomic.AddUint64(&f.highDocNumber, uint64(len(batch.IndexOps)))
	firstDocNumber := lastDocNumber - uint64(len(batch.IndexOps)) + 1

	analysisStart := time.Now()
	resultChan := make(chan *index.AnalysisResult)

	var numUpdates uint64
	var docsDeleted uint64
	for _, doc := range batch.IndexOps {
		if doc != nil {
			doc.Number = firstDocNumber // actually assign doc numbers here
			firstDocNumber++
			numUpdates++
		} else {
			docsDeleted++
		}
	}

	var detectedUnsafeMutex sync.RWMutex
	detectedUnsafe := false

	go func() {
		sofar := uint64(0)
		for _, doc := range batch.IndexOps {
			if doc != nil {
				sofar++
				if sofar > numUpdates {
					detectedUnsafeMutex.Lock()
					detectedUnsafe = true
					detectedUnsafeMutex.Unlock()
					return
				}
				aw := index.NewAnalysisWork(f, doc, resultChan)
				// put the work on the queue
				f.analysisQueue.Queue(aw)
			}
		}
	}()

	allRows := make([]index.IndexRow, 0, 1000)
	// wait for the result
	var itemsDeQueued uint64
	for itemsDeQueued < numUpdates {
		result := <-resultChan
		allRows = append(allRows, result.Rows...)
		itemsDeQueued++
	}
	close(resultChan)

	detectedUnsafeMutex.RLock()
	defer detectedUnsafeMutex.RUnlock()
	if detectedUnsafe {
		return UnsafeBatchUseDetected
	}

	atomic.AddUint64(&f.stats.analysisTime, uint64(time.Since(analysisStart)))

	deleteKeys := make([][]byte, 0)
	// add the internal ops
	for internalKey, internalValue := range batch.InternalOps {
		if internalValue == nil {
			// delete
			deleteInternalRow := NewInternalRow([]byte(internalKey), nil)
			deleteKeys = append(deleteKeys, deleteInternalRow.Key())
		} else {
			updateInternalRow := NewInternalRow([]byte(internalKey), internalValue)
			allRows = append(allRows, updateInternalRow)
		}
	}

	indexStart := time.Now()
	// start a writer for this batch
	var kvwriter store.KVWriter
	kvwriter, err = f.store.Writer()
	if err != nil {
		return
	}

	var dictionaryDeltas map[string]int64
	dictionaryDeltas, err = f.batchRows(kvwriter, allRows, deleteKeys)
	if err != nil {
		_ = kvwriter.Close()
		atomic.AddUint64(&f.stats.errors, 1)
		return
	}

	f.compensator.MutateBatch(batch.IndexOps, lastDocNumber)
	for docID, doc := range batch.IndexOps {
		if doc != nil {
			f.lookuper.Notify(doc.Number, []byte(doc.ID))
		} else {
			f.lookuper.Notify(0, []byte(docID))
		}
	}
	f.dictUpdater.NotifyBatch(dictionaryDeltas)

	err = kvwriter.Close()
	atomic.AddUint64(&f.stats.indexTime, uint64(time.Since(indexStart)))

	if err == nil {
		atomic.AddUint64(&f.stats.updates, numUpdates)
		atomic.AddUint64(&f.stats.deletes, docsDeleted)
		atomic.AddUint64(&f.stats.batches, 1)
	} else {
		atomic.AddUint64(&f.stats.errors, 1)
	}

	return
}

func (f *Firestorm) SetInternal(key, val []byte) (err error) {
	internalRow := NewInternalRow(key, val)
	var writer store.KVWriter
	writer, err = f.store.Writer()
	if err != nil {
		return
	}
	defer func() {
		if cerr := writer.Close(); err == nil && cerr != nil {
			err = cerr
		}
	}()

	wb := writer.NewBatch()
	wb.Set(internalRow.Key(), internalRow.Value())

	return writer.ExecuteBatch(wb)
}

func (f *Firestorm) DeleteInternal(key []byte) (err error) {
	internalRow := NewInternalRow(key, nil)
	var writer store.KVWriter
	writer, err = f.store.Writer()
	if err != nil {
		return
	}
	defer func() {
		if cerr := writer.Close(); err == nil && cerr != nil {
			err = cerr
		}
	}()

	wb := writer.NewBatch()
	wb.Delete(internalRow.Key())

	return writer.ExecuteBatch(wb)
}

func (f *Firestorm) DumpAll() chan interface{} {
	rv := make(chan interface{})
	go func() {
		defer close(rv)

		// start an isolated reader for use during the dump
		kvreader, err := f.store.Reader()
		if err != nil {
			rv <- err
			return
		}
		defer func() {
			cerr := kvreader.Close()
			if cerr != nil {
				rv <- cerr
			}
		}()

		err = f.dumpPrefix(kvreader, rv, nil)
		if err != nil {
			rv <- err
			return
		}
	}()
	return rv
}

func (f *Firestorm) DumpDoc(docID string) chan interface{} {
	rv := make(chan interface{})
	go func() {
		defer close(rv)

		// start an isolated reader for use during the dump
		kvreader, err := f.store.Reader()
		if err != nil {
			rv <- err
			return
		}
		defer func() {
			cerr := kvreader.Close()
			if cerr != nil {
				rv <- cerr
			}
		}()

		err = f.dumpDoc(kvreader, rv, []byte(docID))
		if err != nil {
			rv <- err
			return
		}
	}()
	return rv
}

func (f *Firestorm) DumpFields() chan interface{} {
	rv := make(chan interface{})
	go func() {
		defer close(rv)

		// start an isolated reader for use during the dump
		kvreader, err := f.store.Reader()
		if err != nil {
			rv <- err
			return
		}
		defer func() {
			cerr := kvreader.Close()
			if cerr != nil {
				rv <- cerr
			}
		}()

		err = f.dumpPrefix(kvreader, rv, FieldKeyPrefix)
		if err != nil {
			rv <- err
			return
		}
	}()
	return rv
}

func (f *Firestorm) Reader() (index.IndexReader, error) {
	return newFirestormReader(f)
}

func (f *Firestorm) Stats() json.Marshaler {
	return f.stats

}

func (f *Firestorm) Wait(timeout time.Duration) error {
	return f.dictUpdater.waitTasksDone(timeout)
}

func init() {
	registry.RegisterIndexType(Name, NewFirestorm)
}
<commit_msg>use KeyTo/ValueTo in firestorm.batchRows<commit_after>//  Copyright (c) 2015 Couchbase, Inc.
//  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 firestorm

import (
	"encoding/json"
	"fmt"
	"sync"
	"sync/atomic"
	"time"

	"github.com/blevesearch/bleve/document"
	"github.com/blevesearch/bleve/index"
	"github.com/blevesearch/bleve/index/store"
	"github.com/blevesearch/bleve/registry"
)

const Name = "firestorm"

var UnsafeBatchUseDetected = fmt.Errorf("bleve.Batch is NOT thread-safe, modification after execution detected")

type Firestorm struct {
	storeName        string
	storeConfig      map[string]interface{}
	store            store.KVStore
	compensator      *Compensator
	analysisQueue    *index.AnalysisQueue
	fieldCache       *index.FieldCache
	highDocNumber    uint64
	docCount         *uint64
	garbageCollector *GarbageCollector
	lookuper         *Lookuper
	dictUpdater      *DictUpdater
	stats            *indexStat
}

func NewFirestorm(storeName string, storeConfig map[string]interface{}, analysisQueue *index.AnalysisQueue) (index.Index, error) {
	initialCount := uint64(0)
	rv := Firestorm{
		storeName:     storeName,
		storeConfig:   storeConfig,
		compensator:   NewCompensator(),
		analysisQueue: analysisQueue,
		fieldCache:    index.NewFieldCache(),
		docCount:      &initialCount,
		highDocNumber: 0,
		stats:         &indexStat{},
	}
	rv.stats.f = &rv
	rv.garbageCollector = NewGarbageCollector(&rv)
	rv.lookuper = NewLookuper(&rv)
	rv.dictUpdater = NewDictUpdater(&rv)
	return &rv, nil
}

func (f *Firestorm) Open() (err error) {

	// open the kv store
	storeConstructor := registry.KVStoreConstructorByName(f.storeName)
	if storeConstructor == nil {
		err = index.ErrorUnknownStorageType
		return
	}

	// now open the store
	f.store, err = storeConstructor(&mergeOperator, f.storeConfig)
	if err != nil {
		return
	}

	// start a reader
	var kvreader store.KVReader
	kvreader, err = f.store.Reader()
	if err != nil {
		return
	}

	// assert correct version, and find out if this is new index
	var newIndex bool
	newIndex, err = f.checkVersion(kvreader)
	if err != nil {
		return
	}

	if !newIndex {
		// process existing index before opening
		err = f.warmup(kvreader)
		if err != nil {
			return
		}
	}

	err = kvreader.Close()
	if err != nil {
		return
	}

	if newIndex {
		// prepare a new index
		err = f.bootstrap()
		if err != nil {
			return
		}
	}

	// start the garbage collector
	f.garbageCollector.Start()

	// start the lookuper
	f.lookuper.Start()

	// start the dict updater
	f.dictUpdater.Start()

	return
}

func (f *Firestorm) Close() error {
	f.garbageCollector.Stop()
	f.lookuper.Stop()
	f.dictUpdater.Stop()
	return f.store.Close()
}

func (f *Firestorm) DocCount() (uint64, error) {
	count := atomic.LoadUint64(f.docCount)
	return count, nil

}

func (f *Firestorm) Update(doc *document.Document) (err error) {

	// assign this document a number
	doc.Number = atomic.AddUint64(&f.highDocNumber, 1)

	// do analysis before acquiring write lock
	analysisStart := time.Now()
	resultChan := make(chan *index.AnalysisResult)
	aw := index.NewAnalysisWork(f, doc, resultChan)

	// put the work on the queue
	go f.analysisQueue.Queue(aw)

	// wait for the result
	result := <-resultChan
	close(resultChan)
	atomic.AddUint64(&f.stats.analysisTime, uint64(time.Since(analysisStart)))

	// start a writer for this update
	indexStart := time.Now()
	var kvwriter store.KVWriter
	kvwriter, err = f.store.Writer()
	if err != nil {
		return
	}
	defer func() {
		if cerr := kvwriter.Close(); err == nil && cerr != nil {
			err = cerr
		}
	}()

	var dictionaryDeltas map[string]int64
	dictionaryDeltas, err = f.batchRows(kvwriter, result.Rows, nil)
	if err != nil {
		_ = kvwriter.Close()
		atomic.AddUint64(&f.stats.errors, 1)
		return
	}

	f.compensator.Mutate([]byte(doc.ID), doc.Number)
	f.lookuper.Notify(doc.Number, []byte(doc.ID))
	f.dictUpdater.NotifyBatch(dictionaryDeltas)

	atomic.AddUint64(&f.stats.indexTime, uint64(time.Since(indexStart)))
	return
}

func (f *Firestorm) Delete(id string) error {
	indexStart := time.Now()
	f.compensator.Mutate([]byte(id), 0)
	f.lookuper.Notify(0, []byte(id))
	atomic.AddUint64(&f.stats.indexTime, uint64(time.Since(indexStart)))
	return nil
}

func (f *Firestorm) batchRows(writer store.KVWriter, rows []index.IndexRow, deleteKeys [][]byte) (map[string]int64, error) {

	// prepare batch
	wb := writer.NewBatch()

	var kbuf []byte
	var vbuf []byte

	prepareBuf := func(buf []byte, sizeNeeded int) []byte {
		if cap(buf) < sizeNeeded {
			return make([]byte, sizeNeeded, sizeNeeded+128)
		}
		return buf[0:sizeNeeded]
	}

	dictionaryDeltas := make(map[string]int64)
	for _, row := range rows {
		tfr, ok := row.(*TermFreqRow)
		if ok {
			if tfr.Field() != 0 {
				kbuf = prepareBuf(kbuf, tfr.DictionaryRowKeySize())
				klen, err := tfr.DictionaryRowKeyTo(kbuf)
				if err != nil {
					return nil, err
				}

				dictionaryDeltas[string(kbuf[0:klen])] += 1
			}

			kbuf = prepareBuf(kbuf, tfr.KeySize())
			klen, err := tfr.KeyTo(kbuf)
			if err != nil {
				return nil, err
			}

			vbuf = prepareBuf(vbuf, tfr.ValueSize())
			vlen, err := tfr.ValueTo(vbuf)
			if err != nil {
				return nil, err
			}

			wb.Set(kbuf[0:klen], vbuf[0:vlen])
		} else {
			wb.Set(row.Key(), row.Value())
		}
	}

	for _, dk := range deleteKeys {
		wb.Delete(dk)
	}

	// write out the batch
	err := writer.ExecuteBatch(wb)
	if err != nil {
		return nil, err
	}
	return dictionaryDeltas, nil
}

func (f *Firestorm) Batch(batch *index.Batch) (err error) {

	// acquire enough doc numbers for all updates in the batch
	// FIXME we actually waste doc numbers because deletes are in the
	// same map and we don't need numbers for them
	lastDocNumber := atomic.AddUint64(&f.highDocNumber, uint64(len(batch.IndexOps)))
	firstDocNumber := lastDocNumber - uint64(len(batch.IndexOps)) + 1

	analysisStart := time.Now()
	resultChan := make(chan *index.AnalysisResult)

	var numUpdates uint64
	var docsDeleted uint64
	for _, doc := range batch.IndexOps {
		if doc != nil {
			doc.Number = firstDocNumber // actually assign doc numbers here
			firstDocNumber++
			numUpdates++
		} else {
			docsDeleted++
		}
	}

	var detectedUnsafeMutex sync.RWMutex
	detectedUnsafe := false

	go func() {
		sofar := uint64(0)
		for _, doc := range batch.IndexOps {
			if doc != nil {
				sofar++
				if sofar > numUpdates {
					detectedUnsafeMutex.Lock()
					detectedUnsafe = true
					detectedUnsafeMutex.Unlock()
					return
				}
				aw := index.NewAnalysisWork(f, doc, resultChan)
				// put the work on the queue
				f.analysisQueue.Queue(aw)
			}
		}
	}()

	allRows := make([]index.IndexRow, 0, 1000)
	// wait for the result
	var itemsDeQueued uint64
	for itemsDeQueued < numUpdates {
		result := <-resultChan
		allRows = append(allRows, result.Rows...)
		itemsDeQueued++
	}
	close(resultChan)

	detectedUnsafeMutex.RLock()
	defer detectedUnsafeMutex.RUnlock()
	if detectedUnsafe {
		return UnsafeBatchUseDetected
	}

	atomic.AddUint64(&f.stats.analysisTime, uint64(time.Since(analysisStart)))

	deleteKeys := make([][]byte, 0)
	// add the internal ops
	for internalKey, internalValue := range batch.InternalOps {
		if internalValue == nil {
			// delete
			deleteInternalRow := NewInternalRow([]byte(internalKey), nil)
			deleteKeys = append(deleteKeys, deleteInternalRow.Key())
		} else {
			updateInternalRow := NewInternalRow([]byte(internalKey), internalValue)
			allRows = append(allRows, updateInternalRow)
		}
	}

	indexStart := time.Now()
	// start a writer for this batch
	var kvwriter store.KVWriter
	kvwriter, err = f.store.Writer()
	if err != nil {
		return
	}

	var dictionaryDeltas map[string]int64
	dictionaryDeltas, err = f.batchRows(kvwriter, allRows, deleteKeys)
	if err != nil {
		_ = kvwriter.Close()
		atomic.AddUint64(&f.stats.errors, 1)
		return
	}

	f.compensator.MutateBatch(batch.IndexOps, lastDocNumber)
	for docID, doc := range batch.IndexOps {
		if doc != nil {
			f.lookuper.Notify(doc.Number, []byte(doc.ID))
		} else {
			f.lookuper.Notify(0, []byte(docID))
		}
	}
	f.dictUpdater.NotifyBatch(dictionaryDeltas)

	err = kvwriter.Close()
	atomic.AddUint64(&f.stats.indexTime, uint64(time.Since(indexStart)))

	if err == nil {
		atomic.AddUint64(&f.stats.updates, numUpdates)
		atomic.AddUint64(&f.stats.deletes, docsDeleted)
		atomic.AddUint64(&f.stats.batches, 1)
	} else {
		atomic.AddUint64(&f.stats.errors, 1)
	}

	return
}

func (f *Firestorm) SetInternal(key, val []byte) (err error) {
	internalRow := NewInternalRow(key, val)
	var writer store.KVWriter
	writer, err = f.store.Writer()
	if err != nil {
		return
	}
	defer func() {
		if cerr := writer.Close(); err == nil && cerr != nil {
			err = cerr
		}
	}()

	wb := writer.NewBatch()
	wb.Set(internalRow.Key(), internalRow.Value())

	return writer.ExecuteBatch(wb)
}

func (f *Firestorm) DeleteInternal(key []byte) (err error) {
	internalRow := NewInternalRow(key, nil)
	var writer store.KVWriter
	writer, err = f.store.Writer()
	if err != nil {
		return
	}
	defer func() {
		if cerr := writer.Close(); err == nil && cerr != nil {
			err = cerr
		}
	}()

	wb := writer.NewBatch()
	wb.Delete(internalRow.Key())

	return writer.ExecuteBatch(wb)
}

func (f *Firestorm) DumpAll() chan interface{} {
	rv := make(chan interface{})
	go func() {
		defer close(rv)

		// start an isolated reader for use during the dump
		kvreader, err := f.store.Reader()
		if err != nil {
			rv <- err
			return
		}
		defer func() {
			cerr := kvreader.Close()
			if cerr != nil {
				rv <- cerr
			}
		}()

		err = f.dumpPrefix(kvreader, rv, nil)
		if err != nil {
			rv <- err
			return
		}
	}()
	return rv
}

func (f *Firestorm) DumpDoc(docID string) chan interface{} {
	rv := make(chan interface{})
	go func() {
		defer close(rv)

		// start an isolated reader for use during the dump
		kvreader, err := f.store.Reader()
		if err != nil {
			rv <- err
			return
		}
		defer func() {
			cerr := kvreader.Close()
			if cerr != nil {
				rv <- cerr
			}
		}()

		err = f.dumpDoc(kvreader, rv, []byte(docID))
		if err != nil {
			rv <- err
			return
		}
	}()
	return rv
}

func (f *Firestorm) DumpFields() chan interface{} {
	rv := make(chan interface{})
	go func() {
		defer close(rv)

		// start an isolated reader for use during the dump
		kvreader, err := f.store.Reader()
		if err != nil {
			rv <- err
			return
		}
		defer func() {
			cerr := kvreader.Close()
			if cerr != nil {
				rv <- cerr
			}
		}()

		err = f.dumpPrefix(kvreader, rv, FieldKeyPrefix)
		if err != nil {
			rv <- err
			return
		}
	}()
	return rv
}

func (f *Firestorm) Reader() (index.IndexReader, error) {
	return newFirestormReader(f)
}

func (f *Firestorm) Stats() json.Marshaler {
	return f.stats

}

func (f *Firestorm) Wait(timeout time.Duration) error {
	return f.dictUpdater.waitTasksDone(timeout)
}

func init() {
	registry.RegisterIndexType(Name, NewFirestorm)
}
<|endoftext|>


<commit_before>package nsq

import (
	"encoding/json"
	"net/http"
	"strings"
	"time"

	"github.com/pkg/errors"
)

type ProducerInfo struct {
	BroadcastAddress string `json:"broadcast_address,omitempty"`
	RemoteAddress    string `json:"remote_address,omitempty"`
	Hostname         string `json:"hostname,omitempty"`
	Version          string `json:"version,omitempty"`
	TcpPort          int    `json:"tcp_port"`
	HttpPort         int    `json:"http_port"`
}

type LookupResult struct {
	Channels  []string       `json:"channels"`
	Producers []ProducerInfo `json:"producers"`
}

type lookupResult struct {
	LookupResult
	error
}

func Lookup(topic string, addrs ...string) (resp LookupResult, err error) {
	n := len(addrs)

	if n == 0 {
		err = errors.New("missing addresses when looking up topic " + topic)
		return
	}

	respList := make([]LookupResult, 0, n)
	respChan := make(chan lookupResult, n)
	deadline := time.NewTimer(DefaultLookupTimeout)
	defer deadline.Stop()

	for _, addr := range addrs {
		go lookupAsync(topic, addr, respChan)
	}

respLoop:
	for i := 0; i != n; i++ {
		select {
		case r := <-respChan:
			if r.error != nil {
				err = r.error
			} else {
				respList = append(respList, r.LookupResult)
			}

		case <-deadline.C:
			if i == 0 {
				err = errors.New("no response came back when looking up topic " + topic)
			}
			break respLoop
		}
	}

	resp = mergeLookupResults(respList)
	return
}

func lookupAsync(topic string, addr string, respChan chan<- lookupResult) {
	resp, err := lookup(topic, addr)
	respChan <- lookupResult{resp, err}
}

func lookup(topic string, addr string) (resp LookupResult, err error) {
	var res *http.Response
	var dec *json.Decoder

	if !strings.HasPrefix(addr, "http://") {
		addr = "http://" + addr
	}

	if res, err = http.Get(addr + "/lookup?topic=" + topic); err != nil {
		err = errors.Wrap(err, "looking up topic "+topic+" on http://"+addr)
		return
	}

	defer res.Body.Close()

	v := struct {
		StatusCode int          `json:"status_code"`
		StatusTxt  string       `json:"status_txt"`
		Data       LookupResult `json:"data"`
	}{}

	dec = json.NewDecoder(res.Body)
	err = dec.Decode(&v)

	if v.StatusCode != 200 {
		err = errors.Errorf("looking for topic %s returned %d %s", topic, v.StatusCode, v.StatusTxt)
		return
	}

	resp = v.Data
	return
}

func mergeLookupResults(respList []LookupResult) (resp LookupResult) {
	channels := make(map[string]bool)
	producers := make(map[ProducerInfo]bool)

	for _, r := range respList {
		for _, c := range r.Channels {
			channels[c] = true
		}

		for _, p := range r.Producers {
			producers[p] = true
		}
	}

	resp = LookupResult{
		Channels:  make([]string, 0, len(channels)),
		Producers: make([]ProducerInfo, 0, len(producers)),
	}

	for c := range channels {
		resp.Channels = append(resp.Channels, c)
	}

	for p := range producers {
		resp.Producers = append(resp.Producers, p)
	}

	return
}
<commit_msg>cleanup<commit_after>package nsq

import (
	"encoding/json"
	"net/http"
	"strings"
	"time"

	"github.com/pkg/errors"
)

type ProducerInfo struct {
	BroadcastAddress string `json:"broadcast_address,omitempty"`
	RemoteAddress    string `json:"remote_address,omitempty"`
	Hostname         string `json:"hostname,omitempty"`
	Version          string `json:"version,omitempty"`
	TcpPort          int    `json:"tcp_port"`
	HttpPort         int    `json:"http_port"`
}

type LookupResult struct {
	Channels  []string       `json:"channels"`
	Producers []ProducerInfo `json:"producers"`
}

type lookupResult struct {
	LookupResult
	error
}

func Lookup(topic string, addrs ...string) (result LookupResult, err error) {
	n := len(addrs)

	if n == 0 {
		err = errors.New("missing addresses when looking up topic " + topic)
		return
	}

	resultList := make([]LookupResult, 0, n)
	resultChan := make(chan lookupResult, n)
	deadline := time.NewTimer(DefaultLookupTimeout)
	defer deadline.Stop()

	for _, addr := range addrs {
		go lookupAsync(topic, addr, resultChan)
	}

resultLoop:
	for i := 0; i != n; i++ {
		select {
		case r := <-resultChan:
			if r.error != nil {
				err = r.error
			} else {
				resultList = append(resultList, r.LookupResult)
			}

		case <-deadline.C:
			if i == 0 {
				err = errors.New("no response came back when looking up topic " + topic)
			}
			break resultLoop
		}
	}

	result = mergeLookupResults(resultList)
	return
}

func lookupAsync(topic string, addr string, resultChan chan<- lookupResult) {
	result, err := lookup(topic, addr)
	resultChan <- lookupResult{result, err}
}

func lookup(topic string, addr string) (result LookupResult, err error) {
	var res *http.Response
	var dec *json.Decoder

	if !strings.HasPrefix(addr, "http://") {
		addr = "http://" + addr
	}

	if res, err = http.Get(addr + "/lookup?topic=" + topic); err != nil {
		err = errors.Wrap(err, "looking up topic "+topic+" on http://"+addr)
		return
	}

	defer res.Body.Close()

	v := struct {
		StatusCode int          `json:"status_code"`
		StatusTxt  string       `json:"status_txt"`
		Data       LookupResult `json:"data"`
	}{}

	dec = json.NewDecoder(res.Body)
	err = dec.Decode(&v)

	if v.StatusCode != 200 {
		err = errors.Errorf("looking for topic %s returned %d %s", topic, v.StatusCode, v.StatusTxt)
		return
	}

	result = v.Data
	return
}

func mergeLookupResults(resultList []LookupResult) (result LookupResult) {
	channels := make(map[string]bool)
	producers := make(map[ProducerInfo]bool)

	for _, r := range resultList {
		for _, c := range r.Channels {
			channels[c] = true
		}

		for _, p := range r.Producers {
			producers[p] = true
		}
	}

	result = LookupResult{
		Channels:  make([]string, 0, len(channels)),
		Producers: make([]ProducerInfo, 0, len(producers)),
	}

	for c := range channels {
		result.Channels = append(result.Channels, c)
	}

	for p := range producers {
		result.Producers = append(result.Producers, p)
	}

	return
}
<|endoftext|>


<commit_before>// TODO:
//  - Add an optional prefix to log messages

package lumber

import (
	"fmt"
	"time"
)

const (
	TRACE = iota
	DEBUG
	INFO
	WARN
	ERROR
	FATAL

	TIMEFORMAT = "2006/01/02 15:04:05"
	MSGBUFSIZE = 50
)

var (
	stdLog     = NewConsoleLogger(INFO)
	levels     = [...]string{"TRACE", "DEBUG", "INFO ", "WARN ", "ERROR", "FATAL"}
	timeFormat = "2006/01/02 15:04:05"
	msgBufSize = 50
)

type Logger interface {
	Fatal(format string, v ...interface{})
	Error(format string, v ...interface{})
	Warn(format string, v ...interface{})
	Info(format string, v ...interface{})
	Debug(format string, v ...interface{})
	Trace(format string, v ...interface{})
	output(msg *Message)
	// output(level int, format string, v ...interface{})
	Close() error
}

type Message struct {
	level int
	m     string
	time  time.Time
}

func TimeFormat(f string) {
	stdLog.TimeFormat(f)
}

func LvlStr(l int) string {
	if l >= TRACE && l <= FATAL {
		return levels[l]
	}
	return ""
}

func Level(o int) {
	stdLog.Level(o)
}

func Fatal(format string, v ...interface{}) {
	stdLog.output(&Message{FATAL, fmt.Sprintf(format, v...), time.Now()})
}

func Error(format string, v ...interface{}) {
	stdLog.output(&Message{ERROR, fmt.Sprintf(format, v...), time.Now()})
}

func Warn(format string, v ...interface{}) {
	stdLog.output(&Message{WARN, fmt.Sprintf(format, v...), time.Now()})
}

func Info(format string, v ...interface{}) {
	stdLog.output(&Message{INFO, fmt.Sprintf(format, v...), time.Now()})
}

func Debug(format string, v ...interface{}) {
	stdLog.output(&Message{DEBUG, fmt.Sprintf(format, v...), time.Now()})
}

func Trace(format string, v ...interface{}) {
	stdLog.output(&Message{TRACE, fmt.Sprintf(format, v...), time.Now()})
}
<commit_msg>One more interface fix (added Level)<commit_after>// TODO:
//  - Add an optional prefix to log messages

package lumber

import (
	"fmt"
	"time"
)

const (
	TRACE = iota
	DEBUG
	INFO
	WARN
	ERROR
	FATAL

	TIMEFORMAT = "2006/01/02 15:04:05"
	MSGBUFSIZE = 50
)

var (
	stdLog     = NewConsoleLogger(INFO)
	levels     = [...]string{"TRACE", "DEBUG", "INFO ", "WARN ", "ERROR", "FATAL"}
	timeFormat = "2006/01/02 15:04:05"
	msgBufSize = 50
)

type Logger interface {
	Fatal(string, ...interface{})
	Error(string, ...interface{})
	Warn(string, ...interface{})
	Info(string, ...interface{})
	Debug(string, ...interface{})
	Trace(string, ...interface{})
	Level(int)
	Close() error
	output(msg *Message)
}

type Message struct {
	level int
	m     string
	time  time.Time
}

func TimeFormat(f string) {
	stdLog.TimeFormat(f)
}

func LvlStr(l int) string {
	if l >= TRACE && l <= FATAL {
		return levels[l]
	}
	return ""
}

func Level(o int) {
	stdLog.Level(o)
}

func Fatal(format string, v ...interface{}) {
	stdLog.output(&Message{FATAL, fmt.Sprintf(format, v...), time.Now()})
}

func Error(format string, v ...interface{}) {
	stdLog.output(&Message{ERROR, fmt.Sprintf(format, v...), time.Now()})
}

func Warn(format string, v ...interface{}) {
	stdLog.output(&Message{WARN, fmt.Sprintf(format, v...), time.Now()})
}

func Info(format string, v ...interface{}) {
	stdLog.output(&Message{INFO, fmt.Sprintf(format, v...), time.Now()})
}

func Debug(format string, v ...interface{}) {
	stdLog.output(&Message{DEBUG, fmt.Sprintf(format, v...), time.Now()})
}

func Trace(format string, v ...interface{}) {
	stdLog.output(&Message{TRACE, fmt.Sprintf(format, v...), time.Now()})
}
<|endoftext|>


<commit_before>package exchange

import (
	"testing"

	"github.com/thrasher-/gocryptotrader/config"
)

func TestGetName(t *testing.T) {
	GetName := ExchangeBase{
		Name: "TESTNAME",
	}

	name := GetName.GetName()
	if name != "TESTNAME" {
		t.Error("Test Failed - Exchange getName() returned incorrect name")
	}
}

func TestGetEnabledCurrencies(t *testing.T) {
	enabledPairs := []string{"BTCUSD", "BTCAUD", "LTCUSD", "LTCAUD"}
	GetEnabledCurrencies := ExchangeBase{
		Name:         "TESTNAME",
		EnabledPairs: enabledPairs,
	}

	enCurr := GetEnabledCurrencies.GetEnabledCurrencies()
	if enCurr[0] != "BTCUSD" {
		t.Error("Test Failed - Exchange GetEnabledCurrencies() incorrect string")
	}
}

func TestSetEnabled(t *testing.T) {
	GetEnabledCurrencies := ExchangeBase{
		Name:    "TESTNAME",
		Enabled: false,
	}

	GetEnabledCurrencies.SetEnabled(true)
	if !GetEnabledCurrencies.Enabled {
		t.Error("Test Failed - Exchange SetEnabled(true) did not set boolean")
	}
}

func TestIsEnabled(t *testing.T) {
	GetEnabledCurrencies := ExchangeBase{
		Name:    "TESTNAME",
		Enabled: false,
	}

	if GetEnabledCurrencies.IsEnabled() {
		t.Error("Test Failed - Exchange IsEnabled() did not return correct boolean")
	}
}

func TestSetAPIKeys(t *testing.T) {
	SetAPIKeys := ExchangeBase{
		Name:    "TESTNAME",
		Enabled: false,
	}

	SetAPIKeys.SetAPIKeys("RocketMan", "Digereedoo", "007", false)

	if SetAPIKeys.APIKey != "RocketMan" && SetAPIKeys.APISecret != "Digereedoo" && SetAPIKeys.ClientID != "007" {
		t.Error("Test Failed - Exchange SetAPIKeys() did not set correct values")
	}

}

func TestUpdateAvailableCurrencies(t *testing.T) {
	cfg := config.GetConfig()
	err := cfg.LoadConfig("../testdata/configtest.dat")
	if err != nil {
		t.Log("SOMETHING DONE HAPPENED!")
	}

	UAC := ExchangeBase{
		Name: "ANX",
	}
	exchangeProducts := []string{"ltc", "btc", "usd", "aud"}

	err2 := UAC.UpdateAvailableCurrencies(exchangeProducts)
	if err2 != nil {
		t.Errorf("Test Failed - Exchange UpdateAvailableCurrencies() error: %s", err2)
	}
}
<commit_msg>Changed incorrect names<commit_after>package exchange

import (
	"testing"

	"github.com/thrasher-/gocryptotrader/config"
)

func TestGetName(t *testing.T) {
	GetName := ExchangeBase{
		Name: "TESTNAME",
	}

	name := GetName.GetName()
	if name != "TESTNAME" {
		t.Error("Test Failed - Exchange getName() returned incorrect name")
	}
}

func TestGetEnabledCurrencies(t *testing.T) {
	enabledPairs := []string{"BTCUSD", "BTCAUD", "LTCUSD", "LTCAUD"}
	GetEnabledCurrencies := ExchangeBase{
		Name:         "TESTNAME",
		EnabledPairs: enabledPairs,
	}

	enCurr := GetEnabledCurrencies.GetEnabledCurrencies()
	if enCurr[0] != "BTCUSD" {
		t.Error("Test Failed - Exchange GetEnabledCurrencies() incorrect string")
	}
}

func TestSetEnabled(t *testing.T) {
	SetEnabled := ExchangeBase{
		Name:    "TESTNAME",
		Enabled: false,
	}

	SetEnabled.SetEnabled(true)
	if !SetEnabled.Enabled {
		t.Error("Test Failed - Exchange SetEnabled(true) did not set boolean")
	}
}

func TestIsEnabled(t *testing.T) {
	IsEnabled := ExchangeBase{
		Name:    "TESTNAME",
		Enabled: false,
	}

	if IsEnabled.IsEnabled() {
		t.Error("Test Failed - Exchange IsEnabled() did not return correct boolean")
	}
}

func TestSetAPIKeys(t *testing.T) {
	SetAPIKeys := ExchangeBase{
		Name:    "TESTNAME",
		Enabled: false,
	}

	SetAPIKeys.SetAPIKeys("RocketMan", "Digereedoo", "007", false)

	if SetAPIKeys.APIKey != "RocketMan" && SetAPIKeys.APISecret != "Digereedoo" && SetAPIKeys.ClientID != "007" {
		t.Error("Test Failed - Exchange SetAPIKeys() did not set correct values")
	}

}

func TestUpdateAvailableCurrencies(t *testing.T) {
	cfg := config.GetConfig()
	err := cfg.LoadConfig("../testdata/configtest.dat")
	if err != nil {
		t.Log("SOMETHING DONE HAPPENED!")
	}

	UAC := ExchangeBase{
		Name: "ANX",
	}
	exchangeProducts := []string{"ltc", "btc", "usd", "aud"}

	err2 := UAC.UpdateAvailableCurrencies(exchangeProducts)
	if err2 != nil {
		t.Errorf("Test Failed - Exchange UpdateAvailableCurrencies() error: %s", err2)
	}
}
<|endoftext|>


<commit_before>package imap

import (
	"errors"
	"fmt"
	"regexp"
	"strings"
	"time"

	"github.com/mitchellh/mapstructure"
	"github.com/yesnault/go-imap/imap"

	"github.com/ovh/venom"
	"github.com/ovh/venom/executors"
)

// Name for test imap
const Name = "imap"

var imapLogMask = imap.LogNone
var imapSafeLogMask = imap.LogNone

// New returns a new Test Exec
func New() venom.Executor {
	return &Executor{}
}

// Executor represents a Test Exec
type Executor struct {
	IMAPHost      string `json:"imaphost,omitempty" yaml:"imaphost,omitempty"`
	IMAPPort      string `json:"imapport,omitempty" yaml:"imapport,omitempty"`
	IMAPUser      string `json:"imapuser,omitempty" yaml:"imapuser,omitempty"`
	IMAPPassword  string `json:"imappassword,omitempty" yaml:"imappassword,omitempty"`
	MBox          string `json:"mbox,omitempty" yaml:"mbox,omitempty"`
	MBoxOnSuccess string `json:"mboxonsuccess,omitempty" yaml:"mboxonsuccess,omitempty"`
	SearchFrom    string `json:"searchfrom,omitempty" yaml:"searchfrom,omitempty"`
	SearchSubject string `json:"searchsubject,omitempty" yaml:"searchsubject,omitempty"`
	SearchBody    string `json:"searchbody,omitempty" yaml:"searchbody,omitempty"`
}

// Mail contains an analyzed mail
type Mail struct {
	From    string
	Subject string
	UID     uint32
	Body    string
}

// Result represents a step result
type Result struct {
	Executor    Executor `json:"executor,omitempty" yaml:"executor,omitempty"`
	Err         string   `json:"error" yaml:"error"`
	Subject     string   `json:"subject,omitempty" yaml:"subject,omitempty"`
	Body        string   `json:"body,omitempty" yaml:"body,omitempty"`
	TimeSeconds float64  `json:"timeSeconds,omitempty" yaml:"timeSeconds,omitempty"`
	TimeHuman   string   `json:"timeHuman,omitempty" yaml:"timeHuman,omitempty"`
}

// GetDefaultAssertions return default assertions for type exec
func (Executor) GetDefaultAssertions() *venom.StepAssertions {
	return &venom.StepAssertions{Assertions: []string{"result.err ShouldNotExist"}}
}

// Run execute TestStep of type exec
func (Executor) Run(ctx venom.TestCaseContext, l venom.Logger, step venom.TestStep) (venom.ExecutorResult, error) {
	var t Executor
	if err := mapstructure.Decode(step, &t); err != nil {
		return nil, err
	}

	start := time.Now()

	result := Result{Executor: t}
	find, errs := t.getMail(l)
	if errs != nil {
		result.Err = errs.Error()
	}
	if find != nil {
		result.Subject = find.Subject
		result.Body = find.Body
	} else if result.Err == "" {
		result.Err = "searched mail not found"
	}

	elapsed := time.Since(start)
	result.TimeSeconds = elapsed.Seconds()
	result.TimeHuman = fmt.Sprintf("%s", elapsed)
	result.Executor.IMAPPassword = "****hidden****" // do not output password

	return executors.Dump(result)
}

func (e *Executor) getMail(l venom.Logger) (*Mail, error) {
	if e.SearchFrom == "" && e.SearchSubject == "" && e.SearchBody == "" {
		return nil, fmt.Errorf("You have to use searchfrom and/or searchsubject and/or searchbody")
	}

	c, errc := connect(e.IMAPHost, e.IMAPPort, e.IMAPUser, e.IMAPPassword)
	if errc != nil {
		return nil, fmt.Errorf("Error while connecting:%s", errc.Error())
	}
	defer c.Logout(5 * time.Second)

	var box string

	if e.MBox == "" {
		box = "INBOX"
	} else {
		box = e.MBox
	}

	count, err := queryCount(c, box)
	if err != nil {
		return nil, fmt.Errorf("Error while queryCount:%s", err.Error())
	}

	l.Debugf("count messages:%d", count)

	if count == 0 {
		return nil, errors.New("No message to fetch")
	}

	messages, err := fetch(c, box, count, l)
	if err != nil {
		return nil, fmt.Errorf("Error while feching messages:%s", err.Error())
	}

	for _, msg := range messages {
		m, erre := extract(msg, l)
		if erre != nil {
			return nil, erre
		}

		found, errs := e.isSearched(m)
		if errs != nil {
			return nil, errs
		}

		if found {
			if e.MBoxOnSuccess != "" {
				l.Debugf("Move to %s", e.MBoxOnSuccess)
				if err := m.move(c, e.MBoxOnSuccess); err != nil {
					return nil, err
				}
			}
			return m, nil
		}
	}

	return nil, errors.New("Mail not found")
}

func (e *Executor) isSearched(m *Mail) (bool, error) {
	if e.SearchFrom != "" {
		ma, erra := regexp.MatchString(e.SearchFrom, m.From)
		if erra != nil || !ma {
			return false, erra
		}
	}
	if e.SearchSubject != "" {
		mb, errb := regexp.MatchString(e.SearchSubject, m.Subject)
		if errb != nil || !mb {
			return false, errb
		}
	}
	if e.SearchBody != "" {
		mc, errc := regexp.MatchString(e.SearchBody, m.Body)
		if errc != nil || !mc {
			return false, errc
		}
	}
	return true, nil
}

func (m *Mail) move(c *imap.Client, mbox string) error {
	seq, _ := imap.NewSeqSet("")
	seq.AddNum(m.UID)

	if _, err := c.UIDMove(seq, mbox); err != nil {
		return fmt.Errorf("Error while move msg to %s, err:%s", mbox, err.Error())
	}
	return nil
}

func connect(host, port, imapUsername, imapPassword string) (*imap.Client, error) {
	if !strings.Contains(host, ":") {
		if port == "" {
			port = ":993"
		} else if port != "" && !strings.HasPrefix(port, ":") {
			port = ":" + port
		}
	}

	c, errd := imap.DialTLS(host+port, nil)
	if errd != nil {
		return nil, fmt.Errorf("Unable to dial: %s", errd)
	}

	if c.Caps["STARTTLS"] {
		if _, err := check(c.StartTLS(nil)); err != nil {
			return nil, fmt.Errorf("Unable to start TLS: %s\n", err)
		}
	}

	c.SetLogMask(imapSafeLogMask)
	if _, err := check(c.Login(imapUsername, imapPassword)); err != nil {
		return nil, fmt.Errorf("Unable to login: %s", err)
	}
	c.SetLogMask(imapLogMask)

	return c, nil
}

func fetch(c *imap.Client, box string, nb uint32, l venom.Logger) ([]imap.Response, error) {
	l.Debugf("call Select")
	if _, err := c.Select(box, true); err != nil {
		l.Errorf("Error with select %s", err.Error())
		return []imap.Response{}, err
	}
	defer c.Close(false)

	seqset, _ := imap.NewSeqSet("1:*")

	cmd, err := c.Fetch(seqset, "ENVELOPE", "RFC822.HEADER", "RFC822.TEXT", "UID")
	if err != nil {
		l.Errorf("Error with fetch:%s", err)
		return []imap.Response{}, err
	}

	messages := []imap.Response{}
	for cmd.InProgress() {
		// Wait for the next response (no timeout)
		c.Recv(-1)

		// Process command data
		for _, rsp := range cmd.Data {
			messages = append(messages, *rsp)
		}
		cmd.Data = nil
		c.Data = nil
	}
	l.Debugf("Nb messages fetch:%d", len(messages))
	return messages, nil
}

func queryCount(imapClient *imap.Client, box string) (uint32, error) {
	cmd, errc := check(imapClient.Status(box))
	if errc != nil {
		return 0, errc
	}

	var count uint32
	for _, result := range cmd.Data {
		mailboxStatus := result.MailboxStatus()
		if mailboxStatus != nil {
			count += mailboxStatus.Messages
		}
	}

	return count, nil
}

func check(cmd *imap.Command, erri error) (*imap.Command, error) {
	if erri != nil {
		return nil, erri
	}

	if _, err := cmd.Result(imap.OK); err != nil {
		return nil, err
	}

	return cmd, nil
}
<commit_msg>feat(executor/imap): add deleteonsuccess option to delete a mail once found. (#104)<commit_after>package imap

import (
	"errors"
	"fmt"
	"regexp"
	"strings"
	"time"

	"github.com/mitchellh/mapstructure"
	"github.com/yesnault/go-imap/imap"

	"github.com/ovh/venom"
	"github.com/ovh/venom/executors"
)

// Name for test imap
const Name = "imap"

var imapLogMask = imap.LogNone
var imapSafeLogMask = imap.LogNone

// New returns a new Test Exec
func New() venom.Executor {
	return &Executor{}
}

// Executor represents a Test Exec
type Executor struct {
	IMAPHost        string `json:"imaphost,omitempty" yaml:"imaphost,omitempty"`
	IMAPPort        string `json:"imapport,omitempty" yaml:"imapport,omitempty"`
	IMAPUser        string `json:"imapuser,omitempty" yaml:"imapuser,omitempty"`
	IMAPPassword    string `json:"imappassword,omitempty" yaml:"imappassword,omitempty"`
	MBox            string `json:"mbox,omitempty" yaml:"mbox,omitempty"`
	MBoxOnSuccess   string `json:"mboxonsuccess,omitempty" yaml:"mboxonsuccess,omitempty"`
	DeleteOnSuccess bool   `json:"deleteonsuccess,omitempty" yaml:"deleteonsuccess,omitempty"`
	SearchFrom      string `json:"searchfrom,omitempty" yaml:"searchfrom,omitempty"`
	SearchSubject   string `json:"searchsubject,omitempty" yaml:"searchsubject,omitempty"`
	SearchBody      string `json:"searchbody,omitempty" yaml:"searchbody,omitempty"`
}

// Mail contains an analyzed mail
type Mail struct {
	From    string
	Subject string
	UID     uint32
	Body    string
}

// Result represents a step result
type Result struct {
	Executor    Executor `json:"executor,omitempty" yaml:"executor,omitempty"`
	Err         string   `json:"error" yaml:"error"`
	Subject     string   `json:"subject,omitempty" yaml:"subject,omitempty"`
	Body        string   `json:"body,omitempty" yaml:"body,omitempty"`
	TimeSeconds float64  `json:"timeSeconds,omitempty" yaml:"timeSeconds,omitempty"`
	TimeHuman   string   `json:"timeHuman,omitempty" yaml:"timeHuman,omitempty"`
}

// GetDefaultAssertions return default assertions for type exec
func (Executor) GetDefaultAssertions() *venom.StepAssertions {
	return &venom.StepAssertions{Assertions: []string{"result.err ShouldNotExist"}}
}

// Run execute TestStep of type exec
func (Executor) Run(ctx venom.TestCaseContext, l venom.Logger, step venom.TestStep) (venom.ExecutorResult, error) {
	var t Executor
	if err := mapstructure.Decode(step, &t); err != nil {
		return nil, err
	}

	start := time.Now()

	result := Result{Executor: t}
	find, errs := t.getMail(l)
	if errs != nil {
		result.Err = errs.Error()
	}
	if find != nil {
		result.Subject = find.Subject
		result.Body = find.Body
	} else if result.Err == "" {
		result.Err = "searched mail not found"
	}

	elapsed := time.Since(start)
	result.TimeSeconds = elapsed.Seconds()
	result.TimeHuman = fmt.Sprintf("%s", elapsed)
	result.Executor.IMAPPassword = "****hidden****" // do not output password

	return executors.Dump(result)
}

func (e *Executor) getMail(l venom.Logger) (*Mail, error) {
	if e.SearchFrom == "" && e.SearchSubject == "" && e.SearchBody == "" {
		return nil, fmt.Errorf("You have to use searchfrom and/or searchsubject and/or searchbody")
	}

	c, errc := connect(e.IMAPHost, e.IMAPPort, e.IMAPUser, e.IMAPPassword)
	if errc != nil {
		return nil, fmt.Errorf("Error while connecting:%s", errc.Error())
	}
	defer c.Logout(5 * time.Second)

	var box string

	if e.MBox == "" {
		box = "INBOX"
	} else {
		box = e.MBox
	}

	count, err := queryCount(c, box)
	if err != nil {
		return nil, fmt.Errorf("Error while queryCount:%s", err.Error())
	}

	l.Debugf("count messages:%d", count)

	if count == 0 {
		return nil, errors.New("No message to fetch")
	}

	messages, err := fetch(c, box, count, l)
	if err != nil {
		return nil, fmt.Errorf("Error while feching messages:%s", err.Error())
	}

	for _, msg := range messages {
		m, erre := extract(msg, l)
		if erre != nil {
			return nil, erre
		}

		found, errs := e.isSearched(m)
		if errs != nil {
			return nil, errs
		}

		if found {
			if e.DeleteOnSuccess {
				l.Debugf("Delete message %s", m.UID)
				if err := m.delete(c); err != nil {
					return nil, err
				}
			} else if e.MBoxOnSuccess != "" {
				l.Debugf("Move to %s", e.MBoxOnSuccess)
				if err := m.move(c, e.MBoxOnSuccess); err != nil {
					return nil, err
				}
			}
			return m, nil
		}
	}

	return nil, errors.New("Mail not found")
}

func (e *Executor) isSearched(m *Mail) (bool, error) {
	if e.SearchFrom != "" {
		ma, erra := regexp.MatchString(e.SearchFrom, m.From)
		if erra != nil || !ma {
			return false, erra
		}
	}
	if e.SearchSubject != "" {
		mb, errb := regexp.MatchString(e.SearchSubject, m.Subject)
		if errb != nil || !mb {
			return false, errb
		}
	}
	if e.SearchBody != "" {
		mc, errc := regexp.MatchString(e.SearchBody, m.Body)
		if errc != nil || !mc {
			return false, errc
		}
	}
	return true, nil
}

func (m *Mail) move(c *imap.Client, mbox string) error {
	seq, _ := imap.NewSeqSet("")
	seq.AddNum(m.UID)

	if _, err := c.UIDMove(seq, mbox); err != nil {
		return fmt.Errorf("Error while move msg to %s, err:%s", mbox, err.Error())
	}
	return nil
}

func (m *Mail) delete(c *imap.Client) error {
	seq, _ := imap.NewSeqSet("")
	seq.AddNum(m.UID)

	if _, err := c.UIDStore(seq, "+FLAGS.SILENT", imap.NewFlagSet(`\Deleted`)); err != nil {
		return fmt.Errorf("Error while deleting msg, err: %s", err.Error())
	}
	if _, err := c.Expunge(nil); err != nil {
		return fmt.Errorf("Error while expunging messages: err: %s", err.Error())
	}
	return nil
}

func connect(host, port, imapUsername, imapPassword string) (*imap.Client, error) {
	if !strings.Contains(host, ":") {
		if port == "" {
			port = ":993"
		} else if port != "" && !strings.HasPrefix(port, ":") {
			port = ":" + port
		}
	}

	c, errd := imap.DialTLS(host+port, nil)
	if errd != nil {
		return nil, fmt.Errorf("Unable to dial: %s", errd)
	}

	if c.Caps["STARTTLS"] {
		if _, err := check(c.StartTLS(nil)); err != nil {
			return nil, fmt.Errorf("Unable to start TLS: %s\n", err)
		}
	}

	c.SetLogMask(imapSafeLogMask)
	if _, err := check(c.Login(imapUsername, imapPassword)); err != nil {
		return nil, fmt.Errorf("Unable to login: %s", err)
	}
	c.SetLogMask(imapLogMask)

	return c, nil
}

func fetch(c *imap.Client, box string, nb uint32, l venom.Logger) ([]imap.Response, error) {
	l.Debugf("call Select")
	if _, err := c.Select(box, false); err != nil {
		l.Errorf("Error with select %s", err.Error())
		return []imap.Response{}, err
	}
	defer c.Close(false)

	seqset, _ := imap.NewSeqSet("1:*")

	cmd, err := c.Fetch(seqset, "ENVELOPE", "RFC822.HEADER", "RFC822.TEXT", "UID")
	if err != nil {
		l.Errorf("Error with fetch:%s", err)
		return []imap.Response{}, err
	}

	messages := []imap.Response{}
	for cmd.InProgress() {
		// Wait for the next response (no timeout)
		c.Recv(-1)

		// Process command data
		for _, rsp := range cmd.Data {
			messages = append(messages, *rsp)
		}
		cmd.Data = nil
		c.Data = nil
	}
	l.Debugf("Nb messages fetch:%d", len(messages))
	return messages, nil
}

func queryCount(imapClient *imap.Client, box string) (uint32, error) {
	cmd, errc := check(imapClient.Status(box))
	if errc != nil {
		return 0, errc
	}

	var count uint32
	for _, result := range cmd.Data {
		mailboxStatus := result.MailboxStatus()
		if mailboxStatus != nil {
			count += mailboxStatus.Messages
		}
	}

	return count, nil
}

func check(cmd *imap.Command, erri error) (*imap.Command, error) {
	if erri != nil {
		return nil, erri
	}

	if _, err := cmd.Result(imap.OK); err != nil {
		return nil, err
	}

	return cmd, nil
}
<|endoftext|>


<commit_before>// Package handle provides utilites for working with sandstorm's util.Handle
package handle

import (
	"context"

	"zenhack.net/go/sandstorm/capnp/util"

	"zombiezen.com/go/capnproto2/server"
)

// WithCancel is like context.WithCancel, except that it returns a Handle
// instead of a context.CancelFunc. When the handle is dropped, the context
// is cancelled.
func WithCancel(ctx context.Context) (context.Context, util.Handle) {
	ctx, cancel := context.WithCancel(ctx)

	// The server policy doesn't matter, because handles have no methods:
	return ctx, util.Handle_ServerToClient(&cancelHandle{cancel}, &server.Policy{})
}

type cancelHandle struct {
	cancel context.CancelFunc
}

func (h *cancelHandle) Close() error {
	h.cancel()
	return nil
}
<commit_msg>cancelHandle: implement Ping and rename Close -> Shutdown.<commit_after>// Package handle provides utilites for working with sandstorm's util.Handle
package handle

import (
	"context"

	"zenhack.net/go/sandstorm/capnp/util"

	"zombiezen.com/go/capnproto2/server"
)

// WithCancel is like context.WithCancel, except that it returns a Handle
// instead of a context.CancelFunc. When the handle is dropped, the context
// is cancelled.
func WithCancel(ctx context.Context) (context.Context, util.Handle) {
	ctx, cancel := context.WithCancel(ctx)

	// The server policy doesn't matter, because handles have no methods:
	return ctx, util.Handle_ServerToClient(&cancelHandle{cancel}, &server.Policy{})
}

type cancelHandle struct {
	cancel context.CancelFunc
}

func (h *cancelHandle) Ping(context.Context, util.Handle_ping) error {
	return nil
}

func (h *cancelHandle) Shutdown() error {
	h.cancel()
	return nil
}
<|endoftext|>


<commit_before>/*
Copyright 2021 Google LLC

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

    https://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 docker

import (
	"context"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"

	"github.com/containerd/containerd/containers"
	"github.com/containerd/containerd/metadata"
	"github.com/google/container-explorer/explorers"
	log "github.com/sirupsen/logrus"
	bolt "go.etcd.io/bbolt"
)

const (
	configV1Filename  = "config.json"
	configV2Filename  = "config.v2.json"
	containersDirName = "containers"
)

type explorer struct {
	root          string // docker root directory
	contaierdroot string
	manifest      string
	snapshot      string
	mdb           *bolt.DB // manifest database file
}

// NewExplorer returns a ContainerExplorer interface to explorer docker managed
// containers.
func NewExplorer(root string, containerdroot string, manifest string, snapshot string) (explorers.ContainerExplorer, error) {
	opt := &bolt.Options{
		ReadOnly: true,
	}
	db, err := bolt.Open(manifest, 0444, opt)
	if err != nil {
		return &explorer{}, err
	}

	return &explorer{
		root:          root,
		contaierdroot: containerdroot,
		manifest:      manifest,
		snapshot:      snapshot,
		mdb:           db,
	}, nil
}

// SnapshotRoot returns the snapshot root director for docker managed
// containers.
func (e *explorer) SnapshotRoot(snapshotter string) string {
	// TODO(rmaskey): implement the function
	return ""
}

// ListNamespaces returns namespaces for docker managed containers.
func (e *explorer) ListNamespaces(ctx context.Context) ([]string, error) {
	var nss []string

	// Namespaces in metadata file i.e. meta.db
	// in /var/lib/containerd/io.containerd.metadata.v1.bolt/meta.db
	err := e.mdb.View(func(tx *bolt.Tx) error {
		store := metadata.NewNamespaceStore(tx)
		results, err := store.List(ctx)
		if err != nil {
			return err
		}
		nss = results
		return nil
	})
	if err != nil {
		return nil, err
	}

	// TODO(rmaskey): implement the function

	return nss, nil
}

// GetContainerIDs returns container ID
func (e *explorer) GetContainerIDs(ctx context.Context, containerdir string) ([]string, error) {
	containerpaths, err := filepath.Glob(filepath.Join(containerdir, "*"))
	if err != nil {
		return nil, err
	}

	var containerids []string
	for _, containerpath := range containerpaths {
		_, containerid := filepath.Split(containerpath)
		containerids = append(containerids, containerid)
	}
	return containerids, nil
}

// ListContainers returns container information.
func (e *explorer) ListContainers(ctx context.Context) ([]explorers.Container, error) {
	containersdir := filepath.Join(e.root, containersDirName)
	containerids, err := e.GetContainerIDs(ctx, containersdir)
	if err != nil {
		return nil, err
	}

	var cecontainers []explorers.Container

	for _, containerid := range containerids {
		containerpath := filepath.Join(containersdir, containerid)

		// Read docker config version 2
		configpath := filepath.Join(containerpath, configV2Filename)
		if fileExists(configpath) {
			data, err := ioutil.ReadFile(configpath)
			if err != nil {
				return nil, fmt.Errorf("reading docker config file %s %v", configV2Filename, err)
			}

			var config ConfigFile
			if err := json.Unmarshal(data, &config); err != nil {
				return nil, fmt.Errorf("unmarshalling config file data %v", err)
			}
			cecontainer := convertToContainerExplorerContainer(config)
			cecontainers = append(cecontainers, cecontainer)

			continue
		}

		// Read docker config version 1
		configpath = filepath.Join(containerpath, configV1Filename)
		if fileExists(configpath) {
			// TODO (rmaskey): parse v1 confg
			continue
		}

		log.WithFields(log.Fields{
			"containerid": containerid,
		}).Error("configuration file not found")
	}

	return cecontainers, nil
}

// ListImages returns container images.
func (e *explorer) ListImages(ctx context.Context) ([]explorers.Image, error) {
	// TODO(rmaskey): implement the function
	return nil, nil
}

// ListContent returns content information.
func (e *explorer) ListContent(ctx context.Context) ([]explorers.Content, error) {
	// TODO(rmaskey): implement the function
	return nil, nil
}

// ListSnapshots returns snapshot information.
func (e *explorer) ListSnapshots(ctx context.Context) ([]explorers.SnapshotKeyInfo, error) {
	// TODO(rmaskey): implement the function
	return nil, nil
}

// InfoContainer returns container internal information.
func (e *explorer) InfoContainer(ctx context.Context, containerid string, spec bool) (interface{}, error) {
	// default return
	return nil, nil
}

// MountContainer mounts a container to the specified path
func (e *explorer) MountContainer(ctx context.Context, containerid string, mountpoint string) error {
	return nil
}

// MountAllContainers mounts all the containers
func (e *explorer) MountAllContainers(ctx context.Context, mountpoint string, skipsupportcontainers bool) error {
	// default
	return nil
}

// Close releases internal resources.
func (e *explorer) Close() error {
	return e.mdb.Close()
}

func fileExists(path string) bool {
	_, err := os.Stat(path)
	return err == nil
}

// convertToContainerExplorerContainer maps docker config data to container
// explorer container structure
func convertToContainerExplorerContainer(config ConfigFile) explorers.Container {
	var exposedports []string

	if config.Config.ExposedPorts != nil {
		for k := range config.Config.ExposedPorts {
			exposedports = append(exposedports, k)
		}
	}

	return explorers.Container{
		Hostname: config.Config.Hostname,
		Container: containers.Container{
			ID:          config.ID,
			CreatedAt:   config.Created,
			Image:       config.Image,
			Snapshotter: config.Driver,
			Runtime: containers.RuntimeInfo{
				Name: config.Name,
			},
		},
		Running:      config.State.Running,
		ExposedPorts: exposedports,
	}
}
<commit_msg>Added support to list docker images<commit_after>/*
Copyright 2021 Google LLC

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

    https://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 docker

import (
	"context"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"
	"time"

	"github.com/containerd/containerd/containers"
	"github.com/containerd/containerd/images"
	"github.com/containerd/containerd/metadata"
	"github.com/google/container-explorer/explorers"
	digest "github.com/opencontainers/go-digest"
	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
	log "github.com/sirupsen/logrus"
	bolt "go.etcd.io/bbolt"
)

const (
	configV1Filename     = "config.json"
	configV2Filename     = "config.v2.json"
	containersDirName    = "containers"
	repositoriesDirName  = "image"
	repositoriesFileName = "repositories.json"
	storageOverlay2      = "overlay2"
)

type ImageName map[string]string

type ImageRepository struct {
	Repositories map[string]ImageName
}

type explorer struct {
	root          string // docker root directory
	contaierdroot string
	manifest      string
	snapshot      string
	mdb           *bolt.DB // manifest database file
}

// NewExplorer returns a ContainerExplorer interface to explorer docker managed
// containers.
func NewExplorer(root string, containerdroot string, manifest string, snapshot string) (explorers.ContainerExplorer, error) {
	opt := &bolt.Options{
		ReadOnly: true,
	}
	db, err := bolt.Open(manifest, 0444, opt)
	if err != nil {
		return &explorer{}, err
	}

	return &explorer{
		root:          root,
		contaierdroot: containerdroot,
		manifest:      manifest,
		snapshot:      snapshot,
		mdb:           db,
	}, nil
}

// SnapshotRoot returns the snapshot root director for docker managed
// containers.
func (e *explorer) SnapshotRoot(snapshotter string) string {
	// TODO(rmaskey): implement the function
	return ""
}

// ListNamespaces returns namespaces for docker managed containers.
func (e *explorer) ListNamespaces(ctx context.Context) ([]string, error) {
	var nss []string

	// Namespaces in metadata file i.e. meta.db
	// in /var/lib/containerd/io.containerd.metadata.v1.bolt/meta.db
	err := e.mdb.View(func(tx *bolt.Tx) error {
		store := metadata.NewNamespaceStore(tx)
		results, err := store.List(ctx)
		if err != nil {
			return err
		}
		nss = results
		return nil
	})
	if err != nil {
		return nil, err
	}

	// TODO(rmaskey): implement the function

	return nss, nil
}

// GetContainerIDs returns container ID
func (e *explorer) GetContainerIDs(ctx context.Context, containerdir string) ([]string, error) {
	containerpaths, err := filepath.Glob(filepath.Join(containerdir, "*"))
	if err != nil {
		return nil, err
	}

	var containerids []string
	for _, containerpath := range containerpaths {
		_, containerid := filepath.Split(containerpath)
		containerids = append(containerids, containerid)
	}
	return containerids, nil
}

// ListContainers returns container information.
func (e *explorer) ListContainers(ctx context.Context) ([]explorers.Container, error) {
	containersdir := filepath.Join(e.root, containersDirName)
	containerids, err := e.GetContainerIDs(ctx, containersdir)
	if err != nil {
		return nil, err
	}

	var cecontainers []explorers.Container

	for _, containerid := range containerids {
		containerpath := filepath.Join(containersdir, containerid)

		// Read docker config version 2
		configpath := filepath.Join(containerpath, configV2Filename)
		if fileExists(configpath) {
			data, err := ioutil.ReadFile(configpath)
			if err != nil {
				return nil, fmt.Errorf("reading docker config file %s %v", configV2Filename, err)
			}

			var config ConfigFile
			if err := json.Unmarshal(data, &config); err != nil {
				return nil, fmt.Errorf("unmarshalling config file data %v", err)
			}
			cecontainer := convertToContainerExplorerContainer(config)
			cecontainers = append(cecontainers, cecontainer)

			continue
		}

		// Read docker config version 1
		configpath = filepath.Join(containerpath, configV1Filename)
		if fileExists(configpath) {
			// TODO (rmaskey): parse v1 confg
			continue
		}

		log.WithFields(log.Fields{
			"containerid": containerid,
		}).Error("configuration file not found")
	}

	return cecontainers, nil
}

// structure to hold limited docker image information
//
// The structure hold information from the file
// /var/lib/docker/image/overlay2/imagedb/content/sha256/<imageid>
type rootfs struct {
	Rfstype string   `json:"type"`
	DiffIds []string `json:"diff_ids"`
}

type historyItem map[string]interface{}

type imageContentSummary struct {
	Architecture    string        `json:"architecture"`
	Config          Config        `json:"config"`
	Container       string        `json:"container"`
	ContainerConfig Config        `json:"container_config"`
	Created         time.Time     `json:"created"`
	DockerVersion   string        `json:"docker_version"`
	History         []historyItem `json:"history"`
	Os              string        `json:"os"`
	Rootfs          rootfs        `json:"rootfs"`
}

// ListImages returns information about docker images.
func (e *explorer) ListImages(ctx context.Context) ([]explorers.Image, error) {
	// TODO (rmaskey): Handle docker version 1 images

	// Docker version 2
	//
	// Check for valid image repositories directory
	repositoriesdir := filepath.Join(e.root, repositoriesDirName)
	if !fileExists(repositoriesdir) {
		return nil, fmt.Errorf("valid image repositories directory %s not found", repositoriesdir)
	}

	storagedirs, err := filepath.Glob(filepath.Join(repositoriesdir, "*"))
	if err != nil {
		return nil, fmt.Errorf("listing storage directories %v", err)
	}

	var ceimages []explorers.Image

	for _, storagedir := range storagedirs {
		_, storagename := filepath.Split(storagedir)
		repositoriesfile := filepath.Join(storagedir, repositoriesFileName)

		log.WithFields(log.Fields{
			"storagename":      storagename,
			"storagedir":       storagedir,
			"repositoriesfile": repositoriesfile,
		}).Debug("image repository file")

		data, err := ioutil.ReadFile(repositoriesfile)
		if err != nil {
			return nil, fmt.Errorf("failed read repository file %v. %v", repositoriesfile, err)
		}

		var r ImageRepository
		if err := json.Unmarshal(data, &r); err != nil {
			return nil, fmt.Errorf("unmarshalling image repository file %s. %v", repositoriesfile, err)
		}

		for _, distvalue := range r.Repositories {
			for k, v := range distvalue {
				image := images.Image{
					Name: k,
					Target: ocispec.Descriptor{
						Digest: digest.Digest(v),
					},
				}

				if storagename == storageOverlay2 {
					imagecontent, err := readImageContent(storagename, storagedir, image.Target.Digest)
					if err != nil {
						log.Error("reading image content file ", err)
					} else {
						image.CreatedAt = imagecontent.Created
					}
				}

				ceimages = append(ceimages, explorers.Image{
					Image: image,
				})
			}
		}
	}

	return ceimages, nil
}

// ListContent returns content information.
func (e *explorer) ListContent(ctx context.Context) ([]explorers.Content, error) {
	// TODO(rmaskey): implement the function
	return nil, nil
}

// ListSnapshots returns snapshot information.
func (e *explorer) ListSnapshots(ctx context.Context) ([]explorers.SnapshotKeyInfo, error) {
	// TODO(rmaskey): implement the function
	return nil, nil
}

// InfoContainer returns container internal information.
func (e *explorer) InfoContainer(ctx context.Context, containerid string, spec bool) (interface{}, error) {
	// default return
	return nil, nil
}

// MountContainer mounts a container to the specified path
func (e *explorer) MountContainer(ctx context.Context, containerid string, mountpoint string) error {
	return nil
}

// MountAllContainers mounts all the containers
func (e *explorer) MountAllContainers(ctx context.Context, mountpoint string, skipsupportcontainers bool) error {
	// default
	return nil
}

// Close releases internal resources.
func (e *explorer) Close() error {
	return e.mdb.Close()
}

func fileExists(path string) bool {
	_, err := os.Stat(path)
	return err == nil
}

// convertToContainerExplorerContainer maps docker config data to container
// explorer container structure
func convertToContainerExplorerContainer(config ConfigFile) explorers.Container {
	var exposedports []string

	if config.Config.ExposedPorts != nil {
		for k := range config.Config.ExposedPorts {
			exposedports = append(exposedports, k)
		}
	}

	return explorers.Container{
		Hostname: config.Config.Hostname,
		Container: containers.Container{
			ID:          config.ID,
			CreatedAt:   config.Created,
			Image:       config.Image,
			Snapshotter: config.Driver,
			Runtime: containers.RuntimeInfo{
				Name: config.Name,
			},
		},
		Running:      config.State.Running,
		ExposedPorts: exposedports,
	}
}

// readImageContent reads the content of overlay2 image content
func readImageContent(storagename string, storagepath string, digest digest.Digest) (imageContentSummary, error) {
	m := strings.Split(string(digest), ":")
	if len(m) != 2 {
		return imageContentSummary{}, fmt.Errorf("expecting two colon separated values")
	}
	algo := m[0]
	filename := m[1]

	imagecontentfile := filepath.Join(storagepath, "imagedb", "content", algo, filename)
	log.WithFields(log.Fields{
		"filename": imagecontentfile,
	}).Debug("reading docker image content file")

	data, err := ioutil.ReadFile(imagecontentfile)
	if err != nil {
		log.WithFields(log.Fields{
			"storage name": storagename,
			"algo":         algo,
			"filename":     filename,
		}).Debug("reading docker image content file")

		return imageContentSummary{}, err
	}

	var imagecontent imageContentSummary
	if err := json.Unmarshal(data, &imagecontent); err != nil {
		return imageContentSummary{}, err
	}

	return imagecontent, nil
}
<|endoftext|>


<commit_before>package integration_test

import (
	"fmt"
	"github.com/APTrust/exchange/dpn/network"
	"github.com/APTrust/exchange/util/fileutil"
	"github.com/APTrust/exchange/util/testutil"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"net/url"
	"path/filepath"
	"testing"
)

func TestItemsCopiedToStaging(t *testing.T) {
	// Make sure that each of the expected bags has shown
	// up in our test staging area.
	if !testutil.ShouldRunIntegrationTests() {
		t.Skip("Skipping integration test. Set ENV var RUN_EXCHANGE_INTEGRATION=true if you want to run them.")
    }
	for i := 2; i <= 5; i++ {
		_context, err := testutil.GetContext("integration.json")
		require.Nil(t, err, "Could not create context")
		filename := fmt.Sprintf("00000000-0000-4000-a000-00000000000%d.tar", i)
		path := filepath.Join(_context.Config.DPN.StagingDirectory, filename)
		assert.True(t, fileutil.FileExists(path), "File %s was not copied", path)
	}
}

func TestReplicationDPNWorkItems(t *testing.T) {
	// Make sure that our code updated the DPNWorkItem record for each
	// ReplicationTransfer.
	if !testutil.ShouldRunIntegrationTests() {
		t.Skip("Skipping integration test. Set ENV var RUN_EXCHANGE_INTEGRATION=true if you want to run them.")
    }
	_context, err := testutil.GetContext("integration.json")
	require.Nil(t, err, "Could not create context")
	params := url.Values{}
	params.Set("task", "replication")
	params.Set("page", "1")
	params.Set("per_page", "100")
	resp := _context.PharosClient.DPNWorkItemList(params)
	require.Nil(t, resp.Error)
	for _, item := range resp.DPNWorkItems() {
		require.NotNil(t, item.QueuedAt, "QueuedAt nil for %s", item.Identifier)
		require.NotNil(t, item.Note, "Note nil for %s", item.Identifier)
		require.NotNil(t, item.State, "State nil for %s", item.Identifier)
		assert.False(t, item.QueuedAt.IsZero(), "QueuedAt is zero for %s", item.Identifier)
		assert.NotEmpty(t, item.Note, "Note is empty for %s", item.Identifier)
		assert.True(t, len(*item.State) > 250, "State JSON too short for %s", item.Identifier)
	}
}

func TestCopyResultSentToRemoteNodes(t *testing.T) {
	// Query the FromNode of each copied bag to make sure that
	// we sent a fixity value back to the ingest node.
	if !testutil.ShouldRunIntegrationTests() {
		t.Skip("Skipping integration test. Set ENV var RUN_EXCHANGE_INTEGRATION=true if you want to run them.")
    }
	_context, err := testutil.GetContext("integration.json")
	require.Nil(t, err, "Could not create context")
	dpnClient, err := network.NewDPNRestClient(
		_context.Config.DPN.RestClient.LocalServiceURL,
		_context.Config.DPN.RestClient.LocalAPIRoot,
		_context.Config.DPN.RestClient.LocalAuthToken,
		_context.Config.DPN.LocalNode,
		_context.Config.DPN)
	require.Nil(t, err, "Couldn't build DPN REST client: %v", err)

	// These identifiers are in the fixture data for dpn-server.
	xferIdentifiers := []string {
		"20000000-0000-4000-a000-000000000007", // from chron
		"30000000-0000-4000-a000-000000000001", // from hathi
		"40000000-0000-4000-a000-000000000013", // from sdr
		"50000000-0000-4000-a000-000000000019", // from tdr
	}

	for _, identifier := range xferIdentifiers {
		resp := dpnClient.ReplicationTransferGet(identifier)
		require.Nil(t, resp.Error)
		xfer := resp.ReplicationTransfer()
		require.NotNil(t, xfer, "ReplicationTransfer %s is missing", identifier)
		assert.NotEmpty(t, xfer.FixityValue, "Empty fixity for %s", identifier)
		assert.True(t, xfer.StoreRequested, "StoreRequested should not be false for %s", identifier)
	}
}

func TestValidationQueue(t *testing.T) {
	// Make sure the copied DPN bags made it into the validation queue.
	if !testutil.ShouldRunIntegrationTests() {
		t.Skip("Skipping integration test. Set ENV var RUN_EXCHANGE_INTEGRATION=true if you want to run them.")
    }
	_context, err := testutil.GetContext("integration.json")
	require.Nil(t, err, "Could not create context")
	stats, err := _context.NSQClient.GetStats()
	require.Nil(t, err)
	foundTopic := false
	for _, topic := range stats.Data.Topics {
		if topic.TopicName == _context.Config.DPN.DPNValidationWorker.NsqTopic {
			// Should have 4 items. One from each remote node.
			foundTopic = true
			assert.EqualValues(t, uint64(4), topic.MessageCount)
		}
	}
	assert.True(t, foundTopic, "Nothing was queued in %s",
		_context.Config.DPN.DPNValidationWorker.NsqTopic)
}
<commit_msg>All dpn_copy tests are now passing<commit_after>package integration_test

import (
	"fmt"
	"github.com/APTrust/exchange/dpn/network"
	"github.com/APTrust/exchange/util/fileutil"
	"github.com/APTrust/exchange/util/testutil"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"net/url"
	"path/filepath"
	"testing"
)

func TestItemsCopiedToStaging(t *testing.T) {
	// Make sure that each of the expected bags has shown
	// up in our test staging area.
	if !testutil.ShouldRunIntegrationTests() {
		t.Skip("Skipping integration test. Set ENV var RUN_EXCHANGE_INTEGRATION=true if you want to run them.")
    }
	for i := 2; i <= 5; i++ {
		_context, err := testutil.GetContext("integration.json")
		require.Nil(t, err, "Could not create context")
		filename := fmt.Sprintf("00000000-0000-4000-a000-00000000000%d.tar", i)
		path := filepath.Join(_context.Config.DPN.StagingDirectory, filename)
		assert.True(t, fileutil.FileExists(path), "File %s was not copied", path)
	}
}

func TestReplicationDPNWorkItems(t *testing.T) {
	// Make sure that our code updated the DPNWorkItem record for each
	// ReplicationTransfer.
	if !testutil.ShouldRunIntegrationTests() {
		t.Skip("Skipping integration test. Set ENV var RUN_EXCHANGE_INTEGRATION=true if you want to run them.")
    }
	_context, err := testutil.GetContext("integration.json")
	require.Nil(t, err, "Could not create context")
	params := url.Values{}
	params.Set("task", "replication")
	params.Set("page", "1")
	params.Set("per_page", "100")
	resp := _context.PharosClient.DPNWorkItemList(params)
	require.Nil(t, resp.Error)
	for _, item := range resp.DPNWorkItems() {
		require.NotNil(t, item.QueuedAt, "QueuedAt nil for %s", item.Identifier)
		require.NotNil(t, item.Note, "Note nil for %s", item.Identifier)
		require.NotNil(t, item.State, "State nil for %s", item.Identifier)
		assert.False(t, item.QueuedAt.IsZero(), "QueuedAt is zero for %s", item.Identifier)
		assert.NotEmpty(t, item.Note, "Note is empty for %s", item.Identifier)
		assert.True(t, len(*item.State) > 250, "State JSON too short for %s", item.Identifier)
	}
}

func TestCopyResultSentToRemoteNodes(t *testing.T) {
	// Query the FromNode of each copied bag to make sure that
	// we sent a fixity value back to the ingest node.
	if !testutil.ShouldRunIntegrationTests() {
		t.Skip("Skipping integration test. Set ENV var RUN_EXCHANGE_INTEGRATION=true if you want to run them.")
    }
	_context, err := testutil.GetContext("integration.json")
	require.Nil(t, err, "Could not create context")
	dpnClient, err := network.NewDPNRestClient(
		_context.Config.DPN.RestClient.LocalServiceURL,
		_context.Config.DPN.RestClient.LocalAPIRoot,
		_context.Config.DPN.RestClient.LocalAuthToken,
		_context.Config.DPN.LocalNode,
		_context.Config.DPN)
	require.Nil(t, err, "Couldn't build DPN REST client: %v", err)

	remoteClients, err := dpnClient.GetRemoteClients()
	require.Nil(t, err, "Couldn't build remote DPN clients: %v", err)

	// These identifiers are in the fixture data for dpn-server.
	// Key is the FromNode, value is the ReplicationId
	xferIdentifiers := map[string]string {
		"chron": "20000000-0000-4000-a000-000000000007",
		"hathi": "30000000-0000-4000-a000-000000000001",
		"sdr":   "40000000-0000-4000-a000-000000000013",
		"tdr":   "50000000-0000-4000-a000-000000000019",
	}

	for fromNode, identifier := range xferIdentifiers {
		client := remoteClients[fromNode]
		require.NotNil(t, client, "No DPN REST client for %s", fromNode)
		resp := client.ReplicationTransferGet(identifier)
		require.Nil(t, resp.Error)
		xfer := resp.ReplicationTransfer()
		require.NotNil(t, xfer, "ReplicationTransfer %s is missing", identifier)
		assert.NotEmpty(t, xfer.FixityValue, "Empty fixity for %s", identifier)
		assert.True(t, xfer.StoreRequested, "StoreRequested should not be false for %s", identifier)
	}
}

func TestValidationQueue(t *testing.T) {
	// Make sure the copied DPN bags made it into the validation queue.
	if !testutil.ShouldRunIntegrationTests() {
		t.Skip("Skipping integration test. Set ENV var RUN_EXCHANGE_INTEGRATION=true if you want to run them.")
    }
	_context, err := testutil.GetContext("integration.json")
	require.Nil(t, err, "Could not create context")
	stats, err := _context.NSQClient.GetStats()
	require.Nil(t, err)
	foundTopic := false
	for _, topic := range stats.Data.Topics {
		if topic.TopicName == _context.Config.DPN.DPNValidationWorker.NsqTopic {
			// Should have 4 items. One from each remote node.
			foundTopic = true
			assert.EqualValues(t, uint64(4), topic.MessageCount)
		}
	}
	assert.True(t, foundTopic, "Nothing was queued in %s",
		_context.Config.DPN.DPNValidationWorker.NsqTopic)
}
<|endoftext|>


<commit_before>// Copyright 2020 Google LLC
//
// 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 generate

import (
	"context"
	"fmt"
	"net/http"
	"strings"
	"time"

	"go.opencensus.io/trace"

	"github.com/google/exposure-notifications-server/internal/logging"
	"github.com/google/exposure-notifications-server/internal/publish/database"
	"github.com/google/exposure-notifications-server/internal/publish/model"
	"github.com/google/exposure-notifications-server/internal/serverenv"
	"github.com/google/exposure-notifications-server/internal/util"
	verifyapi "github.com/google/exposure-notifications-server/pkg/api/v1alpha1"
)

func NewHandler(ctx context.Context, config *Config, env *serverenv.ServerEnv) (http.Handler, error) {
	logger := logging.FromContext(ctx)

	if env.Database() == nil {
		return nil, fmt.Errorf("missing database in server environment")
	}

	transformer, err := model.NewTransformer(config.MaxKeysOnPublish, config.MaxIntervalAge, config.TruncateWindow, false)
	if err != nil {
		return nil, fmt.Errorf("model.NewTransformer: %w", err)
	}
	logger.Infof("max keys per upload: %v", config.MaxKeysOnPublish)
	logger.Infof("max interval start age: %v", config.MaxIntervalAge)
	logger.Infof("truncate window: %v", config.TruncateWindow)

	return &generateHandler{
		serverenv:   env,
		transformer: transformer,
		config:      config,
		database:    database.New(env.Database()),
	}, nil
}

type generateHandler struct {
	config      *Config
	serverenv   *serverenv.ServerEnv
	transformer *model.Transformer
	database    *database.PublishDB
}

func (h *generateHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	ctx, span := trace.StartSpan(r.Context(), "(*generate.generateHandler).ServeHTTP")
	defer span.End()
	logger := logging.FromContext(ctx)

	regionStr := h.config.DefaultRegion
	if regionParams, ok := r.URL.Query()["region"]; ok && len(regionParams) > 0 {
		regionStr = regionParams[0]
	}
	regions := strings.Split(regionStr, ",")
	logger.Infof("Rquest to generate data for regions: %v", regions)

	batchTime := time.Now()
	for i := 0; i < h.config.NumExposures; i++ {
		logger.Infof("Generating exposure %v of %v", i+1, h.config.NumExposures)

		tr, err := util.RandomTransmissionRisk()
		if err != nil {
			tr = 0
		}
		publish := verifyapi.Publish{
			Keys:           util.GenerateExposureKeys(h.config.KeysPerExposure, tr, false),
			Regions:        regions,
			AppPackageName: "generated.data",
		}

		exposures, err := h.transformer.TransformPublish(ctx, &publish, batchTime)
		if err != nil {
			message := fmt.Sprintf("Error transofmring generated exposures: %v", err)
			span.SetStatus(trace.Status{Code: trace.StatusCodeInternal, Message: message})
			w.WriteHeader(http.StatusInternalServerError)
			fmt.Fprint(w, message)
			return
		}

		err = h.database.InsertExposures(ctx, exposures)
		if err != nil {
			message := fmt.Sprintf("error writing exposure record: %v", err)
			span.SetStatus(trace.Status{Code: trace.StatusCodeInternal, Message: message})
			w.WriteHeader(http.StatusInternalServerError)
			fmt.Fprint(w, message)
			return
		}
		logger.Infof("Generated %v exposures", len(exposures))
	}
	w.WriteHeader(http.StatusOK)
	fmt.Fprint(w, "Generated exposure keys.")
}
<commit_msg>Update handler.go (#688)<commit_after>// Copyright 2020 Google LLC
//
// 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 generate

import (
	"context"
	"fmt"
	"net/http"
	"strings"
	"time"

	"go.opencensus.io/trace"

	"github.com/google/exposure-notifications-server/internal/logging"
	"github.com/google/exposure-notifications-server/internal/publish/database"
	"github.com/google/exposure-notifications-server/internal/publish/model"
	"github.com/google/exposure-notifications-server/internal/serverenv"
	"github.com/google/exposure-notifications-server/internal/util"
	verifyapi "github.com/google/exposure-notifications-server/pkg/api/v1alpha1"
)

func NewHandler(ctx context.Context, config *Config, env *serverenv.ServerEnv) (http.Handler, error) {
	logger := logging.FromContext(ctx)

	if env.Database() == nil {
		return nil, fmt.Errorf("missing database in server environment")
	}

	transformer, err := model.NewTransformer(config.MaxKeysOnPublish, config.MaxIntervalAge, config.TruncateWindow, false)
	if err != nil {
		return nil, fmt.Errorf("model.NewTransformer: %w", err)
	}
	logger.Infof("max keys per upload: %v", config.MaxKeysOnPublish)
	logger.Infof("max interval start age: %v", config.MaxIntervalAge)
	logger.Infof("truncate window: %v", config.TruncateWindow)

	return &generateHandler{
		serverenv:   env,
		transformer: transformer,
		config:      config,
		database:    database.New(env.Database()),
	}, nil
}

type generateHandler struct {
	config      *Config
	serverenv   *serverenv.ServerEnv
	transformer *model.Transformer
	database    *database.PublishDB
}

func (h *generateHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	ctx, span := trace.StartSpan(r.Context(), "(*generate.generateHandler).ServeHTTP")
	defer span.End()
	logger := logging.FromContext(ctx)

	regionStr := h.config.DefaultRegion
	if regionParams, ok := r.URL.Query()["region"]; ok && len(regionParams) > 0 {
		regionStr = regionParams[0]
	}
	regions := strings.Split(regionStr, ",")
	logger.Infof("Request to generate data for regions: %v", regions)

	batchTime := time.Now()
	for i := 0; i < h.config.NumExposures; i++ {
		logger.Infof("Generating exposure %v of %v", i+1, h.config.NumExposures)

		tr, err := util.RandomTransmissionRisk()
		if err != nil {
			tr = 0
		}
		publish := verifyapi.Publish{
			Keys:           util.GenerateExposureKeys(h.config.KeysPerExposure, tr, false),
			Regions:        regions,
			AppPackageName: "generated.data",
		}

		exposures, err := h.transformer.TransformPublish(ctx, &publish, batchTime)
		if err != nil {
			message := fmt.Sprintf("Error transforming generated exposures: %v", err)
			span.SetStatus(trace.Status{Code: trace.StatusCodeInternal, Message: message})
			w.WriteHeader(http.StatusInternalServerError)
			fmt.Fprint(w, message)
			return
		}

		err = h.database.InsertExposures(ctx, exposures)
		if err != nil {
			message := fmt.Sprintf("error writing exposure record: %v", err)
			span.SetStatus(trace.Status{Code: trace.StatusCodeInternal, Message: message})
			w.WriteHeader(http.StatusInternalServerError)
			fmt.Fprint(w, message)
			return
		}
		logger.Infof("Generated %v exposures", len(exposures))
	}
	w.WriteHeader(http.StatusOK)
	fmt.Fprint(w, "Generated exposure keys.")
}
<|endoftext|>


<commit_before>// Copyright 2019 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

//go:build ignore
// +build ignore

// mkstdlib generates the zstdlib.go file, containing the Go standard
// library API symbols. It's baked into the binary to avoid scanning
// GOPATH in the common case.
package main

import (
	"bufio"
	"bytes"
	"fmt"
	"go/format"
	"io"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"regexp"
	"runtime"
	"sort"

	exec "golang.org/x/sys/execabs"
)

func mustOpen(name string) io.Reader {
	f, err := os.Open(name)
	if err != nil {
		log.Fatal(err)
	}
	return f
}

func api(base string) string {
	return filepath.Join(runtime.GOROOT(), "api", base)
}

var sym = regexp.MustCompile(`^pkg (\S+).*?, (?:var|func|type|const) ([A-Z]\w*)`)

var unsafeSyms = map[string]bool{"Alignof": true, "ArbitraryType": true, "Offsetof": true, "Pointer": true, "Sizeof": true}

func main() {
	var buf bytes.Buffer
	outf := func(format string, args ...interface{}) {
		fmt.Fprintf(&buf, format, args...)
	}
	outf("// Code generated by mkstdlib.go. DO NOT EDIT.\n\n")
	outf("package imports\n")
	outf("var stdlib = map[string][]string{\n")
	f := io.MultiReader(
		mustOpen(api("go1.txt")),
		mustOpen(api("go1.1.txt")),
		mustOpen(api("go1.2.txt")),
		mustOpen(api("go1.3.txt")),
		mustOpen(api("go1.4.txt")),
		mustOpen(api("go1.5.txt")),
		mustOpen(api("go1.6.txt")),
		mustOpen(api("go1.7.txt")),
		mustOpen(api("go1.8.txt")),
		mustOpen(api("go1.9.txt")),
		mustOpen(api("go1.10.txt")),
		mustOpen(api("go1.11.txt")),
		mustOpen(api("go1.12.txt")),
		mustOpen(api("go1.13.txt")),
		mustOpen(api("go1.14.txt")),
		mustOpen(api("go1.15.txt")),
		mustOpen(api("go1.16.txt")),
		mustOpen(api("go1.17.txt")),
		mustOpen(api("go1.18.txt")),
		mustOpen(api("go1.19.txt")),

		// The API of the syscall/js package needs to be computed explicitly,
		// because it's not included in the GOROOT/api/go1.*.txt files at this time.
		syscallJSAPI(),
	)
	sc := bufio.NewScanner(f)

	pkgs := map[string]map[string]bool{
		"unsafe": unsafeSyms,
	}
	paths := []string{"unsafe"}

	for sc.Scan() {
		l := sc.Text()
		if m := sym.FindStringSubmatch(l); m != nil {
			path, sym := m[1], m[2]

			if _, ok := pkgs[path]; !ok {
				pkgs[path] = map[string]bool{}
				paths = append(paths, path)
			}
			pkgs[path][sym] = true
		}
	}
	if err := sc.Err(); err != nil {
		log.Fatal(err)
	}
	sort.Strings(paths)
	for _, path := range paths {
		outf("\t%q: []string{\n", path)
		pkg := pkgs[path]
		var syms []string
		for sym := range pkg {
			syms = append(syms, sym)
		}
		sort.Strings(syms)
		for _, sym := range syms {
			outf("\t\t%q,\n", sym)
		}
		outf("},\n")
	}
	outf("}\n")
	fmtbuf, err := format.Source(buf.Bytes())
	if err != nil {
		log.Fatal(err)
	}
	err = ioutil.WriteFile("zstdlib.go", fmtbuf, 0666)
	if err != nil {
		log.Fatal(err)
	}
}

// syscallJSAPI returns the API of the syscall/js package.
// It's computed from the contents of $(go env GOROOT)/src/syscall/js.
func syscallJSAPI() io.Reader {
	var exeSuffix string
	if runtime.GOOS == "windows" {
		exeSuffix = ".exe"
	}
	cmd := exec.Command("go"+exeSuffix, "run", "cmd/api", "-contexts", "js-wasm", "syscall/js")
	out, err := cmd.Output()
	if err != nil {
		log.Fatalln(err)
	}
	return bytes.NewReader(out)
}
<commit_msg>internal/imports: read entire API dir in mkstdlib.go<commit_after>// Copyright 2019 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

//go:build ignore
// +build ignore

// mkstdlib generates the zstdlib.go file, containing the Go standard
// library API symbols. It's baked into the binary to avoid scanning
// GOPATH in the common case.
package main

import (
	"bufio"
	"bytes"
	"fmt"
	"go/format"
	"io"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"regexp"
	"runtime"
	"sort"
	"strings"

	exec "golang.org/x/sys/execabs"
)

func mustOpen(name string) io.Reader {
	f, err := os.Open(name)
	if err != nil {
		log.Fatal(err)
	}
	return f
}

func api(base string) string {
	return filepath.Join(runtime.GOROOT(), "api", base)
}

var sym = regexp.MustCompile(`^pkg (\S+).*?, (?:var|func|type|const) ([A-Z]\w*)`)

var unsafeSyms = map[string]bool{"Alignof": true, "ArbitraryType": true, "Offsetof": true, "Pointer": true, "Sizeof": true}

func main() {
	var buf bytes.Buffer
	outf := func(format string, args ...interface{}) {
		fmt.Fprintf(&buf, format, args...)
	}
	outf("// Code generated by mkstdlib.go. DO NOT EDIT.\n\n")
	outf("package imports\n")
	outf("var stdlib = map[string][]string{\n")
	f := readAPI()
	sc := bufio.NewScanner(f)

	pkgs := map[string]map[string]bool{
		"unsafe": unsafeSyms,
	}
	paths := []string{"unsafe"}

	for sc.Scan() {
		l := sc.Text()
		if m := sym.FindStringSubmatch(l); m != nil {
			path, sym := m[1], m[2]

			if _, ok := pkgs[path]; !ok {
				pkgs[path] = map[string]bool{}
				paths = append(paths, path)
			}
			pkgs[path][sym] = true
		}
	}
	if err := sc.Err(); err != nil {
		log.Fatal(err)
	}
	sort.Strings(paths)
	for _, path := range paths {
		outf("\t%q: []string{\n", path)
		pkg := pkgs[path]
		var syms []string
		for sym := range pkg {
			syms = append(syms, sym)
		}
		sort.Strings(syms)
		for _, sym := range syms {
			outf("\t\t%q,\n", sym)
		}
		outf("},\n")
	}
	outf("}\n")
	fmtbuf, err := format.Source(buf.Bytes())
	if err != nil {
		log.Fatal(err)
	}
	err = ioutil.WriteFile("zstdlib.go", fmtbuf, 0666)
	if err != nil {
		log.Fatal(err)
	}
}

// readAPI opens an io.Reader that reads all stdlib API content.
func readAPI() io.Reader {
	entries, err := os.ReadDir(filepath.Join(runtime.GOROOT(), "api"))
	if err != nil {
		log.Fatal(err)
	}
	var readers []io.Reader
	for _, entry := range entries {
		name := entry.Name()
		if strings.HasPrefix(name, "go") && strings.HasSuffix(name, ".txt") {
			readers = append(readers, mustOpen(api(name)))
		}
	}
	// The API of the syscall/js package needs to be computed explicitly,
	// because it's not included in the GOROOT/api/go1.*.txt files at this time.
	readers = append(readers, syscallJSAPI())
	return io.MultiReader(readers...)
}

// syscallJSAPI returns the API of the syscall/js package.
// It's computed from the contents of $(go env GOROOT)/src/syscall/js.
func syscallJSAPI() io.Reader {
	var exeSuffix string
	if runtime.GOOS == "windows" {
		exeSuffix = ".exe"
	}
	cmd := exec.Command("go"+exeSuffix, "run", "cmd/api", "-contexts", "js-wasm", "syscall/js")
	out, err := cmd.Output()
	if err != nil {
		log.Fatalln(err)
	}
	return bytes.NewReader(out)
}
<|endoftext|>