<commit_before>package torrent

import (
	"context"
	"errors"
	"io"
	"sync"

	"github.com/anacrolix/log"
	"github.com/anacrolix/missinggo"
)

type Reader interface {
	io.Reader
	io.Seeker
	io.Closer
	missinggo.ReadContexter
	SetReadahead(int64)
	SetResponsive()
}

// Piece range by piece index, [begin, end).
type pieceRange struct {
	begin, end pieceIndex
}

// Accesses Torrent data via a Client. Reads block until the data is
// available. Seeks and readahead also drive Client behaviour.
type reader struct {
	t          *Torrent
	responsive bool
	// Adjust the read/seek window to handle Readers locked to File extents
	// and the like.
	offset, length int64
	// Ensure operations that change the position are exclusive, like Read()
	// and Seek().
	opMu sync.Mutex

	// Required when modifying pos and readahead, or reading them without
	// opMu.
	mu        sync.Locker
	pos       int64
	readahead int64
	// The cached piece range this reader wants downloaded. The zero value
	// corresponds to nothing. We cache this so that changes can be detected,
	// and bubbled up to the Torrent only as required.
	pieces pieceRange
}

var _ io.ReadCloser = &reader{}

// Don't wait for pieces to complete and be verified. Read calls return as
// soon as they can when the underlying chunks become available.
func (r *reader) SetResponsive() {
	r.responsive = true
	r.t.cl.event.Broadcast()
}

// Disable responsive mode. TODO: Remove?
func (r *reader) SetNonResponsive() {
	r.responsive = false
	r.t.cl.event.Broadcast()
}

// Configure the number of bytes ahead of a read that should also be
// prioritized in preparation for further reads.
func (r *reader) SetReadahead(readahead int64) {
	r.mu.Lock()
	r.readahead = readahead
	r.mu.Unlock()
	r.t.cl.lock()
	defer r.t.cl.unlock()
	r.posChanged()
}

// How many bytes are available to read. Max is the most we could require.
func (r *reader) available(off, max int64) (ret int64) {
	off += r.offset
	for max > 0 {
		req, ok := r.t.offsetRequest(off)
		if !ok {
			break
		}
		if !r.responsive && !r.t.pieceComplete(pieceIndex(req.Index)) {
			break
		}
		if !r.t.haveChunk(req) {
			break
		}
		len1 := int64(req.Length) - (off - r.t.requestOffset(req))
		max -= len1
		ret += len1
		off += len1
	}
	// Ensure that ret hasn't exceeded our original max.
	if max < 0 {
		ret += max
	}
	return
}

func (r *reader) waitReadable(off int64) {
	// We may have been sent back here because we were told we could read but
	// it failed.
	r.t.cl.event.Wait()
}

// Calculates the pieces this reader wants downloaded, ignoring the cached
// value at r.pieces.
func (r *reader) piecesUncached() (ret pieceRange) {
	ra := r.readahead
	if ra < 1 {
		// Needs to be at least 1, because [x, x) means we don't want
		// anything.
		ra = 1
	}
	if ra > r.length-r.pos {
		ra = r.length - r.pos
	}
	ret.begin, ret.end = r.t.byteRegionPieces(r.torrentOffset(r.pos), ra)
	return
}

func (r *reader) Read(b []byte) (n int, err error) {
	return r.ReadContext(context.Background(), b)
}

func (r *reader) ReadContext(ctx context.Context, b []byte) (n int, err error) {
	// This is set under the Client lock if the Context is canceled. I think we coordinate on a
	// separate variable so as to avoid false negatives with race conditions due to Contexts being
	// synchronized.
	var ctxErr error
	if ctx.Done() != nil {
		ctx, cancel := context.WithCancel(ctx)
		// Abort the goroutine when the function returns.
		defer cancel()
		go func() {
			<-ctx.Done()
			r.t.cl.lock()
			ctxErr = ctx.Err()
			r.t.tickleReaders()
			r.t.cl.unlock()
		}()
	}
	// Hmmm, if a Read gets stuck, this means you can't change position for
	// other purposes. That seems reasonable, but unusual.
	r.opMu.Lock()
	defer r.opMu.Unlock()
	n, err = r.readOnceAt(b, r.pos, &ctxErr)
	if n == 0 {
		if err == nil {
			panic("expected error")
		} else {
			return
		}
	}

	r.mu.Lock()
	r.pos += int64(n)
	r.posChanged()
	r.mu.Unlock()
	if r.pos >= r.length {
		err = io.EOF
	} else if err == io.EOF {
		err = io.ErrUnexpectedEOF
	}
	return
}

// Wait until some data should be available to read. Tickles the client if it
// isn't. Returns how much should be readable without blocking.
func (r *reader) waitAvailable(pos, wanted int64, ctxErr *error, wait bool) (avail int64, err error) {
	r.t.cl.lock()
	defer r.t.cl.unlock()
	for {
		avail = r.available(pos, wanted)
		if avail != 0 {
			return
		}
		if r.t.closed.IsSet() {
			err = errors.New("torrent closed")
			return
		}
		if *ctxErr != nil {
			err = *ctxErr
			return
		}
		if r.t.dataDownloadDisallowed || !r.t.networkingEnabled {
			err = errors.New("downloading disabled and data not already available")
			return
		}
		if !wait {
			return
		}
		r.waitReadable(pos)
	}
}

// Adds the reader's torrent offset to the reader object offset (for example the reader might be
// constrainted to a particular file within the torrent).
func (r *reader) torrentOffset(readerPos int64) int64 {
	return r.offset + readerPos
}

// Performs at most one successful read to torrent storage.
func (r *reader) readOnceAt(b []byte, pos int64, ctxErr *error) (n int, err error) {
	if pos >= r.length {
		err = io.EOF
		return
	}
	for {
		var avail int64
		avail, err = r.waitAvailable(pos, int64(len(b)), ctxErr, n == 0)
		if avail == 0 {
			return
		}
		firstPieceIndex := pieceIndex(r.torrentOffset(pos) / r.t.info.PieceLength)
		firstPieceOffset := r.torrentOffset(pos) % r.t.info.PieceLength
		b1 := missinggo.LimitLen(b, avail)
		n, err = r.t.readAt(b1, r.torrentOffset(pos))
		if n != 0 {
			err = nil
			return
		}
		r.t.cl.lock()
		// TODO: Just reset pieces in the readahead window. This might help
		// prevent thrashing with small caches and file and piece priorities.
		r.log(log.Fstr("error reading torrent %s piece %d offset %d, %d bytes: %v",
			r.t.infoHash.HexString(), firstPieceIndex, firstPieceOffset, len(b1), err))
		if !r.t.updatePieceCompletion(firstPieceIndex) {
			r.log(log.Fstr("piece %d completion unchanged", firstPieceIndex))
		}
		r.t.cl.unlock()
	}
}

func (r *reader) Close() error {
	r.t.cl.lock()
	defer r.t.cl.unlock()
	r.t.deleteReader(r)
	return nil
}

func (r *reader) posChanged() {
	to := r.piecesUncached()
	from := r.pieces
	if to == from {
		return
	}
	r.pieces = to
	// log.Printf("reader pos changed %v->%v", from, to)
	r.t.readerPosChanged(from, to)
}

func (r *reader) Seek(off int64, whence int) (ret int64, err error) {
	r.opMu.Lock()
	defer r.opMu.Unlock()

	r.mu.Lock()
	defer r.mu.Unlock()
	switch whence {
	case io.SeekStart:
		r.pos = off
	case io.SeekCurrent:
		r.pos += off
	case io.SeekEnd:
		r.pos = r.length + off
	default:
		err = errors.New("bad whence")
	}
	ret = r.pos

	r.posChanged()
	return
}

func (r *reader) log(m log.Msg) {
	r.t.logger.Log(m.Skip(1))
}
<commit_msg>Update the readahead window on read failure<commit_after>package torrent

import (
	"context"
	"errors"
	"fmt"
	"io"
	"sync"

	"github.com/anacrolix/log"
	"github.com/anacrolix/missinggo"
)

type Reader interface {
	io.Reader
	io.Seeker
	io.Closer
	missinggo.ReadContexter
	SetReadahead(int64)
	SetResponsive()
}

// Piece range by piece index, [begin, end).
type pieceRange struct {
	begin, end pieceIndex
}

// Accesses Torrent data via a Client. Reads block until the data is
// available. Seeks and readahead also drive Client behaviour.
type reader struct {
	t          *Torrent
	responsive bool
	// Adjust the read/seek window to handle Readers locked to File extents
	// and the like.
	offset, length int64
	// Ensure operations that change the position are exclusive, like Read()
	// and Seek().
	opMu sync.Mutex

	// Required when modifying pos and readahead, or reading them without
	// opMu.
	mu        sync.Locker
	pos       int64
	readahead int64
	// The cached piece range this reader wants downloaded. The zero value
	// corresponds to nothing. We cache this so that changes can be detected,
	// and bubbled up to the Torrent only as required.
	pieces pieceRange
}

var _ io.ReadCloser = &reader{}

// Don't wait for pieces to complete and be verified. Read calls return as
// soon as they can when the underlying chunks become available.
func (r *reader) SetResponsive() {
	r.responsive = true
	r.t.cl.event.Broadcast()
}

// Disable responsive mode. TODO: Remove?
func (r *reader) SetNonResponsive() {
	r.responsive = false
	r.t.cl.event.Broadcast()
}

// Configure the number of bytes ahead of a read that should also be
// prioritized in preparation for further reads.
func (r *reader) SetReadahead(readahead int64) {
	r.mu.Lock()
	r.readahead = readahead
	r.mu.Unlock()
	r.t.cl.lock()
	defer r.t.cl.unlock()
	r.posChanged()
}

// How many bytes are available to read. Max is the most we could require.
func (r *reader) available(off, max int64) (ret int64) {
	off += r.offset
	for max > 0 {
		req, ok := r.t.offsetRequest(off)
		if !ok {
			break
		}
		if !r.responsive && !r.t.pieceComplete(pieceIndex(req.Index)) {
			break
		}
		if !r.t.haveChunk(req) {
			break
		}
		len1 := int64(req.Length) - (off - r.t.requestOffset(req))
		max -= len1
		ret += len1
		off += len1
	}
	// Ensure that ret hasn't exceeded our original max.
	if max < 0 {
		ret += max
	}
	return
}

func (r *reader) waitReadable(off int64) {
	// We may have been sent back here because we were told we could read but
	// it failed.
	r.t.cl.event.Wait()
}

// Calculates the pieces this reader wants downloaded, ignoring the cached
// value at r.pieces.
func (r *reader) piecesUncached() (ret pieceRange) {
	ra := r.readahead
	if ra < 1 {
		// Needs to be at least 1, because [x, x) means we don't want
		// anything.
		ra = 1
	}
	if ra > r.length-r.pos {
		ra = r.length - r.pos
	}
	ret.begin, ret.end = r.t.byteRegionPieces(r.torrentOffset(r.pos), ra)
	return
}

func (r *reader) Read(b []byte) (n int, err error) {
	return r.ReadContext(context.Background(), b)
}

func (r *reader) ReadContext(ctx context.Context, b []byte) (n int, err error) {
	// This is set under the Client lock if the Context is canceled. I think we coordinate on a
	// separate variable so as to avoid false negatives with race conditions due to Contexts being
	// synchronized.
	var ctxErr error
	if ctx.Done() != nil {
		ctx, cancel := context.WithCancel(ctx)
		// Abort the goroutine when the function returns.
		defer cancel()
		go func() {
			<-ctx.Done()
			r.t.cl.lock()
			ctxErr = ctx.Err()
			r.t.tickleReaders()
			r.t.cl.unlock()
		}()
	}
	// Hmmm, if a Read gets stuck, this means you can't change position for
	// other purposes. That seems reasonable, but unusual.
	r.opMu.Lock()
	defer r.opMu.Unlock()
	n, err = r.readOnceAt(b, r.pos, &ctxErr)
	if n == 0 {
		if err == nil {
			panic("expected error")
		} else {
			return
		}
	}

	r.mu.Lock()
	r.pos += int64(n)
	r.posChanged()
	r.mu.Unlock()
	if r.pos >= r.length {
		err = io.EOF
	} else if err == io.EOF {
		err = io.ErrUnexpectedEOF
	}
	return
}

// Wait until some data should be available to read. Tickles the client if it
// isn't. Returns how much should be readable without blocking.
func (r *reader) waitAvailable(pos, wanted int64, ctxErr *error, wait bool) (avail int64, err error) {
	r.t.cl.lock()
	defer r.t.cl.unlock()
	for {
		avail = r.available(pos, wanted)
		if avail != 0 {
			return
		}
		if r.t.closed.IsSet() {
			err = errors.New("torrent closed")
			return
		}
		if *ctxErr != nil {
			err = *ctxErr
			return
		}
		if r.t.dataDownloadDisallowed || !r.t.networkingEnabled {
			err = errors.New("downloading disabled and data not already available")
			return
		}
		if !wait {
			return
		}
		r.waitReadable(pos)
	}
}

// Adds the reader's torrent offset to the reader object offset (for example the reader might be
// constrainted to a particular file within the torrent).
func (r *reader) torrentOffset(readerPos int64) int64 {
	return r.offset + readerPos
}

// Performs at most one successful read to torrent storage.
func (r *reader) readOnceAt(b []byte, pos int64, ctxErr *error) (n int, err error) {
	if pos >= r.length {
		err = io.EOF
		return
	}
	for {
		var avail int64
		avail, err = r.waitAvailable(pos, int64(len(b)), ctxErr, n == 0)
		if avail == 0 {
			return
		}
		firstPieceIndex := pieceIndex(r.torrentOffset(pos) / r.t.info.PieceLength)
		firstPieceOffset := r.torrentOffset(pos) % r.t.info.PieceLength
		b1 := missinggo.LimitLen(b, avail)
		n, err = r.t.readAt(b1, r.torrentOffset(pos))
		if n != 0 {
			err = nil
			return
		}
		r.t.cl.lock()
		// TODO: Just reset pieces in the readahead window. This might help
		// prevent thrashing with small caches and file and piece priorities.
		r.log(log.Fstr("error reading torrent %s piece %d offset %d, %d bytes: %v",
			r.t.infoHash.HexString(), firstPieceIndex, firstPieceOffset, len(b1), err))
		if !r.t.updatePieceCompletion(firstPieceIndex) {
			r.log(log.Fstr("piece %d completion unchanged", firstPieceIndex))
		}
		// Update the rest of the piece completions in the readahead window, without alerting to
		// changes (since only the first piece, the one above, could have generated the read error
		// we're currently handling).
		if r.pieces.begin != firstPieceIndex {
			panic(fmt.Sprint(r.pieces.begin, firstPieceIndex))
		}
		for index := r.pieces.begin + 1; index < r.pieces.end; index++ {
			r.t.updatePieceCompletion(index)
		}
		r.t.cl.unlock()
	}
}

