<commit_before>// Simple console progress bars
package pb

import (
	"fmt"
	"io"
	"math"
	"strings"
	"sync"
	"sync/atomic"
	"time"
	"unicode/utf8"
)

const (
	// Default refresh rate - 200ms
	DEFAULT_REFRESH_RATE = time.Millisecond * 200
	FORMAT               = "[=>-]"
)

// Create new progress bar object
func New(total int) *ProgressBar {
	return New64(int64(total))
}

// Create new progress bar object uding int64 as total
func New64(total int64) *ProgressBar {
	pb := &ProgressBar{
		Total:         total,
		RefreshRate:   DEFAULT_REFRESH_RATE,
		ShowPercent:   true,
		ShowCounters:  true,
		ShowBar:       true,
		ShowTimeLeft:  true,
		ShowFinalTime: true,
		Units:         U_NO,
		ManualUpdate:  false,

		finish:       make(chan struct{}),
		currentValue: -1,
		scale:        1.0,
		ewma:         &EWMA{},
		mu:           new(sync.Mutex),
	}
	return pb.Format(FORMAT)
}

// Create new object and start
func StartNew(total int) *ProgressBar {
	return New(total).Start()
}

// Callback for custom output
// For example:
// bar.Callback = func(s string) {
//     mySuperPrint(s)
// }
//
type Callback func(out string)

type ProgressBar struct {
	current int64 // current must be first member of struct (https://code.google.com/p/go/issues/detail?id=5278)

	Total                            int64
	RefreshRate                      time.Duration
	ShowPercent, ShowCounters        bool
	ShowSpeed, ShowTimeLeft, ShowBar bool
	ShowFinalTime                    bool
	Output                           io.Writer
	Callback                         Callback
	NotPrint                         bool
	Units                            Units
	Width                            int
	ForceWidth                       bool
	ManualUpdate                     bool

	TimeLeft   time.Duration
	TotalBytes int64

	// default width for unit numbers and time box
	UnitsWidth   int
	TimeBoxWidth int
	BarWidth     int

	finishOnce sync.Once // Guards isFinish
	finish     chan struct{}
	isFinish   bool

	startTime    time.Time
	startValue   int64
	currentValue int64
	scale        float64

	prefix, postfix string

	mu        *sync.Mutex
	ewma      *EWMA
	lastPrint string

	BarStart string
	BarEnd   string
	Empty    string
	Current  string
	CurrentN string

	AlwaysUpdate bool
}

// Start print
func (pb *ProgressBar) Start() *ProgressBar {
	pb.mu.Lock()
	defer pb.mu.Unlock()

	pb.startTime = time.Now()
	pb.startValue = pb.current
	if pb.Total == 0 {
		pb.ShowTimeLeft = false
		pb.ShowPercent = false
	}
	if !pb.ManualUpdate {
		go pb.writer()
	}
	return pb
}

// Set64 sets the current value as int64
func (pb *ProgressBar) Set64(current int64) *ProgressBar {
	pb.mu.Lock()
	defer pb.mu.Unlock()

	pb.current = current
	return pb
}

func (pb *ProgressBar) SetScale(scale float64) *ProgressBar {
	pb.mu.Lock()
	defer pb.mu.Unlock()

	pb.scale = scale
	return pb
}

// Set prefix string
func (pb *ProgressBar) Prefix(prefix string) *ProgressBar {
	pb.mu.Lock()
	defer pb.mu.Unlock()

	pb.prefix = prefix
	return pb
}

// Set postfix string
func (pb *ProgressBar) Postfix(postfix string) *ProgressBar {
	pb.mu.Lock()
	defer pb.mu.Unlock()

	pb.postfix = postfix
	return pb
}

// Set custom format for bar
// Example: bar.Format("[=>_]")
// Example: bar.Format("[\x00=\x00>\x00-\x00]") // \x00 is the delimiter
func (pb *ProgressBar) Format(format string) *ProgressBar {
	var formatEntries []string
	if len(format) == 5 {
		formatEntries = strings.Split(format, "")
	} else {
		formatEntries = strings.Split(format, "\x00")
	}
	if len(formatEntries) == 5 {
		pb.BarStart = formatEntries[0]
		pb.BarEnd = formatEntries[4]
		pb.Empty = formatEntries[3]
		pb.Current = formatEntries[1]
		pb.CurrentN = formatEntries[2]
	}
	return pb
}

// Set bar refresh rate
func (pb *ProgressBar) SetRefreshRate(rate time.Duration) *ProgressBar {
	pb.RefreshRate = rate
	return pb
}

// Set units
// bar.SetUnits(U_NO) - by default
// bar.SetUnits(U_BYTES) - for Mb, Kb, etc
func (pb *ProgressBar) SetUnits(units Units) *ProgressBar {
	pb.Units = units
	return pb
}

// Set max width, if width is bigger than terminal width, will be ignored
func (pb *ProgressBar) SetMaxWidth(width int) *ProgressBar {
	pb.Width = width
	pb.ForceWidth = false
	return pb
}

// Set bar width
func (pb *ProgressBar) SetWidth(width int) *ProgressBar {
	pb.Width = width
	pb.ForceWidth = true
	return pb
}

// End print
func (pb *ProgressBar) Finish() {
	// Protect multiple calls
	pb.finishOnce.Do(func() {
		close(pb.finish)
		pb.write(atomic.LoadInt64(&pb.current))
		if !pb.NotPrint {
			fmt.Printf("\r%s\r", strings.Repeat(" ", pb.Width))
		}
		pb.isFinish = true
	})
}

func (pb *ProgressBar) write(current int64) {
	width := pb.GetWidth() - 1

	var percentBox, countersBox, timeLeftBox, speedBox, barBox, end, out string

	// percents
	if pb.ShowPercent {
		var percent float64
		if pb.Total > 0 {
			percent = float64(current) / (float64(pb.Total) / float64(100))
		} else {
			percent = float64(current) / float64(100)
		}
		percentBox = fmt.Sprintf(" %6.02f%% ", percent)
	} else {
		percentBox = " "
	}

	// counters
	if pb.ShowCounters {
		if pb.Total > 0 {
			countersBox = fmt.Sprintf("%s / %s ", Format(current, pb.Units, pb.UnitsWidth), Format(pb.Total, pb.Units, pb.UnitsWidth))
		} else {
			countersBox = Format(current, pb.Units, pb.UnitsWidth) + " / ? "
		}
	}

	// time left
	fromStart := time.Now().Sub(pb.startTime)
	currentFromStart := current - pb.startValue
	select {
	case <-pb.finish:
		if pb.ShowFinalTime {
			var left time.Duration
			if pb.Total > 0 {
				left = (fromStart / time.Second) * time.Second
			} else {
				left = (time.Duration(currentFromStart) / time.Second) * time.Second
			}
			pb.ewma.Add(left.Seconds())
			timeLeftBox = (time.Duration(pb.ewma.Value()) * time.Second).String()
		}
	default:
		if pb.ShowTimeLeft && currentFromStart > 0 {
			perEntry := fromStart / time.Duration(currentFromStart)
			var left time.Duration
			if pb.Total > 0 {
				left = time.Duration(pb.Total-currentFromStart) * perEntry
				left = (left / time.Second) * time.Second
			} else {
				left = time.Duration(currentFromStart) * perEntry
				left = (left / time.Second) * time.Second
			}
			pb.TimeLeft = left
			timeLeftBox = FormatDuration(left)
		}
	}

	if len(timeLeftBox) < pb.TimeBoxWidth {
		timeLeftBox = fmt.Sprintf("%s%s", strings.Repeat(" ", pb.TimeBoxWidth-len(timeLeftBox)), timeLeftBox)
	}

	// speed
	if pb.ShowSpeed && currentFromStart > 0 {
		fromStart := time.Now().Sub(pb.startTime)
		speed := float64(currentFromStart) / (float64(fromStart) / float64(time.Second))
		speedBox = Format(int64(speed), pb.Units, pb.UnitsWidth) + "/s "
	}

	barWidth := escapeAwareRuneCountInString(countersBox + pb.BarStart + pb.BarEnd + percentBox + timeLeftBox + speedBox + pb.prefix + pb.postfix)
	// bar
	if pb.ShowBar {
		fullSize := min(pb.BarWidth, width-barWidth)
		size := int(math.Ceil(float64(fullSize) * pb.scale))
		padSize := fullSize - size
		if size > 0 {
			if pb.Total > 0 {
				curCount := int(math.Ceil((float64(current) / float64(pb.Total)) * float64(size)))
				emptCount := size - curCount
				barBox = pb.BarStart
				if emptCount < 0 {
					emptCount = 0
				}
				if curCount > size {
					curCount = size
				}
				if emptCount <= 0 {
					barBox += strings.Repeat(pb.Current, curCount)
				} else if curCount > 0 {
					barBox += strings.Repeat(pb.Current, curCount-1) + pb.CurrentN
				}
				barBox += strings.Repeat(pb.Empty, emptCount)
			} else {
				barBox = pb.BarStart
				pos := size - int(current)%int(size)
				if pos-1 > 0 {
					barBox += strings.Repeat(pb.Empty, pos-1)
				}
				barBox += pb.Current
				if size-pos-1 > 0 {
					barBox += strings.Repeat(pb.Empty, size-pos-1)
				}
			}
			if padSize > 0 {
				barBox += strings.Repeat(" ", padSize-1)
			}
			barBox += pb.BarEnd
		} else if padSize > 0 {
			barBox += pb.BarStart + strings.Repeat(" ", padSize-1) + pb.BarEnd
		}
	}

	// check len
	out = pb.prefix + countersBox + barBox + percentBox + speedBox + timeLeftBox + pb.postfix
	if escapeAwareRuneCountInString(out) < width {
		end = strings.Repeat(" ", width-utf8.RuneCountInString(out))
	}

	// and print!
	pb.lastPrint = out + end
	switch {
	case pb.Output != nil:
		fmt.Fprint(pb.Output, "\r"+out+end)
	case pb.Callback != nil:
		pb.Callback(out + end)
	case !pb.NotPrint:
		fmt.Print("\r" + out + end)
	}
}

func min(a, b int) int {
	if a < b {
		return a
	}
	return b
}

func (pb *ProgressBar) GetWidth() int {
	if pb.ForceWidth {
		return pb.Width
	}

	width := pb.Width
	return width
}

// Write the current state of the progressbar
func (pb *ProgressBar) Update() {
	pb.mu.Lock()
	defer pb.mu.Unlock()

	c := atomic.LoadInt64(&pb.current)
	if pb.AlwaysUpdate || c != pb.currentValue {
		pb.write(c)
		pb.currentValue = c
	}
}

func (pb *ProgressBar) GOString() string {
	return pb.lastPrint
}

// Internal loop for writing progressbar
func (pb *ProgressBar) writer() {
	pb.Update()
	for {
		select {
		case <-pb.finish:
			return
		case <-time.After(pb.RefreshRate):
			pb.Update()
		}
	}
}

func (pb *ProgressBar) CurrentValue() int64 {
	return pb.currentValue
}

func (pb *ProgressBar) GetTimeLeft() time.Duration {
	pb.mu.Lock()
	defer pb.mu.Unlock()
	return pb.TimeLeft
}

type window struct {
	Row    uint16
	Col    uint16
	Xpixel uint16
	Ypixel uint16
}
<commit_msg>fix race condition in pb<commit_after>// Simple console progress bars
package pb

import (
	"fmt"
	"io"
	"math"
	"strings"
	"sync"
	"sync/atomic"
	"time"
	"unicode/utf8"
)

const (
	// Default refresh rate - 200ms
	DEFAULT_REFRESH_RATE = time.Millisecond * 200
	FORMAT               = "[=>-]"
)

// Create new progress bar object
func New(total int) *ProgressBar {
	return New64(int64(total))
}

// Create new progress bar object uding int64 as total
func New64(total int64) *ProgressBar {
	pb := &ProgressBar{
		Total:         total,
		RefreshRate:   DEFAULT_REFRESH_RATE,
		ShowPercent:   true,
		ShowCounters:  true,
		ShowBar:       true,
		ShowTimeLeft:  true,
		ShowFinalTime: true,
		Units:         U_NO,
		ManualUpdate:  false,

		finish:       make(chan struct{}),
		currentValue: -1,
		scale:        1.0,
		ewma:         &EWMA{},
		mu:           new(sync.Mutex),
	}
	return pb.Format(FORMAT)
}

// Create new object and start
func StartNew(total int) *ProgressBar {
	return New(total).Start()
}

// Callback for custom output
// For example:
// bar.Callback = func(s string) {
//     mySuperPrint(s)
// }
//
type Callback func(out string)

type ProgressBar struct {
	current int64 // current must be first member of struct (https://code.google.com/p/go/issues/detail?id=5278)

	Total                            int64
	RefreshRate                      time.Duration
	ShowPercent, ShowCounters        bool
	ShowSpeed, ShowTimeLeft, ShowBar bool
	ShowFinalTime                    bool
	Output                           io.Writer
	Callback                         Callback
	NotPrint                         bool
	Units                            Units
	Width                            int
	ForceWidth                       bool
	ManualUpdate                     bool

	TimeLeft   time.Duration
	TotalBytes int64

	// default width for unit numbers and time box
	UnitsWidth   int
	TimeBoxWidth int
	BarWidth     int

	finishOnce sync.Once // Guards isFinish
	finish     chan struct{}
	isFinish   bool

	startTime    time.Time
	startValue   int64
	currentValue int64
	scale        float64

	prefix, postfix string

	mu        *sync.Mutex
	ewma      *EWMA
	lastPrint string

	BarStart string
	BarEnd   string
	Empty    string
	Current  string
	CurrentN string

	AlwaysUpdate bool
}

