<commit_before>package sctp

import (
	"math"
	"sync"

	"github.com/pkg/errors"
)

const (
	// ReliabilityTypeReliable is used for reliable transmission
	ReliabilityTypeReliable byte = 0
	// ReliabilityTypeRexmit is used for partial reliability by retransmission count
	ReliabilityTypeRexmit byte = 1
	// ReliabilityTypeTimed is used for partial reliability by retransmission duration
	ReliabilityTypeTimed byte = 2
)

// Stream represents an SCTP stream
type Stream struct {
	association *Association

	lock sync.RWMutex

	streamIdentifier   uint16
	defaultPayloadType PayloadProtocolIdentifier

	reassemblyQueue *reassemblyQueue
	sequenceNumber  uint16

	readNotifier *sync.Cond

	readErr  error
	writeErr error

	unordered        bool
	reliabilityType  byte
	reliabilityValue uint32
}

// StreamIdentifier returns the Stream identifier associated to the stream.
func (s *Stream) StreamIdentifier() uint16 {
	s.lock.RLock()
	defer s.lock.RUnlock()
	return s.streamIdentifier
}

// SetDefaultPayloadType sets the default payload type used by Write.
func (s *Stream) SetDefaultPayloadType(defaultPayloadType PayloadProtocolIdentifier) {
	s.lock.Lock()
	defer s.lock.Unlock()

	s.setDefaultPayloadType(defaultPayloadType)
}

// setDefaultPayloadType sets the defaultPayloadType. The caller should hold the lock.
func (s *Stream) setDefaultPayloadType(defaultPayloadType PayloadProtocolIdentifier) {
	s.defaultPayloadType = defaultPayloadType
}

// SetReliabilityParams sets reliability parameters for this stream.
func (s *Stream) SetReliabilityParams(unordered bool, relType byte, relVal uint32) {
	s.lock.Lock()
	defer s.lock.Unlock()

	s.setReliabilityParams(unordered, relType, relVal)
}

// setReliabilityParams sets reliability parameters for this stream.
// The caller should hold the lock.
func (s *Stream) setReliabilityParams(unordered bool, relType byte, relVal uint32) {
	s.unordered = unordered
	s.reliabilityType = relType
	s.reliabilityValue = relVal
}

// Read reads a packet of len(p) bytes, dropping the Payload Protocol Identifier.
// Returns EOF when the stream is reset or an error if the stream is closed
// otherwise.
func (s *Stream) Read(p []byte) (int, error) {
	n, _, err := s.ReadSCTP(p)
	return n, err
}

// ReadSCTP reads a packet of len(p) bytes and returns the associated Payload Protocol Identifier.
// Returns EOF when the stream is reset or an error if the stream is closed
// otherwise.
func (s *Stream) ReadSCTP(p []byte) (int, PayloadProtocolIdentifier, error) {
	for {
		s.lock.Lock()
		n, ppi, err := s.reassemblyQueue.read(p)
		s.lock.Unlock()
		if err == nil {
			return n, ppi, nil
		}

		s.lock.RLock()
		err = s.readErr
		if err != nil {
			s.lock.RUnlock()
			return 0, PayloadProtocolIdentifier(0), err
		}

		notifier := s.readNotifier
		s.lock.RUnlock()
		// Wait for read notification
		if notifier != nil {
			notifier.L.Lock()
			notifier.Wait()
			notifier.L.Unlock()
		}
	}
}

func (s *Stream) handleData(pd *chunkPayloadData) {
	var readable bool
	s.lock.Lock()
	//fmt.Printf("stream[%d]: handleData ts=%d ssn=%d\n", s.streamIdentifier, pd.tsn, pd.streamSequenceNumber)
	if s.reassemblyQueue.push(pd) {
		readable = s.reassemblyQueue.isReadable()
	}
	//fmt.Printf("stream[%d]: handleData readable? %v\n", s.streamIdentifier, readable)
	s.lock.Unlock()

	// Notify the reader asynchronously if there's a data chunk to read.
	if readable {
		s.readNotifier.Signal()
	}
}

func (s *Stream) handleForwardTSN(newCumulativeTSN uint32, ssn uint16) {
	var readable bool
	s.lock.Lock()
	// Remove all chunks older than or equal to the new TSN from
	// the reassemblyQueue.
	//fmt.Printf("stream[%d]: handleForwardTSN newTSN=%d ssn=%d\n", s.streamIdentifier, newCumulativeTSN, ssn)
	s.reassemblyQueue.forwardTSN(newCumulativeTSN, s.unordered, ssn)
	readable = s.reassemblyQueue.isReadable()
	//fmt.Printf("stream[%d]: handleForwardTSN readable? %v\n", s.streamIdentifier, readable)
	s.lock.Unlock()

	// Notify the reader asynchronously if there's a data chunk to read.
	if readable {
		s.readNotifier.Signal()
	}
}

// Write writes len(p) bytes from p with the default Payload Protocol Identifier
func (s *Stream) Write(p []byte) (n int, err error) {
	return s.WriteSCTP(p, s.defaultPayloadType)
}

// WriteSCTP writes len(p) bytes from p to the DTLS connection
func (s *Stream) WriteSCTP(p []byte, ppi PayloadProtocolIdentifier) (n int, err error) {
	if len(p) > math.MaxUint16 {
		return 0, errors.Errorf("Outbound packet larger than maximum message size %v", math.MaxUint16)
	}

	s.lock.RLock()
	err = s.writeErr
	s.lock.RUnlock()
	if err != nil {
		return 0, err
	}

	chunks := s.packetize(p, ppi)

	return len(p), s.association.sendPayloadData(chunks)
}

func (s *Stream) packetize(raw []byte, ppi PayloadProtocolIdentifier) []*chunkPayloadData {
	s.lock.Lock()
	defer s.lock.Unlock()

	i := uint16(0)
	remaining := uint16(len(raw))

	// From draft-ietf-rtcweb-data-protocol-09, section 6:
	//   All Data Channel Establishment Protocol messages MUST be sent using
	//   ordered delivery and reliable transmission.
	unordered := ppi != PayloadTypeWebRTCDCEP && s.unordered

	var chunks []*chunkPayloadData
	for remaining != 0 {
		fragmentSize := min(s.association.myMaxMTU, remaining)
		chunk := &chunkPayloadData{
			streamIdentifier:     s.streamIdentifier,
			userData:             raw[i : i+fragmentSize],
			unordered:            unordered,
			beginningFragment:    i == 0,
			endingFragment:       remaining-fragmentSize == 0,
			immediateSack:        false,
			payloadType:          ppi,
			streamSequenceNumber: s.sequenceNumber,
		}

		chunks = append(chunks, chunk)

		remaining -= fragmentSize
		i += fragmentSize
	}

	// RFC 4960 Sec 6.6
	// Note: When transmitting ordered and unordered data, an endpoint does
	// not increment its Stream Sequence Number when transmitting a DATA
	// chunk with U flag set to 1.
	if !unordered {
		s.sequenceNumber++
	}

	return chunks
}

// Close closes the write-direction of the stream.
// Future calls to Write are not permitted after calling Close.
func (s *Stream) Close() error {
	s.lock.Lock()
	if s.writeErr != nil {
		s.lock.Unlock()
		return nil // already closed
	}
	s.writeErr = errors.New("Stream closed")

	a := s.association
	sid := s.streamIdentifier
	s.lock.Unlock()

	// Reset the outgoing stream
	// https://tools.ietf.org/html/rfc6525
	return a.sendResetRequest(sid)
}
<commit_msg>Fix race when send buffer is reused<commit_after>package sctp

import (
	"math"
	"sync"

	"github.com/pkg/errors"
)

const (
	// ReliabilityTypeReliable is used for reliable transmission
	ReliabilityTypeReliable byte = 0
	// ReliabilityTypeRexmit is used for partial reliability by retransmission count
	ReliabilityTypeRexmit byte = 1
	// ReliabilityTypeTimed is used for partial reliability by retransmission duration
	ReliabilityTypeTimed byte = 2
)

// Stream represents an SCTP stream
type Stream struct {
	association *Association

	lock sync.RWMutex

	streamIdentifier   uint16
	defaultPayloadType PayloadProtocolIdentifier

	reassemblyQueue *reassemblyQueue
	sequenceNumber  uint16

	readNotifier *sync.Cond

	readErr  error
	writeErr error

	unordered        bool
	reliabilityType  byte
	reliabilityValue uint32
}

// StreamIdentifier returns the Stream identifier associated to the stream.
func (s *Stream) StreamIdentifier() uint16 {
	s.lock.RLock()
	defer s.lock.RUnlock()
	return s.streamIdentifier
}

// SetDefaultPayloadType sets the default payload type used by Write.
func (s *Stream) SetDefaultPayloadType(defaultPayloadType PayloadProtocolIdentifier) {
	s.lock.Lock()
	defer s.lock.Unlock()

	s.setDefaultPayloadType(defaultPayloadType)
}

// setDefaultPayloadType sets the defaultPayloadType. The caller should hold the lock.
func (s *Stream) setDefaultPayloadType(defaultPayloadType PayloadProtocolIdentifier) {
	s.defaultPayloadType = defaultPayloadType
}

// SetReliabilityParams sets reliability parameters for this stream.
func (s *Stream) SetReliabilityParams(unordered bool, relType byte, relVal uint32) {
	s.lock.Lock()
	defer s.lock.Unlock()

	s.setReliabilityParams(unordered, relType, relVal)
}

// setReliabilityParams sets reliability parameters for this stream.
// The caller should hold the lock.
func (s *Stream) setReliabilityParams(unordered bool, relType byte, relVal uint32) {
	s.unordered = unordered
	s.reliabilityType = relType
	s.reliabilityValue = relVal
}

// Read reads a packet of len(p) bytes, dropping the Payload Protocol Identifier.
// Returns EOF when the stream is reset or an error if the stream is closed
// otherwise.
func (s *Stream) Read(p []byte) (int, error) {
	n, _, err := s.ReadSCTP(p)
	return n, err
}

// ReadSCTP reads a packet of len(p) bytes and returns the associated Payload Protocol Identifier.
// Returns EOF when the stream is reset or an error if the stream is closed
// otherwise.
func (s *Stream) ReadSCTP(p []byte) (int, PayloadProtocolIdentifier, error) {
	for {
		s.lock.Lock()
		n, ppi, err := s.reassemblyQueue.read(p)
		s.lock.Unlock()
		if err == nil {
			return n, ppi, nil
		}

		s.lock.RLock()
		err = s.readErr
		if err != nil {
			s.lock.RUnlock()
			return 0, PayloadProtocolIdentifier(0), err
		}

		notifier := s.readNotifier
		s.lock.RUnlock()
		// Wait for read notification
		if notifier != nil {
			notifier.L.Lock()
			notifier.Wait()
			notifier.L.Unlock()
		}
	}
}

func (s *Stream) handleData(pd *chunkPayloadData) {
	var readable bool
	s.lock.Lock()
	//fmt.Printf("stream[%d]: handleData ts=%d ssn=%d\n", s.streamIdentifier, pd.tsn, pd.streamSequenceNumber)
	if s.reassemblyQueue.push(pd) {
		readable = s.reassemblyQueue.isReadable()
	}
	//fmt.Printf("stream[%d]: handleData readable? %v\n", s.streamIdentifier, readable)
	s.lock.Unlock()

	// Notify the reader asynchronously if there's a data chunk to read.
	if readable {
		s.readNotifier.Signal()
	}
}

func (s *Stream) handleForwardTSN(newCumulativeTSN uint32, ssn uint16) {
	var readable bool
	s.lock.Lock()
	// Remove all chunks older than or equal to the new TSN from
	// the reassemblyQueue.
	//fmt.Printf("stream[%d]: handleForwardTSN newTSN=%d ssn=%d\n", s.streamIdentifier, newCumulativeTSN, ssn)
	s.reassemblyQueue.forwardTSN(newCumulativeTSN, s.unordered, ssn)
	readable = s.reassemblyQueue.isReadable()
	//fmt.Printf("stream[%d]: handleForwardTSN readable? %v\n", s.streamIdentifier, readable)
	s.lock.Unlock()

	// Notify the reader asynchronously if there's a data chunk to read.
	if readable {
		s.readNotifier.Signal()
	}
}

// Write writes len(p) bytes from p with the default Payload Protocol Identifier
func (s *Stream) Write(p []byte) (n int, err error) {
	return s.WriteSCTP(p, s.defaultPayloadType)
}

// WriteSCTP writes len(p) bytes from p to the DTLS connection
func (s *Stream) WriteSCTP(p []byte, ppi PayloadProtocolIdentifier) (n int, err error) {
	if len(p) > math.MaxUint16 {
		return 0, errors.Errorf("Outbound packet larger than maximum message size %v", math.MaxUint16)
	}

	s.lock.RLock()
	err = s.writeErr
	s.lock.RUnlock()
	if err != nil {
		return 0, err
	}

	chunks := s.packetize(p, ppi)

	return len(p), s.association.sendPayloadData(chunks)
}

