<commit_before>package kcp

import (
	"sync"
)

type SendingWindow struct {
	start uint32
	cap   uint32
	len   uint32
	last  uint32

	data  []DataSegment
	inuse []bool
	prev  []uint32
	next  []uint32

	totalInFlightSize uint32
	writer            SegmentWriter
	onPacketLoss      func(uint32)
}

func NewSendingWindow(size uint32, writer SegmentWriter, onPacketLoss func(uint32)) *SendingWindow {
	window := &SendingWindow{
		start:        0,
		cap:          size,
		len:          0,
		last:         0,
		data:         make([]DataSegment, size),
		prev:         make([]uint32, size),
		next:         make([]uint32, size),
		inuse:        make([]bool, size),
		writer:       writer,
		onPacketLoss: onPacketLoss,
	}
	return window
}

func (v *SendingWindow) Release() {
	if v == nil {
		return
	}
	v.len = 0
	for _, seg := range v.data {
		seg.Release()
	}
}

func (v *SendingWindow) Len() int {
	return int(v.len)
}

func (v *SendingWindow) IsEmpty() bool {
	return v.len == 0
}

func (v *SendingWindow) Size() uint32 {
	return v.cap
}

func (v *SendingWindow) IsFull() bool {
	return v.len == v.cap
}

func (v *SendingWindow) Push(number uint32, data []byte) {
	pos := (v.start + v.len) % v.cap
	v.data[pos].SetData(data)
	v.data[pos].Number = number
	v.data[pos].timeout = 0
	v.data[pos].transmit = 0
	v.inuse[pos] = true
	if v.len > 0 {
		v.next[v.last] = pos
		v.prev[pos] = v.last
	}
	v.last = pos
	v.len++
}

func (v *SendingWindow) FirstNumber() uint32 {
	return v.data[v.start].Number
}

func (v *SendingWindow) Clear(una uint32) {
	for !v.IsEmpty() && v.data[v.start].Number < una {
		v.Remove(0)
	}
}

func (v *SendingWindow) Remove(idx uint32) bool {
	if v.IsEmpty() {
		return false
	}

	pos := (v.start + idx) % v.cap
	if !v.inuse[pos] {
		return false
	}
	v.inuse[pos] = false
	v.totalInFlightSize--
	if pos == v.start && pos == v.last {
		v.len = 0
		v.start = 0
		v.last = 0
	} else if pos == v.start {
		delta := v.next[pos] - v.start
		if v.next[pos] < v.start {
			delta = v.next[pos] + v.cap - v.start
		}
		v.start = v.next[pos]
		v.len -= delta
	} else if pos == v.last {
		v.last = v.prev[pos]
	} else {
		v.next[v.prev[pos]] = v.next[pos]
		v.prev[v.next[pos]] = v.prev[pos]
	}
	return true
}

func (v *SendingWindow) HandleFastAck(number uint32, rto uint32) {
	if v.IsEmpty() {
		return
	}

	v.Visit(func(seg *DataSegment) bool {
		if number == seg.Number || number-seg.Number > 0x7FFFFFFF {
			return false
		}

		if seg.transmit > 0 && seg.timeout > rto/3 {
			seg.timeout -= rto / 3
		}
		return true
	})
}

func (v *SendingWindow) Visit(visitor func(seg *DataSegment) bool) {
	if v.IsEmpty() {
		return
	}

	for i := v.start; ; i = v.next[i] {
		if !visitor(&v.data[i]) || i == v.last {
			break
		}
	}
}

func (v *SendingWindow) Flush(current uint32, rto uint32, maxInFlightSize uint32) {
	if v.IsEmpty() {
		return
	}

	var lost uint32
	var inFlightSize uint32

	v.Visit(func(segment *DataSegment) bool {
		if current-segment.timeout >= 0x7FFFFFFF {
			return true
		}
		if segment.transmit == 0 {
			// First time
			v.totalInFlightSize++
		} else {
			lost++
		}
		segment.timeout = current + rto

		segment.Timestamp = current
		segment.transmit++
		v.writer.Write(segment)
		inFlightSize++
		if inFlightSize >= maxInFlightSize {
			return false
		}
		return true
	})

	if v.onPacketLoss != nil && inFlightSize > 0 && v.totalInFlightSize != 0 {
		rate := lost * 100 / v.totalInFlightSize
		v.onPacketLoss(rate)
	}
}

type SendingWorker struct {
	sync.RWMutex
	conn                       *Connection
	window                     *SendingWindow
	firstUnacknowledged        uint32
	firstUnacknowledgedUpdated bool
	nextNumber                 uint32
	remoteNextNumber           uint32
	controlWindow              uint32
	fastResend                 uint32
}

func NewSendingWorker(kcp *Connection) *SendingWorker {
	worker := &SendingWorker{
		conn:             kcp,
		fastResend:       2,
		remoteNextNumber: 32,
		controlWindow:    kcp.Config.GetSendingInFlightSize(),
	}
	worker.window = NewSendingWindow(kcp.Config.GetSendingBufferSize(), worker, worker.OnPacketLoss)
	return worker
}

func (v *SendingWorker) Release() {
	v.Lock()
	v.window.Release()
	v.Unlock()
}

func (v *SendingWorker) ProcessReceivingNext(nextNumber uint32) {
	v.Lock()
	defer v.Unlock()

	v.ProcessReceivingNextWithoutLock(nextNumber)
}

func (v *SendingWorker) ProcessReceivingNextWithoutLock(nextNumber uint32) {
	v.window.Clear(nextNumber)
	v.FindFirstUnacknowledged()
}

// Private: Visible for testing.
func (v *SendingWorker) FindFirstUnacknowledged() {
	first := v.firstUnacknowledged
	if !v.window.IsEmpty() {
		v.firstUnacknowledged = v.window.FirstNumber()
	} else {
		v.firstUnacknowledged = v.nextNumber
	}
	if first != v.firstUnacknowledged {
		v.firstUnacknowledgedUpdated = true
	}
}

// Private: Visible for testing.
func (v *SendingWorker) ProcessAck(number uint32) bool {
	// number < v.firstUnacknowledged || number >= v.nextNumber
	if number-v.firstUnacknowledged > 0x7FFFFFFF || number-v.nextNumber < 0x7FFFFFFF {
		return false
	}

	removed := v.window.Remove(number - v.firstUnacknowledged)
	if removed {
		v.FindFirstUnacknowledged()
	}
	return removed
}

func (v *SendingWorker) ProcessSegment(current uint32, seg *AckSegment, rto uint32) {
	defer seg.Release()

	v.Lock()
	defer v.Unlock()

	if v.remoteNextNumber < seg.ReceivingWindow {
		v.remoteNextNumber = seg.ReceivingWindow
	}
	v.ProcessReceivingNextWithoutLock(seg.ReceivingNext)

	if seg.IsEmpty() {
		return
	}

	var maxack uint32
	var maxackRemoved bool
	for _, number := range seg.NumberList {
		removed := v.ProcessAck(number)
		if maxack < number {
			maxack = number
			maxackRemoved = removed
		}
	}

	if maxackRemoved {
		v.window.HandleFastAck(maxack, rto)
		if current-seg.Timestamp < 10000 {
			v.conn.roundTrip.Update(current-seg.Timestamp, current)
		}
	}
}

func (v *SendingWorker) Push(b []byte) int {
	nBytes := 0
	v.Lock()
	defer v.Unlock()

	for len(b) > 0 && !v.window.IsFull() {
		var size int
		if len(b) > int(v.conn.mss) {
			size = int(v.conn.mss)
		} else {
			size = len(b)
		}
		v.window.Push(v.nextNumber, b[:size])
		v.nextNumber++
		b = b[size:]
		nBytes += size
	}
	return nBytes
}

// Private: Visible for testing.
func (v *SendingWorker) Write(seg Segment) error {
	dataSeg := seg.(*DataSegment)

	dataSeg.Conv = v.conn.conv
	dataSeg.SendingNext = v.firstUnacknowledged
	dataSeg.Option = 0
	if v.conn.State() == StateReadyToClose {
		dataSeg.Option = SegmentOptionClose
	}

	return v.conn.output.Write(dataSeg)
}

func (v *SendingWorker) OnPacketLoss(lossRate uint32) {
	if !v.conn.Config.Congestion || v.conn.roundTrip.Timeout() == 0 {
		return
	}

	if lossRate >= 15 {
		v.controlWindow = 3 * v.controlWindow / 4
	} else if lossRate <= 5 {
		v.controlWindow += v.controlWindow / 4
	}
	if v.controlWindow < 16 {
		v.controlWindow = 16
	}
	if v.controlWindow > 2*v.conn.Config.GetSendingInFlightSize() {
		v.controlWindow = 2 * v.conn.Config.GetSendingInFlightSize()
	}
}

func (v *SendingWorker) Flush(current uint32) {
	v.Lock()

	cwnd := v.firstUnacknowledged + v.conn.Config.GetSendingInFlightSize()
	if cwnd > v.remoteNextNumber {
		cwnd = v.remoteNextNumber
	}
	if v.conn.Config.Congestion && cwnd > v.firstUnacknowledged+v.controlWindow {
		cwnd = v.firstUnacknowledged + v.controlWindow
	}

	if !v.window.IsEmpty() {
		v.window.Flush(current, v.conn.roundTrip.Timeout(), cwnd)
		v.firstUnacknowledgedUpdated = false
	}

	updated := v.firstUnacknowledgedUpdated
	v.firstUnacknowledgedUpdated = false

	v.Unlock()

	if updated {
		v.conn.Ping(current, CommandPing)
	}
}

func (v *SendingWorker) CloseWrite() {
	v.Lock()
	defer v.Unlock()

	v.window.Clear(0xFFFFFFFF)
}

func (v *SendingWorker) IsEmpty() bool {
	v.RLock()
	defer v.RUnlock()

	return v.window.IsEmpty()
}

func (v *SendingWorker) UpdateNecessary() bool {
	return !v.IsEmpty()
}

func (w *SendingWorker) FirstUnacknowledged() uint32 {
	w.RLock()
	defer w.RUnlock()

	return w.firstUnacknowledged
}
<commit_msg>comments<commit_after>package kcp

import (
	"sync"
)

type SendingWindow struct {
	start uint32
	cap   uint32
	len   uint32
	last  uint32

	data  []DataSegment
	inuse []bool
	prev  []uint32
	next  []uint32

	totalInFlightSize uint32
	writer            SegmentWriter
	onPacketLoss      func(uint32)
}

func NewSendingWindow(size uint32, writer SegmentWriter, onPacketLoss func(uint32)) *SendingWindow {
	window := &SendingWindow{
		start:        0,
		cap:          size,
		len:          0,
		last:         0,
		data:         make([]DataSegment, size),
		prev:         make([]uint32, size),
		next:         make([]uint32, size),
		inuse:        make([]bool, size),
		writer:       writer,
		onPacketLoss: onPacketLoss,
	}
	return window
}

func (v *SendingWindow) Release() {
	if v == nil {
		return
	}
	v.len = 0
	for _, seg := range v.data {
		seg.Release()
	}
}

func (v *SendingWindow) Len() int {
	return int(v.len)
}

func (v *SendingWindow) IsEmpty() bool {
	return v.len == 0
}

func (v *SendingWindow) Size() uint32 {
	return v.cap
}

func (v *SendingWindow) IsFull() bool {
	return v.len == v.cap
}

func (v *SendingWindow) Push(number uint32, data []byte) {
	pos := (v.start + v.len) % v.cap
	v.data[pos].SetData(data)
	v.data[pos].Number = number
	v.data[pos].timeout = 0
	v.data[pos].transmit = 0
	v.inuse[pos] = true
	if v.len > 0 {
		v.next[v.last] = pos
		v.prev[pos] = v.last
	}
	v.last = pos
	v.len++
}

func (v *SendingWindow) FirstNumber() uint32 {
	return v.data[v.start].Number
}

func (v *SendingWindow) Clear(una uint32) {
	for !v.IsEmpty() && v.data[v.start].Number < una {
		v.Remove(0)
	}
}

func (v *SendingWindow) Remove(idx uint32) bool {
	if v.IsEmpty() {
		return false
	}

	pos := (v.start + idx) % v.cap
	if !v.inuse[pos] {
		return false
	}
	v.inuse[pos] = false
	v.totalInFlightSize--
	if pos == v.start && pos == v.last {
		v.len = 0
		v.start = 0
		v.last = 0
	} else if pos == v.start {
		delta := v.next[pos] - v.start
		if v.next[pos] < v.start {
			delta = v.next[pos] + v.cap - v.start
		}
		v.start = v.next[pos]
		v.len -= delta
	} else if pos == v.last {
		v.last = v.prev[pos]
	} else {
		v.next[v.prev[pos]] = v.next[pos]
		v.prev[v.next[pos]] = v.prev[pos]
	}
	return true
}

func (v *SendingWindow) HandleFastAck(number uint32, rto uint32) {
	if v.IsEmpty() {
		return
	}

	v.Visit(func(seg *DataSegment) bool {
		if number == seg.Number || number-seg.Number > 0x7FFFFFFF {
			return false
		}

		if seg.transmit > 0 && seg.timeout > rto/3 {
			seg.timeout -= rto / 3
		}
		return true
	})
}

func (v *SendingWindow) Visit(visitor func(seg *DataSegment) bool) {
	if v.IsEmpty() {
		return
	}

	for i := v.start; ; i = v.next[i] {
		if !visitor(&v.data[i]) || i == v.last {
			break
		}
	}
}