// Start print
func (pb *ProgressBar) Start() *ProgressBar {
	pb.mu.Lock()
	defer pb.mu.Unlock()

	pb.startTime = time.Now()
	pb.startValue = pb.current
	if pb.Total == 0 {
		pb.ShowTimeLeft = false
		pb.ShowPercent = false
	}
	if !pb.ManualUpdate {
		go pb.writer()
	}
	return pb
}

// Set64 sets the current value as int64
func (pb *ProgressBar) Set64(current int64) *ProgressBar {
	pb.mu.Lock()
	defer pb.mu.Unlock()

	pb.current = current
	return pb
}

func (pb *ProgressBar) SetScale(scale float64) *ProgressBar {
	pb.mu.Lock()
	defer pb.mu.Unlock()

	pb.scale = scale
	return pb
}

// Set prefix string
func (pb *ProgressBar) Prefix(prefix string) *ProgressBar {
	pb.mu.Lock()
	defer pb.mu.Unlock()

	pb.prefix = prefix
	return pb
}

// Set postfix string
func (pb *ProgressBar) Postfix(postfix string) *ProgressBar {
	pb.mu.Lock()
	defer pb.mu.Unlock()

	pb.postfix = postfix
	return pb
}

// Set custom format for bar
// Example: bar.Format("[=>_]")
// Example: bar.Format("[\x00=\x00>\x00-\x00]") // \x00 is the delimiter
func (pb *ProgressBar) Format(format string) *ProgressBar {
	var formatEntries []string
	if len(format) == 5 {
		formatEntries = strings.Split(format, "")
	} else {
		formatEntries = strings.Split(format, "\x00")
	}
	if len(formatEntries) == 5 {
		pb.BarStart = formatEntries[0]
		pb.BarEnd = formatEntries[4]
		pb.Empty = formatEntries[3]
		pb.Current = formatEntries[1]
		pb.CurrentN = formatEntries[2]
	}
	return pb
}

// Set bar refresh rate
func (pb *ProgressBar) SetRefreshRate(rate time.Duration) *ProgressBar {
	pb.RefreshRate = rate
	return pb
}

// Set units
// bar.SetUnits(U_NO) - by default
// bar.SetUnits(U_BYTES) - for Mb, Kb, etc
func (pb *ProgressBar) SetUnits(units Units) *ProgressBar {
	pb.Units = units
	return pb
}

// Set max width, if width is bigger than terminal width, will be ignored
func (pb *ProgressBar) SetMaxWidth(width int) *ProgressBar {
	pb.Width = width
	pb.ForceWidth = false
	return pb
}

// Set bar width
func (pb *ProgressBar) SetWidth(width int) *ProgressBar {
	pb.Width = width
	pb.ForceWidth = true
	return pb
}

// End print
func (pb *ProgressBar) Finish() {
	// Protect multiple calls
	pb.finishOnce.Do(func() {
		pb.mu.Lock()
		defer pb.mu.Unlock()

		close(pb.finish)
		pb.write(atomic.LoadInt64(&pb.current))
		if !pb.NotPrint {
			fmt.Printf("\r%s\r", strings.Repeat(" ", pb.Width))
		}
		pb.isFinish = true
	})
}

func (pb *ProgressBar) write(current int64) {
	width := pb.GetWidth() - 1

	var percentBox, countersBox, timeLeftBox, speedBox, barBox, end, out string

	// percents
	if pb.ShowPercent {
		var percent float64
		if pb.Total > 0 {
			percent = float64(current) / (float64(pb.Total) / float64(100))
		} else {
			percent = float64(current) / float64(100)
		}
		percentBox = fmt.Sprintf(" %6.02f%% ", percent)
	} else {
		percentBox = " "
	}

	// counters
	if pb.ShowCounters {
		if pb.Total > 0 {
			countersBox = fmt.Sprintf("%s / %s ", Format(current, pb.Units, pb.UnitsWidth), Format(pb.Total, pb.Units, pb.UnitsWidth))
		} else {
			countersBox = Format(current, pb.Units, pb.UnitsWidth) + " / ? "
		}
	}

	// time left
	fromStart := time.Now().Sub(pb.startTime)
	currentFromStart := current - pb.startValue
	select {
	case <-pb.finish:
		if pb.ShowFinalTime {
			var left time.Duration
			if pb.Total > 0 {
				left = (fromStart / time.Second) * time.Second
			} else {
				left = (time.Duration(currentFromStart) / time.Second) * time.Second
			}
			pb.ewma.Add(left.Seconds())
			timeLeftBox = (time.Duration(pb.ewma.Value()) * time.Second).String()
		}
	default:
		if pb.ShowTimeLeft && currentFromStart > 0 {
			perEntry := fromStart / time.Duration(currentFromStart)
			var left time.Duration
			if pb.Total > 0 {
				left = time.Duration(pb.Total-currentFromStart) * perEntry
				left = (left / time.Second) * time.Second
			} else {
				left = time.Duration(currentFromStart) * perEntry
				left = (left / time.Second) * time.Second
			}
			pb.TimeLeft = left
			timeLeftBox = FormatDuration(left)
		}
	}

	if len(timeLeftBox) < pb.TimeBoxWidth {
		timeLeftBox = fmt.Sprintf("%s%s", strings.Repeat(" ", pb.TimeBoxWidth-len(timeLeftBox)), timeLeftBox)
	}

	// speed
	if pb.ShowSpeed && currentFromStart > 0 {
		fromStart := time.Now().Sub(pb.startTime)
		speed := float64(currentFromStart) / (float64(fromStart) / float64(time.Second))
		speedBox = Format(int64(speed), pb.Units, pb.UnitsWidth) + "/s "
	}

	barWidth := escapeAwareRuneCountInString(countersBox + pb.BarStart + pb.BarEnd + percentBox + timeLeftBox + speedBox + pb.prefix + pb.postfix)
	// bar
	if pb.ShowBar {
		fullSize := min(pb.BarWidth, width-barWidth)
		size := int(math.Ceil(float64(fullSize) * pb.scale))
		padSize := fullSize - size
		if size > 0 {
			if pb.Total > 0 {
				curCount := int(math.Ceil((float64(current) / float64(pb.Total)) * float64(size)))
				emptCount := size - curCount
				barBox = pb.BarStart
				if emptCount < 0 {
					emptCount = 0
				}
				if curCount > size {
					curCount = size
				}
				if emptCount <= 0 {
					barBox += strings.Repeat(pb.Current, curCount)
				} else if curCount > 0 {
					barBox += strings.Repeat(pb.Current, curCount-1) + pb.CurrentN
				}
				barBox += strings.Repeat(pb.Empty, emptCount)
			} else {
				barBox = pb.BarStart
				pos := size - int(current)%int(size)
				if pos-1 > 0 {
					barBox += strings.Repeat(pb.Empty, pos-1)
				}
				barBox += pb.Current
				if size-pos-1 > 0 {
					barBox += strings.Repeat(pb.Empty, size-pos-1)
				}
			}
			if padSize > 0 {
				barBox += strings.Repeat(" ", padSize-1)
			}
			barBox += pb.BarEnd
		} else if padSize > 0 {
			barBox += pb.BarStart + strings.Repeat(" ", padSize-1) + pb.BarEnd
		}
	}

	// check len
	out = pb.prefix + countersBox + barBox + percentBox + speedBox + timeLeftBox + pb.postfix
	if escapeAwareRuneCountInString(out) < width {
		end = strings.Repeat(" ", width-utf8.RuneCountInString(out))
	}

	// and print!
	pb.lastPrint = out + end
	switch {
	case pb.Output != nil:
		fmt.Fprint(pb.Output, "\r"+out+end)
	case pb.Callback != nil:
		pb.Callback(out + end)
	case !pb.NotPrint:
		fmt.Print("\r" + out + end)
	}
}

func min(a, b int) int {
	if a < b {
		return a
	}
	return b
}

func (pb *ProgressBar) GetWidth() int {
	if pb.ForceWidth {
		return pb.Width
	}

	width := pb.Width
	return width
}

// Write the current state of the progressbar
func (pb *ProgressBar) Update() {
	pb.mu.Lock()
	defer pb.mu.Unlock()

	c := atomic.LoadInt64(&pb.current)
	if pb.AlwaysUpdate || c != pb.currentValue {
		pb.write(c)
		pb.currentValue = c
	}
}

func (pb *ProgressBar) GOString() string {
	return pb.lastPrint
}

// Internal loop for writing progressbar
func (pb *ProgressBar) writer() {
	pb.Update()
	for {
		select {
		case <-pb.finish:
			return
		case <-time.After(pb.RefreshRate):
			pb.Update()
		}
	}
}

func (pb *ProgressBar) CurrentValue() int64 {
	return pb.currentValue
}

func (pb *ProgressBar) GetTimeLeft() time.Duration {
	pb.mu.Lock()
	defer pb.mu.Unlock()
	return pb.TimeLeft
}

type window struct {
	Row    uint16
	Col    uint16
	Xpixel uint16
	Ypixel uint16
}
<|endoftext|>


<commit_before>package torrent

import (
	"bytes"
	"crypto/sha1" // nolint: gosec
	"errors"
	"fmt"
	"net"

	"github.com/cenkalti/rain/internal/bitfield"
	"github.com/cenkalti/rain/internal/bufferpool"
	"github.com/cenkalti/rain/internal/metainfo"
	"github.com/cenkalti/rain/internal/peer"
	"github.com/cenkalti/rain/internal/peerconn/peerwriter"
	"github.com/cenkalti/rain/internal/peerprotocol"
	"github.com/cenkalti/rain/internal/piecewriter"
	"github.com/cenkalti/rain/internal/tracker"
)

func (t *torrent) handlePieceMessage(pm peer.PieceMessage) {
	msg := pm.Piece
	pe := pm.Peer
	if t.pieces == nil || t.bitfield == nil {
		pe.Logger().Error("piece received but we don't have info")
		t.resumerStats.BytesWasted += int64(len(msg.Buffer.Data))
		t.closePeer(pe)
		return
	}
	if msg.Index >= uint32(len(t.pieces)) {
		pe.Logger().Errorln("invalid piece index:", msg.Index)
		t.resumerStats.BytesWasted += int64(len(msg.Buffer.Data))
		t.closePeer(pe)
		return
	}
	t.downloadSpeed.Update(int64(len(msg.Buffer.Data)))
	t.resumerStats.BytesDownloaded += int64(len(msg.Buffer.Data))
	pd, ok := t.pieceDownloaders[pe]
	if !ok {
		t.resumerStats.BytesWasted += int64(len(msg.Buffer.Data))
		msg.Buffer.Release()
		return
	}
	if pd.Piece.Index != msg.Index {
		t.resumerStats.BytesWasted += int64(len(msg.Buffer.Data))
		msg.Buffer.Release()
		return
	}
	piece := pd.Piece
	block, ok := piece.FindBlock(msg.Begin, uint32(len(msg.Buffer.Data)))
	if !ok {
		pe.Logger().Errorln("invalid piece index:", msg.Index, "begin:", msg.Begin, "length:", len(msg.Buffer.Data))
		t.resumerStats.BytesWasted += int64(len(msg.Buffer.Data))
		t.closePeer(pe)
		msg.Buffer.Release()
		return
	}
	pd.GotBlock(block, msg.Buffer.Data)
	msg.Buffer.Release()
	if !pd.Done() {
		pd.RequestBlocks(t.config.RequestQueueLength)
		return
	}
	// t.log.Debugln("piece download completed. index:", pd.Piece.Index)
	t.closePieceDownloader(pd)
	pe.StopSnubTimer()

	// Request next piece while writing the completed piece, being optimistic about hash check.
	t.startPieceDownloaderFor(pe)

	if piece.Writing {
		panic("piece is already writing")
	}
	piece.Writing = true

	// Prevent receiving piece messages to avoid more than 1 write per torrent.
	t.pieceMessagesC.Suspend()

	pw := piecewriter.New(piece, pe, pd.Buffer)
	go pw.Run(t.pieceWriterResultC, t.doneC)
}