func (r *reader) Close() error {
	r.t.cl.lock()
	defer r.t.cl.unlock()
	r.t.deleteReader(r)
	return nil
}

func (r *reader) posChanged() {
	to := r.piecesUncached()
	from := r.pieces
	if to == from {
		return
	}
	r.pieces = to
	// log.Printf("reader pos changed %v->%v", from, to)
	r.t.readerPosChanged(from, to)
}

func (r *reader) Seek(off int64, whence int) (ret int64, err error) {
	r.opMu.Lock()
	defer r.opMu.Unlock()

	r.mu.Lock()
	defer r.mu.Unlock()
	switch whence {
	case io.SeekStart:
		r.pos = off
	case io.SeekCurrent:
		r.pos += off
	case io.SeekEnd:
		r.pos = r.length + off
	default:
		err = errors.New("bad whence")
	}
	ret = r.pos

	r.posChanged()
	return
}

func (r *reader) log(m log.Msg) {
	r.t.logger.Log(m.Skip(1))
}
<|endoftext|>


<commit_before>// Copyright (C) 2014 Daniel Harrison

package hfile

import (
	"bytes"
	"encoding/binary"
	"encoding/hex"
	"errors"
	"fmt"
	"io"
	"log"
	"os"
	"sort"

	"github.com/edsrzf/mmap-go"
	"github.com/golang/snappy"
)

type Reader struct {
	CollectionConfig

	mmap mmap.MMap

	majorVersion uint32
	minorVersion uint32

	Header
	index []Block

	scannerCache  chan *Scanner
	iteratorCache chan *Iterator
}

type Header struct {
	offset int

	fileInfoOffset             uint64
	dataIndexOffset            uint64
	dataIndexCount             uint32
	metaIndexOffset            uint64
	metaIndexCount             uint32
	totalUncompressedDataBytes uint64
	EntryCount                 uint32
	compressionCodec           uint32
}

type Block struct {
	offset        uint64
	size          uint32
	firstKeyBytes []byte
}

func NewReader(name, path string, lock, debug bool) (*Reader, error) {
	return NewReaderFromConfig(CollectionConfig{name, path, path, lock, debug})
}

func NewReaderFromConfig(cfg CollectionConfig) (*Reader, error) {
	f, err := os.OpenFile(cfg.LocalPath, os.O_RDONLY, 0)

	if err != nil {
		return nil, fmt.Errorf("Error opening file (%s): %v", cfg.LocalPath, err)
	}

	hfile := new(Reader)
	hfile.CollectionConfig = cfg

	hfile.mmap, err = mmap.Map(f, mmap.RDONLY, 0)

	if err != nil {
		return nil, err
	}

	fi, err := f.Stat()
	if err != nil {
		return nil, err
	}

	if hfile.InMem {
		mb := 1024.0 * 1024.0
		log.Printf("[Reader.NewReader] locking %s (%.02fmb)...\n", hfile.Name, float64(fi.Size())/mb)
		if err = hfile.mmap.Lock(); err != nil {
			log.Printf("[Reader.NewReader] error locking %s: %s\n", hfile.Name, err.Error())
			return nil, err
		}
		log.Printf("[Reader.NewReader] locked %s.\n", hfile.Name)
	} else if hfile.Debug {
		log.Printf("[Reader.NewReader] Not locking %s...\n", hfile.Name)
	}

	v := binary.BigEndian.Uint32(hfile.mmap[len(hfile.mmap)-4:])
	hfile.majorVersion = v & 0x00ffffff
	hfile.minorVersion = v >> 24

	err = hfile.readHeader(hfile.mmap)
	if err != nil {
		return nil, err
	}

	err = hfile.loadIndex(hfile.mmap)
	if err != nil {
		return hfile, err
	}
	hfile.scannerCache = make(chan *Scanner, 5)
	hfile.iteratorCache = make(chan *Iterator, 5)
	return hfile, nil
}

func (r *Reader) PrintDebugInfo(out io.Writer, includeStartKeys int) {
	fmt.Fprintln(out, "entries: ", r.EntryCount)
	fmt.Fprintf(out, "compressed: %v (codec: %d)\n", r.compressionCodec != CompressionNone, r.compressionCodec)
	fmt.Fprintln(out, "blocks: ", len(r.index))
	for i, blk := range r.index {
		if i > includeStartKeys {
			fmt.Fprintf(out, "\t... and %d more\n", len(r.index)-i)
			return
		}
		fmt.Fprintf(out, "\t#%d: %s\n", i, hex.EncodeToString(blk.firstKeyBytes))
	}
}

func (r *Reader) readHeader(mmap mmap.MMap) error {
	if r.majorVersion != 1 || r.minorVersion != 0 {
		return fmt.Errorf("wrong version: %d.%d", r.majorVersion, r.minorVersion)
	}

	r.Header.offset = len(mmap) - 60
	buf := bytes.NewReader(mmap[r.Header.offset:])

	headerMagic := make([]byte, 8)
	buf.Read(headerMagic)
	if bytes.Compare(headerMagic, TrailerMagic) != 0 {
		return errors.New("bad header magic")
	}

	binary.Read(buf, binary.BigEndian, &r.fileInfoOffset)
	binary.Read(buf, binary.BigEndian, &r.dataIndexOffset)
	binary.Read(buf, binary.BigEndian, &r.dataIndexCount)
	binary.Read(buf, binary.BigEndian, &r.metaIndexOffset)
	binary.Read(buf, binary.BigEndian, &r.metaIndexCount)
	binary.Read(buf, binary.BigEndian, &r.totalUncompressedDataBytes)
	binary.Read(buf, binary.BigEndian, &r.EntryCount)
	binary.Read(buf, binary.BigEndian, &r.compressionCodec)
	return nil
}

func (r *Reader) loadIndex(mmap mmap.MMap) error {

	dataIndexEnd := r.metaIndexOffset
	if r.metaIndexOffset == 0 {
		dataIndexEnd = uint64(r.Header.offset)
	}
	buf := bytes.NewReader(mmap[r.dataIndexOffset:dataIndexEnd])

	dataIndexMagic := make([]byte, 8)
	buf.Read(dataIndexMagic)
	if bytes.Compare(dataIndexMagic, IndexMagic) != 0 {
		return errors.New("bad data index magic")
	}

	for buf.Len() > 0 {
		dataBlock := Block{}

		binary.Read(buf, binary.BigEndian, &dataBlock.offset)
		binary.Read(buf, binary.BigEndian, &dataBlock.size)

		firstKeyLen, _ := binary.ReadUvarint(buf)
		dataBlock.firstKeyBytes = make([]byte, firstKeyLen)
		buf.Read(dataBlock.firstKeyBytes)

		r.index = append(r.index, dataBlock)
	}

	return nil
}

func After(a, b []byte) bool {
	return bytes.Compare(a, b) > 0
}

func (b *Block) IsAfter(key []byte) bool {
	return After(b.firstKeyBytes, key)
}

func (r *Reader) FindBlock(from int, key []byte) int {
	remaining := len(r.index) - from - 1
	if r.Debug {
		log.Printf("[Reader.findBlock] cur %d, remaining %d\n", from, remaining)
	}

	if remaining <= 0 {
		if r.Debug {
			log.Println("[Reader.findBlock] last block")
		}
		return from // s.cur is the last block, so it is only choice.
	}

	if r.index[from+1].IsAfter(key) {
		if r.Debug {
			log.Println("[Reader.findBlock] next block is past key")
		}
		return from
	}

	offset := sort.Search(remaining, func(i int) bool {
		return r.index[from+i+1].IsAfter(key)
	})

	return from + offset
}

func (r *Reader) GetBlockBuf(i int, dst []byte) ([]byte, error) {
	var err error

	block := r.index[i]

	switch r.compressionCodec {
	case CompressionNone:
		dst = r.mmap[block.offset : block.offset+uint64(block.size)]
	case CompressionSnappy:
		uncompressedByteSize := binary.BigEndian.Uint32(r.mmap[block.offset : block.offset+4])
		if uncompressedByteSize != block.size {
			return nil, errors.New("mismatched uncompressed block size")
		}
		compressedByteSize := binary.BigEndian.Uint32(r.mmap[block.offset+4 : block.offset+8])
		compressedBytes := r.mmap[block.offset+8 : block.offset+8+uint64(compressedByteSize)]
		dst, err = snappy.Decode(dst, compressedBytes)
		if err != nil {
			return nil, err
		}
	default:
		return nil, errors.New("Unsupported compression codec " + string(r.compressionCodec))
	}

	if bytes.Compare(dst[0:8], DataMagic) != 0 {
		return nil, errors.New("bad data block magic")
	}

	return dst, nil
}

func (r *Reader) GetScanner() *Scanner {
	select {
	case s := <-r.scannerCache:
		return s
	default:
		return NewScanner(r)
	}
}

func (r *Reader) GetIterator() *Iterator {
	select {
	case i := <-r.iteratorCache:
		return i
	default:
		return NewIterator(r)
	}
}
<commit_msg>add FirstKey helper<commit_after>// Copyright (C) 2014 Daniel Harrison

package hfile

import (
	"bytes"
	"encoding/binary"
	"encoding/hex"
	"errors"
	"fmt"
	"io"
	"log"
	"os"
	"sort"

	"github.com/edsrzf/mmap-go"
	"github.com/golang/snappy"
)

type Reader struct {
	CollectionConfig

	mmap mmap.MMap

	majorVersion uint32
	minorVersion uint32

	Header
	index []Block

	scannerCache  chan *Scanner
	iteratorCache chan *Iterator
}

type Header struct {
	offset int

	fileInfoOffset             uint64
	dataIndexOffset            uint64
	dataIndexCount             uint32
	metaIndexOffset            uint64
	metaIndexCount             uint32
	totalUncompressedDataBytes uint64
	EntryCount                 uint32
	compressionCodec           uint32
}

type Block struct {
	offset        uint64
	size          uint32
	firstKeyBytes []byte
}

func NewReader(name, path string, lock, debug bool) (*Reader, error) {
	return NewReaderFromConfig(CollectionConfig{name, path, path, lock, debug})
}

func NewReaderFromConfig(cfg CollectionConfig) (*Reader, error) {
	f, err := os.OpenFile(cfg.LocalPath, os.O_RDONLY, 0)

	if err != nil {
		return nil, fmt.Errorf("Error opening file (%s): %v", cfg.LocalPath, err)
	}

	hfile := new(Reader)
	hfile.CollectionConfig = cfg

	hfile.mmap, err = mmap.Map(f, mmap.RDONLY, 0)

	if err != nil {
		return nil, err
	}

	fi, err := f.Stat()
	if err != nil {
		return nil, err
	}

	if hfile.InMem {
		mb := 1024.0 * 1024.0
		log.Printf("[Reader.NewReader] locking %s (%.02fmb)...\n", hfile.Name, float64(fi.Size())/mb)
		if err = hfile.mmap.Lock(); err != nil {
			log.Printf("[Reader.NewReader] error locking %s: %s\n", hfile.Name, err.Error())
			return nil, err
		}
		log.Printf("[Reader.NewReader] locked %s.\n", hfile.Name)
	} else if hfile.Debug {
		log.Printf("[Reader.NewReader] Not locking %s...\n", hfile.Name)
	}

	v := binary.BigEndian.Uint32(hfile.mmap[len(hfile.mmap)-4:])
	hfile.majorVersion = v & 0x00ffffff
	hfile.minorVersion = v >> 24

	err = hfile.readHeader(hfile.mmap)
	if err != nil {
		return nil, err
	}

	err = hfile.loadIndex(hfile.mmap)
	if err != nil {
		return hfile, err
	}
	hfile.scannerCache = make(chan *Scanner, 5)
	hfile.iteratorCache = make(chan *Iterator, 5)
	return hfile, nil
}

func (r *Reader) PrintDebugInfo(out io.Writer, includeStartKeys int) {
	fmt.Fprintln(out, "entries: ", r.EntryCount)
	fmt.Fprintf(out, "compressed: %v (codec: %d)\n", r.compressionCodec != CompressionNone, r.compressionCodec)
	fmt.Fprintln(out, "blocks: ", len(r.index))
	for i, blk := range r.index {
		if i > includeStartKeys {
			fmt.Fprintf(out, "\t... and %d more\n", len(r.index)-i)
			return
		}
		fmt.Fprintf(out, "\t#%d: %s\n", i, hex.EncodeToString(blk.firstKeyBytes))
	}
}

func (r *Reader) readHeader(mmap mmap.MMap) error {
	if r.majorVersion != 1 || r.minorVersion != 0 {
		return fmt.Errorf("wrong version: %d.%d", r.majorVersion, r.minorVersion)
	}

	r.Header.offset = len(mmap) - 60
	buf := bytes.NewReader(mmap[r.Header.offset:])

	headerMagic := make([]byte, 8)
	buf.Read(headerMagic)
	if bytes.Compare(headerMagic, TrailerMagic) != 0 {
		return errors.New("bad header magic")
	}

	binary.Read(buf, binary.BigEndian, &r.fileInfoOffset)
	binary.Read(buf, binary.BigEndian, &r.dataIndexOffset)
	binary.Read(buf, binary.BigEndian, &r.dataIndexCount)
	binary.Read(buf, binary.BigEndian, &r.metaIndexOffset)
	binary.Read(buf, binary.BigEndian, &r.metaIndexCount)
	binary.Read(buf, binary.BigEndian, &r.totalUncompressedDataBytes)
	binary.Read(buf, binary.BigEndian, &r.EntryCount)
	binary.Read(buf, binary.BigEndian, &r.compressionCodec)
	return nil
}