func (s *Stream) packetize(raw []byte, ppi PayloadProtocolIdentifier) []*chunkPayloadData {
	s.lock.Lock()
	defer s.lock.Unlock()

	i := uint16(0)
	remaining := uint16(len(raw))

	// From draft-ietf-rtcweb-data-protocol-09, section 6:
	//   All Data Channel Establishment Protocol messages MUST be sent using
	//   ordered delivery and reliable transmission.
	unordered := ppi != PayloadTypeWebRTCDCEP && s.unordered

	var chunks []*chunkPayloadData
	for remaining != 0 {
		fragmentSize := min(s.association.myMaxMTU, remaining)

		// Copy the userdata since we'll have to store it until acked
		// and the caller may re-use the buffer in the mean time
		userData := make([]byte, fragmentSize)
		copy(userData, raw[i:i+fragmentSize])

		chunk := &chunkPayloadData{
			streamIdentifier:     s.streamIdentifier,
			userData:             userData,
			unordered:            unordered,
			beginningFragment:    i == 0,
			endingFragment:       remaining-fragmentSize == 0,
			immediateSack:        false,
			payloadType:          ppi,
			streamSequenceNumber: s.sequenceNumber,
		}

		chunks = append(chunks, chunk)

		remaining -= fragmentSize
		i += fragmentSize
	}

	// RFC 4960 Sec 6.6
	// Note: When transmitting ordered and unordered data, an endpoint does
	// not increment its Stream Sequence Number when transmitting a DATA
	// chunk with U flag set to 1.
	if !unordered {
		s.sequenceNumber++
	}

	return chunks
}

// Close closes the write-direction of the stream.
// Future calls to Write are not permitted after calling Close.
func (s *Stream) Close() error {
	s.lock.Lock()
	if s.writeErr != nil {
		s.lock.Unlock()
		return nil // already closed
	}
	s.writeErr = errors.New("Stream closed")

	a := s.association
	sid := s.streamIdentifier
	s.lock.Unlock()

	// Reset the outgoing stream
	// https://tools.ietf.org/html/rfc6525
	return a.sendResetRequest(sid)
}
<|endoftext|>


<commit_before>package gofakeit

import (
	"reflect"
)

// Struct fills in exported elements of a struct with random data
// based on the value of `fake` tag of exported elements.
// Use `fake:"skip"` to explicitly skip an element.
// All built-in types are supported, with templating support
// for string types.
func Struct(v interface{}) {
	r(reflect.TypeOf(v), reflect.ValueOf(v), "")
}

func r(t reflect.Type, v reflect.Value, template string) {
	switch t.Kind() {
	case reflect.Ptr:
		rPointer(t, v, template)
	case reflect.Struct:
		rStruct(t, v)
	case reflect.String:
		rString(template, v)
	case reflect.Uint8:
		v.SetUint(uint64(Uint8()))
	case reflect.Uint16:
		v.SetUint(uint64(Uint16()))
	case reflect.Uint32:
		v.SetUint(uint64(Uint32()))
	case reflect.Uint64:
		v.SetUint(Uint64())
	case reflect.Int:
		v.SetInt(Int64())
	case reflect.Int8:
		v.SetInt(int64(Int8()))
	case reflect.Int16:
		v.SetInt(int64(Int16()))
	case reflect.Int32:
		v.SetInt(int64(Int32()))
	case reflect.Int64:
		v.SetInt(Int64())
	case reflect.Float64:
		v.SetFloat(Float64())
	case reflect.Float32:
		v.SetFloat(float64(Float32()))
	case reflect.Bool:
		v.SetBool(Bool())
	}
}

func rString(template string, v reflect.Value) {
	if template != "" {
		r := Generate(template)
		v.SetString(r)
	} else {
		num := Number(4, 10)
		str := ""
		for i := 0; i < num; i++ {
			str = str + "?"
		}
		v.SetString(Generate(str))
	}
}

func rStruct(t reflect.Type, v reflect.Value) {
	n := t.NumField()
	for i := 0; i < n; i++ {
		elementT := t.Field(i)
		elementV := v.Field(i)
		fake := true
		t, ok := elementT.Tag.Lookup("fake")
		if ok && t == "skip" {
			fake = false
		}
		if fake && elementV.CanSet() {
			r(elementT.Type, elementV, t)
		}
	}
}

func rPointer(t reflect.Type, v reflect.Value, template string) {
	elemT := t.Elem()
	if v.IsNil() {
		nv := reflect.New(elemT)
		r(elemT, nv.Elem(), template)
		v.Set(nv)
	} else {
		r(elemT, v.Elem(), template)
	}
}
<commit_msg>struct - remove allocation<commit_after>package gofakeit

import (
	"reflect"
)

// Struct fills in exported elements of a struct with random data
// based on the value of `fake` tag of exported elements.
// Use `fake:"skip"` to explicitly skip an element.
// All built-in types are supported, with templating support
// for string types.
func Struct(v interface{}) {
	r(reflect.TypeOf(v), reflect.ValueOf(v), "")
}

func r(t reflect.Type, v reflect.Value, template string) {
	switch t.Kind() {
	case reflect.Ptr:
		rPointer(t, v, template)
	case reflect.Struct:
		rStruct(t, v)
	case reflect.String:
		rString(template, v)
	case reflect.Uint8:
		v.SetUint(uint64(Uint8()))
	case reflect.Uint16:
		v.SetUint(uint64(Uint16()))
	case reflect.Uint32:
		v.SetUint(uint64(Uint32()))
	case reflect.Uint64:
		v.SetUint(Uint64())
	case reflect.Int:
		v.SetInt(Int64())
	case reflect.Int8:
		v.SetInt(int64(Int8()))
	case reflect.Int16:
		v.SetInt(int64(Int16()))
	case reflect.Int32:
		v.SetInt(int64(Int32()))
	case reflect.Int64:
		v.SetInt(Int64())
	case reflect.Float64:
		v.SetFloat(Float64())
	case reflect.Float32:
		v.SetFloat(float64(Float32()))
	case reflect.Bool:
		v.SetBool(Bool())
	}
}

func rString(template string, v reflect.Value) {
	if template != "" {
		v.SetString(Generate(template))
	} else {
		num := Number(4, 10)
		str := ""
		for i := 0; i < num; i++ {
			str = str + "?"
		}
		v.SetString(Generate(str))
	}
}

func rStruct(t reflect.Type, v reflect.Value) {
	n := t.NumField()
	for i := 0; i < n; i++ {
		elementT := t.Field(i)
		elementV := v.Field(i)
		fake := true
		t, ok := elementT.Tag.Lookup("fake")
		if ok && t == "skip" {
			fake = false
		}
		if fake && elementV.CanSet() {
			r(elementT.Type, elementV, t)
		}
	}
}

func rPointer(t reflect.Type, v reflect.Value, template string) {
	elemT := t.Elem()
	if v.IsNil() {
		nv := reflect.New(elemT)
		r(elemT, nv.Elem(), template)
		v.Set(nv)
	} else {
		r(elemT, v.Elem(), template)
	}
}
<|endoftext|>


<commit_before>package core

import (
	"database/sql"
	"errors"
	"fmt"
	"log"
	"os"
	"path"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"github.com/mdlayher/goset"
	"github.com/wtolson/go-taglib"
)

// validSet is a set of valid file extensions which we should scan as media, as they are the ones
// which TagLib is capable of reading
var validSet = set.New(".ape", ".flac", ".m4a", ".mp3", ".mpc", ".ogg", ".wma", ".wv")

// fsTask is the interface which defines a filesystem task, such as a media scan, or an orphan scan
type fsTask interface {
	Scan(string, string, chan struct{}) error
}

// fsManager handles fsWalker processes, and communicates back and forth with the manager goroutine
func fsManager(mediaFolder string, fsKillChan chan struct{}) {
	log.Println("fs: starting...")

	// Initialize a queue of filesystem tasks
	fsQueue := make(chan fsTask, 10)
	cancelQueue := make(chan chan struct{}, 10)

	// Queue an orphan scan, followed by a media scan
	fsQueue <- new(fsOrphanScan)
	fsQueue <- new(fsMediaScan)

	// Invoke task queue via goroutine, so it can be halted via the manager
	go func() {
		for {
			select {
			// Trigger a fsTask from queue
			case task := <-fsQueue:
				// Create a channel to halt the scan
				cancelChan := make(chan struct{})
				cancelQueue <- cancelChan

				// Start the scan
				// TODO: use actual base folder, actual subfolders from filesystem watcher
				if err := task.Scan(mediaFolder, "", cancelChan); err != nil {
					log.Println(err)
				}

				// On completion, close the cancel channel
				cancelChan = <-cancelQueue
				close(cancelChan)
			}
		}
	}()

	// Trigger manager events via channel
	for {
		select {
		// Stop filesystem manager
		case <-fsKillChan:
			// Halt any in-progress tasks
			log.Println("fs: halting tasks")
			for i := 0; i < len(cancelQueue); i++ {
				// Receive a channel
				f := <-cancelQueue
				if f == nil {
					continue
				}

				// Send termination
				f <- struct{}{}
				log.Println("fs: task halted")
			}

			// Inform manager that shutdown is complete
			log.Println("fs: stopped!")
			fsKillChan <- struct{}{}
			return
		}
	}
}

// fsMediaScan represents a filesystem task which scans the given path for new media
type fsMediaScan struct{}

// Scan scans for media files in a specified path, and queues them up for inclusion
// in the wavepipe database
func (fs *fsMediaScan) Scan(mediaFolder string, subFolder string, walkCancelChan chan struct{}) error {
	// Media scans are comprehensive, so subfolder has no purpose
	if subFolder != "" {
		return errors.New("media scan: subfolder not valid for media scan operation")
	}

	// Halt walk if needed
	var mutex sync.RWMutex
	haltWalk := false
	go func() {
		// Wait for signal
		<-walkCancelChan

		// Halt!
		mutex.Lock()
		haltWalk = true
		mutex.Unlock()
	}()

	// Track metrics about the walk
	artistCount := 0
	albumCount := 0
	songCount := 0
	startTime := time.Now()

	// Invoke a recursive file walk on the given media folder, passing closure variables into
	// walkFunc to enable additional functionality
	log.Println("fs: beginning media scan")
	err := filepath.Walk(mediaFolder, func(currPath string, info os.FileInfo, err error) error {
		// Stop walking immediately if needed
		mutex.RLock()
		if haltWalk {
			return errors.New("media scan: halted by channel")
		}
		mutex.RUnlock()

		// Make sure path is actually valid
		if info == nil {
			return errors.New("media scan: invalid path: " + currPath)
		}

		// Ignore directories for now
		if info.IsDir() {
			return nil
		}

		// Check for a valid media extension
		if !validSet.Has(path.Ext(currPath)) {
			return nil
		}

		// Attempt to scan media file with taglib
		file, err := taglib.Read(currPath)
		if err != nil {
			return fmt.Errorf("%s: %s", currPath, err.Error())
		}
		defer file.Close()

		// Generate a song model from the TagLib file
		song, err := SongFromFile(file)
		if err != nil {
			return err
		}

		// Populate filesystem-related struct fields using OS info
		song.FileName = currPath
		song.FileSize = info.Size()

		// Extract type from the extension, capitalize it, drop the dot
		song.FileType = strings.ToUpper(path.Ext(info.Name()))[1:]
		song.LastModified = info.ModTime().Unix()

		// Generate an artist model from this song's metadata
		artist := ArtistFromSong(song)

		// Check for existing artist
		// Note: if the artist exists, this operation also loads necessary scanning information
		// such as their artist ID, for use in album and song generation
		if err := artist.Load(); err == sql.ErrNoRows {
			// Save new artist
			if err := artist.Save(); err != nil {
				log.Println(err)
			} else if err == nil {
				log.Printf("Artist: [%04d] %s", artist.ID, artist.Title)
				artistCount++
			}
		}

		// Generate the album model from this song's metadata
		album := AlbumFromSong(song)
		album.ArtistID = artist.ID

		// Check for existing album
		// Note: if the album exists, this operation also loads necessary scanning information
		// such as the album ID, for use in song generation
		if err := album.Load(); err == sql.ErrNoRows {
			// Save album
			if err := album.Save(); err != nil {
				log.Println(err)
			} else if err == nil {
				log.Printf("  - Album: [%04d] %s - %d - %s", album.ID, album.Artist, album.Year, album.Title)
				albumCount++
			}
		}

		// Add ID fields to song
		song.ArtistID = artist.ID
		song.AlbumID = album.ID

		// Check for existing song
		if err := song.Load(); err == sql.ErrNoRows {
			// Save song (don't log these because they really slow things down)
			if err := song.Save(); err != nil {
				log.Println(err)
			} else if err == nil {
				songCount++
			}
		}

		// Successful media scan
		return nil
	})

	// Check for filesystem walk errors
	if err != nil {
		return err
	}

	// Print metrics
	log.Printf("fs: media scan complete [time: %s]", time.Since(startTime).String())
	log.Printf("fs: added: [artists: %d] [albums: %d] [songs: %d]", artistCount, albumCount, songCount)

	// No errors
	return nil
}

// fsOrphanScan represents a filesystem task which scans the given path for orphaned media
type fsOrphanScan struct{}