func (t *torrent) handlePeerMessage(pm peer.Message) {
	pe := pm.Peer
	switch msg := pm.Message.(type) {
	case peerprotocol.HaveMessage:
		// Save have messages for processesing later received while we don't have info yet.
		if t.pieces == nil || t.bitfield == nil {
			pe.Messages = append(pe.Messages, msg)
			break
		}
		if msg.Index >= t.info.NumPieces {
			pe.Logger().Errorln("unexpected piece index:", msg.Index)
			t.closePeer(pe)
			break
		}
		pi := &t.pieces[msg.Index]
		// pe.Logger().Debug("Peer ", pe.String(), " has piece #", pi.Index)
		if t.piecePicker != nil {
			t.piecePicker.HandleHave(pe, pi.Index)
		}
		t.updateInterestedState(pe)
		t.startPieceDownloaderFor(pe)
	case peerprotocol.BitfieldMessage:
		// Save bitfield messages while we don't have info yet.
		if t.pieces == nil || t.bitfield == nil {
			pe.Messages = append(pe.Messages, msg)
			break
		}
		bf, err := bitfield.NewBytes(msg.Data, t.info.NumPieces)
		if err != nil {
			pe.Logger().Errorf("%s [len(bitfield)=%d] [numPieces=%d]", err, len(msg.Data), t.info.NumPieces)
			t.closePeer(pe)
			break
		}
		pe.Logger().Debugln("Received bitfield:", bf.Hex())
		if t.piecePicker != nil {
			for i := uint32(0); i < bf.Len(); i++ {
				if bf.Test(i) {
					t.piecePicker.HandleHave(pe, i)
				}
			}
		}
		t.updateInterestedState(pe)
		t.startPieceDownloaderFor(pe)
	case peerprotocol.HaveAllMessage:
		if t.pieces == nil || t.bitfield == nil {
			pe.Messages = append(pe.Messages, msg)
			break
		}
		if t.piecePicker != nil {
			for _, pi := range t.pieces {
				t.piecePicker.HandleHave(pe, pi.Index)
			}
		}
		t.updateInterestedState(pe)
		t.startPieceDownloaderFor(pe)
	case peerprotocol.HaveNoneMessage:
	case peerprotocol.AllowedFastMessage:
		if t.pieces == nil || t.bitfield == nil {
			pe.Messages = append(pe.Messages, msg)
			break
		}
		if msg.Index >= t.info.NumPieces {
			pe.Logger().Errorln("invalid allowed fast piece index:", msg.Index)
			t.closePeer(pe)
			break
		}
		pi := &t.pieces[msg.Index]
		pe.Logger().Debug("Peer ", pe.String(), " has allowed fast for piece #", pi.Index)
		if t.piecePicker != nil {
			t.piecePicker.HandleAllowedFast(pe, msg.Index)
		}
	case peerprotocol.UnchokeMessage:
		pe.PeerChoking = false
		if pd, ok := t.pieceDownloaders[pe]; ok {
			pd.RequestBlocks(t.config.RequestQueueLength)
		}
		t.startPieceDownloaderFor(pe)
	case peerprotocol.ChokeMessage:
		pe.PeerChoking = true
		if pd, ok := t.pieceDownloaders[pe]; ok {
			pd.Choked()
			t.pieceDownloadersChoked[pe] = pd
			if t.piecePicker != nil {
				t.piecePicker.HandleChoke(pe, pd.Piece.Index)
			}
			t.startPieceDownloaderFor(pe)
		}
	case peerprotocol.InterestedMessage:
		pe.PeerInterested = true
	case peerprotocol.NotInterestedMessage:
		pe.PeerInterested = false
	case peerprotocol.RequestMessage:
		if t.pieces == nil || t.bitfield == nil {
			pe.Logger().Error("request received but we don't have info")
			t.closePeer(pe)
			break
		}
		if msg.Index >= t.info.NumPieces {
			pe.Logger().Errorln("invalid request index:", msg.Index)
			t.closePeer(pe)
			break
		}
		if msg.Begin+msg.Length > t.pieces[msg.Index].Length {
			pe.Logger().Errorln("invalid request length:", msg.Length)
			t.closePeer(pe)
			break
		}
		pi := &t.pieces[msg.Index]
		if pe.ClientChoking {
			if pe.FastEnabled {
				m := peerprotocol.RejectMessage{RequestMessage: msg}
				pe.SendMessage(m)
			}
		} else {
			if t.pieceCache != nil {
				pe.SendPiece(msg, t.cachedPiece(pi))
			} else {
				pe.SendPiece(msg, pi.Data)
			}
		}
	case peerprotocol.RejectMessage:
		if t.pieces == nil || t.bitfield == nil {
			pe.Logger().Error("reject received but we don't have info")
			t.closePeer(pe)
			break
		}

		if msg.Index >= t.info.NumPieces {
			pe.Logger().Errorln("invalid reject index:", msg.Index)
			t.closePeer(pe)
			break
		}
		pd, ok := t.pieceDownloaders[pe]
		if !ok {
			break
		}
		if pd.Piece.Index != msg.Index {
			break
		}
		block, ok := pd.Piece.FindBlock(msg.Begin, msg.Length)
		if !ok {
			pe.Logger().Errorln("invalid reject index:", msg.Index, "begin:", msg.Begin, "length:", msg.Length)
			t.closePeer(pe)
			break
		}
		pd.Rejected(block)
	case peerprotocol.CancelMessage:
		if t.pieces == nil || t.bitfield == nil {
			pe.Logger().Error("cancel received but we don't have info")
			t.closePeer(pe)
			break
		}

		if msg.Index >= t.info.NumPieces {
			pe.Logger().Errorln("invalid cancel index:", msg.Index)
			t.closePeer(pe)
			break
		}
		pe.CancelRequest(msg)
	case peerwriter.BlockUploaded:
		t.uploadSpeed.Update(int64(msg.Length))
		t.resumerStats.BytesUploaded += int64(msg.Length)
	case peerprotocol.ExtensionHandshakeMessage:
		pe.Logger().Debugln("extension handshake received:", msg)
		if pe.ExtensionHandshake != nil {
			pe.Logger().Debugln("peer changed extensions")
			break
		}
		pe.ExtensionHandshake = &msg

		if len(msg.YourIP) == 4 {
			t.externalIP = net.IP(msg.YourIP)
		}
		if _, ok := msg.M[peerprotocol.ExtensionKeyMetadata]; ok {
			t.startInfoDownloaders()
		}
		if t.config.PEXEnabled {
			if _, ok := msg.M[peerprotocol.ExtensionKeyPEX]; ok {
				if t.info != nil && t.info.Private != 1 {
					pe.StartPEX(t.peers)
				}
			}
		}
	case peerprotocol.ExtensionMetadataMessage:
		switch msg.Type {
		case peerprotocol.ExtensionMetadataMessageTypeRequest:
			extMsgID, ok := pe.ExtensionHandshake.M[peerprotocol.ExtensionKeyMetadata]
			if !ok {
				break
			}
			if t.info == nil {
				t.sendMetadataReject(pe, msg.Piece, extMsgID)
				break
			}
			start := 16 * 1024 * msg.Piece
			end := start + 16*1024
			totalSize := uint32(len(t.info.Bytes))
			if end > totalSize {
				end = totalSize
			}
			if start >= totalSize {
				t.sendMetadataReject(pe, msg.Piece, extMsgID)
				break
			}
			if end > totalSize {
				t.sendMetadataReject(pe, msg.Piece, extMsgID)
				break
			}
			data := t.info.Bytes[start:end]
			dataMsg := peerprotocol.ExtensionMetadataMessage{
				Type:      peerprotocol.ExtensionMetadataMessageTypeData,
				Piece:     msg.Piece,
				TotalSize: totalSize,
				Data:      data,
			}
			extDataMsg := peerprotocol.ExtensionMessage{
				ExtendedMessageID: extMsgID,
				Payload:           dataMsg,
			}
			pe.SendMessage(extDataMsg)
		case peerprotocol.ExtensionMetadataMessageTypeData:
			id, ok := t.infoDownloaders[pe]
			if !ok {
				break
			}
			err := id.GotBlock(msg.Piece, msg.Data)
			if err != nil {
				pe.Logger().Error(err)
				t.closePeer(pe)
				t.startInfoDownloaders()
				break
			}
			if !id.Done() {
				id.RequestBlocks(t.config.RequestQueueLength)
				break
			}
			pe.StopSnubTimer()

			hash := sha1.New()                              // nolint: gosec
			hash.Write(id.Bytes)                            // nolint: gosec
			if !bytes.Equal(hash.Sum(nil), t.infoHash[:]) { // nolint: gosec
				pe.Logger().Errorln("received info does not match with hash")
				t.closePeer(id.Peer)
				t.startInfoDownloaders()
				break
			}
			t.stopInfoDownloaders()

			info, err := metainfo.NewInfo(id.Bytes)
			if err != nil {
				err = fmt.Errorf("cannot parse info bytes: %s", err)
				t.log.Error(err)
				t.stop(err)
				break
			}
			if info.Private == 1 {
				err = errors.New("private torrent from magnet")
				t.log.Error(err)
				t.stop(err)
				break
			}
			t.info = info
			t.piecePool = bufferpool.New(int(info.PieceLength))
			if t.resume != nil {
				err = t.resume.WriteInfo(t.info.Bytes)
				if err != nil {
					err = fmt.Errorf("cannot write resume info: %s", err)
					t.log.Error(err)
					t.stop(err)
					break
				}
			}
			t.startAllocator()
		case peerprotocol.ExtensionMetadataMessageTypeReject:
			id, ok := t.infoDownloaders[pe]
			if ok {
				t.closePeer(id.Peer)
				t.startInfoDownloaders()
			}
		}
	case peerprotocol.ExtensionPEXMessage:
		if !t.config.PEXEnabled {
			break
		}
		addrs, err := tracker.DecodePeersCompact([]byte(msg.Added))
		if err != nil {
			t.log.Error(err)
			break
		}
		t.handleNewPeers(addrs, peer.SourcePEX)
	default:
		panic(fmt.Sprintf("unhandled peer message type: %T", msg))
	}
}

func (t *torrent) updateInterestedState(pe *peer.Peer) {
	if t.pieces == nil || t.bitfield == nil {
		return
	}
	interested := false
	if !t.completed {
		for i := uint32(0); i < t.bitfield.Len(); i++ {
			weHave := t.bitfield.Test(i)
			peerHave := pe.Bitfield.Test(i)
			if !weHave && peerHave {
				interested = true
				break
			}
		}
	}
	if !pe.ClientInterested && interested {
		pe.ClientInterested = true
		msg := peerprotocol.InterestedMessage{}
		pe.SendMessage(msg)
		return
	}
	if pe.ClientInterested && !interested {
		pe.ClientInterested = false
		msg := peerprotocol.NotInterestedMessage{}
		pe.SendMessage(msg)
		return
	}
}

func (t *torrent) sendMetadataReject(pe *peer.Peer, i uint32, msgID uint8) {
	dataMsg := peerprotocol.ExtensionMetadataMessage{
		Type:  peerprotocol.ExtensionMetadataMessageTypeReject,
		Piece: i,
	}
	extDataMsg := peerprotocol.ExtensionMessage{
		ExtendedMessageID: msgID,
		Payload:           &dataMsg,
	}
	pe.SendMessage(extDataMsg)
}
<commit_msg>release buffer<commit_after>package torrent

import (
	"bytes"
	"crypto/sha1" // nolint: gosec
	"errors"
	"fmt"
	"net"

	"github.com/cenkalti/rain/internal/bitfield"
	"github.com/cenkalti/rain/internal/bufferpool"
	"github.com/cenkalti/rain/internal/metainfo"
	"github.com/cenkalti/rain/internal/peer"
	"github.com/cenkalti/rain/internal/peerconn/peerwriter"
	"github.com/cenkalti/rain/internal/peerprotocol"
	"github.com/cenkalti/rain/internal/piecewriter"
	"github.com/cenkalti/rain/internal/tracker"
)

func (t *torrent) handlePieceMessage(pm peer.PieceMessage) {
	msg := pm.Piece
	pe := pm.Peer
	if t.pieces == nil || t.bitfield == nil {
		pe.Logger().Error("piece received but we don't have info")
		t.resumerStats.BytesWasted += int64(len(msg.Buffer.Data))
		t.closePeer(pe)
		msg.Buffer.Release()
		return
	}
	if msg.Index >= uint32(len(t.pieces)) {
		pe.Logger().Errorln("invalid piece index:", msg.Index)
		t.resumerStats.BytesWasted += int64(len(msg.Buffer.Data))
		t.closePeer(pe)
		msg.Buffer.Release()
		return
	}
	t.downloadSpeed.Update(int64(len(msg.Buffer.Data)))
	t.resumerStats.BytesDownloaded += int64(len(msg.Buffer.Data))
	pd, ok := t.pieceDownloaders[pe]
	if !ok {
		t.resumerStats.BytesWasted += int64(len(msg.Buffer.Data))
		msg.Buffer.Release()
		return
	}
	if pd.Piece.Index != msg.Index {
		t.resumerStats.BytesWasted += int64(len(msg.Buffer.Data))
		msg.Buffer.Release()
		return
	}
	piece := pd.Piece
	block, ok := piece.FindBlock(msg.Begin, uint32(len(msg.Buffer.Data)))
	if !ok {
		pe.Logger().Errorln("invalid piece index:", msg.Index, "begin:", msg.Begin, "length:", len(msg.Buffer.Data))
		t.resumerStats.BytesWasted += int64(len(msg.Buffer.Data))
		t.closePeer(pe)
		msg.Buffer.Release()
		return
	}
	pd.GotBlock(block, msg.Buffer.Data)
	msg.Buffer.Release()
	if !pd.Done() {
		pd.RequestBlocks(t.config.RequestQueueLength)
		return
	}
	// t.log.Debugln("piece download completed. index:", pd.Piece.Index)
	t.closePieceDownloader(pd)
	pe.StopSnubTimer()

	// Request next piece while writing the completed piece, being optimistic about hash check.
	t.startPieceDownloaderFor(pe)

	if piece.Writing {
		panic("piece is already writing")
	}
	piece.Writing = true

	// Prevent receiving piece messages to avoid more than 1 write per torrent.
	t.pieceMessagesC.Suspend()

	pw := piecewriter.New(piece, pe, pd.Buffer)
	go pw.Run(t.pieceWriterResultC, t.doneC)
}