func (r *Reader) loadIndex(mmap mmap.MMap) error {

	dataIndexEnd := r.metaIndexOffset
	if r.metaIndexOffset == 0 {
		dataIndexEnd = uint64(r.Header.offset)
	}
	buf := bytes.NewReader(mmap[r.dataIndexOffset:dataIndexEnd])

	dataIndexMagic := make([]byte, 8)
	buf.Read(dataIndexMagic)
	if bytes.Compare(dataIndexMagic, IndexMagic) != 0 {
		return errors.New("bad data index magic")
	}

	for buf.Len() > 0 {
		dataBlock := Block{}

		binary.Read(buf, binary.BigEndian, &dataBlock.offset)
		binary.Read(buf, binary.BigEndian, &dataBlock.size)

		firstKeyLen, _ := binary.ReadUvarint(buf)
		dataBlock.firstKeyBytes = make([]byte, firstKeyLen)
		buf.Read(dataBlock.firstKeyBytes)

		r.index = append(r.index, dataBlock)
	}

	return nil
}

func After(a, b []byte) bool {
	return bytes.Compare(a, b) > 0
}

func (b *Block) IsAfter(key []byte) bool {
	return After(b.firstKeyBytes, key)
}

func (r *Reader) FirstKey() ([]byte, error) {
	if len(r.index) < 1 {
		return nil, fmt.Errorf("empty collection has no first key")
	}
	return r.index[0].firstKeyBytes, nil
}

func (r *Reader) FindBlock(from int, key []byte) int {
	remaining := len(r.index) - from - 1
	if r.Debug {
		log.Printf("[Reader.findBlock] cur %d, remaining %d\n", from, remaining)
	}

	if remaining <= 0 {
		if r.Debug {
			log.Println("[Reader.findBlock] last block")
		}
		return from // s.cur is the last block, so it is only choice.
	}

	if r.index[from+1].IsAfter(key) {
		if r.Debug {
			log.Println("[Reader.findBlock] next block is past key")
		}
		return from
	}

	offset := sort.Search(remaining, func(i int) bool {
		return r.index[from+i+1].IsAfter(key)
	})

	return from + offset
}

func (r *Reader) GetBlockBuf(i int, dst []byte) ([]byte, error) {
	var err error

	block := r.index[i]

	switch r.compressionCodec {
	case CompressionNone:
		dst = r.mmap[block.offset : block.offset+uint64(block.size)]
	case CompressionSnappy:
		uncompressedByteSize := binary.BigEndian.Uint32(r.mmap[block.offset : block.offset+4])
		if uncompressedByteSize != block.size {
			return nil, errors.New("mismatched uncompressed block size")
		}
		compressedByteSize := binary.BigEndian.Uint32(r.mmap[block.offset+4 : block.offset+8])
		compressedBytes := r.mmap[block.offset+8 : block.offset+8+uint64(compressedByteSize)]
		dst, err = snappy.Decode(dst, compressedBytes)
		if err != nil {
			return nil, err
		}
	default:
		return nil, errors.New("Unsupported compression codec " + string(r.compressionCodec))
	}

	if bytes.Compare(dst[0:8], DataMagic) != 0 {
		return nil, errors.New("bad data block magic")
	}

	return dst, nil
}

func (r *Reader) GetScanner() *Scanner {
	select {
	case s := <-r.scannerCache:
		return s
	default:
		return NewScanner(r)
	}
}

func (r *Reader) GetIterator() *Iterator {
	select {
	case i := <-r.iteratorCache:
		return i
	default:
		return NewIterator(r)
	}
}
<|endoftext|>


<commit_before>package libkbfs

import (
	"fmt"
	"sync"
)

type nodeCacheEntry struct {
	core     *nodeCore
	refCount int
}

// nodeCacheStandard implements the NodeCache interface by tracking
// the reference counts of nodeStandard Nodes, and using their member
// fields to construct paths.
type nodeCacheStandard struct {
	folderBranch FolderBranch
	nodes        map[blockRef]*nodeCacheEntry
	lock         sync.RWMutex
}

var _ NodeCache = (*nodeCacheStandard)(nil)

func newNodeCacheStandard(fb FolderBranch) *nodeCacheStandard {
	return &nodeCacheStandard{
		folderBranch: fb,
		nodes:        make(map[blockRef]*nodeCacheEntry),
	}
}

// lock must be locked for writing by the caller
func (ncs *nodeCacheStandard) forgetLocked(core *nodeCore) {
	ref := core.pathNode.ref()

	entry, ok := ncs.nodes[ref]
	if !ok {
		return
	}
	if entry.core != core {
		return
	}

	entry.refCount--
	if entry.refCount <= 0 {
		delete(ncs.nodes, ref)
	}
}

// should be called only by nodeStandardFinalizer().
func (ncs *nodeCacheStandard) forget(core *nodeCore) {
	ncs.lock.Lock()
	defer ncs.lock.Unlock()
	ncs.forgetLocked(core)
}

// lock must be held for writing by the caller
func (ncs *nodeCacheStandard) newChildForParentLocked(parent Node) (*nodeStandard, error) {
	nodeStandard, ok := parent.(*nodeStandard)
	if !ok {
		return nil, ParentNodeNotFoundError{blockRef{}}
	}

	ref := nodeStandard.core.pathNode.ref()
	entry, ok := ncs.nodes[ref]
	if !ok {
		return nil, ParentNodeNotFoundError{ref}
	}
	if nodeStandard.core != entry.core {
		return nil, ParentNodeNotFoundError{ref}
	}
	return nodeStandard, nil
}

func makeNodeStandardForEntry(entry *nodeCacheEntry) *nodeStandard {
	entry.refCount++
	return makeNodeStandard(entry.core)
}

// GetOrCreate implements the NodeCache interface for nodeCacheStandard.
func (ncs *nodeCacheStandard) GetOrCreate(
	ptr BlockPointer, name string, parent Node) (Node, error) {
	if !ptr.IsValid() {
		// Temporary code to track down bad block
		// pointers. Remove when not needed anymore.
		panic(InvalidBlockRefError{ptr.ref()})
	}

	if name == "" {
		return nil, EmptyNameError{ptr.ref()}
	}

	ncs.lock.Lock()
	defer ncs.lock.Unlock()
	entry, ok := ncs.nodes[ptr.ref()]
	if ok {
		// If the entry happens to be unlinked, we may be in a
		// situation where a node got unlinked and then recreated, but
		// someone held onto a node the whole time and so it never got
		// removed from the cache.  In that case, just stitch it back
		// together.
		if parent != nil && entry.core.parent == nil {
			entry.core.cachedPath = path{}
			entry.core.parent, ok = parent.(*nodeStandard)
			if !ok {
				return nil, ParentNodeNotFoundError{blockRef{}}
			}
			entry.core.pathNode.Name = name
		}
		return makeNodeStandardForEntry(entry), nil
	}

	var parentNS *nodeStandard
	if parent != nil {
		var err error
		parentNS, err = ncs.newChildForParentLocked(parent)
		if err != nil {
			return nil, err
		}
	}

	entry = &nodeCacheEntry{
		core: newNodeCore(ptr, name, parentNS, ncs),
	}
	ncs.nodes[ptr.ref()] = entry
	return makeNodeStandardForEntry(entry), nil
}

// Get implements the NodeCache interface for nodeCacheStandard.
func (ncs *nodeCacheStandard) Get(ref blockRef) Node {
	if ref == (blockRef{}) {
		return nil
	}

	// Temporary code to track down bad block pointers. Remove (or
	// return an error) when not needed anymore.
	if !ref.IsValid() {
		panic(InvalidBlockRefError{ref})
	}

	ncs.lock.Lock()
	defer ncs.lock.Unlock()
	entry, ok := ncs.nodes[ref]
	if !ok {
		return nil
	}
	return makeNodeStandardForEntry(entry)
}

// UpdatePointer implements the NodeCache interface for nodeCacheStandard.
func (ncs *nodeCacheStandard) UpdatePointer(
	oldRef blockRef, newPtr BlockPointer) {
	if oldRef == (blockRef{}) && newPtr == (BlockPointer{}) {
		return
	}

	if !oldRef.IsValid() {
		panic(fmt.Sprintf("invalid oldRef %s with newPtr %s", oldRef, newPtr))
	}

	if !newPtr.IsValid() {
		panic(fmt.Sprintf("invalid newPtr %s with oldRef %s", newPtr, oldRef))
	}

	ncs.lock.Lock()
	defer ncs.lock.Unlock()
	entry, ok := ncs.nodes[oldRef]
	if !ok {
		return
	}

	// Cannot update the pointer for an unlinked node
	if len(entry.core.cachedPath.path) > 0 {
		return
	}

	entry.core.pathNode.BlockPointer = newPtr
	delete(ncs.nodes, oldRef)
	ncs.nodes[newPtr.ref()] = entry
}

// Move implements the NodeCache interface for nodeCacheStandard.
func (ncs *nodeCacheStandard) Move(
	ref blockRef, newParent Node, newName string) error {
	if ref == (blockRef{}) {
		return nil
	}

	// Temporary code to track down bad block pointers. Remove (or
	// return an error) when not needed anymore.
	if !ref.IsValid() {
		panic(InvalidBlockRefError{ref})
	}

	if newName == "" {
		return EmptyNameError{ref}
	}

	ncs.lock.Lock()
	defer ncs.lock.Unlock()
	entry, ok := ncs.nodes[ref]
	if !ok {
		return nil
	}

	newParentNS, err := ncs.newChildForParentLocked(newParent)
	if err != nil {
		return err
	}

	entry.core.parent = newParentNS
	entry.core.pathNode.Name = newName
	return nil
}

// Unlink implements the NodeCache interface for nodeCacheStandard.
func (ncs *nodeCacheStandard) Unlink(ref blockRef, oldPath path) {
	if ref == (blockRef{}) {
		return
	}

	// Temporary code to track down bad block pointers. Remove (or
	// return an error) when not needed anymore.
	if !ref.IsValid() {
		panic(InvalidBlockRefError{ref})
	}

	ncs.lock.Lock()
	defer ncs.lock.Unlock()
	entry, ok := ncs.nodes[ref]
	if !ok {
		return
	}

	entry.core.cachedPath = oldPath
	entry.core.parent = nil
	entry.core.pathNode.Name = ""
	return
}

// PathFromNode implements the NodeCache interface for nodeCacheStandard.
func (ncs *nodeCacheStandard) PathFromNode(node Node) (p path) {
	ncs.lock.RLock()
	defer ncs.lock.RUnlock()

	ns, ok := node.(*nodeStandard)
	if !ok {
		p.path = nil
		return
	}

	for ns != nil {
		core := ns.core
		if core.parent == nil && len(core.cachedPath.path) > 0 {
			// The node was unlinked, but is still in use, so use its
			// cached path.  The path is already reversed, so append
			// it backwards one-by-one to the existing path.  If this
			// is the first node, we can just optimize by returning
			// the complete cached path.
			if len(p.path) == 0 {
				return core.cachedPath
			}
			for i := len(core.cachedPath.path) - 1; i >= 0; i-- {
				p.path = append(p.path, core.cachedPath.path[i])
			}
			break
		}

		p.path = append(p.path, *core.pathNode)
		ns = core.parent
	}

	// need to reverse the path nodes
	for i := len(p.path)/2 - 1; i >= 0; i-- {
		opp := len(p.path) - 1 - i
		p.path[i], p.path[opp] = p.path[opp], p.path[i]
	}

	// TODO: would it make any sense to cache the constructed path?
	p.FolderBranch = ncs.folderBranch
	return
}
<commit_msg>node_cache: don't try to reuse unlinked nodes<commit_after>package libkbfs

import (
	"fmt"
	"sync"
)

type nodeCacheEntry struct {
	core     *nodeCore
	refCount int
}

// nodeCacheStandard implements the NodeCache interface by tracking
// the reference counts of nodeStandard Nodes, and using their member
// fields to construct paths.
type nodeCacheStandard struct {
	folderBranch FolderBranch
	nodes        map[blockRef]*nodeCacheEntry
	lock         sync.RWMutex
}

var _ NodeCache = (*nodeCacheStandard)(nil)

func newNodeCacheStandard(fb FolderBranch) *nodeCacheStandard {
	return &nodeCacheStandard{
		folderBranch: fb,
		nodes:        make(map[blockRef]*nodeCacheEntry),
	}
}

// lock must be locked for writing by the caller
func (ncs *nodeCacheStandard) forgetLocked(core *nodeCore) {
	ref := core.pathNode.ref()

	entry, ok := ncs.nodes[ref]
	if !ok {
		return
	}
	if entry.core != core {
		return
	}

	entry.refCount--
	if entry.refCount <= 0 {
		delete(ncs.nodes, ref)
	}
}

// should be called only by nodeStandardFinalizer().
func (ncs *nodeCacheStandard) forget(core *nodeCore) {
	ncs.lock.Lock()
	defer ncs.lock.Unlock()
	ncs.forgetLocked(core)
}

// lock must be held for writing by the caller
func (ncs *nodeCacheStandard) newChildForParentLocked(parent Node) (*nodeStandard, error) {
	nodeStandard, ok := parent.(*nodeStandard)
	if !ok {
		return nil, ParentNodeNotFoundError{blockRef{}}
	}

	ref := nodeStandard.core.pathNode.ref()
	entry, ok := ncs.nodes[ref]
	if !ok {
		return nil, ParentNodeNotFoundError{ref}
	}
	if nodeStandard.core != entry.core {
		return nil, ParentNodeNotFoundError{ref}
	}
	return nodeStandard, nil
}

func makeNodeStandardForEntry(entry *nodeCacheEntry) *nodeStandard {
	entry.refCount++
	return makeNodeStandard(entry.core)
}