func (v *SendingWindow) Flush(current uint32, rto uint32, maxInFlightSize uint32) {
	if v.IsEmpty() {
		return
	}

	var lost uint32
	var inFlightSize uint32

	v.Visit(func(segment *DataSegment) bool {
		if current-segment.timeout >= 0x7FFFFFFF {
			return true
		}
		if segment.transmit == 0 {
			// First time
			v.totalInFlightSize++
		} else {
			lost++
		}
		segment.timeout = current + rto

		segment.Timestamp = current
		segment.transmit++
		v.writer.Write(segment)
		inFlightSize++
		if inFlightSize >= maxInFlightSize {
			return false
		}
		return true
	})

	if v.onPacketLoss != nil && inFlightSize > 0 && v.totalInFlightSize != 0 {
		rate := lost * 100 / v.totalInFlightSize
		v.onPacketLoss(rate)
	}
}

type SendingWorker struct {
	sync.RWMutex
	conn                       *Connection
	window                     *SendingWindow
	firstUnacknowledged        uint32
	firstUnacknowledgedUpdated bool
	nextNumber                 uint32
	remoteNextNumber           uint32
	controlWindow              uint32
	fastResend                 uint32
}

func NewSendingWorker(kcp *Connection) *SendingWorker {
	worker := &SendingWorker{
		conn:             kcp,
		fastResend:       2,
		remoteNextNumber: 32,
		controlWindow:    kcp.Config.GetSendingInFlightSize(),
	}
	worker.window = NewSendingWindow(kcp.Config.GetSendingBufferSize(), worker, worker.OnPacketLoss)
	return worker
}

func (v *SendingWorker) Release() {
	v.Lock()
	v.window.Release()
	v.Unlock()
}

func (v *SendingWorker) ProcessReceivingNext(nextNumber uint32) {
	v.Lock()
	defer v.Unlock()

	v.ProcessReceivingNextWithoutLock(nextNumber)
}

func (v *SendingWorker) ProcessReceivingNextWithoutLock(nextNumber uint32) {
	v.window.Clear(nextNumber)
	v.FindFirstUnacknowledged()
}

// Private: Visible for testing.
func (v *SendingWorker) FindFirstUnacknowledged() {
	first := v.firstUnacknowledged
	if !v.window.IsEmpty() {
		v.firstUnacknowledged = v.window.FirstNumber()
	} else {
		v.firstUnacknowledged = v.nextNumber
	}
	if first != v.firstUnacknowledged {
		v.firstUnacknowledgedUpdated = true
	}
}

func (v *SendingWorker) processAck(number uint32) bool {
	// number < v.firstUnacknowledged || number >= v.nextNumber
	if number-v.firstUnacknowledged > 0x7FFFFFFF || number-v.nextNumber < 0x7FFFFFFF {
		return false
	}

	removed := v.window.Remove(number - v.firstUnacknowledged)
	if removed {
		v.FindFirstUnacknowledged()
	}
	return removed
}

func (v *SendingWorker) ProcessSegment(current uint32, seg *AckSegment, rto uint32) {
	defer seg.Release()

	v.Lock()
	defer v.Unlock()

	if v.remoteNextNumber < seg.ReceivingWindow {
		v.remoteNextNumber = seg.ReceivingWindow
	}
	v.ProcessReceivingNextWithoutLock(seg.ReceivingNext)

	if seg.IsEmpty() {
		return
	}

	var maxack uint32
	var maxackRemoved bool
	for _, number := range seg.NumberList {
		removed := v.processAck(number)
		if maxack < number {
			maxack = number
			maxackRemoved = removed
		}
	}

	if maxackRemoved {
		v.window.HandleFastAck(maxack, rto)
		if current-seg.Timestamp < 10000 {
			v.conn.roundTrip.Update(current-seg.Timestamp, current)
		}
	}
}

func (v *SendingWorker) Push(b []byte) int {
	nBytes := 0
	v.Lock()
	defer v.Unlock()

	for len(b) > 0 && !v.window.IsFull() {
		var size int
		if len(b) > int(v.conn.mss) {
			size = int(v.conn.mss)
		} else {
			size = len(b)
		}
		v.window.Push(v.nextNumber, b[:size])
		v.nextNumber++
		b = b[size:]
		nBytes += size
	}
	return nBytes
}

// Private: Visible for testing.
func (v *SendingWorker) Write(seg Segment) error {
	dataSeg := seg.(*DataSegment)

	dataSeg.Conv = v.conn.conv
	dataSeg.SendingNext = v.firstUnacknowledged
	dataSeg.Option = 0
	if v.conn.State() == StateReadyToClose {
		dataSeg.Option = SegmentOptionClose
	}

	return v.conn.output.Write(dataSeg)
}

func (v *SendingWorker) OnPacketLoss(lossRate uint32) {
	if !v.conn.Config.Congestion || v.conn.roundTrip.Timeout() == 0 {
		return
	}

	if lossRate >= 15 {
		v.controlWindow = 3 * v.controlWindow / 4
	} else if lossRate <= 5 {
		v.controlWindow += v.controlWindow / 4
	}
	if v.controlWindow < 16 {
		v.controlWindow = 16
	}
	if v.controlWindow > 2*v.conn.Config.GetSendingInFlightSize() {
		v.controlWindow = 2 * v.conn.Config.GetSendingInFlightSize()
	}
}

func (v *SendingWorker) Flush(current uint32) {
	v.Lock()

	cwnd := v.firstUnacknowledged + v.conn.Config.GetSendingInFlightSize()
	if cwnd > v.remoteNextNumber {
		cwnd = v.remoteNextNumber
	}
	if v.conn.Config.Congestion && cwnd > v.firstUnacknowledged+v.controlWindow {
		cwnd = v.firstUnacknowledged + v.controlWindow
	}

	if !v.window.IsEmpty() {
		v.window.Flush(current, v.conn.roundTrip.Timeout(), cwnd)
		v.firstUnacknowledgedUpdated = false
	}

	updated := v.firstUnacknowledgedUpdated
	v.firstUnacknowledgedUpdated = false

	v.Unlock()

	if updated {
		v.conn.Ping(current, CommandPing)
	}
}

func (v *SendingWorker) CloseWrite() {
	v.Lock()
	defer v.Unlock()

	v.window.Clear(0xFFFFFFFF)
}

func (v *SendingWorker) IsEmpty() bool {
	v.RLock()
	defer v.RUnlock()

	return v.window.IsEmpty()
}

func (v *SendingWorker) UpdateNecessary() bool {
	return !v.IsEmpty()
}

func (w *SendingWorker) FirstUnacknowledged() uint32 {
	w.RLock()
	defer w.RUnlock()

	return w.firstUnacknowledged
}
<|endoftext|>


<commit_before>// Copyright (c) 2016 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package transporttest

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"strings"
	"testing"

	"github.com/yarpc/yarpc-go/transport"
)

// RequestMatcher may be used in gomock argument lists to assert that two
// requests match.
//
// Requests are considered to be matching if:
//
// - All their primitive parameters match.
// - The headers of the received request include all the headers from the
//   source request. It may include extra headers.
// - The contents of the request bodies are the same.
type RequestMatcher struct {
	t    *testing.T
	req  *transport.Request
	body []byte
}

// NewRequestMatcher constructs a new RequestMatcher from the given testing.T
// and request.
//
// The request's contents are read in their entirety and replaced with a
// bytes.Reader.
func NewRequestMatcher(t *testing.T, r *transport.Request) RequestMatcher {
	body, err := ioutil.ReadAll(r.Body)
	if err != nil {
		t.Fatalf("failed to read request body: %v", err)
	}

	// restore a copy of the body so that the caller can still use the request
	// object
	r.Body = bytes.NewReader(body)
	return RequestMatcher{t: t, req: r, body: body}
}

// TODO: Headers like User-Agent, Content-Length, etc. make their way to the
// user-level Request object. For now, we're doing the super set check but we
// should do something more specific once yarpc/yarpc.io#2 is resolved.

// Matches checks if the given object matches the Request provided in
// NewRequestMatcher.
func (m RequestMatcher) Matches(got interface{}) bool {
	l := m.req
	r, ok := got.(*transport.Request)
	if !ok {
		panic(fmt.Sprintf("expected *transport.Request, got %v", got))
	}

	if l.Caller != r.Caller {
		m.t.Logf("Caller mismatch: %s != %s", l.Caller, r.Caller)
		return false
	}

	if l.Service != r.Service {
		m.t.Logf("Service mismatch: %s != %s", l.Service, r.Service)
		return false
	}

	if l.Encoding != r.Encoding {
		m.t.Logf("Encoding mismatch: %s != %s", l.Service, r.Service)
		return false
	}

	if l.Procedure != r.Procedure {
		m.t.Logf("Procedure mismatch: %s != %s", l.Procedure, r.Procedure)
		return false
	}

	if l.TTL != r.TTL {
		m.t.Logf("TTL mismatch: %v != %v", l.TTL, r.TTL)
		return false
	}

	if err := checkSuperSet(l.Headers, r.Headers); err != nil {
		m.t.Logf("Headers mismatch: %v != %v\n\t%v", l.Headers, r.Headers, err)
		return false
	}

	rbody, err := ioutil.ReadAll(r.Body)
	if err != nil {
		m.t.Fatalf("failed to read body: %v", err)
	}
	r.Body = bytes.NewReader(rbody) // in case it is reused

	if !bytes.Equal(m.body, rbody) {
		m.t.Logf("Body mismatch: %v != %v", m.body, rbody)
		return false
	}

	return true
}

func (m RequestMatcher) String() string {
	return fmt.Sprintf("matches request %v with body %v", m.req, m.body)
}

// checkSuperSet checks if the items in l are all also present in r.
func checkSuperSet(l, r transport.Headers) error {
	missing := make([]string, 0, len(l))
	for k, vl := range l {
		vr, ok := r.Get(k)
		if !ok || vr != vl {
			missing = append(missing, k)
		}
	}

	if len(missing) > 0 {
		return fmt.Errorf("missing headers: %v", strings.Join(missing, ", "))
	}
	return nil
}

// ResponseMatcher is similar to RequestMatcher but for responses.
type ResponseMatcher struct {
	t    *testing.T
	res  *transport.Response
	body []byte
}

// NewResponseMatcher builds a new ResponseMatcher that verifies that
// responses match the given Response.
func NewResponseMatcher(t *testing.T, r *transport.Response) ResponseMatcher {
	body, err := ioutil.ReadAll(r.Body)
	defer r.Body.Close()
	if err != nil {
		t.Fatalf("failed to read response body: %v", err)
	}

	// restore a copy of the body so that the caller can still use the
	// response object
	r.Body = ioutil.NopCloser(bytes.NewReader(body))
	return ResponseMatcher{t: t, res: r, body: body}
}

// Matches checks if the given object matches the Response provided in
// NewResponseMatcher.
func (m ResponseMatcher) Matches(got interface{}) bool {
	l := m.res
	r, ok := got.(*transport.Response)
	if !ok {
		panic(fmt.Sprintf("expected *transport.Response, got %v", got))
	}

	if err := checkSuperSet(l.Headers, r.Headers); err != nil {
		m.t.Logf("Headers mismatch: %v != %v\n\t%v", l.Headers, r.Headers, err)
		return false
	}

	rbody, err := ioutil.ReadAll(r.Body)
	if err != nil {
		m.t.Fatalf("failed to read body: %v", err)
	}
	r.Body = ioutil.NopCloser(bytes.NewReader(rbody)) // in case it is reused

	if !bytes.Equal(m.body, rbody) {
		m.t.Logf("Body mismatch: %v != %v", m.body, rbody)
		return false
	}

	return true
}

// FakeResponseWriter is a ResponseWriter that records the headers and the body
// written to it.
type FakeResponseWriter struct {
	Headers transport.Headers
	Body    bytes.Buffer
}

// AddHeaders for FakeResponseWriter.
func (fw *FakeResponseWriter) AddHeaders(h transport.Headers) {
	if fw.Headers == nil {
		fw.Headers = make(transport.Headers)
	}
	for k, v := range h {
		fw.Headers.Set(k, v)
	}
}

// Write for FakeResponseWriter.
func (fw *FakeResponseWriter) Write(s []byte) (int, error) {
	return fw.Body.Write(s)
}
<commit_msg>transporttest: Add TTLDelta to RequestMatcher<commit_after>// Copyright (c) 2016 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package transporttest

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"strings"
	"testing"
	"time"

	"github.com/yarpc/yarpc-go/transport"
)

// DefaultTTLDelta is the default TTLDelta used by RequestMatchers.
const DefaultTTLDelta = 5 * time.Millisecond

// RequestMatcher may be used in gomock argument lists to assert that two
// requests match.
//
// Requests are considered to be matching if:
//
// - All their primitive parameters match.
// - The headers of the received request include all the headers from the
//   source request. It may include extra headers.
// - The contents of the request bodies are the same.
type RequestMatcher struct {
	t    *testing.T
	req  *transport.Request
	body []byte

	// Maximum amount of variation allowed when comparing TTLs. Defaults to 5
	// milliseconds.
	TTLDelta time.Duration
}

// NewRequestMatcher constructs a new RequestMatcher from the given testing.T
// and request.
//
// The request's contents are read in their entirety and replaced with a
// bytes.Reader.
func NewRequestMatcher(t *testing.T, r *transport.Request) RequestMatcher {
	body, err := ioutil.ReadAll(r.Body)
	if err != nil {
		t.Fatalf("failed to read request body: %v", err)
	}

	// restore a copy of the body so that the caller can still use the request
	// object
	r.Body = bytes.NewReader(body)
	return RequestMatcher{t: t, req: r, body: body, TTLDelta: DefaultTTLDelta}
}

// TODO: Headers like User-Agent, Content-Length, etc. make their way to the
// user-level Request object. For now, we're doing the super set check but we
// should do something more specific once yarpc/yarpc.io#2 is resolved.