// Scan scans for media files which have been removed from the media directory, and removes
// them as appropriate.  An orphan is defined as follows:
//   - Artist: no more songs contain this artist's ID
//   - Album: no more songs contain this album's ID
//   - Song: song is no longer present in the filesystem
// The baseFolder is the root location of the media folder.  As wavepipe currently supports only
// one media folder, any media which does not reside in this folder is orphaned.  If left blank,
// only the subFolder will be checked.
// The subFolder is the current file location, under the baseFolder.  This is used to allow for
// quick scans of a small subsection of the directory, such as on a filesystem change.  Any files
// which are in the database, but do not exist on disk, will be orphaned and removed.
func (fs *fsOrphanScan) Scan(baseFolder string, subFolder string, orphanCancelChan chan struct{}) error {
	// If both folders are empty, there is nothing to do
	if baseFolder == "" && subFolder == "" {
		return errors.New("orphan scan: no base folder or subfolder")
	}

	// Halt scan if needed
	var mutex sync.RWMutex
	haltOrphanScan := false
	go func() {
		// Wait for signal
		<-orphanCancelChan

		// Halt!
		mutex.Lock()
		haltOrphanScan = true
		mutex.Unlock()
	}()

	// Track metrics about the scan
	songCount := 0
	startTime := time.Now()

	log.Println("fs: beginning orphan scan")

	// Check if a baseFolder is set, meaning remove ANYTHING not under this base
	if baseFolder != "" {
		log.Println("fs: orphan scanning base folder:", baseFolder)

		// Scan for all songs NOT under the base folder
		songs, err := db.SongsNotInPath(baseFolder)
		if err != nil {
			return err
		}

		// Remove all songs which are not in this path
		for _, s := range songs {
			// Remove song from database
			if err := s.Delete(); err != nil {
				return err
			}

			songCount++
		}
	}

	// If no subfolder set, use the base folder to check file existence
	if subFolder == "" {
		subFolder = baseFolder
	}

	// Scan for all songs in subfolder
	log.Println("fs: orphan scanning subfolder:", subFolder)
	songs, err := db.SongsInPath(subFolder)
	if err != nil {
		return err
	}

	// Iterate all songs in this path
	for _, s := range songs {
		// Check that the song still exists in this place
		if _, err := os.Stat(s.FileName); os.IsNotExist(err) {
			// Remove song from database
			if err := s.Delete(); err != nil {
				return err
			}

			songCount++
		}
	}

	// Now that songs have been purged, check for albums
	albumCount, err := db.PurgeOrphanAlbums()
	if err != nil {
		return nil
	}

	// Finally, check for artists
	artistCount, err := db.PurgeOrphanArtists()
	if err != nil {
		return nil
	}

	// Print metrics
	log.Printf("fs: orphan scan complete [time: %s]", time.Since(startTime).String())
	log.Printf("fs: removed: [artists: %d] [albums: %d] [songs: %d]", artistCount, albumCount, songCount)

	// No errors
	return nil
}
<commit_msg>Initial, experimental commit of filesystem watchers<commit_after>package core

import (
	"database/sql"
	"errors"
	"fmt"
	"log"
	"os"
	"path"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"github.com/mdlayher/goset"
	"github.com/romanoff/fsmonitor"
	"github.com/wtolson/go-taglib"
)

// validSet is a set of valid file extensions which we should scan as media, as they are the ones
// which TagLib is capable of reading
var validSet = set.New(".ape", ".flac", ".m4a", ".mp3", ".mpc", ".ogg", ".wma", ".wv")

// fsQueue is a queue of tasks to be performed by the filesystem, such as media and orphan scans
var fsQueue = make(chan fsTask, 10)

// fsTask is the interface which defines a filesystem task, such as a media scan, or an orphan scan
type fsTask interface {
	Folders() (string, string)
	SetFolders(string, string)
	Scan(string, string, chan struct{}) error
}

// fsManager handles fsWalker processes, and communicates back and forth with the manager goroutine
func fsManager(mediaFolder string, fsKillChan chan struct{}) {
	log.Println("fs: starting...")

	// Initialize a queue to cancel filesystem tasks
	cancelQueue := make(chan chan struct{}, 10)

	// Queue an orphan scan
	o := new(fsOrphanScan)
	o.SetFolders(mediaFolder, "")
	fsQueue <- o

	// Queue a media scan
	m := new(fsMediaScan)
	m.SetFolders(mediaFolder, "")
	fsQueue <- m

	// Invoke task queue via goroutine, so it can be halted via the manager
	go func() {
		for {
			select {
			// Trigger a fsTask from queue
			case task := <-fsQueue:
				// Create a channel to halt the scan
				cancelChan := make(chan struct{})
				cancelQueue <- cancelChan

				// Retrieve the folders to use with scan
				baseFolder, subFolder := task.Folders()

				// Start the scan
				if err := task.Scan(baseFolder, subFolder, cancelChan); err != nil {
					log.Println(err)
				}

				// On completion, close the cancel channel
				cancelChan = <-cancelQueue
				close(cancelChan)
			}
		}
	}()

	// Create a filesystem watcher
	watcher, err := fsmonitor.NewWatcher()
	if err != nil {
		log.Println(err)
		return
	}

	// Wait for events on goroutine
	// TODO: for the time being, we ignore IsModify() and IsRename() events, because the
	// IsCreate() and IsDelete() methods, and recurring scans will cover these.  Will re-evalaute
	// this later.
	go func() {
		for {
			select {
			// Event occurred
			case ev := <-watcher.Event:
				switch {
				// On create, trigger a media scan
				case ev.IsCreate():
					// Scan item as the "base folder", so it just adds this item
					m := new(fsMediaScan)
					m.SetFolders(ev.Name, "")
					fsQueue <- m
				// On delete, trigger an orphan scan
				case ev.IsDelete():
					// Scan item as the "subfolder", so it just removes this item
					o := new(fsOrphanScan)
					o.SetFolders("", ev.Name)
					fsQueue <- o
				}
			// Watcher errors
			case err := <-watcher.Error:
				log.Println(err)
				return
			}
		}
	}()

	// Watch media folder
	if err := watcher.Watch(mediaFolder); err != nil {
		log.Println(err)
		return
	}
	log.Println("fs: watching folder:", mediaFolder)

	// Trigger manager events via channel
	for {
		select {
		// Stop filesystem manager
		case <-fsKillChan:
			// Halt any in-progress tasks
			log.Println("fs: halting tasks")
			for i := 0; i < len(cancelQueue); i++ {
				// Receive a channel
				f := <-cancelQueue
				if f == nil {
					continue
				}

				// Send termination
				f <- struct{}{}
				log.Println("fs: task halted")
			}

			// Inform manager that shutdown is complete
			log.Println("fs: stopped!")
			fsKillChan <- struct{}{}
			return
		}
	}
}

// fsMediaScan represents a filesystem task which scans the given path for new media
type fsMediaScan struct {
	baseFolder string
	subFolder  string
}

// Folders returns the base folder and subfolder for use with a scanning task
func (fs *fsMediaScan) Folders() (string, string) {
	return fs.baseFolder, fs.subFolder
}

// SetFolders sets the base folder and subfolder for use with a scanning task
func (fs *fsMediaScan) SetFolders(baseFolder string, subFolder string) {
	fs.baseFolder = baseFolder
	fs.subFolder = subFolder
}

// Scan scans for media files in a specified path, and queues them up for inclusion
// in the wavepipe database
func (fs *fsMediaScan) Scan(mediaFolder string, subFolder string, walkCancelChan chan struct{}) error {
	// Media scans are comprehensive, so subfolder has no purpose
	if subFolder != "" {
		return errors.New("media scan: subfolder not valid for media scan operation")
	}

	// Halt walk if needed
	var mutex sync.RWMutex
	haltWalk := false
	go func() {
		// Wait for signal
		<-walkCancelChan

		// Halt!
		mutex.Lock()
		haltWalk = true
		mutex.Unlock()
	}()

	// Track metrics about the walk
	artistCount := 0
	albumCount := 0
	songCount := 0
	startTime := time.Now()

	// Invoke a recursive file walk on the given media folder, passing closure variables into
	// walkFunc to enable additional functionality
	log.Println("fs: beginning media scan:", mediaFolder)
	err := filepath.Walk(mediaFolder, func(currPath string, info os.FileInfo, err error) error {
		// Stop walking immediately if needed
		mutex.RLock()
		if haltWalk {
			return errors.New("media scan: halted by channel")
		}
		mutex.RUnlock()

		// Make sure path is actually valid
		if info == nil {
			return errors.New("media scan: invalid path: " + currPath)
		}

		// Ignore directories for now
		if info.IsDir() {
			return nil
		}

		// Check for a valid media extension
		if !validSet.Has(path.Ext(currPath)) {
			return nil
		}

		// Attempt to scan media file with taglib
		file, err := taglib.Read(currPath)
		if err != nil {
			return fmt.Errorf("%s: %s", currPath, err.Error())
		}
		defer file.Close()

		// Generate a song model from the TagLib file
		song, err := SongFromFile(file)
		if err != nil {
			return err
		}

		// Populate filesystem-related struct fields using OS info
		song.FileName = currPath
		song.FileSize = info.Size()

		// Extract type from the extension, capitalize it, drop the dot
		song.FileType = strings.ToUpper(path.Ext(info.Name()))[1:]
		song.LastModified = info.ModTime().Unix()

		// Generate an artist model from this song's metadata
		artist := ArtistFromSong(song)

		// Check for existing artist
		// Note: if the artist exists, this operation also loads necessary scanning information
		// such as their artist ID, for use in album and song generation
		if err := artist.Load(); err == sql.ErrNoRows {
			// Save new artist
			if err := artist.Save(); err != nil {
				log.Println(err)
			} else if err == nil {
				log.Printf("Artist: [%04d] %s", artist.ID, artist.Title)
				artistCount++
			}
		}

		// Generate the album model from this song's metadata
		album := AlbumFromSong(song)
		album.ArtistID = artist.ID

		// Check for existing album
		// Note: if the album exists, this operation also loads necessary scanning information
		// such as the album ID, for use in song generation
		if err := album.Load(); err == sql.ErrNoRows {
			// Save album
			if err := album.Save(); err != nil {
				log.Println(err)
			} else if err == nil {
				log.Printf("  - Album: [%04d] %s - %d - %s", album.ID, album.Artist, album.Year, album.Title)
				albumCount++
			}
		}

		// Add ID fields to song
		song.ArtistID = artist.ID
		song.AlbumID = album.ID

		// Check for existing song
		if err := song.Load(); err == sql.ErrNoRows {
			// Save song (don't log these because they really slow things down)
			if err := song.Save(); err != nil {
				log.Println(err)
			} else if err == nil {
				songCount++
			}
		}

		// Successful media scan
		return nil
	})

	// Check for filesystem walk errors
	if err != nil {
		return err
	}

	// Print metrics
	log.Printf("fs: media scan complete [time: %s]", time.Since(startTime).String())
	log.Printf("fs: added: [artists: %d] [albums: %d] [songs: %d]", artistCount, albumCount, songCount)

	// No errors
	return nil
}

// fsOrphanScan represents a filesystem task which scans the given path for orphaned media
type fsOrphanScan struct {
	baseFolder string
	subFolder  string
}

// Folders returns the base folder and subfolder for use with a scanning task
func (fs *fsOrphanScan) Folders() (string, string) {
	return fs.baseFolder, fs.subFolder
}

// SetFolders sets the base folder and subfolder for use with a scanning task
func (fs *fsOrphanScan) SetFolders(baseFolder string, subFolder string) {
	fs.baseFolder = baseFolder
	fs.subFolder = subFolder
}

// Scan scans for media files which have been removed from the media directory, and removes
// them as appropriate.  An orphan is defined as follows:
//   - Artist: no more songs contain this artist's ID
//   - Album: no more songs contain this album's ID
//   - Song: song is no longer present in the filesystem
// The baseFolder is the root location of the media folder.  As wavepipe currently supports only
// one media folder, any media which does not reside in this folder is orphaned.  If left blank,
// only the subFolder will be checked.
// The subFolder is the current file location, under the baseFolder.  This is used to allow for
// quick scans of a small subsection of the directory, such as on a filesystem change.  Any files
// which are in the database, but do not exist on disk, will be orphaned and removed.
func (fs *fsOrphanScan) Scan(baseFolder string, subFolder string, orphanCancelChan chan struct{}) error {
	// If both folders are empty, there is nothing to do
	if baseFolder == "" && subFolder == "" {
		return errors.New("orphan scan: no base folder or subfolder")
	}

	// Halt scan if needed
	var mutex sync.RWMutex
	haltOrphanScan := false
	go func() {
		// Wait for signal
		<-orphanCancelChan

		// Halt!
		mutex.Lock()
		haltOrphanScan = true
		mutex.Unlock()
	}()

	// Track metrics about the scan
	songCount := 0
	startTime := time.Now()

	log.Println("fs: beginning orphan scan")

	// Check if a baseFolder is set, meaning remove ANYTHING not under this base
	if baseFolder != "" {
		log.Println("fs: orphan scanning base folder:", baseFolder)

		// Scan for all songs NOT under the base folder
		songs, err := db.SongsNotInPath(baseFolder)
		if err != nil {
			return err
		}

		// Remove all songs which are not in this path
		for _, s := range songs {
			// Remove song from database
			if err := s.Delete(); err != nil {
				return err
			}

			songCount++
		}
	}

	// If no subfolder set, use the base folder to check file existence
	if subFolder == "" {
		subFolder = baseFolder
	}

	// Scan for all songs in subfolder
	log.Println("fs: orphan scanning subfolder:", subFolder)
	songs, err := db.SongsInPath(subFolder)
	if err != nil {
		return err
	}

	// Iterate all songs in this path
	for _, s := range songs {
		// Check that the song still exists in this place
		if _, err := os.Stat(s.FileName); os.IsNotExist(err) {
			// Remove song from database
			if err := s.Delete(); err != nil {
				return err
			}

			songCount++
		}
	}

	// Now that songs have been purged, check for albums
	albumCount, err := db.PurgeOrphanAlbums()
	if err != nil {
		return nil
	}

	// Finally, check for artists
	artistCount, err := db.PurgeOrphanArtists()
	if err != nil {
		return nil
	}

	// Print metrics
	log.Printf("fs: orphan scan complete [time: %s]", time.Since(startTime).String())
	log.Printf("fs: removed: [artists: %d] [albums: %d] [songs: %d]", artistCount, albumCount, songCount)

	// No errors
	return nil
}
<|endoftext|>


<commit_before>package rtorrent

import (
	"fmt"

	"github.com/tehjojo/go-rtorrent/xmlrpc"
)

// RTorrent is used to communicate with a remote rTorrent instance
type RTorrent struct {
	xmlrpcClient *xmlrpc.Client
}