func (t *torrent) handlePeerMessage(pm peer.Message) {
	pe := pm.Peer
	switch msg := pm.Message.(type) {
	case peerprotocol.HaveMessage:
		// Save have messages for processesing later received while we don't have info yet.
		if t.pieces == nil || t.bitfield == nil {
			pe.Messages = append(pe.Messages, msg)
			break
		}
		if msg.Index >= t.info.NumPieces {
			pe.Logger().Errorln("unexpected piece index:", msg.Index)
			t.closePeer(pe)
			break
		}
		pi := &t.pieces[msg.Index]
		// pe.Logger().Debug("Peer ", pe.String(), " has piece #", pi.Index)
		if t.piecePicker != nil {
			t.piecePicker.HandleHave(pe, pi.Index)
		}
		t.updateInterestedState(pe)
		t.startPieceDownloaderFor(pe)
	case peerprotocol.BitfieldMessage:
		// Save bitfield messages while we don't have info yet.
		if t.pieces == nil || t.bitfield == nil {
			pe.Messages = append(pe.Messages, msg)
			break
		}
		bf, err := bitfield.NewBytes(msg.Data, t.info.NumPieces)
		if err != nil {
			pe.Logger().Errorf("%s [len(bitfield)=%d] [numPieces=%d]", err, len(msg.Data), t.info.NumPieces)
			t.closePeer(pe)
			break
		}
		pe.Logger().Debugln("Received bitfield:", bf.Hex())
		if t.piecePicker != nil {
			for i := uint32(0); i < bf.Len(); i++ {
				if bf.Test(i) {
					t.piecePicker.HandleHave(pe, i)
				}
			}
		}
		t.updateInterestedState(pe)
		t.startPieceDownloaderFor(pe)
	case peerprotocol.HaveAllMessage:
		if t.pieces == nil || t.bitfield == nil {
			pe.Messages = append(pe.Messages, msg)
			break
		}
		if t.piecePicker != nil {
			for _, pi := range t.pieces {
				t.piecePicker.HandleHave(pe, pi.Index)
			}
		}
		t.updateInterestedState(pe)
		t.startPieceDownloaderFor(pe)
	case peerprotocol.HaveNoneMessage:
	case peerprotocol.AllowedFastMessage:
		if t.pieces == nil || t.bitfield == nil {
			pe.Messages = append(pe.Messages, msg)
			break
		}
		if msg.Index >= t.info.NumPieces {
			pe.Logger().Errorln("invalid allowed fast piece index:", msg.Index)
			t.closePeer(pe)
			break
		}
		pi := &t.pieces[msg.Index]
		pe.Logger().Debug("Peer ", pe.String(), " has allowed fast for piece #", pi.Index)
		if t.piecePicker != nil {
			t.piecePicker.HandleAllowedFast(pe, msg.Index)
		}
	case peerprotocol.UnchokeMessage:
		pe.PeerChoking = false
		if pd, ok := t.pieceDownloaders[pe]; ok {
			pd.RequestBlocks(t.config.RequestQueueLength)
		}
		t.startPieceDownloaderFor(pe)
	case peerprotocol.ChokeMessage:
		pe.PeerChoking = true
		if pd, ok := t.pieceDownloaders[pe]; ok {
			pd.Choked()
			t.pieceDownloadersChoked[pe] = pd
			if t.piecePicker != nil {
				t.piecePicker.HandleChoke(pe, pd.Piece.Index)
			}
			t.startPieceDownloaderFor(pe)
		}
	case peerprotocol.InterestedMessage:
		pe.PeerInterested = true
	case peerprotocol.NotInterestedMessage:
		pe.PeerInterested = false
	case peerprotocol.RequestMessage:
		if t.pieces == nil || t.bitfield == nil {
			pe.Logger().Error("request received but we don't have info")
			t.closePeer(pe)
			break
		}
		if msg.Index >= t.info.NumPieces {
			pe.Logger().Errorln("invalid request index:", msg.Index)
			t.closePeer(pe)
			break
		}
		if msg.Begin+msg.Length > t.pieces[msg.Index].Length {
			pe.Logger().Errorln("invalid request length:", msg.Length)
			t.closePeer(pe)
			break
		}
		pi := &t.pieces[msg.Index]
		if pe.ClientChoking {
			if pe.FastEnabled {
				m := peerprotocol.RejectMessage{RequestMessage: msg}
				pe.SendMessage(m)
			}
		} else {
			if t.pieceCache != nil {
				pe.SendPiece(msg, t.cachedPiece(pi))
			} else {
				pe.SendPiece(msg, pi.Data)
			}
		}
	case peerprotocol.RejectMessage:
		if t.pieces == nil || t.bitfield == nil {
			pe.Logger().Error("reject received but we don't have info")
			t.closePeer(pe)
			break
		}

		if msg.Index >= t.info.NumPieces {
			pe.Logger().Errorln("invalid reject index:", msg.Index)
			t.closePeer(pe)
			break
		}
		pd, ok := t.pieceDownloaders[pe]
		if !ok {
			break
		}
		if pd.Piece.Index != msg.Index {
			break
		}
		block, ok := pd.Piece.FindBlock(msg.Begin, msg.Length)
		if !ok {
			pe.Logger().Errorln("invalid reject index:", msg.Index, "begin:", msg.Begin, "length:", msg.Length)
			t.closePeer(pe)
			break
		}
		pd.Rejected(block)
	case peerprotocol.CancelMessage:
		if t.pieces == nil || t.bitfield == nil {
			pe.Logger().Error("cancel received but we don't have info")
			t.closePeer(pe)
			break
		}

		if msg.Index >= t.info.NumPieces {
			pe.Logger().Errorln("invalid cancel index:", msg.Index)
			t.closePeer(pe)
			break
		}
		pe.CancelRequest(msg)
	case peerwriter.BlockUploaded:
		t.uploadSpeed.Update(int64(msg.Length))
		t.resumerStats.BytesUploaded += int64(msg.Length)
	case peerprotocol.ExtensionHandshakeMessage:
		pe.Logger().Debugln("extension handshake received:", msg)
		if pe.ExtensionHandshake != nil {
			pe.Logger().Debugln("peer changed extensions")
			break
		}
		pe.ExtensionHandshake = &msg

		if len(msg.YourIP) == 4 {
			t.externalIP = net.IP(msg.YourIP)
		}
		if _, ok := msg.M[peerprotocol.ExtensionKeyMetadata]; ok {
			t.startInfoDownloaders()
		}
		if t.config.PEXEnabled {
			if _, ok := msg.M[peerprotocol.ExtensionKeyPEX]; ok {
				if t.info != nil && t.info.Private != 1 {
					pe.StartPEX(t.peers)
				}
			}
		}
	case peerprotocol.ExtensionMetadataMessage:
		switch msg.Type {
		case peerprotocol.ExtensionMetadataMessageTypeRequest:
			extMsgID, ok := pe.ExtensionHandshake.M[peerprotocol.ExtensionKeyMetadata]
			if !ok {
				break
			}
			if t.info == nil {
				t.sendMetadataReject(pe, msg.Piece, extMsgID)
				break
			}
			start := 16 * 1024 * msg.Piece
			end := start + 16*1024
			totalSize := uint32(len(t.info.Bytes))
			if end > totalSize {
				end = totalSize
			}
			if start >= totalSize {
				t.sendMetadataReject(pe, msg.Piece, extMsgID)
				break
			}
			if end > totalSize {
				t.sendMetadataReject(pe, msg.Piece, extMsgID)
				break
			}
			data := t.info.Bytes[start:end]
			dataMsg := peerprotocol.ExtensionMetadataMessage{
				Type:      peerprotocol.ExtensionMetadataMessageTypeData,
				Piece:     msg.Piece,
				TotalSize: totalSize,
				Data:      data,
			}
			extDataMsg := peerprotocol.ExtensionMessage{
				ExtendedMessageID: extMsgID,
				Payload:           dataMsg,
			}
			pe.SendMessage(extDataMsg)
		case peerprotocol.ExtensionMetadataMessageTypeData:
			id, ok := t.infoDownloaders[pe]
			if !ok {
				break
			}
			err := id.GotBlock(msg.Piece, msg.Data)
			if err != nil {
				pe.Logger().Error(err)
				t.closePeer(pe)
				t.startInfoDownloaders()
				break
			}
			if !id.Done() {
				id.RequestBlocks(t.config.RequestQueueLength)
				break
			}
			pe.StopSnubTimer()

			hash := sha1.New()                              // nolint: gosec
			hash.Write(id.Bytes)                            // nolint: gosec
			if !bytes.Equal(hash.Sum(nil), t.infoHash[:]) { // nolint: gosec
				pe.Logger().Errorln("received info does not match with hash")
				t.closePeer(id.Peer)
				t.startInfoDownloaders()
				break
			}
			t.stopInfoDownloaders()

			info, err := metainfo.NewInfo(id.Bytes)
			if err != nil {
				err = fmt.Errorf("cannot parse info bytes: %s", err)
				t.log.Error(err)
				t.stop(err)
				break
			}
			if info.Private == 1 {
				err = errors.New("private torrent from magnet")
				t.log.Error(err)
				t.stop(err)
				break
			}
			t.info = info
			t.piecePool = bufferpool.New(int(info.PieceLength))
			if t.resume != nil {
				err = t.resume.WriteInfo(t.info.Bytes)
				if err != nil {
					err = fmt.Errorf("cannot write resume info: %s", err)
					t.log.Error(err)
					t.stop(err)
					break
				}
			}
			t.startAllocator()
		case peerprotocol.ExtensionMetadataMessageTypeReject:
			id, ok := t.infoDownloaders[pe]
			if ok {
				t.closePeer(id.Peer)
				t.startInfoDownloaders()
			}
		}
	case peerprotocol.ExtensionPEXMessage:
		if !t.config.PEXEnabled {
			break
		}
		addrs, err := tracker.DecodePeersCompact([]byte(msg.Added))
		if err != nil {
			t.log.Error(err)
			break
		}
		t.handleNewPeers(addrs, peer.SourcePEX)
	default:
		panic(fmt.Sprintf("unhandled peer message type: %T", msg))
	}
}

func (t *torrent) updateInterestedState(pe *peer.Peer) {
	if t.pieces == nil || t.bitfield == nil {
		return
	}
	interested := false
	if !t.completed {
		for i := uint32(0); i < t.bitfield.Len(); i++ {
			weHave := t.bitfield.Test(i)
			peerHave := pe.Bitfield.Test(i)
			if !weHave && peerHave {
				interested = true
				break
			}
		}
	}
	if !pe.ClientInterested && interested {
		pe.ClientInterested = true
		msg := peerprotocol.InterestedMessage{}
		pe.SendMessage(msg)
		return
	}
	if pe.ClientInterested && !interested {
		pe.ClientInterested = false
		msg := peerprotocol.NotInterestedMessage{}
		pe.SendMessage(msg)
		return
	}
}

func (t *torrent) sendMetadataReject(pe *peer.Peer, i uint32, msgID uint8) {
	dataMsg := peerprotocol.ExtensionMetadataMessage{
		Type:  peerprotocol.ExtensionMetadataMessageTypeReject,
		Piece: i,
	}
	extDataMsg := peerprotocol.ExtensionMessage{
		ExtendedMessageID: msgID,
		Payload:           &dataMsg,
	}
	pe.SendMessage(extDataMsg)
}
<|endoftext|>


<commit_before>package cc

import "fmt"

// Piece is the interface used to be able to use different implementation details
// for the different kinds of pieces.
type Piece interface {
	String() string
	Threatening(b *Board, x, y uint8) []position
}

type blank cell
type dead cell
type king cell
type queen cell
type bishop cell
type rook cell
type knight cell

// The pieces available
// Since they don't carry state use one instace
const (
	Blank  blank  = 0 //"."
	Dead   dead   = 1 //"X"
	King   king   = 5 //"K" //"♚"
	Queen  queen  = 6 //"Q" //"♛"
	Bishop bishop = 7 //"B" //"♝"
	Rook   rook   = 8 //"R" //"♜"
	Knight knight = 9 //"N" //"♞"
)

func (p rook) Threatening(b *Board, x, y uint8) []position {
	if x < 0 || y < 0 || x >= b.columns || y >= b.rows {
		return []position{}
	}
	var ret []position
	for i := uint8(0); i < b.columns; i++ {
		for j := uint8(0); j < b.rows; j++ {
			if i == x && j != y {
				ret = append(ret, position{x: i, y: j})
			}
			if i != x && j == y {
				ret = append(ret, position{x: i, y: j})
			}
		}
	}
	return ret
}

func (p rook) String() string {
	return "R"
}

func (p king) Threatening(b *Board, x, y uint8) []position {
	if x < 0 || y < 0 || x >= b.columns || y >= b.rows {
		return []position{}
	}
	t := []position{
		{x: x - 1, y: y},
		{x: x - 1, y: y - 1},
		{x: x, y: y - 1},
		{x: x + 1, y: y - 1},
		{x: x + 1, y: y},
		{x: x + 1, y: y + 1},
		{x: x, y: y + 1},
		{x: x - 1, y: y + 1},
	}
	var ret []position
	for _, c := range t {
		if c.x < 0 || c.y < 0 || c.x >= b.columns || c.y >= b.rows {
			continue
		}
		ret = append(ret, c)
	}
	return ret
}

func (p king) String() string {
	return "K"
}