// Matches checks if the given object matches the Request provided in
// NewRequestMatcher.
func (m RequestMatcher) Matches(got interface{}) bool {
	l := m.req
	r, ok := got.(*transport.Request)
	if !ok {
		panic(fmt.Sprintf("expected *transport.Request, got %v", got))
	}

	if l.Caller != r.Caller {
		m.t.Logf("Caller mismatch: %s != %s", l.Caller, r.Caller)
		return false
	}

	if l.Service != r.Service {
		m.t.Logf("Service mismatch: %s != %s", l.Service, r.Service)
		return false
	}

	if l.Encoding != r.Encoding {
		m.t.Logf("Encoding mismatch: %s != %s", l.Service, r.Service)
		return false
	}

	if l.Procedure != r.Procedure {
		m.t.Logf("Procedure mismatch: %s != %s", l.Procedure, r.Procedure)
		return false
	}

	ttlDiff := l.TTL - r.TTL
	if ttlDiff < 0 {
		ttlDiff = -ttlDiff
	}
	if ttlDiff > m.TTLDelta {
		m.t.Logf("TTL mismatch: %v != %v", l.TTL, r.TTL)
		return false
	}

	if err := checkSuperSet(l.Headers, r.Headers); err != nil {
		m.t.Logf("Headers mismatch: %v != %v\n\t%v", l.Headers, r.Headers, err)
		return false
	}

	rbody, err := ioutil.ReadAll(r.Body)
	if err != nil {
		m.t.Fatalf("failed to read body: %v", err)
	}
	r.Body = bytes.NewReader(rbody) // in case it is reused

	if !bytes.Equal(m.body, rbody) {
		m.t.Logf("Body mismatch: %v != %v", m.body, rbody)
		return false
	}

	return true
}

func (m RequestMatcher) String() string {
	return fmt.Sprintf("matches request %v with body %v", m.req, m.body)
}

// checkSuperSet checks if the items in l are all also present in r.
func checkSuperSet(l, r transport.Headers) error {
	missing := make([]string, 0, len(l))
	for k, vl := range l {
		vr, ok := r.Get(k)
		if !ok || vr != vl {
			missing = append(missing, k)
		}
	}

	if len(missing) > 0 {
		return fmt.Errorf("missing headers: %v", strings.Join(missing, ", "))
	}
	return nil
}

// ResponseMatcher is similar to RequestMatcher but for responses.
type ResponseMatcher struct {
	t    *testing.T
	res  *transport.Response
	body []byte
}

// NewResponseMatcher builds a new ResponseMatcher that verifies that
// responses match the given Response.
func NewResponseMatcher(t *testing.T, r *transport.Response) ResponseMatcher {
	body, err := ioutil.ReadAll(r.Body)
	defer r.Body.Close()
	if err != nil {
		t.Fatalf("failed to read response body: %v", err)
	}

	// restore a copy of the body so that the caller can still use the
	// response object
	r.Body = ioutil.NopCloser(bytes.NewReader(body))
	return ResponseMatcher{t: t, res: r, body: body}
}

// Matches checks if the given object matches the Response provided in
// NewResponseMatcher.
func (m ResponseMatcher) Matches(got interface{}) bool {
	l := m.res
	r, ok := got.(*transport.Response)
	if !ok {
		panic(fmt.Sprintf("expected *transport.Response, got %v", got))
	}

	if err := checkSuperSet(l.Headers, r.Headers); err != nil {
		m.t.Logf("Headers mismatch: %v != %v\n\t%v", l.Headers, r.Headers, err)
		return false
	}

	rbody, err := ioutil.ReadAll(r.Body)
	if err != nil {
		m.t.Fatalf("failed to read body: %v", err)
	}
	r.Body = ioutil.NopCloser(bytes.NewReader(rbody)) // in case it is reused

	if !bytes.Equal(m.body, rbody) {
		m.t.Logf("Body mismatch: %v != %v", m.body, rbody)
		return false
	}

	return true
}

// FakeResponseWriter is a ResponseWriter that records the headers and the body
// written to it.
type FakeResponseWriter struct {
	Headers transport.Headers
	Body    bytes.Buffer
}

// AddHeaders for FakeResponseWriter.
func (fw *FakeResponseWriter) AddHeaders(h transport.Headers) {
	if fw.Headers == nil {
		fw.Headers = make(transport.Headers)
	}
	for k, v := range h {
		fw.Headers.Set(k, v)
	}
}

// Write for FakeResponseWriter.
func (fw *FakeResponseWriter) Write(s []byte) (int, error) {
	return fw.Body.Write(s)
}
<|endoftext|>


<commit_before>package daemon // import "github.com/docker/docker/daemon"

import (
	"fmt"
	"regexp"
	"sync/atomic"
	"time"

	"github.com/docker/distribution/reference"
	"github.com/docker/docker/api/types"
	"github.com/docker/docker/api/types/filters"
	timetypes "github.com/docker/docker/api/types/time"
	"github.com/docker/docker/image"
	"github.com/docker/docker/layer"
	"github.com/docker/docker/pkg/directory"
	"github.com/docker/docker/runconfig"
	"github.com/docker/docker/volume"
	"github.com/docker/libnetwork"
	digest "github.com/opencontainers/go-digest"
	"github.com/sirupsen/logrus"
	"golang.org/x/net/context"
)

var (
	// errPruneRunning is returned when a prune request is received while
	// one is in progress
	errPruneRunning = fmt.Errorf("a prune operation is already running")

	containersAcceptedFilters = map[string]bool{
		"label":  true,
		"label!": true,
		"until":  true,
	}
	volumesAcceptedFilters = map[string]bool{
		"label":  true,
		"label!": true,
	}
	imagesAcceptedFilters = map[string]bool{
		"dangling": true,
		"label":    true,
		"label!":   true,
		"until":    true,
	}
	networksAcceptedFilters = map[string]bool{
		"label":  true,
		"label!": true,
		"until":  true,
	}
)

// ContainersPrune removes unused containers
func (daemon *Daemon) ContainersPrune(ctx context.Context, pruneFilters filters.Args) (*types.ContainersPruneReport, error) {
	if !atomic.CompareAndSwapInt32(&daemon.pruneRunning, 0, 1) {
		return nil, errPruneRunning
	}
	defer atomic.StoreInt32(&daemon.pruneRunning, 0)

	rep := &types.ContainersPruneReport{}

	// make sure that only accepted filters have been received
	err := pruneFilters.Validate(containersAcceptedFilters)
	if err != nil {
		return nil, err
	}

	until, err := getUntilFromPruneFilters(pruneFilters)
	if err != nil {
		return nil, err
	}

	allContainers := daemon.List()
	for _, c := range allContainers {
		select {
		case <-ctx.Done():
			logrus.Debugf("ContainersPrune operation cancelled: %#v", *rep)
			return rep, nil
		default:
		}

		if !c.IsRunning() {
			if !until.IsZero() && c.Created.After(until) {
				continue
			}
			if !matchLabels(pruneFilters, c.Config.Labels) {
				continue
			}
			cSize, _ := daemon.getSize(c.ID)
			// TODO: sets RmLink to true?
			err := daemon.ContainerRm(c.ID, &types.ContainerRmConfig{})
			if err != nil {
				logrus.Warnf("failed to prune container %s: %v", c.ID, err)
				continue
			}
			if cSize > 0 {
				rep.SpaceReclaimed += uint64(cSize)
			}
			rep.ContainersDeleted = append(rep.ContainersDeleted, c.ID)
		}
	}

	return rep, nil
}

// VolumesPrune removes unused local volumes
func (daemon *Daemon) VolumesPrune(ctx context.Context, pruneFilters filters.Args) (*types.VolumesPruneReport, error) {
	if !atomic.CompareAndSwapInt32(&daemon.pruneRunning, 0, 1) {
		return nil, errPruneRunning
	}
	defer atomic.StoreInt32(&daemon.pruneRunning, 0)

	// make sure that only accepted filters have been received
	err := pruneFilters.Validate(volumesAcceptedFilters)
	if err != nil {
		return nil, err
	}

	rep := &types.VolumesPruneReport{}

	pruneVols := func(v volume.Volume) error {
		select {
		case <-ctx.Done():
			logrus.Debugf("VolumesPrune operation cancelled: %#v", *rep)
			return ctx.Err()
		default:
		}

		name := v.Name()
		refs := daemon.volumes.Refs(v)

		if len(refs) == 0 {
			detailedVolume, ok := v.(volume.DetailedVolume)
			if ok {
				if !matchLabels(pruneFilters, detailedVolume.Labels()) {
					return nil
				}
			}
			vSize, err := directory.Size(v.Path())
			if err != nil {
				logrus.Warnf("could not determine size of volume %s: %v", name, err)
			}
			err = daemon.volumeRm(v)
			if err != nil {
				logrus.Warnf("could not remove volume %s: %v", name, err)
				return nil
			}
			rep.SpaceReclaimed += uint64(vSize)
			rep.VolumesDeleted = append(rep.VolumesDeleted, name)
		}

		return nil
	}

	err = daemon.traverseLocalVolumes(pruneVols)
	if err == context.Canceled {
		return rep, nil
	}

	return rep, err
}

// ImagesPrune removes unused images
func (daemon *Daemon) ImagesPrune(ctx context.Context, pruneFilters filters.Args) (*types.ImagesPruneReport, error) {
	if !atomic.CompareAndSwapInt32(&daemon.pruneRunning, 0, 1) {
		return nil, errPruneRunning
	}
	defer atomic.StoreInt32(&daemon.pruneRunning, 0)

	// make sure that only accepted filters have been received
	err := pruneFilters.Validate(imagesAcceptedFilters)
	if err != nil {
		return nil, err
	}

	rep := &types.ImagesPruneReport{}

	danglingOnly := true
	if pruneFilters.Contains("dangling") {
		if pruneFilters.ExactMatch("dangling", "false") || pruneFilters.ExactMatch("dangling", "0") {
			danglingOnly = false
		} else if !pruneFilters.ExactMatch("dangling", "true") && !pruneFilters.ExactMatch("dangling", "1") {
			return nil, invalidFilter{"dangling", pruneFilters.Get("dangling")}
		}
	}

	until, err := getUntilFromPruneFilters(pruneFilters)
	if err != nil {
		return nil, err
	}

	var allImages map[image.ID]*image.Image
	if danglingOnly {
		allImages = daemon.imageStore.Heads()
	} else {
		allImages = daemon.imageStore.Map()
	}
	allContainers := daemon.List()
	imageRefs := map[string]bool{}
	for _, c := range allContainers {
		select {
		case <-ctx.Done():
			return nil, ctx.Err()
		default:
			imageRefs[c.ID] = true
		}
	}

	// Filter intermediary images and get their unique size
	allLayers := make(map[layer.ChainID]layer.Layer)
	for _, ls := range daemon.layerStores {
		for k, v := range ls.Map() {
			allLayers[k] = v
		}
	}
	topImages := map[image.ID]*image.Image{}
	for id, img := range allImages {
		select {
		case <-ctx.Done():
			return nil, ctx.Err()
		default:
			dgst := digest.Digest(id)
			if len(daemon.referenceStore.References(dgst)) == 0 && len(daemon.imageStore.Children(id)) != 0 {
				continue
			}
			if !until.IsZero() && img.Created.After(until) {
				continue
			}
			if img.Config != nil && !matchLabels(pruneFilters, img.Config.Labels) {
				continue
			}
			topImages[id] = img
		}
	}

	canceled := false
deleteImagesLoop:
	for id := range topImages {
		select {
		case <-ctx.Done():
			// we still want to calculate freed size and return the data
			canceled = true
			break deleteImagesLoop
		default:
		}

		dgst := digest.Digest(id)
		hex := dgst.Hex()
		if _, ok := imageRefs[hex]; ok {
			continue
		}

		deletedImages := []types.ImageDeleteResponseItem{}
		refs := daemon.referenceStore.References(dgst)
		if len(refs) > 0 {
			shouldDelete := !danglingOnly
			if !shouldDelete {
				hasTag := false
				for _, ref := range refs {
					if _, ok := ref.(reference.NamedTagged); ok {
						hasTag = true
						break
					}
				}

				// Only delete if it's untagged (i.e. repo:<none>)
				shouldDelete = !hasTag
			}

			if shouldDelete {
				for _, ref := range refs {
					imgDel, err := daemon.ImageDelete(ref.String(), false, true)
					if err != nil {
						logrus.Warnf("could not delete reference %s: %v", ref.String(), err)
						continue
					}
					deletedImages = append(deletedImages, imgDel...)
				}
			}
		} else {
			imgDel, err := daemon.ImageDelete(hex, false, true)
			if err != nil {
				logrus.Warnf("could not delete image %s: %v", hex, err)
				continue
			}
			deletedImages = append(deletedImages, imgDel...)
		}

		rep.ImagesDeleted = append(rep.ImagesDeleted, deletedImages...)
	}

	// Compute how much space was freed
	for _, d := range rep.ImagesDeleted {
		if d.Deleted != "" {
			chid := layer.ChainID(d.Deleted)
			if l, ok := allLayers[chid]; ok {
				diffSize, err := l.DiffSize()
				if err != nil {
					logrus.Warnf("failed to get layer %s size: %v", chid, err)
					continue
				}
				rep.SpaceReclaimed += uint64(diffSize)
			}
		}
	}

	if canceled {
		logrus.Debugf("ImagesPrune operation cancelled: %#v", *rep)
	}

	return rep, nil
}