// GetOrCreate implements the NodeCache interface for nodeCacheStandard.
func (ncs *nodeCacheStandard) GetOrCreate(
	ptr BlockPointer, name string, parent Node) (Node, error) {
	if !ptr.IsValid() {
		// Temporary code to track down bad block
		// pointers. Remove when not needed anymore.
		panic(InvalidBlockRefError{ptr.ref()})
	}

	if name == "" {
		return nil, EmptyNameError{ptr.ref()}
	}

	ncs.lock.Lock()
	defer ncs.lock.Unlock()
	entry, ok := ncs.nodes[ptr.ref()]
	if ok {
		// If the entry happens to be unlinked, we may be in a
		// situation where a node got unlinked and then recreated, but
		// someone held onto a node the whole time and so it never got
		// removed from the cache.  In that case, forcibly remove it
		// from the cache to make room for the new node.
		if parent != nil && entry.core.parent == nil {
			delete(ncs.nodes, ptr.ref())
		} else {
			return makeNodeStandardForEntry(entry), nil
		}
	}

	var parentNS *nodeStandard
	if parent != nil {
		var err error
		parentNS, err = ncs.newChildForParentLocked(parent)
		if err != nil {
			return nil, err
		}
	}

	entry = &nodeCacheEntry{
		core: newNodeCore(ptr, name, parentNS, ncs),
	}
	ncs.nodes[ptr.ref()] = entry
	return makeNodeStandardForEntry(entry), nil
}

// Get implements the NodeCache interface for nodeCacheStandard.
func (ncs *nodeCacheStandard) Get(ref blockRef) Node {
	if ref == (blockRef{}) {
		return nil
	}

	// Temporary code to track down bad block pointers. Remove (or
	// return an error) when not needed anymore.
	if !ref.IsValid() {
		panic(InvalidBlockRefError{ref})
	}

	ncs.lock.Lock()
	defer ncs.lock.Unlock()
	entry, ok := ncs.nodes[ref]
	if !ok {
		return nil
	}
	return makeNodeStandardForEntry(entry)
}

// UpdatePointer implements the NodeCache interface for nodeCacheStandard.
func (ncs *nodeCacheStandard) UpdatePointer(
	oldRef blockRef, newPtr BlockPointer) {
	if oldRef == (blockRef{}) && newPtr == (BlockPointer{}) {
		return
	}

	if !oldRef.IsValid() {
		panic(fmt.Sprintf("invalid oldRef %s with newPtr %s", oldRef, newPtr))
	}

	if !newPtr.IsValid() {
		panic(fmt.Sprintf("invalid newPtr %s with oldRef %s", newPtr, oldRef))
	}

	ncs.lock.Lock()
	defer ncs.lock.Unlock()
	entry, ok := ncs.nodes[oldRef]
	if !ok {
		return
	}

	// Cannot update the pointer for an unlinked node
	if len(entry.core.cachedPath.path) > 0 {
		return
	}

	entry.core.pathNode.BlockPointer = newPtr
	delete(ncs.nodes, oldRef)
	ncs.nodes[newPtr.ref()] = entry
}

// Move implements the NodeCache interface for nodeCacheStandard.
func (ncs *nodeCacheStandard) Move(
	ref blockRef, newParent Node, newName string) error {
	if ref == (blockRef{}) {
		return nil
	}

	// Temporary code to track down bad block pointers. Remove (or
	// return an error) when not needed anymore.
	if !ref.IsValid() {
		panic(InvalidBlockRefError{ref})
	}

	if newName == "" {
		return EmptyNameError{ref}
	}

	ncs.lock.Lock()
	defer ncs.lock.Unlock()
	entry, ok := ncs.nodes[ref]
	if !ok {
		return nil
	}

	newParentNS, err := ncs.newChildForParentLocked(newParent)
	if err != nil {
		return err
	}

	entry.core.parent = newParentNS
	entry.core.pathNode.Name = newName
	return nil
}

// Unlink implements the NodeCache interface for nodeCacheStandard.
func (ncs *nodeCacheStandard) Unlink(ref blockRef, oldPath path) {
	if ref == (blockRef{}) {
		return
	}

	// Temporary code to track down bad block pointers. Remove (or
	// return an error) when not needed anymore.
	if !ref.IsValid() {
		panic(InvalidBlockRefError{ref})
	}

	ncs.lock.Lock()
	defer ncs.lock.Unlock()
	entry, ok := ncs.nodes[ref]
	if !ok {
		return
	}

	entry.core.cachedPath = oldPath
	entry.core.parent = nil
	entry.core.pathNode.Name = ""
	return
}

// PathFromNode implements the NodeCache interface for nodeCacheStandard.
func (ncs *nodeCacheStandard) PathFromNode(node Node) (p path) {
	ncs.lock.RLock()
	defer ncs.lock.RUnlock()

	ns, ok := node.(*nodeStandard)
	if !ok {
		p.path = nil
		return
	}

	for ns != nil {
		core := ns.core
		if core.parent == nil && len(core.cachedPath.path) > 0 {
			// The node was unlinked, but is still in use, so use its
			// cached path.  The path is already reversed, so append
			// it backwards one-by-one to the existing path.  If this
			// is the first node, we can just optimize by returning
			// the complete cached path.
			if len(p.path) == 0 {
				return core.cachedPath
			}
			for i := len(core.cachedPath.path) - 1; i >= 0; i-- {
				p.path = append(p.path, core.cachedPath.path[i])
			}
			break
		}

		p.path = append(p.path, *core.pathNode)
		ns = core.parent
	}

	// need to reverse the path nodes
	for i := len(p.path)/2 - 1; i >= 0; i-- {
		opp := len(p.path) - 1 - i
		p.path[i], p.path[opp] = p.path[opp], p.path[i]
	}

	// TODO: would it make any sense to cache the constructed path?
	p.FolderBranch = ncs.folderBranch
	return
}
<|endoftext|>


<commit_before>package libmap

import (
	"testing"
)

func TestTSafeMapBytesSetGet(t *testing.T) {
	testData := `{"Free": 1000, "Used": 500}`

	s := NewTSafeMapBytes(nil)
	s.Set("/free", []byte(testData))

	if string(s.Get("/free")) != testData {
		t.Errorf("Failed to test set and get. Actual Data: %v", string(s.Get("/free")))
	}
}

func TestTSafeNestedMapInterfaceInitNestedMap(t *testing.T) {
	m := NewTSafeNestedMapInterface(nil)
	m.initNestedMap("aaa.bbb.ccc")

	if m.Data["aaa"] == nil {
		t.Fatalf("Failed to init nested map")
	}
	if m.Data["aaa"].(map[string]interface{})["bbb"] == nil {
		t.Fatalf("Failed to init nested map")
	}

	m.Data["aaa"].(map[string]interface{})["bbb"].(map[string]interface{})["ccc"] = 42

	val := m.Data["aaa"].(map[string]interface{})["bbb"].(map[string]interface{})["ccc"].(int)
	expected := 42
	if val != expected {
		t.Fatalf("Failed to get value on nested map. Expected: %v, Got: %v", expected, val)
	}
}

func TestTSafeNestedMapInterfaceSetGet(t *testing.T) {
	m := NewTSafeNestedMapInterface(nil)

	m.Set("aaa.bbb.ccc", 42)

	if m.Data["aaa"] == nil {
		t.Fatalf("Failed to init nested map")
	}
	if m.Data["aaa"].(map[string]interface{})["bbb"] == nil {
		t.Fatalf("Failed to init nested map")
	}

	val := m.Data["aaa"].(map[string]interface{})["bbb"].(map[string]interface{})["ccc"].(int)
	expected := 42
	if val != expected {
		t.Fatalf("Failed to get value on nested map. Expected: %v, Got: %v", expected, val)
	}
}
<commit_msg>Basic tests on libmap strings.<commit_after>package libmap

import (
	"testing"
)

func TestTSafeMapBytesSetGet(t *testing.T) {
	testData := `{"Free": 1000, "Used": 500}`

	s := NewTSafeMapBytes(nil)
	s.Set("/free", []byte(testData))

	if string(s.Get("/free")) != testData {
		t.Errorf("Failed to test set and get. Actual Data: %v", string(s.Get("/free")))
	}
}

func TestTSafeNestedMapInterfaceInitNestedMap(t *testing.T) {
	m := NewTSafeNestedMapInterface(nil)
	m.initNestedMap("aaa.bbb.ccc")

	if m.Get("aaa") == nil {
		t.Fatalf("Failed to init nested map")
	}
	if m.Get("aaa").(map[string]interface{})["bbb"] == nil {
		t.Fatalf("Failed to init nested map")
	}

	m.Get("aaa").(map[string]interface{})["bbb"].(map[string]interface{})["ccc"] = 42

	val := m.Get("aaa").(map[string]interface{})["bbb"].(map[string]interface{})["ccc"].(int)
	expected := 42
	if val != expected {
		t.Fatalf("Failed to get value on nested map. Expected: %v, Got: %v", expected, val)
	}
}

func TestTSafeNestedMapInterfaceSetGet(t *testing.T) {
	m := NewTSafeNestedMapInterface(nil)

	m.Set("aaa.bbb.ccc", 42)

	if m.Get("aaa") == nil {
		t.Fatalf("Failed to init nested map")
	}
	if m.Get("aaa").(map[string]interface{})["bbb"] == nil {
		t.Fatalf("Failed to init nested map")
	}

	val := m.Get("aaa").(map[string]interface{})["bbb"].(map[string]interface{})["ccc"].(int)
	expected := 42
	if val != expected {
		t.Fatalf("Failed to get value on nested map. Expected: %v, Got: %v", expected, val)
	}
}

func TestNewTSafeMapStringsBasicFunctionality(t *testing.T) {
	logDB := NewTSafeMapStrings(map[string][]string{
		"Loglines": make([]string, 0),
	})

	logDB.Append("Loglines", "some log")

	logs := logDB.Get("Loglines")
	if len(logs) != 1 {
		t.Fatalf("Failed to get value on string slice. Expected: %v, Got: %v", 1, len(logs))
	}

	logDB.Reset("Loglines")

	logs = logDB.Get("Loglines")
	if len(logs) != 0 {
		t.Fatalf("Failed to get value on string slice. Expected: %v, Got: %v", 0, len(logs))
	}
}
<|endoftext|>


<commit_before>package fq

import (
	"bufio"
	"errors"
	"fmt"
	"io"
	"os"
	"sort"
)

type Reader struct {
	dir          string
	offset       uint64
	r            *bufio.Reader
	file         *os.File
	journalFiles journalFiles
	journalIndex int
}

func NewReader(dir string, offset uint64) (*Reader, error) {
	files, err := getJournalFiles(dir)
	if err != nil {
		return nil, err
	}
	i := sort.Search(len(files), func(i int) bool { return files[i].startOffset > offset })
	if i == 0 {
		return nil, errors.New("offset is too small")
	}

	journalIndex := i - 1
	file := &files[journalIndex]
	reader := Reader{
		dir:          dir,
		journalFiles: files,
		journalIndex: journalIndex,
	}
	if err := reader.openFile(file.fileName); err != nil {
		return nil, err
	}
	reader.r = bufio.NewReader(reader.file)
	reader.offset = file.startOffset
	for reader.offset < offset {
		if _, err := reader.Read(); err != nil {
			return nil, err
		}
	}
	if reader.offset != offset {
		return nil, fmt.Errorf("fail to find offset %d", offset)
	}
	return &reader, nil
}

func (r *Reader) Read() (msg []byte, err error) {
	msg, offset, err := readMessage(r.r)
	if err == io.EOF {
		/*
			watcher, err := fsnotify.NewWatcher()
			if err != nil {
				return nil, err
			}
		*/
		files, err := getJournalFiles(r.dir)
		if err != nil {
			return nil, err
		}
		r.journalFiles = files
		if r.journalIndex < len(r.journalFiles)-1 && r.offset == r.journalFiles[r.journalIndex+1].startOffset {
			r.closeFile()
			r.journalIndex++
			journalFile := &r.journalFiles[r.journalIndex]
			if err := r.openFile(journalFile.fileName); err != nil {
				return nil, err
			}
			r.r = bufio.NewReader(r.file)
			return r.Read()
		}
		if err := r.reopenFile(); err != nil {
			return nil, err
		}
		msg, offset, err := readMessage(r.r)
		if err == io.EOF && offset == r.offset+1 {
			return msg, nil
		} else if err != nil {
			return nil, err
		}
		if offset != r.offset {
			return nil, fmt.Errorf("offset is out of order: %d, %d", offset, r.offset)
		}
		return msg, nil
	}
	if offset != r.offset {
		return nil, fmt.Errorf("offset is out of order: %d, %d", offset, r.offset)
	}
	r.offset++
	return msg, nil
}

func (r *Reader) Offset() uint64 {
	return r.offset
}

func (r *Reader) Close() {
	r.closeFile()
}

func (r *Reader) closeFile() {
	r.file.Close()
}

func (r *Reader) openFile(name string) error {
	var err error
	r.file, err = os.Open(name)
	if err != nil {
		return err
	}
	return nil
}

func (r *Reader) reopenFile() error {
	journalFile := &r.journalFiles[r.journalIndex]
	offset := r.offset
	r.closeFile()
	if err := r.openFile(journalFile.fileName); err != nil {
		return err
	}
	r.r = bufio.NewReader(r.file)

	r.offset = journalFile.startOffset
	for r.offset < offset {
		_, _, err := readMessage(r.r)
		if err != nil {
			return err
		}
		r.offset++
	}
	return nil
}

/*
func (r *Reader) waitForFileAppend() error {
	r.watcher.Add(r.file.Name())
	defer r.watcher.Remove(r.file.Name())
	select {
	case event := <-r.watcher.Events:
		if event.Op&fsnotify.Write == fsnotify.Write {
			return r.reopenFile()
		}
	case err := <-r.watcher.Errors:
		return err
	}
	return nil
}
*/
<commit_msg>refactor<commit_after>package fq

import (
	"bufio"
	"errors"
	"fmt"
	"io"
	"os"
	"sort"
)

type Reader struct {
	dir          string
	offset       uint64
	r            *bufio.Reader
	file         *os.File
	journalFiles journalFiles
	journalIndex int
}