func (p knight) Threatening(b *Board, x, y uint8) []position {
	if x < 0 || y < 0 || x >= b.columns || y >= b.rows {
		return []position{} // TODO: Should be error
	}
	t := []position{
		{x: x - 1, y: y - 2},
		{x: x + 1, y: y - 2},
		{x: x - 2, y: y - 1},
		{x: x + 2, y: y - 1},
		{x: x - 2, y: y + 1},
		{x: x + 2, y: y + 1},
		{x: x - 1, y: y + 2},
		{x: x + 1, y: y + 2},
	}
	var ret []position
	for _, c := range t {
		if c.x < 0 || c.y < 0 || c.x >= b.columns || c.y >= b.rows {
			continue
		}
		ret = append(ret, c)
	}
	return ret
}

func (p knight) String() string {
	return "N"
}

func (p queen) Threatening(b *Board, x, y uint8) []position {
	if x < 0 || y < 0 || x >= b.columns || y >= b.rows {
		return []position{}
	}

	ret := make([]position, 0, b.columns*b.rows)
	ret = append(ret, Rook.Threatening(b, x, y)...)
	ret = append(ret, Bishop.Threatening(b, x, y)...)
	return ret
}

func (p queen) String() string {
	return "Q"
}

func (p bishop) Threatening(b *Board, x, y uint8) []position {
	if x < 0 || y < 0 || x >= b.columns || y >= b.rows {
		return []position{}
	}
	m := make(map[string]position)
	i := x
	j := y
	for {
		if i < 0 || i >= b.columns {
			break
		}
		if j < 0 || j >= b.rows {
			break
		}
		i++
		j++
		m[fmt.Sprintf("%d,%d", i, j)] = position{x: i, y: j}
	}
	i = x
	j = y
	for {
		if i < 0 || i >= b.columns {
			break
		}
		if j < 0 || j >= b.rows {
			break
		}
		i++
		j--
		m[fmt.Sprintf("%d,%d", i, j)] = position{x: i, y: j}
	}
	i = x
	j = y
	for {
		if i < 0 || i >= b.columns {
			break
		}
		if j < 0 || j >= b.rows {
			break
		}
		i--
		j--
		m[fmt.Sprintf("%d,%d", i, j)] = position{x: i, y: j}
	}
	i = x
	j = y
	for {
		if i < 0 || i >= b.columns {
			break
		}
		if j < 0 || j >= b.rows {
			break
		}
		i--
		j++
		m[fmt.Sprintf("%d,%d", i, j)] = position{x: i, y: j}
	}

	ret := make([]position, 0, b.columns*b.rows)
	for _, c := range m {
		if c.x < 0 || c.y < 0 || c.x >= b.columns || c.y >= b.rows {
			continue
		}
		ret = append(ret, c)
	}
	return ret
}

func (p bishop) String() string {
	return "B"
}

func (p blank) String() string {
	return "."
}

func (p blank) Threatening(b *Board, x, y uint8) []position {
	return []position{}
}

func (p dead) String() string {
	return "X"
}

func (p dead) Threatening(b *Board, x, y uint8) []position {
	return []position{}
}
<commit_msg>Improve Bishop threatening<commit_after>package cc

// Piece is the interface used to be able to use different implementation details
// for the different kinds of pieces.
type Piece interface {
	String() string
	Threatening(b *Board, x, y uint8) []position
}

type blank cell
type dead cell
type king cell
type queen cell
type bishop cell
type rook cell
type knight cell

// The pieces available
// Since they don't carry state use one instace
const (
	Blank  blank  = 0 //"."
	Dead   dead   = 1 //"X"
	King   king   = 5 //"K" //"♚"
	Queen  queen  = 6 //"Q" //"♛"
	Bishop bishop = 7 //"B" //"♝"
	Rook   rook   = 8 //"R" //"♜"
	Knight knight = 9 //"N" //"♞"
)

func (p rook) Threatening(b *Board, x, y uint8) []position {
	if x < 0 || y < 0 || x >= b.columns || y >= b.rows {
		return []position{}
	}
	var ret []position
	for i := uint8(0); i < b.columns; i++ {
		for j := uint8(0); j < b.rows; j++ {
			if i == x && j != y {
				ret = append(ret, position{x: i, y: j})
			}
			if i != x && j == y {
				ret = append(ret, position{x: i, y: j})
			}
		}
	}
	return ret
}

func (p rook) String() string {
	return "R"
}

func (p king) Threatening(b *Board, x, y uint8) []position {
	if x < 0 || y < 0 || x >= b.columns || y >= b.rows {
		return []position{}
	}
	t := []position{
		{x: x - 1, y: y},
		{x: x - 1, y: y - 1},
		{x: x, y: y - 1},
		{x: x + 1, y: y - 1},
		{x: x + 1, y: y},
		{x: x + 1, y: y + 1},
		{x: x, y: y + 1},
		{x: x - 1, y: y + 1},
	}
	var ret []position
	for _, c := range t {
		if c.x < 0 || c.y < 0 || c.x >= b.columns || c.y >= b.rows {
			continue
		}
		ret = append(ret, c)
	}
	return ret
}

func (p king) String() string {
	return "K"
}

func (p knight) Threatening(b *Board, x, y uint8) []position {
	if x < 0 || y < 0 || x >= b.columns || y >= b.rows {
		return []position{} // TODO: Should be error
	}
	t := []position{
		{x: x - 1, y: y - 2},
		{x: x + 1, y: y - 2},
		{x: x - 2, y: y - 1},
		{x: x + 2, y: y - 1},
		{x: x - 2, y: y + 1},
		{x: x + 2, y: y + 1},
		{x: x - 1, y: y + 2},
		{x: x + 1, y: y + 2},
	}
	var ret []position
	for _, c := range t {
		if c.x < 0 || c.y < 0 || c.x >= b.columns || c.y >= b.rows {
			continue
		}
		ret = append(ret, c)
	}
	return ret
}

func (p knight) String() string {
	return "N"
}

func (p queen) Threatening(b *Board, x, y uint8) []position {
	if x < 0 || y < 0 || x >= b.columns || y >= b.rows {
		return []position{}
	}

	ret := make([]position, 0, b.columns*b.rows)
	ret = append(ret, Rook.Threatening(b, x, y)...)
	ret = append(ret, Bishop.Threatening(b, x, y)...)
	return ret
}

func (p queen) String() string {
	return "Q"
}

func (p bishop) Threatening(b *Board, x, y uint8) []position {
	if x < 0 || y < 0 || x >= b.columns || y >= b.rows {
		return []position{}
	}
	n := b.columns
	if b.rows > n {
		n = b.rows
	}
	var ret []position
	for step := 1; step < int(n); step++ {
		x1 := uint8(int(x) + (step * -1))
		y1 := uint8(int(y) + (step * -1))
		x2 := uint8(int(x) + (step * 1))
		y2 := uint8(int(y) + (step * 1))
		if x1 < b.columns && y1 < b.rows {
			ret = append(ret, position{x: x1, y: y1})
		}
		if x2 < b.columns && y2 < b.rows {
			ret = append(ret, position{x: x2, y: y2})
		}
		if x1 < b.columns && y2 < b.rows {
			ret = append(ret, position{x: x1, y: y2})
		}
		if x2 < b.columns && y1 < b.rows {
			ret = append(ret, position{x: x2, y: y1})
		}
	}
	return ret
}

func (p bishop) String() string {
	return "B"
}

func (p blank) String() string {
	return "."
}

func (p blank) Threatening(b *Board, x, y uint8) []position {
	return []position{}
}

func (p dead) String() string {
	return "X"
}

func (p dead) Threatening(b *Board, x, y uint8) []position {
	return []position{}
}
<|endoftext|>


<commit_before>/*
Copyright 2016 The Kubernetes 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.
*/

package apiserver

import (
	"net/http"
	"time"

	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime/schema"
	"k8s.io/apimachinery/pkg/util/sets"
	genericapiserver "k8s.io/apiserver/pkg/server"
	serverstorage "k8s.io/apiserver/pkg/server/storage"
	"k8s.io/client-go/pkg/version"

	"k8s.io/kube-aggregator/pkg/apis/apiregistration"
	"k8s.io/kube-aggregator/pkg/apis/apiregistration/v1"
	"k8s.io/kube-aggregator/pkg/apis/apiregistration/v1beta1"
	aggregatorscheme "k8s.io/kube-aggregator/pkg/apiserver/scheme"
	"k8s.io/kube-aggregator/pkg/client/clientset_generated/internalclientset"
	informers "k8s.io/kube-aggregator/pkg/client/informers/internalversion"
	listers "k8s.io/kube-aggregator/pkg/client/listers/apiregistration/internalversion"
	openapicontroller "k8s.io/kube-aggregator/pkg/controllers/openapi"
	statuscontrollers "k8s.io/kube-aggregator/pkg/controllers/status"
	apiservicerest "k8s.io/kube-aggregator/pkg/registry/apiservice/rest"
)

func init() {
	// we need to add the options (like ListOptions) to empty v1
	metav1.AddToGroupVersion(aggregatorscheme.Scheme, schema.GroupVersion{Group: "", Version: "v1"})

	unversioned := schema.GroupVersion{Group: "", Version: "v1"}
	aggregatorscheme.Scheme.AddUnversionedTypes(unversioned,
		&metav1.Status{},
		&metav1.APIVersions{},
		&metav1.APIGroupList{},
		&metav1.APIGroup{},
		&metav1.APIResourceList{},
	)
}

// legacyAPIServiceName is the fixed name of the only non-groupified API version
const legacyAPIServiceName = "v1."

type ExtraConfig struct {
	// ProxyClientCert/Key are the client cert used to identify this proxy. Backing APIServices use
	// this to confirm the proxy's identity
	ProxyClientCert []byte
	ProxyClientKey  []byte

	// If present, the Dial method will be used for dialing out to delegate
	// apiservers.
	ProxyTransport *http.Transport

	// Mechanism by which the Aggregator will resolve services. Required.
	ServiceResolver ServiceResolver
}

type Config struct {
	GenericConfig *genericapiserver.RecommendedConfig
	ExtraConfig   ExtraConfig
}

type completedConfig struct {
	GenericConfig genericapiserver.CompletedConfig
	ExtraConfig   *ExtraConfig
}

type CompletedConfig struct {
	// Embed a private pointer that cannot be instantiated outside of this package.
	*completedConfig
}

// APIAggregator contains state for a Kubernetes cluster master/api server.
type APIAggregator struct {
	GenericAPIServer *genericapiserver.GenericAPIServer

	delegateHandler http.Handler

	// proxyClientCert/Key are the client cert used to identify this proxy. Backing APIServices use
	// this to confirm the proxy's identity
	proxyClientCert []byte
	proxyClientKey  []byte
	proxyTransport  *http.Transport

	// proxyHandlers are the proxy handlers that are currently registered, keyed by apiservice.name
	proxyHandlers map[string]*proxyHandler
	// handledGroups are the groups that already have routes
	handledGroups sets.String

	// lister is used to add group handling for /apis/<group> aggregator lookups based on
	// controller state
	lister listers.APIServiceLister

	// provided for easier embedding
	APIRegistrationInformers informers.SharedInformerFactory

	// Information needed to determine routing for the aggregator
	serviceResolver ServiceResolver

	openAPIAggregationController *openapicontroller.AggregationController
}

// Complete fills in any fields not set that are required to have valid data. It's mutating the receiver.
func (cfg *Config) Complete() CompletedConfig {
	c := completedConfig{
		cfg.GenericConfig.Complete(),
		&cfg.ExtraConfig,
	}

	// the kube aggregator wires its own discovery mechanism
	// TODO eventually collapse this by extracting all of the discovery out
	c.GenericConfig.EnableDiscovery = false
	version := version.Get()
	c.GenericConfig.Version = &version

	return CompletedConfig{&c}
}