// localNetworksPrune removes unused local networks
func (daemon *Daemon) localNetworksPrune(ctx context.Context, pruneFilters filters.Args) *types.NetworksPruneReport {
	rep := &types.NetworksPruneReport{}

	until, _ := getUntilFromPruneFilters(pruneFilters)

	// When the function returns true, the walk will stop.
	l := func(nw libnetwork.Network) bool {
		select {
		case <-ctx.Done():
			// context cancelled
			return true
		default:
		}
		if nw.Info().ConfigOnly() {
			return false
		}
		if !until.IsZero() && nw.Info().Created().After(until) {
			return false
		}
		if !matchLabels(pruneFilters, nw.Info().Labels()) {
			return false
		}
		nwName := nw.Name()
		if runconfig.IsPreDefinedNetwork(nwName) {
			return false
		}
		if len(nw.Endpoints()) > 0 {
			return false
		}
		if err := daemon.DeleteNetwork(nw.ID()); err != nil {
			logrus.Warnf("could not remove local network %s: %v", nwName, err)
			return false
		}
		rep.NetworksDeleted = append(rep.NetworksDeleted, nwName)
		return false
	}
	daemon.netController.WalkNetworks(l)
	return rep
}

// clusterNetworksPrune removes unused cluster networks
func (daemon *Daemon) clusterNetworksPrune(ctx context.Context, pruneFilters filters.Args) (*types.NetworksPruneReport, error) {
	rep := &types.NetworksPruneReport{}

	until, _ := getUntilFromPruneFilters(pruneFilters)

	cluster := daemon.GetCluster()

	if !cluster.IsManager() {
		return rep, nil
	}

	networks, err := cluster.GetNetworks()
	if err != nil {
		return rep, err
	}
	networkIsInUse := regexp.MustCompile(`network ([[:alnum:]]+) is in use`)
	for _, nw := range networks {
		select {
		case <-ctx.Done():
			return rep, nil
		default:
			if nw.Ingress {
				// Routing-mesh network removal has to be explicitly invoked by user
				continue
			}
			if !until.IsZero() && nw.Created.After(until) {
				continue
			}
			if !matchLabels(pruneFilters, nw.Labels) {
				continue
			}
			// https://github.com/docker/docker/issues/24186
			// `docker network inspect` unfortunately displays ONLY those containers that are local to that node.
			// So we try to remove it anyway and check the error
			err = cluster.RemoveNetwork(nw.ID)
			if err != nil {
				// we can safely ignore the "network .. is in use" error
				match := networkIsInUse.FindStringSubmatch(err.Error())
				if len(match) != 2 || match[1] != nw.ID {
					logrus.Warnf("could not remove cluster network %s: %v", nw.Name, err)
				}
				continue
			}
			rep.NetworksDeleted = append(rep.NetworksDeleted, nw.Name)
		}
	}
	return rep, nil
}

// NetworksPrune removes unused networks
func (daemon *Daemon) NetworksPrune(ctx context.Context, pruneFilters filters.Args) (*types.NetworksPruneReport, error) {
	if !atomic.CompareAndSwapInt32(&daemon.pruneRunning, 0, 1) {
		return nil, errPruneRunning
	}
	defer atomic.StoreInt32(&daemon.pruneRunning, 0)

	// make sure that only accepted filters have been received
	err := pruneFilters.Validate(networksAcceptedFilters)
	if err != nil {
		return nil, err
	}

	if _, err := getUntilFromPruneFilters(pruneFilters); err != nil {
		return nil, err
	}

	rep := &types.NetworksPruneReport{}
	if clusterRep, err := daemon.clusterNetworksPrune(ctx, pruneFilters); err == nil {
		rep.NetworksDeleted = append(rep.NetworksDeleted, clusterRep.NetworksDeleted...)
	}

	localRep := daemon.localNetworksPrune(ctx, pruneFilters)
	rep.NetworksDeleted = append(rep.NetworksDeleted, localRep.NetworksDeleted...)

	select {
	case <-ctx.Done():
		logrus.Debugf("NetworksPrune operation cancelled: %#v", *rep)
		return rep, nil
	default:
	}

	return rep, nil
}

func getUntilFromPruneFilters(pruneFilters filters.Args) (time.Time, error) {
	until := time.Time{}
	if !pruneFilters.Contains("until") {
		return until, nil
	}
	untilFilters := pruneFilters.Get("until")
	if len(untilFilters) > 1 {
		return until, fmt.Errorf("more than one until filter specified")
	}
	ts, err := timetypes.GetTimestamp(untilFilters[0], time.Now())
	if err != nil {
		return until, err
	}
	seconds, nanoseconds, err := timetypes.ParseTimestamps(ts, 0)
	if err != nil {
		return until, err
	}
	until = time.Unix(seconds, nanoseconds)
	return until, nil
}

func matchLabels(pruneFilters filters.Args, labels map[string]string) bool {
	if !pruneFilters.MatchKVList("label", labels) {
		return false
	}
	// By default MatchKVList will return true if field (like 'label!') does not exist
	// So we have to add additional Contains("label!") check
	if pruneFilters.Contains("label!") {
		if pruneFilters.MatchKVList("label!", labels) {
			return false
		}
	}
	return true
}
<commit_msg>Remove broken container check from image prune<commit_after>package daemon // import "github.com/docker/docker/daemon"

import (
	"fmt"
	"regexp"
	"sync/atomic"
	"time"

	"github.com/docker/distribution/reference"
	"github.com/docker/docker/api/types"
	"github.com/docker/docker/api/types/filters"
	timetypes "github.com/docker/docker/api/types/time"
	"github.com/docker/docker/errdefs"
	"github.com/docker/docker/image"
	"github.com/docker/docker/layer"
	"github.com/docker/docker/pkg/directory"
	"github.com/docker/docker/runconfig"
	"github.com/docker/docker/volume"
	"github.com/docker/libnetwork"
	digest "github.com/opencontainers/go-digest"
	"github.com/sirupsen/logrus"
	"golang.org/x/net/context"
)

var (
	// errPruneRunning is returned when a prune request is received while
	// one is in progress
	errPruneRunning = fmt.Errorf("a prune operation is already running")

	containersAcceptedFilters = map[string]bool{
		"label":  true,
		"label!": true,
		"until":  true,
	}
	volumesAcceptedFilters = map[string]bool{
		"label":  true,
		"label!": true,
	}
	imagesAcceptedFilters = map[string]bool{
		"dangling": true,
		"label":    true,
		"label!":   true,
		"until":    true,
	}
	networksAcceptedFilters = map[string]bool{
		"label":  true,
		"label!": true,
		"until":  true,
	}
)

// ContainersPrune removes unused containers
func (daemon *Daemon) ContainersPrune(ctx context.Context, pruneFilters filters.Args) (*types.ContainersPruneReport, error) {
	if !atomic.CompareAndSwapInt32(&daemon.pruneRunning, 0, 1) {
		return nil, errPruneRunning
	}
	defer atomic.StoreInt32(&daemon.pruneRunning, 0)

	rep := &types.ContainersPruneReport{}

	// make sure that only accepted filters have been received
	err := pruneFilters.Validate(containersAcceptedFilters)
	if err != nil {
		return nil, err
	}

	until, err := getUntilFromPruneFilters(pruneFilters)
	if err != nil {
		return nil, err
	}

	allContainers := daemon.List()
	for _, c := range allContainers {
		select {
		case <-ctx.Done():
			logrus.Debugf("ContainersPrune operation cancelled: %#v", *rep)
			return rep, nil
		default:
		}

		if !c.IsRunning() {
			if !until.IsZero() && c.Created.After(until) {
				continue
			}
			if !matchLabels(pruneFilters, c.Config.Labels) {
				continue
			}
			cSize, _ := daemon.getSize(c.ID)
			// TODO: sets RmLink to true?
			err := daemon.ContainerRm(c.ID, &types.ContainerRmConfig{})
			if err != nil {
				logrus.Warnf("failed to prune container %s: %v", c.ID, err)
				continue
			}
			if cSize > 0 {
				rep.SpaceReclaimed += uint64(cSize)
			}
			rep.ContainersDeleted = append(rep.ContainersDeleted, c.ID)
		}
	}

	return rep, nil
}

// VolumesPrune removes unused local volumes
func (daemon *Daemon) VolumesPrune(ctx context.Context, pruneFilters filters.Args) (*types.VolumesPruneReport, error) {
	if !atomic.CompareAndSwapInt32(&daemon.pruneRunning, 0, 1) {
		return nil, errPruneRunning
	}
	defer atomic.StoreInt32(&daemon.pruneRunning, 0)

	// make sure that only accepted filters have been received
	err := pruneFilters.Validate(volumesAcceptedFilters)
	if err != nil {
		return nil, err
	}

	rep := &types.VolumesPruneReport{}

	pruneVols := func(v volume.Volume) error {
		select {
		case <-ctx.Done():
			logrus.Debugf("VolumesPrune operation cancelled: %#v", *rep)
			return ctx.Err()
		default:
		}

		name := v.Name()
		refs := daemon.volumes.Refs(v)

		if len(refs) == 0 {
			detailedVolume, ok := v.(volume.DetailedVolume)
			if ok {
				if !matchLabels(pruneFilters, detailedVolume.Labels()) {
					return nil
				}
			}
			vSize, err := directory.Size(v.Path())
			if err != nil {
				logrus.Warnf("could not determine size of volume %s: %v", name, err)
			}
			err = daemon.volumeRm(v)
			if err != nil {
				logrus.Warnf("could not remove volume %s: %v", name, err)
				return nil
			}
			rep.SpaceReclaimed += uint64(vSize)
			rep.VolumesDeleted = append(rep.VolumesDeleted, name)
		}

		return nil
	}

	err = daemon.traverseLocalVolumes(pruneVols)
	if err == context.Canceled {
		return rep, nil
	}

	return rep, err
}

// ImagesPrune removes unused images
func (daemon *Daemon) ImagesPrune(ctx context.Context, pruneFilters filters.Args) (*types.ImagesPruneReport, error) {
	if !atomic.CompareAndSwapInt32(&daemon.pruneRunning, 0, 1) {
		return nil, errPruneRunning
	}
	defer atomic.StoreInt32(&daemon.pruneRunning, 0)

	// make sure that only accepted filters have been received
	err := pruneFilters.Validate(imagesAcceptedFilters)
	if err != nil {
		return nil, err
	}

	rep := &types.ImagesPruneReport{}

	danglingOnly := true
	if pruneFilters.Contains("dangling") {
		if pruneFilters.ExactMatch("dangling", "false") || pruneFilters.ExactMatch("dangling", "0") {
			danglingOnly = false
		} else if !pruneFilters.ExactMatch("dangling", "true") && !pruneFilters.ExactMatch("dangling", "1") {
			return nil, invalidFilter{"dangling", pruneFilters.Get("dangling")}
		}
	}

	until, err := getUntilFromPruneFilters(pruneFilters)
	if err != nil {
		return nil, err
	}

	var allImages map[image.ID]*image.Image
	if danglingOnly {
		allImages = daemon.imageStore.Heads()
	} else {
		allImages = daemon.imageStore.Map()
	}

	// Filter intermediary images and get their unique size
	allLayers := make(map[layer.ChainID]layer.Layer)
	for _, ls := range daemon.layerStores {
		for k, v := range ls.Map() {
			allLayers[k] = v
		}
	}
	topImages := map[image.ID]*image.Image{}
	for id, img := range allImages {
		select {
		case <-ctx.Done():
			return nil, ctx.Err()
		default:
			dgst := digest.Digest(id)
			if len(daemon.referenceStore.References(dgst)) == 0 && len(daemon.imageStore.Children(id)) != 0 {
				continue
			}
			if !until.IsZero() && img.Created.After(until) {
				continue
			}
			if img.Config != nil && !matchLabels(pruneFilters, img.Config.Labels) {
				continue
			}
			topImages[id] = img
		}
	}

	canceled := false
deleteImagesLoop:
	for id := range topImages {
		select {
		case <-ctx.Done():
			// we still want to calculate freed size and return the data
			canceled = true
			break deleteImagesLoop
		default:
		}

		deletedImages := []types.ImageDeleteResponseItem{}
		refs := daemon.referenceStore.References(id.Digest())
		if len(refs) > 0 {
			shouldDelete := !danglingOnly
			if !shouldDelete {
				hasTag := false
				for _, ref := range refs {
					if _, ok := ref.(reference.NamedTagged); ok {
						hasTag = true
						break
					}
				}

				// Only delete if it's untagged (i.e. repo:<none>)
				shouldDelete = !hasTag
			}

			if shouldDelete {
				for _, ref := range refs {
					imgDel, err := daemon.ImageDelete(ref.String(), false, true)
					if imageDeleteFailed(ref.String(), err) {
						continue
					}
					deletedImages = append(deletedImages, imgDel...)
				}
			}
		} else {
			hex := id.Digest().Hex()
			imgDel, err := daemon.ImageDelete(hex, false, true)
			if imageDeleteFailed(hex, err) {
				continue
			}
			deletedImages = append(deletedImages, imgDel...)
		}

		rep.ImagesDeleted = append(rep.ImagesDeleted, deletedImages...)
	}

	// Compute how much space was freed
	for _, d := range rep.ImagesDeleted {
		if d.Deleted != "" {
			chid := layer.ChainID(d.Deleted)
			if l, ok := allLayers[chid]; ok {
				diffSize, err := l.DiffSize()
				if err != nil {
					logrus.Warnf("failed to get layer %s size: %v", chid, err)
					continue
				}
				rep.SpaceReclaimed += uint64(diffSize)
			}
		}
	}

	if canceled {
		logrus.Debugf("ImagesPrune operation cancelled: %#v", *rep)
	}

	return rep, nil
}

func imageDeleteFailed(ref string, err error) bool {
	switch {
	case err == nil:
		return false
	case errdefs.IsConflict(err):
		return true
	default:
		logrus.Warnf("failed to prune image %s: %v", ref, err)
		return true
	}
}