func NewReader(dir string, offset uint64) (*Reader, error) {
	files, err := getJournalFiles(dir)
	if err != nil {
		return nil, err
	}
	i := sort.Search(len(files), func(i int) bool { return files[i].startOffset > offset })
	if i == 0 {
		return nil, errors.New("offset is too small")
	}

	journalIndex := i - 1
	file := &files[journalIndex]
	reader := Reader{
		dir:          dir,
		journalFiles: files,
		journalIndex: journalIndex,
	}
	if err := reader.openFile(file.fileName); err != nil {
		return nil, err
	}
	reader.r = bufio.NewReader(reader.file)
	reader.offset = file.startOffset
	for reader.offset < offset {
		if _, err := reader.Read(); err != nil {
			return nil, err
		}
	}
	if reader.offset != offset {
		return nil, fmt.Errorf("fail to find offset %d", offset)
	}
	return &reader, nil
}

func (r *Reader) Read() (msg []byte, err error) {
	msg, offset, err := readMessage(r.r)
	if err == io.EOF {
		/*
			watcher, err := fsnotify.NewWatcher()
			if err != nil {
				return nil, err
			}
		*/
		files, err := getJournalFiles(r.dir)
		if err != nil {
			return nil, err
		}

		r.journalFiles = files
		if r.journalIndex < len(r.journalFiles)-1 && r.offset == r.journalFiles[r.journalIndex+1].startOffset {
			r.closeFile()
			r.journalIndex++
			journalFile := &r.journalFiles[r.journalIndex]
			if err := r.openFile(journalFile.fileName); err != nil {
				return nil, err
			}
			r.r = bufio.NewReader(r.file)
			return r.Read()
		}
		if err := r.reopenFile(); err != nil {
			return nil, err
		}
		msg, offset, err := readMessage(r.r)
		if err == io.EOF && offset == r.offset+1 {
			return msg, nil
		} else if err != nil {
			return nil, err
		}
		if offset != r.offset {
			return nil, fmt.Errorf("offset is out of order: %d, %d", offset, r.offset)
		}
		return msg, nil
	}
	if offset != r.offset {
		return nil, fmt.Errorf("offset is out of order: %d, %d", offset, r.offset)
	}
	r.offset++
	return msg, nil
}

func (r *Reader) Offset() uint64 {
	return r.offset
}

func (r *Reader) Close() {
	r.closeFile()
}

func (r *Reader) closeFile() {
	r.file.Close()
}

func (r *Reader) openFile(name string) error {
	var err error
	r.file, err = os.Open(name)
	if err != nil {
		return err
	}
	return nil
}

func (r *Reader) reopenFile() error {
	fileName := r.file.Name()
	fileOffset, err := r.file.Seek(0, os.SEEK_CUR)
	if err != nil {
		return err
	}
	r.closeFile()
	if err := r.openFile(fileName); err != nil {
		return err
	}
	if _, err := r.file.Seek(fileOffset, os.SEEK_SET); err != nil {
		return err
	}
	r.r = bufio.NewReader(r.file)
	return nil
}

/*
func (r *Reader) waitForFileAppend() error {
	r.watcher.Add(r.file.Name())
	defer r.watcher.Remove(r.file.Name())
	select {
	case event := <-r.watcher.Events:
		if event.Op&fsnotify.Write == fsnotify.Write {
			return r.reopenFile()
		}
	case err := <-r.watcher.Errors:
		return err
	}
	return nil
}
*/
<|endoftext|>


<commit_before>package operator

import (
	"fmt"
	"github.com/phzfi/RIC/server/images"
	"github.com/phzfi/RIC/server/ops"
	"github.com/phzfi/RIC/server/testutils"
	"sync"
	"testing"
	"time"
)

const cacheFolder = "/tmp/operatortests"

type DummyOperation struct {
	log  *[]int
	name int
}

func (o *DummyOperation) Marshal() string {
	return fmt.Sprintf("test%v", o.name)
}

var logMutex *sync.Mutex = &sync.Mutex{}

func (o *DummyOperation) Apply(img images.Image) error {
	// Take some time for simult opers. tests
	time.Sleep(200 * time.Millisecond)
	logMutex.Lock()
	*(o.log) = append(*(o.log), o.name)
	logMutex.Unlock()
	return nil
}

func TestOperator(t *testing.T) {
	var log []int
	operations := []ops.Operation{
		&DummyOperation{&log, 0},
		&DummyOperation{&log, 1},
		&DummyOperation{&log, 2},
	}

	testutils.RemoveContents(cacheFolder)
	operator := MakeDefault(512*1024*1024, cacheFolder)

	_, err := operator.GetBlob(operations...)
	if err != nil {
		t.Error(err)
	}

	if len(log) != 3 {
		t.Fatal("Too many or too few operations done")
	}
	for i, v := range log {
		if i != v {
			t.Fatal("Wrong operation")
		}
	}
}

func TestDenyIdenticalOperations(t *testing.T) {
	testutils.RemoveContents(cacheFolder)

	var log []int

	// Many identical operations
	operations := [][]ops.Operation{
		{&DummyOperation{&log, 0}, &DummyOperation{&log, 0}},
		{&DummyOperation{&log, 0}, &DummyOperation{&log, 0}},
		{&DummyOperation{&log, 0}, &DummyOperation{&log, 0}},
		{&DummyOperation{&log, 0}, &DummyOperation{&log, 0}},
		{&DummyOperation{&log, 0}, &DummyOperation{&log, 0}},
		{&DummyOperation{&log, 0}, &DummyOperation{&log, 0}},
	}
	operator := MakeDefault(512*1024*1024, cacheFolder)

	// Channel to track amount of completed operations
	c := make(chan bool, len(operations))

	// Launch operations simultaneously
	for _, ops := range operations {
		go func() {
			_, _ = operator.GetBlob(ops...)
			c <- true
		}()
	}

	// Wait for the operations to finish
	for i := 0; i < len(operations); i++ {
		<-c
	}

	// Only 2 operations should've been done - others found from cache
	if len(log) != 2 {
		t.Fatal(fmt.Sprintf("%v operations done. Expected 2", len(log)))
	}
}
<commit_msg>fix test data race when launching go routines inside for loop<commit_after>package operator

import (
	"fmt"
	"github.com/phzfi/RIC/server/images"
	"github.com/phzfi/RIC/server/ops"
	"github.com/phzfi/RIC/server/testutils"
	"sync"
	"testing"
	"time"
)

const cacheFolder = "/tmp/operatortests"

type DummyOperation struct {
	log  *[]int
	name int
}

func (o *DummyOperation) Marshal() string {
	return fmt.Sprintf("test%v", o.name)
}

var logMutex *sync.Mutex = &sync.Mutex{}

func (o *DummyOperation) Apply(img images.Image) error {
	// Take some time for simult opers. tests
	time.Sleep(200 * time.Millisecond)
	logMutex.Lock()
	*(o.log) = append(*(o.log), o.name)
	logMutex.Unlock()
	return nil
}

func TestOperator(t *testing.T) {
	var log []int
	operations := []ops.Operation{
		&DummyOperation{&log, 0},
		&DummyOperation{&log, 1},
		&DummyOperation{&log, 2},
	}

	testutils.RemoveContents(cacheFolder)
	operator := MakeDefault(512*1024*1024, cacheFolder)

	_, err := operator.GetBlob(operations...)
	if err != nil {
		t.Error(err)
	}

	if len(log) != 3 {
		t.Fatal("Too many or too few operations done")
	}
	for i, v := range log {
		if i != v {
			t.Fatal("Wrong operation")
		}
	}
}

func TestDenyIdenticalOperations(t *testing.T) {
	testutils.RemoveContents(cacheFolder)

	var log []int

	// Many identical operations
	operations := [][]ops.Operation{
		{&DummyOperation{&log, 0}, &DummyOperation{&log, 0}},
		{&DummyOperation{&log, 0}, &DummyOperation{&log, 0}},
		{&DummyOperation{&log, 0}, &DummyOperation{&log, 0}},
		{&DummyOperation{&log, 0}, &DummyOperation{&log, 0}},
		{&DummyOperation{&log, 0}, &DummyOperation{&log, 0}},
		{&DummyOperation{&log, 0}, &DummyOperation{&log, 0}},
	}
	operator := MakeDefault(512*1024*1024, cacheFolder)

	// Channel to track amount of completed operations
	c := make(chan bool, len(operations))

	// Launch operations simultaneously
	for i := range operations {
		ops := operations[i]
		go func() {
			_, _ = operator.GetBlob(ops...)
			c <- true
		}()
	}

	// Wait for the operations to finish
	for i := 0; i < len(operations); i++ {
		<-c
	}

	// Only 2 operations should've been done - others found from cache
	if len(log) != 2 {
		t.Fatal(fmt.Sprintf("%v operations done. Expected 2", len(log)))
	}
}
<|endoftext|>


<commit_before>package terraformrules

import (
	"testing"

	hcl "github.com/hashicorp/hcl/v2"
	"github.com/terraform-linters/tflint/tflint"
)

func Test_TerraformUnusedDeclarationsRule(t *testing.T) {
	cases := []struct {
		Name     string
		Content  string
		JSON     bool
		Expected tflint.Issues
	}{
		{
			Name: "unused variable",
			Content: `
variable "not_used" {}

variable "used" {}
output "u" { value = var.used }
`,
			Expected: tflint.Issues{
				{
					Rule:    NewTerraformUnusedDeclarationsRule(),
					Message: `variable "not_used" is declared but not used`,
					Range: hcl.Range{
						Filename: "config.tf",
						Start:    hcl.Pos{Line: 2, Column: 1},
						End:      hcl.Pos{Line: 2, Column: 20},
					},
				},
			},
		},
		{
			Name: "unused data source",
			Content: `
data "null_data_source" "not_used" {}

data "null_data_source" "used" {}
output "u" { value = data.null_data_source.used }
`,
			Expected: tflint.Issues{
				{
					Rule:    NewTerraformUnusedDeclarationsRule(),
					Message: `data "null_data_source" "not_used" is declared but not used`,
					Range: hcl.Range{
						Filename: "config.tf",
						Start:    hcl.Pos{Line: 2, Column: 1},
						End:      hcl.Pos{Line: 2, Column: 35},
					},
				},
			},
		},
		{
			Name: "unused local source",
			Content: `
locals {
	not_used = ""
	used = ""
}

output "u" { value = local.used }
`,
			Expected: tflint.Issues{
				{
					Rule:    NewTerraformUnusedDeclarationsRule(),
					Message: `local.not_used is declared but not used`,
					Range: hcl.Range{
						Filename: "config.tf",
						Start:    hcl.Pos{Line: 3, Column: 2},
						End:      hcl.Pos{Line: 3, Column: 15},
					},
				},
			},
		},
		{
			Name: "variable used in resource",
			Content: `
variable "used" {}
resource "null_resource" "n" {
	triggers = {
		u = var.used
	}
}
`,
			Expected: tflint.Issues{},
		},
		{
			Name: "variable used in module",
			Content: `
variable "used" {}
module "m" {
	source = "."
	u = var.used
}
`,
			Expected: tflint.Issues{},
		},
		{
			Name: "variable used in module",
			Content: `
variable "used" {}
module "m" {
	source = "."
	u = var.used
}
`,
			Expected: tflint.Issues{},
		},
		{
			Name: "variable used in provider",
			Content: `
variable "aws_region" {}
provider "aws" {
	region = var.aws_region
}
`,
			Expected: tflint.Issues{},
		},
		{
			Name: "meta-arguments",
			Content: `
variable "used" {}
resource "null_resource" "n" {
  triggers = {
    u = var.used
	}

  lifecycle {
    ignore_changes = [triggers]
  }

  providers = {
    null = null
  }

  depends_on = [aws_instance.foo]
}
`,
			Expected: tflint.Issues{},
		},
		{
			Name: "additional traversal",
			Content: `
variable "v" {
	type = object({ foo = string })
}
output "v" {
	value = var.v.foo
}

data "terraform_remote_state" "d" {}
output "d" {
	value = data.terraform_remote_state.d.outputs.foo
}
`,
			Expected: tflint.Issues{},
		},
		{
			Name: "json",
			JSON: true,
			Content: `
{
  "resource": {
    "foo": {
      "bar": {
        "nested": [{
          "${var.again}": []
        }]
      }
    }
	},
  "variable": {
    "again": {}
  }
}`,
			Expected: tflint.Issues{},
		},
	}

	rule := NewTerraformUnusedDeclarationsRule()

	for _, tc := range cases {
		filename := "config.tf"
		if tc.JSON {
			filename += ".json"
		}

		t.Run(tc.Name, func(t *testing.T) {
			runner := tflint.TestRunner(t, map[string]string{filename: tc.Content})

			if err := rule.Check(runner); err != nil {
				t.Fatalf("Unexpected error occurred: %s", err)
			}

			tflint.AssertIssues(t, tc.Expected, runner.Issues)
		})
	}
}
<commit_msg>terraform_unused_declarations: add test for local + module<commit_after>package terraformrules

import (
	"testing"

	hcl "github.com/hashicorp/hcl/v2"
	"github.com/terraform-linters/tflint/tflint"
)