// Torrent represents a torrent in rTorrent
type Torrent struct {
	Hash      string
	Name      string
	Path      string
	Size      int
	Label     string
	Completed bool
	Ratio     float64
}

// File represents a file in rTorrent
type File struct {
	Path string
	Size int
}

// View represents a "view" within RTorrent
type View string

const (
	// ViewMain represents the "main" view, containing all torrents
	ViewMain View = "main"
	// ViewStarted represents the "started" view, containing only torrents that have been started
	ViewStarted View = "started"
	// ViewStopped represents the "stopped" view, containing only torrents that have been stopped
	ViewStopped View = "stopped"
	// ViewHashing represents the "hashing" view, containing only torrents that are currently hashing
	ViewHashing View = "hashing"
	// ViewSeeding represents the "seeding" view, containing only torrents that are currently seeding
	ViewSeeding View = "seeding"
)

// Pretty returns a formatted string representing this Torrent
func (t *Torrent) Pretty() string {
	return fmt.Sprintf("Torrent:\n\tHash: %v\n\tName: %v\n\tPath: %v\n\tLabel: %v\n\tSize: %v bytes\n\tCompleted: %v\n\tRatio: %v\n", t.Hash, t.Name, t.Path, t.Label, t.Size, t.Completed, t.Ratio)
}

// Pretty returns a formatted string representing this File
func (f *File) Pretty() string {
	return fmt.Sprintf("File:\n\tPath: %v\n\tSize: %v bytes\n", f.Path, f.Size)
}

// New returns a new instance of `RTorrent`
// Pass in a true value for `insecure` to turn off certificate verification
func New(addr string, insecure bool) *RTorrent {
	return &RTorrent{
		xmlrpcClient: xmlrpc.NewClient(addr, insecure),
	}
}

// IP returns the IP reported by this RTorrent instance
func (r *RTorrent) IP() (string, error) {
	result, err := r.xmlrpcClient.Call("get_ip")
	return result.(string), err
}

// Name returns the name reported by this RTorrent instance
func (r *RTorrent) Name() (string, error) {
	result, err := r.xmlrpcClient.Call("get_name")
	return result.(string), err
}

// DownTotal returns the total downloaded metric reported by this RTorrent instance (bytes)
func (r *RTorrent) DownTotal() (int, error) {
	result, err := r.xmlrpcClient.Call("get_down_total")
	return result.(int), err
}

// UpTotal returns the total uploaded metric reported by this RTorrent instance (bytes)
func (r *RTorrent) UpTotal() (int, error) {
	result, err := r.xmlrpcClient.Call("get_up_total")
	return result.(int), err
}

// GetTorrents returns all of the torrents reported by this RTorrent instance
func (r *RTorrent) GetTorrents(view View) ([]Torrent, error) {
	args := []interface{}{string(view), "d.get_name=", "d.get_size_bytes=", "d.get_hash=", "d.get_custom1=", "d.get_base_path=", "d.is_active=", "d.get_complete=", "d.get_ratio="}
	results, err := r.xmlrpcClient.Call("d.multicall", args...)
	var torrents []Torrent
	if err != nil {
		return torrents, err
	}
	for _, outerResult := range results.([]interface{}) {
		for _, innerResult := range outerResult.([]interface{}) {
			torrentData := innerResult.([]interface{})
			torrents = append(torrents, Torrent{
				Hash:      torrentData[2].(string),
				Name:      torrentData[0].(string),
				Path:      torrentData[4].(string),
				Size:      torrentData[1].(int),
				Label:     torrentData[3].(string),
				Completed: torrentData[6].(int) > 0,
				Ratio:     float64(torrentData[7].(int)) / float64(1000),
			})
		}
	}
	return torrents, nil
}

// GetFiles returns all of the files for a given `Torrent`
func (r *RTorrent) GetFiles(t Torrent) ([]File, error) {
	args := []interface{}{t.Hash, 0, "f.get_path=", "f.get_size_bytes="}
	results, err := r.xmlrpcClient.Call("f.multicall", args...)
	var files []File
	if err != nil {
		return files, err
	}
	for _, outerResult := range results.([]interface{}) {
		for _, innerResult := range outerResult.([]interface{}) {
			fileData := innerResult.([]interface{})
			files = append(files, File{
				Path: fileData[0].(string),
				Size: fileData[1].(int),
			})
		}
	}
	return files, nil
}

// SetLabel sets the label on the given Torrent
func (r *RTorrent) SetLabel(t Torrent, newLabel string) error {
	t.Label = newLabel
	args := []interface{}{t.Hash, newLabel}
	if _, err := r.xmlrpcClient.Call("d.set_custom1", args...); err != nil {
		return err
	}
	return nil
}
<commit_msg>Added method to add torrent by URL<commit_after>package rtorrent

import (
	"fmt"

	"github.com/tehjojo/go-rtorrent/xmlrpc"
)

// RTorrent is used to communicate with a remote rTorrent instance
type RTorrent struct {
	xmlrpcClient *xmlrpc.Client
}

// Torrent represents a torrent in rTorrent
type Torrent struct {
	Hash      string
	Name      string
	Path      string
	Size      int
	Label     string
	Completed bool
	Ratio     float64
}

// File represents a file in rTorrent
type File struct {
	Path string
	Size int
}

// View represents a "view" within RTorrent
type View string

const (
	// ViewMain represents the "main" view, containing all torrents
	ViewMain View = "main"
	// ViewStarted represents the "started" view, containing only torrents that have been started
	ViewStarted View = "started"
	// ViewStopped represents the "stopped" view, containing only torrents that have been stopped
	ViewStopped View = "stopped"
	// ViewHashing represents the "hashing" view, containing only torrents that are currently hashing
	ViewHashing View = "hashing"
	// ViewSeeding represents the "seeding" view, containing only torrents that are currently seeding
	ViewSeeding View = "seeding"
)

// Pretty returns a formatted string representing this Torrent
func (t *Torrent) Pretty() string {
	return fmt.Sprintf("Torrent:\n\tHash: %v\n\tName: %v\n\tPath: %v\n\tLabel: %v\n\tSize: %v bytes\n\tCompleted: %v\n\tRatio: %v\n", t.Hash, t.Name, t.Path, t.Label, t.Size, t.Completed, t.Ratio)
}

// Pretty returns a formatted string representing this File
func (f *File) Pretty() string {
	return fmt.Sprintf("File:\n\tPath: %v\n\tSize: %v bytes\n", f.Path, f.Size)
}

// New returns a new instance of `RTorrent`
// Pass in a true value for `insecure` to turn off certificate verification
func New(addr string, insecure bool) *RTorrent {
	return &RTorrent{
		xmlrpcClient: xmlrpc.NewClient(addr, insecure),
	}
}

// Add adds a new torrent by URL
func (r *RTorrent) Add(url string) error {
	_, err := r.xmlrpcClient.Call("load_start", url)
	if err != nil {
		return err
	}
	return nil
}

// IP returns the IP reported by this RTorrent instance
func (r *RTorrent) IP() (string, error) {
	result, err := r.xmlrpcClient.Call("get_ip")
	return result.(string), err
}

// Name returns the name reported by this RTorrent instance
func (r *RTorrent) Name() (string, error) {
	result, err := r.xmlrpcClient.Call("get_name")
	return result.(string), err
}

// DownTotal returns the total downloaded metric reported by this RTorrent instance (bytes)
func (r *RTorrent) DownTotal() (int, error) {
	result, err := r.xmlrpcClient.Call("get_down_total")
	return result.(int), err
}

// UpTotal returns the total uploaded metric reported by this RTorrent instance (bytes)
func (r *RTorrent) UpTotal() (int, error) {
	result, err := r.xmlrpcClient.Call("get_up_total")
	return result.(int), err
}

// GetTorrents returns all of the torrents reported by this RTorrent instance
func (r *RTorrent) GetTorrents(view View) ([]Torrent, error) {
	args := []interface{}{string(view), "d.get_name=", "d.get_size_bytes=", "d.get_hash=", "d.get_custom1=", "d.get_base_path=", "d.is_active=", "d.get_complete=", "d.get_ratio="}
	results, err := r.xmlrpcClient.Call("d.multicall", args...)
	var torrents []Torrent
	if err != nil {
		return torrents, err
	}
	for _, outerResult := range results.([]interface{}) {
		for _, innerResult := range outerResult.([]interface{}) {
			torrentData := innerResult.([]interface{})
			torrents = append(torrents, Torrent{
				Hash:      torrentData[2].(string),
				Name:      torrentData[0].(string),
				Path:      torrentData[4].(string),
				Size:      torrentData[1].(int),
				Label:     torrentData[3].(string),
				Completed: torrentData[6].(int) > 0,
				Ratio:     float64(torrentData[7].(int)) / float64(1000),
			})
		}
	}
	return torrents, nil
}

// GetFiles returns all of the files for a given `Torrent`
func (r *RTorrent) GetFiles(t Torrent) ([]File, error) {
	args := []interface{}{t.Hash, 0, "f.get_path=", "f.get_size_bytes="}
	results, err := r.xmlrpcClient.Call("f.multicall", args...)
	var files []File
	if err != nil {
		return files, err
	}
	for _, outerResult := range results.([]interface{}) {
		for _, innerResult := range outerResult.([]interface{}) {
			fileData := innerResult.([]interface{})
			files = append(files, File{
				Path: fileData[0].(string),
				Size: fileData[1].(int),
			})
		}
	}
	return files, nil
}

// SetLabel sets the label on the given Torrent
func (r *RTorrent) SetLabel(t Torrent, newLabel string) error {
	t.Label = newLabel
	args := []interface{}{t.Hash, newLabel}
	if _, err := r.xmlrpcClient.Call("d.set_custom1", args...); err != nil {
		return err
	}
	return nil
}
<|endoftext|>


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

package sql_test

import (
	"testing"

	"golang.org/x/net/context"

	"github.com/cockroachdb/cockroach/keys"
	"github.com/cockroachdb/cockroach/roachpb"
	"github.com/cockroachdb/cockroach/server"
	"github.com/cockroachdb/cockroach/sql/sqlbase"
	"github.com/cockroachdb/cockroach/testutils"
	"github.com/cockroachdb/cockroach/testutils/serverutils"
	"github.com/cockroachdb/cockroach/util/leaktest"
)

func TestDatabaseDescriptor(t *testing.T) {
	defer leaktest.AfterTest(t)()
	params, _ := createTestServerParams()
	s, sqlDB, kvDB := serverutils.StartServer(t, params)
	defer s.Stopper().Stop()
	ctx := context.TODO()

	expectedCounter := int64(keys.MaxReservedDescID + 1)

	// Test values before creating the database.
	// descriptor ID counter.
	if ir, err := kvDB.Get(ctx, keys.DescIDGenerator); err != nil {
		t.Fatal(err)
	} else if actual := ir.ValueInt(); actual != expectedCounter {
		t.Fatalf("expected descriptor ID == %d, got %d", expectedCounter, actual)
	}

	// Database name.
	nameKey := sqlbase.MakeNameMetadataKey(keys.RootNamespaceID, "test")
	if gr, err := kvDB.Get(ctx, nameKey); err != nil {
		t.Fatal(err)
	} else if gr.Exists() {
		t.Fatal("expected non-existing key")
	}

	// Write a descriptor key that will interfere with database creation.
	dbDescKey := sqlbase.MakeDescMetadataKey(sqlbase.ID(expectedCounter))
	dbDesc := &sqlbase.Descriptor{
		Union: &sqlbase.Descriptor_Database{
			Database: &sqlbase.DatabaseDescriptor{
				Name:       "sentinel",
				ID:         sqlbase.ID(expectedCounter),
				Privileges: &sqlbase.PrivilegeDescriptor{},
			},
		},
	}
	if err := kvDB.CPut(ctx, dbDescKey, dbDesc, nil); err != nil {
		t.Fatal(err)
	}

	// Database creation should fail, and nothing should have been written.
	if _, err := sqlDB.Exec(`CREATE DATABASE test`); !testutils.IsError(err, "unexpected value") {
		t.Fatalf("unexpected error %v", err)
	}

	if ir, err := kvDB.Get(ctx, keys.DescIDGenerator); err != nil {
		t.Fatal(err)
	} else if actual := ir.ValueInt(); actual != expectedCounter {
		t.Fatalf("expected descriptor ID == %d, got %d", expectedCounter, actual)
	}

	start := roachpb.Key(keys.MakeTablePrefix(uint32(keys.NamespaceTableID)))
	if kvs, err := kvDB.Scan(ctx, start, start.PrefixEnd(), 0); err != nil {
		t.Fatal(err)
	} else {
		if a, e := len(kvs), server.GetBootstrapSchema().SystemDescriptorCount(); a != e {
			t.Fatalf("expected %d keys to have been written, found %d keys", e, a)
		}
	}

	// Remove the junk; allow database creation to proceed.
	if err := kvDB.Del(ctx, dbDescKey); err != nil {
		t.Fatal(err)
	}

	if _, err := sqlDB.Exec(`CREATE DATABASE test`); err != nil {
		t.Fatal(err)
	}
	expectedCounter++

	// Check keys again.
	// descriptor ID counter.
	if ir, err := kvDB.Get(ctx, keys.DescIDGenerator); err != nil {
		t.Fatal(err)
	} else if actual := ir.ValueInt(); actual != expectedCounter {
		t.Fatalf("expected descriptor ID == %d, got %d", expectedCounter, actual)
	}

	// Database name.
	if gr, err := kvDB.Get(ctx, nameKey); err != nil {
		t.Fatal(err)
	} else if !gr.Exists() {
		t.Fatal("key is missing")
	}

	// database descriptor.
	if gr, err := kvDB.Get(ctx, dbDescKey); err != nil {
		t.Fatal(err)
	} else if !gr.Exists() {
		t.Fatal("key is missing")
	}

	// Now try to create it again. We should fail, but not increment the counter.
	if _, err := sqlDB.Exec(`CREATE DATABASE test`); err == nil {
		t.Fatal("failure expected")
	}

	// Check keys again.
	// descriptor ID counter.
	if ir, err := kvDB.Get(ctx, keys.DescIDGenerator); err != nil {
		t.Fatal(err)
	} else if actual := ir.ValueInt(); actual != expectedCounter {
		t.Fatalf("expected descriptor ID == %d, got %d", expectedCounter, actual)
	}

	// Database name.
	if gr, err := kvDB.Get(ctx, nameKey); err != nil {
		t.Fatal(err)
	} else if !gr.Exists() {
		t.Fatal("key is missing")
	}

	// database descriptor.
	if gr, err := kvDB.Get(ctx, dbDescKey); err != nil {
		t.Fatal(err)
	} else if !gr.Exists() {
		t.Fatal("key is missing")
	}
}
<commit_msg>sql: add create table test<commit_after>// Copyright 2015 The Cockroach Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
// implied. See the License for the specific language governing
// permissions and limitations under the License.
//
// Author: Marc Berhault (marc@cockroachlabs.com)