// localNetworksPrune removes unused local networks
func (daemon *Daemon) localNetworksPrune(ctx context.Context, pruneFilters filters.Args) *types.NetworksPruneReport {
	rep := &types.NetworksPruneReport{}

	until, _ := getUntilFromPruneFilters(pruneFilters)

	// When the function returns true, the walk will stop.
	l := func(nw libnetwork.Network) bool {
		select {
		case <-ctx.Done():
			// context cancelled
			return true
		default:
		}
		if nw.Info().ConfigOnly() {
			return false
		}
		if !until.IsZero() && nw.Info().Created().After(until) {
			return false
		}
		if !matchLabels(pruneFilters, nw.Info().Labels()) {
			return false
		}
		nwName := nw.Name()
		if runconfig.IsPreDefinedNetwork(nwName) {
			return false
		}
		if len(nw.Endpoints()) > 0 {
			return false
		}
		if err := daemon.DeleteNetwork(nw.ID()); err != nil {
			logrus.Warnf("could not remove local network %s: %v", nwName, err)
			return false
		}
		rep.NetworksDeleted = append(rep.NetworksDeleted, nwName)
		return false
	}
	daemon.netController.WalkNetworks(l)
	return rep
}

// clusterNetworksPrune removes unused cluster networks
func (daemon *Daemon) clusterNetworksPrune(ctx context.Context, pruneFilters filters.Args) (*types.NetworksPruneReport, error) {
	rep := &types.NetworksPruneReport{}

	until, _ := getUntilFromPruneFilters(pruneFilters)

	cluster := daemon.GetCluster()

	if !cluster.IsManager() {
		return rep, nil
	}

	networks, err := cluster.GetNetworks()
	if err != nil {
		return rep, err
	}
	networkIsInUse := regexp.MustCompile(`network ([[:alnum:]]+) is in use`)
	for _, nw := range networks {
		select {
		case <-ctx.Done():
			return rep, nil
		default:
			if nw.Ingress {
				// Routing-mesh network removal has to be explicitly invoked by user
				continue
			}
			if !until.IsZero() && nw.Created.After(until) {
				continue
			}
			if !matchLabels(pruneFilters, nw.Labels) {
				continue
			}
			// https://github.com/docker/docker/issues/24186
			// `docker network inspect` unfortunately displays ONLY those containers that are local to that node.
			// So we try to remove it anyway and check the error
			err = cluster.RemoveNetwork(nw.ID)
			if err != nil {
				// we can safely ignore the "network .. is in use" error
				match := networkIsInUse.FindStringSubmatch(err.Error())
				if len(match) != 2 || match[1] != nw.ID {
					logrus.Warnf("could not remove cluster network %s: %v", nw.Name, err)
				}
				continue
			}
			rep.NetworksDeleted = append(rep.NetworksDeleted, nw.Name)
		}
	}
	return rep, nil
}

// NetworksPrune removes unused networks
func (daemon *Daemon) NetworksPrune(ctx context.Context, pruneFilters filters.Args) (*types.NetworksPruneReport, error) {
	if !atomic.CompareAndSwapInt32(&daemon.pruneRunning, 0, 1) {
		return nil, errPruneRunning
	}
	defer atomic.StoreInt32(&daemon.pruneRunning, 0)

	// make sure that only accepted filters have been received
	err := pruneFilters.Validate(networksAcceptedFilters)
	if err != nil {
		return nil, err
	}

	if _, err := getUntilFromPruneFilters(pruneFilters); err != nil {
		return nil, err
	}

	rep := &types.NetworksPruneReport{}
	if clusterRep, err := daemon.clusterNetworksPrune(ctx, pruneFilters); err == nil {
		rep.NetworksDeleted = append(rep.NetworksDeleted, clusterRep.NetworksDeleted...)
	}

	localRep := daemon.localNetworksPrune(ctx, pruneFilters)
	rep.NetworksDeleted = append(rep.NetworksDeleted, localRep.NetworksDeleted...)

	select {
	case <-ctx.Done():
		logrus.Debugf("NetworksPrune operation cancelled: %#v", *rep)
		return rep, nil
	default:
	}

	return rep, nil
}

func getUntilFromPruneFilters(pruneFilters filters.Args) (time.Time, error) {
	until := time.Time{}
	if !pruneFilters.Contains("until") {
		return until, nil
	}
	untilFilters := pruneFilters.Get("until")
	if len(untilFilters) > 1 {
		return until, fmt.Errorf("more than one until filter specified")
	}
	ts, err := timetypes.GetTimestamp(untilFilters[0], time.Now())
	if err != nil {
		return until, err
	}
	seconds, nanoseconds, err := timetypes.ParseTimestamps(ts, 0)
	if err != nil {
		return until, err
	}
	until = time.Unix(seconds, nanoseconds)
	return until, nil
}

func matchLabels(pruneFilters filters.Args, labels map[string]string) bool {
	if !pruneFilters.MatchKVList("label", labels) {
		return false
	}
	// By default MatchKVList will return true if field (like 'label!') does not exist
	// So we have to add additional Contains("label!") check
	if pruneFilters.Contains("label!") {
		if pruneFilters.MatchKVList("label!", labels) {
			return false
		}
	}
	return true
}
<|endoftext|>


<commit_before>// +build coprocess

package main

import (
	"github.com/mitchellh/mapstructure"

	"github.com/Sirupsen/logrus"

	"github.com/TykTechnologies/tyk/apidef"
	"github.com/TykTechnologies/tyk/coprocess"

	"bytes"
	"errors"
	"io/ioutil"
	"net/http"
)

var (
	// EnableCoProcess will be overridden by config.EnableCoProcess.
	EnableCoProcess = false

	// GlobalDispatcher will be implemented by the current CoProcess driver.
	GlobalDispatcher coprocess.Dispatcher
)

// CoProcessMiddleware is the basic CP middleware struct.
type CoProcessMiddleware struct {
	*TykMiddleware
	HookType         coprocess.HookType
	HookName         string
	MiddlewareDriver apidef.MiddlewareDriver
}

func (mw *CoProcessMiddleware) GetName() string {
	return "CoProcessMiddleware"
}

// CreateCoProcessMiddleware initializes a new CP middleware, takes hook type (pre, post, etc.), hook name ("my_hook") and driver ("python").
func CreateCoProcessMiddleware(hookName string, hookType coprocess.HookType, mwDriver apidef.MiddlewareDriver, tykMwSuper *TykMiddleware) func(http.Handler) http.Handler {
	dMiddleware := &CoProcessMiddleware{
		TykMiddleware:    tykMwSuper,
		HookType:         hookType,
		HookName:         hookName,
		MiddlewareDriver: mwDriver,
	}

	return CreateMiddleware(dMiddleware, tykMwSuper)
}

func doCoprocessReload() {
	if GlobalDispatcher != nil {
		log.WithFields(logrus.Fields{
			"prefix": "coprocess",
		}).Info("Reloading middlewares")
		GlobalDispatcher.Reload()
	}

}

// CoProcessor represents a CoProcess during the request.
type CoProcessor struct {
	HookType   coprocess.HookType
	Middleware *CoProcessMiddleware
}

// GetObjectFromRequest constructs a CoProcessObject from a given http.Request.
func (c *CoProcessor) GetObjectFromRequest(r *http.Request) *coprocess.Object {

	defer r.Body.Close()
	originalBody, _ := ioutil.ReadAll(r.Body)

	miniRequestObject := &coprocess.MiniRequestObject{
		Headers:        ProtoMap(r.Header),
		SetHeaders:     make(map[string]string),
		DeleteHeaders:  make([]string, 0),
		Body:           string(originalBody),
		Url:            r.URL.Path,
		Params:         ProtoMap(r.URL.Query()),
		AddParams:      make(map[string]string),
		ExtendedParams: ProtoMap(nil),
		DeleteParams:   make([]string, 0),
		ReturnOverrides: &coprocess.ReturnOverrides{
			ResponseCode:  -1,
			ResponseError: "",
		},
	}

	object := &coprocess.Object{
		Request:  miniRequestObject,
		HookName: c.Middleware.HookName,
	}

	// If a middleware is set, take its HookType, otherwise override it with CoProcessor.HookType
	if c.Middleware != nil && c.HookType == 0 {
		c.HookType = c.Middleware.HookType
	}

	object.HookType = c.HookType

	object.Metadata = make(map[string]string)
	object.Spec = make(map[string]string)

	// Append spec data:
	if c.Middleware != nil {
		object.Spec = map[string]string{
			"OrgID": c.Middleware.Spec.OrgID,
			"APIID": c.Middleware.Spec.APIID,
		}
	}

	// Encode the session object (if not a pre-process & not a custom key check):
	if c.HookType != coprocess.HookType_Pre && c.HookType != coprocess.HookType_CustomKeyCheck {
		session := ctxGetSession(r)
		if session != nil {
			object.Session = ProtoSessionState(session)
		}
	}

	return object
}

// ObjectPostProcess does CoProcessObject post-processing (adding/removing headers or params, etc.).
func (c *CoProcessor) ObjectPostProcess(object *coprocess.Object, r *http.Request) {
	r.ContentLength = int64(len(object.Request.Body))
	r.Body = ioutil.NopCloser(bytes.NewBufferString(object.Request.Body))

	for _, dh := range object.Request.DeleteHeaders {
		r.Header.Del(dh)
	}

	for h, v := range object.Request.SetHeaders {
		r.Header.Set(h, v)
	}

	values := r.URL.Query()
	for _, k := range object.Request.DeleteParams {
		values.Del(k)
	}

	for p, v := range object.Request.AddParams {
		values.Set(p, v)
	}

	r.URL.RawQuery = values.Encode()
}

// CoProcessInit creates a new CoProcessDispatcher, it will be called when Tyk starts.
func CoProcessInit() error {
	var err error
	if config.CoProcessOptions.EnableCoProcess {
		GlobalDispatcher, err = NewCoProcessDispatcher()
		EnableCoProcess = true
	}
	return err
}

// CoProcessMiddlewareConfig holds the middleware configuration.
type CoProcessMiddlewareConfig struct {
	ConfigData map[string]string `mapstructure:"config_data" bson:"config_data" json:"config_data"`
}

// New lets you do any initialisations for the object can be done here
func (m *CoProcessMiddleware) New() {}

// GetConfig retrieves the configuration from the API config - we user mapstructure for this for simplicity
func (m *CoProcessMiddleware) GetConfig() (interface{}, error) {
	var moduleConfig CoProcessMiddlewareConfig

	err := mapstructure.Decode(m.Spec.RawData, &moduleConfig)
	if err != nil {
		log.WithFields(logrus.Fields{
			"prefix": "jsvm",
		}).Error(err)
		return nil, err
	}

	return moduleConfig, nil
}

// IsEnabledForSpec checks if this middleware should be enabled for a given API.
func (m *CoProcessMiddleware) IsEnabledForSpec() bool {
	// This flag is true when Tyk has been compiled with CP support and when the configuration enables it.
	enableCoProcess := config.CoProcessOptions.EnableCoProcess && EnableCoProcess
	// This flag indicates if the current spec specifies any CP custom middleware.
	var usesCoProcessMiddleware bool

	supportedDrivers := []apidef.MiddlewareDriver{apidef.PythonDriver, apidef.LuaDriver, apidef.GrpcDriver}

	for _, driver := range supportedDrivers {
		if m.Spec.CustomMiddleware.Driver == driver && CoProcessName == driver {
			usesCoProcessMiddleware = true
			break
		}
	}

	if usesCoProcessMiddleware && enableCoProcess {
		log.WithFields(logrus.Fields{
			"prefix": "coprocess",
		}).Debug("Enabling CP middleware.")
		return true
	}

	if usesCoProcessMiddleware && !enableCoProcess {
		log.WithFields(logrus.Fields{
			"prefix": "coprocess",
		}).Error("Your API specifies a CP custom middleware, either Tyk wasn't build with CP support or CP is not enabled in your Tyk configuration file!")
	}

	if !usesCoProcessMiddleware && m.Spec.CustomMiddleware.Driver != "" {
		log.WithFields(logrus.Fields{
			"prefix": "coprocess",
		}).Error("CP Driver not supported: ", m.Spec.CustomMiddleware.Driver)
	}

	return false
}