func Test_TerraformUnusedDeclarationsRule(t *testing.T) {
	cases := []struct {
		Name     string
		Content  string
		JSON     bool
		Expected tflint.Issues
	}{
		{
			Name: "unused variable",
			Content: `
variable "not_used" {}

variable "used" {}
output "u" { value = var.used }
`,
			Expected: tflint.Issues{
				{
					Rule:    NewTerraformUnusedDeclarationsRule(),
					Message: `variable "not_used" is declared but not used`,
					Range: hcl.Range{
						Filename: "config.tf",
						Start:    hcl.Pos{Line: 2, Column: 1},
						End:      hcl.Pos{Line: 2, Column: 20},
					},
				},
			},
		},
		{
			Name: "unused data source",
			Content: `
data "null_data_source" "not_used" {}

data "null_data_source" "used" {}
output "u" { value = data.null_data_source.used }
`,
			Expected: tflint.Issues{
				{
					Rule:    NewTerraformUnusedDeclarationsRule(),
					Message: `data "null_data_source" "not_used" is declared but not used`,
					Range: hcl.Range{
						Filename: "config.tf",
						Start:    hcl.Pos{Line: 2, Column: 1},
						End:      hcl.Pos{Line: 2, Column: 35},
					},
				},
			},
		},
		{
			Name: "unused local source",
			Content: `
locals {
	not_used = ""
	used = ""
}

output "u" { value = local.used }
`,
			Expected: tflint.Issues{
				{
					Rule:    NewTerraformUnusedDeclarationsRule(),
					Message: `local.not_used is declared but not used`,
					Range: hcl.Range{
						Filename: "config.tf",
						Start:    hcl.Pos{Line: 3, Column: 2},
						End:      hcl.Pos{Line: 3, Column: 15},
					},
				},
			},
		},
		{
			Name: "variable used in resource",
			Content: `
variable "used" {}
resource "null_resource" "n" {
	triggers = {
		u = var.used
	}
}
`,
			Expected: tflint.Issues{},
		},
		{
			Name: "variable used in module",
			Content: `
variable "used" {}
module "m" {
	source = "."
	u = var.used
}
`,
			Expected: tflint.Issues{},
		},
		{
			Name: "variable used in module",
			Content: `
variable "used" {}
module "m" {
	source = "."
	u = var.used
}
`,
			Expected: tflint.Issues{},
		},
		{
			Name: "local used in module",
			Content: `
locals { used = "used" }
module "m" {
	source = "."
	u = local.used
}
`,
			Expected: tflint.Issues{},
		},
		{
			Name: "variable used in provider",
			Content: `
variable "aws_region" {}
provider "aws" {
	region = var.aws_region
}
`,
			Expected: tflint.Issues{},
		},
		{
			Name: "meta-arguments",
			Content: `
variable "used" {}
resource "null_resource" "n" {
  triggers = {
    u = var.used
	}

  lifecycle {
    ignore_changes = [triggers]
  }

  providers = {
    null = null
  }

  depends_on = [aws_instance.foo]
}
`,
			Expected: tflint.Issues{},
		},
		{
			Name: "additional traversal",
			Content: `
variable "v" {
	type = object({ foo = string })
}
output "v" {
	value = var.v.foo
}

data "terraform_remote_state" "d" {}
output "d" {
	value = data.terraform_remote_state.d.outputs.foo
}
`,
			Expected: tflint.Issues{},
		},
		{
			Name: "json",
			JSON: true,
			Content: `
{
  "resource": {
    "foo": {
      "bar": {
        "nested": [{
          "${var.again}": []
        }]
      }
    }
	},
  "variable": {
    "again": {}
  }
}`,
			Expected: tflint.Issues{},
		},
	}

	rule := NewTerraformUnusedDeclarationsRule()

	for _, tc := range cases {
		filename := "config.tf"
		if tc.JSON {
			filename += ".json"
		}

		t.Run(tc.Name, func(t *testing.T) {
			runner := tflint.TestRunner(t, map[string]string{filename: tc.Content})

			if err := rule.Check(runner); err != nil {
				t.Fatalf("Unexpected error occurred: %s", err)
			}

			tflint.AssertIssues(t, tc.Expected, runner.Issues)
		})
	}
}
<|endoftext|>


<commit_before>package dock

import (
	"fmt"
	"io"
	"testing"
	"time"

	"github.com/simulatedsimian/assert"
)

func TestParseRequest(t *testing.T) {
	assert := assert.Make(t)

	assert(msgToRequest("e")).Equal(Request{RequestType: ReqEnquire})
	assert(msgToRequest("r")).Equal(Request{RequestType: ReqResetToBootloader})
	assert(msgToRequest("v")).Equal(Request{RequestType: ReqVersion})
	assert(msgToRequest("d")).Equal(Request{RequestType: ReqDebug})
	assert(msgToRequest("p 0")).Equal(Request{RequestType: ReqPower, Params: []int{0}})
	assert(msgToRequest("p 1")).Equal(Request{RequestType: ReqPower, Params: []int{1}})
	assert(msgToRequest("p")).Equal(Request{})
	assert(msgToRequest("p 3")).Equal(Request{})

	assert(msgToRequest("n d hello")).Equal(Request{RequestType: ReqName, Params: []int{'d'}, ParamStr: "hello"})
	assert(msgToRequest("n u world")).Equal(Request{RequestType: ReqName, Params: []int{'u'}, ParamStr: "world"})

	assert(msgToRequest("s 5 1,2,3")).Equal(Request{RequestType: ReqSet, Channel: 5, Params: []int{1, 2, 3}})
}

func TestSimulatorRequest(t *testing.T) {
	assert := assert.Make(t)

	e1, e2 := NewPipe().Endpoints()
	sim := NewSimulator(e1)

	fmt.Fprint(e2, "e\r")
	assert(<-sim.Requests).Equal(Request{RequestType: ReqEnquire})
	e2.Close()
	assert(<-sim.Requests).Equal(Request{RequestType: ReqError, Error: io.EOF})

}

func TestSimulatorConnectDisconnect(t *testing.T) {
	assert := assert.Make(t)

	buffer := make([]byte, 128)
	e1, e2 := NewPipe().Endpoints()
	sim := NewSimulator(e1)

	assert(sim.modules[2]).Equal(Unknown)
	sim.Connect(Matrix, 3)
	assert(sim.modules[2]).Equal(Matrix)
	time.Sleep(100 * time.Millisecond)
	n, _ := e2.Read(buffer)
	assert(string(buffer[:n])).Equal("c 3/matrix\r\n")

	sim.Disconnect(3)
	assert(sim.modules[2]).Equal(Unknown)
	time.Sleep(100 * time.Millisecond)
	n, _ = e2.Read(buffer)
	assert(string(buffer[:n])).Equal("d 3\r\n")

	e2.Close()
}
<commit_msg>fix test<commit_after>package dock

import (
	"fmt"
	"io"
	"testing"
	"time"

	"github.com/simulatedsimian/assert"
)

func TestParseRequest(t *testing.T) {
	assert := assert.Make(t)

	assert(msgToRequest("e")).Equal(Request{RequestType: ReqEnquire})
	assert(msgToRequest("r")).Equal(Request{RequestType: ReqResetToBootloader})
	assert(msgToRequest("v")).Equal(Request{RequestType: ReqVersion})
	assert(msgToRequest("d")).Equal(Request{RequestType: ReqDebug})
	assert(msgToRequest("p 0")).Equal(Request{RequestType: ReqPower, Params: []int{0}})
	assert(msgToRequest("p 1")).Equal(Request{RequestType: ReqPower, Params: []int{1}})
	assert(msgToRequest("p")).Equal(Request{})
	assert(msgToRequest("p 3")).Equal(Request{})

	assert(msgToRequest("n d hello")).Equal(Request{RequestType: ReqName, Params: []int{'d'}, ParamStr: "hello"})
	assert(msgToRequest("n u world")).Equal(Request{RequestType: ReqName, Params: []int{'u'}, ParamStr: "world"})

	assert(msgToRequest("s 5 1,2,3")).Equal(Request{RequestType: ReqSet, Channel: 5, Params: []int{1, 2, 3}})
}

func TestSimulatorRequest(t *testing.T) {
	assert := assert.Make(t)

	e1, e2 := NewPipe().Endpoints()
	sim := NewSimulator(e1)

	fmt.Fprint(e2, "e\r")
	assert(<-sim.Requests).Equal(Request{RequestType: ReqEnquire})
	e2.Close()
	assert(<-sim.Requests).Equal(Request{RequestType: ReqError, Error: io.EOF})

}

func TestSimulatorConnectDisconnect(t *testing.T) {
	assert := assert.Make(t)

	buffer := make([]byte, 128)
	e1, e2 := NewPipe().Endpoints()
	sim := NewSimulator(e1)

	assert(sim.modules[2]).Equal(Unknown)
	sim.Connect(Matrix, 3)
	assert(sim.modules[2]).Equal(Matrix)
	time.Sleep(100 * time.Millisecond)
	n, _ := e2.Read(buffer)
	assert(string(buffer[:n])).Equal("c 3/matrix\r\n")

	sim.Disconnect(3)
	assert(sim.modules[2]).Equal(Unknown)
	time.Sleep(100 * time.Millisecond)
	n, _ = e2.Read(buffer)
	assert(string(buffer[:n])).Equal("d 3/matrix\r\n")

	e2.Close()
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"net/http"
	"os"

	"github.com/Sirupsen/logrus"
	"github.com/docker/distribution"
	"github.com/docker/distribution/configuration"
	"github.com/docker/distribution/context"
	"github.com/docker/distribution/digest"
	"github.com/docker/distribution/manifest"
	"github.com/docker/distribution/registry/handlers"
	"github.com/docker/distribution/registry/middleware/repository"
	"github.com/flynn/flynn/controller/client"
	ct "github.com/flynn/flynn/controller/types"
	"github.com/flynn/flynn/docker-receive/blobstore"
	"github.com/flynn/flynn/host/types"
	"github.com/flynn/flynn/pkg/status"
	"github.com/flynn/flynn/pkg/version"
)

// main is a modified version of the registry main function:
// https://github.com/docker/distribution/blob/6ba799b/cmd/registry/main.go
func main() {
	logrus.SetLevel(logrus.InfoLevel)

	ctx := context.Background()
	ctx = context.WithValue(ctx, "version", version.String())
	ctx = context.WithLogger(ctx, context.GetLogger(ctx, "version"))

	client, err := controller.NewClient("", os.Getenv("CONTROLLER_KEY"))
	if err != nil {
		context.GetLogger(ctx).Fatalln(err)
	}

	authKey := os.Getenv("AUTH_KEY")

	middleware.Register("flynn", repositoryMiddleware(client, authKey))

	config := configuration.Configuration{
		Version: configuration.CurrentVersion,
		Storage: configuration.Storage{
			blobstore.DriverName: configuration.Parameters{},
		},
		Middleware: map[string][]configuration.Middleware{
			"repository": {
				{Name: "flynn"},
			},
		},
		Auth: configuration.Auth{
			"flynn": configuration.Parameters{
				"auth_key": authKey,
			},
		},
	}
	config.HTTP.Secret = os.Getenv("REGISTRY_HTTP_SECRET")

	status.AddHandler(status.HealthyHandler)

	app := handlers.NewApp(ctx, config)
	http.Handle("/", app)

	addr := ":" + os.Getenv("PORT")
	context.GetLogger(app).Infof("listening on %s", addr)
	if err := http.ListenAndServe(addr, nil); err != nil {
		context.GetLogger(app).Fatalln(err)
	}
}

func repositoryMiddleware(client controller.Client, authKey string) middleware.InitFunc {
	return func(ctx context.Context, r distribution.Repository, _ map[string]interface{}) (distribution.Repository, error) {
		return &repository{
			Repository: r,
			client:     client,
			authKey:    authKey,
		}, nil
	}
}

// repository is a repository middleware which returns a custom ManifestService
// in order to create Flynn artifacts when image manifests are pushed
type repository struct {
	distribution.Repository

	client  controller.Client
	authKey string
}

func (r *repository) Manifests(ctx context.Context, options ...distribution.ManifestServiceOption) (distribution.ManifestService, error) {
	m, err := r.Repository.Manifests(ctx, options...)
	if err != nil {
		return nil, err
	}
	return &manifestService{
		ManifestService: m,
		repository:      r,
		client:          r.client,
		authKey:         r.authKey,
	}, nil
}

type manifestService struct {
	distribution.ManifestService

	repository distribution.Repository
	client     controller.Client
	authKey    string
}

func (m *manifestService) Put(manifest *manifest.SignedManifest) error {
	if err := m.ManifestService.Put(manifest); err != nil {
		return err
	}

	dgst, err := digestManifest(manifest)
	if err != nil {
		return err
	}

	return m.createArtifact(dgst)
}

func (m *manifestService) createArtifact(dgst digest.Digest) error {
	return m.client.CreateArtifact(&ct.Artifact{
		Type: host.ArtifactTypeDocker,
		URI:  fmt.Sprintf("http://flynn:%s@docker-receive.discoverd?name=%s&id=%s", m.authKey, m.repository.Name(), dgst),
		Meta: map[string]string{
			"docker-receive.repository": m.repository.Name(),
			"docker-receive.digest":     string(dgst),
		},
	})
}

// digestManifest is a modified version of:
// https://github.com/docker/distribution/blob/6ba799b/registry/handlers/images.go#L228-L251
func digestManifest(manifest *manifest.SignedManifest) (digest.Digest, error) {
	p, err := manifest.Payload()
	if err != nil {
		return "", err
	}
	return digest.FromBytes(p)
}
<commit_msg>docker-receive: Support deleting images<commit_after>package main

import (
	"fmt"
	"net/http"
	"os"

	"github.com/Sirupsen/logrus"
	"github.com/docker/distribution"
	"github.com/docker/distribution/configuration"
	"github.com/docker/distribution/context"
	"github.com/docker/distribution/digest"
	"github.com/docker/distribution/manifest"
	"github.com/docker/distribution/registry/handlers"
	"github.com/docker/distribution/registry/middleware/repository"
	"github.com/flynn/flynn/controller/client"
	ct "github.com/flynn/flynn/controller/types"
	"github.com/flynn/flynn/docker-receive/blobstore"
	"github.com/flynn/flynn/host/types"
	"github.com/flynn/flynn/pkg/status"
	"github.com/flynn/flynn/pkg/version"
)

// main is a modified version of the registry main function:
// https://github.com/docker/distribution/blob/6ba799b/cmd/registry/main.go
func main() {
	logrus.SetLevel(logrus.InfoLevel)

	ctx := context.Background()
	ctx = context.WithValue(ctx, "version", version.String())
	ctx = context.WithLogger(ctx, context.GetLogger(ctx, "version"))

	client, err := controller.NewClient("", os.Getenv("CONTROLLER_KEY"))
	if err != nil {
		context.GetLogger(ctx).Fatalln(err)
	}

	authKey := os.Getenv("AUTH_KEY")

	middleware.Register("flynn", repositoryMiddleware(client, authKey))

	config := configuration.Configuration{
		Version: configuration.CurrentVersion,
		Storage: configuration.Storage{
			blobstore.DriverName: configuration.Parameters{},
			"delete":             configuration.Parameters{"enabled": true},
		},
		Middleware: map[string][]configuration.Middleware{
			"repository": {
				{Name: "flynn"},
			},
		},
		Auth: configuration.Auth{
			"flynn": configuration.Parameters{
				"auth_key": authKey,
			},
		},
	}
	config.HTTP.Secret = os.Getenv("REGISTRY_HTTP_SECRET")

	status.AddHandler(status.HealthyHandler)

	app := handlers.NewApp(ctx, config)
	http.Handle("/", app)

	addr := ":" + os.Getenv("PORT")
	context.GetLogger(app).Infof("listening on %s", addr)
	if err := http.ListenAndServe(addr, nil); err != nil {
		context.GetLogger(app).Fatalln(err)
	}
}