// New returns a new instance of APIAggregator from the given config.
func (c completedConfig) NewWithDelegate(delegationTarget genericapiserver.DelegationTarget) (*APIAggregator, error) {
	// Prevent generic API server to install OpenAPI handler. Aggregator server
	// has its own customized OpenAPI handler.
	openApiConfig := c.GenericConfig.OpenAPIConfig
	c.GenericConfig.OpenAPIConfig = nil

	genericServer, err := c.GenericConfig.New("kube-aggregator", delegationTarget)
	if err != nil {
		return nil, err
	}

	apiregistrationClient, err := internalclientset.NewForConfig(c.GenericConfig.LoopbackClientConfig)
	if err != nil {
		return nil, err
	}
	informerFactory := informers.NewSharedInformerFactory(
		apiregistrationClient,
		5*time.Minute, // this is effectively used as a refresh interval right now.  Might want to do something nicer later on.
	)

	s := &APIAggregator{
		GenericAPIServer: genericServer,
		delegateHandler:  delegationTarget.UnprotectedHandler(),
		proxyClientCert:  c.ExtraConfig.ProxyClientCert,
		proxyClientKey:   c.ExtraConfig.ProxyClientKey,
		proxyTransport:   c.ExtraConfig.ProxyTransport,
		proxyHandlers:    map[string]*proxyHandler{},
		handledGroups:    sets.String{},
		lister:           informerFactory.Apiregistration().InternalVersion().APIServices().Lister(),
		APIRegistrationInformers: informerFactory,
		serviceResolver:          c.ExtraConfig.ServiceResolver,
	}

	apiGroupInfo := apiservicerest.NewRESTStorage(c.GenericConfig.MergedResourceConfig, c.GenericConfig.RESTOptionsGetter)
	if err := s.GenericAPIServer.InstallAPIGroup(&apiGroupInfo); err != nil {
		return nil, err
	}

	apisHandler := &apisHandler{
		codecs: aggregatorscheme.Codecs,
		lister: s.lister,
	}
	s.GenericAPIServer.Handler.NonGoRestfulMux.Handle("/apis", apisHandler)
	s.GenericAPIServer.Handler.NonGoRestfulMux.UnlistedHandle("/apis/", apisHandler)

	apiserviceRegistrationController := NewAPIServiceRegistrationController(informerFactory.Apiregistration().InternalVersion().APIServices(), s)
	availableController := statuscontrollers.NewAvailableConditionController(
		informerFactory.Apiregistration().InternalVersion().APIServices(),
		c.GenericConfig.SharedInformerFactory.Core().V1().Services(),
		c.GenericConfig.SharedInformerFactory.Core().V1().Endpoints(),
		apiregistrationClient.Apiregistration(),
		c.ExtraConfig.ProxyTransport,
		c.ExtraConfig.ProxyClientCert,
		c.ExtraConfig.ProxyClientKey,
		s.serviceResolver,
	)

	s.GenericAPIServer.AddPostStartHook("start-kube-aggregator-informers", func(context genericapiserver.PostStartHookContext) error {
		informerFactory.Start(context.StopCh)
		c.GenericConfig.SharedInformerFactory.Start(context.StopCh)
		return nil
	})
	s.GenericAPIServer.AddPostStartHook("apiservice-registration-controller", func(context genericapiserver.PostStartHookContext) error {
		go apiserviceRegistrationController.Run(context.StopCh)
		return nil
	})
	s.GenericAPIServer.AddPostStartHook("apiservice-status-available-controller", func(context genericapiserver.PostStartHookContext) error {
		// if we end up blocking for long periods of time, we may need to increase threadiness.
		go availableController.Run(5, context.StopCh)
		return nil
	})

	if openApiConfig != nil {
		specDownloader := openapicontroller.NewDownloader()
		openAPIAggregator, err := openapicontroller.BuildAndRegisterAggregator(
			&specDownloader,
			delegationTarget,
			s.GenericAPIServer.Handler.GoRestfulContainer.RegisteredWebServices(),
			openApiConfig,
			s.GenericAPIServer.Handler.NonGoRestfulMux)
		if err != nil {
			return nil, err
		}
		s.openAPIAggregationController = openapicontroller.NewAggregationController(&specDownloader, openAPIAggregator)

		s.GenericAPIServer.AddPostStartHook("apiservice-openapi-controller", func(context genericapiserver.PostStartHookContext) error {
			go s.openAPIAggregationController.Run(context.StopCh)
			return nil
		})
	}

	return s, nil
}

// AddAPIService adds an API service.  It is not thread-safe, so only call it on one thread at a time please.
// It's a slow moving API, so its ok to run the controller on a single thread
func (s *APIAggregator) AddAPIService(apiService *apiregistration.APIService) error {
	// if the proxyHandler already exists, it needs to be updated. The aggregation bits do not
	// since they are wired against listers because they require multiple resources to respond
	if proxyHandler, exists := s.proxyHandlers[apiService.Name]; exists {
		proxyHandler.updateAPIService(apiService)
		if s.openAPIAggregationController != nil {
			s.openAPIAggregationController.UpdateAPIService(proxyHandler, apiService)
		}
		return nil
	}

	proxyPath := "/apis/" + apiService.Spec.Group + "/" + apiService.Spec.Version
	// v1. is a special case for the legacy API.  It proxies to a wider set of endpoints.
	if apiService.Name == legacyAPIServiceName {
		proxyPath = "/api"
	}

	// register the proxy handler
	proxyHandler := &proxyHandler{
		localDelegate:   s.delegateHandler,
		proxyClientCert: s.proxyClientCert,
		proxyClientKey:  s.proxyClientKey,
		proxyTransport:  s.proxyTransport,
		serviceResolver: s.serviceResolver,
	}
	proxyHandler.updateAPIService(apiService)
	if s.openAPIAggregationController != nil {
		s.openAPIAggregationController.AddAPIService(proxyHandler, apiService)
	}
	s.proxyHandlers[apiService.Name] = proxyHandler
	s.GenericAPIServer.Handler.NonGoRestfulMux.Handle(proxyPath, proxyHandler)
	s.GenericAPIServer.Handler.NonGoRestfulMux.UnlistedHandlePrefix(proxyPath+"/", proxyHandler)

	// if we're dealing with the legacy group, we're done here
	if apiService.Name == legacyAPIServiceName {
		return nil
	}

	// if we've already registered the path with the handler, we don't want to do it again.
	if s.handledGroups.Has(apiService.Spec.Group) {
		return nil
	}

	// it's time to register the group aggregation endpoint
	groupPath := "/apis/" + apiService.Spec.Group
	groupDiscoveryHandler := &apiGroupHandler{
		codecs:    aggregatorscheme.Codecs,
		groupName: apiService.Spec.Group,
		lister:    s.lister,
		delegate:  s.delegateHandler,
	}
	// aggregation is protected
	s.GenericAPIServer.Handler.NonGoRestfulMux.Handle(groupPath, groupDiscoveryHandler)
	s.GenericAPIServer.Handler.NonGoRestfulMux.UnlistedHandle(groupPath+"/", groupDiscoveryHandler)
	s.handledGroups.Insert(apiService.Spec.Group)
	return nil
}

// RemoveAPIService removes the APIService from being handled.  It is not thread-safe, so only call it on one thread at a time please.
// It's a slow moving API, so its ok to run the controller on a single thread.
func (s *APIAggregator) RemoveAPIService(apiServiceName string) {
	version := apiregistration.APIServiceNameToGroupVersion(apiServiceName)

	proxyPath := "/apis/" + version.Group + "/" + version.Version
	// v1. is a special case for the legacy API.  It proxies to a wider set of endpoints.
	if apiServiceName == legacyAPIServiceName {
		proxyPath = "/api"
	}
	s.GenericAPIServer.Handler.NonGoRestfulMux.Unregister(proxyPath)
	s.GenericAPIServer.Handler.NonGoRestfulMux.Unregister(proxyPath + "/")
	if s.openAPIAggregationController != nil {
		s.openAPIAggregationController.RemoveAPIService(apiServiceName)
	}
	delete(s.proxyHandlers, apiServiceName)

	// TODO unregister group level discovery when there are no more versions for the group
	// We don't need this right away because the handler properly delegates when no versions are present
}

func DefaultAPIResourceConfigSource() *serverstorage.ResourceConfig {
	ret := serverstorage.NewResourceConfig()
	// NOTE: GroupVersions listed here will be enabled by default. Don't put alpha versions in the list.
	ret.EnableVersions(
		v1.SchemeGroupVersion,
		v1beta1.SchemeGroupVersion,
	)

	return ret
}
<commit_msg>UPSTREAM: 00000: wait for apiservices on startup<commit_after>/*
Copyright 2016 The Kubernetes 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.
*/

package apiserver

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

	"github.com/golang/glog"

	"k8s.io/apimachinery/pkg/labels"

	"k8s.io/apimachinery/pkg/util/wait"

	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime/schema"
	utilruntime "k8s.io/apimachinery/pkg/util/runtime"
	"k8s.io/apimachinery/pkg/util/sets"
	genericapiserver "k8s.io/apiserver/pkg/server"
	serverstorage "k8s.io/apiserver/pkg/server/storage"
	"k8s.io/client-go/pkg/version"

	"k8s.io/kube-aggregator/pkg/apis/apiregistration"
	apiregistrationapi "k8s.io/kube-aggregator/pkg/apis/apiregistration"
	"k8s.io/kube-aggregator/pkg/apis/apiregistration/v1"
	"k8s.io/kube-aggregator/pkg/apis/apiregistration/v1beta1"
	aggregatorscheme "k8s.io/kube-aggregator/pkg/apiserver/scheme"
	"k8s.io/kube-aggregator/pkg/client/clientset_generated/internalclientset"
	informers "k8s.io/kube-aggregator/pkg/client/informers/internalversion"
	listers "k8s.io/kube-aggregator/pkg/client/listers/apiregistration/internalversion"
	openapicontroller "k8s.io/kube-aggregator/pkg/controllers/openapi"
	statuscontrollers "k8s.io/kube-aggregator/pkg/controllers/status"
	apiservicerest "k8s.io/kube-aggregator/pkg/registry/apiservice/rest"
)

func init() {
	// we need to add the options (like ListOptions) to empty v1
	metav1.AddToGroupVersion(aggregatorscheme.Scheme, schema.GroupVersion{Group: "", Version: "v1"})

	unversioned := schema.GroupVersion{Group: "", Version: "v1"}
	aggregatorscheme.Scheme.AddUnversionedTypes(unversioned,
		&metav1.Status{},
		&metav1.APIVersions{},
		&metav1.APIGroupList{},
		&metav1.APIGroup{},
		&metav1.APIResourceList{},
	)
}

// legacyAPIServiceName is the fixed name of the only non-groupified API version
const legacyAPIServiceName = "v1."

type ExtraConfig struct {
	// ProxyClientCert/Key are the client cert used to identify this proxy. Backing APIServices use
	// this to confirm the proxy's identity
	ProxyClientCert []byte
	ProxyClientKey  []byte

	// If present, the Dial method will be used for dialing out to delegate
	// apiservers.
	ProxyTransport *http.Transport

	// Mechanism by which the Aggregator will resolve services. Required.
	ServiceResolver ServiceResolver
}

type Config struct {
	GenericConfig *genericapiserver.RecommendedConfig
	ExtraConfig   ExtraConfig
}

type completedConfig struct {
	GenericConfig genericapiserver.CompletedConfig
	ExtraConfig   *ExtraConfig
}

type CompletedConfig struct {
	// Embed a private pointer that cannot be instantiated outside of this package.
	*completedConfig
}

// APIAggregator contains state for a Kubernetes cluster master/api server.
type APIAggregator struct {
	GenericAPIServer *genericapiserver.GenericAPIServer

	delegateHandler http.Handler

	// proxyClientCert/Key are the client cert used to identify this proxy. Backing APIServices use
	// this to confirm the proxy's identity
	proxyClientCert []byte
	proxyClientKey  []byte
	proxyTransport  *http.Transport

	// proxyHandlers are the proxy handlers that are currently registered, keyed by apiservice.name
	proxyHandlers map[string]*proxyHandler
	// handledGroups are the groups that already have routes
	handledGroups sets.String

	// lister is used to add group handling for /apis/<group> aggregator lookups based on
	// controller state
	lister listers.APIServiceLister

	// provided for easier embedding
	APIRegistrationInformers informers.SharedInformerFactory

	// Information needed to determine routing for the aggregator
	serviceResolver ServiceResolver

	openAPIAggregationController *openapicontroller.AggregationController
}

// Complete fills in any fields not set that are required to have valid data. It's mutating the receiver.
func (cfg *Config) Complete() CompletedConfig {
	c := completedConfig{
		cfg.GenericConfig.Complete(),
		&cfg.ExtraConfig,
	}

	// the kube aggregator wires its own discovery mechanism
	// TODO eventually collapse this by extracting all of the discovery out
	c.GenericConfig.EnableDiscovery = false
	version := version.Get()
	c.GenericConfig.Version = &version

	return CompletedConfig{&c}
}