// ProcessRequest will run any checks on the request on the way through the system, return an error to have the chain fail
func (m *CoProcessMiddleware) ProcessRequest(w http.ResponseWriter, r *http.Request, configuration interface{}) (error, int) {
	log.WithFields(logrus.Fields{
		"prefix": "coprocess",
	}).Debug("CoProcess Request, HookType: ", m.HookType)

	if !EnableCoProcess {
		return nil, 200
	}

	var extractor IdExtractor
	if m.Spec.EnableCoProcessAuth && m.Spec.CustomMiddleware.IdExtractor.Extractor != nil {
		extractor = m.Spec.CustomMiddleware.IdExtractor.Extractor.(IdExtractor)
	}

	var returnOverrides ReturnOverrides
	var sessionID string

	if m.HookType == coprocess.HookType_CustomKeyCheck && extractor != nil {
		sessionID, returnOverrides = extractor.ExtractAndCheck(r)

		if returnOverrides.ResponseCode != 0 {
			if returnOverrides.ResponseError == "" {
				return nil, returnOverrides.ResponseCode
			}
			err := errors.New(returnOverrides.ResponseError)
			return err, returnOverrides.ResponseCode
		}
	}

	// It's also possible to override the HookType:
	coProcessor := CoProcessor{
		Middleware: m,
		// HookType: coprocess.PreHook,
	}

	object := coProcessor.GetObjectFromRequest(r)

	returnObject, err := coProcessor.Dispatch(object)
	if err != nil {
		if m.HookType == coprocess.HookType_CustomKeyCheck {
			return errors.New("Key not authorised"), 403
		} else {
			return errors.New("Middleware error"), 500
		}
	}

	coProcessor.ObjectPostProcess(returnObject, r)

	token := returnObject.Metadata["token"]

	// The CP middleware indicates this is a bad auth:
	if returnObject.Request.ReturnOverrides.ResponseCode > 400 {

		log.WithFields(logrus.Fields{
			"path":   r.URL.Path,
			"origin": GetIPFromRequest(r),
			"key":    token,
		}).Info("Attempted access with invalid key.")

		// Fire Authfailed Event
		AuthFailed(m.TykMiddleware, r, token)

		// Report in health check
		ReportHealthCheckValue(m.Spec.Health, KeyFailure, "1")

		return errors.New("Key not authorised"), int(returnObject.Request.ReturnOverrides.ResponseCode)
	}

	if returnObject.Request.ReturnOverrides.ResponseCode > 0 {
		for h, v := range returnObject.Request.ReturnOverrides.Headers {
			w.Header().Set(h, v)
		}
		w.WriteHeader(int(returnObject.Request.ReturnOverrides.ResponseCode))
		w.Write([]byte(returnObject.Request.ReturnOverrides.ResponseError))
		return nil, 666
	}

	// Is this a CP authentication middleware?
	if m.Spec.EnableCoProcessAuth && m.HookType == coprocess.HookType_CustomKeyCheck {
		// The CP middleware didn't setup a session:
		if returnObject.Session == nil {
			return errors.New("Key not authorised"), 403
		}

		returnedSession := TykSessionState(returnObject.Session)

		if extractor == nil {
			sessionLifetime := getLifetime(m.Spec, returnedSession)
			// This API is not using the ID extractor, but we've got a session:
			m.Spec.SessionManager.UpdateSession(token, returnedSession, sessionLifetime)
			ctxSetSession(r, returnedSession)
			ctxSetAuthToken(r, token)
		} else {
			// The CP middleware did setup a session, we should pass it to the ID extractor (caching):
			extractor.PostProcess(r, returnedSession, sessionID)
		}
	}

	return nil, 200
}
<commit_msg>coprocess: handle nil request body<commit_after>// +build coprocess

package main

import (
	"github.com/mitchellh/mapstructure"

	"github.com/Sirupsen/logrus"

	"github.com/TykTechnologies/tyk/apidef"
	"github.com/TykTechnologies/tyk/coprocess"

	"bytes"
	"errors"
	"io/ioutil"
	"net/http"
)

var (
	// EnableCoProcess will be overridden by config.EnableCoProcess.
	EnableCoProcess = false

	// GlobalDispatcher will be implemented by the current CoProcess driver.
	GlobalDispatcher coprocess.Dispatcher
)

// CoProcessMiddleware is the basic CP middleware struct.
type CoProcessMiddleware struct {
	*TykMiddleware
	HookType         coprocess.HookType
	HookName         string
	MiddlewareDriver apidef.MiddlewareDriver
}

func (mw *CoProcessMiddleware) GetName() string {
	return "CoProcessMiddleware"
}

// CreateCoProcessMiddleware initializes a new CP middleware, takes hook type (pre, post, etc.), hook name ("my_hook") and driver ("python").
func CreateCoProcessMiddleware(hookName string, hookType coprocess.HookType, mwDriver apidef.MiddlewareDriver, tykMwSuper *TykMiddleware) func(http.Handler) http.Handler {
	dMiddleware := &CoProcessMiddleware{
		TykMiddleware:    tykMwSuper,
		HookType:         hookType,
		HookName:         hookName,
		MiddlewareDriver: mwDriver,
	}

	return CreateMiddleware(dMiddleware, tykMwSuper)
}

func doCoprocessReload() {
	if GlobalDispatcher != nil {
		log.WithFields(logrus.Fields{
			"prefix": "coprocess",
		}).Info("Reloading middlewares")
		GlobalDispatcher.Reload()
	}

}

// CoProcessor represents a CoProcess during the request.
type CoProcessor struct {
	HookType   coprocess.HookType
	Middleware *CoProcessMiddleware
}

// GetObjectFromRequest constructs a CoProcessObject from a given http.Request.
func (c *CoProcessor) GetObjectFromRequest(r *http.Request) *coprocess.Object {
	var body string
	if r.Body == nil {
		body = ""
	} else {
		defer r.Body.Close()
		originalBody, _ := ioutil.ReadAll(r.Body)
		body = string(originalBody)
	}

	miniRequestObject := &coprocess.MiniRequestObject{
		Headers:        ProtoMap(r.Header),
		SetHeaders:     make(map[string]string),
		DeleteHeaders:  make([]string, 0),
		Body:           body,
		Url:            r.URL.Path,
		Params:         ProtoMap(r.URL.Query()),
		AddParams:      make(map[string]string),
		ExtendedParams: ProtoMap(nil),
		DeleteParams:   make([]string, 0),
		ReturnOverrides: &coprocess.ReturnOverrides{
			ResponseCode:  -1,
			ResponseError: "",
		},
	}

	object := &coprocess.Object{
		Request:  miniRequestObject,
		HookName: c.Middleware.HookName,
	}

	// If a middleware is set, take its HookType, otherwise override it with CoProcessor.HookType
	if c.Middleware != nil && c.HookType == 0 {
		c.HookType = c.Middleware.HookType
	}

	object.HookType = c.HookType

	object.Metadata = make(map[string]string)
	object.Spec = make(map[string]string)

	// Append spec data:
	if c.Middleware != nil {
		object.Spec = map[string]string{
			"OrgID": c.Middleware.Spec.OrgID,
			"APIID": c.Middleware.Spec.APIID,
		}
	}

	// Encode the session object (if not a pre-process & not a custom key check):
	if c.HookType != coprocess.HookType_Pre && c.HookType != coprocess.HookType_CustomKeyCheck {
		session := ctxGetSession(r)
		if session != nil {
			object.Session = ProtoSessionState(session)
		}
	}

	return object
}

// ObjectPostProcess does CoProcessObject post-processing (adding/removing headers or params, etc.).
func (c *CoProcessor) ObjectPostProcess(object *coprocess.Object, r *http.Request) {
	r.ContentLength = int64(len(object.Request.Body))
	r.Body = ioutil.NopCloser(bytes.NewBufferString(object.Request.Body))

	for _, dh := range object.Request.DeleteHeaders {
		r.Header.Del(dh)
	}

	for h, v := range object.Request.SetHeaders {
		r.Header.Set(h, v)
	}

	values := r.URL.Query()
	for _, k := range object.Request.DeleteParams {
		values.Del(k)
	}

	for p, v := range object.Request.AddParams {
		values.Set(p, v)
	}

	r.URL.RawQuery = values.Encode()
}

// CoProcessInit creates a new CoProcessDispatcher, it will be called when Tyk starts.
func CoProcessInit() error {
	var err error
	if config.CoProcessOptions.EnableCoProcess {
		GlobalDispatcher, err = NewCoProcessDispatcher()
		EnableCoProcess = true
	}
	return err
}

// CoProcessMiddlewareConfig holds the middleware configuration.
type CoProcessMiddlewareConfig struct {
	ConfigData map[string]string `mapstructure:"config_data" bson:"config_data" json:"config_data"`
}

// New lets you do any initialisations for the object can be done here
func (m *CoProcessMiddleware) New() {}

// GetConfig retrieves the configuration from the API config - we user mapstructure for this for simplicity
func (m *CoProcessMiddleware) GetConfig() (interface{}, error) {
	var moduleConfig CoProcessMiddlewareConfig

	err := mapstructure.Decode(m.Spec.RawData, &moduleConfig)
	if err != nil {
		log.WithFields(logrus.Fields{
			"prefix": "jsvm",
		}).Error(err)
		return nil, err
	}

	return moduleConfig, nil
}

// IsEnabledForSpec checks if this middleware should be enabled for a given API.
func (m *CoProcessMiddleware) IsEnabledForSpec() bool {
	// This flag is true when Tyk has been compiled with CP support and when the configuration enables it.
	enableCoProcess := config.CoProcessOptions.EnableCoProcess && EnableCoProcess
	// This flag indicates if the current spec specifies any CP custom middleware.
	var usesCoProcessMiddleware bool

	supportedDrivers := []apidef.MiddlewareDriver{apidef.PythonDriver, apidef.LuaDriver, apidef.GrpcDriver}

	for _, driver := range supportedDrivers {
		if m.Spec.CustomMiddleware.Driver == driver && CoProcessName == driver {
			usesCoProcessMiddleware = true
			break
		}
	}

	if usesCoProcessMiddleware && enableCoProcess {
		log.WithFields(logrus.Fields{
			"prefix": "coprocess",
		}).Debug("Enabling CP middleware.")
		return true
	}

	if usesCoProcessMiddleware && !enableCoProcess {
		log.WithFields(logrus.Fields{
			"prefix": "coprocess",
		}).Error("Your API specifies a CP custom middleware, either Tyk wasn't build with CP support or CP is not enabled in your Tyk configuration file!")
	}

	if !usesCoProcessMiddleware && m.Spec.CustomMiddleware.Driver != "" {
		log.WithFields(logrus.Fields{
			"prefix": "coprocess",
		}).Error("CP Driver not supported: ", m.Spec.CustomMiddleware.Driver)
	}

	return false
}

// ProcessRequest will run any checks on the request on the way through the system, return an error to have the chain fail
func (m *CoProcessMiddleware) ProcessRequest(w http.ResponseWriter, r *http.Request, configuration interface{}) (error, int) {
	log.WithFields(logrus.Fields{
		"prefix": "coprocess",
	}).Debug("CoProcess Request, HookType: ", m.HookType)

	if !EnableCoProcess {
		return nil, 200
	}

	var extractor IdExtractor
	if m.Spec.EnableCoProcessAuth && m.Spec.CustomMiddleware.IdExtractor.Extractor != nil {
		extractor = m.Spec.CustomMiddleware.IdExtractor.Extractor.(IdExtractor)
	}

	var returnOverrides ReturnOverrides
	var sessionID string

	if m.HookType == coprocess.HookType_CustomKeyCheck && extractor != nil {
		sessionID, returnOverrides = extractor.ExtractAndCheck(r)

		if returnOverrides.ResponseCode != 0 {
			if returnOverrides.ResponseError == "" {
				return nil, returnOverrides.ResponseCode
			}
			err := errors.New(returnOverrides.ResponseError)
			return err, returnOverrides.ResponseCode
		}
	}

	// It's also possible to override the HookType:
	coProcessor := CoProcessor{
		Middleware: m,
		// HookType: coprocess.PreHook,
	}

	object := coProcessor.GetObjectFromRequest(r)

	returnObject, err := coProcessor.Dispatch(object)
	if err != nil {
		if m.HookType == coprocess.HookType_CustomKeyCheck {
			return errors.New("Key not authorised"), 403
		} else {
			return errors.New("Middleware error"), 500
		}
	}

	coProcessor.ObjectPostProcess(returnObject, r)

	token := returnObject.Metadata["token"]

	// The CP middleware indicates this is a bad auth:
	if returnObject.Request.ReturnOverrides.ResponseCode > 400 {

		log.WithFields(logrus.Fields{
			"path":   r.URL.Path,
			"origin": GetIPFromRequest(r),
			"key":    token,
		}).Info("Attempted access with invalid key.")

		// Fire Authfailed Event
		AuthFailed(m.TykMiddleware, r, token)

		// Report in health check
		ReportHealthCheckValue(m.Spec.Health, KeyFailure, "1")

		return errors.New("Key not authorised"), int(returnObject.Request.ReturnOverrides.ResponseCode)
	}

	if returnObject.Request.ReturnOverrides.ResponseCode > 0 {
		for h, v := range returnObject.Request.ReturnOverrides.Headers {
			w.Header().Set(h, v)
		}
		w.WriteHeader(int(returnObject.Request.ReturnOverrides.ResponseCode))
		w.Write([]byte(returnObject.Request.ReturnOverrides.ResponseError))
		return nil, 666
	}

	// Is this a CP authentication middleware?
	if m.Spec.EnableCoProcessAuth && m.HookType == coprocess.HookType_CustomKeyCheck {
		// The CP middleware didn't setup a session:
		if returnObject.Session == nil {
			return errors.New("Key not authorised"), 403
		}

		returnedSession := TykSessionState(returnObject.Session)

		if extractor == nil {
			sessionLifetime := getLifetime(m.Spec, returnedSession)
			// This API is not using the ID extractor, but we've got a session:
			m.Spec.SessionManager.UpdateSession(token, returnedSession, sessionLifetime)
			ctxSetSession(r, returnedSession)
			ctxSetAuthToken(r, token)
		} else {
			// The CP middleware did setup a session, we should pass it to the ID extractor (caching):
			extractor.PostProcess(r, returnedSession, sessionID)
		}
	}

	return nil, 200
}
<|endoftext|>


<commit_before>package ares

import (
	"fmt"
	"log"

	"github.com/nlopes/slack"
)

type Ares struct {
	SlackAppToken string
	SlackBotToken string
	SlackAppID    string
	ImgurClientID string
	BotUserID     string
	// Bot won't be added or re-added to these channels
	ManagedChannels []string
}

func (a *Ares) initBot() {
	api := slack.New(a.SlackAppToken)
	a.getBotUserID()

	channels, err := api.GetChannels(true)

	if err != nil {
		log.Fatal("Failed to get public channels list: ", err.Error())
	}

	for _, channel := range channels {
		a.addBotChannel(channel.ID)
	}

	groups, err := api.GetGroups(true)

	if err != nil {
		log.Fatal("Failed to get private channels list: ", err.Error())
	}

	for _, group := range groups {
		a.addBotGroup(group.ID)
	}
}

func (a *Ares) getBotUserID() {
	api := slack.New(a.SlackAppToken)
	users, err := api.GetUsers()

	if err != nil {
		log.Fatal("Failed to fetch slack users info", err.Error())
	}

	for _, user := range users {

		if user.Profile.ApiAppID == a.SlackAppID {
			a.BotUserID = user.ID
			return
		}
	}
	log.Fatal("Unable to find bot user on the Slack")
}