func repositoryMiddleware(client controller.Client, authKey string) middleware.InitFunc {
	return func(ctx context.Context, r distribution.Repository, _ map[string]interface{}) (distribution.Repository, error) {
		return &repository{
			Repository: r,
			client:     client,
			authKey:    authKey,
		}, nil
	}
}

// repository is a repository middleware which returns a custom ManifestService
// in order to create Flynn artifacts when image manifests are pushed
type repository struct {
	distribution.Repository

	client  controller.Client
	authKey string
}

func (r *repository) Manifests(ctx context.Context, options ...distribution.ManifestServiceOption) (distribution.ManifestService, error) {
	m, err := r.Repository.Manifests(ctx, options...)
	if err != nil {
		return nil, err
	}
	return &manifestService{
		ManifestService: m,
		repository:      r,
		client:          r.client,
		authKey:         r.authKey,
	}, nil
}

type manifestService struct {
	distribution.ManifestService

	repository distribution.Repository
	client     controller.Client
	authKey    string
}

func (m *manifestService) Put(manifest *manifest.SignedManifest) error {
	if err := m.ManifestService.Put(manifest); err != nil {
		return err
	}

	dgst, err := digestManifest(manifest)
	if err != nil {
		return err
	}

	return m.createArtifact(dgst)
}

func (m *manifestService) createArtifact(dgst digest.Digest) error {
	return m.client.CreateArtifact(&ct.Artifact{
		Type: host.ArtifactTypeDocker,
		URI:  fmt.Sprintf("http://flynn:%s@docker-receive.discoverd?name=%s&id=%s", m.authKey, m.repository.Name(), dgst),
		Meta: map[string]string{
			"docker-receive.repository": m.repository.Name(),
			"docker-receive.digest":     string(dgst),
		},
	})
}

// digestManifest is a modified version of:
// https://github.com/docker/distribution/blob/6ba799b/registry/handlers/images.go#L228-L251
func digestManifest(manifest *manifest.SignedManifest) (digest.Digest, error) {
	p, err := manifest.Payload()
	if err != nil {
		return "", err
	}
	return digest.FromBytes(p)
}
<|endoftext|>


<commit_before>// Code generated by protoc-gen-go.
// source: helloworld.proto
// DO NOT EDIT!

/*
Package helloworld is a generated protocol buffer package.

It is generated from these files:
	helloworld.proto

It has these top-level messages:
	HelloRequest
	HelloReply
*/
package helloworld

import proto "github.com/golang/protobuf/proto"

import (
	context "golang.org/x/net/context"
	grpc "google.golang.org/grpc"
)

// Reference imports to suppress errors if they are not otherwise used.
var _ context.Context
var _ grpc.ClientConn

// Reference imports to suppress errors if they are not otherwise used.
var _ = proto.Marshal

// The request message containing the user's name.
type HelloRequest struct {
	Name string `protobuf:"bytes,1,opt,name=name" json:"name,omitempty"`
}

func (m *HelloRequest) Reset()         { *m = HelloRequest{} }
func (m *HelloRequest) String() string { return proto.CompactTextString(m) }
func (*HelloRequest) ProtoMessage()    {}

// The response message containing the greetings
type HelloReply struct {
	Message string `protobuf:"bytes,1,opt,name=message" json:"message,omitempty"`
}

func (m *HelloReply) Reset()         { *m = HelloReply{} }
func (m *HelloReply) String() string { return proto.CompactTextString(m) }
func (*HelloReply) ProtoMessage()    {}

func init() {
}

// Client API for Greeter service

type GreeterClient interface {
	// Sends a greeting
	SayHello(ctx context.Context, in *HelloRequest, opts ...grpc.CallOption) (*HelloReply, error)
}

type greeterClient struct {
	cc *grpc.ClientConn
}

func NewGreeterClient(cc *grpc.ClientConn) GreeterClient {
	return &greeterClient{cc}
}

func (c *greeterClient) SayHello(ctx context.Context, in *HelloRequest, opts ...grpc.CallOption) (*HelloReply, error) {
	out := new(HelloReply)
	err := grpc.Invoke(ctx, "/helloworld.Greeter/SayHello", in, out, c.cc, opts...)
	if err != nil {
		return nil, err
	}
	return out, nil
}

// Server API for Greeter service

type GreeterServer interface {
	// Sends a greeting
	SayHello(context.Context, *HelloRequest) (*HelloReply, error)
}

func RegisterGreeterServer(s *grpc.Server, srv GreeterServer) {
	s.RegisterService(&_Greeter_serviceDesc, srv)
}

func _Greeter_SayHello_Handler(srv interface{}, ctx context.Context, buf []byte) (interface{}, error) {
	in := new(HelloRequest)
	if err := proto.Unmarshal(buf, in); err != nil {
		return nil, err
	}
	out, err := srv.(GreeterServer).SayHello(ctx, in)
	if err != nil {
		return nil, err
	}
	return out, nil
}

var _Greeter_serviceDesc = grpc.ServiceDesc{
	ServiceName: "helloworld.Greeter",
	HandlerType: (*GreeterServer)(nil),
	Methods: []grpc.MethodDesc{
		{
			MethodName: "SayHello",
			Handler:    _Greeter_SayHello_Handler,
		},
	},
	Streams: []grpc.StreamDesc{},
}
<commit_msg>update the generated code<commit_after>// Code generated by protoc-gen-go.
// source: helloworld.proto
// DO NOT EDIT!

/*
Package helloworld is a generated protocol buffer package.

It is generated from these files:
	helloworld.proto

It has these top-level messages:
	HelloRequest
	HelloReply
*/
package helloworld

import proto "github.com/golang/protobuf/proto"

import (
	context "golang.org/x/net/context"
	grpc "google.golang.org/grpc"
)

// Reference imports to suppress errors if they are not otherwise used.
var _ context.Context
var _ grpc.ClientConn

// Reference imports to suppress errors if they are not otherwise used.
var _ = proto.Marshal

// The request message containing the user's name.
type HelloRequest struct {
	Name string `protobuf:"bytes,1,opt,name=name" json:"name,omitempty"`
}

func (m *HelloRequest) Reset()         { *m = HelloRequest{} }
func (m *HelloRequest) String() string { return proto.CompactTextString(m) }
func (*HelloRequest) ProtoMessage()    {}

// The response message containing the greetings
type HelloReply struct {
	Message string `protobuf:"bytes,1,opt,name=message" json:"message,omitempty"`
}

func (m *HelloReply) Reset()         { *m = HelloReply{} }
func (m *HelloReply) String() string { return proto.CompactTextString(m) }
func (*HelloReply) ProtoMessage()    {}

func init() {
}

// Client API for Greeter service

type GreeterClient interface {
	// Sends a greeting
	SayHello(ctx context.Context, in *HelloRequest, opts ...grpc.CallOption) (*HelloReply, error)
}

type greeterClient struct {
	cc *grpc.ClientConn
}

func NewGreeterClient(cc *grpc.ClientConn) GreeterClient {
	return &greeterClient{cc}
}

func (c *greeterClient) SayHello(ctx context.Context, in *HelloRequest, opts ...grpc.CallOption) (*HelloReply, error) {
	out := new(HelloReply)
	err := grpc.Invoke(ctx, "/helloworld.Greeter/SayHello", in, out, c.cc, opts...)
	if err != nil {
		return nil, err
	}
	return out, nil
}

// Server API for Greeter service

type GreeterServer interface {
	// Sends a greeting
	SayHello(context.Context, *HelloRequest) (*HelloReply, error)
}

func RegisterGreeterServer(s *grpc.Server, srv GreeterServer) {
	s.RegisterService(&_Greeter_serviceDesc, srv)
}

func _Greeter_SayHello_Handler(srv interface{}, ctx context.Context, codec grpc.Codec, buf []byte) (interface{}, error) {
	in := new(HelloRequest)
	if err := codec.Unmarshal(buf, in); err != nil {
		return nil, err
	}
	out, err := srv.(GreeterServer).SayHello(ctx, in)
	if err != nil {
		return nil, err
	}
	return out, nil
}

var _Greeter_serviceDesc = grpc.ServiceDesc{
	ServiceName: "helloworld.Greeter",
	HandlerType: (*GreeterServer)(nil),
	Methods: []grpc.MethodDesc{
		{
			MethodName: "SayHello",
			Handler:    _Greeter_SayHello_Handler,
		},
	},
	Streams: []grpc.StreamDesc{},
}
<|endoftext|>


<commit_before>package docker

import (
	"errors"
	"fmt"
	"net/url"
	"os"
	"path"
	"path/filepath"
	"strings"

	"github.com/containers/image/v5/docker/reference"
	"github.com/containers/image/v5/internal/rootless"
	"github.com/containers/image/v5/types"
	"github.com/containers/storage/pkg/homedir"
	"github.com/ghodss/yaml"
	"github.com/opencontainers/go-digest"
	perrors "github.com/pkg/errors"
	"github.com/sirupsen/logrus"
)

// systemRegistriesDirPath is the path to registries.d, used for locating lookaside Docker signature storage.
// You can override this at build time with
// -ldflags '-X github.com/containers/image/v5/docker.systemRegistriesDirPath=$your_path'
var systemRegistriesDirPath = builtinRegistriesDirPath

// builtinRegistriesDirPath is the path to registries.d.
// DO NOT change this, instead see systemRegistriesDirPath above.
const builtinRegistriesDirPath = etcDir + "/containers/registries.d"

// userRegistriesDirPath is the path to the per user registries.d.
var userRegistriesDir = filepath.FromSlash(".config/containers/registries.d")

// defaultUserDockerDir is the default sigstore directory for unprivileged user
var defaultUserDockerDir = filepath.FromSlash(".local/share/containers/sigstore")

// defaultDockerDir is the default sigstore directory for root
var defaultDockerDir = "/var/lib/containers/sigstore"

// registryConfiguration is one of the files in registriesDirPath configuring lookaside locations, or the result of merging them all.
// NOTE: Keep this in sync with docs/registries.d.md!
type registryConfiguration struct {
	DefaultDocker *registryNamespace `json:"default-docker"`
	// The key is a namespace, using fully-expanded Docker reference format or parent namespaces (per dockerReference.PolicyConfiguration*),
	Docker map[string]registryNamespace `json:"docker"`
}

// registryNamespace defines lookaside locations for a single namespace.
type registryNamespace struct {
	SigStore        string `json:"sigstore"`         // For reading, and if SigStoreStaging is not present, for writing.
	SigStoreStaging string `json:"sigstore-staging"` // For writing only.
}

// signatureStorageBase is an "opaque" type representing a lookaside Docker signature storage.
// Users outside of this file should use SignatureStorageBaseURL and signatureStorageURL below.
type signatureStorageBase *url.URL

// SignatureStorageBaseURL reads configuration to find an appropriate signature storage URL for ref, for write access if “write”.
// the usage of the BaseURL is defined under docker/distribution registries—separate storage of docs/signature-protocols.md
// Warning: This function only exposes configuration in registries.d;
// just because this function returns an URL does not mean that the URL will be used by c/image/docker (e.g. if the registry natively supports X-R-S-S).
func SignatureStorageBaseURL(sys *types.SystemContext, ref types.ImageReference, write bool) (*url.URL, error) {
	dr, ok := ref.(dockerReference)
	if !ok {
		return nil, errors.New("ref must be a dockerReference")
	}
	// FIXME? Loading and parsing the config could be cached across calls.
	dirPath := registriesDirPath(sys)
	logrus.Debugf(`Using registries.d directory %s for sigstore configuration`, dirPath)
	config, err := loadAndMergeConfig(dirPath)
	if err != nil {
		return nil, err
	}

	topLevel := config.signatureTopLevel(dr, write)
	var url *url.URL
	if topLevel != "" {
		url, err = url.Parse(topLevel)
		if err != nil {
			return nil, perrors.Wrapf(err, "Invalid signature storage URL %s", topLevel)
		}
	} else {
		// returns default directory if no sigstore specified in configuration file
		url = builtinDefaultSignatureStorageDir(rootless.GetRootlessEUID())
		logrus.Debugf(" No signature storage configuration found for %s, using built-in default %s", dr.PolicyConfigurationIdentity(), url.Redacted())
	}
	// NOTE: Keep this in sync with docs/signature-protocols.md!
	// FIXME? Restrict to explicitly supported schemes?
	repo := reference.Path(dr.ref) // Note that this is without a tag or digest.
	if path.Clean(repo) != repo {  // Coverage: This should not be reachable because /./ and /../ components are not valid in docker references
		return nil, fmt.Errorf("Unexpected path elements in Docker reference %s for signature storage", dr.ref.String())
	}
	url.Path = url.Path + "/" + repo
	return url, nil
}

// registriesDirPath returns a path to registries.d
func registriesDirPath(sys *types.SystemContext) string {
	return registriesDirPathWithHomeDir(sys, homedir.Get())
}

// registriesDirPathWithHomeDir is an internal implementation detail of registriesDirPath,
// it exists only to allow testing it with an artificial home directory.
func registriesDirPathWithHomeDir(sys *types.SystemContext, homeDir string) string {
	if sys != nil && sys.RegistriesDirPath != "" {
		return sys.RegistriesDirPath
	}
	userRegistriesDirPath := filepath.Join(homeDir, userRegistriesDir)
	if _, err := os.Stat(userRegistriesDirPath); err == nil {
		return userRegistriesDirPath
	}
	if sys != nil && sys.RootForImplicitAbsolutePaths != "" {
		return filepath.Join(sys.RootForImplicitAbsolutePaths, systemRegistriesDirPath)
	}

	return systemRegistriesDirPath
}

// builtinDefaultSignatureStorageDir returns default signature storage URL as per euid
func builtinDefaultSignatureStorageDir(euid int) *url.URL {
	if euid != 0 {
		return &url.URL{Scheme: "file", Path: filepath.Join(homedir.Get(), defaultUserDockerDir)}
	}
	return &url.URL{Scheme: "file", Path: defaultDockerDir}
}