package sql_test

import (
	gosql "database/sql"
	"fmt"
	"sync"
	"testing"

	"golang.org/x/net/context"

	"github.com/cockroachdb/cockroach/base"
	"github.com/cockroachdb/cockroach/internal/client"
	"github.com/cockroachdb/cockroach/keys"
	"github.com/cockroachdb/cockroach/roachpb"
	"github.com/cockroachdb/cockroach/server"
	"github.com/cockroachdb/cockroach/sql/sqlbase"
	"github.com/cockroachdb/cockroach/testutils"
	"github.com/cockroachdb/cockroach/testutils/serverutils"
	"github.com/cockroachdb/cockroach/testutils/testcluster"
	"github.com/cockroachdb/cockroach/util/leaktest"
)

func TestDatabaseDescriptor(t *testing.T) {
	defer leaktest.AfterTest(t)()
	params, _ := createTestServerParams()
	s, sqlDB, kvDB := serverutils.StartServer(t, params)
	defer s.Stopper().Stop()
	ctx := context.TODO()

	expectedCounter := int64(keys.MaxReservedDescID + 1)

	// Test values before creating the database.
	// descriptor ID counter.
	if ir, err := kvDB.Get(ctx, keys.DescIDGenerator); err != nil {
		t.Fatal(err)
	} else if actual := ir.ValueInt(); actual != expectedCounter {
		t.Fatalf("expected descriptor ID == %d, got %d", expectedCounter, actual)
	}

	// Database name.
	nameKey := sqlbase.MakeNameMetadataKey(keys.RootNamespaceID, "test")
	if gr, err := kvDB.Get(ctx, nameKey); err != nil {
		t.Fatal(err)
	} else if gr.Exists() {
		t.Fatal("expected non-existing key")
	}

	// Write a descriptor key that will interfere with database creation.
	dbDescKey := sqlbase.MakeDescMetadataKey(sqlbase.ID(expectedCounter))
	dbDesc := &sqlbase.Descriptor{
		Union: &sqlbase.Descriptor_Database{
			Database: &sqlbase.DatabaseDescriptor{
				Name:       "sentinel",
				ID:         sqlbase.ID(expectedCounter),
				Privileges: &sqlbase.PrivilegeDescriptor{},
			},
		},
	}
	if err := kvDB.CPut(ctx, dbDescKey, dbDesc, nil); err != nil {
		t.Fatal(err)
	}

	// Database creation should fail, and nothing should have been written.
	if _, err := sqlDB.Exec(`CREATE DATABASE test`); !testutils.IsError(err, "unexpected value") {
		t.Fatalf("unexpected error %v", err)
	}

	if ir, err := kvDB.Get(ctx, keys.DescIDGenerator); err != nil {
		t.Fatal(err)
	} else if actual := ir.ValueInt(); actual != expectedCounter {
		t.Fatalf("expected descriptor ID == %d, got %d", expectedCounter, actual)
	}

	start := roachpb.Key(keys.MakeTablePrefix(uint32(keys.NamespaceTableID)))
	if kvs, err := kvDB.Scan(ctx, start, start.PrefixEnd(), 0); err != nil {
		t.Fatal(err)
	} else {
		if a, e := len(kvs), server.GetBootstrapSchema().SystemDescriptorCount(); a != e {
			t.Fatalf("expected %d keys to have been written, found %d keys", e, a)
		}
	}

	// Remove the junk; allow database creation to proceed.
	if err := kvDB.Del(ctx, dbDescKey); err != nil {
		t.Fatal(err)
	}

	if _, err := sqlDB.Exec(`CREATE DATABASE test`); err != nil {
		t.Fatal(err)
	}
	expectedCounter++

	// Check keys again.
	// descriptor ID counter.
	if ir, err := kvDB.Get(ctx, keys.DescIDGenerator); err != nil {
		t.Fatal(err)
	} else if actual := ir.ValueInt(); actual != expectedCounter {
		t.Fatalf("expected descriptor ID == %d, got %d", expectedCounter, actual)
	}

	// Database name.
	if gr, err := kvDB.Get(ctx, nameKey); err != nil {
		t.Fatal(err)
	} else if !gr.Exists() {
		t.Fatal("key is missing")
	}

	// database descriptor.
	if gr, err := kvDB.Get(ctx, dbDescKey); err != nil {
		t.Fatal(err)
	} else if !gr.Exists() {
		t.Fatal("key is missing")
	}

	// Now try to create it again. We should fail, but not increment the counter.
	if _, err := sqlDB.Exec(`CREATE DATABASE test`); err == nil {
		t.Fatal("failure expected")
	}

	// Check keys again.
	// descriptor ID counter.
	if ir, err := kvDB.Get(ctx, keys.DescIDGenerator); err != nil {
		t.Fatal(err)
	} else if actual := ir.ValueInt(); actual != expectedCounter {
		t.Fatalf("expected descriptor ID == %d, got %d", expectedCounter, actual)
	}

	// Database name.
	if gr, err := kvDB.Get(ctx, nameKey); err != nil {
		t.Fatal(err)
	} else if !gr.Exists() {
		t.Fatal("key is missing")
	}

	// database descriptor.
	if gr, err := kvDB.Get(ctx, dbDescKey); err != nil {
		t.Fatal(err)
	} else if !gr.Exists() {
		t.Fatal("key is missing")
	}
}

// createTestTable tries to create a new table named based on the passed in id.
// It is designed to be synced with a number of concurrent calls to this
// function. Before starting, it first signals a done on the start waitgroup
// and then will block until the signal channel is closed. Once closed, it will
// proceed to try to create the table. Once the table creation is finished (be
// it successful or not) it signals a done on the end waitgroup.
func createTestTable(
	t *testing.T,
	tc *testcluster.TestCluster,
	id int,
	db *gosql.DB,
	wgStart *sync.WaitGroup,
	wgEnd *sync.WaitGroup,
	signal chan struct{},
	completed chan int,
) {
	defer wgEnd.Done()

	wgStart.Done()
	<-signal

	tableSQL := fmt.Sprintf(`
		CREATE TABLE IF NOT EXISTS "test"."table_%d" (
			id INT PRIMARY KEY,
			val INT
		)`, id)

	if _, err := db.Exec(tableSQL); err != nil {
		t.Errorf("table %d: could not be created: %s", id, err)
	} else {
		completed <- id
	}
}

// verifyTables ensures that the correct number of tables were created and that
// they all correspond to individual table descriptor IDs in the correct range
// of values.
func verifyTables(t *testing.T,
	tc *testcluster.TestCluster,
	completed chan int,
	expectedNumOfTables int,
	descIDStart int64,
) {
	descIDEnd := descIDStart + int64(expectedNumOfTables)
	usedTableIDs := make(map[sqlbase.ID]string)
	var count int
	for id := range completed {
		count++
		tableName := fmt.Sprintf("table_%d", id)
		kvDB := tc.Servers[count%tc.NumServers()].KVClient().(*client.DB)
		tableDesc := sqlbase.GetTableDescriptor(kvDB, "test", tableName)
		if int64(tableDesc.ID) < descIDStart || int64(tableDesc.ID) >= descIDEnd {
			t.Fatalf(
				"table %s's ID %d is not within the expected range of %d to %d",
				tableName,
				tableDesc.ID,
				descIDStart,
				descIDEnd,
			)
		}
		usedTableIDs[tableDesc.ID] = tableName
	}

	if e, a := expectedNumOfTables, len(usedTableIDs); e != a {
		t.Fatalf("expected %d tables created, only got %d", e, a)
	}

	kvDB := tc.Servers[count%tc.NumServers()].KVClient().(*client.DB)
	if descID, err := kvDB.Get(context.Background(), keys.DescIDGenerator); err != nil {
		t.Fatal(err)
	} else {
		if e, a := descIDEnd, descID.ValueInt(); e != a {
			t.Fatalf("expected next descriptor ID to be %d, got %d", e, a)
		}
	}
}

// TestParallelCreateTables tests that concurrent create table requests are
// correctly filled.
func TestParallelCreateTables(t *testing.T) {
	defer leaktest.AfterTest(t)()

	// This number has to be around 10 or else testrace will take too long to
	// finish.
	const numberOfTables = 10
	const numberOfNodes = 3

	tc := testcluster.StartTestCluster(t, numberOfNodes, base.TestClusterArgs{})
	defer tc.Stopper().Stop()
	if err := tc.WaitForFullReplication(); err != nil {
		t.Fatal(err)
	}

	if _, err := tc.ServerConn(0).Exec(`CREATE DATABASE "test"`); err != nil {
		t.Fatal(err)
	}
	// Get the id descriptor generator count.
	kvDB := tc.Servers[0].KVClient().(*client.DB)
	var descIDStart int64
	if descID, err := kvDB.Get(context.Background(), keys.DescIDGenerator); err != nil {
		t.Fatal(err)
	} else {
		descIDStart = descID.ValueInt()
	}

	var wgStart sync.WaitGroup
	var wgEnd sync.WaitGroup
	wgStart.Add(numberOfTables)
	wgEnd.Add(numberOfTables)
	signal := make(chan struct{})
	completed := make(chan int, numberOfTables)
	for i := 0; i < numberOfTables; i++ {
		db := tc.ServerConn(i % numberOfNodes)
		go createTestTable(t, tc, i, db, &wgStart, &wgEnd, signal, completed)
	}

	// Wait until all goroutines are ready.
	wgStart.Wait()
	// Signal the create table goroutines to start.
	close(signal)
	// Wait until all create tables are finished.
	wgEnd.Wait()
	close(completed)

	verifyTables(
		t,
		tc,
		completed,
		numberOfTables,
		descIDStart,
	)
}

// TestParallelCreateConflictingTables tests that concurrent create table
// requests with same name are only filled once. This is the same test as
// TestParallelCreateTables but in this test the tables names are all the same
// and is designed to specifically test the IF NOT EXIST clause.
func TestParallelCreateConflictingTables(t *testing.T) {
	defer leaktest.AfterTest(t)()

	const numberOfTables = 30
	const numberOfNodes = 3

	tc := testcluster.StartTestCluster(t, numberOfNodes, base.TestClusterArgs{})
	defer tc.Stopper().Stop()
	if err := tc.WaitForFullReplication(); err != nil {
		t.Fatal(err)
	}

	if _, err := tc.ServerConn(0).Exec(`CREATE DATABASE "test"`); err != nil {
		t.Fatal(err)
	}

	// Get the id descriptor generator count.
	kvDB := tc.Servers[0].KVClient().(*client.DB)
	var descIDStart int64
	if descID, err := kvDB.Get(context.Background(), keys.DescIDGenerator); err != nil {
		t.Fatal(err)
	} else {
		descIDStart = descID.ValueInt()
	}

	var wgStart sync.WaitGroup
	var wgEnd sync.WaitGroup
	wgStart.Add(numberOfTables)
	wgEnd.Add(numberOfTables)
	signal := make(chan struct{})
	completed := make(chan int, numberOfTables)
	for i := 0; i < numberOfTables; i++ {
		db := tc.ServerConn(i % numberOfNodes)
		go createTestTable(t, tc, 0, db, &wgStart, &wgEnd, signal, completed)
	}

	// Wait until all goroutines are ready.
	wgStart.Wait()
	// Signal the create table goroutines to start.
	close(signal)
	// Wait until all create tables are finished.
	wgEnd.Wait()
	close(completed)

	verifyTables(
		t,
		tc,
		completed,
		1,
		descIDStart,
	)
}
<|endoftext|>


<commit_before>/*
 * Copyright 2012-2016 Aerospike, Inc.
 *
 * Portions may be licensed to Aerospike, Inc. under one or more contributor
 * license agreements WHICH ARE COMPATIBLE WITH THE APACHE LICENSE, VERSION 2.0.
 *
 * Licensed under the Apache License, Version 2.0 (the "License"); you may not
 * use this file except in compliance with the License. You may obtain a copy of
 * the License at http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
 * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
 * License for the specific language governing permissions and limitations under
 * the License.
 */

package aerospike

import (
	"encoding/base64"
	"fmt"
	"strconv"

	. "github.com/aerospike/aerospike-client-go/logger"
	. "github.com/aerospike/aerospike-client-go/types"
)

const (
	_PartitionGeneration = "partition-generation"
	_ReplicasMaster      = "replicas-master"
	_ReplicasAll         = "replicas-all"
)