func (a *Ares) deleteFile(fileId string) {
	api := slack.New(a.SlackAppToken)

	if err := api.DeleteFile(fileId); err != nil {
		fmt.Println("Deletion of the file failed:", err.Error())
	}
}

func (a *Ares) addBotChannel(channelID string) {
	api := slack.New(a.SlackAppToken)
	if _, err := api.InviteUserToChannel(channelID, a.BotUserID); err != nil {
		if err.Error() != "already_in_channel" {
			log.Println(fmt.Sprintf("Failed to add bot to %s: %s", channelID, err.Error()))
		}
	}
}

func (a *Ares) addBotGroup(group string) {
	api := slack.New(a.SlackAppToken)
	if _, _, err := api.InviteUserToGroup(group, a.BotUserID); err != nil {
		log.Println(fmt.Sprintf("Failed to add bot to %s: %s", group, err.Error()))
	}
}

func (a *Ares) notifyUser(user, deleteHash string) {
	api := slack.New(a.SlackBotToken)
	params := slack.PostMessageParameters{}
	attachment := slack.Attachment{
		Color:   "danger",
		Pretext: "Hey there, image uploads to this slack is disabled. Please use external image upload services next time. This incident has been reported :squirrel:",
		Text:    "I have uploaded the image to Imgur for now.",
		Fields: []slack.AttachmentField{
			{
				Title: "If you wish to delete the image uploaded, click on the following link",
				Value: fmt.Sprintf("https://imgur.com/delete/%s", deleteHash),
			},
		},
	}
	params.Attachments = []slack.Attachment{attachment}
	_, _, err := api.PostMessage(user, "", params)
	if err != nil {
		log.Println("Failed to send DM to user", err.Error())
	}
}

func (a *Ares) sendImgToSlack(channel, user, url string) {
	api := slack.New(a.SlackBotToken)
	params := slack.PostMessageParameters{}
	attachment := slack.Attachment{
		Color:    "#D3D3D3",
		Text:     fmt.Sprintf("image originally uploaded by <@%s>", user),
		ImageURL: url,
	}
	params.Attachments = []slack.Attachment{attachment}
	_, _, err := api.PostMessage(channel, "", params)
	if err != nil {
		log.Println("Failed to send DM to user", err.Error())
	}
}

func (a *Ares) handleFile(file *slack.File, channel string) {
	resp := uploadToImgur(file.URLPrivateDownload, a.SlackAppToken, a.ImgurClientID)

	if resp.Status != 200 {
		log.Println("Failed to download/upload")
		return
	}

	a.notifyUser(file.User, resp.Data.Deletehash)
	a.sendImgToSlack(channel, file.User, resp.Data.Link)
	a.deleteFile(file.ID)
}

func isImageFile(fileType string) bool {
	for _, fType := range []string{"jpg", "jpeg", "png", "gif"} {
		if fType == fileType {
			return true
		}
	}
	return false
}

func (a *Ares) Run() {
	api := slack.New(a.SlackBotToken)

	a.initBot()
	log.Println("Bot initialized. Starting moderation duty.")

	rtm := api.NewRTM()
	go rtm.ManageConnection()

	for msg := range rtm.IncomingEvents {
		switch ev := msg.Data.(type) {

		case *slack.MessageEvent:
			if ev.SubType == "file_share" {
				if isImageFile(ev.File.Filetype) {
					a.handleFile(ev.File, ev.Channel)
				}
			}

		case *slack.RTMError:
			fmt.Printf("Error: %s\n", ev.Error())

		case *slack.InvalidAuthEvent:
			fmt.Printf("Invalid credentials")
			return

		default:
			// Ignore other events..
			// fmt.Printf("Unexpected: %v\n", msg.Data)
		}
	}
}
<commit_msg>Re-add bot when it gets removed<commit_after>package ares

import (
	"fmt"
	"log"

	"github.com/nlopes/slack"
)

type Ares struct {
	SlackAppToken string
	SlackBotToken string
	SlackAppID    string
	ImgurClientID string
	BotUserID     string
	// Bot won't be added or re-added to these channels
	ManagedChannels []string
}

func (a *Ares) initBot() {
	api := slack.New(a.SlackAppToken)
	a.getBotUserID()

	channels, err := api.GetChannels(true)

	if err != nil {
		log.Fatal("Failed to get public channels list: ", err.Error())
	}

	for _, channel := range channels {
		a.addBotChannel(channel.ID)
	}

	groups, err := api.GetGroups(true)

	if err != nil {
		log.Fatal("Failed to get private channels list: ", err.Error())
	}

	for _, group := range groups {
		a.addBotGroup(group.ID)
	}
}

func (a *Ares) getBotUserID() {
	api := slack.New(a.SlackAppToken)
	users, err := api.GetUsers()

	if err != nil {
		log.Fatal("Failed to fetch slack users info", err.Error())
	}

	for _, user := range users {

		if user.Profile.ApiAppID == a.SlackAppID {
			a.BotUserID = user.ID
			return
		}
	}
	log.Fatal("Unable to find bot user on the Slack")
}

func (a *Ares) deleteFile(fileId string) {
	api := slack.New(a.SlackAppToken)

	if err := api.DeleteFile(fileId); err != nil {
		fmt.Println("Deletion of the file failed:", err.Error())
	}
}

func (a *Ares) addBotChannel(channelID string) {
	api := slack.New(a.SlackAppToken)
	if _, err := api.InviteUserToChannel(channelID, a.BotUserID); err != nil {
		if err.Error() != "already_in_channel" {
			log.Println(fmt.Sprintf("Failed to add bot to %s: %s", channelID, err.Error()))
		}
	}
}

func (a *Ares) addBotGroup(group string) {
	api := slack.New(a.SlackAppToken)
	if _, _, err := api.InviteUserToGroup(group, a.BotUserID); err != nil {
		log.Println(fmt.Sprintf("Failed to add bot to %s: %s", group, err.Error()))
	}
}

func (a *Ares) notifyUser(user, deleteHash string) {
	api := slack.New(a.SlackBotToken)
	params := slack.PostMessageParameters{}
	attachment := slack.Attachment{
		Color:   "danger",
		Pretext: "Hey there, image uploads to this slack is disabled. Please use external image upload services next time. This incident has been reported :squirrel:",
		Text:    "I have uploaded the image to Imgur for now.",
		Fields: []slack.AttachmentField{
			{
				Title: "If you wish to delete the image uploaded, click on the following link",
				Value: fmt.Sprintf("https://imgur.com/delete/%s", deleteHash),
			},
		},
	}
	params.Attachments = []slack.Attachment{attachment}
	_, _, err := api.PostMessage(user, "", params)
	if err != nil {
		log.Println("Failed to send DM to user", err.Error())
	}
}

func (a *Ares) sendImgToSlack(channel, user, url string) {
	api := slack.New(a.SlackBotToken)
	params := slack.PostMessageParameters{}
	attachment := slack.Attachment{
		Color:    "#D3D3D3",
		Text:     fmt.Sprintf("image originally uploaded by <@%s>", user),
		ImageURL: url,
	}
	params.Attachments = []slack.Attachment{attachment}
	_, _, err := api.PostMessage(channel, "", params)
	if err != nil {
		log.Println("Failed to send DM to user", err.Error())
	}
}

func (a *Ares) handleFile(file *slack.File, channel string) {
	resp := uploadToImgur(file.URLPrivateDownload, a.SlackAppToken, a.ImgurClientID)

	if resp.Status != 200 {
		log.Println("Failed to download/upload")
		return
	}

	a.notifyUser(file.User, resp.Data.Deletehash)
	a.sendImgToSlack(channel, file.User, resp.Data.Link)
	a.deleteFile(file.ID)
}

func isImageFile(fileType string) bool {
	for _, fType := range []string{"jpg", "jpeg", "png", "gif"} {
		if fType == fileType {
			return true
		}
	}
	return false
}

func (a *Ares) Run() {
	api := slack.New(a.SlackBotToken)

	a.initBot()
	log.Println("Bot initialized. Starting moderation duty.")

	rtm := api.NewRTM()
	go rtm.ManageConnection()

	for msg := range rtm.IncomingEvents {
		switch ev := msg.Data.(type) {

		case *slack.MessageEvent:
			if ev.SubType == "file_share" {
				if isImageFile(ev.File.Filetype) {
					a.handleFile(ev.File, ev.Channel)
				}
			}

		case *slack.GroupLeftEvent:
			log.Println("Bot was removed from private channel: ", ev.Channel)
			a.addBotGroup(ev.Channel)

		case *slack.ChannelLeftEvent:
			log.Println("Bot was removed from public channel: ", ev.Channel)
			a.addBotChannel(ev.Channel)

		case *slack.RTMError:
			fmt.Printf("Error: %s\n", ev.Error())

		case *slack.InvalidAuthEvent:
			fmt.Printf("Invalid credentials")
			return

		default:
			// Ignore other events..
			// fmt.Printf("Unexpected: %v\n", msg.Data)
		}
	}
}
<|endoftext|>


<commit_before>package core

import (
	"encoding/base64"
	"fmt"

	context "github.com/jbenet/go-ipfs/Godeps/_workspace/src/code.google.com/p/go.net/context"
	b58 "github.com/jbenet/go-ipfs/Godeps/_workspace/src/github.com/jbenet/go-base58"
	ma "github.com/jbenet/go-ipfs/Godeps/_workspace/src/github.com/jbenet/go-multiaddr"

	bserv "github.com/jbenet/go-ipfs/blockservice"
	config "github.com/jbenet/go-ipfs/config"
	diag "github.com/jbenet/go-ipfs/diagnostics"
	exchange "github.com/jbenet/go-ipfs/exchange"
	bitswap "github.com/jbenet/go-ipfs/exchange/bitswap"
	bsnet "github.com/jbenet/go-ipfs/exchange/bitswap/network"
	mount "github.com/jbenet/go-ipfs/fuse/mount"
	merkledag "github.com/jbenet/go-ipfs/merkledag"
	namesys "github.com/jbenet/go-ipfs/namesys"
	inet "github.com/jbenet/go-ipfs/net"
	mux "github.com/jbenet/go-ipfs/net/mux"
	netservice "github.com/jbenet/go-ipfs/net/service"
	path "github.com/jbenet/go-ipfs/path"
	peer "github.com/jbenet/go-ipfs/peer"
	pin "github.com/jbenet/go-ipfs/pin"
	routing "github.com/jbenet/go-ipfs/routing"
	dht "github.com/jbenet/go-ipfs/routing/dht"
	u "github.com/jbenet/go-ipfs/util"
	ctxc "github.com/jbenet/go-ipfs/util/ctxcloser"
	"github.com/jbenet/go-ipfs/util/debugerror"
)

const IpnsValidatorTag = "ipns"

var log = u.Logger("core")

// IpfsNode is IPFS Core module. It represents an IPFS instance.
type IpfsNode struct {

	// the node's configuration
	Config *config.Config

	// the local node's identity
	Identity peer.Peer

	// storage for other Peer instances
	Peerstore peer.Peerstore

	// the local datastore
	Datastore u.ThreadSafeDatastoreCloser

	// the network message stream
	Network inet.Network

	// the routing system. recommend ipfs-dht
	Routing routing.IpfsRouting

	// the block exchange + strategy (bitswap)
	Exchange exchange.Interface

	// the block service, get/add blocks.
	Blocks *bserv.BlockService

	// the merkle dag service, get/add objects.
	DAG merkledag.DAGService

	// the path resolution system
	Resolver *path.Resolver

	// the name system, resolves paths to hashes
	Namesys namesys.NameSystem

	// the diagnostics service
	Diagnostics *diag.Diagnostics

	// the pinning manager
	Pinning pin.Pinner

	// current mount state, if any.
	Mounts Mounts

	ctxc.ContextCloser

	onlineMode bool // alternatively, offline
}

// Mounts defines what the node's mount state is. This should
// perhaps be moved to the daemon or mount. It's here because
// it needs to be accessible across daemon requests.
type Mounts struct {
	Ipfs mount.Mount
	Ipns mount.Mount
}