// loadAndMergeConfig loads configuration files in dirPath
func loadAndMergeConfig(dirPath string) (*registryConfiguration, error) {
	mergedConfig := registryConfiguration{Docker: map[string]registryNamespace{}}
	dockerDefaultMergedFrom := ""
	nsMergedFrom := map[string]string{}

	dir, err := os.Open(dirPath)
	if err != nil {
		if os.IsNotExist(err) {
			return &mergedConfig, nil
		}
		return nil, err
	}
	configNames, err := dir.Readdirnames(0)
	if err != nil {
		return nil, err
	}
	for _, configName := range configNames {
		if !strings.HasSuffix(configName, ".yaml") {
			continue
		}
		configPath := filepath.Join(dirPath, configName)
		configBytes, err := os.ReadFile(configPath)
		if err != nil {
			return nil, err
		}

		var config registryConfiguration
		err = yaml.Unmarshal(configBytes, &config)
		if err != nil {
			return nil, perrors.Wrapf(err, "parsing %s", configPath)
		}

		if config.DefaultDocker != nil {
			if mergedConfig.DefaultDocker != nil {
				return nil, fmt.Errorf(`Error parsing signature storage configuration: "default-docker" defined both in "%s" and "%s"`,
					dockerDefaultMergedFrom, configPath)
			}
			mergedConfig.DefaultDocker = config.DefaultDocker
			dockerDefaultMergedFrom = configPath
		}

		for nsName, nsConfig := range config.Docker { // includes config.Docker == nil
			if _, ok := mergedConfig.Docker[nsName]; ok {
				return nil, fmt.Errorf(`Error parsing signature storage configuration: "docker" namespace "%s" defined both in "%s" and "%s"`,
					nsName, nsMergedFrom[nsName], configPath)
			}
			mergedConfig.Docker[nsName] = nsConfig
			nsMergedFrom[nsName] = configPath
		}
	}

	return &mergedConfig, nil
}

// config.signatureTopLevel returns an URL string configured in config for ref, for write access if “write”.
// (the top level of the storage, namespaced by repo.FullName etc.), or "" if nothing has been configured.
func (config *registryConfiguration) signatureTopLevel(ref dockerReference, write bool) string {
	if config.Docker != nil {
		// Look for a full match.
		identity := ref.PolicyConfigurationIdentity()
		if ns, ok := config.Docker[identity]; ok {
			logrus.Debugf(` Using "docker" namespace %s`, identity)
			if url := ns.signatureTopLevel(write); url != "" {
				return url
			}
		}

		// Look for a match of the possible parent namespaces.
		for _, name := range ref.PolicyConfigurationNamespaces() {
			if ns, ok := config.Docker[name]; ok {
				logrus.Debugf(` Using "docker" namespace %s`, name)
				if url := ns.signatureTopLevel(write); url != "" {
					return url
				}
			}
		}
	}
	// Look for a default location
	if config.DefaultDocker != nil {
		logrus.Debugf(` Using "default-docker" configuration`)
		if url := config.DefaultDocker.signatureTopLevel(write); url != "" {
			return url
		}
	}
	return ""
}

// ns.signatureTopLevel returns an URL string configured in ns for ref, for write access if “write”.
// or "" if nothing has been configured.
func (ns registryNamespace) signatureTopLevel(write bool) string {
	if write && ns.SigStoreStaging != "" {
		logrus.Debugf(`  Using %s`, ns.SigStoreStaging)
		return ns.SigStoreStaging
	}
	if ns.SigStore != "" {
		logrus.Debugf(`  Using %s`, ns.SigStore)
		return ns.SigStore
	}
	return ""
}

// signatureStorageURL returns an URL usable for accessing signature index in base with known manifestDigest.
// base is not nil from the caller
// NOTE: Keep this in sync with docs/signature-protocols.md!
func signatureStorageURL(base signatureStorageBase, manifestDigest digest.Digest, index int) *url.URL {
	url := *base
	url.Path = fmt.Sprintf("%s@%s=%s/signature-%d", url.Path, manifestDigest.Algorithm(), manifestDigest.Hex(), index+1)
	return &url
}
<commit_msg>Split registryConfiguration.signatureStorageBaseURL from SignatureStorageBaseURL<commit_after>package docker

import (
	"errors"
	"fmt"
	"net/url"
	"os"
	"path"
	"path/filepath"
	"strings"

	"github.com/containers/image/v5/docker/reference"
	"github.com/containers/image/v5/internal/rootless"
	"github.com/containers/image/v5/types"
	"github.com/containers/storage/pkg/homedir"
	"github.com/ghodss/yaml"
	"github.com/opencontainers/go-digest"
	perrors "github.com/pkg/errors"
	"github.com/sirupsen/logrus"
)

// systemRegistriesDirPath is the path to registries.d, used for locating lookaside Docker signature storage.
// You can override this at build time with
// -ldflags '-X github.com/containers/image/v5/docker.systemRegistriesDirPath=$your_path'
var systemRegistriesDirPath = builtinRegistriesDirPath

// builtinRegistriesDirPath is the path to registries.d.
// DO NOT change this, instead see systemRegistriesDirPath above.
const builtinRegistriesDirPath = etcDir + "/containers/registries.d"

// userRegistriesDirPath is the path to the per user registries.d.
var userRegistriesDir = filepath.FromSlash(".config/containers/registries.d")

// defaultUserDockerDir is the default sigstore directory for unprivileged user
var defaultUserDockerDir = filepath.FromSlash(".local/share/containers/sigstore")

// defaultDockerDir is the default sigstore directory for root
var defaultDockerDir = "/var/lib/containers/sigstore"

// registryConfiguration is one of the files in registriesDirPath configuring lookaside locations, or the result of merging them all.
// NOTE: Keep this in sync with docs/registries.d.md!
type registryConfiguration struct {
	DefaultDocker *registryNamespace `json:"default-docker"`
	// The key is a namespace, using fully-expanded Docker reference format or parent namespaces (per dockerReference.PolicyConfiguration*),
	Docker map[string]registryNamespace `json:"docker"`
}

// registryNamespace defines lookaside locations for a single namespace.
type registryNamespace struct {
	SigStore        string `json:"sigstore"`         // For reading, and if SigStoreStaging is not present, for writing.
	SigStoreStaging string `json:"sigstore-staging"` // For writing only.
}

// signatureStorageBase is an "opaque" type representing a lookaside Docker signature storage.
// Users outside of this file should use SignatureStorageBaseURL and signatureStorageURL below.
type signatureStorageBase *url.URL

// SignatureStorageBaseURL reads configuration to find an appropriate signature storage URL for ref, for write access if “write”.
// the usage of the BaseURL is defined under docker/distribution registries—separate storage of docs/signature-protocols.md
// Warning: This function only exposes configuration in registries.d;
// just because this function returns an URL does not mean that the URL will be used by c/image/docker (e.g. if the registry natively supports X-R-S-S).
func SignatureStorageBaseURL(sys *types.SystemContext, ref types.ImageReference, write bool) (*url.URL, error) {
	dr, ok := ref.(dockerReference)
	if !ok {
		return nil, errors.New("ref must be a dockerReference")
	}
	// FIXME? Loading and parsing the config could be cached across calls.
	dirPath := registriesDirPath(sys)
	logrus.Debugf(`Using registries.d directory %s for sigstore configuration`, dirPath)
	config, err := loadAndMergeConfig(dirPath)
	if err != nil {
		return nil, err
	}

	return config.signatureStorageBaseURL(dr, write)
}

// registriesDirPath returns a path to registries.d
func registriesDirPath(sys *types.SystemContext) string {
	return registriesDirPathWithHomeDir(sys, homedir.Get())
}

// registriesDirPathWithHomeDir is an internal implementation detail of registriesDirPath,
// it exists only to allow testing it with an artificial home directory.
func registriesDirPathWithHomeDir(sys *types.SystemContext, homeDir string) string {
	if sys != nil && sys.RegistriesDirPath != "" {
		return sys.RegistriesDirPath
	}
	userRegistriesDirPath := filepath.Join(homeDir, userRegistriesDir)
	if _, err := os.Stat(userRegistriesDirPath); err == nil {
		return userRegistriesDirPath
	}
	if sys != nil && sys.RootForImplicitAbsolutePaths != "" {
		return filepath.Join(sys.RootForImplicitAbsolutePaths, systemRegistriesDirPath)
	}

	return systemRegistriesDirPath
}

// loadAndMergeConfig loads configuration files in dirPath
func loadAndMergeConfig(dirPath string) (*registryConfiguration, error) {
	mergedConfig := registryConfiguration{Docker: map[string]registryNamespace{}}
	dockerDefaultMergedFrom := ""
	nsMergedFrom := map[string]string{}

	dir, err := os.Open(dirPath)
	if err != nil {
		if os.IsNotExist(err) {
			return &mergedConfig, nil
		}
		return nil, err
	}
	configNames, err := dir.Readdirnames(0)
	if err != nil {
		return nil, err
	}
	for _, configName := range configNames {
		if !strings.HasSuffix(configName, ".yaml") {
			continue
		}
		configPath := filepath.Join(dirPath, configName)
		configBytes, err := os.ReadFile(configPath)
		if err != nil {
			return nil, err
		}

		var config registryConfiguration
		err = yaml.Unmarshal(configBytes, &config)
		if err != nil {
			return nil, perrors.Wrapf(err, "parsing %s", configPath)
		}

		if config.DefaultDocker != nil {
			if mergedConfig.DefaultDocker != nil {
				return nil, fmt.Errorf(`Error parsing signature storage configuration: "default-docker" defined both in "%s" and "%s"`,
					dockerDefaultMergedFrom, configPath)
			}
			mergedConfig.DefaultDocker = config.DefaultDocker
			dockerDefaultMergedFrom = configPath
		}

		for nsName, nsConfig := range config.Docker { // includes config.Docker == nil
			if _, ok := mergedConfig.Docker[nsName]; ok {
				return nil, fmt.Errorf(`Error parsing signature storage configuration: "docker" namespace "%s" defined both in "%s" and "%s"`,
					nsName, nsMergedFrom[nsName], configPath)
			}
			mergedConfig.Docker[nsName] = nsConfig
			nsMergedFrom[nsName] = configPath
		}
	}

	return &mergedConfig, nil
}

// signatureStorageBaseURL returns an appropriate signature storage URL for ref, for write access if “write”.
// the usage of the BaseURL is defined under docker/distribution registries—separate storage of docs/signature-protocols.md
func (config *registryConfiguration) signatureStorageBaseURL(dr dockerReference, write bool) (*url.URL, error) {
	topLevel := config.signatureTopLevel(dr, write)
	var url *url.URL
	if topLevel != "" {
		u, err := url.Parse(topLevel)
		if err != nil {
			return nil, perrors.Wrapf(err, "Invalid signature storage URL %s", topLevel)
		}
		url = u
	} else {
		// returns default directory if no sigstore specified in configuration file
		url = builtinDefaultSignatureStorageDir(rootless.GetRootlessEUID())
		logrus.Debugf(" No signature storage configuration found for %s, using built-in default %s", dr.PolicyConfigurationIdentity(), url.Redacted())
	}
	// NOTE: Keep this in sync with docs/signature-protocols.md!
	// FIXME? Restrict to explicitly supported schemes?
	repo := reference.Path(dr.ref) // Note that this is without a tag or digest.
	if path.Clean(repo) != repo {  // Coverage: This should not be reachable because /./ and /../ components are not valid in docker references
		return nil, fmt.Errorf("Unexpected path elements in Docker reference %s for signature storage", dr.ref.String())
	}
	url.Path = url.Path + "/" + repo
	return url, nil
}

// builtinDefaultSignatureStorageDir returns default signature storage URL as per euid
func builtinDefaultSignatureStorageDir(euid int) *url.URL {
	if euid != 0 {
		return &url.URL{Scheme: "file", Path: filepath.Join(homedir.Get(), defaultUserDockerDir)}
	}
	return &url.URL{Scheme: "file", Path: defaultDockerDir}
}

// config.signatureTopLevel returns an URL string configured in config for ref, for write access if “write”.
// (the top level of the storage, namespaced by repo.FullName etc.), or "" if nothing has been configured.
func (config *registryConfiguration) signatureTopLevel(ref dockerReference, write bool) string {
	if config.Docker != nil {
		// Look for a full match.
		identity := ref.PolicyConfigurationIdentity()
		if ns, ok := config.Docker[identity]; ok {
			logrus.Debugf(` Using "docker" namespace %s`, identity)
			if url := ns.signatureTopLevel(write); url != "" {
				return url
			}
		}

		// Look for a match of the possible parent namespaces.
		for _, name := range ref.PolicyConfigurationNamespaces() {
			if ns, ok := config.Docker[name]; ok {
				logrus.Debugf(` Using "docker" namespace %s`, name)
				if url := ns.signatureTopLevel(write); url != "" {
					return url
				}
			}
		}
	}
	// Look for a default location
	if config.DefaultDocker != nil {
		logrus.Debugf(` Using "default-docker" configuration`)
		if url := config.DefaultDocker.signatureTopLevel(write); url != "" {
			return url
		}
	}
	return ""
}

// ns.signatureTopLevel returns an URL string configured in ns for ref, for write access if “write”.
// or "" if nothing has been configured.
func (ns registryNamespace) signatureTopLevel(write bool) string {
	if write && ns.SigStoreStaging != "" {
		logrus.Debugf(`  Using %s`, ns.SigStoreStaging)
		return ns.SigStoreStaging
	}
	if ns.SigStore != "" {
		logrus.Debugf(`  Using %s`, ns.SigStore)
		return ns.SigStore
	}
	return ""
}

// signatureStorageURL returns an URL usable for accessing signature index in base with known manifestDigest.
// base is not nil from the caller
// NOTE: Keep this in sync with docs/signature-protocols.md!
func signatureStorageURL(base signatureStorageBase, manifestDigest digest.Digest, index int) *url.URL {
	url := *base
	url.Path = fmt.Sprintf("%s@%s=%s/signature-%d", url.Path, manifestDigest.Algorithm(), manifestDigest.Hex(), index+1)
	return &url
}
<|endoftext|>