// Parse node's master (and optionally prole) partitions.
type partitionParser struct {
	pmap           partitionMap
	buffer         []byte
	partitionCount int
	generation     int
	length         int
	offset         int
	copied         bool
}

func newPartitionParser(conn *Connection, node *Node, pmap partitionMap, partitionCount int, requestProleReplicas bool) (*partitionParser, error) {
	newPartitionParser := &partitionParser{
		partitionCount: partitionCount,
		pmap:           pmap,
	}

	// Send format 1:  partition-generation\nreplicas-master\n
	// Send format 2:  partition-generation\nreplicas-all\n
	command := _ReplicasMaster
	if requestProleReplicas {
		command = _ReplicasAll
	}
	info, err := newInfo(conn, _PartitionGeneration, command)
	if err != nil {
		return nil, err
	}

	newPartitionParser.buffer = info.msg.Data
	newPartitionParser.length = len(info.msg.Data)
	if newPartitionParser.length == 0 {
		return nil, NewAerospikeError(PARSE_ERROR, fmt.Sprintf("Partition info is empty"))
	}

	newPartitionParser.generation, err = newPartitionParser.parseGeneration()
	if err != nil {
		return nil, err
	}

	if requestProleReplicas {
		err = newPartitionParser.parseReplicasAll(node)
	} else {
		err = newPartitionParser.parseReplicasMaster(node)
	}

	if err != nil {
		return nil, err
	}

	return newPartitionParser, nil
}

func (pp *partitionParser) getGeneration() int {
	return pp.generation
}

func (pp *partitionParser) isPartitionMapCopied() bool {
	return pp.copied
}

func (pp *partitionParser) getPartitionMap() partitionMap {
	return pp.pmap
}

func (pp *partitionParser) parseGeneration() (int, error) {
	if err := pp.expectName(_PartitionGeneration); err != nil {
		return -1, err
	}

	begin := pp.offset
	for pp.offset < pp.length {
		if pp.buffer[pp.offset] == '\n' {
			s := string(pp.buffer[begin:pp.offset])
			pp.offset++
			return strconv.Atoi(s)
		}
		pp.offset++
	}
	return -1, NewAerospikeError(PARSE_ERROR, fmt.Sprintf("Failed to find partition-generation value"))
}

func (pp *partitionParser) parseReplicasMaster(node *Node) error {
	// Use low-level info methods and parse byte array directly for maximum performance.
	// Receive format: replicas-master\t<ns1>:<base 64 encoded bitmap1>;<ns2>:<base 64 encoded bitmap2>...\n
	if err := pp.expectName(_ReplicasMaster); err != nil {
		return err
	}

	begin := pp.offset

	for pp.offset < pp.length {
		if pp.buffer[pp.offset] == ':' {
			// Parse namespace.
			namespace := string(pp.buffer[begin:pp.offset])

			if len(namespace) <= 0 || len(namespace) >= 32 {
				response := pp.getTruncatedResponse()
				return NewAerospikeError(PARSE_ERROR, fmt.Sprintf("Invalid partition namespace `%s` response: `%s`", namespace, response))
			}
			pp.offset++
			begin = pp.offset

			// Parse partition bitmap.
			for pp.offset < pp.length {
				b := pp.buffer[pp.offset]

				if b == ';' || b == '\n' {
					break
				}
				pp.offset++
			}

			if pp.offset == begin {
				response := pp.getTruncatedResponse()
				return NewAerospikeError(PARSE_ERROR, fmt.Sprintf("Empty partition id for namespace `%s` response: `%s`", namespace, response))
			}

			replicaArray := pp.pmap[namespace]

			if replicaArray == nil {
				replicaArray = make([][]*Node, 1)
				replicaArray[0] = make([]*Node, pp.partitionCount)
				pp.copyPartitionMap()
				pp.pmap[namespace] = replicaArray
			}

			if err := pp.decodeBitmap(node, replicaArray[0], begin); err != nil {
				return err
			}
			pp.offset++
			begin = pp.offset
		} else {
			pp.offset++
		}
	}

	return nil
}

func (pp *partitionParser) parseReplicasAll(node *Node) error {
	// Use low-level info methods and parse byte array directly for maximum performance.
	// Receive format: replicas-all\t
	//                 <ns1>:<count>,<base 64 encoded bitmap1>,<base 64 encoded bitmap2>...;
	//                 <ns2>:<count>,<base 64 encoded bitmap1>,<base 64 encoded bitmap2>...;\n
	if err := pp.expectName(_ReplicasAll); err != nil {
		return err
	}

	begin := pp.offset

	for pp.offset < pp.length {
		if pp.buffer[pp.offset] == ':' {
			// Parse namespace.
			namespace := string(pp.buffer[begin:pp.offset])

			if len(namespace) <= 0 || len(namespace) >= 32 {
				response := pp.getTruncatedResponse()
				return NewAerospikeError(PARSE_ERROR, fmt.Sprintf("Invalid partition namespace `%s` response: `%s`", namespace, response))
			}
			pp.offset++
			begin = pp.offset

			// Parse replica count.
			for pp.offset < pp.length {
				b := pp.buffer[pp.offset]

				if b == ',' {
					break
				}
				pp.offset++
			}

			replicaCount, err := strconv.Atoi(string(pp.buffer[begin:pp.offset]))
			if err != nil {
				return err
			}

			// Ensure replicaArray is correct size.
			replicaArray := pp.pmap[namespace]

			if replicaArray == nil {
				// Create new replica array.
				replicaArray = make([][]*Node, replicaCount)

				for i := 0; i < replicaCount; i++ {
					replicaArray[i] = make([]*Node, pp.partitionCount)
				}

				pp.copyPartitionMap()
				pp.pmap[namespace] = replicaArray
			} else if len(replicaArray) != replicaCount {
				Logger.Info("Namespace `%s` replication factor changed from `%d` to `%d` ", namespace, len(replicaArray), replicaCount)

				// Resize replica array.
				replicaTarget := make([][]*Node, replicaCount)

				if len(replicaArray) < replicaCount {
					i := 0

					// Copy existing entries.
					for ; i < len(replicaArray); i++ {
						replicaTarget[i] = replicaArray[i]
					}

					// Create new entries.
					for ; i < replicaCount; i++ {
						replicaTarget[i] = make([]*Node, pp.partitionCount)
					}
				} else {
					// Copy existing entries.
					for i := 0; i < replicaCount; i++ {
						replicaTarget[i] = replicaArray[i]
					}
				}

				pp.copyPartitionMap()
				replicaArray = replicaTarget
				pp.pmap[namespace] = replicaArray
			}

			// Parse partition bitmaps.
			for i := 0; i < replicaCount; i++ {
				pp.offset++
				begin = pp.offset

				// Find bitmap endpoint
				for pp.offset < pp.length {
					b := pp.buffer[pp.offset]

					if b == ',' || b == ';' {
						break
					}
					pp.offset++
				}

				if pp.offset == begin {
					response := pp.getTruncatedResponse()
					return NewAerospikeError(PARSE_ERROR, fmt.Sprintf("Empty partition id for namespace `%s` response: `%s`", namespace, response))
				}

				if err := pp.decodeBitmap(node, replicaArray[i], begin); err != nil {
					return err
				}
			}
			pp.offset++
			begin = pp.offset
		} else {
			pp.offset++
		}
	}

	return nil
}

func (pp *partitionParser) decodeBitmap(node *Node, nodeArray []*Node, begin int) error {
	pp.copyPartitionMap()

	restoreBuffer, err := base64.StdEncoding.DecodeString(string(pp.buffer[begin:pp.offset]))
	if err != nil {
		return err
	}

	for i := 0; i < pp.partitionCount; i++ {
		nodeOld := nodeArray[i]

		if (restoreBuffer[i>>3] & (0x80 >> uint(i&7))) != 0 {
			// Node owns this partition.
			if nodeOld != nil && nodeOld != node {
				// Force previously mapped node to refresh it's partition map on next cluster tend.
				nodeOld.partitionGeneration.Set(-1)
			}

			// Use lazy set because there is only one producer thread. In addition,
			// there is a one second delay due to the cluster tend polling interval.
			// An extra millisecond for a node change will not make a difference and
			// overall performance is improved.
			nodeArray[i] = node
		} else {
			// Node does not own partition.
			if node == nodeOld {
				// Must erase previous map.
				nodeArray[i] = nil
			}
		}
	}

	return nil
}

func (pp *partitionParser) copyPartitionMap() {
	if !pp.copied {
		// Make shallow copy of map.
		pmap := make(partitionMap, len(pp.pmap))
		for ns, replArr := range pp.pmap {
			newReplArr := make([][]*Node, len(replArr))
			for i, nArr := range replArr {
				newNArr := make([]*Node, len(nArr))
				copy(newNArr, nArr)
				newReplArr[i] = newNArr
			}
			pmap[ns] = newReplArr
		}
		pp.pmap = pmap
		pp.copied = true
	}
}

func (pp *partitionParser) expectName(name string) error {
	begin := pp.offset

	for pp.offset < pp.length {
		if pp.buffer[pp.offset] == '\t' {
			s := string(pp.buffer[begin:pp.offset])
			if name == s {
				pp.offset++
				return nil
			}
			break
		}
		pp.offset++
	}

	return NewAerospikeError(PARSE_ERROR, fmt.Sprintf("Failed to find `%s`", name))
}

func (pp *partitionParser) getTruncatedResponse() string {
	max := pp.length
	if max > 200 {
		max = 200
	}
	return string(pp.buffer[0:max])
}
<commit_msg>Always copy the partitionMap in PartitionParser to avoid Data Race<commit_after>/*
 * Copyright 2012-2016 Aerospike, Inc.
 *
 * Portions may be licensed to Aerospike, Inc. under one or more contributor
 * license agreements WHICH ARE COMPATIBLE WITH THE APACHE LICENSE, VERSION 2.0.
 *
 * Licensed under the Apache License, Version 2.0 (the "License"); you may not
 * use this file except in compliance with the License. You may obtain a copy of
 * the License at http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
 * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
 * License for the specific language governing permissions and limitations under
 * the License.
 */

package aerospike

import (
	"encoding/base64"
	"fmt"
	"strconv"

	. "github.com/aerospike/aerospike-client-go/logger"
	. "github.com/aerospike/aerospike-client-go/types"
)

const (
	_PartitionGeneration = "partition-generation"
	_ReplicasMaster      = "replicas-master"
	_ReplicasAll         = "replicas-all"
)

// Parse node's master (and optionally prole) partitions.
type partitionParser struct {
	pmap           partitionMap
	buffer         []byte
	partitionCount int
	generation     int
	length         int
	offset         int
	copied         bool
}

func newPartitionParser(conn *Connection, node *Node, pmap partitionMap, partitionCount int, requestProleReplicas bool) (*partitionParser, error) {
	newPartitionParser := &partitionParser{
		partitionCount: partitionCount,
		pmap:           pmap,
	}

	// Send format 1:  partition-generation\nreplicas-master\n
	// Send format 2:  partition-generation\nreplicas-all\n
	command := _ReplicasMaster
	if requestProleReplicas {
		command = _ReplicasAll
	}
	info, err := newInfo(conn, _PartitionGeneration, command)
	if err != nil {
		return nil, err
	}

	newPartitionParser.buffer = info.msg.Data
	newPartitionParser.length = len(info.msg.Data)
	if newPartitionParser.length == 0 {
		return nil, NewAerospikeError(PARSE_ERROR, fmt.Sprintf("Partition info is empty"))
	}

	newPartitionParser.generation, err = newPartitionParser.parseGeneration()
	if err != nil {
		return nil, err
	}

	// always copy the partition Map
	newPartitionParser.copyPartitionMap()

	if requestProleReplicas {
		err = newPartitionParser.parseReplicasAll(node)
	} else {
		err = newPartitionParser.parseReplicasMaster(node)
	}

	if err != nil {
		return nil, err
	}

	return newPartitionParser, nil
}

func (pp *partitionParser) getGeneration() int {
	return pp.generation
}

func (pp *partitionParser) isPartitionMapCopied() bool {
	return pp.copied
}

func (pp *partitionParser) getPartitionMap() partitionMap {
	return pp.pmap
}

func (pp *partitionParser) parseGeneration() (int, error) {
	if err := pp.expectName(_PartitionGeneration); err != nil {
		return -1, err
	}

	begin := pp.offset
	for pp.offset < pp.length {
		if pp.buffer[pp.offset] == '\n' {
			s := string(pp.buffer[begin:pp.offset])
			pp.offset++
			return strconv.Atoi(s)
		}
		pp.offset++
	}
	return -1, NewAerospikeError(PARSE_ERROR, fmt.Sprintf("Failed to find partition-generation value"))
}

func (pp *partitionParser) parseReplicasMaster(node *Node) error {
	// Use low-level info methods and parse byte array directly for maximum performance.
	// Receive format: replicas-master\t<ns1>:<base 64 encoded bitmap1>;<ns2>:<base 64 encoded bitmap2>...\n
	if err := pp.expectName(_ReplicasMaster); err != nil {
		return err
	}

	begin := pp.offset

	for pp.offset < pp.length {
		if pp.buffer[pp.offset] == ':' {
			// Parse namespace.
			namespace := string(pp.buffer[begin:pp.offset])

			if len(namespace) <= 0 || len(namespace) >= 32 {
				response := pp.getTruncatedResponse()
				return NewAerospikeError(PARSE_ERROR, fmt.Sprintf("Invalid partition namespace `%s` response: `%s`", namespace, response))
			}
			pp.offset++
			begin = pp.offset

			// Parse partition bitmap.
			for pp.offset < pp.length {
				b := pp.buffer[pp.offset]

				if b == ';' || b == '\n' {
					break
				}
				pp.offset++
			}

			if pp.offset == begin {
				response := pp.getTruncatedResponse()
				return NewAerospikeError(PARSE_ERROR, fmt.Sprintf("Empty partition id for namespace `%s` response: `%s`", namespace, response))
			}

			replicaArray := pp.pmap[namespace]

			if replicaArray == nil {
				replicaArray = make([][]*Node, 1)
				replicaArray[0] = make([]*Node, pp.partitionCount)
				pp.pmap[namespace] = replicaArray
			}

			if err := pp.decodeBitmap(node, replicaArray[0], begin); err != nil {
				return err
			}
			pp.offset++
			begin = pp.offset
		} else {
			pp.offset++
		}
	}

	return nil
}