// NewIpfsNode constructs a new IpfsNode based on the given config.
func NewIpfsNode(cfg *config.Config, online bool) (n *IpfsNode, err error) {
	success := false // flip to true after all sub-system inits succeed
	defer func() {
		if !success && n != nil {
			n.Close()
		}
	}()

	if cfg == nil {
		return nil, debugerror.Errorf("configuration required")
	}

	// derive this from a higher context.
	ctx := context.TODO()
	n = &IpfsNode{
		onlineMode: online,
		Config:     cfg,
	}
	n.ContextCloser = ctxc.NewContextCloser(context.TODO(), n.teardown)

	// setup datastore.
	if n.Datastore, err = makeDatastore(cfg.Datastore); err != nil {
		return nil, debugerror.Wrap(err)
	}

	// setup peerstore + local peer identity
	n.Peerstore = peer.NewPeerstore()
	n.Identity, err = initIdentity(&n.Config.Identity, n.Peerstore, online)
	if err != nil {
		return nil, debugerror.Wrap(err)
	}

	// setup online services
	if online {

		dhtService := netservice.NewService(ctx, nil)      // nil handler for now, need to patch it
		exchangeService := netservice.NewService(ctx, nil) // nil handler for now, need to patch it
		diagService := netservice.NewService(ctx, nil)     // nil handler for now, need to patch it

		muxMap := &mux.ProtocolMap{
			mux.ProtocolID_Routing:    dhtService,
			mux.ProtocolID_Exchange:   exchangeService,
			mux.ProtocolID_Diagnostic: diagService,
			// add protocol services here.
		}

		// setup the network
		listenAddrs, err := listenAddresses(cfg)
		if err != nil {
			return nil, debugerror.Wrap(err)
		}

		n.Network, err = inet.NewIpfsNetwork(ctx, listenAddrs, n.Identity, n.Peerstore, muxMap)
		if err != nil {
			return nil, debugerror.Wrap(err)
		}
		n.AddCloserChild(n.Network)

		// setup diagnostics service
		n.Diagnostics = diag.NewDiagnostics(n.Identity, n.Network, diagService)
		diagService.SetHandler(n.Diagnostics)

		// setup routing service
		dhtRouting := dht.NewDHT(ctx, n.Identity, n.Peerstore, n.Network, dhtService, n.Datastore)
		dhtRouting.Validators[IpnsValidatorTag] = namesys.ValidateIpnsRecord

		// TODO(brian): perform this inside NewDHT factory method
		dhtService.SetHandler(dhtRouting) // wire the handler to the service.
		n.Routing = dhtRouting
		n.AddCloserChild(dhtRouting)

		// setup exchange service
		const alwaysSendToPeer = true // use YesManStrategy
		bitswapNetwork := bsnet.NewFromIpfsNetwork(exchangeService, n.Network)
		n.Exchange = bitswap.New(ctx, n.Identity, bitswapNetwork, n.Routing, n.Datastore, alwaysSendToPeer)

		go initConnections(ctx, n.Config, n.Peerstore, dhtRouting)
	}

	// TODO(brian): when offline instantiate the BlockService with a bitswap
	// session that simply doesn't return blocks
	n.Blocks, err = bserv.NewBlockService(n.Datastore, n.Exchange)
	if err != nil {
		return nil, debugerror.Wrap(err)
	}

	n.DAG = merkledag.NewDAGService(n.Blocks)
	n.Namesys = namesys.NewNameSystem(n.Routing)
	n.Pinning, err = pin.LoadPinner(n.Datastore, n.DAG)
	if err != nil {
		n.Pinning = pin.NewPinner(n.Datastore, n.DAG)
	}
	n.Resolver = &path.Resolver{DAG: n.DAG}

	success = true
	return n, nil
}

func (n *IpfsNode) teardown() error {
	if err := n.Datastore.Close(); err != nil {
		return err
	}
	return nil
}

func (n *IpfsNode) OnlineMode() bool {
	return n.onlineMode
}

func initIdentity(cfg *config.Identity, peers peer.Peerstore, online bool) (peer.Peer, error) {
	if cfg.PeerID == "" {
		return nil, debugerror.New("Identity was not set in config (was ipfs init run?)")
	}

	if len(cfg.PeerID) == 0 {
		return nil, debugerror.New("No peer ID in config! (was ipfs init run?)")
	}

	// get peer from peerstore (so it is constructed there)
	id := peer.ID(b58.Decode(cfg.PeerID))
	self, err := peers.Get(id)
	if err != nil {
		return nil, err
	}
	self.SetType(peer.Local)
	self, err = peers.Add(self)
	if err != nil {
		return nil, err
	}

	// when not online, don't need to parse private keys (yet)
	if online {
		skb, err := base64.StdEncoding.DecodeString(cfg.PrivKey)
		if err != nil {
			return nil, err
		}

		if err := self.LoadAndVerifyKeyPair(skb); err != nil {
			return nil, err
		}
	}

	return self, nil
}

func initConnections(ctx context.Context, cfg *config.Config, pstore peer.Peerstore, route *dht.IpfsDHT) {
	for _, p := range cfg.Bootstrap {
		if p.PeerID == "" {
			log.Errorf("error: peer does not include PeerID. %v", p)
		}

		maddr, err := ma.NewMultiaddr(p.Address)
		if err != nil {
			log.Error(err)
			continue
		}

		// setup peer
		npeer, err := pstore.Get(peer.DecodePrettyID(p.PeerID))
		if err != nil {
			log.Errorf("Bootstrapping error: %v", err)
			continue
		}
		npeer.AddAddress(maddr)

		if _, err = route.Connect(ctx, npeer); err != nil {
			log.Errorf("Bootstrapping error: %v", err)
		}
	}
}

func listenAddresses(cfg *config.Config) ([]ma.Multiaddr, error) {
	var listen []ma.Multiaddr

	if len(cfg.Addresses.Swarm) > 0 {
		maddr, err := ma.NewMultiaddr(cfg.Addresses.Swarm)
		if err != nil {
			return nil, fmt.Errorf("Failure to parse config.Addresses.Swarm: %s", cfg.Addresses.Swarm)
		}

		listen = append(listen, maddr)
	}

	return listen, nil
}
<commit_msg>fix node context (tests pass)<commit_after>package core

import (
	"encoding/base64"
	"fmt"

	context "github.com/jbenet/go-ipfs/Godeps/_workspace/src/code.google.com/p/go.net/context"
	b58 "github.com/jbenet/go-ipfs/Godeps/_workspace/src/github.com/jbenet/go-base58"
	ma "github.com/jbenet/go-ipfs/Godeps/_workspace/src/github.com/jbenet/go-multiaddr"

	bserv "github.com/jbenet/go-ipfs/blockservice"
	config "github.com/jbenet/go-ipfs/config"
	diag "github.com/jbenet/go-ipfs/diagnostics"
	exchange "github.com/jbenet/go-ipfs/exchange"
	bitswap "github.com/jbenet/go-ipfs/exchange/bitswap"
	bsnet "github.com/jbenet/go-ipfs/exchange/bitswap/network"
	mount "github.com/jbenet/go-ipfs/fuse/mount"
	merkledag "github.com/jbenet/go-ipfs/merkledag"
	namesys "github.com/jbenet/go-ipfs/namesys"
	inet "github.com/jbenet/go-ipfs/net"
	mux "github.com/jbenet/go-ipfs/net/mux"
	netservice "github.com/jbenet/go-ipfs/net/service"
	path "github.com/jbenet/go-ipfs/path"
	peer "github.com/jbenet/go-ipfs/peer"
	pin "github.com/jbenet/go-ipfs/pin"
	routing "github.com/jbenet/go-ipfs/routing"
	dht "github.com/jbenet/go-ipfs/routing/dht"
	u "github.com/jbenet/go-ipfs/util"
	ctxc "github.com/jbenet/go-ipfs/util/ctxcloser"
	"github.com/jbenet/go-ipfs/util/debugerror"
)

const IpnsValidatorTag = "ipns"

var log = u.Logger("core")

// IpfsNode is IPFS Core module. It represents an IPFS instance.
type IpfsNode struct {

	// the node's configuration
	Config *config.Config

	// the local node's identity
	Identity peer.Peer

	// storage for other Peer instances
	Peerstore peer.Peerstore

	// the local datastore
	Datastore u.ThreadSafeDatastoreCloser

	// the network message stream
	Network inet.Network

	// the routing system. recommend ipfs-dht
	Routing routing.IpfsRouting

	// the block exchange + strategy (bitswap)
	Exchange exchange.Interface

	// the block service, get/add blocks.
	Blocks *bserv.BlockService

	// the merkle dag service, get/add objects.
	DAG merkledag.DAGService

	// the path resolution system
	Resolver *path.Resolver

	// the name system, resolves paths to hashes
	Namesys namesys.NameSystem

	// the diagnostics service
	Diagnostics *diag.Diagnostics

	// the pinning manager
	Pinning pin.Pinner

	// current mount state, if any.
	Mounts Mounts

	ctxc.ContextCloser

	onlineMode bool // alternatively, offline
}

// Mounts defines what the node's mount state is. This should
// perhaps be moved to the daemon or mount. It's here because
// it needs to be accessible across daemon requests.
type Mounts struct {
	Ipfs mount.Mount
	Ipns mount.Mount
}

// NewIpfsNode constructs a new IpfsNode based on the given config.
func NewIpfsNode(cfg *config.Config, online bool) (n *IpfsNode, err error) {
	success := false // flip to true after all sub-system inits succeed
	defer func() {
		if !success && n != nil {
			n.Close()
		}
	}()

	if cfg == nil {
		return nil, debugerror.Errorf("configuration required")
	}

	// derive this from a higher context.
	ctx := context.TODO()
	n = &IpfsNode{
		onlineMode: online,
		Config:     cfg,
	}
	n.ContextCloser = ctxc.NewContextCloser(ctx, n.teardown)

	// setup datastore.
	if n.Datastore, err = makeDatastore(cfg.Datastore); err != nil {
		return nil, debugerror.Wrap(err)
	}

	// setup peerstore + local peer identity
	n.Peerstore = peer.NewPeerstore()
	n.Identity, err = initIdentity(&n.Config.Identity, n.Peerstore, online)
	if err != nil {
		return nil, debugerror.Wrap(err)
	}

	// setup online services
	if online {

		dhtService := netservice.NewService(ctx, nil)      // nil handler for now, need to patch it
		exchangeService := netservice.NewService(ctx, nil) // nil handler for now, need to patch it
		diagService := netservice.NewService(ctx, nil)     // nil handler for now, need to patch it

		muxMap := &mux.ProtocolMap{
			mux.ProtocolID_Routing:    dhtService,
			mux.ProtocolID_Exchange:   exchangeService,
			mux.ProtocolID_Diagnostic: diagService,
			// add protocol services here.
		}

		// setup the network
		listenAddrs, err := listenAddresses(cfg)
		if err != nil {
			return nil, debugerror.Wrap(err)
		}

		n.Network, err = inet.NewIpfsNetwork(ctx, listenAddrs, n.Identity, n.Peerstore, muxMap)
		if err != nil {
			return nil, debugerror.Wrap(err)
		}
		n.AddCloserChild(n.Network)

		// setup diagnostics service
		n.Diagnostics = diag.NewDiagnostics(n.Identity, n.Network, diagService)
		diagService.SetHandler(n.Diagnostics)

		// setup routing service
		dhtRouting := dht.NewDHT(ctx, n.Identity, n.Peerstore, n.Network, dhtService, n.Datastore)
		dhtRouting.Validators[IpnsValidatorTag] = namesys.ValidateIpnsRecord

		// TODO(brian): perform this inside NewDHT factory method
		dhtService.SetHandler(dhtRouting) // wire the handler to the service.
		n.Routing = dhtRouting
		n.AddCloserChild(dhtRouting)

		// setup exchange service
		const alwaysSendToPeer = true // use YesManStrategy
		bitswapNetwork := bsnet.NewFromIpfsNetwork(exchangeService, n.Network)
		n.Exchange = bitswap.New(ctx, n.Identity, bitswapNetwork, n.Routing, n.Datastore, alwaysSendToPeer)

		go initConnections(ctx, n.Config, n.Peerstore, dhtRouting)
	}

	// TODO(brian): when offline instantiate the BlockService with a bitswap
	// session that simply doesn't return blocks
	n.Blocks, err = bserv.NewBlockService(n.Datastore, n.Exchange)
	if err != nil {
		return nil, debugerror.Wrap(err)
	}

	n.DAG = merkledag.NewDAGService(n.Blocks)
	n.Namesys = namesys.NewNameSystem(n.Routing)
	n.Pinning, err = pin.LoadPinner(n.Datastore, n.DAG)
	if err != nil {
		n.Pinning = pin.NewPinner(n.Datastore, n.DAG)
	}
	n.Resolver = &path.Resolver{DAG: n.DAG}

	success = true
	return n, nil
}

func (n *IpfsNode) teardown() error {
	if err := n.Datastore.Close(); err != nil {
		return err
	}
	return nil
}

func (n *IpfsNode) OnlineMode() bool {
	return n.onlineMode
}

func initIdentity(cfg *config.Identity, peers peer.Peerstore, online bool) (peer.Peer, error) {
	if cfg.PeerID == "" {
		return nil, debugerror.New("Identity was not set in config (was ipfs init run?)")
	}

	if len(cfg.PeerID) == 0 {
		return nil, debugerror.New("No peer ID in config! (was ipfs init run?)")
	}

	// get peer from peerstore (so it is constructed there)
	id := peer.ID(b58.Decode(cfg.PeerID))
	self, err := peers.Get(id)
	if err != nil {
		return nil, err
	}
	self.SetType(peer.Local)
	self, err = peers.Add(self)
	if err != nil {
		return nil, err
	}

	// when not online, don't need to parse private keys (yet)
	if online {
		skb, err := base64.StdEncoding.DecodeString(cfg.PrivKey)
		if err != nil {
			return nil, err
		}

		if err := self.LoadAndVerifyKeyPair(skb); err != nil {
			return nil, err
		}
	}

	return self, nil
}

func initConnections(ctx context.Context, cfg *config.Config, pstore peer.Peerstore, route *dht.IpfsDHT) {
	for _, p := range cfg.Bootstrap {
		if p.PeerID == "" {
			log.Errorf("error: peer does not include PeerID. %v", p)
		}

		maddr, err := ma.NewMultiaddr(p.Address)
		if err != nil {
			log.Error(err)
			continue
		}

		// setup peer
		npeer, err := pstore.Get(peer.DecodePrettyID(p.PeerID))
		if err != nil {
			log.Errorf("Bootstrapping error: %v", err)
			continue
		}
		npeer.AddAddress(maddr)

		if _, err = route.Connect(ctx, npeer); err != nil {
			log.Errorf("Bootstrapping error: %v", err)
		}
	}
}

func listenAddresses(cfg *config.Config) ([]ma.Multiaddr, error) {
	var listen []ma.Multiaddr

	if len(cfg.Addresses.Swarm) > 0 {
		maddr, err := ma.NewMultiaddr(cfg.Addresses.Swarm)
		if err != nil {
			return nil, fmt.Errorf("Failure to parse config.Addresses.Swarm: %s", cfg.Addresses.Swarm)
		}

		listen = append(listen, maddr)
	}

	return listen, nil
}
<|endoftext|>