// New returns a new instance of APIAggregator from the given config.
func (c completedConfig) NewWithDelegate(delegationTarget genericapiserver.DelegationTarget) (*APIAggregator, error) {
	// Prevent generic API server to install OpenAPI handler. Aggregator server
	// has its own customized OpenAPI handler.
	openApiConfig := c.GenericConfig.OpenAPIConfig
	c.GenericConfig.OpenAPIConfig = nil

	genericServer, err := c.GenericConfig.New("kube-aggregator", delegationTarget)
	if err != nil {
		return nil, err
	}

	apiregistrationClient, err := internalclientset.NewForConfig(c.GenericConfig.LoopbackClientConfig)
	if err != nil {
		return nil, err
	}
	informerFactory := informers.NewSharedInformerFactory(
		apiregistrationClient,
		5*time.Minute, // this is effectively used as a refresh interval right now.  Might want to do something nicer later on.
	)

	s := &APIAggregator{
		GenericAPIServer: genericServer,
		delegateHandler:  delegationTarget.UnprotectedHandler(),
		proxyClientCert:  c.ExtraConfig.ProxyClientCert,
		proxyClientKey:   c.ExtraConfig.ProxyClientKey,
		proxyTransport:   c.ExtraConfig.ProxyTransport,
		proxyHandlers:    map[string]*proxyHandler{},
		handledGroups:    sets.String{},
		lister:           informerFactory.Apiregistration().InternalVersion().APIServices().Lister(),
		APIRegistrationInformers: informerFactory,
		serviceResolver:          c.ExtraConfig.ServiceResolver,
	}

	apiGroupInfo := apiservicerest.NewRESTStorage(c.GenericConfig.MergedResourceConfig, c.GenericConfig.RESTOptionsGetter)
	if err := s.GenericAPIServer.InstallAPIGroup(&apiGroupInfo); err != nil {
		return nil, err
	}

	apisHandler := &apisHandler{
		codecs: aggregatorscheme.Codecs,
		lister: s.lister,
	}
	s.GenericAPIServer.Handler.NonGoRestfulMux.Handle("/apis", apisHandler)
	s.GenericAPIServer.Handler.NonGoRestfulMux.UnlistedHandle("/apis/", apisHandler)

	apiserviceRegistrationController := NewAPIServiceRegistrationController(informerFactory.Apiregistration().InternalVersion().APIServices(), s)
	availableController := statuscontrollers.NewAvailableConditionController(
		informerFactory.Apiregistration().InternalVersion().APIServices(),
		c.GenericConfig.SharedInformerFactory.Core().V1().Services(),
		c.GenericConfig.SharedInformerFactory.Core().V1().Endpoints(),
		apiregistrationClient.Apiregistration(),
		c.ExtraConfig.ProxyTransport,
		c.ExtraConfig.ProxyClientCert,
		c.ExtraConfig.ProxyClientKey,
		s.serviceResolver,
	)

	s.GenericAPIServer.AddPostStartHook("start-kube-aggregator-informers", func(context genericapiserver.PostStartHookContext) error {
		informerFactory.Start(context.StopCh)
		c.GenericConfig.SharedInformerFactory.Start(context.StopCh)
		return nil
	})
	s.GenericAPIServer.AddPostStartHook("apiservice-registration-controller", func(context genericapiserver.PostStartHookContext) error {
		go apiserviceRegistrationController.Run(context.StopCh)
		return nil
	})
	s.GenericAPIServer.AddPostStartHook("apiservice-status-available-controller", func(context genericapiserver.PostStartHookContext) error {
		// if we end up blocking for long periods of time, we may need to increase threadiness.
		go availableController.Run(5, context.StopCh)
		return nil
	})
	s.GenericAPIServer.AddPostStartHook("apiservice-wait-for-first-sync", func(context genericapiserver.PostStartHookContext) error {
		// when the aggregator first starts, it should make sure that it has proxy handlers for all the known good API services at this time
		// we only need to do this once.
		err := wait.PollImmediateUntil(100*time.Millisecond, func() (bool, error) {
			// fix race
			handledAPIServices := sets.StringKeySet(s.proxyHandlers)
			apiservices, err := s.lister.List(labels.Everything())
			if err != nil {
				return false, err
			}
			expectedAPIServices := sets.NewString()
			for _, apiservice := range apiservices {
				if apiregistrationapi.IsAPIServiceConditionTrue(apiservice, apiregistrationapi.Available) {
					expectedAPIServices.Insert(apiservice.Name)
				}
			}

			notYetHandledAPIServices := expectedAPIServices.Difference(handledAPIServices)
			if len(notYetHandledAPIServices) == 0 {
				return true, nil
			}
			glog.Infof("still waiting on handling APIServices: %v", strings.Join(notYetHandledAPIServices.List(), ","))

			return false, nil
		}, context.StopCh)
		return err
	})

	if openApiConfig != nil {
		specDownloader := openapicontroller.NewDownloader()
		openAPIAggregator, err := openapicontroller.BuildAndRegisterAggregator(
			&specDownloader,
			delegationTarget,
			s.GenericAPIServer.Handler.GoRestfulContainer.RegisteredWebServices(),
			openApiConfig,
			s.GenericAPIServer.Handler.NonGoRestfulMux)
		if err != nil {
			return nil, err
		}
		s.openAPIAggregationController = openapicontroller.NewAggregationController(&specDownloader, openAPIAggregator)

		s.GenericAPIServer.AddPostStartHook("apiservice-openapi-controller", func(context genericapiserver.PostStartHookContext) error {
			go s.openAPIAggregationController.Run(context.StopCh)
			return nil
		})
	}

	return s, nil
}

// AddAPIService adds an API service.  It is not thread-safe, so only call it on one thread at a time please.
// It's a slow moving API, so its ok to run the controller on a single thread
func (s *APIAggregator) AddAPIService(apiService *apiregistration.APIService) error {
	// if the proxyHandler already exists, it needs to be updated. The aggregation bits do not
	// since they are wired against listers because they require multiple resources to respond
	if proxyHandler, exists := s.proxyHandlers[apiService.Name]; exists {
		proxyHandler.updateAPIService(apiService)
		if s.openAPIAggregationController != nil {
			s.openAPIAggregationController.UpdateAPIService(proxyHandler, apiService)
		}
		return nil
	}

	proxyPath := "/apis/" + apiService.Spec.Group + "/" + apiService.Spec.Version
	// v1. is a special case for the legacy API.  It proxies to a wider set of endpoints.
	if apiService.Name == legacyAPIServiceName {
		proxyPath = "/api"
	}

	// register the proxy handler
	proxyHandler := &proxyHandler{
		localDelegate:   s.delegateHandler,
		proxyClientCert: s.proxyClientCert,
		proxyClientKey:  s.proxyClientKey,
		proxyTransport:  s.proxyTransport,
		serviceResolver: s.serviceResolver,
	}
	proxyHandler.updateAPIService(apiService)
	if s.openAPIAggregationController != nil {
		// this is calling a controller.  It should already handle being async.
		go func() {
			defer utilruntime.HandleCrash()
			s.openAPIAggregationController.AddAPIService(proxyHandler, apiService)
		}()
	}
	// we want to update the registration bit last after all the pieces are wired together
	defer func() {
		s.proxyHandlers[apiService.Name] = proxyHandler
	}()
	s.GenericAPIServer.Handler.NonGoRestfulMux.Handle(proxyPath, proxyHandler)
	s.GenericAPIServer.Handler.NonGoRestfulMux.UnlistedHandlePrefix(proxyPath+"/", proxyHandler)

	// if we're dealing with the legacy group, we're done here
	if apiService.Name == legacyAPIServiceName {
		return nil
	}

	// if we've already registered the path with the handler, we don't want to do it again.
	if s.handledGroups.Has(apiService.Spec.Group) {
		return nil
	}

	// it's time to register the group aggregation endpoint
	groupPath := "/apis/" + apiService.Spec.Group
	groupDiscoveryHandler := &apiGroupHandler{
		codecs:    aggregatorscheme.Codecs,
		groupName: apiService.Spec.Group,
		lister:    s.lister,
		delegate:  s.delegateHandler,
	}
	// aggregation is protected
	s.GenericAPIServer.Handler.NonGoRestfulMux.Handle(groupPath, groupDiscoveryHandler)
	s.GenericAPIServer.Handler.NonGoRestfulMux.UnlistedHandle(groupPath+"/", groupDiscoveryHandler)
	s.handledGroups.Insert(apiService.Spec.Group)
	return nil
}

// RemoveAPIService removes the APIService from being handled.  It is not thread-safe, so only call it on one thread at a time please.
// It's a slow moving API, so its ok to run the controller on a single thread.
func (s *APIAggregator) RemoveAPIService(apiServiceName string) {
	version := apiregistration.APIServiceNameToGroupVersion(apiServiceName)

	proxyPath := "/apis/" + version.Group + "/" + version.Version
	// v1. is a special case for the legacy API.  It proxies to a wider set of endpoints.
	if apiServiceName == legacyAPIServiceName {
		proxyPath = "/api"
	}
	s.GenericAPIServer.Handler.NonGoRestfulMux.Unregister(proxyPath)
	s.GenericAPIServer.Handler.NonGoRestfulMux.Unregister(proxyPath + "/")
	if s.openAPIAggregationController != nil {
		s.openAPIAggregationController.RemoveAPIService(apiServiceName)
	}
	delete(s.proxyHandlers, apiServiceName)

	// TODO unregister group level discovery when there are no more versions for the group
	// We don't need this right away because the handler properly delegates when no versions are present
}

func DefaultAPIResourceConfigSource() *serverstorage.ResourceConfig {
	ret := serverstorage.NewResourceConfig()
	// NOTE: GroupVersions listed here will be enabled by default. Don't put alpha versions in the list.
	ret.EnableVersions(
		v1.SchemeGroupVersion,
		v1beta1.SchemeGroupVersion,
	)

	return ret
}
<|endoftext|>


<commit_before>package libvirt

import (
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
	"os"
	"strconv"
	"strings"
	"time"

	libvirt "github.com/libvirt/libvirt-go"
	"github.com/libvirt/libvirt-go-xml"
)

// network transparent image
type image interface {
	Size() (uint64, error)
	Import(func(io.Reader) error, libvirtxml.StorageVolume) error
	String() string
	IsQCOW2() (bool, error)
}

type localImage struct {
	path string
}

func (i *localImage) String() string {
	return i.path
}

func (i *localImage) Size() (uint64, error) {
	file, err := os.Open(i.path)
	if err != nil {
		return 0, err
	}

	fi, err := file.Stat()
	if err != nil {
		return 0, err
	}
	return uint64(fi.Size()), nil
}

func (i *localImage) IsQCOW2() (bool, error) {
	file, err := os.Open(i.path)
	defer file.Close()
	if err != nil {
		return false, fmt.Errorf("Error while opening %s: %s", i.path, err)
	}
	buf := make([]byte, 8)
	_, err = io.ReadAtLeast(file, buf, 8)
	if err != nil {
		return false, err
	}
	return isQCOW2Header(buf)
}

func (i *localImage) Import(copier func(io.Reader) error, vol libvirtxml.StorageVolume) error {
	file, err := os.Open(i.path)
	defer file.Close()
	if err != nil {
		return fmt.Errorf("Error while opening %s: %s", i.path, err)
	}

	fi, err := file.Stat()
	if err != nil {
		return err
	}
	// we can skip the upload if the modification times are the same
	if vol.Target.Timestamps != nil && vol.Target.Timestamps.Mtime != "" {
		if fi.ModTime() == timeFromEpoch(vol.Target.Timestamps.Mtime) {
			log.Printf("Modification time is the same: skipping image copy")
			return nil
		}
	}

	return copier(file)
}

type httpImage struct {
	url *url.URL
}

func (i *httpImage) String() string {
	return i.url.String()
}

func (i *httpImage) Size() (uint64, error) {
	response, err := http.Head(i.url.String())
	if err != nil {
		return 0, err
	}
	if response.StatusCode != 200 {
		return 0,
			fmt.Errorf(
				"Error accessing remote resource: %s - %s",
				i.url.String(),
				response.Status)
	}

	length, err := strconv.Atoi(response.Header.Get("Content-Length"))
	if err != nil {
		err = fmt.Errorf(
			"Error while getting Content-Length of \"%s\": %s - got %s",
			i.url.String(),
			err,
			response.Header.Get("Content-Length"))
		return 0, err
	}
	return uint64(length), nil
}

func (i *httpImage) IsQCOW2() (bool, error) {
	client := &http.Client{}
	req, _ := http.NewRequest("GET", i.url.String(), nil)
	req.Header.Set("Range", "bytes=0-7")
	response, err := client.Do(req)

	if err != nil {
		return false, err
	}
	defer response.Body.Close()

	if response.StatusCode != 206 {
		return false, fmt.Errorf(
			"Can't retrieve partial header of resource to determine file type: %s - %s",
			i.url.String(),
			response.Status)
	}

	header, err := ioutil.ReadAll(response.Body)
	if err != nil {
		return false, err
	}

	if len(header) < 8 {
		return false, fmt.Errorf(
			"Can't retrieve read header of resource to determine file type: %s - %d bytes read",
			i.url.String(),
			len(header))
	}

	return isQCOW2Header(header)
}

// Whether the buffer is the header of a qcow2 file
func isQCOW2Header(buf []byte) (bool, error) {
	if len(buf) < 8 {
		return false, fmt.Errorf("Expected header of 8 bytes. Got %d", len(buf))
	}
	if buf[0] == 'Q' && buf[1] == 'F' && buf[2] == 'I' && buf[3] == 0xfb && buf[4] == 0x00 && buf[5] == 0x00 && buf[6] == 0x00 && buf[7] == 0x03 {
		return true, nil
	}
	return false, nil
}

func (i *httpImage) Import(copier func(io.Reader) error, vol libvirtxml.StorageVolume) error {
	client := &http.Client{}
	req, _ := http.NewRequest("GET", i.url.String(), nil)

	if vol.Target.Timestamps != nil && vol.Target.Timestamps.Mtime != "" {
		req.Header.Set("If-Modified-Since", timeFromEpoch(vol.Target.Timestamps.Mtime).UTC().Format(http.TimeFormat))
	}
	response, err := client.Do(req)

	if err != nil {
		return fmt.Errorf("Error while downloading %s: %s", i.url.String(), err)
	}

	defer response.Body.Close()
	if response.StatusCode == http.StatusNotModified {
		return nil
	}

	return copier(response.Body)
}

func newImage(source string) (image, error) {
	url, err := url.Parse(source)
	if err != nil {
		return nil, fmt.Errorf("Can't parse source '%s' as url: %s", source, err)
	}

	if strings.HasPrefix(url.Scheme, "http") {
		return &httpImage{url: url}, nil
	} else if url.Scheme == "file" || url.Scheme == "" {
		return &localImage{path: url.Path}, nil
	} else {
		return nil, fmt.Errorf("Don't know how to read from '%s': %s", url.String(), err)
	}
}

func newCopier(virConn *libvirt.Connect, volume *libvirt.StorageVol, size uint64) func(src io.Reader) error {
	copier := func(src io.Reader) error {
		var bytesCopied int64

		stream, err := virConn.NewStream(0)
		if err != nil {
			return err
		}

		defer func() {
			if uint64(bytesCopied) != size {
				stream.Abort()
			} else {
				stream.Finish()
			}
			stream.Free()
		}()

		volume.Upload(stream, 0, size, 0)

		sio := NewStreamIO(*stream)

		bytesCopied, err = io.Copy(sio, src)
		if err != nil {
			return err
		}
		log.Printf("%d bytes uploaded\n", bytesCopied)
		return nil
	}
	return copier
}