func (pp *partitionParser) parseReplicasAll(node *Node) error {
	// Use low-level info methods and parse byte array directly for maximum performance.
	// Receive format: replicas-all\t
	//                 <ns1>:<count>,<base 64 encoded bitmap1>,<base 64 encoded bitmap2>...;
	//                 <ns2>:<count>,<base 64 encoded bitmap1>,<base 64 encoded bitmap2>...;\n
	if err := pp.expectName(_ReplicasAll); err != nil {
		return err
	}

	begin := pp.offset

	for pp.offset < pp.length {
		if pp.buffer[pp.offset] == ':' {
			// Parse namespace.
			namespace := string(pp.buffer[begin:pp.offset])

			if len(namespace) <= 0 || len(namespace) >= 32 {
				response := pp.getTruncatedResponse()
				return NewAerospikeError(PARSE_ERROR, fmt.Sprintf("Invalid partition namespace `%s` response: `%s`", namespace, response))
			}
			pp.offset++
			begin = pp.offset

			// Parse replica count.
			for pp.offset < pp.length {
				b := pp.buffer[pp.offset]

				if b == ',' {
					break
				}
				pp.offset++
			}

			replicaCount, err := strconv.Atoi(string(pp.buffer[begin:pp.offset]))
			if err != nil {
				return err
			}

			// Ensure replicaArray is correct size.
			replicaArray := pp.pmap[namespace]

			if replicaArray == nil {
				// Create new replica array.
				replicaArray = make([][]*Node, replicaCount)

				for i := 0; i < replicaCount; i++ {
					replicaArray[i] = make([]*Node, pp.partitionCount)
				}

				pp.pmap[namespace] = replicaArray
			} else if len(replicaArray) != replicaCount {
				Logger.Info("Namespace `%s` replication factor changed from `%d` to `%d` ", namespace, len(replicaArray), replicaCount)

				// Resize replica array.
				replicaTarget := make([][]*Node, replicaCount)

				if len(replicaArray) < replicaCount {
					i := 0

					// Copy existing entries.
					for ; i < len(replicaArray); i++ {
						replicaTarget[i] = replicaArray[i]
					}

					// Create new entries.
					for ; i < replicaCount; i++ {
						replicaTarget[i] = make([]*Node, pp.partitionCount)
					}
				} else {
					// Copy existing entries.
					for i := 0; i < replicaCount; i++ {
						replicaTarget[i] = replicaArray[i]
					}
				}

				replicaArray = replicaTarget
				pp.pmap[namespace] = replicaArray
			}

			// Parse partition bitmaps.
			for i := 0; i < replicaCount; i++ {
				pp.offset++
				begin = pp.offset

				// Find bitmap endpoint
				for pp.offset < pp.length {
					b := pp.buffer[pp.offset]

					if b == ',' || b == ';' {
						break
					}
					pp.offset++
				}

				if pp.offset == begin {
					response := pp.getTruncatedResponse()
					return NewAerospikeError(PARSE_ERROR, fmt.Sprintf("Empty partition id for namespace `%s` response: `%s`", namespace, response))
				}

				if err := pp.decodeBitmap(node, replicaArray[i], begin); err != nil {
					return err
				}
			}
			pp.offset++
			begin = pp.offset
		} else {
			pp.offset++
		}
	}

	return nil
}

func (pp *partitionParser) decodeBitmap(node *Node, nodeArray []*Node, begin int) error {
	restoreBuffer, err := base64.StdEncoding.DecodeString(string(pp.buffer[begin:pp.offset]))
	if err != nil {
		return err
	}

	for i := 0; i < pp.partitionCount; i++ {
		nodeOld := nodeArray[i]

		if (restoreBuffer[i>>3] & (0x80 >> uint(i&7))) != 0 {
			// Node owns this partition.
			if nodeOld != nil && nodeOld != node {
				// Force previously mapped node to refresh it's partition map on next cluster tend.
				nodeOld.partitionGeneration.Set(-1)
			}

			// Use lazy set because there is only one producer thread. In addition,
			// there is a one second delay due to the cluster tend polling interval.
			// An extra millisecond for a node change will not make a difference and
			// overall performance is improved.
			nodeArray[i] = node
		} else {
			// Node does not own partition.
			if node == nodeOld {
				// Must erase previous map.
				nodeArray[i] = nil
			}
		}
	}

	return nil
}

func (pp *partitionParser) copyPartitionMap() {
	if !pp.copied {
		// Make shallow copy of map.
		pmap := make(partitionMap, len(pp.pmap))
		for ns, replArr := range pp.pmap {
			newReplArr := make([][]*Node, len(replArr))
			for i, nArr := range replArr {
				newNArr := make([]*Node, len(nArr))
				copy(newNArr, nArr)
				newReplArr[i] = newNArr
			}
			pmap[ns] = newReplArr
		}
		pp.pmap = pmap
		pp.copied = true
	}
}

func (pp *partitionParser) expectName(name string) error {
	begin := pp.offset

	for pp.offset < pp.length {
		if pp.buffer[pp.offset] == '\t' {
			s := string(pp.buffer[begin:pp.offset])
			if name == s {
				pp.offset++
				return nil
			}
			break
		}
		pp.offset++
	}

	return NewAerospikeError(PARSE_ERROR, fmt.Sprintf("Failed to find `%s`", name))
}

func (pp *partitionParser) getTruncatedResponse() string {
	max := pp.length
	if max > 200 {
		max = 200
	}
	return string(pp.buffer[0:max])
}
<|endoftext|>


<commit_before>package kubernetes

import (
	"encoding/json"
	"reflect"
	"sort"
	"strings"
)

func diffStringMap(pathPrefix string, oldV, newV map[string]interface{}) PatchOperations {
	ops := make([]PatchOperation, 0, 0)

	pathPrefix = strings.TrimRight(pathPrefix, "/")

	// This is suboptimal for adding whole new map from scratch
	// or deleting the whole map, but it's actually intention.
	// There may be some other map items managed outside of TF
	// and we don't want to touch these.

	for k, _ := range oldV {
		if _, ok := newV[k]; ok {
			continue
		}
		ops = append(ops, &RemoveOperation{Path: pathPrefix + "/" + k})
	}

	for k, v := range newV {
		newValue := v.(string)

		if oldValue, ok := oldV[k].(string); ok {
			if oldValue == newValue {
				continue
			}

			ops = append(ops, &ReplaceOperation{
				Path:  pathPrefix + "/" + k,
				Value: newValue,
			})
			continue
		}

		ops = append(ops, &AddOperation{
			Path:  pathPrefix + "/" + k,
			Value: newValue,
		})
	}

	return ops
}

type PatchOperations []PatchOperation

func (po PatchOperations) MarshalJSON() ([]byte, error) {
	var v []PatchOperation = po
	return json.Marshal(v)
}

func (po PatchOperations) Equal(ops []PatchOperation) bool {
	var v []PatchOperation = po

	sort.Slice(v, sortByPathAsc(ops))
	sort.Slice(ops, sortByPathAsc(ops))

	return reflect.DeepEqual(v, ops)
}

func sortByPathAsc(ops []PatchOperation) func(i, j int) bool {
	return func(i, j int) bool {
		return ops[i].GetPath() < ops[j].GetPath()
	}
}

type PatchOperation interface {
	MarshalJSON() ([]byte, error)
	GetPath() string
}

type ReplaceOperation struct {
	Path  string      `json:"path"`
	Value interface{} `json:"value"`
	Op    string      `json:"op"`
}

func (o *ReplaceOperation) GetPath() string {
	return o.Path
}

func (o *ReplaceOperation) MarshalJSON() ([]byte, error) {
	o.Op = "replace"
	return json.Marshal(*o)
}

func (o *ReplaceOperation) String() string {
	b, _ := o.MarshalJSON()
	return string(b)
}

type AddOperation struct {
	Path  string      `json:"path"`
	Value interface{} `json:"value"`
	Op    string      `json:"op"`
}

func (o *AddOperation) GetPath() string {
	return o.Path
}

func (o *AddOperation) MarshalJSON() ([]byte, error) {
	o.Op = "add"
	return json.Marshal(*o)
}

func (o *AddOperation) String() string {
	b, _ := o.MarshalJSON()
	return string(b)
}

type RemoveOperation struct {
	Path string `json:"path"`
	Op   string `json:"op"`
}

func (o *RemoveOperation) GetPath() string {
	return o.Path
}

func (o *RemoveOperation) MarshalJSON() ([]byte, error) {
	o.Op = "remove"
	return json.Marshal(*o)
}

func (o *RemoveOperation) String() string {
	b, _ := o.MarshalJSON()
	return string(b)
}
<commit_msg>provider/kubernetes: Fix sorting of the PatchOperations (#13117)<commit_after>package kubernetes

import (
	"encoding/json"
	"reflect"
	"sort"
	"strings"
)

func diffStringMap(pathPrefix string, oldV, newV map[string]interface{}) PatchOperations {
	ops := make([]PatchOperation, 0, 0)

	pathPrefix = strings.TrimRight(pathPrefix, "/")

	// This is suboptimal for adding whole new map from scratch
	// or deleting the whole map, but it's actually intention.
	// There may be some other map items managed outside of TF
	// and we don't want to touch these.

	for k, _ := range oldV {
		if _, ok := newV[k]; ok {
			continue
		}
		ops = append(ops, &RemoveOperation{Path: pathPrefix + "/" + k})
	}

	for k, v := range newV {
		newValue := v.(string)

		if oldValue, ok := oldV[k].(string); ok {
			if oldValue == newValue {
				continue
			}

			ops = append(ops, &ReplaceOperation{
				Path:  pathPrefix + "/" + k,
				Value: newValue,
			})
			continue
		}

		ops = append(ops, &AddOperation{
			Path:  pathPrefix + "/" + k,
			Value: newValue,
		})
	}

	return ops
}

type PatchOperations []PatchOperation

func (po PatchOperations) MarshalJSON() ([]byte, error) {
	var v []PatchOperation = po
	return json.Marshal(v)
}

func (po PatchOperations) Equal(ops []PatchOperation) bool {
	var v []PatchOperation = po

	sort.Slice(v, sortByPathAsc(v))
	sort.Slice(ops, sortByPathAsc(ops))

	return reflect.DeepEqual(v, ops)
}

func sortByPathAsc(ops []PatchOperation) func(i, j int) bool {
	return func(i, j int) bool {
		return ops[i].GetPath() < ops[j].GetPath()
	}
}

type PatchOperation interface {
	MarshalJSON() ([]byte, error)
	GetPath() string
}

type ReplaceOperation struct {
	Path  string      `json:"path"`
	Value interface{} `json:"value"`
	Op    string      `json:"op"`
}

func (o *ReplaceOperation) GetPath() string {
	return o.Path
}

func (o *ReplaceOperation) MarshalJSON() ([]byte, error) {
	o.Op = "replace"
	return json.Marshal(*o)
}

func (o *ReplaceOperation) String() string {
	b, _ := o.MarshalJSON()
	return string(b)
}

type AddOperation struct {
	Path  string      `json:"path"`
	Value interface{} `json:"value"`
	Op    string      `json:"op"`
}

func (o *AddOperation) GetPath() string {
	return o.Path
}

func (o *AddOperation) MarshalJSON() ([]byte, error) {
	o.Op = "add"
	return json.Marshal(*o)
}

func (o *AddOperation) String() string {
	b, _ := o.MarshalJSON()
	return string(b)
}

type RemoveOperation struct {
	Path string `json:"path"`
	Op   string `json:"op"`
}

func (o *RemoveOperation) GetPath() string {
	return o.Path
}

func (o *RemoveOperation) MarshalJSON() ([]byte, error) {
	o.Op = "remove"
	return json.Marshal(*o)
}

func (o *RemoveOperation) String() string {
	b, _ := o.MarshalJSON()
	return string(b)
}
<|endoftext|>


<commit_before>package etcd

import (
	"log"
	"time"

	"strings"

	"github.com/coreos/etcd/client"
	"github.com/fsouza/go-dockerclient/external/golang.org/x/net/context"
	"github.com/vsco/dcdr/cli/api/stores"
	"github.com/vsco/dcdr/config"
)

type ETCDStore struct {
	kv         client.KeysAPI
	ctx        context.Context
	getOpts    *client.GetOptions
	setOpts    *client.SetOptions
	deleteOpts *client.DeleteOptions
}

var DefaultEndpoints = []string{"http://127.0.0.1:2379"}

func DefaultETCDlStore(cfg *config.Config) (client.KeysAPI, error) {
	endpoints := DefaultEndpoints

	if len(cfg.Etcd.Endpoints) > 0 {
		endpoints = cfg.Etcd.Endpoints
	}

	ecfg := client.Config{
		Endpoints:               endpoints,
		Transport:               client.DefaultTransport,
		HeaderTimeoutPerRequest: time.Second,
	}

	c, err := client.New(ecfg)

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

	return client.NewKeysAPI(c), nil
}

func New(cfg *config.Config) stores.StoreIFace {
	kv, _ := DefaultETCDlStore(cfg)

	es := &ETCDStore{
		kv:  kv,
		ctx: context.Background(),
	}

	return es
}

func (s *ETCDStore) Get(key string) (*stores.KVByte, error) {
	resp, err := s.kv.Get(s.ctx, key, s.getOpts)

	if err != nil {
		return nil, etcdError(err)
	}

	return toKVByte(resp.Node), nil
}

func (s *ETCDStore) Set(key string, bts []byte) error {
	_, err := s.kv.Set(s.ctx, key, string(bts), s.setOpts)

	return err
}

func (s *ETCDStore) Delete(key string) error {
	_, err := s.kv.Delete(s.ctx, key, s.deleteOpts)

	return etcdError(err)
}

func (s *ETCDStore) List(prefix string) (stores.KVBytes, error) {
	opts := &client.GetOptions{
		Recursive: true,
		Sort:      true,
		Quorum:    true,
	}

	resp, err := s.kv.Get(s.ctx, prefix, opts)

	if err != nil {
		return nil, etcdError(err)
	}

	kvbs := FlattenToKVBytes(resp.Node, make(stores.KVBytes, 0))

	return kvbs, nil
}

func FlattenToKVBytes(n *client.Node, nodes stores.KVBytes) stores.KVBytes {
	if n.Dir {
		for _, nd := range n.Nodes {
			nodes = FlattenToKVBytes(nd, nodes)
		}
	} else {
		nodes = append(nodes, toKVByte(n))
	}

	return nodes
}

func toKVByte(n *client.Node) *stores.KVByte {
	return &stores.KVByte{
		// remove leading slash as it adds an empty
		// hash entry when exploded to JSON.
		Key:   strings.TrimPrefix(n.Key, "/"),
		Bytes: []byte(n.Value),
	}
}

func etcdError(err error) error {
	switch err.(type) {
	case client.Error:
		if err.(client.Error).Code == client.ErrorCodeKeyNotFound {
			return nil
		}

		return err
	default:
		return err
	}
}
<commit_msg>vendoring is still terrible<commit_after>package etcd

import (
	"log"
	"time"

	"strings"

	"github.com/coreos/etcd/client"
	"github.com/vsco/dcdr/cli/api/stores"
	"github.com/vsco/dcdr/config"
	"golang.org/x/net/context"
)

type ETCDStore struct {
	kv         client.KeysAPI
	ctx        context.Context
	getOpts    *client.GetOptions
	setOpts    *client.SetOptions
	deleteOpts *client.DeleteOptions
}

var DefaultEndpoints = []string{"http://127.0.0.1:2379"}

func DefaultETCDlStore(cfg *config.Config) (client.KeysAPI, error) {
	endpoints := DefaultEndpoints

	if len(cfg.Etcd.Endpoints) > 0 {
		endpoints = cfg.Etcd.Endpoints
	}

	ecfg := client.Config{
		Endpoints:               endpoints,
		Transport:               client.DefaultTransport,
		HeaderTimeoutPerRequest: time.Second,
	}

	c, err := client.New(ecfg)

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

	return client.NewKeysAPI(c), nil
}

func New(cfg *config.Config) stores.StoreIFace {
	kv, _ := DefaultETCDlStore(cfg)

	es := &ETCDStore{
		kv:  kv,
		ctx: context.Background(),
	}

	return es
}

func (s *ETCDStore) Get(key string) (*stores.KVByte, error) {
	resp, err := s.kv.Get(s.ctx, key, s.getOpts)

	if err != nil {
		return nil, etcdError(err)
	}

	return toKVByte(resp.Node), nil
}

func (s *ETCDStore) Set(key string, bts []byte) error {
	_, err := s.kv.Set(s.ctx, key, string(bts), s.setOpts)

	return err
}

func (s *ETCDStore) Delete(key string) error {
	_, err := s.kv.Delete(s.ctx, key, s.deleteOpts)

	return etcdError(err)
}

func (s *ETCDStore) List(prefix string) (stores.KVBytes, error) {
	opts := &client.GetOptions{
		Recursive: true,
		Sort:      true,
		Quorum:    true,
	}

	resp, err := s.kv.Get(s.ctx, prefix, opts)

	if err != nil {
		return nil, etcdError(err)
	}

	kvbs := FlattenToKVBytes(resp.Node, make(stores.KVBytes, 0))

	return kvbs, nil
}

func FlattenToKVBytes(n *client.Node, nodes stores.KVBytes) stores.KVBytes {
	if n.Dir {
		for _, nd := range n.Nodes {
			nodes = FlattenToKVBytes(nd, nodes)
		}
	} else {
		nodes = append(nodes, toKVByte(n))
	}

	return nodes
}

func toKVByte(n *client.Node) *stores.KVByte {
	return &stores.KVByte{
		// remove leading slash as it adds an empty
		// hash entry when exploded to JSON.
		Key:   strings.TrimPrefix(n.Key, "/"),
		Bytes: []byte(n.Value),
	}
}

func etcdError(err error) error {
	switch err.(type) {
	case client.Error:
		if err.(client.Error).Code == client.ErrorCodeKeyNotFound {
			return nil
		}

		return err
	default:
		return err
	}
}
<|endoftext|>


<commit_before>package sendemail

import "golang.org/x/crypto/openpgp"
import "golang.org/x/crypto/openpgp/armor"
import "golang.org/x/crypto/openpgp/packet"
import "fmt"
import "bytes"
import "net/mail"
import "net/smtp"
import "os"
import "os/exec"
import "net"
import "github.com/hlandau/degoutils/metric"

var cEmailsSent = metric.NewCounter("sendemail.emailsSent")

type Config struct {
	SMTPAddress  string
	SMTPUsername string
	SMTPPassword string

	SendmailPath string
}

type Email struct {
	From             string
	To               []string
	Headers          mail.Header
	Body             string
	OpenPGPEncryptTo []string

	rfc822Message []byte
}

type Sender interface {
	SendEmail(e *Email) error
}

func (cfg *Config) SendEmail(e *Email) error {
	if len(e.From) == 0 {
		return fmt.Errorf("from address must be specified")
	}

	if len(e.To) == 0 {
		return fmt.Errorf("at least one recipient must be specified")
	}

	if e.Headers == nil {
		e.Headers = make(map[string][]string)
	}

	if _, ok := e.Headers["From"]; !ok {
		e.Headers["From"] = []string{e.From}
	}

	if _, ok := e.Headers["To"]; !ok {
		e.Headers["To"] = e.To
	}

	err := encryptEmail(e)
	if err != nil {
		return err
	}

	e.rfc822Message = nil
	e.rfc822Message = append(e.rfc822Message, serializeHeaders(e.Headers)...)
	e.rfc822Message = append(e.rfc822Message, '\n')
	e.rfc822Message = append(e.rfc822Message, e.Body...)

	cEmailsSent.Add(1)

	if cfg.SMTPAddress != "" {
		return cfg.sendViaSMTP(e)
	}

	return cfg.sendViaSendmail(e)
}

func serializeHeaders(headers mail.Header) (b []byte) {
	for k, v := range headers {
		// XXX
		for _, x := range v {
			b = append(b, []byte(fmt.Sprintf("%s: %s\n", k, x))...)
		}
	}
	return
}

func encryptEmail(e *Email) error {
	if len(e.OpenPGPEncryptTo) == 0 {
		return nil
	}

	buf := &bytes.Buffer{}
	var destEntities []*openpgp.Entity
	for _, eto := range e.OpenPGPEncryptTo {
		r := bytes.NewBufferString(eto)
		blk, err := armor.Decode(r)
		if err != nil {
			return err
		}

		rr := packet.NewReader(blk.Body)
		e, err := openpgp.ReadEntity(rr)
		if err != nil {
			return err
		}

		destEntities = append(destEntities, e)
	}

	aew, err := armor.Encode(buf, "PGP MESSAGE", map[string]string{
		"Version": "OpenPGP",
	})
	if err != nil {
		return err
	}

	wr, err := openpgp.Encrypt(aew, destEntities, nil, nil, nil)
	if err != nil {
		return err
	}

	_, err = wr.Write([]byte(e.Body))
	if err != nil {
		wr.Close()
		return err
	}

	wr.Close()
	aew.Close()

	e.Body = string(buf.Bytes())
	return nil
}

func (cfg *Config) sendViaSMTP(e *Email) error {
	var auth smtp.Auth

	if cfg.SMTPUsername != "" {
		host, _, err := net.SplitHostPort(cfg.SMTPAddress)
		if err != nil {
			return err
		}

		auth = smtp.PlainAuth("", cfg.SMTPUsername, cfg.SMTPPassword, host)
	}

	return smtp.SendMail(cfg.SMTPAddress, auth, e.From, e.To, e.rfc822Message)
}

func (cfg *Config) sendViaSendmail(e *Email) error {
	smargs := []string{"-i"}
	smargs = append(smargs, e.To...)

	spath := cfg.SendmailPath
	if spath == "" {
		spath = "/usr/sbin/sendmail"
	}

	sm := exec.Command(spath, smargs...)
	stdin, err := sm.StdinPipe()
	if err != nil {
		return err
	}

	sm.Stdout = os.Stdout
	sm.Stderr = os.Stderr
	err = sm.Start()
	if err != nil {
		return err
	}

	stdin.Write(e.rfc822Message)
	stdin.Close()

	err = sm.Wait()
	if err != nil {
		return err
	}

	return nil
}
<commit_msg>sendemail<commit_after>package sendemail

import "golang.org/x/crypto/openpgp"
import "golang.org/x/crypto/openpgp/armor"
import "golang.org/x/crypto/openpgp/packet"
import "fmt"
import "bytes"
import "net/mail"
import "net/smtp"
import "os"
import "os/exec"
import "net"
import "gopkg.in/hlandau/easymetric.v1/cexp"

var cEmailsSent = cexp.NewCounter("sendemail.emailsSent")

type Config struct {
	SMTPAddress  string
	SMTPUsername string
	SMTPPassword string

	SendmailPath string
}

type Email struct {
	From             string
	To               []string
	Headers          mail.Header
	Body             string
	OpenPGPEncryptTo []string

	rfc822Message []byte
}

type Sender interface {
	SendEmail(e *Email) error
}

func (cfg *Config) SendEmail(e *Email) error {
	if len(e.From) == 0 {
		return fmt.Errorf("from address must be specified")
	}

	if len(e.To) == 0 {
		return fmt.Errorf("at least one recipient must be specified")
	}

	if e.Headers == nil {
		e.Headers = make(map[string][]string)
	}

	if _, ok := e.Headers["From"]; !ok {
		e.Headers["From"] = []string{e.From}
	}

	if _, ok := e.Headers["To"]; !ok {
		e.Headers["To"] = e.To
	}

	err := encryptEmail(e)
	if err != nil {
		return err
	}

	e.rfc822Message = nil
	e.rfc822Message = append(e.rfc822Message, serializeHeaders(e.Headers)...)
	e.rfc822Message = append(e.rfc822Message, '\n')
	e.rfc822Message = append(e.rfc822Message, e.Body...)

	cEmailsSent.Add(1)

	if cfg.SMTPAddress != "" {
		return cfg.sendViaSMTP(e)
	}

	return cfg.sendViaSendmail(e)
}

func serializeHeaders(headers mail.Header) (b []byte) {
	for k, v := range headers {
		// XXX
		for _, x := range v {
			b = append(b, []byte(fmt.Sprintf("%s: %s\n", k, x))...)
		}
	}
	return
}

func encryptEmail(e *Email) error {
	if len(e.OpenPGPEncryptTo) == 0 {
		return nil
	}

	buf := &bytes.Buffer{}
	var destEntities []*openpgp.Entity
	for _, eto := range e.OpenPGPEncryptTo {
		r := bytes.NewBufferString(eto)
		blk, err := armor.Decode(r)
		if err != nil {
			return err
		}

		rr := packet.NewReader(blk.Body)
		e, err := openpgp.ReadEntity(rr)
		if err != nil {
			return err
		}

		destEntities = append(destEntities, e)
	}

	aew, err := armor.Encode(buf, "PGP MESSAGE", map[string]string{
		"Version": "OpenPGP",
	})
	if err != nil {
		return err
	}

	wr, err := openpgp.Encrypt(aew, destEntities, nil, nil, nil)
	if err != nil {
		return err
	}

	_, err = wr.Write([]byte(e.Body))
	if err != nil {
		wr.Close()
		return err
	}

	wr.Close()
	aew.Close()

	e.Body = string(buf.Bytes())
	return nil
}

func (cfg *Config) sendViaSMTP(e *Email) error {
	var auth smtp.Auth

	if cfg.SMTPUsername != "" {
		host, _, err := net.SplitHostPort(cfg.SMTPAddress)
		if err != nil {
			return err
		}

		auth = smtp.PlainAuth("", cfg.SMTPUsername, cfg.SMTPPassword, host)
	}

	return smtp.SendMail(cfg.SMTPAddress, auth, e.From, e.To, e.rfc822Message)
}

func (cfg *Config) sendViaSendmail(e *Email) error {
	smargs := []string{"-i"}
	smargs = append(smargs, e.To...)

	spath := cfg.SendmailPath
	if spath == "" {
		spath = "/usr/sbin/sendmail"
	}

	sm := exec.Command(spath, smargs...)
	stdin, err := sm.StdinPipe()
	if err != nil {
		return err
	}

	sm.Stdout = os.Stdout
	sm.Stderr = os.Stderr
	err = sm.Start()
	if err != nil {
		return err
	}

	stdin.Write(e.rfc822Message)
	stdin.Close()

	err = sm.Wait()
	if err != nil {
		return err
	}

	return nil
}
<|endoftext|>