func timeFromEpoch(str string) time.Time {
	var s, ns int

	ts := strings.Split(str, ".")
	if len(ts) == 2 {
		ns, _ = strconv.Atoi(ts[1])
	}
	s, _ = strconv.Atoi(ts[0])

	return time.Unix(int64(s), int64(ns))
}

// removeVolume removes the volume identified by `key` from libvirt
func removeVolume(client *Client, key string) error {
	volume, err := client.libvirt.LookupStorageVolByKey(key)
	if err != nil {
		return fmt.Errorf("Can't retrieve volume %s: %v", key, err)
	}
	defer volume.Free()

	// Refresh the pool of the volume so that libvirt knows it is
	// not longer in use.
	volPool, err := volume.LookupPoolByVolume()
	if err != nil {
		return fmt.Errorf("Error retrieving pool for volume: %s", err)
	}
	defer volPool.Free()

	poolName, err := volPool.GetName()
	if err != nil {
		return fmt.Errorf("Error retrieving name of volume: %s", err)
	}

	client.poolMutexKV.Lock(poolName)
	defer client.poolMutexKV.Unlock(poolName)

	waitForSuccess("Error refreshing pool for volume", func() error {
		return volPool.Refresh(0)
	})

	// Workaround for redhat#1293804
	// https://bugzilla.redhat.com/show_bug.cgi?id=1293804#c12
	// Does not solve the problem but it makes it happen less often.
	_, err = volume.GetXMLDesc(0)
	if err != nil {
		return fmt.Errorf("Can't retrieve volume %s XML desc: %s", key, err)
	}

	err = volume.Delete(0)
	if err != nil {
		return fmt.Errorf("Can't delete volume %s: %s", key, err)
	}

	return nil
}

// tries really hard to find volume with `key`
// it will try to start the pool if it does not find it
//
// You have to call volume.Free() on the returned volume
func lookupVolumeReallyHard(client *Client, volPoolName string, key string) (*libvirt.StorageVol, error) {
	virConn := client.libvirt
	if virConn == nil {
		return nil, fmt.Errorf(LibVirtConIsNil)
	}

	volume, err := virConn.LookupStorageVolByKey(key)
	if err != nil {
		virErr := err.(libvirt.Error)
		if virErr.Code != libvirt.ERR_NO_STORAGE_VOL {
			return nil, fmt.Errorf("Can't retrieve volume %s", key)
		}
		log.Printf("[INFO] Volume %s not found, attempting to start its pool", key)

		volPool, err := virConn.LookupStoragePoolByName(volPoolName)
		if err != nil {
			return nil, fmt.Errorf("Error retrieving pool %s for volume %s: %s", volPoolName, key, err)
		}
		defer volPool.Free()

		active, err := volPool.IsActive()
		if err != nil {
			return nil, fmt.Errorf("error retrieving status of pool %s for volume %s: %s", volPoolName, key, err)
		}
		if active {
			log.Printf("Can't retrieve volume %s (and pool is active)", key)
			return nil, nil
		}

		err = volPool.Create(0)
		if err != nil {
			return nil, fmt.Errorf("error starting pool %s: %s", volPoolName, err)
		}

		// attempt a new lookup
		volume, err = virConn.LookupStorageVolByKey(key)
		if err != nil {
			virErr := err.(libvirt.Error)
			if virErr.Code != libvirt.ERR_NO_STORAGE_VOL {
				return nil, fmt.Errorf("Can't retrieve volume %s", key)
			}
			// does not exist, but no error
			return nil, nil
		}
	}
	return volume, nil
}
<commit_msg>improve mesg for unEOF so user can check better add more debug infos<commit_after>package libvirt

import (
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
	"os"
	"strconv"
	"strings"
	"time"

	libvirt "github.com/libvirt/libvirt-go"
	"github.com/libvirt/libvirt-go-xml"
)

// network transparent image
type image interface {
	Size() (uint64, error)
	Import(func(io.Reader) error, libvirtxml.StorageVolume) error
	String() string
	IsQCOW2() (bool, error)
}

type localImage struct {
	path string
}

func (i *localImage) String() string {
	return i.path
}

func (i *localImage) Size() (uint64, error) {
	file, err := os.Open(i.path)
	if err != nil {
		return 0, err
	}

	fi, err := file.Stat()
	if err != nil {
		return 0, err
	}
	return uint64(fi.Size()), nil
}

func (i *localImage) IsQCOW2() (bool, error) {
	file, err := os.Open(i.path)
	defer file.Close()
	if err != nil {
		return false, fmt.Errorf("Error while opening %s: %s", i.path, err)
	}
	buf := make([]byte, 8)
	_, err = io.ReadAtLeast(file, buf, 8)
	if err != nil {
		return false, err
	}
	return isQCOW2Header(buf)
}

func (i *localImage) Import(copier func(io.Reader) error, vol libvirtxml.StorageVolume) error {
	file, err := os.Open(i.path)
	defer file.Close()
	if err != nil {
		return fmt.Errorf("Error while opening %s: %s", i.path, err)
	}

	fi, err := file.Stat()
	if err != nil {
		return err
	}
	// we can skip the upload if the modification times are the same
	if vol.Target.Timestamps != nil && vol.Target.Timestamps.Mtime != "" {
		if fi.ModTime() == timeFromEpoch(vol.Target.Timestamps.Mtime) {
			log.Printf("Modification time is the same: skipping image copy")
			return nil
		}
	}

	return copier(file)
}

type httpImage struct {
	url *url.URL
}

func (i *httpImage) String() string {
	return i.url.String()
}

func (i *httpImage) Size() (uint64, error) {
	response, err := http.Head(i.url.String())
	if err != nil {
		return 0, err
	}
	if response.StatusCode != 200 {
		return 0,
			fmt.Errorf(
				"Error accessing remote resource: %s - %s",
				i.url.String(),
				response.Status)
	}

	length, err := strconv.Atoi(response.Header.Get("Content-Length"))
	if err != nil {
		err = fmt.Errorf(
			"Error while getting Content-Length of \"%s\": %s - got %s",
			i.url.String(),
			err,
			response.Header.Get("Content-Length"))
		return 0, err
	}
	return uint64(length), nil
}

func (i *httpImage) IsQCOW2() (bool, error) {
	client := &http.Client{}
	req, _ := http.NewRequest("GET", i.url.String(), nil)
	req.Header.Set("Range", "bytes=0-7")
	response, err := client.Do(req)

	if err != nil {
		return false, err
	}
	defer response.Body.Close()

	if response.StatusCode != 206 {
		return false, fmt.Errorf(
			"Can't retrieve partial header of resource to determine file type: %s - %s",
			i.url.String(),
			response.Status)
	}

	header, err := ioutil.ReadAll(response.Body)
	if err != nil {
		return false, err
	}

	if len(header) < 8 {
		return false, fmt.Errorf(
			"Can't retrieve read header of resource to determine file type: %s - %d bytes read",
			i.url.String(),
			len(header))
	}

	return isQCOW2Header(header)
}

// Whether the buffer is the header of a qcow2 file
func isQCOW2Header(buf []byte) (bool, error) {
	if len(buf) < 8 {
		return false, fmt.Errorf("Expected header of 8 bytes. Got %d", len(buf))
	}
	if buf[0] == 'Q' && buf[1] == 'F' && buf[2] == 'I' && buf[3] == 0xfb && buf[4] == 0x00 && buf[5] == 0x00 && buf[6] == 0x00 && buf[7] == 0x03 {
		return true, nil
	}
	return false, nil
}

func (i *httpImage) Import(copier func(io.Reader) error, vol libvirtxml.StorageVolume) error {
	client := &http.Client{}
	req, err := http.NewRequest("GET", i.url.String(), nil)

	if err != nil {
		log.Printf("[DEBUG:] Error creating new request for source url %s: %s", i.url.String(), err)
		return fmt.Errorf("Error while downloading %s: %s", i.url.String(), err)
	}

	if vol.Target.Timestamps != nil && vol.Target.Timestamps.Mtime != "" {
		req.Header.Set("If-Modified-Since", timeFromEpoch(vol.Target.Timestamps.Mtime).UTC().Format(http.TimeFormat))
	}
	response, err := client.Do(req)

	if err != nil {
		return fmt.Errorf("Error while downloading %s: %s", i.url.String(), err)
	}

	defer response.Body.Close()
	if response.StatusCode == http.StatusNotModified {
		return nil
	}
	log.Printf("[DEBUG]: url resp status code %s\n", response.Status)
	return copier(response.Body)
}

func newImage(source string) (image, error) {
	url, err := url.Parse(source)
	if err != nil {
		return nil, fmt.Errorf("Can't parse source '%s' as url: %s", source, err)
	}

	if strings.HasPrefix(url.Scheme, "http") {
		return &httpImage{url: url}, nil
	} else if url.Scheme == "file" || url.Scheme == "" {
		return &localImage{path: url.Path}, nil
	} else {
		return nil, fmt.Errorf("Don't know how to read from '%s': %s", url.String(), err)
	}
}

func newCopier(virConn *libvirt.Connect, volume *libvirt.StorageVol, size uint64) func(src io.Reader) error {
	copier := func(src io.Reader) error {
		var bytesCopied int64

		stream, err := virConn.NewStream(0)
		if err != nil {
			return err
		}

		defer func() {
			if uint64(bytesCopied) != size {
				stream.Abort()
			} else {
				stream.Finish()
			}
			stream.Free()
		}()

		if err := volume.Upload(stream, 0, size, 0); err != nil {
			stream.Abort()
			return fmt.Errorf("Error while uploading volume %s", err)
		}

		sio := NewStreamIO(*stream)

		bytesCopied, err = io.Copy(sio, src)
		// if we get unexpected EOF this mean that connection was closed suddently from server side
		// the problem is not on the plugin but on server hosting currupted images
		if err == io.ErrUnexpectedEOF {
			return fmt.Errorf("Error: transfer was unexpectedly closed from the server while downloading. Please try again later or check the server hosting sources")
		}
		if err != nil {
			return fmt.Errorf("Error while copying source to volume %s", err)
		}
		log.Printf("%d bytes uploaded\n", bytesCopied)
		return nil
	}
	return copier
}

func timeFromEpoch(str string) time.Time {
	var s, ns int

	ts := strings.Split(str, ".")
	if len(ts) == 2 {
		ns, _ = strconv.Atoi(ts[1])
	}
	s, _ = strconv.Atoi(ts[0])

	return time.Unix(int64(s), int64(ns))
}

// removeVolume removes the volume identified by `key` from libvirt
func removeVolume(client *Client, key string) error {
	volume, err := client.libvirt.LookupStorageVolByKey(key)
	if err != nil {
		return fmt.Errorf("Can't retrieve volume %s: %v", key, err)
	}
	defer volume.Free()

	// Refresh the pool of the volume so that libvirt knows it is
	// not longer in use.
	volPool, err := volume.LookupPoolByVolume()
	if err != nil {
		return fmt.Errorf("Error retrieving pool for volume: %s", err)
	}
	defer volPool.Free()

	poolName, err := volPool.GetName()
	if err != nil {
		return fmt.Errorf("Error retrieving name of volume: %s", err)
	}

	client.poolMutexKV.Lock(poolName)
	defer client.poolMutexKV.Unlock(poolName)

	waitForSuccess("Error refreshing pool for volume", func() error {
		return volPool.Refresh(0)
	})

	// Workaround for redhat#1293804
	// https://bugzilla.redhat.com/show_bug.cgi?id=1293804#c12
	// Does not solve the problem but it makes it happen less often.
	_, err = volume.GetXMLDesc(0)
	if err != nil {
		return fmt.Errorf("Can't retrieve volume %s XML desc: %s", key, err)
	}

	err = volume.Delete(0)
	if err != nil {
		return fmt.Errorf("Can't delete volume %s: %s", key, err)
	}

	return nil
}

// tries really hard to find volume with `key`
// it will try to start the pool if it does not find it
//
// You have to call volume.Free() on the returned volume
func lookupVolumeReallyHard(client *Client, volPoolName string, key string) (*libvirt.StorageVol, error) {
	virConn := client.libvirt
	if virConn == nil {
		return nil, fmt.Errorf(LibVirtConIsNil)
	}

	volume, err := virConn.LookupStorageVolByKey(key)
	if err != nil {
		virErr := err.(libvirt.Error)
		if virErr.Code != libvirt.ERR_NO_STORAGE_VOL {
			return nil, fmt.Errorf("Can't retrieve volume %s", key)
		}
		log.Printf("[INFO] Volume %s not found, attempting to start its pool", key)

		volPool, err := virConn.LookupStoragePoolByName(volPoolName)
		if err != nil {
			return nil, fmt.Errorf("Error retrieving pool %s for volume %s: %s", volPoolName, key, err)
		}
		defer volPool.Free()

		active, err := volPool.IsActive()
		if err != nil {
			return nil, fmt.Errorf("error retrieving status of pool %s for volume %s: %s", volPoolName, key, err)
		}
		if active {
			log.Printf("Can't retrieve volume %s (and pool is active)", key)
			return nil, nil
		}

		err = volPool.Create(0)
		if err != nil {
			return nil, fmt.Errorf("error starting pool %s: %s", volPoolName, err)
		}

		// attempt a new lookup
		volume, err = virConn.LookupStorageVolByKey(key)
		if err != nil {
			virErr := err.(libvirt.Error)
			if virErr.Code != libvirt.ERR_NO_STORAGE_VOL {
				return nil, fmt.Errorf("Can't retrieve volume %s", key)
			}
			// does not exist, but no error
			return nil, nil
		}
	}
	return volume, nil
}
<|endoftext|>